Private Endpoint Configuration for AI Services

import boto3
import json
 
def create_bedrock_vpc_endpoints(
    session: boto3.Session,
    vpc_id: str,
    subnet_ids: list,
    security_group_ids: list,
    region: str = "us-east-1",
) -> dict:
    """Create VPC interface endpoints for Bedrock services."""
    ec2 = session.client("ec2", region_name=region)
    results = {}
 
    endpoints_to_create = {
        "bedrock-management": f"com.amazonaws.{region}.bedrock",
        "bedrock-runtime": f"com.amazonaws.{region}.bedrock-runtime",
    }
 
    for name, service_name in endpoints_to_create.items():
        try:
            response = ec2.create_vpc_endpoint(
                VpcEndpointType="Interface",
                VpcId=vpc_id,
                ServiceName=service_name,
                SubnetIds=subnet_ids,
                SecurityGroupIds=security_group_ids,
                PrivateDnsEnabled=True,
                PolicyDocument=json.dumps({
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Sid": "AllowBedrockAccess",
                            "Effect": "Allow",
                            "Principal": {"AWS": "*"},
                            "Action": [
                                "bedrock:InvokeModel",
                                "bedrock:InvokeModelWithResponseStream",
                            ],
                            "Resource": "*",
                            "Condition": {
                                "StringEquals": {
                                    "aws:PrincipalOrgID": "o-your-org-id"
                                }
                            },
                        }
                    ],
                }),
                TagSpecifications=[
                    {
                        "ResourceType": "vpc-endpoint",
                        "Tags": [
                            {"Key": "Name", "Value": f"bedrock-{name}"},
                            {"Key": "Security", "Value": "ai-private-endpoint"},
                        ],
                    }
                ],
            )
 
            endpoint = response["VpcEndpoint"]
            results[name] = {
                "endpoint_id": endpoint["VpcEndpointId"],
                "service_name": service_name,
                "state": endpoint["State"],
                "dns_entries": [
                    {"dns_name": d["DnsName"], "hosted_zone_id": d["HostedZoneId"]}
                    for d in endpoint.get("DnsEntries", [])
                ],
            }
 
        except Exception as e:
            results[name] = {"error": str(e)}
 
    return results

def generate_restrictive_endpoint_policy(
    org_id: str,
    allowed_model_arns: list,
    allowed_account_ids: list = None,
) -> dict:
    """Generate a restrictive VPC endpoint policy for Bedrock."""
    statements = [
        {
            "Sid": "AllowInvokeSpecificModels",
            "Effect": "Allow",
            "Principal": {"AWS": "*"},
            "Action": [
                "bedrock:InvokeModel",
                "bedrock:InvokeModelWithResponseStream",
            ],
            "Resource": allowed_model_arns,
            "Condition": {
                "StringEquals": {
                    "aws:PrincipalOrgID": org_id,
                },
            },
        },
        {
            "Sid": "AllowListModels",
            "Effect": "Allow",
            "Principal": {"AWS": "*"},
            "Action": [
                "bedrock:ListFoundationModels",
                "bedrock:GetFoundationModel",
            ],
            "Resource": "*",
            "Condition": {
                "StringEquals": {
                    "aws:PrincipalOrgID": org_id,
                },
            },
        },
    ]
 
    if allowed_account_ids:
        # Add account-level restriction
        for statement in statements:
            statement["Condition"]["StringEquals"]["aws:PrincipalAccount"] = allowed_account_ids
 
    return {
        "Version": "2012-10-17",
        "Statement": statements,
    }

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "DenyNonVPCEndpointAccess",
            "Effect": "Deny",
            "Action": "bedrock:*",
            "Resource": "*",
            "Condition": {
                "StringNotEquals": {
                    "aws:sourceVpce": [
                        "vpce-0abc123def456789a",
                        "vpce-0abc123def456789b"
                    ]
                }
            }
        }
    ]
}

from azure.identity import DefaultAzureCredential
from azure.mgmt.network import NetworkManagementClient
from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient
 
def setup_azure_openai_private_link(
    subscription_id: str,
    resource_group: str,
    openai_account_name: str,
    vnet_name: str,
    subnet_name: str,
    location: str = "eastus",
) -> dict:
    """Set up complete private link for Azure OpenAI."""
    credential = DefaultAzureCredential()
    network_client = NetworkManagementClient(credential, subscription_id)
    cognitive_client = CognitiveServicesManagementClient(credential, subscription_id)
 
    results = {"steps": []}
 
    # Step 1: Get the Azure OpenAI resource ID
    account = cognitive_client.accounts.get(resource_group, openai_account_name)
 
    # Step 2: Disable subnet private endpoint policies
    subnet = network_client.subnets.get(resource_group, vnet_name, subnet_name)
    if subnet.private_endpoint_network_policies != "Disabled":
        subnet.private_endpoint_network_policies = "Disabled"
        network_client.subnets.begin_create_or_update(
            resource_group, vnet_name, subnet_name, subnet
        ).result()
        results["steps"].append("Disabled private endpoint network policies on subnet")
 
    # Step 3: Create the private endpoint
    pe_poller = network_client.private_endpoints.begin_create_or_update(
        resource_group,
        f"{openai_account_name}-private-endpoint",
        {
            "location": location,
            "properties": {
                "subnet": {"id": subnet.id},
                "privateLinkServiceConnections": [
                    {
                        "name": f"{openai_account_name}-plsc",
                        "properties": {
                            "privateLinkServiceId": account.id,
                            "groupIds": ["account"],
                        },
                    }
                ],
            },
        },
    )
    pe_result = pe_poller.result()
    results["steps"].append(f"Created private endpoint: {pe_result.id}")
 
    # Step 4: Get the private IP address
    private_ip = None
    for nic in pe_result.network_interfaces:
        nic_detail = network_client.network_interfaces.get(
            resource_group, nic.id.split("/")[-1]
        )
        for ip_config in nic_detail.ip_configurations:
            private_ip = ip_config.private_ip_address
            break
 
    results["private_ip"] = private_ip
    results["steps"].append(f"Private IP assigned: {private_ip}")
 
    # Step 5: Disable public access
    cognitive_client.accounts.begin_update(
        resource_group,
        openai_account_name,
        {
            "properties": {
                "publicNetworkAccess": "Disabled",
            }
        },
    ).result()
    results["steps"].append("Disabled public network access")
 
    return results

from azure.mgmt.privatedns import PrivateDnsManagementClient
 
def configure_azure_openai_private_dns(
    subscription_id: str,
    resource_group: str,
    vnet_name: str,
    private_ip: str,
    openai_account_name: str,
) -> dict:
    """Configure Private DNS Zone for Azure OpenAI private endpoint."""
    credential = DefaultAzureCredential()
    dns_client = PrivateDnsManagementClient(credential, subscription_id)
    network_client = NetworkManagementClient(credential, subscription_id)
 
    zone_name = "privatelink.openai.azure.com"
 
    # Create private DNS zone
    dns_client.private_zones.begin_create_or_update(
        resource_group,
        zone_name,
        {"location": "global"},
    ).result()
 
    # Link zone to VNet
    vnet = network_client.virtual_networks.get(resource_group, vnet_name)
    dns_client.virtual_network_links.begin_create_or_update(
        resource_group,
        zone_name,
        f"{vnet_name}-dns-link",
        {
            "location": "global",
            "properties": {
                "virtualNetwork": {"id": vnet.id},
                "registrationEnabled": False,
            },
        },
    ).result()
 
    # Create A record
    dns_client.record_sets.create_or_update(
        resource_group,
        zone_name,
        openai_account_name,
        "A",
        {
            "properties": {
                "ttl": 300,
                "aRecords": [{"ipv4Address": private_ip}],
            }
        },
    )
 
    return {
        "dns_zone": zone_name,
        "fqdn": f"{openai_account_name}.{zone_name}",
        "resolves_to": private_ip,
    }

from google.cloud import compute_v1
 
def create_vertex_psc_endpoint(
    project_id: str,
    network_name: str,
    subnet_name: str,
    region: str = "us-central1",
) -> dict:
    """Create Private Service Connect endpoint for Vertex AI."""
    # PSC requires a subnet with purpose PRIVATE_SERVICE_CONNECT
    subnets_client = compute_v1.SubnetworksClient()
    addresses_client = compute_v1.GlobalAddressesClient()
    forwarding_client = compute_v1.GlobalForwardingRulesClient()
 
    results = {"steps": []}
 
    # Step 1: Reserve a private IP for PSC
    address_body = compute_v1.Address(
        name="vertex-ai-psc-address",
        address_type="INTERNAL",
        purpose="PRIVATE_SERVICE_CONNECT",
        network=f"projects/{project_id}/global/networks/{network_name}",
    )
 
    try:
        addresses_client.insert(
            project=project_id,
            address_resource=address_body,
        ).result()
        results["steps"].append("Reserved internal address for PSC")
    except Exception as e:
        results["steps"].append(f"Address reservation: {e}")
 
    # Get the reserved address
    address = addresses_client.get(
        project=project_id,
        address="vertex-ai-psc-address",
    )
    private_ip = address.address
    results["private_ip"] = private_ip
 
    # Step 2: Create forwarding rule to Google API bundle
    forwarding_rule = compute_v1.ForwardingRule(
        name="vertex-ai-psc-rule",
        ip_address=address.self_link,
        network=f"projects/{project_id}/global/networks/{network_name}",
        target="all-apis",  # Or "vpc-sc" for VPC-SC compatible bundle
        load_balancing_scheme="",
    )
 
    try:
        forwarding_client.insert(
            project=project_id,
            forwarding_rule_resource=forwarding_rule,
        ).result()
        results["steps"].append("Created PSC forwarding rule")
    except Exception as e:
        results["steps"].append(f"Forwarding rule: {e}")
 
    # Step 3: Configure DNS to route *.googleapis.com to PSC IP
    results["dns_config"] = {
        "zone": "googleapis.com",
        "record": f"*.googleapis.com -> {private_ip}",
        "note": "Configure Cloud DNS private zone to resolve "
                "googleapis.com to the PSC endpoint IP.",
    }
 
    return results

from google.cloud import dns
 
def configure_private_dns_for_vertex(
    project_id: str,
    network_name: str,
    psc_ip: str,
) -> dict:
    """Configure Cloud DNS for private Vertex AI access."""
    dns_client = dns.Client(project=project_id)
 
    # Create private DNS zone for googleapis.com
    zone = dns_client.zone(
        "google-apis-private",
        dns_name="googleapis.com.",
    )
 
    # This creates the zone -- in production, use the resource manager API
    # for private zones with VPC binding
    zone_config = {
        "name": "google-apis-private",
        "dns_name": "googleapis.com.",
        "visibility": "private",
        "private_visibility_config": {
            "networks": [
                {"network_url": f"projects/{project_id}/global/networks/{network_name}"}
            ]
        },
    }
 
    records = [
        {
            "name": "aiplatform.googleapis.com.",
            "type": "A",
            "ttl": 300,
            "rrdatas": [psc_ip],
        },
        {
            "name": "*.aiplatform.googleapis.com.",
            "type": "CNAME",
            "ttl": 300,
            "rrdatas": ["aiplatform.googleapis.com."],
        },
    ]
 
    return {
        "zone_config": zone_config,
        "dns_records": records,
        "detail": "Apply using gcloud dns managed-zones create or Terraform.",
    }

import socket
import subprocess
import json
from typing import Optional
 
def validate_private_endpoint(
    hostname: str,
    expected_private_ip: str,
    test_public_access: bool = True,
) -> dict:
    """Validate private endpoint configuration."""
    results = {
        "hostname": hostname,
        "dns_resolution": {},
        "connectivity": {},
        "public_access_blocked": None,
    }
 
    # Test 1: DNS resolution
    try:
        resolved_ips = socket.getaddrinfo(hostname, 443, socket.AF_INET)
        resolved_ip = resolved_ips[0][4][0]
        results["dns_resolution"] = {
            "resolved_ip": resolved_ip,
            "expected_ip": expected_private_ip,
            "match": resolved_ip == expected_private_ip,
        }
 
        if resolved_ip != expected_private_ip:
            results["dns_resolution"]["warning"] = (
                "DNS resolves to a different IP than expected. "
                "This may indicate DNS is still resolving to the public endpoint."
            )
    except socket.gaierror as e:
        results["dns_resolution"] = {"error": str(e)}
 
    # Test 2: TCP connectivity to private endpoint
    try:
        sock = socket.create_connection((hostname, 443), timeout=5)
        sock.close()
        results["connectivity"]["private"] = {
            "status": "success",
            "detail": f"TCP connection to {hostname}:443 succeeded",
        }
    except (socket.timeout, ConnectionRefusedError, OSError) as e:
        results["connectivity"]["private"] = {
            "status": "failed",
            "error": str(e),
        }
 
    return results
 
 
def validate_endpoint_isolation(
    session,  # boto3.Session or equivalent
    service: str,
    endpoint_id: str,
) -> dict:
    """Validate that the endpoint policy correctly restricts access."""
    results = {"endpoint_id": endpoint_id, "tests": []}
 
    if service == "aws":
        ec2 = session.client("ec2")
        endpoint_detail = ec2.describe_vpc_endpoints(
            VpcEndpointIds=[endpoint_id]
        )["VpcEndpoints"][0]
 
        # Check endpoint state
        results["tests"].append({
            "name": "endpoint_state",
            "pass": endpoint_detail["State"] == "available",
            "state": endpoint_detail["State"],
        })
 
        # Check private DNS enabled
        results["tests"].append({
            "name": "private_dns_enabled",
            "pass": endpoint_detail.get("PrivateDnsEnabled", False),
            "detail": "Private DNS must be enabled to override public resolution",
        })
 
        # Check security groups
        sg_ids = [sg["GroupId"] for sg in endpoint_detail.get("Groups", [])]
        results["tests"].append({
            "name": "security_groups_attached",
            "pass": len(sg_ids) > 0,
            "security_groups": sg_ids,
        })
 
        # Check endpoint policy
        policy = endpoint_detail.get("PolicyDocument", "")
        if isinstance(policy, str):
            policy = json.loads(policy) if policy else {}
 
        has_conditions = False
        for statement in policy.get("Statement", []):
            if statement.get("Condition"):
                has_conditions = True
                break
 
        results["tests"].append({
            "name": "endpoint_policy_has_conditions",
            "pass": has_conditions,
            "detail": "Endpoint policy should include conditions (OrgID, Account) "
                      "to restrict which identities can use the endpoint",
        })
 
    return results

def design_multi_cloud_private_endpoints() -> dict:
    """Design consistent private endpoint patterns across clouds."""
    return {
        "aws": {
            "mechanism": "VPC Interface Endpoints",
            "services": [
                "com.amazonaws.<region>.bedrock",
                "com.amazonaws.<region>.bedrock-runtime",
                "com.amazonaws.<region>.sagemaker.api",
                "com.amazonaws.<region>.sagemaker.runtime",
            ],
            "dns": "Route 53 Private Hosted Zones (auto with PrivateDnsEnabled)",
            "policy": "VPC Endpoint Policy with OrgID condition",
            "enforcement": "SCP with aws:sourceVpce condition",
        },
        "azure": {
            "mechanism": "Azure Private Link / Private Endpoints",
            "services": [
                "Microsoft.CognitiveServices/accounts (groupId: account)",
            ],
            "dns": "Private DNS Zone: privatelink.openai.azure.com",
            "policy": "NSG on PE subnet + Entra ID RBAC",
            "enforcement": "publicNetworkAccess: Disabled on resource",
        },
        "gcp": {
            "mechanism": "Private Service Connect (PSC)",
            "services": [
                "aiplatform.googleapis.com via PSC forwarding rule",
            ],
            "dns": "Cloud DNS private zone for googleapis.com",
            "policy": "VPC-SC perimeter + IAM",
            "enforcement": "VPC-SC + Organization Policy",
        },
        "validation_checklist": [
            "DNS resolves to private IP from all application subnets",
            "DNS resolves to private IP from all peered/spoke networks",
            "Public access returns connection refused",
            "Endpoint policy restricts to organization principals",
            "Security groups/NSGs restrict source subnets",
            "Management and runtime APIs both use private endpoints",
            "Monitoring confirms no traffic over public paths",
        ],
    }