Vertex AI Red Team Walkthrough

# vertex_recon.py
"""Enumerate Vertex AI resources in a GCP project."""
from google.cloud import aiplatform
from google.cloud import resourcemanager_v3
 
def enumerate_vertex_resources(project_id, location="us-central1"):
    """Map all Vertex AI resources in the project."""
    aiplatform.init(project=project_id, location=location)
 
    # List endpoints
    endpoints = aiplatform.Endpoint.list()
    print(f"Prediction Endpoints ({len(endpoints)}):")
    for ep in endpoints:
        print(f"  {ep.display_name} (ID: {ep.name})")
        print(f"    Created: {ep.create_time}")
        print(f"    Traffic split: {ep.traffic_split}")
        for model in ep.list_models():
            print(f"    Model: {model.display_name} "
                  f"(ID: {model.id})")
 
    # List models
    models = aiplatform.Model.list()
    print(f"\nModels ({len(models)}):")
    for m in models:
        print(f"  {m.display_name} (ID: {m.name})")
        print(f"    Framework: {m.framework}")
        print(f"    Container: {m.container_spec}")
 
    # List custom training jobs
    jobs = aiplatform.CustomJob.list()
    print(f"\nTraining Jobs ({len(jobs)}):")
    for j in jobs:
        print(f"  {j.display_name} ({j.state})")
 
    # List pipelines
    pipelines = aiplatform.PipelineJob.list()
    print(f"\nPipelines ({len(pipelines)}):")
    for p in pipelines:
        print(f"  {p.display_name} ({p.state})")
 
    return {
        "endpoints": endpoints,
        "models": models,
        "jobs": jobs,
        "pipelines": pipelines,
    }

def review_vertex_iam(project_id):
    """Review IAM bindings relevant to Vertex AI."""
    from google.cloud import resourcemanager_v3
 
    client = resourcemanager_v3.ProjectsClient()
    policy = client.get_iam_policy(
        request={"resource": f"projects/{project_id}"}
    )
 
    vertex_roles = [
        "roles/aiplatform.admin",
        "roles/aiplatform.user",
        "roles/aiplatform.viewer",
        "roles/aiplatform.customCodeServiceAgent",
        "roles/aiplatform.serviceAgent",
    ]
 
    dangerous_bindings = []
    print("Vertex AI IAM bindings:")
    for binding in policy.bindings:
        if any(role in binding.role for role in vertex_roles):
            print(f"\n  Role: {binding.role}")
            for member in binding.members:
                print(f"    Member: {member}")
                # Flag service accounts with admin roles
                if ("serviceAccount" in member and
                        "admin" in binding.role.lower()):
                    dangerous_bindings.append({
                        "role": binding.role,
                        "member": member,
                        "issue": "Service account with admin role",
                    })
                # Flag allUsers or allAuthenticatedUsers
                if member in ("allUsers", "allAuthenticatedUsers"):
                    dangerous_bindings.append({
                        "role": binding.role,
                        "member": member,
                        "issue": "Public access to Vertex AI",
                    })
 
    if dangerous_bindings:
        print(f"\nFINDINGS ({len(dangerous_bindings)}):")
        for db in dangerous_bindings:
            print(f"  {db['issue']}: {db['member']} -> {db['role']}")
 
    return dangerous_bindings

def assess_model_garden(project_id, location="us-central1"):
    """Assess Model Garden usage and configuration."""
    aiplatform.init(project=project_id, location=location)
 
    models = aiplatform.Model.list()
 
    for model in models:
        print(f"\nModel: {model.display_name}")
        print(f"  Resource name: {model.resource_name}")
        print(f"  Source: {model.model_source_info}")
 
        # Check container image for custom models
        if model.container_spec:
            image = model.container_spec.image_uri
            print(f"  Container: {image}")
 
            # Flag if using latest tag or public registry
            if ":latest" in image:
                print("  FINDING: Using :latest tag -- "
                      "not pinned to specific version")
            if not image.startswith("gcr.io/") and \
               not image.startswith("us-docker.pkg.dev/"):
                print("  FINDING: Container from non-GCP registry")
 
        # Check model artifacts location
        if model.artifact_uri:
            print(f"  Artifacts: {model.artifact_uri}")
            if "gs://" in model.artifact_uri:
                bucket = model.artifact_uri.split("/")[2]
                print(f"  GCS Bucket: {bucket}")
                print("  Check bucket permissions: "
                      f"gsutil iam get gs://{bucket}")

# Check if model artifacts are stored in publicly accessible buckets
gsutil iam get gs://<model-artifacts-bucket>
 
# Check for overly permissive bucket ACLs
gsutil ls -L -b gs://<model-artifacts-bucket>
 
# Verify container image provenance
gcloud artifacts docker images describe <image-uri> \
  --show-provenance

def probe_feature_store(project_id, location="us-central1"):
    """Enumerate and probe Feature Store for data exposure."""
    aiplatform.init(project=project_id, location=location)
 
    # List feature stores
    feature_stores = aiplatform.Featurestore.list()
    print(f"Feature Stores ({len(feature_stores)}):")
 
    for fs in feature_stores:
        print(f"\n  {fs.display_name} (ID: {fs.name})")
 
        # List entity types
        entity_types = fs.list_entity_types()
        for et in entity_types:
            print(f"    Entity Type: {et.display_name}")
 
            # List features
            features = et.list_features()
            for f in features:
                print(f"      Feature: {f.display_name} "
                      f"(type: {f.value_type})")
 
            # Attempt to read feature values
            try:
                # Read a sample of feature values
                entity_ids = ["test_entity_1", "user_1",
                              "admin", "root"]
                for entity_id in entity_ids:
                    try:
                        values = et.read(
                            entity_ids=[entity_id],
                            feature_ids=[f.display_name
                                         for f in features[:5]],
                        )
                        if not values.empty:
                            print(f"      Data retrieved for "
                                  f"entity '{entity_id}':")
                            print(f"        {values.head()}")
                    except Exception:
                        pass  # Entity not found is expected
            except Exception as e:
                print(f"      Read error: {e}")

def test_feature_store_access(project_id, location="us-central1"):
    """Test whether the testing identity has excessive Feature Store access."""
    from google.cloud import aiplatform_v1
 
    client = aiplatform_v1.FeaturestoreServiceClient(
        client_options={"api_endpoint": f"{location}-aiplatform.googleapis.com"}
    )
 
    # Attempt operations that should be restricted
    operations = {
        "list": lambda: client.list_featurestores(
            parent=f"projects/{project_id}/locations/{location}"
        ),
        "search": lambda: client.search_features(
            location=f"projects/{project_id}/locations/{location}",
            query="*",
        ),
    }
 
    for op_name, op_func in operations.items():
        try:
            result = op_func()
            print(f"Operation '{op_name}': PERMITTED")
            # If search returns all features, document the finding
            if op_name == "search":
                features = list(result)
                print(f"  Found {len(features)} features across all stores")
        except Exception as e:
            if "403" in str(e):
                print(f"Operation '{op_name}': DENIED (expected)")
            else:
                print(f"Operation '{op_name}': ERROR ({e})")

# Query Vertex AI prediction logs
gcloud logging read "resource.type=\"aiplatform.googleapis.com/Endpoint\" \
  AND timestamp>=\"2026-03-14T00:00:00Z\"" \
  --project=<project-id> \
  --limit=50 \
  --format=json
 
# Check for prediction request/response logging
gcloud logging read "resource.type=\"aiplatform.googleapis.com/Endpoint\" \
  AND jsonPayload.request IS NOT NULL" \
  --project=<project-id> \
  --limit=10
 
# Check audit logs for Vertex AI admin operations
gcloud logging read "protoPayload.serviceName=\"aiplatform.googleapis.com\" \
  AND logName:\"activity\"" \
  --project=<project-id> \
  --limit=50

def analyze_vertex_logging(project_id):
    """Analyze Vertex AI logging configuration and coverage."""
    from google.cloud import logging_v2
 
    client = logging_v2.Client(project=project_id)
 
    # Check for log sinks that capture Vertex AI logs
    sinks = list(client.list_sinks())
    vertex_sinks = [
        s for s in sinks
        if "aiplatform" in s.filter.lower()
    ]
 
    print(f"Total log sinks: {len(sinks)}")
    print(f"Vertex AI-specific sinks: {len(vertex_sinks)}")
 
    if not vertex_sinks:
        print("FINDING: No dedicated log sinks for Vertex AI. "
              "Prediction logs may be subject to default retention only.")
 
    # Check what is being logged
    entries = client.list_entries(
        filter_="resource.type=\"aiplatform.googleapis.com/Endpoint\"",
        max_results=10,
    )
 
    logged_fields = set()
    for entry in entries:
        if hasattr(entry, "payload"):
            if isinstance(entry.payload, dict):
                logged_fields.update(entry.payload.keys())
 
    print(f"\nLogged fields: {logged_fields}")
 
    # Key gaps to check
    expected_fields = {"request", "response", "model_id",
                       "caller_identity", "latency_ms"}
    missing = expected_fields - logged_fields
    if missing:
        print(f"Missing fields: {missing}")
        print("FINDING: Prediction content may not be logged.")

def assess_pipelines(project_id, location="us-central1"):
    """Assess Vertex AI pipeline security."""
    aiplatform.init(project=project_id, location=location)
 
    pipelines = aiplatform.PipelineJob.list()
 
    for pipeline in pipelines:
        print(f"\nPipeline: {pipeline.display_name}")
        print(f"  State: {pipeline.state}")
        print(f"  Service Account: {pipeline.service_account}")
 
        # Check if pipeline uses default compute service account
        if pipeline.service_account and \
           "compute@developer.gserviceaccount.com" in pipeline.service_account:
            print("  FINDING: Pipeline uses default compute service account. "
                  "This account typically has broad project-level permissions.")
 
        # Check pipeline parameters for hardcoded secrets
        if pipeline.pipeline_spec:
            params = pipeline.pipeline_spec.get("pipelineInfo", {})
            print(f"  Parameters: {list(params.keys())}")