feat: customer managed keys (CMEK) (#249)

* feat: customer managed keys (CMEK)

Implement customer managed keys (CMEK) feature.

WIP. DO NOT MERGE.

* Wrap Status.

* Wrapper for Status, reorganize to avoid circular imports.

* Blacken.

* Make system tests in charge of their own key.

* Consolidate system tests. Get KMS_KEY_NAME from user's environment.

* Fix test.

* Lint.

* Put system tests where nox is expecting to find them.

* Test backup with CMEK.

* Differentiate instance and cluster names for cmek test, so tests aren't stepping on each other's toes. Remove bogus backup with cmek test.

* rename `encryption.py` to `encryption_info.py`

* make sure `kms_key_name` is set to `None` if `encryption_info` is not
PB.

* Fix typo. Use more realistic looking test strings.

google/cloud/bigtable/backup.py

@@ -19,6 +19,7 @@
 from google.cloud._helpers import _datetime_to_pb_timestamp
 from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient
 from google.cloud.bigtable_admin_v2.types import table
+from google.cloud.bigtable.encryption_info import EncryptionInfo
 from google.cloud.bigtable.policy import Policy
 from google.cloud.exceptions import NotFound
 from google.protobuf import field_mask_pb2
@@ -67,13 +68,20 @@ class Backup(object):
     """
 
     def __init__(
-        self, backup_id, instance, cluster_id=None, table_id=None, expire_time=None
+        self,
+        backup_id,
+        instance,
+        cluster_id=None,
+        table_id=None,
+        expire_time=None,
+        encryption_info=None,
     ):
         self.backup_id = backup_id
         self._instance = instance
         self._cluster = cluster_id
         self.table_id = table_id
         self._expire_time = expire_time
+        self._encryption_info = encryption_info
 
         self._parent = None
         self._source_table = None
@@ -176,6 +184,15 @@ def expire_time(self):
     def expire_time(self, new_expire_time):
         self._expire_time = new_expire_time
 
+    @property
+    def encryption_info(self):
+        """Encryption info for this Backup.
+
+        :rtype: :class:`google.cloud.bigtable.encryption.EncryptionInfo`
+        :returns: The encryption information for this backup.
+        """
+        return self._encryption_info
+
     @property
     def start_time(self):
         """The time this Backup was started.
@@ -255,13 +272,15 @@ def from_pb(cls, backup_pb, instance):
         table_id = match.group("table_id") if match else None
 
         expire_time = backup_pb._pb.expire_time
+        encryption_info = EncryptionInfo._from_pb(backup_pb.encryption_info)
 
         backup = cls(
             backup_id,
             instance,
             cluster_id=cluster_id,
             table_id=table_id,
             expire_time=expire_time,
+            encryption_info=encryption_info,
         )
         backup._start_time = backup_pb._pb.start_time
         backup._end_time = backup_pb._pb.end_time

google/cloud/bigtable/cluster.py

@@ -63,6 +63,19 @@ class Cluster(object):
                                  Defaults to
                                  :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`.
 
+    :type kms_key_name: str
+    :param kms_key_name: (Optional, Creation Only) The name of the KMS customer managed
+                         encryption key (CMEK) to use for at-rest encryption of data in
+                         this cluster.  If omitted, Google's default encryption will be
+                         used. If specified, the requirements for this key are:
+
+                         1) The Cloud Bigtable service account associated with the
+                            project that contains the cluster must be granted the
+                            ``cloudkms.cryptoKeyEncrypterDecrypter`` role on the CMEK.
+                         2) Only regional keys can be used and the region of the CMEK
+                            key must match the region of the cluster.
+                         3) All clusters within an instance must use the same CMEK key.
+
     :type _state: int
     :param _state: (`OutputOnly`)
                    The current state of the cluster.
@@ -81,13 +94,15 @@ def __init__(
         location_id=None,
         serve_nodes=None,
         default_storage_type=None,
+        kms_key_name=None,
         _state=None,
     ):
         self.cluster_id = cluster_id
         self._instance = instance
         self.location_id = location_id
         self.serve_nodes = serve_nodes
         self.default_storage_type = default_storage_type
+        self._kms_key_name = kms_key_name
         self._state = _state
 
     @classmethod
@@ -145,6 +160,10 @@ def _update_from_pb(self, cluster_pb):
         self.location_id = cluster_pb.location.split("/")[-1]
         self.serve_nodes = cluster_pb.serve_nodes
         self.default_storage_type = cluster_pb.default_storage_type
+        if cluster_pb.encryption_config:
+            self._kms_key_name = cluster_pb.encryption_config.kms_key_name
+        else:
+            self._kms_key_name = None
         self._state = cluster_pb.state
 
     @property
@@ -187,6 +206,11 @@ def state(self):
         """
         return self._state
 
+    @property
+    def kms_key_name(self):
+        """str: Customer managed encryption key for the cluster."""
+        return self._kms_key_name
+
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             return NotImplemented
@@ -356,4 +380,8 @@ def _to_pb(self):
             serve_nodes=self.serve_nodes,
             default_storage_type=self.default_storage_type,
         )
+        if self._kms_key_name:
+            cluster_pb.encryption_config = instance.Cluster.EncryptionConfig(
+                kms_key_name=self._kms_key_name,
+            )
         return cluster_pb

google/cloud/bigtable/encryption_info.py

@@ -0,0 +1,64 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Class for encryption info for tables and backups."""
+
+from google.cloud.bigtable.error import Status
+
+
+class EncryptionInfo:
+    """Encryption information for a given resource.
+
+    If this resource is protected with customer managed encryption, the in-use Google
+    Cloud Key Management Service (KMS) key versions will be specified along with their
+    status.
+
+    :type encryption_type: int
+    :param encryption_type: See :class:`enums.EncryptionInfo.EncryptionType`
+
+    :type encryption_status: google.cloud.bigtable.encryption.Status
+    :param encryption_status: The encryption status.
+
+    :type kms_key_version: str
+    :param kms_key_version: The key version used for encryption.
+    """
+
+    @classmethod
+    def _from_pb(cls, info_pb):
+        return cls(
+            info_pb.encryption_type,
+            Status(info_pb.encryption_status),
+            info_pb.kms_key_version,
+        )
+
+    def __init__(self, encryption_type, encryption_status, kms_key_version):
+        self.encryption_type = encryption_type
+        self.encryption_status = encryption_status
+        self.kms_key_version = kms_key_version
+
+    def __eq__(self, other):
+        if self is other:
+            return True
+
+        if not isinstance(other, type(self)):
+            return NotImplemented
+
+        return (
+            self.encryption_type == other.encryption_type
+            and self.encryption_status == other.encryption_status
+            and self.kms_key_version == other.kms_key_version
+        )
+
+    def __ne__(self, other):
+        return not self == other

google/cloud/bigtable/enums.py

@@ -156,6 +156,7 @@ class View(object):
         NAME_ONLY = table.Table.View.NAME_ONLY
         SCHEMA_VIEW = table.Table.View.SCHEMA_VIEW
         REPLICATION_VIEW = table.Table.View.REPLICATION_VIEW
+        ENCRYPTION_VIEW = table.Table.View.ENCRYPTION_VIEW
         FULL = table.Table.View.FULL
 
     class ReplicationState(object):
@@ -191,3 +192,32 @@ class ReplicationState(object):
             table.Table.ClusterState.ReplicationState.UNPLANNED_MAINTENANCE
         )
         READY = table.Table.ClusterState.ReplicationState.READY
+
+
+class EncryptionInfo:
+    class EncryptionType:
+        """Possible encryption types for a resource.
+
+        Attributes:
+            ENCRYPTION_TYPE_UNSPECIFIED (int): Encryption type was not specified, though
+                data at rest remains encrypted.
+            GOOGLE_DEFAULT_ENCRYPTION (int): The data backing this resource is encrypted
+                at rest with a key that is fully managed by Google. No key version or
+                status will be populated. This is the default state.
+            CUSTOMER_MANAGED_ENCRYPTION (int): The data backing this resource is
+                encrypted at rest with a key that is managed by the customer. The in-use
+                version of the key and its status are populated for CMEK-protected
+                tables. CMEK-protected backups are pinned to the key version that was in
+                use at the time the backup was taken. This key version is populated but
+                its status is not tracked and is reported as `UNKNOWN`.
+        """
+
+        ENCRYPTION_TYPE_UNSPECIFIED = (
+            table.EncryptionInfo.EncryptionType.ENCRYPTION_TYPE_UNSPECIFIED
+        )
+        GOOGLE_DEFAULT_ENCRYPTION = (
+            table.EncryptionInfo.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION
+        )
+        CUSTOMER_MANAGED_ENCRYPTION = (
+            table.EncryptionInfo.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION
+        )

google/cloud/bigtable/error.py

@@ -0,0 +1,64 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Class for error status."""
+
+
+class Status:
+    """A status, comprising a code and a message.
+
+    See: `Cloud APIs Errors <https://cloud.google.com/apis/design/errors>`_
+
+    This is a thin wrapper for ``google.rpc.status_pb2.Status``.
+
+    :type status_pb: google.rpc.status_pb2.Status
+    :param status_pb: The status protocol buffer.
+    """
+
+    def __init__(self, status_pb):
+        self.status_pb = status_pb
+
+    @property
+    def code(self):
+        """The status code.
+
+        Values are defined in ``google.rpc.code_pb2.Code``.
+
+        See: `google.rpc.Code
+        <https://github.com/googleapis/googleapis/blob/master/google/rpc/code.proto>`_
+
+        :rtype: int
+        :returns: The status code.
+        """
+        return self.status_pb.code
+
+    @property
+    def message(self):
+        """A human readable status message.
+
+        :rypte: str
+        :returns: The status message.
+        """
+        return self.status_pb.message
+
+    def __repr__(self):
+        return repr(self.status_pb)
+
+    def __eq__(self, other):
+        if isinstance(other, type(self)):
+            return self.status_pb == other.status_pb
+        return NotImplemented
+
+    def __ne__(self, other):
+        return not self == other

google/cloud/bigtable/instance.py

@@ -540,7 +540,12 @@ def test_iam_permissions(self, permissions):
         return list(resp.permissions)
 
     def cluster(
-        self, cluster_id, location_id=None, serve_nodes=None, default_storage_type=None
+        self,
+        cluster_id,
+        location_id=None,
+        serve_nodes=None,
+        default_storage_type=None,
+        kms_key_name=None,
     ):
         """Factory to create a cluster associated with this instance.
 
@@ -576,13 +581,30 @@ def cluster(
 
         :rtype: :class:`~google.cloud.bigtable.instance.Cluster`
         :returns: a cluster owned by this instance.
+
+        :type kms_key_name: str
+        :param kms_key_name: (Optional, Creation Only) The name of the KMS customer
+                             managed encryption key (CMEK) to use for at-rest encryption
+                             of data in this cluster.  If omitted, Google's default
+                             encryption will be used. If specified, the requirements for
+                             this key are:
+
+                             1) The Cloud Bigtable service account associated with the
+                                project that contains the cluster must be granted the
+                                ``cloudkms.cryptoKeyEncrypterDecrypter`` role on the
+                                CMEK.
+                             2) Only regional keys can be used and the region of the
+                                CMEK key must match the region of the cluster.
+                             3) All clusters within an instance must use the same CMEK
+                                key.
         """
         return Cluster(
             cluster_id,
             self,
             location_id=location_id,
             serve_nodes=serve_nodes,
             default_storage_type=default_storage_type,
+            kms_key_name=kms_key_name,
         )
 
     def list_clusters(self):

google/cloud/bigtable/table.py

@@ -28,6 +28,7 @@
 from google.cloud.bigtable.column_family import ColumnFamily
 from google.cloud.bigtable.batcher import MutationsBatcher
 from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_ROW_BYTES
+from google.cloud.bigtable.encryption_info import EncryptionInfo
 from google.cloud.bigtable.policy import Policy
 from google.cloud.bigtable.row import AppendRow
 from google.cloud.bigtable.row import ConditionalRow
@@ -484,6 +485,33 @@ def get_cluster_states(self):
             for cluster_id, value_pb in table_pb.cluster_states.items()
         }
 
+    def get_encryption_info(self):
+        """List the encryption info for each cluster owned by this table.
+
+        Gets the current encryption info for the table across all of the clusters.  The
+        returned dict will be keyed by cluster id and contain a status for all of the
+        keys in use.
+
+        :rtype: dict
+        :returns: Dictionary of encryption info for this table. Keys are cluster ids and
+                  values are tuples of :class:`google.cloud.bigtable.encryption.EncryptionInfo` instances.
+        """
+        ENCRYPTION_VIEW = enums.Table.View.ENCRYPTION_VIEW
+        table_client = self._instance._client.table_admin_client
+        table_pb = table_client.get_table(
+            request={"name": self.name, "view": ENCRYPTION_VIEW}
+        )
+
+        return {
+            cluster_id: tuple(
+                (
+                    EncryptionInfo._from_pb(info_pb)
+                    for info_pb in value_pb.encryption_info
+                )
+            )
+            for cluster_id, value_pb in table_pb.cluster_states.items()
+        }
+
     def read_row(self, row_key, filter_=None):
         """Read a single row from this table.