OAST Trim database

Signed-off-by: Simon Bennetts <[email protected]>
zaproxy · Sep 5, 2024 · da4ef96 · da4ef96
1 parent 42fedc8
commit da4ef96
Show file tree

Hide file tree

Showing 13 changed files with 257 additions and 16 deletions.
diff --git a/addOns/database/src/main/java/org/zaproxy/addon/database/Database.java b/addOns/database/src/main/java/org/zaproxy/addon/database/Database.java
@@ -27,6 +27,7 @@
 import javax.jdo.JDOHelper;
 import javax.jdo.PersistenceManager;
 import javax.jdo.PersistenceManagerFactory;
+import javax.jdo.Query;
 import javax.jdo.Transaction;
 import org.datanucleus.PropertyNames;
 import org.datanucleus.api.jdo.JDOPersistenceManagerFactory;
@@ -65,6 +66,9 @@ protected Database(String persistenceUnitName, ClassLoader classLoader) {
         jdoProperties.put(PropertyNames.PROPERTY_SCHEMA_VALIDATE_COLUMNS, false);
         jdoProperties.put(PropertyNames.PROPERTY_SCHEMA_VALIDATE_CONSTRAINTS, false);
 
+        // Required for non "select" SQL statements.
+        jdoProperties.put(PropertyNames.PROPERTY_QUERY_SQL_ALLOWALL, true);
+
         pmf = JDOHelper.getPersistenceManagerFactory(jdoProperties, classLoader);
         this.classLoader = classLoader;
     }
@@ -104,6 +108,17 @@ public void persistEntity(Object entity) {
         }
     }
 
+    public Object runQuery(String sql, Class<?> clazz, boolean unique) {
+        PersistenceManager pm = pmf.getPersistenceManager();
+
+        Query<?> query = pm.newQuery("javax.jdo.query.SQL", sql);
+        if (clazz != null) {
+            query.setResultClass(clazz);
+        }
+        query.setUnique(unique);
+        return query.execute();
+    }
+
     public <T> List<T> getAll(Class<T> clazz) {
         if (clazz == null) {
             throw new IllegalArgumentException("Class cannot be null.");

diff --git a/addOns/oast/CHANGELOG.md b/addOns/oast/CHANGELOG.md
@@ -9,6 +9,7 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 ### Added 
 - API support.
 - Raise alerts for OAST interactions that happened in other sessions.
+- Options to trim the database.
 
 ### Changed
 - Depend on newer version of Database add-on.

diff --git a/addOns/oast/src/main/java/org/zaproxy/addon/oast/ExtensionOast.java b/addOns/oast/src/main/java/org/zaproxy/addon/oast/ExtensionOast.java
@@ -75,6 +75,7 @@ public class ExtensionOast extends ExtensionAdaptor {
     private final Map<String, OastService> services = new HashMap<>();
 
     private OastOptionsPanel oastOptionsPanel;
+    private BoastOptionsPanelTab boastOptionsPanelTab;
     private OastPanel oastPanel;
     private OastParam oastParam;
     private BoastService boastService;
@@ -129,8 +130,8 @@ public void hook(ExtensionHook extensionHook) {
 
         if (hasView()) {
             extensionHook.getHookView().addOptionPanel(getOastOptionsPanel());
-            getOastOptionsPanel().addServicePanel(new GeneralOastOptionsPanelTab());
-            getOastOptionsPanel().addServicePanel(new BoastOptionsPanelTab(boastService));
+            getOastOptionsPanel().addServicePanel(new GeneralOastOptionsPanelTab(this));
+            getOastOptionsPanel().addServicePanel(getBoastOptionsPanelTab());
             getOastOptionsPanel().addServicePanel(new CallbackOptionsPanelTab(callbackService));
             getOastOptionsPanel().addServicePanel(new InteractshOptionsPanelTab(interactshService));
             extensionHook.getHookMenu().addPopupMenuItem(new OastInsertPayloadMenu(this));
@@ -154,13 +155,26 @@ public void optionsLoaded() {
     @Override
     public void postInit() {
         if (oastParam.isUsePermanentDatabase()) {
-            getPermanentDatabase();
+            trimDatabase(oastParam.getDaysToKeepRecords());
         }
+
         boastService.startService();
         callbackService.startService();
         interactshService.startService();
     }
 
+    public void trimDatabase(int days) {
+        getPermanentDatabase().trim(days);
+    }
+
+    public void clearAllRecords() {
+        getPermanentDatabase().clearAllRecords();
+        boastService.clearRegisteredServers();
+        if (hasView()) {
+            ThreadUtils.invokeAndWaitHandled(() -> getBoastOptionsPanelTab().resetBoastServers());
+        }
+    }
+
     private void optionsChanged(OptionsParam optionsParam) {
         getOastServices().values().forEach(OastService::fireOastStateChanged);
         if (!wasUsePermanentDatabase && oastParam.isUsePermanentDatabase()) {
@@ -240,6 +254,13 @@ private OastOptionsPanel getOastOptionsPanel() {
         return oastOptionsPanel;
     }
 
+    private BoastOptionsPanelTab getBoastOptionsPanelTab() {
+        if (boastOptionsPanelTab == null) {
+            boastOptionsPanelTab = new BoastOptionsPanelTab(boastService);
+        }
+        return boastOptionsPanelTab;
+    }
+
     public OastPanel getOastPanel() {
         if (oastPanel == null) {
             oastPanel = new OastPanel(this);

diff --git a/addOns/oast/src/main/java/org/zaproxy/addon/oast/OastApi.java b/addOns/oast/src/main/java/org/zaproxy/addon/oast/OastApi.java
@@ -35,11 +35,14 @@ public class OastApi extends ApiImplementor {
     private static final String PREFIX = "oast";
 
     private static final String ACTION_SET_ACTIVE_SCAN_SERVICE = "setActiveScanService";
+    private static final String ACTION_SET_DAYS_TO_KEEP_RECORDS = "setDaysToKeepRecords";
     private static final String ACTION_SET_BOAST_OPTIONS = "setBoastOptions";
     private static final String ACTION_SET_CALLBACK_OPTIONS = "setCallbackOptions";
     private static final String ACTION_SET_INTERACTSH_OPTIONS = "setInteractshOptions";
+
     private static final String VIEW_GET_ACTIVE_SCAN_SERVICE = "getActiveScanService";
     private static final String VIEW_GET_SERVICES = "getServices";
+    private static final String VIEW_GET_DAYS_TO_KEEP_RECORDS = "getDaysToKeepRecords";
     private static final String VIEW_GET_BOAST_OPTIONS = "getBoastOptions";
     private static final String VIEW_GET_CALLBACK_OPTIONS = "getCallbackOptions";
     private static final String VIEW_GET_INTERACTSH_OPTIONS = "getInteractshOptions";
@@ -51,6 +54,7 @@ public class OastApi extends ApiImplementor {
     private static final String PARAM_SERVER = "server";
     private static final String PARAM_POLL_IN_SECS = "pollInSecs";
     private static final String PARAM_PORT = "port";
+    private static final String PARAM_DAYS = "days";
 
     private ExtensionOast ext;
 
@@ -68,6 +72,8 @@ public OastApi(ExtensionOast ext) {
         this.addApiView(new ApiView(VIEW_GET_BOAST_OPTIONS));
         this.addApiView(new ApiView(VIEW_GET_CALLBACK_OPTIONS));
         this.addApiView(new ApiView(VIEW_GET_INTERACTSH_OPTIONS));
+        this.addApiView(new ApiView(VIEW_GET_DAYS_TO_KEEP_RECORDS));
+
         this.addApiAction(new ApiAction(ACTION_SET_ACTIVE_SCAN_SERVICE, new String[] {PARAM_NAME}));
         this.addApiAction(
                 new ApiAction(
@@ -80,6 +86,8 @@ public OastApi(ExtensionOast ext) {
                 new ApiAction(
                         ACTION_SET_INTERACTSH_OPTIONS,
                         new String[] {PARAM_SERVER, PARAM_POLL_IN_SECS, PARAM_AUTH_TOKEN}));
+        this.addApiAction(
+                new ApiAction(ACTION_SET_DAYS_TO_KEEP_RECORDS, new String[] {PARAM_DAYS}));
     }
 
     @Override
@@ -98,6 +106,12 @@ public ApiResponse handleApiAction(String name, JSONObject params) throws ApiExc
                             ApiException.Type.ILLEGAL_PARAMETER, params.getString(PARAM_NAME));
                 }
                 break;
+            case ACTION_SET_DAYS_TO_KEEP_RECORDS:
+                int days = ApiUtils.getIntParam(params, PARAM_DAYS);
+                ext.getParams().setDaysToKeepRecords(days);
+                ext.trimDatabase(days);
+                break;
+
             case ACTION_SET_BOAST_OPTIONS:
                 ext.getBoastService().getParam().setBoastUri(params.getString(PARAM_SERVER));
                 ext.getBoastService()
@@ -140,6 +154,9 @@ public ApiResponse handleApiView(String name, JSONObject params) throws ApiExcep
             case VIEW_GET_ACTIVE_SCAN_SERVICE:
                 OastService service = ext.getActiveScanOastService();
                 return new ApiResponseElement(name, service != null ? service.getName() : "");
+            case VIEW_GET_DAYS_TO_KEEP_RECORDS:
+                return new ApiResponseElement(
+                        name, Integer.toString(ext.getParams().getDaysToKeepRecords()));
             case VIEW_GET_SERVICES:
                 ApiResponseList servList = new ApiResponseList(name);
 

diff --git a/addOns/oast/src/main/java/org/zaproxy/addon/oast/OastParam.java b/addOns/oast/src/main/java/org/zaproxy/addon/oast/OastParam.java
@@ -30,6 +30,7 @@ public class OastParam extends VersionedAbstractParam {
             PARAM_BASE_KEY + ".activeScanService";
     private static final String PARAM_USE_PERMANENT_DATABASE =
             PARAM_BASE_KEY + ".usePermanentDatabase";
+    private static final String PARAM_DAYS_TO_KEEP_RECORDS = PARAM_BASE_KEY + ".daysToKeepRecords";
 
     public static final String NO_ACTIVE_SCAN_SERVICE_SELECTED_OPTION = "None";
 
@@ -41,6 +42,7 @@ public class OastParam extends VersionedAbstractParam {
 
     private String activeScanServiceName;
     private boolean usePermanentDatabase;
+    private int daysToKeepRecords;
 
     public OastParam() {}
 
@@ -62,11 +64,21 @@ public void setUsePermanentDatabase(boolean usePermanentDatabase) {
         getConfig().setProperty(PARAM_USE_PERMANENT_DATABASE, usePermanentDatabase);
     }
 
+    public int getDaysToKeepRecords() {
+        return daysToKeepRecords;
+    }
+
+    public void setDaysToKeepRecords(int daysToKeepRecords) {
+        this.daysToKeepRecords = daysToKeepRecords;
+        getConfig().setProperty(PARAM_DAYS_TO_KEEP_RECORDS, daysToKeepRecords);
+    }
+
     @Override
     protected void parseImpl() {
         activeScanServiceName =
                 getString(PARAM_ACTIVE_SCAN_SERVICE_NAME, NO_ACTIVE_SCAN_SERVICE_SELECTED_OPTION);
         usePermanentDatabase = getBoolean(PARAM_USE_PERMANENT_DATABASE, true);
+        daysToKeepRecords = getInteger(PARAM_DAYS_TO_KEEP_RECORDS, 45);
     }
 
     @Override

diff --git a/addOns/oast/src/main/java/org/zaproxy/addon/oast/internal/OastPermanentDatabase.java b/addOns/oast/src/main/java/org/zaproxy/addon/oast/internal/OastPermanentDatabase.java
@@ -68,4 +68,64 @@ public AlertEntity getAlertForPayload(String payload) {
         }
         return null;
     }
+
+    public void trim(int days) {
+        if (days <= 0) {
+            return;
+        }
+        LOGGER.debug("Trimming records older than {} days", days);
+        try {
+            // Yes, passing parameters in this way is horrible!
+            // But I couldn't get the parameters working the "correct" way and we know its
+            // definitely a positive int..
+            String dateClause = "DATE_SUB(CURRENT_TIMESTAMP, INTERVAL " + days + " DAY)";
+
+            Object res =
+                    runQuery(
+                            "DELETE FROM BOAST WHERE REGISTERED_TIMESTAMP < " + dateClause,
+                            null,
+                            false);
+            if (Integer.parseInt(res.toString()) > 0) {
+                LOGGER.info("Number of old BOAST records trimmed: {}", res);
+            }
+
+            res = runQuery("DELETE FROM ALERT WHERE CREATETIMESTAMP < " + dateClause, null, false);
+            if (Integer.parseInt(res.toString()) > 0) {
+                LOGGER.info("Number of old ALERT records trimmed: {}", res);
+            }
+
+            res =
+                    runQuery(
+                            "DELETE FROM MESSAGE WHERE CREATETIMESTAMP < " + dateClause,
+                            null,
+                            false);
+            if (Integer.parseInt(res.toString()) > 0) {
+                LOGGER.info("Number of old MESSAGE records trimmed: {}", res);
+            }
+
+        } catch (Exception e) {
+            LOGGER.error("Failed to trim db", e);
+        }
+    }
+
+    public void clearAllRecords() {
+        try {
+            Object res = runQuery("DELETE FROM BOAST", null, false);
+            if (Integer.parseInt(res.toString()) > 0) {
+                LOGGER.info("Number of old BOAST records trimmed: {}", res);
+            }
+
+            res = runQuery("DELETE FROM ALERT", null, false);
+            if (Integer.parseInt(res.toString()) > 0) {
+                LOGGER.info("Number of old ALERT records trimmed: {}", res);
+            }
+
+            res = runQuery("DELETE FROM MESSAGE", null, false);
+            if (Integer.parseInt(res.toString()) > 0) {
+                LOGGER.info("Number of old MESSAGE records trimmed: {}", res);
+            }
+        } catch (Exception e) {
+            LOGGER.error("Failed to trim db", e);
+        }
+    }
 }
diff --git a/addOns/oast/src/main/java/org/zaproxy/addon/oast/services/boast/BoastOptionsPanelTab.java b/addOns/oast/src/main/java/org/zaproxy/addon/oast/services/boast/BoastOptionsPanelTab.java
@@ -140,13 +140,17 @@ private void registerButtonAction() {
         }
     }
 
+    public void resetBoastServers() {
+        boastServersTableModel = null;
+        getBoastServersTable().setModel(getBoastServersTableModel());
+    }
+
     @Override
     public void initParam(OptionsParam options) {
         final BoastParam param = options.getParamSet(BoastParam.class);
         getBoastUri().setText(param.getBoastUri());
         getPollingFrequencySpinner().setValue(param.getPollingFrequency());
-        boastServersTableModel = null;
-        getBoastServersTable().setModel(getBoastServersTableModel());
+        resetBoastServers();
     }
 
     @Override

diff --git a/addOns/oast/src/main/java/org/zaproxy/addon/oast/services/boast/BoastService.java b/addOns/oast/src/main/java/org/zaproxy/addon/oast/services/boast/BoastService.java
@@ -111,6 +111,10 @@ public List<BoastServer> getRegisteredServers() {
         return registeredServers;
     }
 
+    public void clearRegisteredServers() {
+        registeredServers.clear();
+    }
+
     public BoastServer register(String uri) throws IOException {
         getParam().setBoastUri(uri);
         return register();