#!/usr/bin/env python3
"""docker_prescan.py — deterministic Dockerfile / compose security pre-scan.

ARION service artifact. Scans container build and orchestration files
for the misconfiguration classes that are checkable offline:

  * running as root: no ``USER`` instruction, or ``USER root``/``0``
  * mutable base pins: ``:latest``, untagged images, no ``@sha256:``
    digest pin (tags move; digests are immutable)
  * unverified remote fetches: ``ADD <url>``, ``curl|wget ... | sh``
    pipe-to-shell, downloads with no checksum/signature verification
  * secrets in the build: ``ENV``/``ARG``/compose ``environment``
    values matching credential names (password/token/secret/key)
  * broad context copies: ``COPY . /`` / ``ADD .`` — leaks .env,
    .git, and credentials into the layer history
  * dangerous permissions: ``chmod 777``, ``setcap`` grants
  * compose host-equivalence: ``privileged: true``,
    ``network_mode: host``, ``cap_add: SYS_ADMIN|ALL``,
    ``/var/run/docker.sock`` mounts, ``pid: host``/``ipc: host``
  * hygiene signals (INFO): missing HEALTHCHECK, EXPOSE 22,
    ``apt-get`` without cleanup flags, sudo inside RUN

Honest scope: PRE-SCAN, not an audit. Files are scanned
line/indentation-wise — multi-stage dataflow, base-image CVEs, and
runtime behavior are out of scope. A FLAG means "deserves human
review"; a clean report is not proof an image is safe.

Usage:
  docker_prescan.py <Dockerfile-or-dir> [--json]
  docker_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}

SECRET_NAME_RE = re.compile(
    r"(?i)(password|passwd|secret|token|api[_-]?key|private[_-]?key|"
    r"access[_-]?key|auth|credential|mnemonic|seed[_-]?phrase)")
PIPE_SHELL_RE = re.compile(
    r"(curl|wget|fetch)\b[^|\n&;]*\|\s*(sudo\s+)?(bash|sh|zsh|ash)")
URL_RE = re.compile(r"https?://", re.I)
CHECKSUM_RE = re.compile(
    r"(?i)(sha256|sha512|sha-256|md5sum|sha256sum|gpg\s+--verify|"
    r"--check(?!-health)|cosign\s+verify|minisign)")
CHMOD777_RE = re.compile(r"chmod\s+(-R\s+)?(777|a\+rwx|o\+rwx)\b")
COMPOSE_NAME_RE = re.compile(
    r"(?i)(docker-?compose|compose).ya?ml$|^compose$")
DOCKERFILE_NAME_RE = re.compile(r"(?i)dockerfile|\.dockerfile$|\.df$")

KNOWN_REGISTRY_PREFIX = re.compile(r"^[\w.-]+(:\d+)?/")


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


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

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


def is_comment_or_blank(line):
    s = line.strip()
    return not s or s.startswith("#")


def strip_inline_comment(line):
    # '#' inside a quoted RUN string could be real shell; only strip
    # when preceded by whitespace and not inside quotes — cheap rule:
    # Dockerfile comments must start the line anyway after whitespace.
    return line


def parse_instruction(line):
    """Return (instr, rest) for a Dockerfile line, or None."""
    m = re.match(r"^\s*([A-Za-z]+)\s+(.*)$", line)
    if not m:
        return None
    instr = m.group(1).upper()
    if instr not in ("FROM", "RUN", "CMD", "ENTRYPOINT", "COPY", "ADD",
                     "ENV", "ARG", "USER", "WORKDIR", "EXPOSE", "VOLUME",
                     "LABEL", "HEALTHCHECK", "SHELL", "ONBUILD", "STOPSIGNAL"):
        return None
    return instr, m.group(2)


def logical_lines(raw):
    """Join backslash continuations; yields (start_lineno, text)."""
    buf = ""
    start = None
    for i, line in enumerate(raw.splitlines(), 1):
        stripped = line.rstrip()
        if start is None:
            start = i
        if stripped.endswith("\\") and not is_comment_or_blank(stripped):
            buf += stripped[:-1] + " "
            continue
        buf += stripped
        yield start, buf
        buf = ""
        start = None
    if buf:
        yield start or 1, buf


def scan_dockerfile(path, rel):
    findings = []
    raw = open(path, encoding="utf-8", errors="replace").read()
    lines = list(logical_lines(raw))
    if not lines:
        return findings

    stages = []          # per-stage dicts
    cur = {"user": None, "from": None, "healthcheck": False}
    any_user_nonroot = False

    def f(code, sev, ln, msg):
        findings.append(Finding(code, sev, rel, ln, msg))

    for ln, text in lines:
        if is_comment_or_blank(text):
            continue
        ins = parse_instruction(text)
        if not ins:
            continue
        op, arg = ins

        if op == "FROM":
            if cur["from"] is not None:
                stages.append(cur)
            cur = {"user": None, "from": arg.strip(), "healthcheck": False}
            ref = arg.strip().split()[-1]
            if re.search(r"(?i)\s+as\s+\w+$", arg.strip()):
                ref = arg.strip().split()[0]
            if ref.lower() in ("scratch",):
                pass
            elif "@sha256:" in ref:
                pass  # digest-pinned: immutable, best practice
            elif ref.startswith("$") or "${" in ref:
                f("FROM-ARGPIN", "INFO", ln,
                  "base image via ARG variable — pin cannot be "
                  "verified statically (%s)" % ref[:60])
            else:
                name = ref.split(":")[0]
                tag = ref.split(":")[1] if ":" in ref else None
                if tag is None or tag.lower() == "latest":
                    f("BASE-LATEST", "MEDIUM", ln,
                      "base image %s — mutable tag; pin a digest "
                      "(@sha256:...)" % ref)
                else:
                    f("BASE-TAGPIN", "LOW", ln,
                      "base image %s — tag pin is mutable; digest pin "
                      "is stronger" % ref)
        elif op == "USER":
            cur["user"] = arg.strip()
            u = arg.strip().split(":")[0]
            if u in ("0", "root") or u.startswith("${") and "root" in arg:
                f("USER-ROOT", "MEDIUM", ln,
                  "explicit USER %s — container runs as root" % u)
            else:
                any_user_nonroot = True
        elif op == "RUN":
            a = arg
            if PIPE_SHELL_RE.search(a):
                f("PIPE-SHELL", "HIGH", ln,
                  "pipe-to-shell in RUN — remote script executed "
                  "without verification")
            elif URL_RE.search(a) and re.search(r"\b(curl|wget)\b", a) \
                    and not CHECKSUM_RE.search(a):
                f("FETCH-NOCHECKSUM", "MEDIUM", ln,
                  "remote download without checksum/signature "
                  "verification in the same instruction")
            if CHMOD777_RE.search(a):
                f("CHMOD-777", "MEDIUM", ln,
                  "world-writable permissions via chmod")
            if re.search(r"\bsudo\b", a):
                f("SUDO-IN-RUN", "LOW", ln,
                  "sudo inside RUN — usually means the image should "
                  "not need it")
            if SECRET_NAME_RE.search(a) and re.search(
                    r"(?i)(echo|printf|cat)\s+.*(>|>>)", a):
                f("SECRET-TO-FILE", "INFO", ln,
                  "credential-looking value written to a file in RUN "
                  "— persists in layer history")
        elif op in ("ENV", "ARG"):
            # ENV KEY=value  /  ENV KEY value  /  ARG KEY=value
            pairs = re.findall(r"([A-Za-z_][A-Za-z0-9_]*)=", arg)
            single = re.match(r"^([A-Za-z_][A-Za-z0-9_]*)\s+\S", arg)
            bare = re.match(r"^([A-Za-z_][A-Za-z0-9_]*)$", arg.strip())
            names = set(pairs)
            if single:
                names.add(single.group(1))
            if bare:
                names.add(bare.group(1))
            for n in names:
                if SECRET_NAME_RE.search(n):
                    sev = "HIGH" if "=" in arg else "MEDIUM"
                    f("SECRET-IN-" + op, sev, ln,
                      "%s %s — credential name with a literal/default "
                      "value lands in image metadata; use --secret "
                      "mounts or runtime env" % (op, n))
        elif op == "ADD":
            if URL_RE.search(arg):
                f("ADD-REMOTE", "MEDIUM", ln,
                  "ADD fetches a remote URL — no integrity check; "
                  "prefer fetch+verify in RUN")
            elif re.match(r"^\s*(\.|\./)\s", arg):
                f("BROAD-COPY", "LOW", ln,
                  "ADD . copies whole build context — .env/.git can "
                  "leak into layers")
        elif op == "COPY":
            src = arg.split("--")[0] if "--" in arg else arg
            srcs = [t for t in src.split() if not t.startswith("-")]
            # last token is the dest; check sources
            for s in srcs[:-1]:
                if s in (".", "./", "*"):
                    f("BROAD-COPY", "LOW", ln,
                      "COPY . copies whole build context — .env/.git "
                      "can leak into layers")
                    break
                if re.match(r"(?i)\.env", os.path.basename(s)):
                    f("COPY-ENV", "HIGH", ln,
                      "COPY of %s — dotenv file into image; secrets "
                      "persist in layers" % s)
        elif op == "EXPOSE":
            ports = re.findall(r"\d+", arg)
            if "22" in ports:
                f("EXPOSE-SSH", "LOW", ln,
                  "EXPOSE 22 — sshd in a container is usually wrong")
        elif op == "HEALTHCHECK":
            cur["healthcheck"] = True

    if cur["from"] is not None:
        stages.append(cur)

    final = stages[-1]
    if final["user"] is None and not any_user_nonroot:
        findings.append(Finding(
            "NO-USER", "MEDIUM", rel, 0,
            "no USER instruction — final stage runs as root"))
    if not final["healthcheck"]:
        findings.append(Finding(
            "NO-HEALTHCHECK", "INFO", rel, 0,
            "no HEALTHCHECK — orchestrators cannot detect a wedged "
            "container"))
    return findings


def scan_compose(path, rel):
    findings = []
    raw = open(path, encoding="utf-8", errors="replace").read()

    def f(code, sev, ln, msg):
        findings.append(Finding(code, sev, rel, ln, msg))

    for i, line in enumerate(raw.splitlines(), 1):
        s = line.strip()
        if is_comment_or_blank(s):
            continue
        if re.match(r"privileged\s*:\s*true", s, re.I):
            f("PRIVILEGED", "HIGH", i,
              "privileged: true — container has host-equivalent access")
        if re.match(r"network_mode\s*:\s*['\"]?host", s, re.I):
            f("NET-HOST", "MEDIUM", i,
              "network_mode: host — shares the host network namespace")
        if re.match(r"(pid|ipc)\s*:\s*['\"]?host", s, re.I):
            f("NS-HOST", "MEDIUM", i,
              "%s shares a host namespace" % s.split(":")[0])
        if re.match(r"-?\s*cap_add", s, re.I) or re.match(
                r"cap_add\s*:", s, re.I):
            pass  # handled via list items below
        if re.search(r"(?i)(SYS_ADMIN|ALL)\s*$", s) and (
                "cap" in s.lower() or s.startswith("-")):
            f("CAP-ADMIN", "HIGH", i,
              "cap_add %s — near-root capability grant" % s.lstrip("- "))
        if re.match(r"-\s+['\"]?(SYS_ADMIN|ALL)['\"]?\s*$", s):
            f("CAP-ADMIN", "HIGH", i,
              "cap_add %s — near-root capability grant" % s.lstrip("- '\""))
        if "docker.sock" in s:
            f("DOCKER-SOCK", "HIGH", i,
              "docker.sock mount — container controls the host "
              "docker daemon (root-equivalent)")
        if re.match(r"image\s*:\s*(\S+)", s, re.I):
            ref = re.match(r"image\s*:\s*['\"]?(\S+?)['\"]?\s*$", s).group(1)
            if "@sha256:" not in ref:
                tag = ref.split(":")[-1] if ":" in ref.split("/")[-1] \
                    else "latest"
                if tag == "latest":
                    f("IMAGE-LATEST", "LOW", i,
                      "image %s — mutable tag; pin a digest" % ref)
        if re.match(r"(ports|expose)\s*:", s, re.I):
            pass
        m = re.match(r"-\s+['\"]?(\d+):(\d+)", s)
        if m:
            f("PORT-PUBLISH", "INFO", i,
              "published port %s->%s — confirm the bind address is "
              "intended (default 0.0.0.0)" % m.groups())
        m = re.match(r"-\s+['\"]?([A-Za-z_][A-Za-z0-9_]*)=(.+)", s)
        if m and SECRET_NAME_RE.search(m.group(1)):
            if not re.match(r"[\$?{]", m.group(2).strip()):
                f("SECRET-IN-ENV", "HIGH", i,
                  "literal credential value in environment: %s"
                  % m.group(1))
        m2 = re.match(r"([A-Za-z_][A-Za-z0-9_]*)\s*:\s*(.+)", s)
        if m2 and SECRET_NAME_RE.search(m2.group(1)) and not s.startswith("-"):
            v = m2.group(2).strip().strip("'\"")
            if v and not v.startswith("$"):
                f("SECRET-IN-ENV", "HIGH", i,
                  "literal credential value in %s" % m2.group(1))
        if re.match(r"env_file\s*:", s, re.I) or \
                re.match(r"-\s+['\"]?\.env", s):
            f("ENV-FILE", "INFO", i,
              "env_file reference — verify the file is not committed "
              "and holds no production secrets")
    return findings


def find_targets(root):
    if os.path.isfile(root):
        return [root]
    hits = []
    for dp, dn, fn in os.walk(root):
        dn[:] = [d for d in dn if d not in
                 (".git", "node_modules", "vendor", "__pycache__")]
        for name in fn:
            if DOCKERFILE_NAME_RE.search(name) or \
                    COMPOSE_NAME_RE.search(name):
                hits.append(os.path.join(dp, name))
    return sorted(hits)


def render(findings, targets):
    out = ["# Docker / compose security pre-scan", ""]
    if not targets:
        out.append("_No Dockerfile or compose files found._")
        return "\n".join(out)
    out.append("Files scanned: %d" % len(targets))
    out.append("")
    order = sorted(findings, key=lambda x: (sev_rank(x.sev), x.file, x.line))
    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: a clean report is not proof an "
               "image is safe; image CVEs and runtime behavior are out "
               "of scope._")
    return "\n".join(out)


SELFTEST = [
    ("FROM ubuntu:latest\nRUN apt-get update\n",
     ["BASE-LATEST", "NO-USER"]),
    ("FROM python@sha256:" + "a" * 64 + "\nUSER app\nHEALTHCHECK CMD true\n",
     []),
    ("FROM alpine\nUSER root\n", ["USER-ROOT", "NO-HEALTHCHECK"]),
    ("FROM alpine:3.20\nUSER 1000\nHEALTHCHECK CMD true\n",
     ["BASE-TAGPIN"]),
    ("FROM alpine\nRUN curl https://x.sh | bash\nUSER app\n",
     ["PIPE-SHELL"]),
    ("FROM alpine\nRUN wget https://x.tgz && sha256sum -c x.sha && tar xf x.tgz\nUSER app\n",
     []),
    ("FROM alpine\nRUN wget https://x.tgz && tar xf x.tgz\nUSER app\n",
     ["FETCH-NOCHECKSUM"]),
    ("FROM alpine\nENV DB_PASSWORD=hunter2\nUSER app\n", ["SECRET-IN-ENV"]),
    ("FROM alpine\nARG API_KEY\nUSER app\n", ["SECRET-IN-ARG"]),
    ("FROM alpine\nADD https://x/y /opt/y\nUSER app\n", ["ADD-REMOTE"]),
    ("FROM alpine\nCOPY . /app\nUSER app\n", ["BROAD-COPY"]),
    ("FROM alpine\nCOPY .env /app/.env\nUSER app\n", ["COPY-ENV"]),
    ("FROM alpine\nRUN chmod 777 /data\nUSER app\n", ["CHMOD-777"]),
    ("FROM alpine\nUSER app\nEXPOSE 22\n", ["EXPOSE-SSH"]),
]


def run_selftest():
    import tempfile
    fails = []
    for i, (src, expect) in enumerate(SELFTEST):
        with tempfile.TemporaryDirectory() as td:
            p = os.path.join(td, "Dockerfile")
            open(p, "w").write(src)
            got = {f.code for f in scan_dockerfile(p, "Dockerfile")}
            exp = set(expect)
            missing = exp - got
            extra = got - exp - {"NO-USER", "NO-HEALTHCHECK",
                                 "BASE-TAGPIN", "BASE-LATEST",
                                 "FROM-ARGPIN"}
            if missing or extra:
                fails.append((i, "missing=%s extra=%s" % (missing, extra)))
    # compose vectors
    comp = """services:
  web:
    image: nginx:latest
    privileged: true
    network_mode: host
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock
    cap_add:
      - SYS_ADMIN
    environment:
      - DB_PASSWORD=hunter2
      - TOKEN=${TOKEN}
    ports:
      - "8080:80"
"""
    with tempfile.TemporaryDirectory() as td:
        p = os.path.join(td, "docker-compose.yml")
        open(p, "w").write(comp)
        got = {f.code for f in scan_compose(p, "docker-compose.yml")}
        exp = {"PRIVILEGED", "NET-HOST", "DOCKER-SOCK", "CAP-ADMIN",
               "SECRET-IN-ENV", "IMAGE-LATEST", "PORT-PUBLISH"}
        missing = exp - got
        if missing:
            fails.append(("compose", "missing=%s" % missing))
        # literal secret should fire once, ${TOKEN} must not
        sec_lines = [f for f in scan_compose(p, "c.yml")
                     if f.code == "SECRET-IN-ENV"]
        if len(sec_lines) != 1:
            fails.append(("compose-secret-count",
                          "got %d" % len(sec_lines)))
    if fails:
        for f_ in fails:
            print("SELFTEST FAIL", f_)
        return 1
    print("SELFTEST PASS (%d vectors)" % (len(SELFTEST) + 2))
    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 COMPOSE_NAME_RE.search(os.path.basename(p)):
            findings += scan_compose(p, rel)
        else:
            findings += scan_dockerfile(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())
