From 2f32a2df72e4e7614701fcd295a81a15e93a2633 Mon Sep 17 00:00:00 2001
From: Jordan Borean <jborean93@gmail.com>
Date: Fri, 28 Apr 2023 13:41:01 +1000
Subject: [PATCH] More fixes

---
 tests/test_auth_dce.py | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/tests/test_auth_dce.py b/tests/test_auth_dce.py
index 4d21e91..5be8cce 100644
--- a/tests/test_auth_dce.py
+++ b/tests/test_auth_dce.py
@@ -72,29 +72,29 @@ def _ntlm_exchange(
     assert not client.complete
     assert not server.complete
 
-    if client.protocol == "negotiate":
-        authentication = client.step(challenge)
-        assert authentication
-        assert not client.complete
-        assert not server.complete
+    authentication = client.step(challenge)
+    assert authentication
 
-        mech_list = server.step(authentication)
-        assert mech_list
-        assert not client.complete
-        assert server.complete
+    # If it's being wrapped by SPNEGO then an extra msg is expected
+    if authentication.startswith(b"NTLMSSP\x00"):
+        assert client.complete
+        assert not server.complete
 
-        final = client.step(mech_list)
+        final = server.step(authentication)
         assert not final
         assert client.complete
         assert server.complete
 
     else:
-        authentication = client.step(challenge)
-        assert authentication
-        assert client.complete
+        assert not client.complete
         assert not server.complete
 
-        final = server.step(authentication)
+        mech_list = server.step(authentication)
+        assert mech_list
+        assert not client.complete
+        assert server.complete
+
+        final = client.step(mech_list)
         assert not final
         assert client.complete
         assert server.complete