from __future__ import annotations import asyncio import itertools import logging from typing import TYPE_CHECKING import aiofiles from aiohttp import ClientResponseError, ClientTimeout from aiohttp_socks import ProxyType from proxy_scraper_checker.counter import IncrInt from proxy_scraper_checker.http import get_response_text from proxy_scraper_checker.parsers import PROXY_REGEX from proxy_scraper_checker.proxy import Proxy from proxy_scraper_checker.utils import bytes_decode, is_http_url if TYPE_CHECKING: from aiohttp import ClientSession from rich.progress import Progress, TaskID from proxy_scraper_checker.settings import Settings from proxy_scraper_checker.storage import ProxyStorage _logger = logging.getLogger(__name__) async def scrape_one( *, counter: IncrInt, progress: Progress, proto: ProxyType, session: ClientSession, source: str, storage: ProxyStorage, task: TaskID, timeout: ClientTimeout, ) -> None: try: if is_http_url(source): async with session.get(source, timeout=timeout) as response: content = await response.read() text = get_response_text(response=response, content=content) else: async with aiofiles.open(source, "rb") as f: content = await f.read() text = bytes_decode(content) except ClientResponseError as e: _logger.warning( "%s | HTTP status code %d: %s", source, e.status, e.message ) except Exception as e: _logger.warning( "%s | %s.%s: %s", source, e.__class__.__module__, e.__class__.__qualname__, e, ) else: counter.incr() proxies = PROXY_REGEX.finditer(text) try: proxy = next(proxies) except StopIteration: _logger.warning("%s | No proxies found", source) else: for proxy in itertools.chain((proxy,), proxies): # noqa: B020 try: protocol = ProxyType[ proxy.group("protocol").upper().rstrip("S") ] except AttributeError: protocol = proto storage.add( Proxy( protocol=protocol, host=proxy.group("host"), port=int(proxy.group("port")), username=proxy.group("username"), password=proxy.group("password"), ) ) progress.update(task_id=task, advance=1, successful_count=counter.value) async def scrape_all( *, progress: Progress, session: ClientSession, settings: Settings, storage: ProxyStorage, ) -> None: counters = {proto: IncrInt() for proto in settings.sources} progress_tasks = { proto: progress.add_task( description="", total=len(sources), module="Scraper", protocol=proto.name, successful_count=0, ) for proto, sources in settings.sources.items() } timeout = ClientTimeout(total=settings.source_timeout) await asyncio.gather( *( scrape_one( counter=counters[proto], progress=progress, proto=proto, session=session, source=source, storage=storage, task=progress_tasks[proto], timeout=timeout, ) for proto, sources in settings.sources.items() for source in sources ) )