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))
class TestGhettoBitStream(unittest.TestCase):
def setUp(self):
self.bs = GhettoBitStream(tree_data)
def test_pop_byte(self):
length = len(self.bs.bytes)
first = self.bs.pop_byte()
self.assertEqual(first, 5)
self.assertEqual(len(self.bs.bytes), length - 1)
def test_as_iterator(self):
next(self.bs)
self.assertTrue(list(self.bs))
def test_bit_buffer(self):
byte_count = len(self.bs.bytes)
# empty buffer
self.assertEqual(len(self.bs._bit_buffer), 0)
next(self.bs)
# one byte decoded, then one bit consumed
self.assertEqual(len(self.bs._bit_buffer), 7)
# one byte removed
self.assertEqual(len(self.bs.bytes), byte_count - 1)
next(self.bs)
# another bit consumed
self.assertEqual(len(self.bs._bit_buffer), 6)
# remaining bytes still stand
self.assertEqual(len(self.bs.bytes), byte_count - 1)
def test_byte_to_bits(self):
# just spot-checking
self.assertEqual(self.bs._byte_to_bits(0), '00000000')
self.assertEqual(self.bs._byte_to_bits(6), '00000110')
self.assertEqual(self.bs._byte_to_bits(213), '11010101')
def test_bin_backport(self):
# just spot-checking
self.assertEqual(self.bs._bin_backport(0), '00000000')
self.assertEqual(self.bs._bin_backport(6), '00000110')
self.assertEqual(self.bs._bin_backport(213), '11010101')
def test_combine_bytes(self):
# just spot-checking
self.assertEqual(self.bs.combine_bytes([1, 3]), 259)
self.assertEqual(self.bs.combine_bytes([3]), 3)
self.assertEqual(self.bs.combine_bytes([1, 1, 3]), 65795)
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 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)
def setUp(self):
self.bs = GhettoBitStream(tree_data)