#!/usr/bin/env python3
"""R-VULN-01: inventory from recorded API lists."""
import copy
import http.server
import json
import os
import sys
import tempfile
import threading
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parents[3] / "platform/vuln-scanner/src"))
import inventory  # noqa: E402

FAILED = []


def check(name, condition):
    print(("ok   " if condition else "FAIL ") + name)
    if not condition:
        FAILED.append(name)


CLUSTER = json.loads((HERE / "fixtures/cluster.json").read_text())
A1 = "sha256:" + "a1" * 32
B2 = "sha256:" + "b2" * 32
C3 = "sha256:" + "c3" * 32
D4 = "sha256:" + "d4" * 32
E5 = "sha256:" + "e5" * 32
F6 = "sha256:" + "f6" * 32


def by_id(targets):
    return {t["id"]: t for t in targets}


def workloads(target):
    return {(w["namespace"], w["kind"], w["name"]): w["containers"] for w in target["workloads"]}


targets = inventory.build(CLUSTER)
ids = by_id(targets)

nginx = ids["docker.io/library/nginx@" + A1]
check("ReplicaSet resolves to Deployment, not the ReplicaSet", workloads(nginx) == {
    ("apps", "Deployment", "web"): ["app"], ("apps", "Deployment", "api"): ["server"]})
check("docker-pullable:// prefix stripped and digest taken from imageID", nginx["digest"] == A1)
check("init container inventoried", ("apps", "Deployment", "web") in workloads(ids["ghcr.io/org/tool@" + B2]))
check("ephemeral container inventoried", "debug" in workloads(ids["docker.io/library/busybox@" + C3])[("apps", "Deployment", "web")])

backup = ids["ghcr.io/org/backup@" + E5]
check("Job resolves to CronJob", workloads(backup) == {("ops", "CronJob", "backup"): ["backup"]})
check("a CronJob with a running pod is not also inventoried from its template",
      not any(t["id"].startswith("ghcr.io/org/backup:") for t in targets))

nightly = ids["ghcr.io/org/nightly:2.1"]
check("idle CronJob inventoried from template with no digest", nightly["digest"] is None
      and workloads(nightly) == {("ops", "CronJob", "nightly"): ["job"]})

solo = ids["docker.io/library/alpine:3.20"]
check("empty imageID gives a by-tag target from the spec image", solo["digest"] is None)
check("bare Pod is its own workload", ("ops", "Pod", "solo") in workloads(solo))
check("Succeeded pod's image is not inventoried", not any(F6 in t["id"] for t in targets))
check("virt-launcher stays owned by its VirtualMachineInstance",
      ("kubevirt-vms", "VirtualMachineInstance", "vm1") in workloads(ids["quay.io/kubevirt/virt-launcher@sha256:" + "ab" * 32]))
vm, node = ids["vm/kubevirt-vms/vm1"], ids["node/prod-1"]
check("VM inventoried as non-container target", vm["type"] == "vm" and vm["digest"] is None and vm["reason"])
check("node inventoried as non-container target", node["type"] == "node" and node["reason"])

# Non-vacuous: mutating the fixture changes the answer the assertions above rely on.
broken = copy.deepcopy(CLUSTER)
broken["replicasets"] = []
broken["jobs"] = []
broken_ids = by_id(inventory.build(broken))
check("mutation: losing the ReplicaSet record leaves the ReplicaSet as the workload",
      ("apps", "ReplicaSet", "web-5f") in workloads(broken_ids["docker.io/library/nginx@" + A1]))
check("mutation: losing the Job record leaves the Job as the workload",
      ("ops", "Job", "backup-123") in workloads(broken_ids["ghcr.io/org/backup@" + E5]))
check("mutation: a Running phase on the Succeeded pod does inventory its image", any(
    F6 in t["id"] for t in inventory.build({**CLUSTER, "pods": [
        {**p, "status": {**p["status"], "phase": "Running"}} if p["metadata"]["name"] == "done-1" else p
        for p in CLUSTER["pods"]]})))

for reference, expected in (
    ("nginx:1.27", ("docker.io/library/nginx", "1.27", None)),
    ("nginx", ("docker.io/library/nginx", None, None)),
    ("bitnami/redis:7", ("docker.io/bitnami/redis", "7", None)),
    ("ghcr.io/org/app:v1", ("ghcr.io/org/app", "v1", None)),
    ("registry.local:5000/x/y:2", ("registry.local:5000/x/y", "2", None)),
    ("index.docker.io/library/nginx:1", ("docker.io/library/nginx", "1", None)),
    ("ghcr.io/org/app:v1@" + A1, ("ghcr.io/org/app", "v1", A1)),
):
    check("normalise " + reference, inventory.normalise_reference(reference) == expected)
check("bare config id is not a digest", inventory.digest_from_image_id("sha256:" + "a1" * 32) is None)
check("empty imageID has no digest", inventory.digest_from_image_id("") is None)


class FakeListClient:
    """list_all replaced; proves collect() tolerates exactly a 404 on the VMI API."""

    def __init__(self, vmi_error):
        self.vmi_error = vmi_error

    def list_all(self, path):
        if "kubevirt" in path and self.vmi_error:
            raise self.vmi_error
        return []


check("VMI API 404 tolerated as no VMs", inventory.collect(FakeListClient(inventory.NotFound("x")))["vmis"] == [])
try:
    inventory.collect(FakeListClient(inventory.InventoryError("403")))
    check("a non-404 VMI error is not swallowed", False)
except inventory.InventoryError:
    check("a non-404 VMI error is not swallowed", True)


class PagedAPI(http.server.BaseHTTPRequestHandler):
    def log_message(self, *args):
        return

    def do_GET(self):
        if self.headers.get("Authorization") != "Bearer fixture-token":
            self.send_response(403)
            self.end_headers()
            return
        if self.path.startswith("/apis/kubevirt"):
            self.send_response(404)
            self.end_headers()
            return
        if "continue=" in self.path:
            body = {"items": [{"metadata": {"name": "b"}}], "metadata": {}}
        else:
            body = {"items": [{"metadata": {"name": "a"}}], "metadata": {"continue": "page2"}}
        payload = json.dumps(body).encode()
        self.send_response(200)
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)


with tempfile.TemporaryDirectory() as tmp:
    token = os.path.join(tmp, "token")
    Path(token).write_text("fixture-token\n")
    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), PagedAPI)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    client = inventory.KubeClient("http://127.0.0.1:" + str(server.server_address[1]), token, os.devnull)
    items = client.list_all("/api/v1/pods")
    check("pagination follows the continue token with the bearer token", [i["metadata"]["name"] for i in items] == ["a", "b"])
    check("collect over real HTTP tolerates VMI 404", inventory.collect(client)["vmis"] == [])
    server.shutdown()

print("FAILED: " + str(len(FAILED)) if FAILED else "all checks passed")
sys.exit(1 if FAILED else 0)