go-tg-screenshot-bot

1
//go:build cgo && darwin
2

3
package screenshot
4

5
/*
6
#if __ENVIRONMENT_MAC_OS_X_VERSION_MIN_REQUIRED__ > MAC_OS_VERSION_14_4
7
#cgo CFLAGS: -x objective-c
8
#cgo LDFLAGS: -framework CoreGraphics -framework CoreFoundation -framework ScreenCaptureKit
9
#include <ScreenCaptureKit/ScreenCaptureKit.h>
10
#else
11
#cgo LDFLAGS: -framework CoreGraphics -framework CoreFoundation
12
#endif
13
#include <CoreGraphics/CoreGraphics.h>
14

15
static CGImageRef capture(CGDirectDisplayID id, CGRect diIntersectDisplayLocal, CGColorSpaceRef colorSpace) {
16
#if __ENVIRONMENT_MAC_OS_X_VERSION_MIN_REQUIRED__ > MAC_OS_VERSION_14_4
17
    dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
18
    __block CGImageRef result = nil;
19
    [SCShareableContent getShareableContentWithCompletionHandler:^(SCShareableContent* content, NSError* error) {
20
        @autoreleasepool {
21
            if (error) {
22
                dispatch_semaphore_signal(semaphore);
23
                return;
24
            }
25
            SCDisplay* target = nil;
26
            for (SCDisplay *display in content.displays) {
27
                if (display.displayID == id) {
28
                    target = display;
29
                    break;
30
                }
31
            }
32
            if (!target) {
33
                dispatch_semaphore_signal(semaphore);
34
                return;
35
            }
36
            SCContentFilter* filter = [[SCContentFilter alloc] initWithDisplay:target excludingWindows:@[]];
37
            SCStreamConfiguration* config = [[SCStreamConfiguration alloc] init];
38
            config.sourceRect = diIntersectDisplayLocal;
39
            config.width = diIntersectDisplayLocal.size.width;
40
            config.height = diIntersectDisplayLocal.size.height;
41
            [SCScreenshotManager captureImageWithFilter:filter
42
                                          configuration:config
43
                                      completionHandler:^(CGImageRef img, NSError* error) {
44
                if (!error) {
45
                    result = CGImageCreateCopyWithColorSpace(img, colorSpace);
46
                }
47
                dispatch_semaphore_signal(semaphore);
48
            }];
49
        }
50
    }];
51
    dispatch_semaphore_wait(semaphore, DISPATCH_TIME_FOREVER);
52
    dispatch_release(semaphore);
53
    return result;
54
#else
55
    CGImageRef img = CGDisplayCreateImageForRect(id, diIntersectDisplayLocal);
56
    if (!img) {
57
        return nil;
58
    }
59
    CGImageRef copy = CGImageCreateCopyWithColorSpace(img, colorSpace);
60
    CGImageRelease(img);
61
    if (!copy) {
62
        return nil;
63
    }
64
    return copy;
65
#endif
66
}
67
*/
68
import "C"
69

70
import (
71
	"errors"
72
	"image"
73
	"unsafe"
74
)
75

76
func Capture(x, y, width, height int) (*image.RGBA, error) {
77
	if width <= 0 || height <= 0 {
78
		return nil, errors.New("width or height should be > 0")
79
	}
80

81
	rect := image.Rect(0, 0, width, height)
82
	img, err := createImage(rect)
83
	if err != nil {
84
		return nil, err
85
	}
86

87
	// cg: CoreGraphics coordinate (origin: lower-left corner of primary display, x-axis: rightward, y-axis: upward)
88
	// win: Windows coordinate (origin: upper-left corner of primary display, x-axis: rightward, y-axis: downward)
89
	// di: Display local coordinate (origin: upper-left corner of the display, x-axis: rightward, y-axis: downward)
90

91
	cgMainDisplayBounds := getCoreGraphicsCoordinateOfDisplay(C.CGMainDisplayID())
92

93
	winBottomLeft := C.CGPointMake(C.CGFloat(x), C.CGFloat(y+height))
94
	cgBottomLeft := getCoreGraphicsCoordinateFromWindowsCoordinate(winBottomLeft, cgMainDisplayBounds)
95
	cgCaptureBounds := C.CGRectMake(cgBottomLeft.x, cgBottomLeft.y, C.CGFloat(width), C.CGFloat(height))
96

97
	ids := activeDisplayList()
98

99
	ctx := createBitmapContext(width, height, (*C.uint32_t)(unsafe.Pointer(&img.Pix[0])), img.Stride)
100
	if ctx == 0 {
101
		return nil, errors.New("cannot create bitmap context")
102
	}
103

104
	colorSpace := createColorspace()
105
	if colorSpace == 0 {
106
		return nil, errors.New("cannot create colorspace")
107
	}
108
	defer C.CGColorSpaceRelease(colorSpace)
109

110
	for _, id := range ids {
111
		cgBounds := getCoreGraphicsCoordinateOfDisplay(id)
112
		cgIntersect := C.CGRectIntersection(cgBounds, cgCaptureBounds)
113
		if C.CGRectIsNull(cgIntersect) {
114
			continue
115
		}
116
		if cgIntersect.size.width <= 0 || cgIntersect.size.height <= 0 {
117
			continue
118
		}
119

120
		// CGDisplayCreateImageForRect potentially fail in case width/height is odd number.
121
		if int(cgIntersect.size.width)%2 != 0 {
122
			cgIntersect.size.width = C.CGFloat(int(cgIntersect.size.width) + 1)
123
		}
124
		if int(cgIntersect.size.height)%2 != 0 {
125
			cgIntersect.size.height = C.CGFloat(int(cgIntersect.size.height) + 1)
126
		}
127

128
		diIntersectDisplayLocal := C.CGRectMake(cgIntersect.origin.x-cgBounds.origin.x,
129
			cgBounds.origin.y+cgBounds.size.height-(cgIntersect.origin.y+cgIntersect.size.height),
130
			cgIntersect.size.width, cgIntersect.size.height)
131

132
		image := C.capture(id, diIntersectDisplayLocal, colorSpace)
133
		if unsafe.Pointer(image) == nil {
134
			return nil, errors.New("cannot capture display")
135
		}
136
		defer C.CGImageRelease(image)
137

138
		cgDrawRect := C.CGRectMake(cgIntersect.origin.x-cgCaptureBounds.origin.x, cgIntersect.origin.y-cgCaptureBounds.origin.y,
139
			cgIntersect.size.width, cgIntersect.size.height)
140
		C.CGContextDrawImage(ctx, cgDrawRect, image)
141
	}
142

143
	i := 0
144
	for iy := 0; iy < height; iy++ {
145
		j := i
146
		for ix := 0; ix < width; ix++ {
147
			// ARGB => RGBA, and set A to 255
148
			img.Pix[j], img.Pix[j+1], img.Pix[j+2], img.Pix[j+3] = img.Pix[j+1], img.Pix[j+2], img.Pix[j+3], 255
149
			j += 4
150
		}
151
		i += img.Stride
152
	}
153

154
	return img, nil
155
}
156

157
func NumActiveDisplays() int {
158
	var count C.uint32_t = 0
159
	if C.CGGetActiveDisplayList(0, nil, &count) == C.kCGErrorSuccess {
160
		return int(count)
161
	} else {
162
		return 0
163
	}
164
}
165

166
func GetDisplayBounds(displayIndex int) image.Rectangle {
167
	id := getDisplayId(displayIndex)
168
	main := C.CGMainDisplayID()
169

170
	var rect image.Rectangle
171

172
	bounds := getCoreGraphicsCoordinateOfDisplay(id)
173
	rect.Min.X = int(bounds.origin.x)
174
	if main == id {
175
		rect.Min.Y = 0
176
	} else {
177
		mainBounds := getCoreGraphicsCoordinateOfDisplay(main)
178
		mainHeight := mainBounds.size.height
179
		rect.Min.Y = int(mainHeight - (bounds.origin.y + bounds.size.height))
180
	}
181
	rect.Max.X = rect.Min.X + int(bounds.size.width)
182
	rect.Max.Y = rect.Min.Y + int(bounds.size.height)
183

184
	return rect
185
}
186

187
func getDisplayId(displayIndex int) C.CGDirectDisplayID {
188
	main := C.CGMainDisplayID()
189
	if displayIndex == 0 {
190
		return main
191
	} else {
192
		n := NumActiveDisplays()
193
		ids := make([]C.CGDirectDisplayID, n)
194
		if C.CGGetActiveDisplayList(C.uint32_t(n), (*C.CGDirectDisplayID)(unsafe.Pointer(&ids[0])), nil) != C.kCGErrorSuccess {
195
			return 0
196
		}
197
		index := 0
198
		for i := 0; i < n; i++ {
199
			if ids[i] == main {
200
				continue
201
			}
202
			index++
203
			if index == displayIndex {
204
				return ids[i]
205
			}
206
		}
207
	}
208

209
	return 0
210
}
211

212
func getCoreGraphicsCoordinateOfDisplay(id C.CGDirectDisplayID) C.CGRect {
213
	main := C.CGDisplayBounds(C.CGMainDisplayID())
214
	r := C.CGDisplayBounds(id)
215
	return C.CGRectMake(r.origin.x, -r.origin.y-r.size.height+main.size.height,
216
		r.size.width, r.size.height)
217
}
218

219
func getCoreGraphicsCoordinateFromWindowsCoordinate(p C.CGPoint, mainDisplayBounds C.CGRect) C.CGPoint {
220
	return C.CGPointMake(p.x, mainDisplayBounds.size.height-p.y)
221
}
222

223
func createBitmapContext(width int, height int, data *C.uint32_t, bytesPerRow int) C.CGContextRef {
224
	colorSpace := createColorspace()
225
	if colorSpace == 0 {
226
		return 0
227
	}
228
	defer C.CGColorSpaceRelease(colorSpace)
229

230
	return C.CGBitmapContextCreate(unsafe.Pointer(data),
231
		C.size_t(width),
232
		C.size_t(height),
233
		8, // bits per component
234
		C.size_t(bytesPerRow),
235
		colorSpace,
236
		C.kCGImageAlphaNoneSkipFirst)
237
}
238

239
func createColorspace() C.CGColorSpaceRef {
240
	return C.CGColorSpaceCreateWithName(C.kCGColorSpaceSRGB)
241
}
242

243
func activeDisplayList() []C.CGDirectDisplayID {
244
	count := C.uint32_t(NumActiveDisplays())
245
	ret := make([]C.CGDirectDisplayID, count)
246
	if count > 0 && C.CGGetActiveDisplayList(count, (*C.CGDirectDisplayID)(unsafe.Pointer(&ret[0])), nil) == C.kCGErrorSuccess {
247
		return ret
248
	} else {
249
		return make([]C.CGDirectDisplayID, 0)
250
	}
251
}
252