1010import java .time .LocalDateTime ;
1111import java .time .format .DateTimeFormatter ;
1212
13+ import burp .api .montoya .MontoyaApi ;
14+ import burp .api .montoya .http .message .requests .HttpRequest ;
15+ import burp .api .montoya .http .message .responses .HttpResponse ;
16+ import burp .api .montoya .extension .ExtensionUnloadingHandler ;
17+
1318/**
1419 * Handle the recording of the activities into the real storage, SQLite local DB here.
1520 */
16- class ActivityLogger implements IExtensionStateListener {
21+ class ActivityLogger implements ExtensionUnloadingHandler {
1722
1823 /**
1924 * SQL instructions.
@@ -25,12 +30,6 @@ class ActivityLogger implements IExtensionStateListener {
2530 private static final String SQL_BIGGEST_REQUEST_AMOUNT_DATA_SENT = "SELECT MAX(LENGTH(REQUEST_RAW)) FROM ACTIVITY" ;
2631 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" ;
2732
28- /**
29- * Empty string to use when response must be not be logged.
30- */
31- private static final String EMPTY_RESPONSE_CONTENT = "" ;
32-
33-
3433 /**
3534 * Use a single DB connection for performance and to prevent DB file locking issue at filesystem level.
3635 */
@@ -41,11 +40,6 @@ class ActivityLogger implements IExtensionStateListener {
4140 */
4241 private String url ;
4342
44- /**
45- * Ref on Burp tool to manipulate the HTTP requests and have access to API to identify the source of the activity (tool name).
46- */
47- private IBurpExtenderCallbacks callbacks ;
48-
4943 /**
5044 * Ref on project logger.
5145 */
@@ -60,16 +54,13 @@ class ActivityLogger implements IExtensionStateListener {
6054 /**
6155 * Constructor.
6256 *
63- * @param storeName Name of the storage that will be created (file path).
64- * @param callbacks Ref on Burp tool to manipulate the HTTP requests and have access to API to identify the source of the activity (tool name).
65- * @param trace Ref on project logger.
66- * @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.
57+ * @param storeName Name of the storage that will be created (file path).
58+ * @param trace Ref on project logger.
59+ * @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.
6760 */
68- ActivityLogger (String storeName , IBurpExtenderCallbacks callbacks , Trace trace ) throws Exception {
61+ ActivityLogger (String storeName , MontoyaApi api , Trace trace ) throws Exception {
6962 //Load the SQLite driver
7063 Class .forName ("org.sqlite.JDBC" );
71- //Affect the properties
72- this .callbacks = callbacks ;
7364 this .trace = trace ;
7465 updateStoreLocation (storeName );
7566 }
@@ -98,26 +89,32 @@ void updateStoreLocation(String storeName) throws Exception {
9889 /**
9990 * Save an activity event into the storage.
10091 *
101- * @param toolFlag A flag indicating the Burp tool that issued the request.
102- * Burp tool flags are defined in the
103- * <code>IBurpExtenderCallbacks</code> interface .
104- * @param reqInfo Details of the request to be processed .
105- * @param reqContent Raw content of the request.
106- * @throws Exception If event cannot be saved.
92+ * @param request HttpRequest object containing all information about the request
93+ * which was either sent or will be sent out soon.
94+ * @param response HttpResponse object containing all information about the response .
95+ * Is null when only the request ist stored .
96+ * @param tool The name of the tool which was used to issue to request.
97+ * @throws Exception If event cannot be saved.
10798 */
108- void logEvent (int toolFlag , IRequestInfo reqInfo , byte [] reqContent , String statusCode , byte [] resContent ) throws Exception {
99+ void logEvent (HttpRequest request , HttpResponse response , String tool ) throws Exception {
109100 //Verify that the DB connection is still opened
110101 this .ensureDBState ();
111102 //Insert the event into the storage
112103 try (PreparedStatement stmt = this .storageConnection .prepareStatement (SQL_TABLE_INSERT )) {
113104 stmt .setString (1 , InetAddress .getLocalHost ().getHostAddress ());
114- stmt .setString (2 , reqInfo . getUrl (). toString ());
115- stmt .setString (3 , reqInfo . getMethod ());
116- stmt .setString (4 , callbacks . getToolName ( toolFlag ) );
117- stmt .setString (5 , callbacks . getHelpers (). bytesToString ( reqContent ));
105+ stmt .setString (2 , request . url ());
106+ stmt .setString (3 , request . method ());
107+ stmt .setString (4 , tool );
108+ stmt .setString (5 , request . toString ()); //Apparently, bodyToString() does not work..
118109 stmt .setString (6 , LocalDateTime .now ().format (this .datetimeFormatter ));
119- stmt .setString (7 , statusCode );
120- stmt .setString (8 , (resContent != null ) ? callbacks .getHelpers ().bytesToString (resContent ) : EMPTY_RESPONSE_CONTENT );
110+ //Make a distinction if only the request is stored or the response is added as well.
111+ if (response != null ) {
112+ stmt .setString (7 , String .valueOf (response .statusCode ()));
113+ stmt .setString (8 , response .bodyToString ());
114+ } else {
115+ stmt .setString (7 , null );
116+ stmt .setString (8 , null );
117+ }
121118 int count = stmt .executeUpdate ();
122119 if (count != 1 ) {
123120 this .trace .writeLog ("Request was not inserted, no detail available (insertion counter = " + count + ") !" );
@@ -188,8 +185,9 @@ private void ensureDBState() throws Exception {
188185 }
189186
190187 /**
191- * {@inheritDoc}
188+ * Unloads the extension by releasing the DB connection.
192189 */
190+ @ Override
193191 public void extensionUnloaded () {
194192 try {
195193 if (this .storageConnection != null && !this .storageConnection .isClosed ()) {
0 commit comments