# Mobile Verification Toolkit (MVT)
# Copyright (c) 2021-2023 The MVT Authors.
# Use of this software is governed by the MVT License 1.1 that can be found at
#   https://license.mvt.re/1.1/

import json
import logging
import os
import sys
from datetime import datetime
from heapq import heappop, heappush
from typing import Any, Optional

from rich.console import Console
from rich.panel import Panel
from rich.text import Text

from .alerts import AlertLevel, AlertStore
from .config import settings
from .indicators import Indicators
from .module import EncryptedBackupError, MVTModule, run_module, save_timeline
from .module_loader import (
    ModuleOrigin,
    get_module_logger,
    get_module_origin,
    module_supports_command,
)
from .module_types import ModuleTimeline, URLResult
from .utils import (
    CustomJSONEncoder,
    convert_datetime_to_iso,
    generate_hashes_from_path,
    get_sha256_from_file_path,
)
from .version import MVT_VERSION


class Command:
    def __init__(
        self,
        target_path: Optional[str] = None,
        results_path: Optional[str] = None,
        ioc_files: Optional[list] = None,
        iocs: Optional[Indicators] = None,
        module_name: Optional[str] = None,
        serial: Optional[str] = None,
        module_options: Optional[dict] = None,
        hashes: Optional[bool] = False,
        sub_command: Optional[bool] = False,
        log: logging.Logger = logging.getLogger(__name__),
        disable_version_check: bool = False,
        disable_indicator_check: bool = False,
        custom_modules: Optional[list[type[MVTModule]]] = None,
    ) -> None:
        self.name = ""
        self.platform = ""
        self.modules: list[type[MVTModule]] = []
        self.custom_modules = custom_modules if custom_modules else []

        self.target_path = target_path
        self.results_path = results_path
        self.ioc_files = ioc_files if ioc_files else []
        self.module_name = module_name
        self.serial = serial
        self.log = log
        self.sub_command = sub_command
        self.disable_version_check = disable_version_check
        self.disable_indicator_check = disable_indicator_check

        # This dictionary can contain options that will be passed down from
        # the Command to all modules. This can for example be used to pass
        # down a password to decrypt a backup or flags which are need by some modules.
        self.module_options = module_options if module_options else {}

        # This list will contain all executed modules.
        # We can use this to reference e.g. self.executed[0].results.
        self.executed: list[MVTModule] = []
        self.hashes = hashes
        self.hash_values: list[dict[str, Any]] = []
        self.timeline: ModuleTimeline = []
        self.url_results: list[URLResult] = []

        # Load IOCs
        self._create_storage()
        self._setup_logging()

        if iocs is not None:
            self.iocs = iocs
        else:
            self.iocs = Indicators(self.log)
            self.iocs.load_indicators_files(self.ioc_files)

        self.alertstore = AlertStore()

    def _create_storage(self) -> None:
        if self.results_path and not os.path.exists(self.results_path):
            try:
                os.makedirs(self.results_path)
            except Exception as exc:
                self.log.fatal(
                    "Unable to create output folder %s: %s", self.results_path, exc
                )
                sys.exit(1)

    def _setup_logging(self):
        if not self.results_path:
            return

        logger = logging.getLogger("mvt")
        file_handler = logging.FileHandler(
            os.path.join(self.results_path, "command.log")
        )
        formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
        )
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(formatter)

        # MVT can be run in a loop.
        # Old file handlers stick around in subsequent loops.
        # Remove any existing logging.FileHandler instances.
        for handler in logger.handlers:
            if isinstance(handler, logging.FileHandler):
                logger.removeHandler(handler)

        # And finally add the new one.
        logger.addHandler(file_handler)

    def _store_timeline(self) -> None:
        if not self.results_path:
            return

        # We use local timestamps in the timeline on Android as many
        # logs do not contain timezone information.
        if type(self).__name__.startswith("CmdAndroid"):
            is_utc = False
        else:
            is_utc = True

        if len(self.timeline) > 0:
            save_timeline(
                self.timeline,
                os.path.join(self.results_path, "timeline.csv"),
                is_utc=is_utc,
            )

    def _store_alerts(self) -> None:
        if not self.results_path:
            return

        alerts = self.alertstore.as_json()
        if not alerts:
            return

        alerts_path = os.path.join(self.results_path, "alerts.json")
        with open(alerts_path, "w+", encoding="utf-8") as handle:
            json.dump(alerts, handle, indent=4, cls=CustomJSONEncoder)

    def _store_urls(self) -> None:
        if not self.results_path or not self.url_results:
            return

        urls_path = os.path.join(self.results_path, "urls.json")
        with open(urls_path, "w", encoding="utf-8") as handle:
            json.dump(self.url_results, handle, indent=4, cls=CustomJSONEncoder)

    def _store_alerts_timeline(self) -> None:
        if not self.results_path:
            return

        alerts_timeline_path = os.path.join(self.results_path, "alerts_timeline.csv")
        self.alertstore.save_timeline(alerts_timeline_path)

    def _store_info(self) -> None:
        if not self.results_path:
            return

        target_path: Optional[str] = None
        if self.target_path:
            target_path = os.path.abspath(self.target_path)

        info: dict[str, Any] = {
            "target_path": target_path,
            "mvt_version": MVT_VERSION,
            "date": convert_datetime_to_iso(datetime.now()),
            "ioc_files": [],
            "hashes": [],
        }

        for coll in self.iocs.ioc_collections:
            ioc_file_path = coll.get("stix2_file_path", "")
            if ioc_file_path and ioc_file_path not in info["ioc_files"]:
                info["ioc_files"].append(ioc_file_path)

        if self.target_path and (settings.HASH_FILES or self.hashes):
            self.generate_hashes()

        info["hashes"] = self.hash_values

        info_path = os.path.join(self.results_path, "info.json")
        with open(info_path, "w+", encoding="utf-8") as handle:
            json.dump(info, handle, indent=4)

        if self.target_path and (settings.HASH_FILES or self.hashes):
            info_hash = get_sha256_from_file_path(info_path)
            self.log.info('Reference hash of the info.json file: "%s"', info_hash)

    def generate_hashes(self) -> None:
        """
        Compute hashes for files in the target_path
        """
        if not self.target_path:
            return

        for file in generate_hashes_from_path(self.target_path, self.log):
            self.hash_values.append(file)

    @staticmethod
    def _modules_by_origin(
        modules: list[type[MVTModule]],
    ) -> dict[ModuleOrigin, list[str]]:
        origins: dict[ModuleOrigin, list[str]] = {}
        for module in modules:
            origins.setdefault(get_module_origin(module), []).append(module.__name__)
        return origins

    def list_modules(self) -> None:
        self.log.info("Following is the list of available %s modules:", self.name)
        for origin, module_names in self._modules_by_origin(
            self._available_modules()
        ).items():
            self.log.info(
                " - Modules from %s: %s", origin.label, ", ".join(module_names)
            )

    def _log_loaded_modules(self, modules: list[type[MVTModule]]) -> None:
        """Record the loaded modules and their origins for auditability."""
        for origin, module_names in self._modules_by_origin(modules).items():
            self.log.info(
                "Loaded %d %s modules from %s: %s",
                len(module_names),
                self.name,
                origin.label,
                ", ".join(module_names),
            )

    def _available_modules(self) -> list[type[MVTModule]]:
        modules = list(self.modules)
        modules.extend(
            module
            for module in self.custom_modules
            if module_supports_command(module, self.platform, self.name)
        )

        deduplicated = []
        for module in modules:
            if module not in deduplicated:
                deduplicated.append(module)

        return deduplicated

    def init(self) -> None:
        raise NotImplementedError

    def module_init(self, module: MVTModule) -> None:
        raise NotImplementedError

    def finish(self) -> None:
        raise NotImplementedError

    def show_alerts_brief(self) -> None:
        console = Console()

        message = Text()
        for i, level in enumerate(AlertLevel):
            message.append(
                f"MVT produced {self.alertstore.count(level)} {level.name} alerts."
            )
            if i < len(AlertLevel) - 1:
                message.append("\n")

        panel = Panel(
            message, title="ALERTS", style="sandy_brown", border_style="sandy_brown"
        )
        console.print("")
        console.print(panel)

    def show_disable_adb_warning(self) -> None:
        console = Console()
        message = Text(
            "Please disable Developer Options and ADB (Android Debug Bridge) on the device once finished with the acquisition. "
            "ADB is a powerful tool which can allow unauthorized access to the device."
        )
        panel = Panel(message, title="NOTE", style="yellow", border_style="yellow")
        console.print("")
        console.print(panel)

    def show_support_message(self) -> None:
        console = Console()
        message = Text()

        support_message = "Please seek reputable expert help if you have serious concerns about a possible spyware attack. Such support is available to human rights defenders and civil society through Amnesty International's Security Lab at https://securitylab.amnesty.org/get-help/?c=mvt"
        if (
            self.alertstore.count(AlertLevel.HIGH) > 0
            or self.alertstore.count(AlertLevel.CRITICAL) > 0
        ):
            message.append(
                f"MVT produced HIGH or CRITICAL alerts. Only expert review can confirm if the detected indicators are signs of an attack.\n\n{support_message}",
            )
            panel = Panel(message, title="WARNING", style="red", border_style="red")
        else:
            message.append(
                f"The lack of severe alerts does not equate to a clean bill of health.\n\n{support_message}",
            )
            panel = Panel(message, title="NOTE", style="yellow", border_style="yellow")

        console.print("")
        console.print(panel)

    def _skipped_modules(
        self,
        required: list[type[MVTModule]],
        module_indexes: dict[type[MVTModule], int],
    ) -> dict[type[MVTModule], tuple[type[MVTModule], type[MVTModule]]]:
        """Return the modules to drop because a dependency is unavailable.

        A module declaring a dependency this command cannot provide is unable
        to run, and so is every module depending on it. Dropping only those
        keeps a single wrong declaration - in a module scoped to several
        commands, for example - from silencing an entire analysis.

        The returned mapping gives, for each skipped module, the module which
        is missing a dependency and the dependency it is missing.
        """
        skipped: dict[type[MVTModule], tuple[type[MVTModule], type[MVTModule]]] = {}
        # Skipping the one module a run was asked for leaves nothing to run,
        # which the caller reports instead.
        remainder = (
            "" if self.module_name else " The rest of the analysis will still run."
        )

        # Skipping one module can skip the modules depending on it, which the
        # pass over the module list may already have gone past, so repeat the
        # pass until nothing changes.
        changed = True
        while changed:
            changed = False
            for module in required:
                if module in skipped:
                    continue

                for dependency in module.dependencies:
                    if dependency not in module_indexes:
                        skipped[module] = (module, dependency)
                        changed = True
                        self.log.warning(
                            "Module %s will be SKIPPED: it depends on module "
                            "%s, which is not available in this command.%s",
                            module.__name__,
                            dependency.__name__,
                            remainder,
                        )
                        break

                    if dependency in skipped:
                        root, missing = skipped[dependency]
                        skipped[module] = (root, missing)
                        changed = True
                        if dependency is root:
                            self.log.warning(
                                "Module %s will be SKIPPED: it depends on "
                                "module %s, itself skipped for depending on "
                                "unavailable module %s.%s",
                                module.__name__,
                                dependency.__name__,
                                missing.__name__,
                                remainder,
                            )
                        else:
                            self.log.warning(
                                "Module %s will be SKIPPED: it depends on "
                                "skipped module %s, in a chain starting at "
                                "module %s, which depends on unavailable "
                                "module %s.%s",
                                module.__name__,
                                dependency.__name__,
                                root.__name__,
                                missing.__name__,
                                remainder,
                            )
                        break

        return skipped

    def _ordered_modules(self) -> Optional[list[type[MVTModule]]]:
        """Return enabled modules in stable topological order."""
        modules = self._available_modules()
        module_indexes = {module: index for index, module in enumerate(modules)}

        if self.module_name:
            selected = [
                module for module in modules if module.__name__ == self.module_name
            ]
        else:
            selected = [module for module in modules if module.enabled]

        required: set[type[MVTModule]] = set()
        pending = list(selected)
        while pending:
            module = pending.pop()
            if module in required:
                continue
            required.add(module)
            for dependency in module.dependencies:
                # Unavailable dependencies are reported by _skipped_modules().
                if dependency in module_indexes:
                    pending.append(dependency)

        ordered_required = sorted(required, key=lambda module: module_indexes[module])
        skipped = self._skipped_modules(ordered_required, module_indexes)
        runnable = [module for module in ordered_required if module not in skipped]
        if skipped and not runnable:
            self.log.warning(
                "Every selected module was skipped for an unavailable "
                "dependency. No modules will be run."
            )

        dependents: dict[type[MVTModule], list[type[MVTModule]]] = {
            module: [] for module in runnable
        }
        indegree = {module: 0 for module in runnable}
        for module in runnable:
            for dependency in module.dependencies:
                if dependency not in indegree:
                    continue
                dependents[dependency].append(module)
                indegree[module] += 1

        ready: list[tuple[int, type[MVTModule]]] = []
        for module, count in indegree.items():
            if count == 0:
                heappush(ready, (module_indexes[module], module))

        ordered = []
        while ready:
            _, module = heappop(ready)
            ordered.append(module)
            for dependent in dependents[module]:
                indegree[dependent] -= 1
                if indegree[dependent] == 0:
                    heappush(ready, (module_indexes[dependent], dependent))

        if len(ordered) != len(runnable):
            cyclic_modules = sorted(
                (module.__name__ for module, count in indegree.items() if count > 0)
            )
            self.log.warning(
                "Circular module dependency detected involving: %s. "
                "No modules will be run.",
                ", ".join(cyclic_modules),
            )
            return None

        return ordered

    def run(self) -> None:
        ordered_modules = self._ordered_modules()
        if ordered_modules is None:
            return

        self._log_loaded_modules(ordered_modules)

        try:
            self.init()
        except NotImplementedError:
            pass

        executed_by_type: dict[type[MVTModule], MVTModule] = {}
        for module in ordered_modules:

            module_logger = get_module_logger(module)

            m = module(
                target_path=self.target_path,
                results_path=self.results_path,
                module_options=self.module_options,
                log=module_logger,
            )
            m.dependency_modules = {
                dependency: executed_by_type[dependency]
                for dependency in module.dependencies
            }

            if self.iocs.total_ioc_count:
                m.indicators = self.iocs
                m.indicators.log = m.log

            if self.serial:
                m.serial = self.serial

            try:
                self.module_init(m)
            except NotImplementedError:
                pass

            try:
                run_module(m)
            except EncryptedBackupError:
                self.log.critical(
                    "The backup appears to be encrypted. "
                    "Please decrypt it first using `mvt-ios decrypt-backup`."
                )
                return

            self.executed.append(m)
            executed_by_type[module] = m
            self.timeline.extend(m.timeline)
            self.url_results.extend(m.url_results)
            self.alertstore.extend(m.alertstore.alerts)

        try:
            self.finish()
        except NotImplementedError:
            pass

        # We only store the timeline from the parent/main command
        if self.sub_command:
            return

        self._store_timeline()
        self._store_alerts_timeline()
        self._store_alerts()
        self._store_urls()
        self._store_info()
