oh-my-posh/environment_windows_win32.go

// +build windows

package main

import (
	"fmt"
	"regexp"
	"strings"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

// WindowsProcess is an implementation of Process for Windows.
type WindowsProcess struct {
	pid  int
	ppid int
	exe  string
}

// getImagePid returns the
func getImagePid(imageName string) ([]int, error) {
	processes, err := processes()
	if err != nil {
		return nil, err
	}
	var pids []int
	for i := 0; i < len(processes); i++ {
		if strings.ToLower(processes[i].exe) == imageName {
			pids = append(pids, processes[i].pid)
		}
	}
	return pids, nil
}

// getWindowTitle returns the title of a window linked to a process name
func getWindowTitle(imageName string, windowTitleRegex string) (string, error) {
	processPid, err := getImagePid(imageName)
	if err != nil {
		return "", nil
	}
	// returns the first window of the first pid
	_, windowTitle, err := GetWindowTitle(processPid[0], windowTitleRegex)
	if err != nil {
		return "", nil
	}
	return windowTitle, nil
}

func newWindowsProcess(e *windows.ProcessEntry32) *WindowsProcess {
	// Find when the string ends for decoding
	end := 0
	for {
		if e.ExeFile[end] == 0 {
			break
		}
		end++
	}

	return &WindowsProcess{
		pid:  int(e.ProcessID),
		ppid: int(e.ParentProcessID),
		exe:  syscall.UTF16ToString(e.ExeFile[:end]),
	}
}

// Processes returns a snapshot of all the processes
// Taken and adapted from https://github.com/mitchellh/go-ps
func processes() ([]WindowsProcess, error) {
	// get process table snapshot
	handle, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
	if err != nil {
		return nil, syscall.GetLastError()
	}
	defer windows.CloseHandle(handle)

	// get process infor by looping through the snapshot
	var entry windows.ProcessEntry32
	entry.Size = uint32(unsafe.Sizeof(entry))
	err = windows.Process32First(handle, &entry)
	if err != nil {
		return nil, fmt.Errorf("error retrieving process info")
	}

	results := make([]WindowsProcess, 0, 50)
	for {
		results = append(results, *newWindowsProcess(&entry))
		err := windows.Process32Next(handle, &entry)
		if err != nil {
			if err == syscall.ERROR_NO_MORE_FILES {
				break
			}
			return nil, fmt.Errorf("Fail to syscall Process32Next: %v", err)
		}
	}

	return results, nil
}

// win32 specific code

// win32 dll load and function definitions
var (
	user32                       = syscall.NewLazyDLL("user32.dll")
	procEnumWindows              = user32.NewProc("EnumWindows")
	procGetWindowTextW           = user32.NewProc("GetWindowTextW")
	procGetWindowThreadProcessID = user32.NewProc("GetWindowThreadProcessId")
)

// EnumWindows call EnumWindows from user32 and returns all active windows
func EnumWindows(enumFunc uintptr, lparam uintptr) (err error) {
	r1, _, e1 := syscall.Syscall(procEnumWindows.Addr(), 2, uintptr(enumFunc), uintptr(lparam), 0)
	if r1 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

// GetWindowText returns the title and text of a window from a window handle
func GetWindowText(hwnd syscall.Handle, str *uint16, maxCount int32) (len int32, err error) {
	r0, _, e1 := syscall.Syscall(procGetWindowTextW.Addr(), 3, uintptr(hwnd), uintptr(unsafe.Pointer(str)), uintptr(maxCount))
	len = int32(r0)
	if len == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

// GetWindowTitle searchs for a window attached to the pid
func GetWindowTitle(pid int, windowTitleRegex string) (syscall.Handle, string, error) {
	var hwnd syscall.Handle
	var title string
	compiledRegex, err := regexp.Compile(windowTitleRegex)
	if err != nil {
		return 0, "", fmt.Errorf("Error while compiling the regex '%s'", windowTitleRegex)
	}
	// callback fro EnumWindows
	cb := syscall.NewCallback(func(h syscall.Handle, p uintptr) uintptr {
		var prcsID int = 0
		// get pid
		_, _, _ = procGetWindowThreadProcessID.Call(uintptr(h), uintptr(unsafe.Pointer(&prcsID)))
		// check if pid matches spotify pid
		if prcsID == pid {
			b := make([]uint16, 200)
			_, err := GetWindowText(h, &b[0], int32(len(b)))
			if err != nil {
				// ignore the error
				return 1 // continue enumeration
			}
			title = syscall.UTF16ToString(b)
			if compiledRegex.MatchString(title) {
				hwnd = h
				return 0
			}
		}

		return 1 // continue enumeration
	})
	// Enumerates all top-level windows on the screen
	EnumWindows(cb, 0)
	if hwnd == 0 {
		return 0, "", fmt.Errorf("No window with title '%b' found", pid)
	}
	return hwnd, title, nil
}