# 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 base64
import binascii
import hashlib

from .artifact import AndroidArtifact


class DumpsysADBArtifact(AndroidArtifact):
    multiline_fields = ["user_keys", "keystore"]

    def indented_dump_parser(self, dump_data):
        """
        Parse the indented dumpsys output, generated by DualDumpOutputStream in Android.
        """
        res = {}
        stack = [res]
        cur_indent = 0
        in_multiline = False
        for line in dump_data.strip(b"\n").split(b"\n"):
            # Track the level of indentation
            indent = len(line) - len(line.lstrip())
            if indent < cur_indent:
                # If the current line is less indented than the previous one, back out
                stack.pop()
                cur_indent = indent
            else:
                cur_indent = indent

            # Split key and value by '='
            vals = line.lstrip().split(b"=", 1)
            key = vals[0].decode("utf-8")
            current_dict = stack[-1]

            # Annoyingly, some values are multiline and don't have a key on each line
            if in_multiline:
                if key == "":
                    # If the line is empty, it's the terminator for the multiline value
                    in_multiline = False
                    stack.pop()
                else:
                    current_dict.append(line.lstrip())
                continue

            if key == "}":
                stack.pop()
                continue

            if vals[1] == b"{":
                # If the value is a new dictionary, add it to the stack
                current_dict[key] = {}
                stack.append(current_dict[key])

            # Handle continue multiline values
            elif key in self.multiline_fields:
                current_dict[key] = []
                current_dict[key].append(vals[1])

                in_multiline = True
                stack.append(current_dict[key])
            else:
                # If the value something else, store it in the current dictionary
                current_dict[key] = vals[1]

        return res

    def parse_xml(self, xml_data):
        """
        Parse XML data from dumpsys ADB output
        """
        import xml.etree.ElementTree as ET

        keystore = []
        keystore_root = ET.fromstring(xml_data)
        for adb_key in keystore_root.findall("adbKey"):
            key_info = self.calculate_key_info(adb_key.get("key").encode("utf-8"))
            key_info["last_connected"] = adb_key.get("lastConnection")
            keystore.append(key_info)

        return keystore

    @staticmethod
    def calculate_key_info(user_key: bytes) -> dict:
        if b" " in user_key:
            key_base64, user = user_key.split(b" ", 1)
        else:
            key_base64, user = user_key, b""

        try:
            key_raw = base64.b64decode(key_base64)
            key_fingerprint = hashlib.md5(key_raw).hexdigest().upper()
            key_fingerprint_colon = ":".join(
                [key_fingerprint[i : i + 2] for i in range(0, len(key_fingerprint), 2)]
            )
        except binascii.Error:
            # Impossible to parse base64
            key_fingerprint_colon = ""

        return {
            "user": user.decode("utf-8"),
            "fingerprint": key_fingerprint_colon,
            "key": key_base64,
        }

    def check_indicators(self) -> None:
        if not self.results:
            return

        for entry in self.results:
            for user_key in entry.get("user_keys", []):
                self.log.debug(
                    f"Found trusted ADB key for user '{user_key['user']}' with fingerprint "
                    f"'{user_key['fingerprint']}'"
                )

    def parse(self, content: bytes) -> None:
        """
        Parse the Dumpsys ADB section
        Adds results to self.results (List[Dict[str, str]])

        :param content: content of the ADB section (string)
        """
        if not content or b"Can't find service: adb" in content:
            self.log.error(
                "Could not load ADB data from dumpsys. "
                "It may not be supported on this device."
            )
            return

        start_of_json = content.find(b"\n{")
        if start_of_json == -1:
            self.log.error("Unable to find ADB manager state in dumpsys output")
            return

        end_of_json = content.rfind(b"}\n")
        if end_of_json == -1 or end_of_json <= start_of_json:
            self.log.error("Unable to find complete ADB manager state in dumpsys output")
            return

        json_content = content[start_of_json + 2 : end_of_json - 2].rstrip()

        parsed = self.indented_dump_parser(json_content)
        if parsed.get("debugging_manager") is None:
            self.log.error("Unable to find expected ADB entries in dumpsys output")  # noqa
            return

        # Keystore can be in different levels, as the basic parser
        # is not always consistent due to different dumpsys formats.
        if parsed.get("keystore"):
            keystore_data = b"\n".join(parsed["keystore"])
        elif parsed["debugging_manager"].get("keystore"):
            keystore_data = b"\n".join(parsed["debugging_manager"]["keystore"])
        else:
            keystore_data = None

        # Keystore is in XML format on some devices and we need to parse it
        if keystore_data and keystore_data.startswith(b"<?xml"):
            parsed["debugging_manager"]["keystore"] = self.parse_xml(keystore_data)
        else:
            # Keystore is not XML format
            parsed["debugging_manager"]["keystore"] = keystore_data

        parsed = parsed["debugging_manager"]

        # Calculate key fingerprints for better readability
        key_info = []
        for user_key in parsed.get("user_keys", []):
            user_info = self.calculate_key_info(user_key)
            key_info.append(user_info)

        parsed["user_keys"] = key_info
        self.results = [parsed]