summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorVictor Näslund <victor@sunet.se>2022-11-16 19:09:46 +0100
committerVictor Näslund <victor@sunet.se>2022-11-16 19:09:46 +0100
commit43e87d84b15d12d52a4dcde6e80426cbd17e3d6f (patch)
tree18cef29b77973053522a67677121789ebe032285
parent4a56b3aae4114db731eff725e2c6292371a9b8ae (diff)
auth and CLI done
-rw-r--r--data/api_keys.txt2
-rw-r--r--data/collector-dev.soc.sunet.se.crt12
-rw-r--r--data/collector-dev.soc.sunet.se.key3
-rw-r--r--data/collector_container/Dockerfile15
-rw-r--r--data/collector_root_ca.crt13
-rw-r--r--data/mongodb_container/Dockerfile4
-rwxr-xr-xdev-run.sh30
-rw-r--r--pyproject.toml6
-rw-r--r--requirements.txt12
-rw-r--r--setup.cfg8
-rw-r--r--src/soc_collector/__init__.py (renamed from src/collector/__init__.py)0
-rw-r--r--src/soc_collector/auth.py46
-rw-r--r--src/soc_collector/db.py (renamed from src/collector/db.py)8
-rw-r--r--src/soc_collector/healthcheck.py (renamed from src/collector/healthcheck.py)4
-rwxr-xr-xsrc/soc_collector/healthcheck.sh (renamed from src/collector/healthcheck.sh)2
-rwxr-xr-xsrc/soc_collector/main.py (renamed from src/collector/main.py)47
-rw-r--r--src/soc_collector/py.typed (renamed from src/collector/py.typed)0
-rw-r--r--src/soc_collector/schema.py (renamed from src/collector/schema.py)0
-rw-r--r--src/soc_collector/soc_collector_cli.py259
19 files changed, 420 insertions, 51 deletions
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/soc_collector/__init__.py
index 6530fdd..6530fdd 100644
--- a/src/collector/__init__.py
+++ b/src/soc_collector/__init__.py
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/collector/db.py b/src/soc_collector/db.py
index 641e5da..d601a82 100644
--- a/src/collector/db.py
+++ b/src/soc_collector/db.py
@@ -7,6 +7,10 @@ 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,
@@ -66,9 +70,7 @@ class DBClient:
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)
+ cursor.sort({"ip": ASCENDING, "timestamp": DESCENDING}).limit(search_data.limit).skip(search_data.skip)
try:
async for document in cursor:
diff --git a/src/collector/healthcheck.py b/src/soc_collector/healthcheck.py
index 3091d98..e561fa0 100644
--- a/src/collector/healthcheck.py
+++ b/src/soc_collector/healthcheck.py
@@ -16,9 +16,9 @@ def check_collector() -> bool:
time.sleep(2) # Prevent race condition with redis container healthcheck
req = requests.get(
- "http://localhost:8000/info",
+ "https://localhost:8000/info",
timeout=3,
- # TODO: verify="./rootCA.crt",
+ verify="./collector_root_ca.crt",
)
if req.status_code != 200:
diff --git a/src/collector/healthcheck.sh b/src/soc_collector/healthcheck.sh
index aacc906..10d599c 100755
--- a/src/collector/healthcheck.sh
+++ b/src/soc_collector/healthcheck.sh
@@ -8,4 +8,4 @@ then
fi
# If collector
-/usr/bin/python3 ./src/collector/healthcheck.py "$1" || exit 1
+/usr/bin/python3 ./src/soc_collector/healthcheck.py "$1" || exit 1
diff --git a/src/collector/main.py b/src/soc_collector/main.py
index a2d55c7..eb6041f 100755
--- a/src/collector/main.py
+++ b/src/soc_collector/main.py
@@ -6,16 +6,13 @@ from json.decoder import JSONDecodeError
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
-from bson import (
- ObjectId,
- json_util,
-)
+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:
@@ -29,23 +26,25 @@ db = DBClient(environ["MONGODB_USERNAME"], environ["MONGODB_PASSWORD"], environ[
loop = get_running_loop()
startup_task = loop.create_task(db.startup())
-app = FastAPI()
-
+# Load API keys
+API_KEYS = load_api_keys("./api_keys.txt")
-# @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)
+# Disable redoc and swagger endpoints
+app = FastAPI(docs_url=None, redoc_url=None)
@app.post("/sc/v0/search")
-async def search(search_data: SearchInput) -> JSONResponse:
+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})
@@ -59,6 +58,9 @@ async def create(request: Request) -> JSONResponse:
:return: JSONResponse
"""
+ # Ensure authorization
+ authorize_client(request, API_KEYS)
+
try:
json_data = await request.json()
except JSONDecodeError:
@@ -83,6 +85,9 @@ async def replace(request: Request) -> JSONResponse: # pylint: disable=too-many
:return: JSONResponse
"""
+ # Ensure authorization
+ authorize_client(request, API_KEYS)
+
try:
json_data = await request.json()
except JSONDecodeError:
@@ -123,13 +128,16 @@ async def replace(request: Request) -> JSONResponse: # pylint: disable=too-many
@app.get("/sc/v0/{key}")
-async def get(key: str) -> JSONResponse:
+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:
@@ -139,12 +147,16 @@ async def get(key: str) -> JSONResponse:
@app.delete("/sc/v0/{key}")
-async def delete(key: str) -> JSONResponse:
+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:
@@ -154,12 +166,15 @@ async def delete(key: str) -> JSONResponse:
@app.get("/info")
-async def info() -> JSONResponse:
+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:
diff --git a/src/collector/py.typed b/src/soc_collector/py.typed
index e69de29..e69de29 100644
--- a/src/collector/py.typed
+++ b/src/soc_collector/py.typed
diff --git a/src/collector/schema.py b/src/soc_collector/schema.py
index 221990a..221990a 100644
--- a/src/collector/schema.py
+++ b/src/soc_collector/schema.py
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()