diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 8f84c20d9..619211130 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -20,6 +20,7 @@ from mitmproxy.addons import stickycookie from mitmproxy.addons import streambodies from mitmproxy.addons import save from mitmproxy.addons import upstream_auth +from mitmproxy.addons import share def default_addons(): @@ -46,4 +47,5 @@ def default_addons(): streambodies.StreamBodies(), save.Save(), upstream_auth.UpstreamAuth(), + share.Share() ] 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 new file mode 100644 index 000000000..1a234cc94 --- /dev/null +++ b/mitmproxy/addons/share.py @@ -0,0 +1,79 @@ +import typing +import random +import string +import io +import http.client + +from mitmproxy import command +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: 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' + l = [] + for (key, value) in params.items(): + l.append(b'--' + LIMIT) + l.append(b'Content-Disposition: form-data; name="%b"' % key.encode("utf-8")) + l.append(b'') + l.append(value.encode("utf-8")) + l.append(b'--' + LIMIT) + l.append(b'Content-Disposition: form-data; name="file"; filename="%b"' % filename.encode("utf-8")) + l.append(b'Content-Type: application/octet-stream') + l.append(b'') + l.append(content) + l.append(b'--' + LIMIT + b'--') + l.append(b'') + body = CRLF.join(l) + content_type = 'multipart/form-data; boundary=%s' % LIMIT.decode("utf-8") + return content_type, body + + 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) # FIXME: This ultimately needs to be HTTPSConnection + headers = {'content-type': content_type} + try: + conn.request("POST", "", body, headers) + resp = conn.getresponse() + except Exception as v: + raise IOError(v) + finally: + conn.close() + 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: + 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() + 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 new file mode 100644 index 000000000..6c3d6e28f --- /dev/null +++ b/test/mitmproxy/addons/test_share.py @@ -0,0 +1,34 @@ +from unittest import mock +import http.client + +from mitmproxy.test import taddons +from mitmproxy.test import tflow + +from mitmproxy.addons import share +from mitmproxy.addons import view + + +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("https://share.mitmproxy.org/") + + mock_http.return_value.getresponse.return_value = mock.MagicMock(status=403, reason="Forbidden") + sh.share([tflow.tflow(resp=True)]) + assert tctx.master.has_log("Forbidden") + + mock_http.return_value.getresponse.return_value = mock.MagicMock(status=404, reason="") + sh.share([tflow.tflow(resp=True)]) + 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) + tctx.master.addons.add(sh) + tctx.master.commands.call_args("share.flows", ["@shown"])