#!/usr/bin/env python3
"""
Benchmark SSH vs WinRM performance
Compares execution times and capabilities
"""

import time
import subprocess
import os
from winrm import Session


def benchmark_ssh_method():
    """Benchmark SSH-based script execution"""
    print("=== SSH Method Benchmark ===")

    start_time = time.time()

    try:
        # Run a simple recon via SSH
        result = subprocess.run(
            ['timeout', '30', './recon.sh'],
            capture_output=True,
            text=True,
            timeout=35
        )

        end_time = time.time()
        duration = end_time - start_time

        if result.returncode == 0:
            print(f"✅ SSH method successful")
            print(f"⏱️  Execution time: {duration:.2f} seconds")
            return duration, True
        else:
            print(f"❌ SSH method failed: {result.stderr}")
            return duration, False

    except subprocess.TimeoutExpired:
        print(f"❌ SSH method timed out after 35 seconds")
        return 35, False
    except Exception as e:
        end_time = time.time()
        duration = end_time - start_time
        print(f"❌ SSH method error: {e}")
        return duration, False


def benchmark_winrm_method():
    """Benchmark WinRM-based script execution"""
    print("\n=== WinRM Method Benchmark ===")

    host = os.getenv('HOST', '192.168.88.108')
    username = os.getenv('USER_NAME', 'Admin')
    password = os.getenv('PASS', 'P@ssw0rd!')

    start_time = time.time()

    try:
        # Connect to WinRM
        session = Session(
            f'http://{host}:5985/wsman',
            auth=(username, password),
            transport='ntlm'
        )

        # Run equivalent recon commands via WinRM
        recon_script = """
        Write-Host ""
        Write-Host "======== Computer ========"
        Get-ComputerInfo -Property CsName,WindowsProductName,OsVersion,OsInstallDate | Format-List

        Write-Host "======== Boot Time ========"
        (Get-CimInstance Win32_OperatingSystem).LastBootUpTime

        Write-Host "======== Local Users ========"
        Get-LocalUser | Select-Object Name,Enabled,LastLogon | Format-Table -AutoSize

        Write-Host "======== Disks ========"
        Get-Volume | Where-Object DriveLetter | Select-Object DriveLetter,FileSystemLabel,
            @{n='Size(GB)';e={[math]::Round($_.Size / 1GB, 1)}},
            @{n='Free(GB)';e={[math]::Round($_.SizeRemaining / 1GB, 1)}} | Format-Table -AutoSize

        Write-Host "======== IPv4 Addresses ========"
        Get-NetIPAddress -AddressFamily IPv4 | Where-Object { $_.InterfaceAlias -notmatch 'Loopback' } |
            Select-Object InterfaceAlias,IPAddress | Format-Table -AutoSize

        Write-Host "======== WinRM Recon Completed ========"
        """

        result = session.run_ps(recon_script)

        end_time = time.time()
        duration = end_time - start_time

        if result.status_code == 0:
            print(f"✅ WinRM method successful")
            print(f"⏱️  Execution time: {duration:.2f} seconds")

            # Show some output
            output = result.std_out.decode('utf-8', errors='replace')
            lines = [line.strip() for line in output.split('\n') if line.strip()]
            print(f"📊 Output lines: {len(lines)}")

            return duration, True
        else:
            print(f"❌ WinRM method failed: {result.std_err.decode()}")
            return duration, False

    except Exception as e:
        end_time = time.time()
        duration = end_time - start_time
        print(f"❌ WinRM method error: {e}")
        return duration, False


def main():
    print("🚀 SSH vs WinRM Performance Benchmark")
    print("=====================================")

    # Benchmark SSH
    ssh_time, ssh_success = benchmark_ssh_method()

    # Benchmark WinRM
    winrm_time, winrm_success = benchmark_winrm_method()

    # Results comparison
    print("\n📊 RESULTS COMPARISON")
    print("=====================")

    print(f"SSH Method:   {ssh_time:.2f}s  {'✅' if ssh_success else '❌'}")
    print(f"WinRM Method: {winrm_time:.2f}s  {'✅' if winrm_success else '❌'}")

    if ssh_success and winrm_success:
        if ssh_time < winrm_time:
            speedup = winrm_time / ssh_time
            print(f"\n🏃 SSH is {speedup:.1f}x faster than WinRM")
        else:
            speedup = ssh_time / winrm_time
            print(f"\n🏃 WinRM is {speedup:.1f}x faster than SSH")

        print(f"\nTime difference: {abs(ssh_time - winrm_time):.2f} seconds")

    print("\n🎯 RECOMMENDATION")
    print("=================")

    if ssh_success and winrm_success:
        if ssh_time < winrm_time * 1.5:  # If SSH is not significantly slower
            print("✅ Continue using SSH method:")
            print("  - Proven stability")
            print("  - Good performance")
            print("  - Handles large scripts well")
            print("  - Simpler setup")
        else:
            print("✅ Consider WinRM method for:")
            print("  - Better Windows integration")
            print("  - Advanced PowerShell features")
            print("  - Structured error handling")
    elif ssh_success:
        print("✅ SSH method is more reliable for this environment")
    elif winrm_success:
        print("✅ WinRM method works, SSH has issues")
    else:
        print("❌ Both methods have issues, investigate connectivity")

    print("\n💡 HYBRID APPROACH")
    print("==================")
    print("Consider using both methods strategically:")
    print("- SSH: Large scripts, file operations, stability")
    print("- WinRM: Quick commands, Windows-specific tasks, real-time interaction")


if __name__ == "__main__":
    main()