Lifter solution to GoogleCTF 2022 eldar

lifter.py

# Author: hgarrereyn
# Desc: Lifter solution for GoogleCTF 2022 eldar

import lief
from collections import namedtuple
from dataclasses import dataclass
from typing import Any
from capstone import *
from z3 import *
import numpy as np

md = Cs(CS_ARCH_X86, CS_MODE_64)

b = None
try:
    b = lief.ELF.parse('./eldar')
except:
    raise Exception('Must have the ./eldar binary in cwd')

rela = [x for x in b.sections if x.name == '.rela.dyn'][0]
dynsym = [x for x in b.sections if x.name == '.dynsym'][0]


@dataclass
class Symbol(object):
    idx: int
    def __repr__(self):
        return f's{self.idx}'
        
@dataclass
class Reloc(object):
    idx: int
    def __repr__(self):
        return f'r{self.idx}'

@dataclass
class Ref(object):
    val: Any
    def __repr__(self):
        return f'&{self.val}'

@dataclass
class SymAddr(object):
    sym: Symbol
    field: str
    def __repr__(self):
        return f'{self.sym}.{self.field}'

@dataclass
class RelocAddr(object):
    reloc: Reloc
    field: str
    def __repr__(self):
        return f'{self.reloc}.{self.field}'

    def vaddr(self):
        off = 0
        match self.field:
            case 'r_address':off = 0
            case 'r_info': off = 8
            case 'r_addend': off = 16
        
        return (self.reloc.idx * 24) + off + rela.virtual_address
    
@dataclass
class FlagAddr(object):
    idx: int
    def __repr__(self):
        return f'flag[{self.idx}]'

@dataclass
class OutAddr(object):
    idx: int
        
    def __repr__(self):
        return f'out[{self.idx}]'
    
@dataclass
class ArrAddr(object):
    idx: int
        
    def __repr__(self):
        return f'arr[{self.idx}]'

BaseAddr = namedtuple('baseaddr', [])
FailAddr = namedtuple('fail', [])

def format_addr(addr: int):
    if addr >= rela.virtual_address and addr < rela.virtual_address + rela.size:
        offset = addr - rela.virtual_address
        r_offset = (offset // 24)
        r_rem = offset % 24
        
        if r_offset >= 3 and r_offset <= 88:
            arr_idx = (r_offset - 3) * 3 + (r_rem // 8)
            return ArrAddr(arr_idx)
        elif r_offset == 89:
            return OutAddr(r_rem)
        
        match r_rem:
            case 0: return RelocAddr(Reloc(r_offset), 'r_address')
            case 8: return RelocAddr(Reloc(r_offset), 'r_info')
            case 16: return RelocAddr(Reloc(r_offset), 'r_addend')
            case _: return RelocAddr(Reloc(r_offset), r_rem)
    elif addr > dynsym.virtual_address and addr < dynsym.virtual_address + dynsym.size:
        offset = addr - dynsym.virtual_address
        r_offset = (offset // 24)
        r_rem = offset % 24
        
        match r_rem:
            case 0: return SymAddr(Symbol(r_offset), 'st_name')
            case 8: return Symbol(r_offset)
            case 16: return SymAddr(Symbol(r_offset), 'st_size')
            case _: return SymAddr(Symbol(r_offset), r_rem)
    elif addr >= 0x404040 and addr < 0x404040+29:
        off = addr-0x404040
        return FlagAddr(off)
    elif addr == 0x804000:
        return BaseAddr()
    elif addr == 0x404060:
        return FailAddr()
    else:
        return addr

def to_sym(name):
    assert len(name) == 1
    return Symbol(ord(name[0]))

Rel = namedtuple('REL', ['dst','val','ridx'])
Copy = namedtuple('CPY', ['dst', 'symbol', 'ridx'])
R64 = namedtuple('R64', ['dst','symbol','addend','ridx'])
R32 = namedtuple('R32', ['dst','symbol','addend','ridx'])

def parse(b) -> list:
    print('[*] Loading relocations...')
    relocs = list(b.relocations)
    
    print('[*] Parsing...')
    instructions = []
    for i in range(3, len(relocs)):
        r = relocs[i]
        match r.type:
            case 1: # R64
                instructions.append(R64(format_addr(r.address), to_sym(r.symbol.name), format_addr(r.addend), i))
            case 5: # CPY
                instructions.append(Copy(format_addr(r.address), to_sym(r.symbol.name), i))
            case 8: # REL
                instructions.append(Rel(format_addr(r.address), format_addr(r.addend), i))
            case 10: # R32
                instructions.append(R32(format_addr(r.address), to_sym(r.symbol.name), format_addr(r.addend), i))
                
    return instructions

Mov = namedtuple('mov', ['dst', 'src', 'sz', 'ridx'])
Add = namedtuple('add', ['dst', 'src', 'addend', 'ridx'])

def lift_mov_add(instructions):
    idx = 0
    
    sizes = []
    curr = [8] * 8
    sizes.append(curr)

    for instr in instructions:
        c = list(curr)
        match instr:
            case Rel(SymAddr(Symbol(idx), 'st_size'), val, ridx):
                c[idx] = val
        sizes.append(c)
    
    while idx < len(instructions):
        match instructions[idx]:
            case Rel(dst, val, ridx):
                instructions[idx] = Mov(dst, Ref(val), 8, ridx)
            case Copy(dst, sym, ridx):
                instructions[idx] = Mov(dst, sym, sizes[idx][sym.idx], ridx)
            case R64(dst, sym, add, ridx):
                instructions[idx] = Add(dst, sym, add, ridx)
        idx += 1
    return instructions

def remove_sizes(instructions):
    # Sizes are now nops
    idx = 0
    while idx < len(instructions):
        match instructions[idx]:
            case Mov(SymAddr(Symbol(s), 'st_size'), _, _, _) if s != 3:
                instructions[idx:idx+1] = []

        idx += 1
    return instructions

def lift_indirect(instructions):
    # [0349] :: mov r350.r_addend, s2
    # [0350] :: add s4, s4, 0
    # [0351] :: mov r350.r_addend, &0
    idx = 0
    while idx < len(instructions):
        match instructions[idx:idx+3]:
            case [
                Mov(RelocAddr(Reloc(rel_1), 'r_addend'), Symbol(sidx_1), sz_1, ridx_1),
                Add(dst_2, sym_2, _, ridx_2),
                Mov(RelocAddr(Reloc(rel_3), 'r_addend'), Ref(0), sz_3, _),
            ] if (
                (rel_1 == ridx_2) and (rel_3 == ridx_2)
            ):
                instructions[idx:idx+3] = [
                    Add(dst_2, sym_2, Symbol(sidx_1), ridx_1)
                ]

        idx += 1
    return instructions

Block = namedtuple('block', ['arr', 'flag', 'ridx'])
Output = namedtuple('output', ['out', 'arr', 'ridx'])

def lift_block(instructions):
    # [0378] :: mov s2, &arr[1]
    # [0008] :: add s4, s4, s2
    # [0382] :: mov s2, &flag[1]
    # [0384] :: movb s7, s2
    # [0385] :: mov s2, &s7
    # [0008] :: add s4, s4, s2
    # [0390] :: r32 s4.st_value_p1, s1, 0
    # [0391] :: mov s2, &arr[1]
    # [0392] :: mov s6, s2
    # [0393] :: mov s2, &s4
    # [0008] :: add s5, s4, s2
    # [0397] :: mov s2, &s5
    # [0008] :: add s5, s5, s2
    # [0008] :: add s5, s5, s2
    # [0404] :: add s5, s5, arr[0]
    # [0405] :: mov arr[1], s5
    # [0406] :: mov r407.r_address, s2
    # [0407] :: add 0, s6, 0
    idx = 0
    while idx < len(instructions):
        match instructions[idx:idx+18]:
            case [
                Mov(_,arr,_,ridx),
                Add(_,_,_,_),
                Mov(_,flag,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                R32(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                Add(_,_,_,_),
                Add(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
            ]:
                instructions[idx:idx+18] = [
                    Block(arr, flag, ridx)
                ]
        idx += 1
    return instructions

Reset = namedtuple('reset', ['ridx'])
ShuffleBlock = namedtuple('shuffleblock', ['f1', 'f2', 'ridx'])

def lift_reset(instructions):
    idx = 0
    while idx < len(instructions) - 256:
        good = True

        for i in range(256):
            op = instructions[idx+i]
            if type(op) == Mov:
                dst, src, _, _ = op
                if dst != ArrAddr(i) or src != Ref(i):
                    good = False
                    break
            else:
                good = False
                break

        if good:
            instructions[idx:idx+256] = [Reset(instructions[idx].ridx)]

        idx += 1
    return instructions

def lift_shuffle_block(instructions):
    idx = 0
    while idx < len(instructions) - 256:
        good = True

        for i in range(256):
            op = instructions[idx+i]
            if type(op) == Block:
                arr, flag, ridx = op
                if arr != Ref(ArrAddr(i)):
                    good = False
                    break
            else:
                good = False
                break

        if good:
            instructions[idx:idx+256] = [ShuffleBlock(instructions[idx].flag, instructions[idx+1].flag, instructions[idx].ridx)]

        idx += 1
    return instructions

Output = namedtuple('output', ['out', 'arr', 'ridx'])

def lift_output(instructions):
    idx = 0
    while idx < len(instructions):
        match instructions[idx:idx+26]:
            case [
                Mov(_,arr,_,ridx),
                Add(_,_,_,_),
                R32(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                Add(_,_,_,_),
                Add(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                R32(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                Mov(_,_,_,_),
                Add(_,_,_,_),
                Add(_,_,_,_),
                Add(_,_,_,_),
                Mov(out,_,_,_),
            ]:
                instructions[idx:idx+26] = [Output(out, arr, ridx)]
        idx += 1
    return instructions

MultAdd = namedtuple('multadd', ['out', 'val', 'k', 'ridx'])

def lift_multadd(instructions):
    idx = 0
    while idx < len(instructions):
        match instructions[idx:idx+3]:
            # block prefix
            case [
                Mov(Symbol(2), out, _, ridx),
                Mov(Symbol(5), Symbol(2), _, _),
                Mov(Symbol(6), Ref(0), _, _),
            ]:
                k = 0
                double = False

                ptr = idx + 3

                good = True
                while ptr < len(instructions):
                    match instructions[ptr]:
                        case Mov(Symbol(2), Ref(Symbol(6)), _, _):
                            double = True
                        case Mov(Symbol(2), Ref(Symbol(5)), _, _):
                            double = False
                        case Add(Symbol(6), Symbol(6), Symbol(2), _):
                            k = (k * 2) if double else (k + 1)
                        case Add(Symbol(7), Symbol(7), Symbol(2), _):
                            ptr += 1
                            break
                        case _:
                            good = False
                            break

                    ptr += 1
                
                if good:
                    instructions[idx:ptr] = [
                        MultAdd(Symbol(7), out, k, ridx)
                    ]

        idx += 1

    return instructions

Trunc = namedtuple('trunc', ['val', 'k', 'ridx'])

def lift_truncate(instructions):
    idx = 0
    while idx < len(instructions):
        match instructions[idx:idx+2]:
            case [
                Mov(Symbol(2), Ref(SymAddr(Symbol(5), 11)), _, ridx),
                Mov(SymAddr(Symbol(7), 11), Symbol(2), 5, _)
            ]:
                instructions[idx:idx+2] = [
                    Trunc(Symbol(7), 0xffffff, ridx)]
        idx += 1
    return instructions

ArraySlots = namedtuple('arr', ['values', 'ridx'])

def lift_array_slots(instructions):
    idx = 0
    while idx < len(instructions):
        match instructions[idx]:
            case Mov(BaseAddr(), Ref(0), _, ridx):
                ptr = idx+1
                while ptr < len(instructions):
                    op = instructions[ptr]
                    if type(op) != Mov or op.dst != BaseAddr():
                        break
                    ptr += 1

                start = idx
                end = ptr

                data = []

                # Check for movs into array.
                vstart = RelocAddr(Reloc(ridx), 'r_address').vaddr()
                offset = 0
                while end + offset < len(instructions) and offset < ((end - start) * 3):
                    op = instructions[end + offset]
                    if type(op) == Mov and type(op.dst) is RelocAddr and op.dst.vaddr() == vstart + (offset * 8):
                        data.append(op.src.val)
                    else:
                        break
                    offset += 1

                if len(data) > 0:
                    data += [0] * (((end - start) * 3) - len(data))   
                    instructions[idx:end+offset] = [
                        ArraySlots(data, ridx)
                    ]

        idx += 1
    return instructions

Shellcode = namedtuple('shellcode', ['dst', 'code', 'ridx'])

def lift_shellcode(instructions):
    idx = 0
    while idx < len(instructions):
        match instructions[idx:idx+6]:
            case [
                ArraySlots(values, ridx),
                Mov(Symbol(3), Ref(RelocAddr(Reloc(rel2), 'r_address')), _, _),
                Mov(SymAddr(Symbol(3), 'st_name'), _, _, _),
                Add(dst, Symbol(3), _, _),
                Mov(Symbol(2), _, _, _),
                Mov(RelocAddr(Reloc(rel6), 'r_address'), Symbol(2), _, _)
            ] if (rel2 == ridx) and (rel6 == ridx):
                instructions[idx:idx+6] = [
                    Shellcode(dst, b''.join([(x & 0xffffffffffffffff).to_bytes(8, 'little') for x in values]), ridx)
                ]
        idx += 1
    return instructions

Aop = namedtuple('aop', ['dst', 'op', 'val', 'k', 'ridx'])

def lift_aop(instructions):
    idx = 0
    while idx < len(instructions):
        match instructions[idx:idx+5]:
            case [
                Mov(Symbol(2), val, _, ridx),
                Mov(Symbol(5), Symbol(2), _, _),
                Shellcode(_, data, _),
                Mov(Symbol(2), Ref(Symbol(5)), _, _),
                Add(dst, dst2, Symbol(2), _)
            ] if len(data) == 24 and (dst == dst2):
                op = next(md.disasm(data, 0))

                t = op.mnemonic
                k = int(op.op_str.split(', ')[-1], 16)

                instructions[idx:idx+5] = [
                    Aop(dst, t, val, k, ridx)
                ]

        idx += 1
    return instructions

def solve_end_flag(instructions):
    orig = [BitVec('f%d' % i, 8) for i in range(16,28)]
    flag = ([0] * 16) + [ZeroExt(24, x) for x in orig]

    s = Solver()
    for o in orig:
        s.add(UGT(o, 0x20))
        s.add(ULT(o, 0x7f))

    idx = 0
    while idx < len(instructions):
        match instructions[idx]:
            case Mov(Symbol(7), Ref(v), _, ridx) if type(v) is int:
                x = BitVecVal(v, 32)

                ptr = idx+1
                while ptr < len(instructions):
                    match instructions[ptr]:
                        case MultAdd(Symbol(7), Ref(FlagAddr(f)), k, _):
                            x += (flag[f] * k)
                        case Aop(Symbol(7), op, Ref(FlagAddr(f)), k, _):
                            v = None
                            match op:
                                case 'and': v = flag[f] & k
                                case 'xor': v = flag[f] ^ k
                                case 'or': v = flag[f] | k
                                case 'rol': v = ZeroExt(24, RotateLeft(Extract(7,0,flag[f]), k))
                                case 'ror': v = ZeroExt(24, RotateRight(Extract(7,0,flag[f]), k))
                                case 'shl': v = ZeroExt(24, Extract(7,0,flag[f]) << k)
                                case 'shr': v = flag[f] >> k
                                case _:
                                    raise Exception(f'unknown aop: {op}')
                            x += v
                        case Trunc(Symbol(7), k, _):
                            s.add(x == 0)
                        case _:
                            break

                    ptr += 1

        idx += 1

    print('solving...')
    print(s.check())

    m = s.model()
    flag = bytes([m.eval(o).as_long() for o in orig])

    return flag

def solve_out_arr(instructions):
    X = []
    Y = []

    idx = 0
    while idx < len(instructions):
        match instructions[idx]:
            case Mov(Symbol(7), Ref(v), _, ridx) if type(v) is int:
                row = []

                ptr = idx+1
                while ptr < len(instructions):
                    match instructions[ptr]:
                        case MultAdd(Symbol(7), Ref(OutAddr(_)), k, _):
                            row.append(k)
                        case Trunc(Symbol(7), k, _):
                            X.append(row)
                            Y.append(-v)
                        case _:
                            break

                    ptr += 1

        idx += 1

    a = np.array(X, dtype=np.uint32)
    b = np.array(Y, dtype=np.uint32)

    return [int(x) for x in np.linalg.solve(a,b)]

def solve_start_flag(output):
    def sim(a,b):
        arr = list(range(256))

        z = 0
        for i in range(256):
            p = a if i % 2 == 0 else b
            z = (z + arr[i] + p) & 0xff
            arr[i], arr[z] = arr[z], arr[i]

        out = [0,0,0]
        z = 0
        for i in range(3):
            z = (z + arr[i]) & 0xff
            arr[i], arr[z] = arr[z], arr[i]
            out[i] = arr[(arr[i] + arr[z]) & 0xff]
            
        return out

    def solve_chunk(k1,k2,k3):
        # Brute force chunk.
        for a in range(0x20, 0x7f):
            for b in range(0x20, 0x7f):
                out = sim(a,b)
                
                if abs(out[0] - k1) <= 1 and abs(out[1] - k2) <= 1 and abs(out[2] - k3) <= 1:
                    return (a,b)
        
        return None

    f = []
    for i in range(0,len(output),3):
        f += list(solve_chunk(*output[i:i+3]))

    return bytes(f)

def dump(instructions):
    for op in instructions:
        match op:
            case Mov(SymAddr(sym, 'st_name'), Ref(val), 8, ridx) if type(val) is int:
                name = val & 0xffffffff
                info = (val >> 4) & 0xff
                other = (val >> 5) & 0xff
                shndx = (val >> 6) & 0xffff
                print(f'[{ridx:04d}] :: setinfo {sym}, name=0x{name:x}, info=0x{info:x}, other=0x{other:x}, shndx=0x{shndx:x}')
            case Mov(BaseAddr(), Ref(0), _, ridx):
                print(f'[{ridx:04d}] :: [ARRAY SLOT]')
            case Mov(dst, src, 8, ridx):
                print(f'[{ridx:04d}] :: mov {dst}, {src}')
            case Mov(dst, src, sz, ridx):
                print(f'[{ridx:04d}] :: mov({sz}) {dst}, {src}')
            case Add(dst, src, add, ridx):
                print(f'[{ridx:04d}] :: add {dst}, {src}, {add}')
            case R32(dst, src, add, ridx):
                print(f'[{ridx:04d}] :: r32 {dst}, {src}, {add}')
            case Block(arr, flag, ridx):
                print(f'[{ridx:04d}] :: shuffle {arr}, {flag}')
            case Output(out, arr, ridx):
                print(f'[{ridx:04d}] :: output {out}, {arr}')
            case ShuffleBlock(f1, f2, ridx):
                print(f'[{ridx:04d}] :: shuffleblock {f1}, {f2}')
            case MultAdd(dst, val, k, ridx):
                print(f'[{ridx:04d}] :: madd {dst} += ({val} * {k})')
            case Aop(dst, op, val, k, ridx):
                print(f'[{ridx:04d}] :: aop {dst} += ({val} {op} {k})')
            case Reset(ridx):
                print(f'[{ridx:04d}] :: reset')
            case Trunc(val, k, ridx):
                print(f'[{ridx:04d}] :: trunc {val} &= 0x{k:x}')
            case ArraySlots(values, ridx):
                print(f'[{ridx:04d}] :: array [{", ".join([hex(x) for x in values])}]')
            case Shellcode(dst, code, ridx):
                print(f'[{ridx:04d}] :: exec {dst} <- {code.hex()}')
                print('-' * 20)
                for i in md.disasm(code, 0):
                    if i.mnemonic == 'ret':
                        break
                    print("    0x%x:\t%s\t%s" %(i.address, i.mnemonic, i.op_str.replace('0x8040e4', 's5').replace('0x8040cc', 's4')))
                print('-' * 20)
            case _:
                print(op)

LIFTS = [
    lift_mov_add,
    remove_sizes,
    lift_indirect,
    lift_block,
    lift_reset,
    lift_shuffle_block,
    lift_output,
    lift_multadd,
    lift_truncate,
    lift_array_slots,
    lift_shellcode,
    lift_aop,
]

def lift(instructions):
    for lift_fn in LIFTS:
        print(f'[*] {lift_fn.__name__}...')
        instructions = lift_fn(instructions)
    return instructions

instructions = parse(b)
instructions = lift(instructions)
dump(instructions)

out = solve_out_arr(instructions)
start = solve_start_flag(out)
end = solve_end_flag(instructions)

# CTF{H0p3_y0u_l1k3_3LF_m4g1c}
print(start + end)