added support for field access searches

added proper detection/handling for constructors

4 files changed, 254 insertions, 54 deletions

diff --git a/src/cuchaz/enigma/analysis/FieldCallsTreeNode.java b/src/cuchaz/enigma/analysis/FieldCallsTreeNode.java
new file mode 100644
index 0000000..0427b3b
--- /dev/null
+++ b/src/cuchaz/enigma/analysis/FieldCallsTreeNode.java
@@ -0,0 +1,82 @@
+/*******************************************************************************
+ * 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 javax.swing.tree.DefaultMutableTreeNode;
+
+import cuchaz.enigma.mapping.ConstructorEntry;
+import cuchaz.enigma.mapping.Entry;
+import cuchaz.enigma.mapping.FieldEntry;
+import cuchaz.enigma.mapping.MethodEntry;
+import cuchaz.enigma.mapping.Translator;
+
+public class FieldCallsTreeNode extends DefaultMutableTreeNode
+{
+        private static final long serialVersionUID = -7934108091928699835L;
+        
+        private Translator m_deobfuscatingTranslator;
+        private FieldEntry m_entry;
+        
+        public FieldCallsTreeNode( Translator deobfuscatingTranslator, FieldEntry fieldEntry )
+        {
+                m_deobfuscatingTranslator = deobfuscatingTranslator;
+                m_entry = fieldEntry;
+        }
+        
+        public FieldEntry getFieldEntry( )
+        {
+                return m_entry;
+        }
+        
+        @Override
+        public String toString( )
+        {
+                String className = m_deobfuscatingTranslator.translateClass( m_entry.getClassName() );
+                if( className == null )
+                {
+                        className = m_entry.getClassName();
+                }
+                
+                String targetName = m_deobfuscatingTranslator.translate( m_entry );
+                if( targetName == null )
+                {
+                        targetName = m_entry.getName();
+                }
+                return className + "." + targetName;
+        }
+        
+        public void load( JarIndex index, boolean recurse )
+        {
+                // get all the child nodes
+                for( Entry entry : index.getFieldCallers( m_entry ) )
+                {
+                        if( entry instanceof MethodEntry )
+                        {
+                                add( new MethodCallsTreeNode( m_deobfuscatingTranslator, (MethodEntry)entry ) );
+                        }
+                        else if( entry instanceof ConstructorEntry )
+                        {
+                                add( new MethodCallsTreeNode( m_deobfuscatingTranslator, (ConstructorEntry)entry ) );
+                        }
+                }
+                
+                if( recurse && children != null )
+                {
+                        for( Object node : children )
+                        {
+                                if( node instanceof MethodCallsTreeNode )
+                                {
+                                        ((MethodCallsTreeNode)node).load( index, true );
+                                }
+                        }
+                }
+        }
+}
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index 06b0173..96bddc1 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -27,17 +27,26 @@ import javassist.CannotCompileException;
 import javassist.ClassPool;
 import javassist.CtBehavior;
 import javassist.CtClass;
+import javassist.CtConstructor;
+import javassist.CtMethod;
 import javassist.NotFoundException;
 import javassist.bytecode.Descriptor;
+import javassist.expr.ConstructorCall;
 import javassist.expr.ExprEditor;
+import javassist.expr.FieldAccess;
 import javassist.expr.MethodCall;
+import javassist.expr.NewExpr;
 
 import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Lists;
 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.ConstructorEntry;
+import cuchaz.enigma.mapping.Entry;
+import cuchaz.enigma.mapping.FieldEntry;
 import cuchaz.enigma.mapping.MethodEntry;
 import cuchaz.enigma.mapping.Translator;
 
@@ -46,7 +55,8 @@ public class JarIndex
         private Set<String> m_obfClassNames;
         private Ancestries m_ancestries;
         private Multimap<String,MethodEntry> m_methodImplementations;
-        private Multimap<MethodEntry,MethodEntry> m_methodCalls;
+        private Multimap<Entry,Entry> m_methodCalls;
+        private Multimap<FieldEntry,Entry> m_fieldCalls;
         
         public JarIndex( JarFile jar )
         {
@@ -54,6 +64,7 @@ public class JarIndex
                 m_ancestries = new Ancestries();
                 m_methodImplementations = HashMultimap.create();
                 m_methodCalls = HashMultimap.create();
+                m_fieldCalls = HashMultimap.create();
                 
                 // read the class names
                 Enumeration<JarEntry> enumeration = jar.entries();
@@ -133,14 +144,30 @@ public class JarIndex
         {
                 // 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 );
+                final Entry thisEntry;
+                if( behavior instanceof CtMethod )
+                {
+                        MethodEntry methodEntry = new MethodEntry(
+                                new ClassEntry( className ),
+                                behavior.getName(),
+                                behavior.getSignature()
+                        );
+                        thisEntry = methodEntry;
+                        
+                        // index implementation
+                        m_methodImplementations.put( className, methodEntry );
+                }
+                else if( behavior instanceof CtConstructor )
+                {
+                        thisEntry = new ConstructorEntry(
+                                new ClassEntry( className ),
+                                behavior.getSignature()
+                        );
+                }
+                else
+                {
+                        throw new IllegalArgumentException( "behavior must be a method or a constructor!" );
+                }
                 
                 // index method calls
                 try
@@ -150,20 +177,53 @@ public class JarIndex
                                 @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 );
+                                        m_methodCalls.put( calledMethodEntry, thisEntry );
+                                }
+                                
+                                @Override
+                                public void edit( FieldAccess call )
+                                {
+                                        String className = Descriptor.toJvmName( call.getClassName() );
+                                        FieldEntry calledFieldEntry = new FieldEntry(
+                                                new ClassEntry( className ),
+                                                call.getFieldName()
+                                        );
+                                        m_fieldCalls.put( calledFieldEntry, thisEntry );
+                                }
+                                
+                                @Override
+                                public void edit( ConstructorCall call )
+                                {
+                                        String className = Descriptor.toJvmName( call.getClassName() );
+                                        ConstructorEntry calledConstructorEntry = new ConstructorEntry(
+                                                new ClassEntry( className ),
+                                                call.getSignature()
+                                        );
+                                        m_methodCalls.put( calledConstructorEntry, thisEntry );
+                                }
+                                
+                                @Override
+                                public void edit( NewExpr call )
+                                {
+                                        String className = Descriptor.toJvmName( call.getClassName() );
+                                        ConstructorEntry calledConstructorEntry = new ConstructorEntry(
+                                                new ClassEntry( className ),
+                                                call.getSignature()
+                                        );
+                                        
+                                        // TEMP
+                                        if( className.equals( "bgw" ) )
+                                        {
+                                                System.out.println( calledConstructorEntry + " called by " + thisEntry );
+                                        }
+                                        
+                                        m_methodCalls.put( calledConstructorEntry, thisEntry );
                                 }
                         } );
                 }
@@ -202,7 +262,9 @@ public class JarIndex
         public ClassInheritanceTreeNode getClassInheritance( Translator deobfuscatingTranslator, ClassEntry obfClassEntry )
         {
                 // get the root node
-                List<String> ancestry = m_ancestries.getAncestry( obfClassEntry.getName() );
+                List<String> ancestry = Lists.newArrayList();
+                ancestry.add( obfClassEntry.getName() );
+                ancestry.addAll( m_ancestries.getAncestry( obfClassEntry.getName() ) );
                 ClassInheritanceTreeNode rootNode = new ClassInheritanceTreeNode( deobfuscatingTranslator, ancestry.get( ancestry.size() - 1 ) );
                 
                 // expand all children recursively
@@ -241,9 +303,14 @@ public class JarIndex
                 return rootNode;
         }
         
-        public Collection<MethodEntry> getMethodCallers( MethodEntry methodEntry )
+        public Collection<Entry> getFieldCallers( FieldEntry fieldEntry )
+        {
+                return m_fieldCalls.get( fieldEntry );
+        }
+        
+        public Collection<Entry> getMethodCallers( Entry entry )
         {
-                return m_methodCalls.get( methodEntry );
+                return m_methodCalls.get( entry );
         }
         
         private String getMethodKey( String name, String signature )
diff --git a/src/cuchaz/enigma/analysis/MethodCallsTreeNode.java b/src/cuchaz/enigma/analysis/MethodCallsTreeNode.java
index dedfb2e..b5cf4c3 100644
--- a/src/cuchaz/enigma/analysis/MethodCallsTreeNode.java
+++ b/src/cuchaz/enigma/analysis/MethodCallsTreeNode.java
@@ -10,12 +10,15 @@
  ******************************************************************************/
 package cuchaz.enigma.analysis;
 
-import java.util.List;
+import java.util.Set;
 
 import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.TreeNode;
 
-import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
 
+import cuchaz.enigma.mapping.ConstructorEntry;
+import cuchaz.enigma.mapping.Entry;
 import cuchaz.enigma.mapping.MethodEntry;
 import cuchaz.enigma.mapping.Translator;
 
@@ -24,72 +27,117 @@ public class MethodCallsTreeNode extends DefaultMutableTreeNode
         private static final long serialVersionUID = -3658163700783307520L;
         
         private Translator m_deobfuscatingTranslator;
-        private MethodEntry m_entry;
+        private MethodEntry m_methodEntry;
+        private ConstructorEntry m_constructorEntry;
         
         public MethodCallsTreeNode( Translator deobfuscatingTranslator, MethodEntry entry )
         {
                 m_deobfuscatingTranslator = deobfuscatingTranslator;
-                m_entry = entry;
+                m_methodEntry = entry;
+                m_constructorEntry = null;
+        }
+        
+        public MethodCallsTreeNode( Translator deobfuscatingTranslator, ConstructorEntry entry )
+        {
+                m_deobfuscatingTranslator = deobfuscatingTranslator;
+                m_methodEntry = null;
+                m_constructorEntry = entry;
         }
         
         public MethodEntry getMethodEntry( )
         {
-                return m_entry;
+                return m_methodEntry;
         }
         
-        public String getDeobfClassName( )
+        public ConstructorEntry getConstructorEntry( )
         {
-                return m_deobfuscatingTranslator.translateClass( m_entry.getClassName() );
+                return m_constructorEntry;
         }
         
-        public String getDeobfMethodName( )
+        public Entry getEntry( )
         {
-                return m_deobfuscatingTranslator.translate( m_entry );
+                if( m_methodEntry != null )
+                {
+                        return m_methodEntry;
+                }
+                else if( m_constructorEntry != null )
+                {
+                        return m_constructorEntry;
+                }
+                throw new Error( "Illegal state!" );
         }
         
         @Override
         public String toString( )
         {
-                String className = getDeobfClassName();
-                if( className == null )
+                if( m_methodEntry != null )
                 {
-                        className = m_entry.getClassName();
+                        String className = m_deobfuscatingTranslator.translateClass( m_methodEntry.getClassName() );
+                        if( className == null )
+                        {
+                                className = m_methodEntry.getClassName();
+                        }
+                        
+                        String methodName = m_deobfuscatingTranslator.translate( m_methodEntry );
+                        if( methodName == null )
+                        {
+                                methodName = m_methodEntry.getName();
+                        }
+                        return className + "." + methodName + "()";
                 }
-                
-                String methodName = getDeobfMethodName();
-                if( methodName == null )
+                else if( m_constructorEntry != null )
                 {
-                        methodName = m_entry.getName();
+                        String className = m_deobfuscatingTranslator.translateClass( m_constructorEntry.getClassName() );
+                        if( className == null )
+                        {
+                                className = m_constructorEntry.getClassName();
+                        }
+                        return className + "()";
                 }
-                return className + "." + methodName + "()";
+                throw new Error( "Illegal state!" );
         }
         
         public void load( JarIndex index, boolean recurse )
         {
                 // get all the child nodes
-                List<MethodCallsTreeNode> nodes = Lists.newArrayList();
-                for( MethodEntry entry : index.getMethodCallers( m_entry ) )
+                for( Entry entry : index.getMethodCallers( getEntry() ) )
                 {
-                        nodes.add( new MethodCallsTreeNode( m_deobfuscatingTranslator, entry ) );
-                }
-                
-                // add them to this node
-                for( MethodCallsTreeNode node : nodes )
-                {
-                        this.add( node );
+                        if( entry instanceof MethodEntry )
+                        {
+                                add( new MethodCallsTreeNode( m_deobfuscatingTranslator, (MethodEntry)entry ) );
+                        }
+                        else if( entry instanceof ConstructorEntry )
+                        {
+                                add( new MethodCallsTreeNode( m_deobfuscatingTranslator, (ConstructorEntry)entry ) );
+                        }
                 }
                 
-                if( recurse )
+                if( recurse && children != null )
                 {
-                        for( MethodCallsTreeNode node : nodes )
+                        for( Object child : children )
                         {
-                                // don't recurse into self
-                                if( node.getMethodEntry().equals( m_entry ) )
+                                if( child instanceof MethodCallsTreeNode )
                                 {
-                                        continue;
+                                        MethodCallsTreeNode node = (MethodCallsTreeNode)child;
+                                        
+                                        // don't recurse into ancestor
+                                        Set<Entry> ancestors = Sets.newHashSet();
+                                        TreeNode n = (TreeNode)node;
+                                        while( n.getParent() != null )
+                                        {
+                                                n = n.getParent();
+                                                if( n instanceof MethodCallsTreeNode )
+                                                {
+                                                        ancestors.add( ((MethodCallsTreeNode)n).getEntry() );
+                                                }
+                                        }
+                                        if( ancestors.contains( node.getEntry() ) )
+                                        {
+                                                continue;
+                                        }
+                                        
+                                        node.load( index, true );
                                 }
-                                
-                                node.load( index, true );
                         }
                 }
         }
diff --git a/src/cuchaz/enigma/analysis/SourceIndexVisitor.java b/src/cuchaz/enigma/analysis/SourceIndexVisitor.java
index 0ba5996..6c14ee9 100644
--- a/src/cuchaz/enigma/analysis/SourceIndexVisitor.java
+++ b/src/cuchaz/enigma/analysis/SourceIndexVisitor.java
@@ -93,6 +93,7 @@ import com.strobel.decompiler.patterns.Pattern;
 
 import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.ConstructorEntry;
 import cuchaz.enigma.mapping.FieldEntry;
 import cuchaz.enigma.mapping.MethodEntry;
 
@@ -158,7 +159,9 @@ public class SourceIndexVisitor implements IAstVisitor<SourceIndex, Void>
         public Void visitConstructorDeclaration( ConstructorDeclaration node, SourceIndex index )
         {
                 MethodDefinition def = node.getUserData( Keys.METHOD_DEFINITION );
-                index.add( node.getNameToken(), new ClassEntry( def.getDeclaringType().getInternalName() ) );
+                ClassEntry classEntry = new ClassEntry( def.getDeclaringType().getInternalName() );
+                ConstructorEntry constructorEntry = new ConstructorEntry( classEntry, def.getSignature() );
+                index.addDeclaration( node.getNameToken(), constructorEntry );
                 
                 return recurse( node, index );
         }