Improvements for package

client/client.go

@@ -1,139 +1,31 @@
-// client.go
 package client
 
 import (
 	"log"
 	"os"
 
 	"github.com/tomatome/grdp/glog"
-	"github.com/tomatome/grdp/protocol/pdu"
-	"github.com/tomatome/grdp/protocol/rfb"
 )
 
-const (
-	CLIP_OFF = 0
-	CLIP_IN  = 0x1
-	CLIP_OUT = 0x2
-)
-
-const (
-	TC_RDP = 0
-	TC_VNC = 1
-)
+func init() {
+	logger := log.New(os.Stdout, "", 0)
+	glog.SetLogger(logger)
+}
 
-type Control interface {
+type RemoteControl interface {
 	Login(host, user, passwd string, width, height int) error
+	Connect() error
 	KeyUp(sc int, name string)
 	KeyDown(sc int, name string)
 	MouseMove(x, y int)
 	MouseWheel(scroll, x, y int)
 	MouseUp(button int, x, y int)
 	MouseDown(button int, x, y int)
 	On(event string, msg interface{})
+	OnBitmap(handler func([]Bitmap))
 	Close()
 }
 
-type Client struct {
-	host    string
-	user    string
-	passwd  string
-	ctl     Control
-	tc      int
-	setting *Setting
-}
-
-func init() {
-	logger := log.New(os.Stdout, "", 0)
-	glog.SetLogger(logger)
-}
-func NewClient(host, user, passwd string, t int, s *Setting) *Client {
-	if s == nil {
-		s = NewSetting()
-	}
-	c := &Client{
-		host:    host,
-		user:    user,
-		passwd:  passwd,
-		tc:      t,
-		setting: s,
-	}
-
-	switch t {
-	case TC_VNC:
-		c.ctl = newVncClient(s)
-	default:
-		c.ctl = newRdpClient(s)
-	}
-
-	s.SetLogLevel()
-	return c
-}
-
-func (c *Client) Login() error {
-	return c.ctl.Login(c.host, c.user, c.passwd, c.setting.Width, c.setting.Height)
-}
-
-func (c *Client) KeyUp(sc int, name string) {
-	c.ctl.KeyUp(sc, name)
-}
-func (c *Client) KeyDown(sc int, name string) {
-	c.ctl.KeyDown(sc, name)
-}
-func (c *Client) MouseMove(x, y int) {
-	c.ctl.MouseMove(x, y)
-}
-func (c *Client) MouseWheel(scroll, x, y int) {
-	c.ctl.MouseWheel(scroll, x, y)
-}
-func (c *Client) MouseUp(button, x, y int) {
-	c.ctl.MouseUp(button, x, y)
-}
-func (c *Client) MouseDown(button, x, y int) {
-	c.ctl.MouseDown(button, x, y)
-}
-func (c *Client) OnError(f func(e error)) {
-	c.ctl.On("error", f)
-}
-func (c *Client) OnClose(f func()) {
-	c.ctl.On("close", f)
-}
-func (c *Client) OnSuccess(f func()) {
-	c.ctl.On("success", f)
-}
-func (c *Client) OnReady(f func()) {
-	c.ctl.On("ready", f)
-}
-func (c *Client) OnBitmap(f func([]Bitmap)) {
-	f1 := func(data interface{}) {
-		bs := make([]Bitmap, 0, 50)
-		if c.tc == TC_VNC {
-			br := data.(*rfb.BitRect)
-			for _, v := range br.Rects {
-				b := Bitmap{int(v.Rect.X), int(v.Rect.Y), int(v.Rect.X + v.Rect.Width), int(v.Rect.Y + v.Rect.Height),
-					int(v.Rect.Width), int(v.Rect.Height),
-					Bpp(uint16(br.Pf.BitsPerPixel)), false, v.Data}
-				bs = append(bs, b)
-			}
-		} else {
-			for _, v := range data.([]pdu.BitmapData) {
-				IsCompress := v.IsCompress()
-				stream := v.BitmapDataStream
-				if IsCompress {
-					stream = bitmapDecompress(&v)
-					IsCompress = false
-				}
-
-				b := Bitmap{int(v.DestLeft), int(v.DestTop), int(v.DestRight), int(v.DestBottom),
-					int(v.Width), int(v.Height), Bpp(v.BitsPerPixel), IsCompress, stream}
-				bs = append(bs, b)
-			}
-		}
-		f(bs)
-	}
-
-	c.ctl.On("update", f1)
-}
-
 type Bitmap struct {
 	DestLeft     int    `json:"destLeft"`
 	DestTop      int    `json:"destTop"`
@@ -165,24 +57,3 @@ func Bpp(bp uint16) (pixel int) {
 	}
 	return
 }
-
-type Setting struct {
-	Width    int
-	Height   int
-	Protocol string
-	LogLevel glog.LEVEL
-}
-
-func NewSetting() *Setting {
-	return &Setting{
-		Width:    1024,
-		Height:   768,
-		LogLevel: glog.INFO,
-	}
-}
-func (s *Setting) SetLogLevel() {
-	glog.SetLevel(s.LogLevel)
-}
-
-func (s *Setting) SetRequestedProtocol(p uint32) {}
-func (s *Setting) SetClipboard(c int)            {}

client/rdp.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/tomatome/grdp/core"
+	"github.com/tomatome/grdp/glog"
 	"github.com/tomatome/grdp/plugin"
 	"github.com/tomatome/grdp/protocol/nla"
 	"github.com/tomatome/grdp/protocol/pdu"
@@ -16,22 +17,69 @@ import (
 	"github.com/tomatome/grdp/protocol/x224"
 )
 
+const (
+	RdpProtocolRDP      = "rdp"
+	RdpProtocolSSL      = "ssl"
+	RdpProtocolHybrid   = "hybrid"
+	RdpProtocolHybridEx = "hybrid_ex"
+)
+
 type RdpClient struct {
-	tpkt     *tpkt.TPKT
-	x224     *x224.X224
-	mcs      *t125.MCSClient
-	sec      *sec.Client
-	pdu      *pdu.Client
-	channels *plugin.Channels
+	// connection
+	host     string
+	port     string
+	domain   string
+	username string
+	password string
+
+	// window
+	windowHeight uint16
+	windowWidth  uint16
+
+	// protocols
+	tpkt              *tpkt.TPKT
+	x224              *x224.X224
+	mcs               *t125.MCSClient
+	sec               *sec.Client
+	pdu               *pdu.Client
+	channels          *plugin.Channels
+	requestedProtocol string
+}
+
+func (c *RdpClient) setRequestedProtocol() {
+	switch c.requestedProtocol {
+	case RdpProtocolRDP:
+		c.x224.SetRequestedProtocol(x224.PROTOCOL_RDP)
+	case RdpProtocolSSL:
+		c.x224.SetRequestedProtocol(x224.PROTOCOL_SSL)
+	case RdpProtocolHybrid:
+		c.x224.SetRequestedProtocol(x224.PROTOCOL_HYBRID)
+	case RdpProtocolHybridEx:
+		c.x224.SetRequestedProtocol(x224.PROTOCOL_HYBRID_EX)
+	}
+}
+
+func (c *RdpClient) WithWindowSize(height, width uint16) {
+	c.windowWidth = width
+	c.windowHeight = height
 }
 
-func newRdpClient(s *Setting) *RdpClient {
-	return &RdpClient{}
+func (c *RdpClient) WithRequestedProtocol(protocol string) {
+	c.requestedProtocol = protocol
+}
+
+func NewRdpClient(host, port, domain, username, password string) *RdpClient {
+	return &RdpClient{
+		host: host, port: port, domain: domain,
+		username: username, password: password,
+		windowHeight: 600, windowWidth: 800,
+	}
 }
 
 func bitmapDecompress(bitmap *pdu.BitmapData) []byte {
 	return core.Decompress(bitmap.BitmapDataStream, int(bitmap.Width), int(bitmap.Height), Bpp(bitmap.BitsPerPixel))
 }
+
 func split(user string) (domain string, uname string) {
 	if strings.Index(user, "\\") != -1 {
 		t := strings.Split(user, "\\")
@@ -46,12 +94,49 @@ func split(user string) (domain string, uname string) {
 	}
 	return
 }
+
+func (c *RdpClient) Connect() error {
+	addr := net.JoinHostPort(c.host, c.port)
+	conn, err := net.DialTimeout("tcp", addr, 10*time.Second)
+	if err != nil {
+		return err
+	}
+
+	domain := c.domain
+	username := c.username
+	if c.domain == "" {
+		domain, username = split(username)
+	}
+
+	glog.Infof("Connect to %v", addr)
+
+	c.tpkt = tpkt.New(core.NewSocketLayer(conn), nla.NewNTLMv2(domain, username, c.password))
+	c.x224 = x224.New(c.tpkt)
+	c.mcs = t125.NewMCSClient(c.x224)
+	c.sec = sec.NewClient(c.mcs)
+	c.pdu = pdu.NewClient(c.sec)
+	c.channels = plugin.NewChannels(c.sec)
+
+	c.mcs.SetClientCoreData(c.windowWidth, c.windowHeight)
+
+	c.sec.SetUser(username)
+	c.sec.SetPwd(c.password)
+	c.sec.SetDomain(domain)
+
+	c.tpkt.SetFastPathListener(c.sec)
+	c.sec.SetFastPathListener(c.pdu)
+	c.sec.SetChannelSender(c.mcs)
+	c.channels.SetChannelSender(c.sec)
+
+	c.setRequestedProtocol()
+	return c.x224.Connect()
+}
+
 func (c *RdpClient) Login(host, user, pwd string, width, height int) error {
 	conn, err := net.DialTimeout("tcp", host, 3*time.Second)
 	if err != nil {
 		return fmt.Errorf("[dial err] %v", err)
 	}
-
 	domain, user := split(user)
 	c.tpkt = tpkt.New(core.NewSocketLayer(conn), nla.NewNTLMv2(domain, user, pwd))
 	c.x224 = x224.New(c.tpkt)
@@ -80,15 +165,18 @@ func (c *RdpClient) Login(host, user, pwd string, width, height int) error {
 	}
 	return nil
 }
+
 func (c *RdpClient) On(event string, f interface{}) {
 	c.pdu.On(event, f)
 }
+
 func (c *RdpClient) KeyUp(sc int, name string) {
 	p := &pdu.ScancodeKeyEvent{}
 	p.KeyCode = uint16(sc)
 	p.KeyboardFlags |= pdu.KBDFLAGS_RELEASE
 	c.pdu.SendInputEvents(pdu.INPUT_EVENT_SCANCODE, []pdu.InputEventsInterface{p})
 }
+
 func (c *RdpClient) KeyDown(sc int, name string) {
 	p := &pdu.ScancodeKeyEvent{}
 	p.KeyCode = uint16(sc)
@@ -129,6 +217,7 @@ func (c *RdpClient) MouseUp(button int, x, y int) {
 	p.YPos = uint16(y)
 	c.pdu.SendInputEvents(pdu.INPUT_EVENT_MOUSE, []pdu.InputEventsInterface{p})
 }
+
 func (c *RdpClient) MouseDown(button int, x, y int) {
 	p := &pdu.PointerEvent{}
 
@@ -149,8 +238,29 @@ func (c *RdpClient) MouseDown(button int, x, y int) {
 	p.YPos = uint16(y)
 	c.pdu.SendInputEvents(pdu.INPUT_EVENT_MOUSE, []pdu.InputEventsInterface{p})
 }
+
 func (c *RdpClient) Close() {
 	if c != nil && c.tpkt != nil {
 		c.tpkt.Close()
 	}
 }
+
+func (c *RdpClient) OnBitmap(handler func([]Bitmap)) {
+	bitmapsFunc := func(data interface{}) {
+		bs := make([]Bitmap, 0, 50)
+		for _, v := range data.([]pdu.BitmapData) {
+			IsCompress := v.IsCompress()
+			stream := v.BitmapDataStream
+			if IsCompress {
+				stream = bitmapDecompress(&v)
+				IsCompress = false
+			}
+
+			b := Bitmap{int(v.DestLeft), int(v.DestTop), int(v.DestRight), int(v.DestBottom),
+				int(v.Width), int(v.Height), Bpp(v.BitsPerPixel), IsCompress, stream}
+			bs = append(bs, b)
+		}
+		handler(bs)
+	}
+	c.On("update", bitmapsFunc)
+}