Add screenshot example

encoding.go

@@ -1,6 +1,8 @@
 package vnc
 
 import (
+	"bytes"
+	"compress/zlib"
 	"encoding/binary"
 	"io"
 )
@@ -67,3 +69,92 @@ func (*RawEncoding) Read(c *ClientConn, rect *Rectangle, r io.Reader) (Encoding,
 
 	return &RawEncoding{colors}, nil
 }
+
+// ZlibEncoding is raw pixel data sent by the server compressed by Zlib.
+//
+// A single Zlib stream is created. There is only a single header for a framebuffer request response.
+type ZlibEncoding struct {
+	Colors  []Color
+}
+
+func (*ZlibEncoding) Type() int32 {
+	return 6
+}
+
+func (ze *ZlibEncoding) Read(c *ClientConn, rect *Rectangle, r io.Reader) (Encoding, error) {
+	bytesPerPixel := c.PixelFormat.BPP / 8
+	pixelBytes := make([]uint8, bytesPerPixel)
+
+	var byteOrder binary.ByteOrder = binary.LittleEndian
+	if c.PixelFormat.BigEndian {
+		byteOrder = binary.BigEndian
+	}
+
+	// Format
+	// 4 bytes        | uint32 | length
+	// 'length' bytes | []byte | zlibData
+
+	// Read zlib length
+	var zlibLength uint32
+	err := binary.Read(r, binary.BigEndian, &zlibLength)
+	if err != nil {
+		return nil, err
+	}
+
+	// Read all compressed data
+	zBytes := make([]byte, zlibLength)
+	if _, err := io.ReadFull(r, zBytes); err != nil {
+		return nil, err
+	}
+
+	// Create a reader for the buffer
+	// This is needed to avoid 'zlib missing header'
+	zReader, err := zlib.NewReader(bytes.NewReader(zBytes))
+	if err != nil {
+		return nil, err
+	}
+	defer zReader.Close()
+
+	// Create buffer for bytes
+	sizeToRead := int(rect.Height) * int(rect.Width) * int(bytesPerPixel)
+	colorBytes := make([]byte, sizeToRead)
+
+	// Read all data from zlib stream
+	read, err := io.ReadFull(zReader, colorBytes)
+	if read != sizeToRead || err != nil {
+		return nil, err
+	}
+
+	// Create buffer for raw encoding
+	colorReader := bytes.NewReader(colorBytes)
+
+	colors := make([]Color, int(rect.Height)*int(rect.Width))
+
+	for y := uint16(0); y < rect.Height; y++ {
+		for x := uint16(0); x < rect.Width; x++ {
+			if _, err := io.ReadFull(colorReader, pixelBytes); err != nil {
+				return nil, err
+			}
+
+			var rawPixel uint32
+			if c.PixelFormat.BPP == 8 {
+				rawPixel = uint32(pixelBytes[0])
+			} else if c.PixelFormat.BPP == 16 {
+				rawPixel = uint32(byteOrder.Uint16(pixelBytes))
+			} else if c.PixelFormat.BPP == 32 {
+				rawPixel = byteOrder.Uint32(pixelBytes)
+			}
+
+			color := &colors[int(y)*int(rect.Width)+int(x)]
+			if c.PixelFormat.TrueColor {
+				color.R = uint16((rawPixel >> c.PixelFormat.RedShift) & uint32(c.PixelFormat.RedMax))
+				color.G = uint16((rawPixel >> c.PixelFormat.GreenShift) & uint32(c.PixelFormat.GreenMax))
+				color.B = uint16((rawPixel >> c.PixelFormat.BlueShift) & uint32(c.PixelFormat.BlueMax))
+			} else {
+				*color = c.ColorMap[rawPixel]
+			}
+		}
+	}
+
+	return &ZlibEncoding{Colors: colors}, nil
+}

examples/screenshot.go

@@ -0,0 +1,103 @@
+package main
+
+import (
+	"bytes"
+	"image"
+	"image/color"
+	"image/png"
+	"io"
+	"io/ioutil"
+	"net"
+
+	"github.com/tinti/go-vnc"
+)
+
+func TakeScreenshot(address, password string) (*image.RGBA, error) {
+	nc, err := net.Dial("tcp", address)
+	if err != nil {
+		return nil, err
+	}
+	defer nc.Close()
+
+	serverMessageChannel := make(chan vnc.ServerMessage)
+
+	vncClient, err := vnc.Client(nc, &vnc.ClientConfig{
+		Auth: []vnc.ClientAuth{
+			&vnc.PasswordAuth{Password: password},
+		},
+		ServerMessages: []vnc.ServerMessage{
+			&vnc.FramebufferUpdateMessage{},
+		},
+		ServerMessageCh: serverMessageChannel,
+	})
+	if err != nil {
+		return nil, err
+	}
+	defer vncClient.Close()
+
+	err = vncClient.SetEncodings([]vnc.Encoding{
+		&vnc.ZlibEncoding{},
+		&vnc.RawEncoding{},
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	err = vncClient.FramebufferUpdateRequest(false, 0, 0,
+		vncClient.FrameBufferWidth, vncClient.FrameBufferHeight)
+	if err != nil {
+		return nil, err
+	}
+
+	serverMessage := <-serverMessageChannel
+
+	rects := serverMessage.(*vnc.FramebufferUpdateMessage).Rectangles
+	if len(rects) == 0 {
+		panic("vnc: framebuffer rects length")
+	}
+
+	w := int(vncClient.FrameBufferWidth)
+	h := int(vncClient.FrameBufferHeight)
+	img := image.NewRGBA(image.Rect(0, 0, w, h))
+
+	for _, rect := range rects {
+		switch enc := rect.Enc.(type) {
+		case *vnc.RawEncoding:
+			for i, c := range enc.Colors {
+				x, y := i%w, i/w
+				r, g, b := uint8(c.R), uint8(c.G), uint8(c.B)
+
+				img.Set(int(rect.X)+x, int(rect.Y)+y, color.RGBA{r, g, b, 255})
+			}
+		case *vnc.ZlibEncoding:
+			for i, c := range enc.Colors {
+				x, y := i%w, i/w
+				r, g, b := uint8(c.R), uint8(c.G), uint8(c.B)
+
+				img.Set(int(rect.X)+x, int(rect.Y)+y, color.RGBA{r, g, b, 255})
+			}
+		default:
+			panic("vnc: unkown encoding")
+		}
+	}
+
+
+	return img, nil
+}
+
+func main() {
+	img, err := TakeScreenshot("server:port", "password")
+	if err != nil {
+		panic(err)
+	}
+
+	data := bytes.Buffer{}
+	pngEncoder := png.Encoder{CompressionLevel: png.NoCompression}
+
+	err = pngEncoder.Encode(io.Writer(&data), img)
+	if err != nil {
+		panic(err)
+	}
+
+	ioutil.WriteFile("screenshot.png", data.Bytes(), 0600)
+}