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.disassembly = list()

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

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

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

    @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,
                    opcode.size,
                ]
            )

        return opcodes


class X86_intel(_CapstoneBase):
    capstone = Cs(CS_ARCH_X86, CS_MODE_16)
    arch = "x86-16"


class X86(_CapstoneBase):
    capstone = Cs(CS_ARCH_X86, CS_MODE_32)
    arch = "x86-32"


class X86_64(_CapstoneBase):
    capstone = Cs(CS_ARCH_X86, CS_MODE_64)
    arch = "x86-64"


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


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


class ARM_64(_CapstoneBase):
    capstone = Cs(CS_ARCH_ARM64, CS_MODE_ARM)
    arch = "ARM 64"


class MIPS_32_eb(_CapstoneBase):
    capstone = Cs(CS_ARCH_MIPS, CS_MODE_MIPS32 + CS_MODE_BIG_ENDIAN)
    arch = "MIPS-32 (Big-endian)"


class MIPS_64_el(_CapstoneBase):
    capstone = Cs(CS_ARCH_MIPS, CS_MODE_MIPS64 + CS_MODE_LITTLE_ENDIAN)
    arch = "MIPS-64-EL (Little-endian)"


class PPC_64(_CapstoneBase):
    capstone = Cs(CS_ARCH_PPC, CS_MODE_BIG_ENDIAN)
    arch = "PPC-64"


class Sparc(_CapstoneBase):
    capstone = Cs(CS_ARCH_SPARC, CS_MODE_BIG_ENDIAN)
    arch = "Sparc"


class SparcV9(_CapstoneBase):
    capstone = Cs(CS_ARCH_SPARC, CS_MODE_BIG_ENDIAN + CS_MODE_V9)
    arch = "SparcV9"


class SystemZ(_CapstoneBase):
    capstone = Cs(CS_ARCH_SYSZ, 0)
    arch = "SystemZ"


class XCore(_CapstoneBase):
    capstone = Cs(CS_ARCH_XCORE, 0)
    arch = "XCore"