Network Isolation for Cloud AI Workloads

# network_isolation/zones.py
"""Network zone definitions and policy enforcement for AI workloads."""
 
from __future__ import annotations
 
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
 
 
class NetworkZone(Enum):
    """Security zones for AI workload network segmentation."""
    APPLICATION = "application"       # User-facing application layer
    INFERENCE = "inference"           # Model inference and guardrails
    DATA = "data"                     # Vector DBs, document stores, training data
    MANAGEMENT = "management"         # Secrets, IAM, monitoring, deployment
    EXTERNAL_AI = "external_ai"       # Cloud AI service endpoints (Bedrock, etc.)
 
 
class TrafficDirection(Enum):
    INGRESS = "ingress"
    EGRESS = "egress"
 
 
@dataclass
class NetworkPolicy:
    """A network policy rule for AI workload segmentation."""
    policy_id: str
    name: str
    source_zone: NetworkZone
    destination_zone: NetworkZone
    direction: TrafficDirection
    protocol: str                    # tcp, udp, https
    port_range: str                  # "443", "6333-6334", "8080"
    action: str                      # allow, deny
    justification: str
    data_classification: str         # What type of data flows on this path
 
 
# Define the allowed traffic flows between zones
AI_NETWORK_POLICIES: list[NetworkPolicy] = [
    # Application → Inference: App sends prompts to the model proxy
    NetworkPolicy(
        policy_id="NET-001",
        name="App to inference proxy",
        source_zone=NetworkZone.APPLICATION,
        destination_zone=NetworkZone.INFERENCE,
        direction=TrafficDirection.EGRESS,
        protocol="https",
        port_range="443",
        action="allow",
        justification="Application sends user prompts to the inference proxy for model processing",
        data_classification="confidential",
    ),
    # Inference → External AI: Proxy calls cloud AI service via private endpoint
    NetworkPolicy(
        policy_id="NET-002",
        name="Inference to cloud AI (private endpoint)",
        source_zone=NetworkZone.INFERENCE,
        destination_zone=NetworkZone.EXTERNAL_AI,
        direction=TrafficDirection.EGRESS,
        protocol="https",
        port_range="443",
        action="allow",
        justification="Inference proxy forwards validated requests to Bedrock/Azure OpenAI via VPC endpoint",
        data_classification="confidential",
    ),
    # Inference → Data: RAG pipeline retrieves context from vector database
    NetworkPolicy(
        policy_id="NET-003",
        name="Inference to vector database",
        source_zone=NetworkZone.INFERENCE,
        destination_zone=NetworkZone.DATA,
        direction=TrafficDirection.EGRESS,
        protocol="https",
        port_range="6333-6334",
        action="allow",
        justification="RAG pipeline retrieves relevant context from vector database for augmented generation",
        data_classification="confidential",
    ),
    # Inference → Management: Fetch secrets and send logs
    NetworkPolicy(
        policy_id="NET-004",
        name="Inference to secrets manager",
        source_zone=NetworkZone.INFERENCE,
        destination_zone=NetworkZone.MANAGEMENT,
        direction=TrafficDirection.EGRESS,
        protocol="https",
        port_range="443",
        action="allow",
        justification="Inference layer retrieves API credentials from secrets manager",
        data_classification="restricted",
    ),
    # DENY: Application should never directly access cloud AI endpoints
    NetworkPolicy(
        policy_id="NET-010",
        name="Block app direct AI access",
        source_zone=NetworkZone.APPLICATION,
        destination_zone=NetworkZone.EXTERNAL_AI,
        direction=TrafficDirection.EGRESS,
        protocol="https",
        port_range="443",
        action="deny",
        justification="All model API calls must go through the inference proxy for security inspection",
        data_classification="n/a",
    ),
    # DENY: Data zone should never have internet access
    NetworkPolicy(
        policy_id="NET-011",
        name="Block data zone internet",
        source_zone=NetworkZone.DATA,
        destination_zone=NetworkZone.APPLICATION,
        direction=TrafficDirection.EGRESS,
        protocol="tcp",
        port_range="0-65535",
        action="deny",
        justification="Data zone contains sensitive training data and embeddings; no internet access permitted",
        data_classification="n/a",
    ),
    # DENY: Inference zone should not initiate connections to arbitrary internet hosts
    NetworkPolicy(
        policy_id="NET-012",
        name="Block inference internet egress",
        source_zone=NetworkZone.INFERENCE,
        destination_zone=NetworkZone.APPLICATION,  # Represents "internet"
        direction=TrafficDirection.EGRESS,
        protocol="tcp",
        port_range="0-65535",
        action="deny",
        justification="Inference zone should only communicate with approved endpoints via private links",
        data_classification="n/a",
    ),
]
 
 
def validate_traffic_flow(
    source: NetworkZone,
    destination: NetworkZone,
    port: int,
    protocol: str = "https",
) -> dict:
    """
    Validate whether a traffic flow is permitted by the network policy.
 
    Args:
        source: Originating network zone.
        destination: Target network zone.
        port: Destination port number.
        protocol: Network protocol.
 
    Returns:
        Validation result with matching policy and action.
    """
    matching_policies = []
 
    for policy in AI_NETWORK_POLICIES:
        if policy.source_zone != source or policy.destination_zone != destination:
            continue
        if policy.protocol != protocol and policy.protocol != "tcp":
            continue
 
        # Check port range
        if "-" in policy.port_range:
            low, high = policy.port_range.split("-")
            if int(low) <= port <= int(high):
                matching_policies.append(policy)
        elif port == int(policy.port_range):
            matching_policies.append(policy)
 
    if not matching_policies:
        return {
            "allowed": False,
            "reason": "No matching policy — default deny",
            "matching_policies": [],
        }
 
    # If any explicit deny matches, traffic is blocked
    deny_policies = [p for p in matching_policies if p.action == "deny"]
    if deny_policies:
        return {
            "allowed": False,
            "reason": f"Explicitly denied by policy {deny_policies[0].policy_id}",
            "matching_policies": [p.policy_id for p in deny_policies],
        }
 
    allow_policies = [p for p in matching_policies if p.action == "allow"]
    if allow_policies:
        return {
            "allowed": True,
            "reason": f"Allowed by policy {allow_policies[0].policy_id}",
            "matching_policies": [p.policy_id for p in allow_policies],
            "data_classification": allow_policies[0].data_classification,
        }
 
    return {"allowed": False, "reason": "No allow policy matched", "matching_policies": []}

# network_isolation/aws_bedrock_private.py
"""AWS VPC Endpoint configuration for private Bedrock access."""
 
from __future__ import annotations
 
import json
import logging
from typing import Optional
 
import boto3
 
logger = logging.getLogger(__name__)
 
 
class BedrockPrivateEndpointSetup:
    """Configures private VPC endpoint access for AWS Bedrock."""
 
    def __init__(self, region: str = "us-east-1"):
        self.region = region
        self.ec2_client = boto3.client("ec2", region_name=region)
        self.bedrock_client = boto3.client("bedrock", region_name=region)
 
    def create_vpc_endpoint(
        self,
        vpc_id: str,
        subnet_ids: list[str],
        security_group_ids: list[str],
    ) -> dict:
        """
        Create a VPC Interface Endpoint for AWS Bedrock Runtime.
 
        This eliminates the need for Bedrock API traffic to traverse
        the public internet.
 
        Args:
            vpc_id: ID of the VPC where the endpoint will be created.
            subnet_ids: Subnets in the inference zone for ENI placement.
            security_group_ids: Security groups to attach to the endpoint.
 
        Returns:
            Endpoint creation response with endpoint ID and DNS entries.
        """
        try:
            response = self.ec2_client.create_vpc_endpoint(
                VpcEndpointType="Interface",
                ServiceName=f"com.amazonaws.{self.region}.bedrock-runtime",
                VpcId=vpc_id,
                SubnetIds=subnet_ids,
                SecurityGroupIds=security_group_ids,
                PrivateDnsEnabled=True,  # Enables automatic DNS resolution
                TagSpecifications=[{
                    "ResourceType": "vpc-endpoint",
                    "Tags": [
                        {"Key": "Name", "Value": "bedrock-runtime-private"},
                        {"Key": "SecurityZone", "Value": "inference"},
                        {"Key": "ManagedBy", "Value": "ai-security-team"},
                    ],
                }],
            )
 
            endpoint_id = response["VpcEndpoint"]["VpcEndpointId"]
            dns_entries = response["VpcEndpoint"].get("DnsEntries", [])
 
            logger.info(f"Created Bedrock VPC endpoint: {endpoint_id}")
            return {
                "endpoint_id": endpoint_id,
                "dns_entries": dns_entries,
                "status": "pending",
            }
 
        except Exception as e:
            logger.error(f"Failed to create Bedrock VPC endpoint: {e}")
            raise
 
    def create_endpoint_security_group(
        self,
        vpc_id: str,
        inference_subnet_cidr: str,
    ) -> str:
        """
        Create a security group for the Bedrock VPC endpoint.
 
        Only allows HTTPS traffic from the inference subnet.
        """
        try:
            response = self.ec2_client.create_security_group(
                GroupName="bedrock-endpoint-sg",
                Description="Security group for Bedrock VPC endpoint - inference zone only",
                VpcId=vpc_id,
                TagSpecifications=[{
                    "ResourceType": "security-group",
                    "Tags": [
                        {"Key": "Name", "Value": "bedrock-endpoint-sg"},
                        {"Key": "SecurityZone", "Value": "inference"},
                    ],
                }],
            )
            sg_id = response["GroupId"]
 
            # Allow HTTPS from inference subnet only
            self.ec2_client.authorize_security_group_ingress(
                GroupId=sg_id,
                IpPermissions=[{
                    "IpProtocol": "tcp",
                    "FromPort": 443,
                    "ToPort": 443,
                    "IpRanges": [{
                        "CidrIp": inference_subnet_cidr,
                        "Description": "HTTPS from inference subnet to Bedrock endpoint",
                    }],
                }],
            )
 
            logger.info(f"Created endpoint security group: {sg_id}")
            return sg_id
 
        except Exception as e:
            logger.error(f"Failed to create security group: {e}")
            raise
 
    def create_endpoint_policy(self) -> dict:
        """
        Create a VPC endpoint policy that restricts which Bedrock operations
        and models can be accessed through this endpoint.
        """
        policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Sid": "AllowInvokeApprovedModels",
                    "Effect": "Allow",
                    "Principal": {"AWS": "*"},
                    "Action": [
                        "bedrock:InvokeModel",
                        "bedrock:InvokeModelWithResponseStream",
                    ],
                    "Resource": [
                        f"arn:aws:bedrock:{self.region}::foundation-model/anthropic.claude-*",
                        f"arn:aws:bedrock:{self.region}::foundation-model/amazon.titan-*",
                    ],
                    "Condition": {
                        "StringEquals": {
                            "aws:PrincipalOrgID": "o-your-org-id",
                        },
                    },
                },
                {
                    "Sid": "DenyUnapprovedModels",
                    "Effect": "Deny",
                    "Principal": {"AWS": "*"},
                    "Action": "bedrock:InvokeModel",
                    "Resource": "*",
                    "Condition": {
                        "StringNotLike": {
                            "bedrock:ModelId": [
                                "anthropic.claude-*",
                                "amazon.titan-*",
                            ],
                        },
                    },
                },
            ],
        }
        return policy

# network_isolation/azure_openai_private.py
"""Azure Private Link configuration for Azure OpenAI Service."""
 
from __future__ import annotations
 
import logging
from typing import Optional
 
logger = logging.getLogger(__name__)
 
 
def generate_azure_private_endpoint_bicep(
    resource_group: str,
    vnet_name: str,
    subnet_name: str,
    openai_account_name: str,
    location: str = "eastus",
) -> str:
    """
    Generate Azure Bicep template for Azure OpenAI Private Endpoint.
 
    Returns the Bicep template as a string for deployment.
    """
    return f"""
// Azure OpenAI Private Endpoint - Bicep Template
// Deploys a private endpoint for Azure OpenAI Service in the inference subnet
 
param location string = '{location}'
param vnetName string = '{vnet_name}'
param subnetName string = '{subnet_name}'
param openaiAccountName string = '{openai_account_name}'
 
resource vnet 'Microsoft.Network/virtualNetworks@2023-09-01' existing = {{
  name: vnetName
}}
 
resource subnet 'Microsoft.Network/virtualNetworks/subnets@2023-09-01' existing = {{
  parent: vnet
  name: subnetName
}}
 
resource openaiAccount 'Microsoft.CognitiveServices/accounts@2023-10-01-preview' existing = {{
  name: openaiAccountName
}}
 
resource privateEndpoint 'Microsoft.Network/privateEndpoints@2023-09-01' = {{
  name: '${{openaiAccountName}}-pe'
  location: location
  properties: {{
    subnet: {{
      id: subnet.id
    }}
    privateLinkServiceConnections: [
      {{
        name: '${{openaiAccountName}}-connection'
        properties: {{
          privateLinkServiceId: openaiAccount.id
          groupIds: [
            'account'
          ]
        }}
      }}
    ]
  }}
  tags: {{
    SecurityZone: 'inference'
    ManagedBy: 'ai-security-team'
  }}
}}
 
resource privateDnsZone 'Microsoft.Network/privateDnsZones@2020-06-01' = {{
  name: 'privatelink.openai.azure.com'
  location: 'global'
}}
 
resource privateDnsZoneLink 'Microsoft.Network/privateDnsZones/virtualNetworkLinks@2020-06-01' = {{
  parent: privateDnsZone
  name: '${{vnetName}}-link'
  location: 'global'
  properties: {{
    virtualNetwork: {{
      id: vnet.id
    }}
    registrationEnabled: false
  }}
}}
 
resource privateDnsZoneGroup 'Microsoft.Network/privateEndpoints/privateDnsZoneGroups@2023-09-01' = {{
  parent: privateEndpoint
  name: 'default'
  properties: {{
    privateDnsZoneConfigs: [
      {{
        name: 'openai-dns-config'
        properties: {{
          privateDnsZoneId: privateDnsZone.id
        }}
      }}
    ]
  }}
}}
 
// Disable public network access on the OpenAI account
resource openaiNetworkRules 'Microsoft.CognitiveServices/accounts@2023-10-01-preview' = {{
  name: openaiAccountName
  location: location
  properties: {{
    publicNetworkAccess: 'Disabled'
    networkAcls: {{
      defaultAction: 'Deny'
    }}
  }}
}}
 
output privateEndpointId string = privateEndpoint.id
output privateDnsZoneId string = privateDnsZone.id
"""

# network_isolation/monitoring.py
"""Network monitoring and anomaly detection for AI workloads."""
 
from __future__ import annotations
 
from dataclasses import dataclass
from datetime import datetime
from typing import Optional
 
 
@dataclass
class NetworkAnomaly:
    """A detected network anomaly in AI workload traffic."""
    anomaly_id: str
    severity: str            # critical, high, medium, low
    anomaly_type: str
    source_zone: str
    destination: str
    description: str
    detected_at: datetime
    evidence: dict
 
 
class AINetworkMonitor:
    """Monitors network traffic patterns for AI workload anomalies."""
 
    def __init__(self):
        self.anomalies: list[NetworkAnomaly] = []
        self._anomaly_counter = 0
 
    def analyze_flow_log(self, flow_record: dict) -> Optional[NetworkAnomaly]:
        """
        Analyze a VPC flow log record for AI-specific network anomalies.
 
        Checks for:
        - Unexpected egress from inference zone to internet
        - Data zone attempting external connections
        - Unusually large data transfers (potential model/data exfiltration)
        - Connections to unapproved AI service endpoints
        """
        src_ip = flow_record.get("srcaddr", "")
        dst_ip = flow_record.get("dstaddr", "")
        dst_port = flow_record.get("dstport", 0)
        bytes_transferred = flow_record.get("bytes", 0)
        action = flow_record.get("action", "")
 
        # Check for data zone egress attempts
        if self._is_data_zone(src_ip) and not self._is_internal(dst_ip):
            self._anomaly_counter += 1
            return NetworkAnomaly(
                anomaly_id=f"NET-ANOM-{self._anomaly_counter:04d}",
                severity="critical",
                anomaly_type="data_zone_egress",
                source_zone="data",
                destination=dst_ip,
                description=(
                    f"Data zone IP {src_ip} attempted external connection to {dst_ip}:{dst_port}. "
                    "Data zone should have no internet access."
                ),
                detected_at=datetime.utcnow(),
                evidence=flow_record,
            )
 
        # Check for large data transfers from inference zone
        if self._is_inference_zone(src_ip) and bytes_transferred > 100 * 1024 * 1024:
            self._anomaly_counter += 1
            return NetworkAnomaly(
                anomaly_id=f"NET-ANOM-{self._anomaly_counter:04d}",
                severity="high",
                anomaly_type="large_transfer",
                source_zone="inference",
                destination=dst_ip,
                description=(
                    f"Large data transfer ({bytes_transferred / 1024 / 1024:.1f} MB) "
                    f"from inference zone to {dst_ip}. "
                    "Potential model weights or training data exfiltration."
                ),
                detected_at=datetime.utcnow(),
                evidence=flow_record,
            )
 
        # Check for connections to unapproved AI endpoints
        approved_ai_endpoints = [
            "bedrock-runtime", "openai.azure.com",
            "aiplatform.googleapis.com",
        ]
        if (self._is_inference_zone(src_ip)
                and not self._is_internal(dst_ip)
                and not any(ep in str(dst_ip) for ep in approved_ai_endpoints)):
            self._anomaly_counter += 1
            return NetworkAnomaly(
                anomaly_id=f"NET-ANOM-{self._anomaly_counter:04d}",
                severity="medium",
                anomaly_type="unapproved_endpoint",
                source_zone="inference",
                destination=dst_ip,
                description=(
                    f"Inference zone connected to unapproved external endpoint {dst_ip}:{dst_port}. "
                    "Only approved AI service endpoints should be accessible."
                ),
                detected_at=datetime.utcnow(),
                evidence=flow_record,
            )
 
        return None
 
    def _is_data_zone(self, ip: str) -> bool:
        """Check if an IP belongs to the data zone subnet."""
        return ip.startswith("10.0.3.")
 
    def _is_inference_zone(self, ip: str) -> bool:
        """Check if an IP belongs to the inference zone subnet."""
        return ip.startswith("10.0.2.")
 
    def _is_internal(self, ip: str) -> bool:
        """Check if an IP is in a private RFC1918 range."""
        return (ip.startswith("10.") or ip.startswith("172.16.")
                or ip.startswith("192.168."))