#!/usr/bin/env python3

import logging
import os
import re

import xml.etree.cElementTree as Xml
from typing import List, Set, Dict, Union
from xml.etree.cElementTree import Element

from binascii import hexlify
from typing import List, Set

from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Util.Padding import pad

from obfuscator import obfuscator_category
from obfuscator import util
from obfuscator.obfuscation import Obfuscation

class ClassEncryption(obfuscator_category.IEncryptionObfuscator):
    def __init__(self):
        self.logger = logging.getLogger(
            "{0}.{1}".format(__name__, self.__class__.__name__)
        )
        super().__init__()

        self.encryption_secret = "This-key-need-to-be-32-character"

        self.subclass_name_pattern = re.compile(
            r'\s+name\s=\s"(?P<subclass_name>\S+?)"', re.UNICODE
        )

        self.string_pattern = re.compile(r'"(?P<string_value>\S+?)"', re.UNICODE)

        self.split_class_pattern = re.compile(r"[/$]")

        self.package_name: Union[str, None] = None
        self.encrypted_package_name: Union[str, None] = None
        self.ignore_package_names = []

        # Will be populated before running the class rename obfuscator.
        self.class_name_to_smali_file: dict = {}

    def encrypt_identifier(self, identifier: str) -> str:
        identifier_md5 = util.get_string_md5(identifier)
        return "p{0}".format(identifier_md5.lower()[:8])

    def slash_to_dot_notation_for_classes(
        self, rename_transformations: Dict[str, str]
    ) -> Dict[str, str]:
        dot_rename_transformations: Dict[str, str] = {}

        # Remove leading L and trailing ; from class names and replace / and $ with .
        for old_name, new_name in rename_transformations.items():
            dot_rename_transformations[
                old_name[1:-1].replace("/", ".").replace("$", ".")
            ] = (new_name[1:-1].replace("/", ".").replace("$", "."))

        return dot_rename_transformations

    def transform_package_name(self, manifest_xml_root: Element):
        self.encrypted_package_name = ".".join(
            [self.encrypt_identifier(token) for token in self.package_name.split(".")]
        )

        # Rename package name in manifest file.
        manifest_xml_root.set("package", self.encrypted_package_name)
        manifest_xml_root.set(
            "{http://schemas.android.com/apk/res/android}sharedUserId",
            "{0}.uid.shared".format(util.get_random_string(16)),
        )

    def rename_class_declarations(
        self, smali_files: List[str], interactive: bool = False
    ) -> dict:
        renamed_classes = {}

        # Search for class declarations that can be renamed.
        for smali_file in util.show_list_progress(
            smali_files,
            interactive=interactive,
            description="Renaming class declarations",
        ):
            annotation_flag = False
            with util.inplace_edit_file(smali_file) as (in_file, out_file):

                skip_remaining_lines = False
                class_name = None
                r_class = False
                for line in in_file:

                    if skip_remaining_lines:
                        out_file.write(line)
                        continue

                    if not class_name:
                        class_match = util.class_pattern.match(line)
                        if class_match:
                            class_name = class_match.group("class_name")

                            ignore_class = class_name.startswith(
                                tuple(self.ignore_package_names)
                            )

                            # Split class name to its components and encrypt them.
                            class_tokens = self.split_class_pattern.split(
                                class_name[1:-1]
                            )

                            encrypted_class_name = "L"
                            separator_index = 1
                            for token in class_tokens:
                                separator_index += len(token)
                                if token == "R":
                                    r_class = True
                                if token.isdigit():
                                    encrypted_class_name += (
                                        token + class_name[separator_index]
                                    )
                                elif not r_class and not ignore_class:
                                    encrypted_class_name += (
                                        self.encrypt_identifier(token)
                                        + class_name[separator_index]
                                    )
                                else:
                                    encrypted_class_name += (
                                        token + class_name[separator_index]
                                    )
                                separator_index += 1

                            out_file.write(
                                line.replace(class_name, encrypted_class_name)
                            )

                            renamed_classes[class_name] = encrypted_class_name
                            continue

                    if (
                        line.strip()
                        == ".annotation system Ldalvik/annotation/InnerClass;"
                    ):
                        annotation_flag = True
                        out_file.write(line)
                        continue

                    if annotation_flag and 'name = "' in line:
                        # Subclasses have to be renamed as well.
                        subclass_match = self.subclass_name_pattern.match(line)
                        if subclass_match and not r_class:
                            subclass_name = subclass_match.group("subclass_name")
                            out_file.write(
                                line.replace(
                                    subclass_name,
                                    self.encrypt_identifier(subclass_name),
                                )
                            )
                        else:
                            out_file.write(line)
                        continue

                    if line.strip() == ".end annotation":
                        annotation_flag = False
                        out_file.write(line)
                        continue

                    # Method declaration reached, no more class definitions in
                    # this file.
                    if line.startswith(".method "):
                        skip_remaining_lines = True
                        out_file.write(line)
                    else:
                        out_file.write(line)

        return renamed_classes

    def rename_class_usages_in_smali(
        self,
        smali_files: List[str],
        rename_transformations: dict,
        interactive: bool = False,
    ):
        dot_rename_transformations = self.slash_to_dot_notation_for_classes(
            rename_transformations
        )

        # Add package name.
        dot_rename_transformations[self.package_name] = self.encrypted_package_name

        for smali_file in util.show_list_progress(
            smali_files,
            interactive=interactive,
            description="Renaming class usages in smali files",
        ):
            with util.inplace_edit_file(smali_file) as (in_file, out_file):
                for line in in_file:
                    # Rename classes used as strings with . instead of /.
                    string_match = self.string_pattern.search(line)
                    if (
                        string_match
                        and string_match.group("string_value")
                        in dot_rename_transformations
                    ):
                        line = line.replace(
                            string_match.group("string_value"),
                            dot_rename_transformations[
                                string_match.group("string_value")
                            ],
                        )

                    # Sometimes classes are used in annotations as strings
                    # without trailing ;
                    if (
                        string_match
                        and "{0};".format(string_match.group("string_value"))
                        in rename_transformations
                    ):
                        line = line.replace(
                            string_match.group("string_value"),
                            rename_transformations[
                                "{0};".format(string_match.group("string_value"))
                            ][:-1],
                        )

                    # Rename classes used with the "classic" syntax
                    # (leading L and trailing ;).
                    class_names = util.class_name_pattern.findall(line)
                    for class_name in class_names:
                        if class_name in rename_transformations:
                            line = line.replace(
                                class_name, rename_transformations[class_name]
                            )

                    out_file.write(line)

    def rename_class_usages_in_xml(
        self,
        xml_files: List[str],
        rename_transformations: dict,
        interactive: bool = False,
    ):
        dot_rename_transformations = self.slash_to_dot_notation_for_classes(
            rename_transformations
        )

        # Add package name.
        dot_rename_transformations[self.package_name] = self.encrypted_package_name

        for xml_file in util.show_list_progress(
            xml_files,
            interactive=interactive,
            description="Renaming class usages in xml files",
        ):
            with open(xml_file, "r", encoding="utf-8") as current_file:
                file_content = current_file.read()

            # Replace strings from longest to shortest (to avoid replacing
            # partial strings).
            for old_name in sorted(dot_rename_transformations, reverse=True, key=len):
                file_content = file_content.replace(
                    old_name, dot_rename_transformations[old_name]
                )

                # Activity without package name (".ActivityName")
                if (
                    '"{0}"'.format(old_name.replace(self.package_name, ""))
                    in file_content
                ):
                    file_content = file_content.replace(
                        '"{0}"'.format(old_name.replace(self.package_name, "")),
                        '"{0}"'.format(
                            dot_rename_transformations[old_name].replace(
                                self.encrypted_package_name, ""
                            )
                        ),
                    )

            with open(xml_file, "w", encoding="utf-8") as current_file:
                current_file.write(file_content)


    def encrypt_string(self, string_to_encrypt: str) -> str:
            # This is needed to remove the escaping added by Python. For example, if we
            # find in smali the instruction const-string v0, "\"message\"" Android will
            # treat it as "message" while in Python it's \"message\", so we need to encrypt
            # "message" and not \"message\" (we have to remove the unnecessary escaping,
            # otherwise the backslashes would by encrypted as part of the string).
            string_to_encrypt = string_to_encrypt.encode(errors="replace").decode(
                "unicode_escape"
            )

            key = PBKDF2(
                password=self.encryption_secret,
                salt=self.encryption_secret.encode(),
                dkLen=32,
                count=128,
            )
            encrypted_string = hexlify(
                AES.new(key=key, mode=AES.MODE_ECB).encrypt(
                    pad(string_to_encrypt.encode(errors="replace"), AES.block_size)
                )
            ).decode()
            return encrypted_string

                    

    def obfuscate(self, obfuscation_info: Obfuscation):
        self.logger.info('Running "{0}" obfuscator'.format(self.__class__.__name__))

        try:
            encrypted_strings: Set[str] = set()

            # .field <other_optional_stuff> <string_name>:Ljava/lang/String; =
            # "<string_value>"
            static_string_pattern = re.compile(
                r"\.field.+?static.+?(?P<string_name>\S+?):"
                r'Ljava/lang/String;\s=\s"(?P<string_value>.+)"',
                re.UNICODE,
            )

            for smali_file in util.show_list_progress(
                obfuscation_info.get_smali_files(),
                interactive=obfuscation_info.interactive,
                description="Encrypting constant strings",
            ):
                self.logger.debug(
                    'Encrypting constant strings in file "{0}"'.format(smali_file)
                )

                with open(smali_file, "r", encoding="utf-8") as current_file:
                    lines = current_file.readlines()

                class_name = None

                # Line numbers where a static string is declared.
                static_string_index: List[int] = []

                # Names of the static strings.
                static_string_name: List[str] = []

                # Values of the static strings.
                static_string_value: List[str] = []

                direct_methods_line = -1
                static_constructor_line = -1

                # Line numbers where a constant string is declared.
                string_index: List[int] = []

                # Registers containing the constant strings.
                string_register: List[str] = []

                # Values of the constant strings.
                string_value: List[str] = []

                current_local_count = 0
                for line_number, line in enumerate(lines):

                    if not class_name:
                        class_match = util.class_pattern.match(line)
                        if class_match:
                            class_name = class_match.group("class_name")
                            continue

                    if line.startswith("# direct methods"):
                        direct_methods_line = line_number
                        continue

                    if line.startswith(".method static constructor <clinit>()V"):
                        static_constructor_line = line_number
                        continue

                    static_string_match = static_string_pattern.match(line)
                    if static_string_match and static_string_match.group(
                        "string_value"
                    ):
                        # A static non empty string initialization was found.
                        static_string_index.append(line_number)
                        static_string_name.append(
                            static_string_match.group("string_name")
                        )
                        static_string_value.append(
                            static_string_match.group("string_value")
                        )

                    # We are iterating the lines in order, so each time we enter a
                    # method we'll find the declaration with the number of local
                    # registers available. When we'll encounter a constant string later
                    # in the body of the method, we'll look at its register value and if
                    # it's greater than 15 we won't encrypt it (the invoke instruction
                    # that we need later won't take registers with values greater
                    # than 15).
                    match = util.locals_pattern.match(line)
                    if match:
                        current_local_count = int(match.group("local_count"))
                        continue

                    # If the constant string has a register v0-v15 we can proceed with
                    # the encryption, but if it uses a p<number> register, before
                    # encrypting we have to check if <number> + locals <= 15.
                    string_match = util.const_string_pattern.match(line)
                    if string_match and string_match.group("string"):
                        reg_type = string_match.group("register")[:1]
                        reg_number = int(string_match.group("register")[1:])
                        if (reg_type == "v" and reg_number <= 15) or (
                            reg_type == "p" and reg_number + current_local_count <= 15
                        ):
                            # A non empty string was found in a register <= 15.
                            string_index.append(line_number)
                            string_register.append(string_match.group("register"))
                            string_value.append(string_match.group("string"))

                # Const string encryption.

                for string_number, index in enumerate(string_index):
                    lines[index] = (
                        '\tconst-string/jumbo {register}, "{enc_string}"\n'
                        "\n\tinvoke-static {{{register}}}, "
                        "Lcom/decryptstringmanager/DecryptString"
                        ";->decryptString(Ljava/lang/String;)Ljava/lang/String;\n"
                        "\n\tmove-result-object {register}\n".format(
                            register=string_register[string_number],
                            enc_string=self.encrypt_string(string_value[string_number]),
                        )
                    )

                    encrypted_strings.add(string_value[string_number])

                # Static string encryption.

                static_string_encryption_code = ""
                for string_number, index in enumerate(static_string_index):
                    # Remove the original initialization.
                    lines[index] = "{0}\n".format(lines[index].split(" = ")[0])

                    # Initialize the static string from an encrypted string.
                    static_string_encryption_code += (
                        '\tconst-string/jumbo v0, "{enc_string}"\n'
                        "\n\tinvoke-static {{v0}}, "
                        "Lcom/decryptstringmanager/DecryptString"
                        ";->decryptString(Ljava/lang/String;)Ljava/lang/String;\n"
                        "\n\tmove-result-object v0\n"
                        "\n\tsput-object v0, {class_name}->"
                        "{string_name}:Ljava/lang/String;\n\n".format(
                            enc_string=self.encrypt_string(
                                static_string_value[string_number]
                            ),
                            class_name=class_name,
                            string_name=static_string_name[string_number],
                        )
                    )

                    encrypted_strings.add(static_string_value[string_number])

                if static_constructor_line != -1:
                    # Add static string encryption to the existing static constructor.
                    local_match = util.locals_pattern.match(
                        lines[static_constructor_line + 1]
                    )
                    if local_match:
                        # At least one register is needed.
                        local_count = int(local_match.group("local_count"))
                        if local_count == 0:
                            lines[static_constructor_line + 1] = "\t.locals 1\n"
                        lines[static_constructor_line + 2] = "\n{0}".format(
                            static_string_encryption_code
                        )
                else:
                    # Add a new static constructor for the static string encryption.
                    if direct_methods_line != -1:
                        new_constructor_line = direct_methods_line
                    else:
                        new_constructor_line = len(lines) - 1

                    lines[new_constructor_line] = (
                        "{original}"
                        ".method static constructor <clinit>()V\n"
                        "\t.locals 1\n\n"
                        "{encryption_code}"
                        "\treturn-void\n"
                        ".end method\n\n".format(
                            original=lines[new_constructor_line],
                            encryption_code=static_string_encryption_code,
                        )
                    )

                with open(smali_file, "w", encoding="utf-8") as current_file:
                    current_file.writelines(lines)

            if (
                not obfuscation_info.decrypt_string_smali_file_added_flag
                and encrypted_strings
            ):
                # Add to the app the code for decrypting the encrypted strings. The code
                # for decrypting can be put in any smali directory, since it will be
                # moved to the correct directory when rebuilding the application.
                destination_dir = os.path.dirname(obfuscation_info.get_smali_files()[0])
                destination_file = os.path.join(destination_dir, "DecryptString.smali")
                with open(
                    destination_file, "w", encoding="utf-8"
                ) as decrypt_string_smali:
                    decrypt_string_smali.write(
                        util.get_decrypt_string_smali_code(self.encryption_secret)
                    )
                    obfuscation_info.decrypt_string_smali_file_added_flag = True

        except Exception as e:
            self.logger.error(
                'Error during execution of "{0}" obfuscator: {1}'.format(
                    self.__class__.__name__, e
                )
            )
            raise

        finally:
            obfuscation_info.used_obfuscators.append(self.__class__.__name__)


        try:
            Xml.register_namespace(
                "android", "http://schemas.android.com/apk/res/android"
            )

            xml_parser = Xml.XMLParser(encoding="utf-8")
            manifest_tree = Xml.parse(
                obfuscation_info.get_manifest_file(), parser=xml_parser
            )
            manifest_root = manifest_tree.getroot()

            self.package_name = manifest_root.get("package")
            if not self.package_name:
                raise Exception(
                    "Unable to extract package name from application manifest"
                )

            # Get a mapping between class name and smali file path.
            for smali_file in util.show_list_progress(
                obfuscation_info.get_smali_files(),
                interactive=obfuscation_info.interactive,
                description="Class name to smali file mapping",
            ):
                with open(smali_file, "r", encoding="utf-8") as current_file:
                    class_name = None
                    for line in current_file:
                        if not class_name:
                            # Every smali file contains a class.
                            class_match = util.class_pattern.match(line)
                            if class_match:
                                self.class_name_to_smali_file[
                                    class_match.group("class_name")
                                ] = smali_file
                                break

            self.transform_package_name(manifest_root)

            # Write the changes into the manifest file.
            manifest_tree.write(obfuscation_info.get_manifest_file(), encoding="utf-8")

            xml_files: Set[str] = set(
                os.path.join(root, file_name)
                for root, dir_names, file_names in os.walk(
                    obfuscation_info.get_resource_directory()
                )
                for file_name in file_names
                if file_name.endswith(".xml")
                and (
                    "layout" in root or "xml" in root
                )  # Only res/layout-*/ and res/xml-*/ folders.
            )
            xml_files.add(obfuscation_info.get_manifest_file())

            # TODO: use the following code to rename only the classes declared in
            #  application's package.

            # package_smali_files: Set[str] = set(
            #     smali_file
            #     for class_name, smali_file in self.class_name_to_smali_file.items()
            #     if class_name[1:].startswith(self.package_name.replace(".", "/"))
            # )
            #
            # # Rename the classes declared in the application's package.
            # class_rename_transformations = self.rename_class_declarations(
            #     list(package_smali_files), obfuscation_info.interactive
            # )

            # Get user defined ignore package list.
            self.ignore_package_names = obfuscation_info.get_ignore_package_names()

            # Rename all classes declared in smali files.
            class_rename_transformations = self.rename_class_declarations(
                obfuscation_info.get_smali_files(), obfuscation_info.interactive
            )

            # Update renamed classes through all the smali files.
            self.rename_class_usages_in_smali(
                obfuscation_info.get_smali_files(),
                class_rename_transformations,
                obfuscation_info.interactive,
            )

            # Update renamed classes through all the xml files.
            self.rename_class_usages_in_xml(
                list(xml_files),
                class_rename_transformations,
                obfuscation_info.interactive,
            )

        except Exception as e:
            self.logger.error(
                'Error during execution of "{0}" obfuscator: {1}'.format(
                    self.__class__.__name__, e
                )
            )
            raise

        finally:
            obfuscation_info.used_obfuscators.append(self.__class__.__name__)
