Replace custom Singleton and SingletonMeta implementations with a unified design in mongoose.core, streamline imports, and adjust GeoIP configuration defaults.

Author: U039b

mongoose/core/__init__.py

@@ -0,0 +1,48 @@
+import threading
+from typing import Dict, Any
+
+
+class SingletonMeta(type):
+    """
+    Thread-safe metaclass implementing a per-class singleton.
+
+    Each class that uses this metaclass will only ever have a single
+    instance created. The first construction's args/kwargs are used to
+    initialize the singleton; later constructions return the same
+    instance and ignore new args.
+    """
+
+    _instances: Dict[type, Any] = {}
+    _lock: threading.Lock = threading.Lock()
+
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+        # Double-checked locking to avoid acquiring the lock every time
+        if cls not in cls._instances:
+            with cls._lock:
+                if cls not in cls._instances:
+                    instance = super().__call__(*args, **kwargs)
+                    cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class Singleton(type):
+    _instances: Dict[type, type] = {}
+
+    def __call__(cls, *args, **kwargs):
+        """Control instance creation to ensure singleton behavior.
+
+        Args:
+            cls (type): The class being instantiated
+            *args: Positional arguments for class initialization
+            **kwargs: Keyword arguments for class initialization
+
+        Returns:
+            type: The singleton instance of the class
+
+        Note:
+            If an instance already exists, ``__init__`` will still be called with
+            the provided arguments, but no new instance is created.
+        """
+        if cls not in cls._instances:
+            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        return cls._instances[cls]

mongoose/core/cache.py

@@ -1,35 +1,14 @@
-from typing import Optional, Generic, TypeVar, Iterator, Dict, Any, Tuple, List
-from collections import OrderedDict
 import threading
 import time
+from collections import OrderedDict
+from typing import Optional, Generic, TypeVar, Iterator, Dict, Any, List
+
+from mongoose.core import SingletonMeta
 
 K = TypeVar("K")
 V = TypeVar("V")
 
 
-class SingletonMeta(type):
-    """
-    Thread-safe metaclass implementing a per-class singleton.
-
-    Each class that uses this metaclass will only ever have a single
-    instance created. The first construction's args/kwargs are used to
-    initialize the singleton; later constructions return the same
-    instance and ignore new args.
-    """
-
-    _instances: Dict[type, Any] = {}
-    _lock: threading.Lock = threading.Lock()
-
-    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
-        # Double-checked locking to avoid acquiring the lock every time
-        if cls not in cls._instances:
-            with cls._lock:
-                if cls not in cls._instances:
-                    instance = super().__call__(*args, **kwargs)
-                    cls._instances[cls] = instance
-        return cls._instances[cls]
-
-
 class Cache(Generic[K, V], metaclass=SingletonMeta):
     """
     Sharded LRU cache with optional TTL implemented as a singleton per class.

mongoose/core/engine.py

@@ -1,14 +1,16 @@
 import logging
 import threading
-from typing import List, Optional, Type, Dict, Set, Any
+from typing import List, Optional, Type
 
 import yaml
 
 import mongoose.core.cache as cache_module
 from mongoose.collect.nfstream_collector import NFStreamCollector
 from mongoose.collect.suricata_eve_collector import SuricataEveCollector
+from mongoose.core import Singleton
 from mongoose.core.cache import SeverityCache
 from mongoose.core.processing import ProcessingQueue
+from mongoose.core.registry import JobRegistry
 from mongoose.core.sink import Sink
 from mongoose.core.watchdogs import DropInConfigurationWatcher, DropInConfigurationHandler
 from mongoose.enrich.base import Enrich
@@ -21,40 +23,6 @@
 logger = logging.getLogger(__name__)
 
 
-class Singleton(type):
-    _instances: Dict[type, type] = {}
-
-    def __call__(cls, *args, **kwargs):
-        """Control instance creation to ensure singleton behavior.
-
-        Args:
-            cls (type): The class being instantiated
-            *args: Positional arguments for class initialization
-            **kwargs: Keyword arguments for class initialization
-
-        Returns:
-            type: The singleton instance of the class
-
-        Note:
-            If an instance already exists, ``__init__`` will still be called with
-            the provided arguments, but no new instance is created.
-        """
-        if cls not in cls._instances:
-            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
-        return cls._instances[cls]
-
-
-class JobRegistry(metaclass=Singleton):
-    def __init__(self):
-        self.jobs: Set[Any] = set()
-
-    def register(self, job: Any):
-        self.jobs.add(job)
-
-    def clear(self):
-        self.jobs.clear()
-
-
 class Engine(metaclass=Singleton):
     """
     The Engine class is responsible for loading the configuration,
@@ -78,9 +46,13 @@ def __init__(self, config_path: str, interface: str = None, watch_configuration_
         self.started = False
         self.thread_id = threading.get_ident()
         self.watch_configuration_changes = watch_configuration_changes
-        self.webhook_configuration_watcher = DropInConfigurationWatcher(
-            self.config.extra_configuration_dir / "webhook.d",
-            DropInConfigurationHandler(WebhookForwarderConfiguration, self._handle_webhook_configuration_changes),
+        self.webhook_configuration_watcher = (
+            DropInConfigurationWatcher(
+                self.config.extra_configuration_dir / "webhook.d",
+                DropInConfigurationHandler(WebhookForwarderConfiguration, self._handle_webhook_configuration_changes),
+            )
+            if self.watch_configuration_changes
+            else None
         )
 
     def load_extra_config(self, name: str, config_class: Type):

mongoose/core/registry.py

@@ -0,0 +1,14 @@
+from typing import Set, Any
+
+from mongoose.core import Singleton
+
+
+class JobRegistry(metaclass=Singleton):
+    def __init__(self):
+        self.jobs: Set[Any] = set()
+
+    def register(self, job: Any):
+        self.jobs.add(job)
+
+    def clear(self):
+        self.jobs.clear()

mongoose/enrich/__init__.py

@@ -1,14 +0,0 @@
-from .base import Enrich
-from .community_id import CommunityIDEnrichment
-from .direction import DirectionEnrichment
-from .geoip import MaxMindGeoIP, IP66GeoIP
-from .hostname import HostnameEnrichment
-
-__all__ = [
-    "Enrich",
-    "CommunityIDEnrichment",
-    "DirectionEnrichment",
-    "IP66GeoIP",
-    "MaxMindGeoIP",
-    "HostnameEnrichment",
-]

mongoose/enrich/base.py

@@ -5,7 +5,10 @@
 from typing import Optional
 
 from mongoose.core.processing import ProcessingQueue, ProcessingTopic
-from mongoose.enrich import *
+from mongoose.enrich.community_id import CommunityIDEnrichment
+from mongoose.enrich.direction import DirectionEnrichment
+from mongoose.enrich.geoip import MaxMindGeoIP, IP66GeoIP
+from mongoose.enrich.hostname import HostnameEnrichment
 from mongoose.enrich.risk import FlowRiskEnrichment
 from mongoose.enrich.type import EventTypeEnrichment
 from mongoose.models import NetworkDPI, NetworkAlert, NetworkFlow
@@ -34,7 +37,7 @@ def __init__(self, enrichment_configuration: EnrichmentConfiguration):
         )
         if geoip_source and geoip_source.lower() == "maxmind":
             self.geoip_enrichment = MaxMindGeoIP(enrichment_configuration.geoip)
-        else:
+        elif geoip_source and geoip_source.lower() == "ip66":
             self.geoip_enrichment = IP66GeoIP(enrichment_configuration.geoip)
 
     def start(self):

mongoose/enrich/geoip.py

@@ -1,13 +1,13 @@
-import ipaddress
 import logging
+import ipaddress
 from functools import lru_cache
-from typing import Union, Any, Dict, Optional
+from typing import Union, Optional, Dict, Any
+
 import geoip2
 import geoip2.database
-
 import maxminddb
 
-from mongoose.core.engine import JobRegistry
+from mongoose.core.registry import JobRegistry
 from mongoose.models import NetworkDPI, NetworkFlow, NetworkAlert
 from mongoose.models.configuration import GeoIPConfiguration
 from mongoose.utils.exceptions import IgnoreCacheException

mongoose/models/configuration.py

@@ -189,11 +189,11 @@ class GeoIPConfiguration(BaseModel):
     maxmind_db: List[str] = ["GeoLite2-ASN.mmdb", "GeoLite2-City.mmdb", "GeoLite2-Country.mmdb"]
     """The list of GeoIP databases to use."""
 
-    source: str = "ip66"  # maxmind or ip66
+    source: str = "maxmind"  # maxmind or ip66
     """The source to use, either MaxMind or IP66. Defaults to IP66."""
 
-    enable: bool = Field(default=True)
-    """Enable the GeoIP enrichment. Defaults to True."""
+    enable: bool = Field(default=False)
+    """Enable the GeoIP enrichment. Defaults to False."""
 
 
 class EnrichmentConfiguration(BaseModel):

mongoose/utils/jobs.py

@@ -29,7 +29,7 @@ def __init__(self, url, local_path, run_time="09:00"):
     def _parse_run_time(self, run_time_str):
         """Parse run time string to hours and minutes."""
         try:
-            hour, minute = map(int, run_time_str.split(':'))
+            hour, minute = map(int, run_time_str.split(":"))
             if not (0 <= hour <= 23 and 0 <= minute <= 59):
                 raise ValueError("Invalid time range")
             return hour, minute

tests/test_engine.py

@@ -13,7 +13,7 @@ def test_engine_initialization(tmp_path):
                 "suricata": {"socket_path": "/tmp/suricata.socket", "enable": True},
                 "nf_stream": {"interface": "eth0", "enable": False},
             },
-            "enrichment": {"geoip": {"remote_service_url": "http://geoip", "enable": True}},
+            "enrichment": {"geoip": {"enable": False}},
             "forwarder": {
                 "file": {"output_dir": str(tmp_path / "output"), "enable": True},
                 "webhooks": [{"url": "http://webhook", "enable": False}],
@@ -31,17 +31,16 @@ def test_engine_initialization(tmp_path):
 
 
 @patch("mongoose.collect.suricata_eve_collector.SuricataEveCollector.start")
-@patch("mongoose.enrich.Enrich.start")
+@patch("mongoose.enrich.base.Enrich.start")
 @patch("mongoose.forward.file.FileForwarder.start")
-@patch("mongoose.enrich.Enrich.__init__", return_value=None)
-def test_engine_start(mock_enrich_init, mock_file_start, mock_enrich_start, mock_suricata_start, tmp_path):
+def test_engine_start(mock_file_start, mock_enrich_start, mock_suricata_start, tmp_path):
     config_content = {
         "configuration": {
             "collector": {
                 "suricata": {"socket_path": "/tmp/suricata.socket", "enable": True},
                 "nf_stream": {"interface": "eth0", "enable": False},
             },
-            "enrichment": {"geoip": {"enable": True}},
+            "enrichment": {"geoip": {"enable": False}},
             "forwarder": {
                 "file": {"output_dir": str(tmp_path / "output"), "enable": True},
                 "webhooks": [{"url": "http://webhook", "enable": False}],