From e9a93d117b2811a5e5697a9af94f906808e4d68e Mon Sep 17 00:00:00 2001
From: citronneur <citronneur@gmail.com>
Date: Sat, 21 Feb 2015 14:29:17 +0100
Subject: [PATCH] automata modification for NLA

---
 rdpy/core/layer.py            | 20 +++++++++++++++++++-
 rdpy/protocol/rdp/nla/cssp.py |  6 ++++--
 rdpy/protocol/rdp/tpkt.py     | 10 +++++-----
 3 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/rdpy/core/layer.py b/rdpy/core/layer.py
index b4084a6..801a6f6 100644
--- a/rdpy/core/layer.py
+++ b/rdpy/core/layer.py
@@ -200,7 +200,7 @@ class RawLayer(protocol.Protocol, LayerAutomata, IStreamSender):
         #add in buffer
         self._buffer += data
         #while buffer have expected size call local callback
-        while len(self._buffer) >= self._expectedLen:
+        while self._expectedLen > 0 and len(self._buffer) >= self._expectedLen:
             #expected data is first expected bytes
             expectedData = Stream(self._buffer[0:self._expectedLen])
             #rest is for next event of automata
@@ -238,9 +238,27 @@ class RawLayer(protocol.Protocol, LayerAutomata, IStreamSender):
         @param expectedLen: in bytes length use to call next state
         @param callback: callback call when expected length bytes is received
         """
+        self.dataReceived = lambda data:self.__class__.dataReceived(self, data)
         self._expectedLen = expectedLen
         #default callback is recv from LayerAutomata
         self.setNextState(callback)
+    
+    def byPassDataReceived(self, data):
+        """
+        @summary: by pass automata received
+        @warning: ignore rest in buffer
+        @param data: {str} received from twisted
+        """
+        self.recv(Stream(data))
+        
+    def expectImmediately(self, callback = None):
+        """
+        @summary: call immediately when available data is received
+        @param callback: {func} callback called
+        """
+        self.dataReceived = lambda data:self.__class__.byPassDataReceived(self, data)
+        self._expectedLen = 0
+        self.setNextState(callback)
         
     def send(self, message):
         """
diff --git a/rdpy/protocol/rdp/nla/cssp.py b/rdpy/protocol/rdp/nla/cssp.py
index 61ea226..ba09d8c 100644
--- a/rdpy/protocol/rdp/nla/cssp.py
+++ b/rdpy/protocol/rdp/nla/cssp.py
@@ -97,7 +97,7 @@ class TSSmartCardCreds(univ.Sequence):
         namedtype.OptionalNamedType('domainHint', univ.OctetString().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 3)))
         )
 
-def createDERRequest(negoTypes):
+def encodeDERTRequest(negoTypes):
     """
     @summary: create TSRequest from list of Type
     @param negoTypes: {list(Type)}
@@ -119,4 +119,6 @@ def createDERRequest(negoTypes):
     request.setComponentByName("version", univ.Integer(2).subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0)))
     request.setComponentByName("negoTokens", negoData)
     return encoder.encode(request)
-        
\ No newline at end of file
+
+def decodeDERTRequest():
+    pass
\ No newline at end of file
diff --git a/rdpy/protocol/rdp/tpkt.py b/rdpy/protocol/rdp/tpkt.py
index 1b161d6..b40cb2a 100644
--- a/rdpy/protocol/rdp/tpkt.py
+++ b/rdpy/protocol/rdp/tpkt.py
@@ -185,17 +185,17 @@ class TPKT(RawLayer, IFastPathSender):
         @summary: Fast path data
         @param data: {Stream} from twisted layer
         """
-        self._fastPathListener.recvFastPath(self._secFlag, data)
         self.expect(2, self.readHeader)
+        self._fastPathListener.recvFastPath(self._secFlag, data)
     
     def readData(self, data):
         """
         @summary: Read classic TPKT packet, last state in tpkt automata
         @param data: {Stream} with correct size
         """
+        self.expect(2, self.readHeader)
         #next state is pass to 
         self._presentation.recv(data)
-        self.expect(2, self.readHeader)
         
     def send(self, message):
         """
@@ -224,11 +224,11 @@ class TPKT(RawLayer, IFastPathSender):
                     must be called after startTLS function
         """
         #send NTLM negotiate message packet
-        self.transport.write(cssp.createDERRequest( [ ntlm.NegotiateMessage() ] ))
-        
+        self.transport.write(cssp.encodeDERTRequest( [ ntlm.NegotiateMessage() ] ))
+        self.expectImmediately(self.readNTLMChallenge)
         
     def readNTLMChallenge(self, data):
         """
         @summary: server NTLM challenge
         """
-        
\ No newline at end of file
+        print "toto"
\ No newline at end of file