first commit

src/binary_patch.py

@@ -0,0 +1,179 @@
+#!/usr/bin/python3
+import lief
+from pwn import *
+import os
+import sys
+
+global lief_ELF_ALLOC 
+lief_ELF_ALLOC = 2
+global lief_ELF_EXCLUDE 
+lief_ELF_EXECINSTR = 4
+
+def get_binary_file_CLASS(binary, output_info=True):
+  CLASS = CLASS = binary.header.identity_class
+  str_CLASS = ""
+  if CLASS == binary.header.CLASS.ELF32:
+    str_CLASS = "ELF32"
+  elif CLASS == binary.header.CLASS.ELF64:
+    str_CLASS = "ELF64"
+  else:
+    str_CLASS = "UNKNOWN"
+  if output_info:
+    print("[\033[1;34m*\033[0m] CLASS is %s" % (str_CLASS))
+  return (CLASS, str_CLASS)
+
+def get_binary_file_machine_type(binary, output_info=True):
+  machine_type = binary.header.machine_type
+  str_machine_type = ""
+  arch = ""
+  if machine_type == lief._lief.ELF.ARCH.X86_64:
+    str_machine_type = "x86_64"
+    arch = "amd64"
+  else:
+    str_machine_type = "UNKNOWN"
+    arch = "UNKNOWN"
+  
+  print('[\033[1;34m*\033[0m] machine type is %s ==> ARCH : %s' % (str_machine_type, arch))
+  return (str_machine_type, arch)
+
+def load_binary_file_information(path):
+  lief_binary = lief.parse(path)
+  CLASS, str_CLASS = get_binary_file_CLASS(lief_binary)
+  str_machine_type, arch = get_binary_file_machine_type(lief_binary)
+  
+  pwn_binary = ''
+  context.arch = arch
+  if "ELF" in str_CLASS:
+    context.os = "linux"
+    pwn_binary = ELF(path)
+  
+  return (lief_binary, pwn_binary)
+
+# flag : lief.ELF.SEGMENT_FLAGS.PF_R | lief.ELF.SEGMENT_FLAGS.PF_W | lief.ELF.SEGMENT_FLAGS.PF_X
+def add_segment(lief_binary, content, types, flags, base=0x405000):
+  segment              = lief.ELF.Segment()
+  segment.type         = types
+  segment.FLAGS.from_value(flags)
+  segment.content      = list(content)
+  segment.alignment    = 8
+  segment.add(lief._lief.ELF.Segment.FLAGS.R | lief._lief.ELF.Segment.FLAGS.X)
+  segment              = lief_binary.add(segment, base=base)
+  print(segment.FLAGS.value)
+  return segment
+
+'''
+def patch_by_call(start_address_of_call, target_function_address):
+  # caculate the offset
+  jmp_offset = target_function_address - (start_address_of_call + 5)
+  # call + p32(jmp_offset)
+'''
+
+def patch_strcpy(lief_binary, nbytes, save_path, output=True):
+  print("[\033[1;34m*\033[0m] get the length of buffer is 0x%x(%d)" % (nbytes, nbytes))
+  patch_strcpy_code = f"""
+  save_register:
+    push rax;
+    push rcx;
+    push rdx;
+    xor rcx, rcx;
+  loop:
+    mov rdx, {nbytes-1};
+    mov al, [rcx + rsi];
+    
+    cmp rcx, rdx;
+    jge ret_code;
+
+    test al, al;
+    je ret_code;
+    
+    mov [rdi + rcx], al;
+    
+    inc rcx;
+    jmp loop;
+    
+  ret_code:
+    mov [rsi + rcx], al;
+    pop rdx;
+    pop rcx;
+    pop rax;
+    ret;
+  """
+  patch_code = asm(patch_strcpy_code)
+  if output:
+    print("the assmebly code :\n %s" % patch_strcpy_code)
+    print("the machine code :\n %s" % patch_code)
+  
+  new_segment = add_segment(lief_binary, types = lief._lief.ELF.Segment.TYPE.LOAD, flags = 5, content=patch_code)
+  lief_binary.patch_pltgot("strcpy", new_segment.virtual_address)
+  lief_binary.write(save_path)
+  os.system("chmod +x " + save_path)
+
+
+# to do
+def patch_dprintf(lief_binary, save_path, output=True):
+  patch_dprintf_code = f"""
+  save_register:
+    push rdx;
+    push rcx;
+  
+  init_register:
+    push rsi;
+    pop rcx;
+    xor rdx, rdx;
+    
+  get_the_buffer_len:
+    mov al, [rsi+rdx];
+    inc rdx;
+    test al, al;
+    jnz get_the_buffer_len;
+    
+  xor rax, rax;
+  mov al, 1;
+  syscall;
+  
+  ret_code:
+    pop rcx;
+    pop rdx;
+    ret;
+  """
+  patch_code = asm(patch_dprintf_code)
+  if output:
+    print("the assmebly code :\n %s" % patch_dprintf_code)
+    print("the machine code :\n %s" % patch_code)
+  new_segment = add_segment(lief_binary, types = lief._lief.ELF.Segment.TYPE.LOAD, flags = 5, content=patch_code)
+  lief_binary.patch_pltgot("dprintf", new_segment.virtual_address)
+  lief_binary.write(save_path)
+  os.system("chmod +x " + save_path)
+
+
+def patch_recv(lief_binary, nbytes, save_path, output=True):
+  patch_recv_code = f"""
+  mov rdx, {nbytes}
+  mov r10, rcx;
+  xor r8, r8;
+  xor r9, r9;
+  push 45;
+  pop rax;
+  syscall;
+  ret;
+  """
+  patch_code = asm(patch_recv_code)
+  if output:
+    print("the assmebly code :\n %s" % patch_recv_code)
+    print("the machine code :\n %s" % patch_code)
+  new_segment = add_segment(lief_binary, types = lief._lief.ELF.Segment.TYPE.LOAD, flags = 5, content=patch_code)
+  lief_binary.patch_pltgot("recv", new_segment.virtual_address)
+  print(save_path)
+  lief_binary.write(save_path)
+  os.system("chmod +x " + save_path)
+
+'''
+designed for console command
+'''
+# if __name__ == '__main__':
+# # #   argv = sys.argv
+# # #   argc = len(sys.argv)
+#   lief_binary, pwn_binary = load_binary_file_information("../input/edit")
+#   print(lief_binary)
+#   print(pwn_binary)
+#   patch_strcpy(lief_binary, 256, "../output/edit_patch/edit")