<#
.SYNOPSIS
Interactive SSH config manager with cross-shell compatibility and enhanced validation

.DESCRIPTION
This script modifies ~/.ssh/config with safety checks, flexible execution, and robust input validation
Supports both CMD and PowerShell execution environments
#>

$ErrorActionPreference = 'Stop'
$appVersion = "2.5"

#region Helper Functions
function ConvertTo-Base32Hex {
  <#
    .SYNOPSIS
    Converts DateTime to RFC 4648 Base32hex format
    #>
  param([DateTime]$DateTime)
    
  $unixTime = [Math]::Floor(($DateTime - (Get-Date "1970-01-01")).TotalSeconds)
  $base32Chars = "0123456789ABCDEFGHIJKLMNOPQRSTUV"
  $result = ""
    
  if ($unixTime -eq 0) { return "0" }
    
  while ($unixTime -gt 0) {
    $result = $base32Chars[$unixTime % 32] + $result
    $unixTime = [Math]::Floor($unixTime / 32)
  }
    
  # Ensure minimum 8 characters for consistency
  return $result.PadLeft(8, '0')
}

function Get-NormalizedPath {
  <#
    .SYNOPSIS
    Cross-platform path normalization with validation
    #>
  param(
    [string]$Path,
    [string]$BasePath = $env:USERPROFILE
  )
    
  if ([string]::IsNullOrWhiteSpace($Path)) {
    throw "Path cannot be empty"
  }
    
  # Remove quotes and trim
  $cleanPath = $Path.Trim().Trim('"', "'", '`')
  # Replace forward slashes with backslashes for Windows
  $cleanPath = $cleanPath -replace '/', '\'
  # Collapse multiple consecutive separators
  $cleanPath = $cleanPath -replace '\\+', '\'
    
  # Handle tilde expansion
  if ($cleanPath.StartsWith('~')) {
    $cleanPath = $cleanPath -replace '^~', $env:USERPROFILE
  }
    
  # Convert relative to absolute path
  if (-not [System.IO.Path]::IsPathRooted($cleanPath)) {
    $cleanPath = Join-Path $BasePath $cleanPath
  }
    
  # Validate path characters (Windows-specific validation)
  $invalidChars = [System.IO.Path]::GetInvalidPathChars() -join ''
  if ($cleanPath -match "[$([regex]::Escape($invalidChars))]") {
    throw "Path contains invalid characters: $cleanPath"
  }
    
  return [System.IO.Path]::GetFullPath($cleanPath)
}

function Test-ValidHostname {
  <#
    .SYNOPSIS
    Validates hostname or IP address format
    #>
  param([string]$Hostname)
    
  if ([string]::IsNullOrWhiteSpace($Hostname)) {
    return $false
  }
    
  # Check for valid IPv4 address
  $ipv4Pattern = '^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$'
  if ($Hostname -match $ipv4Pattern) {
    return $true
  }
    
  # Check for valid IPv6 address (simplified)
  $ipv6Pattern = '^(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}$|^::1$|^::$'
  if ($Hostname -match $ipv6Pattern) {
    return $true
  }
    
  # Check for valid FQDN/hostname (RFC 1123 compliant)
  $hostnamePattern = '^([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])(\.([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]{0,61}[a-zA-Z0-9]))*$'
  return $Hostname -match $hostnamePattern
}

function Test-ValidUsername {
  <#
    .SYNOPSIS
    Validates SSH username format
    #>
  param([string]$Username)
    
  if ([string]::IsNullOrWhiteSpace($Username)) {
    return $false
  }
    
  # Unix username validation: alphanumeric, underscore, hyphen, dot (no spaces)
  $usernamePattern = '^[a-zA-Z0-9._-]+$'
  return ($Username -match $usernamePattern) -and ($Username.Length -le 32)
}

function Test-ValidHostAlias {
  <#
    .SYNOPSIS
    Validates SSH host alias format
    #>
  param([string]$HostAlias)
    
  if ([string]::IsNullOrWhiteSpace($HostAlias)) {
    return $false
  }
    
  # SSH config Host pattern validation (no spaces, limited special chars)
  $aliasPattern = '^[a-zA-Z0-9._-]+$'
  return ($HostAlias -match $aliasPattern) -and ($HostAlias.Length -le 64)
}

function Test-ValidPort {
  <#
    .SYNOPSIS
    Validates port number range
    #>
  param([string]$Port)
    
  $portNumber = 0
  return ([int]::TryParse($Port, [ref]$portNumber) -and $portNumber -ge 1 -and $portNumber -le 65535)
}

function Get-SafeUserInput {
  <#
    .SYNOPSIS
    Gets user input with validation and retry logic
    #>
  param(
    [string]$Prompt,
    [string]$ValidationMessage,
    [scriptblock]$Validator,
    [string]$DefaultValue = "",
    [int]$MaxAttempts = 3
  )
    
  $attempts = 0
  do {
    if ($attempts -gt 0) {
      Write-Host $ValidationMessage -ForegroundColor Yellow
    }
        
    $value = if ([string]::IsNullOrWhiteSpace($DefaultValue)) {
      Read-Host $Prompt
    }
    else {
      $userInput = Read-Host "$Prompt (default: $DefaultValue)"
      if ([string]::IsNullOrWhiteSpace($userInput)) { $DefaultValue } else { $userInput }
    }
        
    $attempts++
    if ($Validator.Invoke($value)) {
      return $value
    }
        
  } while ($attempts -lt $MaxAttempts)
    
  throw "Maximum validation attempts exceeded for: $Prompt"
}

function Get-SanitizedHostAlias {
  <#
    .SYNOPSIS
    Sanitizes host alias for safe filename usage
    #>
  param([string]$HostAlias)
    
  return $HostAlias -replace '[^\w.-]', '_'
}

function Test-SSHCommand {
  <#
    .SYNOPSIS
    Tests if SSH command exists and is available
    #>
  param([string]$Command)
    
  try {
    $null = Get-Command $Command -ErrorAction Stop
    return $true
  }
  catch {
    return $false
  }
}

function Set-SecureFilePermissions {
  <#
    .SYNOPSIS
    Sets secure permissions on SSH key files (Windows-specific)
    #>
  param([string]$FilePath)
    
  try {
    # Remove inheritance and grant read access only to current user
    $acl = Get-Acl -Path $FilePath
    $acl.SetAccessRuleProtection($true, $false)
    $accessRule = New-Object System.Security.AccessControl.FileSystemAccessRule($env:USERNAME, "Read", "Allow")
    $acl.SetAccessRule($accessRule)
    Set-Acl -Path $FilePath -AclObject $acl
    Write-Host "Secure permissions set on: $FilePath" -ForegroundColor Green
  }
  catch {
    Write-Warning "Failed to set secure permissions on $FilePath : $($_.Exception.Message)"
  }
}

function Copy-SSHKeyToRemote {
  <#
    .SYNOPSIS
    Deploys SSH public key to remote server with fallback methods
    #>
  param(
    [string]$KeyPath,
    [string]$UserName,
    [string]$HostName,
    [string]$Port
  )
    
  $publicKeyPath = "$KeyPath.pub"
  if (-not (Test-Path -Path $publicKeyPath)) {
    throw "Public key file not found: $publicKeyPath"
  }
    
  try {
    if (Test-SSHCommand 'ssh-copy-id') {
      # Method 1: Use ssh-copy-id (preferred)
      Write-Host "Deploying key using ssh-copy-id..." -ForegroundColor Yellow
      ssh-copy-id -i $publicKeyPath -p $Port "$UserName@$HostName"
    }
    else {
      # Method 2: Fallback using ssh and shell commands
      Write-Host "Deploying key using SSH fallback method..." -ForegroundColor Yellow
      $publicKey = Get-Content -Path $publicKeyPath -Raw
      $cleanKey = $publicKey.Trim()
      $remoteCommand = "mkdir -p ~/.ssh && chmod 700 ~/.ssh && echo '$cleanKey' >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && sort ~/.ssh/authorized_keys | uniq > ~/.ssh/authorized_keys.tmp && mv ~/.ssh/authorized_keys.tmp ~/.ssh/authorized_keys"
      ssh -p $Port "$UserName@$HostName" $remoteCommand
    }
        
    Write-Host "SSH key deployed successfully" -ForegroundColor Green
    return $true
  }
  catch {
    Write-Warning "Key deployment failed: $($_.Exception.Message)"
    Write-Host "You may need to manually copy the public key content to ~/.ssh/authorized_keys on the remote server" -ForegroundColor Yellow
    Write-Host "Public key content:" -ForegroundColor Cyan
    Get-Content -Path $publicKeyPath
    return $false
  }
}

function Test-HostExists {
  <#
    .SYNOPSIS
    Checks if host alias already exists in SSH config
    #>
  param(
    [string]$ConfigContent,
    [string]$HostAlias
  )
    
  $pattern = "(?m)^Host\s+$([regex]::Escape($HostAlias))\s*$"
  return $ConfigContent -match $pattern
}

function Backup-SSHConfig {
  <#
    .SYNOPSIS
    Creates timestamped backup of SSH config file
    #>
  param([string]$ConfigPath)
    
  # Proper logical operator syntax
  if ((Test-Path -Path $ConfigPath) -and ((Get-Item -Path $ConfigPath).Length -gt 0)) {
    $timestamp = Get-Date -Format 'yyyyMMddHHmmss'
    $backupPath = "$ConfigPath.backup.$timestamp"
    Copy-Item -Path $ConfigPath -Destination $backupPath
    Write-Host "Backup created: $backupPath" -ForegroundColor Cyan
    return $backupPath
  }
  return $null
}
#endregion

# BEGIN Main Script Execution
Write-Host "==========================================" -ForegroundColor Green
Write-Host "SSH Configuration Manager $appVersion" -ForegroundColor Green
Write-Host "Enhanced with validation and cross-platform support" -ForegroundColor Green
Write-Host "==========================================" -ForegroundColor Green

# BEGIN SSH Directory Setup
$sshDir = Join-Path $env:USERPROFILE ".ssh"
if (-not (Test-Path -Path $sshDir)) {
  Write-Host "Creating SSH directory: $sshDir" -ForegroundColor Yellow
  New-Item -Path $sshDir -ItemType Directory -Force | Out-Null
}

$configPath = Join-Path $sshDir "config"
$configContent = if (Test-Path -Path $configPath) {
  Get-Content -Path $configPath -Raw
}
else {
  Write-Host "Creating new SSH config file: $configPath" -ForegroundColor Yellow
  New-Item -Path $configPath -ItemType File -Force | Out-Null
  ""
}
# END SSH Directory Setup

# BEGIN Validated Input Collection
Write-Host "`nCollecting SSH configuration information..." -ForegroundColor Yellow

# Get and validate host alias
$hostAlias = Get-SafeUserInput -Prompt "Enter host alias (alphanumeric, dots, hyphens only)" -ValidationMessage "Invalid host alias format. Use only letters, numbers, dots, and hyphens." -Validator { param($value) Test-ValidHostAlias $value }

# Check for duplicate host alias
if (Test-HostExists -ConfigContent $configContent -HostAlias $hostAlias) {
  $overwriteChoice = $host.ui.PromptForChoice(
    "Duplicate Host Detected",
    "Host '$hostAlias' already exists in configuration. Do you want to overwrite it?",
    @(
      [System.Management.Automation.Host.ChoiceDescription]::new("&No", "Do not overwrite existing configuration"),
      [System.Management.Automation.Host.ChoiceDescription]::new("&Yes", "Overwrite existing configuration")
    ),
    0
  )
  if ($overwriteChoice -eq 0) {
    Write-Host "Operation cancelled. Exiting..." -ForegroundColor Yellow
    exit 0
  }
}

# Get and validate hostname
$hostName = Get-SafeUserInput -Prompt "Enter target hostname or IP address" -ValidationMessage "Invalid hostname or IP address format." -Validator { param($value) Test-ValidHostname $value }

# Get and validate username
$userName = Get-SafeUserInput -Prompt "Enter SSH username" -ValidationMessage "Invalid username format. Use alphanumeric characters, dots, hyphens, and underscores only." -Validator { param($value) Test-ValidUsername $value }

# Get and validate port
$port = Get-SafeUserInput -Prompt "Enter port number" -ValidationMessage "Invalid port number. Must be between 1 and 65535." -Validator { param($value) Test-ValidPort $value } -DefaultValue "22"
# END Validated Input Collection

# BEGIN SSH Key Authentication Setup
$useKeyAuth = $host.ui.PromptForChoice(
  "SSH Authentication Method",
  "Do you want to use SSH public key authentication?",
  @(
    [System.Management.Automation.Host.ChoiceDescription]::new("&Yes", "Use SSH key authentication (recommended)"),
    [System.Management.Automation.Host.ChoiceDescription]::new("&No", "Use password authentication")
  ),
  0 # Default to Yes
)

$keyPath = $null
if ($useKeyAuth -eq 0) {
  Write-Host "`nConfiguring SSH key authentication..." -ForegroundColor Yellow
    
  # Get key path or generate new key
  $keyPathInput = Read-Host "Enter path to existing private key (leave empty to generate new key)"
    
  if ([string]::IsNullOrWhiteSpace($keyPathInput)) {
    # Generate new SSH key
    if (-not (Test-SSHCommand 'ssh-keygen')) {
      Write-Error "ssh-keygen command not found. Please install OpenSSH client."
      exit 1
    }
        
    # Generate secure filename with Base32hex timestamp
    $sanitizedAlias = Get-SanitizedHostAlias -HostAlias $hostAlias
    $timestamp = ConvertTo-Base32Hex -DateTime (Get-Date)
    $defaultKeyName = "id_ed25519_${sanitizedAlias}_${timestamp}"
    $keyName = Get-SafeUserInput -Prompt "Enter key filename" -ValidationMessage "Invalid filename format." -Validator { param($value) -not [string]::IsNullOrWhiteSpace($value) } -DefaultValue $defaultKeyName
        
    $keyPath = Get-NormalizedPath -Path $keyName -BasePath $sshDir
        
    Write-Host "Generating new SSH key pair..." -ForegroundColor Green
    try {
      # Generate ED25519 key with comment
      ssh-keygen -t ed25519 -f $keyPath -N '""' -C "$userName@$hostName-$(Get-Date -Format 'yyyy-MM-dd')" -q
      Write-Host "SSH key pair generated successfully at: $keyPath" -ForegroundColor Green
            
      # Set secure permissions on private key
      Set-SecureFilePermissions -FilePath $keyPath
    }
    catch {
      Write-Error "Failed to generate SSH key: $($_.Exception.Message)"
      exit 1
    }
  }
  else {
    # Use existing key with enhanced path validation
    try {
      $keyPath = Get-NormalizedPath -Path $keyPathInput -BasePath $sshDir
      if (-not (Test-Path -Path $keyPath)) {
        Write-Error "Private key file not found: $keyPath"
        exit 1
      }
      Write-Host "Using existing private key: $keyPath" -ForegroundColor Green
    }
    catch {
      Write-Error "Invalid key path: $($_.Exception.Message)"
      exit 1
    }
  }
    
  # Deploy SSH key to remote server
  if (Test-Path -Path "$keyPath.pub") {
    $deployKey = $host.ui.PromptForChoice(
      "SSH Key Deployment",
      "Do you want to copy the public key to the remote server?",
      @(
        [System.Management.Automation.Host.ChoiceDescription]::new("&Yes", "Deploy public key to server"),
        [System.Management.Automation.Host.ChoiceDescription]::new("&No", "Skip key deployment")
      ),
      0 # Default to Yes
    )
        
    if ($deployKey -eq 0) {
      Write-Host "`nAttempting to deploy SSH key..." -ForegroundColor Yellow
      Copy-SSHKeyToRemote -KeyPath $keyPath -UserName $userName -HostName $hostName -Port $port
    }
  }
  else {
    Write-Warning "Public key file not found. You may need to generate it manually."
  }
}
# END SSH Key Authentication Setup

# BEGIN Additional Configuration Options
$enableKeepAlive = $host.ui.PromptForChoice(
  "SSH Keep-Alive Settings",
  "Enable SSH zombie session avoidance (ServerAliveInterval and ServerAliveCountMax)?",
  @(
    [System.Management.Automation.Host.ChoiceDescription]::new("&Yes", "Enable keep-alive settings (recommended)"),
    [System.Management.Automation.Host.ChoiceDescription]::new("&No", "Skip keep-alive settings")
  ),
  0 # Default to Yes
)

$configureAdvanced = $host.ui.PromptForChoice(
  "Advanced Configuration",
  "Do you want to configure advanced SSH options?",
  @(
    [System.Management.Automation.Host.ChoiceDescription]::new("&No", "Use basic configuration only"),
    [System.Management.Automation.Host.ChoiceDescription]::new("&Yes", "Configure advanced options")
  ),
  0 # Default to No
)

$advancedSettings = @{}
if ($configureAdvanced -eq 1) {
  Write-Host "`nConfiguring advanced SSH options..." -ForegroundColor Yellow
    
  $advancedOptions = @{
    "AddressFamily"                   = "Address family to use (inet, inet6, any)"
    "ConnectTimeout"                  = "Connection timeout in seconds (default: system default)"
    "ConnectionAttempts"              = "Number of connection attempts (default: 1)"
    "ProxyCommand"                    = "Command to use for proxy connection"
    "ProxyJump"                       = "Jump host (user@host:port)"
    "PreferredAuthentications"        = "Preferred authentication methods (e.g., publickey,password)"
    "ChallengeResponseAuthentication" = "Enable challenge-response authentication (yes/no)"
    "ForwardAgent"                    = "Forward SSH agent (yes/no)"
    "StrictHostKeyChecking"           = "Strict host key checking (yes/no/ask)"
    "UserKnownHostsFile"              = "Path to known hosts file"
    "CheckHostIP"                     = "Check host IP in known_hosts (yes/no)"
    "BatchMode"                       = "Enable batch mode - no prompts (yes/no)"
    "ControlMaster"                   = "Enable connection multiplexing (yes/no/ask/auto/autoask)"
  }
    
  foreach ($option in $advancedOptions.GetEnumerator()) {
    $value = Read-Host "$($option.Key): $($option.Value) (leave empty to skip)"
    if (-not [string]::IsNullOrWhiteSpace($value)) {
      $advancedSettings[$option.Key] = $value.Trim()
    }
  }
}
# END Additional Configuration Options

# BEGIN SSH Configuration File Generation
Write-Host "`nGenerating SSH configuration..." -ForegroundColor Yellow

# Create backup
$backupPath = Backup-SSHConfig -ConfigPath $configPath

# Build configuration block
$configLines = @()
$configLines += "# =========================================="
$configLines += "# SSH CONFIG GENERATED BY PowerShell SSH Config Manager $AppVersion"
$configLines += "# Generated: $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss zzz')"
$configLines += "# Host Alias: $hostAlias"
$configLines += "# Target: $userName@${hostName}:$port"
$configLines += "# Key Authentication: $(if ($keyPath) { 'Enabled' } else { 'Disabled' })"
$configLines += "# =========================================="
$configLines += "# BEGIN Auto-generated SSH config for $hostAlias"
$configLines += "Host $hostAlias"
$configLines += "    # Connection settings - Generated by SSH Config Manager"
$configLines += "    HostName $hostName"
$configLines += "    User $userName"
$configLines += "    Port $port"

if ($keyPath) {
  $configLines += "    # Authentication settings - Generated by SSH Config Manager"
  $configLines += "    IdentityFile `"$keyPath`""
  $configLines += "    IdentitiesOnly yes"
  $configLines += "    PubkeyAuthentication yes"
}

if ($enableKeepAlive -eq 0) {
  $configLines += "    # Keep-alive settings - Generated by SSH Config Manager"
  $configLines += "    ServerAliveInterval 60"
  $configLines += "    ServerAliveCountMax 3"
}

# Add advanced settings with validation
if ($advancedSettings.Count -gt 0) {
  $configLines += "    # Advanced settings - Generated by SSH Config Manager"
  foreach ($setting in $advancedSettings.GetEnumerator()) {
    $configLines += "    $($setting.Key) $($setting.Value)"
  }
}

$configLines += "# END Auto-generated SSH config for $hostAlias"
$configLines += "# =========================================="
$configLines += ""

$newConfigBlock = $configLines -join "`n"
# END SSH Configuration File Generation

# BEGIN Configuration File Update
try {
  # Remove existing configuration for this host (program-generated only)
  $removalPattern = "(?s)# ==========================================.*?# SSH CONFIG GENERATED BY PowerShell SSH Config Manager.*?# END Auto-generated SSH config for $([regex]::Escape($hostAlias)).*?# ==========================================\s*"
  $cleanedContent = $configContent -replace $removalPattern, ""
    
  # Clean up excessive whitespace
  $cleanedContent = $cleanedContent -replace "(`r?`n){3,}", "`n`n"
  $cleanedContent = $cleanedContent.Trim()
    
  # Combine new config with existing content
  $finalContent = if ([string]::IsNullOrWhiteSpace($cleanedContent)) {
    $newConfigBlock.TrimEnd()
  }
  else {
    $newConfigBlock + "`n" + "# ==========================================" + "`n" + "# EXISTING SSH CONFIG (Pre-existing content)" + "`n" + "# ==========================================" + "`n" + $cleanedContent
  }
    
  # Write configuration with UTF-8 encoding (no BOM)
  $utf8NoBom = New-Object System.Text.UTF8Encoding($false)
  [System.IO.File]::WriteAllText($configPath, $finalContent, $utf8NoBom)
    
  Write-Host "`n==========================================" -ForegroundColor Green
  Write-Host "SSH configuration updated successfully!" -ForegroundColor Green
  Write-Host "Configuration file: $configPath" -ForegroundColor Cyan
  if ($backupPath) {
    Write-Host "Backup created: $backupPath" -ForegroundColor Cyan
  }
  Write-Host "You can now connect using: ssh $hostAlias" -ForegroundColor Cyan
  Write-Host "==========================================" -ForegroundColor Green
}
catch {
  Write-Error "Failed to update SSH configuration: $($_.Exception.Message)"
  # Restore from backup if available
  if ($backupPath -and (Test-Path -Path $backupPath)) {
    try {
      Copy-Item -Path $backupPath -Destination $configPath -Force
      Write-Host "Configuration restored from backup" -ForegroundColor Yellow
    }
    catch {
      Write-Error "Failed to restore backup: $($_.Exception.Message)"
    }
  }
  exit 1
}
# END Configuration File Update

Write-Host "`nSSH Configuration Manager completed successfully!" -ForegroundColor Green
# END Main Script Execution