# Check if the script is running with administrative privileges
if (-not (New-Object Security.Principal.WindowsPrincipal([Security.Principal.WindowsIdentity]::GetCurrent())).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
    Write-Host "This script requires administrative privileges..."
    Read-Host "Press Enter to end"
    exit
}

# Function to display instructions
function Show-Instructions {
    Clear-Host
    Write-Host "==== OpenSSH Server Configuration Instructions ===="
    Write-Host ""
    Write-Host "This script configures the OpenSSH server on this Windows machine."
    Write-Host "After installing and configuring the server settings, you need to"
    Write-Host "generate and deploy SSH keys for secure access."
    Write-Host ""
    Write-Host "1. First, complete all necessary server configurations through this script."
    Write-Host "   - Install or reinstall the OpenSSH Server."
    Write-Host "   - Configure SSHD settings, the firewall rule, and SSH agent as needed."
    Write-Host ""
    Write-Host "2. After configuring the server, run the 'SSH_Key_Deployment_and_Connection.ps1'"
    Write-Host "   script on your local machine to generate SSH keys and deploy them to this server."
    Write-Host ""
    Write-Host "3. Finally, use the key-based SSH connection to access this server securely."
    Write-Host ""
    Write-Host "Press Enter to return to the main menu..."
    Read-Host
}

# Function to install OpenSSH Server
function Install-OpenSSHServer {
    <#
    .SYNOPSIS
    Installs the OpenSSH Server capability.

    .DESCRIPTION
    Installs the OpenSSH Server capability if it's available.

    .EXAMPLE
    Install-OpenSSHServer
    #>
    Install-OrReinstallOpenSSHServer -InstallOnly
}

# Function to reinstall OpenSSH Server
function Reinstall-OpenSSHServer {
    <#
    .SYNOPSIS
    Reinstalls the OpenSSH Server capability.

    .DESCRIPTION
    Removes and reinstalls the OpenSSH Server capability if it's available.

    .EXAMPLE
    Reinstall-OpenSSHServer
    #>
    Install-OrReinstallOpenSSHServer -Reinstall
}

# Helper function to install or reinstall OpenSSH Server
function Install-OrReinstallOpenSSHServer {
    param (
        [switch]$InstallOnly,
        [switch]$Reinstall
    )

    # Locate the OpenSSH Server capability
    $opensshCapability = Get-WindowsCapability -Online | Where-Object Name -like 'OpenSSH.Server*'

    if ($opensshCapability) {
        try {
            if ($Reinstall -and $opensshCapability.State -eq 'Installed') {
                # Remove OpenSSH Server capability
                Write-Output "Removing OpenSSH Server..."
                Remove-WindowsCapability -Online -Name $opensshCapability.Name -ErrorAction Stop
                Write-Output "OpenSSH Server removed successfully."
            }

            # Attempt to install OpenSSH Server capability
            Write-Output "Installing OpenSSH Server..."
            Add-WindowsCapability -Online -Name $opensshCapability.Name -ErrorAction Stop
            Write-Output "OpenSSH Server installed successfully."
        }
        catch {
            Write-Error "Failed to $(if ($Reinstall) {"reinstall"} else {"install"}) OpenSSH Server: $($_.Exception.Message)"
        }
    }
    else {
        Write-Error "OpenSSH Server capability not found."
    }
}

# Function to start and set OpenSSH service
function Set-OpenSSHService {
    [string]$serviceName = 'sshd'
    if (Get-Service -Name $serviceName -ErrorAction SilentlyContinue) {
        try {
            Start-Service $serviceName
            Set-Service -Name $serviceName -StartupType 'Automatic'
            Write-Output "OpenSSH service started and set to start automatically."
        }
        catch {
            Write-Error "Failed to configure OpenSSH service: $($_.Exception.Message)"
        }
    } else {
        Write-Error "Service '$serviceName' does not exist."
    }
}

# Function to configure sshd_config file settings
function Configure-SSHDConfig {
    param (
        [string]$Port
    )

    $configPath = "C:\ProgramData\ssh\sshd_config"
    $defaultConfigPath = "C:\Windows\System32\OpenSSH\sshd_config_default"

    # Check if the sshd_config file exists
    if (-not (Test-Path -Path $configPath)) {
        # If it doesn't exist, copy the default config file
        try {
            Copy-Item -Path $defaultConfigPath -Destination $configPath -Force
            Write-Output "Copied default sshd_config file to $configPath"
        }
        catch {
            Write-Error "Failed to copy default sshd_config file: $($_.Exception.Message)"
            return
        }
    }

    # Function to modify or add settings in sshd_config file
	function Set-ConfigLine {
		param (
			[string]$FilePath,
			[string]$Setting,
			[string]$Value
		)

		# Read the file content as a single string
		$content = Get-Content -Path $FilePath -Raw

		# Special handling for 'Port' to ensure it's placed before any 'Match' block
		if ($Setting -eq 'Port') {
			$pattern = "^\s*#?\s*Port\s+.*"  # Match lines with or without comment
			
			# Check if the 'Port' line exists, either commented or uncommented
			if ($content -match $pattern) {
				# Replace the line with the uncommented port setting
				$content = $content -replace "^\s*#?\s*Port\s+.*", "Port $Value"
			} else {
				# If the setting doesn't exist, add it before any Match block
				$matchIndex = $content.IndexOf("Match ")
				if ($matchIndex -ne -1) {
					# Insert before the 'Match' block
					$content = $content.Insert($matchIndex, "`r`nPort $Value`r`n")
				}
				else {
					# If no 'Match' block, append to the end
					$content += "`r`nPort $Value"
				}
			}
		}
		else {
			# Generic pattern for other settings
			$pattern = "^\s*#?\s*$Setting\s+.*"
			
			if ($content -match $pattern) {
				# Replace the matched line with the uncommented version of the setting
				$content = $content -replace "^\s*#?\s*$Setting\s+.*", "$Setting $Value"
			}
			else {
				# If the setting isn't found, add it before any Match block or at the end
				$matchIndex = $content.IndexOf("Match ")
				if ($matchIndex -ne -1) {
					$content = $content.Insert($matchIndex, "$Setting $Value`r`n")
				}
				else {
					$content += "`r`n$Setting $Value"
				}
			}
		}

		# Write the updated content back to the file, trimming extra newlines
		Set-Content -Path $FilePath -Value ($content.TrimEnd() + "`r`n")
	}

    # Configure each setting, replacing any matching commented-out or uncommented lines
    Set-ConfigLine -FilePath $configPath -Setting 'Port' -Value $Port
    Set-ConfigLine -FilePath $configPath -Setting 'PasswordAuthentication' -Value 'no'
    Set-ConfigLine -FilePath $configPath -Setting 'PermitRootLogin' -Value 'no'
    Set-ConfigLine -FilePath $configPath -Setting 'MaxAuthTries' -Value '3'
    Set-ConfigLine -FilePath $configPath -Setting 'LoginGraceTime' -Value '20'
    Set-ConfigLine -FilePath $configPath -Setting 'PubkeyAuthentication' -Value 'yes'
    Set-ConfigLine -FilePath $configPath -Setting 'KbdInteractiveAuthentication' -Value 'no'

    # Restart SSH service after configuration
    function Restart-SSHService {
        try {
            Restart-Service sshd
            Write-Output "SSH service restarted successfully after configuration changes."
        }
        catch {
            Write-Error "Failed to restart SSH service: $($_.Exception.Message)"
        }
    }

    Restart-SSHService
}

# Function to create firewall rule with optional port and force creation
function New-OpenSSHFirewallRule {
    param(
        [int]$Port,
        [switch]$Force
    )

    $ruleName = "OpenSSH-Server-Custom-Port"
    $existingRule = Get-NetFirewallRule -Name $ruleName -ErrorAction SilentlyContinue
    if ($existingRule) {
        if ($Force) {
            # Remove existing rule if -Force is specified
            Remove-NetFirewallRule -Name $ruleName
            Write-Output "Existing firewall rule '$ruleName' was removed due to -Force."
        } else {
            # Notify if rule already exists and -Force is not specified
            Write-Output "Firewall rule '$ruleName' already exists. Use -Force to overwrite it."
            return
        }
    }

    # Attempt to create the new firewall rule
    try {
        New-NetFirewallRule -Name $ruleName -DisplayName 'OpenSSH Server (sshd)' -Enabled True -Direction Inbound -Protocol TCP -Action Allow -LocalPort $Port
        Write-Output "Firewall rule '$ruleName' created successfully on port $Port."
    }
    catch {
        Write-Error "Failed to create firewall rule: $($_.Exception.Message)"
    }
}

# Function to start and configure ssh-agent
function Start-SSHAgent {
    $sshAgentService = Get-Service ssh-agent -ErrorAction SilentlyContinue
    if ($sshAgentService) {
        if ($sshAgentService.Status -eq "Running") {
            Write-Output "SSH Agent is already running."
        }
        else {
            try {
                Start-Service ssh-agent
                Set-Service -Name ssh-agent -StartupType 'Automatic'
                Write-Output "SSH Agent started and set to start automatically."
            }
            catch {
                Write-Error "Failed to start SSH Agent: $($_.Exception.Message)"
            }
        }
    } else {
        Write-Error "SSH Agent service does not exist."
    }
}

# Function to verify SSH setup, including service status check and local connection tests
function Verify-SSHSetup {
    # Path to SSHD configuration
    $sshdConfigPath = "C:\ProgramData\ssh\sshd_config"

    # Display SSHD config port setting
    Write-Host "sshd_config port setting:"
    try {
        $portSetting = (Get-Content -Path $sshdConfigPath | Select-String -Pattern '^Port').Line
        Write-Output "Configured SSH Port: $portSetting"
    }
    catch {
        Write-Error "Could not retrieve Port setting from sshd_config: $($_.Exception.Message)"
    }

    # Display SSH firewall rule ports
    Write-Host "SSH server firewall rule port(s):"
    try {
        $firewallPorts = Get-NetFirewallRule -DisplayName 'OpenSSH Server (sshd)' | Get-NetFirewallPortFilter | ForEach-Object { $_.LocalPort }
        Write-Output "Firewall Rule Ports: $firewallPorts"
    }
    catch {
        Write-Error "Could not retrieve firewall rule ports: $($_.Exception.Message)"
    }

    # Check if sshd service is running
    Write-Host "Checking SSH service status..."
    $sshdService = Get-Service -Name 'sshd' -ErrorAction SilentlyContinue
    if ($sshdService -and $sshdService.Status -eq 'Running') {
        Write-Output "SSH service is running. Proceeding with connection tests..."
        
        # Connection test using password authentication
        Write-Host "Testing SSH connection with password authentication..."
        try {
            $passwordAuthResult = & ssh -o PreferredAuthentications=password -o ConnectTimeout=3 -p 1337 localhost "echo 'Password authentication successful'" 2>&1
            if ($passwordAuthResult -match "Password authentication successful") {
                Write-Host "Password authentication was successful."
            } else {
                Write-Warning "Password authentication failed or is disabled."
            }
        } catch {
            Write-Error "An error occurred during password authentication test: $($_.Exception.Message)"
        }

        # Connection test using key authentication
        Write-Host "Testing SSH connection with key authentication..."
        try {
            $keyAuthResult = & ssh -o PreferredAuthentications=publickey -o ConnectTimeout=3 -p 1337 localhost "echo 'Key authentication successful'" 2>&1
            if ($keyAuthResult -match "Key authentication successful") {
                Write-Host "Key authentication was successful."
            } else {
                Write-Warning "Key authentication failed or is disabled."
            }
        } catch {
            Write-Error "An error occurred during key authentication test: $($_.Exception.Message)"
        }
    }
    elseif ($sshdService) {
        Write-Host "SSH service is installed but currently stopped. Skipping connection tests."
    }
    else {
        Write-Error "SSH service ('sshd') is not installed or cannot be found."
    }

    Write-Host "`nSSH setup verification complete."
}

# Function to set up key-based authentication
function Setup-KeyAuthentication {
    $sshDir = "$env:USERPROFILE\.ssh"
    $authKeysFile = "$sshDir\authorized_keys"

    # Ensure .ssh directory exists
    if (-not (Test-Path -Path $sshDir)) {
        New-Item -ItemType Directory -Path $sshDir -Force | Out-Null
        Write-Output ".ssh directory created at $sshDir"
    }

    # Create authorized_keys file with public key if it doesn't exist
    if (-not (Test-Path -Path $authKeysFile)) {
        # Generate keys if they don't exist
        if (-not (Test-Path -Path "$sshDir\id_ed25519.pub")) {
            ssh-keygen -t ed25519 -f "$sshDir\id_ed25519" -N "" | Out-Null
            Write-Output "SSH key pair generated."
        }

        # Copy public key to authorized_keys
        $publicKey = Get-Content -Path "$sshDir\id_ed25519.pub" -Raw
        Set-Content -Path $authKeysFile -Value $publicKey
        Write-Output "Public key added to authorized_keys."
    }
    else {
        Write-Output "authorized_keys file already exists."
    }

    # Set appropriate permissions on authorized_keys
    icacls $authKeysFile /inheritance:r /grant "$env:USERNAME:F" /grant "SYSTEM:F" | Out-Null
    Write-Output "Permissions set on authorized_keys."

    # Restart SSH service to apply changes
    Restart-Service sshd
    Write-Output "SSH service restarted to apply key-based authentication changes."
}

# Function to set up key-based authentication for administrator access
function Setup-AdminKeyAuthentication {
    $sshDir = "C:\ProgramData\ssh"
    $authKeysFile = "$sshDir\administrators_authorized_keys"

    # Ensure ssh directory exists
    if (-not (Test-Path -Path $sshDir)) {
        New-Item -ItemType Directory -Path $sshDir -Force | Out-Null
        Write-Output "SSH directory created at $sshDir"
    }

    # Define paths for the admin key files
    $adminKeyPath = "$env:USERPROFILE\.ssh\admin_id_ed25519"
    $adminPublicKeyPath = "$adminKeyPath.pub"

    # Generate admin key pair if public key doesn't exist
    if (-not (Test-Path -Path $adminPublicKeyPath)) {
        Start-Process ssh-keygen -ArgumentList "-t", "ed25519", "-f", "$adminKeyPath" -NoNewWindow -Wait
        if (Test-Path -Path $adminPublicKeyPath) {
            Write-Output "Admin SSH key pair generated."
        }
        else {
            Write-Output "Error: Failed to generate admin SSH key pair."
            return  # Exit if key generation failed
        }
    }
    else {
        Write-Output "Admin SSH key pair already exists."
    }

    # Copy public key to administrators_authorized_keys
    $publicKey = Get-Content -Path $adminPublicKeyPath -Raw
    Set-Content -Path $authKeysFile -Value $publicKey
    Write-Output "Admin public key added to administrators_authorized_keys."

    # Set appropriate permissions on administrators_authorized_keys
    icacls $authKeysFile /inheritance:r /grant "Administrators:F" /grant "SYSTEM:F" | Out-Null
    Write-Output "Permissions set on administrators_authorized_keys."

    # Restart SSH service to apply changes
    Restart-Service sshd
    Write-Output "SSH service restarted to apply administrator key-based authentication changes."
}

# Display menu
function Show-Menu {
    Clear-Host
    Write-Host "1. Display Instructions"
    Write-Host "2. Install OpenSSH Server"
    Write-Host "3. Reinstall OpenSSH Server"
    Write-Host "4. Configure OpenSSH Service"
    Write-Host "5. Configure SSHD Settings"
    Write-Host "6. Add Firewall Rule"
    Write-Host "7. Start SSH Agent"
    Write-Host "8. Verify SSH Setup"
    Write-Host "9. Setup Key-Based Authentication for User"
    Write-Host "10. Setup Key-Based Authentication for Administrator"
    Write-Host "11. Exit"
}

# Update the main loop to include the new admin key authentication function
do {
    Show-Menu
    $choice = Read-Host "Select an option [1-11]"
    switch ($choice) {
        "1" { Show-Instructions }
        "2" { Install-OpenSSHServer }
        "3" { Reinstall-OpenSSHServer }
        "4" { Set-OpenSSHService }
        "5" { 
            $Port = Read-Host "Enter the SSH Port (default is 1337)"
            if (-not $Port) { $Port = "1337" }
            Configure-SSHDConfig -Port $Port
        }
        "6" { 
            $Port = Read-Host "Enter the SSH Port for Firewall Rule (default is 1337)"
            if (-not $Port) { $Port = "1337" }
            New-OpenSSHFirewallRule -Port $Port
        }
        "7" { Start-SSHAgent }
        "8" { Verify-SSHSetup }
        "9" { Setup-KeyAuthentication }
        "10" { Setup-AdminKeyAuthentication }
        "11" { Write-Host "Exiting..."; break }
        default { Write-Host "Invalid selection. Please choose a valid option." }
    }
    Read-Host "Press Enter to continue..."
} while ($choice -ne "11")