save context for Insert*Patch

src/patcherex2/components/archinfo/aarch64.py

@@ -4,3 +4,41 @@ class Aarch64Info:
     jmp_asm = "b {dst}"
     jmp_size = 4
     call_asm = "bl {dst}"
+    save_context_asm = """
+        sub sp, sp, #0x1f0
+        stp x0, x1, [sp, #0x0]
+        stp x2, x3, [sp, #0x10]
+        stp x4, x5, [sp, #0x20]
+        stp x6, x7, [sp, #0x30]
+        stp x8, x9, [sp, #0x40]
+        stp x10, x11, [sp, #0x50]
+        stp x12, x13, [sp, #0x60]
+        stp x14, x15, [sp, #0x70]
+        stp x16, x17, [sp, #0x80]
+        stp x18, x19, [sp, #0x90]
+        stp x20, x21, [sp, #0xa0]
+        stp x22, x23, [sp, #0xb0]
+        stp x24, x25, [sp, #0xc0]
+        stp x26, x27, [sp, #0xd0]
+        stp x28, x29, [sp, #0xe0]
+        str x30, [sp, #0xf0]
+    """
+    restore_context_asm = """
+        ldp x0, x1, [sp, #0x0]
+        ldp x2, x3, [sp, #0x10]
+        ldp x4, x5, [sp, #0x20]
+        ldp x6, x7, [sp, #0x30]
+        ldp x8, x9, [sp, #0x40]
+        ldp x10, x11, [sp, #0x50]
+        ldp x12, x13, [sp, #0x60]
+        ldp x14, x15, [sp, #0x70]
+        ldp x16, x17, [sp, #0x80]
+        ldp x18, x19, [sp, #0x90]
+        ldp x20, x21, [sp, #0xa0]
+        ldp x22, x23, [sp, #0xb0]
+        ldp x24, x25, [sp, #0xc0]
+        ldp x26, x27, [sp, #0xd0]
+        ldp x28, x29, [sp, #0xe0]
+        ldr x30, [sp, #0xf0]
+        add sp, sp, #0x1f0
+    """

src/patcherex2/components/archinfo/amd64.py

@@ -4,3 +4,39 @@ class Amd64Info:
     jmp_asm = "jmp {dst}"
     jmp_size = 6
     call_asm = "call {dst}"
+    save_context_asm = """
+    push rax
+    push rbx
+    push rcx
+    push rdx
+    push rsi
+    push rdi
+    push rbp
+    push rsp
+    push r8
+    push r9
+    push r10
+    push r11
+    push r12
+    push r13
+    push r14
+    push r15
+    """
+    restore_context_asm = """
+    pop r15
+    pop r14
+    pop r13
+    pop r12
+    pop r11
+    pop r10
+    pop r9
+    pop r8
+    pop rsp
+    pop rbp
+    pop rdi
+    pop rsi
+    pop rdx
+    pop rcx
+    pop rbx
+    pop rax
+    """

src/patcherex2/components/archinfo/arm.py

@@ -4,3 +4,9 @@ class ArmInfo:
     jmp_asm = "b {dst}"
     jmp_size = 4
     call_asm = "bl {dst}"
+    save_context_asm = """
+        push {r0-r11}
+    """
+    restore_context_asm = """
+        pop {r0-r11}
+    """

src/patcherex2/components/archinfo/mips.py

@@ -5,3 +5,71 @@ class MipsInfo:
     # NOTE: keystone will always add nop for branch delay slot, so include it in size
     jmp_size = 8
     call_asm = "jal {dst}"
+    save_context_asm = """
+    sub $sp, $sp, -124
+    sw $ra, 120($sp)
+    sw $s0, 116($sp)
+    sw $s1, 112($sp)
+    sw $s2, 108($sp)
+    sw $s3, 104($sp)
+    sw $s4, 100($sp)
+    sw $s5, 96($sp)
+    sw $s6, 92($sp)
+    sw $s7, 88($sp)
+    sw $s8, 84($sp)
+    sw $s9, 80($sp)
+    sw $s10, 76($sp)
+    sw $s11, 72($sp)
+    sw $s12, 68($sp)
+    sw $s13, 64($sp)
+    sw $s14, 60($sp)
+    sw $s15, 56($sp)
+    sw $s16, 52($sp)
+    sw $s17, 48($sp)
+    sw $s18, 44($sp)
+    sw $s19, 40($sp)
+    sw $s20, 36($sp)
+    sw $s21, 32($sp)
+    sw $s22, 28($sp)
+    sw $s23, 24($sp)
+    sw $s24, 20($sp)
+    sw $s25, 16($sp)
+    sw $s26, 12($sp)
+    sw $s27, 8($sp)
+    sw $s28, 4($sp)
+    sw $s29, 0($sp)
+    """
+    restore_context_asm = """
+    lw $s29, 0($sp)
+    lw $s28, 4($sp)
+    lw $s27, 8($sp)
+    lw $s26, 12($sp)
+    lw $s25, 16($sp)
+    lw $s24, 20($sp)
+    lw $s23, 24($sp)
+    lw $s22, 28($sp)
+    lw $s21, 32($sp)
+    lw $s20, 36($sp)
+    lw $s19, 40($sp)
+    lw $s18, 44($sp)
+    lw $s17, 48($sp)
+    lw $s16, 52($sp)
+    lw $s15, 56($sp)
+    lw $s14, 60($sp)
+    lw $s13, 64($sp)
+    lw $s12, 68($sp)
+    lw $s11, 72($sp)
+    lw $s10, 76($sp)
+    lw $s9, 80($sp)
+    lw $s8, 84($sp)
+    lw $s7, 88($sp)
+    lw $s6, 92($sp)
+    lw $s5, 96($sp)
+    lw $s4, 100($sp)
+    lw $s3, 104($sp)
+    lw $s2, 108($sp)
+    lw $s1, 112($sp)
+    lw $s0, 116($sp)
+    lw $ra, 120($sp)
+    add $sp, $sp, 124
+    """

src/patcherex2/components/archinfo/mips64.py

@@ -2,6 +2,74 @@ class Mips64Info:
     nop_bytes = b"\x00\x00\x00\x00"
     nop_size = 4
     jmp_asm = "j {dst}"
-    # NOTE: keystone will always add nop for branch delay slot, so include it in size
+    # NOTE: keystone will aldays add nop for branch delay slot, so include it in size
     jmp_size = 8
     call_asm = "jal {dst}"
+    save_context_asm = """
+    sub $sp, $sp, -248
+    sd $ra, 240($sp)
+    sd $s0, 232($sp)
+    sd $s1, 224($sp)
+    sd $s2, 216($sp)
+    sd $s3, 208($sp)
+    sd $s4, 200($sp)
+    sd $s5, 192($sp)
+    sd $s6, 184($sp)
+    sd $s7, 176($sp)
+    sd $s8, 168($sp)
+    sd $s9, 160($sp)
+    sd $s10, 152($sp)
+    sd $s11, 144($sp)
+    sd $s12, 136($sp)
+    sd $s13, 128($sp)
+    sd $s14, 120($sp)
+    sd $s15, 112($sp)
+    sd $s16, 104($sp)
+    sd $s17, 96($sp)
+    sd $s18, 88($sp)
+    sd $s19, 80($sp)
+    sd $s20, 72($sp)
+    sd $s21, 64($sp)
+    sd $s22, 56($sp)
+    sd $s23, 48($sp)
+    sd $s24, 40($sp)
+    sd $s25, 32($sp)
+    sd $s26, 24($sp)
+    sd $s27, 16($sp)
+    sd $s28, 8($sp)
+    sd $s29, 0($sp)
+    """
+    restore_context_asm = """
+    ld $s29, 0($sp)
+    ld $s28, 8($sp)
+    ld $s27, 16($sp)
+    ld $s26, 24($sp)
+    ld $s25, 32($sp)
+    ld $s24, 40($sp)
+    ld $s23, 48($sp)
+    ld $s22, 56($sp)
+    ld $s21, 64($sp)
+    ld $s20, 72($sp)
+    ld $s19, 80($sp)
+    ld $s18, 88($sp)
+    ld $s17, 96($sp)
+    ld $s16, 104($sp)
+    ld $s15, 112($sp)
+    ld $s14, 120($sp)
+    ld $s13, 128($sp)
+    ld $s12, 136($sp)
+    ld $s11, 144($sp)
+    ld $s10, 152($sp)
+    ld $s9, 160($sp)
+    ld $s8, 168($sp)
+    ld $s7, 176($sp)
+    ld $s6, 184($sp)
+    ld $s5, 192($sp)
+    ld $s4, 200($sp)
+    ld $s3, 208($sp)
+    ld $s2, 216($sp)
+    ld $s1, 224($sp)
+    ld $s0, 232($sp)
+    ld $ra, 240($sp)
+    add $sp, $sp, 248
+    """

src/patcherex2/components/archinfo/ppc.py

@@ -4,3 +4,11 @@ class PpcInfo:
     jmp_asm = "b {dst}"
     jmp_size = 4
     call_asm = "bl {dst}"
+    save_context_asm = """
+    stwu r1, -0x80(r1)
+    stmw r3, 0x8(r1)
+    """
+    restore_context_asm = """
+    lmw r3, 0x8(r1)
+    addi r1, r1, 0x80
+    """

src/patcherex2/components/archinfo/ppc64.py

@@ -4,3 +4,11 @@ class Ppc64Info:
     jmp_asm = "b {dst}"
     jmp_size = 4
     call_asm = "bl {dst}"
+    save_context_asm = """
+    stwu r1, -0x80(r1)
+    stmw r3, 0x8(r1)
+    """
+    restore_context_asm = """
+    lmw r3, 0x8(r1)
+    addi r1, r1, 0x80
+    """

src/patcherex2/components/archinfo/ppc_vle.py

@@ -4,3 +4,5 @@ class PpcVleInfo:
     jmp_asm = "b {dst}"
     jmp_size = 4
     call_asm = "bl {dst}"
+    save_context_asm = ""  # TODO
+    restore_context_asm = ""  # TODO

src/patcherex2/components/archinfo/sparc.py

@@ -4,3 +4,5 @@ class SparcInfo:
     jmp_asm = "b {dst}\nnop"  # nop due to delay slot
     jmp_size = 8
     call_asm = "call {dst}"
+    save_context_asm = ""  # TODO
+    restore_context_asm = ""  # TODO

src/patcherex2/components/archinfo/x86.py

@@ -4,3 +4,9 @@ class X86Info:
     jmp_asm = "jmp {dst}"
     jmp_size = 5
     call_asm = "call {dst}"
+    save_context_asm = """
+    pusha
+    """
+    restore_context_asm = """
+    popa
+    """

src/patcherex2/patches/function_patches.py

@@ -93,18 +93,42 @@ def __init__(
         self.prefunc = kwargs["prefunc"] if "prefunc" in kwargs else None
         self.postfunc = kwargs["postfunc"] if "postfunc" in kwargs else None
         self.compile_opts = kwargs["compile_opts"] if "compile_opts" in kwargs else {}
+        self.save_context = (
+            kwargs["save_context"] if "save_context" in kwargs else False
+        )
 
     def apply(self, p) -> None:
         if self.addr:
+            if self.prefunc:
+                if "SAVE_CONTEXT" in self.prefunc:
+                    self.prefunc = self.prefunc.replace(
+                        "SAVE_CONTEXT", f"\n{p.archinfo.save_context_asm}\n"
+                    )
+                if "RESTORE_CONTEXT" in self.prefunc:
+                    self.prefunc = self.prefunc.replace(
+                        "RESTORE_CONTEXT", f"\n{p.archinfo.restore_context_asm}\n"
+                    )
+            if self.postfunc:
+                if "SAVE_CONTEXT" in self.postfunc:
+                    self.postfunc = self.postfunc.replace(
+                        "SAVE_CONTEXT", f"\n{p.archinfo.save_context_asm}\n"
+                    )
+                if "RESTORE_CONTEXT" in self.postfunc:
+                    self.postfunc = self.postfunc.replace(
+                        "RESTORE_CONTEXT", f"\n{p.archinfo.restore_context_asm}\n"
+                    )
             ifp = InsertFunctionPatch(f"__patcherex_{hex(self.addr)}", self.code)
             ifp.apply(p)
-            instrs = self.prefunc if self.prefunc else ""
+            instrs = ""
+            instrs += p.archinfo.save_context_asm if self.save_context else ""
+            instrs += self.prefunc if self.prefunc else ""
             instrs += "\n"
             instrs += p.archinfo.call_asm.format(
                 dst=f"{{__patcherex_{hex(self.addr)}}}"
             )
             instrs += "\n"
             instrs += self.postfunc if self.postfunc else ""
+            instrs += p.archinfo.restore_context_asm if self.save_context else ""
             p.utils.insert_trampoline_code(
                 self.addr,
                 instrs,

src/patcherex2/patches/instruction_patches.py

@@ -36,6 +36,7 @@ def __init__(
         detour_pos=-1,
         symbols: Optional[Dict[str, int]] = None,
         is_thumb=False,
+        **kwargs,
     ) -> None:
         self.addr = None
         self.name = None
@@ -48,9 +49,22 @@ def __init__(
         self.detour_pos = detour_pos
         self.symbols = symbols if symbols else {}
         self.is_thumb = is_thumb
+        self.save_context = (
+            kwargs["save_context"] if "save_context" in kwargs else False
+        )
 
     def apply(self, p) -> None:
         if self.addr:
+            if "SAVE_CONTEXT" in self.instr:
+                self.instr = self.instr.replace(
+                    "SAVE_CONTEXT", f"\n{p.archinfo.save_context_asm}\n"
+                )
+            if "RESTORE_CONTEXT" in self.instr:
+                self.instr = self.instr.replace(
+                    "RESTORE_CONTEXT", f"\n{p.archinfo.restore_context_asm}\n"
+                )
+            if self.save_context:
+                self.instr = f"{p.archinfo.save_context_asm}\n{self.instr}\n{p.archinfo.restore_context_asm}"
             p.utils.insert_trampoline_code(
                 self.addr,
                 self.instr,