import re
from idc import GetMnem, GetOpnd, GetOpType, SetColor, ItemSize, PrevHead, CIC_ITEM, GetDisasm
from random import randint

code = '''
def _sub(a, b):
    return (a - b) & 0xFFFFFFFF

def _add(a, b):
    return (a + b) & 0xFFFFFFFF

def _rol(a, b):
    return ((a << (b & 0x1F)) | (a >> (32 - (b & 0x1F)))) & 0xFFFFFFFF

def _ror(a, b):
    return _rol(a, 32 - (b & 0x1F))

def _shr(a, b):
    return (a >> b) & 0xFFFFFFFF

def _shl(a, b):
    return (a << b) & 0xFFFFFFFF

def _xor(a, b):
    return (a ^ b) & 0xFFFFFFFF

def _and(a, b):
    return (a & b) & 0xFFFFFFFF

def _or(a, b):
    return (a | b) & 0xFFFFFFFF

def _imul(a, b):
    return (a * b) & 0xFFFFFFFF

def _neg(a):
    return (-a) & 0xFFFFFFFF

def _not(a):
    return (~a) & 0xFFFFFFFF

def _dec(a):
    return (a - 1) & 0xFFFFFFFF

def _inc(a):
    return (a + 1) & 0xFFFFFFFF

'''

code_end = """    a2 = eax
    a3 = edi
    a4 = ebx
    return a1, a2, a3, a4

if __name__ == '__main__':
    # from struct import unpack
    # u = lambda x: unpack('I', x)[0]

    # h = '43df6e922f0ad4127ef86627a9156601'.decode('hex')
    # i = 0
    # a = h[i:i+4]
    # i = 4
    # b = h[i:i+4]
    # i = 8
    # c = h[i:i+4]
    # i = 12
    # d = h[i:i+4]
    # print '%08X %08X %08X %08X' % (u(a), u(b), u(c), u(d))
    # a, b, c, d = encrypt(u(a), u(b), u(c), u(d))
    in_ = [0x926EDF43, 0x12D40A2F, 0x2766F87E, 0x016615A9]
    out = [0xAF8FF516, 0x39DA6102, 0x8AD9A164, 0x8B218C70]
    print '%08X %08X %08X %08X' % (in_[0], in_[1], in_[2], in_[3])
    a, b, c, d = encrypt(in_[0], in_[1], in_[2], in_[3])
    print '%08X %08X %08X %08X' % (a, b, c, d)
    print '%08X %08X %08X %08X' % (out[0], out[1], out[2], out[3])
"""

class DecodeRotate:
    def __init__(self, begin, end, pos, used):
        self.begin = begin
        self.end = end
        self.pos = pos
        self.used = used

        self.op1 = GetOpnd(pos, 0)
        self.op2 = GetOpnd(pos, 1)

        self.addesses = []
        self.addesses.append(pos)

        self.count_init_addr = 0
        self.value_init_addr = 0
        self.neg_addr = 0

    def process_shifts(self):
        ops = [self.op1, self.op2]
        ops_o = ops
        pos = self.pos
        shifts_addr = []
        while len(ops) > 0:
            while pos > self.begin:
                pos = PrevHead(pos)
                if GetMnem(pos).startswith('sh') and GetOpnd(pos, 0) in ops and pos not in self.addesses and pos not in self.used:
                    self.addesses.append(pos)
                    ops.remove(GetOpnd(pos, 0))
                    shifts_addr.append(pos)
                    break
                elif GetMnem(pos) == 'mov' and GetOpnd(pos, 0) in ops and pos not in self.addesses and pos not in self.used:
                    self.addesses.append(pos)
                    ops.remove(GetOpnd(pos, 0))
                    ops.append(GetOpnd(pos, 1))
                    ops_o = ops

        and_addr = []
        while len(shifts_addr) > 0:
            pos = max(shifts_addr)
            and_op = 'ecx'
            while pos > self.begin:
                pos = PrevHead(pos)
                if GetMnem(pos) == 'mov' and GetOpnd(pos, 0) == and_op and pos not in self.addesses:
                    and_op = GetOpnd(pos, 1)
                    self.addesses.append(pos)
                elif GetMnem(pos) == 'and' and GetOpnd(pos, 0) == and_op and GetOpnd(pos, 1) == '1Fh' and pos not in self.addesses:
                    self.addesses.append(pos)
                    shifts_addr.remove(max(shifts_addr))
                    and_addr.append(pos)
                    break

        addr = self.find_count_init(and_addr)
        if addr is not None:
            self.find_ops_init(ops_o)
            self.find_neg(addr)

    def find_neg(self, addr):
        op1 = GetOpnd(addr, 0)
        op2 = GetOpnd(addr, 1)
        while addr < self.pos:
            addr += ItemSize(addr)
            mnem = GetMnem(addr)
            if mnem == 'neg' and GetOpnd(addr, 0) in [op1, op2]:
                self.addesses.append(addr)
                break
            elif mnem == 'mov' and GetOpnd(addr, 1) in [op1, op2] and GetOpType(addr, 1) != 5 and addr in self.addesses:
                if GetOpnd(addr, 1) == op1:
                    op1 = GetOpnd(addr, 0)
                elif GetOpnd(addr, 1) == op2:
                    op2 = GetOpnd(addr, 0)
        else:
            print 'find_neg error', hex(self.pos)
            print op1, op2
            pass

    def find_ops_init(self, ops):
        count = 0
        op1 = GetOpnd(self.pos, 0)
        op2 = GetOpnd(self.pos, 1)
        pos1 = self.pos
        while pos1 > self.begin:
            count += 1
            if count > 60:
                self.addesses = []
                print '-', hex(self.pos)
                break
            pos1 = PrevHead(pos1)
            if GetMnem(pos1) == 'mov' and GetOpType(pos1, 1) != 5 and GetOpnd(pos1, 0) in [op1, op2]:
                self.addesses.append(pos1)
                if GetOpnd(pos1, 1) in [op1, op2]:
                    self.value_init_addr = pos1
                    break
                elif GetOpnd(pos1, 0) == op1:
                    op1 = GetOpnd(pos1, 1)
                elif GetOpnd(pos1, 0) == op2:
                    op2 = GetOpnd(pos1, 1)
                else:
                    print('Hmmm')

    def find_count_init(self, starts):
        pos = max(starts)
        op1 = GetOpnd(pos, 0)
        pos1 = min(starts)
        op2 = GetOpnd(pos1, 0)
        count = 0
        while pos > min(starts):
            if count > 60:
                self.addesses = []
                print 'find_count_init (stage #1, count end)', hex(self.pos)
                return None
            pos = PrevHead(pos)
            if GetMnem(pos) in ['mov', 'lea'] and GetOpType(pos, 1) != 5 and GetOpnd(pos, 0) == op1 and pos not in self.addesses:
                op1 = GetOpnd(pos, 1)
                self.addesses.append(pos)
                # if GetOpnd(pos, 0) == op1:
                #     op1 = GetOpnd(pos, 1)
                # elif GetOpnd(pos, 0) == op2:
                #     op2 = GetOpnd(pos, 1)
            # elif GetMnem(pos) == 'neg' and GetOpnd(pos, 0) == op1 and self.neg_addr == 0:
            #     self.neg_addr = pos
            #     self.addesses.append(addesses)
            count += 1
        # else:
        #     self.addesses = []
        #     print 'find_count_init (stage #1, count end)', hex(self.pos)
        #     return None

        count = 0

        while pos1 > self.begin:
            if count > 60:
                self.addesses = []
                print 'find_count_init (stage #2, count end)', hex(self.pos)
                return None
            pos1 = PrevHead(pos1)
            if GetMnem(pos1) in ['mov', 'lea'] and GetOpType(pos1, 1) != 5 and GetOpnd(pos1, 0) in [op1, op2] and pos1 not in self.addesses:
                self.addesses.append(pos1)
                if GetOpnd(pos1, 1) in [op1, op2]:
                    self.count_init_addr = pos1
                    return pos1
                elif GetOpnd(pos1, 0) == op1:
                    op1 = GetOpnd(pos1, 1)
                elif GetOpnd(pos1, 0) == op2:
                    op2 = GetOpnd(pos1, 1)
                else:
                    print('Hmmm')
            # elif GetMnem(pos) == 'neg' and GetOpnd(pos, 0) in [op1, op2] and self.neg_addr == 0:
            #     self.neg_addr = pos
            #     self.addesses.append(addesses)
            count += 1

    def do(self):
        self.process_shifts()

    def get_addresses(self):
        self.addesses.sort()
        return self.addesses

    def get_instruction(self):
        op = None
        self.addesses.sort()
        for addr in self.addesses:
            if GetMnem(addr) == 'neg':
                op = GetOpnd(addr, 0)
            if op is not None and GetMnem(addr) == 'mov' and GetOpnd(addr, 1) == op:
                op = GetOpnd(addr, 0)
            if op is not None and GetMnem(addr).startswith('sh') and op == 'ecx':
                op0 = GetOpnd(self.pos, 0)
                op1 = GetOpnd(self.value_init_addr, 1)
                op2 = GetOpnd(self.count_init_addr, 1)
                mnem = '_rol' if GetMnem(addr) == 'shr' else '_ror'
                return '%s = %s(%s, %s)' % (op0, mnem, op1, op2)


class Decoder:
    def __init__(self, begin, end):
        self.begin = begin
        self.end = end

        self.addesses = []
        self.blocks = []

    def do(self):
        pos = self.begin
        count = 0
        while pos < self.end:
            pos += ItemSize(pos)
            if GetMnem(pos) == 'or' and GetOpType(pos, 1) != 5:
                rot = DecodeRotate(self.begin, self.end, pos, self.addesses)
                rot.do()
                addrs = rot.get_addresses()
                self.addesses.extend(addrs)
                self.blocks.append(addrs)
                print rot.get_instruction()
                # if count > 1:
                #     break
                count += 1

    def colorize(self):
        pos = self.begin
        while pos < self.end:
            SetColor(pos, CIC_ITEM, 0xFFFFFF)
            pos += ItemSize(pos)

        for block in self.blocks:
            color = randint(0x888888, 0xDDDDDD)
            for addr in block:
                SetColor(addr, CIC_ITEM, color)

        return len(set(self.addesses))


class CodeGenerator:
    def __init__(self, begin, end, blocks):
        self.begin = begin
        self.end = end
        self.blocks = blocks

    def decode(self, pos):
        mnem = GetMnem(pos)
        op0 = GetOpnd(pos, 0)
        op1 = op0
        op2 = GetOpnd(pos, 1)
        if mnem in ['push', 'pop']:
            return ''
        if GetOpType(pos, 1) == 5:
            m = re.search('\-?([0-9A-F]+)h?', op2)
            op2 = '0x%s' % m.group(1)
        if mnem == 'lea':
            if GetOpType(pos, 1) == 4:
                m = re.search('\[(e\w\w)(\+|-)([0-9A-F]+)h?\]', op2)
                op1 = m.group(1)
                mnem = m.group(2)
                op2 = '0x%s' % m.group(3)
            elif GetOpType(pos, 1) == 3:
                m = re.search('\[(e\w\w)(\+|-)(.+)\]', op2)
                op1 = m.group(1)
                mnem = m.group(2)
                op2 = m.group(3)
        elif mnem == 'imul':
            op1 = op2
            op2 = GetOpnd(pos, 2)
        else:
            if GetOpType(pos, 1) == 4:
                m = re.search('\[ebp\+(.+)\]', op2)
                op2 = m.group(1)
            if GetOpType(pos, 0) == 4:
                m = re.search('\[ebp\+(.+)\]', op0)
                op0 = m.group(1)
                op1 = op0
            if GetOpType(pos, 1) == 3:
                m = re.search('\[(.+)\]', op2)
                op2 = m.group(1)
            if GetOpType(pos, 0) == 3:
                m = re.search('\[(.+)\]', op0)
                op0 = m.group(1)
                op1 = op0
            if GetOpnd(pos, 1) == 'cl':
                op2 = 'ecx'


        if mnem == 'mov':
            return '%s = %s' % (op0, op2)
        elif GetOpType(pos, 1) == 0:
            return '%s = _%s(%s)' % (op0, mnem, op1)
        elif mnem in ['-', '+']:
            return '%s = _%s(%s, %s)' % (op0, 'sub' if mnem == '-' else 'add', op1, op2)
        elif mnem in ['add', 'sub', 'and', 'xor', 'or', 'shr', 'shl', 'ror', 'rol', 'imul']:
            return '%s = _%s(%s, %s)' % (op0, mnem, op1, op2)
        else:
            Exception('Error [mnem = %s, op0 = %s, op1 = %s, op2 = %s]' % (mnem, op0, op1, op2))

    def generate(self):
        with open('/home/roman/encrypt.py', 'w') as f:
            pos = self.begin
            skip_addr = [0x0040DDA2, 0x0040DDAD, 0x0041B514, 0x0041B548]
            f.write(code)
            f.write('def encrypt(ecx, edx, a3, a4):\n')
            while pos < self.end:
                if pos not in skip_addr:
                    f.write('    ' + self.decode(pos) + '\n')
                pos += ItemSize(pos)
            f.write(code_end)

def main():
    begin = 0x0040DD80
    end = 0x0041B562

    decoder = Decoder(begin, end)
    decoder.do()
    print 22127 - decoder.colorize()
    # gen = CodeGenerator(0x0040DD86, 0x0041B54F, [])
    # gen.generate()

if __name__ == '__main__':
    main()