ShotRDP/client/shotClient.go

package client

import (
	"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"
	"context"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"
)

type RdpClient 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
	screenWidth  int
	screenHeight int
}

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

var (
	timerCycle     = 500 * time.Millisecond
	timer          = time.NewTimer(timerCycle)
	bitmapRevOne   = sync.Once{}
	bitmapRevList  = make([]pdu.BitmapData, 0)
	bitmapRevCount = 0
	bitmapRevSig   = make(chan struct{})
	modifyMux      = sync.Mutex{}
	positionMap    = make(map[Position]*common.Bitmap)
)

func Newclient(host string, logLevel glog.LEVEL, screenWidth, screenHeight int) *RdpClient {
	_logger := log.New(os.Stdout, "", 0)

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

	return &RdpClient{
		Host:         host,
		screenWidth:  screenWidth,
		screenHeight: screenHeight,
	}
}

func (client *RdpClient) Login(domain, username, password string) ([]byte, error) {
	conn, err := net.DialTimeout("tcp", client.Host, 3*time.Second)
	if err != nil {
		return nil, err
	}

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

	// 建立X224连接
	err = client.x224.Connect()
	if err != nil {
		return nil, err
	}

	// 设置事件处理
	client.setEventHandler()
	go func() {
		for {
			select {
			case <-timer.C:
				modifyMux.Lock()
				// 检查是否有新元素添加
				if len(bitmapRevList) == bitmapRevCount {
					bitmapRevOne.Do(func() {
						close(bitmapRevSig)
					})
				}
				modifyMux.Unlock()
			}
		}
	}()

	// 合并登录界面位图
	wg := &sync.WaitGroup{}
	wg.Add(1)

	var buff []byte
	surfaceflingerCtx, _ := context.WithCancel(context.Background())
	go func(ctx context.Context, c *RdpClient) {
		outputName := strings.ReplaceAll(client.Host, ":", "_")
		var bitmapList = make([]*common.Bitmap, 0)

		// TODO: 寻找合适的截图时机
		<-bitmapRevSig

		modifyMux.Lock()
		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)
		}
		modifyMux.Unlock()

		buff = common.DrawFullImage(outputName, bitmapList)
		wg.Done()
	}(surfaceflingerCtx, client)

	wg.Wait()
	return buff, nil
}

// initProtocolStack 初始化RDP协议栈
func (client *RdpClient) 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(client.screenWidth), uint16(client.screenHeight))

	// 设置认证信息
	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 *RdpClient) setEventHandler() {
	client.pdu.On("ready", func() {

	})

	client.pdu.On("success", func() {

	})

	client.pdu.On("close", func() {

	})

	client.pdu.On("done", func() {

	})

	client.pdu.On("error", func(err error) {
		client.pdu.Emit("done")
	})

	client.pdu.On("bitmap", func(rectangles []pdu.BitmapData) {
		modifyMux.Lock()
		defer modifyMux.Unlock()

		// 重置定时器
		if !timer.Stop() {
			<-timer.C
		}
		timer.Reset(timerCycle)

		oldCount := len(bitmapRevList)
		bitmapRevList = append(bitmapRevList, rectangles...)
		newCount := len(bitmapRevList)

		if newCount > oldCount {
			bitmapRevCount = newCount
		}

		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
		}
	})
}

/*
修改部分：grdp/protocol/x224/x224.go#New()
去除requestedProtocol中的PROTOCOL_HYBRID，避免使用NLA
*/

func RealGetScreen(host string, screenWidth, screenHeight int) ([]byte, error) {
	client := Newclient(host, glog.INFO, screenWidth, screenHeight)
	return client.Login("", "", "")
}