# Mobile Verification Toolkit (MVT)
# Copyright (c) 2021-2023 Claudio Guarnieri.
# 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
from typing import Callable, Optional

from rich.progress import track

from mvt.common.module import InsufficientPrivileges

from .modules.adb.base import AndroidExtraction
from .modules.adb.packages import Packages

log = logging.getLogger(__name__)


class DownloadAPKs(AndroidExtraction):
    """DownloadAPKs is the main class operating the download of APKs
    from the device.


    """

    def __init__(
        self,
        results_path: Optional[str] = None,
        all_apks: Optional[bool] = False,
        packages: Optional[list] = None
    ) -> None:
        """Initialize module.
        :param results_path: Path to the folder where data should be stored
        :param all_apks: Boolean indicating whether to download all packages
                         or filter known-goods
        :param packages: Provided list of packages, typically for JSON checks
        """
        super().__init__(results_path=results_path, log=log)

        self.packages = packages
        self.all_apks = all_apks
        self.results_path_apks = None

    @classmethod
    def from_json(cls, json_path: str) -> Callable:
        """Initialize this class from an existing apks.json file.

        :param json_path: Path to the apks.json file to parse.

        """
        with open(json_path, "r", encoding="utf-8") as handle:
            packages = json.load(handle)
            return cls(packages=packages)

    def pull_package_file(self, package_name: str, remote_path: str) -> None:
        """Pull files related to specific package from the device.

        :param package_name: Name of the package to download
        :param remote_path: Path to the file to download
        :returns: Path to the local copy

        """
        log.info("Downloading %s ...", remote_path)

        file_name = ""
        if "==/" in remote_path:
            file_name = "_" + remote_path.split("==/")[1].replace(".apk", "")

        local_path = os.path.join(self.results_path_apks,
                                  f"{package_name}{file_name}.apk")
        name_counter = 0
        while True:
            if not os.path.exists(local_path):
                break

            name_counter += 1
            local_path = os.path.join(self.results_path_apks,
                                      f"{package_name}{file_name}_{name_counter}.apk")

        try:
            self._adb_download(remote_path, local_path)
        except InsufficientPrivileges:
            log.error("Unable to pull package file from %s: insufficient privileges, "
                      "it might be a system app", remote_path)
            self._adb_reconnect()
            return None
        except Exception as exc:
            log.exception("Failed to pull package file from %s: %s",
                          remote_path, exc)
            self._adb_reconnect()
            return None

        return local_path

    def get_packages(self) -> None:
        """Use the Packages adb module to retrieve the list of packages.
        We reuse the same extraction logic to then download the APKs.
        """
        self.log.info("Retrieving list of installed packages...")

        m = Packages()
        m.log = self.log
        m.serial = self.serial
        m.run()

        self.packages = m.results

    def pull_packages(self) -> None:
        """Download all files of all selected packages from the device.
        """
        log.info("Starting extraction of installed APKs at folder %s",
                 self.results_path)

        # If the user provided the flag --all-apks we select all packages.
        packages_selection = []
        if self.all_apks:
            log.info("Selected all %d available packages", len(self.packages))
            packages_selection = self.packages
        else:
            # Otherwise we loop through the packages and get only those that
            # are not marked as system.
            for package in self.packages:
                if not package.get("system", False):
                    packages_selection.append(package)

            log.info("Selected only %d packages which are not marked as \"system\"",
                     len(packages_selection))

        if len(packages_selection) == 0:
            log.info("No packages were selected for download")
            return

        log.info("Downloading packages from device. This might take some time ...")

        self.results_path_apks = os.path.join(self.results_path, "apks")
        if not os.path.exists(self.results_path_apks):
            os.makedirs(self.results_path_apks, exist_ok=True)

        for i in track(range(len(packages_selection)),
                       description=f"Downloading {len(packages_selection)} packages..."):
            package = packages_selection[i]

            log.info("[%d/%d] Package: %s", i, len(packages_selection),
                     package["package_name"])

            # Sometimes the package path contains multiple lines for multiple
            # apks. We loop through each line and download each file.
            for package_file in package["files"]:
                device_path = package_file["path"]
                local_path = self.pull_package_file(package["package_name"],
                                                    device_path)
                if not local_path:
                    continue

                package_file["local_path"] = local_path

        log.info("Download of selected packages completed")

    def save_json(self) -> None:
        json_path = os.path.join(self.results_path, "apks.json")
        with open(json_path, "w", encoding="utf-8") as handle:
            json.dump(self.packages, handle, indent=4)

    def run(self) -> None:
        self.get_packages()
        self._adb_connect()
        self.pull_packages()
        self.save_json()
        self._adb_disconnect()
