Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2025-02-02 00:05:27 +00:00
Code Issues Projects Releases Wiki Activity

session: temporary DB is now stored in temporary dir

Browse Source
This commit is contained in:
madt1m 2018-07-24 15:57:11 +02:00
parent 68eb07b668
commit 8c7793b91a
2 changed files with 12 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
mitmproxy/addons/session.py
View File

@ -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):

8
test/mitmproxy/addons/test_session.py
View File

@ -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'