diff --git a/lemur/common/utils.py b/lemur/common/utils.py index 01cc64ae..283d1eec 100644 --- a/lemur/common/utils.py +++ b/lemur/common/utils.py @@ -71,6 +71,23 @@ def parse_private_key(private_key): ) +def get_key_type_from_certificate(body): + """ + + Helper function to determine key type by pasrding given PEM certificate + + :param body: PEM string + :return: Key type string + """ + parsed_cert = parse_certificate(body) + if isinstance(parsed_cert.public_key(), rsa.RSAPublicKey): + return "RSA{key_size}".format( + key_size=parsed_cert.public_key().key_size + ) + elif isinstance(parsed_cert.public_key(), ec.EllipticCurvePublicKey): + return get_key_type_from_ec_curve(parsed_cert.public_key().curve.name) + + def split_pem(data): """ Split a string of several PEM payloads to a list of strings. diff --git a/lemur/migrations/versions/c301c59688d2_.py b/lemur/migrations/versions/c301c59688d2_.py new file mode 100644 index 00000000..3b0a86f7 --- /dev/null +++ b/lemur/migrations/versions/c301c59688d2_.py @@ -0,0 +1,114 @@ +""" + +This database upgrade updates the key_type information for either +still valid or expired certificates in the last 30 days. For RSA +keys, the algorithm is determined based on the key length. For +the rest of the keys, the certificate body is parsed to determine +the exact key_type information. + +Each individual DB change is explicitly committed, and the respective +log is added to a file named db_upgrade.log in the current working +directory. Any error encountered while parsing a certificate will +also be logged along with the certificate ID. If faced with any issue +while running this upgrade, there is no harm in re-running the upgrade. +Each run processes only rows for which key_type information is not yet +determined. + +A successful complete run will end up updating the Alembic Version to +the new Revision ID c301c59688d2. Currently, Lemur supports only RSA +and ECC certificates. This could be a long-running job depending upon +the number of DB entries it may process. + +Revision ID: c301c59688d2 +Revises: 434c29e40511 +Create Date: 2020-09-21 14:28:50.757998 + +""" + +# revision identifiers, used by Alembic. +revision = 'c301c59688d2' +down_revision = '434c29e40511' + +from alembic import op +from sqlalchemy.sql import text +from lemur.common import utils +import time +import datetime + +log_file = open('db_upgrade.log', 'a') + + +def upgrade(): + log_file.write("\n*** Starting new run(%s) ***\n" % datetime.datetime.now()) + start_time = time.time() + + # Update RSA keys using the key length information + update_key_type_rsa(1024) + update_key_type_rsa(2048) + update_key_type_rsa(4096) + + # Process remaining certificates. Though below method does not make any assumptions, most of the remaining ones should be ECC certs. + update_key_type() + + log_file.write("--- Total %s seconds ---\n" % (time.time() - start_time)) + log_file.close() + + +def downgrade(): + # Change key type column back to null + # Going back 32 days instead of 31 to make sure no certificates are skipped + stmt = text( + "update certificates set key_type=null where not_after > CURRENT_DATE - 32" + ) + op.execute(stmt) + + +""" + Helper methods performing updates for RSA and rest of the keys +""" + + +def update_key_type_rsa(bits): + log_file.write("Processing certificate with key type RSA %s\n" % bits) + + stmt = text( + f"update certificates set key_type='RSA{bits}' where bits={bits} and not_after > CURRENT_DATE - 31 and key_type is null" + ) + log_file.write("Query: %s\n" % stmt) + + start_time = time.time() + op.execute(stmt) + commit() + + log_file.write("--- %s seconds ---\n" % (time.time() - start_time)) + + +def update_key_type(): + conn = op.get_bind() + start_time = time.time() + + # Loop through all certificates that are valid today or expired in the last 30 days. + for cert_id, body in conn.execute( + text( + "select id, body from certificates where bits < 1024 and not_after > CURRENT_DATE - 31 and key_type is null") + ): + try: + cert_key_type = utils.get_key_type_from_certificate(body) + except ValueError as e: + log_file.write("Error in processing certificate - ID: %s Error: %s \n" % (cert_id, str(e))) + else: + log_file.write("Processing certificate - ID: %s key_type: %s\n" % (cert_id, cert_key_type)) + stmt = text( + "update certificates set key_type=:key_type where id=:id" + ) + stmt = stmt.bindparams(key_type=cert_key_type, id=cert_id) + op.execute(stmt) + + commit() + + log_file.write("--- %s seconds ---\n" % (time.time() - start_time)) + + +def commit(): + stmt = text("commit") + op.execute(stmt) diff --git a/lemur/tests/test_utils.py b/lemur/tests/test_utils.py index 1dac39bb..162e53b0 100644 --- a/lemur/tests/test_utils.py +++ b/lemur/tests/test_utils.py @@ -2,11 +2,13 @@ import pytest from lemur.tests.vectors import ( SAN_CERT, + SAN_CERT_STR, INTERMEDIATE_CERT, ROOTCA_CERT, EC_CERT_EXAMPLE, ECDSA_PRIME256V1_CERT, ECDSA_SECP384r1_CERT, + ECDSA_SECP384r1_CERT_STR, DSA_CERT, ) @@ -106,3 +108,9 @@ def test_is_selfsigned(selfsigned_cert): # unsupported algorithm (DSA) with pytest.raises(Exception): is_selfsigned(DSA_CERT) + + +def test_get_key_type_from_certificate(): + from lemur.common.utils import get_key_type_from_certificate + assert (get_key_type_from_certificate(SAN_CERT_STR) == "RSA2048") + assert (get_key_type_from_certificate(ECDSA_SECP384r1_CERT_STR) == "ECCSECP384R1")