From f567555d855522342da072e273a34b8cd42c9d68 Mon Sep 17 00:00:00 2001
From: Richard Downer <richard@cloudsoft.io>
Date: Mon, 27 Sep 2021 16:33:09 +0100
Subject: [PATCH] Bug fixes and updates, mainly around optional parameter
 handling

---
 EC2Access.psm1 | 35 +++++++++++++++++++++++------------
 README.md      |  6 +++---
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/EC2Access.psm1 b/EC2Access.psm1
index 7a74c1f..79fe5c3 100644
--- a/EC2Access.psm1
+++ b/EC2Access.psm1
@@ -71,18 +71,21 @@ function Get-EC2Password {
     [Parameter(Position=2)] [string]$PrivateKeyFile
   )
 
-  $ErrorActionPreference = "Stop"
-
   # Verify the private key files exists
-  if($null -eq $PrivateKeyFile) {
+  if(!($PrivateKeyFile)) {
     $PrivateKeyFile = $HOME + '\.ssh\id_rsa'
+    Write-Verbose "No private key file given - defaulting to $($PrivateKeyFile)"
   }
   if(-not (Test-Path $PrivateKeyFile)) {
     Write-Error "$($PrivateKeyFile) does not exist. Do you need to use -PrivateKeyFile argument?"
   }
 
   Write-Verbose "Requesting password data from AWS"
-  $cipherText = (Get-EC2PasswordData -Region $Region -InstanceId $InstanceId)
+  if ($Region) {
+    $cipherText = Get-EC2PasswordData -InstanceId $InstanceId -Region $Region
+  } else {
+    $cipherText = Get-EC2PasswordData -InstanceId $InstanceId
+  }
 
   Write-Verbose "Decrypting password"
   $password = Convert-RSAEncryptedCipherTextToClearText -PemFile $PrivateKeyFile -CipherText $cipherText
@@ -120,12 +123,16 @@ function Start-DirectEC2RemoteDesktop {
     [Parameter(Position=2)] [string]$PrivateKeyFile
   )
   
-  $password = Get-EC2Password -Instance $InstanceId -Region $Region -PrivateKeyFile $PrivateKeyFile
+  $password = Get-EC2Password -InstanceId $InstanceId -Region $Region -PrivateKeyFile $PrivateKeyFile
   $Credential = New-Object PSCredential "Administrator",$password
 
-  $instance = (Get-EC2Instance -Region $Region -InstanceId $InstanceId).Instances[0]
-  $HostName = $instance.PublicIpAddress
-  Write-Verbose "Instance IP address is $HostName"
+  if($Region) {
+    $response = Get-EC2Instance -InstanceId $InstanceId -Region $Region
+  } else {
+    $response = Get-EC2Instance -InstanceId $InstanceId
+  }
+  $HostName = $response.Instances[0].PublicIpAddress
+  Write-Verbose "Instance address is $HostName"
 
   if ($PSCmdlet.ShouldProcess($InstanceId,'Start remote desktop session')) {
     Start-RemoteDesktop -HostName $HostName -Credential $Credential
@@ -160,12 +167,16 @@ function Start-EC2RemoteDesktopViaSessionManager {
     [Parameter(Position=2)] [string]$PrivateKeyFile
   )
   
-  $password = Get-EC2Password -Instance $InstanceId -Region $Region -PrivateKeyFile $PrivateKeyFile
+  $password = Get-EC2Password -InstanceId $InstanceId -Region $Region -PrivateKeyFile $PrivateKeyFile
   $Credential = New-Object PSCredential "Administrator",$password
 
   $LocalPort = 33389
   $PortForwardParams = @{ portNumber=(,"3389"); localPortNumber=(,$LocalPort.ToString()) }
-  $session = Start-SSMSession -Target $InstanceId -Region $Region -DocumentName AWS-StartPortForwardingSession -Parameters $PortForwardParams
+  if($Region) {
+    $session = Start-SSMSession -Target $InstanceId -DocumentName AWS-StartPortForwardingSession -Parameters $PortForwardParams -Region $Region
+  } else {
+    $session = Start-SSMSession -Target $InstanceId -DocumentName AWS-StartPortForwardingSession -Parameters $PortForwardParams
+  }
 
   # We now need to emulate awscli - it invokes session-manager-plugin with the new session information.
   # AWS Tools for PowerShell don't do this. Also some of the objects seem to look a bit different, and the
@@ -235,8 +246,8 @@ function Start-EC2RemoteDesktopViaSessionManager {
 
 function Start-RemoteDesktop {
   [CmdletBinding(SupportsShouldProcess)] param(
-    [Parameter(Mandatory=$true, Position=0)] [String] [string]$HostName,
-    [Parameter(Mandatory=$true, Position=1)] [PSCredential] [string]$Credential,
+    [Parameter(Mandatory=$true, Position=0)] [String] $HostName,
+    [Parameter(Mandatory=$true, Position=1)] [PSCredential] $Credential,
     [Parameter()] [Int32] [string]$Port
   )
 
diff --git a/README.md b/README.md
index c00c19f..179e2d8 100644
--- a/README.md
+++ b/README.md
@@ -87,9 +87,9 @@ to get a private key.
    same name with a `.pub` extension. Then, using the *Import* button on the EC2 *Key pairs* page to import
    the public key.
 
-By default, the functions in this module will assume that your private key is in the ".ssh\id_rsa" file in
+By default, the functions in this module will assume that your private key is in the `.ssh\id_rsa` file in
 your home directory, which will be the normal situation in method 2 above. If you have used method 1, or have
-your key in any other location, simply pass a "-PrivateKeyFile" parameter to the functions with the path to
+your key in any other location, simply pass a `-PrivateKeyFile` parameter to the functions with the path to
 your private key file.
 
 
@@ -117,7 +117,7 @@ Start-EC2RemoteDesktopViaSessionManager -InstanceId i-12345678abcd `
 ```
 
 If your EC2 instance is reachable on its public IP address, then instead of
-"Start-EC2RemoteDesktopViaSessionManager", you can invoke "Start-DirectEC2RemoteDesktop". This function takes
+`Start-EC2RemoteDesktopViaSessionManager`, you can invoke `Start-DirectEC2RemoteDesktop`. This function takes
 exactly the same parameters but uses the public IP address instead of Session Manager port forwarding.
 
 ```powershell