Improving the logic and test

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))
 

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()

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)