# Gadget Link password-only PowerShell client.
# This file is rendered by the server installer: the placeholder is never a
# secret and the public repository intentionally contains no deployment URL.
[CmdletBinding()]
param()

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

$RelayUrl = 'wss://link.lyutikov.online/client/ws'
$MaxResultBytes = 1024 * 1024
$MaxWebSocketMessageBytes = 6 * (2 * $MaxResultBytes + 4096) + 1024

function ConvertTo-Hex([byte[]]$Bytes) {
    return ([BitConverter]::ToString($Bytes)).Replace('-', '').ToLowerInvariant()
}

function Get-GadgetDeviceId([string]$Name) {
    $machineGuid = (Get-ItemProperty -LiteralPath 'HKLM:\SOFTWARE\Microsoft\Cryptography' -Name MachineGuid -ErrorAction Stop).MachineGuid
    $sha = [Security.Cryptography.SHA256]::Create()
    try {
        $hash = $sha.ComputeHash([Text.Encoding]::UTF8.GetBytes([string]$machineGuid))
        $suffix = (ConvertTo-Hex $hash).Substring(0, 12)
        if ([string]::IsNullOrWhiteSpace($Name)) { return 'pc-' + $suffix }
        $normalized = $Name.Trim()
        if ($normalized -notmatch '^[A-Za-z0-9][A-Za-z0-9._-]{0,47}$') {
            throw 'PC name must use A-Z, a-z, 0-9, dot, underscore, or hyphen (up to 48 characters).'
        }
        return $normalized + '-' + $suffix
    } finally {
        $sha.Dispose()
    }
}

function Get-GadgetSessionToken([Security.SecureString]$Password, [string]$DeviceId) {
    $bstr = [IntPtr]::Zero
    try {
        $bstr = [Runtime.InteropServices.Marshal]::SecureStringToBSTR($Password)
        $plain = [Runtime.InteropServices.Marshal]::PtrToStringBSTR($bstr)
        $hmac = [Security.Cryptography.HMACSHA256]::new([Text.Encoding]::UTF8.GetBytes($plain))
        try {
            $message = [Text.Encoding]::UTF8.GetBytes("gadget-link/shared-token/v1`0$DeviceId")
            return ConvertTo-Hex $hmac.ComputeHash($message)
        } finally {
            $hmac.Dispose()
            $plain = $null
        }
    } finally {
        if ($bstr -ne [IntPtr]::Zero) { [Runtime.InteropServices.Marshal]::ZeroFreeBSTR($bstr) }
    }
}

function Limit-GadgetText([string]$Text, [ref]$Truncated) {
    $bytes = [Text.Encoding]::UTF8.GetBytes($Text)
    if ($bytes.Length -le $MaxResultBytes) { return $Text }
    $Truncated.Value = $true
    return [Text.Encoding]::UTF8.GetString($bytes, 0, $MaxResultBytes)
}

function Invoke-GadgetCommand($Job) {
    $truncated = $false
    $result = [ordered]@{ job_id = $Job.id; exit_code = -1; stdout = ''; stderr = ''; timed_out = $false; truncated = $false; error = '' }
    try {
        if ($Job.shell -cne 'cmd') { throw 'relay sent an unsupported shell' }
        $timeout = [int]$Job.timeout_seconds
        if ($timeout -lt 1 -or $timeout -gt 300) { throw 'relay sent an invalid timeout' }
        $psi = New-Object Diagnostics.ProcessStartInfo
        $psi.FileName = $env:ComSpec
        $psi.Arguments = '/d /s /c "' + [string]$Job.command + '"'
        $psi.UseShellExecute = $false
        $psi.RedirectStandardOutput = $true
        $psi.RedirectStandardError = $true
        $psi.CreateNoWindow = $false
        if (-not [string]::IsNullOrWhiteSpace([string]$Job.cwd)) { $psi.WorkingDirectory = [string]$Job.cwd }
        $process = New-Object Diagnostics.Process
        $process.StartInfo = $psi
        if (-not $process.Start()) { throw 'cmd.exe did not start' }
        $stdoutTask = $process.StandardOutput.ReadToEndAsync()
        $stderrTask = $process.StandardError.ReadToEndAsync()
        if (-not $process.WaitForExit($timeout * 1000)) {
            $result.timed_out = $true
            & "$env:SystemRoot\System32\taskkill.exe" /PID $process.Id /T /F 2>$null | Out-Null
            $process.WaitForExit()
        }
        [Threading.Tasks.Task]::WaitAll(@($stdoutTask, $stderrTask))
        $result.exit_code = $process.ExitCode
        $result.stdout = Limit-GadgetText ([string]$stdoutTask.Result) ([ref]$truncated)
        $result.stderr = Limit-GadgetText ([string]$stderrTask.Result) ([ref]$truncated)
        $result.truncated = $truncated
        $process.Dispose()
    } catch {
        $result.error = Limit-GadgetText $_.Exception.Message ([ref]$truncated)
        $result.truncated = $truncated
    }
    return $result
}

$uri = $null
if (-not [Uri]::TryCreate($RelayUrl, [UriKind]::Absolute, [ref]$uri) -or $uri.Scheme -cne 'wss' -or -not [string]::IsNullOrEmpty($uri.UserInfo)) {
    throw 'This copy is not configured for a relay. Use the /gadget URL from the Gadget Link server.'
}

$password = Read-Host 'Gadget Link password' -AsSecureString
$pcName = Read-Host 'PC name (A-Z, 0-9, dot, underscore, hyphen; optional)'
$deviceId = Get-GadgetDeviceId $pcName
$token = Get-GadgetSessionToken $password $deviceId
$password = $null

function Test-GadgetStopRequest {
    try {
        if (-not [Console]::KeyAvailable) { return $false }
        $key = [Console]::ReadKey($true)
        return $key.Key -eq [ConsoleKey]::C -and ($key.Modifiers -band [ConsoleModifiers]::Control)
    } catch {
        # Some non-interactive hosts expose no console input. The interactive
        # launcher has a console, so Ctrl+C is still the explicit stop action.
        return $false
    }
}

function Wait-GadgetTask([Threading.Tasks.Task]$Task) {
    while (-not $Task.Wait(150)) {
        if (Test-GadgetStopRequest) { return $false }
    }
    return $true
}

function Wait-GadgetReconnectDelay([int]$Seconds) {
    $until = [DateTime]::UtcNow.AddSeconds($Seconds)
    while ([DateTime]::UtcNow -lt $until) {
        if (Test-GadgetStopRequest) { return $false }
        [Threading.Thread]::Sleep(150)
    }
    return $true
}

$MaxReconnectDelaySeconds = 30
$restoreControlCInput = $false
$stopRequested = $false
$reconnectDelaySeconds = 1
try {
    try {
        $previousControlCInput = [Console]::TreatControlCAsInput
        [Console]::TreatControlCAsInput = $true
        $restoreControlCInput = $true
    } catch {}

    while (-not $stopRequested) {
        $socket = $null
        $problem = $null
        try {
            $socket = New-Object Net.WebSockets.ClientWebSocket
            $socket.Options.SetRequestHeader('Authorization', "Bearer $token")
            $socket.Options.SetRequestHeader('X-Gadget-Device-ID', $deviceId)
            $socket.ConnectAsync($uri, [Threading.CancellationToken]::None).GetAwaiter().GetResult() | Out-Null
            $reconnectDelaySeconds = 1
            Write-Host "Gadget Link connected as $deviceId. Press Ctrl+C to disconnect."
            $buffer = New-Object byte[] 65536
            while ($socket.State -eq [Net.WebSockets.WebSocketState]::Open -and -not $stopRequested) {
                $stream = New-Object IO.MemoryStream
                try {
                    do {
                        $segment = New-Object ArraySegment[byte] -ArgumentList (, $buffer)
                        $receiveTask = $socket.ReceiveAsync($segment, [Threading.CancellationToken]::None)
                        if (-not (Wait-GadgetTask $receiveTask)) {
                            $stopRequested = $true
                            break
                        }
                        $received = $receiveTask.GetAwaiter().GetResult()
                        if ($received.MessageType -eq [Net.WebSockets.WebSocketMessageType]::Close) { break }
                        if (($stream.Length + $received.Count) -gt $MaxWebSocketMessageBytes) { throw 'relay WebSocket message exceeds the protocol limit' }
                        $stream.Write($buffer, 0, $received.Count)
                    } while (-not $received.EndOfMessage)
                    if ($stopRequested -or $received.MessageType -eq [Net.WebSockets.WebSocketMessageType]::Close) { break }
                    $job = [Text.Encoding]::UTF8.GetString($stream.ToArray()) | ConvertFrom-Json
                    $result = Invoke-GadgetCommand $job | ConvertTo-Json -Compress
                    $payload = [Text.Encoding]::UTF8.GetBytes($result)
                    $socket.SendAsync((New-Object ArraySegment[byte] -ArgumentList (, $payload)), [Net.WebSockets.WebSocketMessageType]::Text, $true, [Threading.CancellationToken]::None).GetAwaiter().GetResult() | Out-Null
                } finally {
                    $stream.Dispose()
                }
            }
            if (-not $stopRequested) { $problem = 'relay closed the WebSocket session' }
        } catch {
            $problem = $_.Exception.GetBaseException().Message
        } finally {
            if ($null -ne $socket) {
                if ($socket.State -eq [Net.WebSockets.WebSocketState]::Open) {
                    try {
                        $socket.CloseOutputAsync([Net.WebSockets.WebSocketCloseStatus]::NormalClosure, 'client reconnecting', [Threading.CancellationToken]::None).GetAwaiter().GetResult() | Out-Null
                    } catch {}
                }
                $socket.Dispose()
            }
        }

        if ($stopRequested) { break }
        Write-Warning "Gadget Link disconnected: $problem. Reconnecting in $reconnectDelaySeconds second(s). Press Ctrl+C to disconnect."
        if (-not (Wait-GadgetReconnectDelay $reconnectDelaySeconds)) { $stopRequested = $true; break }
        $reconnectDelaySeconds = [Math]::Min($reconnectDelaySeconds * 2, $MaxReconnectDelaySeconds)
    }
} finally {
    if ($restoreControlCInput) { [Console]::TreatControlCAsInput = $previousControlCInput }
    $token = $null
}