--- a/paraspace/injection.py	Tue Aug 02 20:40:48 2011 +0800
+++ b/paraspace/injection.py	Thu Aug 04 15:55:01 2011 +0800
@@ -261,10 +261,9 @@
 
 
 def dexfile_insert_classdefs(dex_dst, dex_src, classdefs):
-    for classdef in classdefs:
-        dexfile_insert_class(dex_dst, classdef)
-        pass
-    pass
+    clones = [dexfile_insert_class(dex_dst, classdef)
+              for classdef in classdefs]
+    return clones
 
 
 ## \brief Clone and insert a _DEX_TypeId to another DEXFile_linked.
@@ -291,10 +290,9 @@
 
 ## \brief Clone and insert a list of _DEX_TypeId objects to a DEXFile_linked.
 def dexfile_insert_typeids(dex_dst, dex_src, typeids):
-    for typeid in typeids:
-        dexfile_insert_typeid(dex_dst, dex_src, typeid)
-        pass
-    pass
+    clones = [dexfile_insert_typeid(dex_dst, dex_src, typeid)
+              for typeid in typeids]
+    return clones
 
 
 ## \brief Collects relative type IDs and classes definition for given class.
@@ -320,6 +318,7 @@
 # \param dex_dst is a DEXFile_linked where the class will be inserted.
 # \param dex_src is a DEXFile_linked where the cloned class is from.
 # \param classdef is a _DEX_ClassDef that will be cloned.
+# \return a vector of a list of classdefs and a list of typeids.
 #
 def dexfile_insert_classdefs_relative(dex_dst, dex_src, classdefs):
     from paraspace.dexfile import DEXFile_linked
@@ -343,22 +342,24 @@
     relative_classdefs, relative_typeids = \
         collect_classdefs_relative(dex_src, classdefs)
     
-    inserting_classdefs = filter(classdef_not_in_dst, relative_classdefs)
+    for classdef in relative_classdefs:
+        if classdef_not_in_dst(classdef):
+            continue
+        raise ValueError, '%s is already in DEX %s: can not insert it' % \
+            (repr(classdef), repr(dex_dst))
+    
     inserting_typeids = filter(typeid_not_in_dst, relative_typeids)
     
-    dexfile_insert_classdefs(dex_dst, dex_src, inserting_classdefs)
-    dexfile_insert_typeids(dex_dst, dex_src, inserting_typeids)
+    cloning_classdefs = \
+        dexfile_insert_classdefs(dex_dst, dex_src, relative_classdefs)
+    cloning_typeids = \
+        dexfile_insert_typeids(dex_dst, dex_src, inserting_typeids)
 
-    classnames = [DEXFile_linked.get_classdef_name(classdef)
-                  for classdef in classdefs]
-    clonings = [dex_dst.find_class_name(classname)
-                for classname in classnames]
-    
-    return clonings
+    return cloning_classdefs, cloning_typeids
 
 
 ## \brief Redirect types and methods for the code of given method.
-def method_redirect_types(dex, method, types_redir, methods_redir):
+def method_redirect_typeidxs(dex, method, typeidxs_redir, methods_redir):
     from paraspace.dalvik_opcodes import decode_insn_blk, all_opcodes
     from paraspace.dalvik_opcodes import encode_opcode_vectors
     from paraspace.dexfile import DEXFile_linked
@@ -373,17 +374,31 @@
     def redirect(opcode, args):
         if opcode == all_opcodes.OP_NEW_INSTANCE:
             typeidx = args[1]
-            if typeidx in types_redir:
-                to_type = types_redir[typeidx]
+            if typeidx in typeidxs_redir:
+                to_type = typeidxs_redir[typeidx]
                 return opcode, (args[0], to_type)
             pass
-        elif opcode == all_opcodes.OP_INVOKE_DIRECT:
+        elif opcode in (all_opcodes.OP_INVOKE_DIRECT,
+                        all_opcodes.OP_INVOKE_VIRTUAL,
+                        all_opcodes.OP_INVOKE_SUPER,
+                        all_opcodes.OP_INVOKE_STATIC,
+                        all_opcodes.OP_INVOKE_INTERFACE):
             methodidx = args[2]
             if methodidx not in methods_redir:
                 return opcode, args
             
             return opcode, (args[0], args[1], methods_redir[methodidx],
                             args[3], args[4], args[5], args[6])
+        elif opcode in (all_opcodes.OP_INVOKE_VIRTUAL_RANGE,
+                        all_opcodes.OP_INVOKE_DIRECT_RANGE,
+                        all_opcodes.OP_INVOKE_SUPER_RANGE,
+                        all_opcodes.OP_INVOKE_STATIC_RANGE,
+                        all_opcodes.OP_INVOKE_INTERFACE_RANGE):
+            methodidx = args[1]
+            if methodidx not in methods_redir:
+                return opcode, args
+            
+            return opcode, (args[0], methodidx, args[2])
         return opcode, args
     
     new_op_vectors = [redirect(opcode, args) for opcode, args in op_vectors]
@@ -425,24 +440,24 @@
 ## \brief Redirect types and methods mentioned in the code of a class.
 #
 # For code of given class definition, Every mentions of types and
-# methods are rewrote to types and methods according types_redir and
+# methods are rewrote to types and methods according typeidxs_redir and
 # methods_redir respectively.
 #
 # \param dex is a DEXFile_linked.
 # \param classdef is a class definition.
-# \param types_redir is a map of types.
+# \param typeidxs_redir is a map of types.
 # \param methods_redir is a map of methods.
 #
-def class_redirect_types(dex, classdef, types_redir, methods_redir):
+def class_redirect_types(dex, classdef, typeidxs_redir, methods_redir):
     if not classdef.classDataOffRef.is_true:
         return
     
     classdata = classdef.classDataOffRef.value
     for method in classdata.directMethods.items:
-        method_redirect_types(dex, method, types_redir, methods_redir)
+        method_redirect_typeidxs(dex, method, typeidxs_redir, methods_redir)
         pass
     for method in classdata.virtualMethods.items:
-        method_redirect_types(dex, method, types_redir, methods_redir)
+        method_redirect_typeidxs(dex, method, typeidxs_redir, methods_redir)
         pass
     pass
 
@@ -450,15 +465,15 @@
 ## \brief Make a map to map methods from source types to destinate types.
 #
 # This function create a map to map methods from source types to
-# methods from destinate types in \ref types_redir.
+# methods from destinate types in \ref typeidxs_redir.
 #
 # \param dex is a DEXFile_linked that owns source and destinate types.
-# \param types_redir is a map of types for redirecting types.
+# \param typeidxs_redir is a map of type indices for redirecting types.
 # \return a map of method indices.
 #
-def make_methods_redir_for_types_redir(dex_src, dex_dst, types_redir):
+def make_methodidxs_redir_map(dex_src, dex_dst, typeidxs_redir):
     methods_map = {}
-    for typeidx_src, typeidx_dst in types_redir.items():
+    for typeidx_src, typeidx_dst in typeidxs_redir.items():
         typeid_src = dex_src.find_typeid_idx(typeidx_src)
         typeid_dst = dex_dst.find_typeid_idx(typeidx_dst)
         class_methods_map =  make_redir_classes_methods_map(dex_src,
@@ -471,13 +486,28 @@
 
 
 ## \biref Redirect types of all code in given DEXFile_linked.
-def dexfile_redirect_types(dex, types_redir, methods_redir, excludes=set([])):
+def dexfile_redirect_types(dex, typeidxs_redir, methods_redir,
+                           excludes=set([])):
     for classdef in dex.classDefs.items:
         typeid = classdef.classIdx
         idx = dex.get_idx_typeid(typeid)
         if idx in excludes:
             continue
-        class_redirect_types(dex, classdef, types_redir, methods_redir)
+        class_redirect_types(dex, classdef, typeidxs_redir, methods_redir)
+        pass
+    pass
+
+
+## \brief Redirect types for code of types specified by given indices.
+def dexfile_redirect_types_typeidxs(dex, typeidxs_redir, methodidxs_redir,
+                                    typeidxs):
+    typeidxs = set(typeidxs)
+    for classdef in dex.classDefs.items:
+        typeid = classdef.classIdx
+        idx = dex.get_idx_typeid(typeid)
+        if idx not in typeidxs:
+            continue
+        class_redirect_types(dex, classdef, typeidxs_redir, methodidxs_redir)
         pass
     pass
 
@@ -549,3 +579,68 @@
         pass
     
     return list(typeidxs)
+
+
+## \brief Make a mapping for type indices of injection.
+def make_typeidxs_map_after_injection(dex_dst, dex_src,
+                                      relative_classdefs,
+                                      relative_typeids):
+    from paraspace.dexfile import DEXFile_linked
+    
+    def map_src_dst_typeid(typeid):
+        idx_src = dex_src.get_idx_typeid(typeid)
+        typename = DEXFile_linked.get_typeid_name(typeid)
+        typeid_dst = dex_dst.find_typeid_name(typename)
+        idx_dst = dex_dst.get_idx_typeid(typeid_dst)
+        dex_dst.find_typeid_idx(idx_dst)
+        return idx_src, idx_dst
+
+    def map_src_dst_classdef(classdef):
+        typeid = classdef.classIdx
+        idx_src, idx_dst = map_src_dst_typeid(typeid)
+        return idx_src, idx_dst
+    
+    typeidxs_classdefs = [map_src_dst_classdef(classdef)
+                          for classdef in relative_classdefs]
+    typeidxs_typeids = [map_src_dst_typeid(typeid)
+                        for typeid in relative_typeids]
+    typeidxs_map = dict(typeidxs_classdefs + typeidxs_typeids)
+    return typeidxs_map
+
+
+## \brief Redirect code for methods of injected classes.
+def redirect_injected_code(dex_dst, dex_src, classdefs):
+    relative_classdefs, relative_typeids = \
+        collect_classdefs_relative(dex_src, classdefs)
+    
+    typeidxs_redir = \
+        make_typeidxs_map_after_injection(dex_dst, dex_src, \
+                                              relative_classdefs, \
+                                              relative_typeids)
+    methodidxs_redir = \
+        make_methodidxs_redir_map(dex_src, dex_dst, typeidxs_redir)
+    
+    dexfile_redirect_types_typeidxs(dex_dst, typeidxs_redir,
+                                    methodidxs_redir,
+                                    typeidxs_redir.values())
+    pass
+
+
+## \brief Inject classes and relative information to a DEX file.
+#
+# \param dex_dst is a DEXFile_linked where to inject classes.
+# \param dex_src is a DEXFile_linked where the class is from.
+# \param classdefs is a list of _DEX_ClassDef instances from \ref dex_src.
+# \return a list of _DEX_ClassDef instances for injection.
+#
+def inject_classdefs(dex_dst, dex_src, classdefs):
+    from paraspace.dexfile import DEXFile_linked
+    
+    assert isinstance(classdefs, (list, tuple))
+    assert isinstance(dex_dst, DEXFile_linked)
+    assert isinstance(dex_src, DEXFile_linked)
+
+    injected_classdefs, injected_typeids = \
+        dexfile_insert_classdefs_relative(dex_dst, dex_src, classdefs)
+    redirect_injected_code(dex_dst, dex_src, classdefs)
+    return injected_classdefs

--- a/paraspace/tests/injection_test.py	Tue Aug 02 20:40:48 2011 +0800
+++ b/paraspace/tests/injection_test.py	Thu Aug 04 15:55:01 2011 +0800
@@ -107,7 +107,7 @@
     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.injection import make_methodidxs_redir_map
     from paraspace import dalvik_opcodes
     
     _install_dexfile_4_deptracker()
@@ -157,9 +157,9 @@
     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)
+    methods_redir = make_methodidxs_redir_map(helloworld_linked,
+                                              helloworld_linked,
+                                              types_redir)
     dexfile_redirect_types(helloworld_linked, types_redir,
                            methods_redir, excludes)
 
@@ -261,12 +261,12 @@
     assert fakefile_dataheader.directMethodsSize == 1
     assert fakefile_dataheader.virtualMethodsSize == 0
 
-    clones = dexfile_insert_classdefs_relative(helloworld_linked,
-                                               fakefile_linked,
-                                               [fakefile_def])
-    assert clones
-    assert len(clones) == 1
-    assert clones[0] != fakefile_def
+    classdefs, typeids = dexfile_insert_classdefs_relative(helloworld_linked,
+                                                           fakefile_linked,
+                                                           [fakefile_def])
+    assert classdefs
+    assert len(classdefs) == 1
+    assert classdefs[0] != fakefile_def
 
     helloworld_unlinked = helloworld_linked.get_unlinked()
     assert helloworld_unlinked
@@ -293,3 +293,44 @@
     codeitems_map = _find_map(helloworld_unlinked, 0x2001)
     assert codeitems_map.size == saved_codeitems_sz + 1
     pass
+
+
+def inject_classdefs_test():
+    from paraspace.dex_deptracker import prepare_dep_decls
+    from paraspace.injection import inject_classdefs
+    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;')
+    clonings = inject_classdefs(helloworld_linked, fakefile_linked,
+                                [fakefile_def])
+    assert clonings
+    assert len(clonings) == 1
+    pass