From ddbf2f5c4291687c53a508dffdd9bc89edd2b242 Mon Sep 17 00:00:00 2001
From: Elliot Chernofsky <35854880+emtuls@users.noreply.github.com>
Date: Thu, 25 Jan 2024 23:59:30 -0500
Subject: [PATCH] Add powershell support in VM-Install-Shortcut

---
 packages/common.vm/common.vm.nuspec           |  2 +-
 .../common.vm/tools/vm.common/vm.common.psm1  | 46 ++++++++-----------
 2 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/packages/common.vm/common.vm.nuspec b/packages/common.vm/common.vm.nuspec
index 13840e194..48acfb9ed 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.20240220</version>
+    <version>0.0.0.20240229</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 64b3daa6a..3ee86d63f 100755
--- a/packages/common.vm/tools/vm.common/vm.common.psm1
+++ b/packages/common.vm/tools/vm.common/vm.common.psm1
@@ -219,7 +219,6 @@ function VM-Install-Raw-GitHub-Repo {
     )
     try {
         $toolDir = Join-Path ${Env:RAW_TOOLS_DIR} $toolName
-        $shortcutDir = Join-Path ${Env:TOOL_LIST_DIR} $category
 
         # Remove files from previous zips for upgrade
         VM-Remove-PreviousZipPackage ${Env:chocolateyPackageFolder}
@@ -242,15 +241,10 @@ function VM-Install-Raw-GitHub-Repo {
         }
 
         if ($powershellCommand) {
-            $executableArgs = "-ExecutionPolicy Bypass -NoExit -Command $powershellCommand"
-            $powershellPath = Join-Path "${Env:WinDir}\system32\WindowsPowerShell\v1.0" "powershell.exe" -Resolve
-            $shortcut = Join-Path $shortcutDir "$toolName.lnk"
-            Install-ChocolateyShortcut -shortcutFilePath $shortcut -targetPath $powershellPath -Arguments $executableArgs -WorkingDirectory $toolDir -IconLocation $powershell
+            VM-Install-Shortcut -toolName $toolName -category $category -arguments $powershellCommand -executableDir $toolDir -powershell $true
         } else {
-            $shortcut = Join-Path $shortcutDir "$toolName.lnk"
-            Install-ChocolateyShortcut -shortcutFilePath $shortcut -targetPath $toolDir
+            VM-Install-Shortcut -toolName $toolName -category $category -executablePath $toolDir
         }
-        VM-Assert-Path $shortcut
 
         return $toolDir
     } catch {
@@ -266,11 +260,13 @@ function VM-Install-Shortcut{
         [string] $toolName,
         [Parameter(Mandatory=$true, Position=1)]
         [string] $category,
-        [Parameter(Mandatory=$true, Position=2)]
+        [Parameter(Mandatory=$false, Position=2)]
         [string] $executablePath,
         [Parameter(Mandatory=$false)]
         [bool] $consoleApp=$false,
         [Parameter(Mandatory=$false)]
+        [bool] $powershell=$false,
+        [Parameter(Mandatory=$false)]
         [switch] $runAsAdmin,
         [Parameter(Mandatory=$false)]
         [string] $executableDir,
@@ -284,15 +280,21 @@ function VM-Install-Shortcut{
 
     if (-Not $iconLocation) {$iconLocation = $executablePath}
 
-    if ($consoleApp) {
-        if (!$executableDir) {
-            $executableDir  = Join-Path ${Env:UserProfile} "Desktop"
+    if ($consoleApp -or $powershell) {
+        if (-not $executableDir) {
+            $executableDir = Join-Path ${Env:UserProfile} "Desktop"
         }
         VM-Assert-Path $executableDir
 
-        $executableCmd  = Join-Path ${Env:WinDir} "system32\cmd.exe" -Resolve
-        # Change to executable dir, print command to execute, and execute command
-        $executableArgs = "/K `"cd `"$executableDir`" && echo $executableDir^> $executablePath $arguments && `"$executablePath`" $arguments`""
+        if ($consoleApp) {
+            $executableCmd = Join-Path ${Env:WinDir} "system32\cmd.exe" -Resolve
+            # Change to executable dir, print command to execute, and execute command
+            $executableArgs = "/K `"cd `"$executableDir`" && echo $executableDir^> $executablePath $arguments && `"$executablePath`" $arguments`""
+        } else {
+            $executableCmd = Join-Path "${PSHome}" "powershell.exe" -Resolve
+            $executableArgs = "-ExecutionPolicy Bypass -NoExit -Command `"`$cmd = '$arguments'; Write-Host `$cmd; Invoke-Expression `$cmd`""
+            $iconLocation = $executableCmd
+        }
 
         $shortcutArgs = @{
             ShortcutFilePath = $shortcut
@@ -460,7 +462,6 @@ function VM-Install-Single-Ps1 {
     )
     try {
         $toolDir = Join-Path ${Env:RAW_TOOLS_DIR} $toolName
-        $shortcutDir = Join-Path ${Env:TOOL_LIST_DIR} $category
 
         # Download and install
         $scriptPath = Join-Path $toolDir "$toolName.ps1"
@@ -477,18 +478,9 @@ function VM-Install-Single-Ps1 {
         Get-ChocolateyWebFile @packageArgs
         VM-Assert-Path $scriptPath
 
-        $shortcut = Join-Path $shortcutDir "$toolName.lnk"
-        $targetCmd = Join-Path ${Env:WinDir} "system32\cmd.exe" -Resolve
-
-        if ($ps1Cmd) {
-            $targetArgs = "/K powershell.exe -ExecutionPolicy Bypass -NoExit -Command `"cd '$toolDir'; $ps1Cmd`""
-        } else {
-            $targetArgs = "/K powershell.exe -ExecutionPolicy Bypass -NoExit -Command `"cd '$toolDir'`""
-        }
-        $targetIcon = Join-Path (Join-Path ${Env:WinDir} "system32\WindowsPowerShell\v1.0") "powershell.exe" -Resolve
+        $powershellPath = Join-Path (Join-Path ${Env:WinDir} "system32\WindowsPowerShell\v1.0") "powershell.exe" -Resolve
+        VM-Install-Shortcut -toolName $toolName -category $category -executablePath $powershellPath -executableDir $toolDir -powershell $true -arguments $ps1Cmd
 
-        Install-ChocolateyShortcut -shortcutFilePath $shortcut -targetPath $targetCmd -arguments $targetArgs -workingDirectory $toolDir -iconLocation $targetIcon
-        VM-Assert-Path $shortcut
         return $scriptPath
     } catch {
         VM-Write-Log-Exception $_