#!/usr/bin/env python3
"""
iam_agent_audit.py - static audit of IAM policy JSON for AI agent roles.

Companion script to "5 IAM Mistakes I Made Giving an AI Agent Access to My
AWS Account" (Part 1) and its continuation (Part 2).

Stdlib only. Reads local JSON files and never calls AWS.

Check IDs
M1 Action "*" grants every action (Part 1, mistakes 1 and 2)
M2 Wildcard action: "svc:*" (HIGH) or "svc:Prefix*" (MEDIUM) (Part 1, mistake 2)
M3 Resource "*" on actions that support resource scoping (Part 1, mistake 3)
M4 Allows IAM writes or PassRole an agent should not hold (Part 1, mistakes 1 and 4)
M5 Boundary does not explicitly deny high-risk IAM actions,
or no boundary was supplied (Part 1, mistake 5)
S3LIST s3:ListBucket not limited by an s3:prefix condition (Part 2, gap 1)

Usage
python3 iam_agent_audit.py POLICY.json [POLICY.json ...] [--boundary BOUNDARY.json]

Exit codes
0 no CRITICAL or HIGH findings
1 at least one CRITICAL or HIGH finding
2 bad input

Simplifications
- Boundary Deny statements are matched on Action only. Resource and Condition
are ignored, so a conditional Deny counts as a full Deny. Review boundaries
by hand as well.
- M5 requires an explicit Deny even when a boundary's Allow list already
excludes IAM. That is deliberate: an explicit Deny keeps blocking IAM if
someone later widens the Allow list.
- Does not evaluate SCPs, resource policies, or session policies.
- Access key age and CloudTrail review need live account data and are out of scope.
"""

import argparse
import fnmatch
import json
import sys

SEVERITY_ORDER = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "INFO": 4}

# IAM actions an AI agent role should not hold. Extend as needed.
HIGH_RISK_IAM = [
"iam:CreateUser",
"iam:CreateRole",
"iam:CreateAccessKey",
"iam:CreateLoginProfile",
"iam:UpdateLoginProfile",
"iam:AttachRolePolicy",
"iam:AttachUserPolicy",
"iam:PutRolePolicy",
"iam:PutUserPolicy",
"iam:UpdateAssumeRolePolicy",
"iam:PassRole",
]

def die(message):
print(f"error: {message}", file=sys.stderr)
sys.exit(2)

def as_list(value):
if value is None:
return []
return value if isinstance(value, list) else [value]

def action_matches(pattern, action):
"""True if the concrete action name is covered by pattern (supports *)."""
return fnmatch.fnmatchcase(action.lower(), pattern.lower())

def load_policy(path):
try:
with open(path, encoding="utf-8") as fh:
doc = json.load(fh)
except (OSError, json.JSONDecodeError) as exc:
die(f"cannot read {path}: {exc}")
if not isinstance(doc, dict) or "Statement" not in doc:
die(f"{path} has no Statement key; not an IAM policy document")
return doc

def statements(doc):
return [s for s in as_list(doc.get("Statement")) if isinstance(s, dict)]

def is_read_only(action):
verb = action.partition(":")[2].lower()
return verb.startswith(("list", "describe"))

def has_prefix_condition(stmt):
conditions = stmt.get("Condition") or {}
for operators in conditions.values():
if isinstance(operators, dict) and any(k.lower() == "s3:prefix" for k in operators):
return True
return False

def denied_actions(boundary):
denied = []
for stmt in statements(boundary):
if stmt.get("Effect") == "Deny":
denied.extend(a for a in as_list(stmt.get("Action")) if isinstance(a, str))
return denied

def audit_policy(doc, boundary=None):
"""Return a list of (severity, check_id, statement_index_or_None, message)."""
findings = []

def add(severity, check, index, message):
findings.append((severity, check, index, message))

for idx, stmt in enumerate(statements(doc)):
if stmt.get("Effect") != "Allow":
continue

actions = [a for a in as_list(stmt.get("Action")) if isinstance(a, str)]
resources = as_list(stmt.get("Resource"))

# M1 and M2: admin and wildcard actions
for action in actions:
if action == "*":
add("CRITICAL", "M1", idx, 'Action "*" grants every action on every service.')
elif "*" in action:
service, _, name = action.partition(":")
if name == "*" or "*" in service:
add("HIGH", "M2", idx,
f'"{action}" is a full wildcard; list the exact actions needed.')
else:
add("MEDIUM", "M2", idx,
f'"{action}" is a prefix wildcard; list the exact actions needed.')

# M4: IAM writes and PassRole granted to this statement
covered = sorted({r for a in actions for r in HIGH_RISK_IAM if action_matches(a, r)})
if covered:
add("CRITICAL", "M4", idx,
"Allows IAM actions an agent should not hold: " + ", ".join(covered) + ".")

# M3: Resource "*"
if "*" in resources:
scoped = [a for a in actions if not is_read_only(a)]
if scoped:
add("MEDIUM", "M3", idx,
'Resource "*" on actions that support resource scoping: ' + ", ".join(scoped) + ".")
elif actions:
add("INFO", "M3", idx,
'Resource "*" on List/Describe actions only; keep these in a read-only group.')

# S3LIST: ListBucket without a prefix condition
if any(action_matches(a, "s3:ListBucket") for a in actions) and not has_prefix_condition(stmt):
add("MEDIUM", "S3LIST", idx,
"s3:ListBucket is not limited by an s3:prefix condition, "
"so any key in the bucket can be listed.")

# M5: permission boundary
if boundary is None:
add("MEDIUM", "M5", None,
"No permission boundary supplied; boundary coverage was not checked.")
else:
denied = denied_actions(boundary)
missing = [r for r in HIGH_RISK_IAM if not any(action_matches(d, r) for d in denied)]
if missing:
add("HIGH", "M5", None,
"Permission boundary does not explicitly deny: " + ", ".join(missing) + ".")

return findings

def main(argv=None):
parser = argparse.ArgumentParser(
description="Audit IAM policy JSON for over-permissive AI agent roles.")
parser.add_argument("policies", nargs="+", help="IAM policy JSON files to audit")
parser.add_argument("--boundary", help="permission boundary JSON to check each policy against")
args = parser.parse_args(argv)

boundary = load_policy(args.boundary) if args.boundary else None

failed = False
counts = {}
for path in args.policies:
findings = audit_policy(load_policy(path), boundary)
findings.sort(key=lambda f: SEVERITY_ORDER[f[0]])

print(f"== {path} ==")
if not findings:
print(" no findings")
for severity, check, index, message in findings:
where = f"Statement {index}: " if index is not None else ""
print(f" [{severity:<8}] {check:<6} {where}{message}")
counts[severity] = counts.get(severity, 0) + 1
if severity in ("CRITICAL", "HIGH"):
failed = True
print()

summary = ", ".join(f"{counts[s]} {s}" for s in SEVERITY_ORDER if s in counts) or "no findings"
print(f"Summary: {summary}")
return 1 if failed else 0

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