/*******************************************************************************
* 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.ConstructorEntry;
import cuchaz.enigma.mapping.Entry;
import cuchaz.enigma.mapping.FieldEntry;
import cuchaz.enigma.mapping.MethodEntry;
import cuchaz.enigma.mapping.SignatureUpdater;
import cuchaz.enigma.mapping.SignatureUpdater.ClassNameUpdater;
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( CtMethod method : c.getDeclaredMethods() )
{
MethodEntry methodEntry = new MethodEntry( m_classEntry, method.getName(), method.getSignature() );
for( EntryReference<BehaviorEntry,BehaviorEntry> reference : index.getBehaviorReferences( methodEntry ) )
{
addReference( reference );
}
}
for( CtConstructor constructor : c.getDeclaredConstructors() )
{
ConstructorEntry constructorEntry = new ConstructorEntry( m_classEntry, constructor.getSignature() );
for( EntryReference<BehaviorEntry,BehaviorEntry> reference : index.getBehaviorReferences( constructorEntry ) )
{
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 SignatureUpdater.update( signature, new ClassNameUpdater( )
{
private Map<String,String> m_classNames = Maps.newHashMap();
@Override
public String update( 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() );
}
} );
}
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;
}
}