SecretManagementArgumentCompleter.psm1
Param( [parameter(Mandatory = $false)] [Switch]$InMemory ) $Script:InMemory = $InMemory $Script:_SecretMgmtCompletion = [System.Collections.Generic.Dictionary[String, String[]]]::new() $Script:_AZkvCompletion = [System.Collections.Generic.Dictionary[String, psobject[]]]::new() $ConfigPaths = @{ Directory = "$env:LOCALAPPDATA\Powershell\SecretManagement" } $ConfigPaths.AzConfig = "$($ConfigPaths.Directory)\AzConfig.json" $ConfigPaths.SecretConfig = "$($ConfigPaths.Directory)\SecretConfig.json" Clear-SecretManagementArgumentCompleterCache if (! $InMemory) { if (Test-Path -Path $ConfigPaths.Directory) { if (Test-Path -Path $ConfigPaths.AzConfig) { $config = Get-Content -Path $ConfigPaths.AzConfig -Raw | ConvertFrom-Json ($Config | get-member -MemberType NoteProperty).Name | foreach { $_AZkvCompletion.Add($_, $Config."$_") } } if (Test-Path -Path $ConfigPaths.SecretConfig) { $Config = Get-Content -Path $ConfigPaths.SecretConfig -Raw | ConvertFrom-Json ($Config | get-member -MemberType NoteProperty).Name | foreach { $_SecretMgmtCompletion.Add($_, $Config."$_") } } } } $AZKVVaultNameArgumentCompletion = { param($commandName, $parameterName, $wordToComplete, $commandAst, $fakeBoundParameter) $Context = [Microsoft.Azure.Commands.Common.Authentication.Abstractions.AzureRmProfileProvider]::Instance.Profile.DefaultContext $Items = $Script:_AZkvCompletion.Keys.Where( { $_ -like "*|$($Context.Subscription.Id)" }) if ($Items.count -eq 0) { $Vnames = (Get-AzKeyVault).VaultName if ($Vnames.count -gt 0 ) { $Vnames | % { $Script:_AZkvCompletion."$_|$($Context.Subscription.Id)" = $null } } } if ($Items.count -gt 0) { return $Items | % { [System.Management.Automation.CompletionResult]::new(($_.Split('|')[0])) } } } $AZKVSecretNameArgumentCompletion = { param($commandName, $parameterName, $wordToComplete, $commandAst, $fakeBoundParameter) $Context = [Microsoft.Azure.Commands.Common.Authentication.Abstractions.AzureRmProfileProvider]::Instance.Profile.DefaultContext if ($fakeBoundParameter.ContainsKey('VaultName')) { if ($fakeBoundParameter.VaultName -is [hashtable]) { $fakeBoundParameter.VaultName = $fakeBoundParameter.VaultName.VaultName } $VaultName = $fakeBoundParameter.item('VaultName') $Key = "$VaultName|$($Context.Subscription.Id)" if (!$Script:_AZkvCompletion.ContainsKey($Key) -and $_AZkvCompletion.Keys.Where( { $_ -like "*|$($Context.Subscription.Id)" }).Count -eq 0) { Set-AZKeyVaultNameCache } if (!$Script:_AZkvCompletion.ContainsKey($Key)) { return $null } $MyVault = $Script:_AZkvCompletion.Item($Key) if ($MyVault -eq $null) { $Secrets = Get-AzKeyVaultSecret -VaultName ($VaultName) $Script:_AZkvCompletion.Item($Key) = $Secrets.Name Save-SecretCache -Az return $Script:_AZkvCompletion.Item($Key) | % { [System.Management.Automation.CompletionResult]::new("'$_'", $_ , [System.Management.Automation.CompletionResultType]::Text, $_) } } else { return $MyVault | % { [System.Management.Automation.CompletionResult]::new("'$_'", $_, [System.Management.Automation.CompletionResultType]::Text, $_) } } } } $RegisterSecretVaultArgCompletion = { param($commandName, $parameterName, $wordToComplete, $commandAst, $fakeBoundParameter) $GetVaultParamStatement = { try { $Context = [Microsoft.Azure.Commands.Common.Authentication.Abstractions.AzureRmProfileProvider]::Instance.Profile.DefaultContext $SubscriptionId = $Context.Subscription.Id } catch { } $Hashtable = '@{ AZKVaultName = ''''; SubscriptionId = ''{0}'' }'.Replace('{0}', $SubscriptionId) return $Hashtable } switch ($parameterName) { 'ModuleName' { $Output = [System.Collections.Generic.List[PSObject]]::new() $Items = (Get-ChildItem -Path $env:PSModulePath.Split(';') -DIrectory).FullName Foreach ($i in $Items) { $Leaf = Split-Path $i -Leaf $Versions = Get-ChildItem -Path "$i/*" -Directory if ($Versions.count -gt 0 ) { $ExtensionPath = (Join-Path -Path $Versions[-1].FullName -ChildPath "$Leaf.Extension") if (Test-Path -Path $ExtensionPath ) { $Leaf } } } } 'VaultParameters' { if ($fakeBoundParameter.ContainsKey('ModuleName') -and $fakeBoundParameter.Item('ModuleName') -eq 'Az.Keyvault') { return [System.Management.Automation.CompletionResult]::new(($GetVaultParamStatement.Invoke())) } } } } $SecretNameArgCompletion = { param($commandName, $parameterName, $wordToComplete, $commandAst, $fakeBoundParameter) $DefaultVault = 'BuiltInLocalVault' $UseCache = $false $VaultNameFilter = @{} if ( $fakeBoundParameter.ContainsKey('Vault')) { if ($fakeBoundParameter.Vault -is [hashtable]) { $fakeBoundParameter.Vault = $fakeBoundParameter.Vault.Vault } if ($fakeBoundParameter.Vault -eq $DefaultVault) { return (Get-SecretInfo -Vault $DefaultVault).Name | % { [System.Management.Automation.CompletionResult]::new("'$_'", $_ , [System.Management.Automation.CompletionResultType]::Text, $_) } } $Vault = $fakeBoundParameter.Vault $VaultNameFilter = @{Vault = $Vault } $UseCache = $Script:_SecretMgmtCompletion.ContainsKey($Vault) if (-not $UseCache) { try { $Secrets = Get-SecretInfo @VaultNameFilter -ErrorAction Stop $Script:_SecretMgmtCompletion.Add($Vault, $Secrets.Name) Save-SecretCache } catch {} } return $Script:_SecretMgmtCompletion.Item($Vault) | % { [System.Management.Automation.CompletionResult]::new("'$_'", $_ , [System.Management.Automation.CompletionResultType]::Text, $_) } } $Vaults = (Get-SecretVault).Name $NeedSave = $false $Output = [System.Collections.Generic.List[Object]]::new() Foreach ($V in $Vaults) { if ($V -eq $DefaultVault) { $Secrets = (Get-SecretInfo -Vault $DefaultVault) | Select VaultName, Name $Output.AddRange($Secrets) Continue } if ($Script:_SecretMgmtCompletion.ContainsKey($V)) { $Secrets = $_SecretMgmtCompletion.Item($V) | % { [PSCustomObject]@{ VaultName = $V Name = $_ } } $Output.AddRange($Secrets) } else { try { $Secrets = Get-SecretInfo -Vault $V -ErrorAction Stop $Script:_SecretMgmtCompletion.Add($V, $Secrets.Name) $Output.AddRange(($Secrets | Select VaultName, Name)) $NeedSave = $true } catch {} } } if ($NeedSave) { Save-SecretCache } $Statement = "" $PName = $commandAst.CommandElements.parameterName if ($null -ne $PName -and $PName.Contains('Name')) { $Statement = "'{1}' -Vault '{0}'" } else { $Statement = "-Vault '{0}' -Name '{1}'" } return $Output | % { $o = $Statement -f $_.VaultName, $_.Name [System.Management.Automation.CompletionResult]::new($o, $_.Name , [System.Management.Automation.CompletionResultType]::Text, $_.Name) } } Function Clear-SecretManagementArgumentCompleterCache() { $Script:_SecretMgmtCompletion = [System.Collections.Generic.Dictionary[String, String[]]]::new() $Script:_AZkvCompletion = [System.Collections.Generic.Dictionary[String, psobject[]]]::new() } function Import-SecretManagementArgumentCompleter { [cmdletBinding()] Param([Switch]$InMemory, [switch]$Force) $Script:InMemory = $InMemory Import-Module -Name SecretManagementArgumentCompleter -ArgumentList $InMemory -Force:$Force } function Save-SecretCache ([Switch]$Az) { if ($Script:InMemory) { return } if (! (Test-Path -Path $ConfigPaths.Directory)) { New-Item $ConfigPaths.Directory -ItemType Directory -Force } if ($az) { $_AZkvCompletion | ConvertTo-Json -Depth 10 | Set-Content -Path $ConfigPaths.AzConfig try { (Get-Item -Path $ConfigPaths.AzConfig).Encrypt() } catch {} } else { $_SecretMgmtCompletion | ConvertTo-Json -Depth 10 | Set-Content -Path $ConfigPaths.SecretConfig try { (Get-Item -Path $ConfigPaths.SecretConfig).Encrypt() } catch {} } } $SecretVaultNameArgCompletion = { Get-SecretVault | Select -ExpandProperty Name | foreach-object { [System.Management.Automation.CompletionResult]::new($_) } } function Set-AZKeyVaultNameCache { $Context = [Microsoft.Azure.Commands.Common.Authentication.Abstractions.AzureRmProfileProvider]::Instance.Profile.DefaultContext $Vaults = Get-AzKeyVault Foreach ($v in $Vaults) { $Key = "$($v.VaultName)|$($Context.Subscription.Id)" if (!$_AZkvCompletion.ContainsKey($Key)) { $_AZkvCompletion.Add($Key, $null) } } Save-SecretCache -Az } Register-ArgumentCompleter -ParameterName Name -ScriptBlock $AZKVSecretNameArgumentCompletion -CommandName Backup-AzKeyVaultSecret, Get-AzKeyVaultSecret, Remove-AzKeyVaultSecret, Set-AzKeyVaultSecret , Update-AzKeyVaultSecret Register-ArgumentCompleter -ParameterName VaultName -ScriptBlock $AZKVVaultNameArgumentCompletion -CommandName Backup-AzKeyVaultSecret, Get-AzKeyVault, Get-AzKeyVaultSecret, Remove-AzKeyVault, Remove-AzKeyVaultSecret, Restore-AzKeyVaultSecret, Set-AzKeyVaultSecret, Update-AzKeyVaultSecret Register-ArgumentCompleter -CommandName Register-SecretVault -ParameterName ModuleName -ScriptBlock $RegisterSecretVaultArgCompletion Register-ArgumentCompleter -CommandName Register-SecretVault -ParameterName VaultParameters -ScriptBlock $RegisterSecretVaultArgCompletion #endregion Register-ArgumentCompleter -CommandName Get-Secret, Get-SecretInfo, Set-Secret, Test-SecretVault, Remove-Secret -ParameterName Vault -ScriptBlock $SecretVaultNameArgCompletion Register-ArgumentCompleter -CommandName Unregister-SecretVault, Test-SecretVault, Get-SecretVault -ParameterName Name -ScriptBlock $SecretVaultNameArgCompletion Register-ArgumentCompleter -CommandName Get-Secret, Get-SecretInfo, Remove-Secret, Set-Secret -ParameterName Name -ScriptBlock $SecretNameArgCompletion |