import logging
from typing import Generator, Iterable, List, Tuple

from volatility3.framework import constants, interfaces, objects, renderers
from volatility3.framework.configuration import requirements
from volatility3.framework.interfaces import configuration
from volatility3.framework.objects.utility import array_to_string
from volatility3.framework.renderers import format_hints
from volatility3.framework.symbols import intermed
from volatility3.framework.symbols.windows.extensions import pe
from volatility3.plugins.windows import modules

vollog = logging.getLogger(__name__)


class Passphrase(interfaces.plugins.PluginInterface):
    """VeraCrypt/TrueCrypt cached passphrase finder (driver .data scan)."""

    _version = (0, 1, 0)
    _required_framework_version = (2, 5, 2)

    @classmethod
    def get_requirements(cls) -> List[configuration.RequirementInterface]:
        return [
            requirements.ModuleRequirement(
                "kernel",
                description="Windows kernel",
                architectures=["Intel32", "Intel64"],
            ),
            requirements.VersionRequirement(
                name="modules", component=modules.Modules, version=(3, 0, 0)
            ),
            requirements.IntRequirement(
                name="min-length",
                description="Minimum length of passphrases to identify",
                default=5,
                optional=True,
            ),
            requirements.StringRequirement(
                name="driver",
                description=(
                    "Driver name substring to scan (case-insensitive), "
                    "e.g. 'veracrypt', 'veracrypt-x64.sys', 'truecrypt.sys'"
                ),
                default="veracrypt",
                optional=True,
            ),
        ]

    def scan_module(
        self, module_base: int, layer_name: str
    ) -> Generator[Tuple[int, str], None, None]:
        pe_table_name = intermed.IntermediateSymbolTable.create(
            self.context, self.config_path, "windows", "pe", class_types=pe.class_types
        )
        dos_header: pe.IMAGE_DOS_HEADER = self.context.object(
            pe_table_name + constants.BANG + "_IMAGE_DOS_HEADER",
            layer_name,
            module_base,
        )

        data_section: objects.StructType = next(
            sec
            for sec in dos_header.get_nt_header().get_sections()
            if array_to_string(sec.Name) == ".data"
        )
        base: int = data_section.VirtualAddress + module_base
        size: int = data_section.Misc.VirtualSize

        # Looking at `Length` in TrueCrypt/Common/Password.h::Password struct
        DWORD_SIZE_BYTES: int = 4
        fmt = objects.DataFormatInfo(
            length=DWORD_SIZE_BYTES, byteorder="little", signed=True
        )
        int32 = objects.templates.ObjectTemplate(
            objects.Integer, pe_table_name + constants.BANG + "int", data_format=fmt
        )
        count, not_aligned = divmod(size, DWORD_SIZE_BYTES)
        if not_aligned:
            raise ValueError("PE data section not DWORD-aligned!")

        lengths = self.context.object(
            pe_table_name + constants.BANG + "array",
            layer_name,
            base,
            count=count,
            subtype=int32,
        )

        min_length = self.config.get("min-length")
        for length in lengths:

            if not min_length <= length <= 64:
                continue

            offset = length.vol["offset"] + DWORD_SIZE_BYTES
            passphrase: objects.Bytes = self.context.object(
                pe_table_name + constants.BANG + "bytes",
                layer_name,
                offset,
                length=length,
            )


            if not all(0x20 <= c < 0x7F for c in passphrase):
                continue

            buf: objects.Bytes = self.context.object(
                pe_table_name + constants.BANG + "bytes",
                layer_name,
                offset + length + 1,
                length=3,
            )
            if any(buf):
                continue

            yield offset, passphrase.decode(encoding="ascii")

    def _find_driver_bases(
        self, mods: Iterable[interfaces.objects.ObjectInterface]
    ) -> List[int]:
        driver_substr = (self.config.get("driver") or "").lower().strip()

        def matches(mod_name: str, needle: str) -> bool:
            return needle and needle in mod_name

        def bases_for(needle: str) -> List[int]:
            out: List[int] = []
            for mod in mods:
                try:
                    name = mod.BaseDllName.get_string().lower()
                except Exception:
                    continue
                if matches(name, needle):
                    out.append(int(mod.DllBase))
            return out

        if driver_substr:
            bases = bases_for(driver_substr)
            if bases:
                return bases

        for needle in ("veracrypt", "truecrypt"):
            bases = bases_for(needle)
            if bases:
                return bases

        return []

    def _generator(self):
        kernel = self.context.modules[self.config["kernel"]]
        mods: Iterable[interfaces.objects.ObjectInterface] = modules.Modules.list_modules(
            self.context, self.config["kernel"]
        )

        driver_bases = self._find_driver_bases(mods)
        if not driver_bases:
            vollog.warning(
                "No VeraCrypt driver module found in the modules list. Unable to proceed."
            )
            return

        seen = set()
        for module_base in driver_bases:
            try:
                for offset, password in self.scan_module(module_base, kernel.layer_name):
                    key = (offset, password)
                    if key in seen:
                        continue
                    seen.add(key)
                    yield (0, (format_hints.Hex(offset), len(password), password))
            except Exception as exc:
                vollog.debug("Failed scanning module at 0x%x: %s", module_base, exc)

    def run(self) -> renderers.TreeGrid:
        return renderers.TreeGrid(
            [
                ("Offset", format_hints.Hex),
                ("Length", int),
                ("Password", str),
            ],
            self._generator(),
        )