"""
AIPTX GCP Key Validator
=======================

Validates Google Cloud Platform service account keys and API keys.

Supports:
- Service Account JSON keys
- API Keys
- OAuth2 tokens
"""

from __future__ import annotations

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

import httpx

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

logger = logging.getLogger(__name__)


class GCPKeyValidator(BaseKeyValidator):
    """
    Validate GCP credentials and enumerate permissions.

    Uses Google Cloud APIs to validate credentials and discover:
    - Service account identity
    - Project access
    - IAM permissions
    """

    TOKEN_URL = "https://oauth2.googleapis.com/token"
    IAM_URL = "https://iam.googleapis.com/v1"
    CRM_URL = "https://cloudresourcemanager.googleapis.com/v1"

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

    def get_key_patterns(self) -> dict[str, re.Pattern]:
        return {
            # Service account JSON key indicators
            "gcp_service_account": re.compile(r'"type"\s*:\s*"service_account"'),
            "gcp_private_key_id": re.compile(r'"private_key_id"\s*:\s*"[a-f0-9]{40}"'),
            "gcp_client_email": re.compile(r'"client_email"\s*:\s*"[^"]+@[^"]+\.iam\.gserviceaccount\.com"'),
            # API Key
            "gcp_api_key": re.compile(r'AIza[0-9A-Za-z_-]{35}'),
        }

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

        Handles both service account JSON and API keys.
        """
        start_time = time.time()

        # Check if it's a service account JSON
        if '"type"' in key_info.key_value and '"service_account"' in key_info.key_value:
            return await self._validate_service_account(key_info, start_time)

        # Check if it's an API key
        if key_info.key_value.startswith("AIza"):
            return await self._validate_api_key(key_info, start_time)

        # Try as OAuth token
        return await self._validate_oauth_token(key_info, start_time)

    async def _validate_service_account(
        self,
        key_info: APIKeyInfo,
        start_time: float,
    ) -> ValidationResult:
        """Validate service account JSON key."""
        try:
            # Parse the JSON key
            try:
                sa_key = json.loads(key_info.key_value)
            except json.JSONDecodeError:
                duration_ms = (time.time() - start_time) * 1000
                return self._create_error_result(
                    key_info,
                    error="Invalid JSON format for service account key",
                    error_code="INVALID_JSON",
                    duration_ms=duration_ms,
                )

            # Extract key components
            client_email = sa_key.get("client_email")
            private_key = sa_key.get("private_key")
            project_id = sa_key.get("project_id")

            if not all([client_email, private_key, project_id]):
                duration_ms = (time.time() - start_time) * 1000
                return self._create_error_result(
                    key_info,
                    error="Missing required fields in service account key",
                    error_code="INCOMPLETE_KEY",
                    duration_ms=duration_ms,
                )

            # Get access token using JWT
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                access_token, error = await self._get_access_token(
                    client, sa_key
                )

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

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

                # Get service account info
                identity_info = {
                    "client_email": client_email,
                    "project_id": project_id,
                    "private_key_id": sa_key.get("private_key_id"),
                    "type": "service_account",
                }

                # Enumerate permissions
                permissions, resources = await self._enumerate_permissions(
                    client, access_token, project_id, client_email
                )

                result = ValidationResult(
                    key_info=key_info,
                    is_valid=True,
                    validation_method="OAuth2 JWT",
                    validation_endpoint=self.TOKEN_URL,
                    identity_info=identity_info,
                    scopes=[],
                    permissions=permissions,
                    resources_accessible=resources,
                    raw_response={"token_acquired": True, "project_id": project_id},
                    validated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    validation_duration_ms=duration_ms,
                )

                result.risk_level = self._assess_gcp_risk(permissions, identity_info)
                result.exploitation_potential = self._get_exploitation_potential(result)
                result.attack_vectors = self._get_attack_vectors(permissions, project_id)

                return result

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

    async def _validate_api_key(
        self,
        key_info: APIKeyInfo,
        start_time: float,
    ) -> ValidationResult:
        """Validate GCP API key."""
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                # Test API key by calling a simple API
                # Using Cloud Resource Manager API
                response = await client.get(
                    f"{self.CRM_URL}/projects",
                    params={"key": key_info.key_value},
                )

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

                if response.status_code == 200:
                    data = response.json()
                    projects = [p.get("projectId") for p in data.get("projects", [])]

                    identity_info = {
                        "type": "api_key",
                        "projects_accessible": len(projects),
                    }

                    result = ValidationResult(
                        key_info=key_info,
                        is_valid=True,
                        validation_method="API Key Test",
                        validation_endpoint=f"{self.CRM_URL}/projects",
                        identity_info=identity_info,
                        scopes=[],
                        permissions=["cloudresourcemanager.projects.list"],
                        resources_accessible=[f"project:{p}" for p in projects[:10]],
                        raw_response=data,
                        validated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                        validation_duration_ms=duration_ms,
                        risk_level=RiskLevel.MEDIUM,
                        exploitation_potential="Medium - API key with project access",
                    )

                    return result

                elif response.status_code == 400:
                    error_data = response.json()
                    error_msg = error_data.get("error", {}).get("message", "Invalid API key")
                    return self._create_error_result(
                        key_info,
                        error=error_msg,
                        error_code="400_BAD_REQUEST",
                        duration_ms=duration_ms,
                    )
                elif response.status_code == 403:
                    # API key is valid but lacks permission
                    return ValidationResult(
                        key_info=key_info,
                        is_valid=True,
                        validation_method="API Key Test",
                        validation_endpoint=f"{self.CRM_URL}/projects",
                        identity_info={"type": "api_key", "restricted": True},
                        scopes=[],
                        permissions=["Restricted API key"],
                        resources_accessible=[],
                        validated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                        validation_duration_ms=duration_ms,
                        risk_level=RiskLevel.LOW,
                        exploitation_potential="Low - Restricted API key",
                    )
                else:
                    return self._create_error_result(
                        key_info,
                        error=f"API returned status {response.status_code}",
                        error_code=str(response.status_code),
                        duration_ms=duration_ms,
                    )

        except Exception as e:
            duration_ms = (time.time() - start_time) * 1000
            return self._create_error_result(
                key_info,
                error=str(e),
                error_code="ERROR",
                duration_ms=duration_ms,
            )

    async def _validate_oauth_token(
        self,
        key_info: APIKeyInfo,
        start_time: float,
    ) -> ValidationResult:
        """Validate OAuth2 access token."""
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                # Validate token using tokeninfo endpoint
                response = await client.get(
                    "https://oauth2.googleapis.com/tokeninfo",
                    params={"access_token": key_info.key_value},
                )

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

                if response.status_code == 200:
                    token_info = response.json()

                    identity_info = {
                        "email": token_info.get("email"),
                        "email_verified": token_info.get("email_verified"),
                        "expires_in": token_info.get("expires_in"),
                        "audience": token_info.get("aud"),
                        "type": "oauth_token",
                    }

                    scopes = token_info.get("scope", "").split()

                    result = ValidationResult(
                        key_info=key_info,
                        is_valid=True,
                        validation_method="OAuth2 tokeninfo",
                        validation_endpoint="https://oauth2.googleapis.com/tokeninfo",
                        identity_info=identity_info,
                        scopes=scopes,
                        permissions=self._scopes_to_permissions(scopes),
                        resources_accessible=[],
                        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_gcp_risk(result.permissions, identity_info)
                    result.exploitation_potential = self._get_exploitation_potential(result)

                    return result

                else:
                    return self._create_error_result(
                        key_info,
                        error="Invalid or expired OAuth token",
                        error_code="INVALID_TOKEN",
                        duration_ms=duration_ms,
                    )

        except Exception as e:
            duration_ms = (time.time() - start_time) * 1000
            return self._create_error_result(
                key_info,
                error=str(e),
                error_code="ERROR",
                duration_ms=duration_ms,
            )

    async def _get_access_token(
        self,
        client: httpx.AsyncClient,
        sa_key: dict,
    ) -> tuple[Optional[str], Optional[str]]:
        """Get access token using service account JWT."""
        try:
            import jwt
            import time as time_module

            # Create JWT
            now = int(time_module.time())
            payload = {
                "iss": sa_key["client_email"],
                "sub": sa_key["client_email"],
                "aud": self.TOKEN_URL,
                "iat": now,
                "exp": now + 3600,
                "scope": "https://www.googleapis.com/auth/cloud-platform",
            }

            # Sign with private key
            assertion = jwt.encode(
                payload,
                sa_key["private_key"],
                algorithm="RS256",
            )

            # Exchange for access token
            response = await client.post(
                self.TOKEN_URL,
                data={
                    "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
                    "assertion": assertion,
                },
            )

            if response.status_code == 200:
                token_data = response.json()
                return token_data.get("access_token"), None
            else:
                error_data = response.json()
                return None, error_data.get("error_description", "Token exchange failed")

        except ImportError:
            # PyJWT not available, try simplified approach
            logger.warning("PyJWT not available, skipping JWT auth")
            return None, "PyJWT library required for service account validation"
        except Exception as e:
            return None, str(e)

    async def _enumerate_permissions(
        self,
        client: httpx.AsyncClient,
        access_token: str,
        project_id: str,
        service_account: str,
    ) -> tuple[list[str], list[str]]:
        """Enumerate GCP permissions for service account."""
        permissions = []
        resources = []

        headers = {
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
        }

        # Get IAM policy for project
        try:
            response = await client.post(
                f"{self.CRM_URL}/projects/{project_id}:getIamPolicy",
                headers=headers,
                json={},
            )
            if response.status_code == 200:
                policy = response.json()
                for binding in policy.get("bindings", []):
                    role = binding.get("role", "")
                    members = binding.get("members", [])
                    if f"serviceAccount:{service_account}" in members:
                        permissions.append(role)
                resources.append(f"project:{project_id}")
        except Exception as e:
            logger.debug(f"Failed to get IAM policy: {e}")

        # List GCS buckets
        try:
            response = await client.get(
                f"https://storage.googleapis.com/storage/v1/b",
                headers=headers,
                params={"project": project_id},
            )
            if response.status_code == 200:
                data = response.json()
                buckets = [b.get("name") for b in data.get("items", [])]
                resources.extend([f"gs://{b}" for b in buckets[:10]])
                permissions.append("storage.buckets.list")
        except Exception:
            pass

        # List Compute instances
        try:
            response = await client.get(
                f"https://compute.googleapis.com/compute/v1/projects/{project_id}/aggregated/instances",
                headers=headers,
            )
            if response.status_code == 200:
                data = response.json()
                instance_count = sum(
                    len(zone.get("instances", []))
                    for zone in data.get("items", {}).values()
                )
                if instance_count:
                    resources.append(f"compute_instances:{instance_count}")
                permissions.append("compute.instances.list")
        except Exception:
            pass

        return permissions, resources

    def _scopes_to_permissions(self, scopes: list[str]) -> list[str]:
        """Convert OAuth scopes to permission descriptions."""
        scope_map = {
            "https://www.googleapis.com/auth/cloud-platform": "Full GCP access",
            "https://www.googleapis.com/auth/compute": "Compute Engine access",
            "https://www.googleapis.com/auth/devstorage.full_control": "Full Cloud Storage access",
            "https://www.googleapis.com/auth/devstorage.read_write": "Cloud Storage read/write",
            "https://www.googleapis.com/auth/devstorage.read_only": "Cloud Storage read-only",
            "https://www.googleapis.com/auth/bigquery": "BigQuery access",
            "https://www.googleapis.com/auth/pubsub": "Pub/Sub access",
            "https://www.googleapis.com/auth/sqlservice.admin": "Cloud SQL admin",
        }

        permissions = []
        for scope in scopes:
            desc = scope_map.get(scope)
            if desc:
                permissions.append(desc)
            else:
                # Extract service name from scope URL
                parts = scope.split("/")
                if parts:
                    permissions.append(parts[-1])

        return permissions

    def _assess_gcp_risk(
        self,
        permissions: list[str],
        identity_info: dict,
    ) -> RiskLevel:
        """Assess risk level based on GCP credentials."""
        # Check for dangerous roles
        critical_roles = ["roles/owner", "roles/editor", "roles/iam.securityAdmin"]
        if any(r in permissions for r in critical_roles):
            return RiskLevel.CRITICAL

        # Check for storage/compute admin
        high_roles = ["roles/storage.admin", "roles/compute.admin", "roles/bigquery.admin"]
        if any(r in permissions for r in high_roles):
            return RiskLevel.HIGH

        # Check for read access
        if any("storage" in p.lower() or "compute" in p.lower() for p in permissions):
            return RiskLevel.MEDIUM

        if permissions or identity_info.get("type") == "service_account":
            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/Editor access to GCP project"
        elif result.risk_level == RiskLevel.HIGH:
            return "High - Admin access to storage/compute resources"
        elif result.risk_level == RiskLevel.MEDIUM:
            return "Medium - Read access to cloud resources"
        elif result.risk_level == RiskLevel.LOW:
            return "Low - Limited GCP permissions"
        return "Minimal - Could not enumerate permissions"

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

        if any("owner" in p.lower() or "editor" in p.lower() for p in permissions):
            vectors.append("Full project control - can modify all resources")
            vectors.append("Create service accounts for persistence")
            vectors.append("Access Secret Manager and KMS")

        if any("storage" in p.lower() for p in permissions):
            vectors.append("Access/exfiltrate Cloud Storage data")
            if any("admin" in p.lower() or "write" in p.lower() for p in permissions):
                vectors.append("Upload malicious files to buckets")

        if any("compute" in p.lower() for p in permissions):
            vectors.append("Enumerate Compute Engine instances")
            if any("admin" in p.lower() for p in permissions):
                vectors.append("Create VMs for crypto mining")
                vectors.append("SSH into existing instances")

        if any("bigquery" in p.lower() for p in permissions):
            vectors.append("Query BigQuery datasets")
            vectors.append("Exfiltrate large data sets")

        if any("iam" in p.lower() for p in permissions):
            vectors.append("Enumerate IAM policies and service accounts")
            if any("admin" in p.lower() for p in permissions):
                vectors.append("Create backdoor service accounts")

        return vectors
