mirror of
https://github.com/nexus-stc/hyperboria
synced 2024-12-04 00:42:53 +01:00
55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
|
from typing import Dict
|
||
|
|
|
||
|
|
import aiohttp.client_exceptions
|
||
|
|
from aiohttp_socks import ProxyError
|
||
|
|
from nexus.pylon.drivers.base import BaseDriver
|
||
|
|
from nexus.pylon.exceptions import (
|
||
|
|
BadResponseError,
|
||
|
|
NotFoundError,
|
||
|
|
)
|
||
|
|
from nexus.pylon.prepared_request import PreparedRequest
|
||
|
|
from nexus.pylon.proto.file_pb2 import Chunk as ChunkPb
|
||
|
|
from nexus.pylon.proto.file_pb2 import FileResponse as FileResponsePb
|
||
|
|
from python_socks import ProxyTimeoutError
|
||
|
|
from tenacity import (
|
||
|
|
retry,
|
||
|
|
retry_if_exception_type,
|
||
|
|
stop_after_attempt,
|
||
|
|
wait_random,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class DirectDriver(BaseDriver):
|
||
|
|
allowed_content_type = None
|
||
|
|
|
||
|
|
@retry(
|
||
|
|
reraise=True,
|
||
|
|
wait=wait_random(min=1, max=2),
|
||
|
|
stop=stop_after_attempt(7),
|
||
|
|
retry=retry_if_exception_type((ProxyError, aiohttp.client_exceptions.ClientPayloadError, ProxyTimeoutError)),
|
||
|
|
)
|
||
|
|
async def execute_prepared_file_request(self, prepared_file_request: PreparedRequest, params: Dict):
|
||
|
|
async with self.get_session() as session:
|
||
|
|
async with prepared_file_request.execute_with(session=session) as resp:
|
||
|
|
if resp.status == 404:
|
||
|
|
raise NotFoundError(url=prepared_file_request.url)
|
||
|
|
elif (
|
||
|
|
resp.status != 200
|
||
|
|
or (
|
||
|
|
self.allowed_content_type
|
||
|
|
and resp.headers.get('Content-Type', '').lower() not in self.allowed_content_type
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise BadResponseError(
|
||
|
|
request_headers=prepared_file_request.headers,
|
||
|
|
url=prepared_file_request.url,
|
||
|
|
status=resp.status,
|
||
|
|
headers=str(resp.headers),
|
||
|
|
)
|
||
|
|
file_validator = self.validator(params)
|
||
|
|
yield FileResponsePb(status=FileResponsePb.Status.BEGIN_TRANSMISSION, source=prepared_file_request.url)
|
||
|
|
async for content, _ in resp.content.iter_chunks():
|
||
|
|
file_validator.update(content)
|
||
|
|
yield FileResponsePb(chunk=ChunkPb(content=content), source=prepared_file_request.url)
|
||
|
|
file_validator.validate()
|