#!/usr/bin/env python3
"""Compare running container image digests for the ready Pods of one Kubernetes Service."""
import argparse
from datetime import datetime, timezone
import json
import re
import subprocess
import sys


DIGEST = r"sha256:[0-9a-f]{64}"
IMAGE_ID = re.compile(rf"(?:[a-z][a-z0-9+.-]*://)?(?:[^\s@]+@)?({DIGEST})")
DNS_NAME = re.compile(r"[a-z0-9](?:[a-z0-9.-]*[a-z0-9])?")


def kubectl_reader(context, namespace):
    base = ["kubectl", f"--context={context}", f"--namespace={namespace}", "--request-timeout=10s", "get"]

    def read(*args):
        try:
            result = subprocess.run([*base, *args, "-o", "json"], capture_output=True,
                                    text=True, encoding="utf-8", errors="replace", timeout=20)
        except (OSError, subprocess.TimeoutExpired) as error:
            raise RuntimeError(f"kubectl read failed: {error}") from error
        if result.returncode:
            raise RuntimeError(result.stderr.strip() or f"kubectl exited {result.returncode}")
        try:
            return json.loads(result.stdout)
        except json.JSONDecodeError as error:
            raise RuntimeError("kubectl did not return valid JSON") from error

    return read


def identity(obj):
    metadata = obj.get("metadata", {})
    uid, version = metadata.get("uid"), metadata.get("resourceVersion")
    if not uid or not version:
        raise RuntimeError("Service or EndpointSlice identity is missing")
    return uid, version


def ready_pods(slices, namespace):
    pods = {}
    for item in slices:
        for endpoint in item.get("endpoints", []):
            conditions = endpoint.get("conditions", {})
            if conditions.get("terminating") is True:
                if conditions.get("serving") is not False or conditions.get("ready") is True:
                    raise RuntimeError("a terminating endpoint may still serve; recheck after it drains")
                continue
            if conditions.get("ready") is False:
                continue
            ref = endpoint.get("targetRef") or {}
            name, uid = ref.get("name", ""), ref.get("uid")
            if (ref.get("kind") != "Pod" or ref.get("namespace", namespace) != namespace
                    or not DNS_NAME.fullmatch(name) or not uid):
                raise RuntimeError("an eligible endpoint has no verifiable local Pod reference")
            if pods.setdefault(name, uid) != uid:
                raise RuntimeError(f"{name} has conflicting Pod identities")
    return pods


def pod_digest(pod, uid, name, container):
    metadata = pod.get("metadata", {})
    if metadata.get("uid") != uid or metadata.get("deletionTimestamp"):
        raise RuntimeError(f"{name} changed or started terminating")
    status = pod.get("status", {})
    if not any(c.get("type") == "Ready" and c.get("status") == "True" for c in status.get("conditions", [])):
        raise RuntimeError(f"{name} no longer reports Ready")
    for current in status.get("containerStatuses", []):
        if current.get("name") != container:
            continue
        match = IMAGE_ID.fullmatch(current.get("imageID") or "")
        if current.get("ready") and "running" in current.get("state", {}) and match:
            return match[1]
        break
    raise RuntimeError(f"{name} has no ready running {container!r} container with a SHA-256 image ID")


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    for name in ("context", "namespace", "service", "container"):
        parser.add_argument(f"--{name}", required=True)
    parser.add_argument("--digest", required=True, action="append",
                        help="approved runtime sha256 digest; repeat for additional approved builds")
    parser.add_argument("--min-pods", type=int, default=1, help="minimum number of distinct ready Pods (default: 1)")
    args = parser.parse_args(argv)
    if (not all(DNS_NAME.fullmatch(value) for value in (args.namespace, args.service, args.container))
            or not args.context.strip() or args.min_pods < 1
            or not all(re.fullmatch(DIGEST, value) for value in args.digest)):
        parser.error("provide valid names, lowercase sha256 digests, and --min-pods of at least 1")

    get = kubectl_reader(args.context, args.namespace)
    service_query = ("service", args.service)
    slices_query = ("endpointslices.discovery.k8s.io", "-l", f"kubernetes.io/service-name={args.service}")
    stamp = datetime.now(timezone.utc).isoformat(timespec="seconds")
    print(f"{stamp} context={args.context} namespace={args.namespace} service={args.service} container={args.container}")
    mismatches = 0
    try:
        service = get(*service_query)
        spec = service.get("spec", {})
        if not spec.get("selector") or spec.get("publishNotReadyAddresses"):
            raise RuntimeError("only selector-backed Services using normal readiness are supported")
        slices = get(*slices_query)["items"]
        pods = ready_pods(slices, args.namespace)
        if len(pods) < args.min_pods:
            raise RuntimeError(f"found {len(pods)} distinct eligible Pods; expected at least {args.min_pods}")
        for name, uid in sorted(pods.items()):
            digest = pod_digest(get("pod", name), uid, name, args.container)
            matched = digest in args.digest
            print(f"{'MATCH' if matched else 'MISMATCH'} {name} {digest}")
            mismatches += not matched
        latest_slices = get(*slices_query)["items"]
        if (sorted(map(identity, slices)) != sorted(map(identity, latest_slices))
                or identity(service) != identity(get(*service_query))):
            raise RuntimeError("the Service or its EndpointSlices changed during the check; rerun it")
    except (RuntimeError, KeyError) as error:
        print(f"INCONCLUSIVE: {error}", file=sys.stderr)
        return 2
    if mismatches:
        print(f"FAIL: {mismatches} of {len(pods)} ready Pods reported an unexpected image digest")
        return 1
    print(f"PASS: {len(pods)} ready Pods reported approved image digests")
    return 0


if __name__ == "__main__":
    sys.exit(main())