1#-------------------------------------------------------------------------
2# Copyright (c) Microsoft.  All rights reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#   http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#--------------------------------------------------------------------------
15
16from .._error import(
17    _ERROR_UNSUPPORTED_ENCRYPTION_VERSION,
18    _ERROR_DECRYPTION_FAILURE,
19    _ERROR_UNSUPPORTED_ENCRYPTION_ALGORITHM,
20    _ERROR_DATA_NOT_ENCRYPTED,
21    _validate_not_none,
22    _validate_key_encryption_key_wrap,
23    _validate_key_encryption_key_unwrap,
24    _validate_kek_id,
25)
26from .._constants import(
27    _ENCRYPTION_PROTOCOL_V1,
28)
29from .._common_conversion import(
30    _decode_base64_to_bytes,
31)
32from .._encryption import(
33    _generate_encryption_data_dict,
34    _dict_to_encryption_data,
35    _generate_AES_CBC_cipher,
36    _validate_and_unwrap_cek,
37    _EncryptionData,
38    _EncryptionAgent,
39    _WrappedContentKey,
40    _EncryptionAlgorithm
41)
42from ._error import(
43    _ERROR_UNSUPPORTED_TYPE_FOR_ENCRYPTION,
44)
45from .models import(
46    Entity,
47    EntityProperty,
48    EdmType,
49)
50from json import(
51    dumps,
52    loads,
53)
54import os
55from copy import deepcopy
56from cryptography.hazmat.backends import default_backend
57from cryptography.hazmat.primitives.padding import PKCS7
58from cryptography.hazmat.primitives.hashes import(
59    Hash,
60    SHA256,
61)
62
63def _encrypt_entity(entity, key_encryption_key, encryption_resolver):
64    '''
65    Encrypts the given entity using AES256 in CBC mode with 128 bit padding.
66    Will generate a content-encryption-key (cek) to encrypt the properties either
67    stored in an EntityProperty with the 'encrypt' flag set or those
68    specified by the encryption resolver. This cek is then wrapped using the
69    provided key_encryption_key (kek). Only strings may be encrypted and the
70    result is stored as binary on the service.
71
72    :param entity:
73        The entity to insert. Could be a dict or an entity object.
74    :param object key_encryption_key:
75        The user-provided key-encryption-key. Must implement the following methods:
76        wrap_key(key)--wraps the specified key using an algorithm of the user's choice.
77        get_key_wrap_algorithm()--returns the algorithm used to wrap the specified symmetric key.
78        get_kid()--returns a string key id for this key-encryption-key.
79    :param function(partition_key, row_key, property_name) encryption_resolver:
80        A function that takes in an entities partition key, row key, and property name and returns
81        a boolean that indicates whether that property should be encrypted.
82    :return: An entity with both the appropriate properties encrypted and the
83        encryption data.
84    :rtype: object
85    '''
86
87    _validate_not_none('entity', entity)
88    _validate_not_none('key_encryption_key', key_encryption_key)
89    _validate_key_encryption_key_wrap(key_encryption_key)
90
91    # AES256 uses 256 bit (32 byte) keys and always with 16 byte blocks
92    content_encryption_key = os.urandom(32)
93    entity_initialization_vector = os.urandom(16)
94
95    encrypted_properties = []
96    encrypted_entity = Entity()
97    for key, value in entity.items():
98        # If the property resolver says it should be encrypted
99        # or it is an EntityProperty with the 'encrypt' property set.
100        if (isinstance(value, EntityProperty) and value.encrypt) or \
101            (encryption_resolver is not None \
102            and encryption_resolver(entity['PartitionKey'], entity['RowKey'], key)):
103
104            # Only strings can be encrypted and None is not an instance of str.
105            if isinstance(value, EntityProperty):
106                if value.type == EdmType.STRING:
107                    value = value.value
108                else:
109                    raise ValueError(_ERROR_UNSUPPORTED_TYPE_FOR_ENCRYPTION)
110            if not isinstance(value, str):
111                raise ValueError(_ERROR_UNSUPPORTED_TYPE_FOR_ENCRYPTION)
112
113            # Value is now confirmed to hold a valid string value to be encrypted
114            # and should be added to the list of encrypted properties.
115            encrypted_properties.append(key)
116
117            propertyIV = _generate_property_iv(entity_initialization_vector,
118                                               entity['PartitionKey'], entity['RowKey'],
119                                               key, False)
120
121            # Encode the strings for encryption.
122            value = value.encode('utf-8')
123
124            cipher = _generate_AES_CBC_cipher(content_encryption_key, propertyIV)
125
126            # PKCS7 with 16 byte blocks ensures compatibility with AES.
127            padder = PKCS7(128).padder()
128            padded_data = padder.update(value) + padder.finalize()
129
130            # Encrypt the data.
131            encryptor = cipher.encryptor()
132            encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
133
134            # Set the new value of this key to be a binary EntityProperty for proper serialization.
135            value = EntityProperty(EdmType.BINARY, encrypted_data)
136
137        encrypted_entity[key] = value
138
139    encrypted_properties = dumps(encrypted_properties)
140
141    # Generate the metadata iv.
142    metadataIV = _generate_property_iv(entity_initialization_vector,
143                          entity['PartitionKey'], entity['RowKey'],
144                          '_ClientEncryptionMetadata2', False)
145
146    encrypted_properties = encrypted_properties.encode('utf-8')
147
148    cipher = _generate_AES_CBC_cipher(content_encryption_key, metadataIV)
149
150    padder = PKCS7(128).padder()
151    padded_data = padder.update(encrypted_properties) + padder.finalize()
152
153    encryptor = cipher.encryptor()
154    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
155
156    encrypted_entity['_ClientEncryptionMetadata2'] = EntityProperty(EdmType.BINARY, encrypted_data)
157
158    encryption_data = _generate_encryption_data_dict(key_encryption_key, content_encryption_key,
159                                                     entity_initialization_vector)
160
161    encrypted_entity['_ClientEncryptionMetadata1'] = dumps(encryption_data)
162    return encrypted_entity
163
164def _decrypt_entity(entity, encrypted_properties_list, content_encryption_key, entityIV, isJavaV1):
165    '''
166    Decrypts the specified entity using AES256 in CBC mode with 128 bit padding. Unwraps the CEK
167    using either the specified KEK or the key returned by the key_resolver. Properties
168    specified in the encrypted_properties_list, will be decrypted and decoded to utf-8 strings.
169
170    :param entity:
171        The entity being retrieved and decrypted. Could be a dict or an entity object.
172    :param list encrypted_properties_list:
173        The encrypted list of all the properties that are encrypted.
174    :param bytes[] content_encryption_key:
175        The key used internally to encrypt the entity. Extrated from the entity metadata.
176    :param bytes[] entityIV:
177        The intialization vector used to seed the encryption algorithm. Extracted from the
178        entity metadata.
179    :return: The decrypted entity
180    :rtype: Entity
181    '''
182
183    _validate_not_none('entity', entity)
184
185    decrypted_entity = deepcopy(entity)
186    try:
187        for property in entity.keys():
188            if property in encrypted_properties_list:
189                value = entity[property]
190
191                propertyIV = _generate_property_iv(entityIV,
192                                                   entity['PartitionKey'], entity['RowKey'],
193                                                   property, isJavaV1)
194                cipher = _generate_AES_CBC_cipher(content_encryption_key,
195                                                  propertyIV)
196
197                # Decrypt the property.
198                decryptor = cipher.decryptor()
199                decrypted_data = (decryptor.update(value.value) + decryptor.finalize())
200
201                # Unpad the data.
202                unpadder = PKCS7(128).unpadder()
203                decrypted_data = (unpadder.update(decrypted_data) + unpadder.finalize())
204
205                decrypted_data = decrypted_data.decode('utf-8')
206
207                decrypted_entity[property] = decrypted_data
208
209        decrypted_entity.pop('_ClientEncryptionMetadata1')
210        decrypted_entity.pop('_ClientEncryptionMetadata2')
211        return decrypted_entity
212    except:
213        raise AzureException(_ERROR_DECRYPTION_FAILURE)
214
215def _extract_encryption_metadata(entity, require_encryption, key_encryption_key, key_resolver):
216    '''
217    Extracts the encryption metadata from the given entity, setting them to be utf-8 strings.
218    If no encryption metadata is present, will return None for all return values unless
219    require_encryption is true, in which case the method will throw.
220
221    :param entity:
222        The entity being retrieved and decrypted. Could be a dict or an entity object.
223    :param bool require_encryption:
224        If set, will enforce that the retrieved entity is encrypted and decrypt it.
225    :param object key_encryption_key:
226        The user-provided key-encryption-key. Must implement the following methods:
227        unwrap_key(key, algorithm)--returns the unwrapped form of the specified symmetric key using the
228        string-specified algorithm.
229        get_kid()--returns a string key id for this key-encryption-key.
230    :param function key_resolver(kid):
231        The user-provided key resolver. Uses the kid string to return a key-encryption-key implementing
232        the interface defined above.
233    :returns: a tuple containing the entity iv, the list of encrypted properties, the entity cek,
234        and whether the entity was encrypted using JavaV1.
235    :rtype: tuple (bytes[], list, bytes[], bool)
236    '''
237    _validate_not_none('entity', entity)
238
239    try:
240        encrypted_properties_list = _decode_base64_to_bytes(entity['_ClientEncryptionMetadata2'])
241        encryption_data = entity['_ClientEncryptionMetadata1']
242        encryption_data = _dict_to_encryption_data(loads(encryption_data))
243    except Exception as e:
244        # Message did not have properly formatted encryption metadata.
245        if require_encryption:
246            raise ValueError(_ERROR_ENTITY_NOT_ENCRYPTED)
247        else:
248            return (None,None,None,None)
249
250    if not(encryption_data.encryption_agent.encryption_algorithm == _EncryptionAlgorithm.AES_CBC_256):
251        raise ValueError(_ERROR_UNSUPPORTED_ENCRYPTION_ALGORITHM)
252
253    content_encryption_key = _validate_and_unwrap_cek(encryption_data, key_encryption_key, key_resolver)
254
255    # Special check for compatibility with Java V1 encryption protocol.
256    isJavaV1 = (encryption_data.key_wrapping_metadata is None) or \
257        ((encryption_data.encryption_agent.protocol == _ENCRYPTION_PROTOCOL_V1) and \
258        'EncryptionLibrary' in encryption_data.key_wrapping_metadata and \
259        'Java' in encryption_data.key_wrapping_metadata['EncryptionLibrary'])
260
261    metadataIV = _generate_property_iv(encryption_data.content_encryption_IV,
262                                      entity['PartitionKey'], entity['RowKey'],
263                                      '_ClientEncryptionMetadata2', isJavaV1)
264
265    cipher = _generate_AES_CBC_cipher(content_encryption_key, metadataIV)
266
267    # Decrypt the data.
268    decryptor = cipher.decryptor()
269    encrypted_properties_list = decryptor.update(encrypted_properties_list) + decryptor.finalize()
270
271    # Unpad the data.
272    unpadder = PKCS7(128).unpadder()
273    encrypted_properties_list = unpadder.update(encrypted_properties_list) + unpadder.finalize()
274
275    encrypted_properties_list = encrypted_properties_list.decode('utf-8')
276
277    if isJavaV1:
278        # Strip the square braces from the ends and split string into list.
279        encrypted_properties_list = encrypted_properties_list[1:-1]
280        encrypted_properties_list = encrypted_properties_list.split(', ')
281    else:
282        encrypted_properties_list = loads(encrypted_properties_list)
283
284    return (encryption_data.content_encryption_IV, encrypted_properties_list, content_encryption_key, isJavaV1)
285
286def _generate_property_iv(entity_iv, pk, rk, property_name, isJavaV1):
287    '''
288    Uses the entity_iv, partition key, and row key to generate and return
289    the iv for the specified property.
290    '''
291    digest = Hash(SHA256(), default_backend())
292    if not isJavaV1:
293        digest.update(entity_iv +
294                      (rk + pk + property_name).encode('utf-8'))
295    else:
296        digest.update(entity_iv +
297                      (pk + rk + property_name).encode('utf-8'))
298    propertyIV = digest.finalize()
299    return propertyIV[:16]