diff options
author | Pietro Francesco Tirenna <pietrotirenna.pt@gmail.com> | 2018-07-24 16:26:10 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2018-07-24 16:26:10 +0200 |
commit | 9c949bd2f852fb5f5a314e6faa45e4869571628c (patch) | |
tree | bf2718f6ca8bd1088ccc9945c7d1b9be4ff62a59 /mitmproxy | |
parent | 16e07996d983483bb732269c921e89e8bda6ee20 (diff) | |
parent | 8c7793b91a2a182ec7d7831f2c03283dad3488cc (diff) | |
download | mitmproxy-9c949bd2f852fb5f5a314e6faa45e4869571628c.tar.gz mitmproxy-9c949bd2f852fb5f5a314e6faa45e4869571628c.tar.bz2 mitmproxy-9c949bd2f852fb5f5a314e6faa45e4869571628c.zip |
Merge pull request #3252 from madt1m/session-db
Session - Hybrid DB
Diffstat (limited to 'mitmproxy')
-rw-r--r-- | mitmproxy/addons/session.py | 73 | ||||
-rw-r--r-- | mitmproxy/exceptions.py | 4 | ||||
-rw-r--r-- | mitmproxy/io/sql/session_create.sql | 20 |
3 files changed, 97 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) +); |