From f43d87e7447bcc8b180d7f0898d4a92fba4da33f Mon Sep 17 00:00:00 2001
From: Carrie Roberts <clr2of8@gmail.com>
Date: Mon, 19 Feb 2024 20:05:42 -0500
Subject: [PATCH] Fix for hung runner when AV kills the process running the
 atomic (#171)

---
 Public/Invoke-AtomicRunner.ps1        | 19 ++++++++++++++-----
 Public/Invoke-KickoffAtomicRunner.ps1 | 12 +++++++++++-
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/Public/Invoke-AtomicRunner.ps1 b/Public/Invoke-AtomicRunner.ps1
index a342516..dbe9da4 100755
--- a/Public/Invoke-AtomicRunner.ps1
+++ b/Public/Invoke-AtomicRunner.ps1
@@ -37,6 +37,9 @@ function Invoke-AtomicRunner {
         [ValidateRange(0, [int]::MaxValue)]
         [int] $PauseBetweenAtomics,
 
+        [parameter(Mandatory = $false)]
+        [switch] $scheduledTaskCleanup,
+
         [Parameter(Mandatory = $false, ValueFromRemainingArguments = $true)]
         $OtherArgs
     )
@@ -178,6 +181,7 @@ function Invoke-AtomicRunner {
         $htvars.Remove('OtherArgs') | Out-Null
         $htvars.Remove('Cleanup') | Out-Null
         $htvars.Remove('PauseBetweenAtomics') | Out-Null
+        $htvars.Remove('scheduledTaskCleanup') | Out-Null
 
         $schedule = Get-Schedule $listOfAtomics
         # If the schedule is empty, end process
@@ -218,11 +222,16 @@ function Invoke-AtomicRunner {
             }
 
             if ($null -ne $tr) {
-                # run the atomic test and exit
-                Invoke-AtomicTestFromScheduleRow $tr
-                # Cleanup after running test
-                Write-Host -Fore cyan "Sleeping for $SleepTillCleanup seconds before cleaning up for $($tr.Technique) $($tr.auto_generated_guid) "; Start-Sleep -Seconds $SleepTillCleanup
-                Invoke-AtomicTestFromScheduleRow $tr $true
+                if ($scheduledTaskCleanup) {
+                    # Cleanup after running test
+                    Write-Host -Fore cyan "Sleeping for $SleepTillCleanup seconds before cleaning up for $($tr.Technique) $($tr.auto_generated_guid) "; Start-Sleep -Seconds $SleepTillCleanup
+                    Invoke-AtomicTestFromScheduleRow $tr $true
+                }
+                else {
+                    # run the atomic test and exit
+                    Invoke-AtomicTestFromScheduleRow $tr
+                    Start-Sleep 3; exit
+                }
             }
             else {
                 LogRunnerMsg "Could not find Test: $guid in schedule. Please update schedule to run this test."
diff --git a/Public/Invoke-KickoffAtomicRunner.ps1 b/Public/Invoke-KickoffAtomicRunner.ps1
index a0b380a..c02ffb0 100644
--- a/Public/Invoke-KickoffAtomicRunner.ps1
+++ b/Public/Invoke-KickoffAtomicRunner.ps1
@@ -25,19 +25,29 @@ function Invoke-KickoffAtomicRunner {
 
     #Create log files as needed
     $all_log_file = Join-Path $artConfig.atomicLogsPath "all-out-$($artConfig.basehostname).txt"
+    $all_log_file_cleanup = Join-Path $artConfig.atomicLogsPath "all-out-$($artConfig.basehostname)-cleanup.txt"
     New-Item $all_log_file -ItemType file -ErrorAction Ignore
+    New-Item $all_log_file_cleanup -ItemType file -ErrorAction Ignore
     New-Item $artConfig.logFile -ItemType File -ErrorAction Ignore
 
     #Rotate logs based on FileSize and Date max_filesize
     $max_filesize = 200 #in MB
     $max_file_age = 30 #in days
     Rotate-Log $all_log_file $max_filesize $max_file_age
+    Rotate-Log $all_log_file_cleanup $max_filesize $max_file_age
+
     Rotate-Log $artConfig.logFile $max_filesize $max_file_age #no need to repeat this. Can reduce further.
 
     # Optional additional delay before starting
     Start-Sleep $artConfig.kickOffDelay.TotalSeconds
 
-    if ($artConfig.debug) { Invoke-AtomicRunner  *>> $all_log_file } else { Invoke-AtomicRunner }
+    $WorkingDirectory = if ($IsLinux -or $IsMacOS) { "/tmp" } else { $env:TEMP }
+    $FileName = if ($IsLinux -or $IsMacOS) { "pwsh" } else { "powershell.exe" }
+    if ($artConfig.debug) { $Arguments = "-Command Invoke-AtomicRunner *>> $all_log_file" } else { $Arguments = "-Command Invoke-AtomicRunner" }
+    # Invoke the atomic as its own process because we don't want to skip the cleanup and rename process in the event that AV kills the process running the atomic
+    Start-Process -FilePath $FileName -ArgumentList $Arguments -WorkingDirectory $WorkingDirectory
+    # Run the cleanup commmands
+    if ($artConfig.debug) { Invoke-AtomicRunner -scheduledTaskCleanup *>> $all_log_file_cleanup } else { Invoke-AtomicRunner -scheduledTaskCleanup }
 }
 
 function LogRunnerMsg ($message) {