#!/usr/bin/env python3
"""
Proxmox VE API 테스트 스크립트
사용법: python test-proxmox-api.py <proxmox-ip> [root-password]
"""

import requests
import json
import sys
import urllib3
from urllib.parse import quote_plus

# SSL 경고 무시
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

def test_api_token_method(host, token):
    """API Token 방식 테스트"""
    print("🔑 API Token 방식 테스트...")
    
    headers = {
        'Authorization': f'PVEAPIToken={token}'
    }
    
    try:
        response = requests.get(
            f'https://{host}:443/api2/json/version',
            headers=headers,
            verify=False,
            timeout=10
        )
        
        if response.status_code == 200:
            version_info = response.json()['data']
            print(f"✅ API Token 인증 성공!")
            print(f"   Proxmox Version: {version_info['version']}")
            print(f"   Release: {version_info['release']}")
            return True
        else:
            print(f"❌ API Token 인증 실패: {response.status_code}")
            print(f"   응답: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ API Token 연결 실패: {e}")
        return False

def test_session_method(host, username, password):
    """세션 쿠키 방식 테스트"""
    print("🍪 세션 쿠키 방식 테스트...")
    
    # 1단계: 로그인하여 티켓 획득
    login_data = {
        'username': username,
        'password': password
    }
    
    try:
        response = requests.post(
            f'https://{host}:443/api2/json/access/ticket',
            data=login_data,
            verify=False,
            timeout=10
        )
        
        if response.status_code != 200:
            print(f"❌ 로그인 실패: {response.status_code}")
            print(f"   응답: {response.text}")
            return False, None, None
        
        data = response.json()['data']
        ticket = data['ticket']
        csrf_token = data['CSRFPreventionToken']
        
        print("✅ 로그인 성공!")
        print(f"   티켓: {ticket[:20]}...")
        print(f"   CSRF: {csrf_token}")
        
        return True, ticket, csrf_token
        
    except Exception as e:
        print(f"❌ 로그인 연결 실패: {e}")
        return False, None, None

def test_vm_list_api(host, ticket=None, csrf_token=None, api_token=None):
    """VM 목록 조회 API 테스트"""
    print("\n📋 VM 목록 조회 API 테스트...")
    
    if api_token:
        headers = {'Authorization': f'PVEAPIToken={api_token}'}
        cookies = None
    else:
        headers = {'CSRFPreventionToken': csrf_token}
        cookies = {'PVEAuthCookie': ticket}
    
    try:
        response = requests.get(
            f'https://{host}:443/api2/json/cluster/resources?type=vm',
            headers=headers,
            cookies=cookies,
            verify=False,
            timeout=10
        )
        
        if response.status_code == 200:
            vms = response.json()['data']
            print(f"✅ VM 목록 조회 성공! (총 {len(vms)}개)")
            
            for vm in vms:
                status_icon = "🟢" if vm.get('status') == 'running' else "🔴"
                print(f"   {status_icon} VM {vm.get('vmid')}: {vm.get('name', 'N/A')} ({vm.get('status', 'unknown')})")
            
            return True, vms
        else:
            print(f"❌ VM 목록 조회 실패: {response.status_code}")
            print(f"   응답: {response.text}")
            return False, []
            
    except Exception as e:
        print(f"❌ VM 목록 API 연결 실패: {e}")
        return False, []

def test_vnc_proxy_api(host, node, vmid, ticket=None, csrf_token=None, api_token=None):
    """VNC 프록시 티켓 생성 테스트"""
    print(f"\n🖥️ VNC 프록시 API 테스트 (VM {vmid})...")
    
    if api_token:
        headers = {'Authorization': f'PVEAPIToken={api_token}'}
        cookies = None
    else:
        headers = {'CSRFPreventionToken': csrf_token}
        cookies = {'PVEAuthCookie': ticket}
    
    data = {'websocket': '1'}
    
    try:
        response = requests.post(
            f'https://{host}:443/api2/json/nodes/{node}/qemu/{vmid}/vncproxy',
            headers=headers,
            cookies=cookies,
            data=data,
            verify=False,
            timeout=10
        )
        
        if response.status_code == 200:
            vnc_data = response.json()['data']
            print("✅ VNC 프록시 티켓 생성 성공!")
            print(f"   포트: {vnc_data['port']}")
            print(f"   티켓: {vnc_data['ticket'][:20]}...")
            
            # WebSocket URL 생성
            encoded_ticket = quote_plus(vnc_data['ticket'])
            ws_url = f"wss://{host}:443/api2/json/nodes/{node}/qemu/{vmid}/vncwebsocket?port={vnc_data['port']}&vncticket={encoded_ticket}"
            print(f"   WebSocket URL: {ws_url[:80]}...")
            
            return True, vnc_data
        else:
            print(f"❌ VNC 프록시 생성 실패: {response.status_code}")
            print(f"   응답: {response.text}")
            return False, None
            
    except Exception as e:
        print(f"❌ VNC 프록시 API 연결 실패: {e}")
        return False, None

def main():
    if len(sys.argv) < 2:
        print("사용법: python test-proxmox-api.py <proxmox-ip> [root-password] [api-token]")
        print()
        print("예시:")
        print("  python test-proxmox-api.py 100.64.0.1 mypassword")
        print("  python test-proxmox-api.py 100.64.0.1 '' root@pam!token=uuid")
        sys.exit(1)
    
    host = sys.argv[1]
    password = sys.argv[2] if len(sys.argv) > 2 else None
    api_token = sys.argv[3] if len(sys.argv) > 3 else None
    
    print("=" * 60)
    print(f"🚀 Proxmox API 테스트 - {host}")
    print("=" * 60)
    
    # API Token 방식 테스트
    if api_token:
        success = test_api_token_method(host, api_token)
        if success:
            # VM 목록 테스트
            vm_success, vms = test_vm_list_api(host, api_token=api_token)
            
            # 실행 중인 VM이 있으면 VNC 테스트
            if vm_success and vms:
                running_vm = next((vm for vm in vms if vm.get('status') == 'running'), None)
                if running_vm:
                    test_vnc_proxy_api(host, running_vm['node'], running_vm['vmid'], api_token=api_token)
                else:
                    print("⚠️ 실행 중인 VM이 없어 VNC 테스트를 건너뜁니다.")
    
    # 패스워드 방식 테스트
    elif password:
        login_success, ticket, csrf_token = test_session_method(host, 'root@pam', password)
        
        if login_success:
            # VM 목록 테스트
            vm_success, vms = test_vm_list_api(host, ticket, csrf_token)
            
            # 실행 중인 VM이 있으면 VNC 테스트
            if vm_success and vms:
                running_vm = next((vm for vm in vms if vm.get('status') == 'running'), None)
                if running_vm:
                    test_vnc_proxy_api(host, running_vm['node'], running_vm['vmid'], ticket, csrf_token)
                else:
                    print("⚠️ 실행 중인 VM이 없어 VNC 테스트를 건너뜁니다.")
    
    else:
        print("❌ 패스워드 또는 API 토큰을 제공해주세요.")
        sys.exit(1)
    
    print("\n" + "=" * 60)
    print("✅ 테스트 완료!")

if __name__ == '__main__':
    main()