Merge pull request #1184 from curiefense/ipinfo-fixes

curiefense · Mar 5, 2023 · e448bf3 · e448bf3
2 parents 4263010 + 25ceadb
commit e448bf3
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 14 deletions.
diff --git a/curiefense/curieconf/client/curieconf/cli/__init__.py b/curiefense/curieconf/client/curieconf/cli/__init__.py
@@ -9,6 +9,7 @@
 import json
 import subprocess
 from enum import Enum
+from pathlib import Path
 from google.cloud import storage
 import base64
 from typing import List, Optional
@@ -496,26 +497,35 @@ def pullipinfo(project: str, bucket: str, ipinfo_dir: str, target_path: str):
     os.makedirs(target_path, exist_ok=True)
 
     for ipinfo_blob in [*(client.list_blobs(bucket_or_name=bucket, prefix=ipinfo_dir))]:
-        file_name = ipinfo_blob.name.split("/")[-1]
+        remote_md5 = base64.b64decode(ipinfo_blob.md5_hash).hex()
+        path = Path(ipinfo_blob.name)
+        file_name = path.name
+        hash_file_name = target_path + path.stem + "_hash"
         if not os.path.isfile(target_path + file_name):
             try:
                 ipinfo_blob.download_to_filename(target_path + file_name)
+                with open(hash_file_name, "w") as hash_file:
+                    hash_file.write(remote_md5)
             except Exception as ex:
-                logger.error(f"failed downloading {file_name} - {ex}")
+                typer.echo(f"failed downloading {file_name} - {ex}", err=True)
                 continue
-            logger.info(f"downloaded {file_name} to {target_path}")
+            typer.echo(f"downloaded {file_name} to {target_path}")
         else:
-            remote_md5 = base64.b64decode(ipinfo_blob.md5_hash).hex()
-            local_md5 = get_md5(target_path + file_name)
+            with open(hash_file_name, "r") as hash_file:
+                local_md5 = hash_file.readline().rstrip()
+
             if not local_md5 == remote_md5:
-                try:
-                    ipinfo_blob.download_to_filename(target_path + file_name)
-                except Exception as ex:
-                    logger.error(f"failed downloading {file_name} - {ex}")
-                    continue
+                with open(hash_file_name, "w") as hash_file:
+                    try:
+                        ipinfo_blob.download_to_filename(target_path + file_name)
+                        hash_file.write(remote_md5)
+                        typer.echo(f"updated {file_name} in {target_path}")
+                    except Exception as ex:
+                        typer.echo(f"failed downloading {file_name} - {ex}", err=True)
+                        continue
 
             else:
-                logger.info(f"{target_path} already contains the updated {file_name}")
+                typer.echo(f"{target_path} already contains the updated {file_name}")
 
 
 @sync.command()

diff --git a/curiefense/images/curiesync/init/pull.sh b/curiefense/images/curiesync/init/pull.sh
@@ -20,8 +20,10 @@ fi
 if [ "$RUN_MODE" = "SYNC_ONCE" ]; then
     info "Synchronizing once"
     curieconfctl sync pull "${CURIE_BUCKET_LINK}" /cf-config
-    curieconfctl sync pullipinfo rbz-internal rbz-dev-auto-acl ipinfo /cf-config/ipinfo/
-
+    if [ "$USE_IPINFO" = "true" ]
+    then
+      curieconfctl sync pullipinfo "${IPINFO_PROJECT}" "${IPINFO_REMOTE_BUCKET}" "${IPINFO_REMOTE_DIR}" /cf-config/ipinfo/
+    fi
     exit 0
 fi
 
@@ -49,8 +51,11 @@ if [ "$RUN_MODE" = "PERIODIC_SYNC" ] || [ -z "$RUN_MODE" ]; then
     while :;
     do
         info "Pulling ${CURIE_BUCKET_LINK}"
-        curieconfctl sync pullipinfo rbz-internal rbz-dev-auto-acl ipinfo /cf-config/ipinfo/
         curieconfctl sync pull "${CURIE_BUCKET_LINK}" /cf-config --on-conf-change "$CONFIGCHANGE"
+        if [ "$USE_IPINFO" = "true" ]
+        then
+          curieconfctl sync pullipinfo "${IPINFO_PROJECT}" "${IPINFO_REMOTE_BUCKET}" "${IPINFO_REMOTE_DIR}" /cf-config/ipinfo/
+        fi
         info "Sleeping"
         sleep $PERIOD
     done