Kill Process Improvements (#545)

Summary:
Pull Request resolved: #545

### Context
Currently the `kill_process` action in TTPForge in cases of failure like `Couldn't find process` and `Couldn't kill process` by default throws an exception which leads to TTP being stopped. While this could be a valid usecase in some situations, allowing user control on whether to throw an error or to print an error and return normally helps make the code versatile.

Following is the decision making in the initial commit of Kill Process Action: [link](https://app.diagrams.net/#G1A4S3RgNffbaMbp40RCpqvmprH1QQvVvi#%7B%22pageId%22%3A%22CNsj_rZLfA5pAP3afiLK%22%7D)
https://pxl.cl/7Lcn3
This flow diagram shows the updated decision making: [link](https://app.diagrams.net/?title=Kill%20Process%20Updated%20Decision%20making.drawio&client=1#G1TePfzl48hKjTKe-K-efOH8F-ASMror7j#%7B%22pageId%22%3A%22CNsj_rZLfA5pAP3afiLK%22%7D)
https://pxl.cl/7Lcmh

### Changes

*   **Code Improvements**: The `kill_process` step has been updated to receive user inputs on actions to take in case of failures like `error_on_find_process_failure` and `error_on_kill_failure`. This facilitates increased control on the code flow by the user to satisfy different needs.
*   **Documentation**: Added an example TTP for reference on the use of the new user inputs for the step `kill_process`.
*   **Testing**: Additional test cases have been added to verify the functionality of the added flags on the `kill_process` step.

### Impact

*   **Improved Usability**: The changes in this diff make it easier for users to control the output of the step on the basis of their needs. (Ignoring or throwing exceptions in cases of failure)

Reviewed By: d0n601

Differential Revision: D78521586

fbshipit-source-id: 4014f0d31ad13fa90c4183d680ca858d40fc56a9

Author: Harshit-Mashru

docs/foundations/actions/kill_process.md

@@ -3,7 +3,8 @@
 The `kill_process` action can be used to terminate processes running on a system.
 Check out the TTP below to see how it works:
 
-<!-- TODO: Link to add for example TTP after code is landed -->
+[Kill Process Unix](https://github.com/facebookincubator/TTPForge/blob/0deabc567751d90078e5db3c2a84574396b43dc1/example-ttps/actions/kill-process/kill-process-unix.yaml)
+[Kill Process Windows](https://github.com/facebookincubator/TTPForge/blob/0deabc567751d90078e5db3c2a84574396b43dc1/example-ttps/actions/kill-process/kill-process-windows.yaml)
 
 You can experiment with the above TTP by installing the `examples` TTP
 repository (skip this if `ttpforge list repos` shows that the `examples` repo is
@@ -27,10 +28,16 @@ You can specify the following YAML fields for the `kill_process:` action:
 you wish to kill
 - `kill_process_name:` (type: `string`) the process name of the process that
 you wish to kill
+- `error_on_find_process_failure:` (type: `bool`) whether to raise an error if
+finding the process name/id fails
+- `error_on_kill_failure:` (type: `bool`) whether to raise an error if killing
+ the process fails
 
 ## Additional Notes
 
 If both `kill_process_id` and `kill_process_name` are specified, the action
 will only consider process ID as long as it is valid.
 If an invalid `kill_process_id` is specified, the action will fall back to
 using `kill_process_name` to kill the processes.
+Both the flags `error_on_find_process_failure` and `error_on_kill_failure` are
+set to `false` by default.

example-ttps/actions/kill-process/kill-process-unix-failure.yaml

@@ -0,0 +1,28 @@
+---
+api_version: 2.0
+uuid: fbf458fb-c351-4633-97b6-67332f94b7c1
+name: "Kill Process Unix Failure"
+description: |
+  "This is an example TTP that kills a ping process on a Unix system that does not exist"
+requirements:
+  platforms:
+    - os: linux
+    - os: darwin
+steps:
+  - name: Process to kill started using python
+    inline: |
+      ps aux | grep ping
+      echo "================="
+      echo "Starting the detached ping process to be killed"
+      python3 -c "import subprocess,os; subprocess.Popen(['ping', '-c', '100', '127.0.0.1'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, preexec_fn=os.setpgrp)"
+      echo "================="
+      ps aux | grep ping | grep -v grep
+      echo "================="
+  - name: Killing the ping process
+    kill_process_id: ""
+    kill_process_name: "ping123"
+    error_on_find_process_failure: false
+    error_on_kill_process_failure: false
+  - name: Show processes
+    inline: |
+      ps aux | grep ping

example-ttps/actions/kill-process/kill-process-unix.yaml

@@ -21,6 +21,7 @@ steps:
   - name: Killing the ping process
     kill_process_id: ""
     kill_process_name: "ping"
+    error_on_find_process_failure: true
   - name: Show processes
     inline: |
       ps aux | grep ping

example-ttps/actions/kill-process/kill-process-windows-failure.yaml

@@ -0,0 +1,27 @@
+---
+api_version: 2.0
+uuid: eb38a84b-837b-47d9-9373-3d35a824b003
+name: "Kill Process Windows Failure"
+description: |
+  "This is an example TTP that kills a ping process on a Windows system that does not exist"
+requirements:
+  platforms:
+    - os: windows
+steps:
+  - name: Process to kill started using python
+    executor: cmd
+    inline: |
+      tasklist | findstr PING.EXE
+      echo "Starting the detached ping process to be killed"
+      python3 -c "import subprocess,os; subprocess.Popen(['ping', '-n', '10', '127.0.0.1'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, creationflags=8)"
+      echo "================="
+      tasklist | findstr PING.EXE
+      echo "================="
+  - name: Killing the ping process
+    kill_process_id: ""
+    kill_process_name: "PING123"
+    error_on_find_process_failure: false
+  - name: Show processes
+    executor: cmd
+    inline: |
+      tasklist | findstr PING.EXE

pkg/blocks/killprocess.go

@@ -32,9 +32,11 @@ import (
 // Its intended use is simulating malicious programs stopping
 // critical applications/processes
 type KillProcessStep struct {
-	actionDefaults `yaml:",inline"`
-	ProcessID      string `yaml:"kill_process_id,omitempty"`
-	ProcessName    string `yaml:"kill_process_name,omitempty"`
+	actionDefaults            `yaml:",inline"`
+	ProcessID                 string `yaml:"kill_process_id,omitempty"`
+	ProcessName               string `yaml:"kill_process_name,omitempty"`
+	ErrorOnFindProcessFailure bool   `yaml:"error_on_find_process_failure,omitempty"`
+	ErrorOnKillFailure        bool   `yaml:"error_on_kill_failure,omitempty"`
 }
 
 // NewKillProcessStep creates a new KillProcessStep instance and returns a pointer to it.
@@ -100,13 +102,26 @@ func (s *KillProcessStep) extractPIDs() ([]int, error) {
 
 	if processID > 0 {
 		logging.L().Infof("Using Process ID: %v", processID)
+
+		err := processutils.VerifyPIDExists(processID)
+		if err != nil {
+			logging.L().Errorf("Error while trying to verify PID exists: %+v", err)
+			if s.ErrorOnFindProcessFailure {
+				return nil, err
+			}
+			return []int{}, nil
+		}
 		return []int{processID}, nil
 	} else if s.ProcessName != "" {
 		logging.L().Infof("Finding processes with name: %v", s.ProcessName)
+
 		processes, err := processutils.GetPIDsByName(s.ProcessName)
 		if err != nil {
 			logging.L().Errorf("Error while trying to get PIDs from name: %+v", err)
-			return nil, err
+			if s.ErrorOnFindProcessFailure {
+				return nil, err
+			}
+			return []int{}, nil
 		}
 
 		pids := make([]int, len(processes))
@@ -126,23 +141,20 @@ func (s *KillProcessStep) extractPIDs() ([]int, error) {
 func (s *KillProcessStep) killProcesses(pids []int) error {
 	logging.L().Infof("Killing the following processes: %v", pids)
 
-	throwError := len(pids) == 1
-
-	// Throwing error if only 1 PID and we fail to kill
 	for _, pid := range pids {
 		proc, err := os.FindProcess(pid)
 		if err != nil {
 			logging.L().Errorf("Error while trying to find process with ID: %v; %+v", pid, err)
-			if throwError {
+			if s.ErrorOnFindProcessFailure {
 				return err
 			}
 			continue
 		}
 
-		logging.L().Infof("Got process handle for PID %d: %+v", pid, proc)
+		logging.L().Infof("Got process handle with PID: %d", pid)
 		if err := proc.Kill(); err != nil {
 			logging.L().Errorf("Failed to kill process with ID: %v; %+v", pid, err)
-			if throwError {
+			if s.ErrorOnKillFailure {
 				return err
 			}
 			continue
@@ -161,6 +173,10 @@ func (s *KillProcessStep) Execute(_ TTPExecutionContext) (*ActResult, error) {
 		return nil, err
 	}
 
+	if len(pids) == 0 {
+		logging.L().Infof("No processes found to kill")
+		return &ActResult{}, nil
+	}
 	if err := s.killProcesses(pids); err != nil {
 		return nil, err
 	}

pkg/blocks/killprocess_test.go

@@ -20,7 +20,6 @@ THE SOFTWARE.
 package blocks
 
 import (
-	"fmt"
 	"os"
 	"strconv"
 	"testing"
@@ -42,33 +41,59 @@ func TestKillProcessExecute(t *testing.T) {
 		expectValidateError bool
 	}{
 		{
-			name:        "Kill non-existent process with process id",
+			name:        "Kill non-existent process with process id - throw error",
 			description: "Trying to kill a process with process id that doesn't exist",
 			step: &KillProcessStep{
-				ProcessID:   "123456789",
-				ProcessName: "ping",
+				ProcessID:                 "123",
+				ProcessName:               "ping",
+				ErrorOnFindProcessFailure: true,
+			},
+			createProcess:      false,
+			expectExecuteError: true,
+		},
+		{
+			name:        "Kill non-existent process with process id - continue on error",
+			description: "Trying to kill a process with process id that doesn't exist",
+			step: &KillProcessStep{
+				ProcessID:                 "123456789",
+				ProcessName:               "ping",
+				ErrorOnFindProcessFailure: false,
+				ErrorOnKillFailure:        false,
+			},
+			createProcess:      false,
+			expectExecuteError: false,
+		},
+		{
+			// Test might throw an error on stress run
+			name:        "Kill non-existent process with process name - throw error",
+			description: "Trying to kill a process with process name that doesn't exist",
+			step: &KillProcessStep{
+				ProcessID:                 "",
+				ProcessName:               "ping",
+				ErrorOnFindProcessFailure: true,
 			},
 			createProcess:      false,
 			expectExecuteError: true,
 		},
 		{
-			name:        "Kill non-existent process with process name",
+			name:        "Kill non-existent process with process name - continue on error",
 			description: "Trying to kill a process with process name that doesn't exist",
 			step: &KillProcessStep{
 				ProcessID:   "",
 				ProcessName: "ping",
 			},
 			createProcess:      false,
-			expectExecuteError: true,
+			expectExecuteError: false,
 		},
 		{
-			name:        "Kill existent process with process id",
+			name:        "Kill non-existent process with process id",
 			description: "Trying to kill a process with process id that exists",
 			step: &KillProcessStep{
 				ProcessID:   "123456789",
 				ProcessName: "ping",
 			},
-			createProcess: true,
+			createProcess:      true,
+			expectExecuteError: false,
 		},
 		{
 			name:        "Kill existent process with process name",
@@ -77,7 +102,8 @@ func TestKillProcessExecute(t *testing.T) {
 				ProcessID:   "",
 				ProcessName: "ping",
 			},
-			createProcess: true,
+			createProcess:      true,
+			expectExecuteError: false,
 		},
 		{
 			name:        "Kill process invalid process id",
@@ -110,8 +136,9 @@ func TestKillProcessExecute(t *testing.T) {
 			name:        "Kill process ID with templating",
 			description: "Trying to kill a process with templating",
 			step: &KillProcessStep{
-				ProcessID:   "{[{.StepVars.pid}]}",
-				ProcessName: "ping",
+				ProcessID:                 "{[{.StepVars.pid}]}",
+				ProcessName:               "ping",
+				ErrorOnFindProcessFailure: true,
 			},
 			stepVars: map[string]string{
 				"pid": "123456789",
@@ -124,8 +151,9 @@ func TestKillProcessExecute(t *testing.T) {
 			name:        "Kill process name with templating",
 			description: "Trying to kill a process with process name and templating",
 			step: &KillProcessStep{
-				ProcessID:   "",
-				ProcessName: "{[{.StepVars.processName}]}",
+				ProcessID:                 "",
+				ProcessName:               "{[{.StepVars.processName}]}",
+				ErrorOnFindProcessFailure: true,
 			},
 			stepVars: map[string]string{
 				"processName": "touch",
@@ -175,7 +203,6 @@ func TestKillProcessExecute(t *testing.T) {
 				require.Error(t, err2)
 				return
 			}
-			fmt.Printf(tc.step.ProcessID)
 			require.NoError(t, err2)
 
 			if tc.createProcess {

pkg/processutils/processutils.go

@@ -42,3 +42,17 @@ func GetPIDsByName(processName string) ([]int32, error) {
 	}
 	return pids, nil
 }
+
+// VerifyPIDExists returns a boolean basis if a process with the input PID exists
+func VerifyPIDExists(pid int) error {
+	processes, err := process.Processes()
+	if err != nil {
+		return err
+	}
+	for _, proc := range processes {
+		if int(proc.Pid) == pid {
+			return nil
+		}
+	}
+	return fmt.Errorf("No process found with PID: %d", pid)
+}

pkg/processutils/processutils_test.go …g/processutils/processutils_unix_test.gopkg/processutils/processutils_test.go renamed to pkg/processutils/processutils_unix_test.go pkg/processutils/processutils_test.go renamed to pkg/processutils/processutils_unix_test.go

@@ -1,3 +1,6 @@
+//go:build !windows
+// +build !windows
+
 /*
 Copyright © 2023-present, Meta Platforms, Inc. and affiliates
 Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -44,7 +47,7 @@ func TestGetPidsByName(t *testing.T) {
 		},
 		{
 			name:                  "Process Does Not Exist",
-			processName:           "ping",
+			processName:           "pingabc",
 			processExists:         false,
 			expectEntriesInResult: false,
 		},
@@ -89,3 +92,38 @@ func TestGetPidsByName(t *testing.T) {
 		})
 	}
 }
+
+// Only PID=1 is always running in Unix & Linux but not windows
+func TestVerifyPIDExists(t *testing.T) {
+
+	testCases := []struct {
+		name        string
+		pid         int
+		expectError bool
+	}{
+		{
+			name:        "Process Exists",
+			pid:         1,
+			expectError: false,
+		},
+		{
+			name:        "Process Does Not Exist",
+			pid:         124567890987654321,
+			expectError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+
+			// Testing
+			err := VerifyPIDExists(tc.pid)
+
+			if tc.expectError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}