#!/usr/bin/python3
import hashlib
import hmac
import json
import os
import secrets
import socket
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

SOCKET_PATH = "{{NEOGUARD_CONTROL_SOCKET_PATH}}"
BASE = "{{DROPS_ORIGIN}}"
KEY_PATH = Path("{{NEOGUARD_KEY_PATH}}")
CLIENTS = set()
CLIENTS_LOCK = threading.Lock()
LAST_EVENT = 0


def signed_request(path, value=None):
    body = b"" if value is None else json.dumps(value, separators=(",", ":")).encode()
    method = "GET" if value is None else "POST"
    for attempt in range(3):
        timestamp = int(time.time())
        nonce = secrets.token_hex(16)
        canonical = "%d\n%s\n%s\n%s\n%s" % (
            timestamp, nonce, method, path, hashlib.sha256(body).hexdigest())
        signature = hmac.new(KEY_PATH.read_bytes(), canonical.encode(), hashlib.sha256).hexdigest()
        request = urllib.request.Request(BASE + path, data=None if value is None else body, method=method, headers={
            "Content-Type": "application/json", "X-Neo-Timestamp": str(timestamp),
            "X-Neo-Nonce": nonce, "X-Neo-Signature": signature,
        })
        try:
            with urllib.request.urlopen(request, timeout=10) as response:
                return json.load(response)
        except urllib.error.HTTPError as error:
            if error.code not in (502, 503, 504) or attempt == 2:
                raise
            time.sleep(0.25 * (attempt + 1))


def broadcast(message):
    payload = (message.replace("\x00", "") + "\n").encode()
    dead = []
    with CLIENTS_LOCK:
        for client in CLIENTS:
            try:
                client.sendall(payload)
            except OSError:
                dead.append(client)
        for client in dead:
            CLIENTS.discard(client)
            client.close()


def attempt_payload(fields):
    if len(fields) == 3:
        address, detail = fields[1], fields[2]
        port = None
    elif len(fields) == 4:
        address, port_text, detail = fields[1], fields[2], fields[3]
        if not port_text.isdecimal() or not 1 <= int(port_text) <= 65535:
            raise ValueError("invalid destination port")
        port = int(port_text)
    else:
        raise ValueError("invalid attempt event")
    payload = {"type": "attempt", "ip": address, "source": "relay", "detail": detail}
    if port is not None:
        payload["destination_port"] = port
    return payload


def handle_client(client):
    buffer = b""
    try:
        client.sendall(b"STATUS\tNeoGuard agent connected\n")
        while True:
            data = client.recv(4096)
            if not data:
                break
            buffer += data
            while b"\n" in buffer:
                line, buffer = buffer.split(b"\n", 1)
                fields = line.decode("utf-8", "replace").split("\t", 3)
                try:
                    if fields[0] == "ATTEMPT":
                        signed_request("/api/v1/neoguard/event", attempt_payload(fields))
                    elif fields[0] == "COMMAND" and len(fields) >= 2:
                        result = signed_request("/api/v1/neoguard/event", {"type":"command","command":fields[1]})
                        client.sendall(("RESPONSE\t" + result.get("response", "No response") + "\n").encode())
                except Exception as error:
                    client.sendall(("ERROR\t" + str(error)[:300] + "\n").encode())
    finally:
        with CLIENTS_LOCK:
            CLIENTS.discard(client)
        client.close()


def socket_server():
    try:
        os.unlink(SOCKET_PATH)
    except FileNotFoundError:
        pass
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(SOCKET_PATH)
    os.chmod(SOCKET_PATH, 0o660)
    server.listen(4)
    while True:
        client, _ = server.accept()
        with CLIENTS_LOCK:
            CLIENTS.add(client)
        threading.Thread(target=handle_client, args=(client,), daemon=True).start()


def poll_events():
    global LAST_EVENT
    while True:
        try:
            result = signed_request("/api/v1/neoguard/poll?after=%d" % LAST_EVENT)
            for event in result.get("events", []):
                LAST_EVENT = max(LAST_EVENT, int(event["id"]))
                broadcast("EVENT\t%s\t%s" % (event["level"], event["message"]))
        except Exception:
            pass
        time.sleep(5)


if __name__ == "__main__":
    threading.Thread(target=poll_events, daemon=True).start()
    socket_server()