added export command with progress bar

5 files changed, 233 insertions, 23 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index 38f7af56..5d87ad0a 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -11,6 +11,7 @@
 package cuchaz.enigma;
 
 import java.io.File;
+import java.io.FileWriter;
 import java.io.IOException;
 import java.io.StringWriter;
 import java.util.List;
@@ -44,6 +45,12 @@ import cuchaz.enigma.mapping.Translator;
 
 public class Deobfuscator
 {
+        public interface ProgressListener
+        {
+                void init( int totalWork );
+                void onProgress( int numDone, String message );
+        }
+        
         private File m_file;
         private JarFile m_jar;
         private DecompilerSettings m_settings;
@@ -137,7 +144,7 @@ public class Deobfuscator
                 }
         }
         
-        public SourceIndex getSource( String className )
+        public CompilationUnit getSourceTree( String className )
         {
                 // is this class deobfuscated?
                 // we need to tell the decompiler the deobfuscated name so it doesn't get freaked out
@@ -156,19 +163,18 @@ public class Deobfuscator
                 AstBuilder builder = new AstBuilder( context );
                 builder.addType( resolvedType );
                 builder.runTransformations( null );
-                CompilationUnit root = builder.getCompilationUnit();
+                return builder.getCompilationUnit();
+        }
+        
+        public SourceIndex getSourceIndex( CompilationUnit sourceTree, String source )
+        {
+                // build the source index
+                SourceIndex index = new SourceIndex( source );
+                sourceTree.acceptVisitor( new SourceIndexVisitor(), index );
                 
-                // render the AST into source
-                StringWriter buf = new StringWriter();
-                root.acceptVisitor( new InsertParenthesesVisitor(), null );
                 // DEBUG
                 //root.acceptVisitor( new TreeDumpVisitor( new File( "tree.txt" ) ), null );
-                root.acceptVisitor( new JavaOutputVisitor( new PlainTextOutput( buf ), m_settings ), null );
-                
-                // build the source index
-                SourceIndex index = new SourceIndex( buf.toString() );
-                root.acceptVisitor( new SourceIndexVisitor(), index );
-                
+
                 /* DEBUG
                 for( Token token : index.referenceTokens() )
                 {
@@ -180,6 +186,56 @@ public class Deobfuscator
                 return index;
         }
         
+        public String getSource( CompilationUnit sourceTree )
+        {
+                // render the AST into source
+                StringWriter buf = new StringWriter();
+                sourceTree.acceptVisitor( new InsertParenthesesVisitor(), null );
+                sourceTree.acceptVisitor( new JavaOutputVisitor( new PlainTextOutput( buf ), m_settings ), null );
+                return buf.toString();
+        }
+        
+        public void writeSources( File dirOut, ProgressListener progress )
+        throws IOException
+        {
+                int numClasses = m_jarIndex.getObfClassNames().size();
+                if( progress != null )
+                {
+                        progress.init( numClasses );
+                }
+                int i = 0;
+                
+                // DEOBFUSCATE ALL THE THINGS!! @_@
+                for( String obfClassName : m_jarIndex.getObfClassNames() )
+                {
+                        // skip inner classes
+                        if( m_jarIndex.getOuterClass( obfClassName ) != null )
+                        {
+                                continue;
+                        }
+                        
+                        ClassEntry deobfClassEntry = deobfuscateEntry( new ClassEntry( obfClassName ) );
+                        if( progress != null )
+                        {
+                                progress.onProgress( i++, deobfClassEntry.toString() );
+                        }
+                        
+                        // get the source
+                        String source = getSource( getSourceTree( obfClassName ) );
+                        
+                        // write the file
+                        File file = new File( dirOut, deobfClassEntry.getName().replace( '.', '/' ) + ".java" );
+                        file.getParentFile().mkdirs();
+                        try( FileWriter out = new FileWriter( file ) )
+                        {
+                                out.write( source );
+                        }
+                }
+                
+                // done!
+                progress.onProgress( numClasses, "Done!" );
+        }
+        
         public <T extends Entry> T obfuscateEntry( T deobfEntry )
         {
                 if( deobfEntry == null )
diff --git a/src/cuchaz/enigma/gui/Gui.java b/src/cuchaz/enigma/gui/Gui.java
index dd82b8e0..febdfd43 100644
--- a/src/cuchaz/enigma/gui/Gui.java
+++ b/src/cuchaz/enigma/gui/Gui.java
@@ -15,7 +15,6 @@ import java.awt.Color;
 import java.awt.Container;
 import java.awt.Dimension;
 import java.awt.FlowLayout;
-import java.awt.Font;
 import java.awt.GridLayout;
 import java.awt.Rectangle;
 import java.awt.event.ActionEvent;
@@ -169,6 +168,7 @@ public class Gui
         private EntryReference<Entry,Entry> m_reference;
         private JFileChooser m_jarFileChooser;
         private JFileChooser m_mappingsFileChooser;
+        private JFileChooser m_exportFileChooser;
         
         public Gui( )
         {
@@ -177,6 +177,8 @@ public class Gui
                 // init file choosers
                 m_jarFileChooser = new JFileChooser();
                 m_mappingsFileChooser = new JFileChooser();
+                m_exportFileChooser = new JFileChooser();
+                m_exportFileChooser.setFileSelectionMode( JFileChooser.DIRECTORIES_ONLY );
                 
                 // init frame
                 m_frame = new JFrame( Constants.Name );
@@ -637,6 +639,22 @@ public class Gui
                         }
                         menu.addSeparator();
                         {
+                                JMenuItem item = new JMenuItem( "Export..." );
+                                menu.add( item );
+                                item.addActionListener( new ActionListener( ) 
+                                {
+                                        @Override
+                                        public void actionPerformed( ActionEvent event )
+                                        {
+                                                if( m_exportFileChooser.showSaveDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
+                                                {
+                                                        m_controller.export( m_exportFileChooser.getSelectedFile() );
+                                                }
+                                        }
+                                } );
+                        }
+                        menu.addSeparator();
+                        {
                                 JMenuItem item = new JMenuItem( "Exit" );
                                 menu.add( item );
                                 item.addActionListener( new ActionListener( )
@@ -686,6 +704,11 @@ public class Gui
                 m_frame.setDefaultCloseOperation( WindowConstants.DO_NOTHING_ON_CLOSE );
         }
         
+        public JFrame getFrame( )
+        {
+                return m_frame;
+        }
+        
         public GuiController getController( )
         {
                 return m_controller;
@@ -887,7 +910,7 @@ public class Gui
         {
                 m_infoPanel.removeAll();
                 JLabel label = new JLabel( "No identifier selected" );
-                unboldLabel( label );
+                GuiTricks.unboldLabel( label );
                 label.setHorizontalAlignment( JLabel.CENTER );
                 m_infoPanel.add( label );
                 
@@ -975,7 +998,7 @@ public class Gui
                 label.setPreferredSize( new Dimension( 100, label.getPreferredSize().height ) );
                 panel.add( label );
                 
-                panel.add( unboldLabel( new JLabel( value, JLabel.LEFT ) ) );
+                panel.add( GuiTricks.unboldLabel( new JLabel( value, JLabel.LEFT ) ) );
         }
         
         private void onCaretMove( int pos )
@@ -1059,7 +1082,7 @@ public class Gui
                 // abort the rename
                 JPanel panel = (JPanel)m_infoPanel.getComponent( 0 );
                 panel.remove( panel.getComponentCount() - 1 );
-                panel.add( unboldLabel( new JLabel( m_reference.entry.getName(), JLabel.LEFT ) ) );
+                panel.add( GuiTricks.unboldLabel( new JLabel( m_reference.entry.getName(), JLabel.LEFT ) ) );
                 
                 m_editor.grabFocus();
                 
@@ -1203,13 +1226,6 @@ public class Gui
                 }
         }
 
-        private JLabel unboldLabel( JLabel label )
-        {
-                Font font = label.getFont();
-                label.setFont( font.deriveFont( font.getStyle() & ~Font.BOLD ) );
-                return label;
-        }
-        
         private void redraw( )
         {
                 m_frame.validate();
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index a35db056..cf4f002e 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -20,6 +20,7 @@ import java.util.List;
 
 import com.google.common.collect.Lists;
 import com.google.common.collect.Queues;
+import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 
 import cuchaz.enigma.Deobfuscator;
 import cuchaz.enigma.analysis.BehaviorReferenceTreeNode;
@@ -106,6 +107,27 @@ public class GuiController
                 refreshCurrentClass();
         }
         
+        public void export( final File dirOut )
+        {
+                new Thread( )
+                {
+                        @Override
+                        public void run( )
+                        {
+                                try
+                                {
+                                        ProgressDialog progress = new ProgressDialog( m_gui.getFrame() );
+                                        m_deobfuscator.writeSources( dirOut, progress );
+                                        progress.close();
+                                }
+                                catch( IOException ex )
+                                {
+                                        throw new Error( ex );
+                                }
+                        }
+                }.start();
+        }
+        
         public Token getToken( int pos )
         {
                 if( m_index == null )
@@ -295,7 +317,9 @@ public class GuiController
                         public void run( )
                         {
                                 // decompile,deobfuscate the bytecode
-                                m_index = m_deobfuscator.getSource( classEntry.getClassName() );
+                                CompilationUnit sourceTree = m_deobfuscator.getSourceTree( classEntry.getClassName() );
+                                String source = m_deobfuscator.getSource( sourceTree );
+                                m_index = m_deobfuscator.getSourceIndex( sourceTree, source );
                                 m_gui.setSource( m_index.getSource() );
                                 if( obfReference != null )
                                 {
diff --git a/src/cuchaz/enigma/gui/GuiTricks.java b/src/cuchaz/enigma/gui/GuiTricks.java
new file mode 100644
index 00000000..c79f4329
--- /dev/null
+++ b/src/cuchaz/enigma/gui/GuiTricks.java
@@ -0,0 +1,25 @@
+/*******************************************************************************
+ * 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.gui;
+
+import java.awt.Font;
+
+import javax.swing.JLabel;
+
+public class GuiTricks
+{
+        public static JLabel unboldLabel( JLabel label )
+        {
+                Font font = label.getFont();
+                label.setFont( font.deriveFont( font.getStyle() & ~Font.BOLD ) );
+                return label;
+        }
+}
diff --git a/src/cuchaz/enigma/gui/ProgressDialog.java b/src/cuchaz/enigma/gui/ProgressDialog.java
new file mode 100644
index 00000000..40ac6a69
--- /dev/null
+++ b/src/cuchaz/enigma/gui/ProgressDialog.java
@@ -0,0 +1,89 @@
+/*******************************************************************************
+ * 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.gui;
+
+import java.awt.BorderLayout;
+import java.awt.Container;
+import java.awt.Dimension;
+import java.awt.FlowLayout;
+
+import javax.swing.BorderFactory;
+import javax.swing.JFrame;
+import javax.swing.JLabel;
+import javax.swing.JPanel;
+import javax.swing.JProgressBar;
+import javax.swing.WindowConstants;
+
+import cuchaz.enigma.Constants;
+import cuchaz.enigma.Deobfuscator.ProgressListener;
+
+public class ProgressDialog implements ProgressListener
+{
+        private JFrame m_frame;
+        private JLabel m_text;
+        private JProgressBar m_progress;
+        
+        public ProgressDialog( JFrame parent )
+        {
+                // init frame
+                m_frame = new JFrame( Constants.Name + " - Export" );
+                final Container pane = m_frame.getContentPane();
+                FlowLayout layout = new FlowLayout();
+                layout.setAlignment( FlowLayout.LEFT );
+                pane.setLayout( layout );
+                
+                pane.add( new JLabel( "Decompiling classes..." ) );
+                
+                // set up the progress bar
+                JPanel panel = new JPanel();
+                pane.add( panel );
+                panel.setLayout( new BorderLayout() );
+                m_text = GuiTricks.unboldLabel( new JLabel() );
+                m_progress = new JProgressBar();
+                m_text.setBorder( BorderFactory.createEmptyBorder( 0, 0, 10, 0 ) );
+                panel.add( m_text, BorderLayout.NORTH );
+                panel.add( m_progress, BorderLayout.CENTER );
+                panel.setPreferredSize( new Dimension( 360, 50 ) );
+                
+                // show the frame
+                pane.doLayout();
+                m_frame.setSize( 400, 120 );
+                m_frame.setResizable( false );
+                m_frame.setLocationRelativeTo( parent );
+                m_frame.setVisible( true );
+                m_frame.setDefaultCloseOperation( WindowConstants.DO_NOTHING_ON_CLOSE );
+        }
+        
+        public void close( )
+        {
+                m_frame.dispose();
+        }
+
+        @Override
+        public void init( int totalWork )
+        {
+                m_text.setText( "Decompiling " + totalWork + " classes..." );
+                m_progress.setMinimum( 0 );
+                m_progress.setMaximum( totalWork );
+                m_progress.setValue( 0 );
+        }
+
+        @Override
+        public void onProgress( int numDone, String message )
+        {
+                m_text.setText( message );
+                m_progress.setValue( numDone );
+                
+                // update the frame
+                m_frame.validate();
+                m_frame.repaint();
+        }
+}