#!/usr/bin/env vpython3
#
# Copyright 2018 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import collections
import functools
import logging
import re
import subprocess
import sys

DEX_CLASS_NAME_RE = re.compile(r'\'L(?P<class_name>[^;]+);\'')
DEX_METHOD_NAME_RE = re.compile(r'\'(?P<method_name>[^\']+)\'')
DEX_METHOD_TYPE_RE = re.compile(  # type descriptor method signature re
    r'\''
    r'\('
    r'(?P<method_params>[^)]*)'
    r'\)'
    r'(?P<method_return_type>[^\']+)'
    r'\''
)
DEX_METHOD_LINE_NR_RE = re.compile(r'line=(?P<line_number>\d+)')

PROFILE_METHOD_RE = re.compile(
    r'(?P<tags>[HSP]+)'  # tags such as H/S/P
    r'(?P<class_name>L[^;]+;)'  # class name in type descriptor format
    r'->(?P<method_name>[^(]+)'
    r'\((?P<method_params>[^)]*)\)'
    r'(?P<method_return_type>.+)'
)

PROGUARD_CLASS_MAPPING_RE = re.compile(
    r'(?P<original_name>[^ ]+)'
    r' -> '
    r'(?P<obfuscated_name>[^:]+):'
)
PROGUARD_METHOD_MAPPING_RE = re.compile(
    # line_start:line_end: (optional)
    r'((?P<line_start>\d+):(?P<line_end>\d+):)?'
    r'(?P<return_type>[^ ]+)'  # original method return type
    # original method class name (if exists)
    r' (?:(?P<original_method_class>[a-zA-Z_\d.$]+)\.)?'
    r'(?P<original_method_name>[^.\(]+)'
    r'\((?P<params>[^\)]*)\)'  # original method params
    r'(?:[^ ]*)'  # original method line numbers (ignored)
    r' -> '
    r'(?P<obfuscated_name>.+)'
)  # obfuscated method name

TYPE_DESCRIPTOR_RE = re.compile(
    r'(?P<brackets>\[*)'
    r'(?:'
    r'(?P<class_name>L[^;]+;)'
    r'|'
    r'[VZBSCIJFD]'
    r')'
)

DOT_NOTATION_MAP = {
    '': '',
    'boolean': 'Z',
    'byte': 'B',
    'void': 'V',
    'short': 'S',
    'char': 'C',
    'int': 'I',
    'long': 'J',
    'float': 'F',
    'double': 'D',
}


@functools.total_ordering
class Method:
    def __init__(self, name, class_name, param_types=None, return_type=None):
        self.name = name
        self.class_name = class_name
        self.param_types = param_types
        self.return_type = return_type

    def __str__(self):
        return '{}->{}({}){}'.format(
            self.class_name,
            self.name,
            self.param_types or '',
            self.return_type or '',
        )

    def __repr__(self):
        return 'Method<{}->{}({}){}>'.format(
            self.class_name,
            self.name,
            self.param_types or '',
            self.return_type or '',
        )

    @staticmethod
    def serialize(method):
        return (
            method.class_name,
            method.name,
            method.param_types,
            method.return_type,
        )

    def __eq__(self, other):
        return self.serialize(self) == self.serialize(other)

    def __lt__(self, other):
        return self.serialize(self) < self.serialize(other)

    def __hash__(self):
        # only hash name and class_name since other fields may not be set yet.
        return hash((self.name, self.class_name))


class Class:
    def __init__(self, name):
        self.name = name
        self._methods = []

    def AddMethod(self, method, line_numbers):
        self._methods.append((method, set(line_numbers)))

    def FindMethodsAtLine(self, method_name, line_start, line_end=None):
        """Searches through dex class for a method given a name and line numbers

        The dex maps methods to line numbers, this method, given the a method name
        in this class as well as a start line and an optional end line (which act as
        hints as to which function in the class is being looked for), returns a list
        of possible matches (or none if none are found).

        Args:
          method_name: name of method being searched for
          line_start: start of hint range for lines in this method
          line_end: end of hint range for lines in this method (optional)

        Returns:
          A list of Method objects that could match the hints given, or None if no
          method is found.
        """
        found_methods = []
        if line_end is None:
            hint_lines = set([line_start])
        else:
            hint_lines = set(range(line_start, line_end + 1))

        named_methods = [
            (method, l)
            for method, l in self._methods
            if method.name == method_name
        ]

        if len(named_methods) == 1:
            return [method for method, l in named_methods]
        if len(named_methods) == 0:
            return None

        for method, line_numbers in named_methods:
            if not hint_lines.isdisjoint(line_numbers):
                found_methods.append(method)

        if len(found_methods) > 0:
            if len(found_methods) > 1:
                logging.warning(
                    'ambigous methods in dex %s at lines %s in class "%s"',
                    found_methods,
                    hint_lines,
                    self.name,
                )
            return found_methods

        for method, line_numbers in named_methods:
            if max(hint_lines) >= min(line_numbers) and min(hint_lines) <= max(
                line_numbers
            ):
                found_methods.append(method)

        if len(found_methods) > 0:
            if len(found_methods) > 1:
                logging.warning(
                    'ambigous methods in dex %s at lines %s in class "%s"',
                    found_methods,
                    hint_lines,
                    self.name,
                )
            return found_methods
        logging.warning(
            'No method named "%s" in class "%s" is mapped to lines %s',
            method_name,
            self.name,
            hint_lines,
        )
        return None


class Profile:
    def __init__(self):
        # {Method: set(char)}
        self._methods = collections.defaultdict(set)
        self._classes = []

    def AddMethod(self, method, tags):
        for tag in tags:
            self._methods[method].add(tag)

    def AddClass(self, cls):
        self._classes.append(cls)

    def WriteToFile(self, path):
        with open(path, 'w') as output_profile:
            for cls in sorted(self._classes):
                output_profile.write(cls + '\n')
            for method in sorted(self._methods):
                tags = sorted(self._methods[method])
                line = '{}{}\n'.format(''.join(tags), str(method))
                output_profile.write(line)


class ProguardMapping:
    def __init__(self):
        # {Method: set(Method)}
        self._method_mapping = collections.defaultdict(set)
        # {String: String} String is class name in type descriptor format
        self._class_mapping = dict()

    def AddMethodMapping(self, from_method, to_method):
        self._method_mapping[from_method].add(to_method)

    def AddClassMapping(self, from_class, to_class):
        self._class_mapping[from_class] = to_class

    def GetMethodMapping(self, from_method):
        return self._method_mapping.get(from_method)

    def GetClassMapping(self, from_class):
        return self._class_mapping.get(from_class, from_class)

    def MapTypeDescriptor(self, type_descriptor):
        match = TYPE_DESCRIPTOR_RE.search(type_descriptor)
        assert match is not None
        class_name = match.group('class_name')
        if class_name is not None:
            return match.group('brackets') + self.GetClassMapping(class_name)
        # just a native type, return as is
        return match.group()

    def MapTypeDescriptorList(self, type_descriptor_list):
        return TYPE_DESCRIPTOR_RE.sub(
            lambda match: self.MapTypeDescriptor(match.group()),
            type_descriptor_list,
        )


class MalformedLineException(Exception):
    def __init__(self, message, line_number):
        super().__init__(message)
        self.message = message
        self.line_number = line_number

    def __str__(self):
        return self.message + ' at line {}'.format(self.line_number)


class MalformedProguardMappingException(MalformedLineException):
    pass


class MalformedProfileException(MalformedLineException):
    pass


def _RunDexDump(dexdump_path, dex_file_path):
    return (
        subprocess.check_output([dexdump_path, dex_file_path])
        .decode('utf-8')
        .splitlines()
    )


def _ReadFile(file_path):
    with open(file_path, 'r') as f:
        return f.readlines()


def _ToTypeDescriptor(dot_notation):
    """Parses a dot notation type and returns it in type descriptor format

    eg:
    org.chromium.browser.ChromeActivity -> Lorg/chromium/browser/ChromeActivity;
    boolean -> Z
    int[] -> [I

    Args:
      dot_notation: trimmed string with a single type in dot notation format

    Returns:
      A string with the type in type descriptor format
    """
    dot_notation = dot_notation.strip()
    prefix = ''
    while dot_notation.endswith('[]'):
        prefix += '['
        dot_notation = dot_notation[:-2]
    if dot_notation in DOT_NOTATION_MAP:
        return prefix + DOT_NOTATION_MAP[dot_notation]
    return prefix + 'L' + dot_notation.replace('.', '/') + ';'


def _DotNotationListToTypeDescriptorList(dot_notation_list_string):
    """Parses a param list of dot notation format and returns it in type
    descriptor format

    eg:
    org.chromium.browser.ChromeActivity,boolean,int[] ->
        Lorg/chromium/browser/ChromeActivity;Z[I

    Args:
      dot_notation_list_string: single string with multiple comma separated types
                                in dot notation format

    Returns:
      A string with the param list in type descriptor format
    """
    return ''.join(
        _ToTypeDescriptor(param)
        for param in dot_notation_list_string.split(',')
    )


def ProcessDex(dex_dump):
    """Parses dexdump output returning a dict of class names to Class objects

    Parses output of the dexdump command on a dex file and extracts information
    about classes and their respective methods and which line numbers a method is
    mapped to.

    Methods that are not mapped to any line number are ignored and not listed
    inside their respective Class objects.

    Args:
      dex_dump: An array of lines of dexdump output

    Returns:
      A dict that maps from class names in type descriptor format (but without the
      surrounding 'L' and ';') to Class objects.
    """
    # class_name: Class
    classes_by_name = {}
    current_class = None
    current_method = None
    reading_positions = False
    reading_methods = False
    method_line_numbers = []
    for line in dex_dump:
        line = line.strip()
        if line.startswith('Class descriptor'):
            # New class started, no longer reading methods.
            reading_methods = False
            current_class = Class(
                DEX_CLASS_NAME_RE.search(line).group('class_name')
            )
            classes_by_name[current_class.name] = current_class
        elif line.startswith('Direct methods') or line.startswith(
            'Virtual methods'
        ):
            reading_methods = True
        elif reading_methods and line.startswith('name'):
            assert current_class is not None
            current_method = Method(
                DEX_METHOD_NAME_RE.search(line).group('method_name'),
                "L" + current_class.name + ";",
            )
        elif reading_methods and line.startswith('type'):
            assert current_method is not None
            match = DEX_METHOD_TYPE_RE.search(line)
            current_method.param_types = match.group('method_params')
            current_method.return_type = match.group('method_return_type')
        elif line.startswith('positions'):
            assert reading_methods
            reading_positions = True
            method_line_numbers = []
        elif reading_positions and line.startswith('0x'):
            line_number = DEX_METHOD_LINE_NR_RE.search(line).group(
                'line_number'
            )
            method_line_numbers.append(int(line_number))
        elif reading_positions and line.startswith('locals'):
            if len(method_line_numbers) > 0:
                current_class.AddMethod(current_method, method_line_numbers)
            # finished reading method line numbers
            reading_positions = False
    return classes_by_name


def ProcessProguardMapping(proguard_mapping_lines, dex):
    """Parses a proguard mapping file

    This takes proguard mapping file lines and then uses the obfuscated dex to
    create a mapping of unobfuscated methods to obfuscated ones and vice versa.

    The dex is used because the proguard mapping file only has the name of the
    obfuscated methods but not their signature, thus the dex is read to look up
    which method with a specific name was mapped to the lines mentioned in the
    proguard mapping file.

    Args:
      proguard_mapping_lines: Array of strings, each is a line from the proguard
                              mapping file (in order).
      dex: a dict of class name (in type descriptor format but without the
           enclosing 'L' and ';') to a Class object.
    Returns:
      Two dicts the first maps from obfuscated methods to a set of non-obfuscated
      ones. It also maps the obfuscated class names to original class names, both
      in type descriptor format (with the enclosing 'L' and ';')
    """
    mapping = ProguardMapping()
    reverse_mapping = ProguardMapping()
    to_be_obfuscated = []
    current_class_orig = None
    current_class_obfs = None
    for index, line in enumerate(proguard_mapping_lines):
        if line.strip() == '':
            continue
        if not line.startswith(' '):
            match = PROGUARD_CLASS_MAPPING_RE.search(line)
            if match is None:
                raise MalformedProguardMappingException(
                    'Malformed class mapping', index
                )
            current_class_orig = match.group('original_name')
            current_class_obfs = match.group('obfuscated_name')
            mapping.AddClassMapping(
                _ToTypeDescriptor(current_class_obfs),
                _ToTypeDescriptor(current_class_orig),
            )
            reverse_mapping.AddClassMapping(
                _ToTypeDescriptor(current_class_orig),
                _ToTypeDescriptor(current_class_obfs),
            )
            continue

        assert current_class_orig is not None
        assert current_class_obfs is not None
        line = line.strip()
        match = PROGUARD_METHOD_MAPPING_RE.search(line)
        # check if is a method mapping (we ignore field mappings)
        if match is not None:
            # check if this line is an inlining by reading ahead 1 line.
            if index + 1 < len(proguard_mapping_lines):
                next_match = PROGUARD_METHOD_MAPPING_RE.search(
                    proguard_mapping_lines[index + 1].strip()
                )
                if (
                    next_match
                    and match.group('line_start') is not None
                    and next_match.group('line_start')
                    == match.group('line_start')
                    and next_match.group('line_end') == match.group('line_end')
                ):
                    continue  # This is an inlining, skip

            original_method = Method(
                match.group('original_method_name'),
                _ToTypeDescriptor(
                    match.group('original_method_class') or current_class_orig
                ),
                _DotNotationListToTypeDescriptorList(match.group('params')),
                _ToTypeDescriptor(match.group('return_type')),
            )

            if match.group('line_start') is not None:
                obfs_methods = dex[
                    current_class_obfs.replace('.', '/')
                ].FindMethodsAtLine(
                    match.group('obfuscated_name'),
                    int(match.group('line_start')),
                    int(match.group('line_end')),
                )

                if obfs_methods is None:
                    continue

                for obfs_method in obfs_methods:
                    mapping.AddMethodMapping(obfs_method, original_method)
                    reverse_mapping.AddMethodMapping(
                        original_method, obfs_method
                    )
            else:
                to_be_obfuscated.append(
                    (original_method, match.group('obfuscated_name'))
                )

    for original_method, obfuscated_name in to_be_obfuscated:
        obfuscated_method = Method(
            obfuscated_name,
            reverse_mapping.GetClassMapping(original_method.class_name),
            reverse_mapping.MapTypeDescriptorList(original_method.param_types),
            reverse_mapping.MapTypeDescriptor(original_method.return_type),
        )
        mapping.AddMethodMapping(obfuscated_method, original_method)
        reverse_mapping.AddMethodMapping(original_method, obfuscated_method)
    return mapping, reverse_mapping


def ProcessProfile(input_profile, proguard_mapping):
    """Parses an android profile and uses the proguard mapping to (de)obfuscate it

    This takes the android profile lines and for each method or class for the
    profile, it uses the mapping to either obfuscate or deobfuscate (based on the
    provided mapping) and returns a Profile object that stores this information.

    Args:
      input_profile: array of lines of the input profile
      proguard_mapping: a proguard mapping that would map from the classes and
                        methods in the input profile to the classes and methods
                        that should be in the output profile.

    Returns:
      A Profile object that stores the information (ie list of mapped classes and
      methods + tags)
    """
    profile = Profile()
    for index, line in enumerate(input_profile):
        line = line.strip()
        if line.startswith('L'):
            profile.AddClass(proguard_mapping.GetClassMapping(line))
            continue
        match = PROFILE_METHOD_RE.search(line)
        if not match:
            raise MalformedProfileException("Malformed line", index)

        method = Method(
            match.group('method_name'),
            match.group('class_name'),
            match.group('method_params'),
            match.group('method_return_type'),
        )

        mapped_methods = proguard_mapping.GetMethodMapping(method)
        if mapped_methods is None:
            logging.warning(
                'No method matching "%s" has been found in the proguard '
                'mapping file',
                method,
            )
            continue

        for original_method in mapped_methods:
            profile.AddMethod(original_method, match.group('tags'))

    return profile


def ObfuscateProfile(
    nonobfuscated_profile,
    dex_file,
    proguard_mapping,
    dexdump_path,
    output_filename,
):
    """Helper method for obfuscating a profile.

    Args:
      nonobfuscated_profile: a profile with nonobfuscated symbols.
      dex_file: path to the dex file matching the mapping.
      proguard_mapping: a mapping from nonobfuscated to obfuscated symbols used
        in the dex file.
      dexdump_path: path to the dexdump utility.
      output_filename: output filename in which to write the obfuscated profile.
    """
    dexinfo = ProcessDex(_RunDexDump(dexdump_path, dex_file))
    _, reverse_mapping = ProcessProguardMapping(
        _ReadFile(proguard_mapping), dexinfo
    )
    obfuscated_profile = ProcessProfile(
        _ReadFile(nonobfuscated_profile), reverse_mapping
    )
    obfuscated_profile.WriteToFile(output_filename)


def main(args):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dexdump-path', required=True, help='Path to dexdump binary.'
    )
    parser.add_argument(
        '--dex-path',
        required=True,
        help='Path to dex file corresponding to the proguard mapping file.',
    )
    parser.add_argument(
        '--proguard-mapping-path',
        required=True,
        help='Path to input proguard mapping file corresponding to the dex file.',
    )
    parser.add_argument(
        '--output-profile-path', required=True, help='Path to output profile.'
    )
    parser.add_argument(
        '--input-profile-path', required=True, help='Path to output profile.'
    )
    parser.add_argument(
        '--verbose',
        action='store_true',
        default=False,
        help='Print verbose output.',
    )
    obfuscation = parser.add_mutually_exclusive_group(required=True)
    obfuscation.add_argument(
        '--obfuscate',
        action='store_true',
        help='Indicates to output an obfuscated profile given a deobfuscated '
        'one.',
    )
    obfuscation.add_argument(
        '--deobfuscate',
        dest='obfuscate',
        action='store_false',
        help='Indicates to output a deobfuscated profile '
        'given an obfuscated one.',
    )
    options = parser.parse_args(args)

    if options.verbose:
        log_level = logging.WARNING
    else:
        log_level = logging.ERROR
    logging.basicConfig(format='%(levelname)s: %(message)s', level=log_level)

    dex = ProcessDex(_RunDexDump(options.dexdump_path, options.dex_path))
    proguard_mapping, reverse_proguard_mapping = ProcessProguardMapping(
        _ReadFile(options.proguard_mapping_path), dex
    )
    if options.obfuscate:
        profile = ProcessProfile(
            _ReadFile(options.input_profile_path), reverse_proguard_mapping
        )
    else:
        profile = ProcessProfile(
            _ReadFile(options.input_profile_path), proguard_mapping
        )
    profile.WriteToFile(options.output_profile_path)


if __name__ == '__main__':
    main(sys.argv[1:])