From eb92e824ac49dd274abb9dd89c97fd7208b2ee9e Mon Sep 17 00:00:00 2001
From: Tal Aloni <tal.aloni.il@gmail.com>
Date: Sat, 23 Sep 2017 16:11:32 +0300
Subject: [PATCH] SPNEGO: Minor code refactoring

---
 .../Authentication/GSSAPI/GSSProvider.cs      |  2 +-
 .../SPNEGO/SimpleProtectedNegotiationToken.cs | 54 +++++++++----------
 .../Helpers/NTLMAuthenticationHelper.cs       |  2 +-
 3 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/SMBLibrary/Authentication/GSSAPI/GSSProvider.cs b/SMBLibrary/Authentication/GSSAPI/GSSProvider.cs
index ac32502..91d01f0 100644
--- a/SMBLibrary/Authentication/GSSAPI/GSSProvider.cs
+++ b/SMBLibrary/Authentication/GSSAPI/GSSProvider.cs
@@ -48,7 +48,7 @@ namespace SMBLibrary.Authentication.GSSAPI
             {
                 token.MechanismTypeList.Add(mechanism.Identifier);
             }
-            return SimpleProtectedNegotiationToken.GetTokenBytes(token);
+            return token.GetBytes(true);
         }
 
         public virtual NTStatus AcceptSecurityContext(ref GSSContext context, byte[] inputToken, out byte[] outputToken)
diff --git a/SMBLibrary/Authentication/GSSAPI/SPNEGO/SimpleProtectedNegotiationToken.cs b/SMBLibrary/Authentication/GSSAPI/SPNEGO/SimpleProtectedNegotiationToken.cs
index 8286d39..c7f5bcd 100644
--- a/SMBLibrary/Authentication/GSSAPI/SPNEGO/SimpleProtectedNegotiationToken.cs
+++ b/SMBLibrary/Authentication/GSSAPI/SPNEGO/SimpleProtectedNegotiationToken.cs
@@ -19,6 +19,32 @@ namespace SMBLibrary.Authentication.GSSAPI
 
         public abstract byte[] GetBytes();
 
+        /// <param name="includeHeader">Prepend the generic GSSAPI header. Required for negTokenInit, optional for negTokenResp.</param>
+        public byte[] GetBytes(bool includeHeader)
+        {
+            byte[] tokenBytes = this.GetBytes();
+            if (includeHeader)
+            {
+                int objectIdentifierFieldSize = DerEncodingHelper.GetLengthFieldSize(SPNEGOIdentifier.Length);
+                int tokenLength = 1 + objectIdentifierFieldSize + SPNEGOIdentifier.Length + tokenBytes.Length;
+                int tokenLengthFieldSize = DerEncodingHelper.GetLengthFieldSize(tokenLength);
+                int headerLength = 1 + tokenLengthFieldSize + 1 + objectIdentifierFieldSize + SPNEGOIdentifier.Length;
+                byte[] buffer = new byte[headerLength + tokenBytes.Length];
+                int offset = 0;
+                ByteWriter.WriteByte(buffer, ref offset, ApplicationTag);
+                DerEncodingHelper.WriteLength(buffer, ref offset, tokenLength);
+                ByteWriter.WriteByte(buffer, ref offset, (byte)DerEncodingTag.ObjectIdentifier);
+                DerEncodingHelper.WriteLength(buffer, ref offset, SPNEGOIdentifier.Length);
+                ByteWriter.WriteBytes(buffer, ref offset, SPNEGOIdentifier);
+                ByteWriter.WriteBytes(buffer, ref offset, tokenBytes);
+                return buffer;
+            }
+            else
+            {
+                return tokenBytes;
+            }
+        }
+
         /// <summary>
         /// https://tools.ietf.org/html/rfc2743
         /// </summary>
@@ -60,33 +86,5 @@ namespace SMBLibrary.Authentication.GSSAPI
             }
             return null;
         }
-
-        /// <summary>
-        /// Will append the generic GSSAPI header.
-        /// </summary>
-        public static byte[] GetTokenBytes(SimpleProtectedNegotiationToken token)
-        {
-            if (token is SimpleProtectedNegotiationTokenInit)
-            {
-                byte[] tokenBytes = token.GetBytes();
-                int objectIdentifierFieldSize = DerEncodingHelper.GetLengthFieldSize(SPNEGOIdentifier.Length);
-                int tokenLength = 1 + objectIdentifierFieldSize + SPNEGOIdentifier.Length + tokenBytes.Length;
-                int tokenLengthFieldSize = DerEncodingHelper.GetLengthFieldSize(tokenLength);
-                int headerLength = 1 + tokenLengthFieldSize + 1 + objectIdentifierFieldSize + SPNEGOIdentifier.Length;
-                byte[] buffer = new byte[headerLength + tokenBytes.Length];
-                int offset = 0;
-                ByteWriter.WriteByte(buffer, ref offset, ApplicationTag);
-                DerEncodingHelper.WriteLength(buffer, ref offset, tokenLength);
-                ByteWriter.WriteByte(buffer, ref offset, (byte)DerEncodingTag.ObjectIdentifier);
-                DerEncodingHelper.WriteLength(buffer, ref offset, SPNEGOIdentifier.Length);
-                ByteWriter.WriteBytes(buffer, ref offset, SPNEGOIdentifier);
-                ByteWriter.WriteBytes(buffer, ref offset, tokenBytes);
-                return buffer;
-            }
-            else
-            {
-                return token.GetBytes();
-            }
-        }
     }
 }
diff --git a/SMBLibrary/Client/Helpers/NTLMAuthenticationHelper.cs b/SMBLibrary/Client/Helpers/NTLMAuthenticationHelper.cs
index c00c4b6..595e93f 100644
--- a/SMBLibrary/Client/Helpers/NTLMAuthenticationHelper.cs
+++ b/SMBLibrary/Client/Helpers/NTLMAuthenticationHelper.cs
@@ -67,7 +67,7 @@ namespace SMBLibrary.Client
                 outputToken.MechanismTypeList = new List<byte[]>();
                 outputToken.MechanismTypeList.Add(GSSProvider.NTLMSSPIdentifier);
                 outputToken.MechanismToken = negotiateMessage.GetBytes();
-                return SimpleProtectedNegotiationToken.GetTokenBytes(outputToken);
+                return outputToken.GetBytes(true);
             }
             else
             {