#!/usr/bin/python3
import base64
import hashlib
import hmac
import ipaddress
import json
import os
import re
import secrets
import shutil
import sqlite3
import threading
import time
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import urlsplit

APP_ROOT = Path("/usr/share/neodrop")
DATA_ROOT = Path("{{NEODROP_DATA_ROOT}}")
STAGING_ROOT = DATA_ROOT / "staging"
OBJECT_ROOT = DATA_ROOT / "objects"
DB_PATH = DATA_ROOT / "neodrop.sqlite3"
MANAGER_TOKEN_HASH_PATH = Path("{{NEODROP_MANAGER_TOKEN_HASH_PATH}}")
LISTEN_ADDRESS = "{{NEODROP_BIND_ADDRESS}}"
LISTEN_PORT = 8080
TRUSTED_PROXY = ipaddress.ip_address("{{TRUSTED_PROXY_ADDRESS}}")
TRUSTED_UPLOADERS = (
    ipaddress.ip_network("{{UPLOAD_CIDR_1}}"),
    ipaddress.ip_network("{{UPLOAD_CIDR_2}}"),
    ipaddress.ip_network("{{UPLOAD_CIDR_3}}"),
)
ALLOWED_ORIGIN = "{{RELAY_ORIGIN}}"
MANAGEMENT_ORIGIN = "{{DROPS_ORIGIN}}"
MANAGEMENT_UPLOADERS = (
    ipaddress.ip_network("{{MANAGEMENT_CIDR_1}}"),
    ipaddress.ip_network("{{MANAGEMENT_CIDR_2}}"),
    ipaddress.ip_network("{{MANAGEMENT_CIDR_3}}"),
)
UPLOAD_ORIGINS = (ALLOWED_ORIGIN, MANAGEMENT_ORIGIN)
CHUNK_SIZE = 4 * 1024 * 1024
PROFILES = {
    0: (50 * 1024 * 1024, 30 * 24 * 60 * 60),
    1: (1024 * 1024 * 1024, 7 * 24 * 60 * 60),
}
UPLOAD_TTL = 24 * 60 * 60
ID_RE = re.compile(r"^[A-Za-z0-9_-]{32}$")
UPLOAD_RE = re.compile(r"^[A-Za-z0-9_-]{22}$")
UNAVAILABLE = b'{"error":"unavailable"}\n'
DB_LOCK = threading.RLock()


def now():
    return int(time.time())


def b64url(data):
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def token_hash(token):
    return hashlib.sha256(token.encode("ascii")).digest()


def valid_manager_token(token):
    if not token or len(token) > 128:
        return False
    try:
        expected = MANAGER_TOKEN_HASH_PATH.read_text().strip()
        actual = hashlib.sha256(token.encode("ascii", "strict")).hexdigest()
    except (OSError, UnicodeEncodeError):
        return False
    return len(expected) == 64 and hmac.compare_digest(actual, expected)


def valid_object_id(value):
    if not ID_RE.fullmatch(value or ""):
        return False
    try:
        return len(base64.urlsafe_b64decode(value + "==")) == 24
    except (ValueError, TypeError):
        return False


def database():
    conn = sqlite3.connect(DB_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA journal_mode=WAL")
    return conn


def initialize_database():
    DATA_ROOT.mkdir(mode=0o700, parents=True, exist_ok=True)
    STAGING_ROOT.mkdir(mode=0o700, exist_ok=True)
    OBJECT_ROOT.mkdir(mode=0o700, exist_ok=True)
    with database() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS tombstones (
                object_id TEXT PRIMARY KEY,
                first_seen INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS uploads (
                upload_id TEXT PRIMARY KEY,
                object_id TEXT UNIQUE NOT NULL,
                token_hash BLOB NOT NULL,
                profile INTEGER NOT NULL,
                size INTEGER NOT NULL,
                chunks INTEGER NOT NULL,
                created_at INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS objects (
                object_id TEXT PRIMARY KEY,
                token_hash BLOB NOT NULL,
                profile INTEGER NOT NULL,
                size INTEGER NOT NULL,
                chunks INTEGER NOT NULL,
                created_at INTEGER NOT NULL,
                expires_at INTEGER NOT NULL
            );
            CREATE INDEX IF NOT EXISTS objects_expiry ON objects(expires_at);
            CREATE INDEX IF NOT EXISTS uploads_created ON uploads(created_at);
            """
        )


def object_path(object_id):
    return OBJECT_ROOT / object_id[:2] / object_id


def report_invalid(client_ip, source, detail, destination_port=None):
    event = {
        "ip": str(client_ip),
        "source": source,
        "detail": detail[:160],
    }
    if destination_port is not None:
        event["destination_port"] = destination_port
    payload = json.dumps(event).encode("utf-8")
    request = urllib.request.Request(
        "{{NEOGUARD_INTERNAL_URL}}",
        data=payload,
        method="POST",
        headers={"Content-Type": "application/json"},
    )
    try:
        urllib.request.urlopen(request, timeout=0.25).close()
    except Exception:
        pass


def cleanup():
    cutoff = now() - UPLOAD_TTL
    expired_paths = []
    object_paths = []
    with DB_LOCK, database() as conn:
        for row in conn.execute("SELECT upload_id FROM uploads WHERE created_at <= ?", (cutoff,)):
            expired_paths.append(STAGING_ROOT / row["upload_id"])
        conn.execute("DELETE FROM uploads WHERE created_at <= ?", (cutoff,))
        for row in conn.execute("SELECT object_id FROM objects WHERE expires_at <= ?", (now(),)):
            object_paths.append((row["object_id"], object_path(row["object_id"])))
    for path in expired_paths:
        shutil.rmtree(path, ignore_errors=True)
    for object_id, path in object_paths:
        try:
            shutil.rmtree(path)
        except FileNotFoundError:
            pass
        except OSError:
            continue
        with DB_LOCK, database() as conn:
            conn.execute("DELETE FROM objects WHERE object_id = ? AND expires_at <= ?", (object_id, now()))


def cleanup_loop():
    while True:
        try:
            cleanup()
        except Exception:
            pass
        time.sleep(3600)


class NeoDropHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "NeoDrop"
    sys_version = ""

    def log_message(self, fmt, *args):
        print("%s %s" % (self.address_string(), fmt % args), flush=True)

    def client_ip(self):
        peer = ipaddress.ip_address(self.client_address[0])
        if peer == TRUSTED_PROXY:
            forwarded = self.headers.get("X-Real-IP", "")
            try:
                return ipaddress.ip_address(forwarded)
            except ValueError:
                return peer
        return peer

    def destination_port(self):
        peer = ipaddress.ip_address(self.client_address[0])
        if peer == TRUSTED_PROXY:
            value = self.headers.get("X-Forwarded-Port", "")
            if value.isdecimal() and 1 <= int(value) <= 65535:
                return int(value)
            return None
        return int(self.server.server_address[1])

    def is_trusted_uploader(self):
        address = self.client_ip()
        return any(address in network for network in TRUSTED_UPLOADERS)

    def send_headers(self, status, content_type, length, extra=None):
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(length))
        self.send_header("Cache-Control", "private, no-store, max-age=0")
        self.send_header("Pragma", "no-cache")
        self.send_header("X-Content-Type-Options", "nosniff")
        if extra:
            for key, value in extra.items():
                self.send_header(key, value)
        self.end_headers()

    def send_json(self, status, value, cors=False):
        body = json.dumps(value, separators=(",", ":")).encode("utf-8") + b"\n"
        extra = {}
        if cors:
            origin = self.headers.get("Origin", "")
            if origin in UPLOAD_ORIGINS:
                extra.update({
                    "Access-Control-Allow-Origin": origin,
                    "Vary": "Origin",
                })
        self.send_headers(status, "application/json", len(body), extra)
        if self.command != "HEAD":
            self.wfile.write(body)

    def unavailable(self, detail="unavailable"):
        report_invalid(self.client_ip(), "drops", detail, self.destination_port())
        self.send_headers(404, "application/json", len(UNAVAILABLE))
        if self.command != "HEAD":
            self.wfile.write(UNAVAILABLE)

    def forbidden_upload(self, detail):
        self.close_connection = True
        report_invalid(self.client_ip(), "drops", detail, self.destination_port())
        self.send_json(403, {"error": "forbidden"}, cors=True)

    def read_json(self, maximum=8192):
        try:
            length = int(self.headers.get("Content-Length", "-1"))
        except ValueError:
            return None
        if length < 0 or length > maximum:
            return None
        try:
            return json.loads(self.rfile.read(length))
        except (json.JSONDecodeError, UnicodeDecodeError):
            return None

    def upload_authorized(self):
        origin = self.headers.get("Origin")
        address = self.client_ip()
        if origin == ALLOWED_ORIGIN:
            source_allowed = any(address in network for network in TRUSTED_UPLOADERS)
        elif origin == MANAGEMENT_ORIGIN:
            source_allowed = any(address in network for network in MANAGEMENT_UPLOADERS)
        else:
            source_allowed = False
        if not source_allowed:
            self.forbidden_upload("external upload")
            return False
        if origin not in UPLOAD_ORIGINS:
            self.forbidden_upload("invalid upload origin")
            return False
        return True

    def bearer(self):
        value = self.headers.get("Authorization", "")
        if not value.startswith("Bearer "):
            return ""
        return value[7:]

    def manager_authorized(self):
        if not self.is_trusted_uploader():
            return False
        return valid_manager_token(self.bearer())

    def lookup_upload(self, upload_id):
        token = self.bearer()
        if not token or len(token) > 128:
            return None
        with database() as conn:
            row = conn.execute("SELECT * FROM uploads WHERE upload_id = ?", (upload_id,)).fetchone()
        if not row or not hmac.compare_digest(row["token_hash"], token_hash(token)):
            return None
        if row["created_at"] + UPLOAD_TTL <= now():
            return None
        return row

    def do_OPTIONS(self):
        if not self.upload_authorized():
            return
        path = urlsplit(self.path).path
        requested = self.headers.get("Access-Control-Request-Method", "")
        valid = (
            (requested == "POST" and (
                path == "/api/v1/uploads"
                or re.fullmatch(r"/api/v1/uploads/[A-Za-z0-9_-]{22}/commit", path)
            ))
            or (requested == "PUT" and (
                re.fullmatch(r"/api/v1/uploads/[A-Za-z0-9_-]{22}/metadata", path)
                or re.fullmatch(r"/api/v1/uploads/[A-Za-z0-9_-]{22}/chunks/\d{1,4}", path)
            ))
            or (requested == "DELETE" and re.fullmatch(
                r"/api/v1/objects/[A-Za-z0-9_-]{32}", path
            ))
        )
        requested_headers = {
            value.strip().lower()
            for value in self.headers.get("Access-Control-Request-Headers", "").split(",")
            if value.strip()
        }
        if not valid or not requested_headers.issubset({"authorization", "content-type"}):
            return self.unavailable("invalid preflight")
        self.send_response(204)
        self.send_header("Content-Length", "0")
        self.send_header("Access-Control-Allow-Origin", self.headers["Origin"])
        self.send_header("Access-Control-Allow-Methods", requested + ", OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Authorization, Content-Type")
        self.send_header("Access-Control-Max-Age", "600")
        self.send_header("Vary", "Origin")
        self.end_headers()

    def do_POST(self):
        path = urlsplit(self.path).path
        if path == "/api/v1/uploads":
            return self.create_upload()
        match = re.fullmatch(r"/api/v1/uploads/([A-Za-z0-9_-]{22})/commit", path)
        if match:
            return self.commit_upload(match.group(1))
        self.unavailable("invalid POST path")

    def create_upload(self):
        if not self.upload_authorized():
            return
        if self.headers.get("Content-Type") != "application/json":
            self.close_connection = True
            return self.send_json(415, {"error": "invalid request"}, cors=True)
        data = self.read_json()
        if not isinstance(data, dict):
            return self.send_json(400, {"error": "invalid request"}, cors=True)
        try:
            object_id = data["id"]
            profile = int(data["profile"])
            size = int(data["size"])
            chunk_size = int(data["chunkSize"])
            chunks = int(data["chunks"])
        except (KeyError, TypeError, ValueError):
            return self.send_json(400, {"error": "invalid request"}, cors=True)
        if not valid_object_id(object_id) or profile not in PROFILES:
            return self.send_json(400, {"error": "invalid request"}, cors=True)
        maximum, _ = PROFILES[profile]
        expected_chunks = (size + CHUNK_SIZE - 1) // CHUNK_SIZE if size else 0
        if size < 0 or size > maximum or chunk_size != CHUNK_SIZE or chunks != expected_chunks:
            return self.send_json(400, {"error": "invalid request"}, cors=True)
        upload_id = b64url(secrets.token_bytes(16))
        token = b64url(secrets.token_bytes(32))
        created = now()
        with DB_LOCK:
            try:
                with database() as conn:
                    conn.execute("INSERT INTO tombstones VALUES (?, ?)", (object_id, created))
                    conn.execute(
                        "INSERT INTO uploads VALUES (?, ?, ?, ?, ?, ?, ?)",
                        (upload_id, object_id, token_hash(token), profile, size, chunks, created),
                    )
                stage = STAGING_ROOT / upload_id
                (stage / "chunks").mkdir(mode=0o700, parents=True)
            except (sqlite3.IntegrityError, FileExistsError):
                return self.send_json(409, {"error": "identifier unavailable"}, cors=True)
        self.send_json(201, {
            "v": 1,
            "uploadId": upload_id,
            "uploadToken": token,
            "sessionExpiresAt": created + UPLOAD_TTL,
        }, cors=True)

    def do_PUT(self):
        path = urlsplit(self.path).path
        metadata = re.fullmatch(r"/api/v1/uploads/([A-Za-z0-9_-]{22})/metadata", path)
        chunk = re.fullmatch(r"/api/v1/uploads/([A-Za-z0-9_-]{22})/chunks/(\d{1,4})", path)
        if metadata:
            return self.put_record(metadata.group(1), None)
        if chunk:
            return self.put_record(chunk.group(1), int(chunk.group(2)))
        self.unavailable("invalid PUT path")

    def put_record(self, upload_id, index):
        if not self.upload_authorized():
            return
        if self.headers.get("Content-Type") != "application/octet-stream":
            self.close_connection = True
            return self.send_json(415, {"error": "invalid record"}, cors=True)
        row = self.lookup_upload(upload_id)
        if not row:
            return self.forbidden_upload("invalid upload authorization")
        try:
            length = int(self.headers.get("Content-Length", "-1"))
        except ValueError:
            length = -1
        if index is None:
            target = STAGING_ROOT / upload_id / "metadata.bin"
            valid_length = 16 <= length <= 4112
        else:
            if index < 0 or index >= row["chunks"]:
                valid_length = False
                expected = -1
            else:
                remaining = row["size"] - index * CHUNK_SIZE
                expected = min(CHUNK_SIZE, remaining) + 16
                valid_length = length == expected
            target = STAGING_ROOT / upload_id / "chunks" / ("%08d.bin" % index)
        if not valid_length:
            return self.send_json(400, {"error": "invalid record"}, cors=True)
        body = self.rfile.read(length)
        if len(body) != length:
            return self.send_json(400, {"error": "incomplete record"}, cors=True)
        if target.exists():
            if target.read_bytes() != body:
                return self.send_json(409, {"error": "immutable record"}, cors=True)
        else:
            temporary = target.with_name(target.name + "." + secrets.token_hex(8) + ".tmp")
            with open(temporary, "xb") as handle:
                handle.write(body)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temporary, target)
        self.send_json(200, {"ok": True}, cors=True)

    def commit_upload(self, upload_id):
        if not self.upload_authorized():
            return
        if self.headers.get("Content-Type") != "application/json":
            self.close_connection = True
            return self.send_json(415, {"error": "invalid request"}, cors=True)
        data = self.read_json(maximum=64)
        if data != {}:
            return self.send_json(400, {"error": "invalid request"}, cors=True)
        row = self.lookup_upload(upload_id)
        if not row:
            return self.forbidden_upload("invalid commit authorization")
        stage = STAGING_ROOT / upload_id
        metadata = stage / "metadata.bin"
        if not metadata.is_file() or not 16 <= metadata.stat().st_size <= 4112:
            return self.send_json(409, {"error": "upload incomplete"}, cors=True)
        for index in range(row["chunks"]):
            path = stage / "chunks" / ("%08d.bin" % index)
            remaining = row["size"] - index * CHUNK_SIZE
            if not path.is_file() or path.stat().st_size != min(CHUNK_SIZE, remaining) + 16:
                return self.send_json(409, {"error": "upload incomplete"}, cors=True)
        created = now()
        expires = created + PROFILES[row["profile"]][1]
        destination = object_path(row["object_id"])
        destination.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
        moved = False
        with DB_LOCK:
            try:
                os.rename(stage, destination)
                moved = True
                with database() as conn:
                    conn.execute(
                        "INSERT INTO objects VALUES (?, ?, ?, ?, ?, ?, ?)",
                        (row["object_id"], row["token_hash"], row["profile"], row["size"],
                         row["chunks"], created, expires),
                    )
                    conn.execute("DELETE FROM uploads WHERE upload_id = ?", (upload_id,))
            except Exception:
                if moved and destination.exists() and not stage.exists():
                    os.rename(destination, stage)
                raise
        self.send_json(200, {"v": 1, "expiresAt": expires}, cors=True)

    def do_DELETE(self):
        path = urlsplit(self.path).path
        match = re.fullmatch(r"/api/v1/objects/([A-Za-z0-9_-]{32})", path)
        if not match or not self.upload_authorized():
            return
        if self.headers.get("Content-Length", "0") != "0":
            self.close_connection = True
            return self.send_json(400, {"error": "invalid request"}, cors=True)
        object_id = match.group(1)
        token = self.bearer()
        target = object_path(object_id)
        with DB_LOCK, database() as conn:
            row = conn.execute("SELECT token_hash FROM objects WHERE object_id = ?", (object_id,)).fetchone()
            if not row or not token or not hmac.compare_digest(row["token_hash"], token_hash(token)):
                return self.forbidden_upload("invalid delete authorization")
            if target.exists():
                shutil.rmtree(target)
            conn.execute("DELETE FROM objects WHERE object_id = ?", (object_id,))
        self.send_json(200, {"deleted": True}, cors=True)

    def do_HEAD(self):
        return self.do_GET()

    def do_GET(self):
        path = urlsplit(self.path).path
        if path == "/api/v1/manager/objects":
            if not self.manager_authorized():
                return self.response_forbidden()
            return self.manager_inventory()
        manifest = re.fullmatch(r"/api/v1/objects/([A-Za-z0-9_-]{32})", path)
        chunk = re.fullmatch(r"/api/v1/objects/([A-Za-z0-9_-]{32})/chunks/(\d{1,4})", path)
        if manifest:
            return self.get_manifest(manifest.group(1))
        if chunk:
            return self.get_chunk(chunk.group(1), int(chunk.group(2)))
        if path == "/" or re.fullmatch(r"/d/[A-Za-z0-9_-]{32}", path):
            return self.static_file("index.html")
        if path in ("/app.js", "/style.css", "/favicon.png", "/avatar-login.jpg"):
            return self.static_file(path[1:])
        self.unavailable("invalid GET path")

    def response_forbidden(self):
        self.send_json(403, {"error": "forbidden"})

    def manager_inventory(self):
        with database() as conn:
            rows = list(conn.execute(
                "SELECT object_id,profile,size,chunks,created_at,expires_at "
                "FROM objects ORDER BY created_at DESC LIMIT 10000"
            ))
        self.send_json(200, {"objects": [
            {
                "id": row["object_id"],
                "profile": row["profile"],
                "size": row["size"],
                "chunks": row["chunks"],
                "createdAt": row["created_at"],
                "expiresAt": row["expires_at"],
            }
            for row in rows
        ]})

    def active_object(self, object_id):
        if not valid_object_id(object_id):
            return None
        with database() as conn:
            return conn.execute(
                "SELECT * FROM objects WHERE object_id = ? AND expires_at > ?",
                (object_id, now()),
            ).fetchone()

    def get_manifest(self, object_id):
        row = self.active_object(object_id)
        metadata = object_path(object_id) / "metadata.bin"
        if not row or not metadata.is_file():
            return self.unavailable("unknown object")
        self.send_json(200, {
            "v": 1,
            "profile": row["profile"],
            "size": row["size"],
            "chunkSize": CHUNK_SIZE,
            "chunks": row["chunks"],
            "metadata": b64url(metadata.read_bytes()),
            "expiresAt": row["expires_at"],
        })

    def get_chunk(self, object_id, index):
        row = self.active_object(object_id)
        if not row or index < 0 or index >= row["chunks"]:
            return self.unavailable("unknown object chunk")
        path = object_path(object_id) / "chunks" / ("%08d.bin" % index)
        if not path.is_file():
            return self.unavailable("unknown object chunk")
        length = path.stat().st_size
        self.send_headers(200, "application/octet-stream", length)
        if self.command != "HEAD":
            with open(path, "rb") as handle:
                shutil.copyfileobj(handle, self.wfile, length=64 * 1024)

    def static_file(self, name):
        path = APP_ROOT / name
        if not path.is_file():
            return self.unavailable("missing application asset")
        types = {"index.html": "text/html; charset=utf-8", "app.js": "application/javascript",
                 "style.css": "text/css", "favicon.png": "image/png",
                 "avatar-login.jpg": "image/jpeg"}
        body = path.read_bytes()
        cached = name in ("favicon.png", "avatar-login.jpg")
        extra = {"Cache-Control": "public, max-age=604800"} if cached else {"Cache-Control": "no-cache"}
        self.send_headers(200, types[name], len(body), extra)
        if self.command != "HEAD":
            self.wfile.write(body)


class NeoDropServer(ThreadingHTTPServer):
    daemon_threads = True
    request_queue_size = 256


if __name__ == "__main__":
    os.umask(0o077)
    initialize_database()
    cleanup()
    threading.Thread(target=cleanup_loop, daemon=True).start()
    server = NeoDropServer((LISTEN_ADDRESS, LISTEN_PORT), NeoDropHandler)
    server.serve_forever()