web: CSP, revert functionality, serve content

1 files changed, 78 insertions, 29 deletions

diff --git a/libmproxy/web/app.py b/libmproxy/web/app.py
index b5ec41c9..27e9aefc 100644
--- a/libmproxy/web/app.py
+++ b/libmproxy/web/app.py
@@ -1,4 +1,5 @@
 import os.path
+import re
 import tornado.web
 import tornado.websocket
 import logging
@@ -9,7 +10,43 @@ from .. import version
 class APIError(tornado.web.HTTPError):
     pass
 
-class IndexHandler(tornado.web.RequestHandler):
+
+class RequestHandler(tornado.web.RequestHandler):
+    def set_default_headers(self):
+        super(RequestHandler, self).set_default_headers()
+        self.set_header("Server", version.NAMEVERSION)
+        self.set_header("X-Frame-Options", "DENY")
+        self.add_header("X-XSS-Protection", "1; mode=block")
+        self.add_header("X-Content-Type-Options", "nosniff")
+        self.add_header("Content-Security-Policy", "default-src 'self'; "
+                                                   "connect-src 'self' ws://* ; "
+                                                   "style-src   'self' 'unsafe-inline'")
+
+    @property
+    def state(self):
+        return self.application.master.state
+
+    @property
+    def master(self):
+        return self.application.master
+
+    @property
+    def flow(self):
+        flow_id = str(self.path_kwargs["flow_id"])
+        flow = self.state.flows.get(flow_id)
+        if flow:
+            return flow
+        else:
+            raise APIError(400, "Flow not found.")
+
+    def write_error(self, status_code, **kwargs):
+        if "exc_info" in kwargs and isinstance(kwargs["exc_info"][1], APIError):
+            self.finish(kwargs["exc_info"][1].log_message)
+        else:
+            super(RequestHandler, self).write_error(status_code, **kwargs)
+
+
+class IndexHandler(RequestHandler):
     def get(self):
         _ = self.xsrf_token  # https://github.com/tornadoweb/tornado/issues/645
         self.render("index.html")
@@ -38,31 +75,6 @@ class ClientConnection(WebSocketEventBroadcaster):
     connections = set()
 
 
-class RequestHandler(tornado.web.RequestHandler):
-    @property
-    def state(self):
-        return self.application.master.state
-
-    @property
-    def master(self):
-        return self.application.master
-
-    @property
-    def flow(self):
-        flow_id = str(self.path_kwargs["flow_id"])
-        flow = self.state.flows.get(flow_id)
-        if flow:
-            return flow
-        else:
-            raise APIError(400, "Flow not found.")
-
-    def write_error(self, status_code, **kwargs):
-        if "exc_info" in kwargs and isinstance(kwargs["exc_info"][1], APIError):
-            self.finish(kwargs["exc_info"][1].log_message)
-        else:
-            super(RequestHandler, self).write_error(status_code, **kwargs)
-
-
 class Flows(RequestHandler):
     def get(self):
         self.write(dict(
@@ -95,13 +107,49 @@ class DuplicateFlow(RequestHandler):
     def post(self, flow_id):
         self.master.duplicate_flow(self.flow)
 
+
+class RevertFlow(RequestHandler):
+    def post(self, flow_id):
+        self.state.revert(self.flow)
+
+
 class ReplayFlow(RequestHandler):
     def post(self, flow_id):
-        self.flow.backup()
         r = self.master.replay_request(self.flow)
         if r:
             raise APIError(400, r)
 
+
+class FlowContent(RequestHandler):
+    def get(self, flow_id, message):
+        message = getattr(self.flow, message)
+
+        if not message.content:
+            raise APIError(400, "No content.")
+
+        content_encoding = message.headers.get_first("Content-Encoding", None)
+        if content_encoding:
+            content_encoding = re.sub(r"[^\w]", "", content_encoding)
+            self.set_header("Content-Encoding", content_encoding)
+
+        original_cd = message.headers.get_first("Content-Disposition", None)
+        filename = None
+        if original_cd:
+            filename = re.search("filename=([\w\" \.\-\(\)]+)", original_cd)
+            if filename:
+                filename = filename.group(1)
+        if not filename:
+            filename = self.flow.request.path.split("?")[0].split("/")[-1]
+
+        filename = re.sub(r"[^\w\" \.\-\(\)]", "", filename)
+        cd = "attachment; filename={}".format(filename)
+        self.set_header("Content-Disposition", cd)
+        self.set_header("Content-Type", "application/text")
+        self.set_header("X-Content-Type-Options", "nosniff")
+        self.set_header("X-Frame-Options", "DENY")
+        self.write(message.content)
+
+
 class Events(RequestHandler):
     def get(self):
         self.write(dict(
@@ -154,6 +202,8 @@ class Application(tornado.web.Application):
             (r"/flows/(?P<flow_id>[0-9a-f\-]+)/accept", AcceptFlow),
             (r"/flows/(?P<flow_id>[0-9a-f\-]+)/duplicate", DuplicateFlow),
             (r"/flows/(?P<flow_id>[0-9a-f\-]+)/replay", ReplayFlow),
+            (r"/flows/(?P<flow_id>[0-9a-f\-]+)/revert", RevertFlow),
+            (r"/flows/(?P<flow_id>[0-9a-f\-]+)/(?P<message>request|response)/content", FlowContent),
             (r"/settings", Settings),
             (r"/clear", ClearAll),
         ]
@@ -164,5 +214,4 @@ class Application(tornado.web.Application):
             cookie_secret=os.urandom(256),
             debug=debug,
         )
-        tornado.web.Application.__init__(self, handlers, **settings)
-
+        super(Application, self).__init__(handlers, **settings)
\ No newline at end of file