seahaven-org-baseline/scripts/check_iam_policies.py

#!/usr/bin/env python3
"""IAM policy checks for seahaven-org-baseline.

Local invariants always run:
- bootstrap trust templates use StringEquals and do not use StringLike
- the plan refresh template grants no lambda write, including lambda:*

When --cdk-out is set, synthesized service control policies are collected.
When AWS credentials can call sts:GetCallerIdentity, each document is sent to
IAM Access Analyzer ValidatePolicy. CheckNoNewAccess compares identity
policies to the base document. It rejects SERVICE_CONTROL_POLICY and any
document with no Allow, so SCP diffs compare Allow actions and Deny
NotAction lists instead. Missing credentials skip the AWS calls and still
pass the local invariants.

The substrate template is not scanned. Its afi plan role still has lambda:*
until the import apply replaces it.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
BOOTSTRAP = ROOT / "lib" / "hcptf-bootstrap"
PLACEHOLDERS = {
    "__ACCOUNT_ID__": "111111111111",
    "__HCP_PROJECT__": "seahaven-prod",
    "__BOOTSTRAP_WORKSPACE__": "iam-bootstrap-prod",
    "__STACK_PREFIX__": "example",
    "__STACK_NAME__": "example",
}
LAMBDA_READ_PREFIXES = ("lambda:Get", "lambda:List", "lambda:Describe")


class CheckFailure(Exception):
    pass


def fail(message: str) -> None:
    raise CheckFailure(message)


def substitute(text: str) -> str:
    for key, value in PLACEHOLDERS.items():
        text = text.replace(key, value)
    return text


def load_json(path: Path) -> dict:
    return json.loads(substitute(path.read_text()))


def statement_actions(statement: dict) -> list[str]:
    action = statement.get("Action", [])
    if isinstance(action, str):
        return [action]
    return list(action)


def lambda_writes(document: dict) -> list[str]:
    found: list[str] = []
    statements = document.get("Statement", [])
    if isinstance(statements, dict):
        statements = [statements]
    for statement in statements:
        if statement.get("Effect") != "Allow":
            continue
        for action in statement_actions(statement):
            if action in {"*", "lambda:*"}:
                found.append(action)
            elif action.startswith("lambda:") and not action.startswith(LAMBDA_READ_PREFIXES):
                found.append(action)
    return found


def check_trust_templates() -> None:
    paths = sorted(BOOTSTRAP.glob("trust-*.json.tmpl"))
    if not paths:
        fail(f"no trust templates in {BOOTSTRAP}")
    for path in paths:
        text = path.read_text()
        if "StringLike" in text:
            fail(f"{path.name}: trust must stay StringEquals")
        if "StringEquals" not in text:
            fail(f"{path.name}: missing StringEquals")
        document = load_json(path)
        condition = document["Statement"][0]["Condition"]
        if "StringLike" in condition or "StringEquals" not in condition:
            fail(f"{path.name}: trust condition must be StringEquals")
        print(f"invariant ok: {path.name} uses StringEquals")


def check_plan_refresh() -> None:
    path = BOOTSTRAP / "plan-refresh-policy.json.tmpl"
    writes = lambda_writes(load_json(path))
    if writes:
        fail(f"{path.name}: plan document grants lambda write {writes}")
    print(f"invariant ok: {path.name} has no lambda write")


def bootstrap_documents() -> list[tuple[str, str, dict]]:
    documents: list[tuple[str, str, dict]] = []
    for path in sorted(BOOTSTRAP.glob("*.json.tmpl")):
        # Trust documents are resource policies without a Resource element.
        # ValidatePolicy rejects that shape. The StringEquals invariant covers them.
        if path.name.startswith("trust-"):
            continue
        documents.append((path.name, "IDENTITY_POLICY", load_json(path)))
    return documents


def synthesized_scps(cdk_out: Path) -> dict[str, dict]:
    found: dict[str, dict] = {}
    if not cdk_out.is_dir():
        fail(f"cdk out directory does not exist: {cdk_out}")
    for path in sorted(cdk_out.glob("*.template.json")):
        template = json.loads(path.read_text())
        for logical, resource in template.get("Resources", {}).items():
            if resource.get("Type") != "AWS::Organizations::Policy":
                continue
            content = resource["Properties"]["Content"]
            if isinstance(content, str):
                content = json.loads(content)
            name = resource["Properties"].get("Name", logical)
            found[f"{path.stem}:{name}"] = content
    if not found:
        fail(f"no service control policies in {cdk_out}")
    return found


def aws_available() -> bool:
    result = subprocess.run(
        ["aws", "sts", "get-caller-identity", "--output", "json"],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        print("AWS checks skipped: sts get-caller-identity failed")
        return False
    identity = json.loads(result.stdout)
    print(f"AWS checks using account {identity.get('Account')}")
    return True


def analyzer(command: list[str], document: dict, extra: list[str]) -> dict:
    with tempfile.NamedTemporaryFile("w", suffix=".json") as handle:
        json.dump(document, handle)
        handle.flush()
        result = subprocess.run(
            ["aws", "accessanalyzer", *command, "--policy-document", f"file://{handle.name}", *extra, "--output", "json"],
            capture_output=True,
            text=True,
        )
    if result.returncode != 0:
        fail(f"{' '.join(command)} failed: {result.stderr.strip()}")
    return json.loads(result.stdout or "{}")


def error_findings(name: str, policy_type: str, document: dict) -> list[tuple[str, str]]:
    payload = analyzer(
        ["validate-policy"],
        document,
        ["--policy-type", policy_type],
    )
    errors = [
        (
            str(finding.get("issueCode")),
            str(finding.get("findingDetails")),
        )
        for finding in payload.get("findings", [])
        if finding.get("findingType") == "ERROR"
    ]
    return errors


def describe_findings(findings: list[tuple[str, str]]) -> str:
    return "; ".join(f"{code}: {details}" for code, details in findings)


def check_no_new_access(existing: dict, new: dict, policy_type: str) -> str:
    with tempfile.NamedTemporaryFile("w", suffix=".json") as new_file:
        json.dump(new, new_file)
        new_file.flush()
        with tempfile.NamedTemporaryFile("w", suffix=".json") as existing_file:
            json.dump(existing, existing_file)
            existing_file.flush()
            result = subprocess.run(
                [
                    "aws",
                    "accessanalyzer",
                    "check-no-new-access",
                    "--policy-type",
                    policy_type,
                    "--existing-policy-document",
                    f"file://{existing_file.name}",
                    "--new-policy-document",
                    f"file://{new_file.name}",
                    "--output",
                    "json",
                ],
                capture_output=True,
                text=True,
            )
    if result.returncode != 0:
        fail(f"CheckNoNewAccess failed: {result.stderr.strip()}")
    payload = json.loads(result.stdout or "{}")
    return payload.get("result", "")


def as_list(value: object) -> list[str]:
    if value is None:
        return []
    if isinstance(value, str):
        return [value]
    return [str(item) for item in value]


def allow_surface(document: dict) -> set[tuple[str, ...]]:
    """Actions a policy newly permits.

    CheckNoNewAccess rejects SERVICE_CONTROL_POLICY and any document with no
    Allow statement. For those documents, new access is an added Allow or a
    larger Deny NotAction list (the deny then skips more actions).
    """
    surface: set[tuple[str, ...]] = set()
    statements = document.get("Statement", [])
    if isinstance(statements, dict):
        statements = [statements]
    for statement in statements:
        sid = str(statement.get("Sid", ""))
        effect = statement.get("Effect")
        if effect == "Allow":
            for action in statement_actions(statement):
                for resource in as_list(statement.get("Resource")) or ["*"]:
                    surface.add(("allow", sid, action, resource))
        elif effect == "Deny" and "NotAction" in statement:
            for action in as_list(statement.get("NotAction")):
                surface.add(("not-action", sid, action))
    return surface


def compare_documents(name: str, policy_type: str, existing: dict, new: dict) -> None:
    if policy_type != "IDENTITY_POLICY":
        added = sorted(allow_surface(new) - allow_surface(existing))
        if added:
            fail(f"SCP allows new access {name}: {added[:8]}")
        print(f"SCP allow check ok: {name}")
        return
    result = check_no_new_access(existing, new, policy_type)
    if result != "PASS":
        fail(f"CheckNoNewAccess {name}: {result or 'empty result'}")
    print(f"CheckNoNewAccess ok: {name}")


def self_test() -> None:
    existing = {
        "Version": "2012-10-17",
        "Statement": [
            {"Effect": "Allow", "Action": "s3:GetObject", "Resource": "*"}
        ],
    }
    widened = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": ["s3:GetObject", "s3:PutObject"],
                "Resource": "*",
            }
        ],
    }
    if check_no_new_access(existing, widened, "IDENTITY_POLICY") != "FAIL":
        fail("self-test expected FAIL when s3:PutObject is added")
    if check_no_new_access(existing, existing, "IDENTITY_POLICY") != "PASS":
        fail("self-test expected PASS for an identical policy")
    print("self-test ok: CheckNoNewAccess distinguishes added access")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--cdk-out", type=Path)
    parser.add_argument("--base-cdk-out", type=Path)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()

    try:
        check_trust_templates()
        check_plan_refresh()
        deny_only = {
            "Statement": [
                {"Sid": "A", "Effect": "Deny", "NotAction": ["iam:*"], "Resource": "*"}
            ]
        }
        widened = {
            "Statement": [
                {
                    "Sid": "A",
                    "Effect": "Deny",
                    "NotAction": ["iam:*", "s3:*"],
                    "Resource": "*",
                },
                {
                    "Sid": "B",
                    "Effect": "Allow",
                    "Action": "s3:GetObject",
                    "Resource": "*",
                },
            ]
        }
        if not (allow_surface(widened) - allow_surface(deny_only)):
            fail("SCP allow check missed a widened NotAction and a new Allow")
        if allow_surface(deny_only) - allow_surface(deny_only):
            fail("SCP allow check flagged an identical document")
        documents: list[tuple[str, str, dict]] = list(bootstrap_documents())
        base_documents: dict[str, dict] = {}
        if args.cdk_out:
            for name, document in synthesized_scps(args.cdk_out).items():
                documents.append((name, "SERVICE_CONTROL_POLICY", document))
        if args.base_cdk_out:
            base_documents = synthesized_scps(args.base_cdk_out)
        if aws_available():
            if args.self_test:
                self_test()
            for name, policy_type, document in documents:
                head_errors = error_findings(name, policy_type, document)
                if name in base_documents:
                    base_errors = set(error_findings(name, policy_type, base_documents[name]))
                    introduced = [item for item in head_errors if item not in base_errors]
                    if introduced:
                        fail(
                            f"ValidatePolicy new ERROR {name}: {describe_findings(introduced)}"
                        )
                    if head_errors:
                        print(
                            f"ValidatePolicy existing ERROR {name}: {describe_findings(head_errors)}"
                        )
                    else:
                        print(f"ValidatePolicy ok: {name} ({policy_type})")
                    compare_documents(
                        name,
                        policy_type,
                        base_documents[name],
                        document,
                    )
                elif head_errors and policy_type == "IDENTITY_POLICY":
                    fail(f"ValidatePolicy ERROR {name}: {describe_findings(head_errors)}")
                elif head_errors:
                    print(
                        f"ValidatePolicy existing ERROR {name}: {describe_findings(head_errors)}"
                    )
                else:
                    print(f"ValidatePolicy ok: {name} ({policy_type})")
        print(f"checked {len(documents)} documents")
    except CheckFailure as exc:
        print(f"FAIL: {exc}", file=sys.stderr)
        return 1
    return 0


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