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):
        self.arch = self.__class__.__name__
        pass

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

    def __len__(self) -> int:
        if not self.disassembly:
            logging.debug(
                f"payload_missing: use {self.__class__}.load(payload=bytes) prior"
            )
            return 0

        return len(self.disassembly)

    def load(self, payload: bytes, offset: int = 0):
        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()

    @property
    def objdump(self) -> str:
        if not self.disassembly:
            logging.debug(
                f"payload_missing: use {self.__class__}.load(payload=bytes) prior"
            )
            return ""

        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:
        if not self.disassembly:
            logging.debug(
                f"payload_missing: use {self.__class__}.load(payload=bytes) prior"
            )
            return []

        opcodes = list()

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

        return opcodes


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


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


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


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


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


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


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


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


class PPC_64(_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)