Write-Up: "xmastree" Day 8 from OverTheWire Advent Bonanza 2018

import r2pipe
import json
from binascii import unhexlify as unhex
import subprocess

r2 = r2pipe.open('./xmastree')

hash_funcs = {
    '0xb53220': 'hash_one',
    '0xb532e0': 'hash_two',
    '0xb533a0': 'hash_three'
}


class State:
    empty = '░'

    def __init__(self, addr, parent=None, constraint=bytearray(b'\x00'*10)):
        self.constraint = constraint
        self.addr = addr
        self.type = classify(addr)
        self.trueconstraint = None
        self.offsets = None
        self.parent = parent

        if self.type == 'hash':
            self.block = parsecheck(addr)
            self.offsets = self.block["input_offsets"]
        elif self.type == 'cmp':
            self.block = parsecmp(addr)
            self.offsets = parent.offsets
            self.constraint = parent.constraint[:]
        elif self.type == 'print':
            self.block = parseprint(addr)
        elif self.type == 'jmp':
            self.block = parsejump(addr)
        else:
            raise Exception("[state] Unknown block type %s" % self.type)

    def true(self):
        if self.type == "print":
            return None
        return self.block["true_addr"]

    def false(self):
        if self.type == "print":
            return None
        return self.block["false_addr"]

    def trueconstrain(self):
        input = crack(self.block["hash"])
        self.trueconstraint = self.constraint[:]
        for k,v in enumerate(self.offsets):
            if self.constraint[v] is not 0x00 and self.constraint[v] is not input[k]:
                return False
            else:
                self.trueconstraint[v] = input[k]
        return True

    def printconstraint(self):
        p = [self.empty if c == 0x00 else chr(c) for c in self.constraint]
        print("".join(p))

    def getconstraint(self):
        p = [self.empty if c == 0x00 else chr(c) for c in self.constraint]
        return "".join(p)

    def printtrueconstraint(self):
        p = [self.empty if c == 0x00 else chr(c) for c in self.trueconstraint]
        print("".join(p))

    def gettrueconstraint(self):
        p = [self.empty if c == 0x00 else chr(c) for c in self.trueconstraint]
        return "".join(p)

def crack(hash):
    try:
        output = subprocess.check_output(['./chash', hash]).split(b"\n")
        assert len(output) == 4
        assert output[2] == b'Found 1'
        return output[1]
    except subprocess.CalledProcessError:
        raise Exception("Unable to crack %s" % hash)


def classify(addr):
    check = json.loads(r2.cmd("pdj 4 @ %s" % addr))
    if check[0]['disasm'] == 'mov eax, 3':
        return "hash"
    elif check[3]['disasm'] == 'call sym.imp.strcmp':
        return "cmp"
    elif check[2]['disasm'] == 'call sym.imp.printf':
        return "print"
    elif check[0]['type'] == 'jmp':
        return "jmp"
    else:
        raise Exception("unknown block @ %s" % addr)


def parsecheck(check_addr):
    #   print("Parsing check @ %s" % check_addr)

    check = json.loads(r2.cmd("pdj 21 @ %s" % check_addr))
    input_offsets = []

    # hash input length
    assert check[0]['disasm'] == 'mov eax, 3'
    checko = {
        "addr": check_addr,
        "type": "hash",
        "input_length": 3
    }

    # input offsets
    for pos in [5,8,11]:
        assert check[pos]['disasm'].startswith('mov r8b, byte [rcx')
        if check[pos]['size'] == 3: # 448a01 < > 448a4103
            offset = 0
        else:
            offset = ord(unhex(check[pos]['bytes'][-2:]))
        assert 0 <= offset < 10
        input_offsets.append(offset)
    checko["input_offsets"] = input_offsets

    # hash func call
    assert check[13]['type'] == 'call'
    hash_addr = check[13]['disasm'].split(" ")[1]
    checko['hash_func'] = hash_funcs[hash_addr]

    # load target hash
    assert check[15]['disasm'].startswith('mov eax, str.')
    str_addr = check[15]['opcode'].split(" ")[2]
    hash_target = r2.cmd("ps @ %s" % str_addr)[:-1]
    checko['hash'] = hash_target

    # fail check address
    assert check[19]['disasm'].startswith('jne')
    checko['false_addr'] = check[19]['disasm'].split(" ")[1]

    # following bb address
    next_bb = hex(check[20]['offset'])
    checko['true_addr'] = next_bb

    return checko


def parsecmp(addr):
    #    print("Parsing cmp @ %s" % addr)
    block = json.loads(r2.cmd("pdj 7 @ %s" % addr))

    assert block[3]['disasm'] == 'call sym.imp.strcmp'
    assert block[1]['disasm'].startswith('mov eax, str.')
    str_addr = block[1]['opcode'].split(" ")[2]
    hash_target = r2.cmd("ps @ %s" % str_addr)[:-1]

    blocko = {
        "addr": addr,
        "type": "cmp",
        "hash": hash_target,
    }

    # fail check address
    assert block[5]['disasm'].startswith('jne')
    blocko['false_addr'] = block[5]['disasm'].split(" ")[1]

    # following bb address
    next_bb = hex(block[6]['offset'])
    blocko['true_addr'] = next_bb
    return blocko


def parseprint(addr):
    #    print("Parsing print @ %s" % addr)
    block = json.loads(r2.cmd("pdj 3 @ %s" % addr))
    assert block[2]['disasm'] == 'call sym.imp.printf'
    assert block[0]['disasm'].startswith('movabs rdi, str.')
    outs = r2.cmd("ps @ %s" % block[0]['opcode'].split(" ")[2])
    blocko = {
        "addr": addr,
        "type": "print",
        "string": outs
    }
    return blocko


def parsejump(addr):
    #    print("Parsing jump @ %s" % addr)
    block = json.loads(r2.cmd("pdj 1 @ %s" % addr))
    assert block['type'] == 'jmp'
    blocko = {
        "addr": addr,
        "type": "jump",
        "dest": hex(block['jump'])
    }
    return blocko


def walk(s):
    stack = []
    stack.append(s)
    while len(stack) > 0:
        w = stack.pop()
        print("%5s @ %s - %s  %s %s %s" % (w.type, w.addr, w.getconstraint(), w.offsets, w.true(), w.false()))
        if w.type == "print":
            print(w.block["string"])
        if w.type == "hash" or w.type == "cmp":
            stack.append(State(w.false(), constraint=w.constraint[:],parent=w))
            if w.trueconstrain() is True:
                stack.append(State(w.true(), constraint=w.trueconstraint[:]))


start = '0x4006cc'
walk(State(start))