#!/usr/bin/env python3
"""
Member-Account Trust Assessment
================================
Extends the org assessment into a member account by assuming a cross-account
role from the management-account profile, then evaluating that account's
IAM role trust relationships and any delegated services it is aware of.

Delegated administrators & trusted-access services are ORG-LEVEL and can only
be read from the management account, so this focuses on what is meaningful
from inside a member account: IAM role trust policies (service / cross-account
/ federated) and account-level org context.

Usage:
  python3 assess_member_account.py --account 547868853286 [--name Lazka]
      [--role OrganizationAccountAccessRole] [--profile PROFILE]
"""
import argparse
import json
import os
import sys
from datetime import datetime, date, timezone

import boto3
from botocore.exceptions import ClientError, BotoCoreError

DEFAULT_PROFILE = "dasnuve-scrivas-louis-impersonation"
FINDINGS_DIR = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "findings")

VENDOR_ACCOUNTS = {
    "123311413059": "Intruder.io (external vulnerability scanning)",
    "728997465891": "Secureframe (SOC 2 / compliance automation)",
}


def _default(o):
    if isinstance(o, (datetime, date)):
        return o.isoformat()
    return str(o)


def paginate(client, op, key, **kwargs):
    out = []
    try:
        for page in client.get_paginator(op).paginate(**kwargs):
            out.extend(page.get(key, []))
    except (ClientError, BotoCoreError) as e:
        return out, str(e)
    return out, None


def assume(session, account_id, role_name, candidate_roles):
    """Try the primary role, then fall back through candidate role names."""
    sts = session.client("sts")
    tried = [role_name] + [r for r in candidate_roles if r != role_name]
    errors = {}
    for rn in tried:
        arn = f"arn:aws:iam::{account_id}:role/{rn}"
        try:
            creds = sts.assume_role(
                RoleArn=arn, RoleSessionName="scrivas-discovery")["Credentials"]
            print(f"[+] Assumed {arn}")
            return boto3.Session(
                aws_access_key_id=creds["AccessKeyId"],
                aws_secret_access_key=creds["SecretAccessKey"],
                aws_session_token=creds["SessionToken"],
            ), rn, None
        except (ClientError, BotoCoreError) as e:
            errors[rn] = str(e)
            print(f"[-] {rn}: {e}")
    return None, None, errors


def assess_iam_trust(member_session):
    iam = member_session.client("iam")
    roles, err = paginate(iam, "list_roles", "Roles")
    if err:
        return {"error": err}
    trust = []
    for r in roles:
        doc = r.get("AssumeRolePolicyDocument", {})
        flat = json.dumps(doc)
        kind = []
        if '"Service"' in flat:
            kind.append("service")
        if '"AWS"' in flat:
            kind.append("cross-account/aws")
        if '"Federated"' in flat:
            kind.append("federated")
        vendor = None
        for stmt in doc.get("Statement", []):
            p = stmt.get("Principal", {})
            aws_p = p.get("AWS") if isinstance(p, dict) else None
            for acct, name in VENDOR_ACCOUNTS.items():
                if aws_p and acct in json.dumps(aws_p):
                    vendor = name
        trust.append({
            "RoleName": r["RoleName"],
            "Path": r.get("Path"),
            "Kind": kind,
            "Vendor": vendor,
            "AssumeRolePolicyDocument": doc,
        })
    return trust


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--account", required=True)
    ap.add_argument("--name", default="member")
    ap.add_argument("--role", default="OrganizationAccountAccessRole")
    ap.add_argument("--profile", default=DEFAULT_PROFILE)
    args = ap.parse_args()

    session = boto3.Session(profile_name=args.profile)

    candidates = ["OrganizationAccountAccessRole", "AWSControlTowerExecution",
                  "OrganizationAccountAccessRole", "AdministratorAccess"]
    member, role_used, err = assume(session, args.account, args.role, candidates)

    result = {
        "generated": datetime.now(timezone.utc).isoformat(),
        "account_id": args.account,
        "account_name": args.name,
    }

    if member is None:
        result["access"] = {"assumed": False, "errors": err}
        print("\n[!] Could not assume any cross-account role into "
              f"{args.name} ({args.account}). This itself is a finding: no "
              "management-account access path exists into the member account.")
    else:
        ident = member.client("sts").get_caller_identity()
        result["access"] = {"assumed": True, "role_used": role_used,
                            "assumed_arn": ident["Arn"]}
        result["iam_role_trust_policies"] = assess_iam_trust(member)

        roles = result["iam_role_trust_policies"]
        if isinstance(roles, list):
            xacct = [r for r in roles if "cross-account/aws" in r["Kind"]]
            fed = [r for r in roles if "federated" in r["Kind"]]
            print(f"\n[{args.name} / {args.account}] IAM roles: {len(roles)} | "
                  f"{len(xacct)} cross-account | {len(fed)} federated")
            for r in xacct:
                tag = f"  [{r['Vendor']}]" if r.get("Vendor") else ""
                print(f"    - {r['RoleName']}{tag}")

    out = os.path.join(FINDINGS_DIR, f"member_{args.name.lower()}_{args.account}.json")
    os.makedirs(FINDINGS_DIR, exist_ok=True)
    with open(out, "w") as f:
        json.dump(result, f, indent=2, default=_default)
    print(f"\nWritten to: {out}")


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