Refactor, improve startup time, fix some bugs

nitro_generator_checker/nitro_checker.py

@@ -4,18 +4,18 @@
 import logging
 from configparser import ConfigParser
 from types import MappingProxyType
-from typing import Optional
+from typing import Optional, Tuple
 
-from aiofile import async_open
 from aiohttp import ClientSession, ClientTimeout, DummyCookieJar
 from aiohttp_socks import ProxyConnector
 from rich.console import Console
 from rich.live import Live
 
-from . import validators
+from . import result_handlers, validators
 from .counter import Counter
 from .nitro_generator import NitroGenerator
 from .proxy_generator import ProxyGenerator
+from .utils import create_background_task
 
 logger = logging.getLogger(__name__)
 HEADERS = MappingProxyType(
@@ -32,13 +32,12 @@ class NitroChecker:
     __slots__ = (
         "console",
         "counter",
-        "file_name",
         "max_connections",
         "nitro_generator",
         "proxy_generator",
         "session",
         "timeout",
-        "webhook_url",
+        "result_handlers",
     )
 
     def __init__(
@@ -57,13 +56,21 @@ def __init__(
 
         self.console = console or Console()
         self.counter = Counter()
-        self.file_name = file_name
         self.max_connections = validators.max_connections(max_connections)
         self.nitro_generator = NitroGenerator()
         self.proxy_generator = ProxyGenerator(session)
         self.session = session
         self.timeout = ClientTimeout(total=timeout, sock_connect=float("inf"))
-        self.webhook_url = webhook_url
+
+        file_handler = result_handlers.FileHandler(file_name)
+        self.result_handlers: Tuple[result_handlers.ABCResultHandler, ...] = (
+            (
+                file_handler,
+                result_handlers.DiscordWebhookHandler(session, webhook_url),
+            )
+            if webhook_url
+            else (file_handler,)
+        )
 
     @classmethod
     async def run_from_configparser(
@@ -84,6 +91,7 @@ async def run_from_configparser(
             await ngc.run()
 
     async def checker(self, live: Live) -> None:
+        await self.proxy_generator.ready_event.wait()
         for code in self.nitro_generator:
             url = (
                 f"https://discord.com/api/v9/entitlements/gift-codes/{code}"
@@ -120,33 +128,19 @@ async def checker(self, live: Live) -> None:
                 else:
                     logger.info("%s | Valid", code)
                     gift_url = f"https://discord.gift/{code}"
-                    await asyncio.gather(
-                        self.save_gift(gift_url),
-                        self.send_webhook_msg(gift_url),
-                    )
+                    for handler in self.result_handlers:
+                        create_background_task(handler.save(gift_url))
                     self.counter.add_valid()
                     self.counter.add_total()
                     live.update(self.counter.as_rich_table())
 
-    async def save_gift(self, gift_url: str) -> None:
-        async with async_open(self.file_name, "a", encoding="utf-8") as f:
-            await f.write(f"\n{gift_url}")
-
-    async def send_webhook_msg(self, gift_url: str) -> None:
-        if not self.webhook_url:
-            return
-        async with self.session.post(
-            self.webhook_url, json={"content": f"@everyone {gift_url}"}
-        ):
-            pass
-
     async def run(self) -> None:
-        set_proxies_task = asyncio.create_task(
-            self.proxy_generator.run_inf_loop()
-        )
-        await self.proxy_generator.wait_for_proxies()
         with Live(self.counter.as_rich_table(), console=self.console) as live:
-            coroutines = (
-                self.checker(live) for _ in range(self.max_connections)
+            await asyncio.gather(
+                self.proxy_generator.run_inf_loop(),
+                *(
+                    result_handler.pre_run()
+                    for result_handler in self.result_handlers
+                ),
+                *(self.checker(live) for _ in range(self.max_connections)),
             )
-            await asyncio.gather(set_proxies_task, *coroutines)

nitro_generator_checker/proxy_generator.py

@@ -11,11 +11,12 @@
 
 
 class ProxyGenerator:
-    __slots__ = ("proxies", "session")
+    __slots__ = ("proxies", "session", "ready_event")
 
     def __init__(self, session: ClientSession) -> None:
         self.session = session
         self.proxies: Tuple[str, ...] = ()
+        self.ready_event = asyncio.Event()
 
     async def set_proxies(self) -> None:
         url = "https://raw.githubusercontent.com/monosans/proxy-list/main/proxies/{}.txt"
@@ -29,16 +30,13 @@ async def set_proxies(self) -> None:
         )
         if proxies:
             self.proxies = proxies
+            self.ready_event.set()
 
     async def run_inf_loop(self) -> NoReturn:
         while True:
             await self.set_proxies()
             await asyncio.sleep(60)
 
-    async def wait_for_proxies(self) -> None:
-        while not self.proxies:
-            await asyncio.sleep(0)
-
     def get_random_proxy(self) -> str:
         return random.choice(self.proxies)
 

nitro_generator_checker/result_handlers.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import asyncio
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from aiofile import async_open
+from aiohttp import ClientSession
+
+from .utils import sync_to_async
+
+
+class ABCResultHandler(metaclass=ABCMeta):
+    __slots__ = ()
+
+    async def pre_run(self) -> None:  # noqa: B027
+        pass
+
+    @abstractmethod
+    async def save(self, gift_url: str) -> None:
+        pass
+
+
+class FileHandler(ABCResultHandler):
+    __slots__ = ("file_path", "_ready_event")
+
+    def __init__(self, file_path: str) -> None:
+        self.file_path = file_path
+        self._ready_event = asyncio.Event()
+
+    async def pre_run(self) -> None:
+        def inner() -> None:
+            path = Path(self.file_path)
+            if path.is_dir():
+                msg = "FileName must be a file, not a directory"
+                raise ValueError(msg)
+            path.parent.mkdir(parents=True, exist_ok=True)
+
+        await sync_to_async(inner)
+        self._ready_event.set()
+
+    async def save(self, gift_url: str) -> None:
+        await self._ready_event.wait()
+        async with async_open(self.file_path, "a", encoding="utf-8") as f:
+            await f.write(f"\n{gift_url}")
+
+
+class DiscordWebhookHandler(ABCResultHandler):
+    __slots__ = ("session", "url")
+
+    def __init__(self, session: ClientSession, url: str) -> None:
+        self.session = session
+        self.url = url
+
+    async def save(self, gift_url: str) -> None:
+        async with self.session.post(
+            self.url, json={"content": f"@everyone {gift_url}"}
+        ):
+            pass

nitro_generator_checker/utils.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable, Coroutine, Set, TypeVar
+
+T = TypeVar("T")
+
+background_tasks: Set[asyncio.Task[Any]] = set()
+
+
+def create_background_task(coro: Coroutine[Any, Any, Any]) -> None:
+    task = asyncio.create_task(coro)
+    background_tasks.add(task)
+    task.add_done_callback(background_tasks.discard)
+
+
+def sync_to_async(func: Callable[[], T]) -> asyncio.Future[T]:
+    loop = asyncio.get_running_loop()
+    return loop.run_in_executor(None, func)