fixit: Clean up Python sample at compute/encryption (#10118)

Co-authored-by: Andrew Ferlitsch <aferlitsch@gmail.com>

Author: m-strzelczyk

compute/encryption/generate_wrapped_rsa_key.py

@@ -1,93 +1,122 @@
 #!/usr/bin/env python
 
-# Copyright 2016 Google Inc. All Rights Reserved.
+#  Copyright 2016 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
+#  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
+#      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.
+#  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.
 
 """Example of authenticating using access tokens directly on Compute Engine.
 
 For more information, see the README.md under /compute.
 """
 
 # [START all]
-
+# [START compute_generate_wrapped_rsa_key]
 import argparse
 import base64
 import os
+from typing import Optional
 
 from cryptography import x509
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
 import requests
 
 
 GOOGLE_PUBLIC_CERT_URL = (
-    'https://cloud-certs.storage.googleapis.com/google-cloud-csek-ingress.pem')
+    "https://cloud-certs.storage.googleapis.com/google-cloud-csek-ingress.pem"
+)
+
 
+def get_google_public_cert_key() -> RSAPublicKey:
+    """
+    Downloads the Google public certificate.
 
-def get_google_public_cert_key():
+    Returns:
+        RSAPublicKey object with the Google public certificate.
+    """
     r = requests.get(GOOGLE_PUBLIC_CERT_URL)
     r.raise_for_status()
 
     # Load the certificate.
-    certificate = x509.load_pem_x509_certificate(
-        r.content, default_backend())
+    certificate = x509.load_pem_x509_certificate(r.content, default_backend())
 
     # Get the certicate's public key.
     public_key = certificate.public_key()
 
     return public_key
 
 
-def wrap_rsa_key(public_key, private_key_bytes):
-    # Use the Google public key to encrypt the customer private key.
-    # This means that only the Google private key is capable of decrypting
-    # the customer private key.
+def wrap_rsa_key(public_key: RSAPublicKey, private_key_bytes: bytes) -> bytes:
+    """
+    Use the Google public key to encrypt the customer private key.
+
+    This means that only the Google private key is capable of decrypting
+    the customer private key.
+
+    Args:
+        public_key: The public key to use for encrypting.
+        private_key_bytes: The private key to be encrypted.
+
+    Returns:
+        private_key_bytes encrypted using the public_key. Encoded using
+        base64.
+    """
     wrapped_key = public_key.encrypt(
         private_key_bytes,
         padding.OAEP(
             mgf=padding.MGF1(algorithm=hashes.SHA1()),
             algorithm=hashes.SHA1(),
-            label=None))
+            label=None,
+        ),
+    )
     encoded_wrapped_key = base64.b64encode(wrapped_key)
     return encoded_wrapped_key
 
 
-def main(key_file):
+def main(key_file: Optional[str]) -> None:
+    """
+    This script will encrypt a private key with Google public key.
+
+    Args:
+        key_file: path to a file containing your private key. If not
+            provided, a new key will be generated (256 bit).
+    """
     # Generate a new 256-bit private key if no key is specified.
     if not key_file:
         customer_key_bytes = os.urandom(32)
     else:
-        with open(key_file, 'rb') as f:
+        with open(key_file, "rb") as f:
             customer_key_bytes = f.read()
 
     google_public_key = get_google_public_cert_key()
     wrapped_rsa_key = wrap_rsa_key(google_public_key, customer_key_bytes)
 
-    print('Base-64 encoded private key: {}'.format(
-        base64.b64encode(customer_key_bytes).decode('utf-8')))
-    print('Wrapped RSA key: {}'.format(wrapped_rsa_key.decode('utf-8')))
+    b64_key = base64.b64encode(customer_key_bytes).decode("utf-8")
+
+    print(f"Base-64 encoded private key: {b64_key}")
+    print(f"Wrapped RSA key: {wrapped_rsa_key.decode('utf-8')}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description=__doc__,
-        formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument(
-        '--key_file', help='File containing your binary private key.')
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--key_file", help="File containing your binary private key.")
 
     args = parser.parse_args()
 
     main(args.key_file)
+# [END compute_generate_wrapped_rsa_key]
 # [END all]

compute/encryption/generate_wrapped_rsa_key_test.py

@@ -19,37 +19,38 @@
 
 import generate_wrapped_rsa_key
 
-PROJECT = os.environ['GOOGLE_CLOUD_PROJECT']
+PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 
 
-def test_main():
+def test_main() -> None:
     generate_wrapped_rsa_key.main(None)
 
 
-def test_create_disk():
-    compute = googleapiclient.discovery.build('compute', 'beta')
+def test_create_disk() -> None:
+    compute = googleapiclient.discovery.build("compute", "beta")
 
     # Generate the key.
     key_bytes = os.urandom(32)
     google_public_key = generate_wrapped_rsa_key.get_google_public_cert_key()
     wrapped_rsa_key = generate_wrapped_rsa_key.wrap_rsa_key(
-        google_public_key, key_bytes)
-    disk_name = f'new-encrypted-disk-{uuid.uuid4().hex}'
+        google_public_key, key_bytes
+    )
+    disk_name = f"new-encrypted-disk-{uuid.uuid4().hex}"
 
     try:
         # Create the disk, if the encryption key is invalid, this will raise.
         compute.disks().insert(
             project=PROJECT,
-            zone='us-central1-f',
+            zone="us-central1-f",
             body={
-                'name': disk_name,
-                'diskEncryptionKey': {
-                    'rsaEncryptedKey': wrapped_rsa_key.decode('utf-8')
-                }
-            }).execute()
+                "name": disk_name,
+                "diskEncryptionKey": {
+                    "rsaEncryptedKey": wrapped_rsa_key.decode("utf-8")
+                },
+            },
+        ).execute()
     finally:
         # Delete the disk.
         compute.disks().delete(
-            project=PROJECT,
-            zone='us-central1-f',
-            disk=disk_name).execute()
+            project=PROJECT, zone="us-central1-f", disk=disk_name
+        ).execute()

compute/encryption/noxfile_config.py

@@ -0,0 +1,42 @@
+# 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.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be imported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    "ignored_versions": ["2.7"],
+    # Old samples are opted out of enforcing Python type hints
+    # All new samples should feature them
+    "enforce_type_hints": True,
+    # An envvar key for determining the project id to use. Change it
+    # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+    # build specific Cloud project. You can also use your own string
+    # to use your own Cloud project.
+    "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+    # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
+    # If you need to use a specific version of pip,
+    # change pip_version_override to the string representation
+    # of the version number, for example, "20.2.4"
+    "pip_version_override": None,
+    # A dictionary you want to inject into your test. Don't put any
+    # secrets here. These values will override predefined values.
+    "envs": {},
+}