Refactor Cloudflare, Redis, and Webhook modules: improve code readability, centralize configurations, and enhance error handling. Upgrade dependencies and update Docker Redis image to Alpine.

Author: Erik05Master

.pylintrc

@@ -2,3 +2,4 @@
 disable=
     missing-module-docstring,
     missing-class-docstring,
+    too-few-public-methods,

compose.yaml

@@ -10,7 +10,7 @@ services:
 
   redis:
     container_name: noattack_redis
-    image: redis
+    image: redis:alpine
     restart: always
     ports:
       - "6379:6379"

main.py

@@ -1,127 +1,113 @@
 import asyncio
+import logging
 import sys
-import time
 
-import aiohttp
 import psutil
 import pyfiglet
 from colorama import Fore
 
-from modules import cloudflare, config, webhook, redis
+from modules.config import get_config, ConfigNotFoundError
+from modules.cloudflare import Cloudflare
+from modules.webhook import Webhook
+from modules.redis import RedisCache
 
-config = config.Config()
-cloudflare = cloudflare.Cloudflare()
-webhook = webhook.Webhook()
-redis = redis.RedisCache()
+logger = logging.getLogger("noattack")
 
 PREFIX = f"{Fore.RED}[\033[38;5;208mNoAttack{Fore.RED}]{Fore.RESET} "
 
 
-def get_network_speed(duration=1):
-    """
-    Measure network speed over a given duration.
-    """
-    io_counters_start = psutil.net_io_counters()
+async def get_network_speed():
+    """Measure incoming and outgoing traffic in MB/s over 1 second."""
+    snap1 = psutil.net_io_counters()
+    await asyncio.sleep(1)
+    snap2 = psutil.net_io_counters()
+    recv = (snap2.bytes_recv - snap1.bytes_recv) / 1024 / 1024
+    sent = (snap2.bytes_sent - snap1.bytes_sent) / 1024 / 1024
+    return recv, sent
 
-    time.sleep(duration)
 
-    io_counters_end = psutil.net_io_counters()
+async def handle_zone(cloudflare, webhook, zone_id, under_attack):
+    """Activate or deactivate Under Attack mode for a zone if the state needs to change."""
+    zone_name = await cloudflare.get_zone(zone_id)
+    if zone_name is None:
+        return
 
-    bytes_received = io_counters_end.bytes_recv - io_counters_start.bytes_recv
-    bytes_sent = io_counters_end.bytes_sent - io_counters_start.bytes_sent
+    currently_attacking = await cloudflare.get_zone_under_attack(zone_id)
+    if currently_attacking is None:
+        return
 
-    mb_received_per_sec = bytes_received / duration / 1024 / 1024
-    mb_sent_per_sec = bytes_sent / duration / 1024 / 1024
-    return mb_received_per_sec, mb_sent_per_sec
+    if under_attack and not currently_attacking:
+        logger.info("Activating Under Attack Mode for %s", zone_name)
+        if await cloudflare.set_zone_under_attack(zone_id, True):
+            await webhook.send("Under Attack Activated", zone_name, color=0xFF0000)
+
+    elif not under_attack and currently_attacking:
+        logger.info("Deactivating Under Attack Mode for %s", zone_name)
+        if await cloudflare.set_zone_under_attack(zone_id, False):
+            await webhook.send("Under Attack Deactivated", zone_name, color=0x00FF00)
 
 
 async def main():
-    """
-    Main function to monitor network speed and manage Cloudflare zones.
-    """
-    if not config.get("CLOUDFLARE", "ZONE_IDS"):
-        print(f"{PREFIX}No Cloudflare zones found.")
+    """Main monitoring loop."""
+    config = get_config()
+    cloudflare = Cloudflare()
+    webhook = Webhook()
+    redis_cache = RedisCache()
+
+    zone_ids = config.get("CLOUDFLARE", "ZONE_IDS")
+    if not zone_ids:
+        logger.critical("No zone IDs configured.")
         sys.exit(1)
 
-    if await redis.check():
-        print(f"{PREFIX}Connected to Redis.")
-    else:
-        print(f"{PREFIX}Failed to connect to Redis.")
+    if not await redis_cache.check():
+        logger.critical("Cannot connect to Redis.")
         sys.exit(1)
 
-    while True:
-        try:
-            if await redis.is_under_attack():
-                print(f"{PREFIX}Under Attack mode is active.")
-            else:
-                mb_received_per_sec, mb_sent_per_sec = get_network_speed()
-                print(
-                    f"{PREFIX}Received: {mb_received_per_sec:.2f} MB/s, "
-                    f"Sent: {mb_sent_per_sec:.2f} MB/s"
-                )
-
-                if mb_received_per_sec > config.get("SETTINGS", "MAX_INCOMING_TRAFFIC_MB"):
-                    print(
-                        f"{PREFIX}Incoming traffic exceeded "
-                        f"{config.get('SETTINGS', 'MAX_INCOMING_TRAFFIC_MB')} MB/s"
-                    )
-
-                    for zone_id in config.get("CLOUDFLARE", "ZONE_IDS"):
-                        await handle_zone(zone_id, True)
-
-                    await redis.set_under_attack()
-
-                else:
-                    print(f"{PREFIX}Incoming traffic is normal.")
-                    for zone_id in config.get("CLOUDFLARE", "ZONE_IDS"):
-                        await handle_zone(zone_id, False)
-
-        except (aiohttp.ClientError, ValueError) as e:
-            print(f"{PREFIX}Error encountered in main loop: {e}")
+    check_interval = config.get("SETTINGS", "CHECK_INTERVAL") or 60
+    threshold = config.get("SETTINGS", "MAX_INCOMING_TRAFFIC_MB")
+    logger.info("Started | interval: %ds | threshold: %s MB/s | zones: %d",
+                check_interval, threshold, len(zone_ids))
 
-        await asyncio.sleep(config.get("SETTINGS", "CHECK_INTERVAL") or 60)
-
-
-async def handle_zone(zone_id, under_attack):
-    """
-    Handle Cloudflare zone based on the traffic condition.
-    """
     try:
-        zone = await cloudflare.get_zone(zone_id)
-        zone_name = zone["result"]["name"]
-
-        under_attack_mode = await cloudflare.get_zone_under_attack(zone_id)
-        if under_attack and under_attack_mode["result"]["value"] != "under_attack":
-            print(f"{PREFIX}Activating Under Attack Mode for {zone_name}")
-            await set_zone_under_attack(zone_id, zone_name, True)
-        elif not under_attack and under_attack_mode["result"]["value"] != "essentially_off":
-            print(f"{PREFIX}Deactivating Under Attack Mode for {zone_name}")
-            await set_zone_under_attack(zone_id, zone_name, False)
-
-    except (aiohttp.ClientError, ValueError) as e:
-        print(f"{PREFIX}Error handling Cloudflare zone {zone_id}: {e}")
+        while True:
+            try:
+                if await redis_cache.is_under_attack():
+                    remaining = await redis_cache.ttl()
+                    logger.info("Cooldown active, %ds remaining", remaining)
+                else:
+                    recv, sent = await get_network_speed()
+                    logger.info("Traffic – recv: %.2f MB/s  sent: %.2f MB/s", recv, sent)
 
+                    if recv > threshold:
+                        logger.warning("Traffic %.2f MB/s exceeded threshold %s MB/s",
+                                       recv, threshold)
+                        for zone_id in zone_ids:
+                            await handle_zone(cloudflare, webhook, zone_id, True)
+                        await redis_cache.set_under_attack()
+                    else:
+                        for zone_id in zone_ids:
+                            await handle_zone(cloudflare, webhook, zone_id, False)
 
-async def set_zone_under_attack(zone_id, zone_name, under_attack):
-    """
-    Set the Cloudflare zone under attack mode.
-    """
-    if config.get("SETTINGS", "LOGGING") and config.get("SETTINGS", "WEBHOOK"):
-        await webhook.send(
-            message=f"{'Activating' if under_attack else 'Deactivating'} "
-                    f"Under Attack Mode for {zone_name}",
-            color=0xFF0000 if under_attack else 0x00FF00
-        )
+            except Exception as exc:  # pylint: disable=broad-except
+                logger.error("Error in main loop: %s", exc)
 
-    await cloudflare.set_zone_under_attack(zone_id, under_attack)
+            await asyncio.sleep(max(0, check_interval - 1))
+    finally:
+        await cloudflare.close()
 
 
 if __name__ == "__main__":
-    print("\033[38;5;208m" + pyfiglet.figlet_format('NoAttack', font='doom') + Fore.RESET)
+    print("\033[38;5;208m" + pyfiglet.figlet_format("NoAttack", font="doom") + Fore.RESET)
 
-    if not config.config_exists():
-        config.create_config()
-        print(f"{PREFIX}Config file created.")
+    try:
+        get_config()
+    except ConfigNotFoundError as config_error:
+        print(f"{PREFIX}{config_error}")
+        sys.exit(1)
 
-    print(f"{PREFIX}Config file loaded.")
+    logging.basicConfig(
+        level=logging.INFO,
+        format="[%(asctime)s] [%(levelname)s] %(name)s: %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
     asyncio.run(main())

modules/cloudflare.py

@@ -1,88 +1,70 @@
+import logging
+
 import aiohttp
 
-from modules import config
+from modules.config import get_config
+
+logger = logging.getLogger("noattack.cloudflare")
 
 
 class Cloudflare:
     def __init__(self):
-        self.config = config.Config()
+        self._cfg = get_config()
+        self._session = None
 
-    async def get_header(self) -> dict:
-        """
-        Get headers for Cloudflare API.
-        :return: A dictionary containing the headers.
-        """
+    def _headers(self):
+        """Build the Cloudflare API auth headers."""
         return {
-            "Authorization": f"Bearer {self.config.get('CLOUDFLARE', 'API_KEY')}",
-            "Content-Type": "application/json"
+            "Authorization": f"Bearer {self._cfg.get('CLOUDFLARE', 'API_KEY')}",
+            "Content-Type": "application/json",
         }
 
-    async def get_zone(self, zone_id: str) -> dict:
-        """
-        Get a zone from Cloudflare.
-        :param zone_id: The ID of the zone.
-        :return: A dictionary containing the zone information.
-        """
-        url = f"https://api.cloudflare.com/client/v4/zones/{zone_id}"
+    async def _get_session(self):
+        """Return the shared aiohttp session, creating it if needed."""
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession()
+        return self._session
 
-        try:
-            async with aiohttp.ClientSession() as session:
-                async with session.get(
-                        url=url,
-                        headers=await self.get_header()
-                ) as response:
-                    return await response.json()
+    async def close(self):
+        """Close the shared HTTP session."""
+        if self._session and not self._session.closed:
+            await self._session.close()
 
-        except aiohttp.ClientConnectorError as e:
-            raise ConnectionError(f"Connection error: {e}") from e
-        except aiohttp.ClientResponseError as e:
-            raise ValueError(f"Response error: {e}") from e
+    async def get_zone(self, zone_id):
+        """Return the zone name for the given zone ID, or None on error."""
+        url = f"https://api.cloudflare.com/client/v4/zones/{zone_id}"
+        try:
+            session = await self._get_session()
+            async with session.get(url, headers=self._headers()) as resp:
+                resp.raise_for_status()
+                data = await resp.json()
+                return data["result"]["name"]
+        except (aiohttp.ClientError, KeyError) as e:
+            logger.error("Failed to fetch zone %s: %s", zone_id, e)
+            return None
 
-    async def get_zone_under_attack(self, zone_id: str) -> dict:
-        """
-        Get the security level of a zone from Cloudflare.
-        :param zone_id: The ID of the zone.
-        :return: A dictionary containing the security level information.
-        """
+    async def get_zone_under_attack(self, zone_id):
+        """Return True if the zone is in under_attack mode, False if not, None on error."""
         url = f"https://api.cloudflare.com/client/v4/zones/{zone_id}/settings/security_level"
-
         try:
-            async with aiohttp.ClientSession() as session:
-                async with session.get(
-                        url=url,
-                        headers=await self.get_header()
-                ) as response:
-                    return await response.json()
+            session = await self._get_session()
+            async with session.get(url, headers=self._headers()) as resp:
+                resp.raise_for_status()
+                data = await resp.json()
+                return data["result"]["value"] == "under_attack"
+        except (aiohttp.ClientError, KeyError) as e:
+            logger.error("Failed to get security level for zone %s: %s", zone_id, e)
+            return None
 
-        except aiohttp.ClientConnectorError as e:
-            raise ConnectionError(f"Connection error: {e}") from e
-
-        except aiohttp.ClientResponseError as e:
-            raise ValueError(f"Response error: {e}") from e
-
-    async def set_zone_under_attack(self, zone_id: str, mode: bool) -> dict:
-        """
-        Set a zone under attack mode.
-        :param zone_id: The ID of the zone.
-        :param mode: True to enable under attack mode, False to disable.
-        :return: A dictionary containing the response from Cloudflare.
-        """
+    async def set_zone_under_attack(self, zone_id, under_attack):
+        """Set the security level for a zone. Returns True on success."""
         url = f"https://api.cloudflare.com/client/v4/zones/{zone_id}/settings/security_level"
-        data = {
-            "value": "under_attack" if mode else "essentially_off"
-        }
-
+        payload = {"value": "under_attack" if under_attack else "essentially_off"}
         try:
-            async with aiohttp.ClientSession() as session:
-                async with session.patch(
-                        url=url,
-                        headers=await self.get_header(),
-                        json=data
-                ) as response:
-                    return await response.json()
-
-        except aiohttp.ClientConnectorError as e:
-            raise ConnectionError(f"Connection error: {e}") from e
-
-        except aiohttp.ClientResponseError as e:
-            raise ValueError(f"Response error: {e}") from e
+            session = await self._get_session()
+            async with session.patch(url, headers=self._headers(), json=payload) as resp:
+                resp.raise_for_status()
+                return True
+        except aiohttp.ClientError as e:
+            logger.error("Failed to set security level for zone %s: %s", zone_id, e)
+            return False