diff --git a/main.go b/main.go index 4f70d35..bffd66f 100644 --- a/main.go +++ b/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,66 @@ 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) + } + } + return + } + } + 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 }