added jar export

5 files changed, 177 insertions, 59 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index ff83d21a..9235cf74 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -12,14 +12,18 @@
 package cuchaz.enigma;
 
 import java.io.File;
+import java.io.FileOutputStream;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.io.StringWriter;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
+import java.util.jar.JarOutputStream;
 
+import javassist.CtClass;
 import javassist.bytecode.Descriptor;
 
 import com.google.common.collect.Lists;
@@ -36,11 +40,11 @@ import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 import com.strobel.decompiler.languages.java.ast.InsertParenthesesVisitor;
 
 import cuchaz.enigma.analysis.EntryReference;
+import cuchaz.enigma.analysis.JarClassIterator;
 import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.analysis.SourceIndexVisitor;
 import cuchaz.enigma.analysis.Token;
-import cuchaz.enigma.analysis.TreeDumpVisitor;
 import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.BehaviorEntry;
 import cuchaz.enigma.mapping.BehaviorEntryFactory;
@@ -61,7 +65,7 @@ public class Deobfuscator
 {
         public interface ProgressListener
         {
-                void init( int totalWork );
+                void init( int totalWork, String title );
                 void onProgress( int numDone, String message );
         }
         
@@ -393,7 +397,7 @@ public class Deobfuscator
                 
                 if( progress != null )
                 {
-                        progress.init( classEntries.size() );
+                        progress.init( classEntries.size(), "Decompiling classes..." );
                 }
                 
                 // DEOBFUSCATE ALL THE THINGS!! @_@
@@ -424,9 +428,60 @@ public class Deobfuscator
                                 throw new Error( "Unable to deobfuscate class " + deobfClassEntry.toString() + " (" + obfClassEntry.toString() + ")", t );
                         }
                 }
-                
-                // done!
-                progress.onProgress( classEntries.size(), "Done!" );
+                if( progress != null )
+                {
+                        progress.onProgress( i, "Done!" );
+                }
+        }
+        
+        public void writeJar( File out, ProgressListener progress )
+        {
+                try( JarOutputStream outJar = new JarOutputStream( new FileOutputStream( out ) ) )
+                {
+                        if( progress != null )
+                        {
+                                progress.init( JarClassIterator.getClassEntries( m_jar ).size(), "Translating classes..." );
+                        }
+                        
+                        // prep the loader
+                        TranslatingTypeLoader loader = new TranslatingTypeLoader(
+                                m_jar,
+                                m_jarIndex,
+                                getTranslator( TranslationDirection.Obfuscating ),
+                                getTranslator( TranslationDirection.Deobfuscating )
+                        );
+                        
+                        int i = 0;
+                        for( CtClass c : JarClassIterator.classes( m_jar ) )
+                        {
+                                if( progress != null )
+                                {
+                                        progress.onProgress( i++, c.getName() );
+                                }
+                                
+                                try
+                                {
+                                        c = loader.transformClass( c );
+                                        outJar.putNextEntry( new JarEntry( c.getName().replace( '.', '/' ) + ".class" ) );
+                                        outJar.write( c.toBytecode() );
+                                        outJar.closeEntry();
+                                }
+                                catch( Throwable t )
+                                {
+                                        throw new Error( "Unable to deobfuscate class " + c.getName(), t );
+                                }
+                        }
+                        if( progress != null )
+                        {
+                                progress.onProgress( i, "Done!" );
+                        }
+                        
+                        outJar.close();
+                }
+                catch( IOException ex )
+                {
+                        throw new Error( "Unable to write to Jar file!" );
+                }
         }
         
         public <T extends Entry> T obfuscateEntry( T deobfEntry )
diff --git a/src/cuchaz/enigma/TranslatingTypeLoader.java b/src/cuchaz/enigma/TranslatingTypeLoader.java
index 8b969857..6179879a 100644
--- a/src/cuchaz/enigma/TranslatingTypeLoader.java
+++ b/src/cuchaz/enigma/TranslatingTypeLoader.java
@@ -13,7 +13,6 @@ package cuchaz.enigma;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.util.Arrays;
 import java.util.Map;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
@@ -157,12 +156,13 @@ public class TranslatingTypeLoader implements ITypeLoader
                         // otherwise, just use the class name (ie for classes in packages)
                         classFileName = obfClassEntry.getName();
                 }
+                
                 JarEntry entry = m_jar.getJarEntry( classFileName + ".class" );
                 if( entry == null )
                 {
                         return null;
                 }
-                
+
                 try
                 {
                         // read the class file into a buffer
@@ -188,33 +188,11 @@ public class TranslatingTypeLoader implements ITypeLoader
                         classPool.insertClassPath( new ByteArrayClassPath( javaClassFileName, buf ) );
                         CtClass c = classPool.get( javaClassFileName );
                         
-                        // we moved a lot of classes out of the default package into the none package
-                        // make sure all the class references are consistent
-                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
-                        
-                        // reconstruct inner classes
-                        new InnerClassWriter( m_jarIndex ).write( c );
-                        
-                        // re-get the javassist handle since we changed class names
-                        String javaClassReconstructedName = Descriptor.toJavaName( obfClassEntry.getName() );
-                        classPool = new ClassPool();
-                        classPool.insertClassPath( new ByteArrayClassPath( javaClassReconstructedName, c.toBytecode() ) );
-                        c = classPool.get( javaClassReconstructedName );
-                        
-                        // check that the file is correct after inner class reconstruction (ie cause Javassist to fail fast if something is wrong)
-                        assertClassName( c, obfClassEntry );
-                        
-                        // do all kinds of deobfuscating transformations on the class
-                        new BridgeFixer( m_jarIndex ).fixBridges( c );
-                        new MethodParameterWriter( m_deobfuscatingTranslator ).writeMethodArguments( c );
-                        new ClassTranslator( m_deobfuscatingTranslator ).translate( c );
+                        c = transformClass( c );
                         
                         // sanity checking
                         assertClassName( c, deobfClassEntry );
                         
-                        // DEBUG
-                        //Util.writeClass( c );
-                        
                         // we have a transformed class!
                         return c.toBytecode();
                 }
@@ -223,6 +201,37 @@ public class TranslatingTypeLoader implements ITypeLoader
                         throw new Error( ex );
                 }
         }
+        
+        public CtClass transformClass( CtClass c )
+        throws IOException, NotFoundException, CannotCompileException
+        {
+                // we moved a lot of classes out of the default package into the none package
+                // make sure all the class references are consistent
+                ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
+                
+                // reconstruct inner classes
+                new InnerClassWriter( m_jarIndex ).write( c );
+                
+                // re-get the javassist handle since we changed class names
+                ClassEntry obfClassEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
+                String javaClassReconstructedName = Descriptor.toJavaName( obfClassEntry.getName() );
+                ClassPool classPool = new ClassPool();
+                classPool.insertClassPath( new ByteArrayClassPath( javaClassReconstructedName, c.toBytecode() ) );
+                c = classPool.get( javaClassReconstructedName );
+                
+                // check that the file is correct after inner class reconstruction (ie cause Javassist to fail fast if something is wrong)
+                assertClassName( c, obfClassEntry );
+                
+                // do all kinds of deobfuscating transformations on the class
+                new BridgeFixer( m_jarIndex ).fixBridges( c );
+                new MethodParameterWriter( m_deobfuscatingTranslator ).writeMethodArguments( c );
+                new ClassTranslator( m_deobfuscatingTranslator ).translate( c );
+                
+                // DEBUG
+                //Util.writeClass( c );
+                
+                return c;
+        }
 
         private void assertClassName( CtClass c, ClassEntry obfClassEntry )
         {
diff --git a/src/cuchaz/enigma/gui/Gui.java b/src/cuchaz/enigma/gui/Gui.java
index 920bc0b9..dbfcba83 100644
--- a/src/cuchaz/enigma/gui/Gui.java
+++ b/src/cuchaz/enigma/gui/Gui.java
@@ -123,12 +123,15 @@ public class Gui
         private JMenuItem m_showCallsMenu;
         private JMenuItem m_showImplementationsMenu;
         private JMenuItem m_toggleMappingMenu;
+        private JMenuItem m_exportSourceMenu;
+        private JMenuItem m_exportJarMenu;
         
         // state
         private EntryReference<Entry,Entry> m_reference;
         private JFileChooser m_jarFileChooser;
         private JFileChooser m_mappingsFileChooser;
-        private JFileChooser m_exportFileChooser;
+        private JFileChooser m_exportSourceFileChooser;
+        private JFileChooser m_exportJarFileChooser;
         
         public Gui( )
         {
@@ -157,8 +160,9 @@ public class Gui
                 // init file choosers
                 m_jarFileChooser = new JFileChooser();
                 m_mappingsFileChooser = new JFileChooser();
-                m_exportFileChooser = new JFileChooser();
-                m_exportFileChooser.setFileSelectionMode( JFileChooser.DIRECTORIES_ONLY );
+                m_exportSourceFileChooser = new JFileChooser();
+                m_exportSourceFileChooser.setFileSelectionMode( JFileChooser.DIRECTORIES_ONLY );
+                m_exportJarFileChooser = new JFileChooser();
                 
                 // init obfuscated classes list
                 m_obfClasses = new ClassSelector( ClassSelector.ObfuscatedClassEntryComparator );
@@ -663,19 +667,36 @@ public class Gui
                         }
                         menu.addSeparator();
                         {
-                                JMenuItem item = new JMenuItem( "Export..." );
+                                JMenuItem item = new JMenuItem( "Export Source..." );
                                 menu.add( item );
                                 item.addActionListener( new ActionListener( ) 
                                 {
                                         @Override
                                         public void actionPerformed( ActionEvent event )
                                         {
-                                                if( m_exportFileChooser.showSaveDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
+                                                if( m_exportSourceFileChooser.showSaveDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
                                                 {
-                                                        m_controller.export( m_exportFileChooser.getSelectedFile() );
+                                                        m_controller.exportSource( m_exportSourceFileChooser.getSelectedFile() );
                                                 }
                                         }
                                 } );
+                                m_exportSourceMenu = item;
+                        }
+                        {
+                                JMenuItem item = new JMenuItem( "Export Jar..." );
+                                menu.add( item );
+                                item.addActionListener( new ActionListener( ) 
+                                {
+                                        @Override
+                                        public void actionPerformed( ActionEvent event )
+                                        {
+                                                if( m_exportJarFileChooser.showSaveDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
+                                                {
+                                                        m_controller.exportJar( m_exportJarFileChooser.getSelectedFile() );
+                                                }
+                                        }
+                                } );
+                                m_exportJarMenu = item;
                         }
                         menu.addSeparator();
                         {
@@ -762,6 +783,8 @@ public class Gui
                 m_saveMappingsMenu.setEnabled( false );
                 m_saveMappingsAsMenu.setEnabled( true );
                 m_closeMappingsMenu.setEnabled( true );
+                m_exportSourceMenu.setEnabled( true );
+                m_exportJarMenu.setEnabled( true );
                 
                 redraw();
         }
@@ -781,6 +804,8 @@ public class Gui
                 m_saveMappingsMenu.setEnabled( false );
                 m_saveMappingsAsMenu.setEnabled( false );
                 m_closeMappingsMenu.setEnabled( false );
+                m_exportSourceMenu.setEnabled( false );
+                m_exportJarMenu.setEnabled( false );
                 
                 redraw();
         }
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index c7efbce6..2862ebed 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -23,6 +23,7 @@ import com.google.common.collect.Queues;
 import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 
 import cuchaz.enigma.Deobfuscator;
+import cuchaz.enigma.Deobfuscator.ProgressListener;
 import cuchaz.enigma.analysis.BehaviorReferenceTreeNode;
 import cuchaz.enigma.analysis.ClassImplementationsTreeNode;
 import cuchaz.enigma.analysis.ClassInheritanceTreeNode;
@@ -32,6 +33,7 @@ import cuchaz.enigma.analysis.MethodImplementationsTreeNode;
 import cuchaz.enigma.analysis.MethodInheritanceTreeNode;
 import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.analysis.Token;
+import cuchaz.enigma.gui.ProgressDialog.ProgressRunnable;
 import cuchaz.enigma.mapping.BehaviorEntry;
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.Entry;
@@ -110,28 +112,29 @@ public class GuiController
                 refreshCurrentClass();
         }
         
-        public void export( final File dirOut )
+        public void exportSource( final File dirOut )
         {
-                new Thread( )
+                ProgressDialog.runInThread( m_gui.getFrame(), new ProgressRunnable( )
                 {
                         @Override
-                        public void run( )
+                        public void run( ProgressListener progress )
+                        throws Exception
                         {
-                                ProgressDialog progress = new ProgressDialog( m_gui.getFrame() );
-                                try
-                                {
-                                        m_deobfuscator.writeSources( dirOut, progress );
-                                }
-                                catch( Exception ex )
-                                {
-                                        throw new Error( ex );
-                                }
-                                finally
-                                {
-                                        progress.close();
-                                }
+                                m_deobfuscator.writeSources( dirOut, progress );
                         }
-                }.start();
+                } );
+        }
+        
+        public void exportJar( final File fileOut )
+        {
+                ProgressDialog.runInThread( m_gui.getFrame(), new ProgressRunnable( )
+                {
+                        @Override
+                        public void run( ProgressListener progress )
+                        {
+                                m_deobfuscator.writeJar( fileOut, progress );
+                        }
+                } );
         }
         
         public Token getToken( int pos )
diff --git a/src/cuchaz/enigma/gui/ProgressDialog.java b/src/cuchaz/enigma/gui/ProgressDialog.java
index 40ac6a69..7f954314 100644
--- a/src/cuchaz/enigma/gui/ProgressDialog.java
+++ b/src/cuchaz/enigma/gui/ProgressDialog.java
@@ -25,22 +25,24 @@ import javax.swing.WindowConstants;
 import cuchaz.enigma.Constants;
 import cuchaz.enigma.Deobfuscator.ProgressListener;
 
-public class ProgressDialog implements ProgressListener
+public class ProgressDialog implements ProgressListener, AutoCloseable
 {
         private JFrame m_frame;
+        private JLabel m_title;
         private JLabel m_text;
         private JProgressBar m_progress;
         
         public ProgressDialog( JFrame parent )
         {
                 // init frame
-                m_frame = new JFrame( Constants.Name + " - Export" );
+                m_frame = new JFrame( Constants.Name + " - Operation in progress" );
                 final Container pane = m_frame.getContentPane();
                 FlowLayout layout = new FlowLayout();
                 layout.setAlignment( FlowLayout.LEFT );
                 pane.setLayout( layout );
                 
-                pane.add( new JLabel( "Decompiling classes..." ) );
+                m_title = new JLabel();
+                pane.add( m_title );
                 
                 // set up the progress bar
                 JPanel panel = new JPanel();
@@ -68,9 +70,9 @@ public class ProgressDialog implements ProgressListener
         }
 
         @Override
-        public void init( int totalWork )
+        public void init( int totalWork, String title )
         {
-                m_text.setText( "Decompiling " + totalWork + " classes..." );
+                m_title.setText( title );
                 m_progress.setMinimum( 0 );
                 m_progress.setMaximum( totalWork );
                 m_progress.setValue( 0 );
@@ -86,4 +88,28 @@ public class ProgressDialog implements ProgressListener
                 m_frame.validate();
                 m_frame.repaint();
         }
+        
+        public static interface ProgressRunnable
+        {
+                void run( ProgressListener listener ) throws Exception;
+        }
+        
+        public static void runInThread( final JFrame parent, final ProgressRunnable runnable )
+        {
+                new Thread( )
+                {
+                        @Override
+                        public void run( )
+                        {
+                                try( ProgressDialog progress = new ProgressDialog( parent ) )
+                                {
+                                        runnable.run( progress );
+                                }
+                                catch( Exception ex )
+                                {
+                                        throw new Error( ex );
+                                }
+                        }
+                }.start();
+        }
 }