netarch

unpack.py

 1#! /usr/bin/python3
 2
 3ENDIAN_LITTLE = 1
 4ENDIAN_BIG = 2
 5ENDIAN_MIDDLE = 3
 6ENDIAN_NETWORK = ENDIAN_BIG
 7
 8class Unpacker:
 9    """Class that lets you peel values off
10    
11    >>> u = Unpacker(bytes((1, 0,2, 0,0,0,3, 0,0,0,0,0,0,0,4)))
12    >>> u.uint8()
13    1
14    >>> u.uint16()
15    2
16    >>> u.uint32()
17    3
18    >>> u.uint64()
19    4
20    
21    >>> u = Unpacker(bytes((1,0, 104,105)), ENDIAN_LITTLE)
22    >>> u.uint16()
23    1
24    >>> u.buf
25    b'hi'
26
27    >>> u = Unpacker(bytes((1,0, 0,2)))
28    >>> u.uint16(ENDIAN_LITTLE)
29    1
30    >>> u.uint(16, ENDIAN_BIG)
31    2
32    
33    >>> u = Unpacker(bytes((0,1,2,3)), ENDIAN_MIDDLE)
34    >>> '%08x' % u.uint32()
35    '01000302'
36    """
37    
38    def __init__(self, buf, endian=ENDIAN_NETWORK):
39        self.endian = endian
40        self.buf = buf
41
42    def uint(self, size, pos=0, endian=None):
43        endian = endian or self.endian
44        if size not in (8, 16, 32, 64):
45            # XXX: I'm pretty sure this can be done, but I don't want to code it up right now.
46            raise ValueError("Can't do weird sizes")
47        noctets = size // 8
48        if endian == ENDIAN_BIG:
49            r = range(0, noctets)
50        elif endian == ENDIAN_LITTLE:
51            r = range(noctets-1, -1, -1)
52        elif endian == ENDIAN_MIDDLE:
53            r = (1, 0, 3, 2,   5, 4, 7, 6)[:noctets]
54        else:
55            raise ValueError("Unsupported byte order")
56        if pos < 0:
57            pos += len(self.buf)
58        pull, self.buf = self.buf[pos:pos+noctets], self.buf[:pos] + self.buf[pos+noctets:]
59        acc = 0
60        for i in r:
61            acc = (acc << 8) | pull[i]
62        return acc
63        
64    def uint8(self, pos=0):
65        return self.uint(8, pos)
66    def uint16(self, pos=0, endian=None):
67        return self.uint(16, pos, endian)
68    def uint32(self, pos=0, endian=None):
69        return self.uint(32, pos, endian)
70    def uint64(self, pos=0, endian=None):
71        return self.uint(64, pos, endian)
72
73if __name__ == "__main__":
74    import doctest
75    doctest.testmod()