Skip to content

Commit 4c7d1c2

Browse files
bytebutcherbytebutcher
committed
Refactored code. Access to IContextMenuInvocation is now more limited/controlled.
1 parent f2f667e commit 4c7d1c2

26 files changed

+147
-138
lines changed

‎burp-send-to-extension/build.gradle‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
apply plugin: 'java'
22

33
group 'net.bytebutcher'
4-
version '1.1'
4+
version '1.2'
55

66
sourceCompatibility = 1.8
77

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/gui/SendToContextMenu.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public SendToContextMenu(BurpExtender burpExtender, SendToTableListener sendToTa
2727

2828
@Override
2929
public List<JMenuItem> createMenuItems(IContextMenuInvocation invocation) {
30-
List<Map<String, IPlaceholder>> placeholders = Placeholders.get(BurpExtender.getCallbacks(), invocation);
30+
List<Map<String, IPlaceholder>> placeholders = Placeholders.get(BurpExtender.getCallbacks(), invocation.getSelectedMessages());
3131
List<CommandObject> commandObjects = BurpExtender.getConfig().getSendToTableData();
3232
if (commandObjects.isEmpty()) {
3333
return Lists.newArrayList();
@@ -64,15 +64,16 @@ public List<JMenuItem> createMenuItems(IContextMenuInvocation invocation) {
6464
private void addMenuItem(JMenu menu, CommandObject commandObject, List<Map<String, IPlaceholder>> placeholders, IContextMenuInvocation invocation) {
6565

6666
JMenuItem item;
67-
if (commandObject.doesRequireRequestResponse(placeholders.get(0)) && Context.getContext(invocation) == Context.UNKNOWN) {
67+
Context context = new Context(invocation);
68+
if (commandObject.doesRequireRequestResponse(placeholders.get(0)) && context.getOrigin() == Context.Origin.UNKNOWN) {
6869
item = new JMenu(commandObject.getName());
69-
SendToContextMenuItem request = new SendToContextMenuItem("request", commandObject, placeholders, invocation, Context.HTTP_REQUEST, sendToTableListener);
70-
SendToContextMenuItem response = new SendToContextMenuItem("response", commandObject, placeholders, invocation, Context.HTTP_RESPONSE, sendToTableListener);
70+
SendToContextMenuItem request = new SendToContextMenuItem("request", commandObject, placeholders, new Context(Context.Origin.HTTP_REQUEST, invocation), sendToTableListener);
71+
SendToContextMenuItem response = new SendToContextMenuItem("response", commandObject, placeholders, new Context(Context.Origin.HTTP_RESPONSE, invocation), sendToTableListener);
7172
item.add(request);
7273
item.add(response);
7374
item.setEnabled(request.isEnabled() || response.isEnabled());
7475
} else {
75-
item = new SendToContextMenuItem(commandObject.getName(), commandObject, placeholders, invocation, Context.getContext(invocation), sendToTableListener);
76+
item = new SendToContextMenuItem(commandObject.getName(), commandObject, placeholders, context, sendToTableListener);
7677
}
7778
menu.add(item);
7879
}

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/gui/SendToContextMenuItem.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package net.bytebutcher.burpsendtoextension.gui;
22

3-
import burp.IContextMenuInvocation;
43
import net.bytebutcher.burpsendtoextension.gui.action.SendToContextMenuItemAction;
54
import net.bytebutcher.burpsendtoextension.models.CommandObject;
65
import net.bytebutcher.burpsendtoextension.models.Context;
@@ -12,15 +11,15 @@
1211

1312
public class SendToContextMenuItem extends JMenuItem {
1413

15-
public SendToContextMenuItem(String title, CommandObject commandObject, List<Map<String, IPlaceholder>> placeholders, IContextMenuInvocation invocation, Context context, SendToTableListener sendToTableListener) {
14+
public SendToContextMenuItem(String title, CommandObject commandObject, List<Map<String, IPlaceholder>> placeholders, Context context, SendToTableListener sendToTableListener) {
1615
String text = "";
1716
List<Map<String, IPlaceholder>> validEntries = commandObject.getValid(placeholders, context);
1817
if (placeholders.size() > 1) {
1918
text = title + " (" + validEntries.size() + "/" + placeholders.size() + ")";
2019
} else {
2120
text = title;
2221
}
23-
this.setAction(new SendToContextMenuItemAction(text, commandObject, placeholders, invocation, sendToTableListener, context));
22+
this.setAction(new SendToContextMenuItemAction(text, commandObject, placeholders, sendToTableListener, context));
2423
this.setEnabled(validEntries.size() > 0);
2524
}
2625

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/gui/action/SendToContextMenuItemAction.java‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package net.bytebutcher.burpsendtoextension.gui.action;
22

33
import burp.BurpExtender;
4-
import burp.IContextMenuInvocation;
4+
import burp.IHttpRequestResponse;
55
import net.bytebutcher.burpsendtoextension.gui.SendToPreviewDialog;
66
import net.bytebutcher.burpsendtoextension.gui.SendToTableListener;
77
import net.bytebutcher.burpsendtoextension.gui.util.DialogUtil;
@@ -13,6 +13,7 @@
1313

1414
import javax.swing.*;
1515
import java.awt.event.ActionEvent;
16+
import java.io.IOException;
1617
import java.io.PrintWriter;
1718
import java.io.StringWriter;
1819
import java.time.ZonedDateTime;
@@ -24,23 +25,39 @@ public class SendToContextMenuItemAction extends AbstractAction {
2425

2526
private final CommandObject commandObject;
2627
private final List<Map<String, IPlaceholder>> placeholders;
27-
private final IContextMenuInvocation invocation;
2828
private final SendToTableListener sendToTableListener;
2929
private final Context context;
3030

31-
public SendToContextMenuItemAction(String title, CommandObject commandObject, List<Map<String, IPlaceholder>> placeholders, IContextMenuInvocation invocation, SendToTableListener sendToTableListener, Context context) {
31+
public SendToContextMenuItemAction(String title, CommandObject commandObject, List<Map<String, IPlaceholder>> placeholders, SendToTableListener sendToTableListener, Context context) {
3232
super(title);
3333
this.commandObject = commandObject;
3434
this.placeholders = placeholders;
35-
this.invocation = invocation;
3635
this.sendToTableListener = sendToTableListener;
3736
this.context = context;
3837
}
3938
@Override
4039
public void actionPerformed(ActionEvent e) {
4140
String result = runCommandObject(commandObject);
42-
if (commandObject.shouldOutputReplaceSelection() && result != null) {
43-
SelectionUtil.replaceSelectedText(invocation, result);
41+
if (result != null && commandObject.shouldOutputReplaceSelection()) {
42+
replaceSelectedText(context, result);
43+
}
44+
}
45+
46+
private void replaceSelectedText(Context context, String replaceText) {
47+
try {
48+
if (context.getSelectedMessages() != null && context.getSelectedMessages().length > 0) {
49+
IHttpRequestResponse message = context.getSelectedMessages()[0];
50+
switch (context.getOrigin()) {
51+
case HTTP_REQUEST:
52+
message.setRequest(SelectionUtil.replaceSelectedText(message.getRequest(), context.getSelectionBounds(), replaceText));
53+
break;
54+
case HTTP_RESPONSE:
55+
message.setResponse(SelectionUtil.replaceSelectedText(message.getResponse(), context.getSelectionBounds(), replaceText));
56+
break;
57+
}
58+
}
59+
} catch (IOException e) {
60+
BurpExtender.getCallbacks().printError("Error during replacing selection with output: " + e.toString());
4461
}
4562
}
4663

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/gui/util/SelectionUtil.java‎

Lines changed: 9 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,18 @@
11
package net.bytebutcher.burpsendtoextension.gui.util;
22

3-
import burp.BurpExtender;
4-
import burp.IContextMenuInvocation;
5-
63
import java.io.ByteArrayOutputStream;
74
import java.io.IOException;
85
import java.util.Arrays;
96

107
public class SelectionUtil {
11-
public static void replaceSelectedText(IContextMenuInvocation invocation, String replaceText) {
12-
try {
13-
if (invocation.getInvocationContext() == IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_REQUEST || invocation.getInvocationContext() == IContextMenuInvocation.CONTEXT_MESSAGE_VIEWER_REQUEST) {
14-
int[] bounds = invocation.getSelectionBounds();
15-
byte[] message = invocation.getSelectedMessages()[0].getRequest();
16-
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
17-
outputStream.write(Arrays.copyOfRange(message, 0, bounds[0]));
18-
outputStream.write(replaceText.getBytes());
19-
outputStream.write(Arrays.copyOfRange(message, bounds[1], message.length));
20-
outputStream.flush();
21-
invocation.getSelectedMessages()[0].setRequest(outputStream.toByteArray());
22-
}
23-
} catch (IOException e) {
24-
BurpExtender.getCallbacks().printError("Error during replacing selection with output: " + e.toString());
25-
}
8+
9+
public static byte[] replaceSelectedText(byte[] message, int[] bounds, String replaceText) throws IOException {
10+
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
11+
outputStream.write(Arrays.copyOfRange(message, 0, bounds[0]));
12+
outputStream.write(replaceText.getBytes());
13+
outputStream.write(Arrays.copyOfRange(message, bounds[1], message.length));
14+
outputStream.flush();
15+
return outputStream.toByteArray();
2616
}
17+
2718
}

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/models/Config.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ public class Config {
1414

1515
private final IBurpExtenderCallbacks callbacks;
1616
private BurpExtender burpExtender;
17-
private String version = "1.1";
17+
private String version = "1.2";
1818

1919
public Config(BurpExtender burpExtender) {
2020
this.burpExtender = burpExtender;

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/models/Context.java‎

Lines changed: 45 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,55 @@
11
package net.bytebutcher.burpsendtoextension.models;
22

33
import burp.IContextMenuInvocation;
4+
import burp.IHttpRequestResponse;
45
import com.google.common.collect.Lists;
56

67
import java.util.ArrayList;
78

8-
public enum Context {
9-
HTTP_RESPONSE, HTTP_REQUEST, UNKNOWN;
10-
11-
public static Context getContext(IContextMenuInvocation contextMenuInvocation) {
12-
ArrayList<Byte> requestContext = Lists.newArrayList(IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_REQUEST, IContextMenuInvocation.CONTEXT_MESSAGE_VIEWER_REQUEST);
13-
ArrayList<Byte> responseContext = Lists.newArrayList(IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_RESPONSE, IContextMenuInvocation.CONTEXT_MESSAGE_VIEWER_RESPONSE);
14-
if (requestContext.contains(contextMenuInvocation.getInvocationContext())) {
15-
return Context.HTTP_REQUEST;
16-
} else if (responseContext.contains(contextMenuInvocation.getInvocationContext())) {
17-
return Context.HTTP_RESPONSE;
18-
} else {
19-
return Context.UNKNOWN;
9+
public class Context {
10+
private final Origin origin;
11+
private final IContextMenuInvocation invocation;
12+
13+
public enum Origin {
14+
HTTP_RESPONSE, HTTP_REQUEST, UNKNOWN;
15+
16+
public static Origin getTarget(IContextMenuInvocation contextMenuInvocation) {
17+
ArrayList<Byte> requestContext = Lists.newArrayList(IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_REQUEST, IContextMenuInvocation.CONTEXT_MESSAGE_VIEWER_REQUEST);
18+
ArrayList<Byte> responseContext = Lists.newArrayList(IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_RESPONSE, IContextMenuInvocation.CONTEXT_MESSAGE_VIEWER_RESPONSE);
19+
if (requestContext.contains(contextMenuInvocation.getInvocationContext())) {
20+
return Origin.HTTP_REQUEST;
21+
} else if (responseContext.contains(contextMenuInvocation.getInvocationContext())) {
22+
return Origin.HTTP_RESPONSE;
23+
} else {
24+
return Origin.UNKNOWN;
25+
}
2026
}
27+
28+
}
29+
30+
public Context(Origin origin, IContextMenuInvocation invocation) {
31+
this.origin = origin;
32+
this.invocation = invocation;
33+
}
34+
35+
public Context(IContextMenuInvocation invocation) {
36+
this(Origin.getTarget(invocation), invocation);
37+
}
38+
39+
public Origin getOrigin() {
40+
return origin;
41+
}
42+
43+
/**
44+
* NOTE: Access to getSelectedMessages should be restricted, since messages should always be accessed via
45+
* RequestResponseHolder. Notably the definition and usage of this function here is a minor design issue.
46+
*/
47+
public IHttpRequestResponse[] getSelectedMessages() {
48+
return invocation.getSelectedMessages();
49+
}
50+
51+
public int[] getSelectionBounds() {
52+
return invocation.getSelectionBounds();
2153
}
2254

23-
}
55+
}

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/models/Placeholders.java‎

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package net.bytebutcher.burpsendtoextension.models;
22

33
import burp.IBurpExtenderCallbacks;
4-
import burp.IContextMenuInvocation;
54
import burp.IHttpRequestResponse;
65
import burp.RequestResponseHolder;
76
import com.google.common.collect.Lists;
@@ -16,25 +15,28 @@ public class Placeholders {
1615
/**
1716
* Initializes the placeholders for each selected message and returns them in a list.
1817
*/
19-
public static List<Map<String, IPlaceholder>> get(IBurpExtenderCallbacks burpExtenderCallbacks, IContextMenuInvocation invocation) {
18+
public static List<Map<String, IPlaceholder>> get(IBurpExtenderCallbacks burpExtenderCallbacks, IHttpRequestResponse[] selectedMessages) {
2019
List<Map<String, IPlaceholder>> result = Lists.newArrayList();
21-
for (IHttpRequestResponse selectedMessage : invocation.getSelectedMessages()) {
20+
if (selectedMessages == null) {
21+
return result;
22+
}
23+
for (IHttpRequestResponse selectedMessage : selectedMessages) {
2224
RequestResponseHolder requestResponseHolder = new RequestResponseHolder(burpExtenderCallbacks, selectedMessage);
2325
Map<String, IPlaceholder> placeholderMap = Maps.newHashMap();
2426
List<IPlaceholder> placeholderList = Lists.newArrayList(
25-
new CookiesPlaceholder(requestResponseHolder, invocation),
26-
new HostPlaceholder(requestResponseHolder, invocation),
27-
new HttpBodyToFilePlaceholder(requestResponseHolder, invocation),
28-
new HttpHeadersToFilePlaceholder(requestResponseHolder, invocation),
29-
new HttpMethodPlaceholder(requestResponseHolder, invocation),
30-
new HttpRequestResponsePlaceholder(requestResponseHolder, invocation),
31-
new PortPlaceholder(requestResponseHolder, invocation),
32-
new ProtocolPlaceholder(requestResponseHolder, invocation),
33-
new SelectedTextPlaceholder(requestResponseHolder, invocation),
34-
new SelectedTextToFilePlaceholder(requestResponseHolder, invocation),
35-
new UrlPathPlaceholder(requestResponseHolder, invocation),
36-
new UrlPlaceholder(requestResponseHolder, invocation),
37-
new UrlQueryPlaceholder(requestResponseHolder, invocation));
27+
new CookiesPlaceholder(requestResponseHolder),
28+
new HostPlaceholder(requestResponseHolder),
29+
new HttpBodyToFilePlaceholder(requestResponseHolder),
30+
new HttpHeadersToFilePlaceholder(requestResponseHolder),
31+
new HttpMethodPlaceholder(requestResponseHolder),
32+
new HttpRequestResponsePlaceholder(requestResponseHolder),
33+
new PortPlaceholder(requestResponseHolder),
34+
new ProtocolPlaceholder(requestResponseHolder),
35+
new SelectedTextPlaceholder(requestResponseHolder),
36+
new SelectedTextToFilePlaceholder(requestResponseHolder),
37+
new UrlPathPlaceholder(requestResponseHolder),
38+
new UrlPlaceholder(requestResponseHolder),
39+
new UrlQueryPlaceholder(requestResponseHolder));
3840
for (IPlaceholder placeholder : placeholderList) {
3941
placeholderMap.put(placeholder.getPlaceholder(), placeholder);
4042
}

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/models/placeholder/AbstractRequestInfoPlaceholder.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package net.bytebutcher.burpsendtoextension.models.placeholder;
22

3-
import burp.IContextMenuInvocation;
43
import burp.IRequestInfoWrapper;
54
import burp.RequestResponseHolder;
65

@@ -11,8 +10,8 @@ public abstract class AbstractRequestInfoPlaceholder extends AbstractRequestResp
1110

1211
private final RequestResponseHolder requestResponseHolder;
1312

14-
public AbstractRequestInfoPlaceholder(String placeholder, boolean doesRequireShellEscape, boolean doWriteToFile, RequestResponseHolder requestResponseHolder, IContextMenuInvocation contextMenuInvocation) {
15-
super(placeholder, doesRequireShellEscape, doWriteToFile, requestResponseHolder, contextMenuInvocation);
13+
public AbstractRequestInfoPlaceholder(String placeholder, boolean doesRequireShellEscape, boolean doWriteToFile, RequestResponseHolder requestResponseHolder) {
14+
super(placeholder, doesRequireShellEscape, doWriteToFile, requestResponseHolder);
1615
this.requestResponseHolder = requestResponseHolder;
1716
}
1817

‎burp-send-to-extension/src/main/java/net/bytebutcher/burpsendtoextension/models/placeholder/AbstractRequestPlaceholder.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,14 @@
11
package net.bytebutcher.burpsendtoextension.models.placeholder;
22

3-
import burp.IContextMenuInvocation;
43
import burp.IHttpRequestResponse;
54
import burp.RequestResponseHolder;
65

76
public abstract class AbstractRequestPlaceholder extends AbstractRequestResponsePlaceholderBase {
87

98
private final RequestResponseHolder requestResponseHolder;
109

11-
public AbstractRequestPlaceholder(String placeholder, boolean doesRequireShellEscape, boolean doWriteToFile, RequestResponseHolder requestResponseHolder, IContextMenuInvocation contextMenuInvocation) {
12-
super(placeholder, doesRequireShellEscape, doWriteToFile, requestResponseHolder, contextMenuInvocation);
10+
public AbstractRequestPlaceholder(String placeholder, boolean doesRequireShellEscape, boolean doWriteToFile, RequestResponseHolder requestResponseHolder) {
11+
super(placeholder, doesRequireShellEscape, doWriteToFile, requestResponseHolder);
1312
this.requestResponseHolder = requestResponseHolder;
1413
}
1514

0 commit comments

Comments
 (0)