keyvault client added

config.py

@@ -3,6 +3,7 @@ import os
 
 from azure.cosmos import PartitionKey
 
+from utils.azure_key_vault_client import AzureKeyVaultClient
 from utils.cosmos_client import CosmosClient
 
 
@@ -24,6 +25,7 @@ class AppConfig:
     """ Bot Configuration """
 
     CLIENT_ID = os.environ.get("CLIENT_ID", None)
+    KEY_VAULT = os.environ.get("KEY_VAULT", '')
     PORT = os.environ.get("HOST_PORT", 8000)
     TENANT_ID = os.environ.get("TENANT_ID",
                                "5df91ebc-64fa-4aa1-862c-bdc0cba3c656")
@@ -36,11 +38,8 @@ class AppConfig:
 
 class CosmosDBConfig:
     """ Cosmos Databases """
-    HOST = os.environ.get('ACCOUNT_HOST',
-                          'https://nancycosomsdb.documents.azure.com:443/')
-    KEY = os.environ.get('COSMOS_KEY',
-                         'fNVRCesO1NAb9MYZNK2rKdAPkY9J4O5ntR8CRuKu6wVGhndiaXch'
-                         'Q6fKwrTTnTbv4tPM8S74YjZsfcX4uAHgiw==')
+    HOST = os.environ.get('ACCOUNT_HOST', '')
+    KEY = os.environ.get('COSMOS_KEY', '')
 
     class Conversations:
         """ Conversation DB """
@@ -72,3 +71,4 @@ class CosmosDBConfig:
 
 
 COSMOS_CLIENT = CosmosClient(CosmosDBConfig.HOST, CosmosDBConfig.KEY)
+KEYVAULT_CLIENT = AzureKeyVaultClient(AppConfig.CLIENT_ID, AppConfig.KEY_VAULT)

requirements.txt

@@ -7,3 +7,6 @@ aiohttp~=3.7.0
 simplejson==3.17.6
 marshmallow-dataclass==8.5.8
 stringcase==1.2.0
+azure-keyvault-secrets==4.2.0
+azure-keyvault-keys==4.3.1
+

utils/azure_key_vault_client.py

+""" AzureVaultClient implementation """
+import asyncio
+import random
+from concurrent.futures.thread import ThreadPoolExecutor
+from typing import Awaitable
+
+# noinspection PyPackageRequirements
+from azure.keyvault.keys import KeyClient, KeyVaultKey
+# noinspection PyPackageRequirements
+from azure.keyvault.keys.crypto import CryptographyClient, EncryptionAlgorithm
+# noinspection PyPackageRequirements
+from azure.keyvault.secrets import SecretClient, KeyVaultSecret
+# noinspection PyPackageRequirements
+from azure.identity import ClientSecretCredential, DefaultAzureCredential, \
+    ManagedIdentityCredential
+
+
+class AzureKeyVaultClient:
+    """ Azure Key Vault Client """
+    def __init__(self, client_id: str, key_vault: str):
+        self.executor = ThreadPoolExecutor(10)
+        self.io_loop = asyncio.get_event_loop()
+        self.key_vault = key_vault
+        self.key_vault_uri = "https://{key_vault}.vault.azure.net".format(
+            key_vault=key_vault
+        )
+        self.credential = ManagedIdentityCredential(client_id=client_id)
+        self.key_client = KeyClient(vault_url=self.key_vault_uri,
+                                    credential=self.credential)
+
+    async def execute_blocking(self, bl, *args):
+        """ Execute blocking code """
+        return await asyncio.get_event_loop().run_in_executor(self.executor,
+                                                              bl,
+                                                              *args)
+
+    def get_key(self, name: str) -> Awaitable["KeyVaultKey"]:
+        """ Async get key """
+        return self.execute_blocking(self.key_client.get_key, name)
+
+    async def get_random_key_bl(self) -> KeyVaultKey:
+        """ Blocking get random key """
+        keys = self.key_client.list_properties_of_keys()
+        await self.execute_blocking(self.key_client.list_properties_of_keys)
+        all_keys = []
+        for key in keys:
+            all_keys.append(key)
+        random_key = random.choice(all_keys)
+        return await self.execute_blocking(self.key_client.get_key,
+                                           random_key.name)
+
+    def get_random_key(self) -> Awaitable["KeyVaultKey"]:
+        """ Async get random key """
+        return self.execute_blocking(self.get_random_key_bl)
+
+    def get_cipher(self, key: KeyVaultKey) -> CryptographyClient:
+        """ Get Cipher """
+        return CryptographyClient(key, self.credential)
+
+    def encrypt_bl(self, key: KeyVaultKey, data: bytes) -> bytes:
+        """ Encrypt data """
+        cipher = CryptographyClient(key, self.credential)
+        result = cipher.encrypt(EncryptionAlgorithm.rsa_oaep, data)
+        return result.ciphertext