got simple method call graph working!

5 files changed, 284 insertions, 51 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index c35a483..33eef08 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -15,12 +15,9 @@ import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.StringWriter;
-import java.util.Enumeration;
 import java.util.List;
-import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 
-import com.google.common.collect.Lists;
 import com.strobel.assembler.metadata.MetadataSystem;
 import com.strobel.assembler.metadata.TypeDefinition;
 import com.strobel.decompiler.DecompilerContext;
@@ -53,7 +50,6 @@ public class Deobfuscator
         private JarIndex m_jarIndex;
         private Mappings m_mappings;
         private Renamer m_renamer;
-        private List<String> m_obfClassNames;
         
         public Deobfuscator( File file )
         throws IOException
@@ -65,7 +61,7 @@ public class Deobfuscator
                 InputStream jarIn = null;
                 try
                 {
-                        m_jarIndex = new JarIndex();
+                        m_jarIndex = new JarIndex( m_jar );
                         jarIn = new FileInputStream( m_file );
                         m_jarIndex.indexJar( jarIn );
                 }
@@ -74,26 +70,6 @@ public class Deobfuscator
                         Util.closeQuietly( jarIn );
                 }
                 
-                // get the obf class names
-                m_obfClassNames = Lists.newArrayList();
-                {
-                        Enumeration<JarEntry> entries = m_jar.entries();
-                        while( entries.hasMoreElements() )
-                        {
-                                JarEntry entry = entries.nextElement();
-                                
-                                // skip everything but class files
-                                if( !entry.getName().endsWith( ".class" ) )
-                                {
-                                        continue;
-                                }
-                                
-                                // get the class name from the file
-                                String className = entry.getName().substring( 0, entry.getName().length() - 6 );
-                                m_obfClassNames.add( className );
-                        }
-                }
-                
                 // config the decompiler
                 m_settings = DecompilerSettings.javaDefaults();
                 m_settings.setShowSyntheticMembers( true );
@@ -140,7 +116,7 @@ public class Deobfuscator
         
         public void getSeparatedClasses( List<String> obfClasses, List<String> deobfClasses )
         {
-                for( String obfClassName : m_obfClassNames )
+                for( String obfClassName : m_jarIndex.getObfClassNames() )
                 {
                         // separate the classes
                         ClassMapping classMapping = m_mappings.getClassByObf( obfClassName );
@@ -302,15 +278,15 @@ public class Deobfuscator
                 if( obfEntry instanceof ClassEntry )
                 {
                         // obf classes must be in the list
-                        return m_obfClassNames.contains( obfEntry.getName() );
+                        return m_jarIndex.getObfClassNames().contains( obfEntry.getName() );
                 }
                 else if( obfEntry instanceof FieldEntry )
                 {
-                        return m_obfClassNames.contains( ((FieldEntry)obfEntry).getClassName() );
+                        return m_jarIndex.getObfClassNames().contains( ((FieldEntry)obfEntry).getClassName() );
                 }
                 else if( obfEntry instanceof MethodEntry )
                 {
-                        return m_obfClassNames.contains( ((MethodEntry)obfEntry).getClassName() );
+                        return m_jarIndex.getObfClassNames().contains( ((MethodEntry)obfEntry).getClassName() );
                 }
                 else if( obfEntry instanceof ArgumentEntry )
                 {
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 );
+                        }
+                }
+        }
+}
diff --git a/src/cuchaz/enigma/gui/Gui.java b/src/cuchaz/enigma/gui/Gui.java
index 3e8a18c..072fb3a 100644
--- a/src/cuchaz/enigma/gui/Gui.java
+++ b/src/cuchaz/enigma/gui/Gui.java
@@ -67,6 +67,7 @@ import com.google.common.collect.Lists;
 
 import cuchaz.enigma.Constants;
 import cuchaz.enigma.analysis.ClassInheritanceTreeNode;
+import cuchaz.enigma.analysis.MethodCallsTreeNode;
 import cuchaz.enigma.analysis.MethodInheritanceTreeNode;
 import cuchaz.enigma.analysis.Token;
 import cuchaz.enigma.mapping.ArgumentEntry;
@@ -139,6 +140,8 @@ public class Gui
         private BoxHighlightPainter m_obfuscatedHighlightPainter;
         private BoxHighlightPainter m_deobfuscatedHighlightPainter;
         private JTree m_inheritanceTree;
+        private JTree m_callsTree;
+        private JTabbedPane m_tabs;
         
         // dynamic menu items
         private JMenuItem m_closeJarMenu;
@@ -147,9 +150,10 @@ public class Gui
         private JMenuItem m_saveMappingsAsMenu;
         private JMenuItem m_closeMappingsMenu;
         private JMenuItem m_renameMenu;
-        private JMenuItem m_inheritanceMenu;
+        private JMenuItem m_showInheritanceMenu;
         private JMenuItem m_openEntryMenu;
         private JMenuItem m_openPreviousMenu;
+        private JMenuItem m_showCallsMenu;
         
         // state
         private EntryPair<Entry> m_selectedEntryPair;
@@ -267,6 +271,10 @@ public class Gui
                                         case KeyEvent.VK_P:
                                                 m_controller.openPreviousEntry();
                                         break;
+                                        
+                                        case KeyEvent.VK_C:
+                                                showCalls();
+                                        break;
                                 }
                         }
                 } );
@@ -306,7 +314,22 @@ public class Gui
                         menu.setAccelerator( KeyStroke.getKeyStroke( KeyEvent.VK_I, 0 ) );
                         menu.setEnabled( false );
                         popupMenu.add( menu );
-                        m_inheritanceMenu = menu;
+                        m_showInheritanceMenu = menu;
+                }
+                {
+                        JMenuItem menu = new JMenuItem( "Show Calls" );
+                        menu.addActionListener( new ActionListener( )
+                        {
+                                @Override
+                                public void actionPerformed( ActionEvent event )
+                                {
+                                        showCalls();
+                                }
+                        } );
+                        menu.setAccelerator( KeyStroke.getKeyStroke( KeyEvent.VK_C, 0 ) );
+                        menu.setEnabled( false );
+                        popupMenu.add( menu );
+                        m_showCallsMenu = menu;
                 }
                 {
                         JMenuItem menu = new JMenuItem( "Go to Declaration" );
@@ -350,12 +373,13 @@ public class Gui
                                 if( event.getClickCount() == 2 )
                                 {
                                         // get the selected node
-                                        Object node = m_inheritanceTree.getSelectionPath().getLastPathComponent();
-                                        if( node == null )
+                                        TreePath path = m_inheritanceTree.getSelectionPath();
+                                        if( path == null )
                                         {
                                                 return;
                                         }
                                         
+                                        Object node = path.getLastPathComponent();
                                         if( node instanceof ClassInheritanceTreeNode )
                                         {
                                                 m_controller.openEntry( new ClassEntry( ((ClassInheritanceTreeNode)node).getObfClassName() ) );
@@ -375,17 +399,47 @@ public class Gui
                 inheritancePanel.setLayout( new BorderLayout() );
                 inheritancePanel.add( new JScrollPane( m_inheritanceTree ) );
                 
+                // init call panel
+                m_callsTree = new JTree();
+                m_callsTree.setModel( null );
+                m_callsTree.addMouseListener( new MouseAdapter( )
+                {
+                        @Override
+                        public void mouseClicked( MouseEvent event )
+                        {
+                                if( event.getClickCount() == 2 )
+                                {
+                                        // get the selected node
+                                        TreePath path = m_callsTree.getSelectionPath();
+                                        if( path == null )
+                                        {
+                                                return;
+                                        }
+                                        
+                                        Object node = path.getLastPathComponent();
+                                        if( node instanceof MethodCallsTreeNode )
+                                        {
+                                                m_controller.openEntry( ((MethodCallsTreeNode)node).getMethodEntry() );
+                                        }
+                                }
+                        }
+                } );
+                JPanel callPanel = new JPanel();
+                callPanel.setLayout( new BorderLayout() );
+                callPanel.add( new JScrollPane( m_callsTree ) );
+                
                 // layout controls
                 JSplitPane splitLeft = new JSplitPane( JSplitPane.VERTICAL_SPLIT, true, obfPanel, deobfPanel );
-                splitLeft.setPreferredSize( new Dimension( 200, 0 ) );
+                splitLeft.setPreferredSize( new Dimension( 250, 0 ) );
                 JPanel centerPanel = new JPanel();
                 centerPanel.setLayout( new BorderLayout() );
                 centerPanel.add( m_infoPanel, BorderLayout.NORTH );
                 centerPanel.add( sourceScroller, BorderLayout.CENTER );
-                JTabbedPane tabbedPane = new JTabbedPane();
-                tabbedPane.setPreferredSize( new Dimension( 200, 0 ) );
-                tabbedPane.addTab( "Inheritance", inheritancePanel );
-                JSplitPane splitRight = new JSplitPane( JSplitPane.HORIZONTAL_SPLIT, true, centerPanel, tabbedPane );
+                m_tabs = new JTabbedPane();
+                m_tabs.setPreferredSize( new Dimension( 250, 0 ) );
+                m_tabs.addTab( "Inheritance", inheritancePanel );
+                m_tabs.addTab( "Method Calls", callPanel );
+                JSplitPane splitRight = new JSplitPane( JSplitPane.HORIZONTAL_SPLIT, true, centerPanel, m_tabs );
                 splitRight.setResizeWeight( 1 ); // let the left side take all the slack
                 splitRight.resetToPreferredSizes();
                 JSplitPane splitCenter = new JSplitPane( JSplitPane.HORIZONTAL_SPLIT, true, splitLeft, splitRight );
@@ -784,7 +838,8 @@ public class Gui
                 
                 showEntryPair( m_selectedEntryPair );
                 
-                m_inheritanceMenu.setEnabled( isClassEntry || isMethodEntry );
+                m_showInheritanceMenu.setEnabled( isClassEntry || isMethodEntry );
+                m_showCallsMenu.setEnabled( isMethodEntry );
                 m_openEntryMenu.setEnabled( isClassEntry || isFieldEntry || isMethodEntry );
                 m_openPreviousMenu.setEnabled( m_controller.hasPreviousEntry() );
         }
@@ -879,6 +934,24 @@ public class Gui
                         m_inheritanceTree.setSelectionRow( m_inheritanceTree.getRowForPath( path ) );
                 }
                 
+                m_tabs.setSelectedIndex( 0 );
+                redraw();
+        }
+        
+        private void showCalls( )
+        {
+                if( m_selectedEntryPair == null )
+                {
+                        return;
+                }
+                
+                if( m_selectedEntryPair.obf instanceof MethodEntry )
+                {
+                        MethodCallsTreeNode node = m_controller.getMethodCalls( (MethodEntry)m_selectedEntryPair.obf );
+                        m_callsTree.setModel( new DefaultTreeModel( node ) );
+                }
+                
+                m_tabs.setSelectedIndex( 1 );
                 redraw();
         }
         
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index 880f001..b54aeba 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -21,6 +21,7 @@ import com.google.common.collect.Lists;
 
 import cuchaz.enigma.Deobfuscator;
 import cuchaz.enigma.analysis.ClassInheritanceTreeNode;
+import cuchaz.enigma.analysis.MethodCallsTreeNode;
 import cuchaz.enigma.analysis.MethodInheritanceTreeNode;
 import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.analysis.Token;
@@ -148,6 +149,16 @@ public class GuiController
                 return MethodInheritanceTreeNode.findNode( rootNode, obfMethodEntry );
         }
         
+        public MethodCallsTreeNode getMethodCalls( MethodEntry obfMethodEntry )
+        {
+                MethodCallsTreeNode rootNode = new MethodCallsTreeNode(
+                        m_deobfuscator.getTranslator( TranslationDirection.Deobfuscating ),
+                        obfMethodEntry
+                );
+                rootNode.load( m_deobfuscator.getJarIndex(), true );
+                return rootNode;
+        }
+        
         public void rename( Entry obfEntry, String newName )
         {
                 m_deobfuscator.rename( obfEntry, newName );