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
16import os
17from json import (
18    dumps,
19    loads,
20)
21
22from azure.common import (
23    AzureException,
24)
25from cryptography.hazmat.primitives.padding import PKCS7
26
27from ..common._common_conversion import (
28    _encode_base64,
29    _decode_base64_to_bytes
30)
31from ..common._encryption import (
32    _generate_encryption_data_dict,
33    _dict_to_encryption_data,
34    _generate_AES_CBC_cipher,
35    _validate_and_unwrap_cek,
36    _EncryptionAlgorithm,
37)
38from ..common._error import (
39    _ERROR_DECRYPTION_FAILURE,
40    _ERROR_UNSUPPORTED_ENCRYPTION_ALGORITHM,
41    _validate_not_none,
42    _validate_key_encryption_key_wrap,
43)
44from ._error import (
45    _ERROR_MESSAGE_NOT_ENCRYPTED
46)
47
48
49def _encrypt_queue_message(message, key_encryption_key):
50    '''
51    Encrypts the given plain text message using AES256 in CBC mode with 128 bit padding.
52    Wraps the generated content-encryption-key using the user-provided key-encryption-key (kek).
53    Returns a json-formatted string containing the encrypted message and the encryption metadata.
54
55    :param object message:
56        The plain text messge to be encrypted.
57    :param object key_encryption_key:
58        The user-provided key-encryption-key. Must implement the following methods:
59        wrap_key(key)--wraps the specified key using an algorithm of the user's choice.
60        get_key_wrap_algorithm()--returns the algorithm used to wrap the specified symmetric key.
61        get_kid()--returns a string key id for this key-encryption-key.
62    :return: A json-formatted string containing the encrypted message and the encryption metadata.
63    :rtype: str
64    '''
65
66    _validate_not_none('message', message)
67    _validate_not_none('key_encryption_key', key_encryption_key)
68    _validate_key_encryption_key_wrap(key_encryption_key)
69
70    # AES256 uses 256 bit (32 byte) keys and always with 16 byte blocks
71    content_encryption_key = os.urandom(32)
72    initialization_vector = os.urandom(16)
73
74    # Queue encoding functions all return unicode strings, and encryption should
75    # operate on binary strings.
76    message = message.encode('utf-8')
77
78    cipher = _generate_AES_CBC_cipher(content_encryption_key, initialization_vector)
79
80    # PKCS7 with 16 byte blocks ensures compatibility with AES.
81    padder = PKCS7(128).padder()
82    padded_data = padder.update(message) + padder.finalize()
83
84    # Encrypt the data.
85    encryptor = cipher.encryptor()
86    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
87
88    # Build the dictionary structure.
89    queue_message = {'EncryptedMessageContents': _encode_base64(encrypted_data),
90                     'EncryptionData': _generate_encryption_data_dict(key_encryption_key,
91                                                                      content_encryption_key,
92                                                                      initialization_vector)}
93
94    return dumps(queue_message)
95
96
97def _decrypt_queue_message(message, require_encryption, key_encryption_key, resolver):
98    '''
99    Returns the decrypted message contents from an EncryptedQueueMessage.
100    If no encryption metadata is present, will return the unaltered message.
101    :param str message:
102        The JSON formatted QueueEncryptedMessage contents with all associated metadata.
103    :param bool require_encryption:
104        If set, will enforce that the retrieved messages are encrypted and decrypt them.
105    :param object key_encryption_key:
106        The user-provided key-encryption-key. Must implement the following methods:
107        unwrap_key(key, algorithm)--returns the unwrapped form of the specified symmetric key using the string-specified algorithm.
108        get_kid()--returns a string key id for this key-encryption-key.
109    :param function resolver(kid):
110        The user-provided key resolver. Uses the kid string to return a key-encryption-key implementing the interface defined above.
111    :return: The plain text message from the queue message.
112    :rtype: str
113    '''
114
115    try:
116        message = loads(message)
117
118        encryption_data = _dict_to_encryption_data(message['EncryptionData'])
119        decoded_data = _decode_base64_to_bytes(message['EncryptedMessageContents'])
120    except (KeyError, ValueError):
121        # Message was not json formatted and so was not encrypted
122        # or the user provided a json formatted message.
123        if require_encryption:
124            raise ValueError(_ERROR_MESSAGE_NOT_ENCRYPTED)
125        else:
126            return message
127    try:
128        return _decrypt(decoded_data, encryption_data, key_encryption_key, resolver).decode('utf-8')
129    except Exception:
130        raise AzureException(_ERROR_DECRYPTION_FAILURE)
131
132
133def _decrypt(message, encryption_data, key_encryption_key=None, resolver=None):
134    '''
135    Decrypts the given ciphertext using AES256 in CBC mode with 128 bit padding.
136    Unwraps the content-encryption-key using the user-provided or resolved key-encryption-key (kek). Returns the original plaintex.
137
138    :param str message:
139        The ciphertext to be decrypted.
140    :param _EncryptionData encryption_data:
141        The metadata associated with this ciphertext.
142    :param object key_encryption_key:
143        The user-provided key-encryption-key. Must implement the following methods:
144        unwrap_key(key, algorithm)--returns the unwrapped form of the specified symmetric key using the string-specified algorithm.
145        get_kid()--returns a string key id for this key-encryption-key.
146    :param function resolver(kid):
147        The user-provided key resolver. Uses the kid string to return a key-encryption-key implementing the interface defined above.
148    :return: The decrypted plaintext.
149    :rtype: str
150    '''
151    _validate_not_none('message', message)
152    content_encryption_key = _validate_and_unwrap_cek(encryption_data, key_encryption_key, resolver)
153
154    if not (_EncryptionAlgorithm.AES_CBC_256 == encryption_data.encryption_agent.encryption_algorithm):
155        raise ValueError(_ERROR_UNSUPPORTED_ENCRYPTION_ALGORITHM)
156
157    cipher = _generate_AES_CBC_cipher(content_encryption_key, encryption_data.content_encryption_IV)
158
159    # decrypt data
160    decrypted_data = message
161    decryptor = cipher.decryptor()
162    decrypted_data = (decryptor.update(decrypted_data) + decryptor.finalize())
163
164    # unpad data
165    unpadder = PKCS7(128).unpadder()
166    decrypted_data = (unpadder.update(decrypted_data) + unpadder.finalize())
167
168    return decrypted_data
169