#!/usr/bin/env python3
"""Agent for auditing Kubernetes Pod Security Standards enforcement."""

import json
import argparse
import subprocess
from datetime import datetime
from collections import Counter


PSS_LEVELS = {
    "privileged": {"order": 0, "description": "Unrestricted, for system workloads"},
    "baseline": {"order": 1, "description": "Minimally restrictive, prevents known escalations"},
    "restricted": {"order": 2, "description": "Heavily restricted, hardened best practices"},
}

BASELINE_VIOLATIONS = [
    "hostNetwork", "hostPID", "hostIPC", "hostPorts",
    "privileged", "allowPrivilegeEscalation",
    "capabilities.add (non-default)", "seccomp (unconfined)",
]

RESTRICTED_VIOLATIONS = BASELINE_VIOLATIONS + [
    "runAsNonRoot not set", "runAsUser=0",
    "seccompProfile not RuntimeDefault/Localhost",
    "capabilities.drop not ALL", "readOnlyRootFilesystem not set",
]


def kubectl_json(args_list):
    """Run kubectl and return JSON."""
    cmd = ["kubectl"] + args_list + ["-o", "json"]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
    if result.returncode != 0:
        return {"error": result.stderr.strip()}
    return json.loads(result.stdout) if result.stdout.strip() else {}


def audit_namespace_labels():
    """Audit PSA labels on all namespaces."""
    ns_data = kubectl_json(["get", "namespaces"])
    if "error" in ns_data:
        return ns_data
    results = []
    for ns in ns_data.get("items", []):
        name = ns["metadata"]["name"]
        labels = ns["metadata"].get("labels", {})
        enforce = labels.get("pod-security.kubernetes.io/enforce", "")
        audit = labels.get("pod-security.kubernetes.io/audit", "")
        warn = labels.get("pod-security.kubernetes.io/warn", "")
        system = name in ("kube-system", "kube-public", "kube-node-lease")
        results.append({
            "namespace": name, "system": system,
            "enforce": enforce, "audit": audit, "warn": warn,
            "protected": bool(enforce),
            "severity": "INFO" if enforce or system else "HIGH",
        })
    return results


def audit_pod_security(pods_path):
    """Audit pods against Pod Security Standards."""
    with open(pods_path) as f:
        data = json.load(f)
    pods = data if isinstance(data, list) else data.get("items", [])
    findings = []

    for pod in pods:
        metadata = pod.get("metadata", {})
        spec = pod.get("spec", {})
        pod_name = metadata.get("name", "")
        ns = metadata.get("namespace", "default")

        if spec.get("hostNetwork"):
            findings.append({"pod": pod_name, "ns": ns, "violation": "hostNetwork",
                             "level": "baseline", "severity": "HIGH"})
        if spec.get("hostPID"):
            findings.append({"pod": pod_name, "ns": ns, "violation": "hostPID",
                             "level": "baseline", "severity": "HIGH"})

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

            if sc.get("privileged"):
                findings.append({"pod": pod_name, "container": c_name, "ns": ns,
                                 "violation": "privileged", "level": "baseline",
                                 "severity": "CRITICAL"})

            if sc.get("allowPrivilegeEscalation", True):
                findings.append({"pod": pod_name, "container": c_name, "ns": ns,
                                 "violation": "allowPrivilegeEscalation",
                                 "level": "restricted", "severity": "MEDIUM"})

            if not sc.get("runAsNonRoot"):
                findings.append({"pod": pod_name, "container": c_name, "ns": ns,
                                 "violation": "runAsNonRoot not set",
                                 "level": "restricted", "severity": "MEDIUM"})

            caps = sc.get("capabilities", {})
            added = caps.get("add", [])
            if added and any(c not in ("NET_BIND_SERVICE",) for c in added):
                findings.append({"pod": pod_name, "container": c_name, "ns": ns,
                                 "violation": f"capabilities.add: {added}",
                                 "level": "baseline", "severity": "HIGH"})

            dropped = caps.get("drop", [])
            if "ALL" not in dropped:
                findings.append({"pod": pod_name, "container": c_name, "ns": ns,
                                 "violation": "capabilities.drop not ALL",
                                 "level": "restricted", "severity": "MEDIUM"})

    return findings


def generate_namespace_labels(namespace, level="restricted"):
    """Generate PSA label patch for a namespace."""
    labels = {
        f"pod-security.kubernetes.io/enforce": level,
        f"pod-security.kubernetes.io/audit": level,
        f"pod-security.kubernetes.io/warn": level,
    }
    return {
        "command": f'kubectl label namespace {namespace} '
                   f'pod-security.kubernetes.io/enforce={level} '
                   f'pod-security.kubernetes.io/audit={level} '
                   f'pod-security.kubernetes.io/warn={level} --overwrite',
        "labels": labels,
    }


def generate_compliance_report(findings):
    """Generate PSS compliance summary."""
    by_level = Counter(f.get("level", "unknown") for f in findings)
    by_severity = Counter(f.get("severity", "unknown") for f in findings)
    by_violation = Counter(f.get("violation", "unknown") for f in findings)
    return {
        "total_violations": len(findings),
        "by_level": dict(by_level),
        "by_severity": dict(by_severity),
        "top_violations": dict(by_violation.most_common(10)),
        "baseline_violations": by_level.get("baseline", 0),
        "restricted_violations": by_level.get("restricted", 0),
    }


def main():
    parser = argparse.ArgumentParser(description="Kubernetes Pod Security Standards Agent")
    parser.add_argument("--pods", help="Pods JSON to audit")
    parser.add_argument("--namespace", help="Namespace for label generation")
    parser.add_argument("--level", choices=["privileged", "baseline", "restricted"],
                        default="restricted")
    parser.add_argument("--action", choices=["audit-ns", "audit-pods", "generate", "full"],
                        default="full")
    parser.add_argument("--output", default="pss_report.json")
    args = parser.parse_args()

    report = {"generated_at": datetime.utcnow().isoformat(), "results": {}}

    if args.action in ("audit-ns", "full"):
        ns_audit = audit_namespace_labels()
        report["results"]["namespace_audit"] = ns_audit
        if isinstance(ns_audit, list):
            unprotected = sum(1 for n in ns_audit if not n["protected"] and not n["system"])
            print(f"[+] Namespaces: {unprotected} unprotected")

    if args.action in ("audit-pods", "full") and args.pods:
        findings = audit_pod_security(args.pods)
        summary = generate_compliance_report(findings)
        report["results"]["pod_audit"] = findings
        report["results"]["summary"] = summary
        print(f"[+] Pod violations: {summary['total_violations']}")

    if args.action in ("generate", "full") and args.namespace:
        labels = generate_namespace_labels(args.namespace, args.level)
        report["results"]["generated"] = labels
        print(f"[+] Labels for {args.namespace}: {args.level}")

    with open(args.output, "w") as f:
        json.dump(report, f, indent=2, default=str)
    print(f"[+] Report saved to {args.output}")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Kubernetes Pod Security Standards Compliance Checker

Audits namespaces and workloads for PSS enforcement levels
and identifies non-compliant pods.
"""

import subprocess
import json
import sys
from dataclasses import dataclass, field


@dataclass
class PSSFinding:
    namespace: str
    resource: str
    level: str  # baseline, restricted
    violation: str
    severity: str  # CRITICAL, HIGH, MEDIUM
    remediation: str


@dataclass
class PSSReport:
    findings: list = field(default_factory=list)
    namespaces_audited: int = 0
    pods_audited: int = 0
    compliant_pods: int = 0
    non_compliant_pods: int = 0


def run_kubectl(args: list) -> tuple:
    """Execute kubectl command and return parsed JSON."""
    cmd = ["kubectl"] + args + ["-o", "json"]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        if result.returncode != 0:
            return None, result.stderr
        return json.loads(result.stdout), None
    except (subprocess.TimeoutExpired, json.JSONDecodeError) as e:
        return None, str(e)


def get_namespace_pss_labels(ns_data: dict) -> dict:
    """Extract PSS labels from namespace metadata."""
    labels = ns_data.get("metadata", {}).get("labels", {})
    return {
        "enforce": labels.get("pod-security.kubernetes.io/enforce", "not-set"),
        "audit": labels.get("pod-security.kubernetes.io/audit", "not-set"),
        "warn": labels.get("pod-security.kubernetes.io/warn", "not-set"),
        "enforce-version": labels.get("pod-security.kubernetes.io/enforce-version", "not-set"),
    }


def check_restricted_compliance(pod_spec: dict, pod_name: str, namespace: str) -> list:
    """Check pod spec against restricted PSS profile."""
    findings = []
    containers = pod_spec.get("containers", []) + pod_spec.get("initContainers", [])
    pod_security_context = pod_spec.get("securityContext", {})

    # Check pod-level runAsNonRoot
    pod_run_as_non_root = pod_security_context.get("runAsNonRoot", False)
    pod_run_as_user = pod_security_context.get("runAsUser")

    # Check pod-level seccomp
    pod_seccomp = pod_security_context.get("seccompProfile", {})
    pod_seccomp_type = pod_seccomp.get("type", "")

    # Check hostNetwork, hostPID, hostIPC
    for host_ns in ["hostNetwork", "hostPID", "hostIPC"]:
        if pod_spec.get(host_ns, False):
            findings.append(PSSFinding(
                namespace=namespace,
                resource=pod_name,
                level="baseline",
                violation=f"{host_ns} is enabled",
                severity="CRITICAL",
                remediation=f"Set {host_ns}: false in pod spec"
            ))

    # Check hostPath volumes
    for vol in pod_spec.get("volumes", []):
        if "hostPath" in vol:
            findings.append(PSSFinding(
                namespace=namespace,
                resource=pod_name,
                level="baseline",
                violation=f"hostPath volume '{vol.get('name')}' detected",
                severity="HIGH",
                remediation="Replace hostPath with emptyDir, PVC, or configMap"
            ))

    for container in containers:
        c_name = container.get("name", "unknown")
        sc = container.get("securityContext", {})

        # Check privileged
        if sc.get("privileged", False):
            findings.append(PSSFinding(
                namespace=namespace,
                resource=f"{pod_name}/{c_name}",
                level="baseline",
                violation="Container runs in privileged mode",
                severity="CRITICAL",
                remediation="Set privileged: false and use specific capabilities"
            ))

        # Check allowPrivilegeEscalation
        if sc.get("allowPrivilegeEscalation", True):
            findings.append(PSSFinding(
                namespace=namespace,
                resource=f"{pod_name}/{c_name}",
                level="restricted",
                violation="allowPrivilegeEscalation is not explicitly false",
                severity="HIGH",
                remediation="Set allowPrivilegeEscalation: false"
            ))

        # Check runAsNonRoot
        c_run_as_non_root = sc.get("runAsNonRoot")
        c_run_as_user = sc.get("runAsUser")
        if not (c_run_as_non_root or pod_run_as_non_root):
            if not (c_run_as_user and c_run_as_user > 0) and not (pod_run_as_user and pod_run_as_user > 0):
                findings.append(PSSFinding(
                    namespace=namespace,
                    resource=f"{pod_name}/{c_name}",
                    level="restricted",
                    violation="Container may run as root (runAsNonRoot not set)",
                    severity="HIGH",
                    remediation="Set runAsNonRoot: true and runAsUser to non-zero value"
                ))

        # Check capabilities
        caps = sc.get("capabilities", {})
        drop_caps = [c.upper() for c in (caps.get("drop") or [])]
        add_caps = [c.upper() for c in (caps.get("add") or [])]

        if "ALL" not in drop_caps:
            findings.append(PSSFinding(
                namespace=namespace,
                resource=f"{pod_name}/{c_name}",
                level="restricted",
                violation="Capabilities not dropped (missing drop: ALL)",
                severity="HIGH",
                remediation="Add capabilities: drop: ['ALL'] to security context"
            ))

        allowed_add = {"NET_BIND_SERVICE"}
        extra_caps = set(add_caps) - allowed_add
        if extra_caps:
            findings.append(PSSFinding(
                namespace=namespace,
                resource=f"{pod_name}/{c_name}",
                level="restricted",
                violation=f"Disallowed capabilities added: {extra_caps}",
                severity="HIGH",
                remediation="Only NET_BIND_SERVICE may be added in restricted profile"
            ))

        # Check seccomp
        c_seccomp = sc.get("seccompProfile", {})
        c_seccomp_type = c_seccomp.get("type", pod_seccomp_type)
        if c_seccomp_type not in ("RuntimeDefault", "Localhost"):
            findings.append(PSSFinding(
                namespace=namespace,
                resource=f"{pod_name}/{c_name}",
                level="restricted",
                violation=f"Seccomp profile not set or is '{c_seccomp_type}'",
                severity="MEDIUM",
                remediation="Set seccompProfile.type to RuntimeDefault or Localhost"
            ))

        # Check readOnlyRootFilesystem (recommended, not required by PSS)
        if not sc.get("readOnlyRootFilesystem", False):
            findings.append(PSSFinding(
                namespace=namespace,
                resource=f"{pod_name}/{c_name}",
                level="restricted",
                violation="Root filesystem is not read-only (recommended)",
                severity="MEDIUM",
                remediation="Set readOnlyRootFilesystem: true and use emptyDir for writable paths"
            ))

    return findings


def audit_cluster(report: PSSReport):
    """Audit all namespaces and pods for PSS compliance."""
    # Get all namespaces
    ns_data, err = run_kubectl(["get", "namespaces"])
    if ns_data is None:
        print(f"[!] Failed to get namespaces: {err}")
        return

    namespaces = ns_data.get("items", [])
    print(f"[*] Found {len(namespaces)} namespaces")

    for ns in namespaces:
        ns_name = ns["metadata"]["name"]
        pss_labels = get_namespace_pss_labels(ns)
        report.namespaces_audited += 1

        enforce_level = pss_labels["enforce"]
        print(f"\n[*] Namespace: {ns_name} (enforce={enforce_level})")

        # Flag namespaces without PSS enforcement
        if enforce_level == "not-set":
            report.findings.append(PSSFinding(
                namespace=ns_name,
                resource="namespace",
                level="baseline",
                violation="No PSS enforcement label set",
                severity="HIGH" if ns_name not in ("kube-system", "kube-public", "kube-node-lease") else "MEDIUM",
                remediation=f"kubectl label namespace {ns_name} pod-security.kubernetes.io/enforce=baseline"
            ))

        # Get pods in namespace
        pods_data, err = run_kubectl(["get", "pods", "-n", ns_name])
        if pods_data is None:
            continue

        pods = pods_data.get("items", [])
        for pod in pods:
            pod_name = pod["metadata"]["name"]
            pod_spec = pod.get("spec", {})
            report.pods_audited += 1

            pod_findings = check_restricted_compliance(pod_spec, pod_name, ns_name)
            if pod_findings:
                report.non_compliant_pods += 1
                report.findings.extend(pod_findings)
            else:
                report.compliant_pods += 1


def print_report(report: PSSReport):
    """Print audit results."""
    print("\n" + "=" * 70)
    print("KUBERNETES POD SECURITY STANDARDS AUDIT REPORT")
    print("=" * 70)
    print(f"Namespaces audited:    {report.namespaces_audited}")
    print(f"Pods audited:          {report.pods_audited}")
    print(f"Compliant pods:        {report.compliant_pods}")
    print(f"Non-compliant pods:    {report.non_compliant_pods}")
    print(f"Total findings:        {len(report.findings)}")

    if report.pods_audited > 0:
        compliance_rate = (report.compliant_pods / report.pods_audited) * 100
        print(f"Compliance rate:       {compliance_rate:.1f}%")

    print("=" * 70)

    # Group by severity
    for severity in ["CRITICAL", "HIGH", "MEDIUM"]:
        severity_findings = [f for f in report.findings if f.severity == severity]
        if severity_findings:
            print(f"\n{severity} FINDINGS ({len(severity_findings)}):")
            print("-" * 70)
            for f in severity_findings:
                print(f"  [{f.namespace}] {f.resource}")
                print(f"    Level: {f.level} | Violation: {f.violation}")
                print(f"    Fix: {f.remediation}")
                print()


def main():
    print("[*] Kubernetes Pod Security Standards Compliance Checker")
    print("[*] Checking against restricted profile\n")

    report = PSSReport()
    audit_cluster(report)
    print_report(report)

    # Save JSON report
    output = {
        "summary": {
            "namespaces": report.namespaces_audited,
            "pods_audited": report.pods_audited,
            "compliant": report.compliant_pods,
            "non_compliant": report.non_compliant_pods,
        },
        "findings": [
            {
                "namespace": f.namespace,
                "resource": f.resource,
                "level": f.level,
                "violation": f.violation,
                "severity": f.severity,
                "remediation": f.remediation,
            }
            for f in report.findings
        ],
    }

    with open("pss_audit_report.json", "w") as f:
        json.dump(output, f, indent=2)
    print("[*] Report saved to pss_audit_report.json")

    critical_high = [f for f in report.findings if f.severity in ("CRITICAL", "HIGH")]
    if critical_high:
        print(f"\n[!] {len(critical_high)} CRITICAL/HIGH findings require attention")
        sys.exit(1)


if __name__ == "__main__":
    main()