Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
// +build windows
package screen
import (
"fmt"
"image"
"reflect"
"syscall"
"unsafe"
"github.com/disintegration/gift"
)
func init() {
// We need to call SetProcessDpiAwareness so that Windows API calls will
// tell us the scale factor for our monitor so that our screenshot works
// on hi-res displays.
procSetProcessDpiAwareness.Call(uintptr(2)) // PROCESS_PER_MONITOR_DPI_AWARE
}
func capture() (image.Image, error) {
// Find the window
handle, err := findWindow()
if err != nil {
return nil, err
}
// Determine the full width and height of the window
rect, err := windowRect(handle)
if err != nil {
return nil, err
}
// Capture!
return captureWindow(handle, rect)
}
var (
modUser32 = syscall.NewLazyDLL("User32.dll")
procFindWindow = modUser32.NewProc("FindWindowW")
procGetClientRect = modUser32.NewProc("GetClientRect")
procGetDC = modUser32.NewProc("GetDC")
procReleaseDC = modUser32.NewProc("ReleaseDC")
modGdi32 = syscall.NewLazyDLL("Gdi32.dll")
procBitBlt = modGdi32.NewProc("BitBlt")
procCreateCompatibleBitmap = modGdi32.NewProc("CreateCompatibleBitmap")
procCreateCompatibleDC = modGdi32.NewProc("CreateCompatibleDC")
procCreateDIBSection = modGdi32.NewProc("CreateDIBSection")
procDeleteDC = modGdi32.NewProc("DeleteDC")
procDeleteObject = modGdi32.NewProc("DeleteObject")
procGetDeviceCaps = modGdi32.NewProc("GetDeviceCaps")
procSelectObject = modGdi32.NewProc("SelectObject")
modShcore = syscall.NewLazyDLL("Shcore.dll")
procSetProcessDpiAwareness = modShcore.NewProc("SetProcessDpiAwareness")
)
const (
// GetDeviceCaps constants from Wingdi.h
deviceCaps_HORZRES = 8
deviceCaps_VERTRES = 10
deviceCaps_LOGPIXELSX = 88
deviceCaps_LOGPIXELSY = 90
// BitBlt constants
bitBlt_SRCCOPY = 0x00CC0020
)
// Windows RECT structure
type win_RECT struct {
Left, Top, Right, Bottom int32
}
// http://msdn.microsoft.com/en-us/library/windows/desktop/dd183375.aspx
type win_BITMAPINFO struct {
BmiHeader win_BITMAPINFOHEADER
BmiColors *win_RGBQUAD
}
// http://msdn.microsoft.com/en-us/library/windows/desktop/dd183376.aspx
type win_BITMAPINFOHEADER struct {
BiSize uint32
BiWidth int32
BiHeight int32
BiPlanes uint16
BiBitCount uint16
BiCompression uint32
BiSizeImage uint32
BiXPelsPerMeter int32
BiYPelsPerMeter int32
BiClrUsed uint32
BiClrImportant uint32
}
// http://msdn.microsoft.com/en-us/library/windows/desktop/dd162938.aspx
type win_RGBQUAD struct {
RgbBlue byte
RgbGreen byte
RgbRed byte
RgbReserved byte
}
// findWindow finds the handle to the window.
func findWindow() (syscall.Handle, error) {
var handle syscall.Handle
// First look for the normal window
ret, _, _ := procFindWindow.Call(
0, uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("AppName"))))
if ret == 0 {
return handle, fmt.Errorf("App not found. Is it running?")
}
handle = syscall.Handle(ret)
return handle, nil
}
// windowRect gets the dimensions for a Window handle.
func windowRect(hwnd syscall.Handle) (image.Rectangle, error) {
var rect win_RECT
ret, _, err := procGetClientRect.Call(uintptr(hwnd), uintptr(unsafe.Pointer(&rect)))
if ret == 0 {
return image.Rectangle{}, fmt.Errorf("Error getting window dimensions: %s", err)
}
return image.Rect(0, 0, int(rect.Right), int(rect.Bottom)), nil
}
// captureWindow captures the desired area from a Window and returns an image.
func captureWindow(handle syscall.Handle, rect image.Rectangle) (image.Image, error) {
// Get the device context for screenshotting
dcSrc, _, err := procGetDC.Call(uintptr(handle))
if dcSrc == 0 {
return nil, fmt.Errorf("Error preparing screen capture: %s", err)
}
defer procReleaseDC.Call(0, dcSrc)
// Grab a compatible DC for drawing
dcDst, _, err := procCreateCompatibleDC.Call(dcSrc)
if dcDst == 0 {
return nil, fmt.Errorf("Error creating DC for drawing: %s", err)
}
defer procDeleteDC.Call(dcDst)
// Determine the width/height of our capture
width := rect.Dx();
height := rect.Dy();
// Get the bitmap we're going to draw onto
var bitmapInfo win_BITMAPINFO
bitmapInfo.BmiHeader = win_BITMAPINFOHEADER{
BiSize: uint32(reflect.TypeOf(bitmapInfo.BmiHeader).Size()),
BiWidth: int32(width),
BiHeight: int32(height),
BiPlanes: 1,
BiBitCount: 32,
BiCompression: 0, // BI_RGB
}
bitmapData := unsafe.Pointer(uintptr(0))
bitmap, _, err := procCreateDIBSection.Call(
dcDst,
uintptr(unsafe.Pointer(&bitmapInfo)),
0,
uintptr(unsafe.Pointer(&bitmapData)), 0, 0)
if bitmap == 0 {
return nil, fmt.Errorf("Error creating bitmap for screen capture: %s", err)
}
defer procDeleteObject.Call(bitmap)
// Select the object and paint it
procSelectObject.Call(dcDst, bitmap)
ret, _, err := procBitBlt.Call(
dcDst, 0, 0, uintptr(width), uintptr(height),
dcSrc, uintptr(rect.Min.X), uintptr(rect.Min.Y), bitBlt_SRCCOPY)
if ret == 0 {
return nil, fmt.Errorf("Error capturing screen: %s", err)
}
// Convert the bitmap to an image.Image. We first start by directly
// creating a slice. This is unsafe but we know the underlying structure
// directly.
var slice []byte
sliceHdr := (*reflect.SliceHeader)(unsafe.Pointer(&slice))
sliceHdr.Data = uintptr(bitmapData)
sliceHdr.Len = width * height * 4
sliceHdr.Cap = sliceHdr.Len
// Using the raw data, grab the RGBA data and transform it into an image.RGBA
imageBytes := make([]byte, len(slice))
for i := 0; i < len(imageBytes); i += 4 {
imageBytes[i], imageBytes[i+2], imageBytes[i+1], imageBytes[i+3] = slice[i+2], slice[i], slice[i+1], slice[i+3]
}
// The window gets screenshotted upside down and I don't know why.
// Flip it.
img := &image.RGBA{imageBytes, 4 * width, image.Rect(0, 0, width, height)}
dst := image.NewRGBA(img.Bounds())
gift.New(gift.FlipVertical()).Draw(dst, img)
return dst, nil
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment