From 43e87d84b15d12d52a4dcde6e80426cbd17e3d6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Victor=20N=C3=A4slund?= Date: Wed, 16 Nov 2022 19:09:46 +0100 Subject: auth and CLI done --- data/api_keys.txt | 2 + data/collector-dev.soc.sunet.se.crt | 12 ++ data/collector-dev.soc.sunet.se.key | 3 + data/collector_container/Dockerfile | 15 +- data/collector_root_ca.crt | 13 ++ data/mongodb_container/Dockerfile | 4 +- dev-run.sh | 30 ++-- pyproject.toml | 6 + requirements.txt | 12 +- setup.cfg | 8 + src/collector/__init__.py | 4 - src/collector/db.py | 176 ---------------------- src/collector/healthcheck.py | 38 ----- src/collector/healthcheck.sh | 11 -- src/collector/main.py | 168 --------------------- src/collector/py.typed | 0 src/collector/schema.py | 107 -------------- src/soc_collector/__init__.py | 4 + src/soc_collector/auth.py | 46 ++++++ src/soc_collector/db.py | 178 ++++++++++++++++++++++ src/soc_collector/healthcheck.py | 38 +++++ src/soc_collector/healthcheck.sh | 11 ++ src/soc_collector/main.py | 183 +++++++++++++++++++++++ src/soc_collector/py.typed | 0 src/soc_collector/schema.py | 107 ++++++++++++++ src/soc_collector/soc_collector_cli.py | 259 +++++++++++++++++++++++++++++++++ 26 files changed, 902 insertions(+), 533 deletions(-) create mode 100644 data/api_keys.txt create mode 100644 data/collector-dev.soc.sunet.se.crt create mode 100644 data/collector-dev.soc.sunet.se.key create mode 100644 data/collector_root_ca.crt create mode 100644 pyproject.toml create mode 100644 setup.cfg delete mode 100644 src/collector/__init__.py delete mode 100644 src/collector/db.py delete mode 100644 src/collector/healthcheck.py delete mode 100755 src/collector/healthcheck.sh delete mode 100755 src/collector/main.py delete mode 100644 src/collector/py.typed delete mode 100644 src/collector/schema.py create mode 100644 src/soc_collector/__init__.py create mode 100644 src/soc_collector/auth.py create mode 100644 src/soc_collector/db.py create mode 100644 src/soc_collector/healthcheck.py create mode 100755 src/soc_collector/healthcheck.sh create mode 100755 src/soc_collector/main.py create mode 100644 src/soc_collector/py.typed create mode 100644 src/soc_collector/schema.py create mode 100644 src/soc_collector/soc_collector_cli.py diff --git a/data/api_keys.txt b/data/api_keys.txt new file mode 100644 index 0000000..8cc5dd6 --- /dev/null +++ b/data/api_keys.txt @@ -0,0 +1,2 @@ +ca7dd92d5a83c9e92b935888d390e919f6b0d0511a569b5373a19c332880fcde;john +bb6b4de1e9839224daf4cdfd6aad5af667655fafc48d7622e80428436ffe0462;victor diff --git a/data/collector-dev.soc.sunet.se.crt b/data/collector-dev.soc.sunet.se.crt new file mode 100644 index 0000000..f75b2a1 --- /dev/null +++ b/data/collector-dev.soc.sunet.se.crt @@ -0,0 +1,12 @@ +-----BEGIN CERTIFICATE----- +MIIBuzCCAW2gAwIBAgIUUARU/fxpFlGCHCwdw6rQS8gDkyAwBQYDK2VwMGsxCzAJ +BgNVBAYTAlNFMRMwEQYDVQQIDApTb21lLVN0YXRlMQ4wDAYDVQQKDAVTVU5FVDES +MBAGA1UECwwJU1VORVQgU09DMSMwIQYDVQQDDBpjb2xsZWN0b3ItZGV2LnNvYy5z +dW5ldC5zZTAeFw0yMjExMTYxNTI5NDVaFw0yNzExMTYxNTI5NDVaMDIxCzAJBgNV +BAYTAlNFMSMwIQYDVQQDDBpjb2xsZWN0b3ItZGV2LnNvYy5zdW5ldC5zZTAqMAUG +AytlcAMhALL1Lx4uRNrYjx3Z/Z41C1BruCOL6slqk2sqz2s0yghIo1wwWjALBgNV +HQ8EBAMCBDAwEwYDVR0lBAwwCgYIKwYBBQUHAwEwNgYDVR0RBC8wLYIaY29sbGVj +dG9yLWRldi5zb2Muc3VuZXQuc2WCCWxvY2FsaG9zdIcEWS+5nzAFBgMrZXADQQBH +g3Yysjjch6OkG/Vo7PyDUg3NlDqbDMktucxDaHLgPkLF508fHNLUhh3LAWn376dr +RULOF42AyfSYQY1WpDgB +-----END CERTIFICATE----- diff --git a/data/collector-dev.soc.sunet.se.key b/data/collector-dev.soc.sunet.se.key new file mode 100644 index 0000000..dcb55bb --- /dev/null +++ b/data/collector-dev.soc.sunet.se.key @@ -0,0 +1,3 @@ +-----BEGIN PRIVATE KEY----- +MC4CAQAwBQYDK2VwBCIEIMbi6GqNxlAvNiRNgYa8hR1F0br8gFhPgJU1hru/n6PO +-----END PRIVATE KEY----- diff --git a/data/collector_container/Dockerfile b/data/collector_container/Dockerfile index 0641587..a35b2cb 100644 --- a/data/collector_container/Dockerfile +++ b/data/collector_container/Dockerfile @@ -1,4 +1,4 @@ -FROM debian:bullseye-20221024-slim@sha256:76cdda8fe5eb597ef5e712e4c9a9f5f1fb119e69f353daaa7bd6d0f6e66e541d +FROM debian:bullseye-20221114-slim@sha256:df172d92d287ec4d4a538e5db8026fcde5f91f5f90061423d69d6148ff05cc47 EXPOSE 8000 @@ -21,21 +21,24 @@ RUN apt-get update \ RUN find / -xdev -perm /6000 -type f -exec chmod a-s {} \; || true # Add user -RUN useradd collector -u 1500 -s /usr/sbin/nologin +RUN useradd soc_collector -u 1500 -s /usr/sbin/nologin COPY ./src /app/src COPY ./data/logging.json /app/logging.json +COPY ./data/collector-dev.soc.sunet.se.crt /app/collector-dev.soc.sunet.se.crt +COPY ./data/collector-dev.soc.sunet.se.key /app/collector-dev.soc.sunet.se.key +COPY ./data/collector_root_ca.crt /app/collector_root_ca.crt +COPY ./data/api_keys.txt /app/api_keys.txt WORKDIR /app/ -USER collector +USER soc_collector # Add healthcheck HEALTHCHECK --interval=2m --timeout=15s --retries=1 --start-period=30s \ - CMD sh ./src/collector/healthcheck.sh COLLECTOR || bash -c 'kill -s 15 1 && (sleep 7; kill -s 9 1)' - -ENTRYPOINT ["uvicorn", "src.collector.main:app", "--log-config", "./logging.json", "--host", "0.0.0.0", "--workers", "1", "--header", "server:collector"] + CMD sh ./src/soc_collector/healthcheck.sh COLLECTOR || bash -c 'kill -s 15 1 && (sleep 7; kill -s 9 1)' +ENTRYPOINT ["uvicorn", "src.soc_collector.main:app", "--log-config", "./logging.json", "--host", "0.0.0.0", "--port", "8000", "--ssl-keyfile", "./collector-dev.soc.sunet.se.key", "--ssl-certfile", "./collector-dev.soc.sunet.se.crt", "--ssl-version", "2", "--workers", "1", "--header", "server:collector"] diff --git a/data/collector_root_ca.crt b/data/collector_root_ca.crt new file mode 100644 index 0000000..b9cddf1 --- /dev/null +++ b/data/collector_root_ca.crt @@ -0,0 +1,13 @@ +-----BEGIN CERTIFICATE----- +MIIB6zCCAZ2gAwIBAgIUXbzkTFrPe+rAKscqzHHB+Kr6oDQwBQYDK2VwMGsxCzAJ +BgNVBAYTAlNFMRMwEQYDVQQIDApTb21lLVN0YXRlMQ4wDAYDVQQKDAVTVU5FVDES +MBAGA1UECwwJU1VORVQgU09DMSMwIQYDVQQDDBpjb2xsZWN0b3ItZGV2LnNvYy5z +dW5ldC5zZTAeFw0yMjExMTYxNTI0NDdaFw0yNzExMTUxNTI0NDdaMGsxCzAJBgNV +BAYTAlNFMRMwEQYDVQQIDApTb21lLVN0YXRlMQ4wDAYDVQQKDAVTVU5FVDESMBAG +A1UECwwJU1VORVQgU09DMSMwIQYDVQQDDBpjb2xsZWN0b3ItZGV2LnNvYy5zdW5l +dC5zZTAqMAUGAytlcAMhAFBwDjb3i5fjxrbcFOMXTZfKnIDx6h9XiojAXPXD/VpD +o1MwUTAdBgNVHQ4EFgQUlY3rq7wlYJwaiTED1AJUTnevhxUwHwYDVR0jBBgwFoAU +lY3rq7wlYJwaiTED1AJUTnevhxUwDwYDVR0TAQH/BAUwAwEB/zAFBgMrZXADQQCL +9S14sR+y0AbHMTmC2BzuLmwPexK23VDgZemCRiBwR1DZ8x5Vzd/IR+WqmUaFhGPv +utaY9PGGHZoZtrbb5WEE +-----END CERTIFICATE----- diff --git a/data/mongodb_container/Dockerfile b/data/mongodb_container/Dockerfile index f37d2a3..e163593 100644 --- a/data/mongodb_container/Dockerfile +++ b/data/mongodb_container/Dockerfile @@ -1,4 +1,4 @@ -FROM debian:bullseye-20221024-slim@sha256:76cdda8fe5eb597ef5e712e4c9a9f5f1fb119e69f353daaa7bd6d0f6e66e541d +FROM debian:bullseye-20221114-slim@sha256:df172d92d287ec4d4a538e5db8026fcde5f91f5f90061423d69d6148ff05cc47 EXPOSE 27017 @@ -20,7 +20,7 @@ RUN find / -xdev -perm /6000 -type f -exec chmod a-s {} \; || true COPY ./data/mongodb_entrypoint.sh /mongodb_entrypoint.sh COPY ./data/init-mongodb.js /init-mongodb.js COPY ./data/healthcheck-mongodb.js /healthcheck-mongodb.js -COPY ./src/collector/healthcheck.sh /healthcheck.sh +COPY ./src/soc_collector/healthcheck.sh /healthcheck.sh USER mongodb diff --git a/dev-run.sh b/dev-run.sh index 084e38b..976b964 100755 --- a/dev-run.sh +++ b/dev-run.sh @@ -1,9 +1,9 @@ #!/bin/bash echo "Checking package" -mypy --strict --namespace-packages --ignore-missing-imports --cache-dir=/tmp/ src/collector/*.py # || exit 1 -black --line-length 120 src/collector/*.py # || exit 1 -pylint --max-line-length 120 src/collector/*.py # || exit 1 +mypy --strict --namespace-packages --ignore-missing-imports --cache-dir=/tmp/ src/soc_collector/*.py # || exit 1 +black --line-length 120 src/soc_collector/*.py # || exit 1 +pylint --max-line-length 120 src/soc_collector/*.py # || exit 1 mkdir -p data/mongodb_data sudo chown -R $USER data/mongodb_data @@ -17,46 +17,46 @@ sleep 3 echo echo -curl -v --data-binary @data/example_data_3.json http://127.0.0.1:8000/sc/v0 +curl -v -k --data-binary @data/example_data_3.json https://127.0.0.1:8000/sc/v0 echo echo -curl -v -X DELETE http://127.0.0.1:8000/sc/v0/63702570e004d2b0b2254d27 +curl -v -k -X DELETE https://127.0.0.1:8000/sc/v0/63702570e004d2b0b2254d27 echo echo -curl -v -X DELETE http://127.0.0.1:8000/sc/v0/63702570e004d2b0b2254d27 +curl -v -k -X DELETE https://127.0.0.1:8000/sc/v0/63702570e004d2b0b2254d27 echo echo -curl -v -d '{"search": {"port": {"$lt": 4}}}' -H 'Content-Type: application/json' http://127.0.0.1:8000/sc/v0/search +curl -v -k -d '{"search": {"port": {"$lt": 4}}}' -H 'Content-Type: application/json' https://127.0.0.1:8000/sc/v0/search echo echo -curl -v -d '{"search": {"port": 112}}' -H 'Content-Type: application/json' http://127.0.0.1:8000/sc/v0/search +curl -v -k -d '{"search": {"port": 112}}' -H 'Content-Type: application/json' https://127.0.0.1:8000/sc/v0/search echo echo -curl -v -d '{"search": {"port": {"$gt": 4}}}' -H 'Content-Type: application/json' http://127.0.0.1:8000/sc/v0/search +curl -v -k -d '{"search": {"port": {"$gt": 4}}}' -H 'Content-Type: application/json' https://127.0.0.1:8000/sc/v0/search echo echo -curl -v -d '{"search": {"port": 111}}' -H 'Content-Type: application/json' http://127.0.0.1:8000/sc/v0/search +curl -v -k -d '{"search": {"port": 111}}' -H 'Content-Type: application/json' https://127.0.0.1:8000/sc/v0/search echo echo -curl -v -d '{"search": {"port": {"sdfsf": 7}}}' -H 'Content-Type: application/json' http://127.0.0.1:8000/sc/v0/search +curl -v -k -d '{"search": {"port": {"sdfsf": 7}}}' -H 'Content-Type: application/json' https://127.0.0.1:8000/sc/v0/search echo echo -curl -v -d '{"search": {"port": {"$sdfsf": 7}}}' -H 'Content-Type: application/json' http://127.0.0.1:8000/sc/v0/search +curl -v -k -d '{"search": {"port": {"$sdfsf": 7}}}' -H 'Content-Type: application/json' https://127.0.0.1:8000/sc/v0/search echo echo -curl -v -d '{"search": {"portfdv": {"$asa": 7}}}' -H 'Content-Type: application/json' http://127.0.0.1:8000/sc/v0/search +curl -v -k -d '{"search": {"portfdv": {"$asa": 7}}}' -H 'Content-Type: application/json' https://127.0.0.1:8000/sc/v0/search echo echo echo echo -curl -v -X PUT --data-binary @data/example_data_3_replace_test.json http://127.0.0.1:8000/sc/v0 +curl -v -k -X PUT --data-binary @data/example_data_3_replace_test.json https://127.0.0.1:8000/sc/v0 echo echo -curl -v http://127.0.0.1:8000/info +curl -v -k https://127.0.0.1:8000/info # bash quickstart.sh -b || exit 1 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8d91941 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,6 @@ +[build-system] +requires = [ + "setuptools", + "wheel" +] +build-backend = "setuptools.build_meta" diff --git a/requirements.txt b/requirements.txt index 664ce71..998c0e0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -30,9 +30,9 @@ dnspython==2.2.1 \ --hash=sha256:0f7569a4a6ff151958b64304071d370daa3243d15941a7beedf0c9fe5105603e \ --hash=sha256:a851e51367fb93e9e1361732c1d60dab63eff98712e503ea7d92e6eccb109b4f # via pymongo -fastapi==0.86.0 \ - --hash=sha256:1020d7ca205d8b95813881fb3282e9c3656e47993531af3aa4ae11065b61dd2c \ - --hash=sha256:cdcaff84ecf7ae939b9579f0c98b0a0989ee3dd855710a32bc985260d92612f6 +fastapi==0.87.0 \ + --hash=sha256:07032e53df9a57165047b4f38731c38bdcc3be5493220471015e2b4b51b486a4 \ + --hash=sha256:254453a2e22f64e2a1b4e1d8baf67d239e55b6c8165c079d25746a5220c81bb4 # via -r requirements.in h11==0.14.0 \ --hash=sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d \ @@ -295,9 +295,9 @@ sniffio==1.3.0 \ --hash=sha256:e60305c5e5d314f5389259b7f22aaa33d8f7dee49763119234af3755c55b9101 \ --hash=sha256:eecefdce1e5bbfb7ad2eeaabf7c1eeb404d7757c379bd1f7e5cce9d8bf425384 # via anyio -starlette==0.20.4 \ - --hash=sha256:42fcf3122f998fefce3e2c5ad7e5edbf0f02cf685d646a83a08d404726af5084 \ - --hash=sha256:c0414d5a56297d37f3db96a84034d61ce29889b9eaccf65eb98a0b39441fcaa3 +starlette==0.21.0 \ + --hash=sha256:0efc058261bbcddeca93cad577efd36d0c8a317e44376bcfc0e097a2b3dc24a7 \ + --hash=sha256:b1b52305ee8f7cfc48cde383496f7c11ab897cd7112b33d998b1317dc8ef9027 # via fastapi typing-extensions==4.4.0 \ --hash=sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa \ diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..8026ee7 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,8 @@ +[options.package_data] +# If any package or subpackage contains *.txt or *.rst files, include them: +soc_collector = py.typed data/collector_root_ca.crt + +[options.entry_points] +# If any package or subpackage contains *.txt or *.rst files, include them: +console_scripts = + soc_collector_cli = soc_collector:soc_collector_cli.main diff --git a/src/collector/__init__.py b/src/collector/__init__.py deleted file mode 100644 index 6530fdd..0000000 --- a/src/collector/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Collector -""" - -__version__ = "1.03" diff --git a/src/collector/db.py b/src/collector/db.py deleted file mode 100644 index 641e5da..0000000 --- a/src/collector/db.py +++ /dev/null @@ -1,176 +0,0 @@ -"""Our database module""" -from typing import List, Dict, Optional, Any -from time import sleep -from sys import exit as app_exit -from dataclasses import dataclass - -from fastapi import HTTPException -from pydantic import BaseModel -from pymongo.errors import OperationFailure -from motor.motor_asyncio import ( - AsyncIOMotorClient, - AsyncIOMotorCollection, -) -from bson import ObjectId - - -class SearchInput(BaseModel): - """Handle search data for HTTP request""" - - search: Optional[Dict[str, Any]] - limit: int = 25 - skip: int = 0 - - -@dataclass() -class DBClient: - """Class to hold database connections for us.""" - - client: AsyncIOMotorClient - collection: AsyncIOMotorCollection - - def __init__(self, username: str, password: str, collection: str) -> None: - self.client = AsyncIOMotorClient( - f"mongodb://{username}:{password}@mongodb:27017/production", - maxConnecting=4, - timeoutMS=3000, - serverSelectionTimeoutMS=3000, - ) - self.collection = self.client["production"][collection] - - async def startup(self) -> None: - """Try query the DB and exit the program if we fail after 5 times. - - :return: None - """ - - for i in range(5): - try: - await self.collection.find_one({"_id": ObjectId("507f1f77bcf86cd799439011")}) - print("Connection to DB - OK") - break - except Exception: # pylint: disable=broad-except - 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") - app_exit(1) - - async def find(self, search_data: SearchInput) -> List[Dict[str, Any]]: - """Wrap the find() method, handling timeouts and return data type. - - :param search_data: Instance of SearchInput. - :return: Optional[List[Dict[str, Any]]] - """ - - data: List[Dict[str, Any]] = [] - cursor = self.collection.find(search_data.search) - - # Sort on timestamp - # TODO: Also sort on IP as well - cursor.sort("timestamp", -1).limit(search_data.limit).skip(search_data.skip) - - try: - async for document in cursor: - if document is not None: - document["_id"] = str(document["_id"]) - data.append(document) - - return data - - except OperationFailure as exc: - print(f"DB failed to process: {exc.details}") - raise HTTPException( - status_code=400, - 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: - print(f"DB connection failed: {exc}") - raise HTTPException(status_code=500, detail="DB connection failed") from exc - - async def find_one(self, object_id: ObjectId) -> Optional[Dict[str, Any]]: - """Wrap the find_one() method, handling timeouts and return data type. - - :param object_id: The object id to find. - :return: Optional[Dict[str, Any]] - """ - - try: - document = await self.collection.find_one({"_id": object_id}) - if isinstance(document, Dict): - document["_id"] = str(document["_id"]) - return document - return None - - except BaseException as exc: - print(f"DB connection failed: {exc}") - raise HTTPException(status_code=500, detail="DB connection failed") from exc - - async def insert_one(self, data: Dict[str, Any]) -> Optional[ObjectId]: - """Wrap the insert_one() method, handling timeouts and return data type. - - :param data: The data to insert into the DB. - :return: Optional[ObjectId] - """ - - try: - result = await self.collection.insert_one(data) - if isinstance(result.inserted_id, ObjectId) and len(str(result.inserted_id)) == 24: - return result.inserted_id - return None - - except BaseException as exc: - print(f"DB connection failed: {exc}") - raise HTTPException(status_code=500, detail="DB connection failed") from exc - - async def replace_one(self, object_id: ObjectId, data: Dict[str, Any]) -> Optional[ObjectId]: - """Wrap the replace_one() method, handling timeouts and return data type. - - :param object_id: The object id to replace. - :param data: The data to replace with. - :return: Optional[ObjectId] - """ - - try: - result = await self.collection.replace_one({"_id": object_id}, data) - if result.matched_count == 1: - return object_id - return None - - except BaseException as exc: - print(f"DB connection failed: {exc}") - raise HTTPException(status_code=500, detail="DB connection failed") from exc - - async def delete_one(self, object_id: ObjectId) -> Optional[ObjectId]: - """Wrap the delete_one() method, handling timeouts and return data type. - - :param object_id: The object id to delete from the DB. - :return: Optional[ObjectId] - """ - - try: - result = await self.collection.delete_one({"_id": object_id}) - if isinstance(result.deleted_count, int) and result.deleted_count == 1: - return object_id - return None - - except BaseException as exc: - print(f"DB connection failed: {exc}") - raise HTTPException(status_code=500, detail="DB connection failed") from exc - - async def estimated_document_count(self) -> Optional[int]: - """Wrap the estimated_document_count() method, handling timeouts and return data type. - - :return: Optional[int] - """ - - try: - result = await self.collection.estimated_document_count() - if isinstance(result, int): - return result - return None - - except BaseException as exc: - print(f"DB connection failed: {exc}") - raise HTTPException(status_code=500, detail="DB connection failed") from exc diff --git a/src/collector/healthcheck.py b/src/collector/healthcheck.py deleted file mode 100644 index 3091d98..0000000 --- a/src/collector/healthcheck.py +++ /dev/null @@ -1,38 +0,0 @@ -""" -Send a healthcheck request -""" -import sys -import time -import json - -import requests - - -def check_collector() -> bool: - """Check our collector using /info - - :return: bool - """ - time.sleep(2) # Prevent race condition with redis container healthcheck - - req = requests.get( - "http://localhost:8000/info", - timeout=3, - # TODO: verify="./rootCA.crt", - ) - - if req.status_code != 200: - return False - - data = json.loads(req.text) - if isinstance(data["Estimated document count"], int) and data["Estimated document count"] >= 0: - return req.status_code == 200 - - return False - - -if __name__ == "__main__": - if sys.argv[1] == "COLLECTOR": - if check_collector(): - sys.exit(0) - sys.exit(1) diff --git a/src/collector/healthcheck.sh b/src/collector/healthcheck.sh deleted file mode 100755 index aacc906..0000000 --- a/src/collector/healthcheck.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash - -# If mongodb container -if [ "$1" = "MONGODB" ] -then - /usr/bin/mongosh --eval 'disableTelemetry()' -u "$MONGODB_USERNAME" -p "$MONGODB_PASSWORD" localhost:27017/production /healthcheck-mongodb.js - exit $? -fi - -# If collector -/usr/bin/python3 ./src/collector/healthcheck.py "$1" || exit 1 diff --git a/src/collector/main.py b/src/collector/main.py deleted file mode 100755 index a2d55c7..0000000 --- a/src/collector/main.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Our main module""" -from os import environ -from asyncio import get_running_loop -from sys import exit as app_exit -from json.decoder import JSONDecodeError - -from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse -from bson import ( - ObjectId, - json_util, -) -from .db import ( - DBClient, - SearchInput, -) -from .schema import valid_schema - - -# Get credentials -if "MONGODB_USERNAME" not in environ or "MONGODB_PASSWORD" not in environ or "MONGODB_COLLECTION" not in environ: - print("Missing MONGODB_USERNAME or MONGODB_PASSWORD or MONGODB_COLLECTION in env") - app_exit(1) - -# Create DB object -db = DBClient(environ["MONGODB_USERNAME"], environ["MONGODB_PASSWORD"], environ["MONGODB_COLLECTION"]) - -# Check DB -loop = get_running_loop() -startup_task = loop.create_task(db.startup()) - -app = FastAPI() - - -# @app.exception_handler(RuntimeError) -# def app_exception_handler(request: Request, exc: RuntimeError) -> JSONResponse: -# print(exc, flush=True) -# return JSONResponse(content={"status": "error", "message": str(exc.with_traceback(None))}, status_code=400) -# return JSONResponse(content={"status": "error", "message": "Error during processing"}, status_code=400) - - -@app.post("/sc/v0/search") -async def search(search_data: SearchInput) -> JSONResponse: - """/sc/v0/search, POST method - - :param search_data: The search data. - :return: JSONResponse - """ - data = await db.find(search_data) - - return JSONResponse(content={"status": "success", "docs": data}) - - -@app.post("/sc/v0") -async def create(request: Request) -> JSONResponse: - """/sc/v0, POST method - - :param request: The request where we get the json body. - :return: JSONResponse - """ - - try: - json_data = await request.json() - except JSONDecodeError: - return JSONResponse(content={"status": "error", "message": "Invalid JSON"}, status_code=400) - - 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) - - if key is None: - return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) - - return JSONResponse(content={"status": "success", "key": str(key)}) - - -@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. - :return: JSONResponse - """ - - try: - json_data = await request.json() - 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 = ObjectId(json_data["_id"]) - elif ( - isinstance(json_data["_id"], dict) and "$oid" in json_data["_id"] and isinstance(json_data["_id"]["$oid"], str) - ): - object_id = ObjectId(json_data["_id"]["$oid"]) - else: - return JSONResponse(content={"status": "error", "message": "Missing key '_id' with valid id"}, status_code=400) - - # Ensure the updating key exist - document = await db.find_one(object_id) - - if document is None: - return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404) - - # Ensure valid schema - del json_data["_id"] - if not valid_schema(json_data): - return JSONResponse(content={"status": "error", "message": "Not our JSON schema"}, status_code=400) - - # Replace the data - json_data["_id"] = object_id - returned_object_id = await db.replace_one(object_id, json_data) - - if returned_object_id is None: - return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) - - return JSONResponse(content={"status": "success", "key": str(object_id)}) - - -@app.get("/sc/v0/{key}") -async def get(key: str) -> JSONResponse: - """/sc/v0/{key}, GET method - - :param key: The document key in the database. - :return: JSONResponse - """ - - document = await db.find_one(ObjectId(key)) - - if document is None: - return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404) - - return JSONResponse(content={"status": "success", "doc": document}) - - -@app.delete("/sc/v0/{key}") -async def delete(key: str) -> JSONResponse: - """/sc/v0/{key}, DELETE method - - :param key: The document key in the database. - :return: JSONResponse - """ - result = await db.delete_one(ObjectId(key)) - - if result is None: - return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404) - - return JSONResponse(content={"status": "success", "key": str(key)}) - - -@app.get("/info") -async def info() -> JSONResponse: - """/info, GET method - - :return: JSONResponse - """ - - count = await db.estimated_document_count() - - if count is None: - return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) - - return JSONResponse(content={"status": "success", "Estimated document count": count}) diff --git a/src/collector/py.typed b/src/collector/py.typed deleted file mode 100644 index e69de29..0000000 diff --git a/src/collector/schema.py b/src/collector/schema.py deleted file mode 100644 index 221990a..0000000 --- a/src/collector/schema.py +++ /dev/null @@ -1,107 +0,0 @@ -"""Our schema module""" -from typing import Any, Dict -import jsonschema - -# fmt:off -# NOTE: Commented out properties are left intentionally, so it is easier to see -# what properties are optional. -schema = { - "$schema": "http://json-schema.org/schema#", - "type": "object", - "properties": { - "document_version": {"type": "integer"}, - "ip": {"type": "string"}, - "port": {"type": "integer"}, - "whois_description": {"type": "string"}, - "asn": {"type": "string"}, - "asn_country_code": {"type": "string"}, - "ptr": {"type": "string"}, - "abuse_mail": {"type": "string"}, - "domain": {"type": "string"}, - "timestamp": {"type": "string", "format": "date-time"}, - "display_name": {"type": "string"}, - "description": {"type": "string"}, - "custom_data": { - "type": "object", - "patternProperties": { - ".*": { - "type": "object", - "properties": { - "display_name": {"type": "string"}, - "data": {"type": ["string", "boolean", "integer"]}, - "description": {"type": "string"}, - }, - "required": [ - "display_name", - "data", - # "description" - ] - }, - }, - }, - "result": { - "type": "object", - "patternProperties": { - ".*": { - "type": "object", - "properties": { - "display_name": {"type": "string"}, - "vulnerable": {"type": "boolean"}, - "investigation_needed": {"type": "boolean"}, - "reliability": {"type": "integer"}, - "description": {"type": "string"}, - }, - "oneOf": [ - { - "required": [ - "display_name", - "vulnerable", - # "reliability", # TODO: reliability is required if vulnerable = true - # "description", - ] - }, - { - "required": - [ - "display_name", - "investigation_needed", - # "reliability", # TODO: reliability is required if investigation_needed = true - # "description", - ] - }, - ] - }, - }, - }, - }, - "required": [ - "document_version", - "ip", - "port", - "whois_description", - "asn", - "asn_country_code", - "ptr", - "abuse_mail", - "domain", - "timestamp", - "display_name", - # "description", - # "custom_data", - "result", - ], -} - - -def valid_schema(json_data: Dict[str, Any]) -> bool: - """Check if json data follows the schema. - - :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 diff --git a/src/soc_collector/__init__.py b/src/soc_collector/__init__.py new file mode 100644 index 0000000..6530fdd --- /dev/null +++ b/src/soc_collector/__init__.py @@ -0,0 +1,4 @@ +"""Collector +""" + +__version__ = "1.03" diff --git a/src/soc_collector/auth.py b/src/soc_collector/auth.py new file mode 100644 index 0000000..4aacb52 --- /dev/null +++ b/src/soc_collector/auth.py @@ -0,0 +1,46 @@ +"""Auth module""" +from typing import List +from fastapi import Request, HTTPException + + +def load_api_keys(path: str) -> List[str]: + """Load API keys from file + + :param path: Path to API keys file. + :return: List[str] + """ + keys: List[str] = [] + + with open(path, encoding="utf-8") as f_data: + key_file = f_data.readlines() + + for line in key_file: + if line[0] == "#" or len(line) < 3: + continue + + key = line.split(";")[0] + keys.append(key) + + return keys + + +def authorize_client(request: Request, api_keys: List[str]) -> None: + """Authorize a client request. + + Parameters: + request (Request): The HTTP request. + api_keys (List[str]): List of accepted api_keys. + + Raises HTTPException with status_code=401 if authorize fail. + + Returns: + None + """ + + if "API-KEY" not in request.headers: + raise HTTPException(status_code=401, detail="API-KEY header missing") + + request_key = request.headers["API-KEY"].strip() + + if request_key not in api_keys: + raise HTTPException(status_code=401, detail="API-KEY header invalid") diff --git a/src/soc_collector/db.py b/src/soc_collector/db.py new file mode 100644 index 0000000..d601a82 --- /dev/null +++ b/src/soc_collector/db.py @@ -0,0 +1,178 @@ +"""Our database module""" +from typing import List, Dict, Optional, Any +from time import sleep +from sys import exit as app_exit +from dataclasses import dataclass + +from fastapi import HTTPException +from pydantic import BaseModel +from pymongo.errors import OperationFailure +from pymongo import ( + ASCENDING, + DESCENDING, +) +from motor.motor_asyncio import ( + AsyncIOMotorClient, + AsyncIOMotorCollection, +) +from bson import ObjectId + + +class SearchInput(BaseModel): + """Handle search data for HTTP request""" + + search: Optional[Dict[str, Any]] + limit: int = 25 + skip: int = 0 + + +@dataclass() +class DBClient: + """Class to hold database connections for us.""" + + client: AsyncIOMotorClient + collection: AsyncIOMotorCollection + + def __init__(self, username: str, password: str, collection: str) -> None: + self.client = AsyncIOMotorClient( + f"mongodb://{username}:{password}@mongodb:27017/production", + maxConnecting=4, + timeoutMS=3000, + serverSelectionTimeoutMS=3000, + ) + self.collection = self.client["production"][collection] + + async def startup(self) -> None: + """Try query the DB and exit the program if we fail after 5 times. + + :return: None + """ + + for i in range(5): + try: + await self.collection.find_one({"_id": ObjectId("507f1f77bcf86cd799439011")}) + print("Connection to DB - OK") + break + except Exception: # pylint: disable=broad-except + 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") + app_exit(1) + + async def find(self, search_data: SearchInput) -> List[Dict[str, Any]]: + """Wrap the find() method, handling timeouts and return data type. + + :param search_data: Instance of SearchInput. + :return: Optional[List[Dict[str, Any]]] + """ + + data: List[Dict[str, Any]] = [] + cursor = self.collection.find(search_data.search) + + cursor.sort({"ip": ASCENDING, "timestamp": DESCENDING}).limit(search_data.limit).skip(search_data.skip) + + try: + async for document in cursor: + if document is not None: + document["_id"] = str(document["_id"]) + data.append(document) + + return data + + except OperationFailure as exc: + print(f"DB failed to process: {exc.details}") + raise HTTPException( + status_code=400, + 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: + print(f"DB connection failed: {exc}") + raise HTTPException(status_code=500, detail="DB connection failed") from exc + + async def find_one(self, object_id: ObjectId) -> Optional[Dict[str, Any]]: + """Wrap the find_one() method, handling timeouts and return data type. + + :param object_id: The object id to find. + :return: Optional[Dict[str, Any]] + """ + + try: + document = await self.collection.find_one({"_id": object_id}) + if isinstance(document, Dict): + document["_id"] = str(document["_id"]) + return document + return None + + except BaseException as exc: + print(f"DB connection failed: {exc}") + raise HTTPException(status_code=500, detail="DB connection failed") from exc + + async def insert_one(self, data: Dict[str, Any]) -> Optional[ObjectId]: + """Wrap the insert_one() method, handling timeouts and return data type. + + :param data: The data to insert into the DB. + :return: Optional[ObjectId] + """ + + try: + result = await self.collection.insert_one(data) + if isinstance(result.inserted_id, ObjectId) and len(str(result.inserted_id)) == 24: + return result.inserted_id + return None + + except BaseException as exc: + print(f"DB connection failed: {exc}") + raise HTTPException(status_code=500, detail="DB connection failed") from exc + + async def replace_one(self, object_id: ObjectId, data: Dict[str, Any]) -> Optional[ObjectId]: + """Wrap the replace_one() method, handling timeouts and return data type. + + :param object_id: The object id to replace. + :param data: The data to replace with. + :return: Optional[ObjectId] + """ + + try: + result = await self.collection.replace_one({"_id": object_id}, data) + if result.matched_count == 1: + return object_id + return None + + except BaseException as exc: + print(f"DB connection failed: {exc}") + raise HTTPException(status_code=500, detail="DB connection failed") from exc + + async def delete_one(self, object_id: ObjectId) -> Optional[ObjectId]: + """Wrap the delete_one() method, handling timeouts and return data type. + + :param object_id: The object id to delete from the DB. + :return: Optional[ObjectId] + """ + + try: + result = await self.collection.delete_one({"_id": object_id}) + if isinstance(result.deleted_count, int) and result.deleted_count == 1: + return object_id + return None + + except BaseException as exc: + print(f"DB connection failed: {exc}") + raise HTTPException(status_code=500, detail="DB connection failed") from exc + + async def estimated_document_count(self) -> Optional[int]: + """Wrap the estimated_document_count() method, handling timeouts and return data type. + + :return: Optional[int] + """ + + try: + result = await self.collection.estimated_document_count() + if isinstance(result, int): + return result + return None + + except BaseException 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/healthcheck.py b/src/soc_collector/healthcheck.py new file mode 100644 index 0000000..e561fa0 --- /dev/null +++ b/src/soc_collector/healthcheck.py @@ -0,0 +1,38 @@ +""" +Send a healthcheck request +""" +import sys +import time +import json + +import requests + + +def check_collector() -> bool: + """Check our collector using /info + + :return: bool + """ + time.sleep(2) # Prevent race condition with redis container healthcheck + + req = requests.get( + "https://localhost:8000/info", + timeout=3, + verify="./collector_root_ca.crt", + ) + + if req.status_code != 200: + return False + + data = json.loads(req.text) + if isinstance(data["Estimated document count"], int) and data["Estimated document count"] >= 0: + return req.status_code == 200 + + return False + + +if __name__ == "__main__": + if sys.argv[1] == "COLLECTOR": + if check_collector(): + sys.exit(0) + sys.exit(1) diff --git a/src/soc_collector/healthcheck.sh b/src/soc_collector/healthcheck.sh new file mode 100755 index 0000000..10d599c --- /dev/null +++ b/src/soc_collector/healthcheck.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# If mongodb container +if [ "$1" = "MONGODB" ] +then + /usr/bin/mongosh --eval 'disableTelemetry()' -u "$MONGODB_USERNAME" -p "$MONGODB_PASSWORD" localhost:27017/production /healthcheck-mongodb.js + exit $? +fi + +# If collector +/usr/bin/python3 ./src/soc_collector/healthcheck.py "$1" || exit 1 diff --git a/src/soc_collector/main.py b/src/soc_collector/main.py new file mode 100755 index 0000000..eb6041f --- /dev/null +++ b/src/soc_collector/main.py @@ -0,0 +1,183 @@ +"""Our main module""" +from os import environ +from asyncio import get_running_loop +from sys import exit as app_exit +from json.decoder import JSONDecodeError + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse +from bson import ObjectId +from .db import ( + DBClient, + SearchInput, +) +from .schema import valid_schema +from .auth import authorize_client, load_api_keys + +# Get credentials +if "MONGODB_USERNAME" not in environ or "MONGODB_PASSWORD" not in environ or "MONGODB_COLLECTION" not in environ: + print("Missing MONGODB_USERNAME or MONGODB_PASSWORD or MONGODB_COLLECTION in env") + app_exit(1) + +# Create DB object +db = DBClient(environ["MONGODB_USERNAME"], environ["MONGODB_PASSWORD"], environ["MONGODB_COLLECTION"]) + +# Check DB +loop = get_running_loop() +startup_task = loop.create_task(db.startup()) + +# Load API keys +API_KEYS = load_api_keys("./api_keys.txt") + +# Disable redoc and swagger endpoints +app = FastAPI(docs_url=None, redoc_url=None) + + +@app.post("/sc/v0/search") +async def search(request: Request, search_data: SearchInput) -> JSONResponse: + """/sc/v0/search, POST method + + :param request: The client request. + :param search_data: The search data. + :return: JSONResponse + """ + + # Ensure authorization + authorize_client(request, API_KEYS) + + data = await db.find(search_data) + + return JSONResponse(content={"status": "success", "docs": data}) + + +@app.post("/sc/v0") +async def create(request: Request) -> JSONResponse: + """/sc/v0, POST method + + :param request: The request where we get the json body. + :return: JSONResponse + """ + + # Ensure authorization + authorize_client(request, API_KEYS) + + try: + json_data = await request.json() + except JSONDecodeError: + return JSONResponse(content={"status": "error", "message": "Invalid JSON"}, status_code=400) + + 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) + + if key is None: + return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) + + return JSONResponse(content={"status": "success", "key": str(key)}) + + +@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. + :return: JSONResponse + """ + + # Ensure authorization + authorize_client(request, API_KEYS) + + try: + json_data = await request.json() + 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 = ObjectId(json_data["_id"]) + elif ( + isinstance(json_data["_id"], dict) and "$oid" in json_data["_id"] and isinstance(json_data["_id"]["$oid"], str) + ): + object_id = ObjectId(json_data["_id"]["$oid"]) + else: + return JSONResponse(content={"status": "error", "message": "Missing key '_id' with valid id"}, status_code=400) + + # Ensure the updating key exist + document = await db.find_one(object_id) + + if document is None: + return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404) + + # Ensure valid schema + del json_data["_id"] + if not valid_schema(json_data): + return JSONResponse(content={"status": "error", "message": "Not our JSON schema"}, status_code=400) + + # Replace the data + json_data["_id"] = object_id + returned_object_id = await db.replace_one(object_id, json_data) + + if returned_object_id is None: + return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) + + return JSONResponse(content={"status": "success", "key": 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. + :return: JSONResponse + """ + + # Ensure authorization + authorize_client(request, API_KEYS) + + document = await db.find_one(ObjectId(key)) + + if document is None: + return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404) + + return JSONResponse(content={"status": "success", "doc": document}) + + +@app.delete("/sc/v0/{key}") +async def delete(request: Request, key: str) -> JSONResponse: + """/sc/v0/{key}, DELETE method + + :param key: The document key in the database. + :return: JSONResponse + """ + + # Ensure authorization + authorize_client(request, API_KEYS) + + result = await db.delete_one(ObjectId(key)) + + if result is None: + return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404) + + return JSONResponse(content={"status": "success", "key": str(key)}) + + +@app.get("/info") +async def info(request: Request) -> JSONResponse: + """/info, GET method + + :return: JSONResponse + """ + + # Ensure authorization + authorize_client(request, API_KEYS) + + count = await db.estimated_document_count() + + if count is None: + return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500) + + return JSONResponse(content={"status": "success", "Estimated document count": count}) diff --git a/src/soc_collector/py.typed b/src/soc_collector/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/src/soc_collector/schema.py b/src/soc_collector/schema.py new file mode 100644 index 0000000..221990a --- /dev/null +++ b/src/soc_collector/schema.py @@ -0,0 +1,107 @@ +"""Our schema module""" +from typing import Any, Dict +import jsonschema + +# fmt:off +# NOTE: Commented out properties are left intentionally, so it is easier to see +# what properties are optional. +schema = { + "$schema": "http://json-schema.org/schema#", + "type": "object", + "properties": { + "document_version": {"type": "integer"}, + "ip": {"type": "string"}, + "port": {"type": "integer"}, + "whois_description": {"type": "string"}, + "asn": {"type": "string"}, + "asn_country_code": {"type": "string"}, + "ptr": {"type": "string"}, + "abuse_mail": {"type": "string"}, + "domain": {"type": "string"}, + "timestamp": {"type": "string", "format": "date-time"}, + "display_name": {"type": "string"}, + "description": {"type": "string"}, + "custom_data": { + "type": "object", + "patternProperties": { + ".*": { + "type": "object", + "properties": { + "display_name": {"type": "string"}, + "data": {"type": ["string", "boolean", "integer"]}, + "description": {"type": "string"}, + }, + "required": [ + "display_name", + "data", + # "description" + ] + }, + }, + }, + "result": { + "type": "object", + "patternProperties": { + ".*": { + "type": "object", + "properties": { + "display_name": {"type": "string"}, + "vulnerable": {"type": "boolean"}, + "investigation_needed": {"type": "boolean"}, + "reliability": {"type": "integer"}, + "description": {"type": "string"}, + }, + "oneOf": [ + { + "required": [ + "display_name", + "vulnerable", + # "reliability", # TODO: reliability is required if vulnerable = true + # "description", + ] + }, + { + "required": + [ + "display_name", + "investigation_needed", + # "reliability", # TODO: reliability is required if investigation_needed = true + # "description", + ] + }, + ] + }, + }, + }, + }, + "required": [ + "document_version", + "ip", + "port", + "whois_description", + "asn", + "asn_country_code", + "ptr", + "abuse_mail", + "domain", + "timestamp", + "display_name", + # "description", + # "custom_data", + "result", + ], +} + + +def valid_schema(json_data: Dict[str, Any]) -> bool: + """Check if json data follows the schema. + + :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 diff --git a/src/soc_collector/soc_collector_cli.py b/src/soc_collector/soc_collector_cli.py new file mode 100644 index 0000000..e32b5ad --- /dev/null +++ b/src/soc_collector/soc_collector_cli.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +"""Our database module""" +from os.path import isfile +from os import environ +from typing import Dict, Any +from argparse import ArgumentParser, RawTextHelpFormatter +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"] + +API_URL = "https://collector-dev.soc.sunet.se:8000" +ROOT_CA_FILE = __file__.replace("soc_collector_cli.py", "data/collector_root_ca.crt") + + +def json_load_data(data: str) -> Dict[str, Any]: + """Load json from argument, json data or path to json file + + :param data: String with either json or path to a json file. + :return: json dict as a Dict[str, Any]. + """ + + ret: Dict[str, Any] + try: + if isfile(data): + with open(data, encoding="utf-8") as f_data: + ret = json.loads(f_data.read()) + else: + ret = json.loads(data) + return ret + + except Exception: # pylint: disable=broad-except + print(f"No such file or invalid json data: {data}") + app_exit(1) + + +def info_action() -> None: + """Get database info, currently number of documents.""" + + req = requests.get(f"{API_URL}/info", headers={"API-KEY": API_KEY}, timeout=5) + + # Ensure ok status + req.raise_for_status() + + # Print data + json_data = json.loads(req.text) + print(f"Estimated document count: {json_data['Estimated document count']}") + + +def search_action(data: str) -> None: + """Search for documents in the database. + + :param data: String with either json or path to a json file. + """ + + 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) + + # Ensure ok status + req.raise_for_status() + + # Print data + json_data = json.loads(req.text) + print(json.dumps(json_data["docs"], indent=4)) + + +def delete_action(data: str) -> None: + """Delete a document in the DB. + + :param data: key or path to a json file containing "_id". + """ + + if isfile(data): + data = json_load_data(data)["_id"] + req = requests.delete(f"{API_URL}/sc/v0/{data}", headers={"API-KEY": API_KEY}, timeout=5) + + # Check status + if req.status_code == 404: + print("ERROR: Document not found") + app_exit(1) + + # Ensure ok status + req.raise_for_status() + + print(f"Deleted data OK - key: {data}") + + +def update_local_action(data: str, update_data: str) -> None: + """Update keys and their data in this document. Does not modify the database. + + :param data: json blob or path to json file. + :param update_data: json blob or path to json file. + """ + + json_data = json_load_data(data) + json_update_data = json_load_data(update_data) + + json_data.update(json_update_data) + + # Print data + print(json.dumps(json_data, indent=4)) + + +def replace_action(data: str) -> 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. + """ + + json_data = json_load_data(data) + req = requests.put(f"{API_URL}/sc/v0", json=json_data, headers={"API-KEY": API_KEY}, timeout=5) + + # Check status + if req.status_code == 404: + print("ERROR: Document not found") + app_exit(1) + + # Ensure ok status + req.raise_for_status() + + json_data = json.loads(req.text) + print(f'Inserted data OK - key: {json_data["_id"]}') + + +def insert_action(data: str) -> 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. + """ + + json_data = json_load_data(data) + req = requests.post(f"{API_URL}/sc/v0", json=json_data, headers={"API-KEY": API_KEY}, timeout=5) + + # Ensure ok status + req.raise_for_status() + + data = json.loads(req.text) + print(f'Inserted data OK - key: {json_data["_id"]}') + + +def get_action(data: str) -> None: + """Get a document from the database. + + :param data: key or path to a json file containing "_id". + """ + if isfile(data): + data = json_load_data(data)["_id"] + req = requests.get(f"{API_URL}/sc/v0/{data}", headers={"API-KEY": API_KEY}, timeout=5) + + # Check status + if req.status_code == 404: + print("ERROR: Document not found") + app_exit(1) + + # Ensure ok status + req.raise_for_status() + + # Print data + json_data = json.loads(req.text) + print(json.dumps(json_data["doc"], indent=4)) + + +def main() -> None: + """Main function.""" + parser = ArgumentParser(formatter_class=RawTextHelpFormatter, description="SOC Collector CLI") + parser.add_argument( + "action", + choices=["info", "search", "get", "insert", "replace", "update_local", "delete"], + help="""Action to take + + info: Show info about the database, currently number of documents. + + search: json blog OR path to file. skip defaults to 0 and limit to 25. + '{"search": {"port": 111, "ip": "192.0.2.28"}, "skip": 0, "limit": 25}' OR ./search_data.json + '{"search": {"asn_country_code": "SE", "result": {"$exists": "cve_2015_0002"}}}' + + get: key OR path to document using its "_id". + 637162378c92893fff92bf7e OR ./data.json + + insert: json blob OR path to file. Document MUST NOT contain "_id". + '{json_blob_here...}' OR ./data.json + + replace: json blob OR path to file. Document MUST contain "_id". + '{json_blob_here...}' OR ./updated_document.json + + update_local: json blob OR path to file json blob or path to file with json to update with. + 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". + 637162378c92893fff92bf7e OR ./data.json + + + Pro tip, in ~/.bashrc: + + _soc_collector_cli() + { + local cur prev words cword + _init_completion || return + + local commands command + + commands='info search get insert replace update_local delete' + + if ((cword == 1)); then + COMPREPLY=($(compgen -W "$commands" -- "$cur")) + else + + case $prev in + search) + COMPREPLY="'{\\"search\\": {\\"asn_country_code\\": \\"SE\\", \\"ip\\": \\"8.8.8.8\\"}}'" + return + ;; + esac + command=${words[1]} + _filedir + return + fi + } + complete -F _soc_collector_cli -o default soc_collector_cli + + """, + ) + parser.add_argument("data", default="info", help="json blob or path to file") + parser.add_argument( + "extra_data", + nargs="?", + default=None, + help="json blob or path to file, only used for local_edit", + ) + + args = parser.parse_args() + + if args.action == "get": + get_action(args.data) + elif args.action == "insert": + insert_action(args.data) + elif args.action == "replace": + replace_action(args.data) + 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) + elif args.action == "search": + search_action(args.data) + elif args.action == "info": + info_action() + else: + print("ERROR: Wrong action") + app_exit(1) + + +if __name__ == "__main__": + main() -- cgit v1.1