subdisassem/subdisassem/rizin_wrapper.py

from pathlib import Path
import logging
import rzpipe
import json


class _RizinBase:
    def __init__(self, path: Path):
        rz_pipe = rzpipe.open(str(path.absolute()))

        for cmd in self.arch_cmds:
            rz_pipe.cmd(cmd)

        rz_pipe.cmd("aa")
        result = rz_pipe.cmd("pdj")
        rz_pipe.quit()
        self.disassembly = json.loads(result)

    def __repr__(self) -> str:
        return self.objdump

    def __len__(self) -> int:
        return len(self.disasm)

    def __lt__(self, other):
        return len(self) < len(other)

    def __contains__(self, name: str):
        return hasattr(self, name)

    @property
    def objdump(self) -> str:
        if "_objdump" in self:
            return self._objdump

        _objdump = str()

        for each in self.disassembly:
            offset = each.get("offset")
            opcode = each.get("opcode")
            if opcode:
                _objdump += f"{offset:#02x}:\t{opcode}\n"

        self._objdump = _objdump
        return self._objdump

    @property
    def disasm(self) -> list:
        if "_disasm" in self:
            return self._disasm

        _disasm = list()

        for each in self.disassembly:
            offset = each.get("offset")
            opcode = each.get("opcode")

            if opcode:
                mnemonic = opcode.split(" ")[0]
                opcode = opcode.split(" ")[1:]
                _disasm.append([offset, mnemonic, opcode])

        self._disasm = _disasm
        return self._disasm

    @property
    def rets(self) -> list:
        if "_rets" in self:
            return self._rets

        _rets = list()

        for each in self.disasm:
            _, mnemonic, _ = each

            if mnemonic and "ret" in mnemonic:
                _rets.append(mnemonic)

        self._rets = _rets
        return self._rets

    @property
    def ret_rates(self) -> list:
        if "_ret_rates" in self:
            return self._ret_rates

        rates = dict()

        for mnemonic in set(self.rets):
            rates[mnemonic] = self.rets.count(mnemonic)

        _ret_rates = sorted(
            ((value, key) for (key, value) in rates.items()), reverse=True
        )

        self._ret_rates = _ret_rates
        return self._ret_rates

    @property
    def mnemonic_rates(self) -> list:
        if "_mnemonic_rates" in self:
            return self._mnemonic_rates

        mnemonics = list()

        for each in self.disasm:
            _, mnemonic, _ = each

            if mnemonic:
                mnemonics.append(mnemonic)

        rates = dict()

        for mnemonic in set(mnemonics):
            rates[mnemonic] = mnemonics.count(mnemonic)

        _mnemonic_rates = sorted(
            ((value, key) for (key, value) in rates.items()), reverse=True
        )

        self._mnemonic_rates = _mnemonic_rates
        return self._mnemonic_rates


class _6502_8(_RizinBase):
    arch_cmds = ["e asm.arch=6502", "e asm.bits=8"]
    name = "6502/NES/C64/Tamagotchi/T-1000 CPU"


class _6502_16(_RizinBase):
    arch_cmds = ["e asm.arch=6502", "e asm.bits=16"]
    name = "6502/NES/C64/Tamagotchi/T-1000 CPU"


class _8051(_RizinBase):
    arch_cmds = ["e asm.arch=8051", "e asm.bits=8"]
    name = "8051 Intel CPU"


class amd29k(_RizinBase):
    arch_cmds = ["e asm.arch=amd29k", "e asm.bits=32"]
    name = "AMD 29k RISC CPU"


class arc_16(_RizinBase):
    arch_cmds = ["e asm.arch=arc", "e asm.bits=16"]
    name = "Argonaut RISC Core"


class arc_32(_RizinBase):
    arch_cmds = ["e asm.arch=arc", "e asm.bits=32"]
    name = "Argonaut RISC Core"


class arm_as_16(_RizinBase):
    arch_cmds = ["e asm.arch=arm.as", "e asm.bits=16"]
    name = "as ARM Assembler (use RZ_ARM32_AS and RZ_ARM64_AS environment)"


class arm_as_32(_RizinBase):
    arch_cmds = ["e asm.arch=arm.as", "e asm.bits=32"]
    name = "as ARM Assembler (use RZ_ARM32_AS and RZ_ARM64_AS environment)"


class arm_as_64(_RizinBase):
    arch_cmds = ["e asm.arch=arm.as", "e asm.bits=64"]
    name = "as ARM Assembler (use RZ_ARM32_AS and RZ_ARM64_AS environment)"


class arm_16(_RizinBase):
    arch_cmds = ["e asm.arch=arm", "e asm.bits=16"]
    name = "Capstone ARM disassembler"


class arm_32(_RizinBase):
    arch_cmds = ["e asm.arch=arm", "e asm.bits=32"]
    name = "Capstone ARM disassembler"


class arm_64(_RizinBase):
    arch_cmds = ["e asm.arch=arm", "e asm.bits=64"]
    name = "Capstone ARM disassembler"


class arm_gnu_16(_RizinBase):
    arch_cmds = ["e asm.arch=arm.gnu", "e asm.bits=16"]
    name = "Acorn RISC Machine CPU"


class arm_gnu_32(_RizinBase):
    arch_cmds = ["e asm.arch=arm.gnu", "e asm.bits=32"]
    name = "Acorn RISC Machine CPU"


class arm_gnu_64(_RizinBase):
    arch_cmds = ["e asm.arch=arm.gnu", "e asm.bits=64"]
    name = "Acorn RISC Machine CPU"


class arm_wine_16(_RizinBase):
    arch_cmds = ["e asm.arch=arm.winedbg", "e asm.bits=16"]
    name = "WineDBG's ARM disassembler"


class arm_wine_32(_RizinBase):
    arch_cmds = ["e asm.arch=arm.winedbg", "e asm.bits=32"]
    name = "WineDBG's ARM disassembler"


class x86_16(_RizinBase):
    arch_cmds = ["e asm.arch=x86", "e asm.bits=16"]


class x86_32(_RizinBase):
    arch_cmds = ["e asm.arch=x86", "e asm.bits=32"]


class x86_64(_RizinBase):
    arch_cmds = ["e asm.arch=x86", "e asm.bits=64"]