From 7860fc1a5c9609e9a6e8ffd986221919bd61bd28 Mon Sep 17 00:00:00 2001
From: Gabriel Claudiu Georgiu <gabriel@claudiugeorgiu.com>
Date: Thu, 21 May 2020 21:44:42 +0200
Subject: [PATCH] Add command line options for the keystore

---
 .idea/Obfuscapk.iml                        |  5 +--
 .idea/webResources.xml                     | 10 +++++
 README.md                                  | 19 +++++++--
 src/obfuscapk/cli.py                       | 44 ++++++++++++++++++++
 src/obfuscapk/main.py                      | 29 ++++++++++---
 src/obfuscapk/obfuscation.py               | 40 +++++++++++++++---
 src/obfuscapk/resources/libs_to_ignore.txt |  1 +
 src/obfuscapk/tool.py                      | 10 ++++-
 src/test/test_cli.py                       | 35 ++++++++++++++++
 src/test/test_obfuscation.py               | 47 +++++++++++++++++++++-
 src/test/test_tool.py                      | 16 ++++++++
 11 files changed, 236 insertions(+), 20 deletions(-)

diff --git a/.idea/Obfuscapk.iml b/.idea/Obfuscapk.iml
index fb8cf84c..5b6fd151 100644
--- a/.idea/Obfuscapk.iml
+++ b/.idea/Obfuscapk.iml
@@ -31,8 +31,7 @@
     <orderEntry type="inheritedJdk" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
-  <component name="PyDocumentationSettings">
-    <option name="format" value="EPYTEXT" />
-    <option name="myDocStringFormat" value="Epytext" />
+  <component name="PackageRequirementsSettings">
+    <option name="requirementsPath" value="" />
   </component>
 </module>
\ No newline at end of file
diff --git a/.idea/webResources.xml b/.idea/webResources.xml
index 0c0a1b30..dae1ef15 100644
--- a/.idea/webResources.xml
+++ b/.idea/webResources.xml
@@ -18,6 +18,16 @@
             <path value="file://$PROJECT_DIR$/src/obfuscapk/resources/smali" />
             <path value="file://$PROJECT_DIR$/src/test/test_resources" />
             <path value="file://$PROJECT_DIR$/src/test/test_resources/v1.0" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/arithmetic_branch/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/asset_encryption/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/call_indirection/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/debug_removal/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/goto/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/lib_encryption/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/method_overload/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/nop/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/random_manifest/example" />
+            <path value="file://$PROJECT_DIR$/src/obfuscapk/obfuscators/reorder/example" />
           </resourceRoots>
         </entryData>
       </entry>
diff --git a/README.md b/README.md
index 28b3305c..b26c79db 100644
--- a/README.md
+++ b/README.md
@@ -248,7 +248,10 @@ Let's start by looking at the help message:
 
 ```Shell
 $ obfuscapk --help
-obfuscapk [-h] -o OBFUSCATOR [-w DIR] [-d OUT_APK] [-i] [-p] [-k VT_API_KEY] <APK_FILE>
+obfuscapk [-h] -o OBFUSCATOR [-w DIR] [-d OUT_APK] [-i] [-p] [-k VT_API_KEY]
+          [--keystore-file KEYSTORE_FILE] [--keystore-password KEYSTORE_PASSWORD]
+          [--key-alias KEY_ALIAS] [--key-password KEY_PASSWORD]
+          <APK_FILE>
 ```
 
 There are two mandatory parameters: `<APK_FILE>`, the path (relative or absolute) to
@@ -283,6 +286,15 @@ key(s) to be used when communicating with Virus Total. Can be set multiple times
 cycle through the API keys during the requests (e.g.,
 `-k VALID_VT_KEY_1 -k VALID_VT_KEY_2`).
 
+* `--keystore-file KEYSTORE_FILE`, `--keystore-password KEYSTORE_PASSWORD`,
+`--key-alias KEY_ALIAS` and `--key-password KEY_PASSWORD` can be used to specify a
+custom keystore (needed for the apk signing). If `--keystore-file` is used,
+`--keystore-password` and `--key-alias` must be specified too, while `--key-password`
+is needed only if the chosen key has a different password from the keystore password.
+By default (when `--keystore-file` is not specified), a
+[keystore bundled with Obfuscapk](https://github.com/ClaudiuGeorgiu/Obfuscapk/blob/master/src/obfuscapk/resources/obfuscation_keystore.jks)
+is used for the signing operations.
+
 Let's consider now a simple working example to see how Obfuscapk works:
 
 ```Shell
@@ -313,8 +325,9 @@ available and ready to be used
     apk file is saved in the working directory created before
 
     - `NewSignature` obfuscator signs the newly created apk file with a custom
-    certificate contained in
-    [this keystore](https://github.com/ClaudiuGeorgiu/Obfuscapk/blob/master/src/obfuscapk/resources/obfuscation_keystore.jks)
+    certificate contained in a
+    [keystore bundled with Obfuscapk](https://github.com/ClaudiuGeorgiu/Obfuscapk/blob/master/src/obfuscapk/resources/obfuscation_keystore.jks)
+    (though a different keystore can be specified with the `--keystore-file` parameter)
 
     - `NewAlignment` obfuscator uses `zipalign` tool to align the resulting apk file
 
diff --git a/src/obfuscapk/cli.py b/src/obfuscapk/cli.py
index 637b83d0..286359ba 100644
--- a/src/obfuscapk/cli.py
+++ b/src/obfuscapk/cli.py
@@ -80,6 +80,34 @@ def get_cmd_args(args: list = None):
         "Can be specified multiple times to use a different API key for each request "
         "(cycling through the API keys)",
     )
+    parser.add_argument(
+        "--keystore-file",
+        type=str,
+        metavar="KEYSTORE_FILE",
+        help="The path to a custom keystore file to be used for signing the obfuscated "
+        ".apk file. By default a keystore bundled with this tool will be used",
+    )
+    parser.add_argument(
+        "--keystore-password",
+        type=str,
+        metavar="KEYSTORE_PASSWORD",
+        help="The password of the custom keystore used for signing the obfuscated .apk "
+        "file (needed only when specifying a custom keystore file)",
+    )
+    parser.add_argument(
+        "--key-alias",
+        type=str,
+        metavar="KEY_ALIAS",
+        help="The key alias for signing the obfuscated .apk file (needed only when "
+        "specifying a custom keystore file)",
+    )
+    parser.add_argument(
+        "--key-password",
+        type=str,
+        metavar="KEY_PASSWORD",
+        help="The key password for signing the obfuscated .apk file (needed only when "
+        "specifying a custom keystore file)",
+    )
     return parser.parse_args(args)
 
 
@@ -118,6 +146,18 @@ def main():
             key.strip(" '\"") for key in arguments.virus_total_key
         ]
 
+    if arguments.keystore_file:
+        arguments.keystore_file = arguments.keystore_file.strip(" '\"")
+
+    if arguments.keystore_password:
+        arguments.keystore_password = arguments.keystore_password.strip(" '\"")
+
+    if arguments.key_alias:
+        arguments.key_alias = arguments.key_alias.strip(" '\"")
+
+    if arguments.key_password:
+        arguments.key_password = arguments.key_password.strip(" '\"")
+
     perform_obfuscation(
         arguments.apk_file,
         arguments.obfuscator,
@@ -126,6 +166,10 @@ def main():
         arguments.interactive,
         arguments.ignore_libs,
         arguments.virus_total_key,
+        arguments.keystore_file,
+        arguments.keystore_password,
+        arguments.key_alias,
+        arguments.key_password,
     )
 
 
diff --git a/src/obfuscapk/main.py b/src/obfuscapk/main.py
index c7916e49..733fda83 100644
--- a/src/obfuscapk/main.py
+++ b/src/obfuscapk/main.py
@@ -50,6 +50,10 @@ def perform_obfuscation(
     interactive: bool = False,
     ignore_libs: bool = False,
     virus_total_api_key: List[str] = None,
+    keystore_file: str = None,
+    keystore_password: str = None,
+    key_alias: str = None,
+    key_password: str = None,
 ):
     """
     Apply the obfuscation techniques to an input application and generate an obfuscated
@@ -69,6 +73,17 @@ def perform_obfuscation(
                         obfuscation operations.
     :param virus_total_api_key: A list containing Virus Total API keys, needed only
                                 when using Virus Total obfuscator.
+    :param keystore_file: The path to a custom keystore file to be used for signing the
+                          resulting obfuscated application. If not provided, a default
+                          keystore bundled with this tool will be used instead.
+    :param keystore_password: The password of the custom keystore used for signing the
+                              resulting obfuscated application (needed only when
+                              specifying a custom keystore file).
+    :param key_alias: The key alias for signing the resulting obfuscated application
+                      (needed only when specifying a custom keystore file).
+    :param key_password: The key password for signing the resulting obfuscated
+                         application (needed only when specifying a custom keystore
+                         file).
     """
 
     check_external_tool_dependencies()
@@ -83,9 +98,13 @@ def perform_obfuscation(
         input_apk_path,
         working_dir_path,
         obfuscated_apk_path,
-        interactive=interactive,
-        ignore_libs=ignore_libs,
-        virus_total_api_key=virus_total_api_key,
+        interactive,
+        ignore_libs,
+        virus_total_api_key,
+        keystore_file,
+        keystore_password,
+        key_alias,
+        key_password,
     )
 
     manager = ObfuscatorManager()
@@ -124,7 +143,5 @@ def perform_obfuscation(
                 )
             (obfuscator_name_to_function[obfuscator_name])(obfuscation)
         except Exception as e:
-            logger.critical(
-                "Error during obfuscation: {0}".format(e), exc_info=True
-            )
+            logger.critical("Error during obfuscation: {0}".format(e), exc_info=True)
             raise
diff --git a/src/obfuscapk/obfuscation.py b/src/obfuscapk/obfuscation.py
index f7df07d2..1ba82a59 100644
--- a/src/obfuscapk/obfuscation.py
+++ b/src/obfuscapk/obfuscation.py
@@ -25,6 +25,10 @@ def __init__(
         ignore_libs: bool = False,
         interactive: bool = False,
         virus_total_api_key: List[str] = None,
+        keystore_file: str = None,
+        keystore_password: str = None,
+        key_alias: str = None,
+        key_password: str = None,
     ):
         self.logger = logging.getLogger(__name__)
 
@@ -34,6 +38,10 @@ def __init__(
         self.ignore_libs: bool = ignore_libs
         self.interactive: bool = interactive
         self.virus_total_api_key: List[str] = virus_total_api_key
+        self.keystore_file: str = keystore_file
+        self.keystore_password: str = keystore_password
+        self.key_alias: str = key_alias
+        self.key_password: str = key_password
 
         # Random string (32 chars long) generation with ASCII letters and digits
         self.encryption_secret = "".join(
@@ -503,14 +511,36 @@ def sign_obfuscated_apk(self) -> None:
         # The obfuscated apk will be signed with jarsigner.
         jarsigner: Jarsigner = Jarsigner()
 
+        # If a custom keystore file is not provided, use the default one bundled with
+        # the tool. Otherwise check that the keystore password and a key alias are
+        # provided along with the custom keystore.
+        if not self.keystore_file:
+            self.keystore_file = os.path.join(
+                os.path.dirname(__file__), "resources", "obfuscation_keystore.jks"
+            )
+            self.keystore_password = "obfuscation_password"
+            self.key_alias = "obfuscation_key"
+        else:
+            if not os.path.isfile(self.keystore_file):
+                self.logger.error(
+                    'Unable to find keystore file "{0}"'.format(self.keystore_file)
+                )
+                raise FileNotFoundError(
+                    'Unable to find keystore file "{0}"'.format(self.keystore_file)
+                )
+            if not self.keystore_password or not self.key_alias:
+                raise ValueError(
+                    "When using a custom keystore file, keystore password and key "
+                    "alias must be provided too"
+                )
+
         try:
             jarsigner.resign(
                 self.obfuscated_apk_path,
-                os.path.join(
-                    os.path.dirname(__file__), "resources", "obfuscation_keystore.jks"
-                ),
-                "obfuscation_password",
-                "obfuscation_key",
+                self.keystore_file,
+                self.keystore_password,
+                self.key_alias,
+                self.key_password,
             )
         except Exception as e:
             self.logger.error("Error during apk signing: {0}".format(e))
diff --git a/src/obfuscapk/resources/libs_to_ignore.txt b/src/obfuscapk/resources/libs_to_ignore.txt
index d3f32111..c1fa9a2d 100644
--- a/src/obfuscapk/resources/libs_to_ignore.txt
+++ b/src/obfuscapk/resources/libs_to_ignore.txt
@@ -3,6 +3,7 @@ android/content/
 android/opengl/
 android/support/
 android/widget/
+androidx/
 anet/channel/
 anetwork/channel/
 anywheresoftware/
diff --git a/src/obfuscapk/tool.py b/src/obfuscapk/tool.py
index 17cda67c..b1e75d67 100644
--- a/src/obfuscapk/tool.py
+++ b/src/obfuscapk/tool.py
@@ -197,6 +197,7 @@ def sign(
         keystore_file_path: str,
         keystore_password: str,
         key_alias: str,
+        key_password: str = None,
     ) -> str:
 
         # Check if the apk file to sign is a valid file.
@@ -220,6 +221,10 @@ def sign(
             key_alias,
         ]
 
+        if key_password:
+            sign_cmd.insert(-2, "-keypass")
+            sign_cmd.insert(-2, key_password)
+
         try:
             self.logger.info('Running sign command "{0}"'.format(" ".join(sign_cmd)))
             output = subprocess.check_output(sign_cmd, stderr=subprocess.STDOUT).strip()
@@ -241,6 +246,7 @@ def resign(
         keystore_file_path: str,
         keystore_password: str,
         key_alias: str,
+        key_password: str = None,
     ) -> str:
 
         # If present, delete the old signature of the apk and then sign it with the
@@ -282,7 +288,9 @@ def resign(
             )
             raise
 
-        return self.sign(apk_path, keystore_file_path, keystore_password, key_alias)
+        return self.sign(
+            apk_path, keystore_file_path, keystore_password, key_alias, key_password
+        )
 
 
 class Zipalign(object):
diff --git a/src/test/test_cli.py b/src/test/test_cli.py
index 052d68c0..2d219c7d 100644
--- a/src/test/test_cli.py
+++ b/src/test/test_cli.py
@@ -60,6 +60,41 @@ def test_valid_basic_command_with_quotes(
 
         assert os.path.isfile(obfuscated_apk_path)
 
+    def test_valid_basic_command_with_custom_keystore(
+        self,
+        tmp_working_directory_path: str,
+        tmp_demo_apk_v10_original_path: str,
+        monkeypatch,
+    ):
+        obfuscated_apk_path = os.path.join(tmp_working_directory_path, "obfuscated.apk")
+
+        # Mock the command line parser.
+        arguments = cli.get_cmd_args(
+            "-w {working_dir} -d {destination} "
+            "-o Rebuild -o NewSignature -o NewAlignment "
+            "--keystore-file {keystore_file} --keystore-password {keystore_password} "
+            "--key-alias {key_alias} --key-password {key_password} {apk_file}".format(
+                working_dir=tmp_working_directory_path,
+                destination=obfuscated_apk_path,
+                apk_file=tmp_demo_apk_v10_original_path,
+                keystore_file=os.path.join(
+                    os.path.dirname(__file__),
+                    os.path.pardir,
+                    "obfuscapk",
+                    "resources",
+                    "obfuscation_keystore.jks",
+                ),
+                keystore_password="obfuscation_password",
+                key_alias="obfuscation_key",
+                key_password="obfuscation_password",
+            ).split()
+        )
+        monkeypatch.setattr(cli, "get_cmd_args", lambda: arguments)
+
+        cli.main()
+
+        assert os.path.isfile(obfuscated_apk_path)
+
     def test_missing_required_parameters(self, monkeypatch):
         # Mock the command line parser.
         original = cli.get_cmd_args
diff --git a/src/test/test_obfuscation.py b/src/test/test_obfuscation.py
index 233c27da..189a859f 100644
--- a/src/test/test_obfuscation.py
+++ b/src/test/test_obfuscation.py
@@ -220,10 +220,10 @@ def test_obfuscation_sign_obfuscated_apk_success(
         )
         obfuscation.obfuscated_apk_path = tmp_demo_apk_v10_rebuild_path
 
-        # In case of errors an exception would be thrown.
+        # In case of errors an exception would be thrown and the test would fail.
         obfuscation.sign_obfuscated_apk()
 
-    def test_obfuscation_sign_obfuscated_apk_error(
+    def test_obfuscation_sign_obfuscated_apk_error_invalid_apk(
         self, tmp_demo_apk_v10_original_path: str
     ):
         obfuscation = Obfuscation(tmp_demo_apk_v10_original_path)
@@ -232,6 +232,49 @@ def test_obfuscation_sign_obfuscated_apk_error(
         with pytest.raises(Exception):
             obfuscation.sign_obfuscated_apk()
 
+    def test_obfuscation_sign_obfuscated_apk_error_missing_keystore_file(
+        self,
+        tmp_working_directory_path: str,
+        tmp_demo_apk_v10_original_path: str,
+        tmp_demo_apk_v10_rebuild_path: str,
+    ):
+        obfuscated_apk_path = os.path.join(tmp_working_directory_path, "obfuscated.apk")
+        obfuscation = Obfuscation(
+            tmp_demo_apk_v10_original_path,
+            tmp_working_directory_path,
+            obfuscated_apk_path,
+            keystore_file="invalid.keystore.path",
+        )
+        obfuscation.obfuscated_apk_path = tmp_demo_apk_v10_rebuild_path
+
+        with pytest.raises(FileNotFoundError):
+            obfuscation.sign_obfuscated_apk()
+
+    def test_obfuscation_sign_obfuscated_apk_error_missing_keystore_password(
+        self,
+        tmp_working_directory_path: str,
+        tmp_demo_apk_v10_original_path: str,
+        tmp_demo_apk_v10_rebuild_path: str,
+    ):
+        obfuscated_apk_path = os.path.join(tmp_working_directory_path, "obfuscated.apk")
+        obfuscation = Obfuscation(
+            tmp_demo_apk_v10_original_path,
+            tmp_working_directory_path,
+            obfuscated_apk_path,
+            keystore_file=os.path.join(
+                os.path.dirname(__file__),
+                os.path.pardir,
+                "obfuscapk",
+                "resources",
+                "obfuscation_keystore.jks",
+            ),
+            key_password=None,
+        )
+        obfuscation.obfuscated_apk_path = tmp_demo_apk_v10_rebuild_path
+
+        with pytest.raises(ValueError):
+            obfuscation.sign_obfuscated_apk()
+
     def test_obfuscation_align_obfuscated_apk_success(
         self,
         tmp_working_directory_path: str,
diff --git a/src/test/test_tool.py b/src/test/test_tool.py
index e2673d90..8c14d67b 100644
--- a/src/test/test_tool.py
+++ b/src/test/test_tool.py
@@ -177,6 +177,22 @@ def test_sign_error_invalid_file(self, tmp_working_directory_path: str):
         with pytest.raises(subprocess.CalledProcessError):
             Jarsigner().sign(invalid_file_path, "ignore", "ignore", "ignore")
 
+    def test_sign_error_invalid_key_password(self, tmp_demo_apk_v10_rebuild_path: str):
+        with pytest.raises(subprocess.CalledProcessError):
+            Jarsigner().sign(
+                tmp_demo_apk_v10_rebuild_path,
+                os.path.join(
+                    os.path.dirname(__file__),
+                    os.path.pardir,
+                    "obfuscapk",
+                    "resources",
+                    "obfuscation_keystore.jks",
+                ),
+                "obfuscation_password",
+                "obfuscation_key",
+                "invalid.key.password",
+            )
+
 
 class TestZipalign(object):
     def test_zipalign_valid_path(self):