def test_generate_path_leaves(self): data = open(DATA, 'rb').read() nodes, bits = PathTree._unpack_data(data) ret = PathTree._generate_path_leaves(GhettoBitStream(bits)) self.assertEqual(len(ret), 4) for node in ret: self.assertTrue(isinstance(node, HuffmanNode))
def __init__(self, data): """ Uncompresses data into a tree that can be traversed for matching paths :param data: binary data as read from a file or pulled directly out of a certificate extension. Data should be compressed with huffman coding as described for v3 entitlement certificates :type data: binary string """ word_leaves, unused_bits = self._unpack_data(data) HuffmanNode.build_tree(word_leaves) word_dict = dict((node.code, node.value) for node in word_leaves) bitstream = GhettoBitStream(unused_bits) path_leaves = self._generate_path_leaves(bitstream) HuffmanNode.build_tree(path_leaves) path_dict = dict((node.code, node) for node in path_leaves) self.path_tree = self._generate_path_tree(path_dict, path_leaves, word_dict, bitstream)
def test_get_node_count_big(self): bs = GhettoBitStream([]) # count bigger than 127, need next 2 bytes to represent it bs.bytes = deque([130, 1, 17]) ret = PathTree._get_node_count(bs) self.assertEqual(ret, 273)
def test_get_node_count_medium(self): bs = GhettoBitStream([]) # count bigger than 127, only need 1 byte to represent it bs.bytes = deque([129, 150]) ret = PathTree._get_node_count(bs) self.assertEqual(ret, 150)
def test_get_node_count_small(self): bs = GhettoBitStream([]) bs.bytes = deque([6]) ret = PathTree._get_node_count(bs) self.assertEqual(ret, 6)
def setUp(self): self.bs = GhettoBitStream(tree_data)