about summary refs log tree commit diff stats
path: root/src/miasm/loader/strpatchwork.py
blob: e1a5de91f12a2d071e71671885ae8583a2556c3f (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
from array import array
import struct
from sys import maxsize

from future.utils import PY3

if PY3:

    def array_frombytes(arr, value):
        return arr.frombytes(value)

    def array_tobytes(arr):
        return arr.tobytes()


else:

    def array_frombytes(arr, value):
        return arr.fromstring(value)

    def array_tobytes(arr):
        return arr.tostring()


class StrPatchwork(object):

    def __init__(self, s=b"", paddingbyte=b"\x00"):
        s_raw = bytes(s)
        val = array("B")
        array_frombytes(val, s_raw)
        self.s = val
        # cache s to avoid rebuilding str after each find
        self.s_cache = s_raw
        self.paddingbyte = paddingbyte

    def __bytes__(self):
        return array_tobytes(self.s)

    def __str__(self):
        if PY3:
            return repr(self)
        return self.__bytes__()

    def __getitem__(self, item):
        s = self.s
        if isinstance(item, slice):
            end = item.stop
            l = len(s)
            if (end is not None and l < end) and end != maxsize:
                # XXX hack [x:] give 2GB limit
                # This is inefficient but avoids complicated maths if step is
                # not 1
                s = s[:]

                tmp = array("B")
                array_frombytes(tmp, self.paddingbyte * (end - l))
                s.extend(tmp)
            r = s[item]
            return array_tobytes(r)

        else:
            if item > len(s):
                return self.paddingbyte
            else:
                return struct.pack("B", s[item])

    def __setitem__(self, item, val):
        if val is None:
            return
        val_array = array("B")
        array_frombytes(val_array, bytes(val))
        if type(item) is not slice:
            item = slice(item, item + len(val_array))
        end = item.stop
        l = len(self.s)
        if l < end:
            tmp = array("B")
            array_frombytes(tmp, self.paddingbyte * (end - l))
            self.s.extend(tmp)
        self.s[item] = val_array
        self.s_cache = None

    def __repr__(self):
        return "<Patchwork %r>" % array_tobytes(self.s)

    def __len__(self):
        return len(self.s)

    def __contains__(self, val):
        return val in bytes(self)

    def __iadd__(self, other):
        tmp = array("B")
        array_frombytes(tmp, bytes(other))
        self.s.extend(tmp)
        return self

    def find(self, pattern, start=0, end=None):
        if not self.s_cache:
            self.s_cache = array_tobytes(self.s)
        return self.s_cache.find(pattern, start, end)

    def rfind(self, pattern, start=0, end=None):
        if not self.s_cache:
            self.s_cache = array_tobytes(self.s)
        return self.s_cache.rfind(pattern, start, end)