#!/usr/bin/env python3
"""
Python WinRM Connection Test
Tests WinRM connectivity using pywinrm library
"""

import os
import sys
from winrm.protocol import Protocol
from winrm import Session
import xml.etree.ElementTree as ET


def test_basic_connection(host, username, password):
    """Test basic WinRM connectivity"""
    print(f"[1/4] Testing basic WinRM protocol connection to {host}...")

    try:
        # Create a basic protocol connection
        endpoint = f"http://{host}:5985/wsman"
        p = Protocol(
            endpoint=endpoint,
            transport='plaintext',
            username=username,
            password=password
        )

        # Try a simple identify operation
        result = p.send_message(
            '<s:Envelope xmlns:s="http://www.w3.org/2003/05/soap-envelope" xmlns:wsmid="http://schemas.dmtf.org/wbem/wsman/identity/1/wsmanidentity.xsd">'
            '<s:Header></s:Header>'
            '<s:Body>'
            '<wsmid:Identify/>'
            '</s:Body>'
            '</s:Envelope>'
        )

        print("✅ Basic WinRM protocol connection successful")

        # Parse response to get some info
        try:
            root = ET.fromstring(result)
            for elem in root.iter():
                if 'ProductVersion' in elem.tag:
                    print(f"   Protocol Version: {elem.text}")
                elif 'ProductVendor' in elem.tag:
                    print(f"   Product Vendor: {elem.text}")
        except:
            pass

        return True

    except Exception as e:
        print(f"❌ Basic connection failed: {e}")
        return False


def test_session_connection(host, username, password):
    """Test WinRM session connectivity"""
    print(f"\n[2/4] Testing WinRM session connection...")

    try:
        session = Session(
            f'http://{host}:5985/wsman',
            auth=(username, password),
            transport='plaintext'
        )

        print("✅ WinRM session created successfully")
        return session

    except Exception as e:
        print(f"❌ Session creation failed: {e}")
        return None


def test_command_execution(session):
    """Test basic command execution"""
    print(f"\n[3/4] Testing command execution...")

    try:
        # Test simple command
        result = session.run_cmd('echo "Hello from WinRM"')

        if result.status_code == 0:
            print("✅ Command execution successful")
            print(f"   Output: {result.std_out.decode().strip()}")
        else:
            print(f"⚠️  Command execution completed with exit code: {result.status_code}")
            print(f"   Error: {result.std_err.decode().strip()}")

        return True

    except Exception as e:
        print(f"❌ Command execution failed: {e}")
        return False


def test_powershell_execution(session):
    """Test PowerShell script execution"""
    print(f"\n[4/4] Testing PowerShell execution...")

    try:
        # Test PowerShell command
        ps_script = """
        $info = @{
            ComputerName = $env:COMPUTERNAME
            UserName = $env:USERNAME
            PowerShellVersion = $PSVersionTable.PSVersion.ToString()
            WindowsVersion = (Get-ComputerInfo -Property WindowsProductName).WindowsProductName
            CurrentTime = Get-Date -Format 'yyyy-MM-dd HH:mm:ss'
        }

        Write-Host "=== Remote System Information ==="
        Write-Host "Computer: $($info.ComputerName)"
        Write-Host "User: $($info.UserName)"
        Write-Host "PowerShell: $($info.PowerShellVersion)"
        Write-Host "Windows: $($info.WindowsVersion)"
        Write-Host "Time: $($info.CurrentTime)"
        Write-Host "=== End Information ==="
        """

        result = session.run_ps(ps_script)

        if result.status_code == 0:
            print("✅ PowerShell execution successful")
            print("Remote system info:")
            output = result.std_out.decode().strip()
            for line in output.split('\n'):
                if line.strip():
                    print(f"   {line.strip()}")
        else:
            print(f"⚠️  PowerShell execution completed with exit code: {result.status_code}")
            if result.std_err:
                print(f"   Error: {result.std_err.decode().strip()}")

        return True

    except Exception as e:
        print(f"❌ PowerShell execution failed: {e}")
        return False


def main():
    # Get connection details from environment or use defaults
    host = os.getenv('HOST', '192.168.88.108')
    username = os.getenv('USER_NAME', 'Admin')
    password = os.getenv('PASS', 'P@ssw0rd!')

    print("=== Python WinRM Connection Test ===")
    print(f"Target: {username}@{host}")
    print("")

    # Test sequence
    success = True

    # Test 1: Basic connection
    if not test_basic_connection(host, username, password):
        success = False
        print("\nℹ️  If basic connection fails, ensure:")
        print("   - WinRM service is running on target")
        print("   - Port 5985 is accessible")
        print("   - AllowUnencrypted=true for HTTP connections")
        sys.exit(1)

    # Test 2: Session creation
    session = test_session_connection(host, username, password)
    if not session:
        success = False
        sys.exit(1)

    # Test 3: Command execution
    if not test_command_execution(session):
        success = False

    # Test 4: PowerShell execution
    if not test_powershell_execution(session):
        success = False

    if success:
        print(f"\n🎉 All tests passed! Python WinRM is working correctly.")
        print(f"\nNext steps:")
        print(f"  1. Create Python WinRM wrapper for existing scripts")
        print(f"  2. Compare performance with SSH method")
        print(f"  3. Implement error handling and logging")
    else:
        print(f"\n⚠️  Some tests failed. Check configuration and try again.")
        sys.exit(1)


if __name__ == "__main__":
    main()