add Log.e for smail script, for android apk reverse engine debugging

smaliinjector.py

"""
Usage:
    smaliinjector [options] inject [<inject_smali_path>] [--exclude=<smali_path>]...
    smaliinjector [options] inject [<smali_path>]... [--verbose | --loglevel=<loglevel>]
    smaliinjector [options] restore [<restore_smali_path>]...

Options:
    -e<smali_path>, --exclude=<smali_path>
                  exclude folders or files for `inject_smali_path`.
                  Only works with `inject [inject_smali_path]`
    --verbose     print more infomation
    --loglevel=<level>
                  log level, values to be: DEBUG, INFO, WARNING, ERROR, CRITICAL [default: INFO]
    -h, --help    print this message and exit

Command:
    inject                 inject debug info into smali files
    restore                remove injected debug info from smali files
    <inject_smali_path>    smail folder that need to be inject [default: ./smali]
    <restore_smali_path>   smail folder that need to be restore [default: ./smali]
"""

import sys
import os
import re
import logging
import json
import atexit


# logging.basicConfig(level=logging.DEBUG)
LOGGER = logging.getLogger('smaliinjector')


def path_under_any(path, super_pathes):
    for super_path in super_pathes:
        if not super_path.endswith(os.sep):
            super_path = super_path + os.sep
        if path.startswith(super_path):
            return True
    return False


class MethodInjector(object):

    CONTINUE = 'continue'
    END = 'end'
    LOGGER = logging.getLogger('smaliinjector.MethodInjector')

    LOG_LINES_NO_BREAK = (
        '# INJECT:BLOCKNEW:START',
        'const-string v{v_tag}, "{tag}"',
        'const-string v{v_payload}, "{payload}"',
        'invoke-static {{v{v_tag}, v{v_payload}}}, Landroid/util/Log;->e(Ljava/lang/String;Ljava/lang/String;)I',
        '# INJECT:BLOCKNEW:END',
    )
    # METHOD_ARGS_RE = re.compile(r'\(([\w\d\[\$\-_;/]*?)\)')
    METHOD_P_RE = re.compile(r'p(\d+)[,}]?')
    PARAM_INLINE_RE = re.compile(r'^.param p\d+,')

    def __init__(self):
        self.lines = []
        self.method_sign = None
        self.status = 'begin'
        self.feed_status = self.CONTINUE
        self.locals_num = None
        self.locals_at = None
        self.dot_chain = []
        self.inject_log_at = None
        self.need_inject = True

    def feed(self, line):
        if self.feed_status == self.END:
            raise RuntimeError('stop feed for {}'.format(self.method_sign))

        line_strip = line.strip()
        status = self.status
        self.lines.append(line)
        if not self.need_inject:
            if line_strip == '.end method':
                self.LOGGER.info('%s end, no inject', self.status)
                self.status = 'end'
                self.feed_status = self.END
                return self.END
            else:
                self.LOGGER.info('%s continue, no inject: %s', self.status, line_strip)
                return self.CONTINUE

        if status == 'begin':
            assert line_strip.startswith('.method')
            self.method_sign = line_strip.split()[-1]
            self.status = 'head.locals'
            self.LOGGER.info('begin with sign %s', self.method_sign)
            return self.CONTINUE
        if status == 'head.locals':
            if line_strip == '.end method':
                self.LOGGER.info('%s with end method, no inject', self.status)
                self.need_inject = False
                self.status = 'end'
                self.feed_status = self.END
                return self.END

            if not line_strip.startswith('.locals ') and line_strip.startswith('.'):
                self.LOGGER.info('%s with no locals, no inject', self.status)
                self.need_inject = False
                self.status = 'head.content'
                self.status = 'end'
                return self.CONTINUE

            locals_num = int(line_strip.split()[-1])
            # if locals_num >= 13:
            #     self.need_inject = False
            #     return self.CONTINUE

            assert self.locals_num is None
            assert self.locals_at is None
            self.locals_num = locals_num
            self.locals_at = len(self.lines) - 1

            self.status = 'head.content'

            self.LOGGER.info('head.locals with locals_num %s locals_at %s', self.locals_num, self.locals_at)
            return self.CONTINUE

        if line_strip.startswith('.end local '):
            self.LOGGER.debug('%s with %s', self.status, line_strip)
            return self.CONTINUE

        if line_strip.startswith('.end '):
            _, end_signer = line_strip.split()
            if end_signer == 'method':
                self.LOGGER.info('%s with method end', self.status)
                assert not self.dot_chain, self.dot_chain
                assert self.status == 'body'
                self.feed_status = self.END
                self.status = 'end'
                return self.END
            else:
                self.LOGGER.debug('%s with %s end', self.status, end_signer)
                assert self.dot_chain, self.dot_chain
                # assert status == 'head.content', status
                dot_signer = end_signer
                last_chain = self.dot_chain.pop(-1)
                assert last_chain == dot_signer, (dot_signer, last_chain)

        elif line_strip.startswith('.annotation '):
            self.LOGGER.debug('%s to head.content: %s', self.status, line_strip)
            self.status = 'head.content'
            self.dot_chain.append('annotation')
            return self.CONTINUE

        elif line_strip.startswith('.param '):
            self.LOGGER.debug('%s to head.content: %s', self.status, line_strip)
            self.status = 'head.content'
            if self.PARAM_INLINE_RE.match(line_strip) is None:  # not one line
                self.dot_chain.append('param')
            return self.CONTINUE

        elif line_strip.startswith('.line ') or line_strip.startswith('.catchall ') or line_strip.startswith('.catch ') or line_strip.startswith('.enum ') or line_strip.startswith('.local ') or line_strip.startswith('.restart '):
            self.LOGGER.debug('%s with %r', self.status, line_strip)
            return self.CONTINUE

        elif line_strip.startswith('.'):
            self.LOGGER.debug('%s with dot %r', self.status, line)
            assert not line_strip.startswith('.locals ')
            dot_signer = line_strip.split()[0][1:]
            self.dot_chain.append(dot_signer)
            return self.CONTINUE

        if line_strip == '':  # check empty line
            if self.dot_chain:  # still in head
                self.LOGGER.debug('%s with empty line in dot_chain: %r', self.status, self.dot_chain)
                assert status == 'head.content'
                return self.CONTINUE
            # no dot chain already, into function body now
            self.LOGGER.debug('%s with empty line out of dot_chain', self.status)
            if status == 'head.content':
                self.inject_log_at = len(self.lines)
            self.status = 'body'
        else:
            self.LOGGER.debug('%s with line %r', self.status, line)

        return self.CONTINUE

    def get_lines(self, file_relpath):
        if not self.need_inject:
            return self.lines

        assert not self.dot_chain, self.dot_chain
        assert self.feed_status == self.END, self.feed_status
        assert self.status == 'end', self.status
        assert self.method_sign is not None
        assert self.inject_log_at is not None
        assert self.locals_at is not None
        assert self.locals_num is not None
        locals_num = self.locals_num
        locals_at = self.locals_at
        locals_old = self.lines[locals_at]
        locals_prepend = locals_old.split('.locals')[0]

        method_sign = self.method_sign

        # method_sign_args = self.METHOD_ARGS_RE.search(method_sign).group(1)
        # method_args_count = self.count_args(method_sign_args) + 1  # + `this`
        method_args_count = self.count_args(''.join(self.lines))

        self.LOGGER.debug('method sign %s; locals_num=%s, method_args_count=%s', method_sign, locals_num, method_args_count)
        # 0-15, total 16
        # local, params, can
        # 0, 16, false
        # 1, 15, false
        # 2, *, true
        # 14, 2, true
        # 14, 1, true
        # 14, 0, true
        # 15, 1, true
        # 16, 0, true
        # NOTE: reserve for param `this`
        # can_inject = method_args_count < 15
        # inject_by_incr = locals_num < 15 and (method_args_count + locals_num) < 15
        #
        # if not can_inject:
        #     return self.lines
        #
        # if not inject_by_incr:  # too many, use v0, v1
        #     v_tag = 0
        #     v_payload = 1
        # else:  # use incr locals
        #     v_tag = locals_num
        #     v_payload = locals_num + 1
        #     locals_new = '{}.locals {}  # INJECT:INLINE:ORIGINAL {}\n'.format(
        #         locals_prepend,
        #         locals_num + 2,
        #         locals_old.rstrip()
        #     )
        #     self.lines[locals_at] = locals_new

        safe_inject = locals_num < 15 and (method_args_count + locals_num) < 15
        if not safe_inject:
            return self.lines
        v_tag = locals_num
        v_payload = locals_num + 1
        locals_new = '{}.locals {}  # INJECT:INLINE:ORIGINAL {}\n'.format(
            locals_prepend,
            locals_num + 2,
            locals_old.rstrip()
        )
        self.lines[locals_at] = locals_new

        log_lines = []
        msg_payload = json.dumps('CALL:{file_relpath}:{method_sign}'.format(
            file_relpath=file_relpath,
            method_sign=self.method_sign
        ), ensure_ascii=False)[1:-1]  # just do some escape here

        for line_formatter in self.LOG_LINES_NO_BREAK:
            line_raw = line_formatter.format(
                v_tag=v_tag,
                v_payload=v_payload,
                tag='INJECT',
                payload=msg_payload
            )
            log_line = '{}{}\n'.format(locals_prepend, line_raw)
            log_lines.append(log_line)
        self.lines[self.inject_log_at: self.inject_log_at] = log_lines

        return self.lines

    @classmethod
    def count_args(klass, method):
        count = 0
        for each in klass.METHOD_P_RE.findall(method):
            count = max(count, int(each) + 1)
        return count

    # @staticmethod
    # def count_args(method_sign, logger=LOGGER):
    #     # e.g.
    #     # ILjava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Landroid/net/Uri;Ljava/lang/String;JLjava/lang/String;Ljava/util/List;Ljava/lang/String;Ljava/lang/String;
    #     count = 0
    #     index = 0
    #     while index <= len(method_sign) - 1:
    #         char = method_sign[index]
    #         if char in (
    #                     'Z',
    #                     'B',
    #                     'S',
    #                     'C',
    #                     'I',
    #                     'J',
    #                     'F',
    #                     'D'
    #                 ):
    #             count += 1
    #             index += 1
    #         elif char == '[':
    #             index += 1
    #         elif char == 'L':
    #             # go find the ;
    #             rest_method_sign = method_sign[index:]
    #             if ';' not in rest_method_sign:
    #                 count += 1
    #                 return count
    #             rest_index = rest_method_sign.index(';')
    #             index += (rest_index + 1)
    #             count += 1
    #     return count


# file_abspath = None


# @atexit.register
# def script_exit():
#     global file_abspath
#     global LOGGER
#     LOGGER.info('process exit at file %s', file_abspath)


def inject(folders, excludes=frozenset()):
    LOGGER = logging.getLogger('smaliinjector.inject')
    LOGGER.info('processing %s, excludes %s', folders, excludes)
    for folder in folders:
        if os.path.isfile(folder):
            LOGGER.debug('%s is file', folder)
            dirpath, filename = os.path.split(folder)
            iter_path = ((
                dirpath,
                [],
                [filename]
            ),)
        else:
            iter_path = os.walk(folder)
        for dirpath, _dirnames, filenames in iter_path:
            relpath = os.path.relpath(dirpath, folder)
            if path_under_any(relpath, excludes):
                LOGGER.info('skip folder %s', dirpath)
                continue
            LOGGER.debug('processing folder %s', relpath)

            for filename in filenames:
                if not filename.endswith('.smali'):
                    continue

                file_abspath = os.path.normpath(os.path.join(folder, dirpath, filename))
                file_relpath = os.path.join(relpath, filename)
                LOGGER.info('processing file %s (%s)', file_relpath, file_abspath)

                with open(file_abspath, 'a+', encoding='utf-8') as f:
                    f.seek(0)
                    lines = []
                    for line in f:
                        if line.strip().startswith('.method '):
                            method_injector = MethodInjector()
                            method_injector.feed(line)
                            while method_injector.feed(next(f)) == method_injector.CONTINUE:
                                pass

                            method_lines = method_injector.get_lines(file_relpath)
                            lines.extend(method_lines)
                            continue
                        lines.append(line)

                    f.seek(0)
                    f.truncate()
                    f.writelines(lines)


def restore(folders):
    LOGGER = logging.getLogger('smaliinjector.restore')
    # {}.locals {}  # INJECT:INLINE:ORIGINAL {}
    LOGGER.info('restore %s', folders)
    # INLINE_RE = re.compile(r'\s*.*?  # INJECT:INLINE:ORIGINAL (?P<ori_line>.*?)\n')
    # BLOCK_RE = re.compile(r'\s*# INJECT:BLOCKNEW:START\n.*?# INJECT:BLOCKNEW:END\n', re.DOTALL)
    for folder in folders:
        if os.path.isfile(folder):
            LOGGER.debug('%s is file', folder)
            dirpath, filename = os.path.split(folder)
            iter_path = ((
                dirpath,
                [],
                [filename]
            ),)
        else:
            iter_path = os.walk(folder)
        for dirpath, _dirnames, filenames in iter_path:
            for filename in filenames:
                if not filename.endswith('.smali'):
                    continue
                file_abspath = os.path.normpath(os.path.join(folder, dirpath, filename))
                LOGGER.info('processing %s', file_abspath)

                with open(file_abspath, 'a+', encoding='utf-8') as f:
                    f.seek(0)
                    lines = []
                    change_count = 0
                    for line in f:
                        if '  # INJECT:INLINE:ORIGINAL ' in line:
                            LOGGER.debug('found line %r', line)
                            _, ori_line = line.split('  # INJECT:INLINE:ORIGINAL ')
                            line = ori_line
                            change_count += 1
                            lines.append(line)
                        elif line.strip() == '# INJECT:BLOCKNEW:START':
                            LOGGER.debug('found line block %r', line)
                            while next(f).strip() != '# INJECT:BLOCKNEW:END':
                                pass
                            else:
                                LOGGER.debug('found line block end')
                            change_count += 1
                        else:
                            # LOGGER.debug('line: %r', line)
                            lines.append(line)

                    if change_count > 0:
                        f.seek(0)
                        f.truncate()
                        f.writelines(lines)
                        LOGGER.info('file restored (%s): %s', change_count, file_abspath)
                    else:
                        LOGGER.debug('not change in file: %s', file_abspath)


if __name__ == '__main__':
    import docpie

    try:
        import colorlog
    except ImportError:
        FORMATTER = logging.Formatter(
            '[%(levelname)1.1s %(lineno)3d %(asctime)s %(funcName)s]'
            ' %(message)s'
        )
    else:
        FORMATTER = colorlog.ColoredFormatter(
            '%(log_color)s'
            '[%(levelname)1.1s %(lineno)3d %(asctime)s %(funcName)s]'
            '%(reset)s'
            ' %(message)s',
            log_colors={
                'DEBUG': 'cyan',
                'INFO': 'green',
                'WARNING': 'yellow',
                'ERROR': 'red',
                'CRITICAL': 'bold_red,bg_black',
        	},
        )

    HDLR = logging.StreamHandler(sys.stdout)
    HDLR.setFormatter(FORMATTER)
    LOGGER.addHandler(HDLR)
    # LOGGER.setLevel(logging.DEBUG)
    # LOGGER.setLevel(logging.WARNING)

    CLIARGS = docpie.docpie(__doc__)
    if CLIARGS['--verbose']:
        LOGLEVEL = logging.DEBUG
    else:
        LOGLEVEL = {
            'DEBUG': logging.DEBUG,
            'INFO': logging.INFO,
            'WARNING': logging.WARNING,
            'ERROR': logging.ERROR,
            'CRITICAL': logging.CRITICAL,
        }[CLIARGS['--loglevel'].upper()]
    LOGGER.setLevel(LOGLEVEL)

    if CLIARGS['inject']:
        excludes = CLIARGS['--exclude']
        if excludes:
            # LOGGER.info('exclude %s', excludes)
            inject_path = os.path.abspath(CLIARGS['<inject_smali_path>'] or 'smali')
            inject_pathes = frozenset((inject_path,))
        else:
            # LOGGER.info('no exclude')
            # print(CLIARGS)
            user_smali_path = (CLIARGS['<inject_smali_path>'],) if CLIARGS['<inject_smali_path>'] else CLIARGS['<smali_path>']
            inject_pathes = frozenset(os.path.abspath(each) for each in (user_smali_path or ['smali']))
        inject(inject_pathes, excludes)
    elif CLIARGS['restore']:
        restore(frozenset(os.path.abspath(each) for each in (CLIARGS['<restore_smali_path>'] or ['smali'])))