# Vaultwarden SSH Agent Script

param(
    [switch]$Debug
)

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

# Constants
$FolderName = "ssh-agent"

function Test-VaultwardenConfig {
    $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 "Vaultwarden server configured to: $ServerUrl"
}

function Get-BWSession {
    # Check for existing permanent session
    $PersistentSession = [Environment]::GetEnvironmentVariable("BW_SESSION", "User")
    if ($PersistentSession) {
        Write-Debug "Existing Bitwarden session found"
        $env:BW_SESSION = $PersistentSession
        
        # Verify session with sync
        $SyncResult = & bw sync --session $PersistentSession | Out-Null
        if ($LASTEXITCODE -eq 0) {
            Write-Host "Using existing session"
            return $PersistentSession
        }
        Write-Host "Existing session invalid: $SyncResult"
    }

    # Get new session
    Write-Debug "No existing Bitwarden session found"
    
    $null = & bw login --check
    if ($LASTEXITCODE -eq 0) { 
        Write-Debug "User logged in, unlocking vault"
        $Session = & bw unlock --raw
    } else { 
        Write-Debug "User not logged in, starting login"
        $Session = & bw login --raw
    }
    
    if (-not $Session) {
        Write-Host ""  # Add line break after failed attempt
        throw "Failed to get Bitwarden session"
    }
    
    Write-Host "Authentication successful"
    [Environment]::SetEnvironmentVariable("BW_SESSION", $Session, "User")
    $env:BW_SESSION = $Session
    & bw sync --session $Session | Out-Null
    return $Session
}

function Clear-SensitiveData {
    param([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 {
    param(
        [string]$KeyContent,
        [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 {
    param($Session)
    
    Write-Debug "Getting folder: $FolderName"
    $Folders = & bw list folders --search $FolderName --session $Session | ConvertFrom-Json
    $Folder = $Folders | Where-Object { $_.name -eq $FolderName }
    
    if (-not $Folder) {
        throw "'$FolderName' folder not found"
    }
    return $Folder.id
}

function Get-FolderItems {
    param(
        $Session,
        $FolderId
    )
    
    Write-Debug "Getting items from folder: $FolderId"
    $Items = & bw list items --folderid $FolderId --session $Session | ConvertFrom-Json
    Write-Debug "Found $($Items.Count) items"
    return $Items
}

function Get-PrivatePublicKey {
    param(
        $Session,
        $Item
    )

    try {
        # Get and validate public key
        $PublicKey = if ($Item.notes -is [array]) { 
            $Item.notes -join "`n" 
        } else { 
            $Item.notes 
        }
        
        if (-not (Test-SSHKey -KeyContent $PublicKey -KeyType 'Public')) {
            throw "Invalid public key format in notes"
        }
        
        Write-Debug "Valid public key found for: $($Item.name)"
        
        # Get and validate private key
        $Attachment = $Item.attachments | Select-Object -First 1
        if (-not $Attachment) {
            throw "No attachment found"
        }  

        $PrivateKey = & bw get attachment $Attachment.id --itemid $Item.id --raw --session $Session
        $PrivateKey = $PrivateKey -join "`n"
        
        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 {
    param(
        [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 ($Process) { 
            $Process.Dispose() 
        }
    }
}

# Main script execution
try {
    Write-Debug "Checking Vaultwarden configuration"
    Test-VaultwardenConfig
    
    Write-Host "Connect to Vaultwarden"
    $session = Get-BWSession
    Write-Debug "Session = $session"
    
    $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
}