Auditing collected data

def decrypt(key, iv, data):
    """
    Decrypts the given data using the provided key and initialization vector (IV).

    This function takes in a key, initialization vector (IV), and encrypted data, 
    and returns the decrypted data as a JSON string.

    Args:
        key (str): base64 encoded encryption key used for decryption (from API response).
        iv (str): base64 encoded initialization vector used for decryption (from API response).
        data (str): data returned by the Credolab SDK.

    Returns:
        str: A JSON string representing the decrypted data.
    """
    # get encrypted data bytes
    encrypted_data_bytes = base64.b64decode(data)

    # get encrypted payload data from the data returned by SDK
    # 292 - magic number defined by Credolab, the bytes before are related to some meta information
    payload_start = 292
    encrypted_payload = encrypted_data_bytes[payload_start:]

    # the data is encrypted using AES encryption algorithm, to decrypt it we need key and iv
    # decode key bytes from base64 encoded string
    key_bytes = base64.b64decode(key)
    iv_bytes = base64.b64decode(iv)

    # create cheaper for encryption
    cipher = AES.new(key_bytes, AES.MODE_CBC, iv_bytes)

    # perform decryption
    compressed_payload =  unpad(cipher.decrypt(encrypted_payload), AES.block_size)

    # the decrypted data is json string compressed with the usage of deflate compression
    # to get the original data we need to decompress data
    payload_bytes = zlib.decompress(compressed_payload, -zlib.MAX_WBITS)

    # decode to json string using utf-8 encoding
    json_data = payload_bytes.decode('utf-8')

    return json_data

# the required dependency - pycryptodome==3.20.0

import base64
import zlib
import requests
import json

from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad


def extract_keys(cl_url, auth_key, data):
    """
    Extracts encryption keys from the collected data using Credolab API.
    
    Args:
        cl_url (str): the URL of the Credolab service.
        auth_key (str): Bearer authentication key received from Credolab CSM.
        data (dict): data collected by the Credolab SDK.
    
    Returns:
        dict: A dictionary containing the extracted AES key and initialization vector required for decryption.
    """
    # set authorization header required for receiving receiving decryption keys for dataset
    headers = {
        'Authorization': f'Bearer {auth_key}'
    }
    
    # build json payload
    payload = {
        'data': data
    }
    
    response = requests.post(
        url=f'{cl_url}/api/datasets/v1/decryptionkeys', 
        headers=headers,
        json=payload)
    
    keys = json.loads(response.content)

    print(keys)
    
    return keys['aes'], keys['iv']


def decrypt(key, iv, data):
    """
    Decrypts the given data using the provided key and initialization vector (IV).

    This function takes in a key, initialization vector (IV), and encrypted data, 
    and returns the decrypted data as a JSON string.

    Args:
        key (str): base64 encoded encryption key used for decryption (from API response).
        iv (str): base64 encoded initialization vector used for decryption (from API response).
        data (str): data returned by the Credolab SDK.

    Returns:
        str: A JSON string representing the decrypted data.
    """
    # get encrypted data bytes
    encrypted_data_bytes = base64.b64decode(data)

    # get encrypted payload data from the data returned by SDK
    # 292 - magic number defined by Credolab, the bytes before are related to some meta information
    payload_start = 292
    encrypted_payload = encrypted_data_bytes[payload_start:]

    # the data is encrypted using AES encryption algorithm, to decrypt it we need key and iv
    # decode key bytes from base64 encoded string
    key_bytes = base64.b64decode(key)
    iv_bytes = base64.b64decode(iv)

    # create cheaper for encryption
    cipher = AES.new(key_bytes, AES.MODE_CBC, iv_bytes)

    # perform decryption
    compressed_payload =  unpad(cipher.decrypt(encrypted_payload), AES.block_size)

    # the decrypted data is json string compressed with the usage of deflate compression
    # to get the original data we need to decompress data
    payload_bytes = zlib.decompress(compressed_payload, -zlib.MAX_WBITS)

    # decode to json string using utf-8 encoding
    json_data = payload_bytes.decode('utf-8')

    return json_data


def main():
    cl_url = 'https://******.credolab.com' # the url to the Credolab server
    auth_key = '*************************' # auth key provided by Credolab CSM
    data = '*****************************' # the data returned by credolab sdk

    aes_key, iv = extract_keys(cl_url, auth_key, data)
    decrypted_data = decrypt(aes_key, iv, data)

    print(decrypted_data)

if __name__ == '__main__':
    main()