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.
SCP diffs run locally. A new or renamed SCP is compared with an empty
baseline, and a base SCP missing from head is treated as deleted. The diff
fails when access widens: a new Allow, a smaller Deny, a larger Deny
NotAction list, or any Condition change. When AWS credentials can call
sts:GetCallerIdentity, each document is sent to IAM Access Analyzer
ValidatePolicy. A new SCP with a ValidatePolicy ERROR fails, because that
error cannot already exist on main. CheckNoNewAccess compares bootstrap
identity policies with the templates in --base-repo. It rejects
SERVICE_CONTROL_POLICY, so SCPs use the local diff. Missing credentials
skip the AWS calls, emit a warning, and still run the local checks.

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(root: Path | None = None) -> dict[str, dict]:
    bootstrap = (root or ROOT) / "lib" / "hcptf-bootstrap"
    if not bootstrap.is_dir():
        fail(f"bootstrap template directory does not exist: {bootstrap}")
    documents: dict[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[path.name] = load_json(path)
    if not documents:
        fail(f"no bootstrap policy templates in {bootstrap}")
    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(
            "::warning title=Access Analyzer skipped::sts get-caller-identity failed. "
            "ValidatePolicy and CheckNoNewAccess did not run."
        )
        print("AWS checks skipped: sts get-caller-identity failed", file=sys.stderr)
        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 statement_list(document: dict) -> list[dict]:
    statements = document.get("Statement", [])
    if isinstance(statements, dict):
        return [statements]
    return list(statements)


def normalize(value: object) -> object:
    if isinstance(value, dict):
        return {key: normalize(value[key]) for key in sorted(value)}
    if isinstance(value, list):
        items = [normalize(item) for item in value]
        return sorted(items, key=lambda item: json.dumps(item, sort_keys=True, default=str))
    return value


def stable(value: object) -> str:
    return json.dumps(normalize(value), sort_keys=True, separators=(",", ":"), default=str)


def name_set(statement: dict, field: str) -> set[str]:
    return set(as_list(statement.get(field)))


def drop_exact(old: list[dict], new: list[dict]) -> tuple[list[dict], list[dict]]:
    used = [False] * len(new)
    new_keys = [stable(item) for item in new]
    remaining_old: list[dict] = []
    for statement in old:
        key = stable(statement)
        matched = False
        for index, candidate in enumerate(new_keys):
            if not used[index] and candidate == key:
                used[index] = True
                matched = True
                break
        if not matched:
            remaining_old.append(statement)
    remaining_new = [item for index, item in enumerate(new) if not used[index]]
    return remaining_old, remaining_new


def pair_widenings(old: dict, new: dict) -> list[str]:
    """Ways a matched statement grants more access than it used to."""
    reasons: list[str] = []
    label = str(new.get("Sid") or old.get("Sid") or "statement")
    if old.get("Effect") != new.get("Effect"):
        return [f"{label}: effect changed"]
    if stable(old.get("Condition")) != stable(new.get("Condition")):
        reasons.append(f"{label}: condition changed")
    if stable(old.get("Principal")) != stable(new.get("Principal")):
        reasons.append(f"{label}: principal changed")
    if ("Action" in old) != ("Action" in new) or ("NotAction" in old) != ("NotAction" in new):
        reasons.append(f"{label}: action form changed")
        return reasons
    if ("Resource" in old) != ("Resource" in new) or ("NotResource" in old) != ("NotResource" in new):
        reasons.append(f"{label}: resource form changed")
        return reasons
    effect = old.get("Effect")
    old_actions, new_actions = name_set(old, "Action"), name_set(new, "Action")
    old_not, new_not = name_set(old, "NotAction"), name_set(new, "NotAction")
    old_resources, new_resources = name_set(old, "Resource"), name_set(new, "Resource")
    old_not_resources = name_set(old, "NotResource")
    new_not_resources = name_set(new, "NotResource")
    if effect == "Allow":
        added = sorted(new_actions - old_actions)
        if added:
            reasons.append(f"{label}: allow actions added {added}")
        removed_exceptions = sorted(old_not - new_not)
        if removed_exceptions:
            reasons.append(f"{label}: allow NotAction shrank {removed_exceptions}")
        if new_resources - old_resources:
            reasons.append(f"{label}: allow resources added")
        if old_not_resources - new_not_resources:
            reasons.append(f"{label}: allow NotResource shrank")
    elif effect == "Deny":
        removed = sorted(old_actions - new_actions)
        if removed:
            reasons.append(f"{label}: deny actions removed {removed}")
        grown = sorted(new_not - old_not)
        if grown:
            reasons.append(f"{label}: deny NotAction grew {grown}")
        if old_resources - new_resources:
            reasons.append(f"{label}: deny resources removed")
        if new_not_resources - old_not_resources:
            reasons.append(f"{label}: deny NotResource grew")
    return reasons


def scp_widenings(existing: dict, new: dict) -> list[str]:
    """Access added relative to existing.

    CheckNoNewAccess rejects SERVICE_CONTROL_POLICY. A new Allow, a removed
    or smaller Deny, a larger Deny NotAction list, or any Condition change
    is new access. A new Deny, or a larger Deny Action list, is not.
    """
    reasons: list[str] = []

    def grouped(document: dict) -> tuple[dict[str, list[dict]], list[dict]]:
        keyed: dict[str, list[dict]] = {}
        loose: list[dict] = []
        for statement in statement_list(document):
            sid = statement.get("Sid")
            if isinstance(sid, str) and sid:
                keyed.setdefault(sid, []).append(statement)
            else:
                loose.append(statement)
        return keyed, loose

    old_keyed, old_loose = grouped(existing)
    new_keyed, new_loose = grouped(new)
    for sid in sorted(set(old_keyed) | set(new_keyed)):
        old_group, new_group = drop_exact(old_keyed.get(sid, []), new_keyed.get(sid, []))
        count = min(len(old_group), len(new_group))
        for old, statement in zip(old_group[:count], new_group[:count]):
            reasons.extend(pair_widenings(old, statement))
        for statement in old_group[count:]:
            if statement.get("Effect") == "Deny":
                reasons.append(f"{sid}: deny statement removed")
        for statement in new_group[count:]:
            if statement.get("Effect") == "Allow":
                reasons.append(f"{sid}: allow statement added")
    old_loose, new_loose = drop_exact(old_loose, new_loose)
    for statement in old_loose:
        if statement.get("Effect") == "Deny":
            reasons.append("deny statement removed")
    for statement in new_loose:
        if statement.get("Effect") == "Allow":
            reasons.append("allow statement added")
    return reasons


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


def expect_widening(existing: dict, new: dict, label: str) -> None:
    if not scp_widenings(existing, new):
        fail(f"SCP allow check missed {label}")


def expect_same(existing: dict, new: dict, label: str) -> None:
    reasons = scp_widenings(existing, new)
    if reasons:
        fail(f"SCP allow check flagged {label}: {reasons}")


def local_self_test() -> None:
    deny_only = {
        "Statement": [
            {"Sid": "A", "Effect": "Deny", "NotAction": ["iam:*"], "Resource": "*"}
        ]
    }
    larger_not_action = {
        "Statement": [
            {
                "Sid": "A",
                "Effect": "Deny",
                "NotAction": ["iam:*", "s3:*"],
                "Resource": "*",
            }
        ]
    }
    expect_widening(deny_only, larger_not_action, "a larger Deny NotAction list")
    expect_widening(deny_only, {"Statement": []}, "a removed Deny")
    expect_widening(
        {
            "Statement": [
                {
                    "Sid": "A",
                    "Effect": "Deny",
                    "Action": ["iam:CreateRole", "iam:DeleteRole"],
                    "Resource": "*",
                }
            ]
        },
        {
            "Statement": [
                {"Sid": "A", "Effect": "Deny", "Action": "iam:CreateRole", "Resource": "*"}
            ]
        },
        "a smaller Deny Action list",
    )
    expect_widening(
        {
            "Statement": [
                {
                    "Sid": "A",
                    "Effect": "Allow",
                    "Action": "s3:GetObject",
                    "Resource": "*",
                    "Condition": {"StringEquals": {"aws:RequestedRegion": "home"}},
                }
            ]
        },
        {
            "Statement": [
                {"Sid": "A", "Effect": "Allow", "Action": "s3:GetObject", "Resource": "*"}
            ]
        },
        "a removed Allow Condition",
    )
    expect_widening(
        {
            "Statement": [
                {
                    "Sid": "A",
                    "Effect": "Deny",
                    "Action": "iam:CreateRole",
                    "Resource": "*",
                    "Condition": {
                        "ArnNotLike": {
                            "aws:PrincipalArn": "arn:aws:iam::*:role/OrganizationAccountAccessRole"
                        }
                    },
                }
            ]
        },
        {
            "Statement": [
                {
                    "Sid": "A",
                    "Effect": "Deny",
                    "Action": "iam:CreateRole",
                    "Resource": "*",
                    "Condition": {
                        "ArnNotLike": {"aws:PrincipalArn": "arn:aws:iam::*:role/other"}
                    },
                }
            ]
        },
        "a changed Deny Condition",
    )
    expect_widening(
        {"Statement": []},
        {
            "Statement": [
                {"Sid": "Wide", "Effect": "Allow", "Action": "*", "Resource": "*"}
            ]
        },
        "a new SCP Allow *",
    )
    expect_widening(
        {
            "Statement": [
                {"Sid": "A", "Effect": "Allow", "Action": "s3:GetObject", "Resource": "*"}
            ]
        },
        {
            "Statement": [
                {"Sid": "A", "Effect": "Allow", "NotAction": "iam:CreateRole", "Resource": "*"}
            ]
        },
        "an Allow rewritten as NotAction",
    )
    expect_same(deny_only, deny_only, "an identical document")
    expect_same(
        {
            "Statement": [
                {"Sid": "A", "Effect": "Deny", "Action": "iam:CreateRole", "Resource": "*"}
            ]
        },
        {
            "Statement": [
                {
                    "Sid": "A",
                    "Effect": "Deny",
                    "Action": ["iam:CreateRole", "iam:DeleteRole"],
                    "Resource": "*",
                }
            ]
        },
        "a larger Deny Action list",
    )
    plan = load_json(BOOTSTRAP / "plan-refresh-policy.json.tmpl")
    widened_plan = json.loads(json.dumps(plan))
    widened_plan["Statement"].append(
        {"Sid": "Wide", "Effect": "Allow", "Action": "s3:PutObject", "Resource": "*"}
    )
    try:
        compare_documents(
            "plan-refresh-policy.json.tmpl",
            "IDENTITY_POLICY",
            plan,
            widened_plan,
            checker=lambda _existing, _new, _policy_type: "FAIL",
        )
    except CheckFailure as exc:
        if "CheckNoNewAccess" not in str(exc):
            raise
    else:
        fail("widened plan-refresh template was accepted")
    compare_documents(
        "plan-refresh-policy.json.tmpl",
        "IDENTITY_POLICY",
        plan,
        plan,
        checker=lambda _existing, _new, _policy_type: "PASS",
    )
    print("self-test ok: SCP widenings and plan-refresh comparison")


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")
    plan = load_json(BOOTSTRAP / "plan-refresh-policy.json.tmpl")
    widened_plan = json.loads(json.dumps(plan))
    widened_plan["Statement"].append(
        {"Effect": "Allow", "Action": "s3:PutObject", "Resource": "*"}
    )
    if check_no_new_access(plan, widened_plan, "IDENTITY_POLICY") != "FAIL":
        fail("self-test expected FAIL when plan refresh gains s3:PutObject")
    print("self-test ok: CheckNoNewAccess distinguishes added access")


def validate_document(
    name: str,
    policy_type: str,
    document: dict,
    base_documents: dict[str, dict],
    require_identity_base: bool,
) -> None:
    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})")
        if policy_type == "IDENTITY_POLICY":
            compare_documents(name, policy_type, base_documents[name], document)
        return
    if head_errors:
        fail(f"ValidatePolicy ERROR {name}: {describe_findings(head_errors)}")
    if policy_type == "IDENTITY_POLICY" and require_identity_base:
        fail(f"new bootstrap policy has no base document: {name}")
    print(f"ValidatePolicy ok: {name} ({policy_type})")


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

    try:
        check_trust_templates()
        check_plan_refresh()
        local_self_test()
        documents: list[tuple[str, str, dict]] = [
            (name, "IDENTITY_POLICY", document)
            for name, document in bootstrap_documents().items()
        ]
        base_documents: dict[str, dict] = {}
        if args.base_repo:
            base_documents.update(bootstrap_documents(args.base_repo))
        head_scps: dict[str, dict] = {}
        if args.cdk_out:
            head_scps = synthesized_scps(args.cdk_out)
            for name, document in head_scps.items():
                documents.append((name, "SERVICE_CONTROL_POLICY", document))
        base_scps: dict[str, dict] = {}
        if args.base_cdk_out:
            base_scps = synthesized_scps(args.base_cdk_out)
            base_documents.update(base_scps)
        empty: dict = {"Statement": []}
        for name, document in head_scps.items():
            compare_documents(
                name,
                "SERVICE_CONTROL_POLICY",
                base_scps.get(name, empty),
                document,
            )
        for name in sorted(set(base_scps) - set(head_scps)):
            compare_documents(name, "SERVICE_CONTROL_POLICY", base_scps[name], empty)
        if aws_available():
            if args.self_test:
                self_test()
            for name, policy_type, document in documents:
                validate_document(
                    name,
                    policy_type,
                    document,
                    base_documents,
                    require_identity_base=args.base_repo is not None,
                )
        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())