#!/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_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_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"""
    mov rdx, {nbytes - 1};
    mov byte ptr [rdx + rsi], 0;
  """
  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)
  new_segment_address = new_segment.virtual_address
  #lief_binary.patch_pltgot("strcpy", new_segment.virtual_address)
  os.system("rm " + save_path)
  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"""
  push rsi;
  pop rdx;
  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)
  new_segment_address = new_segment.virtual_address
  #lief_binary.patch_pltgot("dprintf", new_segment.virtual_address)
  
  os.system("rm " + save_path)
  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};
  """
  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)
  new_segment_address = new_segment.virtual_address
  
  os.system("rm " + save_path)
  lief_binary.write(save_path)
  os.system("chmod +x " + save_path)
  
  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)