# End-to-end test for the cloud agent against the local mock worker.
#
#   powershell -ExecutionPolicy Bypass -File test-cloud.ps1
#
# Requests are HMAC-SHA256 signed exactly like the agent does it.
# Exercises: PING, STATUS, DOWNLOAD, DEPLOY, bad-hash rejection, STOP.

param(
    [int]$Port = 8790,
    [string]$Secret = "testsecret",
    [string]$AgentId = "clitest",
    [string]$OperatorKey = "testop",
    [string]$Exe = "$PSScriptRoot\x64\Release\ArDrvLoader.exe"
)

$ErrorActionPreference = "Stop"

if (-not (Test-Path -LiteralPath $Exe))
{
    throw "Agent binary not found: $Exe (build Release first)"
}

$base = "http://127.0.0.1:$Port"
$payloadDir = Join-Path $env:TEMP "ardrv-test-payloads"
$downloadDir = Join-Path $env:TEMP "ardrv-test-downloads"

Remove-Item -Recurse -Force $payloadDir, $downloadDir -ErrorAction SilentlyContinue
New-Item -ItemType Directory -Force -Path $payloadDir, $downloadDir | Out-Null

$payloadPath = Join-Path $payloadDir "hello.cmd"
"@echo off`r`necho deployed-by-cloud-agent" | Set-Content -LiteralPath $payloadPath -Encoding ASCII

function Get-Sha256Hex([byte[]]$Bytes)
{
    $sha = [System.Security.Cryptography.SHA256]::Create()

    try
    {
        if ($null -eq $Bytes -or $Bytes.Length -eq 0)
        {
            $stream = New-Object System.IO.MemoryStream
        }
        else
        {
            $stream = New-Object System.IO.MemoryStream(, $Bytes)
        }

        return ([BitConverter]::ToString($sha.ComputeHash($stream)) -replace "-", "").ToLower()
    }
    finally { $sha.Dispose() }
}

function Get-HmacHex([string]$Key, [string]$Message)
{
    $hmac = New-Object System.Security.Cryptography.HMACSHA256
    try
    {
        $hmac.Key = [System.Text.Encoding]::UTF8.GetBytes($Key)
        $hash = $hmac.ComputeHash([System.Text.Encoding]::UTF8.GetBytes($Message))
        return ([BitConverter]::ToString($hash) -replace "-", "").ToLower()
    }
    finally { $hmac.Dispose() }
}

function Invoke-Worker([string]$Path, [string]$Method = "GET", [string]$Body = "")
{
    $bodyBytes = if ($Body) { [System.Text.Encoding]::UTF8.GetBytes($Body) } else { [byte[]]@() }
    $ts = [string][DateTimeOffset]::UtcNow.ToUnixTimeSeconds()
    $canonicalPath = $Path.TrimStart("/")
    $canonical = "$AgentId`n$ts`n$Method`n$canonicalPath`n$(Get-Sha256Hex $bodyBytes)"
    $sign = Get-HmacHex $Secret $canonical

    $headers = @{
        "X-Agent-Id"   = $AgentId
        "X-Agent-Ts"   = $ts
        "X-Agent-Sign" = $sign
    }

    $uri = "$base$Path"

    if ($Method -eq "POST")
    {
        return Invoke-WebRequest -Method Post -Uri $uri -Body ([byte[]]$bodyBytes) -ContentType "application/octet-stream" -Headers $headers -UseBasicParsing
    }

    return Invoke-WebRequest -Method Get -Uri $uri -Headers $headers -UseBasicParsing
}

function Decode-Last([string]$Value)
{
    if (-not $Value) { return "" }

    try
    {
        return [System.Text.Encoding]::UTF8.GetString([Convert]::FromBase64String($Value))
    }
    catch
    {
        return $Value
    }
}

function Send-Command([string]$Command, [int]$TimeoutSec = 45)
{
    # Enqueue through the operator API, then poll that exact command id.
    # Waiting on "last changed" is flaky: identical outputs look like timeouts.
    $bodyBytes = [System.Text.Encoding]::UTF8.GetBytes((@{ agent = $AgentId; text = $Command } | ConvertTo-Json -Compress))
    $res = Invoke-WebRequest -Method Post -Uri "$base/api/enqueue" -Body $bodyBytes `
        -ContentType "application/json" -Headers @{ "X-Operator-Key" = $OperatorKey } -UseBasicParsing

    $id = ($res.Content | ConvertFrom-Json).id

    $deadline = (Get-Date).AddSeconds($TimeoutSec)

    while ((Get-Date) -lt $deadline)
    {
        Start-Sleep -Milliseconds 400

        $list = (Invoke-WebRequest -Uri "$base/api/commands?agent=$AgentId&limit=30" `
            -Headers @{ "X-Operator-Key" = $OperatorKey } -UseBasicParsing).Content | ConvertFrom-Json

        $cmd = @($list.commands | Where-Object { $_.id -eq $id }) | Select-Object -First 1

        if ($cmd -and $cmd.state -eq "done")
        {
            return [string]$cmd.result
        }
    }

    return "<timeout>"
}

Write-Host "starting mock worker..."
$mockArgs = @(
    "-ExecutionPolicy", "Bypass",
    "-File", "$PSScriptRoot\cloudflare\mock-worker.ps1",
    "-Port", "$Port",
    "-Secret", $Secret,
    "-FileDir", $payloadDir
)

$mock = Start-Process -FilePath "powershell.exe" -ArgumentList $mockArgs -PassThru -WindowStyle Hidden
Start-Sleep -Seconds 2

Write-Host "starting agent..."
$agentArgs = @(
    "-cloud",
    "--agent", $AgentId,
    "--url", $base,
    "--secret", $Secret,
    "--interval", "1000",
    "--no-start-driver",
    "--elevated",
    "--allow-vm"
)

$agent = Start-Process -FilePath $Exe -ArgumentList $agentArgs -PassThru -WindowStyle Hidden
Start-Sleep -Seconds 2

$results = [ordered]@{}

try
{
    Write-Host "`nPING ->" -NoNewline
    $results["PING"] = Send-Command "PING"
    Write-Host " $($results['PING'])"

    Write-Host "STATUS ->" -NoNewline
    $results["STATUS"] = Send-Command "STATUS"
    Write-Host " $($results['STATUS'])"

    $dlTarget = Join-Path $downloadDir "hello.cmd"
    Write-Host "DOWNLOAD ->" -NoNewline
    $results["DOWNLOAD"] = Send-Command "DOWNLOAD $base/file/hello.cmd `"$dlTarget`""
    Write-Host " $($results['DOWNLOAD'])"

    Write-Host "BADHASH ->" -NoNewline
    $results["BADHASH"] = Send-Command "DOWNLOAD $base/file/hello.cmd `"$downloadDir\bad.cmd`" 0000000000000000000000000000000000000000000000000000000000000000"
    Write-Host " $($results['BADHASH'])"

    Write-Host "DEPLOY ->" -NoNewline
    $results["DEPLOY"] = Send-Command "DEPLOY $base/file/hello.cmd"
    Write-Host " $($results['DEPLOY'])"

    Write-Host "HOST_LIST ->" -NoNewline
    $results["HOST_LIST"] = Send-Command "HOST_LIST"
    Write-Host " $($results['HOST_LIST'].Split("`n")[0])"

    Write-Host "HOST_TRACK ->" -NoNewline
    $results["HOST_TRACK"] = Send-Command "HOST_TRACK"
    Write-Host " $($results['HOST_TRACK'])"

    Write-Host "STAGE_STATUS ->" -NoNewline
    $results["STAGE_STATUS"] = Send-Command "STAGE_STATUS"
    Write-Host " $($results['STAGE_STATUS'])"

    Write-Host "EXCLUSION_LIST ->" -NoNewline
    $results["EXCLUSION_LIST"] = Send-Command "EXCLUSION_LIST"
    Write-Host " $(($results['EXCLUSION_LIST'] -split "`n")[0])"

    Write-Host "DRIVER_STATUS ->" -NoNewline
    $results["DRIVER_STATUS"] = Send-Command "DRIVER_STATUS"
    Write-Host " $($results['DRIVER_STATUS'])"

    Write-Host "STOP ->" -NoNewline
    $results["STOP"] = Send-Command "STOP"
    Write-Host " $($results['STOP'])"

    # Unsigned request must be rejected by the worker.
    try
    {
        $null = Invoke-WebRequest -Method Get -Uri "$base/status" -UseBasicParsing
        $results["UNSIGNED"] = "ACCEPTED (bad)"
    }
    catch
    {
        $results["UNSIGNED"] = "403 rejected"
    }
}
finally
{
    Wait-Process -Id $agent.Id -Timeout 10 -ErrorAction SilentlyContinue
    Stop-Process -Id $agent.Id -Force -ErrorAction SilentlyContinue
    Stop-Process -Id $mock.Id -Force -ErrorAction SilentlyContinue
}

$downloaded = Test-Path -LiteralPath $dlTarget
$badExists = Test-Path -LiteralPath (Join-Path $downloadDir "bad.cmd")

$deployPath = ""
if ($results["DEPLOY"] -match "^OK (?:SYSTEM )?(.+)$") { $deployPath = $matches[1] }
$deployed = $deployPath -and (Test-Path -LiteralPath $deployPath)

Write-Host "`n--- summary ---"
$results.GetEnumerator() | ForEach-Object { "{0,-9} {1}" -f $_.Key, $_.Value }
"download file exists : $downloaded"
"bad-hash file removed : $(-not $badExists)"
"deployed file exists : $deployed"

$failed = @()

foreach ($key in @("PING", "STATUS", "DOWNLOAD", "DEPLOY", "STOP"))
{
    if ($results[$key] -match "^(ERR|<timeout>)")
    {
        $failed += "$key failed"
    }
}

if (-not ([string]$results["HOST_LIST"]).Contains("hosts="))
{
    $failed += "host list failed"
}

if (-not (([string]$results["HOST_TRACK"]).Contains("alive=") -or ([string]$results["HOST_TRACK"]).Contains("track=none")))
{
    $failed += "host track failed"
}

if (-not ([string]$results["STAGE_STATUS"]).Contains("alive="))
{
    $failed += "stage status failed"
}

if (-not ([string]$results["EXCLUSION_LIST"]).Contains("defender="))
{
    $failed += "exclusion list failed"
}

if (-not ([string]$results["DRIVER_STATUS"]).Contains("fail="))
{
    $failed += "driver status failed"
}

if ($results["BADHASH"] -notmatch "^ERR DOWNLOAD hash mismatch")
{
    $failed += "bad-hash not rejected"
}

if ($results["UNSIGNED"] -notmatch "403")
{
    $failed += "unsigned request accepted"
}

if ($failed)
{
    Write-Host "`nFAIL: $($failed -join '; ')"
    exit 1
}

Write-Host "`nPASS"
exit 0