Storage: add KMS samples (GoogleCloudPlatform#1510)

* Storage: add KMS samples

* Add CLOUD_KMS_KEY environment variable

Author: broady

storage/cloud-client/README.rst

@@ -94,7 +94,7 @@ To run this sample:
 
     usage: snippets.py [-h]
                        bucket_name
-                       {create-bucket,delete-bucket,get-bucket-labels,add-bucket-label,remove-bucket-label,list,list-with-prefix,upload,download,delete,metadata,make-public,signed-url,rename,copy}
+                       {create-bucket,delete-bucket,get-bucket-labels,add-bucket-label,remove-bucket-label,list,list-with-prefix,upload,enable-default-kms-key,upload-with-kms-key,download,delete,metadata,make-public,signed-url,rename,copy}
                        ...
 
     This application demonstrates how to perform basic operations on blobs
@@ -105,7 +105,7 @@ To run this sample:
 
     positional arguments:
       bucket_name           Your cloud storage bucket.
-      {create-bucket,delete-bucket,get-bucket-labels,add-bucket-label,remove-bucket-label,list,list-with-prefix,upload,download,delete,metadata,make-public,signed-url,rename,copy}
+      {create-bucket,delete-bucket,get-bucket-labels,add-bucket-label,remove-bucket-label,list,list-with-prefix,upload,enable-default-kms-key,upload-with-kms-key,download,delete,metadata,make-public,signed-url,rename,copy}
         create-bucket       Creates a new bucket.
         delete-bucket       Deletes a bucket. The bucket must be empty.
         get-bucket-labels   Prints out a bucket's labels.
@@ -124,6 +124,11 @@ To run this sample:
                             However, if you specify prefix='/a' and delimiter='/',
                             you'll get back: /a/1.txt
         upload              Uploads a file to the bucket.
+        enable-default-kms-key
+                            Sets a bucket's default KMS key.
+        upload-with-kms-key
+                            Uploads a file to the bucket, encrypting it with the
+                            given KMS key.
         download            Downloads a blob from the bucket.
         delete              Deletes a blob from the bucket.
         metadata            Prints out a blob's metadata.

storage/cloud-client/requirements.txt

@@ -1,2 +1,2 @@
-google-cloud-storage==1.8.0
+google-cloud-storage==1.10.0
 google-cloud-pubsub==0.32.1

storage/cloud-client/snippets.py

@@ -43,6 +43,18 @@ def delete_bucket(bucket_name):
     print('Bucket {} deleted'.format(bucket.name))
 
 
+def enable_default_kms_key(bucket_name, kms_key_name):
+    """Sets a bucket's default KMS key."""
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+    bucket.default_kms_key_name = kms_key_name
+    bucket.patch()
+
+    print('Set default KMS key for bucket {} to {}.'.format(
+        bucket.name,
+        bucket.default_kms_key_name))
+
+
 def get_bucket_labels(bucket_name):
     """Prints out a bucket's labels."""
     storage_client = storage.Client()
@@ -143,6 +155,20 @@ def upload_blob(bucket_name, source_file_name, destination_blob_name):
         destination_blob_name))
 
 
+def upload_blob_with_kms(bucket_name, source_file_name, destination_blob_name,
+                         kms_key_name):
+    """Uploads a file to the bucket, encrypting it with the given KMS key."""
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+    blob = bucket.blob(destination_blob_name, kms_key_name=kms_key_name)
+    blob.upload_from_filename(source_file_name)
+
+    print('File {} uploaded to {} with encryption key {}.'.format(
+        source_file_name,
+        destination_blob_name,
+        kms_key_name))
+
+
 def download_blob(bucket_name, source_blob_name, destination_file_name):
     """Downloads a blob from the bucket."""
     storage_client = storage.Client()
@@ -277,6 +303,16 @@ def copy_blob(bucket_name, blob_name, new_bucket_name, new_blob_name):
     upload_parser.add_argument('source_file_name')
     upload_parser.add_argument('destination_blob_name')
 
+    enable_default_kms_parser = subparsers.add_parser(
+        'enable-default-kms-key', help=enable_default_kms_key.__doc__)
+    enable_default_kms_parser.add_argument('kms_key_name')
+
+    upload_kms_parser = subparsers.add_parser(
+        'upload-with-kms-key', help=upload_blob_with_kms.__doc__)
+    upload_kms_parser.add_argument('source_file_name')
+    upload_kms_parser.add_argument('destination_blob_name')
+    upload_kms_parser.add_argument('kms_key_name')
+
     download_parser = subparsers.add_parser(
         'download', help=download_blob.__doc__)
     download_parser.add_argument('source_blob_name')
@@ -310,6 +346,8 @@ def copy_blob(bucket_name, blob_name, new_bucket_name, new_blob_name):
 
     if args.command == 'create-bucket':
         create_bucket(args.bucket_name)
+    if args.command == 'enable-default-kms-key':
+        enable_default_kms_key(args.bucket_name, args.kms_key_name)
     elif args.command == 'delete-bucket':
         delete_bucket(args.bucket_name)
     if args.command == 'get-bucket-labels':
@@ -327,6 +365,12 @@ def copy_blob(bucket_name, blob_name, new_bucket_name, new_blob_name):
             args.bucket_name,
             args.source_file_name,
             args.destination_blob_name)
+    elif args.command == 'upload-with-kms-key':
+        upload_blob_with_kms(
+            args.bucket_name,
+            args.source_file_name,
+            args.destination_blob_name,
+            args.kms_key_name)
     elif args.command == 'download':
         download_blob(
             args.bucket_name,

storage/cloud-client/snippets_test.py

@@ -23,6 +23,7 @@
 import snippets
 
 BUCKET = os.environ['CLOUD_STORAGE_BUCKET']
+KMS_KEY = os.environ['CLOUD_KMS_KEY']
 
 
 def test_get_bucket_labels():
@@ -79,6 +80,17 @@ def test_upload_blob():
             'test_upload_blob')
 
 
+def test_upload_blob_with_kms():
+    with tempfile.NamedTemporaryFile() as source_file:
+        source_file.write(b'test')
+
+        snippets.upload_blob_with_kms(
+            BUCKET,
+            source_file.name,
+            'test_upload_blob_encrypted',
+            KMS_KEY)
+
+
 def test_download_blob(test_blob):
     with tempfile.NamedTemporaryFile() as dest_file:
         snippets.download_blob(