striga/container.py at master · LLVMParty/striga · GitHub

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
from pefile import PE


class Container:
    @property
    def image_base(self) -> int:
        raise NotImplementedError()

    @property
    def image_size(self) -> int:
        raise NotImplementedError()

    def in_range(self, va: int):
        return va >= self.image_base and va < self.image_base + self.image_size

    def get_data(self, va: int, size: int) -> bytes:
        raise NotImplementedError()


class RawContainer(Container):
    def __init__(self, data: bytes, image_base=0):
        self.data = data
        self._image_base = image_base
        self._image_size = len(data)

    @property
    def image_base(self) -> int:
        return self._image_base

    @property
    def image_size(self) -> int:
        return self._image_size

    def get_data(self, va: int, size: int) -> bytes:
        assert self.in_range(va)
        rva = va - self.image_base
        return self.data[rva : rva + size]


class PEContainer(Container):
    def __init__(self, pe: str | PE) -> None:
        if isinstance(pe, str):
            pe = PE(pe)
        self.pe = pe
        self._image_base = pe.OPTIONAL_HEADER.ImageBase  # pyright: ignore
        self._image_size = pe.OPTIONAL_HEADER.SizeOfImage  # pyright: ignore

    @property
    def image_base(self) -> int:
        return self._image_base

    @property
    def image_size(self) -> int:
        return self._image_size

    def get_data(self, va: int, size: int) -> bytes:
        assert self.in_range(va)
        rva = va - self.image_base
        return self.pe.get_data(rva, size)