diff --git a/.gitignore b/.gitignore index 9f97022..ff3432e 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ -target/ \ No newline at end of file +target/ +mise.toml \ No newline at end of file diff --git a/README.md b/README.md index 50c4106..dd041c9 100644 --- a/README.md +++ b/README.md @@ -183,3 +183,99 @@ Procedure kindly provided by the PortSwigger support: # SQLite client Cross-platform: https://github.com/sqlitebrowser/sqlitebrowser + +# Burp Suite Activity Logger - PostgreSQL Database Setup + +This Docker Compose configuration sets up a PostgreSQL database for the Burp Suite Activity Logger extension. + +## Quick Start + +1. **Start the database:** + ```bash + docker-compose up -d + ``` + +2. **Configure your Burp Suite extension with these connection parameters:** + - **Host:** `localhost` + - **Port:** `5432` + - **Database:** `burp_activity` + - **Username:** `burp_user` + - **Password:** `burp_password` + +3. **Stop the database:** + ```bash + docker-compose down + ``` + +## Services Included + +### PostgreSQL Database +- **Container:** `burp-activity-db` +- **Port:** 5432 (exposed to host) +- **Database:** `burp_activity` +- **User:** `burp_user` +- **Password:** `burp_password` + +### pgAdmin (Optional) +- **Container:** `burp-pgadmin` +- **Port:** 8080 (web interface) +- **Email:** `admin@example.com` +- **Password:** `admin` + +To start with pgAdmin included: +```bash +docker-compose --profile admin up -d +``` + +## Data Persistence + +Database data is stored in a Docker volume named `postgres_data`, so your data will persist between container restarts. + +## Database Schema + +The database automatically creates an `ACTIVITY` table with the following structure: +- `id` - Primary key (auto-increment) +- `local_source_ip` - Source IP address +- `target_url` - Target URL of the request +- `http_method` - HTTP method (GET, POST, etc.) +- `burp_tool` - Burp Suite tool that generated the request +- `request_raw` - Raw HTTP request +- `send_datetime` - When the request was sent +- `http_status_code` - HTTP response status code +- `response_raw` - Raw HTTP response +- `created_at` - When the record was created + +## Performance Optimizations + +The setup includes several performance optimizations: +- Indexes on commonly queried columns +- Proper user permissions +- Health checks for container monitoring + +## Customization + +You can modify the following in `docker-compose.yml`: +- Database name, username, and password in the `environment` section +- Port mappings if you need different ports +- Volume configurations for data storage + +## Troubleshooting + +1. **Connection refused errors:** + - Ensure the container is running: `docker-compose ps` + - Check container logs: `docker-compose logs postgres` + +2. **Permission errors:** + - The init script sets up proper permissions automatically + - If issues persist, check the logs: `docker-compose logs postgres` + +3. **Data not persisting:** + - Ensure the volume is properly created: `docker volume ls` + - Check that the container has write permissions + +## Security Notes + +- The default credentials are for development only +- For production use, change the default passwords +- Consider using Docker secrets for sensitive information +- Restrict network access as needed diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..589181e --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,42 @@ +version: '3.8' + +services: + postgres: + image: postgres:15-alpine + container_name: burp-activity-db + environment: + POSTGRES_DB: burp_activity + POSTGRES_USER: burp_user + POSTGRES_PASSWORD: burp_password + POSTGRES_INITDB_ARGS: "--encoding=UTF8" + ports: + - "5432:5432" + volumes: + - postgres_data:/var/lib/postgresql/data + - ./init.sql:/docker-entrypoint-initdb.d/init.sql + restart: unless-stopped + healthcheck: + test: ["CMD-SHELL", "pg_isready -U burp_user -d burp_activity"] + interval: 10s + timeout: 5s + retries: 5 + + # Optional: pgAdmin for database management + pgadmin: + image: dpage/pgadmin4:latest + container_name: burp-pgadmin + environment: + PGADMIN_DEFAULT_EMAIL: admin@example.com + PGADMIN_DEFAULT_PASSWORD: admin + PGADMIN_CONFIG_SERVER_MODE: 'False' + ports: + - "8080:80" + depends_on: + - postgres + restart: unless-stopped + profiles: + - admin + +volumes: + postgres_data: + driver: local \ No newline at end of file diff --git a/init.sql b/init.sql new file mode 100644 index 0000000..dec3d11 --- /dev/null +++ b/init.sql @@ -0,0 +1,35 @@ +-- Initialize the database for Burp Suite Activity Logger +-- This script runs automatically when the container starts for the first time + +-- Grant necessary permissions to the burp_user +GRANT ALL PRIVILEGES ON DATABASE burp_activity TO burp_user; + +-- Create the activity table (this will also be created by the Java extension, but having it here ensures it exists) +CREATE TABLE IF NOT EXISTS ACTIVITY ( + id SERIAL PRIMARY KEY, + local_source_ip TEXT, + target_url TEXT, + http_method TEXT, + burp_tool TEXT, + request_raw TEXT, + send_datetime TIMESTAMP, + http_status_code TEXT, + response_raw TEXT, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP +); + +-- Grant permissions on the table +GRANT ALL PRIVILEGES ON TABLE ACTIVITY TO burp_user; +GRANT USAGE, SELECT ON SEQUENCE activity_id_seq TO burp_user; + +-- Create indexes for better performance on common queries +CREATE INDEX IF NOT EXISTS idx_activity_send_datetime ON ACTIVITY(send_datetime); +CREATE INDEX IF NOT EXISTS idx_activity_http_method ON ACTIVITY(http_method); +CREATE INDEX IF NOT EXISTS idx_activity_burp_tool ON ACTIVITY(burp_tool); +CREATE INDEX IF NOT EXISTS idx_activity_target_url ON ACTIVITY(target_url); + +-- Log initialization completion +DO $$ +BEGIN + RAISE NOTICE 'Burp Activity Logger database initialized successfully'; +END $$; \ No newline at end of file diff --git a/pom.xml b/pom.xml index 13f2fcf..8e2c620 100644 --- a/pom.xml +++ b/pom.xml @@ -20,6 +20,11 @@ sqlite-jdbc 3.49.1.0 + + org.postgresql + postgresql + 42.7.4 + diff --git a/src/main/java/burp/ActivityHttpListener.java b/src/main/java/burp/ActivityHttpListener.java index 63f7c8b..c0e8e52 100644 --- a/src/main/java/burp/ActivityHttpListener.java +++ b/src/main/java/burp/ActivityHttpListener.java @@ -13,7 +13,7 @@ class ActivityHttpListener implements HttpHandler { /** * Ref on handler that will store the activity information into the activity log storage. */ - private ActivityLogger activityLogger; + private ActivityStorage activityStorage; /** * Ref on project logger. @@ -23,22 +23,34 @@ class ActivityHttpListener implements HttpHandler { /** * Constructor. * - * @param activityLogger Ref on handler that will store the activity information into the activity log storage. + * @param activityStorage Ref on handler that will store the activity information into the activity log storage. * @param trace Ref on project logger. */ - ActivityHttpListener(ActivityLogger activityLogger, Trace trace) { - this.activityLogger = activityLogger; + ActivityHttpListener(ActivityStorage activityStorage, Trace trace) { + this.activityStorage = activityStorage; this.trace = trace; } + /** + * Replace the current activity storage with a new one. + * This allows switching between storage backends without restarting Burp Suite. + * + * @param newStorage The new storage instance to use + */ + void replaceStorage(ActivityStorage newStorage) { + this.activityStorage = newStorage; + this.trace.writeLog("HTTP listener activity storage replaced."); + } + @Override public RequestToBeSentAction handleHttpRequestToBeSent(HttpRequestToBeSent requestToBeSent) { //Check if the response will be logged as well. If yes, wait until response is received. if (!ConfigMenu.INCLUDE_HTTP_RESPONSE_CONTENT) { try { - if (this.mustLogRequest(requestToBeSent)) { - this.activityLogger.logEvent(requestToBeSent, null, requestToBeSent.toolSource().toolType().toolName()); + String toolName = requestToBeSent.toolSource().toolType().toolName(); + if (this.mustLogRequest(requestToBeSent, toolName)) { + this.activityStorage.logEvent(requestToBeSent, null, toolName); } } catch (Exception e) { this.trace.writeLog("Cannot save request: " + e.getMessage()); @@ -56,8 +68,9 @@ public ResponseReceivedAction handleHttpResponseReceived(HttpResponseReceived re if (ConfigMenu.INCLUDE_HTTP_RESPONSE_CONTENT) { try { //Save the information of the current request if the message is an HTTP response and according to the restriction options - if (this.mustLogRequest(responseReceived.initiatingRequest())) { - this.activityLogger.logEvent(responseReceived.initiatingRequest(), responseReceived, responseReceived.toolSource().toolType().toolName()); + String toolName = responseReceived.toolSource().toolType().toolName(); + if (this.mustLogRequest(responseReceived.initiatingRequest(), toolName)) { + this.activityStorage.logEvent(responseReceived.initiatingRequest(), responseReceived, toolName); } } catch (Exception e) { this.trace.writeLog("Cannot save response: " + e.getMessage()); @@ -70,18 +83,40 @@ public ResponseReceivedAction handleHttpResponseReceived(HttpResponseReceived re * Determine if the current request must be logged according to the configuration options selected by the users. * * @param request HttpRequest object containing all the information about the request + * @param toolName Name of the tool that generated this request (e.g., "Repeater", "Intruder", "Proxy") * @return TRUE if the request must be logged, FALSE otherwise */ - private boolean mustLogRequest(HttpRequest request) { + private boolean mustLogRequest(HttpRequest request, String toolName) { //By default: Request is logged boolean mustLogRequest = true; + String url = request.url(); + + //this.trace.writeLog("DEBUG: Checking request from " + toolName + " to " + url); //Initially we check the pause state if (ConfigMenu.IS_LOGGING_PAUSED) { mustLogRequest = false; + //this.trace.writeLog("DEBUG: Request filtered out - logging is paused"); } else { - //First: We check if we must apply restriction about image resource - if (ConfigMenu.EXCLUDE_IMAGE_RESOURCE_REQUESTS) { + //First: We check if we must apply restriction about tool source + if (ConfigMenu.FILTER_BY_TOOL_SOURCE) { + // Debug logging to see actual tool names + //this.trace.writeLog("Received tool name: '" + toolName + "', Included tools: " + ConfigMenu.INCLUDED_TOOL_SOURCES.toString()); + + // Check if any of the included tool sources match (case-insensitive) + boolean toolMatches = ConfigMenu.INCLUDED_TOOL_SOURCES.stream() + .anyMatch(includedTool -> includedTool.equalsIgnoreCase(toolName)); + + if (!toolMatches) { + mustLogRequest = false; + //this.trace.writeLog("DEBUG: Request from tool '" + toolName + "' filtered out by tool source filter."); + } else { + //this.trace.writeLog("DEBUG: Tool '" + toolName + "' passed tool source filter."); + } + } + //Second: We check if we must apply restriction about image resource + //Configuration restrictions options are applied in sequence so we only work here if the request is marked to be logged + if (mustLogRequest && ConfigMenu.EXCLUDE_IMAGE_RESOURCE_REQUESTS) { //Get the file extension of the current URL and remove the parameters from the URL String filename = request.url(); if (filename != null && filename.indexOf('?') != -1) { @@ -94,16 +129,19 @@ private boolean mustLogRequest(HttpRequest request) { String extension = filename.substring(filename.lastIndexOf('.') + 1).trim().toLowerCase(Locale.US); if (ConfigMenu.IMAGE_RESOURCE_EXTENSIONS.contains(extension)) { mustLogRequest = false; + //this.trace.writeLog("DEBUG: Request filtered out - image resource with extension: " + extension); } } } - //Secondly: We check if we must apply restriction about the URL scope + //Finally: We check if we must apply restriction about the URL scope //Configuration restrictions options are applied in sequence so we only work here if the request is marked to be logged if (mustLogRequest && ConfigMenu.ONLY_INCLUDE_REQUESTS_FROM_SCOPE && ! request.isInScope()) { mustLogRequest = false; + //this.trace.writeLog("DEBUG: Request filtered out - not in scope: " + url); } } + //this.trace.writeLog("DEBUG: Final decision for " + toolName + " request to " + url + ": " + (mustLogRequest ? "LOGGED" : "FILTERED")); return mustLogRequest; } diff --git a/src/main/java/burp/ActivityLogger.java b/src/main/java/burp/ActivityLogger.java index 8871019..7e5c4f9 100644 --- a/src/main/java/burp/ActivityLogger.java +++ b/src/main/java/burp/ActivityLogger.java @@ -9,6 +9,10 @@ import java.sql.Statement; import java.time.LocalDateTime; import java.time.format.DateTimeFormatter; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import burp.api.montoya.MontoyaApi; import burp.api.montoya.http.message.requests.HttpRequest; @@ -17,8 +21,9 @@ /** * Handle the recording of the activities into the real storage, SQLite local DB here. + * Now uses async writes with a background thread for improved performance. */ -class ActivityLogger implements ExtensionUnloadingHandler { +class ActivityLogger implements ActivityStorage { /** * SQL instructions. @@ -30,6 +35,21 @@ class ActivityLogger implements ExtensionUnloadingHandler { private static final String SQL_BIGGEST_REQUEST_AMOUNT_DATA_SENT = "SELECT MAX(LENGTH(REQUEST_RAW)) FROM ACTIVITY"; private static final String SQL_MAX_HITS_BY_SECOND = "SELECT COUNT(REQUEST_RAW) AS HITS, SEND_DATETIME FROM ACTIVITY GROUP BY SEND_DATETIME ORDER BY HITS DESC"; + /** + * Maximum queue size to prevent memory issues + */ + private static final int MAX_QUEUE_SIZE = 10000; + + /** + * Batch size for database writes + */ + private static final int BATCH_SIZE = 100; + + /** + * Maximum wait time for batch processing (milliseconds) + */ + private static final long BATCH_TIMEOUT_MS = 1000; + /** * Use a single DB connection for performance and to prevent DB file locking issue at filesystem level. */ @@ -50,6 +70,20 @@ class ActivityLogger implements ExtensionUnloadingHandler { */ private DateTimeFormatter datetimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"); + /** + * Queue for async event processing + */ + private final BlockingQueue eventQueue = new LinkedBlockingQueue<>(MAX_QUEUE_SIZE); + + /** + * Background thread for database writes + */ + private Thread writerThread; + + /** + * Flag to control the writer thread lifecycle + */ + private final AtomicBoolean running = new AtomicBoolean(true); /** * Constructor. @@ -63,6 +97,163 @@ class ActivityLogger implements ExtensionUnloadingHandler { Class.forName("org.sqlite.JDBC"); this.trace = trace; updateStoreLocation(storeName); + startWriterThread(); + } + + /** + * Start the background writer thread + */ + private void startWriterThread() { + writerThread = new Thread(this::processEventQueue, "ActivityLogger-Writer"); + writerThread.setDaemon(true); + writerThread.start(); + this.trace.writeLog("Async writer thread started."); + } + + /** + * Background thread that processes the event queue + */ + private void processEventQueue() { + LogEvent[] batch = new LogEvent[BATCH_SIZE]; + + while (running.get() || !eventQueue.isEmpty()) { + try { + int batchCount = 0; + long batchStartTime = System.currentTimeMillis(); + + // Collect events for batching + while (batchCount < BATCH_SIZE && running.get()) { + LogEvent event = eventQueue.poll(BATCH_TIMEOUT_MS, TimeUnit.MILLISECONDS); + if (event == null) { + break; // Timeout reached + } + batch[batchCount++] = event; + + // Check if we should flush early due to timeout + if (System.currentTimeMillis() - batchStartTime >= BATCH_TIMEOUT_MS) { + break; + } + } + + // Process the batch if we have events + if (batchCount > 0) { + writeBatch(batch, batchCount); + } + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } catch (Exception e) { + this.trace.writeLog("Error in writer thread: " + e.getMessage()); + } + } + + // Flush remaining events when shutting down + flushRemainingEvents(); + } + + /** + * Write a batch of events to the database + */ + private void writeBatch(LogEvent[] batch, int count) throws Exception { + ensureDBState(); + + // Use batch inserts for better performance + this.storageConnection.setAutoCommit(false); + + try (PreparedStatement stmt = this.storageConnection.prepareStatement(SQL_TABLE_INSERT)) { + for (int i = 0; i < count; i++) { + LogEvent event = batch[i]; + stmt.setString(1, event.localSourceIp); + stmt.setString(2, event.targetUrl); + stmt.setString(3, event.httpMethod); + stmt.setString(4, event.tool); + stmt.setString(5, event.requestRaw); + stmt.setString(6, event.sendDateTime); + stmt.setString(7, event.httpStatusCode); + stmt.setString(8, event.responseRaw); + stmt.addBatch(); + } + + int[] results = stmt.executeBatch(); + this.storageConnection.commit(); + + // Log any failed inserts + int successCount = 0; + for (int result : results) { + if (result > 0) successCount++; + } + + if (successCount != count) { + this.trace.writeLog("Batch insert: " + successCount + "/" + count + " events inserted successfully"); + } + + } catch (Exception e) { + this.storageConnection.rollback(); + throw e; + } finally { + this.storageConnection.setAutoCommit(true); + } + } + + /** + * Flush any remaining events in the queue (used during shutdown) + */ + private void flushRemainingEvents() { + LogEvent[] batch = new LogEvent[BATCH_SIZE]; + int count = 0; + + while (!eventQueue.isEmpty() && count < BATCH_SIZE) { + LogEvent event = eventQueue.poll(); + if (event != null) { + batch[count++] = event; + } + } + + if (count > 0) { + try { + writeBatch(batch, count); + this.trace.writeLog("Flushed " + count + " remaining events during shutdown."); + } catch (Exception e) { + this.trace.writeLog("Error flushing remaining events: " + e.getMessage()); + } + } + } + + /** + * Save an activity event into the storage (now async). + * + * @param request HttpRequest object containing all information about the request + * which was either sent or will be sent out soon. + * @param response HttpResponse object containing all information about the response. + * Is null when only the request ist stored. + * @param tool The name of the tool which was used to issue to request. + * @throws Exception If event cannot be saved. + */ + public void logEvent(HttpRequest request, HttpResponse response, String tool) throws Exception { + try { + // Create event object with pre-computed values + LogEvent event = new LogEvent( + InetAddress.getLocalHost().getHostAddress(), + request.url(), + request.method(), + tool, + request.toString(), + LocalDateTime.now().format(this.datetimeFormatter), + response != null ? String.valueOf(response.statusCode()) : null, + response != null ? response.bodyToString() : null + ); + + // Add to queue (non-blocking) + if (!eventQueue.offer(event)) { + // Queue is full - could log a warning or implement backpressure + this.trace.writeLog("Event queue full, dropping event. Consider adjusting MAX_QUEUE_SIZE."); + } + + } catch (Exception e) { + this.trace.writeLog("Error queueing event: " + e.getMessage()); + // Could fallback to synchronous write in critical cases + } } /** @@ -74,61 +265,23 @@ class ActivityLogger implements ExtensionUnloadingHandler { void updateStoreLocation(String storeName) throws Exception { String newUrl = "jdbc:sqlite:" + storeName; this.url = newUrl; - //Open the connection to the DB this.trace.writeLog("Activity information will be stored in database file '" + storeName + "'."); this.storageConnection = DriverManager.getConnection(newUrl); this.storageConnection.setAutoCommit(true); this.trace.writeLog("Open new connection to the storage."); - //Create the table try (Statement stmt = this.storageConnection.createStatement()) { stmt.execute(SQL_TABLE_CREATE); this.trace.writeLog("Recording table initialized."); } } - /** - * Save an activity event into the storage. - * - * @param request HttpRequest object containing all information about the request - * which was either sent or will be sent out soon. - * @param response HttpResponse object containing all information about the response. - * Is null when only the request ist stored. - * @param tool The name of the tool which was used to issue to request. - * @throws Exception If event cannot be saved. - */ - void logEvent(HttpRequest request, HttpResponse response, String tool) throws Exception { - //Verify that the DB connection is still opened - this.ensureDBState(); - //Insert the event into the storage - try (PreparedStatement stmt = this.storageConnection.prepareStatement(SQL_TABLE_INSERT)) { - stmt.setString(1, InetAddress.getLocalHost().getHostAddress()); - stmt.setString(2, request.url()); - stmt.setString(3, request.method()); - stmt.setString(4, tool); - stmt.setString(5, request.toString()); //Apparently, bodyToString() does not work.. - stmt.setString(6, LocalDateTime.now().format(this.datetimeFormatter)); - //Make a distinction if only the request is stored or the response is added as well. - if (response != null) { - stmt.setString(7, String.valueOf(response.statusCode())); - stmt.setString(8, response.bodyToString()); - } else { - stmt.setString(7, null); - stmt.setString(8, null); - } - int count = stmt.executeUpdate(); - if (count != 1) { - this.trace.writeLog("Request was not inserted, no detail available (insertion counter = " + count + ") !"); - } - } - } - /** * Extract and compute statistics about the DB. * * @return A VO object containing the statistics. * @throws Exception If computation meet and error. */ - DBStats getEventsStats() throws Exception { + public DBStats getEventsStats() throws Exception { //Verify that the DB connection is still opened this.ensureDBState(); //Get the total of the records in the activity table @@ -189,6 +342,22 @@ private void ensureDBState() throws Exception { */ @Override public void extensionUnloaded() { + // Signal the writer thread to stop + running.set(false); + + // Wait for writer thread to finish processing + if (writerThread != null) { + try { + writerThread.interrupt(); + writerThread.join(5000); // Wait up to 5 seconds + this.trace.writeLog("Writer thread stopped."); + } catch (InterruptedException e) { + this.trace.writeLog("Interrupted while waiting for writer thread to finish."); + Thread.currentThread().interrupt(); + } + } + + // Close database connection try { if (this.storageConnection != null && !this.storageConnection.isClosed()) { this.storageConnection.close(); @@ -198,4 +367,30 @@ public void extensionUnloaded() { this.trace.writeLog("Cannot close the connection to the storage: " + e.getMessage()); } } + + /** + * Value object to hold event data for async processing + */ + private static class LogEvent { + final String localSourceIp; + final String targetUrl; + final String httpMethod; + final String tool; + final String requestRaw; + final String sendDateTime; + final String httpStatusCode; + final String responseRaw; + + LogEvent(String localSourceIp, String targetUrl, String httpMethod, String tool, + String requestRaw, String sendDateTime, String httpStatusCode, String responseRaw) { + this.localSourceIp = localSourceIp; + this.targetUrl = targetUrl; + this.httpMethod = httpMethod; + this.tool = tool; + this.requestRaw = requestRaw; + this.sendDateTime = sendDateTime; + this.httpStatusCode = httpStatusCode; + this.responseRaw = responseRaw; + } + } } diff --git a/src/main/java/burp/ActivityStorage.java b/src/main/java/burp/ActivityStorage.java new file mode 100644 index 0000000..b441682 --- /dev/null +++ b/src/main/java/burp/ActivityStorage.java @@ -0,0 +1,33 @@ +package burp; + +import burp.api.montoya.http.message.requests.HttpRequest; +import burp.api.montoya.http.message.responses.HttpResponse; +import burp.api.montoya.extension.ExtensionUnloadingHandler; + +/** + * Interface for activity storage implementations. + * This allows for different storage backends (SQLite, PostgreSQL, etc.) + * while maintaining a consistent API. + */ +interface ActivityStorage extends ExtensionUnloadingHandler { + + /** + * Save an activity event into the storage. + * + * @param request HttpRequest object containing all information about the request + * which was either sent or will be sent out soon. + * @param response HttpResponse object containing all information about the response. + * Is null when only the request is stored. + * @param tool The name of the tool which was used to issue the request. + * @throws Exception If event cannot be saved. + */ + void logEvent(HttpRequest request, HttpResponse response, String tool) throws Exception; + + /** + * Extract and compute statistics about the storage. + * + * @return A VO object containing the statistics. + * @throws Exception If computation meets an error. + */ + DBStats getEventsStats() throws Exception; +} \ No newline at end of file diff --git a/src/main/java/burp/ActivityStorageFactory.java b/src/main/java/burp/ActivityStorageFactory.java new file mode 100644 index 0000000..ed398b3 --- /dev/null +++ b/src/main/java/burp/ActivityStorageFactory.java @@ -0,0 +1,137 @@ +package burp; + +import burp.api.montoya.MontoyaApi; +import burp.api.montoya.persistence.Preferences; +import javax.swing.JFrame; +import javax.swing.JOptionPane; +import javax.swing.JTextField; +import javax.swing.JPasswordField; +import javax.swing.JPanel; +import javax.swing.JLabel; +import java.awt.GridLayout; + +/** + * Factory class to create the appropriate activity storage instance based on configuration. + */ +class ActivityStorageFactory { + + /** + * Create an activity storage instance based on configuration preferences. + * + * @param preferences Preferences for configuration + * @param storeName SQLite database file path (used if SQLite is selected) + * @param api MontoyaApi instance + * @param trace Trace logger + * @return ActivityStorage instance + * @throws Exception If storage creation fails + */ + static ActivityStorage createStorage(Preferences preferences, String storeName, MontoyaApi api, Trace trace) throws Exception { + boolean usePostgreSQL = Boolean.TRUE.equals(preferences.getBoolean(ConfigMenu.USE_POSTGRESQL_CFG_KEY)); + + if (usePostgreSQL) { + return createPostgreSQLStorage(preferences, api, trace); + } else { + return createSQLiteStorage(storeName, api, trace); + } + } + + /** + * Create SQLite storage instance. + */ + private static ActivityStorage createSQLiteStorage(String storeName, MontoyaApi api, Trace trace) throws Exception { + return new ActivityLogger(storeName, api, trace); + } + + /** + * Create PostgreSQL storage instance. + */ + private static ActivityStorage createPostgreSQLStorage(Preferences preferences, MontoyaApi api, Trace trace) throws Exception { + String host = preferences.getString(ConfigMenu.POSTGRESQL_HOST_CFG_KEY); + String portStr = preferences.getString(ConfigMenu.POSTGRESQL_PORT_CFG_KEY); + String database = preferences.getString(ConfigMenu.POSTGRESQL_DATABASE_CFG_KEY); + String username = preferences.getString(ConfigMenu.POSTGRESQL_USERNAME_CFG_KEY); + String password = preferences.getString(ConfigMenu.POSTGRESQL_PASSWORD_CFG_KEY); + + // Set defaults if not configured + if (host == null || host.trim().isEmpty()) { + host = "localhost"; + } + + int port = 5432; // Default PostgreSQL port + if (portStr != null && !portStr.trim().isEmpty()) { + try { + port = Integer.parseInt(portStr.trim()); + } catch (NumberFormatException e) { + trace.writeLog("Invalid PostgreSQL port, using default 5432"); + } + } + + if (database == null || database.trim().isEmpty()) { + database = "burp_requests"; + } + + if (username == null || username.trim().isEmpty()) { + username = "postgres"; + } + + if (password == null) { + password = ""; + } + + return new PostgreSQLActivityLogger(host, port, database, username, password, api, trace); + } + + /** + * Show PostgreSQL connection configuration dialog. + * + * @param preferences Preferences to save configuration + * @param parentFrame Parent frame for dialog + * @return true if user confirmed, false if cancelled + */ + static boolean showPostgreSQLConfigDialog(Preferences preferences, JFrame parentFrame) { + JPanel panel = new JPanel(new GridLayout(5, 2, 5, 5)); + + String currentHost = preferences.getString(ConfigMenu.POSTGRESQL_HOST_CFG_KEY); + String currentPort = preferences.getString(ConfigMenu.POSTGRESQL_PORT_CFG_KEY); + String currentDatabase = preferences.getString(ConfigMenu.POSTGRESQL_DATABASE_CFG_KEY); + String currentUsername = preferences.getString(ConfigMenu.POSTGRESQL_USERNAME_CFG_KEY); + String currentPassword = preferences.getString(ConfigMenu.POSTGRESQL_PASSWORD_CFG_KEY); + + JTextField hostField = new JTextField(currentHost != null ? currentHost : "localhost"); + JTextField portField = new JTextField(currentPort != null ? currentPort : "5432"); + JTextField databaseField = new JTextField(currentDatabase != null ? currentDatabase : "burp_requests"); + JTextField usernameField = new JTextField(currentUsername != null ? currentUsername : "postgres"); + JPasswordField passwordField = new JPasswordField(currentPassword != null ? currentPassword : ""); + + panel.add(new JLabel("Host:")); + panel.add(hostField); + panel.add(new JLabel("Port:")); + panel.add(portField); + panel.add(new JLabel("Database:")); + panel.add(databaseField); + panel.add(new JLabel("Username:")); + panel.add(usernameField); + panel.add(new JLabel("Password:")); + panel.add(passwordField); + + int result = JOptionPane.showConfirmDialog( + parentFrame, + panel, + "PostgreSQL Connection Configuration", + JOptionPane.OK_CANCEL_OPTION, + JOptionPane.PLAIN_MESSAGE + ); + + if (result == JOptionPane.OK_OPTION) { + // Save configuration + preferences.setString(ConfigMenu.POSTGRESQL_HOST_CFG_KEY, hostField.getText().trim()); + preferences.setString(ConfigMenu.POSTGRESQL_PORT_CFG_KEY, portField.getText().trim()); + preferences.setString(ConfigMenu.POSTGRESQL_DATABASE_CFG_KEY, databaseField.getText().trim()); + preferences.setString(ConfigMenu.POSTGRESQL_USERNAME_CFG_KEY, usernameField.getText().trim()); + preferences.setString(ConfigMenu.POSTGRESQL_PASSWORD_CFG_KEY, new String(passwordField.getPassword())); + return true; + } + + return false; + } +} \ No newline at end of file diff --git a/src/main/java/burp/BurpExtender.java b/src/main/java/burp/BurpExtender.java index bed63bb..6b4d01e 100644 --- a/src/main/java/burp/BurpExtender.java +++ b/src/main/java/burp/BurpExtender.java @@ -45,47 +45,114 @@ public void initialize(MontoyaApi api) { customStoreFileName = defaultStoreFileName; } Boolean isLoggingPaused = Boolean.TRUE.equals(preferences.getBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY)); + Boolean usePostgreSQL = Boolean.TRUE.equals(preferences.getBoolean(ConfigMenu.USE_POSTGRESQL_CFG_KEY)); + if (!isLoggingPaused) { - Object[] options = {"Keep the DB file", "Change the DB file", "Pause the logging"}; - String msg = "Continue to log events into the following database file?\n\r" + customStoreFileName; - //Mapping of the buttons with the dialog: options[0] => YES / options[1] => NO / options[2] => CANCEL - int loggingQuestionReply = JOptionPane.showOptionDialog(burpFrame, msg, extensionName, JOptionPane.YES_NO_CANCEL_OPTION, JOptionPane.QUESTION_MESSAGE, null, options, null); - //Case for YES is already handled, use the stored file - if (loggingQuestionReply == JOptionPane.YES_OPTION) { - preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); - this.api.logging().logToOutput("Logging is enabled."); - } - //Case for the NO => Change DB file - if (loggingQuestionReply == JOptionPane.NO_OPTION) { - JFileChooser customStoreFileNameFileChooser = Utilities.createDBFileChooser(); - int dbFileSelectionReply = customStoreFileNameFileChooser.showDialog(burpFrame, "Use"); - if (dbFileSelectionReply == JFileChooser.APPROVE_OPTION) { - customStoreFileName = customStoreFileNameFileChooser.getSelectedFile().getAbsolutePath().replaceAll("\\\\", "/"); + String msg; + Object[] options; + + if (usePostgreSQL) { + // Check if PostgreSQL is configured + String pgHost = preferences.getString(ConfigMenu.POSTGRESQL_HOST_CFG_KEY); + String pgDb = preferences.getString(ConfigMenu.POSTGRESQL_DATABASE_CFG_KEY); + if (pgHost == null || pgHost.trim().isEmpty() || pgDb == null || pgDb.trim().isEmpty()) { + // PostgreSQL not configured, ask user to configure + msg = "PostgreSQL storage is enabled but not configured. Would you like to configure it now?"; + options = new Object[]{"Configure PostgreSQL", "Use SQLite instead", "Pause the logging"}; } else { - JOptionPane.showMessageDialog(burpFrame, "The following database file will continue to be used:\n\r" + customStoreFileName, extensionName, JOptionPane.INFORMATION_MESSAGE); + msg = "Continue to log events into PostgreSQL database?\n\rHost: " + pgHost + "\n\rDatabase: " + pgDb; + options = new Object[]{"Continue", "Reconfigure PostgreSQL", "Pause the logging"}; } - preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); - this.api.logging().logToOutput("Logging is enabled."); + } else { + msg = "Continue to log events into the following SQLite database file?\n\r" + customStoreFileName; + options = new Object[]{"Keep the DB file", "Change the DB file", "Pause the logging"}; } - //Case for the CANCEL => Pause the logging - if (loggingQuestionReply == JOptionPane.CANCEL_OPTION) { - preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.TRUE); - this.api.logging().logToOutput("Logging is paused."); + //Mapping of the buttons with the dialog: options[0] => YES / options[1] => NO / options[2] => CANCEL + int loggingQuestionReply = JOptionPane.showOptionDialog(burpFrame, msg, extensionName, JOptionPane.YES_NO_CANCEL_OPTION, JOptionPane.QUESTION_MESSAGE, null, options, null); + + if (usePostgreSQL) { + //Case for YES (Continue/Configure PostgreSQL) + if (loggingQuestionReply == JOptionPane.YES_OPTION) { + String pgHost = preferences.getString(ConfigMenu.POSTGRESQL_HOST_CFG_KEY); + String pgDb = preferences.getString(ConfigMenu.POSTGRESQL_DATABASE_CFG_KEY); + if (pgHost == null || pgHost.trim().isEmpty() || pgDb == null || pgDb.trim().isEmpty()) { + // Configure PostgreSQL + if (ActivityStorageFactory.showPostgreSQLConfigDialog(preferences, burpFrame)) { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); + this.api.logging().logToOutput("PostgreSQL storage configured and logging is enabled."); + } else { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.TRUE); + this.api.logging().logToOutput("PostgreSQL configuration cancelled. Logging is paused."); + } + } else { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); + this.api.logging().logToOutput("PostgreSQL logging is enabled."); + } + } + //Case for NO (Use SQLite instead/Reconfigure PostgreSQL) + if (loggingQuestionReply == JOptionPane.NO_OPTION) { + String pgHost = preferences.getString(ConfigMenu.POSTGRESQL_HOST_CFG_KEY); + String pgDb = preferences.getString(ConfigMenu.POSTGRESQL_DATABASE_CFG_KEY); + if (pgHost == null || pgHost.trim().isEmpty() || pgDb == null || pgDb.trim().isEmpty()) { + // Use SQLite instead + preferences.setBoolean(ConfigMenu.USE_POSTGRESQL_CFG_KEY, Boolean.FALSE); + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); + this.api.logging().logToOutput("Switched to SQLite storage and logging is enabled."); + } else { + // Reconfigure PostgreSQL + if (ActivityStorageFactory.showPostgreSQLConfigDialog(preferences, burpFrame)) { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); + this.api.logging().logToOutput("PostgreSQL storage reconfigured and logging is enabled."); + } else { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); + this.api.logging().logToOutput("PostgreSQL configuration unchanged. Logging is enabled."); + } + } + } + //Case for CANCEL => Pause the logging + if (loggingQuestionReply == JOptionPane.CANCEL_OPTION) { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.TRUE); + this.api.logging().logToOutput("Logging is paused."); + } + } else { + //SQLite mode + //Case for YES is already handled, use the stored file + if (loggingQuestionReply == JOptionPane.YES_OPTION) { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); + this.api.logging().logToOutput("SQLite logging is enabled."); + } + //Case for the NO => Change DB file + if (loggingQuestionReply == JOptionPane.NO_OPTION) { + JFileChooser customStoreFileNameFileChooser = Utilities.createDBFileChooser(); + int dbFileSelectionReply = customStoreFileNameFileChooser.showDialog(burpFrame, "Use"); + if (dbFileSelectionReply == JFileChooser.APPROVE_OPTION) { + customStoreFileName = customStoreFileNameFileChooser.getSelectedFile().getAbsolutePath().replaceAll("\\\\", "/"); + } else { + JOptionPane.showMessageDialog(burpFrame, "The following database file will continue to be used:\n\r" + customStoreFileName, extensionName, JOptionPane.INFORMATION_MESSAGE); + } + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.FALSE); + this.api.logging().logToOutput("SQLite logging is enabled."); + } + //Case for the CANCEL => Pause the logging + if (loggingQuestionReply == JOptionPane.CANCEL_OPTION) { + preferences.setBoolean(ConfigMenu.PAUSE_LOGGING_CFG_KEY, Boolean.TRUE); + this.api.logging().logToOutput("Logging is paused."); + } } //Save the location of the database file chosen by the user preferences.setString(ConfigMenu.DB_FILE_CUSTOM_LOCATION_CFG_KEY, customStoreFileName); } else { this.api.logging().logToOutput("Logging is paused."); } - //Init logger and HTTP listener - ActivityLogger activityLogger = new ActivityLogger(customStoreFileName, this.api, trace); - ActivityHttpListener activityHttpListener = new ActivityHttpListener(activityLogger, trace); + //Init storage and HTTP listener + ActivityStorage activityStorage = ActivityStorageFactory.createStorage(preferences, customStoreFileName, this.api, trace); + ActivityHttpListener activityHttpListener = new ActivityHttpListener(activityStorage, trace); //Setup the configuration menu - configMenu = new ConfigMenu(this.api, trace, activityLogger); + configMenu = new ConfigMenu(this.api, trace, activityStorage, activityHttpListener); SwingUtilities.invokeLater(configMenu); //Register all listeners this.api.http().registerHttpHandler(activityHttpListener); - this.api.extension().registerUnloadingHandler(activityLogger); + this.api.extension().registerUnloadingHandler(activityStorage); } catch (Exception e) { String errMsg = "Cannot start the extension due to the following reason:\n\r" + e.getMessage(); //Notification of the error in the dashboard tab diff --git a/src/main/java/burp/ConfigMenu.java b/src/main/java/burp/ConfigMenu.java index d5749e6..1dd6b7a 100644 --- a/src/main/java/burp/ConfigMenu.java +++ b/src/main/java/burp/ConfigMenu.java @@ -7,6 +7,9 @@ import javax.swing.JMenu; import javax.swing.JMenuItem; import javax.swing.JOptionPane; +import javax.swing.JPanel; +import javax.swing.JLabel; +import javax.swing.JCheckBox; import java.awt.Frame; import java.awt.event.ActionEvent; import java.util.ArrayList; @@ -14,6 +17,8 @@ import java.util.List; import java.util.Locale; import java.util.ResourceBundle; +import java.awt.GridLayout; +import java.io.File; import burp.api.montoya.MontoyaApi; import burp.api.montoya.persistence.Preferences; @@ -49,6 +54,21 @@ public class ConfigMenu implements Runnable { */ static volatile boolean IS_LOGGING_PAUSED = Boolean.FALSE; + /** + * Expose the configuration option to choose storage type (SQLite or PostgreSQL). + */ + static volatile boolean USE_POSTGRESQL = Boolean.FALSE; + + /** + * Expose the configuration option for filtering requests by tool source. + */ + static volatile boolean FILTER_BY_TOOL_SOURCE = Boolean.FALSE; + + /** + * Expose the list of included tools when tool source filtering is enabled. + */ + static final List INCLUDED_TOOL_SOURCES = new ArrayList<>(); + /** * Option configuration key for the restriction of the logging of requests in defined target scope. */ @@ -74,6 +94,30 @@ public class ConfigMenu implements Runnable { */ public static final String INCLUDE_HTTP_RESPONSE_CONTENT_CFG_KEY = "INCLUDE_HTTP_RESPONSE_CONTENT"; + /** + * Option configuration key for using PostgreSQL instead of SQLite. + */ + public static final String USE_POSTGRESQL_CFG_KEY = "USE_POSTGRESQL"; + + /** + * PostgreSQL configuration keys. + */ + public static final String POSTGRESQL_HOST_CFG_KEY = "POSTGRESQL_HOST"; + public static final String POSTGRESQL_PORT_CFG_KEY = "POSTGRESQL_PORT"; + public static final String POSTGRESQL_DATABASE_CFG_KEY = "POSTGRESQL_DATABASE"; + public static final String POSTGRESQL_USERNAME_CFG_KEY = "POSTGRESQL_USERNAME"; + public static final String POSTGRESQL_PASSWORD_CFG_KEY = "POSTGRESQL_PASSWORD"; + + /** + * Option configuration key for filtering requests by tool source. + */ + public static final String FILTER_BY_TOOL_SOURCE_CFG_KEY = "FILTER_BY_TOOL_SOURCE"; + + /** + * Option configuration key for storing included tool sources. + */ + public static final String INCLUDED_TOOL_SOURCES_CFG_KEY = "INCLUDED_TOOL_SOURCES"; + /** * Extension root configuration menu. */ @@ -95,21 +139,28 @@ public class ConfigMenu implements Runnable { private Trace trace; /** - * Ref on activity logger in order to enable the access to the DB statistics. + * Ref on activity storage in order to enable the access to the DB statistics. */ - private ActivityLogger activityLogger; + private ActivityStorage activityStorage; + + /** + * Ref on activity HTTP listener to enable storage replacement. + */ + private ActivityHttpListener activityHttpListener; /** * Constructor. * - * @param api The MontoyaAPI object used for accessing all the Burp features and ressources such as requests and responses. - * @param trace Ref on project logger. - * @param activityLogger Ref on activity logger in order to enable the access to the DB statistics. + * @param api The MontoyaAPI object used for accessing all the Burp features and ressources such as requests and responses. + * @param trace Ref on project logger. + * @param activityStorage Ref on activity storage in order to enable the access to the DB statistics. + * @param activityHttpListener Ref on activity HTTP listener to enable storage replacement. */ - ConfigMenu(MontoyaApi api, Trace trace, ActivityLogger activityLogger) { + ConfigMenu(MontoyaApi api, Trace trace, ActivityStorage activityStorage, ActivityHttpListener activityHttpListener) { this.api = api; this.trace = trace; - this.activityLogger = activityLogger; + this.activityStorage = activityStorage; + this.activityHttpListener = activityHttpListener; this.preferences = this.api.persistence().preferences(); String value; @@ -126,6 +177,23 @@ public class ConfigMenu implements Runnable { EXCLUDE_IMAGE_RESOURCE_REQUESTS = Boolean.TRUE.equals(this.preferences.getBoolean(EXCLUDE_IMAGE_RESOURCE_REQUESTS_CFG_KEY)); IS_LOGGING_PAUSED = Boolean.TRUE.equals(this.preferences.getBoolean(PAUSE_LOGGING_CFG_KEY)); INCLUDE_HTTP_RESPONSE_CONTENT = Boolean.TRUE.equals(this.preferences.getBoolean(INCLUDE_HTTP_RESPONSE_CONTENT_CFG_KEY)); + USE_POSTGRESQL = Boolean.TRUE.equals(this.preferences.getBoolean(USE_POSTGRESQL_CFG_KEY)); + FILTER_BY_TOOL_SOURCE = Boolean.TRUE.equals(this.preferences.getBoolean(FILTER_BY_TOOL_SOURCE_CFG_KEY)); + + // Load included tool sources from preferences + String includedToolsStr = this.preferences.getString(INCLUDED_TOOL_SOURCES_CFG_KEY); + if (includedToolsStr != null && !includedToolsStr.trim().isEmpty()) { + String[] includedTools = includedToolsStr.split(","); + for (String tool : includedTools) { + String trimmedTool = tool.trim(); + if (!trimmedTool.isEmpty()) { + INCLUDED_TOOL_SOURCES.add(trimmedTool); + } + } + } else { + // Default: include all tools if no preference is set + Collections.addAll(INCLUDED_TOOL_SOURCES, "Proxy", "Repeater", "Intruder", "Scanner", "Sequencer", "Spider", "Target", "Extender"); + } } /** @@ -134,7 +202,7 @@ public class ConfigMenu implements Runnable { @Override public void run() { //Build the menu - this.cfgMenu = new JMenu("Log Requests to SQLite"); + this.cfgMenu = new JMenu("Log Requests to Database"); //Add the sub menu to restrict the logging of requests in defined target scope String menuText = "Log only requests from defined target scope"; final JCheckBoxMenuItem subMenuRestrictToScope = new JCheckBoxMenuItem(menuText, ONLY_INCLUDE_REQUESTS_FROM_SCOPE); @@ -186,6 +254,175 @@ public void actionPerformed(ActionEvent e) { } }); this.cfgMenu.add(subMenuIncludeHttpResponseContent); + //Add the menu to choose storage type + menuText = "Use PostgreSQL instead of SQLite"; + final JCheckBoxMenuItem subMenuUsePostgreSQL = new JCheckBoxMenuItem(menuText, USE_POSTGRESQL); + subMenuUsePostgreSQL.addActionListener(new AbstractAction(menuText) { + public void actionPerformed(ActionEvent e) { + if (subMenuUsePostgreSQL.isSelected()) { + // Show PostgreSQL configuration dialog + if (ActivityStorageFactory.showPostgreSQLConfigDialog(ConfigMenu.this.preferences, ConfigMenu.getBurpFrame())) { + try { + ConfigMenu.this.preferences.setBoolean(USE_POSTGRESQL_CFG_KEY, Boolean.TRUE); + ConfigMenu.USE_POSTGRESQL = Boolean.TRUE; + + // Create new PostgreSQL storage + String defaultStoreFileName = new File(System.getProperty("user.home"), "LogRequestsToSQLite.db").getAbsolutePath().replaceAll("\\\\", "/"); + String customStoreFileName = ConfigMenu.this.preferences.getString(ConfigMenu.DB_FILE_CUSTOM_LOCATION_CFG_KEY); + if (customStoreFileName == null) { + customStoreFileName = defaultStoreFileName; + } + + ActivityStorage newStorage = ActivityStorageFactory.createStorage( + ConfigMenu.this.preferences, + customStoreFileName, + ConfigMenu.this.api, + ConfigMenu.this.trace + ); + + // Replace storage without restart + ConfigMenu.this.replaceActivityStorage(newStorage); + + ConfigMenu.this.trace.writeLog("PostgreSQL storage enabled and active."); + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), + "PostgreSQL storage is now active. No restart required!", + "Configuration Updated", + JOptionPane.INFORMATION_MESSAGE); + } catch (Exception ex) { + ConfigMenu.this.trace.writeLog("Failed to switch to PostgreSQL storage: " + ex.getMessage()); + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), + "Failed to switch to PostgreSQL storage: " + ex.getMessage() + + "\nPlease check your PostgreSQL configuration and try again.", + "Configuration Error", + JOptionPane.ERROR_MESSAGE); + // Revert checkbox state + subMenuUsePostgreSQL.setSelected(false); + ConfigMenu.this.preferences.setBoolean(USE_POSTGRESQL_CFG_KEY, Boolean.FALSE); + ConfigMenu.USE_POSTGRESQL = Boolean.FALSE; + } + } else { + // User cancelled, revert checkbox + subMenuUsePostgreSQL.setSelected(false); + } + } else { + try { + ConfigMenu.this.preferences.setBoolean(USE_POSTGRESQL_CFG_KEY, Boolean.FALSE); + ConfigMenu.USE_POSTGRESQL = Boolean.FALSE; + + // Create new SQLite storage + String defaultStoreFileName = new File(System.getProperty("user.home"), "LogRequestsToSQLite.db").getAbsolutePath().replaceAll("\\\\", "/"); + String customStoreFileName = ConfigMenu.this.preferences.getString(ConfigMenu.DB_FILE_CUSTOM_LOCATION_CFG_KEY); + if (customStoreFileName == null) { + customStoreFileName = defaultStoreFileName; + } + + ActivityStorage newStorage = ActivityStorageFactory.createStorage( + ConfigMenu.this.preferences, + customStoreFileName, + ConfigMenu.this.api, + ConfigMenu.this.trace + ); + + // Replace storage without restart + ConfigMenu.this.replaceActivityStorage(newStorage); + + ConfigMenu.this.trace.writeLog("SQLite storage enabled and active."); + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), + "SQLite storage is now active. No restart required!", + "Configuration Updated", + JOptionPane.INFORMATION_MESSAGE); + } catch (Exception ex) { + ConfigMenu.this.trace.writeLog("Failed to switch to SQLite storage: " + ex.getMessage()); + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), + "Failed to switch to SQLite storage: " + ex.getMessage(), + "Configuration Error", + JOptionPane.ERROR_MESSAGE); + // Revert checkbox state + subMenuUsePostgreSQL.setSelected(true); + ConfigMenu.this.preferences.setBoolean(USE_POSTGRESQL_CFG_KEY, Boolean.TRUE); + ConfigMenu.USE_POSTGRESQL = Boolean.TRUE; + } + } + } + }); + this.cfgMenu.add(subMenuUsePostgreSQL); + //Add the menu to configure PostgreSQL connection + menuText = "Configure PostgreSQL Connection"; + final JMenuItem subMenuConfigurePostgreSQL = new JMenuItem(menuText); + subMenuConfigurePostgreSQL.addActionListener(new AbstractAction(menuText) { + public void actionPerformed(ActionEvent e) { + if (ActivityStorageFactory.showPostgreSQLConfigDialog(ConfigMenu.this.preferences, ConfigMenu.getBurpFrame())) { + try { + // If currently using PostgreSQL, update the connection + if (ConfigMenu.USE_POSTGRESQL) { + String defaultStoreFileName = new File(System.getProperty("user.home"), "LogRequestsToSQLite.db").getAbsolutePath().replaceAll("\\\\", "/"); + String customStoreFileName = ConfigMenu.this.preferences.getString(ConfigMenu.DB_FILE_CUSTOM_LOCATION_CFG_KEY); + if (customStoreFileName == null) { + customStoreFileName = defaultStoreFileName; + } + + ActivityStorage newStorage = ActivityStorageFactory.createStorage( + ConfigMenu.this.preferences, + customStoreFileName, + ConfigMenu.this.api, + ConfigMenu.this.trace + ); + + // Replace storage without restart + ConfigMenu.this.replaceActivityStorage(newStorage); + + ConfigMenu.this.trace.writeLog("PostgreSQL connection parameters updated and reconnected."); + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), + "PostgreSQL connection updated and active. No restart required!", + "Configuration Updated", + JOptionPane.INFORMATION_MESSAGE); + } else { + ConfigMenu.this.trace.writeLog("PostgreSQL connection parameters updated."); + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), + "PostgreSQL connection parameters saved. Enable PostgreSQL storage to use the new settings.", + "Configuration Updated", + JOptionPane.INFORMATION_MESSAGE); + } + } catch (Exception ex) { + ConfigMenu.this.trace.writeLog("Failed to update PostgreSQL connection: " + ex.getMessage()); + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), + "Failed to update PostgreSQL connection: " + ex.getMessage() + + "\nPlease check your PostgreSQL configuration and try again.", + "Configuration Error", + JOptionPane.ERROR_MESSAGE); + } + } + } + }); + this.cfgMenu.add(subMenuConfigurePostgreSQL); + //Add the menu to filter by tool source + menuText = "Select tools to log"; + final JCheckBoxMenuItem subMenuFilterByToolSource = new JCheckBoxMenuItem(menuText, FILTER_BY_TOOL_SOURCE); + subMenuFilterByToolSource.addActionListener(new AbstractAction(menuText) { + public void actionPerformed(ActionEvent e) { + if (subMenuFilterByToolSource.isSelected()) { + // Show dialog with checkboxes for each tool + if (showToolSelectionDialog()) { + ConfigMenu.this.preferences.setBoolean(FILTER_BY_TOOL_SOURCE_CFG_KEY, Boolean.TRUE); + ConfigMenu.this.preferences.setString(INCLUDED_TOOL_SOURCES_CFG_KEY, String.join(",", INCLUDED_TOOL_SOURCES)); + FILTER_BY_TOOL_SOURCE = Boolean.TRUE; + + ConfigMenu.this.trace.writeLog("Tool source filtering enabled. Included tools: " + INCLUDED_TOOL_SOURCES.toString()); + } else { + // User cancelled, uncheck the menu item + subMenuFilterByToolSource.setSelected(false); + } + } else { + ConfigMenu.this.preferences.setBoolean(FILTER_BY_TOOL_SOURCE_CFG_KEY, Boolean.FALSE); + FILTER_BY_TOOL_SOURCE = Boolean.FALSE; + // Reset to include all tools + INCLUDED_TOOL_SOURCES.clear(); + Collections.addAll(INCLUDED_TOOL_SOURCES, "Proxy", "Repeater", "Intruder", "Scanner", "Sequencer", "Spider", "Target", "Extender"); + ConfigMenu.this.trace.writeLog("Tool source filtering disabled. All tools will be logged."); + } + } + }); + this.cfgMenu.add(subMenuFilterByToolSource); //Add the menu to pause the logging menuText = "Pause the logging"; final JCheckBoxMenuItem subMenuPauseTheLogging = new JCheckBoxMenuItem(menuText, IS_LOGGING_PAUSED); @@ -205,14 +442,18 @@ public void actionPerformed(ActionEvent e) { } }); this.cfgMenu.add(subMenuPauseTheLogging); - //Add the menu to change the DB file - menuText = "Change the DB file"; + //Add the menu to change the DB file (SQLite only) + menuText = "Change the SQLite DB file"; final JMenuItem subMenuDBFileLocationMenuItem = new JMenuItem(menuText); subMenuDBFileLocationMenuItem.addActionListener( new AbstractAction(menuText) { public void actionPerformed(ActionEvent e) { try { - String title = "Change the DB file"; + String title = "Change the SQLite DB file"; + if (ConfigMenu.USE_POSTGRESQL) { + JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), "This option is only available when using SQLite storage.", title, JOptionPane.WARNING_MESSAGE); + return; + } if (!ConfigMenu.IS_LOGGING_PAUSED) { JOptionPane.showMessageDialog(ConfigMenu.getBurpFrame(), "Logging must be paused prior to update the DB file location!", title, JOptionPane.WARNING_MESSAGE); } else { @@ -221,9 +462,14 @@ public void actionPerformed(ActionEvent e) { int dbFileSelectionReply = customStoreFileNameFileChooser.showDialog(getBurpFrame(), "Use"); if (dbFileSelectionReply == JFileChooser.APPROVE_OPTION) { customStoreFileName = customStoreFileNameFileChooser.getSelectedFile().getAbsolutePath().replaceAll("\\\\", "/"); - activityLogger.updateStoreLocation(customStoreFileName); - ConfigMenu.this.preferences.setString(ConfigMenu.DB_FILE_CUSTOM_LOCATION_CFG_KEY, customStoreFileName); - JOptionPane.showMessageDialog(getBurpFrame(), "DB file updated to use:\n\r" + customStoreFileName, title, JOptionPane.INFORMATION_MESSAGE); + // Only works with SQLite ActivityLogger + if (ConfigMenu.this.activityStorage instanceof ActivityLogger) { + ((ActivityLogger) ConfigMenu.this.activityStorage).updateStoreLocation(customStoreFileName); + ConfigMenu.this.preferences.setString(ConfigMenu.DB_FILE_CUSTOM_LOCATION_CFG_KEY, customStoreFileName); + JOptionPane.showMessageDialog(getBurpFrame(), "DB file updated to use:\n\r" + customStoreFileName, title, JOptionPane.INFORMATION_MESSAGE); + } else { + JOptionPane.showMessageDialog(getBurpFrame(), "This feature is only available with SQLite storage.", title, JOptionPane.WARNING_MESSAGE); + } } else { JOptionPane.showMessageDialog(getBurpFrame(), "The following database file will continue to be used:\n\r" + customStoreFileName, title, JOptionPane.INFORMATION_MESSAGE); } @@ -243,7 +489,7 @@ public void actionPerformed(ActionEvent e) { public void actionPerformed(ActionEvent e) { try { //Get the data - DBStats stats = ConfigMenu.this.activityLogger.getEventsStats(); + DBStats stats = ConfigMenu.this.activityStorage.getEventsStats(); //Build the message String buffer = "Size of the database file on the disk: \n\r" + formatStat(stats.getSizeOnDisk()) + ".\n\r"; buffer += "Amount of data sent by the biggest HTTP request: \n\r" + formatStat(stats.getBiggestRequestSize()) + ".\n\r"; @@ -307,4 +553,92 @@ static String formatStat(long stat) { double amount = stat / unit; return String.format("%.2f %s", amount, unitLabel); } + + /** + * Replace the current activity storage with a new one. + * This allows switching between storage types without restarting Burp Suite. + * + * @param newStorage The new storage instance to use + */ + private void replaceActivityStorage(ActivityStorage newStorage) { + try { + // Pause logging temporarily + boolean wasLoggingPaused = IS_LOGGING_PAUSED; + IS_LOGGING_PAUSED = true; + + // Clean up old storage + if (this.activityStorage != null) { + this.activityStorage.extensionUnloaded(); + this.trace.writeLog("Old activity storage cleaned up."); + } + + // Replace storage instances + this.activityStorage = newStorage; + this.activityHttpListener.replaceStorage(newStorage); + + // Register new storage as unloading handler + this.api.extension().registerUnloadingHandler(newStorage); + + // Restore logging state + IS_LOGGING_PAUSED = wasLoggingPaused; + + this.trace.writeLog("Activity storage successfully replaced."); + } catch (Exception e) { + this.trace.writeLog("Error replacing activity storage: " + e.getMessage()); + throw new RuntimeException("Failed to replace activity storage", e); + } + } + + /** + * Show a dialog with checkboxes for tool selection. + * + * @return true if user confirmed the selection, false if cancelled + */ + private boolean showToolSelectionDialog() { + // Include both lowercase and uppercase versions to handle API variations + String[] allTools = {"Proxy", "Repeater", "Intruder", "Scanner", "Sequencer", "Spider", "Target", "Extender"}; + JCheckBox[] checkBoxes = new JCheckBox[allTools.length]; + + // Create checkboxes and set their initial state based on current included tools + for (int i = 0; i < allTools.length; i++) { + // Check for case-insensitive match + final String currentTool = allTools[i]; // Make it final for lambda + boolean isChecked = INCLUDED_TOOL_SOURCES.stream() + .anyMatch(includedTool -> includedTool.equalsIgnoreCase(currentTool)); + checkBoxes[i] = new JCheckBox(currentTool, isChecked); + } + + // Create a panel to hold the checkboxes + JPanel panel = new JPanel(); + panel.setLayout(new GridLayout(0, 2)); // 2 columns + panel.add(new JLabel("Select tools to log:")); + panel.add(new JLabel("")); // Empty cell for spacing + + for (JCheckBox checkBox : checkBoxes) { + panel.add(checkBox); + } + + // Show the dialog + int result = JOptionPane.showConfirmDialog( + getBurpFrame(), + panel, + "Tool Selection", + JOptionPane.OK_CANCEL_OPTION, + JOptionPane.PLAIN_MESSAGE + ); + + if (result == JOptionPane.OK_OPTION) { + // Update the included tools list based on checkbox selections + INCLUDED_TOOL_SOURCES.clear(); + for (int i = 0; i < checkBoxes.length; i++) { + if (checkBoxes[i].isSelected()) { + INCLUDED_TOOL_SOURCES.add(allTools[i]); + } + } + this.trace.writeLog("Tool selection updated. New included tools: " + INCLUDED_TOOL_SOURCES.toString()); + return true; + } + + return false; + } } diff --git a/src/main/java/burp/PostgreSQLActivityLogger.java b/src/main/java/burp/PostgreSQLActivityLogger.java new file mode 100644 index 0000000..14ba5c7 --- /dev/null +++ b/src/main/java/burp/PostgreSQLActivityLogger.java @@ -0,0 +1,470 @@ +package burp; + +import java.net.InetAddress; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.time.LocalDateTime; +import java.time.format.DateTimeFormatter; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import burp.api.montoya.MontoyaApi; +import burp.api.montoya.http.message.requests.HttpRequest; +import burp.api.montoya.http.message.responses.HttpResponse; +import burp.api.montoya.extension.ExtensionUnloadingHandler; + +/** + * Handle the recording of the activities into PostgreSQL database. + * Uses async writes with a background thread for improved performance. + */ +class PostgreSQLActivityLogger implements ActivityStorage { + + /** + * SQL instructions for PostgreSQL. + */ + private static final String SQL_TABLE_CREATE = "CREATE TABLE IF NOT EXISTS ACTIVITY (" + + "id SERIAL PRIMARY KEY, " + + "local_source_ip TEXT, " + + "target_url TEXT, " + + "http_method TEXT, " + + "burp_tool TEXT, " + + "request_raw TEXT, " + + "send_datetime TIMESTAMP, " + + "http_status_code TEXT, " + + "response_raw TEXT, " + + "created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)"; + + private static final String SQL_TABLE_INSERT = "INSERT INTO ACTIVITY " + + "(local_source_ip, target_url, http_method, burp_tool, request_raw, send_datetime, http_status_code, response_raw) " + + "VALUES(?,?,?,?,?,?::timestamp,?,?)"; + + private static final String SQL_COUNT_RECORDS = "SELECT COUNT(http_method) FROM ACTIVITY"; + private static final String SQL_TOTAL_AMOUNT_DATA_SENT = "SELECT SUM(LENGTH(request_raw)) FROM ACTIVITY"; + private static final String SQL_BIGGEST_REQUEST_AMOUNT_DATA_SENT = "SELECT MAX(LENGTH(request_raw)) FROM ACTIVITY"; + private static final String SQL_MAX_HITS_BY_SECOND = "SELECT COUNT(request_raw) AS hits, " + + "DATE_TRUNC('second', send_datetime) as second_bucket " + + "FROM ACTIVITY GROUP BY second_bucket ORDER BY hits DESC LIMIT 1"; + + /** + * Maximum queue size to prevent memory issues + */ + private static final int MAX_QUEUE_SIZE = 10000; + + /** + * Batch size for database writes + */ + private static final int BATCH_SIZE = 100; + + /** + * Maximum wait time for batch processing (milliseconds) + */ + private static final long BATCH_TIMEOUT_MS = 1000; + + /** + * Use a single DB connection for performance. + */ + private Connection storageConnection; + + /** + * Database connection parameters + */ + private String host; + private int port; + private String database; + private String username; + private String password; + + /** + * Ref on project logger. + */ + private Trace trace; + + /** + * Formatter for date/time. + */ + private DateTimeFormatter datetimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"); + + /** + * Queue for async event processing + */ + private final BlockingQueue eventQueue = new LinkedBlockingQueue<>(MAX_QUEUE_SIZE); + + /** + * Background thread for database writes + */ + private Thread writerThread; + + /** + * Flag to control the writer thread lifecycle + */ + private final AtomicBoolean running = new AtomicBoolean(true); + + /** + * Constructor. + * + * @param host PostgreSQL host + * @param port PostgreSQL port + * @param database Database name + * @param username Database username + * @param password Database password + * @param api Montoya API reference + * @param trace Ref on project logger. + * @throws Exception If connection with the DB cannot be opened or if the DB cannot be created or if the JDBC driver cannot be loaded. + */ + PostgreSQLActivityLogger(String host, int port, String database, String username, String password, MontoyaApi api, Trace trace) throws Exception { + //Load the PostgreSQL driver + Class.forName("org.postgresql.Driver"); + this.trace = trace; + this.host = host; + this.port = port; + this.database = database; + this.username = username; + this.password = password; + + initializeConnection(); + startWriterThread(); + } + + /** + * Initialize the database connection and create table if needed + */ + private void initializeConnection() throws Exception { + String url = String.format("jdbc:postgresql://%s:%d/%s", host, port, database); + this.storageConnection = DriverManager.getConnection(url, username, password); + this.storageConnection.setAutoCommit(true); + this.trace.writeLog("Connected to PostgreSQL database at " + host + ":" + port + "/" + database); + + try (Statement stmt = this.storageConnection.createStatement()) { + stmt.execute(SQL_TABLE_CREATE); + this.trace.writeLog("PostgreSQL recording table initialized."); + } + } + + /** + * Start the background writer thread + */ + private void startWriterThread() { + writerThread = new Thread(this::processEventQueue, "PostgreSQLActivityLogger-Writer"); + writerThread.setDaemon(true); + writerThread.start(); + this.trace.writeLog("PostgreSQL async writer thread started."); + } + + /** + * Background thread that processes the event queue + */ + private void processEventQueue() { + LogEvent[] batch = new LogEvent[BATCH_SIZE]; + + while (running.get() || !eventQueue.isEmpty()) { + try { + int batchCount = 0; + long batchStartTime = System.currentTimeMillis(); + + // Collect events for batching + while (batchCount < BATCH_SIZE && running.get()) { + LogEvent event = eventQueue.poll(BATCH_TIMEOUT_MS, TimeUnit.MILLISECONDS); + if (event == null) { + break; // Timeout reached + } + batch[batchCount++] = event; + + // Check if we should flush early due to timeout + if (System.currentTimeMillis() - batchStartTime >= BATCH_TIMEOUT_MS) { + break; + } + } + + // Process the batch if we have events + if (batchCount > 0) { + writeBatch(batch, batchCount); + } + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } catch (Exception e) { + this.trace.writeLog("Error in PostgreSQL writer thread: " + e.getMessage()); + // Try to reconnect if connection issues + try { + if (this.storageConnection.isClosed()) { + initializeConnection(); + this.trace.writeLog("Reconnected to PostgreSQL database."); + } + } catch (Exception reconnectEx) { + this.trace.writeLog("Failed to reconnect to PostgreSQL: " + reconnectEx.getMessage()); + } + } + } + + // Flush remaining events when shutting down + flushRemainingEvents(); + } + + /** + * Write a batch of events to the database + */ + private void writeBatch(LogEvent[] batch, int count) throws Exception { + ensureDBState(); + + // Use batch inserts for better performance + this.storageConnection.setAutoCommit(false); + + try (PreparedStatement stmt = this.storageConnection.prepareStatement(SQL_TABLE_INSERT)) { + for (int i = 0; i < count; i++) { + LogEvent event = batch[i]; + stmt.setString(1, event.localSourceIp); + stmt.setString(2, event.targetUrl); + stmt.setString(3, event.httpMethod); + stmt.setString(4, event.tool); + stmt.setString(5, event.requestRaw); + stmt.setString(6, event.sendDateTime); + stmt.setString(7, event.httpStatusCode); + stmt.setString(8, event.responseRaw); + stmt.addBatch(); + } + + int[] results = stmt.executeBatch(); + this.storageConnection.commit(); + + // Log any failed inserts + int successCount = 0; + for (int result : results) { + if (result > 0) successCount++; + } + + if (successCount != count) { + this.trace.writeLog("PostgreSQL batch insert: " + successCount + "/" + count + " events inserted successfully"); + } + + } catch (Exception e) { + this.storageConnection.rollback(); + throw e; + } finally { + this.storageConnection.setAutoCommit(true); + } + } + + /** + * Flush any remaining events in the queue (used during shutdown) + */ + private void flushRemainingEvents() { + LogEvent[] batch = new LogEvent[BATCH_SIZE]; + int count = 0; + + while (!eventQueue.isEmpty() && count < BATCH_SIZE) { + LogEvent event = eventQueue.poll(); + if (event != null) { + batch[count++] = event; + } + } + + if (count > 0) { + try { + writeBatch(batch, count); + this.trace.writeLog("PostgreSQL flushed " + count + " remaining events during shutdown."); + } catch (Exception e) { + this.trace.writeLog("Error flushing remaining PostgreSQL events: " + e.getMessage()); + } + } + } + + /** + * Save an activity event into the storage (now async). + * + * @param request HttpRequest object containing all information about the request + * which was either sent or will be sent out soon. + * @param response HttpResponse object containing all information about the response. + * Is null when only the request ist stored. + * @param tool The name of the tool which was used to issue to request. + * @throws Exception If event cannot be saved. + */ + public void logEvent(HttpRequest request, HttpResponse response, String tool) throws Exception { + try { + // Create event object with pre-computed values + LogEvent event = new LogEvent( + InetAddress.getLocalHost().getHostAddress(), + request.url(), + request.method(), + tool, + request.toString(), + LocalDateTime.now().format(this.datetimeFormatter), + response != null ? String.valueOf(response.statusCode()) : null, + response != null ? response.bodyToString() : null + ); + + // Add to queue (non-blocking) + if (!eventQueue.offer(event)) { + // Queue is full - could log a warning or implement backpressure + this.trace.writeLog("PostgreSQL event queue full, dropping event. Consider adjusting MAX_QUEUE_SIZE."); + } + + } catch (Exception e) { + this.trace.writeLog("Error queueing PostgreSQL event: " + e.getMessage()); + // Could fallback to synchronous write in critical cases + } + } + + /** + * Update database connection parameters and reconnect. + * + * @param host PostgreSQL host + * @param port PostgreSQL port + * @param database Database name + * @param username Database username + * @param password Database password + * @throws Exception If connection with the DB cannot be opened or if the DB cannot be created. + */ + void updateConnectionParameters(String host, int port, String database, String username, String password) throws Exception { + this.host = host; + this.port = port; + this.database = database; + this.username = username; + this.password = password; + + // Close existing connection + if (this.storageConnection != null && !this.storageConnection.isClosed()) { + this.storageConnection.close(); + } + + // Create new connection + initializeConnection(); + this.trace.writeLog("PostgreSQL connection parameters updated and reconnected."); + } + + /** + * Extract and compute statistics about the DB. + * + * @return A VO object containing the statistics. + * @throws Exception If computation meets an error. + */ + public DBStats getEventsStats() throws Exception { + //Verify that the DB connection is still opened + this.ensureDBState(); + + //Get the total of the records in the activity table + long recordsCount; + try (PreparedStatement stmt = this.storageConnection.prepareStatement(SQL_COUNT_RECORDS)) { + try (ResultSet rst = stmt.executeQuery()) { + recordsCount = rst.next() ? rst.getLong(1) : 0; + } + } + + //Get data amount if the DB is not empty + long totalAmountDataSent = 0; + long biggestRequestAmountDataSent = 0; + long maxHitsBySecond = 0; + + if (recordsCount > 0) { + //Get the total amount of data sent, we assume here that 1 character = 1 byte + try (PreparedStatement stmt = this.storageConnection.prepareStatement(SQL_TOTAL_AMOUNT_DATA_SENT)) { + try (ResultSet rst = stmt.executeQuery()) { + if (rst.next()) { + totalAmountDataSent = rst.getLong(1); + } + } + } + + //Get the amount of data sent by the biggest request, we assume here that 1 character = 1 byte + try (PreparedStatement stmt = this.storageConnection.prepareStatement(SQL_BIGGEST_REQUEST_AMOUNT_DATA_SENT)) { + try (ResultSet rst = stmt.executeQuery()) { + if (rst.next()) { + biggestRequestAmountDataSent = rst.getLong(1); + } + } + } + + //Get the maximum number of hits sent in a second + try (PreparedStatement stmt = this.storageConnection.prepareStatement(SQL_MAX_HITS_BY_SECOND)) { + try (ResultSet rst = stmt.executeQuery()) { + if (rst.next()) { + maxHitsBySecond = rst.getLong(1); + } + } + } + } + + //For PostgreSQL, we'll estimate DB size based on table statistics + //This is an approximation since PostgreSQL doesn't have a simple file size equivalent + long estimatedSize = recordsCount * 1024; // Rough estimate + + //Build the VO and return it + return new DBStats(estimatedSize, recordsCount, totalAmountDataSent, biggestRequestAmountDataSent, maxHitsBySecond); + } + + /** + * Ensure the connection to the DB is valid. + * + * @throws Exception If connection cannot be verified or opened. + */ + private void ensureDBState() throws Exception { + //Verify that the DB connection is still opened + if (this.storageConnection.isClosed()) { + //Get new one + this.trace.writeLog("PostgreSQL connection lost, reconnecting..."); + initializeConnection(); + } + } + + /** + * Unloads the extension by releasing the DB connection. + */ + @Override + public void extensionUnloaded() { + // Signal the writer thread to stop + running.set(false); + + // Wait for writer thread to finish processing + if (writerThread != null) { + try { + writerThread.interrupt(); + writerThread.join(5000); // Wait up to 5 seconds + this.trace.writeLog("PostgreSQL writer thread stopped."); + } catch (InterruptedException e) { + this.trace.writeLog("Interrupted while waiting for PostgreSQL writer thread to finish."); + Thread.currentThread().interrupt(); + } + } + + // Close database connection + try { + if (this.storageConnection != null && !this.storageConnection.isClosed()) { + this.storageConnection.close(); + this.trace.writeLog("PostgreSQL connection released."); + } + } catch (Exception e) { + this.trace.writeLog("Cannot close the PostgreSQL connection: " + e.getMessage()); + } + } + + /** + * Value object to hold event data for async processing + */ + private static class LogEvent { + final String localSourceIp; + final String targetUrl; + final String httpMethod; + final String tool; + final String requestRaw; + final String sendDateTime; + final String httpStatusCode; + final String responseRaw; + + LogEvent(String localSourceIp, String targetUrl, String httpMethod, String tool, + String requestRaw, String sendDateTime, String httpStatusCode, String responseRaw) { + this.localSourceIp = localSourceIp; + this.targetUrl = targetUrl; + this.httpMethod = httpMethod; + this.tool = tool; + this.requestRaw = requestRaw; + this.sendDateTime = sendDateTime; + this.httpStatusCode = httpStatusCode; + this.responseRaw = responseRaw; + } + } +} \ No newline at end of file