Merge pull request #3150 from charhate/key_type_column

Backfill the key_type column
This commit is contained in:
charhate 2020-09-23 12:11:20 -07:00 committed by GitHub
commit 59bfcec808
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 139 additions and 0 deletions

View File

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

View File

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

View File

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