"Matriochka step 4 (I did it again)" solution in the form of IDA script — Nuit du Hack Quals 2017

matriochka4_ida_solve.py

from __future__ import print_function

from idautils import *
from idaapi import *
import hashlib


def step_forward(addr, n=1, checks=None):
    if checks:
        assert(n == len(checks))
    for i in range(n):
        q = DecodeInstruction(addr)
        if checks:
            assert(q.get_canon_mnem() == checks[i])
        addr += q.size
    return addr


def process_one(addr):
    start = GetOperandValue(addr, 1)
    addr = step_forward(addr, 5, ['lea', 'mov', 'jmp', 'mov', 'movzx'])
    xor_value = GetOperandValue(addr, 1) & 0xFF
    addr = step_forward(addr, 6, ['xor', 'mov', 'mov', 'mov', 'add', 'mov'])
    end = GetOperandValue(addr, 1)
    addr = step_forward(addr, 3, ['lea', 'cmp', 'jb'])
    return (addr, start, end, xor_value)


def make_magic(addr=None):
    if not addr:
        addr = ScreenEA()
        assert(addr != BADADDR)
        _ = step_forward(addr, 1, ['jmp'])
        addr = GetOperandValue(addr, 0)
    addr, start1, end1, xor_value1 = process_one(addr)
    print('[*] XORing from {} to {} with value {}'.format(hex(start1), hex(end1), hex(xor_value1)))
    addr, start2, end2, xor_value2 = process_one(addr)
    print('[*] XORing from {} to {} with value {}'.format(hex(start2), hex(end2), hex(xor_value2)))
    # for i in range(start1, end1, 1):
    #     PatchByte(i, Byte(i) ^ xor_value1)
    for i in range(start2, end2, 1):
        PatchByte(i, Byte(i) ^ xor_value2)


def make_more_magic(addr=None):
    if not addr:
        addr = ScreenEA()
    print('[*] Starting from {}'.format(hex(addr)))
    try:
        while addr < 0x40E000:
            q = DecodeInstruction(addr)
            if q.get_canon_mnem() == 'jmp' and GetOperandValue(addr, 0) > addr:
                print('[-] Found candidate at {}'.format(hex(addr)))
                try:
                    make_magic(GetOperandValue(addr, 0))
                    if addr != 0x40bd06:
                        PatchByte(addr + 0, 0x31)
                        PatchByte(addr + 1, 0xC0)
                        PatchByte(addr + 2, 0xC9)
                        PatchByte(addr + 3, 0xC3)
                        PatchByte(addr + 4, 0xCC)
                        assert(Byte(addr - 2) == 0x74)
                        PatchByte(addr - 2, 0x75)
                        # for j in range(q.size):
                        #     PatchByte(addr + j, 0x90)
                except:
                    # print('[!] Failed')
                    pass          
            addr += q.size
    except:
        pass
    print('[!] Ended at {}'.format(hex(addr)))


def find_constraints():
    addr = 0x40b000
    constraints, temp = [], []
    while addr < 0x425000:
        try:
            q = DecodeInstruction(addr)
            if q.get_canon_mnem() == 'mov' and GetOperandValue(addr, 1) == 0:
                insns = ['mov', 'mov', 'mov', 'mov', 'jmp', 'mov', 'mov', 'mov', 'add', 'movzx', 'movsx', 'mov', 'mov', 'movzx', 'xor', 'mov', 'jmp', 'mov', 'and', 'neg', 'mov', 'mov', 'shr']
                _ = step_forward(addr, len(insns), insns)
                print('[*] Found constraint candidate at {}'.format(hex(addr)))
                addr = step_forward(addr, 1)
                cons0 = GetOperandValue(addr, 1)
                addr = step_forward(addr, 1)
                cons1 = GetOperandValue(addr, 1) & 0xFFFFFFFF
                addr = step_forward(addr, 34)
                cons2 = GetOperandValue(addr, 1) & 0xFFFFFFFF
                temp.append((cons0, cons1, cons2))
            elif q.get_canon_mnem() == 'retn' and len(temp) > 0:
                print('[-] {} constraints found in function'.format(len(temp)))
                constraints = temp[::-1] + constraints
                temp = []
        except:
            pass
        finally:
            addr += 1
    print('[!] Found constraints for {} bytes totally'.format(sum(x[0] for x in constraints)))
    constraints = constraints[::-1]
    ans = b''
    for con_i, con in enumerate(constraints):
        print('[-] Bruteforcing constraint {}/{}'.format(con_i, len(constraints)))
        results = []
        for a1 in range(128):
            for a2 in range(128):
                t = -1 & 0xFFFFFFFF
                t ^= a1
                for _ in range(8):
                    t = (t >> 1) ^ (-(t & 1) & con[1])
                    t &= 0xFFFFFFFF
                t ^= a2
                for _ in range(8):
                    t = (t >> 1) ^ (-(t & 1) & con[1])
                    t &= 0xFFFFFFFF
                if (~t) & 0xFFFFFFFF == con[2]:
                    results.append([a1, a2])
        if len(results) > 1:
            print('[!] More than 1 result found')
            break
        elif len(results) == 0:
            print('[!] No result found')
            break
        else:
            ans += ''.join(map(chr, results[0]))
            print(ans)
    print(ans)
    if len(ans) == 334:
        m = hashlib.md5()
        m.update(ans)
        print('[!] Answer:', m.hexdigest())
    # with open('C:\\Users\\vient\\Desktop\\ctf\\nuitduhack\\step4_cons.txt', 'w') as f:
    #     _ = f.write(repr(constraints[::-1]))


def main():
    make_more_magic(0x400900)
    find_constraints()


if __name__ == '__main__':
    main()

'''
___  ___      _        _            _     _         
|  \/  |     | |      (_)          | |   | |        
| .  . | __ _| |_ _ __ _  ___   ___| |__ | | ____ _ 
| |\/| |/ _` | __| `__| |/ _ \ / __| `_ \| |/ / _` |
| |  | | (_| | |_| |  | | (_) | (__| | | |   < (_| |
\_|  |_/\__,_|\__|_|  |_|\___/ \___|_| |_|_|\_\__,_|
                  I did it again.
'''