<#
    LXS - Remote Access (Windows)
    Description: Remote Desktop and the OpenSSH server - status, enable,
                 disable. The Windows counterpart of root-ssh-login.sh.
    Repo: https://git.hyko.cx/hykocx/lxs
#>

# Load LXS common library (colors, UI helpers, spinner, loggers, guards).
# Prefers the sibling ..\lib\common.ps1 (repo checkout or installed layout) and
# only hits the network when this script is run standalone.
$LxsRawBase = if ($env:LXS_RAW_BASE) { $env:LXS_RAW_BASE } else { 'https://git.hyko.cx/hykocx/lxs/raw/branch/main' }
if (-not $env:LXS_RAW_PLATFORM_BASE) { $env:LXS_RAW_PLATFORM_BASE = "$LxsRawBase/windows" }
$LxsLibPath = if ($PSScriptRoot) { Join-Path $PSScriptRoot '..\lib\common.ps1' } else { $null }
if ($LxsLibPath -and (Test-Path $LxsLibPath)) {
    . $LxsLibPath
} else {
    try {
        $LxsLibSource = Invoke-RestMethod -Uri "$env:LXS_RAW_PLATFORM_BASE/lib/common.ps1" -UseBasicParsing -ErrorAction Stop
    } catch {
        Write-Error 'Failed to fetch lib/common.ps1'
        exit 1
    }
    # The library ships with a UTF-8 BOM (Windows PowerShell 5.1 needs it to
    # read the file as UTF-8); over HTTP that BOM arrives as a leading U+FEFF
    # character, which the parser will not accept.
    . ([scriptblock]::Create(($LxsLibSource -replace '^\uFEFF', '')))
}

$env:LXS_LOG_FILE = Join-Path (Get-LxsTempDir) 'lxs_remote_access.log'

if (-not (Assert-LxsWindows)) { exit 1 }
Assert-LxsAdmin -ScriptPath $PSCommandPath -Arguments $args

$LxsRdpKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server'
$LxsRdpTcpKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp'

function Get-LxsRdpState {
    try {
        # fDenyTSConnections 0 = RDP allowed, 1 = denied.
        $deny = (Get-ItemProperty -Path $LxsRdpKey -Name 'fDenyTSConnections' -ErrorAction Stop).fDenyTSConnections
        return ($deny -eq 0)
    } catch {
        return $false
    }
}

function Get-LxsNlaState {
    try {
        return ((Get-ItemProperty -Path $LxsRdpTcpKey -Name 'UserAuthentication' -ErrorAction Stop).UserAuthentication -eq 1)
    } catch {
        return $false
    }
}

function Show-LxsRemoteStatus {
    Clear-Host
    Show-LxsBoxTop -Title 'REMOTE ACCESS STATUS'
    Write-Host ''

    Write-Host "$($script:Cyan)$($script:Bold)Remote Desktop$($script:NC)"
    if (Get-LxsRdpState) {
        Write-LxsOk 'RDP is ENABLED'
        if (Get-LxsNlaState) {
            Write-LxsOk 'Network Level Authentication is on'
        } else {
            Write-LxsWarn 'Network Level Authentication is OFF - turn it on (option 2).'
        }
        try {
            $port = (Get-ItemProperty -Path $LxsRdpTcpKey -Name 'PortNumber' -ErrorAction Stop).PortNumber
            Write-Host "  Listening port: $port"
        } catch {
            Write-Host '  Listening port: 3389 (default)'
        }
    } else {
        Write-Host "  $($script:Gray)RDP is disabled$($script:NC)"
    }

    try {
        $rules = @(Get-NetFirewallRule -DisplayGroup 'Remote Desktop' -ErrorAction Stop | Where-Object { $_.Enabled -eq 'True' })
        Write-Host "  Firewall rules enabled: $($rules.Count)"
    } catch {
        # Group may not exist on Server Core images.
    }

    Write-Host ''
    Write-Host "$($script:Cyan)$($script:Bold)OpenSSH Server$($script:NC)"
    $svc = Get-Service -Name sshd -ErrorAction SilentlyContinue
    if (-not $svc) {
        Write-Host "  $($script:Gray)Not installed$($script:NC)"
    } else {
        if ($svc.Status -eq 'Running') { Write-LxsOk "sshd is running (startup: $($svc.StartType))" }
        else { Write-LxsWarn "sshd is installed but $($svc.Status) (startup: $($svc.StartType))" }
        try {
            $shell = (Get-ItemProperty -Path 'HKLM:\SOFTWARE\OpenSSH' -Name 'DefaultShell' -ErrorAction Stop).DefaultShell
            Write-Host "  Default shell : $shell"
        } catch {
            Write-Host "  Default shell : cmd.exe (Windows default)"
        }
    }
    Write-Host ''
}

function Enable-LxsRdp {
    Clear-Host
    Show-LxsBoxTop -Title 'ENABLE REMOTE DESKTOP'
    Write-Host ''
    Write-Host 'This allows inbound RDP connections and opens the matching'
    Write-Host 'firewall rules. Network Level Authentication will be required,'
    Write-Host 'so clients must authenticate before a session is created.'
    Write-Host ''
    if (-not (Confirm-LxsAction -Question 'Enable Remote Desktop?')) {
        Write-LxsInfo 'Cancelled.'
        return
    }
    Write-Host ''
    if (Set-LxsRegistryValue -Path $LxsRdpKey -Name 'fDenyTSConnections' -Value 0) {
        Write-LxsOk 'RDP connections allowed'
    }
    if (Set-LxsRegistryValue -Path $LxsRdpTcpKey -Name 'UserAuthentication' -Value 1) {
        Write-LxsOk 'Network Level Authentication required'
    }
    try {
        Enable-NetFirewallRule -DisplayGroup 'Remote Desktop' -ErrorAction Stop
        Write-LxsOk 'Firewall rules for Remote Desktop enabled'
    } catch {
        Write-LxsWarn "Could not enable the firewall rules: $($_.Exception.Message)"
    }
    Write-Host ''
    Write-LxsWarn 'Only expose RDP to the internet behind a VPN - it is a constant brute-force target.'
}

function Disable-LxsRdp {
    Clear-Host
    Show-LxsBoxTop -Title 'DISABLE REMOTE DESKTOP'
    Write-Host ''
    if (-not (Confirm-LxsAction -Question 'Disable Remote Desktop and its firewall rules?')) {
        Write-LxsInfo 'Cancelled.'
        return
    }
    Write-Host ''
    if (Set-LxsRegistryValue -Path $LxsRdpKey -Name 'fDenyTSConnections' -Value 1) {
        Write-LxsOk 'RDP connections denied'
    }
    try {
        Disable-NetFirewallRule -DisplayGroup 'Remote Desktop' -ErrorAction Stop
        Write-LxsOk 'Firewall rules for Remote Desktop disabled'
    } catch {
        Write-LxsWarn "Could not disable the firewall rules: $($_.Exception.Message)"
    }
}

function Install-LxsSshServer {
    Clear-Host
    Show-LxsBoxTop -Title 'OPENSSH SERVER'
    Write-Host ''
    Write-Host 'Installs the OpenSSH Server capability, starts sshd, sets it to'
    Write-Host 'start automatically and opens TCP/22 in the firewall.'
    Write-Host ''
    if (-not (Confirm-LxsAction -Question 'Install and enable the OpenSSH server?')) {
        Write-LxsInfo 'Cancelled.'
        return
    }
    Write-Host ''

    try {
        $cap = Get-WindowsCapability -Online -Name 'OpenSSH.Server*' -ErrorAction Stop |
            Select-Object -First 1
    } catch {
        Write-LxsErr "Could not query Windows capabilities: $($_.Exception.Message)"
        return
    }

    if (-not $cap) {
        Write-LxsErr 'The OpenSSH Server capability is not available on this edition.'
        return
    }

    if ($cap.State -ne 'Installed') {
        Write-LxsInfo "Installing $($cap.Name) (this can take a few minutes)..."
        try {
            Add-WindowsCapability -Online -Name $cap.Name -ErrorAction Stop | Out-Null
            Write-LxsOk 'OpenSSH Server installed'
        } catch {
            Write-LxsErr "Installation failed: $($_.Exception.Message)"
            return
        }
    } else {
        Write-LxsOk 'OpenSSH Server was already installed'
    }

    try {
        Set-Service -Name sshd -StartupType Automatic -ErrorAction Stop
        Start-Service -Name sshd -ErrorAction Stop
        Write-LxsOk 'sshd is running and set to start automatically'
    } catch {
        Write-LxsErr "Could not start sshd: $($_.Exception.Message)"
        return
    }

    # The capability normally creates this rule; create it when it is missing.
    try {
        if (-not (Get-NetFirewallRule -Name 'OpenSSH-Server-In-TCP' -ErrorAction SilentlyContinue)) {
            New-NetFirewallRule -Name 'OpenSSH-Server-In-TCP' -DisplayName 'OpenSSH Server (sshd)' `
                -Enabled True -Direction Inbound -Protocol TCP -Action Allow -LocalPort 22 -ErrorAction Stop | Out-Null
            Write-LxsOk 'Firewall rule created for TCP/22'
        } else {
            Write-LxsOk 'Firewall rule for TCP/22 already present'
        }
    } catch {
        Write-LxsWarn "Could not create the firewall rule: $($_.Exception.Message)"
    }

    Write-Host ''
    Write-Host "$($script:Gray)Connect with: ssh $env:USERNAME@$(Get-LxsPublicIP)$($script:NC)"
    Write-Host "$($script:Gray)Key-based auth: put your public key in %ProgramData%\ssh\administrators_authorized_keys$($script:NC)"
    Write-Host "$($script:Gray)for admin accounts, or in %USERPROFILE%\.ssh\authorized_keys otherwise.$($script:NC)"
}

function Disable-LxsSshServer {
    Clear-Host
    Show-LxsBoxTop -Title 'DISABLE OPENSSH SERVER'
    Write-Host ''
    $svc = Get-Service -Name sshd -ErrorAction SilentlyContinue
    if (-not $svc) {
        Write-LxsWarn 'The OpenSSH server is not installed.'
        return
    }
    if (-not (Confirm-LxsAction -Question 'Stop sshd and set it to Disabled?')) {
        Write-LxsInfo 'Cancelled.'
        return
    }
    try {
        Stop-Service -Name sshd -Force -ErrorAction Stop
        Set-Service -Name sshd -StartupType Disabled -ErrorAction Stop
        Write-LxsOk 'sshd stopped and disabled'
    } catch {
        Write-LxsErr "Failed: $($_.Exception.Message)"
    }
}

function Set-LxsSshDefaultShell {
    Clear-Host
    Show-LxsBoxTop -Title 'SSH DEFAULT SHELL'
    Write-Host ''
    Write-Host 'Which shell should SSH sessions land in?'
    Write-Host ''
    Show-LxsMenuItem '1' 'PowerShell' 'powershell.exe'
    Show-LxsMenuItem '2' 'PowerShell 7' 'pwsh.exe (must be installed)'
    Show-LxsMenuItem '3' 'Command Prompt' 'cmd.exe (Windows default)'
    Show-LxsMenuItem '0' 'Cancel' '' -Exit
    Write-Host ''
    $choice = Read-LxsChoice

    $shell = switch ($choice) {
        '1' { "$env:SystemRoot\System32\WindowsPowerShell\v1.0\powershell.exe" }
        '2' { (Get-Command pwsh -ErrorAction SilentlyContinue).Source }
        '3' { "$env:SystemRoot\System32\cmd.exe" }
        default { $null }
    }
    if (-not $shell) {
        if ($choice -eq '2') { Write-LxsErr 'pwsh.exe was not found - install PowerShell 7 first.' }
        else { Write-LxsInfo 'Cancelled.' }
        return
    }
    Write-Host ''
    if (Set-LxsRegistryValue -Path 'HKLM:\SOFTWARE\OpenSSH' -Name 'DefaultShell' -Value $shell -Type String) {
        Write-LxsOk "SSH sessions will start: $shell"
    }
}

function Show-LxsRemoteMenu {
    while ($true) {
        Clear-Host
        Show-LxsBoxTop -Title 'REMOTE ACCESS' -Right 'ADMIN'
        Write-Host ''
        Show-LxsMenuItem '1' 'Show status'
        Show-LxsMenuItem '2' 'Enable Remote Desktop'  'with NLA + firewall'
        Show-LxsMenuItem '3' 'Disable Remote Desktop'
        Show-LxsMenuItem '4' 'Install OpenSSH Server' 'service + firewall'
        Show-LxsMenuItem '5' 'Disable OpenSSH Server'
        Show-LxsMenuItem '6' 'Set SSH default shell'
        Show-LxsMenuItem '0' 'Back' '' -Exit
        Write-Host ''
        Show-LxsBoxBottom
        Write-Host ''
        $choice = Read-LxsChoice
        Write-Host ''

        switch ($choice) {
            '1' { Show-LxsRemoteStatus;    Read-LxsEnter }
            '2' { Enable-LxsRdp;           Read-LxsEnter }
            '3' { Disable-LxsRdp;          Read-LxsEnter }
            '4' { Install-LxsSshServer;    Read-LxsEnter }
            '5' { Disable-LxsSshServer;    Read-LxsEnter }
            '6' { Set-LxsSshDefaultShell;  Read-LxsEnter }
            '0' { return }
            default { Write-LxsErr 'Invalid protocol. Select 0-6.'; Start-Sleep -Seconds 1 }
        }
    }
}

Show-LxsRemoteMenu
exit 75