az keyvault key sign/verify: Fix --digest to accept base64 encoded string

Author: evelyn-ys

src/azure-cli/azure/cli/command_modules/keyvault/custom.py

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 
 # pylint: disable=too-many-lines
+import base64
 import codecs
 import hashlib
 import json
@@ -1127,19 +1128,31 @@ def decrypt_key(cmd, client, algorithm, value, iv=None, tag=None, aad=None,
 
 
 def sign_key(cmd, client, algorithm, digest, name=None, version=None):
+    if '256' in algorithm:
+        hash_func = hashlib.sha256
+    elif '384' in algorithm:
+        hash_func = hashlib.sha384
+    else:
+        hash_func = hashlib.sha512
     SignatureAlgorithm = cmd.loader.get_sdk('SignatureAlgorithm', mod='crypto._enums',
                                             resource_type=ResourceType.DATA_KEYVAULT_KEYS)
     crypto_client = client.get_cryptography_client(name, key_version=version)
-    return crypto_client.sign(SignatureAlgorithm(algorithm), digest.encode('utf-8'))
+    return crypto_client.sign(SignatureAlgorithm(algorithm),
+                              hash_func(base64.b64decode(digest.encode('utf-8'))).digest())
 
 
 def verify_key(cmd, client, algorithm, digest, signature, name=None, version=None):
-    import base64
+    if '256' in algorithm:
+        hash_func = hashlib.sha256
+    elif '384' in algorithm:
+        hash_func = hashlib.sha384
+    else:
+        hash_func = hashlib.sha512
     SignatureAlgorithm = cmd.loader.get_sdk('SignatureAlgorithm', mod='crypto._enums',
                                             resource_type=ResourceType.DATA_KEYVAULT_KEYS)
     crypto_client = client.get_cryptography_client(name, key_version=version)
     return crypto_client.verify(SignatureAlgorithm(algorithm),
-                                digest.encode('utf-8'),
+                                hash_func(base64.b64decode(digest.encode('utf-8'))).digest(),
                                 base64.b64decode(signature.encode('utf-8')))
 
 
@@ -1493,7 +1506,6 @@ def download_secret(client, file_path, name=None, encoding=None, version=''):  #
                 f.write(secret_value)
         else:
             if encoding == 'base64':
-                import base64
                 decoded = base64.b64decode(secret_value)
             elif encoding == 'hex':
                 import binascii
@@ -1558,7 +1570,6 @@ def download_certificate(client, file_path, certificate_name=None, encoding='PEM
             if encoding == 'DER':
                 f.write(cert)
             else:
-                import base64
                 encoded = base64.encodebytes(cert)
                 if isinstance(encoded, bytes):
                     encoded = encoded.decode("utf-8")