diff options
author | Pietro Francesco Tirenna <pietrotirenna.pt@gmail.com> | 2018-07-28 17:44:02 +0200 |
---|---|---|
committer | Pietro Francesco Tirenna <pietrotirenna.pt@gmail.com> | 2018-07-28 17:44:02 +0200 |
commit | 53b85d23605409a46754f217d061f7d8a1d3cb6d (patch) | |
tree | c053bd80198d92556f0dbd8966317da5d54a10c0 | |
parent | 9c949bd2f852fb5f5a314e6faa45e4869571628c (diff) | |
download | mitmproxy-53b85d23605409a46754f217d061f7d8a1d3cb6d.tar.gz mitmproxy-53b85d23605409a46754f217d061f7d8a1d3cb6d.tar.bz2 mitmproxy-53b85d23605409a46754f217d061f7d8a1d3cb6d.zip |
session: adding methods to capture and store flows
-rw-r--r-- | mitmproxy/addons/session.py | 107 |
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]) |