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

import ctypes
import logging
import multiprocessing
import os
import signal
import subprocess
import sys
import threading
import time

from picosnitch.config import Config
from picosnitch.constants import FD_CACHE
from picosnitch.process_manager import ProcessManager
from picosnitch.subprocesses.fuse import run_fuse
from picosnitch.subprocesses.monitor import run_monitor
from picosnitch.subprocesses.notifications import run_notifications
from picosnitch.subprocesses.primary import run_primary
from picosnitch.subprocesses.remote_sql import run_remote_sql
from picosnitch.subprocesses.secondary import run_secondary
from picosnitch.subprocesses.virustotal import run_virustotal
from picosnitch.types import State
from picosnitch.utils import relaunch_argv


def run_main_loop(config: Config, state: State) -> int:
    """coordinates all picosnitch subprocesses"""
    # init fanotify; CDLL(None) binds fanotify_* from the already-loaded libc, avoiding
    # find_library's gcc/objdump $PATH fallback (like bpf_wrapper._resolve_library)
    libc = ctypes.CDLL(None)
    _FAN_CLOEXEC = 0x1
    _FAN_CLASS_CONTENT = 0x4
    _FAN_UNLIMITED_MARKS = 0x20
    flags = _FAN_CLOEXEC | _FAN_CLASS_CONTENT
    if FD_CACHE >= 8192:
        flags |= _FAN_UNLIMITED_MARKS
    fan_fd = libc.fanotify_init(flags, os.O_RDONLY)
    if fan_fd < 0:
        logging.error("fanotify_init() failed")
        sys.exit(1)
    # start subprocesses
    event_pipes = [multiprocessing.Pipe(duplex=False) for _ in range(5)]
    event_recv_pipes, event_send_pipes = zip(*event_pipes)
    secondary_recv_pipe, secondary_send_pipe = multiprocessing.Pipe(duplex=False)
    q_error: multiprocessing.Queue[str] = multiprocessing.Queue()
    p_notifications = ProcessManager(
        name="snitchnotify",
        target=run_notifications,
        init_args=(
            config,
            fan_fd,
            q_error,
        ),
    )
    p_monitor = ProcessManager(
        name="snitchmonitor",
        target=run_monitor,
        init_args=(
            config,
            fan_fd,
            event_send_pipes,
            q_error,
        ),
    )
    p_fuse = ProcessManager(
        name="snitchrfuse",
        target=run_fuse,
        init_args=(
            config,
            fan_fd,
            q_error,
        ),
    )
    p_virustotal = ProcessManager(
        name="snitchvirustotal",
        target=run_virustotal,
        init_args=(
            config,
            fan_fd,
            q_error,
        ),
    )
    p_remote_sql = ProcessManager(
        name="snitchremotesql",
        target=run_remote_sql,
        init_args=(
            config,
            fan_fd,
            q_error,
        ),
    )
    p_primary = ProcessManager(
        name="snitchprimary",
        target=run_primary,
        init_args=(
            config,
            state,
            event_recv_pipes,
            secondary_send_pipe,
            q_error,
            p_notifications.q_in,
        ),
    )
    p_secondary = ProcessManager(
        name="snitchsecondary",
        target=run_secondary,
        init_args=(
            config,
            state,
            fan_fd,
            p_fuse,
            p_virustotal,
            secondary_recv_pipe,
            p_primary.q_in,
            p_remote_sql.q_in,
            q_error,
        ),
    )
    # set signals
    subprocesses = [p_monitor, p_fuse, p_virustotal, p_remote_sql, p_primary, p_secondary, p_notifications]
    shutdown_event = threading.Event()
    resume_event = threading.Event()

    def _request_shutdown(signum, frame):
        # also wake the resume_event.wait() below so stop/restart tears down promptly
        # instead of after the 5s poll timeout (the loop breaks on shutdown before resuming)
        shutdown_event.set()
        resume_event.set()

    signal.signal(signal.SIGINT, _request_shutdown)
    signal.signal(signal.SIGTERM, _request_shutdown)
    signal.signal(signal.SIGUSR1, lambda signum, frame: resume_event.set())
    # watch subprocesses; detect a real suspend/resume by the gap between CLOCK_BOOTTIME
    # (counts time asleep) and CLOCK_MONOTONIC (does not). a plain wall-clock gap can be
    # forged by cpu starvation or SIGSTOP, which must not trigger a monitor restart -- that
    # would open a bpf re-attach window an attacker could slip traffic through.
    mono_last = time.monotonic()
    boot_last = time.clock_gettime(time.CLOCK_BOOTTIME)
    try:
        while not shutdown_event.is_set():
            resume_event.wait(timeout=5)
            if shutdown_event.is_set():
                break
            if not all(p.is_alive() for p in subprocesses):
                dead = " ".join([p.name for p in subprocesses if not p.is_alive()])
                q_error.put(f"picosnitch subprocess died, attempting restart, terminate by running `picosnitch stop` ({dead})")
                break
            if any(p.is_zombie() for p in subprocesses):
                zombies = " ".join([p.name for p in subprocesses if p.is_zombie()])
                q_error.put(f"picosnitch subprocess became a zombie, attempting restart ({zombies})")
                break
            if sum(p.memory() for p in subprocesses) > 4096000000:
                q_error.put("picosnitch memory usage exceeded 4096 MB, attempting restart")
                break
            mono_now = time.monotonic()
            boot_now = time.clock_gettime(time.CLOCK_BOOTTIME)
            suspended_for = (boot_now - boot_last) - (mono_now - mono_last)
            # skip the resume/suspend restart once shutdown is requested: the shutdown handler
            # also sets resume_event (to break the wait() above), and a spurious monitor re-attach
            # here would only add teardown latency before the clean exit below
            if not shutdown_event.is_set() and (resume_event.is_set() or suspended_for > 5):
                resume_event.clear()
                logging.info("restarting picosnitch monitor after resuming from suspend")
                p_monitor.q_in.put("terminate")
                # let the monitor drain the conn maps and exit before SIGTERM aborts it
                p_monitor.join(timeout=10)
                p_monitor.terminate()
                p_monitor.q_in.get()
                p_monitor.start()
            mono_last, boot_last = mono_now, boot_now
    except Exception as e:
        tb = sys.exc_info()[2]
        q_error.put("picosnitch subprocess exception: %s%s on line %s" % (type(e).__name__, str(e.args), tb.tb_lineno if tb else "?"))
    # clean exit on signal (no restart)
    if shutdown_event.is_set():
        for p in subprocesses:
            p.terminate()
        sys.exit(0)
    if sys.argv[1] == "start-no-daemon":
        return 1
    # attempt to restart picosnitch (terminate by running `picosnitch stop`)
    time.sleep(5)
    for p in subprocesses:
        p.terminate()
    # cwd="/" so the -m fallback doesn't put an inherited cwd on sys.path[0]
    subprocess.Popen(relaunch_argv("restart"), cwd="/")
    return 0
