From 4cebfbabadfb6f500cf762923187017c47295961 Mon Sep 17 00:00:00 2001 From: madt1m Date: Fri, 20 Jul 2018 16:56:26 +0200 Subject: session: Implemented SessionDB Schema --- mitmproxy/io/sql/session_create.sql | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 mitmproxy/io/sql/session_create.sql diff --git a/mitmproxy/io/sql/session_create.sql b/mitmproxy/io/sql/session_create.sql new file mode 100644 index 00000000..8387981c --- /dev/null +++ b/mitmproxy/io/sql/session_create.sql @@ -0,0 +1,26 @@ +CREATE TABLE FLOW ( +FID INTEGER PRIMARY KEY, +MID INTEGER, +BID INTEGER, +CONTENT BLOB +); + +CREATE TABLE META ( +MID INTEGER PRIMARY KEY, +INTERCEPTED INTEGER, +MARKED INTEGER, +MODE VARCHAR(20) +); + +CREATE TABLE BODY ( +BID INTEGER, +BREQ BLOB, +BRES BLOB +); + +CREATE TABLE ANNOTATION ( +AID INTEGER PRIMARY KEY, +FID INTEGER, +TYPE VARCHAR(20), +CONTENT BLOB +); \ No newline at end of file -- cgit v1.2.3 From 04b9555db8c821c053e156e65cb97d24eaf134b8 Mon Sep 17 00:00:00 2001 From: madt1m Date: Fri, 20 Jul 2018 16:56:51 +0200 Subject: session: load/create wrapper methods for SessionDB --- mitmproxy/addons/session.py | 59 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 mitmproxy/addons/session.py diff --git a/mitmproxy/addons/session.py b/mitmproxy/addons/session.py new file mode 100644 index 00000000..d065ade2 --- /dev/null +++ b/mitmproxy/addons/session.py @@ -0,0 +1,59 @@ +import os +import sqlite3 + +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: + """ + 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) + 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) + + 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) + + @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',), ('META',), ('ANNOTATION',)] + if all(elem in rows for elem in tables): + c.close() + return True + except: + if c: + c.close() + return False -- cgit v1.2.3 From 1a584420276e4fab7997797ba14c004396924771 Mon Sep 17 00:00:00 2001 From: madt1m Date: Fri, 20 Jul 2018 16:57:23 +0200 Subject: exceptions: new SessionLoad exception for load errors --- mitmproxy/exceptions.py | 4 ++++ 1 file changed, 4 insertions(+) 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""" -- cgit v1.2.3 From fad8e7c99b5db84f6e03da674453c28e47e97d9f Mon Sep 17 00:00:00 2001 From: madt1m Date: Fri, 20 Jul 2018 16:58:11 +0200 Subject: tests: SessionDB fully tested --- test/mitmproxy/addons/test_session.py | 52 +++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 test/mitmproxy/addons/test_session.py diff --git a/test/mitmproxy/addons/test_session.py b/test/mitmproxy/addons/test_session.py new file mode 100644 index 00000000..2424bd28 --- /dev/null +++ b/test/mitmproxy/addons/test_session.py @@ -0,0 +1,52 @@ +import sqlite3 +import os +import pytest + +from mitmproxy.exceptions import SessionLoadException +from mitmproxy.addons import session +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') + + def test_session_not_valid(self, tdata): + path = tdata.path('mitmproxy/data/') + 'test.sqlite' + if os.path.isfile(path): + os.remove(path) + with open(path, 'w') as handle: + handle.write("Not valid data") + with pytest.raises(SessionLoadException): + s = session.SessionDB(path) + + def test_session_new_persistent(self, tdata): + path = tdata.path('mitmproxy/data/') + 'test.sqlite' + if os.path.isfile(path): + os.remove(path) + s = session.SessionDB(path) + assert session.SessionDB.is_session_db(path) + + def test_session_load_existing(self, tdata): + path = tdata.path('mitmproxy/data/') + 'test.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, 1, 1, "{blob}");') + con.close() + s = session.SessionDB(path) + con = sqlite3.connect(path) + with con: + cur = con.cursor() + cur.execute('SELECT * FROM FLOW;') + rows = cur.fetchall() + assert len(rows) == 1 -- cgit v1.2.3 From 3b5cdf7f67e030fa0887bc309068b727f20918e3 Mon Sep 17 00:00:00 2001 From: madt1m Date: Fri, 20 Jul 2018 18:15:27 +0200 Subject: test_linting: removed some unused local variables --- test/mitmproxy/addons/test_session.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/test/mitmproxy/addons/test_session.py b/test/mitmproxy/addons/test_session.py index 2424bd28..7089b89e 100644 --- a/test/mitmproxy/addons/test_session.py +++ b/test/mitmproxy/addons/test_session.py @@ -16,23 +16,23 @@ class TestSession: os.remove('tmp.sqlite') def test_session_not_valid(self, tdata): - path = tdata.path('mitmproxy/data/') + 'test.sqlite' + path = tdata.path('mitmproxy/data/') + '/test.sqlite' if os.path.isfile(path): os.remove(path) with open(path, 'w') as handle: handle.write("Not valid data") with pytest.raises(SessionLoadException): - s = session.SessionDB(path) + session.SessionDB(path) def test_session_new_persistent(self, tdata): - path = tdata.path('mitmproxy/data/') + 'test.sqlite' + path = tdata.path('mitmproxy/data/') + '/test.sqlite' if os.path.isfile(path): os.remove(path) - s = session.SessionDB(path) + session.SessionDB(path) assert session.SessionDB.is_session_db(path) def test_session_load_existing(self, tdata): - path = tdata.path('mitmproxy/data/') + 'test.sqlite' + path = tdata.path('mitmproxy/data/') + '/test.sqlite' if os.path.isfile(path): os.remove(path) con = sqlite3.connect(path) @@ -43,7 +43,7 @@ class TestSession: blob = b'blob_of_data' con.execute(f'INSERT INTO FLOW VALUES(1, 1, 1, "{blob}");') con.close() - s = session.SessionDB(path) + session.SessionDB(path) con = sqlite3.connect(path) with con: cur = con.cursor() -- cgit v1.2.3 From 68eb07b668b44074271b358b25e48ffd2b4726e0 Mon Sep 17 00:00:00 2001 From: madt1m Date: Tue, 24 Jul 2018 11:14:39 +0200 Subject: session: modified schema. Now SessionDB uses tempfile module for temp session --- mitmproxy/addons/session.py | 36 +++++++++++++++++++++++++---------- mitmproxy/io/sql/session_create.sql | 36 +++++++++++++++-------------------- 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) -- cgit v1.2.3 From 8c7793b91a2a182ec7d7831f2c03283dad3488cc Mon Sep 17 00:00:00 2001 From: madt1m Date: Tue, 24 Jul 2018 15:57:11 +0200 Subject: session: temporary DB is now stored in temporary dir --- mitmproxy/addons/session.py | 14 ++++++-------- test/mitmproxy/addons/test_session.py | 8 ++++++-- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/mitmproxy/addons/session.py b/mitmproxy/addons/session.py index 7f1c0025..c49b95c4 100644 --- a/mitmproxy/addons/session.py +++ b/mitmproxy/addons/session.py @@ -1,4 +1,5 @@ import tempfile +import shutil import sqlite3 import os @@ -20,7 +21,7 @@ class SessionDB: or create a new one with optional path. :param db_path: """ - self.temp = None + self.tempdir = None self.con = None if db_path is not None and os.path.isfile(db_path): self._load_session(db_path) @@ -28,19 +29,16 @@ class SessionDB: 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.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.temp: - # This is a workaround to ensure portability - os.remove(self.temp.name) + if self.tempdir: + shutil.rmtree(self.tempdir) def _load_session(self, path): if not self.is_session_db(path): diff --git a/test/mitmproxy/addons/test_session.py b/test/mitmproxy/addons/test_session.py index cb36e283..d4b1109b 100644 --- a/test/mitmproxy/addons/test_session.py +++ b/test/mitmproxy/addons/test_session.py @@ -8,10 +8,14 @@ from mitmproxy.utils.data import pkg_data class TestSession: - def test_session_temporary(self, tdata): + def test_session_temporary(self): s = session.SessionDB() - filename = s.temp.name + 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' -- cgit v1.2.3