Networking and Security

def list_by_database_account(
    self,
    resource_group_name: str,
    account_name: str
) -> ItemPaged[PrivateEndpointConnection]:
    """
    List private endpoint connections for a database account.
    
    Parameters:
    - resource_group_name: Name of the resource group
    - account_name: Name of the Cosmos DB account
    
    Returns:
    ItemPaged[PrivateEndpointConnection]: Paginated list of private endpoint connections
    """

def get(
    self,
    resource_group_name: str,
    account_name: str,
    private_endpoint_connection_name: str
) -> PrivateEndpointConnection:
    """
    Get properties of a private endpoint connection.
    
    Parameters:
    - resource_group_name: Name of the resource group
    - account_name: Name of the Cosmos DB account
    - private_endpoint_connection_name: Name of the private endpoint connection
    
    Returns:
    PrivateEndpointConnection: Private endpoint connection properties and status
    """

def begin_create_or_update(
    self,
    resource_group_name: str,
    account_name: str,
    private_endpoint_connection_name: str,
    parameters: PrivateEndpointConnection
) -> LROPoller[PrivateEndpointConnection]:
    """
    Create or update a private endpoint connection (Long Running Operation).
    
    Parameters:
    - resource_group_name: Name of the resource group
    - account_name: Name of the Cosmos DB account
    - private_endpoint_connection_name: Name of the connection to create/update
    - parameters: Private endpoint connection configuration
    
    Returns:
    LROPoller[PrivateEndpointConnection]: Poller for monitoring operation progress
    """

def begin_delete(
    self,
    resource_group_name: str,
    account_name: str,
    private_endpoint_connection_name: str
) -> LROPoller[None]:
    """
    Delete a private endpoint connection (Long Running Operation).
    
    Parameters:
    - resource_group_name: Name of the resource group
    - account_name: Name of the Cosmos DB account
    - private_endpoint_connection_name: Name of the connection to delete
    
    Returns:
    LROPoller[None]: Poller for monitoring operation progress
    """

def list_by_database_account(
    self,
    resource_group_name: str,
    account_name: str
) -> ItemPaged[PrivateLinkResource]:
    """
    List private link resources for a database account.
    
    Parameters:
    - resource_group_name: Name of the resource group
    - account_name: Name of the Cosmos DB account
    
    Returns:
    ItemPaged[PrivateLinkResource]: Paginated list of private link resources
    """

def get(
    self,
    resource_group_name: str,
    account_name: str,
    group_name: str
) -> PrivateLinkResource:
    """
    Get properties of a private link resource.
    
    Parameters:
    - resource_group_name: Name of the resource group
    - account_name: Name of the Cosmos DB account
    - group_name: Name of the private link resource group
    
    Returns:
    PrivateLinkResource: Private link resource properties and supported sub-resources
    """

from azure.mgmt.cosmosdb import CosmosDBManagementClient
from azure.mgmt.cosmosdb.models import (
    PrivateEndpointConnection,
    PrivateLinkServiceConnectionState,
    PrivateEndpoint
)
from azure.identity import DefaultAzureCredential

client = CosmosDBManagementClient(DefaultAzureCredential(), "subscription-id")

# Create private endpoint connection
private_endpoint_config = PrivateEndpointConnection(
    private_endpoint=PrivateEndpoint(
        id="/subscriptions/sub-id/resourceGroups/network-rg/providers/Microsoft.Network/privateEndpoints/my-private-endpoint"
    ),
    private_link_service_connection_state=PrivateLinkServiceConnectionState(
        status="Approved",
        description="Approved for production use",
        actions_required="None"
    )
)

# Create the connection
poller = client.private_endpoint_connections.begin_create_or_update(
    "my-resource-group",
    "my-cosmos-account",
    "my-private-endpoint-connection",
    private_endpoint_config
)

connection = poller.result()
print(f"Private endpoint connection status: {connection.private_link_service_connection_state.status}")

# List available private link resources
private_link_resources = client.private_link_resources.list_by_database_account(
    "my-resource-group",
    "my-cosmos-account"
)

for resource in private_link_resources:
    print(f"Group ID: {resource.group_id}")
    print(f"Required Members: {resource.required_members}")
    print(f"Required Zone Names: {resource.required_zone_names}")
    print("---")

# Get specific private link resource details
sql_resource = client.private_link_resources.get(
    "my-resource-group",
    "my-cosmos-account", 
    "Sql"  # Common group IDs: "Sql", "MongoDB", "Cassandra", "Gremlin", "Table"
)

print(f"SQL Private Link Resource:")
print(f"  Group ID: {sql_resource.group_id}")
print(f"  Required Members: {sql_resource.required_members}")
print(f"  Required DNS Zone Names: {sql_resource.required_zone_names}")

# List all private endpoint connections
connections = client.private_endpoint_connections.list_by_database_account(
    "my-resource-group",
    "my-cosmos-account"
)

for connection in connections:
    print(f"Connection: {connection.name}")
    print(f"Status: {connection.private_link_service_connection_state.status}")
    print(f"Provisioning State: {connection.provisioning_state}")
    
    # Approve pending connections
    if connection.private_link_service_connection_state.status == "Pending":
        print(f"Approving connection: {connection.name}")
        
        # Update connection state to approved
        connection.private_link_service_connection_state.status = "Approved"
        connection.private_link_service_connection_state.description = "Auto-approved by management script"
        
        poller = client.private_endpoint_connections.begin_create_or_update(
            "my-resource-group",
            "my-cosmos-account",
            connection.name.split('/')[-1],  # Extract connection name from full resource path
            connection
        )
        
        updated_connection = poller.result()
        print(f"Updated status: {updated_connection.private_link_service_connection_state.status}")

print("---")

# Reject a specific private endpoint connection
connection_name = "unwanted-private-endpoint-connection"

try:
    connection = client.private_endpoint_connections.get(
        "my-resource-group",
        "my-cosmos-account",
        connection_name
    )
    
    # Reject the connection
    connection.private_link_service_connection_state.status = "Rejected"
    connection.private_link_service_connection_state.description = "Connection rejected due to security policy"
    
    poller = client.private_endpoint_connections.begin_create_or_update(
        "my-resource-group",
        "my-cosmos-account",
        connection_name,
        connection
    )
    
    rejected_connection = poller.result()
    print(f"Connection rejected: {rejected_connection.private_link_service_connection_state.status}")
    
    # Optionally delete the rejected connection
    delete_poller = client.private_endpoint_connections.begin_delete(
        "my-resource-group",
        "my-cosmos-account", 
        connection_name
    )
    
    delete_poller.result()
    print(f"Connection {connection_name} deleted")
    
except Exception as e:
    print(f"Connection {connection_name} not found or error occurred: {e}")

# Monitor connection health and troubleshoot issues
connections = client.private_endpoint_connections.list_by_database_account(
    "my-resource-group",
    "my-cosmos-account"
)

print("Private Endpoint Connection Health Report:")
print("=" * 50)

for connection in connections:
    print(f"Connection: {connection.name}")
    print(f"  Private Endpoint ID: {connection.private_endpoint.id if connection.private_endpoint else 'None'}")
    print(f"  Connection Status: {connection.private_link_service_connection_state.status}")
    print(f"  Provisioning State: {connection.provisioning_state}")
    print(f"  Description: {connection.private_link_service_connection_state.description}")
    print(f"  Actions Required: {connection.private_link_service_connection_state.actions_required}")
    
    # Check for any issues
    if connection.provisioning_state != "Succeeded":
        print(f"  ⚠️  WARNING: Provisioning state is {connection.provisioning_state}")
    
    if connection.private_link_service_connection_state.status == "Rejected":
        print(f"  ❌ ERROR: Connection is rejected")
    elif connection.private_link_service_connection_state.status == "Pending":
        print(f"  ⏳ PENDING: Connection requires approval")
    elif connection.private_link_service_connection_state.status == "Approved":
        print(f"  ✅ SUCCESS: Connection is approved and active")
    
    print("---")

class PrivateEndpointConnection:
    """Private endpoint connection configuration and status."""
    id: str  # Resource ID of the connection
    name: str  # Connection name
    type: str  # Resource type
    private_endpoint: PrivateEndpoint  # Private endpoint details
    private_link_service_connection_state: PrivateLinkServiceConnectionState  # Connection state
    group_id: str  # Private link sub-resource group ID
    provisioning_state: str  # Provisioning state ("Succeeded", "Failed", "Creating", etc.)

class PrivateEndpoint:
    """Private endpoint resource reference."""
    id: str  # Resource ID of the private endpoint

class PrivateLinkServiceConnectionState:
    """Private link service connection state and status."""
    status: str  # Connection status ("Pending", "Approved", "Rejected")
    description: str  # Human-readable description of the connection state
    actions_required: str  # Required actions (if any) to complete the connection

class PrivateLinkResource:
    """Private link resource definition and capabilities."""
    id: str  # Resource ID
    name: str  # Resource name
    type: str  # Resource type
    group_id: str  # Private link resource group identifier
    required_members: List[str]  # Required members for the private link resource
    required_zone_names: List[str]  # Required private DNS zone names for the resource

# Common Private Link Group IDs
class PrivateLinkGroupIds:
    """Common private link group identifiers for different Cosmos DB APIs."""
    SQL = "Sql"  # SQL (Core) API
    MONGODB = "MongoDB"  # MongoDB API
    CASSANDRA = "Cassandra"  # Cassandra API
    GREMLIN = "Gremlin"  # Gremlin (Graph) API
    TABLE = "Table"  # Table API
    ANALYTICAL_STORAGE = "AnalyticalStorage"  # Azure Synapse Link analytical storage

# Connection Status Values
class ConnectionStatus:
    """Private endpoint connection status values."""
    PENDING = "Pending"  # Connection awaiting approval
    APPROVED = "Approved"  # Connection approved and active
    REJECTED = "Rejected"  # Connection rejected

# Provisioning States
class ProvisioningState:
    """Private endpoint connection provisioning states."""
    SUCCEEDED = "Succeeded"  # Successfully provisioned
    FAILED = "Failed"  # Provisioning failed
    CREATING = "Creating"  # Currently being created
    UPDATING = "Updating"  # Currently being updated
    DELETING = "Deleting"  # Currently being deleted
    CANCELED = "Canceled"  # Operation was canceled

# Configure Cosmos DB account for private endpoint only access
from azure.mgmt.cosmosdb.models import DatabaseAccountUpdateParameters

# Disable public network access
update_params = DatabaseAccountUpdateParameters(
    public_network_access="Disabled",
    network_acl_bypass="None"  # Disable bypass for Azure services
)

poller = client.database_accounts.begin_update(
    "my-resource-group",
    "my-cosmos-account",
    update_params
)

updated_account = poller.result()
print(f"Public network access: {updated_account.public_network_access}")

def validate_private_endpoint_setup(resource_group: str, account_name: str):
    """Validate private endpoint configuration for a Cosmos DB account."""
    
    # Check account network access settings
    account = client.database_accounts.get(resource_group, account_name)
    
    if account.public_network_access == "Enabled":
        print("⚠️  WARNING: Public network access is enabled")
    else:
        print("✅ Public network access is disabled")
    
    # Check private endpoint connections
    connections = client.private_endpoint_connections.list_by_database_account(
        resource_group, account_name
    )
    
    approved_connections = []
    for connection in connections:
        if connection.private_link_service_connection_state.status == "Approved":
            approved_connections.append(connection)
    
    if len(approved_connections) == 0:
        print("❌ ERROR: No approved private endpoint connections found")
    else:
        print(f"✅ Found {len(approved_connections)} approved private endpoint connections")
    
    # Check available private link resources
    private_link_resources = client.private_link_resources.list_by_database_account(
        resource_group, account_name
    )
    
    print(f"Available private link resources: {[r.group_id for r in private_link_resources]}")

# Run validation
validate_private_endpoint_setup("my-resource-group", "my-cosmos-account")