Merge pull request #3252 from madt1m/session-db

Session - Hybrid DB

4 files changed, 155 insertions, 0 deletions

diff --git a/mitmproxy/addons/session.py b/mitmproxy/addons/session.py
new file mode 100644
index 00000000..c49b95c4
--- /dev/null
+++ b/mitmproxy/addons/session.py
@@ -0,0 +1,73 @@
+import tempfile
+import shutil
+import sqlite3
+import os
+
+from mitmproxy.exceptions import SessionLoadException
+from mitmproxy.utils.data import pkg_data
+
+
+# Could be implemented using async libraries
+class SessionDB:
+    """
+    This class wraps connection to DB
+    for Sessions and handles creation,
+    retrieving and insertion in tables.
+    """
+
+    def __init__(self, db_path=None):
+        """
+        Connect to an already existing database,
+        or create a new one with optional path.
+        :param db_path:
+        """
+        self.tempdir = None
+        self.con = None
+        if db_path is not None and os.path.isfile(db_path):
+            self._load_session(db_path)
+        else:
+            if db_path:
+                path = db_path
+            else:
+                self.tempdir = tempfile.mkdtemp()
+                path = os.path.join(self.tempdir, 'tmp.sqlite')
+            self.con = sqlite3.connect(path)
+            self._create_session()
+
+    def __del__(self):
+        if self.con:
+            self.con.close()
+        if self.tempdir:
+            shutil.rmtree(self.tempdir)
+
+    def _load_session(self, path):
+        if not self.is_session_db(path):
+            raise SessionLoadException('Given path does not point to a valid Session')
+        self.con = sqlite3.connect(path)
+
+    def _create_session(self):
+        script_path = pkg_data.path("io/sql/session_create.sql")
+        qry = open(script_path, 'r').read()
+        with self.con:
+            self.con.executescript(qry)
+
+    @staticmethod
+    def is_session_db(path):
+        """
+        Check if database entered from user
+        is a valid Session SQLite DB.
+        :return: True if valid, False if invalid.
+        """
+        try:
+            c = sqlite3.connect(f'file:{path}?mode=rw', uri=True)
+            cursor = c.cursor()
+            cursor.execute("SELECT NAME FROM sqlite_master WHERE type='table';")
+            rows = cursor.fetchall()
+            tables = [('flow',), ('body',), ('annotation',)]
+            if all(elem in rows for elem in tables):
+                c.close()
+                return True
+        except:
+            if c:
+                c.close()
+        return False
diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py
index d568898b..9f0a8c30 100644
--- a/mitmproxy/exceptions.py
+++ b/mitmproxy/exceptions.py
@@ -129,6 +129,10 @@ class NetlibException(MitmproxyException):
         super().__init__(message)
 
 
+class SessionLoadException(MitmproxyException):
+    pass
+
+
 class Disconnect:
     """Immediate EOF"""
 
diff --git a/mitmproxy/io/sql/session_create.sql b/mitmproxy/io/sql/session_create.sql
new file mode 100644
index 00000000..bfc98b94
--- /dev/null
+++ b/mitmproxy/io/sql/session_create.sql
@@ -0,0 +1,20 @@
+CREATE TABLE flow (
+id VARCHAR(36) PRIMARY KEY,
+content BLOB
+);
+
+CREATE TABLE body (
+id INTEGER PRIMARY KEY,
+flow_id VARCHAR(36),
+type_id INTEGER,
+content BLOB,
+FOREIGN KEY(flow_id) REFERENCES flow(id)
+);
+
+CREATE TABLE annotation (
+id INTEGER PRIMARY KEY,
+flow_id VARCHAR(36),
+type VARCHAR(16),
+content BLOB,
+FOREIGN KEY(flow_id) REFERENCES flow(id)
+);
diff --git a/test/mitmproxy/addons/test_session.py b/test/mitmproxy/addons/test_session.py
new file mode 100644
index 00000000..d4b1109b
--- /dev/null
+++ b/test/mitmproxy/addons/test_session.py
@@ -0,0 +1,58 @@
+import sqlite3
+import pytest
+import os
+
+from mitmproxy.addons import session
+from mitmproxy.exceptions import SessionLoadException
+from mitmproxy.utils.data import pkg_data
+
+
+class TestSession:
+    def test_session_temporary(self):
+        s = session.SessionDB()
+        td = s.tempdir
+        filename = os.path.join(td, 'tmp.sqlite')
+        assert session.SessionDB.is_session_db(filename)
+        assert os.path.isdir(td)
+        del s
+        assert not os.path.isdir(td)
+
+    def test_session_not_valid(self, tdata):
+        path = tdata.path('mitmproxy/data/') + '/test_snv.sqlite'
+        if os.path.isfile(path):
+            os.remove(path)
+        with open(path, 'w') as handle:
+            handle.write("Not valid data")
+        with pytest.raises(SessionLoadException):
+            session.SessionDB(path)
+        os.remove(path)
+
+    def test_session_new_persistent(self, tdata):
+        path = tdata.path('mitmproxy/data/') + '/test_np.sqlite'
+        if os.path.isfile(path):
+            os.remove(path)
+        session.SessionDB(path)
+        assert session.SessionDB.is_session_db(path)
+        os.remove(path)
+
+    def test_session_load_existing(self, tdata):
+        path = tdata.path('mitmproxy/data/') + '/test_le.sqlite'
+        if os.path.isfile(path):
+            os.remove(path)
+        con = sqlite3.connect(path)
+        script_path = pkg_data.path("io/sql/session_create.sql")
+        qry = open(script_path, 'r').read()
+        with con:
+            con.executescript(qry)
+            blob = b'blob_of_data'
+            con.execute(f'INSERT INTO FLOW VALUES(1, "{blob}");')
+        con.close()
+        session.SessionDB(path)
+        con = sqlite3.connect(path)
+        with con:
+            cur = con.cursor()
+            cur.execute('SELECT * FROM FLOW;')
+            rows = cur.fetchall()
+            assert len(rows) == 1
+        con.close()
+        os.remove(path)