session: modified schema. Now SessionDB uses tempfile module for temp session

3 files changed, 53 insertions, 41 deletions

diff --git a/mitmproxy/addons/session.py b/mitmproxy/addons/session.py
index d065ade2..7f1c0025 100644
--- a/mitmproxy/addons/session.py
+++ b/mitmproxy/addons/session.py
@@ -1,5 +1,6 @@
-import os
+import tempfile
 import sqlite3
+import os
 
 from mitmproxy.exceptions import SessionLoadException
 from mitmproxy.utils.data import pkg_data
@@ -19,24 +20,39 @@ class SessionDB:
         or create a new one with optional path.
         :param db_path:
         """
+        self.temp = None
+        self.con = None
         if db_path is not None and os.path.isfile(db_path):
             self._load_session(db_path)
         else:
-            path = db_path or 'tmp.sqlite'
-            # in case tmp.sqlite already exists in FS
-            if os.path.isfile(path):
-                os.remove(path)
+            if db_path:
+                path = db_path
+            else:
+                # We use tempfile only to generate a path, since we demand file creation to sqlite, and removal to os.
+                self.temp = tempfile.NamedTemporaryFile()
+                path = self.temp.name
+                self.temp.close()
             self.con = sqlite3.connect(path)
-            script_path = pkg_data.path("io/sql/session_create.sql")
-            qry = open(script_path, 'r').read()
-            with self.con:
-                self.con.executescript(qry)
+            self._create_session()
+
+    def __del__(self):
+        if self.con:
+            self.con.close()
+        if self.temp:
+            # This is a workaround to ensure portability
+            os.remove(self.temp.name)
 
     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):
         """
@@ -49,7 +65,7 @@ class SessionDB:
             cursor = c.cursor()
             cursor.execute("SELECT NAME FROM sqlite_master WHERE type='table';")
             rows = cursor.fetchall()
-            tables = [('FLOW',), ('BODY',), ('META',), ('ANNOTATION',)]
+            tables = [('flow',), ('body',), ('annotation',)]
             if all(elem in rows for elem in tables):
                 c.close()
                 return True
diff --git a/mitmproxy/io/sql/session_create.sql b/mitmproxy/io/sql/session_create.sql
index 8387981c..bfc98b94 100644
--- a/mitmproxy/io/sql/session_create.sql
+++ b/mitmproxy/io/sql/session_create.sql
@@ -1,26 +1,20 @@
-CREATE TABLE FLOW (
-FID INTEGER PRIMARY KEY,
-MID INTEGER,
-BID INTEGER,
-CONTENT BLOB
+CREATE TABLE flow (
+id VARCHAR(36) PRIMARY KEY,
+content BLOB
 );
 
-CREATE TABLE META (
-MID INTEGER PRIMARY KEY,
-INTERCEPTED INTEGER,
-MARKED INTEGER,
-MODE VARCHAR(20)
+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 BODY (
-BID INTEGER,
-BREQ BLOB,
-BRES BLOB
+CREATE TABLE annotation (
+id INTEGER PRIMARY KEY,
+flow_id VARCHAR(36),
+type VARCHAR(16),
+content BLOB,
+FOREIGN KEY(flow_id) REFERENCES flow(id)
 );
-
-CREATE TABLE ANNOTATION (
-AID INTEGER PRIMARY KEY,
-FID INTEGER,
-TYPE VARCHAR(20),
-CONTENT BLOB
-);
\ No newline at end of file
diff --git a/test/mitmproxy/addons/test_session.py b/test/mitmproxy/addons/test_session.py
index 7089b89e..cb36e283 100644
--- a/test/mitmproxy/addons/test_session.py
+++ b/test/mitmproxy/addons/test_session.py
@@ -1,38 +1,38 @@
 import sqlite3
-import os
 import pytest
+import os
 
-from mitmproxy.exceptions import SessionLoadException
 from mitmproxy.addons import session
+from mitmproxy.exceptions import SessionLoadException
 from mitmproxy.utils.data import pkg_data
 
 
 class TestSession:
     def test_session_temporary(self, tdata):
-        open('tmp.sqlite', 'w')
         s = session.SessionDB()
-        assert session.SessionDB.is_session_db('tmp.sqlite')
-        s.con.close()
-        os.remove('tmp.sqlite')
+        filename = s.temp.name
+        assert session.SessionDB.is_session_db(filename)
 
     def test_session_not_valid(self, tdata):
-        path = tdata.path('mitmproxy/data/') + '/test.sqlite'
+        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.sqlite'
+        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.sqlite'
+        path = tdata.path('mitmproxy/data/') + '/test_le.sqlite'
         if os.path.isfile(path):
             os.remove(path)
         con = sqlite3.connect(path)
@@ -41,7 +41,7 @@ class TestSession:
         with con:
             con.executescript(qry)
             blob = b'blob_of_data'
-            con.execute(f'INSERT INTO FLOW VALUES(1, 1, 1, "{blob}");')
+            con.execute(f'INSERT INTO FLOW VALUES(1, "{blob}");')
         con.close()
         session.SessionDB(path)
         con = sqlite3.connect(path)
@@ -50,3 +50,5 @@ class TestSession:
             cur.execute('SELECT * FROM FLOW;')
             rows = cur.fetchall()
             assert len(rows) == 1
+        con.close()
+        os.remove(path)