From 52b113a52a7102d643fd17f411638a148883484a Mon Sep 17 00:00:00 2001 From: MalwareMechanic Date: Tue, 26 Oct 2021 21:07:30 -0400 Subject: [PATCH] Update error handling, add context menu functions Change `$ErrorActionPreference` from `Continue` to 'Stop' to properly handle errors in installers. Noticed when an error occurred due to a failed `Join-Path` in `VM-Install-From-Zip` the installer continued when it should have stopped. Add new `-zipFolder_64` parameter to `VM-Install-From-Zip` to handle cases where different architecture URLs may have different ZIP folder names, closes #36. Add functions `VM-Add-To-Right-Click-Menu` and `VM-Remove-From-Right-Click-Menu` (helps #33). --- .../common.vm/tools/VM.Common/VM.Common.psm1 | 130 +++++++++++++++--- 1 file changed, 114 insertions(+), 16 deletions(-) diff --git a/packages/common.vm/tools/VM.Common/VM.Common.psm1 b/packages/common.vm/tools/VM.Common/VM.Common.psm1 index 297557295..734440514 100755 --- a/packages/common.vm/tools/VM.Common/VM.Common.psm1 +++ b/packages/common.vm/tools/VM.Common/VM.Common.psm1 @@ -1,4 +1,5 @@ -$ErrorActionPreference = 'Continue' +# Setting this to "Stop". Functions should properly handle errors or throw to calling function. +$ErrorActionPreference = 'Stop' # ################################################################################################ # @@ -187,10 +188,11 @@ function VM-Remove-PreviousZipPackage { [OPTIONAL] A wildcard expression for a file type containing a list of files to delete. #> param( + [Parameter(Mandatory=$true, Position=0)] [string] $packagePath, + [Parameter(Mandatory=$false)] [string] $expression=$null ) - if ($expression) { $previousZipFiles = Get-ChildItem -Path (Join-Path $packagePath $expression) } else { @@ -199,10 +201,12 @@ function VM-Remove-PreviousZipPackage { foreach ($zipFileName in $previousZipFiles) { if ((Test-Path -Path $zipFileName)) { - $zipContents = Get-Content $zipFileName -Force - foreach ($fileInZip in $zipContents) { - if ($fileInZip -ne $null -and $fileInZip.Trim() -ne '' -and (Test-Path "$fileInZip")) { - Remove-Item -Path "$fileInZip" -ErrorAction SilentlyContinue -Recurse -Force + $zipContents = @(Get-Content $zipFileName -Force) + if ($zipContents) { + foreach ($fileInZip in $zipContents) { + if (($null -ne $fileInZip) -AND ($fileInZip.Trim() -ne '') -AND (Test-Path $fileInZip)) { + Remove-Item -Path $fileInZip -Recurse -Force -ea 0 + } } } } @@ -403,7 +407,7 @@ function VM-Install-From-Zip { [Parameter(Mandatory=$false)] [bool] $consoleApp=$false, [Parameter(Mandatory=$false)] - [string] $zipFolder # subfolder in zip with the app files + [bool] $innerFolder=$false # subfolder in zip with the app files ) try { $toolDir = Join-Path ${Env:RAW_TOOLS_DIR} $toolName @@ -412,21 +416,42 @@ function VM-Install-From-Zip { # Remove files from previous zips for upgrade VM-Remove-PreviousZipPackage ${Env:chocolateyPackageFolder} + # Snapshot all folders in $toolDir + $oldDirList = @() + if (Test-Path $toolDir) { + $oldDirList = @(Get-ChildItem $toolDir | Where-Object {$_.PSIsContainer}) + } + # Download and unzip $packageArgs = @{ - packageName = ${Env:ChocolateyPackageName} - unzipLocation = $toolDir - url = $zipUrl - checksum = $zipSha256 - url64bit = $zipUrl_64 - checksum64 = $zipSha256_64 - checksumType = 'sha256' + packageName = ${Env:ChocolateyPackageName} + unzipLocation = $toolDir + url = $zipUrl + checksum = $zipSha256 + checksumType = 'sha256' + url64bit = $zipUrl_64 + checksum64 = $zipSha256_64 } Install-ChocolateyZipPackage @packageArgs VM-Assert-Path $toolDir - if ($zipFolder) { - $toolDir = Join-Path $toolDir $zipFolder -Resolve + # Diff and find new folders in $toolDir + $newDirList = @(Get-ChildItem $toolDir | Where-Object {$_.PSIsContainer}) + $diffDirs = Compare-Object -ReferenceObject $oldDirList -DifferenceObject $newDirList -PassThru + + # If $innerFolder is set to $true, after unzipping only a single folder should be new. + # GitHub ZIP files typically unzip to a single folder that contains the tools. + if ($innerFolder) { + # First time install, use the single resulting folder name from Install-ChocolateyZipPackage. + if ($diffDirs.Count -eq 1) { + # Save the "new tool directory" to assist with upgrading. + $newToolDir = Join-Path $toolDir $diffDirs[0].Name -Resolve + Set-Content (Join-Path ${Env:chocolateyPackageFolder} "innerFolder.txt") $newToolDir + $toolDir = $newToolDir + } else { + # On upgrade there may be no new directory, in this case retrieve previous "new tool directory" from saved file. + $toolDir = Get-Content (Join-Path ${Env:chocolateyPackageFolder} "innerFolder.txt") + } } $executablePath = Join-Path $toolDir "$toolName.exe" -Resolve @@ -481,4 +506,77 @@ function VM-Write-Log-Exception { $position_msg = $error_record.InvocationInfo.PositionMessage VM-Write-Log "ERROR" "[ERR] $msg`r`n$position_msg" throw $error_record +} + +function VM-Add-To-Right-Click-Menu { + Param + ( + [Parameter(Mandatory=$true, Position=0)] + [String] $menuKey, # name of registry key + [Parameter(Mandatory=$true, Position=1)] + [string] $menuLabel, # value displayed in right-click menu + [Parameter(Mandatory=$true, Position=2)] + [string] $command, + [Parameter(Mandatory=$true, Position=3)] + [ValidateSet("file", "directory")] + [string] $type + ) + try { + # Determine if file or directory should show item in right-click menu + if ($type -eq "file") { + $key = "*" + } else { + $key = "directory" + } + + # Check and map "HKCR" to correct drive + if (-NOT (Test-Path -path 'HKCR:')) { + New-PSDrive -Name HKCR -PSProvider Registry -Root HKEY_CLASSES_ROOT | Out-Null + } + + # Add right-click menu display name + if (-NOT (Test-Path -LiteralPath "HKCR:\$key\shell\$menuKey")) { + New-Item -Path "HKCR:\$key\shell\$menuKey" | Out-Null + } + Set-ItemProperty -LiteralPath "HKCR:\$key\shell\$menuKey" -Name '(Default)' -Value "$menuLabel" -Type String + + # Add command to run when executed from right-click menu + if(-NOT (Test-Path -LiteralPath "HKCR:\$key\shell\$menuKey\command")) { + New-Item -Path "HKCR:\$key\shell\$menuKey\command" | Out-Null + } + Set-ItemProperty -LiteralPath "HKCR:\$key\shell\$menuKey\command" -Name '(Default)' -Value $command -Type String + } catch { + FE-Write-Log "ERROR" "Failed to add $menuKey to right-click menu" + } +} + +function VM-Remove-From-Right-Click-Menu { + Param + ( + [Parameter(Mandatory=$true, Position=0)] + [String] $menuKey, # name of registry key + [Parameter(Mandatory=$true, Position=1)] + [ValidateSet("file", "directory")] + [string] $type + ) + try { + # Determine if file or directory should show item in right-click menu + if ($type -eq "file") { + $key = "*" + } else { + $key = "directory" + } + + # Check and map "HKCR" to correct drive + if (-NOT (Test-Path -path 'HKCR:')) { + New-PSDrive -Name HKCR -PSProvider Registry -Root HKEY_CLASSES_ROOT | Out-Null + } + + # Remove right-click menu settings from registry + if (Test-Path -LiteralPath "HKCR:\$key\shell\$menuKey") { + Remove-Item -LiteralPath "HKCR:\$key\shell\$menuKey" -Recurse + } + } catch { + FE-Write-Log "ERROR" "Failed to remove $menuKey from right-click menu" + } } \ No newline at end of file