diff --git a/reciever/main.py b/reciever/main.py new file mode 100644 index 0000000..82ceace --- /dev/null +++ b/reciever/main.py @@ -0,0 +1,109 @@ +import sys +#!/usr/bin/env python3 +import os +import os.path +import sqlite3 +import uuid +from os import makedirs +import time +from typing import Annotated + +from argon2 import PasswordHasher +from argon2.exceptions import VerifyMismatchError +from fastapi import Depends, FastAPI, Response, UploadFile, status +from fastapi.security import HTTPBasic, HTTPBasicCredentials + +# Security singleton +security = HTTPBasic() + +class Inventory: + + def __init__(self): + + try: + self.host_dir = os.environ["INVENT_HOST_DIR"] + except KeyError: + self.host_dir = '/var/cache/invent/hosts' + try: + self.image_dir = os.environ["INVENT_IMAGE_DIR"] + except KeyError: + self.image_dir = '/var/cache/invent/images' + try: + self.db_dir = os.environ["INVENT_DB_DIR"] + except KeyError: + self.db_dir = '/etc/invent/db' + if not os.path.isdir(self.host_dir): + makedirs(self.host_dir) + if not os.path.isdir(self.image_dir): + makedirs(self.image_dir) + if not os.path.isdir(self.db_dir): + makedirs(self.db_dir) + + self.ph = PasswordHasher() + self.db = sqlite3.connect(os.path.join(self.db_dir, "users.db")) + self.cursor = self.db.cursor() + self.cursor.execute("CREATE TABLE IF NOT EXISTS users(username, salt, hash)") + self.cursor.execute("CREATE TABLE IF NOT EXISTS reports(username, timestamp, endpoint)") + + + + def get_or_create_user(self, credentials: Annotated[HTTPBasicCredentials, Depends(security)]) -> tuple[str,str]: + username = credentials.username + query = self.cursor.execute(f"SELECT salt, hash FROM users WHERE username='{username}'") + result = query.fetchone() + # If the user is not in the database, we will trust the user and add it to the database + # TOFU: https://developer.mozilla.org/en-US/docs/Glossary/TOFU + if result is None: + salt = uuid.uuid4().hex + hash = self.ph.hash(salt + credentials.password) + self.cursor.execute(f"INSERT INTO users (username, salt, hash) values('{username}', '{salt}', '{hash}')") + self.db.commit() + return (salt, hash) + else: + return result + + def validate_credentials(self,credentials: Annotated[HTTPBasicCredentials, Depends(security)], salt: str, hash: str) -> bool: + password = credentials.password + try: + self.ph.verify(hash, salt + password) + except VerifyMismatchError: + return False + return True + + async def upload(self, endpoint: str, file: UploadFile, name: str, credentials: Annotated[HTTPBasicCredentials, Depends(security)], response: Response): + dir = self.host_dir + if endpoint == 'image': + dir = self.image_dir + + if credentials.username != name: + response.status_code = status.HTTP_403_FORBIDDEN + return {"ERROR": "Username and andpoint does not match"} + salt, hash = self.get_or_create_user(credentials) + if self.validate_credentials(credentials, salt, hash): + filename = os.path.join(dir, name + '.json') + with open(filename, 'w') as fh: + contents: bytes = await file.read() + fh.write(contents.decode('utf-8')) + + return {"SUCCESS": f"File: {filename} saved"} + else: + query = self.cursor.execute(f"SELECT * FROM reports WHERE username='{credentials.username}' and endpoint={endpoint}") + result = query.fetchone() + if result is None: + timestamp = int(time.time()) + self.cursor.execute(f"INSERT INTO reports (username, timestamp) values('{credentials.username}','{endpoint}', '{timestamp}')") + self.db.commit() + response.status_code = status.HTTP_401_UNAUTHORIZED + return {"ERROR": "Invalid password, this incident will be reported"} + + +app = FastAPI() +inventory = Inventory() + +@app.post("/host/{hostname}", status_code=status.HTTP_201_CREATED) +async def upload_host(file: UploadFile, hostname: str, credentials: Annotated[HTTPBasicCredentials, Depends(security)], response: Response): + return await inventory.upload('host', file, hostname, credentials, response) + +@app.post("/image/{imagename}", status_code=status.HTTP_201_CREATED) +async def upload_image(file: UploadFile, imagename: str, credentials: Annotated[HTTPBasicCredentials, Depends(security)], response: Response): + return await inventory.upload('image', file, imagename, credentials, response) diff --git a/reciever/requirements.txt b/reciever/requirements.txt new file mode 100644 index 0000000..f07f9aa --- /dev/null +++ b/reciever/requirements.txt @@ -0,0 +1,5 @@ +fastapi[all] +uvicorn[standard] +python-multipart +argon2-cffi +uuid diff --git a/scanner.py b/scanner/scanner.py similarity index 100% rename from scanner.py rename to scanner/scanner.py