# 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 csv
import json
import logging
import os
import re
from dataclasses import asdict, is_dataclass
from typing import Any, Dict, Optional, Sequence

from .alerts import AlertStore
from .indicators import Indicators
from .module_types import (
    ModuleAtomicResult,
    ModuleResults,
    ModuleSerializedResult,
    ModuleTimeline,
    URLResult,
)
from .utils import CustomJSONEncoder, exec_or_profile


class DatabaseNotFoundError(Exception):
    pass


class DatabaseCorruptedError(Exception):
    pass


class EncryptedBackupError(Exception):
    pass


class InsufficientPrivileges(Exception):
    pass


class MVTModule:
    """This class provides a base for all extraction modules."""

    enabled: bool = True
    slug: Optional[str] = None
    dependencies: Sequence[type["MVTModule"]] = ()
    supported_commands: Sequence[tuple[str, str]] = ()
    # A custom module can name a module class it supersedes, usually a
    # built-in one. When both are available to a command, the named module is
    # dropped from the run and this module takes its place, including in the
    # dependencies of any other module.
    replaces: Optional[type["MVTModule"]] = None

    def __init__(
        self,
        file_path: Optional[str] = None,
        target_path: Optional[str] = None,
        results_path: Optional[str] = None,
        module_options: Optional[Dict[str, Any]] = None,
        log: logging.Logger = logging.getLogger(__name__),
        results: Optional[ModuleResults] = None,
    ) -> None:
        """Initialize module.

        :param file_path: Path to the module's database file, if there is any
        :type file_path: str
        :param target_path: Path to the target folder (backup or filesystem
                            dump)
        :type target_path: str
        :param results_path: Folder where results will be stored
        :type results_path: str
        :param fast_mode: Flag to enable or disable slow modules
        :type fast_mode: bool
        :param log: Handle to logger
        :param results: Provided list of results entries
        :type results: list
        """
        self.file_path: Optional[str] = file_path
        self.target_path: Optional[str] = target_path
        self.results_path: Optional[str] = results_path
        self.serial: Optional[str] = None
        self.module_options: Dict[str, Any] = module_options if module_options else {}

        self.log = log
        self.indicators: Optional[Indicators] = None
        self.alertstore: AlertStore = AlertStore(log=log)

        self.results: ModuleResults = results if results is not None else []
        self.timeline: ModuleTimeline = []
        self.url_results: list[URLResult] = []
        self.dependency_modules: Dict[type["MVTModule"], "MVTModule"] = {}

    def get_dependency_results(
        self, module_class: type["MVTModule"]
    ) -> ModuleResults:
        """Return the results produced by a prerequisite module."""
        return self.dependency_modules[module_class].results

    @classmethod
    def from_json(cls, json_path: str, log: logging.Logger):
        with open(json_path, "r", encoding="utf-8") as handle:
            results = json.load(handle)
            if log:
                log.info('Loaded %d results from "%s"', len(results), json_path)

            return cls(results=results, log=log)

    @classmethod
    def get_slug(cls) -> str:
        if cls.slug:
            return cls.slug

        sub = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", cls.__name__)
        return re.sub("([a-z0-9])([A-Z])", r"\1_\2", sub).lower()

    def check_indicators(self) -> None:
        raise NotImplementedError

    def collect_url_results(self) -> None:
        """Collect URL records exposed by this module."""

    def add_url_result(self, url: str, timestamp: Optional[str], source: str) -> None:
        expanded_url = None
        if self.indicators:
            expanded_url = self.indicators.get_expanded_url(url)

        self.url_results.append(
            {
                "url": url,
                "expanded_url": expanded_url,
                "timestamp": timestamp,
                "source": source,
            }
        )

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

        name = self.get_slug()

        if self.results:
            converted_results: Any
            if isinstance(self.results, dict):
                converted_results = self.results
            else:
                converted_results = [
                    asdict(result)
                    if is_dataclass(result) and not isinstance(result, type)
                    else result
                    for result in self.results
                ]
            results_file_name = f"{name}.json"
            results_json_path = os.path.join(self.results_path, results_file_name)
            with open(results_json_path, "w", encoding="utf-8") as handle:
                try:
                    json.dump(
                        converted_results, handle, indent=4, cls=CustomJSONEncoder
                    )
                except Exception as exc:
                    self.log.error(
                        "Unable to store results of module %s to file %s: %s",
                        self.__class__.__name__,
                        results_file_name,
                        exc,
                    )

        if self.alertstore.alerts:
            detected_file_name = f"{name}_detected.json"
            detected_json_path = os.path.join(self.results_path, detected_file_name)
            with open(detected_json_path, "w", encoding="utf-8") as handle:
                json.dump(
                    self.alertstore.as_json(), handle, indent=4, cls=CustomJSONEncoder
                )

    def serialize(self, result: ModuleAtomicResult) -> ModuleSerializedResult:
        raise NotImplementedError

    @staticmethod
    def _deduplicate_timeline(timeline: list) -> list:
        """Serialize entry as JSON to deduplicate repeated entries

        :param timeline: List of entries from timeline to deduplicate

        """
        timeline_set = set()
        for record in timeline:
            timeline_set.add(
                json.dumps(
                    asdict(record)
                    if is_dataclass(record) and not isinstance(record, type)
                    else record,
                    sort_keys=True,
                )
            )

        return [json.loads(record) for record in timeline_set]

    def to_timeline(self) -> None:
        """Convert results into a timeline."""
        if not self.results:
            return

        for result in self.results:
            record: ModuleSerializedResult = self.serialize(result)
            if record:
                if isinstance(record, list):
                    self.timeline.extend(record)  # type: ignore[arg-type]
                else:
                    self.timeline.append(record)  # type: ignore[arg-type]

        # De-duplicate timeline entries.
        self.timeline = self._deduplicate_timeline(self.timeline)

    def run(self) -> None:
        """Run the main module procedure."""
        raise NotImplementedError


def run_module(module: MVTModule) -> None:
    module.log.info("Running module %s...", module.__class__.__name__)

    try:
        exec_or_profile("module.run()", globals(), locals())
    except EncryptedBackupError:
        raise
    except NotImplementedError:
        module.log.exception(
            "The run() procedure of module %s was not implemented yet!",
            module.__class__.__name__,
        )
    except InsufficientPrivileges as exc:
        module.log.info(
            "Insufficient privileges for module %s: %s", module.__class__.__name__, exc
        )
    except DatabaseNotFoundError as exc:
        module.log.info(
            "There might be no data to extract by module %s: %s",
            module.__class__.__name__,
            exc,
        )
    except DatabaseCorruptedError as exc:
        module.log.error(
            "The %s module database seems to be corrupted: %s",
            module.__class__.__name__,
            exc,
        )
    except Exception as exc:
        module.log.exception(
            "Error in running extraction from module %s: %s",
            module.__class__.__name__,
            exc,
        )
    else:
        try:
            exec_or_profile("module.check_indicators()", globals(), locals())
        except NotImplementedError:
            module.log.info(
                "The %s module does not support checking for indicators",
                module.__class__.__name__,
            )
        except Exception as exc:
            module.log.exception(
                "Error when checking indicators from module %s: %s",
                module.__class__.__name__,
                exc,
            )

        else:
            if module.indicators and not module.alertstore.alerts:
                module.log.info(
                    "The %s module produced no detections!", module.__class__.__name__
                )

        try:
            module.collect_url_results()
        except Exception as exc:
            module.log.exception(
                "Error when collecting URLs from module %s: %s",
                module.__class__.__name__,
                exc,
            )

        try:
            module.to_timeline()
        except NotImplementedError:
            pass
        except Exception as exc:
            module.log.exception(
                "Error when serializing data from module %s: %s",
                module.__class__.__name__,
                exc,
            )

        module.save_to_json()


def save_timeline(timeline: list, timeline_path: str, is_utc: bool = True) -> None:
    """Save the timeline in a csv file.

    :param timeline: List of records to order and store
    :param timeline_path: Path to the csv file to store the timeline to

    """
    with open(timeline_path, "w", encoding="utf-8") as handle:
        csvoutput = csv.writer(
            handle, delimiter=",", quotechar='"', quoting=csv.QUOTE_ALL, escapechar="\\"
        )

        if is_utc:
            timestamp_header = "UTC Timestamp"
        else:
            timestamp_header = "Device Local Timestamp"
        csvoutput.writerow([timestamp_header, "Plugin", "Event", "Description"])

        for event in sorted(
            timeline, key=lambda x: x["timestamp"] if x["timestamp"] is not None else ""
        ):
            csvoutput.writerow(
                [
                    event.get("timestamp"),
                    event.get("module"),
                    event.get("event"),
                    event.get("data"),
                ]
            )
