Upload files to "src"

src/binary_patch.py

@@ -61,42 +61,24 @@ def add_segment(lief_binary, content, types, flags, base=0x405000):
   print(segment.FLAGS.value)
   return segment
 
-'''
-def patch_by_call(start_address_of_call, target_function_address):
+
+def patch_by_pltsec_jmp(elf_file, symbol, start_address_of_pltsec_jmp, target_function_address, target_function_len, save_path):
   # caculate the offset
-  jmp_offset = target_function_address - (start_address_of_call + 5)
-  # call + p32(jmp_offset)
-'''
+  jmp_offset = target_function_address - (start_address_of_pltsec_jmp + 5)
+  shellcode = b'\xe9' + p32(jmp_offset & 0xffffffff)
+  elf_file.write(start_address_of_pltsec_jmp, shellcode)
+  
+  jmp_offset = elf_file.got[symbol] - (target_function_address + target_function_len + 7)
+  shellcode = b'\xf2\xff\x25' + p32(jmp_offset & 0xffffffff)
+  elf_file.write(target_function_address + target_function_len, shellcode)
+  elf_file.save(save_path)
+
 
 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;
+    mov byte ptr [rdx + rsi], 0;
   """
   patch_code = asm(patch_strcpy_code)
   if output:
@@ -104,76 +86,65 @@ def patch_strcpy(lief_binary, nbytes, save_path, output=True):
     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)
+  new_segment_address = new_segment.virtual_address
+  #lief_binary.patch_pltgot("strcpy", new_segment.virtual_address)
   lief_binary.write(save_path)
   os.system("chmod +x " + save_path)
   
+  elf_patch = ELF(save_path)
+  patch_by_pltsec_jmp(elf_patch, 'strcpy', elf_patch.plt['strcpy'], new_segment_address, len(patch_code), 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;
+  mov rax, [rsp];
+  push 0x7325;
+  push rsp
+  pop rsi;
+  push rax;
   """
   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)
+  new_segment_address = new_segment.virtual_address
+  #lief_binary.patch_pltgot("dprintf", new_segment.virtual_address)
   lief_binary.write(save_path)
   os.system("chmod +x " + save_path)
   
+  elf_patch = ELF(save_path)
+  patch_by_pltsec_jmp(elf_patch, 'dprintf', elf_patch.plt['dprintf'], new_segment_address, len(patch_code), 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;
+  mov rdx, {nbytes};
   """
   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)
+  new_segment_address = new_segment.virtual_address
+  
   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")
+  elf_patch = ELF(save_path)
+  patch_by_pltsec_jmp(elf_patch, 'recv', elf_patch.plt['recv'], new_segment_address, len(patch_code), save_path)
+
+if __name__ == '__main__':
+  argv = sys.argv
+  argc = len(sys.argv)
+  path = sys.argv[1]
+  save_path = path + "_patch" 
+  lief_binary, pwn_binary = load_binary_file_information(path)
+  if sys.argv[2] == 'dprintf':
+    patch_dprintf(lief_binary, save_path)
+  elif sys.argv[2] == 'strcpy':
+    patch_strcpy(lief_binary, int(sys.argv[3]), save_path)
+  elif sys.argv[2] == 'recv':
+    patch_recv(lief_binary, int(sys.argv[3]), save_path)