mitmweb: 100% app test coverage, numerous fixes

5 files changed, 101 insertions, 99 deletions

diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 395a9d52..d7c312e5 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -303,8 +303,8 @@ class FlowListWalker(urwid.ListWalker):
 
 class FlowListBox(urwid.ListBox):
 
-    def __init__(self, master):
-        self.master = master
+    def __init__(self, master: "mitmproxy.tools.console.master.ConsoleMaster"):
+        self.master = master  # type: "mitmproxy.tools.console.master.ConsoleMaster"
         super().__init__(FlowListWalker(master))
 
     def get_method_raw(self, k):
@@ -348,7 +348,7 @@ class FlowListBox(urwid.ListBox):
         if key == "A":
             for f in self.master.view:
                 if f.intercepted:
-                    f.resume()
+                    f.resume(self.master)
             signals.flowlist_change.send(self)
         elif key == "z":
             self.master.view.clear()
diff --git a/mitmproxy/tools/console/flowview.py b/mitmproxy/tools/console/flowview.py
index e2b24fab..ecb070d8 100644
--- a/mitmproxy/tools/console/flowview.py
+++ b/mitmproxy/tools/console/flowview.py
@@ -510,8 +510,10 @@ class FlowView(tabs.Tabs):
             self.flow.resume(self.master)
             signals.flow_change.send(self, flow = self.flow)
         elif key == "A":
-            self.master.accept_all()
-            signals.flow_change.send(self, flow = self.flow)
+            for f in self.view:
+                if f.intercepted:
+                    f.resume(self.master)
+                    signals.flow_change.send(self, flow=f)
         elif key == "d":
             if self.flow.killable:
                 self.flow.kill(self.master)
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index f8850404..834d6ead 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -71,7 +71,7 @@ class ConsoleMaster(master.Master):
 
     def __init__(self, options, server):
         super().__init__(options, server)
-        self.view = view.View()
+        self.view = view.View()  # type: view.View
         self.stream_path = None
         # This line is just for type hinting
         self.options = self.options  # type: Options
diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py
index fa0b702f..f617bd08 100644
--- a/mitmproxy/tools/web/app.py
+++ b/mitmproxy/tools/web/app.py
@@ -11,6 +11,7 @@ import tornado.escape
 import tornado.web
 import tornado.websocket
 from mitmproxy import contentviews
+from mitmproxy import exceptions
 from mitmproxy import flowfilter
 from mitmproxy import http
 from mitmproxy import io
@@ -108,9 +109,23 @@ class RequestHandler(tornado.web.RequestHandler):
 
     @property
     def json(self):
-        if not self.request.headers.get("Content-Type").startswith("application/json"):
-            return None
-        return json.loads(self.request.body.decode())
+        if not self.request.headers.get("Content-Type", "").startswith("application/json"):
+            raise APIError(400, "Invalid Content-Type, expected application/json.")
+        try:
+            return json.loads(self.request.body.decode())
+        except Exception as e:
+            raise APIError(400, "Malformed JSON: {}".format(str(e)))
+
+    @property
+    def filecontents(self):
+        """
+        Accept either a multipart/form file upload or just take the plain request body.
+
+        """
+        if self.request.files:
+            return next(iter(self.request.files.values()))[0].body
+        else:
+            return self.request.body
 
     @property
     def view(self) -> mitmproxy.addons.view.View:
@@ -124,11 +139,11 @@ class RequestHandler(tornado.web.RequestHandler):
     def flow(self) -> mitmproxy.flow.Flow:
         flow_id = str(self.path_kwargs["flow_id"])
         # FIXME: Add a facility to addon.view to safely access the store
-        flow = self.view._store.get(flow_id)
+        flow = self.view.get_by_id(flow_id)
         if flow:
             return flow
         else:
-            raise APIError(400, "Flow not found.")
+            raise APIError(404, "Flow not found.")
 
     def write_error(self, status_code: int, **kwargs):
         if "exc_info" in kwargs and isinstance(kwargs["exc_info"][1], APIError):
@@ -168,7 +183,7 @@ class WebSocketEventBroadcaster(tornado.websocket.WebSocketHandler):
         for conn in cls.connections:
             try:
                 conn.write_message(message)
-            except Exception:
+            except Exception:  # pragma: no cover
                 logging.error("Error sending message", exc_info=True)
 
 
@@ -196,10 +211,8 @@ class DumpFlows(RequestHandler):
 
     def post(self):
         self.view.clear()
-
-        content = self.request.files.values()[0][0].body
-        bio = BytesIO(content)
-        self.master.load_flows(io.FlowReader(bio).stream())
+        bio = BytesIO(self.filecontents)
+        self.master.load_flows(io.FlowReader(bio))
         bio.close()
 
 
@@ -211,7 +224,8 @@ class ClearAll(RequestHandler):
 
 class AcceptFlows(RequestHandler):
     def post(self):
-        self.master.accept_all(self.master)
+        for f in self.view:
+            f.resume(self.master)
 
 
 class AcceptFlow(RequestHandler):
@@ -228,53 +242,59 @@ class FlowHandler(RequestHandler):
     def put(self, flow_id):
         flow = self.flow
         flow.backup()
-        for a, b in self.json.items():
-            if a == "request" and hasattr(flow, "request"):
-                request = flow.request
-                for k, v in b.items():
-                    if k in ["method", "scheme", "host", "path", "http_version"]:
-                        setattr(request, k, str(v))
-                    elif k == "port":
-                        request.port = int(v)
-                    elif k == "headers":
-                        request.headers.clear()
-                        for header in v:
-                            request.headers.add(*header)
-                    elif k == "content":
-                        request.text = v
-                    else:
-                        print("Warning: Unknown update {}.{}: {}".format(a, k, v))
-
-            elif a == "response" and hasattr(flow, "response"):
-                response = flow.response
-                for k, v in b.items():
-                    if k == "msg":
-                        response.msg = str(v)
-                    elif k == "code":
-                        response.status_code = int(v)
-                    elif k == "http_version":
-                        response.http_version = str(v)
-                    elif k == "headers":
-                        response.headers.clear()
-                        for header in v:
-                            response.headers.add(*header)
-                    elif k == "content":
-                        response.text = v
-                    else:
-                        print("Warning: Unknown update {}.{}: {}".format(a, k, v))
-            else:
-                print("Warning: Unknown update {}: {}".format(a, b))
+        try:
+            for a, b in self.json.items():
+                if a == "request" and hasattr(flow, "request"):
+                    request = flow.request
+                    for k, v in b.items():
+                        if k in ["method", "scheme", "host", "path", "http_version"]:
+                            setattr(request, k, str(v))
+                        elif k == "port":
+                            request.port = int(v)
+                        elif k == "headers":
+                            request.headers.clear()
+                            for header in v:
+                                request.headers.add(*header)
+                        elif k == "content":
+                            request.text = v
+                        else:
+                            raise APIError(400, "Unknown update request.{}: {}".format(k, v))
+
+                elif a == "response" and hasattr(flow, "response"):
+                    response = flow.response
+                    for k, v in b.items():
+                        if k in ["msg", "http_version"]:
+                            setattr(response, k, str(v))
+                        elif k == "code":
+                            response.status_code = int(v)
+                        elif k == "headers":
+                            response.headers.clear()
+                            for header in v:
+                                response.headers.add(*header)
+                        elif k == "content":
+                            response.text = v
+                        else:
+                            raise APIError(400, "Unknown update response.{}: {}".format(k, v))
+                else:
+                    raise APIError(400, "Unknown update {}: {}".format(a, b))
+        except APIError:
+            flow.revert()
+            raise
         self.view.update(flow)
 
 
 class DuplicateFlow(RequestHandler):
     def post(self, flow_id):
-        self.master.view.duplicate_flow(self.flow)
+        f = self.flow.copy()
+        self.view.add(f)
+        self.write(f.id)
 
 
 class RevertFlow(RequestHandler):
     def post(self, flow_id):
-        self.flow.revert()
+        if self.flow.modified():
+            self.flow.revert()
+            self.view.update(self.flow)
 
 
 class ReplayFlow(RequestHandler):
@@ -283,16 +303,17 @@ class ReplayFlow(RequestHandler):
         self.flow.response = None
         self.view.update(self.flow)
 
-        r = self.master.replay_request(self.flow)
-        if r:
-            raise APIError(400, r)
+        try:
+            self.master.replay_request(self.flow)
+        except exceptions.ReplayException as e:
+            raise APIError(400, str(e))
 
 
 class FlowContent(RequestHandler):
     def post(self, flow_id, message):
         self.flow.backup()
         message = getattr(self.flow, message)
-        message.content = self.request.files.values()[0][0].body
+        message.content = self.filecontents
         self.view.update(self.flow)
 
     def get(self, flow_id, message):
@@ -364,46 +385,16 @@ class Settings(RequestHandler):
         ))
 
     def put(self):
-        update = {}
-        for k, v in self.json.items():
-            if k == "intercept":
-                self.master.options.intercept = v
-                update[k] = v
-            elif k == "showhost":
-                self.master.options.showhost = v
-                update[k] = v
-            elif k == "no_upstream_cert":
-                self.master.options.no_upstream_cert = v
-                update[k] = v
-            elif k == "rawtcp":
-                self.master.options.rawtcp = v
-                update[k] = v
-            elif k == "http2":
-                self.master.options.http2 = v
-                update[k] = v
-            elif k == "anticache":
-                self.master.options.anticache = v
-                update[k] = v
-            elif k == "anticomp":
-                self.master.options.anticomp = v
-                update[k] = v
-            elif k == "stickycookie":
-                self.master.options.stickycookie = v
-                update[k] = v
-            elif k == "stickyauth":
-                self.master.options.stickyauth = v
-                update[k] = v
-            elif k == "stream":
-                self.master.options.stream_large_bodies = v
-                update[k] = v
-            else:
-                print("Warning: Unknown setting {}: {}".format(k, v))
-
-        ClientConnection.broadcast(
-            resource="settings",
-            cmd="update",
-            data=update
-        )
+        update = self.json
+        option_whitelist = {
+            "intercept", "showhost", "no_upstream_cert",
+            "rawtcp", "http2", "anticache", "anticomp",
+            "stickycookie", "stickyauth", "stream_large_bodies"
+        }
+        for k in update:
+            if k not in option_whitelist:
+                raise APIError(400, "Unknown setting {}".format(k))
+        self.master.options.update(**update)
 
 
 class Application(tornado.web.Application):
diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py
index 453455f0..5457fb40 100644
--- a/mitmproxy/tools/web/master.py
+++ b/mitmproxy/tools/web/master.py
@@ -45,6 +45,8 @@ class WebMaster(master.Master):
         self.events.sig_add.connect(self._sig_events_add)
         self.events.sig_refresh.connect(self._sig_events_refresh)
 
+        self.options.changed.connect(self._sig_options_update)
+
         self.addons.add(*addons.default_addons())
         self.addons.add(self.view, self.events, intercept.Intercept())
         self.app = app.Application(
@@ -101,6 +103,13 @@ class WebMaster(master.Master):
             cmd="reset"
         )
 
+    def _sig_options_update(self, options, updated):
+        app.ClientConnection.broadcast(
+            resource="settings",
+            cmd="update",
+            data={k: getattr(options, k) for k in updated}
+        )
+
     def run(self):  # pragma: no cover
 
         iol = tornado.ioloop.IOLoop.instance()