<#
.SYNOPSIS
    Adds SSH keys from Vaultwarden to the SSH agent.
.DESCRIPTION
    Retrieves SSH keys stored in Vaultwarden and adds them to the running SSH agent.
.PARAMETER Debug
    Enables debug output.
.PARAMETER FolderName
    The name of the Vaultwarden folder containing SSH keys. Defaults to "ssh-agent".
.PARAMETER ClearSession
    Removes the stored Vaultwarden session from the Windows Credential Manager.
.EXAMPLE
    .\Vaultwarden_ssh-agent.ps1
    Imports all SSH keys from Vaultwarden to the SSH agent.
.EXAMPLE
    .\Vaultwarden_ssh-agent.ps1 -FolderName "my-ssh-keys"
    Imports SSH keys from a custom folder name.
.EXAMPLE
    .\Vaultwarden_ssh-agent.ps1 -ClearSession
    Removes the stored Vaultwarden session from the Windows Credential Manager.
#>

param(
    [switch]$Debug,
    [string]$FolderName = "ssh-agent",
    [switch]$ClearSession
)

# Set debug preference
if ($Debug) {
    $DebugPreference = 'Continue'
} else {
    $DebugPreference = 'SilentlyContinue'
}

# Add CredManager type definition at script start
Add-Type -TypeDefinition @"
    using System;
    using System.Runtime.InteropServices;
    
    public class CredManager {
        [DllImport("advapi32.dll", SetLastError=true, CharSet=CharSet.Unicode)]
        public static extern bool CredRead(string target, int type, int reservedFlag, out IntPtr credentialPtr);
        
        [DllImport("advapi32.dll", SetLastError=true)]
        public static extern void CredFree(IntPtr buffer);
        
        [StructLayout(LayoutKind.Sequential, CharSet=CharSet.Unicode)]
        public struct CREDENTIAL {
            public int Flags;
            public int Type;
            public string TargetName;
            public string Comment;
            public System.Runtime.InteropServices.ComTypes.FILETIME LastWritten;
            public int CredentialBlobSize;
            public IntPtr CredentialBlob;
            public int Persist;
            public int AttributeCount;
            public IntPtr Attributes;
            public string TargetAlias;
            public string UserName;
        }
    }
"@

function Clear-BWSession {
    [CmdletBinding()]
    param()

    Write-Debug "Removing stored session from Windows Credential Manager"
    & cmdkey /delete:"Vaultwarden_Session" | Out-Null
}

function Test-Prerequisites {
    # Check if ssh-add exists
    if (-not (Get-Command "ssh-add" -ErrorAction SilentlyContinue)) {
        throw "ssh-add command not found. Please ensure OpenSSH is installed."
    }

    # Check if SSH agent is running
    $process = Get-Process ssh-agent -ErrorAction SilentlyContinue
    if (-not $process) {
        throw "SSH agent is not running. Please start the SSH agent service."
    }

    # Check if bw CLI is available and version
    try {
        $bwVersion = & bw --version
        if ($bwVersion -match '(\d{4})\.(\d{1,2})' -and 
            ("$($matches[1])$($matches[2].PadLeft(2,'0'))" -lt "202412")) {
            throw "Bitwarden CLI version $bwVersion is not supported. Please upgrade to version 2024.12.0 or above to use SSH key features."
        }
    } catch {
        throw "Bitwarden CLI not found or version check failed. Please install Bitwarden CLI 2024.12.0 or above."
    }
}

function Test-VaultwardenConfig {
    [CmdletBinding()]
    param()

    $ConfigPath = "$env:APPDATA\Bitwarden CLI\data.json"
    if (Test-Path $ConfigPath) {
        $Config = Get-Content $ConfigPath -Raw | ConvertFrom-Json
        if ($Config.global_environment_environment.region -eq "Self-hosted" -and 
            $Config.global_environment_environment.urls.base) {
            Write-Debug "Vaultwarden server is configured: $($Config.global_environment_environment.urls.base)"
            return
        }
    }
    
    Write-Host "Vaultwarden server is not configured. Please enter the server URL:"
    $ServerUrl = Read-Host
    & bw config server $ServerUrl
    Write-Host ""  # Add line break after failed attempt
    Write-Host "Vaultwarden server configured to: $ServerUrl"
}

function Get-BWSession {
    [CmdletBinding()]
    param()

    Write-Debug "Retrieving session from Windows Credential Manager"
    
    $CredPtr = [IntPtr]::Zero
    $Result = [CredManager]::CredRead("Vaultwarden_Session", 1, 0, [ref]$CredPtr)
    
    if ($Result) {
        try {
            $Cred = [System.Runtime.InteropServices.Marshal]::PtrToStructure($CredPtr, [Type][CredManager+CREDENTIAL])
            
            if (-not $Cred.CredentialBlobSize) {
                Write-Debug "No stored session found"
                return $null
            }
            
            $SecretBytes = New-Object byte[] $Cred.CredentialBlobSize
            [System.Runtime.InteropServices.Marshal]::Copy($Cred.CredentialBlob, $SecretBytes, 0, $Cred.CredentialBlobSize)
            $Session = [System.Text.Encoding]::Unicode.GetString($SecretBytes)
            
            # Check session status
            Write-Debug "Stored session found"
            # bw status is always locked when using --session
            # https://github.com/bitwarden/clients/issues/9254
            [Environment]::SetEnvironmentVariable("BW_SESSION", $Session, [System.EnvironmentVariableTarget]::Process)
            $Status = & bw status | ConvertFrom-Json
            [Environment]::SetEnvironmentVariable("BW_SESSION", $null, [System.EnvironmentVariableTarget]::Process)
            if ($LASTEXITCODE -eq 0) {
                if ($Status.status -eq "unlocked") {
                    Write-Debug "Stored session is valid"
                    Write-Host "Using stored session"
                    return $Session
                }
            }
            
            Write-Debug "Stored session is invalid, removing"
            & cmdkey /delete:"Vaultwarden_Session" | Out-Null
        }
        finally {
            [CredManager]::CredFree($CredPtr)
            if ($Session) {
                Clear-SensitiveData $Session
            }
        }
    }

    Write-Host "Getting new session"
    
    # Check current login status
    Write-Debug "Checking login status"
    $Status = & bw status | ConvertFrom-Json
    
    if ($Status.status -eq "unauthenticated") {
        Write-Debug "Not logged in, attempting API key login"
        & bw login --apikey | Out-Null
        if ($LASTEXITCODE -ne 0) {
            throw "Login failed"
        }
        # Get status again after login to retrieve user email
        Write-Debug "Getting user info after login"
        $Status = & bw status | ConvertFrom-Json
    }
    else {
        Write-Debug "Already logged in as $($Status.userEmail)"
    }

    Write-Debug "Unlocking vault"
    $NewSession = & bw unlock --raw
    if (-not $NewSession) {
        throw "Failed to unlock vault"
    }

    # Store session in Windows Credential Manager
    Write-Debug "Storing session in Windows Credential Manager for user: $($Status.userEmail)"
    $Result = & cmdkey /generic:"Vaultwarden_Session" /user:"$($Status.userEmail)" /pass:"$NewSession" | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to store session: $Result"
    }

    Write-Host "New session established"
    return $NewSession
}

function Clear-SensitiveData {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [System.Object]$Variable
    )
    
    if ($Variable -is [SecureString]) {
        $Variable.Dispose()
    }
    elseif ($Variable -is [string]) {
        [Runtime.InteropServices.Marshal]::ZeroFreeBSTR(
            [Runtime.InteropServices.Marshal]::StringToBSTR($Variable)
        )
    }
}

function Test-SSHKey {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [string]$KeyContent,
        
        [Parameter(Mandatory=$true)]
        [ValidateSet('Public', 'Private')]
        [string]$KeyType
    )
    
    $Patterns = @{
        Public = '^(ssh-rsa|ssh-dss|ssh-ed25519|ecdsa-sha2-nistp\d+)\s+[A-Za-z0-9+/=]+\s*.*$'
        Private = '(?sm)^-----BEGIN\s+(OPENSSH|RSA|EC|DSA)\s+PRIVATE\s+KEY-----.*-----END\s+(OPENSSH|RSA|EC|DSA)\s+PRIVATE\s+KEY-----$'
    }
    
    return $KeyContent -match $Patterns[$KeyType]
}

function Get-FolderId {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [string]$Session
    )
    
    Write-Debug "Getting folder: $FolderName"
    try {
        [Environment]::SetEnvironmentVariable("BW_SESSION", $Session, [System.EnvironmentVariableTarget]::Process)
        $Folders = & bw list folders --search $FolderName | ConvertFrom-Json
    }
    finally {
        [Environment]::SetEnvironmentVariable("BW_SESSION", $null, [System.EnvironmentVariableTarget]::Process)
    }
    $Folder = $Folders | Where-Object { $_.name -eq $FolderName }
    
    if (-not $Folder) {
        throw "'$FolderName' folder not found"
    }
    return $Folder.id
}

function Get-FolderItems {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [string]$Session,
        [Parameter(Mandatory=$true)]
        [string]$FolderId
    )
    
    Write-Debug "Getting items from folder: $FolderId"
    try {
        [Environment]::SetEnvironmentVariable("BW_SESSION", $Session, [System.EnvironmentVariableTarget]::Process)
        $Items = & bw list items --folderid $FolderId | ConvertFrom-Json
    }
    finally {
        [Environment]::SetEnvironmentVariable("BW_SESSION", $null, [System.EnvironmentVariableTarget]::Process)
    }
    # Add filter for SSH key type (type=5)
    $SshKeyItems = $Items | Where-Object { $_.type -eq 5 }
    Write-Debug "Found $($SshKeyItems.Count) SSH key items"
    return $SshKeyItems
}

function Get-PrivatePublicKey {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [string]$Session,
        [Parameter(Mandatory=$true)]
        [PSObject]$Item
    )

    try {
        if ($Item.type -ne 5) {
            throw "Item is not an SSH key type"
        }

        # Get public key from the sshKey property
        $PublicKey = $Item.sshKey.publicKey
        
        if (-not (Test-SSHKey -KeyContent $PublicKey -KeyType 'Public')) {
            throw "Invalid public key format"
        }
        
        Write-Debug "Valid public key found for: $($Item.name)"
        
        # Get private key from the sshKey property
        $PrivateKey = $Item.sshKey.privateKey
        
        if (-not (Test-SSHKey -KeyContent $PrivateKey -KeyType 'Private')) {
            throw "Invalid private key format"
        }
        
        Write-Debug "Valid private key found for: $($Item.name)"
        
        return @{
            PublicKey = $PublicKey
            PrivateKey = ConvertTo-SecureString $PrivateKey -AsPlainText -Force
            Name = $Item.name
        }
    }
    finally {
        if ($PrivateKey) {
            Clear-SensitiveData $PrivateKey
        }
    }
}

function Add-PrivateKeyToSSHAgent {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [ValidateNotNull()]
        [PSCustomObject]$SSHKey
    )
    
    try {
        Write-Debug "Adding key: $($SSHKey.Name)"
        
        $PrivateKeyPlain = [Runtime.InteropServices.Marshal]::PtrToStringAuto(
            [Runtime.InteropServices.Marshal]::SecureStringToBSTR($SSHKey.PrivateKey)
        )
        
        $ProcessInfo = New-Object System.Diagnostics.ProcessStartInfo
        $ProcessInfo.FileName = "ssh-add"
        $ProcessInfo.Arguments = "-"
        $ProcessInfo.RedirectStandardInput = $true
        $ProcessInfo.RedirectStandardOutput = $true
        $ProcessInfo.RedirectStandardError = $true
        $ProcessInfo.UseShellExecute = $false
        
        $Process = New-Object System.Diagnostics.Process
        $Process.StartInfo = $ProcessInfo
        
        $Process.Start() | Out-Null
        $Process.StandardInput.WriteLine($PrivateKeyPlain)
        $Process.StandardInput.Close()
        $Process.WaitForExit()

        if ($Process.ExitCode -eq 0) {
            Write-Debug "Key added successfully: $($SSHKey.Name)"
            return $true
        } else {
            Write-Warning "Failed to add key: $($SSHKey.Name)"
            return $false
        }
    }
    catch {
        Write-Error "Exception adding key $($SSHKey.Name): $_"
        return $false
    }
    finally {
        if ($PrivateKeyPlain) { 
            Clear-SensitiveData $PrivateKeyPlain 
        }
        if ($SSHKey.PrivateKey) {
            Clear-SensitiveData $SSHKey.PrivateKey
        }
        if ($Process) { 
            $Process.Dispose() 
        }
    }
}

# Main script execution
try {
    if ($ClearSession) {
        Clear-BWSession
        Write-Host "Stored session has been cleared"
        return
    }

    Test-Prerequisites
    
    Write-Debug "Checking Vaultwarden configuration"
    Test-VaultwardenConfig

    # Get session
    $Session = Get-BWSession

    # Sync vault
    Write-Host "Syncing vault"
    try {
        [Environment]::SetEnvironmentVariable("BW_SESSION", $Session, [System.EnvironmentVariableTarget]::Process)
        & bw sync | Out-Null
        $SyncExitCode = $LASTEXITCODE
    }
    finally {
        [Environment]::SetEnvironmentVariable("BW_SESSION", $null, [System.EnvironmentVariableTarget]::Process)
    }
    if ($SyncExitCode -ne 0) {
        throw "Failed to sync vault"
    }

    $FolderId = Get-FolderId -Session $Session
    $Items = Get-FolderItems -Session $Session -FolderId $FolderId
    
    $Results = @{
        Success = @()
        Failed = @()
    }

    Write-Host "Adding SSH keys to agent..."
    foreach ($Item in $Items) {
        try {
            $SSHKey = Get-PrivatePublicKey -Session $Session -Item $Item
            if (Add-PrivateKeyToSSHAgent -SSHKey $SSHKey) {
                $Results.Success += $Item.name
            } else {
                $Results.Failed += $Item.name
            }
        }
        catch {
            Write-Warning "Failed to process $($Item.name): $_"
            $Results.Failed += $Item.name
        }
    }

    # Report results
    Write-Host "`nResults:"
    Write-Host "Successfully added keys: $($Results.Success.Count)"
    $Results.Success | ForEach-Object { Write-Host "  - $_" }
    
    if ($Results.Failed.Count -gt 0) {
        Write-Host "`nFailed to add keys: $($Results.Failed.Count)"
        $Results.Failed | ForEach-Object { Write-Host "  - $_" }
    }
}
catch {
    Write-Error "Critical error: $_"
    exit 1
}
finally {
    if ($Session) {
        Clear-SensitiveData $Session
    }
}