1#!/usr/bin/env python
2
3import string
4import struct
5import sys
6
7
8class FileExtract:
9    '''Decode binary data from a file'''
10
11    def __init__(self, f, b='='):
12        '''Initialize with an open binary file and optional byte order'''
13
14        self.file = f
15        self.byte_order = b
16        self.offsets = list()
17
18    def set_byte_order(self, b):
19        '''Set the byte order, valid values are "big", "little", "swap", "native", "<", ">", "@", "="'''
20        if b == 'big':
21            self.byte_order = '>'
22        elif b == 'little':
23            self.byte_order = '<'
24        elif b == 'swap':
25            # swap what ever the current byte order is
26            self.byte_order = swap_unpack_char()
27        elif b == 'native':
28            self.byte_order = '='
29        elif b == '<' or b == '>' or b == '@' or b == '=':
30            self.byte_order = b
31        else:
32            print("error: invalid byte order specified: '%s'" % b)
33
34    def is_in_memory(self):
35        return False
36
37    def seek(self, offset, whence=0):
38        if self.file:
39            return self.file.seek(offset, whence)
40        raise ValueError
41
42    def tell(self):
43        if self.file:
44            return self.file.tell()
45        raise ValueError
46
47    def read_size(self, byte_size):
48        s = self.file.read(byte_size)
49        if len(s) != byte_size:
50            return None
51        return s
52
53    def push_offset_and_seek(self, offset):
54        '''Push the current file offset and seek to "offset"'''
55        self.offsets.append(self.file.tell())
56        self.file.seek(offset, 0)
57
58    def pop_offset_and_seek(self):
59        '''Pop a previously pushed file offset, or do nothing if there were no previously pushed offsets'''
60        if len(self.offsets) > 0:
61            self.file.seek(self.offsets.pop())
62
63    def get_sint8(self, fail_value=0):
64        '''Extract a single int8_t from the binary file at the current file position, returns a single integer'''
65        s = self.read_size(1)
66        if s:
67            v, = struct.unpack(self.byte_order + 'b', s)
68            return v
69        else:
70            return fail_value
71
72    def get_uint8(self, fail_value=0):
73        '''Extract a single uint8_t from the binary file at the current file position, returns a single integer'''
74        s = self.read_size(1)
75        if s:
76            v, = struct.unpack(self.byte_order + 'B', s)
77            return v
78        else:
79            return fail_value
80
81    def get_sint16(self, fail_value=0):
82        '''Extract a single int16_t from the binary file at the current file position, returns a single integer'''
83        s = self.read_size(2)
84        if s:
85            v, = struct.unpack(self.byte_order + 'h', s)
86            return v
87        else:
88            return fail_value
89
90    def get_uint16(self, fail_value=0):
91        '''Extract a single uint16_t from the binary file at the current file position, returns a single integer'''
92        s = self.read_size(2)
93        if s:
94            v, = struct.unpack(self.byte_order + 'H', s)
95            return v
96        else:
97            return fail_value
98
99    def get_sint32(self, fail_value=0):
100        '''Extract a single int32_t from the binary file at the current file position, returns a single integer'''
101        s = self.read_size(4)
102        if s:
103            v, = struct.unpack(self.byte_order + 'i', s)
104            return v
105        else:
106            return fail_value
107
108    def get_uint32(self, fail_value=0):
109        '''Extract a single uint32_t from the binary file at the current file position, returns a single integer'''
110        s = self.read_size(4)
111        if s:
112            v, = struct.unpack(self.byte_order + 'I', s)
113            return v
114        else:
115            return fail_value
116
117    def get_sint64(self, fail_value=0):
118        '''Extract a single int64_t from the binary file at the current file position, returns a single integer'''
119        s = self.read_size(8)
120        if s:
121            v, = struct.unpack(self.byte_order + 'q', s)
122            return v
123        else:
124            return fail_value
125
126    def get_uint64(self, fail_value=0):
127        '''Extract a single uint64_t from the binary file at the current file position, returns a single integer'''
128        s = self.read_size(8)
129        if s:
130            v, = struct.unpack(self.byte_order + 'Q', s)
131            return v
132        else:
133            return fail_value
134
135    def get_fixed_length_c_string(
136            self,
137            n,
138            fail_value='',
139            isprint_only_with_space_padding=False):
140        '''Extract a single fixed length C string from the binary file at the current file position, returns a single C string'''
141        s = self.read_size(n)
142        if s:
143            cstr, = struct.unpack(self.byte_order + ("%i" % n) + 's', s)
144            # Strip trialing NULLs
145            cstr = string.strip(cstr, "\0")
146            if isprint_only_with_space_padding:
147                for c in cstr:
148                    if c in string.printable or ord(c) == 0:
149                        continue
150                    return fail_value
151            return cstr
152        else:
153            return fail_value
154
155    def get_c_string(self):
156        '''Extract a single NULL terminated C string from the binary file at the current file position, returns a single C string'''
157        cstr = ''
158        byte = self.get_uint8()
159        while byte != 0:
160            cstr += "%c" % byte
161            byte = self.get_uint8()
162        return cstr
163
164    def get_n_sint8(self, n, fail_value=0):
165        '''Extract "n" int8_t integers from the binary file at the current file position, returns a list of integers'''
166        s = self.read_size(n)
167        if s:
168            return struct.unpack(self.byte_order + ("%u" % n) + 'b', s)
169        else:
170            return (fail_value,) * n
171
172    def get_n_uint8(self, n, fail_value=0):
173        '''Extract "n" uint8_t integers from the binary file at the current file position, returns a list of integers'''
174        s = self.read_size(n)
175        if s:
176            return struct.unpack(self.byte_order + ("%u" % n) + 'B', s)
177        else:
178            return (fail_value,) * n
179
180    def get_n_sint16(self, n, fail_value=0):
181        '''Extract "n" int16_t integers from the binary file at the current file position, returns a list of integers'''
182        s = self.read_size(2 * n)
183        if s:
184            return struct.unpack(self.byte_order + ("%u" % n) + 'h', s)
185        else:
186            return (fail_value,) * n
187
188    def get_n_uint16(self, n, fail_value=0):
189        '''Extract "n" uint16_t integers from the binary file at the current file position, returns a list of integers'''
190        s = self.read_size(2 * n)
191        if s:
192            return struct.unpack(self.byte_order + ("%u" % n) + 'H', s)
193        else:
194            return (fail_value,) * n
195
196    def get_n_sint32(self, n, fail_value=0):
197        '''Extract "n" int32_t integers from the binary file at the current file position, returns a list of integers'''
198        s = self.read_size(4 * n)
199        if s:
200            return struct.unpack(self.byte_order + ("%u" % n) + 'i', s)
201        else:
202            return (fail_value,) * n
203
204    def get_n_uint32(self, n, fail_value=0):
205        '''Extract "n" uint32_t integers from the binary file at the current file position, returns a list of integers'''
206        s = self.read_size(4 * n)
207        if s:
208            return struct.unpack(self.byte_order + ("%u" % n) + 'I', s)
209        else:
210            return (fail_value,) * n
211
212    def get_n_sint64(self, n, fail_value=0):
213        '''Extract "n" int64_t integers from the binary file at the current file position, returns a list of integers'''
214        s = self.read_size(8 * n)
215        if s:
216            return struct.unpack(self.byte_order + ("%u" % n) + 'q', s)
217        else:
218            return (fail_value,) * n
219
220    def get_n_uint64(self, n, fail_value=0):
221        '''Extract "n" uint64_t integers from the binary file at the current file position, returns a list of integers'''
222        s = self.read_size(8 * n)
223        if s:
224            return struct.unpack(self.byte_order + ("%u" % n) + 'Q', s)
225        else:
226            return (fail_value,) * n
227