115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
"""
|
|
|
|
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, only RSA and ECC
|
|
certificates are supported by Lemur. 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(
|
|
"update certificates set key_type='RSA{0}' where bits={0} and not_after > CURRENT_DATE - 31 and key_type is null".format(bits)
|
|
)
|
|
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)
|