Add support for per-tool signing configuration, configuration saving + minor bug fixes

src/main/java/burp/BurpExtender.java

@@ -8,19 +8,39 @@
 import com.netspi.awssigner.signing.DelegatingAwsRequestSigner;
 import com.netspi.awssigner.signing.ParsedAuthHeader;
 import com.netspi.awssigner.signing.SigningException;
+import com.netspi.awssigner.utils.AWSSignerUtils;
 import com.netspi.awssigner.view.BurpUIComponentCustomizer;
 import com.netspi.awssigner.view.BurpTabPanel;
 import java.awt.Component;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import javax.swing.JMenuItem;
 import javax.swing.SwingUtilities;
 
 //This is the Burp primary class. It needs to live in this package and have this name
 public class BurpExtender implements IBurpExtender, ITab, IHttpListener, IContextMenuFactory {
 
+    static {
+        final Map<Integer, String> tempMap = new HashMap<Integer, String>();
+
+        tempMap.put(IBurpExtenderCallbacks.TOOL_SUITE, "All");
+        tempMap.put(IBurpExtenderCallbacks.TOOL_EXTENDER, "Extensions");
+        tempMap.put(IBurpExtenderCallbacks.TOOL_INTRUDER, "Intruder");
+        tempMap.put(IBurpExtenderCallbacks.TOOL_PROXY, "Proxy");
+        tempMap.put(IBurpExtenderCallbacks.TOOL_REPEATER, "Repeater");
+        tempMap.put(IBurpExtenderCallbacks.TOOL_SCANNER, "Scanner");
+        tempMap.put(IBurpExtenderCallbacks.TOOL_SEQUENCER, "Sequencer");
+        tempMap.put(IBurpExtenderCallbacks.TOOL_TARGET, "Target");
+
+        TOOL_FLAG_TRANSLATION_MAP = Collections.unmodifiableMap(tempMap);
+    }
     public static final String EXTENSION_NAME = "AWS Signer";
 
+    private static final Map<Integer, String> TOOL_FLAG_TRANSLATION_MAP;
+
     private BurpTabPanel view;
     private AWSSignerConfiguration model;
     private AWSSignerController controller;
@@ -33,8 +53,14 @@ public void registerExtenderCallbacks(IBurpExtenderCallbacks callbacks) {
 
         //Save callbacks and helpers for later reference
         this.callbacks = callbacks;
+        AWSSignerUtils.setBurpExtenderCallbacks(callbacks);
         helpers = callbacks.getHelpers();
 
+        // Extension unload/shutdown callback
+        this.callbacks.registerExtensionStateListener(()-> {
+            this.model.persist();
+        });
+
         //Setup styling
         BurpUIComponentCustomizer.setBurpStyler((Component component) -> {
             callbacks.customizeUiComponent(component);
@@ -47,7 +73,7 @@ public void registerExtenderCallbacks(IBurpExtenderCallbacks callbacks) {
         //Create the view
         view = new BurpTabPanel();
         //Create the model
-        model = new AWSSignerConfiguration();
+        model = AWSSignerConfiguration.getOrCreateProjectConfiguration();
         //Create controller to keep them in sync
         controller = new AWSSignerController(view, model);
 
@@ -99,6 +125,12 @@ public void processHttpMessage(int toolFlag, boolean messageIsRequest, IHttpRequ
             return;
         }
 
+        // Is the request from a tool we want to sign?
+        if ((model.signForTools & toolFlag) == 0 && (model.signForTools & IBurpExtenderCallbacks.TOOL_SUITE) == 0) {
+            LogWriter.logDebug("Signing for requests from " + TOOL_FLAG_TRANSLATION_MAP.get(toolFlag) + " is not enabled. Ignoring Message.");
+            return;
+        }
+
         //Could be a request we want to sign. Let's analyze it
         IRequestInfo request = helpers.analyzeRequest(messageInfo);
 

src/main/java/com/netspi/awssigner/controller/AWSSignerController.java

@@ -1,6 +1,7 @@
 package com.netspi.awssigner.controller;
 
 import burp.BurpExtender;
+import burp.IBurpExtenderCallbacks;
 import burp.IContextMenuInvocation;
 import com.netspi.awssigner.credentials.SigningCredentials;
 import com.google.gson.GsonBuilder;
@@ -33,7 +34,6 @@
 import java.awt.Component;
 import java.awt.event.MouseAdapter;
 import java.awt.event.MouseEvent;
-import java.awt.event.MouseListener;
 import java.io.IOException;
 import java.nio.file.Path;
 import java.util.ArrayList;
@@ -66,8 +66,42 @@ public AWSSignerController(BurpTabPanel view, AWSSignerConfiguration model) {
 
         initListeners();
 
-        //Add the initial profile
-        addProfile(new StaticCredentialsProfile(INIT_PROFILE_NAME));
+        //Add the initial profile if no profiles exist already
+
+        if (this.model.profiles.isEmpty()) {
+            addProfile(new StaticCredentialsProfile(INIT_PROFILE_NAME));
+        } else {
+            this.view.profileList.setSelectedIndex(0);
+        }
+
+        syncViewWithModel();
+    }
+
+    private void syncViewWithModel(){
+
+        //Setup profile list
+        DefaultListModel listModel = resetProfileList();
+        view.profileList.setSelectedIndex(0);
+
+        //Reset "Always Sign With" combo box
+        resetAlwaysSignWithProfileComboBox();
+
+        initializeProfileConfigurationTab(model.profiles.get(0));
+
+        // Hack to preserve the state of the flag bits, since the handlers for setSelected flip the bits currently.
+
+        int tmp = model.signForTools;
+        this.view.persistConfigurationCheckbox.setSelected(model.shouldPersist);
+        this.view.signingEnabledCheckbox.setSelected(model.isEnabled);
+        this.view.signForAllCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_SUITE) != 0);
+        this.view.signForProxyCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_PROXY) != 0);
+        this.view.signForIntruderCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_INTRUDER) != 0);
+        this.view.signForExtensionsCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_EXTENDER) != 0);
+        this.view.signForRepeaterCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_REPEATER) != 0);
+        this.view.signForTargetCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_TARGET) != 0);
+        this.view.signForScannerCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_SCANNER) != 0);
+        this.view.signForSequencerCheckbox.setSelected((model.signForTools & IBurpExtenderCallbacks.TOOL_SEQUENCER) != 0);
+        model.signForTools = tmp;
     }
 
     private void initListeners() {
@@ -78,6 +112,13 @@ private void initListeners() {
             logInfo("Signing Enabled: " + model.isEnabled);
         });
 
+        //Configuration persistence checkbox
+        view.persistConfigurationCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("Persist Configuration Checkbox State Change.");
+            model.shouldPersist = (e.getStateChange() == ItemEvent.SELECTED);
+            logInfo("Persist Configuration Enabled: " + model.shouldPersist);
+        });
+
         //"Always Sign With" profile selection combox box
         alwaysSignWithComboBoxProfileSelectionListener = new ComboBoxProfileSelectionListener(model, "Always Sign With", (Profile profile) -> {
             logDebug("Setting \"Always Sign With\" Profile to: " + profile);
@@ -96,6 +137,62 @@ private void initListeners() {
             }
         }));
 
+        //Sign for all check box
+        view.signForAllCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_SUITE;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
+        //Sign for repeater check box
+        view.signForRepeaterCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_REPEATER;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
+        //Sign for proxy check box
+        view.signForProxyCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_PROXY;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
+        //Sign for intruder check box
+        view.signForIntruderCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_INTRUDER;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
+        //Sign target check box
+        view.signForTargetCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_TARGET;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
+        //Sign for extensions check box
+        view.signForExtensionsCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_EXTENDER;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
+        //Sign for sequencer check box
+        view.signForSequencerCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_SEQUENCER;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
+        //Sign for scanner check box
+        view.signForScannerCheckbox.addItemListener((ItemEvent e) -> {
+            logDebug("signForTools before: " + Integer.toBinaryString(model.signForTools));
+            model.signForTools = model.signForTools ^ IBurpExtenderCallbacks.TOOL_SCANNER;
+            logDebug("signForTools after: " + Integer.toBinaryString(model.signForTools));
+        });
+
         //Add button
         view.addProfileButton.addActionListener(((ActionEvent e) -> {
             logDebug("Add Profile Button Clicked.");
@@ -268,6 +365,31 @@ private void initListeners() {
             }
 
             Profile profile = currentProfileOptional.get();
+
+            // Bug fix to use the updated config if the input text field still has focus
+            if (profile instanceof StaticCredentialsProfile) {
+                if (view.staticAccessKeyTextField.hasFocus()) {
+                    ((StaticCredentialsProfile) profile).setAccessKey(view.staticAccessKeyTextField.getText());
+                }
+                if (view.staticSecretKeyTextField.hasFocus()) {
+                    ((StaticCredentialsProfile) profile).setSecretKey(view.staticSecretKeyTextField.getText());
+                }
+                if (view.staticSessionTokenTextField.hasFocus()) {
+                    ((StaticCredentialsProfile) profile).setSessionToken(view.staticSessionTokenTextField.getText());
+                }
+            } else if (profile instanceof CommandProfile) {
+                if (view.commandCommandTextField.hasFocus()) {
+                    ((CommandProfile) profile).setCommand(view.commandCommandTextField.getText());
+                }
+                if (view.commandDurationTextField.hasFocus()) {
+                    ((CommandProfile) profile).setDurationSecondsFromText(view.commandDurationTextField.getText());
+                }
+            } else if ( profile instanceof AssumeRoleProfile) {
+                if (view.assumeRoleRoleArnTextField.hasFocus()) {
+                    ((AssumeRoleProfile) profile).setRoleArn(view.assumeRoleRoleArnTextField.getText());
+                }
+            }
+
             //Check if we even have enough information to test this profile
             if (!profile.requiredFieldsAreSet()) {
                 logDebug("Profile " + profile.getName() + " does not have all required fields");

src/main/java/com/netspi/awssigner/model/AWSSignerConfiguration.java

@@ -1,14 +1,42 @@
 package com.netspi.awssigner.model;
 
+import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
 
-public class AWSSignerConfiguration {
+import com.netspi.awssigner.utils.AWSSignerUtils;
+
+import burp.IBurpExtenderCallbacks;
+
+public class AWSSignerConfiguration implements Serializable {
     public volatile boolean isEnabled =  true;
     public volatile List<Profile> profiles = new ArrayList<>();
     public volatile Profile alwaysSignWithProfile;
+    public volatile int signForTools = IBurpExtenderCallbacks.TOOL_SUITE;
+    public volatile boolean shouldPersist = true;
+
+    public static final String STORED_CONFIG_OBJECT_KEY = "aws-signer-configuration";
+
+    public synchronized void persist() {
+        if (shouldPersist) {
+            AWSSignerUtils.storeObjectForCurrentProject(STORED_CONFIG_OBJECT_KEY, this);
+        }
+    }
+
+    public static AWSSignerConfiguration getOrCreateProjectConfiguration() {
+
+        AWSSignerConfiguration config = new AWSSignerConfiguration();
+
+        Object o = AWSSignerUtils.getStoredObjectForCurrentProject(STORED_CONFIG_OBJECT_KEY);
+
+        if ( o != null && o instanceof AWSSignerConfiguration) {
+            config = (AWSSignerConfiguration) o;
+        }
+
+        return config;
+    }
 
     public List<String> getProfileNames() {
         return profiles.stream().map(Profile::getName).collect(Collectors.toList());

src/main/java/com/netspi/awssigner/model/Profile.java

@@ -2,10 +2,12 @@
 
 import com.netspi.awssigner.credentials.SignerCredentialException;
 import com.netspi.awssigner.credentials.SigningCredentials;
+
+import java.io.Serializable;
 import java.util.Objects;
 import java.util.Optional;
 
-public abstract class Profile {
+public abstract class Profile implements Serializable {
 
     protected String name;
     protected boolean isEnabled = true;