feat: adopt structured logging, add config via env (FORBIDDEN, SCAN_INTERVAL, TIMEOUT, WORKERS), introduce worker pool and improved kill logic with timeouts

2025-08-07 11:40:00 +02:00
parent 33a65ffbff
commit 4fb751a268
3 changed files with 264 additions and 101 deletions
--- a/main.go
+++ b/main.go
@@ -2,105 +2,175 @@ package main

 import (
 	"context"
-	"fmt"
-	"io"
-	"log"
 	"os"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
+
+	logger "git.site.quack-lab.dev/dave/cylogger"
+	utils "git.site.quack-lab.dev/dave/cyutils"
 )

-var Error *log.Logger
-var Warning *log.Logger
-
-func init() {
-	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
-	logFile, err := os.Create("main.log")
-	if err != nil {
-		log.Printf("Error creating log file: %v", err)
-		os.Exit(1)
-	}
-	logger := io.MultiWriter(os.Stdout, logFile)
-	log.SetOutput(logger)
-
-	Error = log.New(io.MultiWriter(logFile, os.Stderr, os.Stdout),
-		fmt.Sprintf("%sERROR:%s ", "\033[0;101m", "\033[0m"),
-		log.Lmicroseconds|log.Lshortfile)
-	Warning = log.New(io.MultiWriter(logFile, os.Stdout),
-		fmt.Sprintf("%sWarning:%s ", "\033[0;93m", "\033[0m"),
-		log.Lmicroseconds|log.Lshortfile)
+type Config struct {
+	Forbidden    []string
+	ScanInterval time.Duration
+	Timeout      time.Duration
+	Workers      int
 }

-func main() {
-	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
+func getenv(key, def string) string {
+	if v, ok := os.LookupEnv(key); ok {
+		return v
+	}
+	return def
+}
+
+var timeUnits = map[string]int64{
+	"ms": 1,
+	"s":  1000,
+	"m":  60_000,
+	"h":  3_600_000,
+	"d":  86_400_000,
+	"M":  2_592_000_000,
+	"y":  31_536_000_000,
+}
+
+// parseDurationMS supports "1s", "500ms", and compound "1s_500ms"
+func parseDurationMS(expr string) int64 {
+	expr = strings.TrimSpace(expr)
+	if expr == "" {
+		return 0
+	}
+	var total int64
+	var val strings.Builder
+	var unit strings.Builder
+
+	flush := func() {
+		if val.Len() == 0 || unit.Len() == 0 {
+			return
+		}
+		v, err := strconv.ParseInt(val.String(), 10, 64)
+		if err != nil {
+			logger.Warning("Invalid duration value: %q: %v", val.String(), err)
+			val.Reset()
+			unit.Reset()
+			return
+		}
+		u := unit.String()
+		mul, ok := timeUnits[u]
+		if !ok {
+			logger.Warning("Invalid duration unit: %q", u)
+			val.Reset()
+			unit.Reset()
+			return
+		}
+		total += v * mul
+		val.Reset()
+		unit.Reset()
 	}

-	delay := time.Duration(2) * 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")
+	for _, part := range strings.Split(expr, "_") {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+		val.Reset()
+		unit.Reset()
+		for _, r := range part {
+			if r >= '0' && r <= '9' {
+				val.WriteRune(r)
+			} else if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') {
+				unit.WriteRune(r)
+			}
+		}
+		flush()
+	}
+	return total
+}
+
+func loadConfig() Config {
+	logger.InitFlag()
+	log := logger.Default.WithPrefix("loadConfig")
+
+	// Forbidden names
+	forbidden := []string{}
+	if env := strings.TrimSpace(getenv("FORBIDDEN", "")); env != "" {
+		for _, p := range strings.Split(env, ",") {
+			p = strings.TrimSpace(p)
+			if p != "" {
+				forbidden = append(forbidden, p)
+			}
+		}
+	}
+
+	// Default scan interval: 2s
+	scan := time.Duration(parseDurationMS(getenv("SCAN_INTERVAL", "2s"))) * time.Millisecond
+	if scan <= 0 {
+		scan = 2 * time.Second
+	}
+
+	// Default timeout: half the scan interval
+	timeout := time.Duration(parseDurationMS(getenv("TIMEOUT", ""))) * time.Millisecond
+	if timeout <= 0 {
+		timeout = scan / 2
+	}
+
+	workers := 10
+	if w := strings.TrimSpace(getenv("WORKERS", "")); w != "" {
+		if n, err := strconv.Atoi(w); err == nil && n > 0 {
+			workers = n
+		} else if err != nil {
+			log.Warning("Invalid WORKERS value %q: %v (using default %d)", w, err, workers)
+		}
+	}
+
+	cfg := Config{
+		Forbidden:    forbidden,
+		ScanInterval: scan,
+		Timeout:      timeout,
+		Workers:      workers,
+	}
+
+	// Config dump
+	log.Info("Configuration loaded")
+	log.Info("SCAN_INTERVAL(ms): %d", cfg.ScanInterval.Milliseconds())
+	log.Info("TIMEOUT(ms): %d", cfg.Timeout.Milliseconds())
+	log.Info("WORKERS: %d", cfg.Workers)
+	if len(cfg.Forbidden) == 0 {
+		log.Warning("FORBIDDEN is empty - nothing to kill")
 	} else {
-		var err error
-		delay, err = time.ParseDuration(scanDelay)
-		if err != nil {
-			Error.Printf("Error parsing scan delay: %v", err)
-			return
-		}
+		log.Info("Forbidden process names: %d", len(cfg.Forbidden))
+		log.Trace("Forbidden list: %v", cfg.Forbidden)
 	}

-	timeout := delay / 2
-	etimeout, exists := os.LookupEnv("HITMAN_TIMEOUT")
-	if !exists {
-		log.Printf("No timeout is set, defaulting to %vs", timeout.Seconds())
-		log.Printf("Set HITMAN_TIMEOUT to change this")
-	} else {
-		var err error
-		timeout, err = time.ParseDuration(etimeout)
-		if err != nil {
-			Error.Printf("Error parsing timeout: %v", err)
-			return
-		}
+	return cfg
+}
+
+// Kill sends SIGKILL (-9) immediately and waits until the process is reaped or timeout elapses.
+func Kill(ctx context.Context, pid int) error {
+	proc, err := os.FindProcess(pid)
+	if err != nil {
+		return err
 	}

-	procs := strings.Split(forbidden, ",")
-	workers := make(chan struct{}, 10)
+	// Unconditional -9
+	if err := proc.Signal(syscall.SIGKILL); err != nil {
+		return err
+	}

-	for {
-		log.Printf("Running")
-		procmap, err := BuildProcessMap()
-		if err != nil {
-			Error.Printf("Error building process map: %v", err)
-			return
-		}
+	done := make(chan error, 1)
+	go func() {
+		_, err := proc.Wait()
+		done <- err
+	}()

-		for _, proc := range procs {
-			workers <- struct{}{}
-			go func(proc string) {
-				defer func() { <-workers }()
-				proc = strings.Trim(proc, " ")
-				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 := KillWithTimeout(int(node.Proc.ProcessID), timeout)
-						if err != nil {
-							Error.Printf("Error terminating process %d: %v", node.Proc.ProcessID, err)
-						}
-					}
-				} else {
-					log.Printf("No forbidden process %s found", proc)
-				}
-			}(proc)
-		}
-		time.Sleep(delay)
+	select {
+	case err := <-done:
+		return err
+	case <-ctx.Done():
+		_ = proc.Release()
+		return ctx.Err()
 	}
 }

@@ -110,28 +180,78 @@ func KillWithTimeout(pid int, timeout time.Duration) error {
 	return Kill(ctx, pid)
 }

-func Kill(ctx context.Context, pid int) error {
-	process, err := os.FindProcess(pid)
+func processCycle(cfg Config) {
+	log := logger.Default.WithPrefix("processCycle")
+	log.Info("Starting process cleanup cycle")
+	log.Debug("Timeout(ms): %d", cfg.Timeout.Milliseconds())
+
+	procmap, err := BuildProcessMap()
 	if err != nil {
-		return err
+		log.Error("BuildProcessMap failed: %v", err)
+		return
 	}

-	err = process.Signal(syscall.SIGKILL)
-	if err != nil {
-		return err
+	if len(cfg.Forbidden) == 0 {
+		log.Warning("No forbidden processes defined; skipping")
+		return
 	}

-	done := make(chan error, 1)
-	go func() {
-		_, err := process.Wait()
-		done <- err
-	}()
+	// Parallel pass through forbidden names
+	utils.WithWorkers(cfg.Workers, cfg.Forbidden, func(worker int, name string) {
+		name = strings.TrimSpace(name)
+		ilog := log.WithPrefix("forbidden").WithPrefix(name).WithField("worker", worker)

-	select {
-	case err := <-done:
-		return err
-	case <-ctx.Done():
-		process.Release()
-		return ctx.Err()
+		if name == "" {
+			ilog.Warning("Empty name, skipping")
+			return
+		}
+
+		ilog.Debug("Searching processes by name")
+		ilog.Trace("Query: %s", name)
+
+		results, ok := procmap.findByName(name)
+		if !ok || len(results) == 0 {
+			ilog.Info("No matching processes found")
+			return
+		}
+
+		ilog.Info("Found %d matching processes", len(results))
+		for _, node := range results {
+			pid := int(node.Proc.ProcessID)
+			plog := ilog.WithPrefix("pid").WithPrefix(strconv.Itoa(pid))
+
+			plog.Debug("Killing with SIGKILL (-9)")
+			plog.Trace("PID: %d", pid)
+
+			if err := KillWithTimeout(pid, cfg.Timeout); err != nil {
+				plog.Error("Kill failed: %v", err)
+			} else {
+				plog.Info("Process killed")
+			}
+		}
+	})
+	log.Info("Process cleanup cycle complete")
+}
+
+func main() {
+	app := logger.Default.WithPrefix("main")
+	app.Info("Starting hitman (no-questions-asked)")
+
+	cfg := loadConfig()
+
+	// Initial cycle
+	processCycle(cfg)
+
+	// Ticker loop
+	t := time.NewTicker(cfg.ScanInterval)
+	defer t.Stop()
+
+	for {
+		ts := <-t.C
+		tlog := app.WithPrefix("tick").WithPrefix(strconv.FormatInt(ts.UnixMilli(), 10))
+		tlog.Info("Timer tick")
+		tlog.Trace("Timestamp(ms): %d", ts.UnixMilli())
+
+		processCycle(cfg)
 	}
 }