fix(jans-pycloudlib): split google secrets when payload is larger than 65536 bytes

jans-pycloudlib/jans/pycloudlib/secret/google_secret.py

@@ -11,7 +11,9 @@
 import os
 import typing as _t
 import zlib
+from contextlib import suppress
 from functools import cached_property
+from math import ceil
 
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.exceptions import InvalidTag
@@ -82,6 +84,15 @@ def __init__(self) -> None:
         self.salt = os.urandom(16)
         self.key = self._set_key()
 
+        # max payload size (currently 64K)
+        self.max_payload_size = 65_536
+
+        # max multiparts for given prefix (currenty unused)
+        self.max_multiparts = ceil(1_000_000 / self.max_payload_size)
+
+        # iterable contains multipart secret names
+        self.multiparts: list[str] = []
+
     @cached_property
     def client(self) -> secretmanager.SecretManagerServiceClient:
         """Create the Secret Manager client."""
@@ -125,50 +136,38 @@ def get_all(self) -> dict[str, _t.Any]:
         Returns:
             A mapping of secrets (if any)
         """
-        # Try to get the latest resource name. Used in initialization. If the latest version doesn't exist
-        # its a state where the secret and initial version must be created
-        name = f"projects/{self.project_id}/secrets/{self.google_secret_name}/versions/latest"
+        # get a list of secrets with prefixed name
+        resp = self.client.list_secrets(
+            request={
+                "parent": f"projects/{self.project_id}",
+                "filter": f"name:{self.google_secret_name}",
+            }
+        )
 
-        try:
-            self.client.access_secret_version(request={"name": name})
-        except NotFound:
-            logger.warning("Secret may not exist or have any versions created yet")
-            self.create_secret()
-            self.add_secret_version(safe_value({}))
+        # collect all secret names (if any) for further request
+        names = [f"{scr.name}/versions/{self.version_id}" for scr in resp]
+
+        if not names:
+            return {}
 
-        # Build the resource name of the secret version.
-        name = f"projects/{self.project_id}/secrets/{self.google_secret_name}/versions/{self.version_id}"
         data = {}
+        payload = b""
 
-        try:
-            # Access the secret version.
-            response = self.client.access_secret_version(request={"name": name})
+        for name in names:
+            # the secret with given name may not exist or have any versions created yet
+            with suppress(NotFound):
+                fragment = self.client.access_secret_version(request={"name": name})
+                payload = payload + fragment.payload.data
 
-        except NotFound:
-            logger.warning(
-                "Secret may not exist or have any versions created. Make sure "
-                "CN_GOOGLE_SECRET_VERSION_ID and CN_GOOGLE_SECRET_NAME_PREFIX "
-                "environment variables are set correctly. In Google secrets manager, "
-                "a secret with the name jans-secret would have "
-                "CN_GOOGLE_SECRET_NAME_PREFIX set to jans."
-            )
+        if not payload:
+            return {}
 
-        else:
-            # logger.info(f"Secret {self.google_secret_name} has been found. Accessing version {self.version_id}.")
-            # backward-compat checks
-            try:
-                # previously data is compressed using zlib
-                payload = zlib.decompress(response.payload.data).decode("UTF-8")
-                logger.warning("Decompressed legacy data.")
-            except zlib.error:
-                payload = lzma.decompress(response.payload.data).decode("UTF-8")
-
-            try:
-                # previously data is double-encrypted
-                data = json.loads(self._decrypt(payload))
-                logger.warning("Loaded legacy data.")
-            except binascii.Error:
-                data = json.loads(payload)
+        try:
+            data = self._maybe_legacy_payload(payload)
+        except lzma.LZMAError:
+            data = json.loads(payload)
+
+        # decoded payload
         return data
 
     def get(self, key: str, default: _t.Any = "") -> _t.Any:
@@ -196,13 +195,10 @@ def set(self, key: str, value: _t.Any) -> bool:
         """
         all_ = self.get_all()
         all_[key] = safe_value(value)
+        logger.info(f"Adding key {key}.")
 
-        self.create_secret()
-
-        logger.info(f'Adding key {key} to google secret manager')
-        logger.info(f'Size of secret payload : {sys.getsizeof(safe_value(all_))} bytes')
-        secret_version_bool = self.add_secret_version(safe_value(all_))
-        return secret_version_bool
+        payload = safe_value(all_)
+        return self._add_secret_version_multipart(payload)
 
     def set_all(self, data: dict[str, _t.Any]) -> bool:
         """Push a full dictionary to secrets.
@@ -220,70 +216,111 @@ def set_all(self, data: dict[str, _t.Any]) -> bool:
         for k, v in data.items():
             all_[k] = safe_value(v)
 
-        self.create_secret()
+        payload = safe_value(all_)
+        return self._add_secret_version_multipart(payload)
 
-        logger.info(f'Size of secret payload : {sys.getsizeof(safe_value(all_))} bytes')
-        secret_version_bool = self.add_secret_version(safe_value(all_))
-        return secret_version_bool
-
-    def create_secret(self) -> None:
-        """Create a new secret with the given name.
-
-        A secret is a logical wrapper around a collection of secret versions.
-        Secret versions hold the actual secret material.
-        """
-        # Build the resource name of the parent project.
-        parent = f"projects/{self.project_id}"
+    def delete(self) -> None:
+        """Delete the secret with the given name and all of its versions."""
+        # Build the resource name of the secret.
+        name = self.client.secret_path(self.project_id, self.google_secret_name)
 
         try:
-            # Create the secret.
-            response = self.client.create_secret(
-                request={
-                    "parent": parent,
-                    "secret_id": self.google_secret_name,
-                    "secret": {"replication": {"automatic": {}}},
-                }
-            )
-            logger.info(f"Created secret: {response.name}")
-        except AlreadyExists:
-            logger.warning(f'Secret {self.google_secret_name} already exists. A new version will be created.')
+            # Delete the secret.
+            self.client.delete_secret(request={"name": name})
+        except NotFound:
+            logger.warning(f'Secret {self.google_secret_name} does not exist in the secret manager.')
 
-    def add_secret_version(self, payload: _t.AnyStr) -> bool:
+    def _add_secret_version_multipart(self, payload: _t.AnyStr) -> bool:
         """Add a new secret version to the given secret with the provided payload.
 
         Args:
-            payload: encrypted payload
+            payload: secret's payload
         """
-        # Build the resource name of the parent secret.
-        parent = self.client.secret_path(self.project_id, self.google_secret_name)
-
         if isinstance(payload, str):
             # Convert the string payload into a bytes. This step can be omitted if you
             # pass in bytes instead of a str for the payload argument.
             payload_bytes = payload.encode("UTF-8")
         else:
             payload_bytes = payload
 
-        # compress the payload
-        payload_bytes = lzma.compress(payload_bytes)
+        data_length = sys.getsizeof(payload_bytes)
+        parts = ceil(data_length / self.max_payload_size)
 
-        logger.info(f'Size of final compressed secret payload : {sys.getsizeof(payload_bytes)} bytes')
+        if parts > 1:
+            logger.warning(
+                f"The secret payload size is {data_length} bytes and is exceeding max. size of {self.max_payload_size} bytes. "
+                f"It will be splitted into {parts} parts."
+            )
 
-        # Add the secret version.
-        response = self.client.add_secret_version(
-            request={"parent": parent, "payload": {"data": payload_bytes}}
-        )
+        for part in range(0, parts):
+            name = self._prepare_secret_multipart(part)
 
-        logger.info(f"Added secret version: {response.name}")
-        return bool(response)
+            start_bytes = part * self.max_payload_size
+            stop_bytes = (part + 1) * self.max_payload_size
+            fragment = payload_bytes[start_bytes:stop_bytes]
 
-    def delete(self) -> None:
-        """Delete the secret with the given name and all of its versions."""
-        # Build the resource name of the secret.
-        name = self.client.secret_path(self.project_id, self.google_secret_name)
+            # Build the resource name of the parent secret.
+            parent = self.client.secret_path(self.project_id, name)
+
+            # Add the secret version.
+            response = self.client.add_secret_version(
+                request={"parent": parent, "payload": {"data": fragment}}
+            )
+            logger.info(f"Added secret version: {response.name}")
+        return True
+
+    def _prepare_secret_multipart(self, part: int) -> str:
+        """Create a new secret with the given name.
+
+        A secret is a logical wrapper around a collection of secret versions.
+        Secret versions hold the actual secret material.
+
+        Args:
+            part: multipart number.
+
+        Returns:
+            Newly created secret's name.
+        """
+        name = self.google_secret_name
+
+        if part > 0:
+            name = f"{self.google_secret_name}-{part}"
 
+        if name in self.multiparts:
+            return name
+
+        # Build the resource name of the parent project.
+        parent = f"projects/{self.project_id}"
+
+        # Secret with given name may already exists
+        with suppress(AlreadyExists):
+            # Create the secret.
+            response = self.client.create_secret(
+                request={
+                    "parent": parent,
+                    "secret_id": name,
+                    "secret": {
+                        "replication": {"automatic": {}},
+                        "labels": {"multipart_enabled": "true"},
+                    },
+                }
+            )
+            logger.info(f"Created secret: {response.name}")
+            self.multiparts.append(name)
+        return name
+
+    def _maybe_legacy_payload(self, payload: bytes) -> dict[str, _t.Any]:
         try:
-            # Delete the secret.
-            self.client.delete_secret(request={"name": name})
-        except NotFound:
-            logger.warning(f'Secret {self.google_secret_name} does not exist in the secret manager.')
+            # previously data is compressed using zlib
+            payload_str = zlib.decompress(payload).decode("UTF-8")
+            logger.warning("Decompressed legacy data.")
+        except zlib.error:
+            payload_str = lzma.decompress(payload).decode("UTF-8")
+
+        try:
+            # previously data is double-encrypted
+            data: dict[str, _t.Any] = json.loads(self._decrypt(payload_str))
+            logger.warning("Loaded legacy data.")
+        except binascii.Error:
+            data = json.loads(payload_str)
+        return data