"""Scanner entrypoint and the scan loop. Import has no side effects; main() runs the service."""
import collections
import logging
import os
import signal
import sys
import threading
import time
from datetime import timedelta

import acceptances
import inventory
import metrics
import page
import report as report_mod
import server
import state as st
import trivy
from state import fmt_ts, parse_ts

log = logging.getLogger("vuln-scanner")

FULL_RUN_INTERVAL = timedelta(hours=20)
DB_MAX_AGE = timedelta(hours=48)
DB_REFRESH_AGE = timedelta(hours=24)
KILLED_REASON = "scanner killed during scan (OOM/eviction?)"


class Engine:
    def __init__(self, data_dir, client, runner, acceptances_path, clock=st.utcnow, tick_seconds=900, publish_seconds=120,
                 private_registries=()):
        self.store = st.Store(data_dir)
        self.data_dir = data_dir
        self.client = client
        self.runner = runner
        self.acceptances_path = acceptances_path
        # Registries known to need credentials the scanner deliberately does not hold (R-VULN-09) and,
        # being on the LAN, not reachable from its egress policy: a scan would end in a connection
        # error (`failed`) rather than a 401 (`credential-required`), so it is never attempted.
        self.private_registries = {name.strip().lower() for name in private_registries if name.strip()}
        self.clock = clock
        self.tick_seconds = tick_seconds
        self.publish_seconds = publish_seconds
        self.lock = threading.RLock()
        self.stop = threading.Event()
        self.state = self.store.load()
        self.report = report_mod.load_published(data_dir)
        self.report_json = report_mod.serialise(self.report) if self.report else None
        self.loop_thread = None
        self._last_publish = 0.0
        self._last_inventory = None
        self._recover_interrupted_scan()

    def _recover_interrupted_scan(self):
        """An inFlight key at startup means the last process died inside that scan."""
        key = self.state.get("inFlight")
        if not key:
            return
        log.error("the previous process was killed while scanning %s; marking it failed", key)
        record = self.state["scans"].get(key)
        if record is not None:
            record["killedAt"] = fmt_ts(self.clock())
        self._record_failure(key, "failed", KILLED_REASON)

    # --- views used by the HTTP servers; none of them waits for a scan --------------------------
    def report_bytes(self):
        return self.report_json

    def render_page(self, query):
        return page.render(self.report, query, self.clock())

    def metrics_text(self):
        with self.lock:
            return metrics.render(self.report, self.state, self.clock())

    def healthy(self):
        return self.loop_thread is None or self.loop_thread.is_alive()

    def record_triage(self, document):
        with self.lock:
            self.state["remediation"] = document["remediation"]
            self.state["triageHeartbeatAt"] = fmt_ts(self.clock())  # server receipt time, per contract
            self.store.save(self.state)
            self.publish()

    # --- persistence -----------------------------------------------------------------------------
    def publish(self):
        """True when the report was written; False when it was refused and the previous one stands."""
        with self.lock:
            now = self.clock()
            entries, errors = acceptances.load(self.acceptances_path, now.date())
            for problem in errors:
                log.error("invalid acceptance ignored: %s", problem)
            built = report_mod.build(self.state, entries, errors, now)
            published = report_mod.publish(built, self.data_dir)
            if published:
                self.report = built
                self.report_json = report_mod.serialise(built)
            self._last_publish = time.monotonic()
            return published

    # --- inventory -------------------------------------------------------------------------------
    def refresh_inventory(self):
        """Install a fresh inventory. A failure keeps the previous one and degrades the report."""
        now = self.clock()
        self._last_inventory = time.monotonic()
        try:
            targets = inventory.build(inventory.collect(self.client))
        except Exception as exc:  # pylint: disable=broad-except
            log.exception("inventory failed; keeping the previous inventory")
            with self.lock:
                self.state["inventoryError"] = type(exc).__name__
                self.store.save(self.state)
            return False
        with self.lock:
            st.merge_inventory(self.state, targets, now)
            self.store.save(self.state)
        return True

    def _units(self):
        """{scan key: reference}. Two targets sharing a digest are one unit, scanned once."""
        units = {}
        for target in sorted(self.state["inventory"], key=lambda t: t["id"]):
            key = st.scan_key(target)
            if key is not None and key not in units:
                units[key] = target["image"] + "@" + target["digest"] if target["digest"] else target["ref"]
        return units

    def _never_attempted(self):
        units = self._units()
        return [k for k in units if self.state["scans"][k]["lastAttemptAt"] is None]

    def _full_run_due(self):
        """Keyed on the last ATTEMPT, not the last success: a permanently failing image would
        otherwise look 'older than 20h' forever and restart a full run every tick."""
        if self.state["run"] is None:
            return True
        now = self.clock()
        for key in self._units():
            attempted = parse_ts(self.state["scans"][key]["lastAttemptAt"])
            if attempted is not None and now - attempted >= FULL_RUN_INTERVAL:
                return True
        return False

    # --- scanning --------------------------------------------------------------------------------
    def _refresh_database(self, always):
        """True when a usable database is present. Failures are logged, never raised."""
        info = self._db_info()
        if always or info is None or self._db_age(info) >= DB_REFRESH_AGE:
            try:
                self.runner.download_db()
            except Exception:  # pylint: disable=broad-except
                log.exception("vulnerability database update failed")
                return False
            info = self._db_info()
        if info is None:
            return False
        with self.lock:
            self.state["database"] = info
            self.store.save(self.state)
        return True

    def _db_info(self):
        try:
            info = self.runner.db_info()
        except Exception:  # pylint: disable=broad-except
            log.exception("could not read vulnerability database information")
            return None
        return dict(info) if info else None

    def _db_age(self, info):
        moment = parse_ts(info.get("updatedAt"))
        return timedelta.max if moment is None else self.clock() - moment

    def _scan_unit(self, key, reference):
        registry = trivy.registry_of(reference)
        if registry in self.private_registries or registry.split(":", 1)[0] in self.private_registries:
            self._record_failure(key, "credential-required",
                                 "registry " + registry + " is private and the scanner holds no credential for it")
            return
        with self.lock:
            self.state["inFlight"] = key
            self.store.save(self.state)
        try:
            found = trivy.parse_results(self.runner.scan(reference))
        except trivy.ScanFailure as exc:
            self._record_failure(key, exc.kind, exc.reason)
        except Exception as exc:  # pylint: disable=broad-except
            log.exception("scan of %s failed unexpectedly", reference)
            self._record_failure(key, "failed", "scan error: " + type(exc).__name__)
        else:
            self._record_success(key, found)
        try:
            self.runner.clean_cache()
        except Exception:  # pylint: disable=broad-except
            log.exception("scan cache cleanup failed")

    def _finish_in_flight(self, key):
        """Cleared in the same write as the scan's outcome, so no crash can separate the two."""
        if self.state.get("inFlight") == key:
            self.state["inFlight"] = None

    def _killed_recently(self, key):
        """A unit that killed the scanner sits out full runs for a day: retrying it first on every
        restart would OOM again before any later image, which is the wedge this exists to prevent."""
        killed = parse_ts(self.state["scans"][key].get("killedAt"))
        return killed is not None and self.clock() - killed < FULL_RUN_INTERVAL

    def _record_failure(self, key, kind, reason):
        now = fmt_ts(self.clock())
        with self.lock:
            self._finish_in_flight(key)
            record = self.state["scans"].get(key)
            if record is None:
                self.store.save(self.state)
                return
            record["lastAttemptAt"] = now
            record["outcome"] = kind
            record["reason"] = reason
            record["failedSince"] = (record["failedSince"] or now) if kind == "failed" else None
            self.store.save(self.state)

    def _record_success(self, key, found):
        now = fmt_ts(self.clock())
        with self.lock:
            self._finish_in_flight(key)
            record = self.state["scans"].get(key)
            if record is None:
                self.store.save(self.state)
                return
            record["lastAttemptAt"] = now
            record["killedAt"] = None
            record["lastSuccessAt"] = now
            record["outcome"] = "ok"
            record["reason"] = None
            record["failedSince"] = None
            record["findings"] = found
            record["firstSeen"] = {f["id"]: record["firstSeen"].get(f["id"], now) for f in found}
            self.store.save(self.state)

    def _scan_queue(self, keys):
        """Scan sequentially, keeping the inventory and report fresh while a long run proceeds."""
        queue = collections.deque(keys)
        while queue and not self.stop.is_set():
            if self._last_inventory is None or time.monotonic() - self._last_inventory >= self.tick_seconds:
                if self.refresh_inventory():
                    queue.extendleft(k for k in self._never_attempted() if k not in queue)
            key = queue.popleft()
            reference = self._units().get(key)
            if reference is None:
                continue  # it left the inventory while queued
            try:
                self._scan_unit(key, reference)
            except Exception:  # pylint: disable=broad-except
                log.exception("scan bookkeeping failed for %s", key)
            if time.monotonic() - self._last_publish >= self.publish_seconds:
                self.publish()

    def run_full(self):
        started = self.clock()
        log.info("full run starting")
        database_ok = self._refresh_database(always=True)
        self._scan_queue([key for key in sorted(self._units()) if not self._killed_recently(key)])
        if self.stop.is_set():
            log.info("full run interrupted by shutdown; not recorded")
            return
        info = self.state["database"]
        now = self.clock()
        degraded = not database_ok or info is None or self._db_age(info) > DB_MAX_AGE
        _, errors = acceptances.load(self.acceptances_path, now.date())
        if errors or self.state.get("inventoryError") or self.state.get("loadError"):
            degraded = True
        status = "degraded" if degraded else "complete"
        with self.lock:
            self.state["run"] = {"startedAt": fmt_ts(started), "finishedAt": fmt_ts(now), "status": status}
            self.store.save(self.state)
            # Only a report that actually reached the store counts as a completed run. Otherwise a
            # refused report would freeze the page while VulnScanStale stayed quiet.
            if self.publish() and not degraded and self.report["run"]["status"] == "complete":
                self.state["lastCompleteRunAt"] = fmt_ts(now)
            # A corrupt state.json degrades the run that rebuilds from it, and no further: once a full
            # run has completed against a good inventory, the state is whole again.
            if self.state.get("loadError") and self._last_inventory is not None and not self.state.get("inventoryError"):
                self.state["loadError"] = None
            self.store.save(self.state)
        log.info("full run finished: %s", status)

    def tick(self):
        self.refresh_inventory()
        self.publish()
        if self._full_run_due():
            self.run_full()
        else:
            never = self._never_attempted()
            if never:
                self._refresh_database(always=False)
                self._scan_queue(never)
        self.publish()

    def run_forever(self):
        while not self.stop.is_set():
            try:
                self.tick()
            except Exception:  # pylint: disable=broad-except
                log.exception("scan cycle crashed; continuing with the next cycle")
            self.stop.wait(self.tick_seconds)

    def start(self):
        self.loop_thread = threading.Thread(target=self.run_forever, name="scan-loop", daemon=True)
        self.loop_thread.start()


def main():
    logging.basicConfig(stream=sys.stdout, level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
    data_dir = os.environ.get("VULN_DATA_DIR", "/data")
    cache_dir = os.environ.get("VULN_CACHE_DIR", "/cache")
    tools_dir = os.environ.get("VULN_TOOLS_DIR", "/tools")
    trivy_bin = os.environ.get("TRIVY_BIN", os.path.join(tools_dir, "trivy"))
    runner = trivy.Runner(trivy_bin, cache_dir, int(os.environ.get("VULN_SCAN_TIMEOUT_SECONDS", "600")))
    private_registries = os.environ.get("VULN_PRIVATE_REGISTRIES", "").split(",")
    engine = Engine(data_dir, inventory.KubeClient(), runner, os.environ.get("VULN_ACCEPTANCES", "/config/acceptances.yaml"),
                    private_registries=private_registries)
    server.start_all(engine, ("127.0.0.1", 8080), ("0.0.0.0", 8081), ("0.0.0.0", 9090))
    finished = threading.Event()
    signal.signal(signal.SIGTERM, lambda *_: finished.set())
    signal.signal(signal.SIGINT, lambda *_: finished.set())
    engine.start()
    finished.wait()
    engine.stop.set()


if __name__ == "__main__":
    main()