decode_automaton.py

# (c) 2024 Michael Karcher
# Distributable according to the 2-clause BSD license
# No warranty!

import struct
from dataclasses import dataclass
from typing import List

@dataclass
class next_state:
    delta_actions: int
    delta_mask: int

@dataclass
class final_state:
    result_code: int

def parse_action(statedesc: bytes):
    (dm_flags, da) = struct.unpack("<HH", statedesc)
    dm = dm_flags & ~7
    flags = dm_flags & 7
    if flags & 4:
        return final_state(da)
    if flags & 2:
        pass
    else:
        da = -da
    if flags & 1:
        pass
    else:
        dm = -dm
    return next_state(delta_actions = da-4, delta_mask = dm-16)

def parse_actiontable(serialized: bytes):
    while True:
        state = serialized[:4]
        if not state:
            return
        serialized = serialized[4:]
        yield parse_action(state)

def parse_mask(serialized: bytes) -> str:
    binmask = int.from_bytes(serialized, "little")
    result = ""
    for i in range(128):
        if binmask & (1 << i):
            result += chr(i)
    return result

class automaton:
    def __init__(self, maskdata : bytes, actiondata : bytes):
        self.maskdata = maskdata
        self.actiondata = actiondata

    def all_strings(self):
        init_mask_cursor = len(self.maskdata)-16
        init_mask = parse_mask(self.maskdata[-16:])
        init_action_cursor = len(self.actiondata)-4*len(init_mask)
        yield from self._all_recursive("", init_mask_cursor, init_action_cursor)

    def _all_recursive(self, prefix, mask_cursor, action_cursor):
        mask = parse_mask(self.maskdata[mask_cursor:mask_cursor+16])
        actions = parse_actiontable(self.actiondata[action_cursor:action_cursor + 4*len(mask)])
        for (c, action) in zip(mask, actions):
            string = prefix + c
            if isinstance(action, final_state):
                yield (action.result_code, string)
            else:
                yield from self._all_recursive(string, mask_cursor + action.delta_mask, action_cursor + action.delta_actions)

def load_malware_sample_as_elf():
    from elftools.elf.elffile import ELFFile
    with open("liblzma_la-crc64-fast.o.this-is-malware", "rb") as malware_object:
        malware_elf = ELFFile(malware_object)
        mask_data = malware_elf.get_section_by_name(".rodata.crc64_clmul1").data()
        action_data = malware_elf.get_section_by_name(".rodata.lzip_decode0").data()
        return automaton(mask_data, action_data)

def ida_get_seg_date(segname):
    import ida_segment
    seg=ida_segment.get_segm_by_name(segname)
    if seg:
        return idc.get_bytes(seg.start_ea, seg.size())
    else:
        return None

def load_malware_sample_inside_ida():
    mask_data  = ida_get_seg_date(".rodata.crc64_clmul1")
    action_data = ida_get_seg_date(".rodata.lzip_decode0")
    return automaton(mask_data, action_data)

def load_malware_sample():
    try:
        # check [Is it running in the IDAPython]
        import idc
        return load_malware_sample_inside_ida()
    except Exception as e:
        # not-inside IDAPython
        return load_malware_sample_as_elf()

def main():
    a = load_malware_sample()
    for (id, string) in a.all_strings():
        print(f"{id:4x}: {repr(string)}")

if __name__ == "__main__":
    main()