ShotRDP/client.go

package main

import (
	"context"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"ShotRDP/common"
	"ShotRDP/grdp/core"
	"ShotRDP/grdp/glog"
	"ShotRDP/grdp/plugin"
	"ShotRDP/grdp/protocol/nla"
	"ShotRDP/grdp/protocol/pdu"
	"ShotRDP/grdp/protocol/sec"
	"ShotRDP/grdp/protocol/t125"
	"ShotRDP/grdp/protocol/tpkt"
	"ShotRDP/grdp/protocol/x224"
)

type Client struct {
	Host string          // 服务地址(ip:port)
	tpkt *tpkt.TPKT      // TPKT协议层
	x224 *x224.X224      // X224协议层
	mcs  *t125.MCSClient // MCS协议层
	sec  *sec.Client     // 安全层
	pdu  *pdu.Client     // PDU协议层

	channels *plugin.Channels
}

type Position struct {
	Top    int
	Left   int
	Right  int
	Bottom int
}

var (
	modifyMux   = sync.Mutex{}
	positionMap = make(map[Position]*common.Bitmap)
)

func NewClient(host string, logLevel glog.LEVEL) *Client {
	_logger := log.New(os.Stdout, "", 0)

	glog.SetLogger(_logger)
	glog.SetLevel(logLevel)

	return &Client{Host: host}
}

func (client *Client) Login(domain, username, password string) error {
	conn, err := net.Dial("tcp", client.Host)
	if err != nil {
		glog.Errorf("TCP连接失败：%v", err)
		return err
	}

	// 初始化协议栈
	client.initProtocolStack(conn, domain, username, password)

	// 建立X224连接
	err = client.x224.Connect()
	if err != nil {
		glog.Errorf("建立X224连接失败：%v", err)
		return err
	}

	wg := &sync.WaitGroup{}
	wg.Add(1)

	// 设置事件处理
	client.setEventHandler(wg)

	wg.Wait()
	return nil
}

// initProtocolStack 初始化RDP协议栈
func (client *Client) initProtocolStack(conn net.Conn, domain, username, password string) {
	// 创建协议层实例
	client.tpkt = tpkt.New(core.NewSocketLayer(conn), nla.NewNTLMv2(domain, username, password))
	client.x224 = x224.New(client.tpkt)
	client.mcs = t125.NewMCSClient(client.x224)
	client.sec = sec.NewClient(client.mcs)
	client.pdu = pdu.NewClient(client.sec)

	// 配置桌面信息
	client.channels = plugin.NewChannels(client.sec)
	client.mcs.SetClientDesktop(uint16(1920), uint16(1080))

	// 设置认证信息
	client.sec.SetDomain(domain)
	client.sec.SetUsername(username)
	client.sec.SetPassword(password)

	// 配置协议层关联
	client.tpkt.SetFastPathListener(client.sec)
	client.sec.SetFastPathListener(client.pdu)

	client.sec.SetChannelSender(client.mcs)
}

// setEventHandler 设置PDU事件处理器
func (client *Client) setEventHandler(wg *sync.WaitGroup) {
	client.pdu.On("ready", func() {
		glog.Info("PDU连接就绪")

		// Test
		//glog.Infof("服务端核心信息: %+v", *client.sec.ServerCoreData())
		//glog.Infof("服务端安全信息: %+v", *client.sec.ServerSecurityData())
	})

	client.pdu.On("success", func() {
		glog.Info("PDU连接成功")
	})

	client.pdu.On("close", func() {
		glog.Info("PDU连接关闭")
	})

	client.pdu.On("done", func() {
		glog.Info("PDU处理完成")
	})

	client.pdu.On("error", func(err error) {
		glog.Errorf("PDU错误事件：%v", err)
		client.pdu.Emit("done")
	})

	client.pdu.On("bitmap", func(rectangles []pdu.BitmapData) {
		modifyMux.Lock()
		for _, rectangle := range rectangles {
			pos := Position{Left: int(rectangle.DestLeft / rectangle.Width), Top: int(rectangle.DestTop / rectangle.Height), Right: int(rectangle.DestRight / rectangle.Width), Bottom: int(rectangle.DestBottom / rectangle.Height)}
			isCompress := rectangle.IsCompress()
			data := rectangle.BitmapDataStream
			if isCompress {
				data = common.BitmapDecompress(&rectangle)
				isCompress = false
			}

			bitmap := &common.Bitmap{
				DestLeft:     int(rectangle.DestLeft),
				DestTop:      int(rectangle.DestTop),
				DestRight:    int(rectangle.DestRight),
				DestBottom:   int(rectangle.DestBottom),
				Width:        int(rectangle.Width),
				Height:       int(rectangle.Height),
				BitsPerPixel: common.Bpp(rectangle.BitsPerPixel),
				IsCompress:   isCompress,
				Data:         data,
			}
			positionMap[pos] = bitmap
		}
		modifyMux.Unlock()
	})
 
	// 合并登录界面位图
	surfaceflingerCtx, _ := context.WithCancel(context.Background())
	go func(ctx context.Context, c *Client) {
		outputName := strings.ReplaceAll(client.Host, ":", "_")

		for common.Opened(ctx) {
			// TODO: 寻找合适的截图时机
			time.Sleep(5 * time.Second)
			modifyMux.Lock()

			var bitmapList = make([]*common.Bitmap, 0)
			for _, bitmap := range positionMap {
				//common.DrawOneBitmap(bitmap, fmt.Sprintf("%d-%d-%d-%d.png", pos.Left, pos.Top, pos.Right, pos.Bottom))
				bitmapList = append(bitmapList, bitmap)
			}

			common.DrawFullImage(outputName, bitmapList)
			modifyMux.Unlock()

			wg.Done()
		}
	}(surfaceflingerCtx, client)
}