#!/usr/bin/env python3
"""
Parallel Scanning Example

Demonstrates:
- Running multiple scans simultaneously
- Background execution
- Session monitoring
- Resource management

Usage:
    python parallel_scans.py <target1> <target2> <target3> ...
    
Example:
    python parallel_scans.py 192.168.1.1 192.168.1.2 192.168.1.3
"""

import asyncio
import sys
from pentest_mcp_server import PentestMCPServer


async def monitor_session(server, session_id, interval=5, max_checks=10):
    """Monitor a session and print periodic updates."""
    for i in range(max_checks):
        await asyncio.sleep(interval)
        
        result = await server._handle_read_output({
            "session_id": session_id,
            "lines": 10
        })
        
        if result["status"] == "success":
            output = result["output"]
            # Check if scan is complete
            if "Nmap done" in output or "finished" in output.lower():
                print(f"  [{session_id}] Scan completed!")
                return True
            else:
                # Show progress
                lines = [l for l in output.split('\n') if l.strip()]
                if lines:
                    print(f"  [{session_id}] Progress: {lines[-1][:60]}...")
    
    return False


async def main():
    if len(sys.argv) < 2:
        print("Usage: python parallel_scans.py <target1> <target2> <target3> ...")
        print("Example: python parallel_scans.py 192.168.1.1 192.168.1.2 192.168.1.3")
        sys.exit(1)
    
    targets = sys.argv[1:]
    server = PentestMCPServer()
    
    try:
        print(f"[*] Starting parallel scans of {len(targets)} target(s)")
        print("[*] This demonstrates concurrent operations\n")
        
        # Create sessions for each target
        sessions = {}
        print("[+] Creating sessions...")
        
        for i, target in enumerate(targets, 1):
            session_id = f"scan_{i}"
            result = await server._handle_create_session({
                "session_id": session_id
            })
            
            if result["status"] == "created":
                sessions[session_id] = target
                print(f"  ✓ {session_id}: {target}")
            else:
                print(f"  ✗ Failed to create session for {target}")
        
        print()
        
        # Start all scans in background
        print("[*] Starting scans in background...")
        for session_id, target in sessions.items():
            result = await server._handle_execute({
                "session_id": session_id,
                "command": f"nmap -sV {target} -oX /tmp/{session_id}.xml",
                "background": True
            })
            
            if result["status"] == "background":
                print(f"  ✓ {session_id}: Scan started for {target}")
            else:
                print(f"  ✗ {session_id}: Failed to start scan")
        
        print()
        
        # Monitor all sessions
        print("[*] Monitoring scan progress...")
        print("[*] Checking every 5 seconds (max 10 checks = 50 seconds)")
        print("-" * 60)
        
        # Create monitoring tasks for all sessions
        monitor_tasks = [
            monitor_session(server, session_id, interval=5, max_checks=10)
            for session_id in sessions.keys()
        ]
        
        # Wait for all monitors to complete
        results = await asyncio.gather(*monitor_tasks)
        
        print("-" * 60)
        print()
        
        # Check which scans completed
        completed = sum(results)
        print(f"[+] {completed}/{len(sessions)} scans completed within monitoring period")
        
        if completed < len(sessions):
            print("[!] Some scans are still running. They will continue in the background.")
        
        print()
        
        # Get system status
        print("[*] System status:")
        status_result = await server._handle_get_system_status({})
        
        if status_result["status"] == "success":
            print(f"  Active sessions: {status_result['active_sessions']}")
            print(f"  Connection: {status_result['connection_status']['host']}:{status_result['connection_status']['port']}")
            
            if "system_info" in status_result:
                sys_info = status_result["system_info"]
                print(f"  Hostname: {sys_info.get('hostname', 'unknown')}")
                print(f"  CPU cores: {sys_info.get('cpu_info', 'unknown')}")
        
        print()
        
        # Parse completed scans
        print("[*] Parsing completed scan results...")
        for i, (session_id, target) in enumerate(sessions.items(), 1):
            if results[i-1]:  # If scan completed
                parse_result = await server._handle_parse_tool_output({
                    "tool": "nmap",
                    "file_path": f"/tmp/{session_id}.xml",
                    "format": "xml"
                })
                
                if parse_result["status"] == "success":
                    hosts = parse_result["parsed_data"]["hosts"]
                    print(f"\n[+] Results for {target}:")
                    
                    if hosts:
                        for host in hosts:
                            ports = host.get('ports', [])
                            print(f"  Host: {host.get('address', 'unknown')}")
                            print(f"  Open ports: {len(ports)}")
                            
                            for port in ports[:5]:  # Show first 5 ports
                                print(f"    - {port['port']}/{port['protocol']}: {port['service']}")
                            
                            if len(ports) > 5:
                                print(f"    ... and {len(ports) - 5} more")
                    else:
                        print("  No hosts found or scan incomplete")
        
    except Exception as e:
        print(f"[-] Error: {e}")
        import traceback
        traceback.print_exc()
    
    finally:
        # Cleanup all sessions
        print("\n[*] Cleaning up sessions...")
        for session_id in sessions.keys():
            cleanup_result = await server._handle_kill_session({
                "session_id": session_id
            })
            if cleanup_result["status"] == "killed":
                print(f"  ✓ {session_id} terminated")
        
        await server.shutdown()
        print("[+] Done!")


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