From c9bbce088e1f8afcc111fb0d311c3e59cde434c3 Mon Sep 17 00:00:00 2001 From: kira0204 Date: Fri, 19 Jan 2018 02:13:00 +0530 Subject: [PATCH 1/7] Adding 'share' feature --- mitmproxy/addons/__init__.py | 2 + mitmproxy/addons/multipart.py | 43 ++++++++++++ mitmproxy/addons/upload.py | 124 ++++++++++++++++++++++++++++++++++ 3 files changed, 169 insertions(+) create mode 100644 mitmproxy/addons/multipart.py create mode 100644 mitmproxy/addons/upload.py diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 8f84c20d9..803409123 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 upload def default_addons(): @@ -46,4 +47,5 @@ def default_addons(): streambodies.StreamBodies(), save.Save(), upstream_auth.UpstreamAuth(), + upload.Upload() ] diff --git a/mitmproxy/addons/multipart.py b/mitmproxy/addons/multipart.py new file mode 100644 index 000000000..c29626edd --- /dev/null +++ b/mitmproxy/addons/multipart.py @@ -0,0 +1,43 @@ +from urllib.request import Request, urlopen +from urllib.error import HTTPError, URLError +import mimetypes + + +def get_content_type(filename): + return mimetypes.guess_type(filename)[0] or 'application/octet-stream' + + +def encode_multipart_formdata(filename, content): + 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: %b' % get_content_type(filename).encode("utf-8")) + L.append(b'') + L.append(content) + L.append(b'--' + LIMIT + b'--') + L.append(b'') + body = CRLF.join(L) + content_type = b'multipart/form-data; boundary=%b' % LIMIT + return content_type, body + +def post_multipart(host, filename, content): + content_type, body = encode_multipart_formdata(filename, content) + req = Request("http://upload.share.mitmproxy.org.s3.amazonaws.com", data=body, method='POST') + req.add_header('content-type', content_type) + req.add_header('content-length', str(len(body))) + try: + response = urlopen(req) + except HTTPError as e: + return ('The server couldn\'t fulfill the request.') + except URLError as e: + return ('We failed to reach a server.') + else: + return ('URL: share.mitmproxy.org/%s' % filename) diff --git a/mitmproxy/addons/upload.py b/mitmproxy/addons/upload.py new file mode 100644 index 000000000..660f1d5ac --- /dev/null +++ b/mitmproxy/addons/upload.py @@ -0,0 +1,124 @@ +import os.path +import typing +import random +import datetime +import time +from urllib.request import Request, urlopen +from urllib.error import HTTPError, URLError +import mimetypes + +from mitmproxy import command +from mitmproxy import exceptions +from mitmproxy import flowfilter +from mitmproxy import io +from mitmproxy import ctx +from mitmproxy import flow + +import mitmproxy.types + +from mitmproxy.addons import multipart + + +class Upload: + def __init__(self): + self.stream = None + self.filt = None + self.active_flows = set() # type: Set[flow.Flow] + + def open_file(self, path): + mode = "wb" + path = os.path.expanduser(path) + return open(path, mode) + + def read_file(self, path): + path = os.path.expanduser(path) + with open(path, "rb") as f: + content = f.read() + return content + + def start_stream_to_path(self, path, flt): + try: + f = self.open_file(path) + except IOError as v: + raise exceptions.OptionsError(str(v)) + self.stream = io.FilteredFlowWriter(f, flt) + self.active_flows = set() + + def configure(self, updated): + # We're already streaming - stop the previous stream and restart + if "save_stream_filter" in updated: + if ctx.options.save_stream_filter: + self.filt = flowfilter.parse(ctx.options.save_stream_filter) + if not self.filt: + raise exceptions.OptionsError( + "Invalid filter specification: %s" % ctx.options.save_stream_filter + ) + else: + self.filt = None + if "save_stream_file" in updated or "save_stream_filter" in updated: + if self.stream: + self.done() + if ctx.options.save_stream_file: + self.start_stream_to_path(ctx.options.save_stream_file, self.filt) + + def base36encode(self, integer): + chars, encoded = "0123456789abcdefghijklmnopqrstuvwxyz", "" + + while integer > 0: + integer, remainder = divmod(integer, 36) + encoded = chars[remainder] + encoded + + return encoded + + @command.command("upload.file") + def upload(self, flows: typing.Sequence[flow.Flow]) -> None: + + d = datetime.datetime.utcnow() + id = self.base36encode(int(time.mktime(d.timetuple()) * 1000 * random.random()))[0:7] + try: + f = self.open_file(id) # Making temporary file to store the flows + except IOError as v: + raise exceptions.CommandError(v) from v + stream = io.FlowWriter(f) + for i in flows: + stream.add(i) + f.close() + content = self.read_file(id) + res = multipart.post_multipart('http://upload.share.mitmproxy.org.s3.amazonaws.com', id, content) + ctx.log.alert("%s" % res) + os.remove(os.path.expanduser(id)) # Deleting the temporary file + + def tcp_start(self, flow): + if self.stream: + self.active_flows.add(flow) + + def tcp_end(self, flow): + if self.stream: + self.stream.add(flow) + self.active_flows.discard(flow) + + def websocket_start(self, flow): + if self.stream: + self.active_flows.add(flow) + + def websocket_end(self, flow): + if self.stream: + self.stream.add(flow) + self.active_flows.discard(flow) + + def response(self, flow): + if self.stream: + self.stream.add(flow) + self.active_flows.discard(flow) + + def request(self, flow): + if self.stream: + self.active_flows.add(flow) + + def done(self): + if self.stream: + for f in self.active_flows: + self.stream.add(f) + self.active_flows = set([]) + self.stream.fo.close() + self.stream = None From 8a48cea50252ac4cde65a853cf50c92fa541f52f Mon Sep 17 00:00:00 2001 From: kira0204 Date: Fri, 19 Jan 2018 13:31:48 +0530 Subject: [PATCH 2/7] upload<->share,multipart.py optimised --- mitmproxy/addons/__init__.py | 4 ++-- mitmproxy/addons/multipart.py | 10 +++------- mitmproxy/addons/{upload.py => share.py} | 9 +++------ 3 files changed, 8 insertions(+), 15 deletions(-) rename mitmproxy/addons/{upload.py => share.py} (94%) diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 803409123..619211130 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -20,7 +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 upload +from mitmproxy.addons import share def default_addons(): @@ -47,5 +47,5 @@ def default_addons(): streambodies.StreamBodies(), save.Save(), upstream_auth.UpstreamAuth(), - upload.Upload() + share.Share() ] diff --git a/mitmproxy/addons/multipart.py b/mitmproxy/addons/multipart.py index c29626edd..2d3001d14 100644 --- a/mitmproxy/addons/multipart.py +++ b/mitmproxy/addons/multipart.py @@ -1,10 +1,5 @@ from urllib.request import Request, urlopen from urllib.error import HTTPError, URLError -import mimetypes - - -def get_content_type(filename): - return mimetypes.guess_type(filename)[0] or 'application/octet-stream' def encode_multipart_formdata(filename, content): @@ -19,7 +14,7 @@ def encode_multipart_formdata(filename, content): 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: %b' % get_content_type(filename).encode("utf-8")) + L.append(b'Content-Type: application/octet-stream') L.append(b'') L.append(content) L.append(b'--' + LIMIT + b'--') @@ -28,9 +23,10 @@ def encode_multipart_formdata(filename, content): content_type = b'multipart/form-data; boundary=%b' % LIMIT return content_type, body + def post_multipart(host, filename, content): content_type, body = encode_multipart_formdata(filename, content) - req = Request("http://upload.share.mitmproxy.org.s3.amazonaws.com", data=body, method='POST') + req = Request(host, data=body, method='POST') req.add_header('content-type', content_type) req.add_header('content-length', str(len(body))) try: diff --git a/mitmproxy/addons/upload.py b/mitmproxy/addons/share.py similarity index 94% rename from mitmproxy/addons/upload.py rename to mitmproxy/addons/share.py index 660f1d5ac..e4d85c47d 100644 --- a/mitmproxy/addons/upload.py +++ b/mitmproxy/addons/share.py @@ -3,9 +3,6 @@ import typing import random import datetime import time -from urllib.request import Request, urlopen -from urllib.error import HTTPError, URLError -import mimetypes from mitmproxy import command from mitmproxy import exceptions @@ -19,7 +16,7 @@ import mitmproxy.types from mitmproxy.addons import multipart -class Upload: +class Share: def __init__(self): self.stream = None self.filt = None @@ -70,8 +67,8 @@ class Upload: return encoded - @command.command("upload.file") - def upload(self, flows: typing.Sequence[flow.Flow]) -> None: + @command.command("share.file") + def share(self, flows: typing.Sequence[flow.Flow]) -> None: d = datetime.datetime.utcnow() id = self.base36encode(int(time.mktime(d.timetuple()) * 1000 * random.random()))[0:7] From 12399e9c1253c776e21c46877a554c98e7049924 Mon Sep 17 00:00:00 2001 From: kira0204 Date: Fri, 19 Jan 2018 21:05:50 +0530 Subject: [PATCH 3/7] changing to http.client and BytesIO object --- mitmproxy/addons/multipart.py | 49 ++++++++++++++++++----------------- mitmproxy/addons/share.py | 21 +++++---------- 2 files changed, 31 insertions(+), 39 deletions(-) diff --git a/mitmproxy/addons/multipart.py b/mitmproxy/addons/multipart.py index 2d3001d14..827ca4ef0 100644 --- a/mitmproxy/addons/multipart.py +++ b/mitmproxy/addons/multipart.py @@ -1,39 +1,40 @@ -from urllib.request import Request, urlopen -from urllib.error import HTTPError, URLError +import http.client def encode_multipart_formdata(filename, content): params = {"key": filename, "acl": "bucket-owner-full-control", "Content-Type": "application/octet-stream"} LIMIT = b'---------------------------198495659117975628761412556003' CRLF = b'\r\n' - L = [] + 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) + 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 = b'multipart/form-data; boundary=%b' % LIMIT return content_type, body def post_multipart(host, filename, content): content_type, body = encode_multipart_formdata(filename, content) - req = Request(host, data=body, method='POST') - req.add_header('content-type', content_type) - req.add_header('content-length', str(len(body))) + conn = http.client.HTTPConnection(host, 80) + headers = {'content-type': content_type, 'content-length': str(len(body))} try: - response = urlopen(req) - except HTTPError as e: - return ('The server couldn\'t fulfill the request.') - except URLError as e: - return ('We failed to reach a server.') + 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: - return ('URL: share.mitmproxy.org/%s' % filename) + conn.close() + return 'URL: share.mitmproxy.org/%s' % filename diff --git a/mitmproxy/addons/share.py b/mitmproxy/addons/share.py index e4d85c47d..f109adc59 100644 --- a/mitmproxy/addons/share.py +++ b/mitmproxy/addons/share.py @@ -3,6 +3,7 @@ import typing import random import datetime import time +import _io from mitmproxy import command from mitmproxy import exceptions @@ -27,12 +28,6 @@ class Share: path = os.path.expanduser(path) return open(path, mode) - def read_file(self, path): - path = os.path.expanduser(path) - with open(path, "rb") as f: - content = f.read() - return content - def start_stream_to_path(self, path, flt): try: f = self.open_file(path) @@ -69,21 +64,17 @@ class Share: @command.command("share.file") def share(self, flows: typing.Sequence[flow.Flow]) -> None: - d = datetime.datetime.utcnow() - id = self.base36encode(int(time.mktime(d.timetuple()) * 1000 * random.random()))[0:7] - try: - f = self.open_file(id) # Making temporary file to store the flows - except IOError as v: - raise exceptions.CommandError(v) from v + 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) + f.seek(0) + content = f.read() + res = multipart.post_multipart('upload.share.mitmproxy.org.s3.amazonaws.com', u_id, content) f.close() - content = self.read_file(id) - res = multipart.post_multipart('http://upload.share.mitmproxy.org.s3.amazonaws.com', id, content) ctx.log.alert("%s" % res) - os.remove(os.path.expanduser(id)) # Deleting the temporary file def tcp_start(self, flow): if self.stream: From 3b51f6127232fc8269f6821b7450af0954eff478 Mon Sep 17 00:00:00 2001 From: kira0204 Date: Wed, 24 Jan 2018 22:30:44 +0530 Subject: [PATCH 4/7] removed unnecessary methods --- mitmproxy/addons/multipart.py | 40 ---------- mitmproxy/addons/share.py | 113 +++++++++------------------- test/mitmproxy/addons/test_share.py | 1 + 3 files changed, 38 insertions(+), 116 deletions(-) delete mode 100644 mitmproxy/addons/multipart.py create mode 100644 test/mitmproxy/addons/test_share.py diff --git a/mitmproxy/addons/multipart.py b/mitmproxy/addons/multipart.py deleted file mode 100644 index 827ca4ef0..000000000 --- a/mitmproxy/addons/multipart.py +++ /dev/null @@ -1,40 +0,0 @@ -import http.client - - -def encode_multipart_formdata(filename, content): - 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 = b'multipart/form-data; boundary=%b' % LIMIT - return content_type, body - - -def post_multipart(host, filename, content): - content_type, body = encode_multipart_formdata(filename, content) - conn = http.client.HTTPConnection(host, 80) - headers = {'content-type': content_type, 'content-length': str(len(body))} - 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: - conn.close() - return 'URL: share.mitmproxy.org/%s' % filename diff --git a/mitmproxy/addons/share.py b/mitmproxy/addons/share.py index f109adc59..17eefadbd 100644 --- a/mitmproxy/addons/share.py +++ b/mitmproxy/addons/share.py @@ -1,57 +1,53 @@ -import os.path import typing import random import datetime import time import _io +import http.client from mitmproxy import command -from mitmproxy import exceptions -from mitmproxy import flowfilter from mitmproxy import io from mitmproxy import ctx from mitmproxy import flow -import mitmproxy.types - -from mitmproxy.addons import multipart - class Share: - def __init__(self): - self.stream = None - self.filt = None - self.active_flows = set() # type: Set[flow.Flow] + def encode_multipart_formdata(self, filename, content): + 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 = b'multipart/form-data; boundary=%b' % LIMIT + return content_type, body - def open_file(self, path): - mode = "wb" - path = os.path.expanduser(path) - return open(path, mode) - - def start_stream_to_path(self, path, flt): + def post_multipart(self, host, filename, content): + 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))} try: - f = self.open_file(path) - except IOError as v: - raise exceptions.OptionsError(str(v)) - self.stream = io.FilteredFlowWriter(f, flt) - self.active_flows = set() - - def configure(self, updated): - # We're already streaming - stop the previous stream and restart - if "save_stream_filter" in updated: - if ctx.options.save_stream_filter: - self.filt = flowfilter.parse(ctx.options.save_stream_filter) - if not self.filt: - raise exceptions.OptionsError( - "Invalid filter specification: %s" % ctx.options.save_stream_filter - ) - else: - self.filt = None - if "save_stream_file" in updated or "save_stream_filter" in updated: - if self.stream: - self.done() - if ctx.options.save_stream_file: - self.start_stream_to_path(ctx.options.save_stream_file, self.filt) + 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: + conn.close() + return 'URL: share.mitmproxy.org/%s' % filename def base36encode(self, integer): chars, encoded = "0123456789abcdefghijklmnopqrstuvwxyz", "" @@ -62,7 +58,7 @@ class Share: return encoded - @command.command("share.file") + @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] @@ -72,41 +68,6 @@ class Share: stream.add(i) f.seek(0) content = f.read() - res = multipart.post_multipart('upload.share.mitmproxy.org.s3.amazonaws.com', u_id, content) + res = self.post_multipart('upload.share.mitmproxy.org.s3.amazonaws.com', u_id, content) f.close() ctx.log.alert("%s" % res) - - def tcp_start(self, flow): - if self.stream: - self.active_flows.add(flow) - - def tcp_end(self, flow): - if self.stream: - self.stream.add(flow) - self.active_flows.discard(flow) - - def websocket_start(self, flow): - if self.stream: - self.active_flows.add(flow) - - def websocket_end(self, flow): - if self.stream: - self.stream.add(flow) - self.active_flows.discard(flow) - - def response(self, flow): - if self.stream: - self.stream.add(flow) - self.active_flows.discard(flow) - - def request(self, flow): - if self.stream: - self.active_flows.add(flow) - - def done(self): - if self.stream: - for f in self.active_flows: - self.stream.add(f) - self.active_flows = set([]) - self.stream.fo.close() - self.stream = None diff --git a/test/mitmproxy/addons/test_share.py b/test/mitmproxy/addons/test_share.py new file mode 100644 index 000000000..e09696a72 --- /dev/null +++ b/test/mitmproxy/addons/test_share.py @@ -0,0 +1 @@ +# TODO: Write test From 2bcbaaece2b37e13698ce02a0c38f9c331d80842 Mon Sep 17 00:00:00 2001 From: kira0204 Date: Sat, 27 Jan 2018 04:39:37 +0530 Subject: [PATCH 5/7] Added test --- mitmproxy/addons/share.py | 6 ++++++ test/mitmproxy/addons/test_share.py | 30 ++++++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/mitmproxy/addons/share.py b/mitmproxy/addons/share.py index 17eefadbd..81d0fc022 100644 --- a/mitmproxy/addons/share.py +++ b/mitmproxy/addons/share.py @@ -12,6 +12,10 @@ from mitmproxy import flow class Share: + def __init__(self): + self.res = None + self.u_id = None + def encode_multipart_formdata(self, filename, content): params = {"key": filename, "acl": "bucket-owner-full-control", "Content-Type": "application/octet-stream"} LIMIT = b'---------------------------198495659117975628761412556003' @@ -69,5 +73,7 @@ class Share: f.seek(0) content = f.read() res = self.post_multipart('upload.share.mitmproxy.org.s3.amazonaws.com', u_id, content) + self.res = res + self.u_id = u_id f.close() ctx.log.alert("%s" % res) diff --git a/test/mitmproxy/addons/test_share.py b/test/mitmproxy/addons/test_share.py index e09696a72..106e8f041 100644 --- a/test/mitmproxy/addons/test_share.py +++ b/test/mitmproxy/addons/test_share.py @@ -1 +1,29 @@ -# TODO: Write test +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: + sh.share([tflow.tflow(resp=True)]) + assert sh.res == "URL: share.mitmproxy.org/%s" % sh.u_id + + mock_http.return_value.getresponse.side_effect = http.client.RemoteDisconnected + sh.share([tflow.tflow(resp=True)]) + assert sh.res == "The server couldn\'t fulfill the request." + + mock_http.return_value.request.side_effect = http.client.CannotSendRequest + sh.share([tflow.tflow(resp=True)]) + assert sh.res == "We failed to reach a server." + + v = view.View() + tctx.master.addons.add(v) + tctx.master.addons.add(sh) + tctx.master.commands.call_args("share.flows", ["@shown"]) \ No newline at end of file From c551b4aac72e4ea93c22358bd8009a743a517e1e Mon Sep 17 00:00:00 2001 From: kira0204 Date: Sun, 28 Jan 2018 22:56:05 +0530 Subject: [PATCH 6/7] Better approach to testing --- mitmproxy/addons/share.py | 6 ------ test/mitmproxy/addons/test_share.py | 8 ++++---- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/mitmproxy/addons/share.py b/mitmproxy/addons/share.py index 81d0fc022..17eefadbd 100644 --- a/mitmproxy/addons/share.py +++ b/mitmproxy/addons/share.py @@ -12,10 +12,6 @@ from mitmproxy import flow class Share: - def __init__(self): - self.res = None - self.u_id = None - def encode_multipart_formdata(self, filename, content): params = {"key": filename, "acl": "bucket-owner-full-control", "Content-Type": "application/octet-stream"} LIMIT = b'---------------------------198495659117975628761412556003' @@ -73,7 +69,5 @@ class Share: f.seek(0) content = f.read() res = self.post_multipart('upload.share.mitmproxy.org.s3.amazonaws.com', u_id, content) - self.res = res - self.u_id = u_id f.close() ctx.log.alert("%s" % res) diff --git a/test/mitmproxy/addons/test_share.py b/test/mitmproxy/addons/test_share.py index 106e8f041..a85c20b7b 100644 --- a/test/mitmproxy/addons/test_share.py +++ b/test/mitmproxy/addons/test_share.py @@ -13,17 +13,17 @@ def test_share_command(): sh = share.Share() with taddons.context() as tctx: sh.share([tflow.tflow(resp=True)]) - assert sh.res == "URL: share.mitmproxy.org/%s" % sh.u_id + assert tctx.master.has_log("URL: share.mitmproxy.org/") mock_http.return_value.getresponse.side_effect = http.client.RemoteDisconnected sh.share([tflow.tflow(resp=True)]) - assert sh.res == "The server couldn\'t fulfill the request." + assert tctx.master.has_log("The server couldn\'t fulfill the request.") mock_http.return_value.request.side_effect = http.client.CannotSendRequest sh.share([tflow.tflow(resp=True)]) - assert sh.res == "We failed to reach a server." + assert tctx.master.has_log("We failed to reach a server.") v = view.View() tctx.master.addons.add(v) tctx.master.addons.add(sh) - tctx.master.commands.call_args("share.flows", ["@shown"]) \ No newline at end of file + tctx.master.commands.call_args("share.flows", ["@shown"]) From de06a157ea862f87550c755358f0463eb12ec1cd Mon Sep 17 00:00:00 2001 From: kira0204 Date: Thu, 1 Feb 2018 05:01:19 +0530 Subject: [PATCH 7/7] 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)