Merge branch 'develop' of https://github.com/homotechsual/NinjaOne in…

…to develop
homotechsual · May 24, 2024 · 0eb5e6c · 0eb5e6c
2 parents 0d0ed08 + a3f6bc2
commit 0eb5e6c
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 15 deletions.
diff --git a/Source/Private/Get-NinjaOneSecrets.ps1 b/Source/Private/Get-NinjaOneSecrets.ps1
@@ -46,7 +46,8 @@ function Get-NinjaOneSecrets {
     foreach ($ConnectionSecret in $Secrets.ConnectionInfo.GetEnumerator()) {
         Write-Verbose ('Processing secret {0} for vault retrieval.' -f $ConnectionSecret.Key)
         $SecretName = $ConnectionSecret.Key
-        $SecretValue = Get-Secret -Name $SecretName -Vault $VaultName -ErrorAction SilentlyContinue
+        $VaultSecretName = $ConnectionSecret.Value
+        $SecretValue = Get-Secret -Name $VaultSecretName -Vault $VaultName -AsPlainText -ErrorAction SilentlyContinue
         if ($null -eq $SecretValue) {
             Write-Verbose ('Secret {0} is null. Skipping.' -f $SecretName)
             continue
@@ -58,7 +59,8 @@ function Get-NinjaOneSecrets {
     foreach ($AuthenticationSecret in $Secrets.AuthenticationInfo.GetEnumerator()) {
         Write-Verbose ('Processing secret {0} for vault retrieval.' -f $AuthenticationSecret.Key)
         $SecretName = $AuthenticationSecret.Key
-        $SecretValue = Get-Secret -Name $SecretName -Vault $VaultName -ErrorAction SilentlyContinue
+        $VaultSecretName = $AuthenticationSecret.Value
+        $SecretValue = Get-Secret -Name $VaultSecretName -Vault $VaultName -AsPlainText -ErrorAction SilentlyContinue
         if ($null -eq $SecretValue) {
             Write-Verbose ('Secret {0} is null. Skipping.' -f $SecretName)
             continue
@@ -125,4 +127,4 @@ function Get-NinjaOneSecrets {
     $Script:NRAPIConnectionInformation.WriteToSecretVault = $true
     $Script:NRAPIConnectionInformation.VaultName = $VaultName
     $Script:NRAPIConnectionInformation.ReadFromSecretVault = $true
-}
+}
diff --git a/Source/Private/Start-OAuthHTTPListener.ps1 b/Source/Private/Start-OAuthHTTPListener.ps1
@@ -4,7 +4,7 @@ function Start-OAuthHTTPListener {
         .SYNOPSIS
             Instantiates and starts a .NET HTTP listener to handle OAuth authorization code responses.
         .DESCRIPTION
-            Utilises the `System.Net.HttpListener` class to create a simple HTTP listener on a user-defined port 
+            Utilises the `System.Net.HttpListener` class to create a simple HTTP listener on a user-defined port
         .EXAMPLE
             PS C:\> New-NinjaOnePATCHRequest -OpenURI 'http://localhost:9090'
         .OUTPUTS
@@ -14,31 +14,55 @@ function Start-OAuthHTTPListener {
     [Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSUseShouldProcessForStateChangingFunctions', '', Justification = 'Private function - no need to support.')]
     param (
         [Parameter(Mandatory)]
-        [System.UriBuilder]$OpenURI
+        [System.UriBuilder]$OpenURI,
+        [Parameter()]
+        [int] $TimeoutSeconds = 15
     )
     Write-Verbose 'Opening browser to authenticate.'
     Write-Verbose "Authentication URL: $($OpenURI.ToString())"
     $HTTP = [System.Net.HttpListener]::new()
     $HTTP.Prefixes.Add("http://localhost:$Port/")
     $HTTP.Start()
     Start-Process $OpenURI.ToString()
+    $Timeout = [System.TimeSpan]::FromSeconds($TimeoutSeconds)
+    $ContextTask = $HTTP.GetContextAsync()
     $Result = @{}
-    while ($HTTP.IsListening) {
-        $Context = $HTTP.GetContext()
+    while ($ContextTask.AsyncWaitHandle.WaitOne($Timeout)) {
+        $Context = $ContextTask.GetAwaiter().GetResult()
+
+        [string]$HTML = '<h1>NinjaOne PowerShell Module</h1><br /><p>An authorisation code has been received. The HTTP listener will stop in 5 seconds.</p><p>Please close this tab / window.</p>'
+        [string]$HTMLError = '<h1>NinjaOne PowerShell Module</h1><br /><p>An error occured. The HTTP listener will stop in 5 seconds.</p><p>Please close this tab / window and try again.</p>'
+
         if ($Context.Request.QueryString -and $Context.Request.QueryString['Code']) {
             $Result.Code = $Context.Request.QueryString['Code']
             Write-Verbose "Authorisation code received: $($Result.Code)"
             if ($null -ne $Result.Code) {
                 $Result.GotAuthorisationCode = $True
             }
-            [string]$HTML = '<h1>NinjaOne PowerShell Module</h1><br /><p>An authorisation code has been received. The HTTP listener will stop in 5 seconds.</p><p>Please close this tab / window.</p>'
-            $Response = [System.Text.Encoding]::UTF8.GetBytes($HTML)
-            $Context.Response.ContentLength64 = $Response.Length
-            $Context.Response.OutputStream.Write($Response, 0, $Response.Length)
-            $Context.Response.OutputStream.Close()
-            Start-Sleep -Seconds 5
-            $HTTP.Stop()
+        } else {
+            $HTML = $HTMLError
         }
+
+        $Response = [System.Text.Encoding]::UTF8.GetBytes($HTML)
+        $Context.Response.ContentLength64 = $Response.Length
+        $Context.Response.OutputStream.Write($Response, 0, $Response.Length)
+        $Context.Response.OutputStream.Close()
+        Start-Sleep -Seconds 5
+        $HTTP.Stop()
+        $HTTP.Dispose()
+        break
+    }
+
+    if ($HTTP.IsListening) {
+        $HTTP.Stop()
+        $HTTP.Dispose()
+    }
+
+    if (!$Result.GotAuthorisationCode) {
+        Remove-Variable -Name 'NRAPIConnectionInformation' -Scope 'Script' -Force
+        Remove-Variable -Name 'NRAPIAuthenticationInformation' -Scope 'Script' -Force
+        throw 'Authorisation failed, please try again.'
     }
+
     return $Result
-}
+}