Small clean ups and type hints of mvt-android

mvt/android/cli.py

@@ -71,15 +71,7 @@ def download_apks(ctx, all_apks, virustotal, output, from_file, serial):
                 log.critical("You need to specify an output folder with --output!")
                 ctx.exit(1)
 
-            if not os.path.exists(output):
+            download = DownloadAPKs(results_path=output, all_apks=all_apks)
-                try:
-                    os.makedirs(output)
-                except Exception as e:
-                    log.critical("Unable to create output folder %s: %s", output, e)
-                    ctx.exit(1)
-
-            download = DownloadAPKs(output_folder=output, all_apks=all_apks,
-                                    log=logging.getLogger(DownloadAPKs.__module__))
             if serial:
                 download.serial = serial
             download.run()

mvt/android/cmd_check_backup.py

@@ -9,6 +9,7 @@ import os
 import sys
 import tarfile
 from pathlib import Path
+from typing import Callable
 
 from rich.prompt import Prompt
 
@@ -38,7 +39,7 @@ class CmdAndroidCheckBackup(Command):
         self.backup_archive = None
         self.backup_files = []
 
-    def init(self):
+    def init(self) -> None:
         if os.path.isfile(self.target_path):
             self.backup_type = "ab"
             with open(self.target_path, "rb") as handle:
@@ -77,7 +78,7 @@ class CmdAndroidCheckBackup(Command):
             log.critical("Invalid backup path, path should be a folder or an Android Backup (.ab) file")
             sys.exit(1)
 
-    def module_init(self, module):
+    def module_init(self, module: Callable) -> None:
         if self.backup_type == "folder":
             module.from_folder(self.target_path, self.backup_files)
         else:

mvt/android/cmd_check_bugreport.py

@@ -6,6 +6,7 @@
 import logging
 import os
 from pathlib import Path
+from typing import Callable
 from zipfile import ZipFile
 
 from mvt.common.command import Command
@@ -31,7 +32,7 @@ class CmdAndroidCheckBugreport(Command):
         self.bugreport_archive = None
         self.bugreport_files = []
 
-    def init(self):
+    def init(self) -> None:
         if os.path.isfile(self.target_path):
             self.bugreport_format = "zip"
             self.bugreport_archive = ZipFile(self.target_path)
@@ -44,7 +45,7 @@ class CmdAndroidCheckBugreport(Command):
                 for file_name in subfiles:
                     self.bugreport_files.append(os.path.relpath(os.path.join(root, file_name), parent_path))
 
-    def module_init(self, module):
+    def module_init(self, module: Callable) -> None:
         if self.bugreport_format == "zip":
             module.from_zip(self.bugreport_archive, self.bugreport_files)
         else:

mvt/android/cmd_download_apks.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import os
+from typing import Callable
 
 from tqdm import tqdm
 
@@ -36,23 +37,22 @@ class DownloadAPKs(AndroidExtraction):
 
     """
 
-    def __init__(self, output_folder=None, all_apks=False, log=None,
+    def __init__(self, results_path: str = "", all_apks: bool = False,
-                 packages=None):
+                 packages: list = []):
         """Initialize module.
-        :param output_folder: Path to the folder where data should be stored
+        :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__(log=log)
+        super().__init__(results_path=results_path, log=log)
 
         self.packages = packages
         self.all_apks = all_apks
-        self.output_folder_apk = None
+        self.results_path_apks = None
-        self.output_folder = output_folder
 
     @classmethod
-    def from_json(cls, json_path):
+    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.
@@ -62,7 +62,7 @@ class DownloadAPKs(AndroidExtraction):
             packages = json.load(handle)
             return cls(packages=packages)
 
-    def pull_package_file(self, package_name, remote_path):
+    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
@@ -76,7 +76,7 @@ class DownloadAPKs(AndroidExtraction):
         if "==/" in remote_path:
             file_name = "_" + remote_path.split("==/")[1].replace(".apk", "")
 
-        local_path = os.path.join(self.output_folder_apk,
+        local_path = os.path.join(self.results_path_apks,
                                   f"{package_name}{file_name}.apk")
         name_counter = 0
         while True:
@@ -84,7 +84,7 @@ class DownloadAPKs(AndroidExtraction):
                 break
 
             name_counter += 1
-            local_path = os.path.join(self.output_folder_apk,
+            local_path = os.path.join(self.results_path_apks,
                                       f"{package_name}{file_name}_{name_counter}.apk")
 
         try:
@@ -105,11 +105,9 @@ class DownloadAPKs(AndroidExtraction):
 
         return local_path
 
-    def get_packages(self):
+    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...")
 
@@ -120,12 +118,11 @@ class DownloadAPKs(AndroidExtraction):
 
         self.packages = m.results
 
-    def pull_packages(self):
+    def pull_packages(self) -> None:
-        """Download all files of all selected packages from the device."""
+        """Download all files of all selected packages from the device.
-        log.info("Starting extraction of installed APKs at folder %s", self.output_folder)
+        """
-
+        log.info("Starting extraction of installed APKs at folder %s",
-        if not os.path.exists(self.output_folder):
+                 self.results_path)
-            os.mkdir(self.output_folder)
 
         # If the user provided the flag --all-apks we select all packages.
         packages_selection = []
@@ -139,7 +136,7 @@ class DownloadAPKs(AndroidExtraction):
                 if not package.get("system", False):
                     packages_selection.append(package)
 
-            log.info("Selected only %d packages which are not marked as system",
+            log.info("Selected only %d packages which are not marked as \"system\"",
                      len(packages_selection))
 
         if len(packages_selection) == 0:
@@ -148,9 +145,9 @@ class DownloadAPKs(AndroidExtraction):
 
         log.info("Downloading packages from device. This might take some time ...")
 
-        self.output_folder_apk = os.path.join(self.output_folder, "apks")
+        self.results_path_apks = os.path.join(self.results_path, "apks")
-        if not os.path.exists(self.output_folder_apk):
+        if not os.path.exists(self.results_path_apks):
-            os.mkdir(self.output_folder_apk)
+            os.mkdirs(self.results_path_apks)
 
         counter = 0
         for package in packages_selection:
@@ -172,14 +169,12 @@ class DownloadAPKs(AndroidExtraction):
 
         log.info("Download of selected packages completed")
 
-    def save_json(self):
+    def save_json(self) -> None:
-        """Save the results to the package.json file."""
+        json_path = os.path.join(self.results_path, "apks.json")
-        json_path = os.path.join(self.output_folder, "apks.json")
         with open(json_path, "w", encoding="utf-8") as handle:
             json.dump(self.packages, handle, indent=4)
 
     def run(self) -> None:
-        """Run all steps of fetch-apk."""
         self.get_packages()
         self._adb_connect()
         self.pull_packages()