Skip to content

Commit 39661fc

Browse files
bit4woobit4woo
committed
tuning the logic of "Update Cookie" & "Set Cookie"
优化更新cookie和set cookie的逻辑 单独一个类来管理cookie
1 parent 8a07435 commit 39661fc

File tree

10 files changed

+180
-226
lines changed

10 files changed

+180
-226
lines changed

‎src/burp/BurpExtender.java‎

Lines changed: 11 additions & 149 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ public List<JMenuItem> createMenuItems(IContextMenuInvocation invocation) {
101101
if (context == IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_REQUEST) {
102102

103103
menu_list.add(new UpdateCookieMenu(this));
104-
if (this.config.getTmpMap().containsKey("UsedCookie")){
104+
if (this.config.getUsedCookie()!=null){
105105
menu_list.add(new UpdateCookieWithHistoryMenu(this));
106106
}
107107

@@ -114,7 +114,7 @@ public List<JMenuItem> createMenuItems(IContextMenuInvocation invocation) {
114114
}
115115

116116
menu_list.add(new SetCookieMenu(this));
117-
if (this.config.getTmpMap().containsKey("cookieToSetHistory")){
117+
if (this.config.getUsedCookie() != null){
118118
menu_list.add(new SetCookieWithHistoryMenu(this));
119119
}
120120

@@ -195,17 +195,19 @@ public void processProxyMessage(boolean messageIsRequest, IInterceptedProxyMessa
195195
}
196196

197197
//当函数第一次被调用时，还没来得及设置cookie，获取到的cookieToSet必然为空。
198-
String cookieToSet = config.getTmpMap().get("cookieToSet");
198+
HashMap<String, HeaderEntry> cookieToSetMap = config.getSetCookieMap();
199199
//stderr.println("called"+cookieToSet);
200-
if (cookieToSet != null){//第二次调用如果cookie不为空，就走到这里
201-
String targetUrl = cookieToSet.split(CookieUtils.SPLITER)[0];
202-
String originUrl = cookieToSet.split(CookieUtils.SPLITER)[1];
203-
String cookieValue = cookieToSet.split(CookieUtils.SPLITER)[2];
200+
if (cookieToSetMap != null && !cookieToSetMap.isEmpty()){//第二次调用如果cookie不为空，就走到这里
204201

205202
IHttpRequestResponse messageInfo = message.getMessageInfo();
206203
String CurrentUrl = messageInfo.getHttpService().toString();
207204
//stderr.println(CurrentUrl+" "+targetUrl);
208-
if (targetUrl.equalsIgnoreCase(CurrentUrl)){
205+
HeaderEntry cookieToSet = cookieToSetMap.get(CurrentUrl);
206+
if (cookieToSet != null){
207+
208+
String targetUrl = cookieToSet.getTargetUrl();
209+
String cookieValue = cookieToSet.getHeaderValue();
210+
209211
if (messageIsRequest) {
210212
byte[] newRequest = CookieUtils.updateCookie(messageInfo,cookieValue);
211213
messageInfo.setRequest(newRequest);
@@ -219,9 +221,7 @@ public void processProxyMessage(boolean messageIsRequest, IInterceptedProxyMessa
219221
byte[] response = helpers.buildHttpMessage(responseHeaders,responseBody);
220222

221223
messageInfo.setResponse(response);
222-
config.getTmpMap().remove("cookieToSet");//only need to set once
223-
config.getTmpMap().put("cookieToSetHistory",cookieToSet);//store used cookie, change name to void change every request of host
224-
//临时换名称存储，避免这个参数影响这里的逻辑，导致域名下的每个请求都会进行该操作。
224+
cookieToSetMap.remove(CurrentUrl);//only need to set once
225225
}
226226
}
227227

@@ -363,144 +363,6 @@ public void processHttpMessage(int toolFlag, boolean messageIsRequest, IHttpRequ
363363
}
364364
}
365365

366-
@Deprecated
367-
public void processHttpMessageWithEditor(int toolFlag, boolean messageIsRequest, IHttpRequestResponse messageInfo) {
368-
//messageeditor
369-
synchronized (messageInfo) {
370-
if (messageIsRequest) {
371-
372-
boolean isRequestChanged = false;
373-
MessageEditor editer = new MessageEditor(messageIsRequest, messageInfo, helpers);
374-
375-
URL url = editer.getURL();
376-
String path = url.getPath();
377-
String host = editer.getHost();
378-
byte[] body = editer.getBody();
379-
LinkedHashMap<String, String> headers = editer.getHeaderMap();//this will lost the first line
380-
381-
382-
//remove header
383-
List<ConfigEntry> configEntries = tableModel.getConfigByType(ConfigEntry.Action_Remove_From_Headers);
384-
for (ConfigEntry entry : configEntries) {
385-
String key = entry.getKey();
386-
if (headers.remove(key) != null) {
387-
isRequestChanged = true;
388-
}
389-
}
390-
391-
if (config.getTmpMap().containsKey(host)) {//自动更新cookie
392-
String cookieValue = config.getTmpMap().get(host);
393-
String[] values = cookieValue.split("::::");
394-
String trueCookie = values[1];
395-
headers.put("Cookie", trueCookie);
396-
isRequestChanged = true;
397-
}
398-
399-
//add/update/append header
400-
if (toolFlag == (toolFlag & checkEnabledFor())) {
401-
//if ((config.isOnlyForScope() && callbacks.isInScope(url))|| !config.isOnlyForScope()) {
402-
if (!config.isOnlyForScope()||callbacks.isInScope(url)){
403-
try {
404-
List<ConfigEntry> updateOrAddEntries = tableModel.getConfigEntries();
405-
for (ConfigEntry entry : updateOrAddEntries) {
406-
String key = entry.getKey();
407-
String value = entry.getValue();
408-
409-
if (value.contains("%host")) {
410-
value = value.replaceAll("%host", host);
411-
//stdout.println("3333"+value);
412-
}
413-
414-
if (value.toLowerCase().contains("%dnslogserver")) {
415-
String dnslog = tableModel.getConfigByKey("DNSlogServer");
416-
Pattern p = Pattern.compile("(?u)%dnslogserver");
417-
Matcher m = p.matcher(value);
418-
419-
while (m.find()) {
420-
String found = m.group(0);
421-
value = value.replaceAll(found, dnslog);
422-
}
423-
}
424-
425-
if (entry.getType().equals(ConfigEntry.Action_Add_Or_Replace_Header) && entry.isEnable()) {
426-
headers.put(key, value);
427-
isRequestChanged = true;
428-
429-
} else if (entry.getType().equals(ConfigEntry.Action_Append_To_header_value) && entry.isEnable()) {
430-
value = headers.get(key) + value;
431-
headers.put(key, value);
432-
isRequestChanged = true;
433-
//stdout.println("2222"+value);
434-
} else if (entry.getKey().equalsIgnoreCase("Chunked-AutoEnable") && entry.isEnable()) {
435-
headers.put("Transfer-Encoding", "chunked");
436-
isRequestChanged = true;
437-
438-
try {
439-
boolean useComment = false;
440-
if (this.tableModel.getConfigByKey("Chunked-UseComment") != null) {
441-
useComment = true;
442-
}
443-
String lenStr = this.tableModel.getConfigByKey("Chunked-Length");
444-
int len = 10;
445-
if (lenStr != null) {
446-
len = Integer.parseInt(lenStr);
447-
}
448-
body = Methods.encoding(body, len, useComment);
449-
editer.setBody(body);
450-
} catch (UnsupportedEncodingException e) {
451-
e.printStackTrace(stderr);
452-
}
453-
}
454-
}
455-
456-
457-
///proxy function should be here
458-
//reference https://support.portswigger.net/customer/portal/questions/17350102-burp-upstream-proxy-settings-and-sethttpservice
459-
String proxy = this.tableModel.getConfigByKey("Proxy-ServerList");
460-
String mode = this.tableModel.getConfigByKey("Proxy-UseRandomMode");
461-
462-
if (proxy != null) {//if enable is false, will return null.
463-
List<String> proxyList = Arrays.asList(proxy.split(";"));//如果字符串是以;结尾，会被自动丢弃
464-
465-
if (mode != null) {//random mode
466-
proxyServerIndex = (int) (Math.random() * proxyList.size());
467-
//proxyServerIndex = new Random().nextInt(proxyList.size());
468-
} else {
469-
proxyServerIndex = (proxyServerIndex + 1) % proxyList.size();
470-
}
471-
String proxyhost = proxyList.get(proxyServerIndex).split(":")[0].trim();
472-
int port = Integer.parseInt(proxyList.get(proxyServerIndex).split(":")[1].trim());
473-
editer.setService(
474-
helpers.buildHttpService(proxyhost, port, messageInfo.getHttpService().getProtocol()));
475-
String firstrline = editer.getFirstLineOfHeader().replaceFirst(path, url.toString().split("\\?",0)[0]);
476-
editer.setFirstLineOfHeader(firstrline);
477-
isRequestChanged = true;
478-
//success or failed,need to check?
479-
}
480-
} catch (Exception e) {
481-
e.printStackTrace(stderr);
482-
}
483-
}
484-
}
485-
//set final request
486-
editer.setHeaderMap(headers);
487-
messageInfo = editer.getMessageInfo();
488-
489-
if (isRequestChanged) {
490-
//debug
491-
List<String> finalheaders = helpers.analyzeRequest(messageInfo).getHeaders();
492-
//List<String> finalheaders = editer.getHeaderList();//error here:bodyOffset getted twice are different
493-
stdout.println(System.lineSeparator() + "//////////edited request by knife//////////////" + System.lineSeparator());
494-
for (String entry : finalheaders) {
495-
stdout.println(entry);
496-
}
497-
}
498-
}
499-
}//sync
500-
}
501-
502-
503-
504366
public List<String> GetSetCookieHeaders(String cookies){
505367
if (cookies.startsWith("Cookie: ")){
506368
cookies = cookies.replaceFirst("Cookie: ","");

‎src/config/Config.java‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,16 @@
99
import com.alibaba.fastjson.annotation.JSONField;
1010

1111
import burp.IBurpExtenderCallbacks;
12+
import knife.HeaderEntry;
1213

1314
public class Config {
1415

1516
private String ConfigName = "";
1617
private List<String> stringConfigEntries = new ArrayList<String>();// get from configTableModel
1718
private int enableStatus = IBurpExtenderCallbacks.TOOL_PROXY;
1819
private boolean onlyForScope = true;
19-
private HashMap<String,String> tmpMap = new HashMap<String,String>();
20+
private HashMap<String,HeaderEntry> setCookieMap = new HashMap<String,HeaderEntry>();
21+
private HeaderEntry usedCookie = null;
2022

2123
Config(){
2224
//to resolve "default constructor not found" error
@@ -58,14 +60,23 @@ public void setOnlyForScope(boolean onlyForScope) {
5860
this.onlyForScope = onlyForScope;
5961
}
6062

61-
public HashMap<String, String> getTmpMap() {
62-
return tmpMap;
63+
@JSONField(serialize=false)//表明不序列号该字段
64+
public HashMap<String, HeaderEntry> getSetCookieMap() {
65+
return setCookieMap;
66+
}
67+
@JSONField(serialize=false)//表明不序列号该字段
68+
public void setSetCookieMap(HashMap<String, HeaderEntry> setCookieMap) {
69+
this.setCookieMap = setCookieMap;
70+
}
71+
@JSONField(serialize=false)//表明不序列号该字段
72+
public HeaderEntry getUsedCookie() {
73+
return usedCookie;
74+
}
75+
@JSONField(serialize=false)//表明不序列号该字段
76+
public void setUsedCookie(HeaderEntry usedCookie) {
77+
this.usedCookie = usedCookie;
6378
}
6479

65-
public void setTmpMap(HashMap<String, String> tmpMap) {
66-
this.tmpMap = tmpMap;
67-
}
68-
6980
@JSONField(serialize=false)//表明不序列号该字段
7081
public String ToJson(){//注意函数名称，如果是get set开头，会被认为是Getter和Setter函数，会在序列化过程中被调用。
7182
return JSONObject.toJSONString(this);

‎src/knife/Cookie.java‎

Lines changed: 0 additions & 11 deletions
This file was deleted.

‎src/knife/CookieUtils.java‎

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
package knife;
22

3-
import burp.*;
4-
53
import java.util.LinkedHashMap;
64
import java.util.List;
75

6+
import burp.BurpExtender;
7+
import burp.Getter;
8+
import burp.IHttpRequestResponse;
9+
import burp.Methods;
10+
811
public class CookieUtils {
912

1013
//////////////////////////////////////////common methods for cookie handle///////////////////////////////
@@ -25,7 +28,7 @@ public static IHttpRequestResponse[] Reverse(IHttpRequestResponse[] input){
2528
return a String url_which_cookie_from+SPLITER+cookievalue
2629
*/
2730

28-
public static String getLatestHeaderFromHistory(String shortUrl,String headerName){
31+
public static HeaderEntry getLatestHeaderFromHistory(String shortUrl,String headerName){
2932
//还是草粉师傅说得对，直接从history里面拿最好
3033

3134
IHttpRequestResponse[] historyMessages = Reverse(BurpExtender.callbacks.getProxyHistory());
@@ -53,36 +56,40 @@ public static String getLatestHeaderFromHistory(String shortUrl,String headerNam
5356
if (hisShortUrl.equalsIgnoreCase(shortUrl)) {
5457
String cookieValue = getter.getHeaderValueOf(true,historyMessage,headerName);
5558
if (cookieValue != null){
56-
return shortUrl+SPLITER+cookieValue;
59+
HeaderEntry entry = new HeaderEntry(shortUrl,headerName,cookieValue, null);
60+
return entry;
61+
//return shortUrl+SPLITER+cookieValue;
5762
}
5863
}
5964
}
6065

6166
return null;
6267
}
6368

64-
public static String getLatestCookieFromHistory(String shortUrl){
69+
public static HeaderEntry getLatestCookieFromHistory(String shortUrl){
6570
return getLatestHeaderFromHistory(shortUrl,"Cookie");
6671
}
6772

6873
/*
6974
return a String url_which_cookie_from+SPLITER+cookievalue
7075
*/
71-
public static String getLatestCookieFromSpeicified() {
72-
String latestCookie = null;
76+
public static HeaderEntry getLatestCookieFromSpeicified() {
77+
HeaderEntry latestCookie = null;
7378
String domainOrCookie = Methods.prompt_and_validate_input("cookie OR cookie of ", null);
7479
String url1 = "";
7580
String url2 = "";
7681
try{
7782
if (domainOrCookie == null){
7883
return null;
7984
}else if (domainOrCookie.contains("=") && !domainOrCookie.contains("?") && !domainOrCookie.contains("/")){//直接是cookie
80-
latestCookie = domainOrCookie.trim();
81-
if (latestCookie.startsWith("Cookie:")){
82-
latestCookie = latestCookie.replaceFirst("Cookie:","").trim();
85+
String cookieValue = domainOrCookie.trim();
86+
87+
if (cookieValue.startsWith("Cookie:")){
88+
cookieValue = cookieValue.replaceFirst("Cookie:","").trim();
8389
}
84-
String tips = "Cookie: "+latestCookie.substring(0,latestCookie.indexOf("="))+"...";
85-
latestCookie = tips+SPLITER+latestCookie;
90+
String tips = "Cookie: "+cookieValue.substring(0,cookieValue.indexOf("="))+"...";
91+
latestCookie = new HeaderEntry(tips,"Cookie",cookieValue, null);
92+
8693
return latestCookie;
8794
}else if (domainOrCookie.startsWith("http://") || domainOrCookie.startsWith("https://")) {//不包含协议头的域名或url
8895
url1 = domainOrCookie;

0 commit comments

Comments
 (0)