Add lazy KMS creation in AthenianKMS

Signed-off-by: Lou Marvin Caraig <[email protected]>

Author: se7entyse7en

server/athenian/api/kms.py

@@ -38,19 +38,32 @@ def __init__(self):
                     "%s must be defined, see https://cloud.google.com/kms/docs/reference/rest"
                     % env_names[0],
                 )
-        service_file_inline = os.getenv("GOOGLE_KMS_SERVICE_ACCOUNT_JSON_INLINE")
+
+        self._evars = evars
+
+        service_file_inline = os.getenv(
+            "GOOGLE_KMS_SERVICE_ACCOUNT_JSON_INLINE")
         if service_file_inline is not None:
-            service_file = io.StringIO(service_file_inline)
+            self._service_file = io.StringIO(service_file_inline)
         else:
-            service_file = os.getenv("GOOGLE_KMS_SERVICE_ACCOUNT_JSON")
-        self._kms = KMS(**evars, service_file=service_file)
-        self.log.info("Using Google KMS %(keyproject)s/%(keyring)s/%(keyname)s", evars)
+            self._service_file = os.getenv("GOOGLE_KMS_SERVICE_ACCOUNT_JSON")
+
+        self._kms = None
+        self.log.info(
+            "Using Google KMS %(keyproject)s/%(keyring)s/%(keyname)s", evars)
+
+    async def _get_kms(self) -> KMS:
+        if self._kms is None:
+            self._kms = KMS(**self._evars, service_file=self._service_file)
+
+        return self._kms
 
     async def encrypt(self, plaintext: Union[bytes, str]) -> str:
         """Encrypt text using Google KMS."""
+        kms = await self._get_kms()
         for attempt in range(self.timeout_retries):
             try:
-                return await self._kms.encrypt(encode(plaintext))
+                return await kms.encrypt(encode(plaintext))
             except asyncio.TimeoutError as e:
                 self.log.warning("encrypt attempt %d", attempt + 1)
                 if attempt == self.timeout_retries - 1:
@@ -59,9 +72,10 @@ async def encrypt(self, plaintext: Union[bytes, str]) -> str:
     async def decrypt(self, ciphertext: str) -> bytes:
         """Decrypt text using Google KMS."""
         # we cannot use gcloud.aio.kms.decode because it converts bytes to string with str.decode()
+        kms = await self._get_kms()
         for attempt in range(self.timeout_retries):
             try:
-                payload = await self._kms.decrypt(ciphertext)
+                payload = await kms.decrypt(ciphertext)
             except asyncio.TimeoutError as e:
                 self.log.warning("decrypt attempt %d", attempt + 1)
                 if attempt == self.timeout_retries - 1:
@@ -72,4 +86,5 @@ async def decrypt(self, ciphertext: str) -> bytes:
 
     async def close(self):
         """Close the underlying HTTPS session."""
-        await self._kms.close()
+        if self._kms is not None:
+            await self._kms.close()