From 8de9842092d073852907a0242503e941e2d7c0e3 Mon Sep 17 00:00:00 2001
From: sayali <scharhate@netflix.com>
Date: Tue, 22 Sep 2020 18:22:45 -0700
Subject: [PATCH] Backfill the key_type column: DB Upgrade

---
 lemur/common/utils.py                      |  17 ++++
 lemur/migrations/versions/c301c59688d2_.py | 108 +++++++++++++++++++++
 2 files changed, 125 insertions(+)
 create mode 100644 lemur/migrations/versions/c301c59688d2_.py

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..6bd94cfb
--- /dev/null
+++ b/lemur/migrations/versions/c301c59688d2_.py
@@ -0,0 +1,108 @@
+"""
+
+This upgrade of database updates the key_type information for certificates
+that are either still valid or have expired in last 30 days. For RSA keys,
+the algorithm is determined based on the key length. For rest of the keys,
+the certificate body is parsed to determine the exact key type information.
+
+Each individual change is explicitly committed. The logs are added to file
+named upgrade_logs in current working directory. If faced any issue while
+running this upgrade, there is no harm in re-running the upgrade. Each run
+processes only the keys for which key type information is not yet determined.
+A successful end to end run will end up updating the Alembic Version to 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
+keys 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
+
+log_file = open('upgrade_logs', 'a')
+
+
+def upgrade():
+    log_file.write("\n*** Starting new run ***\n")
+    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 are valid today or expired in 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:
+            log_file.write("Error in processing certificate. ID: %s\n" % cert_id)
+        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)