Skip to content

Commit 507219d

Browse files
committed
Add support for IP66 GeoIP integration and GeoIP source selection
- Introduce `IP66GeoIP` with daily file download capability. - Add `source` field in `GeoIPConfiguration` for dynamic source selection between MaxMind and IP66. - Refactor GeoIP initialization to support multiple enrichment sources. - Implement `JobRegistry` to manage background jobs in the Engine. - Enhance GeoIP handling for better configurability and performance.
1 parent a71d766 commit 507219d

6 files changed

Lines changed: 235 additions & 9 deletions

File tree

mongoose/core/engine.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import logging
22
import threading
3-
from typing import List, Optional, Type, Dict
3+
from typing import List, Optional, Type, Dict, Set, Any
44

55
import yaml
66

@@ -44,6 +44,17 @@ def __call__(cls, *args, **kwargs):
4444
return cls._instances[cls]
4545

4646

47+
class JobRegistry(metaclass=Singleton):
48+
def __init__(self):
49+
self.jobs: Set[Any] = set()
50+
51+
def register(self, job: Any):
52+
self.jobs.add(job)
53+
54+
def clear(self):
55+
self.jobs.clear()
56+
57+
4758
class Engine(metaclass=Singleton):
4859
"""
4960
The Engine class is responsible for loading the configuration,
@@ -57,6 +68,7 @@ def __init__(self, config_path: str, interface: str = None, watch_configuration_
5768
self.processing_queue = ProcessingQueue()
5869
self.collectors: List = []
5970
self.forwarders: List = []
71+
self.job_registry = JobRegistry()
6072
self.enrichment: Optional[Enrich] = None
6173
self.database_storage = None
6274
self.sink: Sink = Sink()
@@ -193,6 +205,9 @@ def start(self):
193205
if self.watch_configuration_changes:
194206
self.webhook_configuration_watcher.run()
195207

208+
for job in self.job_registry.jobs:
209+
job.start()
210+
196211
self.sink.start()
197212
self.database_storage.start()
198213

@@ -218,6 +233,9 @@ def stop(self):
218233

219234
self.processing_queue.stop_processing()
220235

236+
for job in self.job_registry.jobs:
237+
job.stop()
238+
221239
if self.watch_configuration_changes:
222240
self.webhook_configuration_watcher.stop()
223241

mongoose/enrich/__init__.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
1+
from .base import Enrich
12
from .community_id import CommunityIDEnrichment
23
from .direction import DirectionEnrichment
3-
from .geoip import GeoIP
4+
from .geoip import MaxMindGeoIP, IP66GeoIP
45
from .hostname import HostnameEnrichment
5-
from .base import Enrich
66

77
__all__ = [
88
"Enrich",
99
"CommunityIDEnrichment",
1010
"DirectionEnrichment",
11-
"GeoIP",
11+
"IP66GeoIP",
12+
"MaxMindGeoIP",
1213
"HostnameEnrichment",
1314
]

mongoose/enrich/base.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from typing import Optional
66

77
from mongoose.core.processing import ProcessingQueue, ProcessingTopic
8-
from mongoose.enrich import DirectionEnrichment, CommunityIDEnrichment, HostnameEnrichment, GeoIP
8+
from mongoose.enrich import *
99
from mongoose.enrich.risk import FlowRiskEnrichment
1010
from mongoose.enrich.type import EventTypeEnrichment
1111
from mongoose.models import NetworkDPI, NetworkAlert, NetworkFlow
@@ -27,11 +27,15 @@ def __init__(self, enrichment_configuration: EnrichmentConfiguration):
2727
EventTypeEnrichment(),
2828
FlowRiskEnrichment(),
2929
]
30-
self.geoip_enrichment = (
31-
GeoIP(enrichment_configuration.geoip)
30+
geoip_source = (
31+
enrichment_configuration.geoip.source or "maxmind"
3232
if enrichment_configuration.geoip and enrichment_configuration.geoip.enable
3333
else None
3434
)
35+
if geoip_source and geoip_source.lower() == "maxmind":
36+
self.geoip_enrichment = MaxMindGeoIP(enrichment_configuration.geoip)
37+
else:
38+
self.geoip_enrichment = IP66GeoIP(enrichment_configuration.geoip)
3539

3640
def start(self):
3741
if self.thread and self.thread.is_alive():

mongoose/enrich/geoip.py

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,74 @@
22
import logging
33
from functools import lru_cache
44
from typing import Union, Any, Dict, Optional
5+
import geoip2
56
import geoip2.database
67

7-
import geoip2
8+
import maxminddb
89

10+
from mongoose.core.engine import JobRegistry
911
from mongoose.models import NetworkDPI, NetworkFlow, NetworkAlert
1012
from mongoose.models.configuration import GeoIPConfiguration
1113
from mongoose.utils.exceptions import IgnoreCacheException
14+
from mongoose.utils.jobs import DailyFileDownloadJob
1215

1316
logger = logging.getLogger(__name__)
1417

1518

16-
class GeoIP:
19+
class IP66GeoIP:
20+
database_filename = "ip66.mmdb"
21+
database_url = "https://downloads.ip66.dev/db/ip66.mmdb"
22+
download_job = None
23+
24+
def __init__(self, geoip_configuration: GeoIPConfiguration):
25+
self.geoip_configuration = geoip_configuration
26+
self.database_path = self.geoip_configuration.maxmind_db_path / self.database_filename
27+
if not self.download_job:
28+
self.download_job = DailyFileDownloadJob(url=self.database_url, local_path=self.database_path)
29+
JobRegistry().register(self.download_job)
30+
31+
@lru_cache(maxsize=512)
32+
def request_geoip(self, ip_address: str) -> Optional[Dict[Any, Any]]:
33+
reader = maxminddb.open_database(str(self.database_path))
34+
geoip_data = {}
35+
record = reader.get(ip_address)
36+
if not record:
37+
raise IgnoreCacheException() # prevents caching
38+
39+
geoip_data["details"] = record.get("anonymous_ip", None)
40+
geoip_data["traits"] = record.get("traits", None)
41+
geoip_data["asn"] = record.get("autonomous_system_number")
42+
geoip_data["organization"] = record.get("autonomous_system_organization")
43+
geoip_data["country"] = record.get("country", {}).get("iso_code")
44+
geoip_data["country_name"] = record.get("country", {}).get("names", {}).get("en")
45+
geoip_data["continent"] = record.get("continent", {}).get("code")
46+
geoip_data["continent_name"] = record.get("continent", {}).get("names", {}).get("en")
47+
48+
if geoip_data:
49+
return geoip_data
50+
raise IgnoreCacheException() # prevents caching
51+
52+
def enrich_network_event(self, event: Union[NetworkDPI, NetworkFlow, NetworkAlert]):
53+
if not hasattr(event, "src_ip") or not hasattr(event, "dst_ip"):
54+
return
55+
56+
src_ip = ipaddress.ip_address(event.src_ip)
57+
dst_ip = ipaddress.ip_address(event.dst_ip)
58+
if src_ip.is_global:
59+
try:
60+
event.enrichment["geoip"] = self.request_geoip(event.src_ip)
61+
event.enrichment["geoip"]["ip"] = event.src_ip
62+
except (Exception,):
63+
pass
64+
elif dst_ip.is_global:
65+
try:
66+
event.enrichment["geoip"] = self.request_geoip(event.dst_ip)
67+
event.enrichment["geoip"]["ip"] = event.dst_ip
68+
except (Exception,):
69+
pass
70+
71+
72+
class MaxMindGeoIP:
1773
def __init__(self, geoip_configuration: GeoIPConfiguration):
1874
self.geoip_configuration = geoip_configuration
1975
self.databases = []

mongoose/models/configuration.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,9 @@ class GeoIPConfiguration(BaseModel):
189189
maxmind_db: List[str] = ["GeoLite2-ASN.mmdb", "GeoLite2-City.mmdb", "GeoLite2-Country.mmdb"]
190190
"""The list of GeoIP databases to use."""
191191

192+
source: str = "ip66" # maxmind or ip66
193+
"""The source to use, either MaxMind or IP66. Defaults to IP66."""
194+
192195
enable: bool = Field(default=True)
193196
"""Enable the GeoIP enrichment. Defaults to True."""
194197

mongoose/utils/jobs.py

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
import logging
2+
import os
3+
import threading
4+
import urllib.error
5+
import urllib.request
6+
from datetime import datetime, timedelta
7+
from pathlib import Path
8+
9+
logger = logging.getLogger(__name__)
10+
11+
12+
class DailyFileDownloadJob:
13+
def __init__(self, url, local_path, run_time="09:00"):
14+
"""
15+
Initialize a daily file download job.
16+
17+
Args:
18+
url: URL to download the file from
19+
local_path: Local path to save the file
20+
run_time: Time to run daily download in "HH:MM" format (24-hour)
21+
"""
22+
self.url = url
23+
self.local_path: Path = Path(local_path)
24+
self.local_path.parent.mkdir(parents=True, exist_ok=True)
25+
self.run_time = self._parse_run_time(run_time)
26+
self.stop_event = threading.Event()
27+
self.thread = None
28+
29+
def _parse_run_time(self, run_time_str):
30+
"""Parse run time string to hours and minutes."""
31+
try:
32+
hour, minute = map(int, run_time_str.split(':'))
33+
if not (0 <= hour <= 23 and 0 <= minute <= 59):
34+
raise ValueError("Invalid time range")
35+
return hour, minute
36+
except (ValueError, AttributeError):
37+
logger.error(f"Invalid run_time format: {run_time_str}. Must be in 'HH:MM' format")
38+
raise ValueError("run_time must be in 'HH:MM' format")
39+
40+
def _get_next_run_time(self):
41+
"""Calculate the next run time based on the current time."""
42+
now = datetime.now()
43+
run_hour, run_minute = self.run_time
44+
next_run = now.replace(hour=run_hour, minute=run_minute, second=0, microsecond=0)
45+
46+
# If today's run time has passed, schedule for tomorrow
47+
if next_run <= now:
48+
next_run += timedelta(days=1)
49+
50+
return next_run
51+
52+
def _file_needs_update(self):
53+
"""Check if the file needs to be downloaded (doesn't exist or is older than a day)."""
54+
try:
55+
file_mod_time = datetime.fromtimestamp(os.path.getmtime(self.local_path))
56+
one_day_ago = datetime.now() - timedelta(days=1)
57+
needs_update = file_mod_time < one_day_ago
58+
if needs_update:
59+
logger.info(f"File {self.local_path} is older than 24 hours, needs update")
60+
else:
61+
logger.debug(f"File {self.local_path} is up to date")
62+
return needs_update
63+
except OSError as e:
64+
logger.warning(f"Could not get file modification time for {self.local_path}: {e}")
65+
# If we can't get file info, assume it needs to be updated
66+
return True
67+
68+
def _download_file(self):
69+
"""Download the file from URL to the local path."""
70+
try:
71+
logger.info(f"Starting download from {self.url} to {self.local_path}")
72+
urllib.request.urlretrieve(self.url, self.local_path)
73+
logger.info(f"Download completed successfully to {self.local_path}")
74+
return True
75+
except urllib.error.URLError as e:
76+
logger.error(f"Network error during download from {self.url}: {e}")
77+
return False
78+
except Exception as e:
79+
logger.error(f"Unexpected error during download from {self.url}: {e}")
80+
return False
81+
82+
def _run_immediate_download_if_needed(self):
83+
"""Check and perform immediate download if the file needs update."""
84+
if self._file_needs_update():
85+
logger.info("Performing immediate download due to outdated or missing file")
86+
return self._download_file()
87+
return True
88+
89+
def _run_loop(self):
90+
"""Main execution loop."""
91+
logger.info("Daily file download job started")
92+
93+
# First, check if the immediate download is needed
94+
self._run_immediate_download_if_needed()
95+
96+
while not self.stop_event.is_set():
97+
next_run = self._get_next_run_time()
98+
sleep_time = (next_run - datetime.now()).total_seconds()
99+
100+
logger.info(f"Next scheduled download: {next_run}")
101+
102+
# Wait until the next run time or the stop event
103+
if self.stop_event.wait(timeout=sleep_time):
104+
logger.info("Stop event received, exiting download loop")
105+
break
106+
107+
# Execute daily download if not stopping
108+
if not self.stop_event.is_set():
109+
self._download_file()
110+
111+
def start(self):
112+
"""Start the daily download job in a background thread."""
113+
if self.thread is not None and self.thread.is_alive():
114+
logger.warning("Download job is already running")
115+
return
116+
117+
self.stop_event.clear()
118+
self.thread = threading.Thread(target=self._run_loop, daemon=False)
119+
self.thread.start()
120+
logger.info(f"Daily file download job started for {self.url}")
121+
122+
def stop(self, timeout=30):
123+
"""
124+
Stop the download job gracefully.
125+
126+
Args:
127+
timeout: Maximum time to wait for graceful shutdown (seconds)
128+
"""
129+
if self.thread is None or not self.thread.is_alive():
130+
logger.debug("Download job is not running")
131+
return True
132+
133+
logger.info("Stopping daily download job...")
134+
self.stop_event.set()
135+
136+
# Wait for the thread to finish gracefully
137+
self.thread.join(timeout=timeout)
138+
139+
if self.thread.is_alive():
140+
logger.warning("Download job did not stop gracefully within timeout")
141+
return False
142+
else:
143+
logger.info("Daily download job stopped successfully")
144+
return True

0 commit comments

Comments
 (0)