Add API endpoints for getting and setting UI options

zaproxy · Dec 28, 2018 · b1ae3e8 · b1ae3e8
1 parent 29613bf
commit b1ae3e8
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 0 deletions.
diff --git a/src/main/java/org/zaproxy/zap/extension/hud/HudAPI.java b/src/main/java/org/zaproxy/zap/extension/hud/HudAPI.java
@@ -76,20 +76,26 @@ public class HudAPI extends ApiImplementor {
 
     private static final String PREFIX = "hud";
 
+    private static final int MAX_KEY_LENGTH = 50;
+
     private Map<String, String> siteUrls = new HashMap<String, String>();
     private ExtensionHUD extension;
 
     private static final String ACTION_LOG = "log";
     private static final String ACTION_RECORD_REQUEST = "recordRequest";
     private static final String ACTION_RESET_TUTORIAL_TASKS = "resetTutorialTasks";
+    private static final String ACTION_SET_UI_OPTION = "setUiOption";
 
+    private static final String VIEW_GET_UI_OPTION = "getUiOption";
     private static final String VIEW_HUD_ALERT_DATA = "hudAlertData";
     private static final String VIEW_HEARTBEAT = "heartbeat";
 
     private static final String PARAM_RECORD = "record";
     private static final String PARAM_HEADER = "header";
     private static final String PARAM_BODY = "body";
     private static final String PARAM_URL = "url";
+    private static final String PARAM_KEY = "key";
+    private static final String PARAM_VALUE = "value";
 
     /** The only files that can be included on domain */
     private static final List<String> DOMAIN_FILE_WHITELIST =
@@ -121,9 +127,15 @@ public HudAPI(ExtensionHUD extension) {
         this.addApiAction(
                 new ApiAction(ACTION_RECORD_REQUEST, new String[] {PARAM_HEADER, PARAM_BODY}));
         this.addApiAction(new ApiAction(ACTION_RESET_TUTORIAL_TASKS));
+        this.addApiAction(
+                new ApiAction(
+                        ACTION_SET_UI_OPTION,
+                        new String[] {PARAM_KEY},
+                        new String[] {PARAM_VALUE}));
 
         this.addApiView(new ApiView(VIEW_HUD_ALERT_DATA, new String[] {PARAM_URL}));
         this.addApiView(new ApiView(VIEW_HEARTBEAT));
+        this.addApiView(new ApiView(VIEW_GET_UI_OPTION, new String[] {PARAM_KEY}));
 
         hudFileProxy = new HudFileProxy(this);
         hudFileUrl = API.getInstance().getCallBackUrl(hudFileProxy, API.API_URL_S);
@@ -247,13 +259,28 @@ public ApiResponse handleApiAction(String name, JSONObject params) throws ApiExc
                 this.extension.resetTutorialTasks();
                 break;
 
+            case ACTION_SET_UI_OPTION:
+                String key = params.getString(PARAM_KEY);
+                String value = params.optString(PARAM_VALUE, "");
+
+                if (!isValidKey(key)) {
+                    throw new ApiException(ApiException.Type.ILLEGAL_PARAMETER, PARAM_KEY);
+                }
+
+                this.extension.getHudParam().setUiOption(key, value);
+                break;
+
             default:
                 throw new ApiException(ApiException.Type.BAD_ACTION);
         }
 
         return ApiResponseElement.OK;
     }
 
+    private boolean isValidKey(String key) {
+        return key.length() > 0 && key.length() < MAX_KEY_LENGTH && key.matches("[a-zA-Z0-9]+");
+    }
+
     @Override
     public ApiResponse handleApiView(String name, JSONObject params) throws ApiException {
 
@@ -267,6 +294,12 @@ public ApiResponse handleApiView(String name, JSONObject params) throws ApiExcep
             case VIEW_HEARTBEAT:
                 logger.debug("Received heartbeat");
                 return ApiResponseElement.OK;
+            case VIEW_GET_UI_OPTION:
+                String key = params.getString(PARAM_KEY);
+                if (!isValidKey(key)) {
+                    throw new ApiException(ApiException.Type.ILLEGAL_PARAMETER, PARAM_KEY);
+                }
+                return new ApiResponseElement(key, this.extension.getHudParam().getUiOption(key));
             default:
                 throw new ApiException(ApiException.Type.BAD_VIEW);
         }

diff --git a/src/main/java/org/zaproxy/zap/extension/hud/HudParam.java b/src/main/java/org/zaproxy/zap/extension/hud/HudParam.java
@@ -47,6 +47,7 @@ public class HudParam extends VersionedAbstractParam {
     private static final String PARAM_SHOW_WELCOME_SCREEN = PARAM_BASE_KEY + ".showWelcomeScreen";
     private static final String PARAM_ENABLE_ON_DOMAIN_MSGS =
             PARAM_BASE_KEY + ".enableOnDomainMsgs";
+    private static final String PARAM_UI_OPTION_PREFIX = PARAM_BASE_KEY + ".uiOption.";
 
     /**
      * The version of the configurations. Used to keep track of configurations changes between
@@ -281,4 +282,17 @@ public void resetTutorialTasks() {
             log.error(e.getMessage(), e);
         }
     }
+
+    public void setUiOption(String key, String value) {
+        getConfig().setProperty(PARAM_UI_OPTION_PREFIX + key, value);
+        try {
+            this.getConfig().save();
+        } catch (ConfigurationException e) {
+            log.error(e.getMessage(), e);
+        }
+    }
+
+    public String getUiOption(String key) {
+        return getConfig().getString(PARAM_UI_OPTION_PREFIX + key, "");
+    }
 }
diff --git a/src/main/resources/org/zaproxy/zap/extension/hud/resources/Messages.properties b/src/main/resources/org/zaproxy/zap/extension/hud/resources/Messages.properties
@@ -8,6 +8,8 @@ hud.api.action.setOptionDevelopmentMode = Sets the boolean option 'Development m
 hud.api.action.setOptionEnabled = Sets whether the HUD is enabled or not
 hud.api.action.setOptionInScopeOnly = Sets the boolean option 'In scope only'
 hud.api.action.setOptionRemoveCSP = Sets the boolean option 'Remove CSP'
+hud.api.action.setUiOption = Sets a UI option with the given key and value. The key must be 50 alphanumeric characters or less
+hud.api.view.getUiOption = Gets a UI option set using setUiOption. If the supplied key is not set then an empty string is returned
 hud.api.view.optionAllowUnsafeEval = Returns true if the 'Allow unsafe eval' option is set
 hud.api.view.optionBaseDirectory = Returns the base directory from which the HUD files are loaded
 hud.api.view.optionDevelopmentMode = Returns true if the 'Development mode' option is set