diff --git a/encoding.go b/encoding.go index 7198595..2c4b2e0 100644 --- a/encoding.go +++ b/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 +} diff --git a/examples/screenshot.go b/examples/screenshot.go new file mode 100644 index 0000000..c1e47e0 --- /dev/null +++ b/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) +}