Hooking the view to load from new DB handler

4 files changed, 23 insertions, 20 deletions

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])))