[CSAW CTF Finals 2017] rabbithole

rabbithole.py

from pwn import *

class Node(object):
    def __init__(self, addr):
        self.addr = addr

e = ELF("./rabbithole")

name_addr = {}
for i in e.symbols:
    if "node" in i:
        name_addr[e.symbols[i]] = i
name_addr[0] = "leaf"
roots = []
idx = 0

for idx in xrange(59):
    addr = u64(e.read(e.symbols['roots']+(idx*8),8))
    root = Node(addr)
    roots.append(root)

tree = {}

def form_dict(root):
    if not root:
        return
    chk = ord(e.read(root,1))
    l = u64(e.read(root+(8*2),8))
    r = u64(e.read(root+(8*3),8))
    a,b = ord(e.read(root+8,1)), ord(e.read(root+9,1))
    tree[root] = (l,r,a,b,chk)
    if not all((l,r)):
        return
    form_dict(r)
    form_dict(l)

for idx in xrange(59):
    form_dict(roots[idx].addr)

success("Trees formed")

def print_tree(head):
    lvl = [head]
    while lvl:
        nxt = []
        for node in lvl:
            print name_addr[node],
            if node in tree:
                nxt.append(tree[node][0])
                nxt.append(tree[node][1])
        print
        lvl = nxt

# print_tree(roots[0].addr)

from z3 import *

for i in xrange(59):
    globals()['a%i' % i]=BitVec('a%i' %i,32)

solver = Solver()

def parse_forz3(idx, root, s):
    for i in s:
        if i=="L":
            solver.add(globals()['a%i' % idx] >= tree[root][2])
            solver.add(globals()['a%i' % idx] < tree[root][3]) 
            root = tree[root][0]
        if i=="R":
            solver.add(Or(globals()['a%i' % idx] < tree[root][2],(globals()['a%i' % idx] >= tree[root][3])))
            root = tree[root][1]

def preorder_solve(idx, root, gs):
    if  not root:
        return
    if tree[root][2] == 1 and tree[root][4] == 1: 
        parse_forz3(idx, roots[idx].addr, gs)
    if tree[root][0]:
        preorder_solve(idx, tree[root][0], gs+"L")
    if tree[root][1]:
        preorder_solve(idx, tree[root][1], gs+"R")

success("Traversing and adding constraints")

for i in xrange(59):
    preorder_solve(i, roots[i].addr,"")

print solver.check()
if solver.check() != sat:
    sys.exit(-1)

modl=solver.model()
flag=""
for i in xrange(59):
    obj=globals()['a%i' % i]
    c=modl[obj].as_long()
    flag += chr(c)

print flag