From 74d659d795f9ef680cfc75338eb50abe8d4771c9 Mon Sep 17 00:00:00 2001
From: justdan96 <justdan96@gmail.com>
Date: Sun, 15 Apr 2018 03:31:17 +0100
Subject: [PATCH] use SplitN for domain parsing, moved to separate function
 added battery of tests for username, domain, workstation and expected
 negotiate message bytes

---
 negotiator.go | 21 +++++++++++++-------
 nlmp_test.go  | 54 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/negotiator.go b/negotiator.go
index f61521e..72e668f 100644
--- a/negotiator.go
+++ b/negotiator.go
@@ -9,6 +9,18 @@ import (
 	"strings"
 )
 
+func GetDomain(user string) (string, string) {
+	// parse domain name from based on slashes in the input
+	domain := ""
+
+	if strings.Contains(user, "\\") {
+		ucomponents := strings.SplitN(user, "\\", 2)
+		domain = ucomponents[0]
+		user = ucomponents[1]
+	}
+	return user, domain
+}
+
 //Negotiator is a http.Roundtripper decorator that automatically
 //converts basic authentication to NTLM/Negotiate authentication when appropriate.
 type Negotiator struct{ http.RoundTripper }
@@ -77,14 +89,9 @@ func (l Negotiator) RoundTrip(req *http.Request) (res *http.Response, err error)
 			return nil, err
 		}
 
-		// parse domain name from username
+		// get domain from username
 		domain := ""
-
-		if strings.Contains(u, "\\") {
-			ucomponents := strings.Split(u, "\\")
-			domain = ucomponents[0]
-			u = ucomponents[1]
-		}
+		u, domain = GetDomain(u)
 
 		// send negotiate
 		negotiateMessage, err := NewNegotiateMessage(domain, "")
diff --git a/nlmp_test.go b/nlmp_test.go
index dea5aaf..b48af8d 100644
--- a/nlmp_test.go
+++ b/nlmp_test.go
@@ -10,8 +10,62 @@ import (
 var username = "user"
 var password = "SecREt01"
 var target = "DOMAIN"
+var domain = "MYDOMAIN"
+var workstation = "MYPC"
 var challenge = []byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef}
 
+func TestUsernameDomainWorkstation(t *testing.T) {
+	// taking a username and workstation as input, check that the username, domain, workstation
+	// and negotiate message bytes all match their expected values
+	tables := []struct {
+		u  string
+		w  string
+		xu string
+		xd string
+		xb []byte
+	}{
+		{username, "", username, "", []byte{
+			0x4e, 0x54, 0x4c, 0x4d, 0x53, 0x53, 0x50, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 
+			0x80, 0xa0, 0x00, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
+			0x28, 0x00, 0x00, 0x00, 0x06, 0x01, 0xb1, 0x1d, 0x00, 0x00, 0x00, 0x0f }},
+		{domain+"\\"+username, "", username, domain, []byte{
+			0x4e, 0x54, 0x4c, 0x4d, 0x53, 0x53, 0x50, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x10, 
+			0x80, 0xa0, 0x08, 0x00, 0x08, 0x00, 0x28, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
+			0x30, 0x00, 0x00, 0x00, 0x06, 0x01, 0xb1, 0x1d, 0x00, 0x00, 0x00, 0x0f, 0x4d, 0x59, 
+			0x44, 0x4f, 0x4d, 0x41, 0x49, 0x4e }},
+		{domain+"\\"+username, workstation, username, domain, []byte{
+			0x4e, 0x54, 0x4c, 0x4d, 0x53, 0x53, 0x50, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x30, 
+			0x80, 0xa0, 0x08, 0x00, 0x08, 0x00, 0x28, 0x00, 0x00, 0x00, 0x04, 0x00, 0x04, 0x00, 
+			0x30, 0x00, 0x00, 0x00, 0x06, 0x01, 0xb1, 0x1d, 0x00, 0x00, 0x00, 0x0f, 0x4d, 0x59, 
+			0x44, 0x4f, 0x4d, 0x41, 0x49, 0x4e, 0x4d, 0x59, 0x50, 0x43 }},
+		{username, workstation, username, "", []byte{
+			0x4e, 0x54, 0x4c, 0x4d, 0x53, 0x53, 0x50, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x20, 
+			0x80, 0xa0, 0x00, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x04, 0x00, 0x04, 0x00, 
+			0x28, 0x00, 0x00, 0x00, 0x06, 0x01, 0xb1, 0x1d, 0x00, 0x00, 0x00, 0x0f, 0x4d, 0x59, 
+			0x50, 0x43 }},
+	}
+
+	for _, table := range tables {
+		tuser, tdomain := GetDomain(table.u)
+		if tuser != table.xu {
+			t.Fatalf("username not correct, expected %v got %v", tuser, table.xu)
+		}
+		if tdomain != table.xd {
+			t.Fatalf("domain not correct, expected %v got %v", tdomain, table.xd)
+		}
+
+		tb, err := NewNegotiateMessage(tdomain, table.w)
+		if err != nil {
+			t.Fatalf("error creating new negotiate message with domain '%v' and workstation '%v'", tdomain, table.w)
+		}
+
+		if !bytes.Equal(tb, table.xb) {
+			t.Fatalf("negotiate message bytes not correct, expected %v got %v", tb, table.xb)
+		}
+	}
+
+}
+
 func TestCalculateNTLMv2Response(t *testing.T) {
 	NTLMv2Hash := getNtlmV2Hash(password, username, target)
 	ClientChallenge := []byte{0xff, 0xff, 0xff, 0x00, 0x11, 0x22, 0x33, 0x44}