Fix immediatly automata

rdpy/core/layer.py

@@ -183,6 +183,7 @@ class RawLayer(protocol.Protocol, LayerAutomata, IStreamSender):
         #len of next packet pass to next state function
         self._expectedLen = 0
         self._factory = None
+        self._immediateCallback = None
         
     def setFactory(self, factory):
         """
@@ -238,27 +239,31 @@ 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):
+    def dataReceivedAll(self, data):
         """
         @summary: by pass automata received
         @warning: ignore rest in buffer
         @param data: {str} received from twisted
         """
-        self.recv(Stream(data))
+        self._immediateCallback(Stream(data))
         
-    def expectImmediately(self, callback = None):
+    def expectImmediately(self, callback):
         """
         @summary: call immediately when available data is received
         @param callback: {func} callback called
         """
-        self.dataReceived = lambda data:self.__class__.byPassDataReceived(self, data)
+        self.dataReceived = lambda data:self.__class__.dataReceivedAll(self, data)
-        self._expectedLen = 0
+        self._immediateCallback = callback
-        self.setNextState(callback)
+        
+    def restartAutomata(self):
+        """
+        @summary: restart automata
+        """
+        self.dataReceived = lambda data:self.__class__.dataReceived(self, data)
         
     def send(self, message):
         """

rdpy/protocol/rdp/nla/cssp.py

@@ -23,8 +23,8 @@
 """
 
 from pyasn1.type import namedtype, univ, tag
-from pyasn1.codec.der import encoder
+from pyasn1.codec.der import encoder, decoder
-from rdpy.core.type import Stream
+from rdpy.core.type import Stream, String
 
 class NegoToken(univ.Sequence):
     componentType = namedtype.NamedTypes(
@@ -101,7 +101,7 @@ def encodeDERTRequest(negoTypes):
     """
     @summary: create TSRequest from list of Type
     @param negoTypes: {list(Type)}
-    @return: {str}
+    @return: {String}
     """
     negoData = NegoData().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 1))
     
@@ -118,7 +118,12 @@ def encodeDERTRequest(negoTypes):
     request = TSRequest()
     request.setComponentByName("version", univ.Integer(2).subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0)))
     request.setComponentByName("negoTokens", negoData)
-    return encoder.encode(request)
+    return String(encoder.encode(request))
 
-def decodeDERTRequest():
+def decodeDERTRequest(s):
-    pass
+    """
+    @summary: Decode the stream as 
+    @param s: {Stream}
+    """
+    tRequest = decoder.decode(s.getvalue(), asn1Spec=TSRequest())
+    print tRequest

rdpy/protocol/rdp/tpkt.py

@@ -185,17 +185,17 @@ class TPKT(RawLayer, IFastPathSender):
         @summary: Fast path data
         @param data: {Stream} from twisted layer
         """
-        self.expect(2, self.readHeader)
         self._fastPathListener.recvFastPath(self._secFlag, data)
+        self.expect(2, self.readHeader)
     
     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,7 +224,7 @@ class TPKT(RawLayer, IFastPathSender):
                     must be called after startTLS function
         """
         #send NTLM negotiate message packet
-        self.transport.write(cssp.encodeDERTRequest( [ ntlm.NegotiateMessage() ] ))
+        RawLayer.send(self, cssp.encodeDERTRequest( [ ntlm.NegotiateMessage() ] ))
         self.expectImmediately(self.readNTLMChallenge)
         
     def readNTLMChallenge(self, data):