1 files changed, 589 insertions, 0 deletions

diff --git a/src/cuchaz/enigma/gui/ClassMatchingGui.java b/src/cuchaz/enigma/gui/ClassMatchingGui.java
new file mode 100644
index 0000000..89b19c3
--- /dev/null
+++ b/src/cuchaz/enigma/gui/ClassMatchingGui.java
@@ -0,0 +1,589 @@
+/*******************************************************************************
+ * Copyright (c) 2015 Jeff Martin.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the GNU Lesser General Public
+ * License v3.0 which accompanies this distribution, and is available at
+ * http://www.gnu.org/licenses/lgpl.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 java.awt.event.ActionEvent;
+import java.awt.event.ActionListener;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+
+import javax.swing.BoxLayout;
+import javax.swing.ButtonGroup;
+import javax.swing.JButton;
+import javax.swing.JCheckBox;
+import javax.swing.JFrame;
+import javax.swing.JLabel;
+import javax.swing.JPanel;
+import javax.swing.JRadioButton;
+import javax.swing.JScrollPane;
+import javax.swing.JSplitPane;
+import javax.swing.SwingConstants;
+import javax.swing.WindowConstants;
+
+import com.google.common.collect.BiMap;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+
+import cuchaz.enigma.Constants;
+import cuchaz.enigma.Deobfuscator;
+import cuchaz.enigma.convert.ClassIdentifier;
+import cuchaz.enigma.convert.ClassIdentity;
+import cuchaz.enigma.convert.ClassMatch;
+import cuchaz.enigma.convert.ClassMatches;
+import cuchaz.enigma.convert.ClassMatching;
+import cuchaz.enigma.convert.ClassNamer;
+import cuchaz.enigma.convert.MappingsConverter;
+import cuchaz.enigma.gui.ClassSelector.ClassSelectionListener;
+import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.Mappings;
+import cuchaz.enigma.mapping.MappingsChecker;
+import de.sciss.syntaxpane.DefaultSyntaxKit;
+
+
+public class ClassMatchingGui {
+        
+        private static enum SourceType {
+                Matched {
+                        
+                        @Override
+                        public Collection<ClassEntry> getSourceClasses(ClassMatches matches) {
+                                return matches.getUniqueMatches().keySet();
+                        }
+                },
+                Unmatched {
+                        
+                        @Override
+                        public Collection<ClassEntry> getSourceClasses(ClassMatches matches) {
+                                return matches.getUnmatchedSourceClasses();
+                        }
+                },
+                Ambiguous {
+                        
+                        @Override
+                        public Collection<ClassEntry> getSourceClasses(ClassMatches matches) {
+                                return matches.getAmbiguouslyMatchedSourceClasses();
+                        }
+                };
+                
+                public JRadioButton newRadio(ActionListener listener, ButtonGroup group) {
+                        JRadioButton button = new JRadioButton(name(), this == getDefault());
+                        button.setActionCommand(name());
+                        button.addActionListener(listener);
+                        group.add(button);
+                        return button;
+                }
+                
+                public abstract Collection<ClassEntry> getSourceClasses(ClassMatches matches);
+                
+                public static SourceType getDefault() {
+                        return values()[0];
+                }
+        }
+        
+        public static interface SaveListener {
+                public void save(ClassMatches matches);
+        }
+        
+        // controls
+        private JFrame m_frame;
+        private ClassSelector m_sourceClasses;
+        private ClassSelector m_destClasses;
+        private CodeReader m_sourceReader;
+        private CodeReader m_destReader;
+        private JLabel m_sourceClassLabel;
+        private JLabel m_destClassLabel;
+        private JButton m_matchButton;
+        private Map<SourceType,JRadioButton> m_sourceTypeButtons;
+        private JCheckBox m_advanceCheck;
+        
+        private ClassMatches m_classMatches;
+        private Deobfuscator m_sourceDeobfuscator;
+        private Deobfuscator m_destDeobfuscator;
+        private ClassEntry m_sourceClass;
+        private ClassEntry m_destClass;
+        private SourceType m_sourceType;
+        private SaveListener m_saveListener;
+
+        public ClassMatchingGui(ClassMatches matches, Deobfuscator sourceDeobfuscator, Deobfuscator destDeobfuscator) {
+                
+                m_classMatches = matches;
+                m_sourceDeobfuscator = sourceDeobfuscator;
+                m_destDeobfuscator = destDeobfuscator;
+                
+                // init frame
+                m_frame = new JFrame(Constants.Name + " - Class Matcher");
+                final Container pane = m_frame.getContentPane();
+                pane.setLayout(new BorderLayout());
+                
+                // init source side
+                JPanel sourcePanel = new JPanel();
+                sourcePanel.setLayout(new BoxLayout(sourcePanel, BoxLayout.PAGE_AXIS));
+                sourcePanel.setPreferredSize(new Dimension(200, 0));
+                pane.add(sourcePanel, BorderLayout.WEST);
+                sourcePanel.add(new JLabel("Source Classes"));
+                
+                // init source type radios
+                JPanel sourceTypePanel = new JPanel();
+                sourcePanel.add(sourceTypePanel);
+                sourceTypePanel.setLayout(new BoxLayout(sourceTypePanel, BoxLayout.PAGE_AXIS));
+                ActionListener sourceTypeListener = new ActionListener() {
+                        @Override
+                        public void actionPerformed(ActionEvent event) {
+                                setSourceType(SourceType.valueOf(event.getActionCommand()));
+                        }
+                };
+                ButtonGroup sourceTypeButtons = new ButtonGroup();
+                m_sourceTypeButtons = Maps.newHashMap();
+                for (SourceType sourceType : SourceType.values()) {
+                        JRadioButton button = sourceType.newRadio(sourceTypeListener, sourceTypeButtons);
+                        m_sourceTypeButtons.put(sourceType, button);
+                        sourceTypePanel.add(button);
+                }
+                
+                m_sourceClasses = new ClassSelector(ClassSelector.DeobfuscatedClassEntryComparator);
+                m_sourceClasses.setListener(new ClassSelectionListener() {
+                        @Override
+                        public void onSelectClass(ClassEntry classEntry) {
+                                setSourceClass(classEntry);
+                        }
+                });
+                JScrollPane sourceScroller = new JScrollPane(m_sourceClasses);
+                sourcePanel.add(sourceScroller);
+                
+                // init dest side
+                JPanel destPanel = new JPanel();
+                destPanel.setLayout(new BoxLayout(destPanel, BoxLayout.PAGE_AXIS));
+                destPanel.setPreferredSize(new Dimension(200, 0));
+                pane.add(destPanel, BorderLayout.WEST);
+                destPanel.add(new JLabel("Destination Classes"));
+                
+                m_destClasses = new ClassSelector(ClassSelector.DeobfuscatedClassEntryComparator);
+                m_destClasses.setListener(new ClassSelectionListener() {
+                        @Override
+                        public void onSelectClass(ClassEntry classEntry) {
+                                setDestClass(classEntry);
+                        }
+                });
+                JScrollPane destScroller = new JScrollPane(m_destClasses);
+                destPanel.add(destScroller);
+                
+                JButton autoMatchButton = new JButton("AutoMatch");
+                autoMatchButton.addActionListener(new ActionListener() {
+                        @Override
+                        public void actionPerformed(ActionEvent event) {
+                                autoMatch();
+                        }
+                });
+                destPanel.add(autoMatchButton);
+                
+                // init source panels
+                DefaultSyntaxKit.initKit();
+                m_sourceReader = new CodeReader();
+                m_destReader = new CodeReader();
+                
+                // init all the splits
+                JSplitPane splitLeft = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, sourcePanel, new JScrollPane(m_sourceReader));
+                splitLeft.setResizeWeight(0); // let the right side take all the slack
+                JSplitPane splitRight = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, new JScrollPane(m_destReader), destPanel);
+                splitRight.setResizeWeight(1); // let the left side take all the slack
+                JSplitPane splitCenter = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, splitLeft, splitRight);
+                splitCenter.setResizeWeight(0.5); // resize 50:50
+                pane.add(splitCenter, BorderLayout.CENTER);
+                splitCenter.resetToPreferredSizes();
+                
+                // init bottom panel
+                JPanel bottomPanel = new JPanel();
+                bottomPanel.setLayout(new FlowLayout());
+                
+                m_sourceClassLabel = new JLabel();
+                m_sourceClassLabel.setHorizontalAlignment(SwingConstants.RIGHT);
+                m_destClassLabel = new JLabel();
+                m_destClassLabel.setHorizontalAlignment(SwingConstants.LEFT);
+                
+                m_matchButton = new JButton();
+                
+                m_advanceCheck = new JCheckBox("Advance to next likely match");
+                m_advanceCheck.addActionListener(new ActionListener() {
+                        @Override
+                        public void actionPerformed(ActionEvent event) {
+                                if (m_advanceCheck.isSelected()) {
+                                        advance();
+                                }
+                        }
+                });
+                
+                bottomPanel.add(m_sourceClassLabel);
+                bottomPanel.add(m_matchButton);
+                bottomPanel.add(m_destClassLabel);
+                bottomPanel.add(m_advanceCheck);
+                pane.add(bottomPanel, BorderLayout.SOUTH);
+                
+                // show the frame
+                pane.doLayout();
+                m_frame.setSize(1024, 576);
+                m_frame.setMinimumSize(new Dimension(640, 480));
+                m_frame.setVisible(true);
+                m_frame.setDefaultCloseOperation(WindowConstants.DISPOSE_ON_CLOSE);
+                
+                // init state
+                updateDestMappings();
+                setSourceType(SourceType.getDefault());
+                updateMatchButton();
+                m_saveListener = null;
+        }
+        
+        public void setSaveListener(SaveListener val) {
+                m_saveListener = val;
+        }
+
+        private void updateDestMappings() {
+                
+                Mappings newMappings = MappingsConverter.newMappings(
+                        m_classMatches,
+                        m_sourceDeobfuscator.getMappings(),
+                        m_sourceDeobfuscator,
+                        m_destDeobfuscator
+                );
+                
+                // look for dropped mappings
+                MappingsChecker checker = new MappingsChecker(m_destDeobfuscator.getJarIndex());
+                checker.dropBrokenMappings(newMappings);
+                
+                // count them
+                int numDroppedFields = checker.getDroppedFieldMappings().size();
+                int numDroppedMethods = checker.getDroppedMethodMappings().size();
+                System.out.println(String.format(
+                        "%d mappings from matched classes don't match the dest jar:\n\t%5d fields\n\t%5d methods",
+                        numDroppedFields + numDroppedMethods,
+                        numDroppedFields,
+                        numDroppedMethods
+                ));
+                
+                m_destDeobfuscator.setMappings(newMappings);
+        }
+
+        protected void setSourceType(SourceType val) {
+                
+                // show the source classes
+                m_sourceType = val;
+                m_sourceClasses.setClasses(deobfuscateClasses(m_sourceType.getSourceClasses(m_classMatches), m_sourceDeobfuscator));
+                
+                // update counts
+                for (SourceType sourceType : SourceType.values()) {
+                        m_sourceTypeButtons.get(sourceType).setText(String.format("%s (%d)",
+                                sourceType.name(),
+                                sourceType.getSourceClasses(m_classMatches).size()
+                        ));
+                }
+        }
+        
+        private Collection<ClassEntry> deobfuscateClasses(Collection<ClassEntry> in, Deobfuscator deobfuscator) {
+                List<ClassEntry> out = Lists.newArrayList();
+                for (ClassEntry entry : in) {
+                        
+                        ClassEntry deobf = deobfuscator.deobfuscateEntry(entry);
+                        
+                        // make sure we preserve any scores
+                        if (entry instanceof ScoredClassEntry) {
+                                deobf = new ScoredClassEntry(deobf, ((ScoredClassEntry)entry).getScore());
+                        }
+                        
+                        out.add(deobf);
+                }
+                return out;
+        }
+
+        protected void setSourceClass(ClassEntry classEntry) {
+                
+                Runnable onGetDestClasses = null;
+                if (m_advanceCheck.isSelected()) {
+                        onGetDestClasses = new Runnable() {
+                                @Override
+                                public void run() {
+                                        pickBestDestClass();
+                                }
+                        };
+                }
+                
+                setSourceClass(classEntry, onGetDestClasses);
+        }
+        
+        protected void setSourceClass(ClassEntry classEntry, final Runnable onGetDestClasses) {
+
+                // update the current source class
+                m_sourceClass = classEntry;
+                m_sourceClassLabel.setText(m_sourceClass != null ? m_sourceClass.getName() : "");
+                
+                if (m_sourceClass != null) {
+                        
+                        // show the dest class(es)
+                        ClassMatch match = m_classMatches.getMatchBySource(m_sourceDeobfuscator.obfuscateEntry(m_sourceClass));
+                        assert(match != null);
+                        if (match.destClasses.isEmpty()) {
+                                
+                                m_destClasses.setClasses(null);
+                                
+                                // run in a separate thread to keep ui responsive
+                                new Thread() {
+                                        @Override
+                                        public void run() {
+                                                m_destClasses.setClasses(deobfuscateClasses(getLikelyMatches(m_sourceClass), m_destDeobfuscator));
+                                                m_destClasses.expandAll();
+                                                
+                                                if (onGetDestClasses != null) {
+                                                        onGetDestClasses.run();
+                                                }
+                                        }
+                                }.start();
+                                
+                        } else {
+                                
+                                m_destClasses.setClasses(deobfuscateClasses(match.destClasses, m_destDeobfuscator));
+                                m_destClasses.expandAll();
+                                
+                                if (onGetDestClasses != null) {
+                                        onGetDestClasses.run();
+                                }
+                        }
+                }
+                
+                setDestClass(null);
+                m_sourceReader.decompileClass(m_sourceClass, m_sourceDeobfuscator, new Runnable() {
+                        @Override
+                        public void run() {
+                                m_sourceReader.navigateToClassDeclaration(m_sourceClass);
+                        }
+                });
+                
+                updateMatchButton();
+        }
+
+        private Collection<ClassEntry> getLikelyMatches(ClassEntry sourceClass) {
+                
+                ClassEntry obfSourceClass = m_sourceDeobfuscator.obfuscateEntry(sourceClass);
+                
+                // set up identifiers
+                ClassNamer namer = new ClassNamer(m_classMatches.getUniqueMatches());
+                ClassIdentifier sourceIdentifier = new ClassIdentifier(
+                        m_sourceDeobfuscator.getJar(), m_sourceDeobfuscator.getJarIndex(),
+                        namer.getSourceNamer(), true
+                );
+                ClassIdentifier destIdentifier = new ClassIdentifier(
+                        m_destDeobfuscator.getJar(), m_destDeobfuscator.getJarIndex(),
+                        namer.getDestNamer(), true
+                );
+                
+                try {
+                        
+                        // rank all the unmatched dest classes against the source class
+                        ClassIdentity sourceIdentity = sourceIdentifier.identify(obfSourceClass);
+                        List<ClassEntry> scoredDestClasses = Lists.newArrayList();
+                        for (ClassEntry unmatchedDestClass : m_classMatches.getUnmatchedDestClasses()) {
+                                ClassIdentity destIdentity = destIdentifier.identify(unmatchedDestClass);
+                                float score = 100.0f*(sourceIdentity.getMatchScore(destIdentity) + destIdentity.getMatchScore(sourceIdentity))
+                                        /(sourceIdentity.getMaxMatchScore() + destIdentity.getMaxMatchScore());
+                                scoredDestClasses.add(new ScoredClassEntry(unmatchedDestClass, score));
+                        }
+                        return scoredDestClasses;
+                        
+                } catch (ClassNotFoundException ex) {
+                        throw new Error("Unable to find class " + ex.getMessage());
+                }
+        }
+        
+        protected void setDestClass(ClassEntry classEntry) {
+                
+                // update the current source class
+                m_destClass = classEntry;
+                m_destClassLabel.setText(m_destClass != null ? m_destClass.getName() : "");
+                
+                m_destReader.decompileClass(m_destClass, m_destDeobfuscator, new Runnable() {
+                        @Override
+                        public void run() {
+                                m_destReader.navigateToClassDeclaration(m_destClass);
+                        }
+                });
+                
+                updateMatchButton();
+        }
+        
+        private void updateMatchButton() {
+                
+                ClassEntry obfSource = m_sourceDeobfuscator.obfuscateEntry(m_sourceClass);
+                ClassEntry obfDest = m_destDeobfuscator.obfuscateEntry(m_destClass);
+                
+                BiMap<ClassEntry,ClassEntry> uniqueMatches = m_classMatches.getUniqueMatches();
+                boolean twoSelected = m_sourceClass != null && m_destClass != null;
+                boolean isMatched = uniqueMatches.containsKey(obfSource) && uniqueMatches.containsValue(obfDest);
+                boolean canMatch = !uniqueMatches.containsKey(obfSource) && ! uniqueMatches.containsValue(obfDest);
+                
+                GuiTricks.deactivateButton(m_matchButton);
+                if (twoSelected) {
+                        if (isMatched) {
+                                GuiTricks.activateButton(m_matchButton, "Unmatch", new ActionListener() {
+                                        @Override
+                                        public void actionPerformed(ActionEvent event) {
+                                                onUnmatchClick();
+                                        }
+                                });
+                        } else if (canMatch) {
+                                GuiTricks.activateButton(m_matchButton, "Match", new ActionListener() {
+                                        @Override
+                                        public void actionPerformed(ActionEvent event) {
+                                                onMatchClick();
+                                        }
+                                });
+                        }
+                }
+        }
+
+        private void onMatchClick() {
+                // precondition: source and dest classes are set correctly
+                
+                ClassEntry obfSource = m_sourceDeobfuscator.obfuscateEntry(m_sourceClass);
+                ClassEntry obfDest = m_destDeobfuscator.obfuscateEntry(m_destClass);
+                
+                // remove the classes from their match
+                m_classMatches.removeSource(obfSource);
+                m_classMatches.removeDest(obfDest);
+                
+                // add them as matched classes
+                m_classMatches.add(new ClassMatch(obfSource, obfDest));
+                
+                ClassEntry nextClass = null;
+                if (m_advanceCheck.isSelected()) {
+                        nextClass = m_sourceClasses.getNextClass(m_sourceClass);
+                }
+                
+                save();
+                updateMatches();
+                
+                if (nextClass != null) {
+                        advance(nextClass);
+                }
+        }
+        
+        private void onUnmatchClick() {
+                // precondition: source and dest classes are set to a unique match
+                
+                ClassEntry obfSource = m_sourceDeobfuscator.obfuscateEntry(m_sourceClass);
+                
+                // remove the source to break the match, then add the source back as unmatched
+                m_classMatches.removeSource(obfSource);
+                m_classMatches.add(new ClassMatch(obfSource, null));
+                
+                save();
+                updateMatches();
+        }
+        
+        private void updateMatches() {
+                updateDestMappings();
+                setDestClass(null);
+                m_destClasses.setClasses(null);
+                updateMatchButton();
+                
+                // remember where we were in the source tree
+                String packageName = m_sourceClasses.getSelectedPackage();
+                
+                setSourceType(m_sourceType);
+                
+                m_sourceClasses.expandPackage(packageName);
+        }
+        
+        private void save() {
+                if (m_saveListener != null) {
+                        m_saveListener.save(m_classMatches);
+                }
+        }
+        
+        private void autoMatch() {
+                
+                System.out.println("Automatching...");
+                
+                // compute a new matching
+                ClassMatching matching = MappingsConverter.computeMatching(
+                        m_sourceDeobfuscator.getJar(), m_sourceDeobfuscator.getJarIndex(),
+                        m_destDeobfuscator.getJar(), m_destDeobfuscator.getJarIndex(),
+                        m_classMatches.getUniqueMatches()
+                );
+                ClassMatches newMatches = new ClassMatches(matching.matches());
+                System.out.println(String.format("Automatch found %d new matches",
+                        newMatches.getUniqueMatches().size() - m_classMatches.getUniqueMatches().size()
+                ));
+                
+                // update the current matches
+                m_classMatches = newMatches;
+                save();
+                updateMatches();
+        }
+        
+        private void advance() {
+                advance(null);
+        }
+        
+        private void advance(ClassEntry sourceClass) {
+
+                // make sure we have a source class
+                if (sourceClass == null) {
+                        sourceClass = m_sourceClasses.getSelectedClass();
+                        if (sourceClass != null) {
+                                sourceClass = m_sourceClasses.getNextClass(sourceClass);
+                        } else {
+                                sourceClass = m_sourceClasses.getFirstClass();
+                        }
+                }
+                
+                // set the source class
+                setSourceClass(sourceClass, new Runnable() {
+                        @Override
+                        public void run() {
+                                pickBestDestClass();
+                        }
+                });
+                m_sourceClasses.setSelectionClass(sourceClass);
+        }
+        
+        private void pickBestDestClass() {
+                
+                // then, pick the best dest class
+                ClassEntry firstClass = null;
+                ScoredClassEntry bestDestClass = null;
+                for (ClassSelectorPackageNode packageNode : m_destClasses.packageNodes()) {
+                        for (ClassSelectorClassNode classNode : m_destClasses.classNodes(packageNode)) {
+                                if (firstClass == null) {
+                                        firstClass = classNode.getClassEntry();
+                                }
+                                if (classNode.getClassEntry() instanceof ScoredClassEntry) {
+                                        ScoredClassEntry scoredClass = (ScoredClassEntry)classNode.getClassEntry();
+                                        if (bestDestClass == null || bestDestClass.getScore() < scoredClass.getScore()) {
+                                                bestDestClass = scoredClass;
+                                        }
+                                }
+                        }
+                }
+                
+                // pick the entry to show
+                ClassEntry destClass = null;
+                if (bestDestClass != null) {
+                        destClass = bestDestClass;
+                } else if (firstClass != null) {
+                        destClass = firstClass;
+                }
+                
+                setDestClass(destClass);
+                m_destClasses.setSelectionClass(destClass);
+        }
+}