diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/soc_collector/db.py | 39 | ||||
| -rwxr-xr-x | src/soc_collector/main.py | 46 | ||||
| -rw-r--r-- | src/soc_collector/schema.py | 28 | ||||
| -rw-r--r-- | src/soc_collector/soc_collector_cli.py | 130 |
4 files changed, 131 insertions, 112 deletions
diff --git a/src/soc_collector/db.py b/src/soc_collector/db.py index b10d865..b1501d8 100644 --- a/src/soc_collector/db.py +++ b/src/soc_collector/db.py @@ -7,7 +7,6 @@ from dataclasses import dataclass from fastapi import HTTPException from pydantic import BaseModel from bson import ObjectId -from bson.errors import InvalidId from pymongo.errors import OperationFailure from pymongo import ( ASCENDING, @@ -19,19 +18,6 @@ from motor.motor_asyncio import ( ) -def object_id_from_key(key: str) -> ObjectId: - """Get ObjectId from key, 400 if invalid ObjectId - - :param key: Key. - :return: ObjectId - """ - - try: - return ObjectId(key) - except InvalidId as exc: - raise HTTPException(status_code=400, detail="Invalid key/object id") from exc - - class SearchInput(BaseModel): """Handle search data for HTTP request""" @@ -71,7 +57,10 @@ class DBClient: print(f"WARNING failed to connect to DB - {i} / 4", flush=True) sleep(1) else: - print("Could not connect to DB - mongodb://REDACTED_USERNAME:REDACTED_PASSWORD@mongodb:27017/production") + print( + "Could not connect to DB - mongodb://REDACTED_USERNAME:REDACTED_PASSWORD@mongodb:27017/production", + flush=True, + ) app_exit(1) async def find(self, search_data: SearchInput) -> List[Dict[str, Any]]: @@ -82,11 +71,11 @@ class DBClient: """ data: List[Dict[str, Any]] = [] - cursor = self.collection.find(search_data.filter) - - cursor.sort([("ip", ASCENDING), ("timestamp", DESCENDING)]).limit(search_data.limit).skip(search_data.skip) try: + cursor = self.collection.find(search_data.filter) + cursor.sort([("ip", ASCENDING), ("timestamp", DESCENDING)]).limit(search_data.limit).skip(search_data.skip) + async for document in cursor: if document is not None: document["_id"] = str(document["_id"]) @@ -101,7 +90,7 @@ class DBClient: detail="Probably wrong syntax, note the dictionary for find: " + "https://motor.readthedocs.io/en/stable/tutorial-asyncio.html#async-for", ) from exc - except BaseException as exc: + except Exception as exc: print(f"DB connection failed: {exc}") raise HTTPException(status_code=500, detail="DB connection failed") from exc @@ -119,7 +108,7 @@ class DBClient: return document return None - except BaseException as exc: + except Exception as exc: print(f"DB connection failed: {exc}") raise HTTPException(status_code=500, detail="DB connection failed") from exc @@ -136,7 +125,7 @@ class DBClient: return result.inserted_id return None - except BaseException as exc: + except Exception as exc: print(f"DB connection failed: {exc}") raise HTTPException(status_code=500, detail="DB connection failed") from exc @@ -154,7 +143,7 @@ class DBClient: return object_id return None - except BaseException as exc: + except Exception as exc: print(f"DB connection failed: {exc}") raise HTTPException(status_code=500, detail="DB connection failed") from exc @@ -171,7 +160,7 @@ class DBClient: return object_id return None - except BaseException as exc: + except Exception as exc: print(f"DB connection failed: {exc}") raise HTTPException(status_code=500, detail="DB connection failed") from exc @@ -182,11 +171,11 @@ class DBClient: """ try: - result = await self.collection.estimated_document_count() + result = await self.collection.estimated_document_count(maxTimeMS=4000) if isinstance(result, int): return result return None - except BaseException as exc: + except Exception as exc: print(f"DB connection failed: {exc}") raise HTTPException(status_code=500, detail="DB connection failed") from exc diff --git a/src/soc_collector/main.py b/src/soc_collector/main.py index e70199b..fd1bded 100755 --- a/src/soc_collector/main.py +++ b/src/soc_collector/main.py @@ -10,9 +10,8 @@ from fastapi.responses import JSONResponse from .db import ( DBClient, SearchInput, - object_id_from_key, ) -from .schema import valid_schema +from .schema import valid_schema, object_id_from_data from .auth import authorize_client, load_api_keys # Get credentials @@ -52,10 +51,10 @@ async def search(request: Request, search_data: SearchInput) -> JSONResponse: @app.post("/sc/v0") -async def create(request: Request) -> JSONResponse: +async def insert(request: Request) -> JSONResponse: """/sc/v0, POST method - :param request: The request where we get the json body. + :param request: The client request. :return: JSONResponse """ @@ -75,19 +74,19 @@ async def create(request: Request) -> JSONResponse: if not valid_schema(json_data): return JSONResponse(content={"status": "error", "message": "Not our JSON schema"}, status_code=400) - key = await db.insert_one(json_data) + object_id = await db.insert_one(json_data) - if key is None: + if object_id is None: return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) - return JSONResponse(content={"status": "success", "key": str(key)}) + return JSONResponse(content={"status": "success", "_id": str(object_id)}) @app.put("/sc/v0") async def replace(request: Request) -> JSONResponse: # pylint: disable=too-many-return-statements """/sc/v0, PUT method - :param request: The request where we get the json body. + :param request: The client request. :return: JSONResponse """ @@ -99,13 +98,9 @@ async def replace(request: Request) -> JSONResponse: # pylint: disable=too-many except JSONDecodeError: return JSONResponse(content={"status": "error", "message": "Invalid JSON"}, status_code=400) - if "_id" not in json_data: - return JSONResponse(content={"status": "error", "message": "Missing key '_id'"}, status_code=400) - # Get the key - if isinstance(json_data["_id"], str): - object_id = object_id_from_key(json_data["_id"]) - else: + object_id = object_id_from_data(json_data) + if object_id is None: return JSONResponse(content={"status": "error", "message": "Missing key '_id' with valid id"}, status_code=400) # Ensure the updating key exist @@ -123,24 +118,28 @@ async def replace(request: Request) -> JSONResponse: # pylint: disable=too-many json_data["_id"] = object_id returned_object_id = await db.replace_one(object_id, json_data) - if returned_object_id is None: + if returned_object_id is None or returned_object_id != object_id: return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) - return JSONResponse(content={"status": "success", "key": str(object_id)}) + return JSONResponse(content={"status": "success", "_id": str(object_id)}) @app.get("/sc/v0/{key}") async def get(request: Request, key: str) -> JSONResponse: """/sc/v0/{key}, GET method - :param key: The document key in the database. + :param request: The client request. + :param key: The document id in the database. :return: JSONResponse """ # Ensure authorization authorize_client(request, API_KEYS) - object_id = object_id_from_key(key) + # Get the id + object_id = object_id_from_data(key) + if object_id is None: + return JSONResponse(content={"status": "error", "message": "Invalid id"}, status_code=400) document = await db.find_one(object_id) @@ -154,27 +153,32 @@ async def get(request: Request, key: str) -> JSONResponse: async def delete(request: Request, key: str) -> JSONResponse: """/sc/v0/{key}, DELETE method - :param key: The document key in the database. + :param request: The client request. + :param key: The document id in the database. :return: JSONResponse """ # Ensure authorization authorize_client(request, API_KEYS) - object_id = object_id_from_key(key) + # Get the id + object_id = object_id_from_data(key) + if object_id is None: + return JSONResponse(content={"status": "error", "message": "Invalid id"}, status_code=400) result = await db.delete_one(object_id) if result is None: return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404) - return JSONResponse(content={"status": "success", "key": str(object_id)}) + return JSONResponse(content={"status": "success", "_id": str(object_id)}) @app.get("/info") async def info(request: Request) -> JSONResponse: """/info, GET method + :param request: The client request. :return: JSONResponse """ diff --git a/src/soc_collector/schema.py b/src/soc_collector/schema.py index 221990a..2c2dfb9 100644 --- a/src/soc_collector/schema.py +++ b/src/soc_collector/schema.py @@ -1,6 +1,8 @@ """Our schema module""" -from typing import Any, Dict +from typing import Any, Dict, Optional, Union import jsonschema +from bson import ObjectId +from bson.errors import InvalidId # fmt:off # NOTE: Commented out properties are left intentionally, so it is easier to see @@ -99,9 +101,33 @@ def valid_schema(json_data: Dict[str, Any]) -> bool: :param json_data: Json object :return: bool """ + try: jsonschema.validate(json_data, schema, format_checker=jsonschema.FormatChecker()) except jsonschema.exceptions.ValidationError as exc: print(f"Validation failed with error: {exc.message}") return False return True + + +def object_id_from_data(data: Union[str, Dict[str, Any]]) -> Optional[ObjectId]: + """Get ObjectId from key. None if invalid. + + :param data: Key. + :return: Optional[ObjectId] + """ + + if isinstance(data, str): + try: + return ObjectId(data) + except InvalidId: + return None + + elif isinstance(data, Dict): + if "_id" in data and isinstance(data["_id"], str): + try: + return ObjectId(data["_id"]) + except InvalidId: + return None + + return None diff --git a/src/soc_collector/soc_collector_cli.py b/src/soc_collector/soc_collector_cli.py index d7add30..4929655 100644 --- a/src/soc_collector/soc_collector_cli.py +++ b/src/soc_collector/soc_collector_cli.py @@ -8,31 +8,11 @@ from sys import exit as app_exit import json import requests -if "COLLECTOR_API_KEY" not in environ: - print("Missing 'COLLECTOR_API_KEY' in environment") - app_exit(1) -API_KEY = environ["COLLECTOR_API_KEY"] +from .schema import object_id_from_data -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 @@ -54,10 +34,14 @@ def json_load_data(data: str) -> Dict[str, Any]: app_exit(1) -def info_action() -> None: - """Get database info, currently number of documents.""" +def info_action(api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: + """Get database info, currently number of documents. + + :param api_key: The API key. + :param base_url: URL to the API. + """ - req = requests.get(f"{API_URL}/info", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE) + req = requests.get(f"{base_url}/info", headers={"API-KEY": api_key}, timeout=5, verify=ROOT_CA_FILE) # Ensure ok status req.raise_for_status() @@ -67,16 +51,18 @@ def info_action() -> None: print(f"Estimated document count: {json_data['Estimated document count']}") -def search_action(data: str) -> None: +def search_action(data: str, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Search for documents in the database. :param data: String with either json or path to a json file. + :param api_key: The API key. + :param base_url: URL to the API. """ search_data = json_load_data(data) req = requests.post( - f"{API_URL}/sc/v0/search", headers={"API-KEY": API_KEY}, json=search_data, timeout=5, verify=ROOT_CA_FILE + f"{base_url}/sc/v0/search", headers={"API-KEY": api_key}, json=search_data, timeout=5, verify=ROOT_CA_FILE ) # Ensure ok status @@ -87,24 +73,26 @@ def search_action(data: str) -> None: print(json.dumps(json_data["docs"], indent=4)) -def delete_action(data: str) -> None: +def delete_action(data: str, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Delete a document in the DB. - :param data: key or path to a json file containing "_id". + :param data: id or path to a json file containing "_id". + :param api_key: The API key. + :param base_url: URL to the API. """ if data and isfile(data): 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"] + object_id = object_id_from_data(json_data) else: - key = data + object_id = object_id_from_data(data) - valid_key(key) + if object_id is None: + print("ERROR: id is not valid") + app_exit(1) - req = requests.delete(f"{API_URL}/sc/v0/{key}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE) + req = requests.delete( + f"{base_url}/sc/v0/{str(object_id)}", headers={"API-KEY": api_key}, timeout=5, verify=ROOT_CA_FILE + ) # Check status if req.status_code == 404: @@ -114,7 +102,7 @@ def delete_action(data: str) -> None: # Ensure ok status req.raise_for_status() - print(f"Deleted data OK - key: {key}") + print(f"Deleted data OK - key: {str(object_id)}") def update_local_action(data: str, update_data: str) -> None: @@ -133,21 +121,24 @@ def update_local_action(data: str, update_data: str) -> None: print(json.dumps(json_data, indent=4)) -def replace_action(data: str) -> None: +def replace_action(data: str, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Replace the entire document in the database with this document, "_id" must exist as a key. :param data: json blob or path to json file, "_id" key must exist. + :param api_key: The API key. + :param base_url: URL to the API. """ json_data = json_load_data(data) - if "_id" not in json_data or not isinstance(json_data["_id"], str): + object_id = object_id_from_data(json_data) + if object_id is None: 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) + req = requests.put( + f"{base_url}/sc/v0", json=json_data, headers={"API-KEY": api_key}, timeout=5, verify=ROOT_CA_FILE + ) # Check status if req.status_code == 404: @@ -158,13 +149,15 @@ def replace_action(data: str) -> None: req.raise_for_status() json_data = json.loads(req.text) - print(f'Replaced data OK - key: {json_data["key"]}') + print(f'Replaced data OK - key: {json_data["_id"]}') -def insert_action(data: str) -> None: +def insert_action(data: str, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Insert a new document into the database, "_id" must not exist in the document. :param data: json blob or path to json file, "_id" key must not exist. + :param api_key: The API key. + :param base_url: URL to the API. """ json_data = json_load_data(data) @@ -174,34 +167,35 @@ def insert_action(data: str) -> None: app_exit(1) req = requests.post( - f"{API_URL}/sc/v0", json=json_data, headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE + f"{base_url}/sc/v0", json=json_data, headers={"API-KEY": api_key}, timeout=5, verify=ROOT_CA_FILE ) # Ensure ok status req.raise_for_status() json_data = json.loads(req.text) - print(f'Inserted data OK - key: {json_data["key"]}') + print(f'Inserted data OK - key: {json_data["_id"]}') -def get_action(data: str) -> None: +def get_action(data: str, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Get a document from the database. - :param data: key or path to a json file containing "_id". + :param data: id or path to a json file containing "_id". + :param api_key: The API key. + :param base_url: URL to the API. """ if data and isfile(data): 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"] + object_id = object_id_from_data(json_data) else: - key = data - - valid_key(key) + object_id = object_id_from_data(data) + if object_id is None: + print("ERROR: Invalid id") + app_exit(1) - req = requests.get(f"{API_URL}/sc/v0/{key}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE) + req = requests.get( + f"{base_url}/sc/v0/{str(object_id)}", headers={"API-KEY": api_key}, timeout=5, verify=ROOT_CA_FILE + ) # Check status if req.status_code == 404: @@ -218,6 +212,7 @@ def get_action(data: str) -> None: def main() -> None: """Main function.""" + parser = ArgumentParser(formatter_class=RawTextHelpFormatter, description="SOC Collector CLI") parser.add_argument( "action", @@ -231,7 +226,7 @@ def main() -> None: '{"filter": {"asn_country_code": "SE", "result": {"$exists": "cve_2015_0002"}}}' '{"filter": {}}' - get: key OR path to document using its "_id". + get: id OR path to document using its "_id". 637162378c92893fff92bf7e OR ./data.json insert: json blob OR path to file. Document MUST NOT contain "_id". @@ -244,7 +239,7 @@ def main() -> None: This does NOT send data to the database, use replace for that. 1st ARG: '{json_blob_here...}' OR ./data.json 2th ARG:'{"port": 555, "some_key": "some_data"}' OR ./data.json - delete: key OR path to file using its "_id". + delete: id OR path to file using its "_id". 637162378c92893fff92bf7e OR ./data.json @@ -288,20 +283,25 @@ def main() -> None: args = parser.parse_args() + if "COLLECTOR_API_KEY" not in environ: + print("Missing 'COLLECTOR_API_KEY' in environment") + app_exit(1) + api_key = environ["COLLECTOR_API_KEY"] + if args.action == "get": - get_action(args.data) + get_action(args.data, api_key) elif args.action == "insert": - insert_action(args.data) + insert_action(args.data, api_key) elif args.action == "replace": - replace_action(args.data) + replace_action(args.data, api_key) elif args.action == "update_local" and args.extra_data is not None: update_local_action(args.data, args.extra_data) elif args.action == "delete": - delete_action(args.data) + delete_action(args.data, api_key) elif args.action == "search": - search_action(args.data) + search_action(args.data, api_key) elif args.action == "info": - info_action() + info_action(api_key) else: print("ERROR: Wrong action") app_exit(1) |