Cleanup

main.go

@@ -5,10 +5,14 @@ import (
 	"io"
 	"log"
 	"os"
+	"strings"
+	"syscall"
+	"time"
 )
 
 var Error *log.Logger
 var Warning *log.Logger
+
 func init() {
 	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
 	logFile, err := os.Create("main.log")
@@ -28,11 +32,67 @@ func init() {
 }
 
 func main() {
-	procmap, err := BuildProcessMap()
-	if err != nil {
-		Error.Printf("Error building process map: %v", err)
+	forbidden, exists := os.LookupEnv("HITMAN_FORBIDDEN")
+	if !exists {
+		Error.Println("HITMAN_FORBIDDEN environment variable not set")
+		log.Printf("Please set to a comma separated list of process names to forbid")
 		return
 	}
+	delay := time.Duration(3) * time.Second
+	scanDelay, exists := os.LookupEnv("HITMAN_SCAN_DELAY")
+	if !exists {
+		log.Printf("No scan delay is set, defaulting to %vs", delay.Seconds())
+		log.Printf("Set HITMAN_SCAN_DELAY to change this")
+	} else {
+		var err error
+		delay, err = time.ParseDuration(scanDelay)
+		if err != nil {
+			Error.Printf("Error parsing scan delay: %v", err)
+			return
+		}
+	}
 
-	log.Printf("%#v", procmap)
+	procs := strings.Split(forbidden, ",")
+
+	for {
+		log.Printf("Running")
+		procmap, err := BuildProcessMap()
+		if err != nil {
+			Error.Printf("Error building process map: %v", err)
+			return
+		}
+
+		for _, proc := range procs {
+			log.Printf("Checking %s", proc)
+			res, ok := procmap.findByName(proc)
+			if ok {
+				log.Printf("Forbidden process %s found (x%d)", proc, len(res))
+				for _, node := range res {
+					log.Printf("Killing forbidden process %d", node.Proc.ProcessID)
+					err := Kill(node.Proc.ProcessID)
+					if err != nil {
+						Error.Printf("Error terminating process %d: %v", node.Proc.ProcessID, err)
+					}
+				}
+			} else {
+				log.Printf("No forbidden process %s found", proc)
+			}
+		}
+		time.Sleep(delay)
+	}
+}
+
+func Kill(pid uint32) error {
+	handle, err := syscall.OpenProcess(syscall.PROCESS_TERMINATE, false, uint32(pid))
+	if err != nil {
+		return fmt.Errorf("error opening process: %v", err)
+	}
+	defer syscall.CloseHandle(handle)
+
+	err = syscall.TerminateProcess(handle, 7172)
+	if err != nil {
+		return fmt.Errorf("error terminating process: %v", err)
+	}
+
+	return nil
 }