Fix immediatly automata

This commit is contained in:
citronneur
2015-02-22 14:04:09 +01:00
parent e9a93d117b
commit 8cd789480f
3 changed files with 26 additions and 16 deletions

View File

@@ -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._expectedLen = 0
self.setNextState(callback)
self.dataReceived = lambda data:self.__class__.dataReceivedAll(self, data)
self._immediateCallback = callback
def restartAutomata(self):
"""
@summary: restart automata
"""
self.dataReceived = lambda data:self.__class__.dataReceived(self, data)
def send(self, message):
"""

View File

@@ -23,8 +23,8 @@
"""
from pyasn1.type import namedtype, univ, tag
from pyasn1.codec.der import encoder
from rdpy.core.type import Stream
from pyasn1.codec.der import encoder, decoder
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():
pass
def decodeDERTRequest(s):
"""
@summary: Decode the stream as
@param s: {Stream}
"""
tRequest = decoder.decode(s.getvalue(), asn1Spec=TSRequest())
print tRequest

View File

@@ -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):