Amazon AWS Amplify CLI before 12

#!/usr/bin/env python3
# check_cve_2024_28056.py
# Exit codes:
#   0 = PATCHED (no vulnerable roles found)
#   1 = VULNERABLE (vulnerable role(s) + same-account classic-flow identity pool(s) found)
#   2 = UNKNOWN (insufficient permissions/errors, or vulnerable roles found but no classic-flow pool found)

import argparse
import sys
import json
from botocore.exceptions import BotoCoreError, ClientError

try:
    import boto3
except Exception:
    print('UNKNOWN - boto3 is not installed')
    sys.exit(2)


def make_session(profile=None):
    try:
        return boto3.Session(profile_name=profile) if profile else boto3.Session()
    except Exception as e:
        print(f'UNKNOWN - unable to create boto3 session: {e}')
        sys.exit(2)


def normalize_actions(action):
    if isinstance(action, list):
        return action
    if isinstance(action, str):
        return [action]
    return []


def is_cognito_trust_vulnerable(stmt):
    if stmt.get('Effect') != 'Allow':
        return False

    principal = stmt.get('Principal', {})
    federated = principal.get('Federated')
    if federated != 'cognito-identity.amazonaws.com':
        return False

    actions = normalize_actions(stmt.get('Action'))
    if 'sts:AssumeRoleWithWebIdentity' not in actions:
        return False

    # Vulnerable pattern for this CVE family: missing aud scoping condition.
    cond = stmt.get('Condition')
    if not cond:
        return True

    string_equals = cond.get('StringEquals', {}) if isinstance(cond, dict) else {}
    aud = string_equals.get('cognito-identity.amazonaws.com:aud')
    if not aud:
        return True

    return False


def list_vulnerable_roles(session):
    iam = session.client('iam')
    vulnerable = []
    paginator = iam.get_paginator('list_roles')
    for page in paginator.paginate():
        for role in page.get('Roles', []):
            doc = role.get('AssumeRolePolicyDocument', {})
            statements = doc.get('Statement', [])
            if isinstance(statements, dict):
                statements = [statements]
            for stmt in statements:
                if is_cognito_trust_vulnerable(stmt):
                    vulnerable.append({
                        'RoleName': role.get('RoleName'),
                        'Arn': role.get('Arn')
                    })
                    break
    return vulnerable


def list_classic_flow_pools(session):
    regions = session.get_available_regions('cognito-identity')
    found = []
    errors = []

    for region in regions:
        try:
            client = session.client('cognito-identity', region_name=region)
            token = None
            while True:
                kwargs = {'MaxResults': 60}
                if token:
                    kwargs['NextToken'] = token
                resp = client.list_identity_pools(**kwargs)
                for pool in resp.get('IdentityPools', []):
                    pool_id = pool['IdentityPoolId']
                    desc = client.describe_identity_pool(IdentityPoolId=pool_id)
                    if desc.get('AllowClassicFlow') is True:
                        found.append({
                            'Region': region,
                            'IdentityPoolId': pool_id,
                            'IdentityPoolName': desc.get('IdentityPoolName', '')
                        })
                token = resp.get('NextToken')
                if not token:
                    break
        except (ClientError, BotoCoreError) as e:
            errors.append(f'{region}: {e}')
        except Exception as e:
            errors.append(f'{region}: {e}')

    return found, errors


def main():
    parser = argparse.ArgumentParser(description='Assess residual exposure for CVE-2024-28056 in an AWS account')
    parser.add_argument('--profile', help='AWS profile name', default=None)
    args = parser.parse_args()

    session = make_session(args.profile)

    try:
        vulnerable_roles = list_vulnerable_roles(session)
    except (ClientError, BotoCoreError) as e:
        print(f'UNKNOWN - IAM enumeration failed: {e}')
        sys.exit(2)
    except Exception as e:
        print(f'UNKNOWN - unexpected IAM error: {e}')
        sys.exit(2)

    try:
        classic_pools, pool_errors = list_classic_flow_pools(session)
    except Exception as e:
        print(f'UNKNOWN - Cognito enumeration failed: {e}')
        sys.exit(2)

    summary = {
        'vulnerable_role_count': len(vulnerable_roles),
        'classic_flow_pool_count': len(classic_pools),
        'example_roles': vulnerable_roles[:5],
        'example_classic_flow_pools': classic_pools[:5],
        'regional_errors': pool_errors[:10],
    }

    if vulnerable_roles and classic_pools:
        print('VULNERABLE - permissive Cognito-trusted role(s) and same-account classic-flow identity pool(s) found')
        print(json.dumps(summary, indent=2, sort_keys=True))
        sys.exit(1)

    if vulnerable_roles and not classic_pools:
        print('UNKNOWN - vulnerable role(s) found, but no classic-flow identity pools were discovered')
        print('UNKNOWN - cross-account abuse is blocked by AWS, but stale role trust policies still need review')
        print(json.dumps(summary, indent=2, sort_keys=True))
        sys.exit(2)

    if not vulnerable_roles:
        print('PATCHED - no vulnerable Cognito-trusted IAM role patterns found for CVE-2024-28056')
        print(json.dumps(summary, indent=2, sort_keys=True))
        sys.exit(0)

    print('UNKNOWN - unable to determine state conclusively')
    print(json.dumps(summary, indent=2, sort_keys=True))
    sys.exit(2)


if __name__ == '__main__':
    main()