diff --git a/lib/core/testing.py b/lib/core/testing.py index 4b74b909b94..2923314c864 100644 --- a/lib/core/testing.py +++ b/lib/core/testing.py @@ -5,6 +5,7 @@ See the file 'LICENSE' for copying permission """ +import contextlib import doctest import json import logging @@ -206,9 +207,11 @@ def _thread(): os.close(handle) cleanups.append(database) - with sqlite3.connect(database) as conn: + # Note: sqlite3's own context manager only commits (it does not close the connection) + with contextlib.closing(sqlite3.connect(database)) as conn: c = conn.cursor() c.executescript(vulnserver.SCHEMA) + conn.commit() handle, request = tempfile.mkstemp(suffix=".req") os.close(handle) diff --git a/lib/utils/sqlalchemy.py b/lib/utils/sqlalchemy.py index 34c377c5a96..9ca71c2e0f1 100644 --- a/lib/utils/sqlalchemy.py +++ b/lib/utils/sqlalchemy.py @@ -50,6 +50,7 @@ def __init__(self, dialect=None): self.dialect = dialect self.address = conf.direct + self.engine = None if conf.dbmsUser: self.address = self.address.replace("'%s':" % conf.dbmsUser, "%s:" % _urllib.parse.quote(conf.dbmsUser)) @@ -84,6 +85,7 @@ def connect(self): else: engine = _sqlalchemy.create_engine(self.address, connect_args={}) + self.engine = engine self.connector = engine.connect() except (TypeError, ValueError) as ex: if "_get_server_version_info" in traceback.format_exc(): @@ -110,6 +112,19 @@ def connect(self): else: raise SqlmapMissingDependence("SQLAlchemy not available (e.g. 'pip%s install SQLAlchemy')" % ('3' if six.PY3 else "")) + def close(self): + try: + GenericConnector.close(self) + finally: + # Note: closing the connection only returns it to the engine's pool (the underlying DBAPI + # connection, e.g. sqlite3, stays open) - disposing the engine is what actually releases it + if self.engine is not None: + try: + self.engine.dispose() + except Exception as ex: + logger.debug(getSafeExString(ex)) + self.engine = None + def fetchall(self): try: retVal = [] diff --git a/tests/test_multibit.py b/tests/test_multibit.py index c69d6b6f544..b51cf5c41e7 100644 --- a/tests/test_multibit.py +++ b/tests/test_multibit.py @@ -140,8 +140,14 @@ def tearDown(self): multibit._page = self._savedPage multibit.singleTimeWarnMessage, multibit.singleTimeLogMessage = _SAID - def _attempt(self, secret, length=None, **kwargs): + def _target(self, secret, **kwargs): target = _Target(secret, **kwargs) + self.addCleanup(target.db.close) + + return target + + def _attempt(self, secret, length=None, **kwargs): + target = self._target(secret, **kwargs) multibit._page = target.page return self._read(target, secret, length), target @@ -231,7 +237,7 @@ def test_first_value_that_cannot_settle_the_cross_check_keeps_the_channel(self): # non-ASCII one, and one whose codepoint has an ASCII low byte (U+0141 -> 0x41 'A') all settle # nothing, and none of them may condemn a channel that calibrated cleanly for first in ("", u"\u017eaba pod vodom", u"\u0141ukasz Nowak here"): - target = _Target(first) + target = self._target(first) self.assertIsNone(self._read(target, first, len(first) or None)) good = target.secret("plain ascii value that must still come back") @@ -242,7 +248,7 @@ def test_null_values_do_not_condemn_the_channel(self): # what the channel did - i.e. exactly the way a channel dropping rows fails it. That is the # data's doing, and a handful of NULLs in one dump must not retire a working channel first = "first value is perfectly readable ascii" - target = _Target(first) + target = self._target(first) self.assertEqual(self._read(target, first, len(first)), first) for _ in range(MULTIBIT_MAX_FAILURES + 1): @@ -253,7 +259,7 @@ def test_null_values_do_not_condemn_the_channel(self): self.assertEqual(self._read(target, good, len(good)), good) def test_unreadable_value_mid_run_keeps_the_channel(self): - target = _Target("first value is perfectly readable ascii") + target = self._target("first value is perfectly readable ascii") self.assertEqual(self._read(target, "first value is perfectly readable ascii", 39), "first value is perfectly readable ascii") bad = target.secret(u"\u0141ukasz") # reads back as "Aukasz" - right length, wrong char @@ -295,7 +301,7 @@ def test_nothing_runs_without_the_switch(self): def test_resumed_partial_value_is_left_to_bisection(self): # a resumed value starts mid-string, and multi-bit reads from the first character: it would # re-read what is already cached and then fail its own length read-back - target = _Target("resuming this one is bisection's job") + target = self._target("resuming this one is bisection's job") multibit._page = target.page self.assertIsNone(multibit.attempt("SELECT s FROM v", lambda idx: 'x', 35, None, True)) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 768c1241119..a531c5219f7 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -58,6 +58,7 @@ def test_absolute_sqlite_path_opens_correct_file(self): sa = SQLAlchemy(dialect="sqlite") sa.connect() + self.addCleanup(sa.close) # the reformatted URL must resolve to the exact absolute file (not '//...path') self.assertEqual(_sa.engine.url.make_url(sa.address).database, os.path.abspath(dbfile))