from capstone import Cs
from capstone import (
    CS_ARCH_ARM,
    CS_ARCH_ARM64,
    CS_ARCH_MIPS,
    CS_ARCH_PPC,
    CS_ARCH_SPARC,
    CS_ARCH_SYSZ,
    CS_ARCH_X86,
    CS_ARCH_XCORE,
)
from capstone import (
    CS_MODE_16,
    CS_MODE_32,
    CS_MODE_64,
    CS_MODE_ARM,
    CS_MODE_BIG_ENDIAN,
    CS_MODE_LITTLE_ENDIAN,
    CS_MODE_MCLASS,
    CS_MODE_MICRO,
    CS_MODE_MIPS3,
    CS_MODE_MIPS32,
    CS_MODE_MIPS32R6,
    CS_MODE_MIPS64,
    CS_MODE_THUMB,
    CS_MODE_V8,
    CS_MODE_V9,
)
import logging


class _CapstoneBase:
    def __init__(self, payload: bytes, offset: int = 0):
        self.arch = self.__class__.__name__
        self.capstone.skipdata = True
        self.capstone.skipdata_setup = ("unknown", None, None)

        disassembly = list()

        for opcode in self.capstone.disasm(payload, offset):
            disassembly.append(opcode)

        if disassembly:
            self.disassembly = disassembly
        else:
            logging.debug("disassembly_empty")
            self.disassembly = list()

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

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

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

    @property
    def objdump(self) -> str:
        opcodes = str()

        for opcode in self.disassembly:
            opcodes += f"{opcode.address:#02x}:\t{opcode.mnemonic}\t{opcode.op_str}\n"

        return opcodes

    @property
    def disasm(self) -> list:
        opcodes = list()

        for opcode in self.disassembly:
            opcodes.append([opcode.address, opcode.mnemonic, opcode.op_str])

        return opcodes

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

        self._rets = list()

        for opcode in self.disassembly:
            if "ret" in opcode.mnemonic:
                self._rets.append(opcode.mnemonic)

        return self._rets

    @property
    def ret_rates(self) -> list:
        rates = dict()

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

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

        return listed

    @property
    def mnemonic_rates(self) -> list:
        mnemonics = list()

        for opcode in self.disassembly:
            mnemonics.append(opcode.mnemonic)

        rates = dict()

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

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

        return listed


class x86_16(_CapstoneBase):
    capstone = Cs(CS_ARCH_X86, CS_MODE_16)


class x86_32(_CapstoneBase):
    capstone = Cs(CS_ARCH_X86, CS_MODE_32)


class x86_64(_CapstoneBase):
    capstone = Cs(CS_ARCH_X86, CS_MODE_64)


class armv7(_CapstoneBase):
    capstone = Cs(CS_ARCH_ARM, CS_MODE_ARM)


class thumb2(_CapstoneBase):
    capstone = Cs(CS_ARCH_ARM, CS_MODE_THUMB)


class aarch64(_CapstoneBase):
    capstone = Cs(CS_ARCH_ARM64, CS_MODE_ARM)


class mips32(_CapstoneBase):
    capstone = Cs(CS_ARCH_MIPS, CS_MODE_MIPS32 + CS_MODE_BIG_ENDIAN)


class mips64_el(_CapstoneBase):
    capstone = Cs(CS_ARCH_MIPS, CS_MODE_MIPS64 + CS_MODE_LITTLE_ENDIAN)


class ppc64(_CapstoneBase):
    capstone = Cs(CS_ARCH_PPC, CS_MODE_BIG_ENDIAN)


class sparc(_CapstoneBase):
    capstone = Cs(CS_ARCH_SPARC, CS_MODE_BIG_ENDIAN)


class sparcv9(_CapstoneBase):
    capstone = Cs(CS_ARCH_SPARC, CS_MODE_BIG_ENDIAN + CS_MODE_V9)


class systemz(_CapstoneBase):
    capstone = Cs(CS_ARCH_SYSZ, 0)


class xcore(_CapstoneBase):
    capstone = Cs(CS_ARCH_XCORE, 0)