#!/usr/bin/env python3
"""iac_prescan.py — deterministic Terraform / Kubernetes security pre-scan.

ARION service artifact. Scans infrastructure-as-code files for the
misconfiguration classes that are checkable offline, line/indentation-
wise with lightweight block tracking — no HCL/YAML parser, no cloud
APIs, no network.

Terraform (.tf / .tfvars):
  * world-open ingress: security-group / firewall rules exposing
    0.0.0.0/0 (or Azure "*") to admin/datastore ports (22, 3389, 3306,
    5432, 6379, 9200, 27017, 2375, 10250, ...) — SG-WORLD-ADMIN (HIGH);
    world-open on other ports — SG-WORLD (MEDIUM)
  * public datastores / buckets: publicly_accessible = true
    (DB-PUBLIC HIGH), S3 acl public-read[-write] (S3-PUBLIC-ACL HIGH),
    associate_public_ip_address (PUB-IP LOW)
  * encryption off: storage_encrypted/encrypted = false (ENC-OFF
    MEDIUM), s3 backend without encrypt = true (BACKEND-ENC INFO)
  * hardcoded credentials: literal values for password/secret/token/
    *_key assignments (TF-SECRET HIGH), AWS access-key shapes
    (AWS-KEY HIGH), PEM private-key blocks (PRIVKEY HIGH)
  * hygiene: versioning enabled = false (S3-NO-VERSIONING LOW),
    skip_final_snapshot = true (LOW), tfsec/checkov skip comments
    (SKIP-COMMENT INFO)

Kubernetes (.yaml/.yml docs carrying apiVersion+kind):
  * privilege & escape: privileged: true (K8S-PRIV HIGH),
    allowPrivilegeEscalation: true (K8S-PRIV-ESC MEDIUM),
    runAsUser: 0 / runAsNonRoot: false (K8S-ROOT MEDIUM),
    added capabilities SYS_ADMIN/ALL (K8S-CAP HIGH; NET_ADMIN,
    SYS_MODULE, SYS_PTRACE, DAC_* MEDIUM)
  * host sharing: hostNetwork/hostPID/hostIPC: true (K8S-HOST-NS
    MEDIUM), hostPath volumes (K8S-HOSTPATH MEDIUM)
  * secrets: env entries named *PASSWORD/*SECRET/*TOKEN/*KEY with a
    literal value: (K8S-SECRET-ENV HIGH; valueFrom is fine)
  * RBAC: "*" wildcard verbs/resources in Role/ClusterRole rules
    (K8S-RBAC-WILDCARD MEDIUM)
  * hygiene: :latest / untagged images (K8S-LATEST LOW),
    readOnlyRootFilesystem: false (LOW), Ingress without tls: (LOW),
    workload without resources.limits (K8S-NO-LIMITS INFO),
    LoadBalancer/NodePort services (INFO)

Honest scope: PRE-SCAN, not an audit. It does not evaluate module
composition, variable resolution, plan output, or cloud state —
a FLAG means "deserves human review"; a clean report is not proof
the infrastructure is safe.

Usage:
  iac_prescan.py <file-or-dir> [--json]
  iac_prescan.py --selftest
Stdlib only. Exit 0 = no HIGH/MEDIUM findings, 1 = findings, 2 = usage error.
"""
import json
import os
import re
import sys

SEV_ORDER = {"HIGH": 0, "MEDIUM": 1, "LOW": 2, "INFO": 3}

SENSITIVE_PORTS = {
    22, 23, 135, 139, 445, 1433, 1521, 2375, 2376, 2379, 2380, 3000,
    3306, 3389, 5432, 5601, 5672, 5900, 5985, 5986, 6379, 7000, 7199,
    8000, 8080, 8500, 8888, 9000, 9042, 9092, 9200, 9300, 10250,
    11211, 15672, 15675, 2049, 27017, 27018, 28015, 50000, 61616,
}

SECRET_NAME_RE = re.compile(
    r"(?i)(password|passwd|secret|token|api[_-]?key|private[_-]?key|"
    r"access[_-]?key|secret[_-]?key|client[_-]?secret|ssh[_-]?key|"
    r"credential|mnemonic|seed[_-]?phrase)")
PRIVKEY_RE = re.compile(r"BEGIN [A-Z0-9 ]*PRIVATE KEY")
AWS_KEY_RE = re.compile(r"\b(?:AKIA|ASIA|ABIA|ACCA)[0-9A-Z]{16}\b")
SKIP_COMMENT_RE = re.compile(
    r"(?i)(tfsec:ignore|checkov:skip|bridgecrew:skip|kics-scan\s+ignore)")

RESOURCE_RE = re.compile(r'^\s*resource\s+"([^"]+)"\s+"([^"]+)"')
DATA_RE = re.compile(r'^\s*data\s+"([^"]+)"')
BACKEND_RE = re.compile(r'^\s*backend\s+"([^"]+)"')
SUBBLOCK_RE = re.compile(r'^\s*(ingress|egress|allow|deny|rules|'
                         r'internal|external|ebs_block_device|'
                         r'root_block_device|versioning|logging)\b')

# resource types treated as ingress rules
SG_RULE_TYPES = {
    "aws_security_group", "aws_security_group_rule",
    "aws_vpc_security_group_ingress_rule",
    "azurerm_network_security_rule", "azurerm_network_security_group",
    "azurerm_firewall_network_rule_collection",
    "google_compute_firewall", "google_compute_firewall_rule",
    "openstack_compute_secgroup_v2", "openstack_networking_secgroup_rule_v2",
    "digitalocean_firewall", "oci_core_security_list",
    "hcloud_firewall", "linode_firewall",
}
WORLD_CIDRS = re.compile(
    r'(?:"|\s)(0\.0\.0\.0/0|::/0|\*|Internet|any|world)(?:"|\s|,|\])')
CIDR_ATTR_RE = re.compile(
    r"(?i)(cidr_blocks|ipv6_cidr_blocks|cidr_ipv4|cidr_ipv6|"
    r"source_ranges|"
    r"source_address_prefix(?:es)?|source_cidr|remote_ip_prefix|"
    r"source_ip_ranges|allowed_cidrs|cidr)\s*=")
PORT_ATTR_RE = re.compile(
    r"(?i)(from_port|to_port|port|ports|destination_port_range|"
    r"destination_port_ranges|port_range|dst_ports)\s*=\s*"
    r"\"?([0-9\-*,\s\"\[\]]+)")
DIRECTION_IN_RE = re.compile(
    r'(?i)(direction\s*=\s*"?(ingress|inbound|in|INGRESS)|type\s*=\s*"ingress)')


def sev_rank(s):
    return SEV_ORDER.get(s, 9)


class Finding:
    def __init__(self, code, sev, file, line, msg):
        self.code, self.sev, self.file, self.line, self.msg = \
            code, sev, file, line, msg

    def as_dict(self):
        return {"code": self.code, "severity": self.sev,
                "file": self.file, "line": self.line, "msg": self.msg}


def find_targets(root):
    out = []
    SKIP = {".git", "node_modules", "vendor", ".terraform", "dist",
            "build", "__pycache__", ".github"}
    if os.path.isfile(root):
        return [root]
    for base, dirs, files in os.walk(root):
        dirs[:] = [d for d in dirs if d not in SKIP]
        for f in files:
            lf = f.lower()
            if lf.endswith((".tf", ".tfvars", ".yaml", ".yml")):
                out.append(os.path.join(base, f))
    return sorted(out)


# ---------------------------------------------------------------- terraform

def _ports_in(text):
    """Extract integer port set from a rule block's text."""
    ports = set()
    for m in PORT_ATTR_RE.finditer(text):
        body = m.group(2)
        if m.group(1).lower() == "ports" or "[" in body:
            for num in re.findall(r"\d+", body):
                ports.add(int(num))
        else:
            for tok in re.findall(r"\d+-\d+|\d+", body):
                if "-" in tok:
                    a, b = tok.split("-", 1)
                    ports.update(range(int(a), min(int(b), int(a) + 512) + 1))
                else:
                    ports.add(int(tok))
    return ports


def _analyze_rule_block(buf, base_line, rel, out):
    """buf = list of (lineno, text) for one ingress rule body."""
    text = "\n".join(t for _, t in buf)
    if not (CIDR_ATTR_RE.search(text) and WORLD_CIDRS.search(text)):
        return
    # egress-only rules are not ingress findings
    if re.search(r'(?i)type\s*=\s*"egress"', text) and \
            not DIRECTION_IN_RE.search(text):
        return
    world_line = base_line
    for ln, t in buf:
        if CIDR_ATTR_RE.search(t) and WORLD_CIDRS.search(t):
            world_line = ln
            break
    ports = _ports_in(text)
    hit = sorted(ports & SENSITIVE_PORTS)
    if hit:
        out.append(Finding(
            "SG-WORLD-ADMIN", "HIGH", rel, world_line,
            "world-open ingress (0.0.0.0/0) to sensitive port(s) %s — "
            "restrict source CIDRs" % ",".join(map(str, hit[:8]))))
    else:
        out.append(Finding(
            "SG-WORLD", "MEDIUM", rel, world_line,
            "world-open ingress (0.0.0.0/0)%s — restrict source CIDRs"
            % (" on port(s) %s" % ",".join(map(str, sorted(ports)[:8]))
               if ports else "")))


def scan_tf(path, rel):
    out = []
    try:
        lines = open(path, encoding="utf-8", errors="replace")\
            .read().splitlines()
    except OSError:
        return out

    # --- pass 1: capture resource/rule blocks for SG analysis ---
    i = 0
    n = len(lines)
    while i < n:
        line = lines[i]
        m = RESOURCE_RE.match(line)
        if m and m.group(1) in SG_RULE_TYPES:
            rtype = m.group(1)
            start = i
            depth = line.count("{") - line.count("}")
            buf = [(i + 1, line)]
            i += 1
            # capture until block closes; split aws_security_group
            # into per-direction sub-buffers
            sub_name, sub_buf, sub_start = None, [], 0
            while i < n and depth > 0:
                l = lines[i]
                depth += l.count("{") - l.count("}")
                sb = SUBBLOCK_RE.match(l)
                if rtype == "aws_security_group" and sb and \
                        l.rstrip().endswith("{"):
                    if sub_buf:
                        _flush_sg_sub(sub_name, sub_buf, sub_start,
                                      rel, out)
                    sub_name, sub_buf, sub_start = \
                        sb.group(1), [(i + 1, l)], i + 1
                elif sub_buf is not None and sub_name:
                    sub_buf.append((i + 1, l))
                buf.append((i + 1, l))
                i += 1
            if sub_buf:
                _flush_sg_sub(sub_name, sub_buf, sub_start, rel, out)
            if rtype != "aws_security_group":
                _analyze_rule_block(buf, start + 1, rel, out)
            continue
        i += 1

    # --- pass 2: line-wise checks ---
    in_heredoc = False
    cur_res = ""
    depth = 0
    backend_stack = []  # (type, saw_encrypt)
    for idx, line in enumerate(lines):
        ln = idx + 1
        rm = RESOURCE_RE.match(line)
        if rm:
            cur_res = rm.group(1)
        bm = BACKEND_RE.match(line)
        if bm:
            backend_stack.append([bm.group(1), False, ln])
        if re.search(r"<<-?\s*[A-Z_]+", line):
            in_heredoc = True
        if in_heredoc and re.match(r"^\s*[A-Z_]+\s*$", line):
            in_heredoc = False
        if backend_stack:
            if "{" in line:
                pass
            if re.search(r"\bencrypt\s*=\s*true", line):
                backend_stack[-1][1] = True
            if re.match(r"^\s*}\s*$", line):
                b = backend_stack.pop()
                if b[0] in ("s3", "gcs", "azurerm") and not b[1]:
                    out.append(Finding(
                        "BACKEND-ENC", "INFO", rel, b[2],
                        "tfstate backend %r without encrypt = true" % b[0]))

        s = line.strip()
        if s.startswith("#") or s.startswith("//"):
            if SKIP_COMMENT_RE.search(s):
                out.append(Finding(
                    "SKIP-COMMENT", "INFO", rel, ln,
                    "linter skip directive — verify the suppressed "
                    "finding was triaged"))
            continue
        if SKIP_COMMENT_RE.search(s):
            out.append(Finding(
                "SKIP-COMMENT", "INFO", rel, ln,
                "linter skip directive — verify the suppressed "
                "finding was triaged"))
        if PRIVKEY_RE.search(line):
            out.append(Finding(
                "PRIVKEY", "HIGH", rel, ln,
                "private-key material embedded in IaC — move to a "
                "secret manager"))
        if AWS_KEY_RE.search(line):
            out.append(Finding(
                "AWS-KEY", "HIGH", rel, ln,
                "AWS access-key-shaped literal committed to IaC"))
        m = re.search(
            r"(?i)\b(password|passwd|secret|token|api_key|private_key|"
            r"access_key|secret_key|client_secret)\s*=\s*"
            r"\"([^\"]*)\"", line)
        if m:
            val = m.group(2)
            # skip references/interpolation and placeholders
            if val and not re.search(
                    r"(var\.|local\.|each\.|data\.|module\.|\$\{|"
                    r"self\.|^\*+$|^x+$|^<|^CHANGEME|^EXAMPLE|^TODO)",
                    val, re.I):
                out.append(Finding(
                    "TF-SECRET", "HIGH", rel, ln,
                    "literal credential assigned to %r — move to "
                    "tfvars outside VCS or a secret manager"
                    % m.group(1)))
        else:
            # env-style unquoted assignments (user_data heredocs, .env
            # embeds): NAME=value where the name carries a secret word
            m2 = re.search(
                r"\b([A-Za-z_][A-Za-z0-9_]*(?:PASSWORD|PASSWD|SECRET|"
                r"TOKEN|API_KEY|PRIVATE_KEY|ACCESS_KEY|SECRET_KEY|"
                r"CLIENT_SECRET|SSH_KEY)[A-Za-z0-9_]*)\s*=\s*"
                r"([^\s\"'\[{][^\s\"']{3,})", line)
            if m2 and not re.match(
                    r"(?i)(var\.|local\.|each\.|data\.|module\.|file\(|"
                    r"templatefile\(|base64|true$|false$|null$|\d+$)",
                    m2.group(2)):
                out.append(Finding(
                    "TF-SECRET", "HIGH", rel, ln,
                    "literal credential assigned to %r — move to "
                    "a secret manager" % m2.group(1)))
        if re.search(r"(?i)publicly_accessible\s*=\s*true", line):
            out.append(Finding(
                "DB-PUBLIC", "HIGH", rel, ln,
                "datastore publicly_accessible = true"))
        m = re.search(r'(?i)\bacl\s*=\s*"([a-z-]+)"', line)
        if m and m.group(1) in ("public-read", "public-read-write"):
            out.append(Finding(
                "S3-PUBLIC-ACL", "HIGH", rel, ln,
                "bucket acl %r grants public access" % m.group(1)))
        elif m and m.group(1) == "authenticated-read":
            out.append(Finding(
                "S3-AUTH-READ", "MEDIUM", rel, ln,
                "acl authenticated-read exposes bucket to every AWS "
                "account, not just yours"))
        if re.search(r"(?i)\b(storage_encrypted|encrypted|"
                     r"encryption_enabled|disk_encryption)\s*=\s*false",
                     line):
            out.append(Finding(
                "ENC-OFF", "MEDIUM", rel, ln,
                "encryption explicitly disabled"))
        if re.search(r"(?i)associate_public_ip_address\s*=\s*true",
                     line):
            out.append(Finding(
                "PUB-IP", "LOW", rel, ln,
                "instance gets a public IP — confirm intended exposure"))
        if re.search(r"(?i)enabled\s*=\s*false", line) and \
                cur_res in ("aws_s3_bucket",
                            "aws_s3_bucket_versioning") or \
                re.search(r"(?i)versioning\s*=\s*false", line):
            if cur_res.startswith("aws_s3") or \
                    "versioning" in line.lower():
                out.append(Finding(
                    "S3-NO-VERSIONING", "LOW", rel, ln,
                    "bucket versioning disabled — no recovery from "
                    "overwrite/delete"))
        if re.search(r"(?i)skip_final_snapshot\s*=\s*true", line):
            out.append(Finding(
                "SKIP-SNAPSHOT", "LOW", rel, ln,
                "DB deletion skips final snapshot"))
        if re.search(r"(?i)deletion_protection\s*=\s*false", line):
            out.append(Finding(
                "NO-DEL-PROT", "INFO", rel, ln,
                "deletion protection off"))
    return out


def _flush_sg_sub(name, buf, start, rel, out):
    if name in ("ingress", "allow"):
        _analyze_rule_block(buf, start, rel, out)


# ------------------------------------------------------------- kubernetes

K8S_WORKLOAD = {
    "Pod", "Deployment", "DaemonSet", "StatefulSet", "Job", "CronJob",
    "ReplicaSet", "ReplicationController",
}
K8S_CAP_HIGH = {"SYS_ADMIN", "ALL"}
K8S_CAP_MED = {"NET_ADMIN", "SYS_MODULE", "SYS_PTRACE", "SYS_RAWIO",
               "DAC_OVERRIDE", "DAC_READ_SEARCH", "SYS_BOOT",
               "SYS_TIME", "MKNOD"}


def scan_k8s(path, rel):
    try:
        text = open(path, encoding="utf-8", errors="replace").read()
    except OSError:
        return []
    if "apiVersion:" not in text or "kind:" not in text:
        return []
    out = []
    kind = ""
    saw_limits = False
    saw_tls = False
    env_name_pending = None
    doc_workload = False

    def doc_end():
        nonlocal saw_limits, saw_tls
        if doc_workload and not saw_limits:
            out.append(Finding(
                "K8S-NO-LIMITS", "INFO", rel, doc_start,
                "workload %s has no resources.limits — unbounded "
                "CPU/memory" % (kind or "?")))
        if kind == "Ingress" and not saw_tls:
            out.append(Finding(
                "K8S-ING-NOTLS", "LOW", rel, doc_start,
                "Ingress without a tls: block — plaintext only"))

    doc_start = 1
    for idx, line in enumerate(text.splitlines()):
        ln = idx + 1
        s = line.strip()
        if s.startswith("---"):
            doc_end()
            kind, saw_limits, saw_tls = "", False, False
            doc_workload = False
            env_name_pending = None
            doc_start = ln
            continue
        m = re.match(r"kind:\s*(\S+)", s)
        if m:
            kind = m.group(1)
            doc_workload = kind in K8S_WORKLOAD
            continue
        if re.match(r"tls\s*:", s):
            saw_tls = True
        if re.match(r"limits\s*:", s):
            saw_limits = True
        if PRIVKEY_RE.search(line):
            out.append(Finding(
                "PRIVKEY", "HIGH", rel, ln,
                "private-key material embedded in manifest"))
        if AWS_KEY_RE.search(line):
            out.append(Finding(
                "AWS-KEY", "HIGH", rel, ln,
                "AWS access-key-shaped literal in manifest"))
        if SKIP_COMMENT_RE.search(s):
            out.append(Finding(
                "SKIP-COMMENT", "INFO", rel, ln,
                "linter skip directive — verify triage"))
        if re.search(r"privileged\s*:\s*true", s):
            out.append(Finding(
                "K8S-PRIV", "HIGH", rel, ln,
                "privileged: true — container = host root"))
        if re.search(r"allowPrivilegeEscalation\s*:\s*true", s):
            out.append(Finding(
                "K8S-PRIV-ESC", "MEDIUM", rel, ln,
                "allowPrivilegeEscalation: true"))
        if re.search(r"runAsUser\s*:\s*0\b", s) or \
                re.search(r"runAsNonRoot\s*:\s*false", s):
            out.append(Finding(
                "K8S-ROOT", "MEDIUM", rel, ln,
                "container runs as root"))
        if re.search(r"host(Network|PID|IPC)\s*:\s*true", s):
            out.append(Finding(
                "K8S-HOST-NS", "MEDIUM", rel, ln,
                "shares host %s namespace"
                % re.search(r"host(Network|PID|IPC)", s).group(1)))
        if re.match(r"hostPath\s*:", s):
            out.append(Finding(
                "K8S-HOSTPATH", "MEDIUM", rel, ln,
                "hostPath volume mounts host filesystem"))
        m = re.search(r"-\s*\"?([A-Z_]+)\"?\s*$", s)
        if m and m.group(1) in K8S_CAP_HIGH:
            out.append(Finding(
                "K8S-CAP", "HIGH", rel, ln,
                "added capability %s — near-root equivalence"
                % m.group(1)))
        elif m and m.group(1) in K8S_CAP_MED:
            out.append(Finding(
                "K8S-CAP", "MEDIUM", rel, ln,
                "added capability %s" % m.group(1)))
        if kind in ("Role", "ClusterRole") and re.search(
                r'("|\')\*("|\')|^\s*-\s*\*\s*$', s):
            out.append(Finding(
                "K8S-RBAC-WILDCARD", "MEDIUM", rel, ln,
                "wildcard * in %s rule" % kind))
        m = re.match(r"image\s*:\s*(\S+)", s)
        if m:
            img = m.group(1).strip("'\"")
            if img.endswith(":latest") or (":" not in img and
                                          "@" not in img):
                out.append(Finding(
                    "K8S-LATEST", "LOW", rel, ln,
                    "mutable image ref %r" % img))
        if re.search(r"readOnlyRootFilesystem\s*:\s*false", s):
            out.append(Finding(
                "K8S-RW-ROOT", "LOW", rel, ln,
                "writable root filesystem"))
        if re.search(r"type\s*:\s*LoadBalancer", s):
            out.append(Finding(
                "K8S-LB", "INFO", rel, ln,
                "Service type LoadBalancer — public endpoint"))
        if re.search(r"type\s*:\s*NodePort", s):
            out.append(Finding(
                "K8S-NODEPORT", "INFO", rel, ln,
                "NodePort exposes the service on every node"))
        if re.search(r"automountServiceAccountToken\s*:\s*true", s):
            out.append(Finding(
                "K8S-SATOKEN", "INFO", rel, ln,
                "service-account token mounted — set false if unused"))
        # env literal secrets: `- name: X` then `value: literal`
        m = re.match(r"-\s*name\s*:\s*([A-Za-z0-9_]+)", s)
        if m:
            env_name_pending = m.group(1) if SECRET_NAME_RE.search(
                m.group(1)) else None
            continue
        m = re.match(r"value\s*:\s*(.+)", s)
        if m and env_name_pending:
            out.append(Finding(
                "K8S-SECRET-ENV", "HIGH", rel, ln,
                "env %s carries a literal value — use a Secret/"
                "valueFrom" % env_name_pending))
            env_name_pending = None
        elif re.match(r"valueFrom\s*:", s) or re.match(r"-\s", s):
            env_name_pending = None
    doc_end()
    return out


# ------------------------------------------------------------------ driver

def render(findings, targets):
    order = sorted(findings,
                   key=lambda f: (sev_rank(f.sev), f.file, f.line))
    out = ["# IaC security pre-scan report", "",
           "Scanned %d file(s)." % len(targets), ""]
    counts = {}
    for x in findings:
        counts[x.sev] = counts.get(x.sev, 0) + 1
    out.append("Findings: " + (", ".join(
        "%d %s" % (counts[s], s) for s in ("HIGH", "MEDIUM", "LOW", "INFO")
        if s in counts) or "none"))
    out.append("")
    for x in order:
        out.append("- **%s** `%s` %s:%d — %s" % (x.sev, x.code, x.file,
                                                x.line, x.msg))
    out.append("")
    out.append("_Pre-scan, not an audit: variable resolution, module "
               "composition, and cloud state are out of scope; a clean "
               "report is not proof the infrastructure is safe._")
    return "\n".join(out)


SELFTEST_TF = [
    ('resource "aws_security_group" "sg" {\n'
     '  ingress {\n'
     '    from_port = 22\n    to_port = 22\n'
     '    cidr_blocks = ["0.0.0.0/0"]\n  }\n}\n',
     ["SG-WORLD-ADMIN"]),
    ('resource "aws_vpc_security_group_ingress_rule" "r" {\n'
     '  from_port = 7777\n  to_port = 7777\n'
     '  cidr_ipv4 = "0.0.0.0/0"\n}\n',
     ["SG-WORLD"]),
    ('resource "azurerm_network_security_rule" "r" {\n'
     '  direction = "Inbound"\n  destination_port_range = "3389"\n'
     '  source_address_prefix = "*"\n}\n',
     ["SG-WORLD-ADMIN"]),
    ('resource "aws_security_group" "sg" {\n'
     '  egress {\n    from_port = 0\n    to_port = 0\n'
     '    cidr_blocks = ["0.0.0.0/0"]\n  }\n}\n',
     []),
    ('resource "aws_db_instance" "db" {\n'
     '  publicly_accessible = true\n  skip_final_snapshot = true\n}\n',
     ["DB-PUBLIC", "SKIP-SNAPSHOT"]),
    ('resource "aws_s3_bucket" "b" {\n  acl = "public-read"\n}\n',
     ["S3-PUBLIC-ACL"]),
    ('resource "aws_instance" "i" {\n'
     '  associate_public_ip_address = true\n}\n', ["PUB-IP"]),
    ('variable "x" {\n  default = "ok"\n}\n'
     'resource "aws_db_instance" "d" {\n'
     '  password = "hunter2plain"\n  storage_encrypted = false\n}\n',
     ["TF-SECRET", "ENC-OFF"]),
    ('resource "aws_db_instance" "d" {\n'
     '  password = var.db_password\n}\n', []),
    ('terraform {\n  backend "s3" {\n    bucket = "st"\n'
     '    key = "x.tfstate"\n  }\n}\n', ["BACKEND-ENC"]),
]

SELFTEST_K8S = [
    ("""apiVersion: apps/v1
kind: Deployment
spec:
  template:
    spec:
      containers:
      - name: a
        image: nginx:latest
        securityContext:
          privileged: true
          allowPrivilegeEscalation: true
          runAsUser: 0
""",
     {"K8S-PRIV", "K8S-PRIV-ESC", "K8S-ROOT", "K8S-LATEST",
      "K8S-NO-LIMITS"}),
    ("""apiVersion: v1
kind: Pod
spec:
  hostNetwork: true
  containers:
  - name: a
    image: repo/x@sha256:abc
    env:
    - name: DB_PASSWORD
      value: hunter2
    resources:
      limits:
        cpu: "1"
  volumes:
  - name: h
    hostPath:
      path: /etc
""",
     {"K8S-HOST-NS", "K8S-SECRET-ENV", "K8S-HOSTPATH"}),
    ("""apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
rules:
- apiGroups: ["*"]
  resources: ["*"]
  verbs: ["*"]
""", {"K8S-RBAC-WILDCARD"}),
    ("""apiVersion: v1
kind: Pod
spec:
  containers:
  - name: a
    image: nginx:1.27
    env:
    - name: API_TOKEN
      valueFrom:
        secretKeyRef:
          name: s
    resources:
      limits: {cpu: "1"}
""", set()),
]


def run_selftest():
    import tempfile
    fails = []
    for i, (src, expect) in enumerate(SELFTEST_TF):
        with tempfile.TemporaryDirectory() as td:
            p = os.path.join(td, "main.tf")
            open(p, "w").write(src)
            got = {f.code for f in scan_tf(p, "main.tf")}
            missing = set(expect) - got
            if missing:
                fails.append(("tf%d" % i, "missing=%s" % missing))
    for i, (src, expect) in enumerate(SELFTEST_K8S):
        with tempfile.TemporaryDirectory() as td:
            p = os.path.join(td, "m.yaml")
            open(p, "w").write(src)
            got = {f.code for f in scan_k8s(p, "m.yaml")}
            missing = expect - got
            extra = got - expect - {"K8S-NO-LIMITS", "K8S-LATEST"}
            if missing or extra:
                fails.append(("k8s%d" % i, "missing=%s extra=%s"
                              % (missing, extra)))
    if fails:
        for f_ in fails:
            print("SELFTEST FAIL", f_)
        return 1
    print("SELFTEST PASS (%d vectors)"
          % (len(SELFTEST_TF) + len(SELFTEST_K8S)))
    return 0


def main():
    args = [a for a in sys.argv[1:] if a != "--json"]
    want_json = "--json" in sys.argv
    if "--selftest" in sys.argv:
        return run_selftest()
    if not args:
        print(__doc__)
        return 2
    root = args[0]
    targets = find_targets(root)
    findings = []
    for p in targets:
        rel = os.path.relpath(p, root) if os.path.isdir(root) \
            else os.path.basename(p)
        if p.lower().endswith((".yaml", ".yml")):
            findings += scan_k8s(p, rel)
        else:
            findings += scan_tf(p, rel)
    if want_json:
        print(json.dumps([f.as_dict() for f in findings], indent=1))
    else:
        print(render(findings, targets))
    return 1 if any(f.sev in ("HIGH", "MEDIUM") for f in findings) else 0


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