fixed healthcheck and client CLI small stuff

2 files changed, 48 insertions, 9 deletions

diff --git a/src/soc_collector/healthcheck.py b/src/soc_collector/healthcheck.py
index e561fa0..1ccd8f3 100644
--- a/src/soc_collector/healthcheck.py
+++ b/src/soc_collector/healthcheck.py
@@ -7,16 +7,18 @@ import json
 
 import requests
 
+from src.soc_collector.auth import load_api_keys
 
 def check_collector() -> bool:
     """Check our collector using /info
 
     :return: bool
     """
-    time.sleep(2)  # Prevent race condition with redis container healthcheck
+    time.sleep(2)  # Prevent race condition with mongodb container healthcheck
 
     req = requests.get(
         "https://localhost:8000/info",
+        headers={"API-KEY": load_api_keys("./api_keys.txt")[-1]},
         timeout=3,
         verify="./collector_root_ca.crt",
     )
diff --git a/src/soc_collector/soc_collector_cli.py b/src/soc_collector/soc_collector_cli.py
index 2398e4f..d7add30 100644
--- a/src/soc_collector/soc_collector_cli.py
+++ b/src/soc_collector/soc_collector_cli.py
@@ -17,6 +17,22 @@ API_URL = "https://collector-dev.soc.sunet.se:8000"
 ROOT_CA_FILE = __file__.replace("soc_collector_cli.py", "data/collector_root_ca.crt")
 
 
+def valid_key(key: str) -> None:
+    """Ensure the document key is valid. exit(1) otherwise.
+
+    :param key: The key.
+    """
+    valid_chars = ["a", "b", "c", "d", "e", "f", "1", "2", "3", "4", "5", "6", "7", "8", "9", "0"]
+    if len(key) != 24:
+        print(f"ERROR: Invalid key '{key}'")
+        app_exit(1)
+
+    for char in key:
+        if char not in valid_chars:
+            print(f"ERROR: Invalid key '{key}'")
+            app_exit(1)
+
+
 def json_load_data(data: str) -> Dict[str, Any]:
     """Load json from argument, json data or path to json file
 
@@ -76,10 +92,19 @@ def delete_action(data: str) -> None:
 
     :param data: key or path to a json file containing "_id".
     """
-
     if data and isfile(data):
-        data = json_load_data(data)["_id"]
-    req = requests.delete(f"{API_URL}/sc/v0/{data}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)
+        json_data = json_load_data(data)
+
+        if "_id" not in json_data or not isinstance(json_data["_id"], str):
+            print("ERROR: Valid '_id' key not in data")
+            app_exit(1)
+        key: str = json_data["_id"]
+    else:
+        key = data
+
+    valid_key(key)
+
+    req = requests.delete(f"{API_URL}/sc/v0/{key}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)
 
     # Check status
     if req.status_code == 404:
@@ -89,7 +114,7 @@ def delete_action(data: str) -> None:
     # Ensure ok status
     req.raise_for_status()
 
-    print(f"Deleted data OK - key: {data}")
+    print(f"Deleted data OK - key: {key}")
 
 
 def update_local_action(data: str, update_data: str) -> None:
@@ -116,10 +141,12 @@ def replace_action(data: str) -> None:
 
     json_data = json_load_data(data)
 
-    if "_id" not in json_data:
-        print("ERROR: '_id' key not in data")
+    if "_id" not in json_data or not isinstance(json_data["_id"], str):
+        print("ERROR: Valid '_id' key not in data")
         app_exit(1)
 
+    valid_key(json_data["_id"])
+
     req = requests.put(f"{API_URL}/sc/v0", json=json_data, headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)
 
     # Check status
@@ -163,8 +190,18 @@ def get_action(data: str) -> None:
     :param data: key or path to a json file containing "_id".
     """
     if data and isfile(data):
-        data = json_load_data(data)["_id"]
-    req = requests.get(f"{API_URL}/sc/v0/{data}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)
+        json_data = json_load_data(data)
+
+        if "_id" not in json_data or not isinstance(json_data["_id"], str):
+            print("ERROR: Valid '_id' key not in data")
+            app_exit(1)
+        key: str = json_data["_id"]
+    else:
+        key = data
+
+    valid_key(key)
+
+    req = requests.get(f"{API_URL}/sc/v0/{key}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)
 
     # Check status
     if req.status_code == 404: