got simple method call graph working!

2 files changed, 185 insertions, 12 deletions

diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index 8a8384c..06b0173 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -14,19 +14,28 @@ import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Collection;
+import java.util.Enumeration;
 import java.util.List;
+import java.util.Set;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.Multimap;
-
 import javassist.ByteArrayClassPath;
+import javassist.CannotCompileException;
 import javassist.ClassPool;
+import javassist.CtBehavior;
 import javassist.CtClass;
 import javassist.NotFoundException;
 import javassist.bytecode.Descriptor;
-import javassist.bytecode.MethodInfo;
+import javassist.expr.ExprEditor;
+import javassist.expr.MethodCall;
+
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Sets;
+
 import cuchaz.enigma.Constants;
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.MethodEntry;
@@ -34,16 +43,35 @@ import cuchaz.enigma.mapping.Translator;
 
 public class JarIndex
 {
+        private Set<String> m_obfClassNames;
         private Ancestries m_ancestries;
-        private Multimap<String,String> m_methodImplementations;
+        private Multimap<String,MethodEntry> m_methodImplementations;
+        private Multimap<MethodEntry,MethodEntry> m_methodCalls;
         
-        public JarIndex( )
+        public JarIndex( JarFile jar )
         {
+                m_obfClassNames = Sets.newHashSet();
                 m_ancestries = new Ancestries();
                 m_methodImplementations = HashMultimap.create();
+                m_methodCalls = HashMultimap.create();
+                
+                // read the class names
+                Enumeration<JarEntry> enumeration = jar.entries();
+                while( enumeration.hasMoreElements() )
+                {
+                        JarEntry entry = enumeration.nextElement();
+                        
+                        // filter out non-classes
+                        if( entry.isDirectory() || !entry.getName().endsWith( ".class" ) )
+                        {
+                                continue;
+                        }
+                        
+                        String className = entry.getName().substring( 0, entry.getName().length() - 6 );
+                        m_obfClassNames.add( Descriptor.toJvmName( className ) );
+                }
         }
         
-        @SuppressWarnings( "unchecked" )
         public void indexJar( InputStream in )
         throws IOException
         {
@@ -89,7 +117,10 @@ public class JarIndex
                         {
                                 CtClass c = classPool.get( className );
                                 m_ancestries.addSuperclass( c.getName(), c.getClassFile().getSuperclass() );
-                                addMethodImplementations( c.getName(), (List<MethodInfo>)c.getClassFile().getMethods() );
+                                for( CtBehavior behavior : c.getDeclaredBehaviors() )
+                                {
+                                        indexBehavior( behavior );
+                                }
                         }
                         catch( NotFoundException ex )
                         {
@@ -98,14 +129,55 @@ public class JarIndex
                 }
         }
         
-        private void addMethodImplementations( String name, List<MethodInfo> methods )
+        private void indexBehavior( CtBehavior behavior )
         {
-                for( MethodInfo method : methods )
+                // get the method entry
+                String className = Descriptor.toJvmName( behavior.getDeclaringClass().getName() );
+                final MethodEntry methodEntry = new MethodEntry(
+                        new ClassEntry( className ),
+                        behavior.getName(),
+                        behavior.getSignature()
+                );
+                
+                // index implementation
+                m_methodImplementations.put( className, methodEntry );
+                
+                // index method calls
+                try
+                {
+                        behavior.instrument( new ExprEditor( )
+                        {
+                                @Override
+                                public void edit( MethodCall call )
+                                {
+                                        // is this a jar class?
+                                        String className = Descriptor.toJvmName( call.getClassName() );
+                                        if( !m_obfClassNames.contains( className ) )
+                                        {
+                                                return;
+                                        }
+                                        
+                                        // make entry for the called method
+                                        MethodEntry calledMethodEntry = new MethodEntry(
+                                                new ClassEntry( className ),
+                                                call.getMethodName(),
+                                                call.getSignature()
+                                        );
+                                        m_methodCalls.put( calledMethodEntry, methodEntry );
+                                }
+                        } );
+                }
+                catch( CannotCompileException ex )
                 {
-                        m_methodImplementations.put( name, getMethodKey( method.getName(), method.getDescriptor() ) );
+                        throw new Error( ex );
                 }
         }
         
+        public Set<String> getObfClassNames( )
+        {
+                return m_obfClassNames;
+        }
+        
         public Ancestries getAncestries( )
         {
                 return m_ancestries;
@@ -118,7 +190,7 @@ public class JarIndex
         
         public boolean isMethodImplemented( String className, String methodName, String methodSignature )
         {
-                Collection<String> implementations = m_methodImplementations.get( className );
+                Collection<MethodEntry> implementations = m_methodImplementations.get( className );
                 if( implementations == null )
                 {
                         return false;
@@ -169,6 +241,11 @@ public class JarIndex
                 return rootNode;
         }
         
+        public Collection<MethodEntry> getMethodCallers( MethodEntry methodEntry )
+        {
+                return m_methodCalls.get( methodEntry );
+        }
+        
         private String getMethodKey( String name, String signature )
         {
                 return name + signature;
diff --git a/src/cuchaz/enigma/analysis/MethodCallsTreeNode.java b/src/cuchaz/enigma/analysis/MethodCallsTreeNode.java
new file mode 100644
index 0000000..dedfb2e
--- /dev/null
+++ b/src/cuchaz/enigma/analysis/MethodCallsTreeNode.java
@@ -0,0 +1,96 @@
+/*******************************************************************************
+ * Copyright (c) 2014 Jeff Martin.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the GNU Public License v3.0
+ * which accompanies this distribution, and is available at
+ * http://www.gnu.org/licenses/gpl.html
+ * 
+ * Contributors:
+ *     Jeff Martin - initial API and implementation
+ ******************************************************************************/
+package cuchaz.enigma.analysis;
+
+import java.util.List;
+
+import javax.swing.tree.DefaultMutableTreeNode;
+
+import com.google.common.collect.Lists;
+
+import cuchaz.enigma.mapping.MethodEntry;
+import cuchaz.enigma.mapping.Translator;
+
+public class MethodCallsTreeNode extends DefaultMutableTreeNode
+{
+        private static final long serialVersionUID = -3658163700783307520L;
+        
+        private Translator m_deobfuscatingTranslator;
+        private MethodEntry m_entry;
+        
+        public MethodCallsTreeNode( Translator deobfuscatingTranslator, MethodEntry entry )
+        {
+                m_deobfuscatingTranslator = deobfuscatingTranslator;
+                m_entry = entry;
+        }
+        
+        public MethodEntry getMethodEntry( )
+        {
+                return m_entry;
+        }
+        
+        public String getDeobfClassName( )
+        {
+                return m_deobfuscatingTranslator.translateClass( m_entry.getClassName() );
+        }
+        
+        public String getDeobfMethodName( )
+        {
+                return m_deobfuscatingTranslator.translate( m_entry );
+        }
+        
+        @Override
+        public String toString( )
+        {
+                String className = getDeobfClassName();
+                if( className == null )
+                {
+                        className = m_entry.getClassName();
+                }
+                
+                String methodName = getDeobfMethodName();
+                if( methodName == null )
+                {
+                        methodName = m_entry.getName();
+                }
+                return className + "." + methodName + "()";
+        }
+        
+        public void load( JarIndex index, boolean recurse )
+        {
+                // get all the child nodes
+                List<MethodCallsTreeNode> nodes = Lists.newArrayList();
+                for( MethodEntry entry : index.getMethodCallers( m_entry ) )
+                {
+                        nodes.add( new MethodCallsTreeNode( m_deobfuscatingTranslator, entry ) );
+                }
+                
+                // add them to this node
+                for( MethodCallsTreeNode node : nodes )
+                {
+                        this.add( node );
+                }
+                
+                if( recurse )
+                {
+                        for( MethodCallsTreeNode node : nodes )
+                        {
+                                // don't recurse into self
+                                if( node.getMethodEntry().equals( m_entry ) )
+                                {
+                                        continue;
+                                }
+                                
+                                node.load( index, true );
+                        }
+                }
+        }
+}