From fdf1255bb7557e0d26b66a4ac5efe3b8f7f973b9 Mon Sep 17 00:00:00 2001
From: Elliot Chernofsky <echernofsky@google.com>
Date: Thu, 12 Dec 2024 20:27:27 -0500
Subject: [PATCH] Add signature verify arg to VM-Install-From-Zip

---
 packages/common.vm/common.vm.nuspec           |  2 +-
 .../common.vm/tools/vm.common/vm.common.psm1  | 53 +++++++++++++++----
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/packages/common.vm/common.vm.nuspec b/packages/common.vm/common.vm.nuspec
index 79076e2e..d76e3da3 100755
--- a/packages/common.vm/common.vm.nuspec
+++ b/packages/common.vm/common.vm.nuspec
@@ -2,7 +2,7 @@
 <package xmlns="http://schemas.microsoft.com/packaging/2015/06/nuspec.xsd">
   <metadata>
     <id>common.vm</id>
-    <version>0.0.0.20241106</version>
+    <version>0.0.0.20241212</version>
     <description>Common libraries for VM-packages</description>
     <authors>Mandiant</authors>
   </metadata>
diff --git a/packages/common.vm/tools/vm.common/vm.common.psm1 b/packages/common.vm/tools/vm.common/vm.common.psm1
index 6a54eb2a..252bf200 100755
--- a/packages/common.vm/tools/vm.common/vm.common.psm1
+++ b/packages/common.vm/tools/vm.common/vm.common.psm1
@@ -397,6 +397,8 @@ function VM-Install-From-Zip {
         [Parameter(Mandatory=$false)]
         [string] $executableName, # Executable name, needed if different from "$toolName.exe"
         [Parameter(Mandatory=$false)]
+        [bool] $verifySignature=$false,
+        [Parameter(Mandatory=$false)]
         [switch] $withoutBinFile, # Tool should not be installed as a bin file
         # Examples:
         # $powershellCommand = "Get-Content README.md"
@@ -410,17 +412,31 @@ function VM-Install-From-Zip {
         # Remove files from previous zips for upgrade
         VM-Remove-PreviousZipPackage ${Env:chocolateyPackageFolder}
 
-        # Download and unzip
-        $packageArgs = @{
-            packageName    = ${Env:ChocolateyPackageName}
-            unzipLocation  = $toolDir
-            url            = $zipUrl
-            checksum       = $zipSha256
-            checksumType   = 'sha256'
-            url64bit       = $zipUrl_64
-            checksum64     = $zipSha256_64
-        }
-        Install-ChocolateyZipPackage @packageArgs | Out-Null
+        # We do not check hashes for tools that we use signature verification for
+        if ($verifySignature) {
+            # Download zip
+            $packageArgs      = @{
+                packageName     = $env:ChocolateyPackageName
+                file            = Join-Path ${Env:TEMP} $toolName
+                url             = $zipUrl
+            }
+            $filePath = Get-ChocolateyWebFile @packageArgs
+            # Extract zip
+            Get-ChocolateyUnzip -FileFullPath $filePath -Destination $toolDir
+        }
+        else {  # Not verifying signature, so check if hash is as expected
+            # Download and unzip
+            $packageArgs = @{
+                packageName    = ${Env:ChocolateyPackageName}
+                unzipLocation  = $toolDir
+                url            = $zipUrl
+                checksum       = $zipSha256
+                checksumType   = 'sha256'
+                url64bit       = $zipUrl_64
+                checksum64     = $zipSha256_64
+            }
+            Install-ChocolateyZipPackage @packageArgs | Out-Null
+        }
         VM-Assert-Path $toolDir
 
         # If $innerFolder is set to $true, after unzipping there should be only one folder
@@ -430,6 +446,21 @@ function VM-Install-From-Zip {
             $toolDir = Join-Path $toolDir $dirList[0].Name -Resolve
         }
 
+        if ($verifySignature) {
+            # Check signature of all executable files individually
+            Get-ChildItem -Path "$toolDir\*.exe" | ForEach-Object {
+                try {
+                    # Check signature for each file
+                    VM-Assert-Signature $_.FullName
+                } catch {
+                    # Remove the file with invalid signature
+                    Write-Warning "Removing file '$($_.FullName)' due to invalid signature"
+                    Remove-Item $_.FullName -Force -ea 0 | Out-Null
+                    VM-Write-Log-Exception $_
+                }
+            }
+        }
+
         if ($powershellCommand) {
             $executablePath = $toolDir
             VM-Install-Shortcut -toolName $toolName -category $category -arguments $powershellCommand -executableDir $executablePath -powershell