#!/usr/bin/env python3
"""Local, fail-closed bridge between Message Safety and host KESL."""

from __future__ import annotations

import asyncio
import hashlib
import json
import os
import re
import socket
import struct
import sys
import tempfile
from pathlib import Path
from typing import Any

MAX_HEADER_BYTES = 4096
MAX_FILE_BYTES = 5 * 1024 * 1024
MAX_RESPONSE_BYTES = 16 * 1024
DEFAULT_TIMEOUT_SECONDS = 60.0
KESL_CONTROL = Path("/opt/kaspersky/kesl/bin/kesl-control")
STAGING_DIR = Path("/var/lib/han-kesl-scan/staging")

_SUMMARY_PATTERNS = {
    "scanned": re.compile(
        r"(?im)^\s*(?:Scanned objects|Objects scanned|Проверенные объекты)"
        r"\s*:\s*(\d+)\s*$"
    ),
    "detected": re.compile(
        r"(?im)^\s*(?:Total detected objects|Всего обнаружено объектов)"
        r"\s*:\s*(\d+)\s*$"
    ),
    "errors": re.compile(
        r"(?im)^\s*(?:Scan errors|Ошибки проверки)\s*:\s*(\d+)\s*$"
    ),
    "skipped": re.compile(
        r"(?im)^\s*(?:Skipped objects|Objects skipped|Пропущено объектов)"
        r"\s*:\s*(\d+)\s*$"
    ),
}
_THREAT_RE = re.compile(r"(?im)^\s*(?:Threat|Detect name)\s*:\s*(.{1,256})\s*$")


class ProtocolError(ValueError):
    pass


class KeslError(RuntimeError):
    pass


def _json_bytes(payload: dict[str, Any]) -> bytes:
    body = json.dumps(payload, ensure_ascii=True, separators=(",", ":")).encode()
    if len(body) > MAX_RESPONSE_BYTES:
        raise KeslError("response exceeds protocol limit")
    return struct.pack(">I", len(body)) + body


async def _read_frame(reader: asyncio.StreamReader) -> tuple[dict[str, Any], bytes]:
    (header_size,) = struct.unpack(">I", await reader.readexactly(4))
    if header_size < 2 or header_size > MAX_HEADER_BYTES:
        raise ProtocolError("invalid header size")
    try:
        header = json.loads((await reader.readexactly(header_size)).decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise ProtocolError("invalid header") from exc
    if not isinstance(header, dict) or set(header) - {"op", "size"}:
        raise ProtocolError("invalid header fields")
    operation = header.get("op")
    if operation == "status":
        if "size" in header:
            raise ProtocolError("status request cannot contain a body")
        return header, b""
    if operation != "scan" or not isinstance(header.get("size"), int):
        raise ProtocolError("unsupported operation")
    size = header["size"]
    if isinstance(size, bool) or size < 1 or size > MAX_FILE_BYTES:
        raise ProtocolError("invalid file size")
    return header, await reader.readexactly(size)


def parse_scan_output(output: str, *, exit_code: int = 0) -> tuple[str, str | None]:
    """Parse the documented KESL Scan_File summary, rejecting format drift."""
    if exit_code not in (0, 72):
        raise KeslError(f"unexpected KESL scan exit code: {exit_code}")
    summary: dict[str, int] = {}
    for name, pattern in _SUMMARY_PATTERNS.items():
        match = pattern.search(output)
        if not match:
            raise KeslError(f"KESL scan summary field is missing: {name}")
        summary[name] = int(match.group(1))
    if summary["scanned"] < 1:
        raise KeslError("KESL did not scan the submitted object")
    if summary["errors"] or summary["skipped"]:
        raise KeslError("KESL scan completed with errors or skipped objects")
    detected = summary["detected"]
    if (exit_code == 72) != (detected > 0):
        raise KeslError("KESL exit code contradicts scan summary")
    if detected == 0:
        return "clean", None
    threat = _THREAT_RE.search(output)
    return "infected", threat.group(1).strip() if threat else "detected"


def _flatten_json(value: Any, prefix: str = "") -> dict[str, Any]:
    result: dict[str, Any] = {}
    if isinstance(value, dict):
        for key, item in value.items():
            normalized = re.sub(r"[^\w]+", "_", str(key).casefold()).strip("_")
            result.update(_flatten_json(item, f"{prefix}_{normalized}".strip("_")))
    else:
        result[prefix] = value
    return result


def parse_app_info(output: str) -> dict[str, Any]:
    try:
        flattened = _flatten_json(json.loads(output))
    except json.JSONDecodeError as exc:
        raise KeslError("invalid KESL app-info JSON") from exc

    def find(*suffixes: str) -> Any:
        for suffix in suffixes:
            for key, value in flattened.items():
                if key == suffix or key.endswith(f"_{suffix}"):
                    return value
        raise KeslError(f"KESL app-info field is missing: {suffixes[0]}")

    version = str(find("version", "application_version", "версия")).strip()
    databases_loaded = find(
        "databases_loaded",
        "application_databases_loaded",
        "базы_приложения_загружены",
    )
    databases_date = str(
        find(
            "databases_date",
            "last_release_date_of_databases",
            "database_date",
            "дата_последнего_выпуска_баз_приложения",
        )
    ).strip()
    license_info = str(
        find(
            "license_info",
            "license_status",
            "key_status",
            "license_information",
            "информация_о_лицензии_приложения",
        )
    ).strip()
    if databases_loaded not in (True, "Yes", "yes", "true", "Да", "да", 1):
        raise KeslError("KESL databases are not loaded")
    if not version or not databases_date or databases_date.upper() == "N/A":
        raise KeslError("KESL version or database date is unavailable")
    if "valid" not in license_info.lower() and "действ" not in license_info.lower():
        raise KeslError("KESL license is not valid")
    signature_source = f"{version}\0{databases_date}".encode()
    return {
        "status": "ready",
        "engine_version": version,
        "databases_date": databases_date,
        "signatures_version": "sha256:" + hashlib.sha256(signature_source).hexdigest(),
    }


async def _run_kesl(
    *arguments: str,
    timeout: float,
    allowed_returncodes: frozenset[int] = frozenset({0}),
) -> tuple[str, int]:
    if not KESL_CONTROL.is_file():
        raise KeslError("KESL control utility is unavailable")
    process = await asyncio.create_subprocess_exec(
        str(KESL_CONTROL),
        *arguments,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env={**os.environ, "LC_ALL": "C", "LANG": "C", "LANGUAGE": "C"},
    )
    try:
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout)
    except TimeoutError:
        process.kill()
        await process.wait()
        raise KeslError("KESL command timed out") from None
    if process.returncode not in allowed_returncodes:
        detail = stderr.decode("utf-8", "replace").strip()[:256]
        raise KeslError(f"KESL command failed: {detail or process.returncode}")
    return stdout.decode("utf-8", "replace"), process.returncode


class Broker:
    def __init__(self, *, timeout: float, concurrency: int) -> None:
        self.timeout = timeout
        self.slots = asyncio.Semaphore(concurrency)

    async def status(self) -> dict[str, Any]:
        output, _ = await _run_kesl(
            "--app-info", "--json", timeout=min(self.timeout, 10.0)
        )
        return parse_app_info(output)

    async def scan(self, body: bytes) -> dict[str, Any]:
        async with self.slots:
            status_before = await self.status()
            STAGING_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
            descriptor, name = tempfile.mkstemp(prefix="scan-", dir=STAGING_DIR)
            path = Path(name)
            try:
                os.fchmod(descriptor, 0o600)
                with os.fdopen(descriptor, "wb", closefd=True) as target:
                    target.write(body)
                    target.flush()
                    os.fsync(target.fileno())
                output, exit_code = await _run_kesl(
                    "--scan-file",
                    str(path),
                    "--action",
                    "Inform",
                    timeout=self.timeout,
                    allowed_returncodes=frozenset({0, 72}),
                )
                verdict, threat = parse_scan_output(output, exit_code=exit_code)
                status_after = await self.status()
                if (
                    status_before["signatures_version"]
                    != status_after["signatures_version"]
                ):
                    raise KeslError("KESL databases changed during scan")
                return {
                    "status": "completed",
                    "verdict": verdict,
                    "threat": threat,
                    "engine_version": status_after["engine_version"],
                    "signatures_version": status_after["signatures_version"],
                }
            finally:
                path.unlink(missing_ok=True)

    async def handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            header, body = await asyncio.wait_for(_read_frame(reader), 10.0)
            result = await (self.status() if header["op"] == "status" else self.scan(body))
        except (ProtocolError, asyncio.IncompleteReadError, TimeoutError):
            result = {"status": "error", "error": "invalid_request"}
        except KeslError as exc:
            print(f"KESL scanner unavailable: {exc}", file=sys.stderr, flush=True)
            result = {"status": "error", "error": "scanner_unavailable"}
        except Exception:
            result = {"status": "error", "error": "internal_error"}
        writer.write(_json_bytes(result))
        await writer.drain()
        writer.close()
        await writer.wait_closed()


def _activation_socket() -> socket.socket:
    if int(os.environ.get("LISTEN_FDS", "0")) != 1 or os.getpid() != int(
        os.environ.get("LISTEN_PID", "0")
    ):
        raise SystemExit("exactly one systemd activation socket is required")
    descriptor = socket.fromfd(3, socket.AF_UNIX, socket.SOCK_STREAM)
    descriptor.setblocking(False)
    return descriptor


async def serve() -> None:
    timeout = float(os.environ.get("HAN_KESL_SCAN_TIMEOUT_SEC", DEFAULT_TIMEOUT_SECONDS))
    concurrency = int(os.environ.get("HAN_KESL_SCAN_CONCURRENCY", "5"))
    if not 1 <= concurrency <= 16 or not 1 <= timeout <= 300:
        raise SystemExit("invalid broker limits")
    server = await asyncio.start_unix_server(
        Broker(timeout=timeout, concurrency=concurrency).handle,
        sock=_activation_socket(),
        limit=MAX_FILE_BYTES + MAX_HEADER_BYTES + 4,
    )
    async with server:
        await server.serve_forever()


async def probe() -> None:
    reader, writer = await asyncio.wait_for(
        asyncio.open_unix_connection("/run/han-kesl/scan.sock"), 5.0
    )
    try:
        header = json.dumps({"op": "status"}, separators=(",", ":")).encode()
        writer.write(struct.pack(">I", len(header)) + header)
        await writer.drain()
        (size,) = struct.unpack(">I", await asyncio.wait_for(reader.readexactly(4), 10.0))
        if size < 2 or size > MAX_RESPONSE_BYTES:
            raise SystemExit("invalid broker response")
        response = json.loads((await asyncio.wait_for(reader.readexactly(size), 10.0)).decode())
        if response.get("status") != "ready":
            print(json.dumps(response, ensure_ascii=True, sort_keys=True), file=sys.stderr)
            raise SystemExit("KESL broker is not ready")
        print(json.dumps(response, ensure_ascii=True, sort_keys=True))
    finally:
        writer.close()
        await writer.wait_closed()


if __name__ == "__main__":
    if not sys.argv[1:]:
        asyncio.run(serve())
    elif sys.argv[1:] == ["--probe"]:
        asyncio.run(probe())
    else:
        raise SystemExit("usage: han-kesl-scan-broker [--probe]")