PyGuardian/tests/integration/test_api_integration.py

#!/usr/bin/env python3
"""
Integration tests for PyGuardian API and cluster management.
"""

import unittest
import tempfile
import os
import sys
import json
import asyncio
import aiohttp
from unittest.mock import Mock, patch, AsyncMock
import sqlite3
from datetime import datetime

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


class TestAPIServer(unittest.TestCase):
    """Integration tests for API server."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, 'test_guardian.db')
        
    def tearDown(self):
        """Clean up test fixtures."""
        if os.path.exists(self.db_path):
            os.remove(self.db_path)
        os.rmdir(self.temp_dir)
    
    def test_api_health_endpoint(self):
        """Test API health check endpoint."""
        # This would be an actual HTTP test
        # For now, just test that we can import the module
        try:
            from api_server import PyGuardianAPI
            self.assertTrue(True)
        except ImportError:
            self.fail("Could not import API server module")
    
    def test_agent_registration_flow(self):
        """Test agent registration API flow."""
        # Mock test for agent registration
        test_data = {
            'agent_name': 'test_agent',
            'host_info': {
                'hostname': 'test-host',
                'os': 'linux',
                'arch': 'x86_64'
            }
        }
        
        # This would test the actual API endpoint
        self.assertIsNotNone(test_data)
    
    def test_jwt_authentication_middleware(self):
        """Test JWT authentication middleware."""
        # Test JWT authentication in API requests
        test_token = "Bearer test.jwt.token"
        
        # Mock authorization header validation
        self.assertTrue(test_token.startswith("Bearer "))


class TestClusterManager(unittest.TestCase):
    """Integration tests for cluster management."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        
    def tearDown(self):
        """Clean up test fixtures."""
        os.rmdir(self.temp_dir)
    
    def test_cluster_manager_import(self):
        """Test cluster manager module import."""
        try:
            from cluster_manager import ClusterManager
            self.assertTrue(True)
        except ImportError:
            self.fail("Could not import ClusterManager")
    
    def test_agent_registration(self):
        """Test agent registration in cluster."""
        # Mock agent registration
        agent_data = {
            'agent_id': 'agent_test123',
            'hostname': 'test-agent',
            'ip_address': '192.168.1.100',
            'status': 'active'
        }
        
        self.assertEqual(agent_data['status'], 'active')
    
    def test_agent_health_check(self):
        """Test agent health monitoring."""
        # Mock health check
        health_data = {
            'agent_id': 'agent_test123',
            'last_seen': datetime.now().isoformat(),
            'status': 'healthy',
            'cpu_usage': 25.5,
            'memory_usage': 60.2,
            'disk_usage': 45.0
        }
        
        self.assertEqual(health_data['status'], 'healthy')
        self.assertLess(health_data['cpu_usage'], 100)
        self.assertLess(health_data['memory_usage'], 100)


class TestTelegramBot(unittest.TestCase):
    """Integration tests for Telegram bot."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        
    def tearDown(self):
        """Clean up test fixtures."""
        os.rmdir(self.temp_dir)
    
    def test_bot_import(self):
        """Test Telegram bot module import."""
        try:
            from bot import TelegramBot
            self.assertTrue(True)
        except ImportError:
            self.fail("Could not import TelegramBot")
    
    def test_command_parsing(self):
        """Test bot command parsing."""
        # Mock command parsing
        test_commands = [
            '/start',
            '/status',
            '/cluster',
            '/agents',
            '/help'
        ]
        
        for cmd in test_commands:
            self.assertTrue(cmd.startswith('/'))
    
    def test_authentication_commands(self):
        """Test authentication-related bot commands."""
        # Mock authentication commands
        auth_commands = [
            '/generate_agent',
            '/revoke_token',
            '/list_agents',
            '/agent_status'
        ]
        
        for cmd in auth_commands:
            self.assertTrue(isinstance(cmd, str))


class TestSecurityMonitor(unittest.TestCase):
    """Integration tests for security monitoring."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        
    def tearDown(self):
        """Clean up test fixtures."""
        os.rmdir(self.temp_dir)
    
    def test_security_monitor_import(self):
        """Test security monitor import."""
        try:
            from monitor import SecurityMonitor
            self.assertTrue(True)
        except ImportError:
            self.fail("Could not import SecurityMonitor")
    
    def test_threat_detection(self):
        """Test threat detection logic."""
        # Mock threat detection
        threat_events = [
            {
                'type': 'brute_force',
                'source_ip': '192.168.1.100',
                'attempts': 5,
                'timestamp': datetime.now().isoformat()
            },
            {
                'type': 'port_scan',
                'source_ip': '10.0.0.50',
                'ports': [22, 80, 443],
                'timestamp': datetime.now().isoformat()
            }
        ]
        
        for event in threat_events:
            self.assertIn('type', event)
            self.assertIn('source_ip', event)
            self.assertIn('timestamp', event)


class TestFirewallManager(unittest.TestCase):
    """Integration tests for firewall management."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        
    def tearDown(self):
        """Clean up test fixtures."""
        os.rmdir(self.temp_dir)
    
    def test_firewall_import(self):
        """Test firewall module import."""
        try:
            from firewall import FirewallManager
            self.assertTrue(True)
        except ImportError:
            self.fail("Could not import FirewallManager")
    
    def test_ip_blocking(self):
        """Test IP address blocking."""
        # Mock IP blocking
        blocked_ips = [
            '192.168.1.100',
            '10.0.0.50',
            '203.0.113.1'
        ]
        
        for ip in blocked_ips:
            # Validate IP format (basic check)
            parts = ip.split('.')
            self.assertEqual(len(parts), 4)
            for part in parts:
                self.assertTrue(0 <= int(part) <= 255)
    
    def test_whitelist_management(self):
        """Test IP whitelist management."""
        # Mock whitelist
        whitelist = [
            '127.0.0.1',
            '192.168.1.0/24',
            '10.0.0.0/8'
        ]
        
        for entry in whitelist:
            self.assertIsInstance(entry, str)
            self.assertTrue('.' in entry)


class TestDatabaseOperations(unittest.TestCase):
    """Integration tests for database operations."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, 'test_integration.db')
        
    def tearDown(self):
        """Clean up test fixtures."""
        if os.path.exists(self.db_path):
            os.remove(self.db_path)
        os.rmdir(self.temp_dir)
    
    def test_database_creation(self):
        """Test database creation and schema."""
        # Create SQLite database
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Create a test table
        cursor.execute('''
            CREATE TABLE test_agents (
                id INTEGER PRIMARY KEY,
                agent_id TEXT UNIQUE,
                status TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # Insert test data
        cursor.execute('''
            INSERT INTO test_agents (agent_id, status)
            VALUES (?, ?)
        ''', ('agent_test123', 'active'))
        
        conn.commit()
        
        # Verify data
        cursor.execute('SELECT * FROM test_agents')
        results = cursor.fetchall()
        
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0][1], 'agent_test123')
        self.assertEqual(results[0][2], 'active')
        
        conn.close()
    
    def test_agent_authentication_tables(self):
        """Test agent authentication tables."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Create authentication tables
        cursor.execute('''
            CREATE TABLE agent_auth (
                id INTEGER PRIMARY KEY,
                agent_id TEXT UNIQUE NOT NULL,
                key_hash TEXT NOT NULL,
                encrypted_key TEXT NOT NULL,
                is_active BOOLEAN DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        cursor.execute('''
            CREATE TABLE agent_tokens (
                id INTEGER PRIMARY KEY,
                agent_id TEXT NOT NULL,
                token TEXT NOT NULL,
                expires_at TIMESTAMP NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (agent_id) REFERENCES agent_auth (agent_id)
            )
        ''')
        
        # Test data insertion
        cursor.execute('''
            INSERT INTO agent_auth (agent_id, key_hash, encrypted_key)
            VALUES (?, ?, ?)
        ''', ('agent_test123', 'test_hash', 'encrypted_key'))
        
        conn.commit()
        
        # Verify tables exist and have data
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
        tables = [row[0] for row in cursor.fetchall()]
        
        self.assertIn('agent_auth', tables)
        self.assertIn('agent_tokens', tables)
        
        cursor.execute('SELECT COUNT(*) FROM agent_auth')
        count = cursor.fetchone()[0]
        self.assertEqual(count, 1)
        
        conn.close()


def run_integration_tests():
    """Run all integration tests."""
    print("🔄 Running PyGuardian Integration Tests...")
    print("=" * 50)
    
    # Create test suite
    test_suite = unittest.TestSuite()
    
    # Add test classes
    test_classes = [
        TestAPIServer,
        TestClusterManager,
        TestTelegramBot,
        TestSecurityMonitor,
        TestFirewallManager,
        TestDatabaseOperations
    ]
    
    for test_class in test_classes:
        tests = unittest.TestLoader().loadTestsFromTestCase(test_class)
        test_suite.addTests(tests)
    
    # Run tests
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(test_suite)
    
    # Print summary
    print("\n" + "=" * 50)
    print(f"🏁 Integration Tests completed:")
    print(f"   ✅ Passed: {result.testsRun - len(result.failures) - len(result.errors)}")
    print(f"   ❌ Failed: {len(result.failures)}")
    print(f"   💥 Errors: {len(result.errors)}")
    
    return 0 if result.wasSuccessful() else 1


if __name__ == '__main__':
    sys.exit(run_integration_tests())