samples: Minor fixes for importing-a-key snippets (#68)

Author: lwolfowitz-google

kms/snippets/create_import_job.py

@@ -35,7 +35,10 @@ def create_import_job(project_id, location_id, key_ring_id, import_job_id):
 
     # Set paramaters for the import job, allowed values for ImportMethod and ProtectionLevel found here:
     # https://googleapis.dev/python/cloudkms/latest/_modules/google/cloud/kms_v1/types/resources.html
-    import_job_params = {"import_method": kms.ImportJob.ImportMethod.RSA_OAEP_3072_SHA1_AES_256, "protection_level": kms.ProtectionLevel.HSM}
+
+    import_method = kms.ImportJob.ImportMethod.RSA_OAEP_3072_SHA1_AES_256
+    protection_level = kms.ProtectionLevel.HSM
+    import_job_params = {"import_method": import_method, "protection_level": protection_level}
 
     # Call the client to create a new import job.
     import_job = client.create_import_job({"parent": key_ring_name, "import_job_id": import_job_id, "import_job": import_job_params})

kms/snippets/create_key_for_import.py

@@ -15,7 +15,9 @@
 # [START kms_create_key_for_import]
 def create_key_for_import(project_id, location_id, key_ring_id, crypto_key_id):
     """
-    Generate Cloud KMS-compatible key material locally and sets up an empty CryptoKey within a KeyRing for import.
+
+    Sets up an empty CryptoKey within a KeyRing for import.
+
 
     Args:
         project_id (string): Google Cloud project ID (e.g. 'my-project').
@@ -24,24 +26,9 @@ def create_key_for_import(project_id, location_id, key_ring_id, crypto_key_id):
         crypto_key_id (string): ID of the key to import (e.g. 'my-asymmetric-signing-key').
     """
 
-    # Import Python standard cryptographic libraries.
-    from cryptography.hazmat.backends import default_backend
-    from cryptography.hazmat.primitives import serialization
-    from cryptography.hazmat.primitives.asymmetric import ec
-
     # Import the client library.
     from google.cloud import kms
 
-    # Generate some key material in Python and format it in PKCS #8 DER as
-    # required by Google Cloud KMS.
-    key = ec.generate_private_key(ec.SECP256R1, default_backend())
-    formatted_key = key.private_bytes(
-        serialization.Encoding.DER,
-        serialization.PrivateFormat.PKCS8,
-        serialization.NoEncryption())
-
-    print('Generated key bytes: {}'.format(formatted_key))
-
     # Create the client.
     client = kms.KeyManagementServiceClient()
 

kms/snippets/import_manually_wrapped_key.py

@@ -13,28 +13,35 @@
 
 
 # [START kms_import_manually_wrapped_key]
-def import_manually_wrapped_key(project_id, location_id, key_ring_id, crypto_key_id, import_job_id, key_material):
+def import_manually_wrapped_key(project_id, location_id, key_ring_id, crypto_key_id, import_job_id):
     """
-    Imports local key material to Cloud KMS.
+    Generates and imports local key material to Cloud KMS.
 
     Args:
         project_id (string): Google Cloud project ID (e.g. 'my-project').
         location_id (string): Cloud KMS location (e.g. 'us-east1').
         key_ring_id (string): ID of the Cloud KMS key ring (e.g. 'my-key-ring').
         crypto_key_id (string): ID of the key to import (e.g. 'my-asymmetric-signing-key').
         import_job_id (string): ID of the import job (e.g. 'my-import-job').
-        key_material (bytes): Locally generated key material in PKCS #8 DER format.
-    Returns:
-        CryptoKeyVersion: An instance of the imported key in Cloud KMS.
     """
 
     # Import the client library and Python standard cryptographic libraries.
     import os
-    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat import backends
     from cryptography.hazmat.primitives import hashes, keywrap, serialization
-    from cryptography.hazmat.primitives.asymmetric import padding
+    from cryptography.hazmat.primitives.asymmetric import ec, padding
     from google.cloud import kms
 
+    # Generate some key material in Python and format it in PKCS #8 DER as
+    # required by Google Cloud KMS.
+    key = ec.generate_private_key(ec.SECP256R1, backends.default_backend())
+    formatted_key = key.private_bytes(
+        serialization.Encoding.DER,
+        serialization.PrivateFormat.PKCS8,
+        serialization.NoEncryption())
+
+    print('Generated key bytes: {}'.format(formatted_key))
+
     # Create the client.
     client = kms.KeyManagementServiceClient()
 
@@ -47,12 +54,12 @@ def import_manually_wrapped_key(project_id, location_id, key_ring_id, crypto_key
     # Generate a temporary 32-byte key for AES-KWP and wrap the key material.
     kwp_key = os.urandom(32)
     wrapped_target_key = keywrap.aes_key_wrap_with_padding(
-        kwp_key, key_material, default_backend())
+        kwp_key, formatted_key, backends.default_backend())
 
     # Retrieve the public key from the import job.
     import_job = client.get_import_job(name=import_job_name)
     import_job_pub = serialization.load_pem_public_key(
-        bytes(import_job.public_key.pem, 'UTF-8'), default_backend())
+        bytes(import_job.public_key.pem, 'UTF-8'), backends.default_backend())
 
     # Wrap the KWP key using the import job key.
     wrapped_kwp_key = import_job_pub.encrypt(

kms/snippets/snippets_test.py

@@ -20,7 +20,7 @@
 from cryptography.exceptions import InvalidSignature
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes, serialization
-from cryptography.hazmat.primitives.asymmetric import ec, padding, utils
+from cryptography.hazmat.primitives.asymmetric import padding, utils
 from google.cloud import kms
 import pytest
 
@@ -226,7 +226,7 @@ def test_create_key_asymmetric_sign(project_id, location_id, key_ring_id):
 def test_create_key_for_import(project_id, location_id, key_ring_id, import_tests_key_id, capsys):
     create_key_for_import(project_id, location_id, key_ring_id, import_tests_key_id)
     out, _ = capsys.readouterr()
-    assert "Generated key" in out
+    assert "Created hsm key" in out
 
 
 def test_create_key_hsm(project_id, location_id, key_ring_id):
@@ -387,12 +387,7 @@ def test_iam_remove_member(client, project_id, location_id, key_ring_id, asymmet
 
 
 def test_import_manually_wrapped_key(project_id, location_id, key_ring_id, import_job_id, import_tests_key_id, capsys):
-    key = ec.generate_private_key(ec.SECP256R1, default_backend())
-    formatted_key = key.private_bytes(
-        serialization.Encoding.DER,
-        serialization.PrivateFormat.PKCS8,
-        serialization.NoEncryption())
-    import_manually_wrapped_key(project_id, location_id, key_ring_id, import_tests_key_id, import_job_id, formatted_key)
+    import_manually_wrapped_key(project_id, location_id, key_ring_id, import_tests_key_id, import_job_id)
     out, _ = capsys.readouterr()
     assert "Imported" in out