93 lines
2.0 KiB
Go
93 lines
2.0 KiB
Go
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 > 5000 {
|
|
break
|
|
}
|
|
err = windows.Process32Next(snapshot, &pe32)
|
|
}
|
|
return tree, nil
|
|
}
|