1 files changed, 183 insertions, 0 deletions
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})