diff options
| -rw-r--r-- | mitmproxy/addons/session.py | 36 | ||||
| -rw-r--r-- | mitmproxy/io/sql/session_create.sql | 36 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_session.py | 22 |
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) |