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

DATA_ROOT = Path("{{NEOGUARD_DATA_ROOT}}")
DB_PATH = DATA_ROOT / "neoguard.sqlite3"
KEY_PATH = Path("{{NEOGUARD_KEY_PATH}}")
LISTEN_ADDRESS = "{{NEOGUARD_BIND_ADDRESS}}"
LISTEN_PORT = 8081
TRUSTED_PROXY = ipaddress.ip_address("{{TRUSTED_PROXY_ADDRESS}}")
EXEMPT_NETWORKS = (
    ipaddress.ip_network("{{UPLOAD_CIDR_1}}"),
    ipaddress.ip_network("{{UPLOAD_CIDR_2}}"),
    ipaddress.ip_network("{{UPLOAD_CIDR_3}}"),
)
WINDOW = 600
MAX_FAILURES = 21
BAN_SECONDS = 24 * 60 * 60
DEFAULT_REDIRECT = "{{BLOCKED_REDIRECT_URL}}"
DB_LOCK = threading.RLock()
NONCES = {}
URL_RE = re.compile(r"^https://[^\s\\\"'{};$]{1,2039}$")


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


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


def initialize_database():
    DATA_ROOT.mkdir(mode=0o700, parents=True, exist_ok=True)
    with database() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS attempts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                occurred_at INTEGER NOT NULL,
                address TEXT NOT NULL,
                source TEXT NOT NULL,
                detail TEXT NOT NULL,
                destination_port INTEGER
            );
            CREATE INDEX IF NOT EXISTS attempts_window ON attempts(address, occurred_at);
            CREATE TABLE IF NOT EXISTS offenses (
                address TEXT PRIMARY KEY,
                count INTEGER NOT NULL,
                last_at INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS bans (
                address TEXT PRIMARY KEY,
                started_at INTEGER NOT NULL,
                expires_at INTEGER,
                offense INTEGER NOT NULL,
                action TEXT NOT NULL,
                target TEXT NOT NULL,
                message TEXT NOT NULL,
                reason TEXT NOT NULL,
                rdns TEXT NOT NULL DEFAULT '',
                active INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                occurred_at INTEGER NOT NULL,
                level TEXT NOT NULL,
                message TEXT NOT NULL,
                destination_port INTEGER
            );
            """
        )
        for table in ("attempts", "events"):
            columns = {row["name"] for row in conn.execute("PRAGMA table_info(%s)" % table)}
            if "destination_port" not in columns:
                conn.execute("ALTER TABLE %s ADD COLUMN destination_port INTEGER" % table)


def normalize_address(value):
    address = ipaddress.ip_address(value)
    if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped:
        address = address.ipv4_mapped
    return address


def exempt(address):
    return any(address in network for network in EXEMPT_NETWORKS)


def add_event(conn, level, message, destination_port=None):
    conn.execute("INSERT INTO events(occurred_at, level, message, destination_port) VALUES (?, ?, ?, ?)",
                 (now(), level, message[:2000], destination_port))


def confirmed_rdns(address):
    try:
        hostname = socket.gethostbyaddr(str(address))[0].lower().rstrip(".")
        if not hostname or len(hostname) > 253 or any(ord(c) < 33 for c in hostname):
            return ""
        results = socket.getaddrinfo(hostname, None)
        returned = {normalize_address(item[4][0]) for item in results}
        return hostname if address in returned else ""
    except (OSError, ValueError):
        return ""


def resolve_ban(address_text):
    address = normalize_address(address_text)
    hostname = confirmed_rdns(address)
    with DB_LOCK, database() as conn:
        conn.execute("UPDATE bans SET rdns = ? WHERE address = ?", (hostname, str(address)))
        if hostname:
            add_event(conn, "INFO", "[RDNS] ip=%s rdns=%s" % (address, hostname))


def expire_bans(conn):
    timestamp = now()
    rows = list(conn.execute("SELECT address FROM bans WHERE active = 1 AND expires_at IS NOT NULL AND expires_at <= ?", (timestamp,)))
    for row in rows:
        conn.execute("UPDATE bans SET active = 0 WHERE address = ?", (row["address"],))
        add_event(conn, "INFO", "[UNBAN] ip=%s reason=expired" % row["address"])


def normalize_destination_port(value):
    if value is None or value == "":
        return None
    if isinstance(value, bool):
        raise ValueError("invalid destination port")
    if isinstance(value, int):
        port = value
    elif isinstance(value, str) and value.isdecimal():
        port = int(value)
    else:
        raise ValueError("invalid destination port")
    if not 1 <= port <= 65535:
        raise ValueError("invalid destination port")
    return port


def record_attempt(address_value, source, detail, destination_port=None):
    destination_port = normalize_destination_port(destination_port)
    try:
        address = normalize_address(address_value)
    except ValueError:
        return
    if exempt(address):
        return
    address_text = str(address)
    timestamp = now()
    created_ban = False
    with DB_LOCK, database() as conn:
        expire_bans(conn)
        active = conn.execute("SELECT 1 FROM bans WHERE address = ? AND active = 1", (address_text,)).fetchone()
        conn.execute("DELETE FROM attempts WHERE occurred_at < ?", (timestamp - WINDOW,))
        if active:
            return
        conn.execute("INSERT INTO attempts(occurred_at,address,source,detail,destination_port) VALUES (?,?,?,?,?)",
                     (timestamp, address_text, source[:32], detail[:300], destination_port))
        count = conn.execute("SELECT COUNT(*) AS n FROM attempts WHERE address = ? AND occurred_at >= ?",
                             (address_text, timestamp - WINDOW)).fetchone()["n"]
        if count in (1, 10, 20):
            port_text = " port=%d" % destination_port if destination_port else ""
            add_event(conn, "WARN", "[WARN] ip=%s%s rdns=pending count=%d/%d window=10m source=%s" %
                       (address_text, port_text, count, MAX_FAILURES, source[:32]), destination_port)
        if count >= MAX_FAILURES:
            offense_row = conn.execute("SELECT count FROM offenses WHERE address = ?", (address_text,)).fetchone()
            offense = (offense_row["count"] if offense_row else 0) + 1
            conn.execute("INSERT INTO offenses VALUES (?,?,?) ON CONFLICT(address) DO UPDATE SET count=excluded.count,last_at=excluded.last_at",
                         (address_text, offense, timestamp))
            expires = None if offense >= 3 else timestamp + BAN_SECONDS
            conn.execute("""
                INSERT INTO bans(address,started_at,expires_at,offense,action,target,message,reason,rdns,active)
                VALUES (?,?,?,?,?,?,?,?,?,1)
                ON CONFLICT(address) DO UPDATE SET started_at=excluded.started_at,expires_at=excluded.expires_at,
                  offense=excluded.offense,action=excluded.action,target=excluded.target,message=excluded.message,
                  reason=excluded.reason,rdns='',active=1
            """, (address_text, timestamp, expires, offense, "redirect", DEFAULT_REDIRECT,
                    "Access denied by NeoGuard.", "automatic threshold", ""))
            conn.execute("DELETE FROM attempts WHERE address = ?", (address_text,))
            level = "PERM" if expires is None else "BAN"
            until = "permanent" if expires is None else time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(expires))
            port_text = " port=%d" % destination_port if destination_port else ""
            add_event(conn, level, "[%s] ip=%s%s rdns=pending offense=%d failures=%d until=%s action=redirect" %
                       (level, address_text, port_text, offense, count, until), destination_port)
            created_ban = True
    if created_ban:
        threading.Thread(target=resolve_ban, args=(address_text,), daemon=True).start()


def active_ban(address_value):
    try:
        address = normalize_address(address_value)
    except ValueError:
        return None
    if exempt(address):
        return None
    with DB_LOCK, database() as conn:
        expire_bans(conn)
        return conn.execute("SELECT * FROM bans WHERE address = ? AND active = 1", (str(address),)).fetchone()


def duration_value(value):
    if value == "permanent":
        return None
    match = re.fullmatch(r"(\d+)(m|h|d)", value)
    if not match:
        raise ValueError("duration must be Nm, Nh, Nd, or permanent")
    factor = {"m": 60, "h": 3600, "d": 86400}[match.group(2)]
    duration = int(match.group(1)) * factor
    if duration < 60 or duration > 10 * 365 * 86400:
        raise ValueError("duration out of range")
    return now() + duration


def format_ban(row):
    until = "permanent" if row["expires_at"] is None else time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(row["expires_at"]))
    return "ip=%s rdns=%s offense=%d until=%s action=%s target=%s reason=%s" % (
        row["address"], row["rdns"] or "-", row["offense"], until, row["action"],
        row["target"] if row["action"] == "redirect" else "-", row["reason"])


def run_command(command):
    try:
        args = shlex.split(command)
    except ValueError as error:
        return "ERROR " + str(error)
    if not args:
        return "Commands: list, show IP, add IP DURATION, remove IP, extend IP DURATION, action IP redirect URL|message TEXT, status"
    operation = args[0].lower()
    with DB_LOCK, database() as conn:
        expire_bans(conn)
        if operation == "list":
            rows = list(conn.execute("SELECT * FROM bans WHERE active=1 ORDER BY started_at DESC LIMIT 25"))
            return "No active bans." if not rows else "\n".join(format_ban(row) for row in rows)
        if operation == "status":
            active = conn.execute("SELECT COUNT(*) n FROM bans WHERE active=1").fetchone()["n"]
            attempts = conn.execute("SELECT COUNT(*) n FROM attempts WHERE occurred_at>=?", (now()-WINDOW,)).fetchone()["n"]
            return "NeoGuard active_bans=%d recent_attempts=%d threshold=%d/10m" % (active, attempts, MAX_FAILURES)
        if len(args) < 2:
            return "ERROR missing IP address"
        try:
            address = normalize_address(args[1])
        except ValueError:
            return "ERROR invalid IP address"
        if exempt(address):
            return "ERROR protected internal/trusted address"
        address_text = str(address)
        if operation == "show":
            row = conn.execute("SELECT * FROM bans WHERE address=?", (address_text,)).fetchone()
            offense = conn.execute("SELECT count FROM offenses WHERE address=?", (address_text,)).fetchone()
            return (format_ban(row) if row else "No ban for %s" % address_text) + " offenses=%d" % (offense["count"] if offense else 0)
        if operation == "remove":
            conn.execute("UPDATE bans SET active=0 WHERE address=?", (address_text,))
            conn.execute("DELETE FROM attempts WHERE address=?", (address_text,))
            add_event(conn, "ADMIN", "[ADMIN] operation=remove ip=%s" % address_text)
            return "Removed active ban for %s; offense history preserved." % address_text
        if operation == "add":
            expiry = duration_value(args[2] if len(args) > 2 else "24h")
            offense_row = conn.execute("SELECT count FROM offenses WHERE address=?", (address_text,)).fetchone()
            offense = offense_row["count"] if offense_row else 0
            conn.execute("""
                INSERT INTO bans VALUES (?,?,?,?,?,?,?,?,?,1)
                ON CONFLICT(address) DO UPDATE SET started_at=excluded.started_at,expires_at=excluded.expires_at,
                  action=excluded.action,target=excluded.target,message=excluded.message,reason=excluded.reason,active=1
            """, (address_text, now(), expiry, offense, "redirect", DEFAULT_REDIRECT,
                    "Access denied by NeoGuard.", "manual", confirmed_rdns(address)))
            add_event(conn, "ADMIN", "[ADMIN] operation=add ip=%s" % address_text)
            return "Added " + address_text
        row = conn.execute("SELECT * FROM bans WHERE address=? AND active=1", (address_text,)).fetchone()
        if not row:
            return "ERROR no active ban for %s" % address_text
        if operation == "extend":
            if len(args) < 3:
                return "ERROR missing duration"
            expiry = duration_value(args[2])
            conn.execute("UPDATE bans SET expires_at=? WHERE address=?", (expiry, address_text))
            add_event(conn, "ADMIN", "[ADMIN] operation=extend ip=%s duration=%s" % (address_text, args[2]))
            return "Extended " + address_text
        if operation == "action":
            if len(args) < 4 or args[2] not in ("redirect", "message"):
                return "ERROR usage: action IP redirect URL | action IP message TEXT"
            action = args[2]
            value = " ".join(args[3:])[:1000]
            if action == "redirect" and not URL_RE.fullmatch(value):
                return "ERROR redirect must be a safe https:// URL"
            if action == "message" and any(ord(c) < 32 and c not in "\t" for c in value):
                return "ERROR invalid message"
            conn.execute("UPDATE bans SET action=?,target=?,message=? WHERE address=?",
                         (action, value if action == "redirect" else "", value if action == "message" else "", address_text))
            add_event(conn, "ADMIN", "[ADMIN] operation=action ip=%s action=%s" % (address_text, action))
            return "Updated action for " + address_text
    return "ERROR unknown command"


def verify_hmac(handler, body):
    try:
        timestamp = int(handler.headers.get("X-Neo-Timestamp", "0"))
        nonce = handler.headers.get("X-Neo-Nonce", "")
        signature = handler.headers.get("X-Neo-Signature", "")
        if abs(now() - timestamp) > 60 or not re.fullmatch(r"[0-9a-f]{32}", nonce):
            return False
        with DB_LOCK:
            cutoff = now() - 120
            for key in list(NONCES):
                if NONCES[key] < cutoff:
                    del NONCES[key]
            if nonce in NONCES:
                return False
            path = handler.path
            digest = hashlib.sha256(body).hexdigest()
            canonical = "%d\n%s\n%s\n%s\n%s" % (timestamp, nonce, handler.command, path, digest)
            expected = hmac.new(KEY_PATH.read_bytes(), canonical.encode(), hashlib.sha256).hexdigest()
            if not hmac.compare_digest(signature, expected):
                return False
            NONCES[nonce] = now()
            return True
    except (ValueError, OSError):
        return False


class NeoGuardHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "NeoGuard"
    sys_version = ""

    def log_message(self, fmt, *args):
        pass

    def peer(self):
        return normalize_address(self.client_address[0])

    def body(self, maximum=16384):
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            return None
        return self.rfile.read(length) if 0 <= length <= maximum else None

    def response(self, status, body=b"", content_type="application/json", headers=None):
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.send_header("X-Content-Type-Options", "nosniff")
        if headers:
            for key, value in headers.items():
                self.send_header(key, value)
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(body)

    def json_response(self, status, value):
        self.response(status, json.dumps(value, separators=(",", ":")).encode() + b"\n")

    def proxy_client(self):
        if self.peer() != TRUSTED_PROXY:
            return None
        try:
            return normalize_address(self.headers.get("X-Real-IP", ""))
        except ValueError:
            return None

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

    def do_GET(self):
        path = urlsplit(self.path).path
        if path == "/check":
            client = self.proxy_client()
            if not client:
                return self.response(403)
            return self.response(403 if active_ban(client) else 204)
        if path == "/blocked":
            client = self.proxy_client()
            row = active_ban(client) if client else None
            if not row:
                return self.response(404)
            if row["action"] == "redirect":
                return self.response(302, headers={"Location": row["target"]})
            message = html.escape(row["message"] or "Access denied by NeoGuard.")
            page = ("<!doctype html><meta charset=utf-8><meta name=robots content=noindex>"
                    "<title>NeoGuard</title><style>body{margin:0;background:#07030d;color:#f1eaff;font:18px monospace}"
                    "main{max-width:760px;margin:12vh auto;padding:48px;border:1px solid #9d5cff;background:#150a24;"
                    "box-shadow:14px 14px #351360}h1{color:#ff65c8}b{color:#b9ff4f}</style>"
                    "<main><b>NEOGUARD // ACCESS DENIED</b><h1>Connection blocked</h1><p>%s</p></main>" % message).encode()
            return self.response(403, page, "text/html; charset=utf-8")
        if path == "/api/v1/neoguard/poll":
            if not verify_hmac(self, b""):
                return self.response(403)
            try:
                after = int(parse_qs(urlsplit(self.path).query).get("after", ["0"])[0])
            except ValueError:
                after = 0
            with database() as conn:
                rows = list(conn.execute("SELECT * FROM events WHERE id>? ORDER BY id LIMIT 100", (after,)))
            return self.json_response(200, {"events": [dict(row) for row in rows]})
        return self.response(404)

    def do_POST(self):
        path = urlsplit(self.path).path
        body = self.body()
        if body is None:
            return self.response(400)
        if path == "/internal/attempt" and self.peer().is_loopback:
            try:
                data = json.loads(body)
                record_attempt(data["ip"], data.get("source", "drops"), data.get("detail", "invalid request"),
                               data.get("destination_port"))
                return self.response(204)
            except (ValueError, KeyError, json.JSONDecodeError):
                return self.response(400)
        if path == "/api/v1/neoguard/event":
            if not verify_hmac(self, body):
                return self.response(403)
            try:
                data = json.loads(body)
                if data.get("type") == "attempt":
                    record_attempt(data["ip"], data.get("source", "relay"), data.get("detail", "authentication failed"),
                                   data.get("destination_port"))
                    return self.json_response(200, {"ok": True})
                if data.get("type") == "command":
                    return self.json_response(200, {"response": run_command(data.get("command", ""))})
            except (ValueError, KeyError, json.JSONDecodeError):
                pass
            return self.response(400)
        return self.response(404)


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


if __name__ == "__main__":
    os.umask(0o077)
    initialize_database()
    NeoGuardServer((LISTEN_ADDRESS, LISTEN_PORT), NeoGuardHandler).serve_forever()