go-tg-screenshot-bot

1
//go:build !s390x && !ppc64le && !darwin && !windows && (linux || freebsd || openbsd || netbsd)
2

3
package screenshot
4

5
import (
6
	"fmt"
7
	"github.com/gen2brain/shm"
8
	"github.com/jezek/xgb"
9
	mshm "github.com/jezek/xgb/shm"
10
	"github.com/jezek/xgb/xinerama"
11
	"github.com/jezek/xgb/xproto"
12
	"image"
13
	"image/color"
14
)
15

16
func captureXinerama(x, y, width, height int) (img *image.RGBA, e error) {
17
	defer func() {
18
		err := recover()
19
		if err != nil {
20
			img = nil
21
			e = fmt.Errorf("%v", err)
22
		}
23
	}()
24
	c, err := xgb.NewConn()
25
	if err != nil {
26
		return nil, err
27
	}
28
	defer c.Close()
29

30
	err = xinerama.Init(c)
31
	if err != nil {
32
		return nil, err
33
	}
34

35
	reply, err := xinerama.QueryScreens(c).Reply()
36
	if err != nil {
37
		return nil, err
38
	}
39

40
	primary := reply.ScreenInfo[0]
41
	x0 := int(primary.XOrg)
42
	y0 := int(primary.YOrg)
43

44
	useShm := true
45
	err = mshm.Init(c)
46
	if err != nil {
47
		useShm = false
48
	}
49

50
	screen := xproto.Setup(c).DefaultScreen(c)
51
	wholeScreenBounds := image.Rect(0, 0, int(screen.WidthInPixels), int(screen.HeightInPixels))
52
	targetBounds := image.Rect(x+x0, y+y0, x+x0+width, y+y0+height)
53
	intersect := wholeScreenBounds.Intersect(targetBounds)
54

55
	rect := image.Rect(0, 0, width, height)
56
	img, err = createImage(rect)
57
	if err != nil {
58
		return nil, err
59
	}
60

61
	// Paint with opaque black
62
	index := 0
63
	for iy := 0; iy < height; iy++ {
64
		j := index
65
		for ix := 0; ix < width; ix++ {
66
			img.Pix[j+3] = 255
67
			j += 4
68
		}
69
		index += img.Stride
70
	}
71

72
	if !intersect.Empty() {
73
		var data []byte
74

75
		if useShm {
76
			shmSize := intersect.Dx() * intersect.Dy() * 4
77
			shmId, err := shm.Get(shm.IPC_PRIVATE, shmSize, shm.IPC_CREAT|0777)
78
			if err != nil {
79
				return nil, err
80
			}
81

82
			seg, err := mshm.NewSegId(c)
83
			if err != nil {
84
				return nil, err
85
			}
86

87
			data, err = shm.At(shmId, 0, 0)
88
			if err != nil {
89
				return nil, err
90
			}
91

92
			mshm.Attach(c, seg, uint32(shmId), false)
93

94
			defer mshm.Detach(c, seg)
95
			defer func() {
96
				_ = shm.Rm(shmId)
97
			}()
98
			defer func() {
99
				_ = shm.Dt(data)
100
			}()
101

102
			_, err = mshm.GetImage(c, xproto.Drawable(screen.Root),
103
				int16(intersect.Min.X), int16(intersect.Min.Y),
104
				uint16(intersect.Dx()), uint16(intersect.Dy()), 0xffffffff,
105
				byte(xproto.ImageFormatZPixmap), seg, 0).Reply()
106
			if err != nil {
107
				return nil, err
108
			}
109
		} else {
110
			xImg, err := xproto.GetImage(c, xproto.ImageFormatZPixmap, xproto.Drawable(screen.Root),
111
				int16(intersect.Min.X), int16(intersect.Min.Y),
112
				uint16(intersect.Dx()), uint16(intersect.Dy()), 0xffffffff).Reply()
113
			if err != nil {
114
				return nil, err
115
			}
116

117
			data = xImg.Data
118
		}
119

120
		// BitBlt by hand
121
		offset := 0
122
		for iy := intersect.Min.Y; iy < intersect.Max.Y; iy++ {
123
			for ix := intersect.Min.X; ix < intersect.Max.X; ix++ {
124
				r := data[offset+2]
125
				g := data[offset+1]
126
				b := data[offset]
127
				img.SetRGBA(ix-(x+x0), iy-(y+y0), color.RGBA{r, g, b, 255})
128
				offset += 4
129
			}
130
		}
131
	}
132

133
	return img, e
134
}
135