From bfe300c3dc3a085400886fbd06b6e3c7d9939877 Mon Sep 17 00:00:00 2001
From: Bolke de Bruin <bolke@xs4all.nl>
Date: Tue, 28 Jul 2020 08:59:19 +0200
Subject: [PATCH] Add tests WIP

---
 protocol/client.go       | 46 ++++++++++++++++++++++++++++++++++++++++
 protocol/handler.go      |  2 +-
 protocol/handler_test.go | 44 ++++++++++++++++++++++++++++++++++++++
 protocol/rdpgw_test.go   |  2 ++
 4 files changed, 93 insertions(+), 1 deletion(-)
 create mode 100644 protocol/client.go
 create mode 100644 protocol/handler_test.go
 create mode 100644 protocol/rdpgw_test.go

diff --git a/protocol/client.go b/protocol/client.go
new file mode 100644
index 0000000..a15d1cb
--- /dev/null
+++ b/protocol/client.go
@@ -0,0 +1,46 @@
+package protocol
+
+import (
+	"bytes"
+	"encoding/binary"
+)
+
+const (
+	MajorVersion = 0x0
+	MinorVersion = 0x0
+	Version      = 0x00
+)
+
+type ClientConfig struct {
+	SmartCardAuth bool
+	PAAToken	  string
+	NTLMAuth	  bool
+}
+
+func (c *ClientConfig) handshakeRequest() []byte {
+	var caps uint16
+
+	if c.SmartCardAuth {
+		caps = caps | HTTP_EXTENDED_AUTH_SC
+	}
+
+	if len(c.PAAToken) > 0 {
+		caps = caps | HTTP_EXTENDED_AUTH_PAA
+	}
+
+	if c.NTLMAuth {
+		caps = caps | HTTP_EXTENDED_AUTH_SSPI_NTLM
+	}
+
+	buf := new(bytes.Buffer)
+
+	binary.Write(buf, binary.LittleEndian, byte(MajorVersion))
+	binary.Write(buf, binary.LittleEndian, byte(MinorVersion))
+	binary.Write(buf, binary.LittleEndian, uint16(Version))
+
+	binary.Write(buf, binary.LittleEndian, uint16(caps))
+
+	return createPacket(PKT_TYPE_HANDSHAKE_REQUEST, buf.Bytes())
+}
+
+func (c *ClientConfig) readServerHandshakeResponse(data []byte) ()
diff --git a/protocol/handler.go b/protocol/handler.go
index 4330c56..d4fdb16 100644
--- a/protocol/handler.go
+++ b/protocol/handler.go
@@ -225,7 +225,7 @@ func (h *Handler) ReadMessage() (pt int, n int, msg []byte, err error) {
 func (h *Handler) handshakeResponse(major byte, minor byte) []byte {
 	var caps uint16
 	if h.SmartCardAuth {
-		caps = caps | HTTP_EXTENDED_AUTH_PAA
+		caps = caps | HTTP_EXTENDED_AUTH_SC
 	}
 	if h.TokenAuth {
 		caps = caps | HTTP_EXTENDED_AUTH_PAA
diff --git a/protocol/handler_test.go b/protocol/handler_test.go
new file mode 100644
index 0000000..c938cad
--- /dev/null
+++ b/protocol/handler_test.go
@@ -0,0 +1,44 @@
+package protocol
+
+import (
+	"log"
+	"testing"
+)
+
+const (
+	HeaderLen = 8
+	HandshakeRequestLen = HeaderLen + 6
+)
+
+func TestHandshake(t *testing.T) {
+	client := ClientConfig{
+		PAAToken: "abab",
+	}
+
+	data := client.handshakeRequest()
+	pt, size, pkt, err := readHeader(data)
+
+	if pt != PKT_TYPE_HANDSHAKE_REQUEST {
+		t.Fatalf("readHeader failed, expected packet type %d got %d", PKT_TYPE_HANDSHAKE_REQUEST, pt)
+	}
+
+	if size != HandshakeRequestLen {
+		t.Fatalf("readHeader failed, expected size %d, got %d", HandshakeRequestLen, size)
+	}
+
+	if err != nil {
+		t.Fatalf("readHeader failed got error %s", err)
+	}
+
+	log.Printf("pkt: %x", pkt)
+
+	major, minor, version, extAuth := readHandshake(pkt)
+	if major != MajorVersion || minor != MinorVersion || version != Version {
+		t.Fatalf("readHandshake failed got version %d.%d protocol %d, expected %d.%d protocol %d",
+			major, minor, version, MajorVersion, MinorVersion, Version)
+	}
+
+	if !((extAuth & HTTP_EXTENDED_AUTH_PAA) == HTTP_EXTENDED_AUTH_PAA) {
+		t.Fatalf("readHandshake failed got ext auth %d, expected %d", extAuth, extAuth | HTTP_EXTENDED_AUTH_PAA)
+	}
+}
diff --git a/protocol/rdpgw_test.go b/protocol/rdpgw_test.go
new file mode 100644
index 0000000..dc6bdf6
--- /dev/null
+++ b/protocol/rdpgw_test.go
@@ -0,0 +1,2 @@
+package protocol
+
-- 
GitLab