#!/usr/bin/env python3
"""
Agent Readiness Check, by Abacross.

Runs READ-ONLY against the AWS account your current credentials point at and reports
whether an autonomous agent running there can spend without limit, reach further than
intended, or act without a record you could prove afterwards. Nothing is changed and
nothing leaves your machine unless you choose to send the digest at the end.

    python3 agent_readiness_check.py                 # uses your default credentials
    python3 agent_readiness_check.py --regions us-east-1,eu-west-1
    python3 agent_readiness_check.py --json report.json --html report.html

Needs boto3 (pip install boto3) and credentials that can read IAM, CloudTrail, Budgets,
Bedrock, GuardDuty, Config, CloudWatch and S3 bucket configuration. The ReadOnlyAccess
managed policy is enough. It never calls anything that writes.

Every API call this makes is listed by --calls, --policy prints an IAM policy allowing
exactly those and nothing else, and --audit-log records each call as it is made, so the
read-only claim can be checked from outside rather than taken on trust.
"""
import argparse
import base64
import datetime as dt
import json
import sys

try:
    import boto3
    from botocore.exceptions import BotoCoreError, ClientError
except ImportError:
    sys.exit("boto3 is not installed: pip install boto3")

VERSION = "0.1"
DEFAULT_REGIONS = ["us-east-1", "us-west-2"]
AGENTISH = ("agent", "bedrock", "lambda", "ecs", "task", "bot", "assistant", "copilot", "worker", "runner")
KEY_STALE_DAYS = 90

# Severity weights for the score, and the questions each check answers. The three
# questions are the three ways agents go wrong; every finding belongs to one of them.
WEIGHT = {"critical": 20, "high": 10, "medium": 5, "low": 2}
QUESTION = {
    "spend": "Can an agent here spend without a limit that actually stops it?",
    "reach": "Can an agent here reach further than it was meant to?",
    "proof": "Could you prove afterwards what an agent here did?",
}


class Finding:
    def __init__(self, id, severity, question, title, evidence, blast_radius, fix):
        self.id, self.severity, self.question = id, severity, question
        self.title, self.evidence, self.blast_radius, self.fix = title, evidence, blast_radius, fix

    def as_dict(self):
        return dict(id=self.id, severity=self.severity, question=self.question, title=self.title,
                    evidence=self.evidence, blast_radius=self.blast_radius, fix=self.fix)


def _safe(fn, *a, **kw):
    """Call an API; a permission or availability error is a note, not a crash."""
    try:
        return fn(*a, **kw), None
    except (ClientError, BotoCoreError) as e:
        code = getattr(e, "response", {}).get("Error", {}).get("Code", type(e).__name__)
        return None, code


def _age_days(when, now):
    if when is None:
        return None
    if when.tzinfo is None:
        when = when.replace(tzinfo=dt.timezone.utc)
    return (now - when).days


# ---------------------------------------------------------------- checks ---------
# Each check takes a client factory `c(service, region=None)` and returns findings plus
# a list of notes about what it could not see. Keeping them separate keeps them testable.

def check_spend(c, account_id, regions, notes):
    out = []
    budgets, err = _safe(lambda: c("budgets", "us-east-1").describe_budgets(AccountId=account_id))
    if err:
        notes.append(f"budgets: could not read ({err})")
    else:
        items = budgets.get("Budgets") or []
        if not items:
            out.append(Finding("budget.none", "critical", "spend", "No AWS Budget exists",
                               "describe_budgets returned nothing",
                               "A looping agent bills until somebody notices. There is no limit of any kind.",
                               "Create a monthly cost budget with an action that stops the agent's execution role or scales its compute to zero, not only an email."))
        else:
            with_actions = 0
            for b in items:
                acts, e2 = _safe(lambda: c("budgets", "us-east-1").describe_budget_actions_for_budget(
                    AccountId=account_id, BudgetName=b["BudgetName"]))
                if not e2 and (acts.get("Actions") or []):
                    with_actions += 1
            if with_actions == 0:
                out.append(Finding("budget.no-action", "high", "spend", "Budgets alert but do not act",
                                   f"{len(items)} budget(s), none with a budget action",
                                   "An alert at 2am is read at 9am. The seven hours between are billed.",
                                   "Attach a budget action that applies a deny policy to the agent's role or stops the service when the threshold is crossed."))
    alarms, err = _safe(lambda: c("cloudwatch", "us-east-1").describe_alarms_for_metric(
        MetricName="EstimatedCharges", Namespace="AWS/Billing", Dimensions=[{"Name": "Currency", "Value": "USD"}]))
    if err:
        notes.append(f"billing alarms: could not read ({err})")
    elif not (alarms.get("MetricAlarms") or []):
        out.append(Finding("billing.no-alarm", "medium", "spend", "No billing alarm",
                           "no CloudWatch alarm on AWS/Billing EstimatedCharges",
                           "Even the cheapest early warning is absent; the first signal is the invoice.",
                           "Enable billing alerts and create an EstimatedCharges alarm below your monthly expectation."))
    return out


def check_reach(c, now, notes):
    out = []
    iam = c("iam")
    summary, err = _safe(lambda: iam.get_account_summary())
    if err:
        notes.append(f"iam summary: could not read ({err})")
    elif summary["SummaryMap"].get("AccountMFAEnabled", 1) == 0:
        out.append(Finding("iam.root-no-mfa", "critical", "reach", "Root has no MFA",
                           "AccountMFAEnabled = 0",
                           "Whoever holds the root password holds the account, every agent in it included.",
                           "Enable a hardware or virtual MFA device on the root user today."))

    # Long-lived keys: the credential an agent is most often given, and the one that
    # outlives every rotation policy nobody wrote down.
    users, err = _safe(lambda: iam.list_users(MaxItems=200))
    if err:
        notes.append(f"iam users: could not read ({err})")
    else:
        stale, unused = [], []
        for u in users.get("Users", []):
            keys, e2 = _safe(lambda: iam.list_access_keys(UserName=u["UserName"]))
            if e2:
                continue
            for k in keys.get("AccessKeyMetadata", []):
                if k.get("Status") != "Active":
                    continue
                age = _age_days(k.get("CreateDate"), now)
                if age is not None and age > KEY_STALE_DAYS:
                    stale.append(f"{u['UserName']} ({age}d)")
                last, e3 = _safe(lambda: iam.get_access_key_last_used(AccessKeyId=k["AccessKeyId"]))
                if not e3:
                    lu = last.get("AccessKeyLastUsed", {}).get("LastUsedDate")
                    lu_age = _age_days(lu, now)
                    if lu is None or (lu_age is not None and lu_age > KEY_STALE_DAYS):
                        unused.append(u["UserName"])
        if stale:
            out.append(Finding("iam.key-stale", "high", "reach", "Long-lived access keys not rotated",
                               f"{len(stale)} active key(s) older than {KEY_STALE_DAYS} days: " + ", ".join(stale[:5]),
                               "A leaked key from months ago still works today. Agents copy keys into places you did not intend.",
                               "Rotate, then move agents to roles with short-lived credentials so there is nothing to leak."))
        if unused:
            out.append(Finding("iam.key-unused", "medium", "reach", "Active keys nobody uses",
                               f"{len(unused)} active key(s) unused for {KEY_STALE_DAYS}+ days or never: " + ", ".join(sorted(set(unused))[:5]),
                               "A live credential with no owner is a credential nobody will notice being used.",
                               "Deactivate, wait a week, delete."))

    # Wildcard power on roles that look like they belong to agents or their runtimes.
    roles, err = _safe(lambda: iam.list_roles(MaxItems=200))
    if err:
        notes.append(f"iam roles: could not read ({err})")
    else:
        wide = []
        for r in roles.get("Roles", []):
            name = r["RoleName"]
            if not any(t in name.lower() for t in AGENTISH):
                continue
            att, e2 = _safe(lambda: iam.list_attached_role_policies(RoleName=name))
            if not e2 and any(p["PolicyArn"].endswith("/AdministratorAccess") or p["PolicyArn"].endswith("/PowerUserAccess")
                              for p in att.get("AttachedPolicies", [])):
                wide.append(name)
                continue
            inl, e3 = _safe(lambda: iam.list_role_policies(RoleName=name))
            if e3:
                continue
            for pn in inl.get("PolicyNames", []):
                doc, e4 = _safe(lambda: iam.get_role_policy(RoleName=name, PolicyName=pn))
                if e4:
                    continue
                stmts = doc["PolicyDocument"].get("Statement", [])
                if isinstance(stmts, dict):
                    stmts = [stmts]
                for s in stmts:
                    acts = s.get("Action", [])
                    acts = [acts] if isinstance(acts, str) else acts
                    res = s.get("Resource", [])
                    res = [res] if isinstance(res, str) else res
                    if s.get("Effect") == "Allow" and "*" in acts and "*" in res:
                        wide.append(name)
                        break
        if wide:
            out.append(Finding("iam.wildcard-admin", "critical", "reach", "Agent-like roles with administrator power",
                               f"{len(wide)} role(s) with Action * on Resource * or AdministratorAccess: " + ", ".join(sorted(set(wide))[:5]),
                               "An agent that can do anything will, eventually, do something you did not mean. The limit of the damage is the account.",
                               "Scope each agent role to the actions and resources it uses, and put a permissions boundary on it so a helpful widening cannot exceed the ceiling."))
    return out


def check_proof(c, regions, notes):
    out = []
    ct = c("cloudtrail", regions[0])
    trails, err = _safe(lambda: ct.describe_trails(includeShadowTrails=True))
    if err:
        notes.append(f"cloudtrail: could not read ({err})")
    else:
        items = trails.get("trailList") or []
        if not items:
            out.append(Finding("trail.none", "critical", "proof", "No CloudTrail trail",
                               "describe_trails returned nothing",
                               "After an incident there is no record of which API calls were made, by what, when. The question 'what did the agent do' has no answer.",
                               "Create an organization or multi-region trail to an S3 bucket with log file validation on."))
        else:
            if not any(t.get("IsMultiRegionTrail") for t in items):
                out.append(Finding("trail.single-region", "medium", "proof", "No multi-region trail",
                                   f"{len(items)} trail(s), none multi-region",
                                   "An agent that acts in a region you did not expect leaves no record.",
                                   "Make one trail multi-region."))
            if not any(t.get("LogFileValidationEnabled") for t in items):
                out.append(Finding("trail.no-validation", "high", "proof", "Trail logs are not integrity-validated",
                                   "no trail with LogFileValidationEnabled",
                                   "A log that can be edited proves only that someone wrote a log. Without digests you cannot show it was not changed.",
                                   "Enable log file validation; it costs nothing."))
            for t in items:
                b = t.get("S3BucketName")
                if not b:
                    continue
                s3 = c("s3")
                ver, e2 = _safe(lambda: s3.get_bucket_versioning(Bucket=b))
                lock, e3 = _safe(lambda: s3.get_object_lock_configuration(Bucket=b))
                versioned = not e2 and ver.get("Status") == "Enabled"
                locked = not e3 and bool(lock.get("ObjectLockConfiguration"))
                if not versioned and not locked:
                    out.append(Finding("s3.log-bucket-unprotected", "high", "proof", "Trail bucket allows silent deletion",
                                       f"bucket {b}: versioning {'on' if versioned else 'off'}, object lock {'on' if locked else 'off'}",
                                       "Whoever can write to the bucket can delete the evidence and leave no trace.",
                                       "Enable versioning and Object Lock in compliance mode on the log bucket, with a retention that outlasts any investigation."))
                    break
    for region in regions:
        # Bedrock findings only where Bedrock is actually used. Event history is readable
        # for 90 days with or without a trail, so a single InvokeModel event is the
        # evidence. Without it, a team calling another provider directly from ECS would
        # be scored, and priced, for guardrails on a service they do not touch.
        ev, err = _safe(lambda: c("cloudtrail", region).lookup_events(
            LookupAttributes=[{"AttributeKey": "EventSource", "AttributeValue": "bedrock-runtime.amazonaws.com"}], MaxResults=1))
        if err:
            notes.append(f"cloudtrail {region}: could not read event history, Bedrock findings skipped ({err})")
            bedrock_used = False
        else:
            bedrock_used = bool(ev.get("Events"))
        if not bedrock_used:
            notes.append(f"bedrock {region}: no invocations in the last 90 days, so not assessed")
        br = c("bedrock", region)
        cfg, err = _safe(lambda: br.get_model_invocation_logging_configuration()) if bedrock_used else (None, "skipped")
        if err == "skipped":
            pass
        elif err:
            notes.append(f"bedrock {region}: could not read logging config ({err})")
        elif not (cfg.get("loggingConfig") or {}):
            out.append(Finding(f"bedrock.no-invocation-logging.{region}", "high", "proof", f"Bedrock invocation logging off in {region}",
                               "get_model_invocation_logging_configuration has no loggingConfig",
                               "Prompts and completions are not recorded. What the model was asked and what it answered is unrecoverable.",
                               "Enable model invocation logging to S3 or CloudWatch Logs, in every region agents run."))
        gr, err = _safe(lambda: br.list_guardrails(maxResults=5)) if bedrock_used else (None, "skipped")
        if err == "skipped":
            pass
        elif err:
            if err not in ("AccessDeniedException",):
                notes.append(f"bedrock {region}: could not list guardrails ({err})")
        elif not (gr.get("guardrails") or []):
            out.append(Finding(f"bedrock.no-guardrails.{region}", "medium", "reach", f"No Bedrock guardrail in {region}",
                               "list_guardrails returned nothing",
                               "Nothing between the model and a prompt that asks it to do something it should not.",
                               "Define at least one guardrail with denied topics and sensitive-information filters, and attach it to every agent."))
        gd = c("guardduty", region)
        dets, err = _safe(lambda: gd.list_detectors())
        if err:
            notes.append(f"guardduty {region}: could not read ({err})")
        elif not (dets.get("DetectorIds") or []):
            out.append(Finding(f"guardduty.off.{region}", "medium", "proof", f"GuardDuty off in {region}",
                               "no detector", "Credential misuse and unusual API patterns, which is what a compromised agent looks like, go undetected.",
                               "Enable GuardDuty; it reads CloudTrail and costs little at agent scale."))
        cf = c("config", region)
        rec, err = _safe(lambda: cf.describe_configuration_recorder_status())
        if err:
            notes.append(f"config {region}: could not read ({err})")
        elif not any(r.get("recording") for r in rec.get("ConfigurationRecordersStatus") or []):
            out.append(Finding(f"config.off.{region}", "low", "proof", f"AWS Config not recording in {region}",
                               "no recorder recording", "Resource changes an agent makes have no history to diff against.",
                               "Turn on the configuration recorder for the resource types agents touch."))
    return out


# ---------------------------------------------------------------- scoring --------

def score(findings):
    s = 100
    for f in findings:
        s -= WEIGHT[f.severity]
    return max(0, s)


def account_class(role_count):
    if role_count < 25:
        return "small"
    if role_count < 100:
        return "medium"
    return "large"


def digest(findings, role_count, regions, now):
    counts = {k: 0 for k in WEIGHT}
    for f in findings:
        counts[f.severity] += 1
    d = {"v": 1, "score": score(findings), "class": account_class(role_count), "counts": counts,
         "ids": sorted({f.id.split(".")[0] + "." + f.id.split(".")[1] for f in findings}),
         "regions": len(regions), "at": now.strftime("%Y-%m-%d")}
    return d


def encode_digest(d):
    return base64.urlsafe_b64encode(json.dumps(d, separators=(",", ":")).encode()).decode().rstrip("=")


# ---------------------------------------------------------------- run ------------

def run(session, regions, now=None, audit_log=None):
    now = now or dt.datetime.now(dt.timezone.utc)
    cache = {}
    log = open(audit_log, "a") if audit_log else None

    def c(service, region=None):
        key = (service, region)
        if key not in cache:
            client = session.client(service, region_name=region) if region else session.client(service)
            if log:
                # botocore emits before-call for every request; record service, operation
                # and region, never parameters or responses.
                def record(model, **kw):
                    log.write(f"{dt.datetime.now(dt.timezone.utc).isoformat()} {model.service_model.service_name}:{model.name} {region or '-'}\n")
                    log.flush()
                client.meta.events.register("before-call.*.*", record)
            cache[key] = client
        return cache[key]

    notes = []
    ident, err = _safe(lambda: c("sts", regions[0]).get_caller_identity())
    if err:
        sys.exit(f"cannot identify the account with these credentials ({err})")
    account_id = ident["Account"]
    roles, err = _safe(lambda: c("iam").list_roles(MaxItems=200))
    role_count = len(roles.get("Roles", [])) if not err else 0

    findings = check_spend(c, account_id, regions, notes) + check_reach(c, now, notes) + check_proof(c, regions, notes)
    order = {"critical": 0, "high": 1, "medium": 2, "low": 3}
    findings.sort(key=lambda f: (order[f.severity], f.id))
    return {"version": VERSION, "at": now.isoformat(), "regions": regions, "roles": role_count,
            "score": score(findings), "class": account_class(role_count),
            "findings": [f.as_dict() for f in findings], "could_not_check": notes,
            "digest": digest(findings, role_count, regions, now)}


def print_report(rep):
    print(f"\n  Agent Readiness Check v{VERSION}   score {rep['score']}/100   account class: {rep['class']}   regions: {', '.join(rep['regions'])}\n")
    for q, text in QUESTION.items():
        fs = [f for f in rep["findings"] if f["question"] == q]
        print(f"  {text}")
        if not fs:
            print("     nothing found\n")
            continue
        for f in fs:
            print(f"     [{f['severity']:<8}] {f['title']}")
            print(f"                evidence: {f['evidence']}")
            print(f"                blast radius: {f['blast_radius']}")
        print()
    if rep["could_not_check"]:
        print("  Could not check (missing permission or service not available):")
        for n in rep["could_not_check"]:
            print(f"     {n}")
        print()
    code = encode_digest(rep["digest"])
    print("  Nothing above has left your machine. To see a fixed price for fixing it, open:")
    print(f"     https://abacross.com/readiness/?d={code}")
    print("  That link carries only the score, the account size class and the finding ids, no names or ARNs.\n")


def html_report(rep):
    rows = "".join(
        f"<tr><td class='{f['severity']}'>{f['severity']}</td><td><strong>{f['title']}</strong><br><small>{f['evidence']}</small></td>"
        f"<td>{f['blast_radius']}</td><td>{f['fix']}</td></tr>" for f in rep["findings"])
    notes = "".join(f"<li>{n}</li>" for n in rep["could_not_check"])
    return f"""<!doctype html><meta charset=utf-8><title>Agent Readiness Check</title>
<style>body{{font:15px/1.5 -apple-system,Segoe UI,Roboto,sans-serif;max-width:1100px;margin:2rem auto;padding:0 1rem;color:#0E1B2A}}
table{{border-collapse:collapse;width:100%}}td,th{{border-top:1px solid #E2E8EE;padding:.6rem;vertical-align:top}}
td.critical{{color:#B3261E;font-weight:700}}td.high{{color:#8A5B00;font-weight:700}}.score{{font-size:2.4rem;font-weight:800}}</style>
<h1>Agent Readiness Check</h1><p class=score>{rep['score']}/100</p>
<p>{rep['at']} · account class {rep['class']} · regions {', '.join(rep['regions'])} · {len(rep['findings'])} finding(s)</p>
<table><tr><th>severity</th><th>finding</th><th>blast radius</th><th>fix</th></tr>{rows}</table>
{('<h2>Could not check</h2><ul>' + notes + '</ul>') if notes else ''}
<p><a href="https://abacross.com/readiness/?d={encode_digest(rep['digest'])}">See a fixed price to fix these</a> (the link carries only the score, size class and finding ids).</p>"""


CALLS = """sts:GetCallerIdentity budgets:DescribeBudgets budgets:DescribeBudgetActionsForBudget
cloudwatch:DescribeAlarmsForMetric iam:GetAccountSummary iam:ListUsers iam:ListAccessKeys
iam:GetAccessKeyLastUsed iam:ListRoles iam:ListAttachedRolePolicies iam:ListRolePolicies
iam:GetRolePolicy cloudtrail:DescribeTrails cloudtrail:LookupEvents s3:GetBucketVersioning s3:GetObjectLockConfiguration
bedrock:GetModelInvocationLoggingConfiguration bedrock:ListGuardrails guardduty:ListDetectors
config:DescribeConfigurationRecorderStatus"""


def main():
    ap = argparse.ArgumentParser(description="Agent Readiness Check (read-only)")
    ap.add_argument("--regions", default=",".join(DEFAULT_REGIONS), help="comma-separated regions agents run in")
    ap.add_argument("--profile", help="AWS profile to use")
    ap.add_argument("--json", help="write the full report here")
    ap.add_argument("--html", help="write an HTML report here")
    ap.add_argument("--calls", action="store_true", help="list every API call this makes and exit")
    ap.add_argument("--audit-log", help="append every API call actually made (service:operation region) to this file")
    ap.add_argument("--policy", action="store_true", help="print an IAM policy allowing exactly these calls and nothing else")
    a = ap.parse_args()
    if a.calls:
        print(CALLS)
        return 0
    if a.policy:
        print(json.dumps({"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Action": sorted(CALLS.split()), "Resource": "*"}]}, indent=2))
        return 0
    regions = [r.strip() for r in a.regions.split(",") if r.strip()]
    session = boto3.Session(profile_name=a.profile) if a.profile else boto3.Session()
    rep = run(session, regions, audit_log=a.audit_log)
    print_report(rep)
    if a.json:
        with open(a.json, "w") as fh:
            json.dump(rep, fh, indent=2, default=str)
        print(f"  full report: {a.json}")
    if a.html:
        with open(a.html, "w") as fh:
            fh.write(html_report(rep))
        print(f"  html report: {a.html}")
    return 0


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