Google KMS support and audit_callback

Author: rog555

CHANGELOG.md

@@ -1,6 +1,9 @@
 Changelog
 =========
 
+## [v0.0.24] - 2024-05-15
+- Google KMS support and audit_callback
+
 ## [v0.0.23] - 2024-04-26
 - Google Firestore support
 

README.md

@@ -52,6 +52,7 @@ For optional [client side](#client-side-encryption) field level envelope encrypt
 ```
 pip install 'abnosql[aws-kms]'
 pip install 'abnosql[azure-kms]'
+pip install 'abnosql[gcp-kms]'
 ```
 
 By default, abnosql does not include database dependencies.  This is to facilitate packaging
@@ -195,17 +196,27 @@ This works for AWS DyanmoDB & Firestore, however Azure Cosmos has a limitation w
 
 ## Audit
 
-`put_item()` and `put_items()` take an optional `audit_user` kwarg.  If supplied, absnosql will add the following to the item:
+Table config attribute `audit_user` will add the following to the item being written to database:
 
 - `createdBy` - value of `audit_user`, added if does not exist in item supplied to put_item()
 - `createdDate` - UTC ISO timestamp string, added if does not exist
 - `modifiedBy` - value of `audit_user` always added
 - `modifiedDate` - UTC ISO timestamp string, always added
 
-You can also specify `audit_user` as config attribute to table.  If you prefer snake_case over CamelCase, you can set env var `ABNOSQL_CAMELCASE` = `FALSE`
+If snake_case over CamelCase is preferred, set env var `ABNOSQL_CAMELCASE` = `FALSE`
 
 NOTE: created* will only be added if `update` is not True in a `put_item()` operation
 
+Table config attribute `audit_callback` with value as a function callback can be used to hook into additional audit stores.
+
+Callback function must accept the following positional args:
+
+- `table_name` - table name
+- `dt_iso` - ISO date timestamp
+- `operation` - `create`, `update`, `get` or `delete`
+- `key` - key of item serialised in <key>=<value>; format
+- `audit_user` - user performing the operation
+
 ## Change Feed / Stream Support
 
 **AWS DynamoDB** [Streams](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Streams.html) allow Lambda functions to be triggered upon create, update and delete table operations.  The event sent to the lambda (see [aws docs](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Streams.Lambda.Tutorial2.html)) contains `eventName` and `eventSourceARN`, where:
@@ -240,14 +251,15 @@ To write an Azure Function / AWS Lambda that is able to process both DynamoDB an
 
 ## Client Side Encryption
 
-If configured in table config with `kms` attribute, abnosql will perform client side encryption using AWS KMS or Azure KeyVault
+If configured in table config with `kms` attribute, abnosql will perform client side encryption using AWS KMS, Azure KeyVault or Google KMS
 
 Each attribute value defined in the config is encrypted with a 256-bit AES-GCM data key generated for each attribute value:
 
 - `aws` uses [AWS Encryption SDK for Python](https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/python.html)
-- `azure` uses [python cryptography](https://cryptography.io/en/latest/hazmat/primitives/aead/#cryptography.hazmat.primitives.ciphers.aead.AESGCM.generate_key) to generate AES-GCM data key, encrypt the attribute value and then uses an RSA CMK in Azure Keyvault to wrap/unwrap (envelope encryption) the AES-GCM data key.  The module uses the [azure-keyvaults-keys](https://learn.microsoft.com/en-us/python/api/overview/azure/keyvault-keys-readme?view=azure-python) python SDK for wrap/unrap functionality of the generated data key (Azure doesnt support generate data key as AWS does)
+- `azure` uses [python cryptography](https://cryptography.io/en/latest/hazmat/primitives/aead/#cryptography.hazmat.primitives.ciphers.aead.AESGCM.generate_key) to generate AES-GCM data key, encrypt the attribute value and then uses an RSA CMK in Azure Keyvault to wrap/unwrap (envelope encryption) the AES-GCM data key.  The plugin uses the [azure-keyvault-keys](https://learn.microsoft.com/en-us/python/api/overview/azure/keyvault-keys-readme?view=azure-python) python SDK for wrap/unrap functionality of the generated data key (Azure doesnt support generate data key as AWS does - see also [tink issue](https://github.com/tink-crypto/tink/issues/158#issuecomment-1382589658))
+- `gcp` uses [Google Tink](https://developers.google.com/tink/client-side-encryption)
 
-Both providers use a [256-bit AES-GCM](https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/supported-algorithms.html) generated data key with AAD/encryption context (Azure provider uses a 96-nonce).  AES-GCM is an Authenticated symmetric encryption scheme used by both AWS and Azure (and [Hashicorp Vault](https://developer.hashicorp.com/vault/docs/secrets/transit#aes256-gcm96))
+All providers use a [256-bit AES-GCM](https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/supported-algorithms.html) generated data key with AAD/encryption context (Azure provider uses a 96-nonce).  AES-GCM is an Authenticated symmetric encryption scheme used by AWS, Azure & Google (and [Hashicorp Vault](https://developer.hashicorp.com/vault/docs/secrets/transit#aes256-gcm96))
 
 See also [AWS Encryption Best Practices](https://docs.aws.amazon.com/prescriptive-guidance/latest/encryption-best-practices/welcome.html)
 
@@ -256,15 +268,23 @@ Example config:
 ```
 {
     'kms': {
+        # Azure example
         'key_ids': ['https://foo.vault.azure.net/keys/bar/45e36a1024a04062bd489db0d9004d09'],
+
+        # AWS example
+        # 'key_ids': ['arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab'],
+
+        # Google Example
+        # 'key_ids': ['gcp-kms://projects/p1/locations/global/keyRings/kr1/cryptoKeys/ck1'],
+
         'key_attrs': ['hk', 'rk'],
         'attrs': ['obj', 'str']
     }
 }
 ```
 
 Where:
-- `key_ids`: list of AWS KMS Key ARNs or Azure KeyVault identifier (URL to RSA CMK).  This is picked up via `ABNOSQL_KMS_KEYS` env var as a comma separated list (*NOTE: env var recommended to avoid provider specific code*)
+- `key_ids`: list of AWS KMS Key ARNs, Azure KeyVault identifier (URL to RSA CMK) or Google KMS URI.  This is picked up via `ABNOSQL_KMS_KEYS` env var as a comma separated list (*NOTE: env var recommended to avoid provider specific code*)
 - `key_attrs`: list of key attributes in the item from which the AAD/encryption context is set.  Taken from `ABNOSQL_KEY_ATTRS` env var or table `key_attrs` if defined there
 - `attrs`: list of attributes keys to encrypt
 - `key_bytes`: optional for azure, use your own AESGCM key if specified, otherwise generate one
@@ -383,6 +403,7 @@ abnosql uses pluggy and registers in the `abnosql.table` namespace
 The following hooks are available
 
 - `set_config` - set config
+- `get_item_pre`
 - `get_item_post` - called after `get_item()`, can return modified data
 - `put_item_pre`
 - `put_item_post`
@@ -434,20 +455,25 @@ More examples in [tests/test_cosmos.py](./tests/test_cosmos.py)
 
 ## Google Firestore
 
-Use [python-mock-firestore](https://github.com/mdowds/python-mock-firestore) and pass `MockFirestore()` to table config as `client` attribute
+Use [python-mock-firestore](https://github.com/mdowds/python-mock-firestore) and pass `MockFirestore()` to table config as `client` attribute, or patch get_client()
 
 Example:
 
 ```
+from unittest.mock import patch
 from mockfirestore import MockFirestore
+from abnosql.plugins.table.firestore import Table as FirestoreTable
 
 
+@patch.object(FirestoreTable, 'get_client', MockFirestore)
 def test_something():
-    tb = table('mytable', {'client': MockFirestore()})
+    tb = table('mytable', {})
     item = tb.get_item(foo='bar')
 
 ```
 
+More examples in [tests/test_firestore.py](./tests/test_firestore.py)
+
 # CLI
 
 Small abnosql CLI installed with few of the commands above
@@ -489,7 +515,7 @@ p2           p2.2      5  {'foo': 'bar', 'num': 5, 'list': [1, 2, 3]}  [1, 2, 3]
 - [x] client side encryption
 - [x] test pagination & exception handling
 - [x] [Google Firestore](https://cloud.google.com/python/docs/reference/firestore/latest) support, ideally in the core library (though could be added outside via use of the plugin system).  Would need something like [FireSQL](https://firebaseopensource.com/projects/jsayol/firesql/) implemented for python, maybe via sqlglot
-- [ ] [Google Vault](https://cloud.google.com/python/docs/reference/cloudkms/latest/) KMS support
+- [x] [Google Vault](https://cloud.google.com/python/docs/reference/cloudkms/latest/) KMS support
 - [ ] [Hashicorp Vault](https://github.com/hashicorp/vault-examples/blob/main/examples/_quick-start/python/example.py) KMS support
 - [ ] Simple caching (maybe) using globals (used for AWS Lambda / Azure Functions)
 - [ ] PostgresSQL support using JSONB column (see [here](https://medium.com/geekculture/json-and-postgresql-using-json-to-mimic-nosqls-storage-benefits-1564c69f61fc) for example).  Would be nice to avoid an ORM and having to define a model for each table...

abnosql/plugins/kms/aws.py

@@ -32,14 +32,14 @@ def wrapper(*args, **kwargs):
             except ClientError as e:
                 code = e.response['Error']['Code']
                 if raise_not_found and code in ['ResourceNotFoundException']:
-                    raise ex.NotFoundException(e) from None
+                    raise ex.NotFoundException(detail=e) from None
                 elif code == 'UnrecognizedClientException':
-                    raise ex.ConfigException(e) from None
-                raise ex.ValidationException(e) from None
+                    raise ex.ConfigException(detail=e) from None
+                raise ex.ValidationException(detail=e) from None
             except NoCredentialsError as e:
-                raise ex.ConfigException(e) from None
+                raise ex.ConfigException(detail=e) from None
             except Exception as e:
-                raise ex.PluginException(e)
+                raise ex.PluginException(detail=e)
         return wrapper
     return decorator
 
@@ -51,6 +51,7 @@ def __init__(
     ) -> None:
         self.pm = pm
         self.config = config or {}
+        self.provider = 'aws'
         self.session = self.config.get(
             'session', Session()
         )

abnosql/plugins/kms/azure.py

@@ -32,10 +32,10 @@ def wrapper(*args, **kwargs):
                 return func(*args, **kwargs)
             except azex.ResourceNotFoundError as e:
                 if raise_not_found:
-                    raise ex.NotFoundException(e.message) from None
+                    raise ex.NotFoundException(detail=e.message) from None
                 return None
             except Exception as e:
-                raise ex.PluginException(e)
+                raise ex.PluginException(detail=e)
         return wrapper
     return decorator
 
@@ -47,6 +47,7 @@ def __init__(
     ) -> None:
         self.pm = pm
         self.config = config or {}
+        self.provider = 'azure'
         key_ids = self.config.get('key_ids', get_keys())
         if not isinstance(key_ids, list) or len(key_ids) == 0:
             raise ex.ConfigException('kms key_ids required')
@@ -56,6 +57,9 @@ def __init__(
                 'credential', DefaultAzureCredential()
             )
         )
+        self.pack_bytes_maxlen = self.config.get(
+            'pack_bytes_maxlen', 10000
+        )
 
     @kms_ex_handler()
     def encrypt(
@@ -64,46 +68,58 @@ def encrypt(
         # azure doesnt have GenerateDataKey equivilent
         # as AWS does, and its encrypt/decrypt APIs
         # are only for use against CMKs not data keys
-        # so we must do our own AESGCM key to encrypt/decrypt
-        # the plaintext and then use azure to wrap/unwrap
-        # this with CMK.
-        # This follows similar pattern to aws-encryption-sdk
+        # so generate own AESGCM DEK to encrypt/decrypt
+        # the plaintext and then use azure to wrap/unwrap this with CMK.
+        # This follows similar pattern to aws-encryption-sdk and google tink
         # and the wrapped/encrypted AES key lives with the data
+        # see https://developers.google.com/tink/client-side-encryption
+        # https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html  # noqa
         context = dict(sorted(context.items()))
         aad = json.dumps(context).encode()
+
+        # 1) generate random Data Encryption Key (DEK)
         # 256-bit AES-GCM key with 96-bit nonce
         nonce = os.urandom(96)
-        key = key or AESGCM.generate_key(bit_length=256)
-        aesgcm = AESGCM(key)
-        # encrypt the key using Azure Key Vault CMK
-        enc_key = self.crypto_client.wrap_key(
-            KeyWrapAlgorithm.rsa_oaep_256, key
+        dek = key or AESGCM.generate_key(bit_length=256)
+        dek_aesgcm = AESGCM(dek)
+
+        # 2) The DEK is encrypted by a Key Encryption Key (KEK)
+        # that is stored in a cloud KMS (Azure Key Vault CMK)
+        enc_dek = self.crypto_client.wrap_key(
+            KeyWrapAlgorithm.rsa_oaep_256, dek
         ).encrypted_key
-        # delete unencrypted key from memory asap
-        del key
-        # encrypt
-        ct = aesgcm.encrypt(nonce, plaintext.encode(), aad)
-        del aesgcm
-        # byte packing is smaller than json and what aws-encryption-sdk does
+        del dek  # delete unencrypted DEK from memory asap
+
+        # 3) Data is encrypted using the DEK by the client.
+        ct = dek_aesgcm.encrypt(nonce, plaintext.encode(), aad)
+        del dek_aesgcm
+
+        # 4) Concatenates the KEK-encrypted encryption DEK with the encrypted
+        # data (byte packing is what aws-encryption-sdk and google tink do)
         serialized = b64encode(
-            pack_bytes([ct, nonce, enc_key], 10000)
+            pack_bytes([ct, nonce, enc_dek], self.pack_bytes_maxlen)
         ).decode()
         return serialized
 
     @kms_ex_handler()
     def decrypt(self, serialized: str, context: t.Dict) -> str:
         context = dict(sorted(context.items()))
         aad = json.dumps(context).encode()
+
+        # 1) Extracts the KEK-encrypted DEK key.
         unpacked = unpack_bytes(b64decode(serialized.encode()))
         if len(unpacked) != 3:
             raise ValueError('invalid serialization')
-        (ct, nonce, enc_key) = unpacked
+        (ct, nonce, enc_dek) = unpacked
+
+        # 2) Makes a request to your KMS to decrypt the KEK-encrypted DEK.
         # decrypt the key using Azure Key Vault CMK
-        key = self.crypto_client.unwrap_key(
-            KeyWrapAlgorithm.rsa_oaep_256, enc_key  # obj['key']
+        dek = self.crypto_client.unwrap_key(
+            KeyWrapAlgorithm.rsa_oaep_256, enc_dek
         ).key
-        aesgcm = AESGCM(key)
-        del key
-        plaintext = aesgcm.decrypt(nonce, ct, aad).decode()
-        # plaintext = aesgcm.decrypt(obj['nonce'], obj['ct'], aad).decode()
+
+        # 3) Decrypts the ciphertext locally using the DEK.
+        dek_aesgcm = AESGCM(dek)
+        del dek
+        plaintext = dek_aesgcm.decrypt(nonce, ct, aad).decode()
         return plaintext

abnosql/plugins/kms/gcp.py

@@ -0,0 +1,85 @@
+from base64 import b64decode
+from base64 import b64encode
+import functools
+import json
+import os
+import typing as t
+
+import abnosql.exceptions as ex
+from abnosql.kms import get_keys
+from abnosql.kms import KmsBase
+from abnosql.plugin import PM
+
+try:
+    from tink import aead  # type: ignore
+    from tink import core  # type: ignore
+    from tink.integration import gcpkms  # type: ignore
+    from tink import new_keyset_handle  # type: ignore
+except ImportError:
+    MISSING_DEPS = True
+
+
+def kms_ex_handler(raise_not_found: t.Optional[bool] = True):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except core.TinkError as e:
+                raise ex.ConfigException(detail=str(e))
+            except Exception as e:
+                raise ex.PluginException(detail=e)
+        return wrapper
+    return decorator
+
+
+def mock_remote_aead(*args, **kwargs):
+    # used for patching during tests
+    # see https://github.com/tink-crypto/tink-py/blob/main/tink/aead/_kms_envelope_aead_test.py  # noqa
+    keyset_handle = new_keyset_handle(aead.aead_key_templates.AES256_GCM)
+    return keyset_handle.primitive(aead.Aead)
+
+
+class Kms(KmsBase):
+
+    @kms_ex_handler()
+    def __init__(
+        self, pm: PM, config: t.Optional[dict] = None
+    ) -> None:
+        self.pm = pm
+        self.config = config or {}
+        self.provider = 'gcp'
+
+        self.key_ids = self.config.get('key_ids', get_keys())
+        if not isinstance(self.key_ids, list) or len(self.key_ids) == 0:
+            raise ex.ConfigException('kms key_ids required')
+        self.kek_uri = self.key_ids[0]
+        self.credentials = self.config.get(
+            'credentials', os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
+        )
+        # see https://developers.google.com/tink/client-side-encryption
+        aead.register()
+        self.client = gcpkms.GcpKmsClient(
+            self.kek_uri,
+            self.credentials
+        )
+        remote_aead = self.client.get_aead(self.kek_uri)
+        self.env_aead = aead.KmsEnvelopeAead(
+            aead.aead_key_templates.AES256_GCM, remote_aead
+        )
+
+    @kms_ex_handler()
+    def encrypt(
+        self, plaintext: str, context: t.Dict, key: t.Optional[bytes] = None
+    ) -> str:
+        ciphertext = self.env_aead.encrypt(
+            plaintext.encode(), json.dumps(context).encode()
+        )
+        return b64encode(ciphertext).decode()
+
+    @kms_ex_handler()
+    def decrypt(self, serialized: str, context: t.Dict) -> str:
+        plaintext = self.env_aead.decrypt(
+            b64decode(serialized), json.dumps(context).encode()
+        )
+        return plaintext.decode()