From 1ad23f9db0328cb87bc23b5b0bf58657e811a8af Mon Sep 17 00:00:00 2001 From: Joona Hoikkala Date: Tue, 4 Feb 2020 13:13:04 +0200 Subject: [PATCH] Move DBM handling to a context manager --- .../certbot_apache/_internal/prefetch_ocsp.py | 107 +++++++++++------- certbot-apache/tests/ocsp_prefetch_test.py | 47 ++++---- 2 files changed, 93 insertions(+), 61 deletions(-) diff --git a/certbot-apache/certbot_apache/_internal/prefetch_ocsp.py b/certbot-apache/certbot_apache/_internal/prefetch_ocsp.py index 13b0945ae..ad27c6cab 100644 --- a/certbot-apache/certbot_apache/_internal/prefetch_ocsp.py +++ b/certbot-apache/certbot_apache/_internal/prefetch_ocsp.py @@ -18,6 +18,46 @@ from certbot_apache._internal import constants logger = logging.getLogger(__name__) +class DBMHandler(object): + """Context manager to handle DBM file reads and writes""" + + def __init__(self, filename, mode): + self.filename = filename + self.filemode = mode + self.bsddb = False + self.database = None + + def __enter__(self): + """Open the DBM file and return the filehandle""" + + if not os.path.isfile(self.filename + ".db"): + raise errors.PluginError( + "The OCSP stapling cache DBM file wasn't created by Apache.") + try: + import bsddb + self.bsddb = True + cache_path = self.filename + ".db" + try: + self.database = bsddb.hashopen(cache_path, self.filemode) + except Exception: + raise errors.PluginError("Unable to open dbm database file.") + except ImportError: + # Python3 doesn't have bsddb module, so we use dbm.ndbm instead + import dbm + try: + self.database = dbm.ndbm.open(self.filename, self.filemode) # pylint: disable=no-member + except Exception: + # This is raised if a file cannot be found + raise errors.PluginError("Unable to open dbm database file.") + return self.database + + def __exit__(self, *args): + """Close the DBM file""" + if self.bsddb: + self.database.sync() + self.database.close() + + class OCSPPrefetchMixin(object): """OCSPPrefetchMixin implements OCSP response prefetching""" @@ -53,40 +93,6 @@ class OCSPPrefetchMixin(object): "Apache OCSP stapling cache database.") raise errors.NotSupportedError(msg) - def _ocsp_dbm_open(self, filepath): - """Helper method to open an DBM file in a way that depends on the platform - that Certbot is run on. Returns an open database structure.""" - - if not os.path.isfile(filepath+".db"): - raise errors.PluginError( - "The OCSP stapling cache DBM file wasn't created by Apache.") - try: - import bsddb - self._ocsp_dbm_bsddb = True - cache_path = filepath + ".db" - try: - database = bsddb.hashopen(cache_path, 'w') - except Exception: - raise errors.PluginError("Unable to open dbm database file.") - except ImportError: - # Python3 doesn't have bsddb module, so we use dbm.ndbm instead - import dbm - try: - database = dbm.ndbm.open(filepath, 'w') # pylint: disable=no-member - except Exception: - # This is raised if a file cannot be found - raise errors.PluginError("Unable to open dbm database file.") - return database - - def _ocsp_dbm_close(self, database): - """Helper method to sync and close a DBM file, in a way required by the - used dbm implementation.""" - if self._ocsp_dbm_bsddb: - database.sync() - database.close() - else: - database.close() - def _ocsp_refresh_needed(self, pf_obj): """Refreshes OCSP response for a certiifcate if it's due @@ -119,11 +125,9 @@ class OCSPPrefetchMixin(object): if not handler.ocsp_revoked_by_paths(cert_path, chain_path, ocsp_workfile): # Guaranteed good response cache_path = os.path.join(self.config.work_dir, "ocsp", "ocsp_cache") - # dbm.open automatically adds the file extension, it will be - db = self._ocsp_dbm_open(cache_path) cert_sha = apache_util.certid_sha1(cert_path) - db[cert_sha] = self._ocsp_response_dbm(ocsp_workfile) - self._ocsp_dbm_close(db) + # dbm.open automatically adds the file extension + self._write_to_dbm(cache_path, cert_sha, self._ocsp_response_dbm(ocsp_workfile)) else: logger.warning("Encountered an issue while trying to prefetch OCSP " "response for certificate: %s", cert_path) @@ -134,6 +138,33 @@ class OCSPPrefetchMixin(object): # The OCSP workfile did not exist because of an OCSP response fetching error return + def _write_to_dbm(self, filename, key, value): + """Helper method to write an OCSP response cache value to DBM + + :param filename: DBM database filename + :param bytes key: Database key name + :param bytes value: Database entry value + """ + + with DBMHandler(filename, 'w') as db: + db[key] = value + + def _read_dbm(self, filename): + """Helper method for reading the dbm using context manager. + Used for tests. + + :param str filename: DBM database filename + + :returns: Dictionary of database keys and values + :rtype: dict + """ + + ret = dict() + with DBMHandler(filename, 'r') as db: + for k in db.keys(): + ret[k] = db[k] + return ret + def _ocsp_ttl(self, next_update): """Calculates Apache internal TTL for the next OCSP staple update. diff --git a/certbot-apache/tests/ocsp_prefetch_test.py b/certbot-apache/tests/ocsp_prefetch_test.py index 0f8cfa60b..6e2f2b19e 100644 --- a/certbot-apache/tests/ocsp_prefetch_test.py +++ b/certbot-apache/tests/ocsp_prefetch_test.py @@ -198,11 +198,10 @@ class OCSPPrefetchTest(util.ApacheTest): mock_times.return_value = produced_at, this_update, next_update self.call_mocked_py2(self.config.enable_ocsp_prefetch, self.lineage, ["ocspvhost.com"]) - odbm = self.config._ocsp_dbm_open(self.db_path) + odbm = self.config._read_dbm(self.db_path) self.assertEqual(len(odbm.keys()), 1) # The actual response data is prepended by Apache timestamp self.assertTrue(odbm[list(odbm.keys())[0]].endswith(b'MOCKRESPONSE')) - self.config._ocsp_dbm_close(odbm) with mock.patch(ocsp_path, side_effect=ocsp_req_mock) as mock_ocsp: with mock.patch('certbot._internal.ocsp.RevocationChecker.ocsp_times') as mock_times: @@ -240,9 +239,7 @@ class OCSPPrefetchTest(util.ApacheTest): # Make sure that the db file exists open(self.db_fullpath, 'a').close() - odbm = self.call_mocked_py2(self.config._ocsp_dbm_open, self.db_path) - odbm[b'mock_key'] = b'mock_value' - self.config._ocsp_dbm_close(odbm) + self.call_mocked_py2(self.config._write_to_dbm, self.db_path, b'mock_key', b'mock_value') # Mock OCSP prefetch dict to signify that there should be a db self.config._ocsp_prefetch = {"mock": "value"} @@ -250,9 +247,8 @@ class OCSPPrefetchTest(util.ApacheTest): with mock.patch(rel_path, side_effect=ocsp_del_db): self.config.restart() - odbm = self.config._ocsp_dbm_open(self.db_path) + odbm = self.config._read_dbm(self.db_path) self.assertEqual(odbm[b'mock_key'], b'mock_value') - self.config._ocsp_dbm_close(odbm) @mock.patch("certbot_apache._internal.configurator.ApacheConfigurator.config_test") @mock.patch("certbot_apache._internal.configurator.ApacheConfigurator._reload") @@ -306,9 +302,12 @@ class OCSPPrefetchTest(util.ApacheTest): def test_ocsp_prefetch_open_dbm_no_file(self): open(self.db_fullpath, 'a').close() db_not_exists = self.db_path+"nonsense" - self.call_mocked_py2(self.config._ocsp_dbm_open, self.db_path) + self.call_mocked_py2(self.config._write_to_dbm, self.db_path, b'k', b'v') self.assertRaises(errors.PluginError, - self.call_mocked_py2, self.config._ocsp_dbm_open, db_not_exists) + self.call_mocked_py2, + self.config._write_to_dbm, + db_not_exists, + b'k', b'v') def test_ocsp_prefetch_py2_open_file_error(self): open(self.db_fullpath, 'a').close() @@ -316,8 +315,9 @@ class OCSPPrefetchTest(util.ApacheTest): mock_db.hashopen.side_effect = Exception("error") sys.modules["bsddb"] = mock_db self.assertRaises(errors.PluginError, - self.config._ocsp_dbm_open, - self.db_path) + self.config._write_to_dbm, + self.db_path, + b'k', b'v') def test_ocsp_prefetch_py3_open_file_error(self): open(self.db_fullpath, 'a').close() @@ -326,27 +326,28 @@ class OCSPPrefetchTest(util.ApacheTest): sys.modules["dbm"] = mock_db sys.modules["bsddb"] = None self.assertRaises(errors.PluginError, - self.config._ocsp_dbm_open, - self.db_path) + self.config._write_to_dbm, + self.db_path, + b'k', b'v') def test_ocsp_prefetch_open_close_py2_noerror(self): expected_val = b'whatever_value' open(self.db_fullpath, 'a').close() - db = self.call_mocked_py2( - self.config._ocsp_dbm_open, self.db_path) - db[b'key'] = expected_val - self.call_mocked_py2(self.config._ocsp_dbm_close, db) - db2 = self.call_mocked_py2(self.config._ocsp_dbm_open, self.db_path) + self.call_mocked_py2( + self.config._write_to_dbm, self.db_path, + b'key', expected_val + ) + db2 = self.call_mocked_py2(self.config._read_dbm, self.db_path) self.assertEqual(db2[b'key'], expected_val) def test_ocsp_prefetch_open_close_py3_noerror(self): expected_val = b'whatever_value' open(self.db_fullpath, 'a').close() - db = self.call_mocked_py3( - self.config._ocsp_dbm_open, self.db_path) - db[b'key'] = expected_val - self.call_mocked_py3(self.config._ocsp_dbm_close, db) - db2 = self.call_mocked_py3(self.config._ocsp_dbm_open, self.db_path) + self.call_mocked_py3( + self.config._write_to_dbm, self.db_path, + b'key', expected_val + ) + db2 = self.call_mocked_py3(self.config._read_dbm, self.db_path) self.assertEqual(db2[b'key'], expected_val) @mock.patch("certbot_apache._internal.constants.OCSP_APACHE_TTL", 1234)