session: adding methods to capture and store flows

1 files changed, 107 insertions, 0 deletions

diff --git a/mitmproxy/addons/session.py b/mitmproxy/addons/session.py
index c49b95c4..6176bd5b 100644
--- a/mitmproxy/addons/session.py
+++ b/mitmproxy/addons/session.py
@@ -1,12 +1,30 @@
 import tempfile
+import asyncio
+import typing
+import bisect
 import shutil
 import sqlite3
 import os
 
+from mitmproxy import types
+from mitmproxy import http
+from mitmproxy import ctx
 from mitmproxy.exceptions import SessionLoadException
 from mitmproxy.utils.data import pkg_data
 
 
+class KeyifyList(object):
+    def __init__(self, inner, key):
+        self.inner = inner
+        self.key = key
+
+    def __len__(self):
+        return len(self.inner)
+
+    def __getitem__(self, k):
+        return self.key(self.inner[k])
+
+
 # Could be implemented using async libraries
 class SessionDB:
     """
@@ -71,3 +89,92 @@ class SessionDB:
             if c:
                 c.close()
         return False
+
+
+orders = [
+    ("t", "time"),
+    ("m", "method"),
+    ("u", "url"),
+    ("z", "size")
+]
+
+
+class Session:
+    def __init__(self):
+        self.sdb = SessionDB(ctx.options.session_path)
+        self._hot_store = []
+        self._view = []
+        self.order = orders[0]
+        self._flush_period = 3.0
+        self._flush_rate = 150
+
+    def load(self, loader):
+        loader.add_option(
+            "session_path", typing.Optional[types.Path], None,
+            "Path of session to load or to create."
+        )
+        loader.add_option(
+            "view_order", str, "time",
+            "Flow sort order.",
+            choices=list(map(lambda c: c[1], orders))
+        )
+
+    def _generate_order(self, f: http.HTTPFlow) -> typing.Union[str, int, float]:
+        o = self.order
+        if o == "time":
+            return f.request.timestamp_start or 0
+        if o == "method":
+            return f.request.method
+        if o == "url":
+            return f.request.url
+        if o == "size":
+            s = 0
+            if f.request.raw_content:
+                s += len(f.request.raw_content)
+            if f.response and f.response.raw_content:
+                s += len(f.response.raw_content)
+            return s
+
+    async def _writer(self):
+        while True:
+            await asyncio.sleep(self._flush_period)
+            tof = []
+            to_dump = min(self._flush_rate, len(self._hot_store))
+            for _ in range(to_dump):
+                tof.append(self._hot_store.pop())
+            self.store(tof)
+
+    def store(self, flows: typing.Sequence[http.HTTPFlow]):
+        pass
+
+    def running(self):
+        pass
+
+    def add(self, flows: typing.Sequence[http.HTTPFlow]) -> None:
+        for f in flows:
+            if f.id not in [f.id for f in self._hot_store] and f.id not in self.sdb:
+                # Flow has to be filtered here before adding to view. Later
+                o = self._generate_order(f)
+                self._view.insert(bisect.bisect_left(KeyifyList(self._view, lambda x: x[0]), o), (o, f.id))
+                self._hot_store.append(f)
+
+    def update(self, flow):
+        pass
+
+    def request(self, f):
+        self.add([f])
+
+    def error(self, f):
+        self.update([f])
+
+    def response(self, f):
+        self.update([f])
+
+    def intercept(self, f):
+        self.update([f])
+
+    def resume(self, f):
+        self.update([f])
+
+    def kill(self, f):
+        self.update([f])