Package tlslite :: Module defragmenter
[hide private]
[frames] | no frames]

Source Code for Module tlslite.defragmenter

  1  # Copyright (c) 2015, Hubert Kario 
  2  # 
  3  # See the LICENSE file for legal information regarding use of this file. 
  4   
  5  """ Helper package for handling fragmentation of messages """ 
  6   
  7  from __future__ import generators 
  8   
  9  from .utils.codec import Parser 
 10   
11 -class Defragmenter(object):
12 13 """ 14 Class for demultiplexing TLS messages. 15 16 Since the messages can be interleaved and fragmented between each other 17 we need to cache not complete ones and return in order of urgency. 18 19 Supports messages with given size (like Alerts) or with a length header 20 in specific place (like Handshake messages). 21 22 @ivar priorities: order in which messages from given types should be 23 returned. 24 @ivar buffers: data buffers for message types 25 @ivar decoders: functions which check buffers if a message of given type 26 is complete 27 """ 28
29 - def __init__(self):
30 """Set up empty defregmenter""" 31 self.priorities = [] 32 self.buffers = {} 33 self.decoders = {}
34
35 - def addStaticSize(self, msgType, size):
36 """Add a message type which all messages are of same length""" 37 if msgType in self.priorities: 38 raise ValueError("Message type already defined") 39 if size < 1: 40 raise ValueError("Message size must be positive integer") 41 42 self.priorities += [msgType] 43 44 self.buffers[msgType] = bytearray(0) 45 def sizeHandler(data): 46 """ 47 Size of message in parameter 48 49 If complete message is present in parameter returns its size, 50 None otherwise. 51 """ 52 if len(data) < size: 53 return None 54 else: 55 return size
56 self.decoders[msgType] = sizeHandler
57
58 - def addDynamicSize(self, msgType, sizeOffset, sizeOfSize):
59 """Add a message type which has a dynamic size set in a header""" 60 if msgType in self.priorities: 61 raise ValueError("Message type already defined") 62 if sizeOfSize < 1: 63 raise ValueError("Size of size must be positive integer") 64 if sizeOffset < 0: 65 raise ValueError("Offset can't be negative") 66 67 self.priorities += [msgType] 68 self.buffers[msgType] = bytearray(0) 69 70 def sizeHandler(data): 71 """ 72 Size of message in parameter 73 74 If complete message is present in parameter returns its size, 75 None otherwise. 76 """ 77 if len(data) < sizeOffset+sizeOfSize: 78 return None 79 else: 80 parser = Parser(data) 81 # skip the header 82 parser.getFixBytes(sizeOffset) 83 84 payloadLength = parser.get(sizeOfSize) 85 if parser.getRemainingLength() < payloadLength: 86 # not enough bytes in buffer 87 return None 88 return sizeOffset + sizeOfSize + payloadLength
89 90 self.decoders[msgType] = sizeHandler 91
92 - def addData(self, msgType, data):
93 """Adds data to buffers""" 94 if msgType not in self.priorities: 95 raise ValueError("Message type not defined") 96 97 self.buffers[msgType] += data
98
99 - def getMessage(self):
100 """Extract the highest priority complete message from buffer""" 101 for msgType in self.priorities: 102 length = self.decoders[msgType](self.buffers[msgType]) 103 if length is None: 104 continue 105 106 # extract message 107 data = self.buffers[msgType][:length] 108 # remove it from buffer 109 self.buffers[msgType] = self.buffers[msgType][length:] 110 return (msgType, data) 111 return None
112
113 - def clearBuffers(self):
114 """Remove all data from buffers""" 115 for key in self.buffers.keys(): 116 self.buffers[key] = bytearray(0)
117