Initial commit

procmap.go

@@ -0,0 +1,92 @@
+package main
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+type (
+	ProcessMap struct {
+		Map       map[uint32]*ProcessMapNode
+		NameIndex map[string][]*ProcessMapNode
+	}
+	ProcessMapNode struct {
+		Proc     windows.ProcessEntry32
+		Name     string
+		Children map[uint32]*ProcessMapNode
+	}
+)
+
+func (pt *ProcessMap) add(proc *windows.ProcessEntry32) error {
+	if pt.Map == nil {
+		pt.Map = make(map[uint32]*ProcessMapNode)
+	}
+	if pt.NameIndex == nil {
+		pt.NameIndex = make(map[string][]*ProcessMapNode)
+	}
+
+	procNode := &ProcessMapNode{
+		Proc: *proc,
+		Name: windows.UTF16ToString(proc.ExeFile[:]),
+	}
+
+	_, ok := pt.NameIndex[procNode.Name]
+	if !ok {
+		pt.NameIndex[procNode.Name] = make([]*ProcessMapNode, 0)
+	}
+	pt.NameIndex[procNode.Name] = append(pt.NameIndex[procNode.Name], procNode)
+
+	parent, parentExists := pt.Map[proc.ParentProcessID]
+	if !parentExists {
+		parent = procNode
+		pt.Map[proc.ParentProcessID] = parent
+	}
+
+	child, childExists := pt.Map[proc.ProcessID]
+	if !childExists {
+		child = procNode
+		pt.Map[proc.ProcessID] = child
+	}
+
+	if parent.Children == nil {
+		parent.Children = make(map[uint32]*ProcessMapNode)
+	}
+	parent.Children[proc.ProcessID] = child
+
+	return nil
+}
+
+func (pt *ProcessMap) findByName(name string) ([]*ProcessMapNode, bool) {
+	val, ok := pt.NameIndex[name]
+	return val, ok
+}
+func (pt *ProcessMap) findByPid(pid uint32) (*ProcessMapNode, bool) {
+	val, ok := pt.Map[pid]
+	return val, ok
+}
+
+func BuildProcessMap() (*ProcessMap, error) {
+	tree := &ProcessMap{}
+
+	snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
+	if err != nil {
+		return nil, err
+	}
+	defer windows.CloseHandle(snapshot)
+
+	var pe32 windows.ProcessEntry32
+	pe32.Size = uint32(unsafe.Sizeof(pe32))
+
+	var i int
+	err = windows.Process32First(snapshot, &pe32)
+	for err == nil {
+		tree.add(&pe32)
+		i++
+		if i > 500 {
+			break
+		}
+		err = windows.Process32Next(snapshot, &pe32)
+	}
+	return tree, nil
+}