Merge pull request #3150 from charhate/key_type_column

Backfill the key_type column

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.

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)

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