dedicated cleanup task for runner

redcanaryco · Nov 9, 2023 · 932c181 · 932c181
1 parent b7863c2
commit 932c181
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 11 deletions.
diff --git a/Public/Invoke-AtomicRunner.ps1 b/Public/Invoke-AtomicRunner.ps1
@@ -37,12 +37,14 @@ function Invoke-AtomicRunner {
         [ValidateRange(0, [int]::MaxValue)]
         [int] $PauseBetweenAtomics,
 
+        [parameter(Mandatory = $false)]
+        [switch] $scheduledTaskCleanup,
+
         [Parameter(Mandatory = $false, ValueFromRemainingArguments = $true)]
         $OtherArgs
     )
     Begin { }
     Process {       
-
         function Get-GuidFromHostName( $basehostname ) {
             $guid = [System.Net.Dns]::GetHostName() -replace $($basehostname + "-"), ""
 
@@ -181,6 +183,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
@@ -221,11 +224,14 @@ function Invoke-AtomicRunner {
             }
 
             if ($null -ne $tr) {
-                Invoke-AtomicTestFromScheduleRow $tr
-                Write-Host -Fore cyan "Sleeping for $SleepTillCleanup seconds before cleaning up"; Start-Sleep -Seconds $SleepTillCleanup
-
-                # Cleanup after running test
-                Invoke-AtomicTestFromScheduleRow $tr $true
+                if (-not $scheduledTaskCleanup) {
+                    Invoke-AtomicTestFromScheduleRow $tr
+                }
+                else {
+                    # 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 {
                 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
@@ -1,5 +1,10 @@
 function Invoke-KickoffAtomicRunner {
 
+    Param(
+        [parameter(Mandatory = $false, Position = 0)]
+        [switch] $scheduledTaskCleanup
+    )
+
     #log rotation function
     function Rotate-Log {
         Param ($log, $max_filesize, $max_age)
@@ -45,10 +50,10 @@ function Invoke-KickoffAtomicRunner {
 
     # Invoke the Runner Script
     if ($artConfig.debug) {
-        Invoke-AtomicRunner *>> $all_log_file
+        Invoke-AtomicRunner -scheduledTaskCleanup:$scheduledTaskCleanup *>> $all_log_file
     }
     else {
-        Invoke-AtomicRunner
+        Invoke-AtomicRunner -scheduledTaskCleanup:$scheduledTaskCleanup
     }
 }
 

diff --git a/Public/Invoke-SetupAtomicRunner.ps1 b/Public/Invoke-SetupAtomicRunner.ps1
@@ -66,8 +66,15 @@ function Invoke-SetupAtomicRunner {
             $triggers += $trigger
         }
         $task = New-ScheduledTask -Action $taskAction -Principal $taskPrincipal -Trigger $triggers -Description "A task that runs 1 minute or later after boot to start the atomic test runner script"
+        # setup scheduled task that will start the runner cleanup task after restart
+        $taskName2 = "KickOff-AtomicRunnerScheduledTask"
+        Unregister-ScheduledTask $taskName -confirm:$false -ErrorAction Ignore
+        $taskAction2 = New-ScheduledTaskAction -Execute "powershell.exe" -Argument "-exec bypass -Command Invoke-KickoffAtomicRunner -scheduledTaskCleanup"
+        $taskPrincipal = New-ScheduledTaskPrincipal -UserId $artConfig.user
+        $task2 = New-ScheduledTask -Action $taskAction2 -Principal $taskPrincipal -Trigger $triggers -Description "A task that runs 1 minute or later after boot to start the atomic test runner cleanup script"
         try {
             $null = Register-ScheduledTask -TaskName $taskName -InputObject $task -User $artConfig.user -Password $($cred.GetNetworkCredential().password) -ErrorAction Stop
+            $null = Register-ScheduledTask -TaskName $taskName2 -InputObject $task2 -User $artConfig.user -Password $($cred.GetNetworkCredential().password) -ErrorAction Stop
         }
         catch {
             if ($_.CategoryInfo.Category -eq "AuthenticationError") {
@@ -84,14 +91,23 @@ function Invoke-SetupAtomicRunner {
         # sets cronjob string using basepath from config.ps1
         $pwshPath = which pwsh
         $job = "@reboot root sleep 60;$pwshPath -Command Invoke-KickoffAtomicRunner"
-        $exists = cat /etc/crontab | Select-String -Quiet "KickoffAtomicRunner"
-        #checks if the Kickoff-AtomicRunner job exists. If not appends it to the system crontab.
+        $job2 = "@reboot root sleep 60;$pwshPath -Command Invoke-KickoffAtomicRunner -scheduledTaskCleanup"
+        $exists = cat /etc/crontab | Select-String -Quiet "KickoffAtomicRunner$"
+        $exists2 = cat /etc/crontab | Select-String -Quiet "Invoke-KickoffAtomicRunner -scheduledTaskCleanup"
+        #checks if the Kickoff-AtomicRunner jobs exist. If not appends it to the system crontab.
         if ($null -eq $exists) {
             $(echo "$job" >> /etc/crontab)
             write-host "setting cronjob"
         }
         else {
-            write-host "cronjob already exists"
+            write-host "Invoke-KickoffAtomicRunner cronjob already exists"
+        }
+        if ($null -eq $exists2) {
+            $(echo "$job2" >> /etc/crontab)
+            write-host "setting cronjob for cleanup task"
+        }
+        else {
+            write-host "Invoke-KickoffAtomicRunner -scheduledTaskCleanup cronjob already exists"
         }
     }