From df1f808002e542c886b8b48c6eefc8a1e284ae43 Mon Sep 17 00:00:00 2001
From: Dan <justdan96@gmail.com>
Date: Thu, 12 Apr 2018 18:05:24 +0100
Subject: [PATCH] add version, domain and workstation fields

Add version support from jh125486
Add workstation and domain support from jh125486
Enable to extract domain from username
---
 negotiate_message.go | 64 +++++++++++++++++++++++++++++++++-----------
 negotiator.go        | 19 ++++++++++---
 version.go           | 20 ++++++++++++++
 3 files changed, 84 insertions(+), 19 deletions(-)
 create mode 100644 version.go

diff --git a/negotiate_message.go b/negotiate_message.go
index 97aa07e..151a91c 100644
--- a/negotiate_message.go
+++ b/negotiate_message.go
@@ -3,29 +3,61 @@ package ntlmssp
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
+	"strings"
 )
 
+const expMsgBodyLen = 40
+
 type negotiateMessageFields struct {
-	messageHeader
-	NegotiateFlags negotiateFlags
+        messageHeader
+        NegotiateFlags negotiateFlags
+
+        Domain      varField
+        Workstation varField
+
+        Version
 }
 
+var defaultFlags = negotiateFlagNTLMSSPNEGOTIATETARGETINFO |
+        negotiateFlagNTLMSSPNEGOTIATE56 |
+        negotiateFlagNTLMSSPNEGOTIATE128 |
+        negotiateFlagNTLMSSPNEGOTIATEUNICODE
+
 //NewNegotiateMessage creates a new NEGOTIATE message with the
 //flags that this package supports.
-func NewNegotiateMessage() []byte {
-	m := negotiateMessageFields{
-		messageHeader: newMessageHeader(1),
-	}
+func NewNegotiateMessage(domainName, workstationName string) ([]byte, error) {
+        payloadOffset := expMsgBodyLen
+        flags := defaultFlags
 
-	m.NegotiateFlags = negotiateFlagNTLMSSPREQUESTTARGET |
-		negotiateFlagNTLMSSPNEGOTIATENTLM |
-		negotiateFlagNTLMSSPNEGOTIATEALWAYSSIGN |
-		negotiateFlagNTLMSSPNEGOTIATEUNICODE
+        if domainName != "" {
+                flags |= negotiateFlagNTLMSSPNEGOTIATEOEMDOMAINSUPPLIED
+        }
 
-	b := bytes.Buffer{}
-	err := binary.Write(&b, binary.LittleEndian, &m)
-	if err != nil {
-		panic(err)
-	}
-	return b.Bytes()
+        if workstationName != "" {
+                flags |= negotiateFlagNTLMSSPNEGOTIATEOEMWORKSTATIONSUPPLIED
+        }
+
+        msg := negotiateMessageFields{
+                messageHeader:  newMessageHeader(1),
+                NegotiateFlags: flags,
+                Domain:         newVarField(&payloadOffset, len(domainName)),
+                Workstation:    newVarField(&payloadOffset, len(workstationName)),
+                Version:        DefaultVersion(),
+        }
+
+        b := bytes.Buffer{}
+        if err := binary.Write(&b, binary.LittleEndian, &msg); err != nil {
+                return nil, err
+        }
+        if b.Len() != expMsgBodyLen {
+                return nil, errors.New("incorrect body length")
+        }
+
+        payload := strings.ToUpper(domainName + workstationName)
+        if _, err := b.WriteString(payload); err != nil {
+                return nil, err
+        }
+
+        return b.Bytes(), nil
 }
diff --git a/negotiator.go b/negotiator.go
index 9ae4586..ac40e16 100644
--- a/negotiator.go
+++ b/negotiator.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"strings"
 )
 
 //Negotiator is a http.Roundtripper decorator that automatically
@@ -76,9 +77,21 @@ func (l Negotiator) RoundTrip(req *http.Request) (res *http.Response, err error)
 			return nil, err
 		}
 
-		// send negotiate
-		negotiateMessage := NewNegotiateMessage()
-		if resauth.IsNTLM() {
+       // parse domain name from username
+       domain := ""
+
+       if strings.Contains(u, "\\") {
+            u_components := strings.Split(u, "\\")
+            domain = u_components[0]
+            u = u_components[1]
+       }
+
+       // send negotiate
+       negotiateMessage, err := NewNegotiateMessage(domain, "")
+       if err != nil {
+           return nil, err
+       }
+       if resauth.IsNTLM() {
 			req.Header.Set("Authorization", "NTLM "+base64.StdEncoding.EncodeToString(negotiateMessage))
 		} else {
 			req.Header.Set("Authorization", "Negotiate "+base64.StdEncoding.EncodeToString(negotiateMessage))
diff --git a/version.go b/version.go
new file mode 100644
index 0000000..6d84892
--- /dev/null
+++ b/version.go
@@ -0,0 +1,20 @@
+package ntlmssp
+
+// Version is a struct representing https://msdn.microsoft.com/en-us/library/cc236654.aspx
+type Version struct {
+	ProductMajorVersion uint8
+	ProductMinorVersion uint8
+	ProductBuild        uint16
+	_                   [3]byte
+	NTLMRevisionCurrent uint8
+}
+
+// DefaultVersion returns a Version with "sensible" defaults (Windows 7)
+func DefaultVersion() Version {
+	return Version{
+		ProductMajorVersion: 6,
+		ProductMinorVersion: 1,
+		ProductBuild:        7601,
+		NTLMRevisionCurrent: 15,
+	}
+}