Adds snippets for enabling and restoring a key version (#1196)

* Adds snippets for enabling and restoring a key version

* Fixed lint issues

Author: WalterHub

kms/api-client/snippets.py

@@ -155,6 +155,33 @@ def disable_crypto_key_version(project_id, location_id, key_ring_id,
 # [END kms_disable_cryptokey_version]
 
 
+# [START kms_enable_cryptokey_version]
+def enable_crypto_key_version(project_id, location_id, key_ring_id,
+                              crypto_key_id, version_id):
+    """Enables a CryptoKeyVersion associated with a given CryptoKey and
+    KeyRing."""
+
+    # Creates an API client for the KMS API.
+    kms_client = googleapiclient.discovery.build('cloudkms', 'v1')
+
+    # Construct the resource name of the CryptoKeyVersion.
+    name = (
+        'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}/'
+        'cryptoKeyVersions/{}'
+        .format(
+            project_id, location_id, key_ring_id, crypto_key_id, version_id))
+
+    # Use the KMS API to enable the CryptoKeyVersion.
+    crypto_keys = kms_client.projects().locations().keyRings().cryptoKeys()
+    request = crypto_keys.cryptoKeyVersions().patch(
+        name=name, body={'state': 'ENABLED'}, updateMask='state')
+    response = request.execute()
+
+    print('CryptoKeyVersion {}\'s state has been set to {}.'.format(
+        name, response['state']))
+# [END kms_enable_cryptokey_version]
+
+
 # [START kms_destroy_cryptokey_version]
 def destroy_crypto_key_version(
         project_id, location_id, key_ring_id, crypto_key_id, version_id):
@@ -181,6 +208,31 @@ def destroy_crypto_key_version(
 # [END kms_destroy_cryptokey_version]
 
 
+# [START kms_restore_cryptokey_version]
+def restore_crypto_key_version(
+        project_id, location_id, key_ring_id, crypto_key_id, version_id):
+    """Restores a CryptoKeyVersion that is scheduled for destruction."""
+
+    # Creates an API client for the KMS API.
+    kms_client = googleapiclient.discovery.build('cloudkms', 'v1')
+
+    # Construct the resource name of the CryptoKeyVersion.
+    name = (
+        'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}/'
+        'cryptoKeyVersions/{}'
+        .format(
+            project_id, location_id, key_ring_id, crypto_key_id, version_id))
+
+    # Use the KMS API to restore the CryptoKeyVersion.
+    crypto_keys = kms_client.projects().locations().keyRings().cryptoKeys()
+    request = crypto_keys.cryptoKeyVersions().restore(name=name, body={})
+    response = request.execute()
+
+    print('CryptoKeyVersion {}\'s state has been set to {}.'.format(
+        name, response['state']))
+# [END kms_restore_cryptokey_version]
+
+
 # [START kms_add_member_to_cryptokey_policy]
 def add_member_to_crypto_key_policy(
         project_id, location_id, key_ring_id, crypto_key_id, member, role):
@@ -294,6 +346,14 @@ def get_key_ring_policy(project_id, location_id, key_ring_id):
     disable_crypto_key_version_parser.add_argument('crypto_key')
     disable_crypto_key_version_parser.add_argument('version')
 
+    enable_crypto_key_version_parser = subparsers.add_parser(
+        'enable_crypto_key_version')
+    enable_crypto_key_version_parser.add_argument('project')
+    enable_crypto_key_version_parser.add_argument('location')
+    enable_crypto_key_version_parser.add_argument('key_ring')
+    enable_crypto_key_version_parser.add_argument('crypto_key')
+    enable_crypto_key_version_parser.add_argument('version')
+
     destroy_crypto_key_version_parser = subparsers.add_parser(
         'destroy_crypto_key_version')
     destroy_crypto_key_version_parser.add_argument('project')
@@ -302,6 +362,14 @@ def get_key_ring_policy(project_id, location_id, key_ring_id):
     destroy_crypto_key_version_parser.add_argument('crypto_key')
     destroy_crypto_key_version_parser.add_argument('version')
 
+    restore_crypto_key_version_parser = subparsers.add_parser(
+        'restore_crypto_key_version')
+    restore_crypto_key_version_parser.add_argument('project')
+    restore_crypto_key_version_parser.add_argument('location')
+    restore_crypto_key_version_parser.add_argument('key_ring')
+    restore_crypto_key_version_parser.add_argument('crypto_key')
+    restore_crypto_key_version_parser.add_argument('version')
+
     add_member_to_crypto_key_policy_parser = subparsers.add_parser(
         'add_member_to_crypto_key_policy')
     add_member_to_crypto_key_policy_parser.add_argument('project')
@@ -352,13 +420,27 @@ def get_key_ring_policy(project_id, location_id, key_ring_id):
             args.key_ring,
             args.crypto_key,
             args.version)
+    elif args.command == 'enable_crypto_key_version':
+        enable_crypto_key_version(
+            args.project,
+            args.location,
+            args.key_ring,
+            args.crypto_key,
+            args.version)
     elif args.command == 'destroy_crypto_key_version':
         destroy_crypto_key_version(
             args.project,
             args.location,
             args.key_ring,
             args.crypto_key,
             args.version)
+    elif args.command == 'restore_crypto_key_version':
+        restore_crypto_key_version(
+            args.project,
+            args.location,
+            args.key_ring,
+            args.crypto_key,
+            args.version)
     elif args.command == 'add_member_to_crypto_key_policy':
         add_member_to_crypto_key_policy(
             args.project,

kms/api-client/snippets_test.py

@@ -102,6 +102,19 @@ def test_disable_crypto_key_version(capsys):
     assert expected in out
 
 
+def test_enable_crypto_key_version(capsys):
+    snippets.enable_crypto_key_version(
+        PROJECT, LOCATION, KEY_RING, CRYPTO_KEY, VERSION)
+    out, _ = capsys.readouterr()
+    expected = (
+        'CryptoKeyVersion projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}/'
+        'cryptoKeyVersions/{}\'s state has been set to {}.'
+        .format(
+            PROJECT, LOCATION, KEY_RING, CRYPTO_KEY, VERSION,
+            'ENABLED'))
+    assert expected in out
+
+
 def test_destroy_crypto_key_version(capsys):
     snippets.destroy_crypto_key_version(
         PROJECT, LOCATION, KEY_RING, CRYPTO_KEY, VERSION)
@@ -115,6 +128,19 @@ def test_destroy_crypto_key_version(capsys):
     assert expected in out
 
 
+def test_restore_crypto_key_version(capsys):
+    snippets.restore_crypto_key_version(
+        PROJECT, LOCATION, KEY_RING, CRYPTO_KEY, VERSION)
+    out, _ = capsys.readouterr()
+    expected = (
+        'CryptoKeyVersion projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}/'
+        'cryptoKeyVersions/{}\'s state has been set to {}.'
+        .format(
+            PROJECT, LOCATION, KEY_RING, CRYPTO_KEY, VERSION,
+            'DISABLED'))
+    assert expected in out
+
+
 def test_add_member_to_crypto_key_policy(capsys):
     snippets.add_member_to_crypto_key_policy(
         PROJECT, LOCATION, KEY_RING, CRYPTO_KEY, MEMBER, ROLE)