upstream/certbot
1
0
Fork 0
mirror of https://github.com/certbot/certbot.git synced 2026-06-06 23:32:06 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Move DBM handling to a context manager

Browse source
This commit is contained in:
Joona Hoikkala 2020-02-04 13:13:04 +02:00
parent a5d739f8ff
commit 1ad23f9db0
No known key found for this signature in database
GPG key ID: D5AA86BBF9B29A5C
2 changed files with 93 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

107
certbot-apache/certbot_apache/_internal/prefetch_ocsp.py
View file

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

47
certbot-apache/tests/ocsp_prefetch_test.py
View file

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