from paraspace import dexfile
import os

def _install_dexfile_4_deptracker():
    global dexfile
    import imp, sys
    from paraspace import dex_deptracker
    from paraspace import injection
    
    try:
        new_dexfile = imp.load_compiled('dexfile', dexfile.__file__)
    except ImportError:
        new_dexfile = imp.load_source('dexfile', dexfile.__file__)
        pass
    dex_deptracker.dexfile = new_dexfile
    dexfile = new_dexfile
    dex_deptracker._nest_types = (dexfile.array, dexfile.cond, dexfile.switch)

    injection.dex_type_2_array_attr_map = \
        injection._saved_dex_type_2_array_attr_map
    
    sys.modules['paraspace.dexfile'] = new_dexfile
    pass


def _find_map(dex, map_type):
    for map in dex.maps.items.items:
        if map.type == map_type:
            return map
        pass
    pass


def inject_fakefile_to_helloworld_test():
    from paraspace.dex_deptracker import prepare_dep_decls
    from paraspace.injection import dexfile_insert_class
    
    _install_dexfile_4_deptracker()
    
    all_dep_decls = prepare_dep_decls()
    
    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    
    helloworld_fn = os.path.join(srcroot, 'data', 'helloworld.dex')
    helloworld_dex = dexfile.DEXFile.open(helloworld_fn)

    classdef_map = _find_map(helloworld_dex, 0x0006)
    saved_classdef_map_sz = classdef_map.size
    saved_methodids_sz = len(helloworld_dex.methodIds.items)
    
    codeitems_map = _find_map(helloworld_dex, 0x2001)
    saved_codeitems_sz = codeitems_map.size

    helloworld_linked = \
        dexfile.DEXFile_linked.build_dependencies(helloworld_dex,
                                                  all_dep_decls)

    fakefile_fn = os.path.join(srcroot, 'data', 'fakefile.dex')
    fakefile_dex = dexfile.DEXFile.open(fakefile_fn)
    fakefile_linked = \
        dexfile.DEXFile_linked. \
        build_dependencies(fakefile_dex, all_dep_decls)
    
    fakefile_def = fakefile_linked. \
        find_class_name('Lcom/codemud/fakefile/fakefile;')

    fakefile_data = fakefile_def.classDataOffRef.value
    assert len(fakefile_data.directMethods.items) == 1
    assert len(fakefile_data.virtualMethods.items) == 0
    fakefile_dataheader = fakefile_data.header
    assert fakefile_dataheader.directMethodsSize == 1
    assert fakefile_dataheader.virtualMethodsSize == 0

    clone = dexfile_insert_class(helloworld_linked, fakefile_def)
    assert clone
    assert clone != fakefile_def

    helloworld_unlinked = helloworld_linked.get_unlinked()
    assert helloworld_unlinked

    # map size for classdef must be increased by 1
    classdef_map = _find_map(helloworld_unlinked, 0x0006)
    assert classdef_map.size == saved_classdef_map_sz + 1

    classdata_map = _find_map(helloworld_unlinked, 0x2000)
    assert classdata_map.size == classdef_map.size

    # Check strings
    strdatas = helloworld_unlinked.stringDataItems.items
    strs = sorted([strdata.data.data for strdata in strdatas])
    assert len(strs) == len(set(strs)) # uniquely
    assert 'Lcom/codemud/fakefile/fakefile;' in strs

    # Check Method List
    methodids_map = _find_map(helloworld_unlinked, 0x0005) # method ids
    assert methodids_map.size == len(helloworld_unlinked.methodIds.items)
    assert methodids_map.size == saved_methodids_sz + 1

    # Check Code item List
    codeitems_map = _find_map(helloworld_unlinked, 0x2001)
    assert codeitems_map.size == saved_codeitems_sz + 1
    pass


def redirect_types_test():
    from paraspace.dex_deptracker import prepare_dep_decls
    from paraspace.injection import dexfile_insert_class
    from paraspace.injection import dexfile_redirect_types
    from paraspace.injection import make_methods_redir_for_types_redir
    from paraspace import dalvik_opcodes
    
    _install_dexfile_4_deptracker()
    
    all_dep_decls = prepare_dep_decls()
    
    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    
    helloworld_fn = os.path.join(srcroot, 'data', 'helloworld.dex')
    helloworld_dex = dexfile.DEXFile.open(helloworld_fn)

    classdef_map = _find_map(helloworld_dex, 0x0006)
    saved_classdef_map_sz = classdef_map.size
    saved_methodids_sz = len(helloworld_dex.methodIds.items)
    
    codeitems_map = _find_map(helloworld_dex, 0x2001)
    saved_codeitems_sz = codeitems_map.size

    helloworld_linked = \
        dexfile.DEXFile_linked.build_dependencies(helloworld_dex,
                                                  all_dep_decls)

    fakefile_fn = os.path.join(srcroot, 'data', 'fakefile.dex')
    fakefile_dex = dexfile.DEXFile.open(fakefile_fn)
    fakefile_linked = \
        dexfile.DEXFile_linked. \
        build_dependencies(fakefile_dex, all_dep_decls)
    
    fakefile_def = \
        fakefile_linked.find_class_name('Lcom/codemud/fakefile/fakefile;')
    clone = dexfile_insert_class(helloworld_linked, fakefile_def)

    init_method = fakefile_linked.find_method_name('<init>', fakefile_def)
    init_protoid = init_method.methodIdx.protoIdx

    File_typeid = helloworld_linked.find_typeid_name('Ljava/io/File;')
    assert File_typeid
    File_typeidx = helloworld_linked.get_idx_typeid(File_typeid)

    clone_typeidx = helloworld_linked.get_idx_typeid(clone.classIdx)

    types_redir = {File_typeidx: clone_typeidx}
    excludes = set([clone_typeidx])
    
    fakefile_data = clone.classDataOffRef.value
    fakefile_methods = fakefile_data.directMethods.items
    fakefile_codes = [method.codeOffRef.value for method in fakefile_methods]
    
    methods_redir = make_methods_redir_for_types_redir(helloworld_linked,
                                                       helloworld_linked,
                                                       types_redir)
    dexfile_redirect_types(helloworld_linked, types_redir,
                           methods_redir, excludes)

    for code in helloworld_linked.codeItems.items:
        op_vectors = dalvik_opcodes.decode_insn_blk(code.insns.data)
        for opcode, args in op_vectors:
            if opcode == dalvik_opcodes.all_opcodes.OP_NEW_INSTANCE:
                assert args[1] != File_typeidx
                pass
            elif opcode == dalvik_opcodes.all_opcodes.OP_INVOKE_DIRECT:
                methodid = helloworld_linked.find_methodid_idx(args[2])
                assert methodid.classIdx != File_typeid or \
                    code in fakefile_codes
                pass
            pass
        pass
    pass


def collect_types_in_method_test():
    from paraspace.dex_deptracker import prepare_dep_decls
    from paraspace.injection import collect_typeidxs_in_method
    from paraspace.injection import collect_typeidxs_mentioned_by_class
    from paraspace.dexfile import DEXFile_linked
    
    _install_dexfile_4_deptracker()
    
    all_dep_decls = prepare_dep_decls()
    
    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    
    fakefile_fn = os.path.join(srcroot, 'data', 'fakefile.dex')
    fakefile_dex = dexfile.DEXFile.open(fakefile_fn)
    fakefile_linked = \
        dexfile.DEXFile_linked. \
        build_dependencies(fakefile_dex, all_dep_decls)
    
    fakefile_def = \
        fakefile_linked.find_class_name('Lcom/codemud/fakefile/fakefile;')

    init_method = fakefile_linked.find_method_name('<init>', fakefile_def)

    typeidxs = collect_typeidxs_in_method(fakefile_linked, init_method)
    assert len(typeidxs) == 1

    typeid = fakefile_linked.find_typeid_idx(typeidxs[0])
    typeid_name = DEXFile_linked.get_typeid_name(typeid)
    assert typeid_name == 'Ljava/io/File;'
    
    typeidxs = collect_typeidxs_mentioned_by_class(fakefile_linked,
                                                   fakefile_def)
    assert len(typeidxs) == 1

    typeid = fakefile_linked.find_typeid_idx(typeidxs[0])
    typeid_name = DEXFile_linked.get_typeid_name(typeid)
    assert typeid_name == 'Ljava/io/File;'
    pass


def dexfile_insert_class_relative_test():
    from paraspace.dex_deptracker import prepare_dep_decls
    from paraspace.injection import dexfile_insert_class_relative
    
    _install_dexfile_4_deptracker()
    
    all_dep_decls = prepare_dep_decls()
    
    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    
    helloworld_fn = os.path.join(srcroot, 'data', 'helloworld.dex')
    helloworld_dex = dexfile.DEXFile.open(helloworld_fn)

    classdef_map = _find_map(helloworld_dex, 0x0006)
    saved_classdef_map_sz = classdef_map.size
    saved_methodids_sz = len(helloworld_dex.methodIds.items)
    
    codeitems_map = _find_map(helloworld_dex, 0x2001)
    saved_codeitems_sz = codeitems_map.size

    helloworld_linked = \
        dexfile.DEXFile_linked.build_dependencies(helloworld_dex,
                                                  all_dep_decls)

    fakefile_fn = os.path.join(srcroot, 'data', 'fakefile.dex')
    fakefile_dex = dexfile.DEXFile.open(fakefile_fn)
    fakefile_linked = \
        dexfile.DEXFile_linked. \
        build_dependencies(fakefile_dex, all_dep_decls)
    
    fakefile_def = fakefile_linked. \
        find_class_name('Lcom/codemud/fakefile/fakefile;')

    fakefile_data = fakefile_def.classDataOffRef.value
    assert len(fakefile_data.directMethods.items) == 1
    assert len(fakefile_data.virtualMethods.items) == 0
    fakefile_dataheader = fakefile_data.header
    assert fakefile_dataheader.directMethodsSize == 1
    assert fakefile_dataheader.virtualMethodsSize == 0

    clone = dexfile_insert_class_relative(helloworld_linked,
                                          fakefile_linked, fakefile_def)
    assert clone
    assert clone != fakefile_def

    helloworld_unlinked = helloworld_linked.get_unlinked()
    assert helloworld_unlinked

    # map size for classdef must be increased by 1
    classdef_map = _find_map(helloworld_unlinked, 0x0006)
    assert classdef_map.size == saved_classdef_map_sz + 1

    classdata_map = _find_map(helloworld_unlinked, 0x2000)
    assert classdata_map.size == classdef_map.size

    # Check strings
    strdatas = helloworld_unlinked.stringDataItems.items
    strs = sorted([strdata.data.data for strdata in strdatas])
    assert len(strs) == len(set(strs)) # uniquely
    assert 'Lcom/codemud/fakefile/fakefile;' in strs

    # Check Method List
    methodids_map = _find_map(helloworld_unlinked, 0x0005) # method ids
    assert methodids_map.size == len(helloworld_unlinked.methodIds.items)
    assert methodids_map.size == saved_methodids_sz + 1

    # Check Code item List
    codeitems_map = _find_map(helloworld_unlinked, 0x2001)
    assert codeitems_map.size == saved_codeitems_sz + 1
    pass