def _flush_info_headers(info_data, write_headers, type): if (len(write_headers) > 0): info_data.write(getVarint(type)) info_data.write(getVarint(len(write_headers))) for str_key, str_value in write_headers.iteritems(): info_data.write(THeaderTransport._serialize_string(str_key)) info_data.write(THeaderTransport._serialize_string(str_value)) write_headers.clear()
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for CLIENT_TYPE.HEADER @param buf(StringIO): Buffer to write message to @param wout(str): Payload @param wsz(int): Payload length """ transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity info_data = StringIO() # Write persistent kv-headers _flush_info_headers(info_data, self.get_write_persistent_headers(), INFO.PERSISTENT) # Write non-persistent kv-headers _flush_info_headers(info_data, self.__write_headers, INFO.NORMAL) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + \ info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size # MAGIC(2) | FLAGS(2) + SEQ_ID(4) + HEADER_SIZE(2) wsz += header_size + 10 if wsz > MAX_FRAME_SIZE: buf.write(pack("!I", BIG_FRAME_MAGIC)) buf.write(pack("!Q", wsz)) else: buf.write(pack("!I", wsz)) buf.write(pack("!HH", HEADER_MAGIC >> 16, self.__flags)) buf.write(pack("!I", self.seq_id)) buf.write(pack("!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) buf.write(info_data.getvalue()) # Pad out the header with 0x00 for x in range(0, padding_size, 1): buf.write(pack("!c", b'\0')) # Send data section buf.write(wout)
def _flush_info_headers(info_data, write_headers, type): if (len(write_headers) > 0): info_data.write(getVarint(type)) info_data.write(getVarint(len(write_headers))) write_headers_iter = write_headers.items() for str_key, str_value in write_headers_iter: info_data.write(_serialize_string(str_key)) info_data.write(_serialize_string(str_value)) write_headers.clear()
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for self.HEADERS_CLIENT_TYPE @param buf(StringIO): Buffer to write message to @param wout(str): Payload @param wsz(int): Payload length """ transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity info_data = StringIO() # Write persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.get_write_persistent_headers(), self.INFO_PKEYVALUE) # Write non-persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.__write_headers, self.INFO_KEYVALUE) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + \ info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size wsz += header_size + 10 buf.write(pack(b"!I", wsz)) buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags)) buf.write(pack(b"!I", self.__seq_id)) buf.write(pack(b"!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) buf.write(info_data.getvalue()) # Pad out the header with 0x00 for x in range(0, padding_size, 1): buf.write(pack(b"!c", b'\0')) # Send data section buf.write(wout)
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for self.HEADERS_CLIENT_TYPE @param buf(StringIO): Buffer to write message to @param wout(str): Payload @param wsz(int): Payload length """ transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity info_data = StringIO() # Write persistent kv-headers THeaderTransport._flush_info_headers( info_data, self.get_write_persistent_headers(), self.INFO_PKEYVALUE) # Write non-persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.__write_headers, self.INFO_KEYVALUE) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + \ info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size wsz += header_size + 10 buf.write(pack(b"!I", wsz)) buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags)) buf.write(pack(b"!I", self.__seq_id)) buf.write(pack(b"!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) buf.write(info_data.getvalue()) # Pad out the header with 0x00 for x in range(0, padding_size, 1): buf.write(pack(b"!c", b'\0')) # Send data section buf.write(wout)
def _flush_info_headers(info_data, write_headers, type): if (len(write_headers) > 0): info_data.write(getVarint(type)) info_data.write(getVarint(len(write_headers))) if sys.version_info[0] >= 3: write_headers_iter = write_headers.items() else: write_headers_iter = write_headers.iteritems() for str_key, str_value in write_headers_iter: info_data.write(THeaderTransport._serialize_string(str_key)) info_data.write(THeaderTransport._serialize_string(str_value)) write_headers.clear()
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for CLIENT_TYPE.HEADER This method writes a message using the same logic as THeaderTransport._flushHeaderMessage but mutates fields included in self._fuzz_fields """ cls = self.__class__ transform_data = StringIO() num_transforms = len(self._THeaderTransport__write_transforms) for trans_id in self._THeaderTransport__write_transforms: trans_id = self._get_fuzzy_field('transform_id', trans_id, 'i32') transform_data.write(getVarint(trans_id)) # Add in special flags. if self._THeaderTransport__identity: id_version = self._get_fuzzy_field( 'version', self.ID_VERSION, 'str') self._THeaderTransport__write_headers[self.ID_VERSION_HEADER] = ( id_version) identity = self._get_fuzzy_field( 'identity', self._THeaderTransport__identity, 'str') self._THeaderTransport__write_headers[self.IDENTITY_HEADER] = ( identity) info_data = StringIO() # Write persistent kv-headers _flush_info_headers( info_data, self._THeaderTransport__write_persistent_headers, INFO.PERSISTENT) # Write non-persistent kv-headers _flush_info_headers( info_data, self._THeaderTransport__write_headers, INFO.NORMAL) header_data = StringIO() proto_id = self._get_fuzzy_field( 'proto_id', self._THeaderTransport__proto_id, 'i32') header_data.write(getVarint(proto_id)) num_transforms = self._get_fuzzy_field( 'num_transforms', num_transforms, 'i32') header_data.write(getVarint(num_transforms)) header_size = (transform_data.tell() + header_data.tell() + info_data.tell()) padding_size = 4 - (header_size % 4) # Fuzz padding size, but do not let total header size exceed 2**16 - 1 padding_size = min((2 ** 16 - 1) - header_size, self._get_fuzzy_field('padding', padding_size, 'i16')) header_size = header_size + padding_size wsz += header_size + 10 self._write_fuzzy_field(buf, 'length', wsz, 'i32') self._write_fuzzy_field(buf, 'magic', HEADER_MAGIC >> 16, 'i16') self._write_fuzzy_field( buf, 'flags', self.header_flags(), 'i16') self._write_fuzzy_field( buf, 'seq_id', self.seq_id, 'i32') self._write_fuzzy_field(buf, 'header_size', header_size // 4, 'i16') buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) buf.write(info_data.getvalue()) # Pad out the header with 0x00 if 'padding_bytes' in self._fuzz_fields: # Print that padding bytes are being fuzzed in favor of printing # the value of each individual padding byte self._print("Fuzzing %d padding bytes" % padding_size) old_verbose, self._verbose = self._verbose, False for _ in sm.xrange(padding_size): self._write_fuzzy_field( buf, 'padding_bytes', six.int2byte(0), 'char') self._verbose = old_verbose self._write_fuzzy_field(buf, 'payload', wout, 'str')
def _serialize_string(str_): if sys.version_info[0] >= 3 and not isinstance(str_, bytes): str_ = str_.encode() return getVarint(len(str_)) + str_
def _serialize_string(str_): if PY3 and not isinstance(str_, bytes): str_ = str_.encode() return getVarint(len(str_)) + str_
def flush(self): wout = self.__wbuf.getvalue() wout = self.transform(wout) wsz = len(wout) # reset wbuf before write/flush to preserve state on underlying failure self.__wbuf.seek(0) self.__wbuf.truncate() if self.__proto_id == 1 and self.__client_type != self.HTTP_CLIENT_TYPE: raise TTransportException(TTransportException.INVALID_CLIENT_TYPE, "Trying to send JSON encoding over binary") buf = StringIO() if self.__client_type == self.HEADERS_CLIENT_TYPE: transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) if self.__hmac_func: num_transforms += 1 transform_data.write(getVarint(self.HMAC_TRANSFORM)) transform_data.write(b'\0') # size of hmac, fixup later. # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity info_data = StringIO() # Write persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.__write_persistent_headers, self.INFO_PKEYVALUE) # Write non-persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.__write_headers, self.INFO_KEYVALUE) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + \ info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size wsz += header_size + 10 buf.write(pack(b"!I", wsz)) buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags)) buf.write(pack(b"!I", self.__seq_id)) buf.write(pack(b"!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) hmac_loc = buf.tell() - 1 # Fixup hmac size later buf.write(info_data.getvalue()) # Pad out the header with 0x00 for x in range(0, padding_size, 1): buf.write(pack(b"!c", b'\0')) # Send data section buf.write(wout) # HMAC calculation should always be last. if self.__hmac_func: hmac_data = buf.getvalue()[4:] hmac = self.__hmac_func(hmac_data) # Fill in hmac size. buf.seek(hmac_loc) buf.write(chr(len(hmac))) buf.seek(0, os.SEEK_END) buf.write(hmac) # Fix packet size since we appended data. new_sz = buf.tell() - 4 buf.seek(0) buf.write(pack(b"!I", new_sz)) elif self.__client_type == self.FRAMED_DEPRECATED: buf.write(pack(b"!I", wsz)) buf.write(wout) elif self.__client_type == self.UNFRAMED_DEPRECATED: buf.write(wout) elif self.__client_type == self.HTTP_CLIENT_TYPE: # Reset the client type if we sent something - # oneway calls via HTTP expect a status response otherwise buf.write(self.header.getvalue()) buf.write(wout) self.__client_type == self.HEADERS_CLIENT_TYPE elif self.__client_type == self.UNKNOWN_CLIENT_TYPE: raise TTransportException(TTransportException.INVALID_CLIENT_TYPE, "Unknown client type") # We don't include the framing bytes as part of the frame size check # (so -4) if buf.tell() - 4 > self.__max_frame_size: raise TTransportException(TTransportException.INVALID_FRAME_SIZE, "Attempting to send frame that is too large") self.__trans.write(buf.getvalue()) self.__trans.flush()
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for CLIENT_TYPE.HEADER This method writes a message using the same logic as THeaderTransport._flushHeaderMessage but mutates fields included in self._fuzz_fields """ transform_data = StringIO() num_transforms = len(self._THeaderTransport__write_transforms) for trans_id in self._THeaderTransport__write_transforms: trans_id = self._get_fuzzy_field('transform_id', trans_id, 'i32') transform_data.write(getVarint(trans_id)) # Add in special flags. if self._THeaderTransport__identity: id_version = self._get_fuzzy_field( 'version', self.ID_VERSION, 'str') self._THeaderTransport__write_headers[self.ID_VERSION_HEADER] = ( id_version) identity = self._get_fuzzy_field( 'identity', self._THeaderTransport__identity, 'str') self._THeaderTransport__write_headers[self.IDENTITY_HEADER] = ( identity) info_data = StringIO() # Write persistent kv-headers _flush_info_headers( info_data, self._THeaderTransport__write_persistent_headers, INFO.PERSISTENT) # Write non-persistent kv-headers _flush_info_headers( info_data, self._THeaderTransport__write_headers, INFO.NORMAL) header_data = StringIO() proto_id = self._get_fuzzy_field( 'proto_id', self._THeaderTransport__proto_id, 'i32') header_data.write(getVarint(proto_id)) num_transforms = self._get_fuzzy_field( 'num_transforms', num_transforms, 'i32') header_data.write(getVarint(num_transforms)) header_size = (transform_data.tell() + header_data.tell() + info_data.tell()) padding_size = 4 - (header_size % 4) # Fuzz padding size, but do not let total header size exceed 2**16 - 1 padding_size = min((2 ** 16 - 1) - header_size, self._get_fuzzy_field('padding', padding_size, 'i16')) header_size = header_size + padding_size wsz += header_size + 10 self._write_fuzzy_field(buf, 'length', wsz, 'i32') self._write_fuzzy_field(buf, 'magic', HEADER_MAGIC >> 16, 'i16') self._write_fuzzy_field( buf, 'flags', self.header_flags(), 'i16') self._write_fuzzy_field( buf, 'seq_id', self.seq_id, 'i32') self._write_fuzzy_field(buf, 'header_size', header_size // 4, 'i16') buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) buf.write(info_data.getvalue()) # Pad out the header with 0x00 if 'padding_bytes' in self._fuzz_fields: # Print that padding bytes are being fuzzed in favor of printing # the value of each individual padding byte self._print("Fuzzing %d padding bytes" % padding_size) old_verbose, self._verbose = self._verbose, False for _ in sm.xrange(padding_size): self._write_fuzzy_field( buf, 'padding_bytes', six.int2byte(0), 'char') self._verbose = old_verbose self._write_fuzzy_field(buf, 'payload', wout, 'str')
def _serialize_string(str_): return getVarint(len(str_)) + str_
def flush(self): wout = self.__wbuf.getvalue() wout = self.transform(wout) wsz = len(wout) # reset wbuf before write/flush to preserve state on underlying failure self.__wbuf.seek(0) self.__wbuf.truncate() if self.__proto_id == 1 and self.__client_type != self.HTTP_CLIENT_TYPE: raise TTransportException( TTransportException.INVALID_CLIENT_TYPE, "Trying to send JSON encoding over binary") buf = StringIO() if self.__client_type == self.HEADERS_CLIENT_TYPE: transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) if self.__hmac_func: num_transforms += 1 transform_data.write(getVarint(self.HMAC_TRANSFORM)) transform_data.write(b'\0') # size of hmac, fixup later. # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity info_data = StringIO() # Write persistent kv-headers THeaderTransport._flush_info_headers( info_data, self.__write_persistent_headers, self.INFO_PKEYVALUE) # Write non-persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.__write_headers, self.INFO_KEYVALUE) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + \ info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size wsz += header_size + 10 buf.write(pack(b"!I", wsz)) buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags)) buf.write(pack(b"!I", self.__seq_id)) buf.write(pack(b"!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) hmac_loc = buf.tell() - 1 # Fixup hmac size later buf.write(info_data.getvalue()) # Pad out the header with 0x00 for x in range(0, padding_size, 1): buf.write(pack(b"!c", b'\0')) # Send data section buf.write(wout) # HMAC calculation should always be last. if self.__hmac_func: hmac_data = buf.getvalue()[4:] hmac = self.__hmac_func(hmac_data) # Fill in hmac size. buf.seek(hmac_loc) buf.write(chr(len(hmac))) buf.seek(0, os.SEEK_END) buf.write(hmac) # Fix packet size since we appended data. new_sz = buf.tell() - 4 buf.seek(0) buf.write(pack(b"!I", new_sz)) elif self.__client_type == self.FRAMED_DEPRECATED: buf.write(pack(b"!I", wsz)) buf.write(wout) elif self.__client_type == self.UNFRAMED_DEPRECATED: buf.write(wout) elif self.__client_type == self.HTTP_CLIENT_TYPE: # Reset the client type if we sent something - # oneway calls via HTTP expect a status response otherwise buf.write(self.header.getvalue()) buf.write(wout) self.__client_type == self.HEADERS_CLIENT_TYPE elif self.__client_type == self.UNKNOWN_CLIENT_TYPE: raise TTransportException(TTransportException.INVALID_CLIENT_TYPE, "Unknown client type") # We don't include the framing bytes as part of the frame size check # (so -4) if buf.tell() - 4 > self.__max_frame_size: raise TTransportException( TTransportException.INVALID_FRAME_SIZE, "Attempting to send frame that is too large") self.__trans.write(buf.getvalue()) self.__trans.flush()
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for self.HEADERS_CLIENT_TYPE @param buf(StringIO): Buffer to write message to @param wout(str): Payload @param wsz(int): Payload length """ transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) if self.__hmac_func: num_transforms += 1 transform_data.write(getVarint(self.HMAC_TRANSFORM)) transform_data.write(b'\0') # size of hmac, fixup later. # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity info_data = StringIO() # Write persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.__write_persistent_headers, self.INFO_PKEYVALUE) # Write non-persistent kv-headers THeaderTransport._flush_info_headers(info_data, self.__write_headers, self.INFO_KEYVALUE) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + \ info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size wsz += header_size + 10 buf.write(pack(b"!I", wsz)) buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags)) buf.write(pack(b"!I", self.__seq_id)) buf.write(pack(b"!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) hmac_loc = buf.tell() - 1 # Fixup hmac size later buf.write(info_data.getvalue()) # Pad out the header with 0x00 for x in range(0, padding_size, 1): buf.write(pack(b"!c", b'\0')) # Send data section buf.write(wout) # HMAC calculation should always be last. if self.__hmac_func: hmac_data = buf.getvalue()[4:] hmac = self.__hmac_func(hmac_data) # Fill in hmac size. buf.seek(hmac_loc) buf.write(chr(len(hmac))) buf.seek(0, os.SEEK_END) buf.write(hmac) # Fix packet size since we appended data. new_sz = buf.tell() - 4 buf.seek(0) buf.write(pack(b"!I", new_sz))
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for self.HEADERS_CLIENT_TYPE This method writes a message using the same logic as THeaderTransport._flushHeaderMessage but mutates fields included in self._fuzz_fields """ cls = self.__class__ transform_data = StringIO() num_transforms = len(self._THeaderTransport__write_transforms) for trans_id in self._THeaderTransport__write_transforms: trans_id = self._get_fuzzy_field('transform_id', trans_id, 'i32') transform_data.write(getVarint(trans_id)) if self._THeaderTransport__hmac_func: num_transforms += 1 transform_data.write(getVarint(self.HMAC_TRANSFORM)) transform_data.write(b'\0') # size of hmac, fixup later. # Add in special flags. if self._THeaderTransport__identity: id_version = self._get_fuzzy_field( 'version', self.ID_VERSION, 'str') self._THeaderTransport__write_headers[self.ID_VERSION_HEADER] = ( id_version) identity = self._get_fuzzy_field( 'identity', self._THeaderTransport__identity, 'str') self._THeaderTransport__write_headers[self.IDENTITY_HEADER] = ( identity) info_data = StringIO() # Write persistent kv-headers cls._flush_info_headers( info_data, self._THeaderTransport__write_persistent_headers, self.INFO_PKEYVALUE) # Write non-persistent kv-headers cls._flush_info_headers( info_data, self._THeaderTransport__write_headers, self.INFO_KEYVALUE) header_data = StringIO() proto_id = self._get_fuzzy_field( 'proto_id', self._THeaderTransport__proto_id, 'i32') header_data.write(getVarint(proto_id)) num_transforms = self._get_fuzzy_field( 'num_transforms', num_transforms, 'i32') header_data.write(getVarint(num_transforms)) header_size = (transform_data.tell() + header_data.tell() + info_data.tell()) padding_size = 4 - (header_size % 4) # Fuzz padding size, but do not let total header size exceed 2**16 - 1 padding_size = min((2 ** 16 - 1) - header_size, self._get_fuzzy_field('padding', padding_size, 'i16')) header_size = header_size + padding_size wsz += header_size + 10 self._write_fuzzy_field(buf, 'length', wsz, 'i32') self._write_fuzzy_field(buf, 'magic', self.HEADER_MAGIC, 'i16') self._write_fuzzy_field( buf, 'flags', self._THeaderTransport__flags, 'i16') self._write_fuzzy_field( buf, 'seq_id', self._THeaderTransport__seq_id, 'i32') self._write_fuzzy_field(buf, 'header_size', header_size // 4, 'i16') buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) hmac_loc = buf.tell() - 1 # Fixup hmac size later buf.write(info_data.getvalue()) # Pad out the header with 0x00 if 'padding_bytes' in self._fuzz_fields: # Print that padding bytes are being fuzzed in favor of printing # the value of each individual padding byte self._print("Fuzzing %d padding bytes" % padding_size) old_verbose, self._verbose = self._verbose, False for _ in sm.xrange(padding_size): self._write_fuzzy_field( buf, 'padding_bytes', six.int2byte(0), 'char') self._verbose = old_verbose self._write_fuzzy_field(buf, 'payload', wout, 'str') # HMAC calculation should always be last. if self._THeaderTransport__hmac_func: hmac_data = buf.getvalue()[4:] hmac = self._THeaderTransport__hmac_func(hmac_data) # Fill in hmac size. buf.seek(hmac_loc) self._write_hmac_size(buf, hmac) buf.seek(0, os.SEEK_END) self._write_hmac(buf, hmac) # Fix packet size since we appended data. new_sz = buf.tell() - 4 buf.seek(0) self._write_wsz(buf, new_sz)
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for self.HEADERS_CLIENT_TYPE This method writes a message using the same logic as THeaderTransport._flushHeaderMessage but mutates fields included in self._fuzz_fields """ cls = self.__class__ transform_data = StringIO() num_transforms = len(self._THeaderTransport__write_transforms) for trans_id in self._THeaderTransport__write_transforms: trans_id = self._get_fuzzy_field('transform_id', trans_id, 'i32') transform_data.write(getVarint(trans_id)) if self._THeaderTransport__hmac_func: num_transforms += 1 transform_data.write(getVarint(self.HMAC_TRANSFORM)) transform_data.write(b'\0') # size of hmac, fixup later. # Add in special flags. if self._THeaderTransport__identity: id_version = self._get_fuzzy_field('version', self.ID_VERSION, 'str') self._THeaderTransport__write_headers[self.ID_VERSION_HEADER] = ( id_version) identity = self._get_fuzzy_field('identity', self._THeaderTransport__identity, 'str') self._THeaderTransport__write_headers[self.IDENTITY_HEADER] = ( identity) info_data = StringIO() # Write persistent kv-headers cls._flush_info_headers( info_data, self._THeaderTransport__write_persistent_headers, self.INFO_PKEYVALUE) # Write non-persistent kv-headers cls._flush_info_headers(info_data, self._THeaderTransport__write_headers, self.INFO_KEYVALUE) header_data = StringIO() proto_id = self._get_fuzzy_field('proto_id', self._THeaderTransport__proto_id, 'i32') header_data.write(getVarint(proto_id)) num_transforms = self._get_fuzzy_field('num_transforms', num_transforms, 'i32') header_data.write(getVarint(num_transforms)) header_size = (transform_data.tell() + header_data.tell() + info_data.tell()) padding_size = 4 - (header_size % 4) # Fuzz padding size, but do not let total header size exceed 2**16 - 1 padding_size = min((2**16 - 1) - header_size, self._get_fuzzy_field('padding', padding_size, 'i16')) header_size = header_size + padding_size wsz += header_size + 10 self._write_fuzzy_field(buf, 'length', wsz, 'i32') self._write_fuzzy_field(buf, 'magic', self.HEADER_MAGIC, 'i16') self._write_fuzzy_field(buf, 'flags', self._THeaderTransport__flags, 'i16') self._write_fuzzy_field(buf, 'seq_id', self._THeaderTransport__seq_id, 'i32') self._write_fuzzy_field(buf, 'header_size', header_size // 4, 'i16') buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) hmac_loc = buf.tell() - 1 # Fixup hmac size later buf.write(info_data.getvalue()) # Pad out the header with 0x00 if 'padding_bytes' in self._fuzz_fields: # Print that padding bytes are being fuzzed in favor of printing # the value of each individual padding byte self._print("Fuzzing %d padding bytes" % padding_size) old_verbose, self._verbose = self._verbose, False for _ in sm.xrange(padding_size): self._write_fuzzy_field(buf, 'padding_bytes', six.int2byte(0), 'char') self._verbose = old_verbose self._write_fuzzy_field(buf, 'payload', wout, 'str') # HMAC calculation should always be last. if self._THeaderTransport__hmac_func: hmac_data = buf.getvalue()[4:] hmac = self._THeaderTransport__hmac_func(hmac_data) # Fill in hmac size. buf.seek(hmac_loc) self._write_hmac_size(buf, hmac) buf.seek(0, os.SEEK_END) self._write_hmac(buf, hmac) # Fix packet size since we appended data. new_sz = buf.tell() - 4 buf.seek(0) self._write_wsz(buf, new_sz)
def _flushHeaderMessage(self, buf, wout, wsz): """Write a message for CLIENT_TYPE.HEADER @param buf(StringIO): Buffer to write message to @param wout(str): Payload @param wsz(int): Payload length """ transform_data = StringIO() # For now, all transforms don't require data. num_transforms = len(self.__write_transforms) for trans_id in self.__write_transforms: transform_data.write(getVarint(trans_id)) # Add in special flags. if self.__identity: self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION self.__write_headers[self.IDENTITY_HEADER] = self.__identity info_data = StringIO() # Write persistent kv-headers _flush_info_headers(info_data, self.get_write_persistent_headers(), INFO.PERSISTENT) # Write non-persistent kv-headers _flush_info_headers(info_data, self.__write_headers, INFO.NORMAL) header_data = StringIO() header_data.write(getVarint(self.__proto_id)) header_data.write(getVarint(num_transforms)) header_size = transform_data.tell() + header_data.tell() + \ info_data.tell() padding_size = 4 - (header_size % 4) header_size = header_size + padding_size # MAGIC(2) | FLAGS(2) + SEQ_ID(4) + HEADER_SIZE(2) wsz += header_size + 10 if wsz > MAX_FRAME_SIZE: buf.write(pack("!I", BIG_FRAME_MAGIC)) buf.write(pack("!Q", wsz)) else: buf.write(pack("!I", wsz)) buf.write(pack("!HH", HEADER_MAGIC >> 16, self.__flags)) buf.write(pack("!I", self.seq_id)) buf.write(pack("!H", header_size // 4)) buf.write(header_data.getvalue()) buf.write(transform_data.getvalue()) buf.write(info_data.getvalue()) # Pad out the header with 0x00 for _ in range(0, padding_size, 1): buf.write(pack("!c", b'\0')) # Send data section buf.write(wout)