# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright (C) 2020 Eric Lesiuta
from __future__ import annotations

import contextlib
import hashlib
import importlib
import multiprocessing
import os
import pickle
import queue
import re

from picosnitch.config import Config
from picosnitch.constants import SCHEMA_CONNECTIONS


def remote_dialect(client: str) -> tuple[bool, str]:
    """(is_postgres, placeholder) for a validated remote client name
    - psycopg / psycopg2 -> PostgreSQL; pymysql / mariadb -> MySQL/MariaDB
    - the mariadb connector uses qmark (?) paramstyle, the others use %s
    """
    is_postgres = client in ("psycopg", "psycopg2")
    placeholder = "?" if client == "mariadb" else "%s"
    return is_postgres, placeholder


def executables_key_hash(exe_key: tuple) -> str:
    """dedup key for an executables row: sha256 over (exe, name, cmdline, sha256).
    uniqueness on the raw TEXT columns fails for real cmdlines: MySQL unique
    prefix indexes dedup on the first 191 chars (id lookup then finds no row),
    Postgres btree rejects index rows > 2704 bytes (whole batch lost)"""
    # Length prefixes keep field boundaries unambiguous before hashing.
    return hashlib.sha256("".join(f"{len(f)}\x1f{f}" for f in exe_key).encode("utf-8", "surrogatepass")).hexdigest()


def create_tables(con, conn_table: str, is_postgres: bool) -> None:
    """create the remote reference + connections tables if missing
    - every column is BIGINT: remote INTEGER is 32-bit, but sqlite's is dynamic
      and holds u64 send/recv byte counts and netns inodes well past 2**31
    - executables dedups on key_hash (see executables_key_hash), never the
      unbounded TEXT columns
    """
    if is_postgres:
        pk = "BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY"
        exe_schema = f"id {pk}, exe TEXT NOT NULL, name TEXT NOT NULL, cmdline TEXT NOT NULL, sha256 TEXT NOT NULL, key_hash CHAR(64) NOT NULL UNIQUE"
        dom_schema = f"id {pk}, domain TEXT NOT NULL UNIQUE"
        addr_schema = f"id {pk}, addr TEXT NOT NULL UNIQUE"
    else:
        pk = "BIGINT PRIMARY KEY AUTO_INCREMENT"
        exe_schema = f"id {pk}, exe TEXT NOT NULL, name TEXT NOT NULL, cmdline TEXT NOT NULL, sha256 TEXT NOT NULL, key_hash CHAR(64) NOT NULL UNIQUE"
        dom_schema = f"id {pk}, domain VARCHAR(255) NOT NULL UNIQUE"
        addr_schema = f"id {pk}, addr VARCHAR(64) NOT NULL UNIQUE"
    conn_schema = SCHEMA_CONNECTIONS.replace("INTEGER", "BIGINT")
    cur = con.cursor()
    cur.execute(f"CREATE TABLE IF NOT EXISTS executables ({exe_schema})")
    cur.execute(f"CREATE TABLE IF NOT EXISTS domains ({dom_schema})")
    cur.execute(f"CREATE TABLE IF NOT EXISTS addresses ({addr_schema})")
    cur.execute(f"CREATE TABLE IF NOT EXISTS {conn_table} ({conn_schema})")
    con.commit()
    # CREATE TABLE IF NOT EXISTS won't migrate an existing remote 'executables' table from an
    # older schema; probe for the current schema and fail loudly with the drop-and-recreate remedy
    # instead of spamming opaque errors every batch while silently dropping all remote inserts
    try:
        cur.execute("SELECT key_hash FROM executables LIMIT 0")
        cur.fetchall()
    except Exception as e:
        raise RuntimeError(
            "remote 'executables' table has an outdated schema; drop the old remote tables (executables, domains, addresses, and the connections table) so picosnitch can recreate them"
        ) from e


def insert_entries(con, conn_table: str, entries: list[tuple], is_postgres: bool, ph: str) -> None:
    """resolve normalized ids and insert each connection entry (INSERT only)
    - ph is the driver's placeholder (%s or ?); dedup via INSERT IGNORE (MySQL)
      or INSERT ... ON CONFLICT DO NOTHING (Postgres)
    """
    ignore = " ON CONFLICT DO NOTHING" if is_postgres else ""
    insert = "INSERT INTO" if is_postgres else "INSERT IGNORE INTO"
    sql_insert_exe = f"{insert} executables(exe, name, cmdline, sha256, key_hash) VALUES ({ph}, {ph}, {ph}, {ph}, {ph}){ignore}"
    sql_select_exe = f"SELECT id FROM executables WHERE key_hash = {ph}"
    sql_insert_dom = f"{insert} domains(domain) VALUES ({ph}){ignore}"
    sql_select_dom = f"SELECT id FROM domains WHERE domain = {ph}"
    sql_insert_addr = f"{insert} addresses(addr) VALUES ({ph}){ignore}"
    sql_select_addr = f"SELECT id FROM addresses WHERE addr = {ph}"
    sql_insert_conn = f"INSERT INTO {conn_table} VALUES ({', '.join([ph] * 16)})"
    with con.cursor() as cur:
        conn_rows = []
        for contime, send, recv, n_events, exe_key, pexe_key, gpexe_key, uid, family, protocol, lport, rport, laddr, raddr, domain, netns in entries:
            exe_ids = []
            for key in (exe_key, pexe_key, gpexe_key):
                key_hash = executables_key_hash(key)
                cur.execute(sql_insert_exe, (*key, key_hash))
                cur.execute(sql_select_exe, (key_hash,))
                exe_ids.append(cur.fetchone()[0])
            exe_id, pexe_id, gpexe_id = exe_ids
            cur.execute(sql_insert_addr, (laddr,))
            cur.execute(sql_select_addr, (laddr,))
            laddr_id = cur.fetchone()[0]
            cur.execute(sql_insert_addr, (raddr,))
            cur.execute(sql_select_addr, (raddr,))
            raddr_id = cur.fetchone()[0]
            cur.execute(sql_insert_dom, (domain,))
            cur.execute(sql_select_dom, (domain,))
            domain_id = cur.fetchone()[0]
            conn_rows.append((contime, send, recv, n_events, exe_id, pexe_id, gpexe_id, uid, family, protocol, lport, rport, laddr_id, raddr_id, domain_id, netns))
        cur.executemany(sql_insert_conn, conn_rows)
    con.commit()


def run_remote_sql(config: Config, fan_fd: int, q_error: multiprocessing.Queue[str], q_in: multiprocessing.Queue[bytes], _q_out: multiprocessing.Queue) -> int:
    """writes connection log entries from secondary to the remote SQL server ([database.remote]),
    runs without root so third-party database drivers are never imported by a root process"""
    parent_process = multiprocessing.parent_process()
    assert parent_process is not None
    # drop root (desktop.user, or nobody) before importing or running any third-party code
    from ..utils import drop_root_permanent, resolve_unprivileged_user

    # close the privileged fanotify handle before dropping: never used here, must not persist into
    # the dropped-privilege domain
    try:
        os.close(fan_fd)
    except OSError:
        pass
    uid, gid = resolve_unprivileged_user(config.desktop.user)
    drop_root_permanent(uid, gid)
    # config was validated loudly by check_remote_config at start; revalidate before use
    # (defense in depth) but only report and idle, never crash-loop the daemon
    sql = None
    sql_kwargs = dict(config.database.remote)
    sql_client = sql_kwargs.pop("client", "no client error")
    conn_table = sql_kwargs.pop("connections_table", "connections")
    is_postgres, ph = False, "%s"
    if config.database.remote:
        if sql_client not in ["mariadb", "psycopg", "psycopg2", "pymysql"]:
            q_error.put(f'unsupported database.remote "client": {sql_client}')
        elif not isinstance(conn_table, str) or not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", conn_table):
            q_error.put(f"invalid remote table name: {conn_table!r}")
        else:
            is_postgres, ph = remote_dialect(sql_client)
            try:
                sql = importlib.import_module(sql_client)
            except ImportError:
                q_error.put(f'database.remote "client" {sql_client} is not installed, remote logging disabled')
        if sql is not None and not any(k in sql_kwargs for k in ("ssl", "ssl_context", "sslmode", "ssl_mode")):
            q_error.put("warning: remote database connection has no SSL/TLS parameters configured")
    tables_ready = False
    while True:
        if not parent_process.is_alive():
            return 0
        entries: list[tuple] = []
        try:
            # timeout so the parent-alive check above runs even when idle
            entries = pickle.loads(q_in.get(block=True, timeout=15))
            if sql is None:
                continue
            # closing() so a create_tables/insert_entries error doesn't leak the driver socket
            with contextlib.closing(sql.connect(**sql_kwargs)) as con:
                if not tables_ready:
                    create_tables(con, conn_table, is_postgres)
                    tables_ready = True
                insert_entries(con, conn_table, entries, is_postgres, ph)
        except queue.Empty:
            pass
        except Exception as e:
            q_error.put("SQL server execute %s%s on line %s, lost %s entries" % (type(e).__name__, str(e.args), e.__traceback__.tb_lineno if e.__traceback__ else "?", len(entries)))
