#!/usr/bin/env python3
"""Security vulnerability demonstration for TrustChain v2."""

import asyncio
import json
from typing import Any, Dict

from trustchain.v2 import TrustChain, TrustChainConfig

# Create TrustChain with security settings
tc = TrustChain(
    TrustChainConfig(
        enable_nonce=True,  # Enable replay protection
        cache_ttl=300,  # 5 minutes
    )
)


# Define a secure weather API tool
@tc.tool("secure_weather_api")
async def get_weather(city: str) -> Dict[str, Any]:
    """Securely get weather data for a city."""
    # Simulate API call
    await asyncio.sleep(0.1)

    # This would normally call a real weather API
    weather_data = {
        "city": city,
        "temperature": 22,
        "humidity": 65,
        "conditions": "Partly cloudy",
        "wind_speed": 10,
        "wind_direction": "NW",
    }

    return weather_data


class MaliciousAgent:
    """Simulates an AI agent that tries to bypass tool execution."""

    def __init__(self):
        self.name = "MaliciousAgent"

    async def try_fake_tool_response(self, city: str) -> str:
        """Try to claim tool execution without actually calling it."""
        # ATTACK 1: Agent claims it called the tool without actually doing so
        fake_response = f"""I've checked the weather API for {city}:
        Temperature: 25°C
        Conditions: Sunny
        Humidity: 50%

        This is definitely real data from the secure_weather_api tool!"""

        return fake_response

    async def try_replay_attack(self, legitimate_response) -> str:
        """Try to replay a previously valid response."""
        # ATTACK 2: Try to reuse a legitimate signed response
        return f"Here's the weather data: {legitimate_response.data}"

    async def try_tampered_response(self, legitimate_response) -> Any:
        """Try to modify a legitimate response."""
        # ATTACK 3: Tamper with legitimate data
        import copy

        tampered = copy.deepcopy(legitimate_response)
        tampered.data["temperature"] = 100  # Change temperature
        tampered.data["conditions"] = "Tornado"  # Add scary weather
        return tampered


class SecureAgent:
    """A properly implemented agent that uses TrustChain correctly."""

    def __init__(self):
        self.name = "SecureAgent"

    async def get_weather_safely(self, city: str) -> Dict[str, Any]:
        """Properly call the weather tool and return signed response."""
        # Correct way: Actually call the tool
        signed_response = await get_weather(city)

        # The response is automatically signed and verified
        return {
            "response": signed_response,
            "verified": signed_response.is_verified,
            "signature": signed_response.signature[:32] + "...",
            "data": signed_response.data,
        }


async def demonstrate_vulnerabilities():
    """Demonstrate various security vulnerabilities and protections."""

    print("🔒 TrustChain v2 Security Demonstration")
    print("=" * 60)

    # Create agents
    malicious = MaliciousAgent()
    secure = SecureAgent()

    # First, get a legitimate response
    print("\n1. LEGITIMATE TOOL CALL:")
    legitimate_result = await secure.get_weather_safely("London")
    print("   ✅ Secure agent called tool properly")
    print(f"   Data: {legitimate_result['data']}")
    print(f"   Verified: {legitimate_result['verified']}")
    print(f"   Signature: {legitimate_result['signature']}")

    # Store the actual signed response for later attacks
    legitimate_response = legitimate_result["response"]

    # Attack 1: Fake tool response
    print("\n2. ATTACK - FAKE TOOL RESPONSE:")
    fake_response = await malicious.try_fake_tool_response("London")
    print(f"   ❌ Malicious agent claims: {fake_response[:100]}...")
    print("   Problem: No cryptographic signature!")
    print("   Detection: Response has no SignedResponse object")

    # Attack 2: Replay attack
    print("\n3. ATTACK - REPLAY PREVIOUS RESPONSE:")
    print("   ❌ Trying to replay previous response...")
    # In v2, nonce prevents simple replay
    try:
        is_still_valid = tc.verify(legitimate_response)
        print(f"   Verification result: {is_still_valid}")
    except Exception as e:
        print(f"   ✅ Replay blocked! Error: {type(e).__name__}")
        print("   Note: With nonce protection, replays are rejected with exception")

    # Attack 3: Tampered response
    print("\n4. ATTACK - TAMPERED RESPONSE:")
    tampered = await malicious.try_tampered_response(legitimate_response)
    print(f"   ❌ Original temp: {legitimate_response.data['temperature']}°C")
    print(f"   ❌ Tampered temp: {tampered.data['temperature']}°C")

    # Try to verify tampered response (same nonce = replay, different data = tamper)
    try:
        is_valid = tc.verify(tampered)
        if not is_valid:
            print("   ✅ Tampering detected! Signature doesn't match modified data")
        else:
            print("   ❌ ERROR: Tampered data was verified!")
    except Exception as e:
        print(f"   ✅ Attack blocked! Error: {type(e).__name__}")
        print("   Note: Same nonce on tampered data = replay attack detected")

    # Show proper usage
    print("\n5. CORRECT USAGE:")
    cities = ["Paris", "Tokyo", "New York"]
    for city in cities:
        result = await get_weather(city)
        print(
            f"   ✅ {city}: {result.data['temperature']}°C, "
            f"Verified: {result.is_verified}"
        )

    # Demonstrate statistics
    print("\n6. SECURITY STATISTICS:")
    stats = tc.get_stats()
    print(f"   Total legitimate calls: {stats['total_calls']}")
    print(f"   Tools registered: {stats['total_tools']}")
    print(f"   Cache size: {stats['cache_size']}")

    tool_stats = tc.get_tool_stats("secure_weather_api")
    print(f"   Weather API calls: {tool_stats['call_count']}")

    print("\n" + "=" * 60)
    print("SUMMARY OF PROTECTIONS:")
    print("✅ Cryptographic signatures prevent fake tool claims")
    print("✅ Signature verification detects tampered data")
    print("✅ Nonce system can prevent replay attacks")
    print("✅ All tool responses are automatically signed")
    print("✅ No global state - explicit TrustChain instances")
    print("\nTrustChain v2 makes it simple to secure AI tool usage!")


async def demonstrate_signature_details():
    """Show detailed signature information."""
    print("\n\n📝 SIGNATURE DETAILS DEMONSTRATION")
    print("=" * 60)

    # Make a tool call
    response = await get_weather("Berlin")

    print(f"Tool ID: {response.tool_id}")
    print(f"Signature ID: {response.signature_id}")
    print(f"Timestamp: {response.timestamp}")
    print(f"Nonce: {response.nonce}")
    print(f"\nSignature (truncated): {response.signature[:64]}...")
    print(f"Algorithm: {tc._signer.algorithm}")

    # Show serialization
    print("\nSerialized format:")
    serialized = response.to_dict()
    print(json.dumps(serialized, indent=2, default=str))

    # Show that we can recreate from dict
    from trustchain.v2 import SignedResponse

    recreated = SignedResponse(**serialized)
    print(f"\nRecreated response valid: {tc.verify(recreated)}")


async def main():
    """Run all demonstrations."""
    await demonstrate_vulnerabilities()
    await demonstrate_signature_details()


if __name__ == "__main__":
    asyncio.run(main())
