"""
AIPTX Azure Key Validator
=========================

Validates Azure credentials and service principal keys.

Supports:
- Service Principal (client_id + client_secret + tenant_id)
- Managed Identity tokens
- Azure CLI tokens
- Storage Account keys
"""

from __future__ import annotations

import re
import time
import logging
from typing import Any, Optional

import httpx

from ..base import (
    APIKeyInfo,
    ValidationResult,
    BaseKeyValidator,
    KeyType,
    RiskLevel,
)

logger = logging.getLogger(__name__)


class AzureKeyValidator(BaseKeyValidator):
    """
    Validate Azure credentials and enumerate permissions.

    Uses Azure Resource Manager and Graph APIs to validate and discover:
    - Service principal identity
    - Subscription access
    - Role assignments
    """

    LOGIN_URL = "https://login.microsoftonline.com"
    MANAGEMENT_URL = "https://management.azure.com"
    GRAPH_URL = "https://graph.microsoft.com"

    def get_key_type(self) -> KeyType:
        return KeyType.AZURE

    def get_key_patterns(self) -> dict[str, re.Pattern]:
        return {
            # Azure client secret (various formats)
            "azure_client_secret": re.compile(r'[A-Za-z0-9~._-]{34,}'),
            # Azure tenant ID (GUID)
            "azure_tenant_id": re.compile(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}'),
            # Azure client ID (GUID)
            "azure_client_id": re.compile(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}'),
            # Azure storage account key
            "azure_storage_key": re.compile(r'[A-Za-z0-9+/]{86}=='),
            # Azure SAS token
            "azure_sas_token": re.compile(r'sv=\d{4}-\d{2}-\d{2}&[^"\']+sig=[A-Za-z0-9%]+'),
        }

    async def validate(self, key_info: APIKeyInfo) -> ValidationResult:
        """
        Validate Azure credentials.

        Requires client_id, client_secret, and tenant_id for service principal.
        """
        start_time = time.time()

        # Get credentials from key_info
        client_id = key_info.key_id or key_info.metadata.get("client_id")
        client_secret = key_info.secret_value or key_info.key_value
        tenant_id = key_info.metadata.get("tenant_id")

        # Check for storage key
        if key_info.key_value.endswith("==") and len(key_info.key_value) == 88:
            return await self._validate_storage_key(key_info, start_time)

        # Check for SAS token
        if "sv=" in key_info.key_value and "sig=" in key_info.key_value:
            return await self._validate_sas_token(key_info, start_time)

        # Service principal requires all three components
        if not all([client_id, client_secret, tenant_id]):
            duration_ms = (time.time() - start_time) * 1000
            return self._create_error_result(
                key_info,
                error="Service principal requires client_id, client_secret, and tenant_id",
                error_code="INCOMPLETE_CREDENTIALS",
                duration_ms=duration_ms,
            )

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                # Get access token
                access_token, token_info, error = await self._get_access_token(
                    client, client_id, client_secret, tenant_id
                )

                duration_ms = (time.time() - start_time) * 1000

                if error:
                    return self._create_error_result(
                        key_info,
                        error=error,
                        error_code="AUTH_ERROR",
                        duration_ms=duration_ms,
                    )

                # Build identity info
                identity_info = {
                    "client_id": client_id,
                    "tenant_id": tenant_id,
                    "type": "service_principal",
                    "token_type": token_info.get("token_type"),
                    "expires_in": token_info.get("expires_in"),
                }

                # Get service principal details
                sp_info = await self._get_service_principal_info(
                    client, access_token, client_id
                )
                if sp_info:
                    identity_info["display_name"] = sp_info.get("displayName")
                    identity_info["app_id"] = sp_info.get("appId")

                # Enumerate permissions
                permissions, resources, subscriptions = await self._enumerate_permissions(
                    client, access_token, client_id, tenant_id
                )

                if subscriptions:
                    identity_info["subscriptions"] = subscriptions[:5]

                result = ValidationResult(
                    key_info=key_info,
                    is_valid=True,
                    validation_method="OAuth2 Client Credentials",
                    validation_endpoint=f"{self.LOGIN_URL}/{tenant_id}/oauth2/v2.0/token",
                    identity_info=identity_info,
                    scopes=[],
                    permissions=permissions,
                    resources_accessible=resources,
                    raw_response=token_info,
                    validated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    validation_duration_ms=duration_ms,
                )

                result.risk_level = self._assess_azure_risk(permissions, subscriptions, identity_info)
                result.exploitation_potential = self._get_exploitation_potential(result)
                result.attack_vectors = self._get_attack_vectors(permissions, subscriptions)

                return result

        except httpx.TimeoutException:
            duration_ms = (time.time() - start_time) * 1000
            return self._create_error_result(
                key_info,
                error="Request timeout",
                error_code="TIMEOUT",
                duration_ms=duration_ms,
            )
        except Exception as e:
            duration_ms = (time.time() - start_time) * 1000
            logger.error(f"Azure validation error: {e}")
            return self._create_error_result(
                key_info,
                error=str(e),
                error_code="ERROR",
                duration_ms=duration_ms,
            )

    async def _validate_storage_key(
        self,
        key_info: APIKeyInfo,
        start_time: float,
    ) -> ValidationResult:
        """Validate Azure Storage Account key."""
        duration_ms = (time.time() - start_time) * 1000

        storage_account = key_info.metadata.get("storage_account")
        if not storage_account:
            return self._create_error_result(
                key_info,
                error="Storage account name required for storage key validation",
                error_code="MISSING_ACCOUNT_NAME",
                duration_ms=duration_ms,
            )

        # Storage keys are valid if they match the format
        # Full validation would require making a storage request
        return ValidationResult(
            key_info=key_info,
            is_valid=True,  # Assumed valid based on format
            validation_method="Format validation",
            validation_endpoint="N/A",
            identity_info={
                "storage_account": storage_account,
                "type": "storage_account_key",
            },
            scopes=[],
            permissions=["Storage Account access"],
            resources_accessible=[f"storage:{storage_account}"],
            validated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            validation_duration_ms=duration_ms,
            risk_level=RiskLevel.HIGH,
            exploitation_potential="High - Full access to storage account data",
            attack_vectors=[
                "Access all blobs, tables, queues, and files",
                "Modify or delete storage data",
                "Upload malicious content",
            ],
        )

    async def _validate_sas_token(
        self,
        key_info: APIKeyInfo,
        start_time: float,
    ) -> ValidationResult:
        """Validate Azure SAS token."""
        duration_ms = (time.time() - start_time) * 1000

        # Parse SAS token parameters
        params = {}
        for part in key_info.key_value.split("&"):
            if "=" in part:
                k, v = part.split("=", 1)
                params[k] = v

        # Determine permissions from 'sp' parameter
        permissions = []
        sp = params.get("sp", "")
        perm_map = {
            "r": "Read",
            "w": "Write",
            "d": "Delete",
            "l": "List",
            "a": "Add",
            "c": "Create",
        }
        for char in sp:
            if char in perm_map:
                permissions.append(perm_map[char])

        # Check expiry
        se = params.get("se", "")

        return ValidationResult(
            key_info=key_info,
            is_valid=True,
            validation_method="SAS token parsing",
            validation_endpoint="N/A",
            identity_info={
                "type": "sas_token",
                "expiry": se,
                "signed_resource": params.get("sr", "unknown"),
            },
            scopes=[],
            permissions=permissions,
            resources_accessible=[f"resource:{params.get('sr', 'unknown')}"],
            validated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            validation_duration_ms=duration_ms,
            risk_level=RiskLevel.MEDIUM if "w" in sp or "d" in sp else RiskLevel.LOW,
            exploitation_potential="Depends on SAS permissions",
        )

    async def _get_access_token(
        self,
        client: httpx.AsyncClient,
        client_id: str,
        client_secret: str,
        tenant_id: str,
    ) -> tuple[Optional[str], dict, Optional[str]]:
        """Get Azure access token using client credentials."""
        try:
            response = await client.post(
                f"{self.LOGIN_URL}/{tenant_id}/oauth2/v2.0/token",
                data={
                    "grant_type": "client_credentials",
                    "client_id": client_id,
                    "client_secret": client_secret,
                    "scope": f"{self.MANAGEMENT_URL}/.default",
                },
            )

            if response.status_code == 200:
                token_data = response.json()
                return token_data.get("access_token"), token_data, None
            else:
                error_data = response.json()
                error_desc = error_data.get("error_description", "Authentication failed")
                return None, {}, error_desc

        except Exception as e:
            return None, {}, str(e)

    async def _get_service_principal_info(
        self,
        client: httpx.AsyncClient,
        access_token: str,
        client_id: str,
    ) -> Optional[dict]:
        """Get service principal details from Graph API."""
        try:
            # Need Graph API token
            response = await client.get(
                f"{self.GRAPH_URL}/v1.0/servicePrincipals",
                headers={"Authorization": f"Bearer {access_token}"},
                params={"$filter": f"appId eq '{client_id}'"},
            )

            if response.status_code == 200:
                data = response.json()
                principals = data.get("value", [])
                if principals:
                    return principals[0]
        except Exception as e:
            logger.debug(f"Failed to get SP info: {e}")

        return None

    async def _enumerate_permissions(
        self,
        client: httpx.AsyncClient,
        access_token: str,
        client_id: str,
        tenant_id: str,
    ) -> tuple[list[str], list[str], list[str]]:
        """Enumerate Azure permissions and resources."""
        permissions = []
        resources = []
        subscriptions = []

        headers = {"Authorization": f"Bearer {access_token}"}

        # List subscriptions
        try:
            response = await client.get(
                f"{self.MANAGEMENT_URL}/subscriptions?api-version=2020-01-01",
                headers=headers,
            )
            if response.status_code == 200:
                data = response.json()
                for sub in data.get("value", []):
                    sub_id = sub.get("subscriptionId")
                    sub_name = sub.get("displayName")
                    subscriptions.append(f"{sub_name} ({sub_id})")
                    resources.append(f"subscription:{sub_id}")

                    # Get role assignments for this subscription
                    await self._get_role_assignments(
                        client, headers, sub_id, permissions
                    )
        except Exception as e:
            logger.debug(f"Failed to list subscriptions: {e}")

        # List resource groups (if we have subscriptions)
        for sub in subscriptions[:3]:  # Limit to first 3
            sub_id = sub.split("(")[-1].rstrip(")")
            try:
                response = await client.get(
                    f"{self.MANAGEMENT_URL}/subscriptions/{sub_id}/resourcegroups?api-version=2021-04-01",
                    headers=headers,
                )
                if response.status_code == 200:
                    data = response.json()
                    for rg in data.get("value", [])[:5]:
                        resources.append(f"resourceGroup:{rg.get('name')}")
            except Exception:
                pass

        return permissions, resources, subscriptions

    async def _get_role_assignments(
        self,
        client: httpx.AsyncClient,
        headers: dict,
        subscription_id: str,
        permissions: list[str],
    ) -> None:
        """Get role assignments for subscription."""
        try:
            response = await client.get(
                f"{self.MANAGEMENT_URL}/subscriptions/{subscription_id}/providers/Microsoft.Authorization/roleAssignments?api-version=2022-04-01",
                headers=headers,
            )
            if response.status_code == 200:
                data = response.json()
                role_ids = set()
                for assignment in data.get("value", []):
                    role_id = assignment.get("properties", {}).get("roleDefinitionId", "")
                    if role_id:
                        role_ids.add(role_id.split("/")[-1])

                # Map common role IDs to names
                role_map = {
                    "8e3af657-a8ff-443c-a75c-2fe8c4bcb635": "Owner",
                    "b24988ac-6180-42a0-ab88-20f7382dd24c": "Contributor",
                    "acdd72a7-3385-48ef-bd42-f606fba81ae7": "Reader",
                    "ba92f5b4-2d11-453d-a403-e96b0029c9fe": "Storage Blob Data Contributor",
                    "2a2b9908-6ea1-4ae2-8e65-a410df84e7d1": "Storage Blob Data Reader",
                }

                for role_id in role_ids:
                    role_name = role_map.get(role_id, f"Role:{role_id[:8]}")
                    if role_name not in permissions:
                        permissions.append(role_name)
        except Exception as e:
            logger.debug(f"Failed to get role assignments: {e}")

    def _assess_azure_risk(
        self,
        permissions: list[str],
        subscriptions: list[str],
        identity_info: dict,
    ) -> RiskLevel:
        """Assess risk level based on Azure credentials."""
        # Critical: Owner or Contributor on subscriptions
        if "Owner" in permissions:
            return RiskLevel.CRITICAL

        if "Contributor" in permissions and subscriptions:
            return RiskLevel.CRITICAL

        # High: Multiple subscriptions or storage access
        if len(subscriptions) > 1 or any("Storage" in p for p in permissions):
            return RiskLevel.HIGH

        # Medium: Reader access to resources
        if "Reader" in permissions or subscriptions:
            return RiskLevel.MEDIUM

        # Low: Limited access
        if permissions:
            return RiskLevel.LOW

        return RiskLevel.INFO

    def _get_exploitation_potential(self, result: ValidationResult) -> str:
        """Get exploitation potential description."""
        if result.risk_level == RiskLevel.CRITICAL:
            return "Critical - Owner/Contributor access to Azure subscriptions"
        elif result.risk_level == RiskLevel.HIGH:
            return "High - Access to multiple subscriptions or storage"
        elif result.risk_level == RiskLevel.MEDIUM:
            return "Medium - Reader access to Azure resources"
        elif result.risk_level == RiskLevel.LOW:
            return "Low - Limited Azure permissions"
        return "Minimal - Could not enumerate permissions"

    def _get_attack_vectors(
        self,
        permissions: list[str],
        subscriptions: list[str],
    ) -> list[str]:
        """Identify potential attack vectors."""
        vectors = []

        if "Owner" in permissions:
            vectors.append("Full control over Azure resources")
            vectors.append("Create VMs for crypto mining")
            vectors.append("Access Key Vault secrets")
            vectors.append("Modify IAM for persistence")

        if "Contributor" in permissions:
            vectors.append("Modify most Azure resources")
            vectors.append("Deploy malicious applications")
            vectors.append("Access and modify storage")

        if any("Storage" in p for p in permissions):
            vectors.append("Access Azure Blob Storage")
            if "Contributor" in str(permissions):
                vectors.append("Upload malicious blobs")

        if "Reader" in permissions:
            vectors.append("Enumerate Azure resources")
            vectors.append("Discover sensitive configurations")

        if subscriptions:
            vectors.append(f"Access to {len(subscriptions)} subscription(s)")

        return vectors
