From de06a157ea862f87550c755358f0463eb12ec1cd Mon Sep 17 00:00:00 2001 From: kira0204 Date: Thu, 1 Feb 2018 05:01:19 +0530 Subject: [PATCH] Improving the logic and test --- mitmproxy/addons/save.py | 4 +- mitmproxy/addons/share.py | 76 ++++++++++++++++------------- test/mitmproxy/addons/test_share.py | 15 ++++-- 3 files changed, 53 insertions(+), 42 deletions(-) diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py index 44afef686..47da29b25 100644 --- a/mitmproxy/addons/save.py +++ b/mitmproxy/addons/save.py @@ -61,8 +61,8 @@ class Save: except IOError as v: raise exceptions.CommandError(v) from v stream = io.FlowWriter(f) - for i in flows: - stream.add(i) + for x in flows: + stream.add(x) f.close() ctx.log.alert("Saved %s flows." % len(flows)) diff --git a/mitmproxy/addons/share.py b/mitmproxy/addons/share.py index 17eefadbd..1a234cc94 100644 --- a/mitmproxy/addons/share.py +++ b/mitmproxy/addons/share.py @@ -1,18 +1,18 @@ import typing import random -import datetime -import time -import _io +import string +import io import http.client from mitmproxy import command -from mitmproxy import io +import mitmproxy.io from mitmproxy import ctx from mitmproxy import flow +from mitmproxy.net.http import status_codes class Share: - def encode_multipart_formdata(self, filename, content): + def encode_multipart_formdata(self, filename: str, content: bytes) -> typing.Tuple[str, bytes]: params = {"key": filename, "acl": "bucket-owner-full-control", "Content-Type": "application/octet-stream"} LIMIT = b'---------------------------198495659117975628761412556003' CRLF = b'\r\n' @@ -30,44 +30,50 @@ class Share: l.append(b'--' + LIMIT + b'--') l.append(b'') body = CRLF.join(l) - content_type = b'multipart/form-data; boundary=%b' % LIMIT + content_type = 'multipart/form-data; boundary=%s' % LIMIT.decode("utf-8") return content_type, body - def post_multipart(self, host, filename, content): + def post_multipart(self, host: str, filename: str, content: bytes) -> str: + """ + Upload flows to the specified S3 server. + + Returns: + - The share URL, if upload is successful. + Raises: + - IOError, otherwise. + """ content_type, body = self.encode_multipart_formdata(filename, content) - conn = http.client.HTTPConnection(host, 80) - headers = {'content-type': content_type, 'content-length': str(len(body))} + conn = http.client.HTTPConnection(host) # FIXME: This ultimately needs to be HTTPSConnection + headers = {'content-type': content_type} try: conn.request("POST", "", body, headers) - except http.client.CannotSendRequest: - return 'We failed to reach a server.' - try: - conn.getresponse() - except http.client.RemoteDisconnected: - return 'The server couldn\'t fulfill the request.' - else: + resp = conn.getresponse() + except Exception as v: + raise IOError(v) + finally: conn.close() - return 'URL: share.mitmproxy.org/%s' % filename - - def base36encode(self, integer): - chars, encoded = "0123456789abcdefghijklmnopqrstuvwxyz", "" - - while integer > 0: - integer, remainder = divmod(integer, 36) - encoded = chars[remainder] + encoded - - return encoded + if resp.status != 204: + if resp.reason: + reason = resp.reason + else: + reason = status_codes.RESPONSES.get(resp.status, str(resp.status)) + raise IOError(reason) + return "https://share.mitmproxy.org/%s" % filename @command.command("share.flows") def share(self, flows: typing.Sequence[flow.Flow]) -> None: - d = datetime.datetime.utcnow() - u_id = self.base36encode(int(time.mktime(d.timetuple()) * 1000 * random.random()))[0:7] - f = _io.BytesIO() - stream = io.FlowWriter(f) - for i in flows: - stream.add(i) + u_id = "".join(random.choice(string.ascii_lowercase + string.digits)for _ in range(7)) + f = io.BytesIO() + stream = mitmproxy.io.FlowWriter(f) + for x in flows: + stream.add(x) f.seek(0) content = f.read() - res = self.post_multipart('upload.share.mitmproxy.org.s3.amazonaws.com', u_id, content) - f.close() - ctx.log.alert("%s" % res) + try: + res = self.post_multipart('upload.share.mitmproxy.org.s3.amazonaws.com', u_id, content) + except IOError as v: + ctx.log.warn("%s" % v) + else: + ctx.log.alert("%s" % res) + finally: + f.close() \ No newline at end of file diff --git a/test/mitmproxy/addons/test_share.py b/test/mitmproxy/addons/test_share.py index a85c20b7b..6c3d6e28f 100644 --- a/test/mitmproxy/addons/test_share.py +++ b/test/mitmproxy/addons/test_share.py @@ -12,16 +12,21 @@ def test_share_command(): with mock.patch('mitmproxy.addons.share.http.client.HTTPConnection') as mock_http: sh = share.Share() with taddons.context() as tctx: + mock_http.return_value.getresponse.return_value = mock.MagicMock(status=204, reason="No Content") sh.share([tflow.tflow(resp=True)]) - assert tctx.master.has_log("URL: share.mitmproxy.org/") + assert tctx.master.has_log("https://share.mitmproxy.org/") - mock_http.return_value.getresponse.side_effect = http.client.RemoteDisconnected + mock_http.return_value.getresponse.return_value = mock.MagicMock(status=403, reason="Forbidden") sh.share([tflow.tflow(resp=True)]) - assert tctx.master.has_log("The server couldn\'t fulfill the request.") + assert tctx.master.has_log("Forbidden") - mock_http.return_value.request.side_effect = http.client.CannotSendRequest + mock_http.return_value.getresponse.return_value = mock.MagicMock(status=404, reason="") sh.share([tflow.tflow(resp=True)]) - assert tctx.master.has_log("We failed to reach a server.") + assert tctx.master.has_log("Not Found") + + mock_http.return_value.request.side_effect = http.client.CannotSendRequest("Error in sending req") + sh.share([tflow.tflow(resp=True)]) + assert tctx.master.has_log("Error in sending req") v = view.View() tctx.master.addons.add(v)