ida-genpatch/patchgen.py at main · frasten/ida-genpatch · 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
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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
import re
import ida_idaapi
import ida_ida
import idc
import ida_bytes
import ida_funcs
# RunPlugin("patchgen", 0)

class PatchGen(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_UNL | ida_idaapi.PLUGIN_MULTI
    comment = "Lists the patched addressess, source bytes and patches bytes to use them in custom-made patchers."
    help = "Press Alt-F8 to generate the patch code."
    wanted_name = "patchgen"
    wanted_hotkey = "Alt-F8"

    def init(self):
        print("[+] PatchGen plugin loaded. Press %s to generate the patch code." % PatchGen.wanted_hotkey)
        return PatchGenMod()


class PatchGenMod(ida_idaapi.plugmod_t):
    def run(self, arg):
        self.print_patches(self.get_patched_bytes())

    def get_patched_bytes(self, start=None, end=None):
        if start is None:
            start = ida_ida.inf_get_min_ea()
        if end is None:
            end = ida_ida.inf_get_max_ea()

        patched_bytes = dict()

        def collector(ea, fpos, original, patched):
            patched_bytes[ea] = PatchData(ea, fpos, original, patched)
            return 0

        ida_bytes.visit_patched_bytes(start, end, collector)
        return patched_bytes

    def h(self, num):
        # print('0x{0:08X}'.format(ida_loader.get_fileregion_offset(ScreenEA())))
        return '0x{:08X}'.format(num)

    def h_bytes(self, bytes):
        return ", ".join("0x{:02X}".format(x) for x in bytes)

    def group_patches(self, patches):
        last_addr = 0
        current_group = PatchGroup()

        patch_groups = []
        sorted_addresses = sorted(patches.keys()) # Dictionaries aren't sorted by default
        for addr in sorted_addresses:
            if addr > last_addr + 1 and current_group.length() > 0:
                # Not in sequence:
                patch_groups.append(current_group)
                current_group = PatchGroup()

            bytePatch = patches[addr]
            current_group.append(bytePatch)
            last_addr = addr

        if current_group.length() > 0:
            # The last group:
            patch_groups.append(current_group)
        return patch_groups

    def print_patches(self, patches):
        patch_groups = self.group_patches(patches)

        print()
        print("/" * 20 + " %d Patches " % (len(patch_groups)) + "/" * 20)
        print()
        for chunk in patch_groups:
            patched_instructions = " | ".join(chunk.disasm_patched())
            original_instructions = " | ".join(chunk.disasm_original())

            patch_descr = "// %s(): [%s]  ==>  [%s]" % (chunk.func_name(), original_instructions, patched_instructions)
            print(patch_descr)

            fpos_str = self.h(chunk.fpos())
            orig_str = self.h_bytes(chunk.original())
            patch_str = self.h_bytes(chunk.patched())

            # print "%s: %s => %s" % (fpos_str, orig_str, patch_str)

            if chunk.length() == 1:
                print("Hunks.Add(new SinglePatchHunk(%s, %s, %s));" % (fpos_str, orig_str, patch_str))
            else:
                print("Hunks.Add(new SinglePatchHunk(%s, [%s], [%s]));" % (fpos_str, orig_str, patch_str))
            print()


class PatchData:
    def __init__(self, ea, fpos, original, patched):
        self.ea = ea
        self.fpos = fpos
        self.original = original
        self.patched = patched


class PatchGroup:
    def __init__(self):
        self.bytes = []

    def append(self, byte_patch):
        self.bytes.append(byte_patch)

    def ea(self):
        return self.bytes[0].ea

    def fpos(self):
        return self.bytes[0].fpos

    def func_name(self):
        return ida_funcs.get_func_name(self.ea())

    def length(self):
        return len(self.bytes)

    def disasm_patched(self):
        result = []

        offset = 0
        while True:
            if offset >= self.length():
                break
            instruction_addr = self.ea() + offset
            disasm = idc.GetDisasm(instruction_addr)
            disasm = re.sub("\\s+", " ", disasm)
            result.append(disasm)

            offset += ida_bytes.get_item_size(instruction_addr)
        return result

    def disasm_original(self):
        self.revert_patch()
        disasm = self.disasm_patched()
        self.apply_patch()
        return disasm

    def revert_patch(self):
        for b in self.bytes:
            ida_bytes.patch_byte(b.ea, b.original)

    def apply_patch(self):
        for b in self.bytes:
            ida_bytes.patch_byte(b.ea, b.patched)

    def original(self):
        return map(lambda b: b.original, self.bytes)

    def patched(self):
        return map(lambda b: b.patched, self.bytes)


def PLUGIN_ENTRY():
    return PatchGen()