1 files changed, 411 insertions, 0 deletions

diff --git a/src/cuchaz/enigma/convert/ClassIdentity.java b/src/cuchaz/enigma/convert/ClassIdentity.java
new file mode 100644
index 00000000..1be61239
--- /dev/null
+++ b/src/cuchaz/enigma/convert/ClassIdentity.java
@@ -0,0 +1,411 @@
+/*******************************************************************************
+ * 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.convert;
+
+import java.io.UnsupportedEncodingException;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.Enumeration;
+import java.util.List;
+import java.util.Map;
+
+import javassist.CannotCompileException;
+import javassist.CtBehavior;
+import javassist.CtClass;
+import javassist.CtConstructor;
+import javassist.CtField;
+import javassist.CtMethod;
+import javassist.bytecode.BadBytecode;
+import javassist.bytecode.CodeIterator;
+import javassist.bytecode.ConstPool;
+import javassist.bytecode.Descriptor;
+import javassist.bytecode.Opcode;
+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.HashMultiset;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Multiset;
+
+import cuchaz.enigma.Constants;
+import cuchaz.enigma.Util;
+import cuchaz.enigma.analysis.ClassImplementationsTreeNode;
+import cuchaz.enigma.analysis.EntryReference;
+import cuchaz.enigma.analysis.JarIndex;
+import cuchaz.enigma.bytecode.ConstPoolEditor;
+import cuchaz.enigma.bytecode.InfoType;
+import cuchaz.enigma.bytecode.accessors.ConstInfoAccessor;
+import cuchaz.enigma.convert.ClassNamer.SidedClassNamer;
+import cuchaz.enigma.mapping.BehaviorEntry;
+import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.ClassNameReplacer;
+import cuchaz.enigma.mapping.Entry;
+import cuchaz.enigma.mapping.FieldEntry;
+import cuchaz.enigma.mapping.JavassistUtil;
+import cuchaz.enigma.mapping.Signature;
+
+public class ClassIdentity {
+        
+        private ClassEntry m_classEntry;
+        private SidedClassNamer m_namer;
+        private Multiset<String> m_fields;
+        private Multiset<String> m_methods;
+        private Multiset<String> m_constructors;
+        private String m_staticInitializer;
+        private String m_extends;
+        private Multiset<String> m_implements;
+        private Multiset<String> m_implementations;
+        private Multiset<String> m_references;
+        
+        public ClassIdentity(CtClass c, SidedClassNamer namer, JarIndex index, boolean useReferences) {
+                m_namer = namer;
+                
+                // stuff from the bytecode
+                
+                m_classEntry = new ClassEntry(Descriptor.toJvmName(c.getName()));
+                m_fields = HashMultiset.create();
+                for (CtField field : c.getDeclaredFields()) {
+                        m_fields.add(scrubSignature(field.getSignature()));
+                }
+                m_methods = HashMultiset.create();
+                for (CtMethod method : c.getDeclaredMethods()) {
+                        m_methods.add(scrubSignature(method.getSignature()) + "0x" + getBehaviorSignature(method));
+                }
+                m_constructors = HashMultiset.create();
+                for (CtConstructor constructor : c.getDeclaredConstructors()) {
+                        m_constructors.add(scrubSignature(constructor.getSignature()) + "0x" + getBehaviorSignature(constructor));
+                }
+                m_staticInitializer = "";
+                if (c.getClassInitializer() != null) {
+                        m_staticInitializer = getBehaviorSignature(c.getClassInitializer());
+                }
+                m_extends = "";
+                if (c.getClassFile().getSuperclass() != null) {
+                        m_extends = scrubClassName(c.getClassFile().getSuperclass());
+                }
+                m_implements = HashMultiset.create();
+                for (String interfaceName : c.getClassFile().getInterfaces()) {
+                        m_implements.add(scrubClassName(interfaceName));
+                }
+                
+                // stuff from the jar index
+                
+                m_implementations = HashMultiset.create();
+                ClassImplementationsTreeNode implementationsNode = index.getClassImplementations(null, m_classEntry);
+                if (implementationsNode != null) {
+                        @SuppressWarnings("unchecked")
+                        Enumeration<ClassImplementationsTreeNode> implementations = implementationsNode.children();
+                        while (implementations.hasMoreElements()) {
+                                ClassImplementationsTreeNode node = implementations.nextElement();
+                                m_implementations.add(scrubClassName(node.getClassEntry().getName()));
+                        }
+                }
+                
+                m_references = HashMultiset.create();
+                if (useReferences) {
+                        for (CtField field : c.getDeclaredFields()) {
+                                FieldEntry fieldEntry = new FieldEntry(m_classEntry, field.getName());
+                                for (EntryReference<FieldEntry,BehaviorEntry> reference : index.getFieldReferences(fieldEntry)) {
+                                        addReference(reference);
+                                }
+                        }
+                        for (CtBehavior behavior : c.getDeclaredBehaviors()) {
+                                BehaviorEntry behaviorEntry = JavassistUtil.getBehaviorEntry(behavior);
+                                for (EntryReference<BehaviorEntry,BehaviorEntry> reference : index.getBehaviorReferences(behaviorEntry)) {
+                                        addReference(reference);
+                                }
+                        }
+                }
+        }
+        
+        private void addReference(EntryReference<? extends Entry,BehaviorEntry> reference) {
+                if (reference.context.getSignature() != null) {
+                        m_references.add(String.format("%s_%s", scrubClassName(reference.context.getClassName()), scrubSignature(reference.context.getSignature())));
+                } else {
+                        m_references.add(String.format("%s_<clinit>", scrubClassName(reference.context.getClassName())));
+                }
+        }
+        
+        public ClassEntry getClassEntry() {
+                return m_classEntry;
+        }
+        
+        @Override
+        public String toString() {
+                StringBuilder buf = new StringBuilder();
+                buf.append("class: ");
+                buf.append(m_classEntry.getName());
+                buf.append(" ");
+                buf.append(hashCode());
+                buf.append("\n");
+                for (String field : m_fields) {
+                        buf.append("\tfield ");
+                        buf.append(field);
+                        buf.append("\n");
+                }
+                for (String method : m_methods) {
+                        buf.append("\tmethod ");
+                        buf.append(method);
+                        buf.append("\n");
+                }
+                for (String constructor : m_constructors) {
+                        buf.append("\tconstructor ");
+                        buf.append(constructor);
+                        buf.append("\n");
+                }
+                if (m_staticInitializer.length() > 0) {
+                        buf.append("\tinitializer ");
+                        buf.append(m_staticInitializer);
+                        buf.append("\n");
+                }
+                if (m_extends.length() > 0) {
+                        buf.append("\textends ");
+                        buf.append(m_extends);
+                        buf.append("\n");
+                }
+                for (String interfaceName : m_implements) {
+                        buf.append("\timplements ");
+                        buf.append(interfaceName);
+                        buf.append("\n");
+                }
+                for (String implementation : m_implementations) {
+                        buf.append("\timplemented by ");
+                        buf.append(implementation);
+                        buf.append("\n");
+                }
+                for (String reference : m_references) {
+                        buf.append("\treference ");
+                        buf.append(reference);
+                        buf.append("\n");
+                }
+                return buf.toString();
+        }
+        
+        private String scrubClassName(String className) {
+                return scrubSignature("L" + Descriptor.toJvmName(className) + ";");
+        }
+        
+        private String scrubSignature(String signature) {
+                return scrubSignature(new Signature(signature));
+        }
+        
+        private String scrubSignature(Signature signature) {
+                
+                return new Signature(signature, new ClassNameReplacer() {
+                        
+                        private Map<String,String> m_classNames = Maps.newHashMap();
+                        
+                        @Override
+                        public String replace(String className) {
+                                
+                                // classes not in the none package can be passed through
+                                ClassEntry classEntry = new ClassEntry(className);
+                                if (!classEntry.getPackageName().equals(Constants.NonePackage)) {
+                                        return className;
+                                }
+                                
+                                // is this class ourself?
+                                if (className.equals(m_classEntry.getName())) {
+                                        return "CSelf";
+                                }
+                                
+                                // try the namer
+                                if (m_namer != null) {
+                                        String newName = m_namer.getName(className);
+                                        if (newName != null) {
+                                                return newName;
+                                        }
+                                }
+                                
+                                // otherwise, use local naming
+                                if (!m_classNames.containsKey(className)) {
+                                        m_classNames.put(className, getNewClassName());
+                                }
+                                return m_classNames.get(className);
+                        }
+                        
+                        private String getNewClassName() {
+                                return String.format("C%03d", m_classNames.size());
+                        }
+                }).toString();
+        }
+        
+        private boolean isClassMatchedUniquely(String className) {
+                return m_namer != null && m_namer.getName(Descriptor.toJvmName(className)) != null;
+        }
+        
+        private String getBehaviorSignature(CtBehavior behavior) {
+                try {
+                        // does this method have an implementation?
+                        if (behavior.getMethodInfo().getCodeAttribute() == null) {
+                                return "(none)";
+                        }
+                        
+                        // compute the hash from the opcodes
+                        ConstPool constants = behavior.getMethodInfo().getConstPool();
+                        final MessageDigest digest = MessageDigest.getInstance("MD5");
+                        CodeIterator iter = behavior.getMethodInfo().getCodeAttribute().iterator();
+                        while (iter.hasNext()) {
+                                int pos = iter.next();
+                                
+                                // update the hash with the opcode
+                                int opcode = iter.byteAt(pos);
+                                digest.update((byte)opcode);
+                                
+                                switch (opcode) {
+                                        case Opcode.LDC: {
+                                                int constIndex = iter.byteAt(pos + 1);
+                                                updateHashWithConstant(digest, constants, constIndex);
+                                        }
+                                        break;
+                                        
+                                        case Opcode.LDC_W:
+                                        case Opcode.LDC2_W: {
+                                                int constIndex = (iter.byteAt(pos + 1) << 8) | iter.byteAt(pos + 2);
+                                                updateHashWithConstant(digest, constants, constIndex);
+                                        }
+                                        break;
+                                }
+                        }
+                        
+                        // update hash with method and field accesses
+                        behavior.instrument(new ExprEditor() {
+                                @Override
+                                public void edit(MethodCall call) {
+                                        updateHashWithString(digest, scrubClassName(call.getClassName()));
+                                        updateHashWithString(digest, scrubSignature(call.getSignature()));
+                                        if (isClassMatchedUniquely(call.getClassName())) {
+                                                updateHashWithString(digest, call.getMethodName());
+                                        }
+                                }
+                                
+                                @Override
+                                public void edit(FieldAccess access) {
+                                        updateHashWithString(digest, scrubClassName(access.getClassName()));
+                                        updateHashWithString(digest, scrubSignature(access.getSignature()));
+                                        if (isClassMatchedUniquely(access.getClassName())) {
+                                                updateHashWithString(digest, access.getFieldName());
+                                        }
+                                }
+                                
+                                @Override
+                                public void edit(ConstructorCall call) {
+                                        updateHashWithString(digest, scrubClassName(call.getClassName()));
+                                        updateHashWithString(digest, scrubSignature(call.getSignature()));
+                                }
+                                
+                                @Override
+                                public void edit(NewExpr expr) {
+                                        updateHashWithString(digest, scrubClassName(expr.getClassName()));
+                                }
+                        });
+                        
+                        // convert the hash to a hex string
+                        return toHex(digest.digest());
+                } catch (BadBytecode | NoSuchAlgorithmException | CannotCompileException ex) {
+                        throw new Error(ex);
+                }
+        }
+        
+        private void updateHashWithConstant(MessageDigest digest, ConstPool constants, int index) {
+                ConstPoolEditor editor = new ConstPoolEditor(constants);
+                ConstInfoAccessor item = editor.getItem(index);
+                if (item.getType() == InfoType.StringInfo) {
+                        updateHashWithString(digest, constants.getStringInfo(index));
+                }
+                // TODO: other constants
+        }
+        
+        private void updateHashWithString(MessageDigest digest, String val) {
+                try {
+                        digest.update(val.getBytes("UTF8"));
+                } catch (UnsupportedEncodingException ex) {
+                        throw new Error(ex);
+                }
+        }
+        
+        private String toHex(byte[] bytes) {
+                // function taken from:
+                // http://stackoverflow.com/questions/9655181/convert-from-byte-array-to-hex-string-in-java
+                final char[] hexArray = "0123456789ABCDEF".toCharArray();
+                char[] hexChars = new char[bytes.length * 2];
+                for (int j = 0; j < bytes.length; j++) {
+                        int v = bytes[j] & 0xFF;
+                        hexChars[j * 2] = hexArray[v >>> 4];
+                        hexChars[j * 2 + 1] = hexArray[v & 0x0F];
+                }
+                return new String(hexChars);
+        }
+        
+        @Override
+        public boolean equals(Object other) {
+                if (other instanceof ClassIdentity) {
+                        return equals((ClassIdentity)other);
+                }
+                return false;
+        }
+        
+        public boolean equals(ClassIdentity other) {
+                return m_fields.equals(other.m_fields)
+                        && m_methods.equals(other.m_methods)
+                        && m_constructors.equals(other.m_constructors)
+                        && m_staticInitializer.equals(other.m_staticInitializer)
+                        && m_extends.equals(other.m_extends)
+                        && m_implements.equals(other.m_implements)
+                        && m_implementations.equals(other.m_implementations)
+                        && m_references.equals(other.m_references);
+        }
+        
+        @Override
+        public int hashCode() {
+                List<Object> objs = Lists.newArrayList();
+                objs.addAll(m_fields);
+                objs.addAll(m_methods);
+                objs.addAll(m_constructors);
+                objs.add(m_staticInitializer);
+                objs.add(m_extends);
+                objs.addAll(m_implements);
+                objs.addAll(m_implementations);
+                objs.addAll(m_references);
+                return Util.combineHashesOrdered(objs);
+        }
+        
+        public int getMatchScore(ClassIdentity other) {
+                return getNumMatches(m_fields, other.m_fields) 
+                        + getNumMatches(m_methods, other.m_methods)
+                        + getNumMatches(m_constructors, other.m_constructors);
+        }
+        
+        public int getMaxMatchScore() {
+                return m_fields.size() + m_methods.size() + m_constructors.size();
+        }
+        
+        public boolean matches(CtClass c) {
+                // just compare declaration counts
+                return m_fields.size() == c.getDeclaredFields().length
+                        && m_methods.size() == c.getDeclaredMethods().length
+                        && m_constructors.size() == c.getDeclaredConstructors().length;
+        }
+        
+        private int getNumMatches(Multiset<String> a, Multiset<String> b) {
+                int numMatches = 0;
+                for (String val : a) {
+                        if (b.contains(val)) {
+                                numMatches++;
+                        }
+                }
+                return numMatches;
+        }
+}