ShotRDP/common/bitmap.go

package common

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"os"
	"sort"

	"ShotRDP/grdp/core"
	"ShotRDP/grdp/glog"
	"ShotRDP/grdp/protocol/pdu"
)

type Bitmap struct {
	// 位图坐标信息
	DestTop    int
	DestLeft   int
	DestRight  int
	DestBottom int

	// 位图尺寸信息
	Width  int
	Height int

	Data         []byte
	IsCompress   bool
	BitsPerPixel int
}

func BitmapDecompress(bitmap *pdu.BitmapData) []byte {
	return core.Decompress(bitmap.BitmapDataStream, int(bitmap.Width), int(bitmap.Height), Bpp(bitmap.BitsPerPixel))
}

// Bpp 计算每像素所需的字节数
func Bpp(BitsPerPixel uint16) (pixel int) {
	switch BitsPerPixel {
	case 15:
		pixel = 1

	case 16:
		pixel = 2

	case 24:
		pixel = 3

	case 32:
		pixel = 4

	default:
		glog.Error("invalid BitsPerPixel")

	}
	return
}

func ToRGBA(pixel int, i int, data []byte) (r, g, b, a uint8) {
	a = 255
	switch pixel {
	case 1:
		rgb555 := core.Uint16BE(data[i], data[i+1])
		r, g, b = core.RGB555ToRGB(rgb555)
	case 2:
		rgb565 := core.Uint16BE(data[i], data[i+1])
		r, g, b = core.RGB565ToRGB(rgb565)
	case 3, 4:
		fallthrough
	default:
		r, g, b = data[i+2], data[i+1], data[i]
	}

	return
}

func DrawOneBitmap(bitmap *Bitmap, name string) {
	finalImg := image.NewRGBA(image.Rect(0, 0, bitmap.Width, bitmap.Height))

	var (
		i, p       int
		r, g, b, a uint8
	)

	i = 0
	p = bitmap.BitsPerPixel
	m := image.NewRGBA(image.Rect(0, 0, bitmap.Width, bitmap.Height))
	for y := 0; y < bitmap.Height; y++ {
		for x := 0; x < bitmap.Width; x++ {
			r, g, b, a = ToRGBA(p, i, bitmap.Data)
			c := color.RGBA{R: r, G: g, B: b, A: a}
			m.Set(x, y, c)

			i += p
		}
	}

	draw.Draw(finalImg, finalImg.Bounds(), m, m.Bounds().Min, draw.Src)

	outFile, _ := os.Create(fmt.Sprintf("%s.png", name))
	defer outFile.Close()

	// 保存图像为PNG文件
	_ = png.Encode(outFile, finalImg)
}

func DrawFullImage(outputName string, bitmapList []*Bitmap) []byte {
	// 对位图按 DestTop 和 DestLeft 排序
	sort.Slice(bitmapList, func(i, j int) bool {
		if bitmapList[i].DestTop == bitmapList[j].DestTop {
			return bitmapList[i].DestLeft < bitmapList[j].DestLeft
		}
		return bitmapList[i].DestTop < bitmapList[j].DestTop
	})

	// 计算合并后的总尺寸
	var totalWidth, totalHeight int
	for _, bmp := range bitmapList {
		if bmp.DestLeft+bmp.Width > totalWidth {
			totalWidth = bmp.DestLeft + bmp.Width
		}
		if bmp.DestTop+bmp.Height > totalHeight {
			totalHeight = bmp.DestTop + bmp.Height
		}
	}

	// 创建一个空白的画布，大小为合并后的总尺寸
	finalImg := image.NewRGBA(image.Rect(0, 0, totalWidth, totalHeight))

	var (
		i, p       int
		r, g, b, a uint8
	)

	for _, bmp := range bitmapList {
		i = 0
		p = bmp.BitsPerPixel
		m := image.NewRGBA(image.Rect(0, 0, bmp.Width, bmp.Height))
		for y := 0; y < bmp.Height; y++ {
			for x := 0; x < bmp.Width; x++ {
				r, g, b, a = ToRGBA(p, i, bmp.Data)
				c := color.RGBA{R: r, G: g, B: b, A: a}
				m.Set(x, y, c)

				i += p
			}
		}

		draw.Draw(finalImg, finalImg.Bounds().Add(image.Pt(bmp.DestLeft, bmp.DestTop)), m, m.Bounds().Min, draw.Src)
	}

	// 创建保存PNG文件
	//outFile, _ := os.Create(fmt.Sprintf("%s.png", outputName))
	//defer outFile.Close()
	//_ = png.Encode(outFile, finalImg)

	var buff bytes.Buffer
	err := png.Encode(&buff, finalImg)
	if err != nil {
		return nil
	}

	return buff.Bytes()
}