add key provisioning for postgres

VAuthenticator · Oct 27, 2024 · 3e40c90 · 3e40c90
1 parent c1cf4b8
commit 3e40c90
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 12 deletions.
diff --git a/local-environment/local-initializer/database_setup.py b/local-environment/local-initializer/database_setup.py
@@ -101,12 +101,10 @@ def pass_encoded(password):
                                 password="postgres",
                                 port="5432")
         cur = conn.cursor()
-
         create_schema()
 
         store_roles()
         store_account()
-
         store_client_applications()
 
         cur.close()

diff --git a/local-environment/local-initializer/key_setup.py b/local-environment/local-initializer/key_setup.py
@@ -3,6 +3,7 @@
 import os
 import sys
 import uuid
+import psycopg2
 from dotenv import load_dotenv
 
 load_dotenv(dotenv_path="env")
@@ -32,21 +33,40 @@ def kmsClient():
 
 
 def store_key(key_table_name, master_key):
-    table = dynamodb.Table(key_table_name)
     key_pair = kms_client.generate_data_key_pair(KeyId=master_key, KeyPairSpec='RSA_2048')
-    table.put_item(Item={
-        "master_key_id": key_pair["KeyId"].split("/")[1],
-        "key_id": str(uuid.uuid4()),
-        "encrypted_private_key": base64.b64encode(key_pair["PrivateKeyCiphertextBlob"]).decode(),
-        "public_key": base64.b64encode(key_pair["PublicKey"]).decode(),
-        "key_purpose": "SIGNATURE",
-        "key_type": "ASYMMETRIC",
-        "enabled": True
-    })
+    if os.getenv("experimental_database_persistence"):
+        cur.execute(
+            f" INSERT INTO KEYS(master_key_id, key_id, key_purpose, key_type, encrypted_private_key, public_key, enabled, key_expiration_date_timestamp)  VALUES('{key_pair["KeyId"].split("/")[1]}', '{str(uuid.uuid4()),}', 'SIGNATURE', 'ASYMMETRIC', '{base64.b64encode(key_pair["PrivateKeyCiphertextBlob"]).decode()}','{base64.b64encode(key_pair["PublicKey"]).decode()}' True, 0)")
+        conn.commit()
+
+    else:
+        table = dynamodb.Table(key_table_name)
+        table.put_item(Item={
+            "master_key_id": key_pair["KeyId"].split("/")[1],
+            "key_id": str(uuid.uuid4()),
+            "encrypted_private_key": base64.b64encode(key_pair["PrivateKeyCiphertextBlob"]).decode(),
+            "public_key": base64.b64encode(key_pair["PublicKey"]).decode(),
+            "key_purpose": "SIGNATURE",
+            "key_type": "ASYMMETRIC",
+            "enabled": True
+        })
 
 
 if __name__ == '__main__':
     input_master_key = sys.argv[1]
     input_key_table_name = f'VAuthenticator_Signature_Keys{sys.argv[2]}'
 
     store_key(input_key_table_name, input_master_key)
+
+    if os.getenv("experimental_database_persistence"):
+        conn = psycopg2.connect(database="postgres",
+                                host=database_host,
+                                user="postgres",
+                                password="postgres",
+                                port="5432")
+        cur = conn.cursor()
+
+        store_key(input_key_table_name, input_master_key)
+
+        cur.close()
+        conn.close()