#!/usr/bin/python3
# -*- coding: utf-8 -*-

import argparse
import sys
import pychrysalide
from pychrysalide.analysis.contents import FileContent
from pychrysalide.format.dex import DexFormat
from pychrysalide.analysis import LoadedBinary



_access_flags = [
    [ DexFormat.ACC_PUBLIC, 'PUBLIC' ],
    [ DexFormat.ACC_PRIVATE, 'PRIVATE' ],
    [ DexFormat.ACC_PROTECTED, 'PROTECTED' ],
    [ DexFormat.ACC_STATIC, 'STATIC' ],
    [ DexFormat.ACC_FINAL, 'FINAL' ],
    [ DexFormat.ACC_SYNCHRONIZED, 'SYNCHRONIZED' ],
    [ DexFormat.ACC_VOLATILE, 'VOLATILE' ],
    [ DexFormat.ACC_BRIDGE, 'BRIDGE' ],
    [ DexFormat.ACC_TRANSIENT, 'TRANSIENT' ],
    [ DexFormat.ACC_VARARGS, 'VARARGS' ],
    [ DexFormat.ACC_NATIVE, 'NATIVE' ],
    [ DexFormat.ACC_INTERFACE, 'INTERFACE' ],
    [ DexFormat.ACC_ABSTRACT, 'ABSTRACT' ],
    [ DexFormat.ACC_STRICT, 'STRICT' ],
    [ DexFormat.ACC_SYNTHETIC, 'SYNTHETIC' ],
    [ DexFormat.ACC_ANNOTATION, 'ANNOTATION' ],
    [ DexFormat.ACC_ENUM, 'ENUM' ],
    [ DexFormat.ACC_CONSTRUCTOR, 'CONSTRUCTOR' ],
    [ DexFormat.ACC_DECLARED_SYNCHRONIZED, 'DECLARED_SYNCHRONIZED' ]
]


def _build_format_string(level, suffix = None):
    """Build the format string for the dump tree."""

    level *= 2

    string = ' ' * level;

    string += '{0:<%d}' % (20 - level)

    string += ':' if suffix else '-'

    string += ' '

    if suffix:
        string += suffix

    return string


def _translate_access_flags(mask):
    """Build a description of access flags."""

    access_desc = ''

    for acc in _access_flags:
        if mask & acc[0]:
            if len(access_desc) > 0:
                access_desc += ' ' + acc[1]
            else:
                access_desc += acc[1]

    return access_desc


def dump_dex_class(cls, idx):
    """Output the content of a Dex class."""

    print(_build_format_string(0).format('Class #%d' % idx))

    print(_build_format_string(1, '{1}').format('Class descriptor', cls.type))

    access_desc = _translate_access_flags(cls.definition.access_flags)

    print(_build_format_string(1, '0x{1:x} ({2})').format('Access flags', cls.definition.access_flags, access_desc))

    print(_build_format_string(1, '{1}').format('Superclass', cls.super))

    print(_build_format_string(1).format('Interfaces'))

    if cls.interfaces:

        counter = 0

        for ifc in cls.interfaces:
            print(_build_format_string(2, '{1}').format('#%u' % counter, ifc))
            counter += 1

    print(_build_format_string(1).format('Static fields'))

    if cls.static_fields:

        counter = 0

        for fld in cls.static_fields:
            dump_dex_field(fld, counter)
            counter += 1

    print(_build_format_string(1).format('Instance fields'))

    if cls.instance_fields:

        counter = 0

        for fld in cls.instance_fields:
            dump_dex_field(fld, counter)
            counter += 1

    print(_build_format_string(1).format('Direct methods'))

    if cls.direct_methods:

        counter = 0

        for meth in cls.direct_methods:
            dump_dex_method(meth, counter)
            counter += 1

    print(_build_format_string(1).format('Virtual methods'))

    if cls.virtual_methods:

        counter = 0

        for meth in cls.virtual_methods:
            dump_dex_method(meth, counter)
            counter += 1

    print(_build_format_string(1, '0x{1:x} ({2})').format('Source', cls.definition.source_file_idx, cls.source_file))


def dump_dex_field(fld, idx):
    """Output the description of a Dex class field."""

    print(_build_format_string(2).format('#%u' % idx))

    print(_build_format_string(3, '{1:s}').format('name', fld.variable.name))

    print(_build_format_string(3, '{1}').format('type', fld.variable.type))

    access_desc = _translate_access_flags(fld.encoded.access_flags)

    print(_build_format_string(3, '0x{1:x} ({2})').format('access flags', fld.encoded.access_flags, access_desc))


def dump_dex_method(meth, idx):
    """Output the content of a Dex method."""

    print(_build_format_string(2).format('#%u' % idx))

    print(_build_format_string(3, '{1:s}').format('name', meth.routine.name))

    print(_build_format_string(3, '{1}').format('prototype', meth.routine))

    access_desc = _translate_access_flags(meth.encoded.access_flags)

    print(_build_format_string(3, '0x{1:x} ({2})').format('access flags', meth.encoded.access_flags, access_desc))

    if meth.code_item:

        print(_build_format_string(3).format('code'))

        print(_build_format_string(4, '{1}').format('registers', meth.code_item.registers_size))

        print(_build_format_string(4, '{1}').format('ins', meth.code_item.ins_size))

        print(_build_format_string(4, '{1}').format('outs', meth.code_item.outs_size))

        print(_build_format_string(4, '{1} 16-bit code unit{2}').format('insns size', meth.code_item.insns_size, 's' if meth.code_item.insns_size > 1 else ''))

    else:

        print(_build_format_string(3, '{1}').format('code', '(none)'))


def dump_dex_pool(fmt):
    """Dump the Dex pool content."""

    print(_build_format_string(0).format('Strings'))

    counter = 0

    for s in fmt.pool.strings:
        print(_build_format_string(1, '{1}').format('#%u' % counter, s))
        counter += 1

    print()

    print(_build_format_string(0).format('Types'))

    counter = 0

    for t in fmt.pool.types:
        print(_build_format_string(1, '{1}').format('#%u' % counter, t))
        counter += 1

    print()

    print(_build_format_string(0).format('Prototypes'))

    counter = 0

    for p in fmt.pool.prototypes:
        print(_build_format_string(1, '{1}').format('#%u' % counter, p))
        counter += 1

    print()

    print(_build_format_string(0).format('Fields'))

    counter = 0

    for f in fmt.pool.fields:
        print(_build_format_string(1, '{1}').format('#%u' % counter, f))
        counter += 1

    print()

    print(_build_format_string(0).format('Methods'))

    counter = 0

    for m in fmt.pool.methods:
        print(_build_format_string(1, '{1}').format('#%u' % counter, m.routine))
        counter += 1

    print()

    print(_build_format_string(0).format('Classes'))

    counter = 0

    for c in fmt.pool.classes:
        print(_build_format_string(1, '{1}').format('#%u' % counter, c.type))
        counter += 1


if __name__ == '__main__':
    """Script entry point."""

    parser = argparse.ArgumentParser()

    parser.add_argument('-p', '--pool', help='dump the Dex pool', action='store_true')
    parser.add_argument('-d', '--dump', help='dump the Dex classes (default action)', action='store_true')
    parser.add_argument('classes',      help='path to the Dex file to analyze', metavar='classes.dex')

    args = parser.parse_args()

    cnt = FileContent(args.classes)

    if cnt is None:
        sys.exit('No content to load!')

    fmt = DexFormat(cnt)

    if fmt is None:
        sys.exit('Failed to load the Dex format!')

    print('%s: %s' % (args.classes, fmt.description))
    print()

    binary = LoadedBinary(fmt)

    if binary is None:
        sys.exit('Failed to load the binary!')

    status = binary.analyze_and_wait()

    if not(status):
        sys.exit('Failed to analyze the binary!')

    if args.pool:

        dump_dex_pool(fmt)

    else:

        counter = 0

        for c in fmt.classes:

            dump_dex_class(c, counter)
            print()

            counter += 1