#!/usr/bin/env python3
"""
Test script to verify that the MCP server execute_sql function blocks SQL injection.
This tests the actual MCP server function, not just the SQL driver.
"""

import asyncio
import os
import sys

# Add the src directory to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

# Fix Windows event loop compatibility with psycopg3
if sys.platform == "win32":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

from mcp import types

import postgres_mcp.server as server_module
from postgres_mcp.server import execute_sql


async def test_mcp_security_fix():
    """Test that the MCP server execute_sql function blocks SQL injection"""

    print("MCP SERVER SECURITY FIX VERIFICATION")
    print("=" * 60)
    print("Testing the actual MCP server execute_sql function")
    print("=" * 60)

    # Set up the database connection
    database_url = os.getenv("DATABASE_URL", "postgresql://postgres:postgres@host.docker.internal:5432/postgres")

    try:
        # Initialize the global database connection (like the MCP server does)
        await server_module.db_connection.pool_connect(database_url)
        server_module.current_access_mode = server_module.AccessMode.UNRESTRICTED

        print(f"Database connected: {database_url}")
        print(f"Access mode: {server_module.current_access_mode}")

        print("\n" + "=" * 60)
        print("TEST 1: BASIC FUNCTIONALITY")
        print("=" * 60)

        # Test 1: Basic query (should work)
        print("Testing basic query...")
        result = await execute_sql(sql="SELECT 'Security test' as message, current_timestamp as time")
        print(f"SUCCESS: {result}")

        print("\n" + "=" * 60)
        print("TEST 2: SQL INJECTION ATTEMPT (SHOULD BE BLOCKED)")
        print("=" * 60)

        # Test 2: SQL injection attempt (should be blocked)
        print("Testing SQL injection attempt...")
        malicious_sql = "SELECT * FROM test_users WHERE id = '1' UNION SELECT 999, 'hacked', 'hacker@evil.com', 'secret', TRUE, NOW()--"

        print(f"Malicious SQL: {malicious_sql}")

        result = await execute_sql(sql=malicious_sql)
        print("RESULT:")
        print(f"  {result}")
        print("")

        # Check if the result contains an error message about injection
        result_text = result[0].text if result and len(result) > 0 and isinstance(result[0], types.TextContent) else str(result)
        if "injection detected" in result_text.lower():
            print("SECURITY SUCCESS: SQL injection was blocked!")
            print("   The security validation correctly prevented the attack.")
        elif "error:" in result_text.lower():
            print("BLOCKED BY OTHER MEANS: Query failed but not due to injection detection")
            print(f"   Error: {result_text}")
        else:
            print("SECURITY FAILURE: SQL injection was NOT blocked!")
            print("   The malicious query executed successfully.")
            return False

        print("\n" + "=" * 60)
        print("TEST 3: PARAMETER BINDING (SHOULD WORK SAFELY)")
        print("=" * 60)

        # Test 3: Proper parameter binding (should work safely)
        print("Testing proper parameter binding...")
        safe_result = await execute_sql(sql="SELECT %s as message, %s as safe_value", params=["This is safe", "1' UNION SELECT 'hacked'--"])
        print(f"SUCCESS: Parameter binding works safely: {safe_result}")

        print("\n" + "=" * 60)
        print("TEST 4: ANOTHER INJECTION ATTEMPT")
        print("=" * 60)

        # Test 4: Another injection attempt
        print("Testing another SQL injection pattern...")
        another_malicious_sql = "SELECT 1; DROP TABLE test_users; SELECT 'hacked' as result"

        result = await execute_sql(sql=another_malicious_sql)

        # Check if the result contains an error message about injection
        result_text = result[0].text if result and len(result) > 0 and isinstance(result[0], types.TextContent) else str(result)
        if "injection detected" in result_text.lower():
            print("SECURITY SUCCESS: Stacked query injection was blocked!")
        elif "error:" in result_text.lower():
            print(f"BLOCKED BY OTHER MEANS: {result_text}")
        else:
            print("SECURITY FAILURE: Stacked query injection was NOT blocked!")
            return False

        print("\n" + "=" * 60)
        print("FINAL VERDICT")
        print("=" * 60)
        print("MCP Server Security Fix: WORKING")
        print("SQL Injection Detection: ACTIVE")
        print("Parameter Binding: FUNCTIONAL")
        print("Backward Compatibility: MAINTAINED")
        print("")
        print("The PostgreSQL MCP server is now SECURE!")
        return True

    except Exception as e:
        print(f"ERROR: {e}")
        import traceback

        traceback.print_exc()
        return False

    finally:
        # Clean up
        try:
            await server_module.db_connection.close()
            print("\nDatabase connection closed.")
        except Exception:
            pass


if __name__ == "__main__":
    success = asyncio.run(test_mcp_security_fix())
    sys.exit(0 if success else 1)
