# AgentX Fleet Client Runner for Windows (PowerShell 3.0 - 7+)
param (
    [string]$ServerUrl = "wss://agentx.toplevel.group",
    [string]$Token = "AgentX_Sec7788_9vKxM2qLpWz",
    [string]$NodeId = $env:COMPUTERNAME,
    [switch]$Install,
    [switch]$Uninstall
)

$ErrorActionPreference = "Continue"

# UTF-8 Encoding support for Windows Console and PowerShell output
try {
    [Console]::OutputEncoding = [System.Text.Encoding]::UTF8
    $OutputEncoding = [System.Text.Encoding]::UTF8
} catch {}

# Fallback for script path
if (-not $PSCommandPath) {
    if ($MyInvocation.MyCommand.Path) { $PSCommandPath = $MyInvocation.MyCommand.Path }
    elseif ($MyInvocation.MyCommand.Definition) { $PSCommandPath = $MyInvocation.MyCommand.Definition }
}

# Check Administrator Privileges
function Test-IsAdmin {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = New-Object Security.Principal.WindowsPrincipal($identity)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

# 1. Handle Uninstall
if ($Uninstall) {
    if (-not (Test-IsAdmin)) {
        Write-Host "[!] Elevating privileges for uninstall..." -ForegroundColor Yellow
        Start-Process powershell.exe -ArgumentList "-NoProfile -ExecutionPolicy Bypass -File `"$PSCommandPath`" -Uninstall" -Verb RunAs
        exit
    }
    Write-Host "[*] Removing AgentX Scheduled Task..." -ForegroundColor Yellow
    schtasks.exe /Delete /TN "AgentX" /F 2>$null
    Write-Host "[+] AgentX task removed." -ForegroundColor Green
    Add-Type -AssemblyName System.Windows.Forms -ErrorAction SilentlyContinue
    [System.Windows.Forms.MessageBox]::Show("AgentX успешно удален из системы.", "AgentX Fleet Runner", 0, 64) 2>$null
    exit
}

# 2. Handle Auto-Install as System Task (Runs at boot, survives reboots)
if ($Install) {
    if (-not (Test-IsAdmin)) {
        Write-Host "[!] Elevating privileges for installation..." -ForegroundColor Yellow
        Start-Process powershell.exe -ArgumentList "-NoProfile -ExecutionPolicy Bypass -File `"$PSCommandPath`" -Install" -Verb RunAs
        exit
    }

    $targetDir = "$env:ProgramData\AgentX"
    if (-not (Test-Path $targetDir)) {
        New-Item -ItemType Directory -Path $targetDir -Force | Out-Null
    }

    $targetFile = "$targetDir\runner.ps1"
    Copy-Item -Path $PSCommandPath -Destination $targetFile -Force

    $taskCmd = "powershell.exe -NoProfile -NonInteractive -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$targetFile`" -ServerUrl `"$ServerUrl`" -Token `"$Token`" -NodeId `"$NodeId`""
    
    Write-Host "[*] Registering AgentX scheduled task (runs on boot as SYSTEM)..." -ForegroundColor Yellow
    schtasks.exe /Create /TN "AgentX" /TR $taskCmd /SC ONSTART /RU "NT AUTHORITY\SYSTEM" /RL HIGHEST /F

    Write-Host "[*] Starting AgentX service..." -ForegroundColor Yellow
    schtasks.exe /Run /TN "AgentX"

    Write-Host "[+] AgentX successfully installed! It will now auto-start after every reboot." -ForegroundColor Green
    Add-Type -AssemblyName System.Windows.Forms -ErrorAction SilentlyContinue
    [System.Windows.Forms.MessageBox]::Show("AgentX успешно установлен и запущен как фоновая системная задача!", "AgentX Fleet Runner", 0, 64) 2>$null
    exit
}

Write-Host "======================================" -ForegroundColor Cyan
Write-Host "   AgentX Fleet Runner v1.2" -ForegroundColor Cyan
Write-Host "   Connecting to: $ServerUrl" -ForegroundColor Cyan
Write-Host "   Node ID: $NodeId" -ForegroundColor Cyan
Write-Host "======================================" -ForegroundColor Cyan

# Fallback check for .NET WebSocket support (Requires .NET 4.5+)
Add-Type -AssemblyName System.Net.Http -ErrorAction SilentlyContinue

function Run-AgentLoop {
    while ($true) {
        try {
            Write-Host "[*] Establishing connection to $ServerUrl..." -ForegroundColor Yellow
            $wsUri = [System.Uri]"$ServerUrl/?token=$Token"
            $ws = New-Object System.Net.WebSockets.ClientWebSocket
            $ws.Options.KeepAliveInterval = [System.TimeSpan]::FromSeconds(20)
            $cts = New-Object System.Threading.CancellationTokenSource

            $connectTask = $ws.ConnectAsync($wsUri, $cts.Token)
            $connectTask.Wait(10000)

            if ($ws.State -ne [System.Net.WebSockets.WebSocketState]::Open) {
                Write-Host "[-] Failed to connect. Retrying in 10s..." -ForegroundColor Red
                Start-Sleep -Seconds 10
                continue
            }

            Write-Host "[+] Connected successfully!" -ForegroundColor Green

            # Register
            $osDesc = (Get-WmiObject Win32_OperatingSystem -ErrorAction SilentlyContinue).Caption
            if (-not $osDesc) { $osDesc = [System.Environment]::OSVersion.ToString() }
            $regMsg = @{
                type = "register"
                nodeId = $NodeId
                hostname = $env:COMPUTERNAME
                os = $osDesc
                arch = $env:PROCESSOR_ARCHITECTURE
            } | ConvertTo-Json -Compress

            $regBytes = [System.Text.Encoding]::UTF8.GetBytes($regMsg)
            $regSeg = New-Object System.ArraySegment[byte] -ArgumentList @(,$regBytes)
            $ws.SendAsync($regSeg, [System.Net.WebSockets.WebSocketMessageType]::Text, $true, $cts.Token).Wait()

            # Message receiving loop
            $buffer = New-Object byte[] 65536
            while ($ws.State -eq [System.Net.WebSockets.WebSocketState]::Open) {
                $seg = New-Object System.ArraySegment[byte] -ArgumentList @(,$buffer)
                $recvTask = $ws.ReceiveAsync($seg, $cts.Token)
                $recvTask.Wait()

                if ($recvTask.Result.MessageType -eq [System.Net.WebSockets.WebSocketMessageType]::Close) {
                    Write-Host "[-] Server closed connection." -ForegroundColor Yellow
                    break
                }

                $msgStr = [System.Text.Encoding]::UTF8.GetString($buffer, 0, $recvTask.Result.Count)
                $json = $msgStr | ConvertFrom-Json

                if ($json.type -eq "exec") {
                    Write-Host "[>] Executing command ($($json.reqId)): $($json.command)" -ForegroundColor Cyan
                    
                    $outText = ""
                    $errText = ""
                    $exitCode = 0

                    try {
                        # Safe non-blocking execution via native pipeline (avoids buffer deadlocks)
                        $outText = & powershell.exe -NoProfile -NonInteractive -ExecutionPolicy Bypass -Command $json.command 2>&1 | Out-String
                        $exitCode = $LASTEXITCODE
                        if ($null -eq $exitCode) { $exitCode = 0 }
                    } catch {
                        $errText = $_.Exception.Message
                        $exitCode = 1
                    }

                    $resMsg = @{
                        type = "exec_result"
                        reqId = $json.reqId
                        exitCode = $exitCode
                        stdout = $outText
                        stderr = $errText
                    } | ConvertTo-Json -Compress

                    $resBytes = [System.Text.Encoding]::UTF8.GetBytes($resMsg)
                    $resSeg = New-Object System.ArraySegment[byte] -ArgumentList @(,$resBytes)
                    $ws.SendAsync($resSeg, [System.Net.WebSockets.WebSocketMessageType]::Text, $true, $cts.Token).Wait()
                    Write-Host "[<] Completed ($($json.reqId)) with code $exitCode" -ForegroundColor Green
                }
            }
        } catch {
            Write-Host "[-] Connection error: $($_.Exception.Message)" -ForegroundColor Red
        }

        Write-Host "[*] Reconnecting in 10 seconds..." -ForegroundColor DarkGray
        Start-Sleep -Seconds 10
    }
}

Run-AgentLoop