#!/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(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(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(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)
  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)