#!/usr/bin/env python3
"""Kubernetes Pod Security Admission audit agent.

Audits Kubernetes namespaces and workloads for Pod Security Standards (PSS)
compliance using kubectl. Checks namespace labels for enforce/audit/warn
modes, analyzes pod specs against Baseline and Restricted profiles, and
reports violations.
"""
import argparse
import json
import os
import subprocess
import sys
from datetime import datetime, timezone


def run_kubectl(args_list, timeout=60):
    """Execute a kubectl command and return parsed JSON output."""
    cmd = ["kubectl"] + args_list + ["-o", "json"]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    if result.returncode != 0:
        print(f"[!] kubectl error: {result.stderr[:200]}", file=sys.stderr)
        return None
    try:
        return json.loads(result.stdout)
    except json.JSONDecodeError:
        return None


def get_namespaces():
    """Get all namespaces with their PSA labels."""
    print("[*] Fetching namespaces...")
    data = run_kubectl(["get", "namespaces"])
    if not data:
        return []
    namespaces = []
    for ns in data.get("items", []):
        name = ns.get("metadata", {}).get("name", "")
        labels = ns.get("metadata", {}).get("labels", {})
        psa_enforce = labels.get("pod-security.kubernetes.io/enforce", "")
        psa_audit = labels.get("pod-security.kubernetes.io/audit", "")
        psa_warn = labels.get("pod-security.kubernetes.io/warn", "")
        psa_enforce_version = labels.get("pod-security.kubernetes.io/enforce-version", "")
        namespaces.append({
            "name": name,
            "enforce": psa_enforce,
            "audit": psa_audit,
            "warn": psa_warn,
            "enforce_version": psa_enforce_version,
            "has_psa": bool(psa_enforce or psa_audit or psa_warn),
        })
    print(f"[+] Found {len(namespaces)} namespaces")
    return namespaces


def audit_namespace_psa(namespaces):
    """Audit PSA label configuration across namespaces."""
    findings = []
    system_ns = {"kube-system", "kube-public", "kube-node-lease", "default"}

    for ns in namespaces:
        name = ns["name"]
        if name in system_ns:
            continue

        if not ns["has_psa"]:
            findings.append({
                "namespace": name,
                "check": "No PSA labels configured",
                "severity": "HIGH",
                "recommendation": (
                    f"Apply PSA labels: kubectl label namespace {name} "
                    f"pod-security.kubernetes.io/enforce=baseline "
                    f"pod-security.kubernetes.io/warn=restricted"
                ),
            })
        elif ns["enforce"] == "privileged":
            findings.append({
                "namespace": name,
                "check": "PSA enforce set to privileged (permissive)",
                "severity": "HIGH",
                "recommendation": "Upgrade to baseline or restricted enforcement",
            })
        elif ns["enforce"] == "baseline" and not ns["warn"]:
            findings.append({
                "namespace": name,
                "check": "Enforce baseline but no warn=restricted",
                "severity": "MEDIUM",
                "recommendation": (
                    f"Add warn label: kubectl label namespace {name} "
                    f"pod-security.kubernetes.io/warn=restricted"
                ),
            })
        elif ns["enforce"] == "restricted":
            findings.append({
                "namespace": name,
                "check": "PSA enforce=restricted (most secure)",
                "severity": "INFO",
            })

    return findings


def audit_pod_security(namespace=None):
    """Audit pods for security spec violations against PSS profiles."""
    findings = []
    cmd = ["get", "pods", "--all-namespaces"] if not namespace else ["get", "pods", "-n", namespace]
    data = run_kubectl(cmd)
    if not data:
        return findings

    for pod in data.get("items", []):
        pod_name = pod.get("metadata", {}).get("name", "")
        pod_ns = pod.get("metadata", {}).get("namespace", "")
        spec = pod.get("spec", {})

        # Check host namespaces
        if spec.get("hostNetwork"):
            findings.append({
                "namespace": pod_ns, "pod": pod_name,
                "check": "hostNetwork enabled",
                "severity": "CRITICAL", "profile": "baseline",
            })
        if spec.get("hostPID"):
            findings.append({
                "namespace": pod_ns, "pod": pod_name,
                "check": "hostPID enabled",
                "severity": "CRITICAL", "profile": "baseline",
            })
        if spec.get("hostIPC"):
            findings.append({
                "namespace": pod_ns, "pod": pod_name,
                "check": "hostIPC enabled",
                "severity": "HIGH", "profile": "baseline",
            })

        containers = spec.get("containers", []) + spec.get("initContainers", [])
        for container in containers:
            c_name = container.get("name", "")
            sec_ctx = container.get("securityContext", {})

            # Privileged container
            if sec_ctx.get("privileged"):
                findings.append({
                    "namespace": pod_ns, "pod": pod_name, "container": c_name,
                    "check": "Privileged container",
                    "severity": "CRITICAL", "profile": "baseline",
                })

            # Run as root
            if sec_ctx.get("runAsUser") == 0:
                findings.append({
                    "namespace": pod_ns, "pod": pod_name, "container": c_name,
                    "check": "Running as root (UID 0)",
                    "severity": "HIGH", "profile": "restricted",
                })

            # Missing runAsNonRoot
            if not sec_ctx.get("runAsNonRoot"):
                findings.append({
                    "namespace": pod_ns, "pod": pod_name, "container": c_name,
                    "check": "runAsNonRoot not set",
                    "severity": "MEDIUM", "profile": "restricted",
                })

            # Writable root filesystem
            if not sec_ctx.get("readOnlyRootFilesystem"):
                findings.append({
                    "namespace": pod_ns, "pod": pod_name, "container": c_name,
                    "check": "Root filesystem not read-only",
                    "severity": "MEDIUM", "profile": "restricted",
                })

            # Privilege escalation
            if sec_ctx.get("allowPrivilegeEscalation", True):
                findings.append({
                    "namespace": pod_ns, "pod": pod_name, "container": c_name,
                    "check": "allowPrivilegeEscalation not disabled",
                    "severity": "MEDIUM", "profile": "restricted",
                })

            # Dangerous capabilities
            caps = sec_ctx.get("capabilities", {})
            added = caps.get("add", [])
            dangerous_caps = {"SYS_ADMIN", "NET_ADMIN", "SYS_PTRACE", "ALL",
                              "NET_RAW", "SYS_RAWIO", "SYS_MODULE"}
            for cap in added:
                if cap.upper() in dangerous_caps:
                    findings.append({
                        "namespace": pod_ns, "pod": pod_name, "container": c_name,
                        "check": f"Dangerous capability added: {cap}",
                        "severity": "CRITICAL" if cap.upper() in ("SYS_ADMIN", "ALL") else "HIGH",
                        "profile": "baseline",
                    })

            # Drop ALL capabilities check (restricted)
            dropped = [c.upper() for c in caps.get("drop", [])]
            if "ALL" not in dropped:
                findings.append({
                    "namespace": pod_ns, "pod": pod_name, "container": c_name,
                    "check": "Capabilities not dropping ALL",
                    "severity": "LOW", "profile": "restricted",
                })

            # Host ports
            for port in container.get("ports", []):
                if port.get("hostPort"):
                    findings.append({
                        "namespace": pod_ns, "pod": pod_name, "container": c_name,
                        "check": f"hostPort exposed: {port['hostPort']}",
                        "severity": "HIGH", "profile": "baseline",
                    })

    return findings


def format_summary(ns_findings, pod_findings, namespaces):
    """Print audit summary."""
    print(f"\n{'='*60}")
    print(f"  Pod Security Admission Audit Report")
    print(f"{'='*60}")
    print(f"  Namespaces      : {len(namespaces)}")
    psa_configured = sum(1 for ns in namespaces if ns["has_psa"])
    print(f"  PSA Configured  : {psa_configured}/{len(namespaces)}")
    print(f"  NS Findings     : {len(ns_findings)}")
    print(f"  Pod Findings    : {len(pod_findings)}")

    all_findings = ns_findings + pod_findings
    severity_counts = {}
    for f in all_findings:
        sev = f.get("severity", "INFO")
        severity_counts[sev] = severity_counts.get(sev, 0) + 1

    print(f"\n  By Severity:")
    for sev in ["CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"]:
        count = severity_counts.get(sev, 0)
        if count > 0:
            print(f"    {sev:10s}: {count}")

    print(f"\n  Namespace PSA Status:")
    for ns in namespaces:
        if ns["name"] in ("kube-system", "kube-public", "kube-node-lease"):
            continue
        status = ns["enforce"] or "none"
        print(f"    {ns['name']:30s}: enforce={status:12s} audit={ns['audit'] or 'none':12s}")

    if pod_findings:
        print(f"\n  Top Pod Violations:")
        for f in pod_findings[:15]:
            if f["severity"] in ("CRITICAL", "HIGH"):
                print(f"    [{f['severity']:8s}] {f.get('namespace', '')}/"
                      f"{f.get('pod', '')}: {f['check']}")

    return severity_counts


def main():
    parser = argparse.ArgumentParser(
        description="Kubernetes Pod Security Admission audit agent"
    )
    parser.add_argument("--namespace", "-n", help="Specific namespace to audit (default: all)")
    parser.add_argument("--skip-pods", action="store_true", help="Only audit namespace labels")
    parser.add_argument("--kubeconfig", help="Path to kubeconfig file")
    parser.add_argument("--context", help="Kubernetes context to use")
    parser.add_argument("--output", "-o", help="Output JSON report path")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    if args.kubeconfig:
        os.environ["KUBECONFIG"] = args.kubeconfig

    namespaces = get_namespaces()
    ns_findings = audit_namespace_psa(namespaces)

    pod_findings = []
    if not args.skip_pods:
        pod_findings = audit_pod_security(args.namespace)

    severity_counts = format_summary(ns_findings, pod_findings, namespaces)

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "tool": "PSA Audit",
        "namespaces": namespaces,
        "namespace_findings": ns_findings,
        "pod_findings": pod_findings,
        "severity_counts": severity_counts,
        "risk_level": (
            "CRITICAL" if severity_counts.get("CRITICAL", 0) > 0
            else "HIGH" if severity_counts.get("HIGH", 0) > 0
            else "MEDIUM" if severity_counts.get("MEDIUM", 0) > 0
            else "LOW"
        ),
    }

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"\n[+] Report saved to {args.output}")
    elif args.verbose:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Pod Security Admission Manager - Audit namespaces for PSA compliance,
apply labels, and generate migration reports.
"""

import json
import subprocess
import sys
import argparse


PSA_LABELS = {
    "enforce": "pod-security.kubernetes.io/enforce",
    "enforce-version": "pod-security.kubernetes.io/enforce-version",
    "audit": "pod-security.kubernetes.io/audit",
    "audit-version": "pod-security.kubernetes.io/audit-version",
    "warn": "pod-security.kubernetes.io/warn",
    "warn-version": "pod-security.kubernetes.io/warn-version",
}

SYSTEM_NAMESPACES = {
    "kube-system", "kube-public", "kube-node-lease",
    "calico-system", "tigera-operator", "gatekeeper-system", "falco",
}


def run_kubectl(args: list) -> str:
    """Execute kubectl command."""
    cmd = ["kubectl"] + args
    result = subprocess.run(cmd, capture_output=True, text=True)
    return result.stdout, result.stderr, result.returncode


def get_namespace_psa_labels() -> list:
    """Get PSA labels for all namespaces."""
    stdout, _, _ = run_kubectl(["get", "namespaces", "-o", "json"])
    data = json.loads(stdout)
    results = []
    for ns in data.get("items", []):
        name = ns["metadata"]["name"]
        labels = ns["metadata"].get("labels", {})
        psa_info = {
            "namespace": name,
            "enforce": labels.get(PSA_LABELS["enforce"], "none"),
            "audit": labels.get(PSA_LABELS["audit"], "none"),
            "warn": labels.get(PSA_LABELS["warn"], "none"),
            "is_system": name in SYSTEM_NAMESPACES,
        }
        results.append(psa_info)
    return results


def dry_run_enforcement(namespace: str, level: str) -> dict:
    """Test what would happen if PSA enforcement was applied."""
    stdout, stderr, rc = run_kubectl([
        "label", "--dry-run=server", "--overwrite", "namespace", namespace,
        f"pod-security.kubernetes.io/enforce={level}"
    ])
    violations = []
    if "Warning" in stderr or "Warning" in stdout:
        for line in (stderr + stdout).split("\n"):
            if "Warning" in line or "violate" in line:
                violations.append(line.strip())

    return {
        "namespace": namespace,
        "level": level,
        "would_pass": rc == 0 and len(violations) == 0,
        "violations": violations,
    }


def audit_all_namespaces() -> list:
    """Audit all namespaces for PSA configuration."""
    ns_info = get_namespace_psa_labels()

    print("\n=== Pod Security Admission Audit ===\n")
    print(f"{'Namespace':<30} {'Enforce':<12} {'Audit':<12} {'Warn':<12} {'System'}")
    print("-" * 85)

    compliant = 0
    total = 0

    for ns in ns_info:
        if ns["is_system"]:
            status = "EXEMPT"
        elif ns["enforce"] == "restricted":
            status = "COMPLIANT"
            compliant += 1
        elif ns["enforce"] == "baseline":
            status = "PARTIAL"
        else:
            status = "NON-COMPLIANT"

        if not ns["is_system"]:
            total += 1

        system = "Yes" if ns["is_system"] else "No"
        print(f"{ns['namespace']:<30} {ns['enforce']:<12} {ns['audit']:<12} {ns['warn']:<12} {system}")

    print(f"\n{compliant}/{total} non-system namespaces at restricted level")
    return ns_info


def apply_psa_labels(namespace: str, level: str, version: str = "latest"):
    """Apply PSA labels to a namespace."""
    labels = [
        f"pod-security.kubernetes.io/enforce={level}",
        f"pod-security.kubernetes.io/enforce-version={version}",
        f"pod-security.kubernetes.io/audit=restricted",
        f"pod-security.kubernetes.io/audit-version={version}",
        f"pod-security.kubernetes.io/warn=restricted",
        f"pod-security.kubernetes.io/warn-version={version}",
    ]
    cmd = ["label", "--overwrite", "namespace", namespace] + labels
    stdout, stderr, rc = run_kubectl(cmd)
    if rc == 0:
        print(f"Applied {level} PSA labels to {namespace}")
    else:
        print(f"Failed: {stderr}", file=sys.stderr)


def main():
    parser = argparse.ArgumentParser(description="Pod Security Admission Manager")
    subparsers = parser.add_subparsers(dest="command")

    subparsers.add_parser("audit", help="Audit namespace PSA labels")

    test_cmd = subparsers.add_parser("test", help="Dry-run PSA enforcement")
    test_cmd.add_argument("--namespace", "-n", required=True)
    test_cmd.add_argument("--level", default="restricted", choices=["baseline", "restricted"])

    apply_cmd = subparsers.add_parser("apply", help="Apply PSA labels")
    apply_cmd.add_argument("--namespace", "-n", required=True)
    apply_cmd.add_argument("--level", required=True, choices=["privileged", "baseline", "restricted"])
    apply_cmd.add_argument("--version", default="latest")

    args = parser.parse_args()

    if args.command == "audit":
        audit_all_namespaces()
    elif args.command == "test":
        result = dry_run_enforcement(args.namespace, args.level)
        print(json.dumps(result, indent=2))
    elif args.command == "apply":
        apply_psa_labels(args.namespace, args.level, args.version)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()