Optimization on the download of mmdb files

src/common/core/jobs/jobs/mmdb-asn.py

@@ -35,13 +35,25 @@ try:
     # Compute the mmdb URL
     mmdb_url = f"https://download.db-ip.com/free/dbip-asn-lite-{date.today().strftime('%Y-%m')}.mmdb.gz"
 
-    # Download the mmdb file
+    # Download the mmdb file and save it to tmp
     logger.info(f"Downloading mmdb file from url {mmdb_url} ...")
-    resp = get(mmdb_url)
+    file_content = b""
+    with get(mmdb_url, stream=True) as resp:
+        resp.raise_for_status()
+        for chunk in resp.iter_content(chunk_size=4 * 1024):
+            if chunk:
+                file_content += chunk
 
-    # Save it to temp
-    logger.info("Saving mmdb file to tmp ...")
-    Path(f"/var/tmp/bunkerweb/asn.mmdb").write_bytes(decompress(resp.content))
+    try:
+        assert file_content
+    except AssertionError:
+        logger.error(f"Error while downloading mmdb file from {mmdb_url}")
+        _exit(2)
+
+    # Decompress it
+    logger.info("Decompressing mmdb file ...")
+    file_content = decompress(file_content)
+    Path(f"/var/tmp/bunkerweb/asn.mmdb").write_bytes(file_content)
 
     # Try to load it
     logger.info("Checking if mmdb file is valid ...")
@@ -73,7 +85,7 @@ try:
     # Update db
     with lock:
         err = db.update_job_cache(
-            "mmdb-asn", None, "asn.mmdb", resp.content, checksum=new_hash
+            "mmdb-asn", None, "asn.mmdb", file_content, checksum=new_hash
         )
 
     if err:

src/common/core/jobs/jobs/mmdb-country.py

@@ -16,8 +16,8 @@ sys_path.extend(
     )
 )
 
-from requests import get
 from maxminddb import open_database
+from requests import get
 
 from Database import Database
 from logger import setup_logger
@@ -27,28 +27,6 @@ logger = setup_logger("JOBS.mmdb-country", getenv("LOG_LEVEL", "INFO"))
 status = 0
 
 try:
-    # Only download mmdb if the country blacklist or whitelist is enabled
-    dl_mmdb = False
-    # Multisite case
-    if getenv("MULTISITE", "no") == "yes":
-        for first_server in getenv("SERVER_NAME", "").split(" "):
-            if getenv(
-                f"{first_server}_BLACKLIST_COUNTRY", getenv("BLACKLIST_COUNTRY")
-            ) or getenv(
-                f"{first_server}_WHITELIST_COUNTRY", getenv("WHITELIST_COUNTRY")
-            ):
-                dl_mmdb = True
-                break
-    # Singlesite case
-    elif getenv("BLACKLIST_COUNTRY") or getenv("WHITELIST_COUNTRY"):
-        dl_mmdb = True
-
-    if not dl_mmdb:
-        logger.info(
-            "Country blacklist or whitelist is not enabled, skipping download..."
-        )
-        _exit(0)
-
     # Don't go further if the cache is fresh
     if is_cached_file("/var/cache/bunkerweb/country.mmdb", "month"):
         logger.info("country.mmdb is already in cache, skipping download...")
@@ -57,13 +35,25 @@ try:
     # Compute the mmdb URL
     mmdb_url = f"https://download.db-ip.com/free/dbip-country-lite-{date.today().strftime('%Y-%m')}.mmdb.gz"
 
-    # Download the mmdb file
+    # Download the mmdb file and save it to tmp
     logger.info(f"Downloading mmdb file from url {mmdb_url} ...")
-    resp = get(mmdb_url)
+    file_content = b""
+    with get(mmdb_url, stream=True) as resp:
+        resp.raise_for_status()
+        for chunk in resp.iter_content(chunk_size=4 * 1024):
+            if chunk:
+                file_content += chunk
 
-    # Save it to temp
-    logger.info("Saving mmdb file to tmp ...")
-    Path(f"/var/tmp/bunkerweb/country.mmdb").write_bytes(decompress(resp.content))
+    try:
+        assert file_content
+    except AssertionError:
+        logger.error(f"Error while downloading mmdb file from {mmdb_url}")
+        _exit(2)
+
+    # Decompress it
+    logger.info("Decompressing mmdb file ...")
+    file_content = decompress(file_content)
+    Path(f"/var/tmp/bunkerweb/country.mmdb").write_bytes(file_content)
 
     # Try to load it
     logger.info("Checking if mmdb file is valid ...")
@@ -80,9 +70,7 @@ try:
     # Move it to cache folder
     logger.info("Moving mmdb file to cache ...")
     cached, err = cache_file(
-        "/var/tmp/bunkerweb/country.mmdb",
-        "/var/cache/bunkerweb/country.mmdb",
-        new_hash,
+        "/var/tmp/bunkerweb/country.mmdb", "/var/cache/bunkerweb/country.mmdb", new_hash
     )
     if not cached:
         logger.error(f"Error while caching mmdb file : {err}")
@@ -97,7 +85,7 @@ try:
     # Update db
     with lock:
         err = db.update_job_cache(
-            "mmdb-country", None, "country.mmdb", resp.content, checksum=new_hash
+            "mmdb-country", None, "country.mmdb", file_content, checksum=new_hash
         )
 
     if err: