From 2aacf94a63c7e0c396d6254356ba9739c7319077 Mon Sep 17 00:00:00 2001
From: Chase Kelley <32003197+wchasekelley@users.noreply.github.com>
Date: Fri, 28 Aug 2020 06:41:22 -0400
Subject: [PATCH] update raw export to not remove headers (#4181)

---
 CHANGELOG                            | 1 +
 mitmproxy/addons/export.py           | 8 +++++++-
 test/mitmproxy/addons/test_export.py | 2 ++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG b/CHANGELOG
index 75fef3b52..78dd1788c 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -11,6 +11,7 @@ Unreleased: mitmproxy next
     * Display HTTP trailers in mitmweb (@sanlengjingvv)
     * Revamp onboarding app (@mhils)
     * Add ASGI support for embedded apps (@mhils)
+    * Updated raw exports to not remove headers (@wchasekelley)
 
     * --- TODO: add new PRs above this line ---
 
diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py
index d874c95a7..81409e440 100644
--- a/mitmproxy/addons/export.py
+++ b/mitmproxy/addons/export.py
@@ -18,7 +18,11 @@ def cleanup_request(f: flow.Flow) -> http.HTTPRequest:
     assert isinstance(f, http.HTTPFlow)
     request = f.request.copy()
     request.decode(strict=False)
-    # a bit of clean-up - these headers should be automatically set by curl/httpie
+    return request
+
+
+def pop_headers(request: http.HTTPRequest) -> http.HTTPRequest:
+    # Remove some headers that are redundant for curl/httpie export
     request.headers.pop('content-length')
     if request.headers.get("host", "") == request.host:
         request.headers.pop("host")
@@ -53,6 +57,7 @@ def request_content_for_console(request: http.HTTPRequest) -> str:
 
 def curl_command(f: flow.Flow) -> str:
     request = cleanup_request(f)
+    request = pop_headers(request)
     args = ["curl"]
     for k, v in request.headers.items(multi=True):
         if k.lower() == "accept-encoding":
@@ -70,6 +75,7 @@ def curl_command(f: flow.Flow) -> str:
 
 def httpie_command(f: flow.Flow) -> str:
     request = cleanup_request(f)
+    request = pop_headers(request)
     args = ["http", request.method, request.url]
     for k, v in request.headers.items(multi=True):
         args.append(f"{k}: {v}")
diff --git a/test/mitmproxy/addons/test_export.py b/test/mitmproxy/addons/test_export.py
index b0e5e47e9..333addda7 100644
--- a/test/mitmproxy/addons/test_export.py
+++ b/test/mitmproxy/addons/test_export.py
@@ -144,6 +144,7 @@ class TestRaw:
 
     def test_get_request_present(self, get_request):
         assert b"header: qvalue" in export.raw(get_request)
+        assert b"content-length: 0" in export.raw_request(get_request)
 
     def test_get_response_present(self, get_response):
         delattr(get_response, 'request')
@@ -163,6 +164,7 @@ class TestRaw:
 class TestRawRequest:
     def test_get(self, get_request):
         assert b"header: qvalue" in export.raw_request(get_request)
+        assert b"content-length: 0" in export.raw_request(get_request)
 
     def test_no_request(self, get_response):
         delattr(get_response, 'request')