basic_usage.py

#!/usr/bin/env python3
"""
Basic SchemaPin Usage Example

This example demonstrates how to use SchemaPin integration with MockLoop MCP
for basic schema verification and key pinning scenarios.
"""

import asyncio
import json
import sys
from pathlib import Path

# Add src to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))

from mockloop_mcp.schemapin import (
    SchemaPinConfig,
    SchemaVerificationInterceptor,
    KeyPinningManager,
    PolicyHandler,
    SchemaPinAuditLogger,
    VerificationResult,
    PolicyAction
)


async def basic_verification_example():
    """Demonstrate basic schema verification workflow."""
    print("=== Basic Schema Verification Example ===\n")

    # 1. Configure SchemaPin
    config = SchemaPinConfig(
        enabled=True,
        policy_mode="warn",  # Options: enforce, warn, log
        auto_pin_keys=False,
        key_pin_storage_path="example_keys.db",
        trusted_domains=["api.example.com"],
        interactive_mode=False
    )

    print(f"✓ SchemaPin configured: policy_mode={config.policy_mode}")

    # 2. Initialize verification interceptor
    interceptor = SchemaVerificationInterceptor(config)
    print("✓ Verification interceptor initialized")

    # 3. Define a tool schema (what you'd get from an MCP tool)
    tool_schema = {
        "name": "database_query",
        "description": "Execute SQL queries against database",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "SQL query to execute"
                },
                "database": {
                    "type": "string",
                    "description": "Target database name",
                    "default": "main"
                }
            },
            "required": ["query"]
        }
    }

    # 4. Simulate verification scenarios
    print("\n--- Scenario 1: Unsigned tool (no signature) ---")
    result = await interceptor.verify_tool_schema(
        tool_name="database_query",
        schema=tool_schema,
        signature=None,  # No signature provided
        domain="api.example.com"
    )

    print(f"Valid: {result.valid}")
    print(f"Error: {result.error}")
    print(f"Tool ID: {result.tool_id}")

    print("\n--- Scenario 2: Tool with signature but no pinned key ---")
    # In real usage, this would be a cryptographic signature
    mock_signature = "eyJhbGciOiJFUzI1NiJ9.mock_signature_data"

    result = await interceptor.verify_tool_schema(
        tool_name="database_query",
        schema=tool_schema,
        signature=mock_signature,
        domain="api.example.com"
    )

    print(f"Valid: {result.valid}")
    print(f"Domain: {result.domain}")
    print(f"Key pinned: {result.key_pinned}")
    if result.error:
        print(f"Error: {result.error}")


async def key_management_example():
    """Demonstrate key pinning and management."""
    print("\n=== Key Management Example ===\n")

    # Initialize key manager
    key_manager = KeyPinningManager("example_keys.db")
    print("✓ Key manager initialized")

    # Example public key (in real usage, this would be discovered or provided)
    example_public_key = """-----BEGIN PUBLIC KEY-----
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEExample1234567890abcdefghijklmn
opqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890example
-----END PUBLIC KEY-----"""

    # Pin a key for a tool
    tool_id = "api.example.com/database_query"
    domain = "api.example.com"
    metadata = {
        "developer": "Example Corp",
        "version": "1.0.0",
        "description": "Database query tool"
    }

    success = key_manager.pin_key(tool_id, domain, example_public_key, metadata)
    print(f"✓ Key pinned successfully: {success}")

    # Retrieve pinned key
    retrieved_key = key_manager.get_pinned_key(tool_id)
    print(f"✓ Key retrieved: {retrieved_key is not None}")

    # Check if key is pinned
    is_pinned = key_manager.is_key_pinned(tool_id)
    print(f"✓ Key is pinned: {is_pinned}")

    # Get detailed key information
    key_info = key_manager.get_key_info(tool_id)
    if key_info:
        print(f"✓ Key info - Domain: {key_info['domain']}")
        print(f"✓ Key info - Pinned at: {key_info['pinned_at']}")
        print(f"✓ Key info - Verification count: {key_info['verification_count']}")
        print(f"✓ Key info - Developer: {key_info['metadata']['developer']}")

    # List all pinned keys
    pinned_keys = key_manager.list_pinned_keys()
    print(f"✓ Total pinned keys: {len(pinned_keys)}")

    # Update verification stats (simulating successful verification)
    key_manager.update_verification_stats(tool_id)
    print("✓ Verification stats updated")

    # Get updated info
    updated_info = key_manager.get_key_info(tool_id)
    if updated_info:
        print(f"✓ Updated verification count: {updated_info['verification_count']}")


async def policy_enforcement_example():
    """Demonstrate policy enforcement scenarios."""
    print("\n=== Policy Enforcement Example ===\n")

    # Test different policy modes
    policy_modes = ["enforce", "warn", "log"]

    for mode in policy_modes:
        print(f"--- Testing {mode.upper()} mode ---")

        config = SchemaPinConfig(policy_mode=mode)
        policy_handler = PolicyHandler(config)

        # Simulate failed verification
        failed_result = VerificationResult(
            valid=False,
            tool_id="untrusted.com/suspicious_tool",
            domain="untrusted.com",
            error="Signature verification failed"
        )

        decision = await policy_handler.evaluate_verification_result(
            failed_result, "suspicious_tool"
        )

        print(f"  Action: {decision.action.value}")
        print(f"  Reason: {decision.reason}")
        print(f"  Policy mode: {decision.policy_mode}")

        # Show what each action means
        if decision.action == PolicyAction.BLOCK:
            print("  → Tool execution would be BLOCKED")
        elif decision.action == PolicyAction.WARN:
            print("  → Tool execution would proceed with WARNING")
        elif decision.action == PolicyAction.LOG:
            print("  → Tool execution would proceed with LOGGING only")

        print()


async def audit_logging_example():
    """Demonstrate audit logging capabilities."""
    print("\n=== Audit Logging Example ===\n")

    # Initialize audit logger
    audit_logger = SchemaPinAuditLogger("example_audit.db")
    print("✓ Audit logger initialized")

    # Log various events
    print("Logging verification events...")

    # Successful verification
    success_result = VerificationResult(
        valid=True,
        tool_id="api.example.com/secure_tool",
        domain="api.example.com",
        key_pinned=True
    )

    await audit_logger.log_verification_attempt(
        "api.example.com/secure_tool",
        "api.example.com",
        success_result,
        execution_time_ms=125.5
    )

    # Failed verification
    await audit_logger.log_verification_error(
        "malicious.com/bad_tool",
        "malicious.com",
        "Invalid signature detected"
    )

    # Key pinning event
    await audit_logger.log_key_pinning_event(
        "new.com/new_tool",
        "new.com",
        "new_public_key",
        "pin"
    )

    # Policy decision
    await audit_logger.log_policy_decision(
        "questionable.com/tool",
        "warn",
        "Unsigned tool execution",
        "warn"
    )

    print("✓ Events logged")

    # Get verification statistics
    stats = audit_logger.get_verification_stats()
    print("\n--- Audit Statistics ---")
    print(f"Total verifications: {stats.get('total_verifications', 0)}")
    print(f"Successful verifications: {stats.get('successful_verifications', 0)}")
    print(f"Failed verifications: {stats.get('failed_verifications', 0)}")
    print(f"Unique tools: {stats.get('unique_tools', 0)}")
    print(f"Unique domains: {stats.get('unique_domains', 0)}")

    if 'policy_breakdown' in stats:
        print("\nPolicy action breakdown:")
        for action, count in stats['policy_breakdown'].items():
            print(f"  {action}: {count}")


async def configuration_example():
    """Demonstrate configuration management."""
    print("\n=== Configuration Management Example ===\n")

    # Create custom configuration
    config = SchemaPinConfig(
        enabled=True,
        policy_mode="enforce",
        auto_pin_keys=True,
        key_pin_storage_path="production_keys.db",
        discovery_timeout=60,
        cache_ttl=7200,
        well_known_endpoints={
            "api.example.com": "https://api.example.com/.well-known/schemapin.json",
            "tools.corp.com": "https://tools.corp.com/security/schemapin.json"
        },
        trusted_domains=[
            "api.example.com",
            "tools.corp.com",
            "internal.company.com"
        ],
        revocation_check=True,
        interactive_mode=False
    )

    print("✓ Custom configuration created")

    # Save configuration to file
    config.save_to_file("schemapin_config.json")
    print("✓ Configuration saved to file")

    # Load configuration from file
    loaded_config = SchemaPinConfig.load_from_file("schemapin_config.json")
    print("✓ Configuration loaded from file")

    # Verify configuration
    print(f"Policy mode: {loaded_config.policy_mode}")
    print(f"Auto-pin keys: {loaded_config.auto_pin_keys}")
    print(f"Trusted domains: {len(loaded_config.trusted_domains)}")
    print(f"Discovery timeout: {loaded_config.discovery_timeout}s")

    # Convert to dictionary for inspection
    config_dict = loaded_config.to_dict()
    print("\nConfiguration as dictionary:")
    print(json.dumps(config_dict, indent=2))


async def main():
    """Run all examples."""
    print("SchemaPin Integration Examples")
    print("=" * 50)

    try:
        await basic_verification_example()
        await key_management_example()
        await policy_enforcement_example()
        await audit_logging_example()
        await configuration_example()

        print("\n" + "=" * 50)
        print("✓ All examples completed successfully!")
        print("\nNext steps:")
        print("1. Review the generated database files (example_keys.db, example_audit.db)")
        print("2. Examine the configuration file (schemapin_config.json)")
        print("3. Integrate SchemaPin verification into your MCP tools")
        print("4. Set up proper key discovery endpoints for your domains")

    except Exception as e:
        print(f"\n❌ Error running examples: {e}")
        import traceback
        traceback.print_exc()

    finally:
        # Clean up example files
        cleanup_files = [
            "example_keys.db",
            "example_audit.db",
            "schemapin_config.json"
        ]

        print("\nCleaning up example files...")
        for file_path in cleanup_files:
            try:
                Path(file_path).unlink(missing_ok=True)
                print(f"✓ Removed {file_path}")
            except Exception as e:
                print(f"⚠ Could not remove {file_path}: {e}")


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