Refactor Set HTTP Cookie Operation

usdAG · Nov 15, 2024 · 5266e74 · 5266e74
1 parent 4887756
commit 5266e74
Showing 1 changed file with 55 additions and 27 deletions.
diff --git a/src/main/java/de/usd/cstchef/operations/setter/HttpSetCookie.java b/src/main/java/de/usd/cstchef/operations/setter/HttpSetCookie.java
@@ -4,16 +4,10 @@
 
 import javax.swing.JCheckBox;
 
-import burp.BurpUtils;
-import burp.api.montoya.MontoyaApi;
 import burp.api.montoya.core.ByteArray;
-import burp.api.montoya.http.message.Cookie;
 import burp.api.montoya.http.message.HttpHeader;
-import burp.api.montoya.http.message.params.HttpParameter;
-import burp.api.montoya.http.message.params.HttpParameterType;
 import burp.api.montoya.http.message.requests.HttpRequest;
 import burp.api.montoya.http.message.responses.HttpResponse;
-import de.usd.cstchef.Utils;
 import de.usd.cstchef.Utils.MessageType;
 import de.usd.cstchef.operations.Operation.OperationInfos;
 import de.usd.cstchef.operations.OperationCategory;
@@ -31,34 +25,68 @@ protected ByteArray perform(ByteArray input, MessageType messageType) throws Exc
         if (getWhat().equals(""))
             return input;
 
-        if (messageType == MessageType.REQUEST) {
+
+        if(messageType == MessageType.REQUEST) {
             HttpRequest request = HttpRequest.httpRequest(input);
-            if (!Utils.httpRequestCookieExtractor(request, cookieName).equals(ByteArray.byteArray(0))
-                    || addIfNotPresent.isSelected()) {
-                return Utils.addCookieToHttpRequest(request, new Utils.CSTCCookie(cookieName, cookieValue))
-                        .toByteArray();
-            } else {
-                return input;
+
+            // has Cookie header
+            if(request.hasHeader("Cookie")) {
+                String cookies = request.header("Cookie").value();
+                // has this particular cookie set
+                if(cookies.contains(cookieName + "=")) {
+                    String[] c = cookies.split("; ");
+                    cookies = "";
+                    for(String cookie : c) {
+                        cookie = cookie.replaceAll(cookieName + "=\\S*", cookieName + "=" + cookieValue);
+                        cookies = cookies.concat(cookie + "; ");
+                    }
+                    cookies = cookies.replaceAll(";\s$", "");
+                    return request.withUpdatedHeader("Cookie", cookies).toByteArray();
+                }
+                // has this particular cookie not set
+                else {
+                    cookies = cookies.concat("; " + cookieName + "=" + cookieValue);
+                        return addIfNotPresent.isSelected() ? request.withUpdatedHeader("Cookie", cookies).toByteArray() : input;
+                }
+            }
+            // has no Cookie header
+            else {
+                return addIfNotPresent.isSelected() ? request.withAddedHeader("Cookie", cookieName + "=" + cookieValue).toByteArray() : input;
             }
-        } else if (messageType == MessageType.RESPONSE) {
+        }
+
+        else if (messageType == MessageType.RESPONSE) {
             HttpResponse response = HttpResponse.httpResponse(input);
-            List<HttpHeader> headers = response.headers();
-            for (HttpHeader h : headers) {
-                if (h.name().equals("Set-Cookie")) {
-                    if (h.value().contains(cookieName)) {
-                        return response.withRemovedHeader(h)
-                                .withAddedHeader(HttpHeader.httpHeader("Set-Cookie", cookieName + "=" + cookieValue))
-                                .toByteArray();
+            List<HttpHeader> httpHeader = response.headers();
+
+            // has Set-Cookie header
+            if(response.hasCookie(cookieName)) {
+                response = response.withRemovedHeaders(httpHeader);
+                for(int i = 0; i < httpHeader.size(); i++) {
+                    if(httpHeader.get(i).name().equals("Set-Cookie")) {
+                        // has this particular cookie set
+                        if(httpHeader.get(i).value().contains(cookieName + "=")) {
+                            response = response.withAddedHeader("Set-Cookie", cookieName + "=" + cookieValue);
+                        }
+                        else{
+                            response = response.withAddedHeader(httpHeader.get(i));
+                        }
+                    }
+                    else {
+                        response = response.withAddedHeader(httpHeader.get(i));
                     }
                 }
+
+                return response.toByteArray();
+
             }
-            if (addIfNotPresent.isSelected()) {
-                return response.withAddedHeader(HttpHeader.httpHeader("Set-Cookie", cookieName + "=" + cookieValue))
-                        .toByteArray();
-            } else {
-                return input;
+            // has no Set-Cookie header
+            else {
+                return addIfNotPresent.isSelected() ? response.withAddedHeader("Set-Cookie", cookieName + "=" + cookieValue).toByteArray(): input;
             }
-        } else {
+        }
+
+        else {
             return parseRawMessage(input);
         }
     }