0vercl0k chall """solution"""

# ugly poc written for 0vercl0k[] chall
# hardcoded addresses, slow, but you'll have the idea :)
# IL + data tainting + symbolic execution + some trickz
# it could be a lot faster/cleaner but I don't have time to work on it, it was just a fast PoC
# thks to 0vercl0k for taking the time to write a clean writeup ;)
# -> http://doar-e.github.io/blog/2013/09/16/breaking-kryptonites-obfuscation-with-symbolic-execution/

import sys
sys.path.append("C:\\Program Files (x86)\\z3-4.3.0-x86\\bin")
from idc import *
from collections import defaultdict
from z3 import *
from pprint import pprint
import re
import PySide
from PySide.QtGui import QMessageBox

addr_check = 0x140F19996

class BVRegs(dict) :
    def __init__(self) :
        dict.__init__(self, {r:BitVec(r, 64) for r in Reg.regs})

    def __setitem__(self, item, value) :
        dict.__setitem__(self, item, simplify(value))

    def __str__(self) :
        return "\n".join("%s : %s"%(r,self[r]) for r in Reg.regs)

class BVContext(object) :
    def __init__(self) :
        self.sym_count = 0
        self.regs = BVRegs()
        self.memory = {}
        self.used_addr = set()

    def deref_read(self, addr) :
        # fucking ugly but works ... just consider that != str => != addr
        return self.memory.setdefault(str(addr), BitVec("[%s]"%addr, 64))

    def deref_write(self, addr, value) :
        self.memory[str(addr)] = simplify(value)

    def int2abstract(self, i) :
        return BitVecVal(i, 64)

    def __getitem__(self, item) :
        if isinstance(item, Expr) :
            return item.get_value(self)
        elif isinstance(item, int) or isinstance(item, long) :
            return BitVecVal(item, 64)
        raise Exception("invalid item :\n\t%r"%item)

    def __setitem__(self, item, value) :
        if isinstance(item, Expr) :
            return item.set_value(self, value)
        raise Exception("invalid item :\n\t%r"%item)

    def __str__(self) :
        ret_val = "%s\n"%self.regs
        mem = ["%s : %s"%(r,v) for r,v in self.memory.items()]
        mem.sort()
        ret_val += "\n".join(mem)
        return ret_val

class Singleton(type):
    def __init__(self, *args, **kwargs):
        super(Singleton, self).__init__(*args, **kwargs)
        self.__instances = {}

    def __call__(self, *args, **kwargs):
        arg = args[0] if len(args) else None
        return self.__instances.setdefault(arg, super(Singleton, self).__call__(*args, **kwargs))

class Taint(frozenset) :
    __metaclass__ = Singleton

    def __init__(self, taints=None) :
        if taints :
            if len(taints) > 1 :
                raise DecodeException("toto")
            frozenset.__init__(self, taints)
        else :
            frozenset.__init__(self)

    def __add__(self, b) :
        return Taint(self.union(b))

    def __sub__(self, b) :
        return Taint(self.union(b))

    def __and__(self, b) :
        return Taint(self.union(b))

    def __or__(self, b) :
        return Taint(self.union(b))

    def __xor__(self, b) :
        return Taint(self.union(b))

    def __neg__(self) :
        return self

    def __invert__(self) :
        return self

    def __lshift__(self, b) :
        return Taint(self.union(b))

    def __rshift__(self, b) :
        return Taint(self.union(b))

    def __str__(self)  :
        return ", ".join(self)

class TaintedRegs(dict) :
    def __init__(self) :
        dict.__init__(self, {r:Taint(frozenset((r,))) for r in Reg.regs})

    def __str__(self) :
        return "\n".join("%s : %s"%(r,self[r]) for r in Reg.regs)

class TaintedContext(object) :
    def __init__(self) :
        self.sym_count = 0
        self.regs = TaintedRegs()
        self.memory = {}
        self.used_addr = set()

    def deref_read(self, addr) :
        # fucking ugly but works ... just consider that != str => != addr
        v = self.memory.setdefault(str(addr), Taint(frozenset(("[%s]"%addr,))))
        return v

    def deref_write(self, addr, value) :
        self.memory[str(addr)] = value

    def int2abstract(self, i) :
        return Taint()

    def __getitem__(self, item) :
        if isinstance(item, Expr) :
            return item.get_value(self)
        elif isinstance(item, int) or isinstance(item, long) :
            return Taint()
        raise Exception("invalid item :\n\t%r"%item)

    def __setitem__(self, item, value) :
        if isinstance(item, Expr) :
            return item.set_value(self, value)
        raise Exception("invalid item :\n\t%r"%item)


    def __str__(self) :
        ret_val = "%s\n"%self.regs
        mem = ["%s : %s"%(r,v) for r,v in self.memory.items()]
        mem.sort()
        ret_val += "\n".join(mem)
        return ret_val

class Mix(object) :
    def __init__(self, bv, taint) :
        self.bv = bv
        self.taint = taint

    @staticmethod
    def wrapp(name):
        bv_attr = getattr(BitVecRef, name)
        taint_attr = getattr(Taint, name)
        if hasattr(bv_attr, '__call__') and hasattr(taint_attr, '__call__') :
            def newfunc(*args):
                bv_args = tuple(arg.bv for arg in args)
                taint_args = tuple(arg.taint for arg in args)
                result = Mix(bv_attr(*bv_args), taint_attr(*taint_args))
                return result
            return newfunc
        raise Exception("unkown attribute")

    def __str__(self) :
        return "\tbv: %s\n\ttaint: %s"%(self.bv, self.taint)

for n in ("add", "sub", "and", "or", "xor", "neg", "invert", "lshift", "rshift") :
    setattr(Mix, "__%s__"%n, Mix.wrapp("__%s__"%n))

class MixedRegs(object) :

    def __init__(self, mixedContext) :
        self.bv_regs = mixedContext.bv_ctx.regs
        self.tainted_regs = mixedContext.tainted_ctx.regs

    def __getitem__(self, item) :
        return Mix(self.bv_regs[item],
                self.tainted_regs[item])

    def __setitem__(self, item, value) :
        self.bv_regs[item] = value.bv
        self.tainted_regs[item] = value.taint


    def __str__(self) :
        return "\n".join("%s : %s"%(r, self[r]) for r in Reg.regs)

class MixedContext(object) :

    def __init__(self) :
        self.bv_ctx = BVContext()
        self.tainted_ctx = TaintedContext()
        self.regs = MixedRegs(self)

    def deref_read(self, addr) :
        # use bv address
        return Mix(self.bv_ctx.deref_read(addr.bv), self.tainted_ctx.deref_read(addr.bv))

    def deref_write(self, addr, value) :
        self.bv_ctx.deref_write(addr.bv, value.bv)
        self.tainted_ctx.deref_write(addr.bv, value.taint)

    def int2abstract(self, i) :
        return Mix(self.bv_ctx.int2abstract(i),
                self.tainted_ctx.int2abstract(i))

    def __getitem__(self, item) :
        if isinstance(item, Expr) :
            return item.get_value(self)
        elif isinstance(item, int) or isinstance(item, long) :
            return self.int2abstract(item)
        raise Exception("invalid item :\n\t%r"%item)

    def __setitem__(self, item, value) :
        if isinstance(item, Expr) :
            return item.set_value(self, value)
        raise Exception("invalid item :\n\t%r"%item)

    def __str__(self) :
        ret_val = "%s\n"%self.regs
        mem = ["%s :\n\tbv: %s\n\ttaint: %s)"%(r,v,self.tainted_ctx.memory[r]) for r,v in self.bv_ctx.memory.items()]
        mem.sort()
        ret_val += "\n".join(mem)
        return ret_val


class Expr(object) :
    size = None
    funs = {
        "+" : (lambda x, y : x + y),
        "-" : (lambda x, y : x - y),
        "&" : (lambda x, y : x & y),
        "|" : (lambda x, y : x | y),
        "^" : (lambda x, y : x ^ y),
        "<<" : (lambda x, y : x << y),
        ">>" : (lambda x, y : x >> y),
        "u-" : (lambda x : -x),
        "~" : (lambda x : ~y),
    }
    def __init__(self, operator, *operands) :
        self.opr = str(operator)
        self.ops = operands
        self.fun = Expr.funs[self.opr]

        if len(self.ops) > 1 :
            self._str = self.opr.join("%s"%op if isinstance(op, Operand) else "(%s)"%op for op in self.ops)
        else :
            op = self.ops[0]
            self._str = self.opr + ("%s"%op if isinstance(op, Operand) else "(%s)"%op)

    def __str__(self) :
        return self._str

    def __add__(self, b) :
        return Expr("+", self, b)

    def __sub__(self, b) :
        return Expr("-", self, b)

    def __and__(self, b) :
        return Expr("&", self, b)

    def __or__(self, b) :
        return Expr("|", self, b)

    def __xor__(self, b) :
        return Expr("^", self, b)

    def __neg__(self) :
        return Expr("u-", self)

    def __invert__(self) :
        return Expr("~", self)

    def __lshift__(self, b) :
        return Expr("<<", self, b)

    def __rshift__(self, b) :
        return Expr(">>", self, b)

    def get_value(self, context) :
        return self.fun(*tuple(op.get_value(context) for op in self.ops))

class Operand(Expr) :
    pass

class Reg(Operand) :
    regs = []
    _str_to_rms = {}
    for s in ("a", "b", "c", "d") :
        r = "r" + s + "x"
        regs.append(r)
        _str_to_rms[s+"l"] = (r, 0xFF, 0, 1)
        _str_to_rms[s+"h"] = (r, 0xFF00, 8, 1)
        _str_to_rms[s+"x"] = (r, 0xFFFF, 0, 2)
        _str_to_rms["e"+s+"x"] = (r, 0xFFFFFFFF, 0, 4)
        _str_to_rms["r"+s+"x"] = (r, 0xFFFFFFFFFFFFFFFF, 0, 8)

    for s in ("di", "si", "sp", "bp") :
        r = "r" + s
        regs.append(r)
        _str_to_rms[s+"l"] = (r, 0xFF, 0, 1)
        _str_to_rms[s] = (r, 0xFFFF, 0, 2)
        _str_to_rms["e"+s] = (r, 0xFFFFFFFF, 0, 4)
        _str_to_rms["r"+s] = (r, 0xFFFFFFFFFFFFFFFF, 0, 8)

    for s in xrange(8, 16) :
        s = "r%d"%s
        regs.append(s)
        _str_to_rms[s+"b"] = (s, 0xFF, 0, 1)
        _str_to_rms[s+"w"] = (s, 0xFFFF, 0, 2)
        _str_to_rms[s+"d"] = (s, 0xFFFFFFFF, 0, 4)
        _str_to_rms[s] = (s, 0xFFFFFFFFFFFFFFFF, 0, 8)
    regs = tuple(regs)

    def __init__(self, s) :
        self._str = str(s)
        self.reg, self.mask, self.shift, self.size = Reg._str_to_rms[self._str]

    def __str__(self) :
        return self._str

    def get_value(self, context) :
        v = context.regs[self.reg]
        if self.mask != 0xFFFFFFFFFFFFFFFF :
            v &= context[self.mask]
        if self.shift != 0 :
            v >>= context[self.shift]
        return v

    def set_value(self, context, value) :
        if self.mask >= 0xFFFFFFFF :
            context.regs[self.reg] = value
            return
        if self.shift != 0 :
            value <<= self.shift
        context.regs[self.reg] = (context.regs[self.reg] & ~context[self.mask]) | value

class Mem(Operand) :
    magic_regex = re.compile(r'\[((?P<base>([abcdspirexhl]{1,3}[0-9]{0,2}(?!\*)))?(((?P<index_sign>[+-]))?(?P<index>[abcdspirexhl]{1,3}[0-9]{0,2})(\*(?P<index_mult>[2-8]))?)?((?P<disp>[+-][0-9A-F]+)h?)?|((?P<addr>-?[0-9A-F])+h?))\]')

    def __init__(self, expr) :
        self.expr = expr
        self._str = "[%s]"%expr

    def __str__(self) :
        return self._str

    @staticmethod
    def str2mem(s) :
        s = str(s)
        try :
            base, index_sign, index, _index_mult, _disp, addr = Mem.magic_regex.match(s).group("base", "index_sign", "index", "index_mult", "disp", "addr")
        except :
            raise Exception("Unkown memory address : %s"%s)
        exp = None
        if _disp :
            _disp = int(_disp, 16)
            if _disp :
                exp = Cst(_disp)
        if addr :
            exp = Cst(int(addr, 16))
        index_mult = None
        if index_mult :
            _index_mult = int(_index_mult)
            if _index_mult != 1 :
                index_mult = _index_mult

        if base :
            base = Reg(base)
            exp = base + exp if exp is not None else base
        if index :
            index = Reg(index)
            if index_mult :
                index = index * Cst(index_mult)
            if index_sign == "-" :
                index = -index
            exp = index + exp if exp is not None else index
        return Mem(exp)

    def get_value(self, context) :
        addr = self.expr.get_value(context)
        return context.deref_read(addr)

    def set_value(self, context, value) :
        addr = context[self.expr]
        return context.deref_write(addr, value)

class Cst(Operand) :

    def __init__(self, v) :
        self._str = "0x%08X"%v
        self.v = int(v)
        self._size = None

    @property
    def size(self):
        return self._size

    @size.setter
    def size(self, value):
        assert (1 << (value*8)) > abs(self.v)
        self._size = value

    def get_value(self, context) :
        return context[self.v]

    def __str__(self) :
        return self._str

def normalize_sizes(*operands) :
    '''set operand size so every sizes are equals, if there is any non matching size : raise Exc'''
    for op in operands :
        if op.size :
            size = op.size
            break
    else :
        raise Exception("no known size")

    if any(op.size != size for op in operands if op.size is not None) :
        raise Exception("sizes are not matching")

    for op in operands :
        op.size = size

def get_op(addr, i) :
    inst = GetOpType(addr, i)
    if inst < 1 :
        return None
    elif inst == 1 : # reg
        return Reg(GetOpnd(addr, i))
    elif 2 <= inst <= 4 :
        return Mem.str2mem(GetOpnd(addr, i))
    elif 5 <= inst <= 7 :
        return Cst(GetOperandValue(addr, i))
    else :
        raise Exception("unkown operand (%X %d)"%(addr, i))

ins2mc_funs = {}
def ins2mc_add(dst, src) :
    normalize_sizes(dst, src)
    return [(dst, dst + src)]
ins2mc_funs["add"] = ins2mc_add

def ins2mc_and(dst, src) :
    normalize_sizes(dst, src)
    return [(dst, dst & src)]
ins2mc_funs["and"] = ins2mc_and

def ins2mc_shr(dst, src) :
    return [(dst, dst >> src)]
ins2mc_funs["shr"] = ins2mc_shr

def ins2mc_shl(dst, src) :
    return [(dst, dst << src)]
ins2mc_funs["shl"] = ins2mc_shl

def ins2mc_xor(dst, src) :
    normalize_sizes(dst, src)
    return [(dst, dst ^ src)]
ins2mc_funs["xor"] = ins2mc_xor

def ins2mc_or(dst, src) :
    normalize_sizes(dst, src)
    return [(dst, dst | src)]
ins2mc_funs["or"] = ins2mc_or

def ins2mc_mov(dst, src) :
    normalize_sizes(dst, src)
    return [(dst, src)]
ins2mc_funs["mov"] = ins2mc_mov


class DecodeException(Exception) :
    pass

def ins2mc(addr) :
    mnem = GetMnem(addr)
    if mnem == "" :
        raise DecodeException("Unkown instruction @ %X"%(mnem, addr))
    ops = [op for op in [get_op(addr, i) for i in xrange(3)] if op]
    if mnem not in ins2mc_funs :
        raise DecodeException("Unkown instruction %s @ %X"%(mnem, addr))
    return ins2mc_funs[mnem](*ops)


def simplify_linear(v, symbol) :
    s = Solver()

    a = simplify(substitute(v, (BitVec(symbol, 64), BitVecVal(0, 64))))
    s.add(a + BitVec(symbol, 64)  & 0xFFFFFFFF != v & 0xFFFFFFFF)
    print a
    print s.check()
    if s.check() == unsat :
        return a + BitVec(symbol, 64)
    print s.model()
    return v


eip = 0x00140F199E9
prev = eip
context = MixedContext()

msgBox = QMessageBox()
msgBox.setText("Break")
msgBox.setInformativeText("Break")
msgBox.setStandardButtons(QMessageBox.Cancel | QMessageBox.Ok | QMessageBox.Close)
msgBox.setDefaultButton(QMessageBox.Ok)
stepinto = False
try :
    while True :
        if eip == 0x0140F1BD06   :
            stepinto = True
        if stepinto :
            # print context.bv_ctx
            # print context.tainted_ctx
            # print context
            syms = set(context.tainted_ctx.regs["rdi"])
            sym = syms.pop()
            print context.tainted_ctx.regs["rdi"]
            print simplify_linear(context.bv_ctx.regs["rdi"], sym)
            ret = msgBox.exec_()
            if ret == QMessageBox.Cancel:
                sys.exit()
            if ret == QMessageBox.Close :
                stepinto = False
        for dst, expr in ins2mc(eip) :
            v = context[expr]
            # if len(v.taint) > 1 :
            #   print hex(eip)
            #   print "bv:", str(simplify(v.bv))
            #   print "taint:", str(v.taint)
            #   stepinto = True
            context[dst] = v
        eip += ItemSize(eip)
        if eip - prev >= 0x1000 :
            print hex(eip)
            prev = eip
except DecodeException as e :
    pass

print "prout"