Merge pull request #22 from nccgroup/zehuan-dev

Export, import and clear tabs

src/trust/nccgroup/decoderimproved/ConfigPanel.java

@@ -0,0 +1,82 @@
+package trust.nccgroup.decoderimproved;
+
+import javax.swing.*;
+import java.awt.*;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+
+public class ConfigPanel extends JPanel {
+    JButton exportButton;
+    JButton loadButton;
+    JToggleButton clearButton;
+
+    public ConfigPanel(ExtensionRoot extensionRoot) {
+        this.setLayout(new FlowLayout(FlowLayout.LEFT));
+        exportButton = new JButton("Export all tabs to file");
+        loadButton = new JButton("Load tabs from file");
+        String clearButtonText = "Clear all tabs on exit";
+        clearButton = new JToggleButton(clearButtonText);
+
+        this.add(exportButton);
+        this.add(loadButton);
+        this.add(clearButton);
+
+        // Listeners
+        exportButton.addActionListener((e) -> {
+            try {
+                JFileChooser fileChooser = new JFileChooser();
+                fileChooser.setDialogTitle("Save all data to...");
+                // Grab focus to save file dialog
+                fileChooser.addHierarchyListener((_event) -> {
+                    grabFocus();
+                });
+                if (fileChooser.showSaveDialog(extensionRoot.multiDecoderTab) == JFileChooser.APPROVE_OPTION) {
+                    FileOutputStream fileOutputStream = new FileOutputStream(fileChooser.getSelectedFile());
+                    // Get state and write to file
+                    fileOutputStream.write(extensionRoot.multiDecoderTab.getState().getBytes());
+                    fileOutputStream.close();
+                }
+            } catch (Exception ee) {
+                Logger.printErrorFromException(ee);
+            }
+        });
+
+        loadButton.addActionListener((e) -> {
+            try {
+                JFileChooser fileChooser = new JFileChooser();
+                fileChooser.setDialogTitle("Load data from...");
+                // Grab focus to load file dialog
+                fileChooser.addHierarchyListener((_event) -> {
+                    grabFocus();
+                });
+                if (fileChooser.showOpenDialog(extensionRoot.multiDecoderTab) == JFileChooser.APPROVE_OPTION) {
+                    // Read file content
+                    File selectedFile = fileChooser.getSelectedFile();
+                    FileInputStream fileInputStream = new FileInputStream(selectedFile);
+                    byte[] fileContent = new byte[(int) selectedFile.length()];
+                    fileInputStream.read(fileContent);
+                    fileInputStream.close();
+                    extensionRoot.multiDecoderTab.setState(new String(fileContent), false);
+                }
+            } catch (Exception ee) {
+                Logger.printErrorFromException(ee);
+                JOptionPane.showMessageDialog(extensionRoot.multiDecoderTab, ee.getClass().getName() + ", please check extension errors for details", "Error loading file", JOptionPane.ERROR_MESSAGE);
+            }
+        });
+
+        clearButton.addItemListener((e) -> {
+            try {
+                if (clearButton.isSelected()) {
+                    extensionRoot.setClearState();
+                    clearButton.setText("ALL TABS will be CLEARED on exit");
+                } else {
+                    extensionRoot.setSaveFullState();
+                    clearButton.setText(clearButtonText);
+                }
+            } catch (Exception ee) {
+                Logger.printErrorFromException(ee);
+            }
+        });
+    }
+}

src/trust/nccgroup/decoderimproved/ExtensionRoot.java

@@ -7,6 +7,8 @@ public class ExtensionRoot implements IBurpExtender {
     private IBurpExtenderCallbacks callbacks;
     private IExtensionHelpers helpers;
 
+    MultiDecoderTab multiDecoderTab;
+
     public void registerExtenderCallbacks(IBurpExtenderCallbacks _callbacks) {
 
         callbacks = _callbacks;
@@ -16,14 +18,25 @@ public void registerExtenderCallbacks(IBurpExtenderCallbacks _callbacks) {
 
         callbacks.setExtensionName("Decoder Improved");
 
-        MultiDecoderTab multiDecoderTab = new MultiDecoderTab(callbacks);
+        multiDecoderTab = new MultiDecoderTab(this);
         //callbacks.customizeUiComponent(multiDecoderTab);
         callbacks.addSuiteTab(multiDecoderTab);
-        callbacks.registerContextMenuFactory(new SendToDecoderImprovedContextMenuFactory(callbacks, multiDecoderTab));
+        callbacks.registerContextMenuFactory(new SendToDecoderImprovedContextMenuFactory(multiDecoderTab));
 
         String savedSettings = callbacks.loadExtensionSetting(multiDecoderTab.getTabCaption());
-        multiDecoderTab.setState(savedSettings);
+        // null state will be handled in MultiDecoderTab
+        multiDecoderTab.setState(savedSettings, true);
+
+        setSaveFullState();
+    }
 
+    public void setSaveFullState() {
+        callbacks.getExtensionStateListeners().forEach((x) -> callbacks.removeExtensionStateListener(x));
         callbacks.registerExtensionStateListener(() -> callbacks.saveExtensionSetting(multiDecoderTab.getTabCaption(), multiDecoderTab.getState()));
     }
+
+    public void setClearState() {
+        callbacks.getExtensionStateListeners().forEach((x) -> callbacks.removeExtensionStateListener(x));
+        callbacks.registerExtensionStateListener(() -> callbacks.saveExtensionSetting(multiDecoderTab.getTabCaption(), null));
+    }
 }

src/trust/nccgroup/decoderimproved/Logger.java

@@ -2,6 +2,7 @@
 
 import burp.IBurpExtenderCallbacks;
 
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.time.LocalDateTime;
@@ -29,10 +30,16 @@ public static void printError(String errorString){
 
     public static void printErrorFromException(Exception e) {
         if (callbacks != null) {
-            StringWriter stringWriter = new StringWriter();
-            PrintWriter printWriter = new PrintWriter(stringWriter);
-            e.printStackTrace(printWriter);
-            callbacks.printError(stringWriter.toString());
+            try {
+                StringWriter stringWriter = new StringWriter();
+                PrintWriter printWriter = new PrintWriter(stringWriter);
+                e.printStackTrace(printWriter);
+                callbacks.printError(stringWriter.toString());
+                printWriter.close();
+                stringWriter.close();
+            } catch (IOException ee){
+                printError(ee.getMessage());
+            }
         }
     }
 

src/trust/nccgroup/decoderimproved/MultiDecoderTab.java

@@ -1,7 +1,5 @@
 package trust.nccgroup.decoderimproved;
 
-import burp.IBurpExtenderCallbacks;
-import burp.IExtensionHelpers;
 import burp.ITab;
 import com.google.common.collect.Lists;
 import com.google.gson.*;
@@ -24,13 +22,11 @@
 
 public class MultiDecoderTab extends JPanel implements ITab {
 
-    private IBurpExtenderCallbacks callbacks;
-
-    private IExtensionHelpers helpers;
-
     private JTabbedPane main;
     private JPanel newTabButton;
 
+    private ConfigPanel configPanel;
+
     boolean tabChangeListenerLock = false;
 
     //Plugin starts with one decoder tab open and the "new tab" tab
@@ -44,12 +40,10 @@ public void setTabChangeListenerLock(boolean tabChangeListenerLock) {
         this.tabChangeListenerLock = tabChangeListenerLock;
     }
 
-    public MultiDecoderTab(IBurpExtenderCallbacks _callbacks) {
+    public MultiDecoderTab(ExtensionRoot extensionRoot) {
         // Set main tab layout
         setLayout(new BorderLayout());
         //initialize ui elements
-        callbacks = _callbacks;
-        helpers = callbacks.getHelpers();
         main = new JTabbedPane();
 
         // Add "new tab" tab
@@ -74,6 +68,9 @@ public MultiDecoderTab(IBurpExtenderCallbacks _callbacks) {
             }
         });
         add(main, BorderLayout.CENTER);
+
+        configPanel = new ConfigPanel(extensionRoot);
+        add(configPanel, BorderLayout.SOUTH);
     }
 
     // Logic for adding new tabs
@@ -82,7 +79,6 @@ public void addTab() {
         // Add a new tab
         overallCount += 1;
         DecoderTab mt2 = new DecoderTab(Integer.toString(overallCount, 10), this);
-        //callbacks.customizeUiComponent(mt2);
         main.add(mt2);
         main.setTabComponentAt(main.indexOfComponent(mt2), mt2.getTabHandleElement());
         main.setSelectedComponent(mt2);
@@ -175,25 +171,33 @@ public String getState() {
     }
 
     // Decode the saved extension setting string and recover all tabs
-    public void setState(String stateString) {
+    public void setState(String stateString, boolean initial) {
         if (stateString == null || stateString.isEmpty()) {
-            addTab();
-            main.setSelectedIndex(0);
-            return;
+            if (initial) {
+                addTab();
+                main.setSelectedIndex(0);
+                return;
+            } else {
+                throw new IllegalArgumentException("Error reading file or file is empty");
+            }
         }
         try {
+            int originalIndex = initial ? 0 : main.getSelectedIndex();
+            int originalTabCount = main.getTabCount();
             JsonObject extensionStateObject = JsonParser.parseString(stateString).getAsJsonObject();
             JsonArray tabStateArray = extensionStateObject.get("t").getAsJsonArray();
             if (tabStateArray.size() == 0) {
-                addTab();
-                main.setSelectedIndex(0);
+                if (initial) {
+                    addTab();
+                    main.setSelectedIndex(0);
+                }
                 return;
             }
             for (int i = 0; i < tabStateArray.size(); i++) {
                 JsonObject tabStateObject = tabStateArray.get(i).getAsJsonObject();
                 // Build a new tab for each tab object
                 addTab();
-                DecoderTab dt = (DecoderTab) main.getComponentAt(i);
+                DecoderTab dt = (DecoderTab) main.getComponentAt(originalTabCount + i - 1);
                 dt.decoderTabHandle.tabName.setText(tabStateObject.get("n").getAsString());
                 DecoderSegmentState dsState = dt.getDecoderSegments().get(0).dsState;
                 dsState.setByteArrayList(Base64.getDecoder().decode(tabStateObject.get("b").getAsString()));
@@ -224,11 +228,15 @@ public void setState(String stateString) {
                     }
                 }
             }
-            main.setSelectedIndex(extensionStateObject.get("i").getAsInt());
+            main.setSelectedIndex(initial ? extensionStateObject.get("i").getAsInt() : originalIndex);
         } catch (Exception e) {
-            addTab();
-            main.setSelectedIndex(0);
-            Logger.printErrorFromException(e);
+            if (initial) {
+                addTab();
+                main.setSelectedIndex(0);
+                Logger.printErrorFromException(e);
+            } else {
+                throw e;
+            }
         }
     }
 
@@ -240,7 +248,6 @@ private static class DecoderTab extends JPanel {
 
         public DecoderTab(String _title, MultiDecoderTab _parent) {
             decoderTabHandle = new DecoderTabHandle(_title, _parent, this);
-            //_parent.callbacks.customizeUiComponent(decoderTabHandle);
             setupComponents();
         }
 

src/trust/nccgroup/decoderimproved/SendToDecoderImprovedContextMenuFactory.java

@@ -12,13 +12,9 @@
  * Created by j on 12/9/16.
  */
 class SendToDecoderImprovedContextMenuFactory implements IContextMenuFactory {
-    private IBurpExtenderCallbacks callbacks;
-    private IExtensionHelpers helpers;
     private MultiDecoderTab tab;
 
-    public SendToDecoderImprovedContextMenuFactory(IBurpExtenderCallbacks _callbacks, MultiDecoderTab _tab) {
-        callbacks = _callbacks;
-        helpers = callbacks.getHelpers();
+    public SendToDecoderImprovedContextMenuFactory(MultiDecoderTab _tab) {
         tab = _tab;
     }