129 lines
4.0 KiB
Python
129 lines
4.0 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 configured in LOG_UPGRADE_FILE or, by default,
|
|
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
|
|
import time
|
|
import datetime
|
|
from flask import current_app
|
|
|
|
from logging import Formatter, FileHandler, getLogger
|
|
|
|
from lemur.common import utils
|
|
|
|
log = getLogger(__name__)
|
|
handler = FileHandler(current_app.config.get("LOG_UPGRADE_FILE", "db_upgrade.log"))
|
|
handler.setFormatter(
|
|
Formatter(
|
|
"%(asctime)s %(levelname)s: %(message)s " "[in %(pathname)s:%(lineno)d]"
|
|
)
|
|
)
|
|
handler.setLevel(current_app.config.get("LOG_LEVEL", "DEBUG"))
|
|
log.setLevel(current_app.config.get("LOG_LEVEL", "DEBUG"))
|
|
log.addHandler(handler)
|
|
|
|
|
|
def upgrade():
|
|
log.info("\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.info("--- Total %s seconds ---\n" % (time.time() - start_time))
|
|
|
|
|
|
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)
|
|
commit()
|
|
|
|
|
|
"""
|
|
Helper methods performing updates for RSA and rest of the keys
|
|
"""
|
|
|
|
|
|
def update_key_type_rsa(bits):
|
|
log.info("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.info("Query: %s\n" % stmt)
|
|
|
|
start_time = time.time()
|
|
op.execute(stmt)
|
|
commit()
|
|
|
|
log.info("--- %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 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.error("Error in processing certificate - ID: %s Error: %s \n" % (cert_id, str(e)))
|
|
else:
|
|
log.info("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.info("--- %s seconds ---\n" % (time.time() - start_time))
|
|
|
|
|
|
def commit():
|
|
stmt = text("commit")
|
|
op.execute(stmt)
|