go-tg-screenshot-bot

1
//go:build windows
2

3
package screenshot
4

5
import (
6
	"errors"
7
	"github.com/lxn/win"
8
	"image"
9
	"syscall"
10
	"unsafe"
11
)
12

13
var (
14
	libUser32, _               = syscall.LoadLibrary("user32.dll")
15
	funcGetDesktopWindow, _    = syscall.GetProcAddress(syscall.Handle(libUser32), "GetDesktopWindow")
16
	funcEnumDisplayMonitors, _ = syscall.GetProcAddress(syscall.Handle(libUser32), "EnumDisplayMonitors")
17
	funcGetMonitorInfo, _      = syscall.GetProcAddress(syscall.Handle(libUser32), "GetMonitorInfoW")
18
	funcEnumDisplaySettings, _ = syscall.GetProcAddress(syscall.Handle(libUser32), "EnumDisplaySettingsW")
19
)
20

21
func Capture(x, y, width, height int) (*image.RGBA, error) {
22
	rect := image.Rect(0, 0, width, height)
23
	img, err := createImage(rect)
24
	if err != nil {
25
		return nil, err
26
	}
27

28
	hwnd := getDesktopWindow()
29
	hdc := win.GetDC(hwnd)
30
	if hdc == 0 {
31
		return nil, errors.New("GetDC failed")
32
	}
33
	defer win.ReleaseDC(hwnd, hdc)
34

35
	memory_device := win.CreateCompatibleDC(hdc)
36
	if memory_device == 0 {
37
		return nil, errors.New("CreateCompatibleDC failed")
38
	}
39
	defer win.DeleteDC(memory_device)
40

41
	bitmap := win.CreateCompatibleBitmap(hdc, int32(width), int32(height))
42
	if bitmap == 0 {
43
		return nil, errors.New("CreateCompatibleBitmap failed")
44
	}
45
	defer win.DeleteObject(win.HGDIOBJ(bitmap))
46

47
	var header win.BITMAPINFOHEADER
48
	header.BiSize = uint32(unsafe.Sizeof(header))
49
	header.BiPlanes = 1
50
	header.BiBitCount = 32
51
	header.BiWidth = int32(width)
52
	header.BiHeight = int32(-height)
53
	header.BiCompression = win.BI_RGB
54
	header.BiSizeImage = 0
55

56
	// GetDIBits balks at using Go memory on some systems. The MSDN example uses
57
	// GlobalAlloc, so we'll do that too. See:
58
	// https://docs.microsoft.com/en-gb/windows/desktop/gdi/capturing-an-image
59
	bitmapDataSize := uintptr(((int64(width)*int64(header.BiBitCount) + 31) / 32) * 4 * int64(height))
60
	hmem := win.GlobalAlloc(win.GMEM_MOVEABLE, bitmapDataSize)
61
	defer win.GlobalFree(hmem)
62
	memptr := win.GlobalLock(hmem)
63
	defer win.GlobalUnlock(hmem)
64

65
	old := win.SelectObject(memory_device, win.HGDIOBJ(bitmap))
66
	if old == 0 {
67
		return nil, errors.New("SelectObject failed")
68
	}
69
	defer win.SelectObject(memory_device, old)
70

71
	if !win.BitBlt(memory_device, 0, 0, int32(width), int32(height), hdc, int32(x), int32(y), win.SRCCOPY) {
72
		return nil, errors.New("BitBlt failed")
73
	}
74

75
	if win.GetDIBits(hdc, bitmap, 0, uint32(height), (*uint8)(memptr), (*win.BITMAPINFO)(unsafe.Pointer(&header)), win.DIB_RGB_COLORS) == 0 {
76
		return nil, errors.New("GetDIBits failed")
77
	}
78

79
	i := 0
80
	src := uintptr(memptr)
81
	for y := 0; y < height; y++ {
82
		for x := 0; x < width; x++ {
83
			v0 := *(*uint8)(unsafe.Pointer(src))
84
			v1 := *(*uint8)(unsafe.Pointer(src + 1))
85
			v2 := *(*uint8)(unsafe.Pointer(src + 2))
86

87
			// BGRA => RGBA, and set A to 255
88
			img.Pix[i], img.Pix[i+1], img.Pix[i+2], img.Pix[i+3] = v2, v1, v0, 255
89

90
			i += 4
91
			src += 4
92
		}
93
	}
94

95
	return img, nil
96
}
97

98
func NumActiveDisplays() int {
99
	var count int = 0
100
	enumDisplayMonitors(win.HDC(0), nil, syscall.NewCallback(countupMonitorCallback), uintptr(unsafe.Pointer(&count)))
101
	return count
102
}
103

104
func GetDisplayBounds(displayIndex int) image.Rectangle {
105
	var ctx getMonitorBoundsContext
106
	ctx.Index = displayIndex
107
	ctx.Count = 0
108
	enumDisplayMonitors(win.HDC(0), nil, syscall.NewCallback(getMonitorBoundsCallback), uintptr(unsafe.Pointer(&ctx)))
109
	return image.Rect(
110
		int(ctx.Rect.Left), int(ctx.Rect.Top),
111
		int(ctx.Rect.Right), int(ctx.Rect.Bottom))
112
}
113

114
func getDesktopWindow() win.HWND {
115
	ret, _, _ := syscall.Syscall(funcGetDesktopWindow, 0, 0, 0, 0)
116
	return win.HWND(ret)
117
}
118

119
func enumDisplayMonitors(hdc win.HDC, lprcClip *win.RECT, lpfnEnum uintptr, dwData uintptr) bool {
120
	ret, _, _ := syscall.Syscall6(funcEnumDisplayMonitors, 4,
121
		uintptr(hdc),
122
		uintptr(unsafe.Pointer(lprcClip)),
123
		lpfnEnum,
124
		dwData,
125
		0,
126
		0)
127
	return int(ret) != 0
128
}
129

130
func countupMonitorCallback(hMonitor win.HMONITOR, hdcMonitor win.HDC, lprcMonitor *win.RECT, dwData uintptr) uintptr {
131
	var count *int
132
	count = (*int)(unsafe.Pointer(dwData))
133
	*count = *count + 1
134
	return uintptr(1)
135
}
136

137
type getMonitorBoundsContext struct {
138
	Index int
139
	Rect  win.RECT
140
	Count int
141
}
142

143
func getMonitorBoundsCallback(hMonitor win.HMONITOR, hdcMonitor win.HDC, lprcMonitor *win.RECT, dwData uintptr) uintptr {
144
	var ctx *getMonitorBoundsContext
145
	ctx = (*getMonitorBoundsContext)(unsafe.Pointer(dwData))
146
	if ctx.Count != ctx.Index {
147
		ctx.Count = ctx.Count + 1
148
		return uintptr(1)
149
	}
150

151
	if realSize := getMonitorRealSize(hMonitor); realSize != nil {
152
		ctx.Rect = *realSize
153
	} else {
154
		ctx.Rect = *lprcMonitor
155
	}
156

157
	return uintptr(0)
158
}
159

160
type _MONITORINFOEX struct {
161
	win.MONITORINFO
162
	DeviceName [win.CCHDEVICENAME]uint16
163
}
164

165
const _ENUM_CURRENT_SETTINGS = 0xFFFFFFFF
166

167
type _DEVMODE struct {
168
	_            [68]byte
169
	DmSize       uint16
170
	_            [6]byte
171
	DmPosition   win.POINT
172
	_            [86]byte
173
	DmPelsWidth  uint32
174
	DmPelsHeight uint32
175
	_            [40]byte
176
}
177

178
// getMonitorRealSize makes a call to GetMonitorInfo
179
// to obtain the device name for the monitor handle
180
// provided to the method.
181
//
182
// With the device name, EnumDisplaySettings is called to
183
// obtain the current configuration for the monitor, this
184
// information includes the real resolution of the monitor
185
// rather than the scaled version based on DPI.
186
//
187
// If either handle calls fail, it will return a nil
188
// allowing the calling method to use the bounds information
189
// returned by EnumDisplayMonitors which may be affected
190
// by DPI.
191
func getMonitorRealSize(hMonitor win.HMONITOR) *win.RECT {
192
	info := _MONITORINFOEX{}
193
	info.CbSize = uint32(unsafe.Sizeof(info))
194

195
	ret, _, _ := syscall.Syscall(funcGetMonitorInfo, 2, uintptr(hMonitor), uintptr(unsafe.Pointer(&info)), 0)
196
	if ret == 0 {
197
		return nil
198
	}
199

200
	devMode := _DEVMODE{}
201
	devMode.DmSize = uint16(unsafe.Sizeof(devMode))
202

203
	if ret, _, _ := syscall.Syscall(funcEnumDisplaySettings, 3, uintptr(unsafe.Pointer(&info.DeviceName[0])), _ENUM_CURRENT_SETTINGS, uintptr(unsafe.Pointer(&devMode))); ret == 0 {
204
		return nil
205
	}
206

207
	return &win.RECT{
208
		Left:   devMode.DmPosition.X,
209
		Right:  devMode.DmPosition.X + int32(devMode.DmPelsWidth),
210
		Top:    devMode.DmPosition.Y,
211
		Bottom: devMode.DmPosition.Y + int32(devMode.DmPelsHeight),
212
	}
213
}
214