From b6969df52ac13cace7112ff53156171f7563acfd Mon Sep 17 00:00:00 2001 From: Louis-Philippe Huberdeau Date: Thu, 29 Jun 2017 10:14:20 -0400 Subject: [PATCH] Add missing httpVersion in request render, avoid encoding to base64 unless binary data is included --- lib/utils/collect.py | 33 ++++++++++++++++++++++++++------- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/lib/utils/collect.py b/lib/utils/collect.py index b88ef4111..fc6915061 100644 --- a/lib/utils/collect.py +++ b/lib/utils/collect.py @@ -108,6 +108,7 @@ class Request: def toDict(self): out = { + "httpVersion": self.httpVersion, "method": self.method, "url": self.url, "headers": [dict(name=key, value=value) for key, value in self.headers.items()], @@ -160,16 +161,22 @@ class Response: raw=raw) def toDict(self): + content = { + "mimeType": self.headers.get('Content-Type'), + "text": self.content, + } + + binary = set(['\0', '\1']) + if any(c in binary for c in self.content): + content["encoding"] = "base64" + content["text"] = base64.b64encode(self.content) + return { "httpVersion": self.httpVersion, "status": self.status, "statusText": self.statusText, "headers": [dict(name=key, value=value) for key, value in self.headers.items()], - "content": { - "mimeType": self.headers.get('Content-Type'), - "encoding": "base64", - "text": base64.b64encode(self.content), - }, + "content": content, "comment": self.comment, } @@ -241,6 +248,7 @@ if __name__ == '__main__': self.assertEqual("http://example.com/test.php", out["url"]) self.assertIn({"name": "Host", "value": "example.com"}, out["headers"]) self.assertEqual("Hello World", out["comment"]) + self.assertEqual("HTTP/1.1", out["httpVersion"]) def test_render_with_post_body(self): req = Request(method="POST", @@ -276,15 +284,26 @@ if __name__ == '__main__': resp = Response(status=200, statusText="OK", httpVersion="HTTP/1.1", headers={"Content-Type": "text/html"}, - content="Hello\n") + content="\nHello\n") out = resp.toDict() self.assertEqual(200, out["status"]) self.assertEqual("OK", out["statusText"]) self.assertIn({"name": "Content-Type", "value": "text/html"}, out["headers"]) self.assertEqual(out["content"], { "mimeType": "text/html", + "text": "\nHello\n", + }) + + def test_simple_body_contains_binary_data(self): + resp = Response(status=200, statusText="OK", + httpVersion="HTTP/1.1", + headers={"Content-Type": "application/octet-stream"}, + content="test\0abc") + out = resp.toDict() + self.assertEqual(out["content"], { "encoding": "base64", - "text": "PGh0bWw+PGJvZHk+SGVsbG88L2JvZHk+PC9odG1sPgo=", + "mimeType": "application/octet-stream", + "text": "dGVzdABhYmM=", }) unittest.main(buffer=False)