### Global Varaibles ###

$global:windowsTopMost = $false
$global:returnCode = 0
$global:privateKeyPath = $args[0]
$global:passphraseEnvName = $args[1]
$global:sshJob = $null
$global:sshPassEnvVarName = $null

### Main Functions ###

function OpenSSHKeyFileDialog {
  param(
    $InitialDirectory
  )
  [System.Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms") | Out-Null
  $openFileDialog = New-Object System.Windows.Forms.OpenFileDialog
  #$openFileDialog = (New-Object Microsoft.Win32.OpenFileDialog)
  if (Test-Path $InitialDirectory -PathType container) {
    $openFileDialog.InitialDirectory = $InitialDirectory
  }
  $openFileDialog.Title = "Select Private Key File"
  $openFileDialog.Filter = "All Files (*.*)| *.*"
  $openFileDialog.Multiselect = $false
  $openFileDialog.AutoUpgradeEnabled = $true

  $openFileDialogWindow = $null
  if ($global:windowsTopMost) {
    $openFileDialogWindow = (New-Object System.Windows.Forms.Form -Property @{TopMost = $true })
  }

  if ($openFileDialog.ShowDialog($openFileDialogWindow)) {
    return $openFileDialog.FileName
  }
  return $null
}

function SSHAddKey {
  param(
    $PrivateKeyPath,
    $Passphrase
  )
  $return=$false
  #Invoke-Expression 'ssh-agent'
  Start-Process -WindowStyle Hidden ssh-agent
  $global:sshJob = Start-Job -ScriptBlock {
    $encodedPassphrase = [Convert]::ToBase64String([System.Text.Encoding]::Unicode.GetBytes($Using:Passphrase))
    $global:sshPassEnvVarName = "SSH_TEMP_PASS_$((( 1..16 | foreach-object {'{0:x}' -f (Get-Random -Maximum 16)}) -join '').toUpper())"
    Set-Item -Path "env:$global:sshPassEnvVarName" -Value $encodedPassphrase

    $Env:DISPLAY="0"
    $Env:SSH_ASKPASS="$Env:COMSPEC /c powershell.exe -NoLogo -NonInteractive -F `"$Using:PSCommandPath`" `"$Using:PrivateKeyPath`" `"$global:sshPassEnvVarName`""

    $global:LASTEXITCODE = 0
    Invoke-Expression 'ssh-add "$Using:PrivateKeyPath" 2> $null > $null'

    Write-Output $global:LASTEXITCODE
    if (Test-Path -Path "env:$global:sshPassEnvVarName"){
      Remove-Item -Path "env:$global:sshPassEnvVarName"
    }
    $global:sshPassEnvVarName = $null
  }
  $global:sshJob | Wait-Job -Timeout 3 # Unfortunately, 1 second is too low and might returns false-positive
  if ($global:sshJob.JobStateInfo.State -Eq "Running"){
    Remove-Job -Force -InstanceId $global:sshJob.InstanceId > $null
    return $false
  } else {
    Receive-Job $global:sshJob -OutVariable sshJobOutput > $null
    if ([String]$sshJobOutput -Eq "0"){
      $return=$true
    } else {
      return $false
    }
  }
  return $return
}

### Main Process ###

if (([string]::IsNullOrEmpty($global:privateKeyPath)) -Or -Not (Test-Path $global:privateKeyPath -PathType leaf)){
  if (Test-Path "$env:USERPROFILE\.ssh" -PathType container){
    $global:privateKeyPath = (OpenSSHKeyFileDialog -InitialDirectory "$env:USERPROFILE\.ssh")
  } else {
    $global:privateKeyPath = ({ OpenSSHKeyFileDialog })
  }
  if (([string]::IsNullOrEmpty($global:privateKeyPath)) -Or -Not (Test-Path $global:privateKeyPath -PathType leaf)){
    Exit 1
  }
}

if (-Not ([string]::IsNullOrEmpty($global:passphraseEnvName))){
  $decodedPassphrase = [System.Text.Encoding]::Unicode.GetString([System.Convert]::FromBase64String((Get-Item -Path "env:$global:passphraseEnvName").Value))
  Write-Output $decodedPassphrase
  Exit $global:returnCode
}

Add-Type -AssemblyName System.Windows.Forms
[System.Windows.Forms.Application]::EnableVisualStyles()

$frmPassphrase                   = New-Object system.Windows.Forms.Form
$frmPassphrase.ClientSize        = New-Object System.Drawing.Point(318,120)
$frmPassphrase.text              = "OpenSSH Agent Passphrase"
$frmPassphrase.TopMost           = $global:windowsTopMost
$frmPassphrase.MaximizeBox       = $false
$frmPassphrase.MinimizeBox       = $false
$frmPassphrase.FormBorderStyle   = 'Fixed3D'
$frmPassphrase.StartPosition     = "CenterScreen"
$frmPassphrase.KeyPreview        = $true
$frmPassphrase.ShowIcon          = $false

$frmPassphrase.Add_KeyDown({if($_.KeyCode -eq "Enter"){ok($txtPassphrase.Text)}})
$frmPassphrase.Add_KeyDown({if($_.KeyCode -eq "Escape"){cancel}})

$lblPassphrase                   = New-Object system.Windows.Forms.Label
$lblPassphrase.text              = "Passphare for `"" + (Split-Path $global:privateKeyPath -leaf) + "`":"
$lblPassphrase.AutoSize          = $true
$lblPassphrase.width             = 25
$lblPassphrase.height            = 10
$lblPassphrase.location          = New-Object System.Drawing.Point(10,20)

$txtPassphrase                    = New-Object system.Windows.Forms.MaskedTextBox
$txtPassphrase.multiline          = $false
$txtPassphrase.width              = 298
$txtPassphrase.height             = 20
$txtPassphrase.location           = New-Object System.Drawing.Point(10,42)
$txtPassphrase.PasswordChar       = '•'

$btnOk                           = New-Object system.Windows.Forms.Button
$btnOk.text                      = "OK"
$btnOk.width                     = 74
$btnOk.height                    = 30
$btnOk.location                  = New-Object System.Drawing.Point(157,77)

$btnCancel                       = New-Object system.Windows.Forms.Button
$btnCancel.text                  = "Cancel"
$btnCancel.width                 = 73
$btnCancel.height                = 30
$btnCancel.location              = New-Object System.Drawing.Point(235,77)

$frmPassphrase.controls.AddRange(@($btnOk,$btnCancel,$lblPassphrase,$txtPassphrase))

$btnOk.Add_Click({ ok($txtPassphrase.Text) })
$btnCancel.Add_Click({ cancel })

function ok ($text) {
  $btnOk.enabled = $false
  $btnCancel.enabled = $false
  $txtPassphrase.enabled = $false
  $sshPassphraseSuccess = (SSHAddKey -PrivateKeyPath $global:privateKeyPath -Passphrase $text)
  if ( $sshPassphraseSuccess -Eq $true){
    # Good!
    $frmPassphrase.Close()
  }
  $btnOk.enabled = $true
  $btnCancel.enabled = $true
  $txtPassphrase.enabled = $true
  $txtPassphrase.text = ""
  $frmPassphrase.Activate()
  $txtPassphrase.Focus()
}

function cancel() {
  $global:returnCode = 1
  $frmPassphrase.Close()
}

$frmPassphrase.Add_Shown({$frmPassphrase.Activate(); $txtPassphrase.Focus()})
[void]$frmPassphrase.ShowDialog()

Exit $global:returnCode