escrowai-encrypt 0.0.1__py3-none-any.whl

escrowai_encrypt/encryption.py

@@ -0,0 +1,249 @@
+import os
+import uuid
+import yaml
+import pathlib
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.serialization import load_pem_public_key
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.backends import default_backend
+from zipfile import ZipFile
+from azure.storage.blob import ContainerClient
+from escrowai_encrypt.utils import encrypt_data
+import shutil
+
+
+def generate_content_encryption_key(filename=""):
+    """
+    Generate a Content Encryption Key (CEK), the key used to encrypt an algorithm or dataset.
+    """
+
+    if filename == "":
+        filename = f"cek_{uuid.uuid4()}.key"
+
+    with open(filename, "wb") as file:
+        file.write(os.urandom(32))
+
+    print(f"Content Encryption Key (CEK) downloaded successfully to {filename}.")
+
+
+def generate_wrapped_content_encryption_key(
+    content_encryption_key: str, key_encryption_key: str, filename=""
+):
+    """
+    Generate a Wrapped Content Encryption Key (WCEK), a Content Encryption Key (CEK) that has been encrypted
+    using a Key Encryption Key (KEK).
+    """
+
+    if filename == "":
+        filename = f"wcek_{uuid.uuid4()}.key.bkenc"
+
+    with open(content_encryption_key, "rb") as file:
+        cek = file.read()
+    with open(key_encryption_key, "rb") as file:
+        kek = file.read()
+
+    print(
+        f"Wrapping CEK {content_encryption_key} with KEK {key_encryption_key} to file {filename}..."
+    )
+
+    public_key = load_pem_public_key(kek, backend=default_backend())
+    encrypted = public_key.encrypt(
+        cek,
+        padding=padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA1()),
+            algorithm=hashes.SHA1(),
+            label=None,
+        ),
+    )
+
+    with open(filename, "wb") as file:
+        file.write(encrypted)
+
+    print(
+        f"Wrapped Content Encryption Key (WCEK) downloaded successfully to {filename}."
+    )
+
+
+def encrypt_algorithm(
+    algorithm_directory: str, content_encryption_key: str, filename=""
+):
+    """
+    Encrypt an Algorithm. Algorithm encryption safeguards your algorithm during storage, transmission,
+    and execution. The algorithm is encrypted using a Content Encryption Key (CEK).
+    """
+
+    if filename == "":
+        filename = f"{algorithm_directory}_encrypted_{uuid.uuid4()}.zip"
+
+    with open(content_encryption_key, "rb") as file:
+        cek = file.read()
+
+    original_files = []
+    new_files = []
+    with open(f"{algorithm_directory}/secrets.yaml", "r") as f:
+        files = yaml.safe_load(f)
+
+    count = str(len(files["secretFiles"]))
+    n = 1
+
+    print(f"{count} secret(s) to be encrypted.")
+    for f in files["secretFiles"]:
+        print(f"Encrypting secret {f[1]} ({n}/{count})...")
+        newpaths = ["", ""]
+        newpaths[0] = f"{algorithm_directory}/{f[0]}"
+        newpaths[1] = f"{algorithm_directory}/{f[1]}"
+
+        salt = os.urandom(8)
+        pbkdf2_hash = PBKDF2HMAC(
+            algorithm=hashes.SHA256(),
+            salt=salt,
+            length=32 + 12,
+            iterations=10000,
+            backend=default_backend(),
+        )
+
+        derived_password = pbkdf2_hash.derive(cek)
+        key = derived_password[0:32]
+        iv = derived_password[32 : (32 + 12)]
+
+        with open(newpaths[1], "rb") as encrypt:
+            data = encrypt.read()
+
+        encrypted = AESGCM(key).encrypt(iv, data, None)
+        encrypted = b"Salted__" + salt + encrypted
+
+        with open(newpaths[0], "wb") as write:
+            write.write(encrypted)
+
+        original_files.append(f[1])
+        new_files.append(newpaths[0])
+
+        print(f"Secret {f[0]} encrypted ({n}/{count}).")
+
+    shutil.make_archive("temp", "zip", algorithm_directory)
+    for i in new_files:
+        pathlib.Path(i).unlink()
+
+    zip_in = ZipFile("temp.zip", "r")
+    zip_out = ZipFile(filename, "w")
+
+    for item in zip_in.infolist():
+        buffer = zip_in.read(item.filename)
+        if item.filename not in original_files:
+            zip_out.writestr(item, buffer)
+    zip_in.close()
+    zip_out.close()
+
+    pathlib.Path("temp.zip").unlink()
+
+    print(f"Algorithm successfully encrypted as {filename}.")
+
+
+def encrypt_upload_dataset(
+    dataset_directory: str, content_encryption_key: str, dataset_sas_uri: str
+):
+    """
+    Encrypt and upload a Dataset to Azure Blob Storage. Dataset encryption safeguards your dataset during storage,
+    transmission, and execution. The dataset is encrypted using a Content Encryption Key (CEK) and uploaded using
+    a Shared Access Signature (SAS) URL.
+    """
+
+    with open(content_encryption_key, "rb") as file:
+        cek = file.read()
+
+    salt = os.urandom(8)
+    pbkdf2_hash = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        salt=salt,
+        length=32 + 12,
+        iterations=10000,
+        backend=default_backend(),
+    )
+
+    derived_password = pbkdf2_hash.derive(cek)
+    key = derived_password[0:32]
+    iv = derived_password[32 : (32 + 12)]
+
+    count = sum(
+        len([x for x in files if x[0] != "."])
+        for _, _, files in os.walk(dataset_directory)
+    )
+    n = 1
+
+    print(f"{count} files(s) to be encrypted and uploaded.")
+    for root, _, files in os.walk(dataset_directory):
+        for filename in files:
+            file_path = os.path.join(root, filename)
+            dirname = os.path.dirname(file_path)
+            n += encrypt_data(
+                filename,
+                salt,
+                key,
+                dataset_directory,
+                dirname,
+                iv,
+                dataset_sas_uri,
+                n,
+                count,
+            )
+
+    print(f"Dataset successfully encrypted and uploaded to SAS URL {dataset_sas_uri}.")
+
+
+def encrypt_upload_dataset_from_blob(
+    dataset_sas_uri_unencrypted: str, content_encryption_key: str, dataset_sas_uri: str
+):
+    """
+    Encrypt and upload a Dataset to Azure Blob Storage. Dataset encryption safeguards your dataset during storage,
+    transmission, and execution. The dataset is encrypted using a Content Encryption Key (CEK) and uploaded using
+    a Shared Access Signature (SAS) URL.
+    """
+
+    with open(content_encryption_key, "rb") as file:
+        cek = file.read()
+
+    salt = os.urandom(8)
+    pbkdf2_hash = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        salt=salt,
+        length=32 + 12,
+        iterations=10000,
+        backend=default_backend(),
+    )
+
+    derived_password = pbkdf2_hash.derive(cek)
+    key = derived_password[0:32]
+    iv = derived_password[32 : (32 + 12)]
+
+    unencrypted_client = ContainerClient.from_container_url(dataset_sas_uri_unencrypted)
+    encrypted_client = ContainerClient.from_container_url(dataset_sas_uri)
+    blobs = unencrypted_client.list_blob_names()
+
+    count = 0
+    for blob in blobs:
+        count += 1
+
+    blobs = unencrypted_client.list_blob_names()  # reset page
+
+    n = 1
+    print(f"{str(count)} blob(s) to be encrypted and uploaded.")
+    for blob in blobs:
+        unencrypted_client = ContainerClient.from_container_url(
+            dataset_sas_uri_unencrypted
+        )
+        encrypted_client = ContainerClient.from_container_url(dataset_sas_uri)
+        print(f"Beginning download of blob {blob} ({n}/{str(count)})...")
+        data = unencrypted_client.download_blob(blob).readall()
+        print(f"Encrypting blob {blob}...")
+        encrypted = AESGCM(key).encrypt(iv, data, None)
+        encrypted = b"Salted__" + salt + encrypted
+        print(f"Uploading blob {blob}...")
+        encrypted_client.upload_blob(blob + ".bkenc", encrypted, overwrite=True)
+        print(f"Blob {blob + '.bkenc'} uploaded ({n}/{str(count)}).")
+        n += 1
+
+    print(
+        f"Dataset {dataset_sas_uri_unencrypted} successfully encrypted and uploaded to SAS URL {dataset_sas_uri}."
+    )

escrowai_encrypt/utils.py

@@ -0,0 +1,39 @@
+import pathlib
+from azure.storage.blob import BlobClient
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+
+def encrypt_data(file, salt, key, base_folder, folder, iv, url, n, count):
+    if file[0] == ".":
+        return 0
+    if file.endswith(".bkenc"):
+        raise Exception(f"Error: cannot encrypt file {file}.")
+
+    print(f"Encrypting file {file} ({n}/{count})...")
+
+    with open(folder + "/" + file, "rb") as encrypt:
+        data = encrypt.read()
+
+    encrypted = AESGCM(key).encrypt(iv, data, None)
+    encrypted = b"Salted__" + salt + encrypted
+
+    # write new encrypted file, delete old
+    with open(folder + "/" + file + ".bkenc", "wb") as write:
+        write.write(encrypted)
+
+    upload_to_blob(url, folder + "/" + file + ".bkenc", base_folder)
+
+    pathlib.Path(folder + "/" + file + ".bkenc").unlink()
+
+    print(f"Uploaded file {file + '.bkenc'} ({n}/{count}).")
+
+    return 1
+
+
+def upload_to_blob(url, file, folder):
+    uri = url.partition("?")
+    new_uri = uri[0] + file.split(folder, 1)[1] + uri[1] + uri[2]
+    client = BlobClient.from_blob_url(new_uri)
+
+    with open(file, "rb") as upload:
+        client.upload_blob(upload, overwrite=True)

escrowai_encrypt-0.0.1.dist-info/METADATA

@@ -0,0 +1,29 @@
+Metadata-Version: 2.2
+Name: escrowai-encrypt
+Version: 0.0.1
+Summary: This package allows EscrowAI users to programmatically encrypt their artifacts and add encryption to their workflows.
+Home-page: https://github.com/BeeKeeperAI/escrowai-encrypt
+Author: BeeKeeperAI
+Author-email: engineering@beekeeperai.com
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+Requires-Dist: azure-storage-blob
+Requires-Dist: cryptography
+Requires-Dist: pyyaml
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# escrowai-encrypt
+A module for encryption of artifacts and keys to be uploaded to EscrowAI.

escrowai_encrypt-0.0.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+escrowai_encrypt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+escrowai_encrypt/encryption.py,sha256=HGP2AH5wEchaRIGGwFNoV0zGJMDDjvyHH7L5t-DnRqo,7901
+escrowai_encrypt/main.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+escrowai_encrypt/utils.py,sha256=JQeozbUkwoSEnHG4JWFqPsJXC17Kv4lhzFJog5Lj0gA,1185
+escrowai_encrypt-0.0.1.dist-info/METADATA,sha256=ePf7jOzp-IGMiPX53q1GDcyFjQ8OXyiykDTU2lYe9_I,988
+escrowai_encrypt-0.0.1.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+escrowai_encrypt-0.0.1.dist-info/entry_points.txt,sha256=mjB_P4wzCRMYNUOBCdP6rz8xEyuTVwU7qZM4F5VWKBE,55
+escrowai_encrypt-0.0.1.dist-info/top_level.txt,sha256=R8QP9gXiLcbs1sVwwxPHtfEgMf5Ww0eqBhnUQkD318I,17
+escrowai_encrypt-0.0.1.dist-info/RECORD,,

escrowai_encrypt-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

escrowai_encrypt-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+encrypt = escrowai_encrypt.main:main

escrowai_encrypt-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+escrowai_encrypt