New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Session - Hybrid DB #3252
Merged
Merged
Session - Hybrid DB #3252
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit
Hold shift + click to select a range
4cebfba
session: Implemented SessionDB Schema
04b9555
session: load/create wrapper methods for SessionDB
1a58442
exceptions: new SessionLoad exception for load errors
fad8e7c
tests: SessionDB fully tested
3b5cdf7
test_linting: removed some unused local variables
suidpit 68eb07b
session: modified schema. Now SessionDB uses tempfile module for temp…
suidpit 8c7793b
session: temporary DB is now stored in temporary dir
suidpit File filter
Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
); | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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): | ||
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) | ||
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() | ||
session.SessionDB(path) | ||
con = sqlite3.connect(path) | ||
with con: | ||
cur = con.cursor() | ||
cur.execute('SELECT * FROM FLOW;') | ||
rows = cur.fetchall() | ||
assert len(rows) == 1 |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I have some comments on the schema design here, some trivial/aesthetic, some substantive. I'm going to put them all down here for discussion, and let's chat about this on Slack when we're both on, or here if you prefer.
AID
,FID
, etc. There are two schools of thought here, and I'm OK with both. The first is to just use "id" for the primary key in all tables. In practice this is fine, because they will be referred to along with the table name in almost all circumstances. The second is to prefix the ID with the table name, soflow_id
orflowid
. Pick whichever is more aesthetic to you. :)