# -------------------------------------------------------------------------
# Copyright (c) Microsoft.  All rights reserved.
#
# 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.
# --------------------------------------------------------------------------

import uuid
from os import urandom

from azure.common import AzureException
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric.padding import (
    OAEP,
    MGF1,
)
from cryptography.hazmat.primitives.asymmetric.rsa import generate_private_key
from cryptography.hazmat.primitives.hashes import SHA1
from cryptography.hazmat.primitives.keywrap import (
    aes_key_wrap,
    aes_key_unwrap,
)

from azure.cosmosdb.table import (
    Entity,
    TableBatch,
    EdmType,
    EntityProperty,
    TablePayloadFormat,
)


# Sample implementations of the encryption-related interfaces.
class KeyWrapper:
    def __init__(self, kid):
        self.kek = urandom(32)
        self.backend = default_backend()
        self.kid = 'local:' + kid

    def wrap_key(self, key, algorithm='A256KW'):
        if algorithm == 'A256KW':
            return aes_key_wrap(self.kek, key, self.backend)
        else:
            raise ValueError(_ERROR_UNKNOWN_KEY_WRAP_ALGORITHM)

    def unwrap_key(self, key, algorithm):
        if algorithm == 'A256KW':
            return aes_key_unwrap(self.kek, key, self.backend)
        else:
            raise ValueError(_ERROR_UNKNOWN_KEY_WRAP_ALGORITHM)

    def get_key_wrap_algorithm(self):
        return 'A256KW'

    def get_kid(self):
        return self.kid


class KeyResolver:
    def __init__(self):
        self.keys = {}

    def put_key(self, key):
        self.keys[key.get_kid()] = key

    def resolve_key(self, kid):
        return self.keys[kid]


class RSAKeyWrapper:
    def __init__(self, kid):
        self.private_key = generate_private_key(public_exponent=65537,
                                                key_size=2048,
                                                backend=default_backend())
        self.public_key = self.private_key.public_key()
        self.kid = 'local:' + kid

    def wrap_key(self, key, algorithm='RSA'):
        if algorithm == 'RSA':
            return self.public_key.encrypt(key,
                                           OAEP(
                                               mgf=MGF1(algorithm=SHA1()),
                                               algorithm=SHA1(),
                                               label=None)
                                           )
        else:
            raise ValueError(_ERROR_UNKNOWN_KEY_WRAP_ALGORITHM)

    def unwrap_key(self, key, algorithm):
        if algorithm == 'RSA':
            return self.private_key.decrypt(key,
                                            OAEP(
                                                mgf=MGF1(algorithm=SHA1()),
                                                algorithm=SHA1(),
                                                label=None)
                                            )
        else:
            raise ValueError(_ERROR_UNKNOWN_KEY_WRAP_ALGORITHM)

    def get_key_wrap_algorithm(self):
        return 'RSA'

    def get_kid(self):
        return self.kid


class TableEncryptionSamples():
    def __init__(self, account):
        self.account = account

    def run_all_samples(self):
        self.service = self.account.create_table_service()

        self.put_encrypted_entity_properties()
        self.put_encrypted_entity_encryption_resolver()
        self.get_encrypted_entity()
        self.get_encrypted_entity_key_encryption_key()
        self.replace_encrypted_entity()
        self.query_encrypted_entities()
        self.batch_encrypted_entities()
        self.require_encryption()
        self.alternate_key_encryption_algorithms()
        self.merge_not_supported()

    def _get_table_reference(self, prefix='table'):
        table_name = '{}{}'.format(prefix, str(uuid.uuid4()).replace('-', ''))
        return table_name

    def _create_table(self, prefix='table'):
        table_name = self._get_table_reference(prefix)
        self.service.create_table(table_name)
        return table_name

    def _create_base_entity_dict(self):
        entity = {}
        # Partition key and row key must be strings and are required
        entity['PartitionKey'] = 'pk{}'.format(str(uuid.uuid4()).replace('-', ''))
        entity['RowKey'] = 'rk{}'.format(str(uuid.uuid4()).replace('-', ''))
        return entity

    def _create_base_entity_class(self):
        # Partition key and row key must be strings and are required
        entity = Entity()
        entity['PartitionKey'] = 'pk{}'.format(str(uuid.uuid4()).replace('-', ''))
        entity['RowKey'] = 'rk{}'.format(str(uuid.uuid4()).replace('-', ''))
        return entity

    def _create_entity_for_encryption(self):
        entity = self._create_base_entity_class()
        entity['foo'] = EntityProperty(EdmType.STRING, 'bar', True)
        return entity

    def _create_query_table_encrypted(self, entity_count):
        '''
        Creates a table with the specified name and adds entities with the
        default set of values. PartitionKey is set to 'MyPartition' and RowKey
        is set to a unique counter value starting at 1 (as a string). The
        'foo' attribute is set to be encrypted.
        '''
        table_name = self._create_table(prefix='querytable')
        self.service.require_encryption = True

        entity = self._create_entity_for_encryption()
        with self.service.batch(table_name) as batch:
            for i in range(1, entity_count + 1):
                entity['RowKey'] = entity['RowKey'] + str(i)
                batch.insert_entity(entity)
        return table_name

    # A sample encryption resolver. This resolver is a simple case that will mark
    # any property named 'foo' for encryption, regardless of the partition or row 
    # it is in.
    def encryption_resolver(self, pk, rk, property):
        return property == 'foo'

    def put_encrypted_entity_properties(self):
        table_name = self._create_table()

        # Can use a dict or the Entity class to encrypt entities.
        # The EntityProperty object takes an optional parameteter, 'encrypt'
        # that marks the property for encryption when set to true.
        entity1 = self._create_base_entity_dict()
        entity1['foo'] = EntityProperty(EdmType.STRING, 'bar', True)
        entity2 = self._create_base_entity_class()
        entity2.foo = EntityProperty(EdmType.STRING, 'bar', True)
        entity3 = self._create_base_entity_class()
        entity3['badValue'] = EntityProperty(EdmType.INT64, 12, True)
        entity4 = self._create_base_entity_class()

        # KeyWrapper implements the key encryption key interface outlined
        # in the insert/get entity documentation.
        # Setting this property will tell these APIs to encrypt the entity.
        self.service.key_encryption_key = KeyWrapper('key1')
        self.service.insert_entity(table_name, entity1)
        self.service.insert_entity(table_name, entity2)

        # Note: The internal encryption process requires two properties, so there
        # are only 250 custom properties available when encrypting.
        # Note: str is the only type valid for encryption. Trying to encrypt other
        # properties will throw.

        self.service.delete_table(table_name)

    def put_encrypted_entity_encryption_resolver(self):
        table_name = self._create_table()

        entity = self._create_base_entity_class()
        entity['foo'] = 'bar'
        self.service.key_encryption_key = KeyWrapper('key1')

        # An encryption resolver is a function that takes in the Partition Key,
        # Row Key, and property name and returns true if the property should be 
        # encrypted and false otherwise. This can be used in place of explictly
        # setting each property to be encrypted through the EntityProperty class.
        self.service.encryption_resolver_function = self.encryption_resolver

        self.service.insert_entity(table_name, entity)

        self.service.delete_table(table_name)

    def get_encrypted_entity(self):
        table_name = self._create_table()
        entity = self._create_entity_for_encryption()
        self.service.key_encryption_key = KeyWrapper('key1')
        self.service.insert_entity(table_name, entity)

        # Entities can be decrypted by setting a key_resolver function on the service
        # without directly setting the key_encryption_key property itself. The function takes 
        # in the key_id (retrieved from the encrypted entity metadata) and returns the 
        # corresponding key_encryption_key. 
        key_resolver = KeyResolver()
        key_resolver.put_key(self.service.key_encryption_key)
        self.service.key_resolver_function = key_resolver.resolve_key
        self.service.key_encryption_key = None

        # Decrypted entities are stored in their raw string form, regardless of whether
        # they were stored in an EntityProperty when encrypted.

        # Retrieving and decrypting an encrypted entity works regardless of the accepted
        # payload format.
        entity_full = self.service.get_entity(table_name, entity['PartitionKey'], entity['RowKey'],
                                              accept=TablePayloadFormat.JSON_FULL_METADATA)
        entity_none = self.service.get_entity(table_name, entity['PartitionKey'], entity['RowKey'],
                                              accept=TablePayloadFormat.JSON_NO_METADATA)

        # Note: Properties that are encrypted on upload but not decrypted on download due to lack
        # of an encryption policy are stored in an EntityProperty with as an EdmBinary type.
        # Note: The encryption metadata headers are preserved on the entity if 
        # it is not decrypted when downloaded.
        # Note: Decrypted entities are stored in their raw string form, regardless of whether
        # they were stored in an EntityProperty when encrypted.

        self.service.key_resolver_function = None
        self.service.delete_table(table_name)

    def query_encrypted_entities(self):
        self.service.key_encryption_key = KeyWrapper('key1')
        key_resolver = KeyResolver()
        key_resolver.put_key(self.service.key_encryption_key)
        self.service.key_resolver_function = key_resolver.resolve_key
        table_name = self._create_query_table_encrypted(5)

        # Querying for entire entities will transparently decrypt retrieved entities.
        response = self.service.query_entities(table_name, num_results=5)

        # Performing a projection on a subset of properties will also implicilty
        # retrieve the encryption metatdata properties when an encryption policy is set.
        response = self.service.query_entities(table_name, num_results=5,
                                               select='PartitionKey,RowKey,foo')

        self.service.delete_table(table_name)

    def batch_encrypted_entities(self):
        table_name = self._create_table()
        entity1 = self._create_entity_for_encryption()
        entity2 = self._create_entity_for_encryption()
        entity2['PartitionKey'] = entity1['PartitionKey']

        # Batches will encrypt the entities at the time of inserting into the batch, not
        # committing the batch to the service, so the encryption policy must be 
        # passed in at the time of batch creation.
        kek = KeyWrapper('key1')
        batch = TableBatch(require_encryption=True, key_encryption_key=kek)
        batch.insert_entity(entity1)
        batch.insert_entity(entity2)
        self.service.commit_batch(table_name, batch)

        # When using the batch as a context manager, the tableservice object will
        # automatically apply its encryption policy to the batch. 
        entity3 = self._create_entity_for_encryption()
        entity4 = self._create_entity_for_encryption()
        entity4['PartitionKey'] = entity3['PartitionKey']
        self.service.key_encryption_key = KeyWrapper('key1')
        with self.service.batch(table_name) as batch:
            batch.insert_entity(entity3)
            batch.insert_entity(entity4)

        # Note that batches follow all the same client-side-encryption behavior as
        # the corresponding individual table operations.

        self.service.delete_table(table_name)

    def require_encryption(self):
        self.service.key_encryption_key = None
        self.service.key_resolver_function = None
        self.service.require_encryption = False
        table_name = self._create_table()
        entity_unencrypted = self._create_base_entity_class()
        entity_unencrypted['foo'] = 'bar'
        self.service.insert_entity(table_name, entity_unencrypted)

        # If the require_encryption flag is set, the service object will throw if there
        # is no encryption policy set on upload.
        self.service.key_encryption_key = None
        self.service.require_encryption = True
        try:
            self.service.insert_entity(table_name, entity_unencrypted)
            raise Exception
        except ValueError:
            pass

        # If the require_encryption flag is set, the service object will throw if there
        # is no encryption policy set on download.
        kek = KeyWrapper('key1')
        self.service.key_encryption_key = kek

        key_resolver = KeyResolver()
        key_resolver.put_key(self.service.key_encryption_key)
        self.service.key_resolver_function = key_resolver.resolve_key

        entity_encrypted = self._create_entity_for_encryption()
        self.service.insert_entity(table_name, entity_encrypted)

        self.service.key_encryption_key = None
        self.service.key_resolver_function = None
        try:
            self.service.get_entity(table_name, entity_encrypted['PartitionKey'],
                                    entity_encrypted['RowKey'])
            raise Exception
        except ValueError:
            pass

        # If the require_encryption flag is set, but the retrieved object is not encrypted,
        # the service object will throw.
        self.service.key_resolver_function = key_resolver.resolve_key
        try:
            self.service.get_entity(table_name, entity_unencrypted['PartitionKey'],
                                    entity_unencrypted['RowKey'])
            raise Exception
        except AzureException:
            pass

        self.service.delete_table(table_name)

    def alternate_key_encryption_algorithms(self):
        table_name = self._create_table()
        entity = self._create_entity_for_encryption()

        # The key wrapping algorithm used by the key_encryption_key is entirely
        # up to the choice of the user. For instance, RSA may be used.
        self.service.key_encryption_key = RSAKeyWrapper('key2')
        self.service.insert_entity(table_name, entity)

        key_resolver = KeyResolver()
        key_resolver.put_key(self.service.key_encryption_key)
        self.service.key_resolver_function = key_resolver.resolve_key
        entity = self.service.get_entity(table_name, entity['PartitionKey'], entity['RowKey'])

        self.service.delete_table(table_name)

    def merge_not_supported(self):
        table_name = self._create_table()
        entity = self._create_entity_for_encryption()
        self.service.key_encryption_key = KeyWrapper('key1')
        self.service.insert_entity(table_name, entity)

        # Merging encrypted entities is not supported. Calling merge with 
        # an encryption policy set will cause merge entities to fail.
        # If the require_encryption flag is set, merge_entities will fail.
        # Note that insert_or_merge exhibits the same encryption behavior.
        self.service.require_encryption = True
        try:
            self.service.merge_entity(table_name, entity)
            raise Exception
        except ValueError:
            pass

        self.service.require_encryption = False
        try:
            self.service.merge_entity(table_name, entity)
            raise Exception
        except ValueError:
            pass

        self.service.delete_table(table_name)

    def get_encrypted_entity_key_encryption_key(self):
        table_name = self._create_table()
        entity = self._create_entity_for_encryption()
        kek = KeyWrapper('key1')
        self.service.key_encryption_key = kek
        self.service.insert_entity(table_name, entity)

        # If the key_encryption_key property is set, the tableservice object will
        # try to decrypt entities using that key. If both the key_resolver and key_encryption_key
        # properties are set, the result of the key_resolver will take precedence and the decryption
        # will fail if that key is not successful.
        entity = self.service.get_entity(table_name, entity['PartitionKey'], entity['RowKey'])

        self.service.delete_table(table_name)

    def replace_encrypted_entity(self):
        table_name = self._create_table()
        entity = self._create_entity_for_encryption()
        self.service.key_encryption_key = KeyWrapper('key1')
        self.service.insert_entity(table_name, entity)

        # An entity, encrypted or decrypted, may be replaced by an encrypted entity.
        entity['foo'].value = 'updated'
        self.service.update_entity(table_name, entity)

        self.service.delete_table(table_name)