1# ------------------------------------
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4# ------------------------------------
5from azure.keyvault.keys.crypto.aio import CryptographyClient
6
7from _shared.test_case_async import KeyVaultTestCase
8from _test_case import client_setup, get_decorator, KeysTestCase
9
10
11all_api_versions = get_decorator(is_async=True, vault_only=True)
12
13
14class TestCryptoExamples(KeysTestCase, KeyVaultTestCase):
15    def __init__(self, *args, **kwargs):
16        kwargs["match_body"] = False
17        super(TestCryptoExamples, self).__init__(*args, **kwargs)
18
19    @all_api_versions()
20    @client_setup
21    async def test_encrypt_decrypt_async(self, key_client, **kwargs):
22        credential = self.get_credential(CryptographyClient, is_async=True)
23        key_name = self.get_resource_name("crypto-test-encrypt-key")
24        await key_client.create_rsa_key(key_name)
25
26        # [START create_client]
27        # create a CryptographyClient using a KeyVaultKey instance
28        key = await key_client.get_key(key_name)
29        crypto_client = CryptographyClient(key, credential)
30
31        # or a key's id, which must include a version
32        key_id = "https://<your vault>.vault.azure.net/keys/<key name>/fe4fdcab688c479a9aa80f01ffeac26"
33        crypto_client = CryptographyClient(key_id, credential)
34
35        # the client and credential should be closed when no longer needed
36        # (both are also async context managers)
37        await crypto_client.close()
38        await credential.close()
39        # [END create_client]
40
41        client = CryptographyClient(key, credential, api_version=key_client.api_version)
42
43        # [START encrypt]
44        from azure.keyvault.keys.crypto import EncryptionAlgorithm
45
46        # the result holds the ciphertext and identifies the encryption key and algorithm used
47        result = await client.encrypt(EncryptionAlgorithm.rsa_oaep, b"plaintext")
48        print(result.key_id)
49        print(result.algorithm)
50        ciphertext = result.ciphertext
51        # [END encrypt]
52
53        # [START decrypt]
54        from azure.keyvault.keys.crypto import EncryptionAlgorithm
55
56        result = await client.decrypt(EncryptionAlgorithm.rsa_oaep, ciphertext)
57        print(result.plaintext)
58        # [END decrypt]
59
60    @all_api_versions()
61    @client_setup
62    async def test_wrap_unwrap_async(self, key_client, **kwargs):
63        credential = self.get_credential(CryptographyClient, is_async=True)
64        key_name = self.get_resource_name("crypto-test-wrapping-key")
65        key = await key_client.create_rsa_key(key_name)
66        client = CryptographyClient(key, credential, api_version=key_client.api_version)
67
68        key_bytes = b"5063e6aaa845f150200547944fd199679c98ed6f99da0a0b2dafeaf1f4684496fd532c1c229968cb9dee44957fcef7ccef59ceda0b362e56bcd78fd3faee5781c623c0bb22b35beabde0664fd30e0e824aba3dd1b0afffc4a3d955ede20cf6a854d52cfd"
69
70        # [START wrap_key]
71        from azure.keyvault.keys.crypto import KeyWrapAlgorithm
72
73        # wrap returns a tuple with the wrapped bytes and the metadata required to unwrap the key
74        result = await client.wrap_key(KeyWrapAlgorithm.rsa_oaep, key_bytes)
75        print(result.key_id)
76        print(result.algorithm)
77        encrypted_key = result.encrypted_key
78        # [END wrap_key]
79
80        # [START unwrap_key]
81        from azure.keyvault.keys.crypto import KeyWrapAlgorithm
82
83        result = await client.unwrap_key(KeyWrapAlgorithm.rsa_oaep, encrypted_key)
84        # [END unwrap_key]
85
86    @all_api_versions()
87    @client_setup
88    async def test_sign_verify_async(self, key_client, **kwargs):
89        credential = self.get_credential(CryptographyClient, is_async=True)
90        key_name = self.get_resource_name("crypto-test-wrapping-key")
91        key = await key_client.create_rsa_key(key_name)
92        client = CryptographyClient(key, credential, api_version=key_client.api_version)
93
94        # [START sign]
95        import hashlib
96        from azure.keyvault.keys.crypto import SignatureAlgorithm
97
98        digest = hashlib.sha256(b"plaintext").digest()
99
100        # sign returns the signature and the metadata required to verify it
101        result = await client.sign(SignatureAlgorithm.rs256, digest)
102        print(result.key_id)
103        print(result.algorithm)
104        signature = result.signature
105        # [END sign]
106
107        # [START verify]
108        from azure.keyvault.keys.crypto import SignatureAlgorithm
109
110        verified = await client.verify(SignatureAlgorithm.rs256, digest, signature)
111        assert verified.is_valid
112        # [END verify]
113