diff options
| -rw-r--r-- | src/soc_collector/healthcheck.py | 4 | ||||
| -rw-r--r-- | src/soc_collector/soc_collector_cli.py | 53 |
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: |