From b686073f56667f7caf11c9013d592a8f091ccd9d Mon Sep 17 00:00:00 2001 From: madt1m Date: Sat, 23 Jun 2018 03:36:58 +0200 Subject: Hooking the view to load from new DB handler --- examples/addons/protodumper.py | 5 ----- mitmproxy/addons/view.py | 13 +++++-------- mitmproxy/io/__init__.py | 3 ++- mitmproxy/io/db.py | 22 ++++++++++++++++------ 4 files changed, 23 insertions(+), 20 deletions(-) delete mode 100644 examples/addons/protodumper.py diff --git a/examples/addons/protodumper.py b/examples/addons/protodumper.py deleted file mode 100644 index 6a1d676c..00000000 --- a/examples/addons/protodumper.py +++ /dev/null @@ -1,5 +0,0 @@ -from mitmproxy.io import protobuf - - -def response(f): - protobuf.loads(protobuf.dumps(f)) diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py index 1c8bd0ce..3730aa8d 100644 --- a/mitmproxy/addons/view.py +++ b/mitmproxy/addons/view.py @@ -447,16 +447,13 @@ class View(collections.Sequence): Load flows into the view, without processing them with addons. """ try: - with open(path, "rb") as f: - for i in io.FlowReader(f).stream(): - # Do this to get a new ID, so we can load the same file N times and - # get new flows each time. It would be more efficient to just have a - # .newid() method or something. - self.add([i.copy()]) + dh = io.DbHandler(path) + for f in dh.load(): + self.add([f.copy()]) + except exceptions.TypeError as e: + ctx.log.error(str(e)) except IOError as e: ctx.log.error(e.strerror) - except exceptions.FlowReadException as e: - ctx.log.error(str(e)) def add(self, flows: typing.Sequence[mitmproxy.flow.Flow]) -> None: """ diff --git a/mitmproxy/io/__init__.py b/mitmproxy/io/__init__.py index 540e6871..854d2505 100644 --- a/mitmproxy/io/__init__.py +++ b/mitmproxy/io/__init__.py @@ -1,7 +1,8 @@ from .io import FlowWriter, FlowReader, FilteredFlowWriter, read_flows_from_paths +from .db import DbHandler __all__ = [ - "FlowWriter", "FlowReader", "FilteredFlowWriter", "read_flows_from_paths" + "FlowWriter", "FlowReader", "FilteredFlowWriter", "read_flows_from_paths", "DbHandler" ] diff --git a/mitmproxy/io/db.py b/mitmproxy/io/db.py index f7e13e91..b59650ab 100644 --- a/mitmproxy/io/db.py +++ b/mitmproxy/io/db.py @@ -1,5 +1,9 @@ import sqlite3 +import os + from mitmproxy.io import protobuf +from mitmproxy.http import HTTPFlow +from mitmproxy import exceptions class DbHandler: @@ -8,7 +12,11 @@ class DbHandler: This class is wrapping up connection to SQLITE DB. """ - def __init__(self, db_path="tmp.sqlite"): + def __init__(self, db_path="/tmp/tmp.sqlite"): + if os.path.isfile(db_path): + self.db_path = db_path + else: + raise IOError("Invalid path!") self.db_path = db_path self._con = sqlite3.connect(self.db_path) self._c = self._con.cursor() @@ -17,16 +25,18 @@ class DbHandler: def _create_db(self): with self._con: self._con.execute('CREATE TABLE IF NOT EXISTS FLOWS(' - 'id INTEGER PRIMARY KEY AUTOINCREMENT,' + 'id INTEGER PRIMARY KEY,' 'pbuf_blob BLOB)') def store(self, flows): blobs = [] for flow in flows: - blobs.append(protobuf.dumps(flow)) + blobs.append((protobuf.dumps(flow),)) with self._con: - self._con.executemany('INSERT INTO FLOWS values (?)', blobs) + self._con.executemany('INSERT INTO FLOWS (pbuf_blob) values (?)', blobs) def load(self): - self._c.execute('SELECT * FROM FLOWS') - return self._c.fetchall() + flows = [] + self._c.execute('SELECT pbuf_blob FROM FLOWS') + for row in self._c.fetchall(): + flows.append(HTTPFlow.from_state(protobuf.loads(row[0]))) -- cgit v1.2.3