Add files via upload

heshan-basnayaka · web-flow · commit ba934d2cacd8 · 2024-05-18T21:26:22.000+05:30
diff --git a/Final_CPU_GPU_Combined_Multiprocess_AES_Algorithm.py b/Final_CPU_GPU_Combined_Multiprocess_AES_Algorithm.py
@@ -0,0 +1,272 @@
+import hashlib
+import hmac
+import scrypt
+import secrets
+import os
+import time
+import re
+import psutil
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+import base64
+import torch
+import torchcsprng as csprng
+import numpy as np
+
+def hkdf_scrypt(password, salt, length, n, r, p):
+    # HKDF extraction step
+    prk = hmac.new(salt.encode(), password.encode(), hashlib.sha256).digest()
+
+    # HKDF expansion step
+    info = b'Scrypt key derivation'
+    t = b''
+    okm = b''
+    while len(okm) < length:
+        t = hmac.new(prk, t + info + bytes([len(t) + 1]), hashlib.sha256).digest()
+        okm += scrypt.hash(t, salt.encode(), n, r, p, length)
+
+    return okm[:length]
+
+
+def encrypt_file(file_path, key, ratio, device):
+    # Read the file contents
+    with open(file_path, 'rb') as file:
+        plaintext = file.read()
+
+    # Workload division according to the ratio between cpu and gpu
+    workload_division = int(len(plaintext) * ratio)
+
+    # plaintext for cpu
+    cpu_plaintext = plaintext[:workload_division]
+
+    # Generate a random IV (Initialization Vector)
+    iv = secrets.token_bytes(AES.block_size)
+
+    # Create AES cipher object with key and mode (ECB)
+    cipher = AES.new(key, AES.MODE_ECB)
+
+    # Pad the cpu plaintext to match AES block size
+    padded_cpu_plaintext = pad(cpu_plaintext, AES.block_size)
+
+    # Encrypt the padded cpu plaintext
+    cpu_ciphertext = cipher.encrypt(padded_cpu_plaintext)
+
+    if(ratio != 1): # if GPU is now available, only use CPU
+        # Encrypt in GPU
+
+        # transfer key into tensor to use in GPU
+        key_array = np.frombuffer(key, dtype=np.uint8)
+        gpu_key = torch.from_numpy(key_array)
+        gpu_key = gpu_key.to(device)
+
+        # plaintext for gpu
+        gpu_plaintext = plaintext[workload_division:]
+
+        # transfer gpu_plaintext into tensor to use in GPU
+        gpu_plaintext_array = np.frombuffer(gpu_plaintext, dtype=np.uint8)
+        gpu_plaintext_tensor = torch.from_numpy(gpu_plaintext_array)
+        gpu_plaintext_tensor = gpu_plaintext_tensor.to(device)
+
+        # Pad the gpu_plaintext_tensor to match AES block size
+        padded_gpu_plaintext = pad(gpu_plaintext, 16)
+
+        padded_gpu_plaintext_array = np.frombuffer(padded_gpu_plaintext, dtype=np.uint8)
+        padded_gpu_plaintext_tensor = torch.from_numpy(padded_gpu_plaintext_array)
+        padded_gpu_plaintext_tensor = padded_gpu_plaintext_tensor.to(device)
+
+        # encrypt padded_gpu_plaintext_tensor in GPU
+        gpu_encrypted = torch.empty(len(padded_gpu_plaintext_tensor), dtype=torch.int8, device=device)
+        csprng.encrypt(padded_gpu_plaintext_tensor, gpu_encrypted, gpu_key, "aes128", "ecb")
+
+        # transfer gpu_encrypted results into bytes to use in cpu
+        gpu_encrypted_numpy = gpu_encrypted.cpu().numpy()
+        gpu_ciphertext = gpu_encrypted_numpy.tobytes()
+
+        # merge iv, cpu_ciphertext, gpu_ciphertext
+        results = iv + cpu_ciphertext + gpu_ciphertext
+    else:
+        results = iv + cpu_ciphertext
+
+    # Write the IV and encrypted data back to the file
+    with open(file_path, 'wb') as file:
+        file.write(results)
+
+def decrypt_file(file_path, key, ratio, device):
+    # Read the file contents
+    with open(file_path, 'rb') as file:
+        ciphertext = file.read()
+
+    # Extract the IV and ciphertext from the file
+    iv = ciphertext[:AES.block_size]
+    ciphertext = ciphertext[AES.block_size:]
+
+    # Workload division according to the ratio between cpu and gpu
+    workload_division = 16 + (int((len(ciphertext) - 16) * ratio / 16)) * 16
+
+    # ciphertext for cput
+    cpu_ciphertext = ciphertext[:workload_division]
+
+    # Create AES cipher object with key and mode (ECB)
+    cipher = AES.new(key, AES.MODE_ECB)
+
+    # Decrypt the cpu_ciphertext
+    cpu_plaintext = cipher.decrypt(cpu_ciphertext)
+
+    # Unpad the decrypted cpu_plaintext
+    cpu_plaintext = unpad(cpu_plaintext, AES.block_size)
+
+
+    if(ratio != 1): # if GPU is now available, only use CPU
+        # Decrypt in GPU
+
+        # transfer key into tensor to use in GPU 
+        key_array = np.frombuffer(key, dtype=np.uint8)
+        gpu_key = torch.from_numpy(key_array)
+        gpu_key = gpu_key.to(device)
+
+        # cipertext for gpu
+        gpu_ciphertext = ciphertext[workload_division:]
+
+        # transfer gpu_ciphertext into tensor to use in GPU
+        encrypted_input_array = np.frombuffer(gpu_ciphertext, dtype=np.uint8)
+        encrypted_input_tensor = torch.from_numpy(encrypted_input_array)
+        encrypted_input_tensor = encrypted_input_tensor.to(device)
+
+        # Decrypt encrypted_input_tensor in GPU
+        gpu_decrypted = torch.empty_like(encrypted_input_tensor)
+        csprng.decrypt(encrypted_input_tensor, gpu_decrypted, gpu_key, "aes128", "ecb")
+
+        # transfer gpu_encrypted results into bytes to use in cpu
+        gpu_decrypted_numpy = gpu_decrypted.cpu().numpy()
+        gpu_plaintext = gpu_decrypted_numpy.tobytes()
+
+        # Unpad the decrypted gpu_plaintext
+        gpu_plaintext = unpad(gpu_plaintext, 16)
+
+        # Merge cpu_plaintext and gpu_plaintext
+        results = cpu_plaintext + gpu_plaintext
+    else:
+        results = cpu_plaintext
+
+    # Write the decrypted data back to the file
+    with open(file_path, 'wb') as file:
+        file.write(results)
+
+
+def cpu_gpu_ratio():
+
+    #Define the ration between cpu and gpu
+    gpu_total = torch.cuda.device_count()
+
+    # Usage CPU
+    cpu_percent = psutil.cpu_percent(interval=0.1)
+
+    # Usage RAM
+    memory_percent = psutil.virtual_memory().percent
+
+    print(f"CPU Usage: {cpu_percent}% - Memory Usage: {memory_percent}")
+
+    # Set the Ratio between CPU and GPU according to the CPU and RAM Usage
+    if(gpu_total == 0):
+        ratio = 1
+    elif(memory_percent > 90 or cpu_percent > 90):
+        ratio = 0.1
+    elif(memory_percent > 70 or cpu_percent > 70):
+        ratio = 0.3
+    elif(memory_percent > 40 or cpu_percent > 40):
+        ratio = 0.5
+    else:
+        ratio = 0.7
+    return ratio
+
+    
+def encrypt_folder(folder_path, key, mode, ratio, device):
+    start_time = time.time()
+    for root, dirs, files in os.walk(folder_path):
+        for file in files:
+            file_path = os.path.join(root, file)
+            if mode == 'encrypt':
+                encrypt_file(file_path, key, ratio, device)
+                print(f'Encrypted: {file_path}')
+            elif mode == 'decrypt':
+                decrypt_file(file_path, key, ratio, device)
+                print(f'Decrypted: {file_path}')
+
+            # Monitor system resources while processing files
+            cpu_percent = psutil.cpu_percent(interval=0.1)  
+            memory_percent = psutil.virtual_memory().percent
+            print(f"CPU Usage: {cpu_percent}% - Memory Usage: {memory_percent}% -Ratio between CPU and GPU: {ratio}")
+
+    end_time = time.time()
+    execution_time = end_time - start_time
+    print(f"Execution Time: {execution_time:.3f} seconds ({execution_time * 1000:.3f} milliseconds)")
+
+if __name__=='__main__':
+    # Define regular expressions for input validation
+    password_regex = re.compile(r'^.{8,}$')  # Minimum 8 characters
+    salt_regex = re.compile(r'^[a-fA-F0-9]{32}$')  # 32 hexadecimal characters
+    folder_path_regex = re.compile(r'^[a-zA-Z0-9_./\\-]+$')  # Alphanumeric, underscore, dot, forward slash, backslash, and hyphen
+
+    # Prompt the user for inputs
+    while True:
+        try:
+            password = "1234567890" # password = input("Enter the password: ")
+            if not password_regex.match(password):
+                print("Invalid password. Password must be at least 8 characters long.")
+                continue
+
+            salt = "12345678901234567890123456789012" # salt = input("Enter the salt: ")
+            if not salt_regex.match(salt):
+                print("Invalid salt. Salt must be a 32-character hexadecimal string.")
+                continue
+
+            length = 16 # length = int(input("Enter the desired key length in bytes: "))
+            n = 2^10 # n = int(input("Enter the value for 'n': "))
+            r = 8 # r = int(input("Enter the value for 'r': "))
+            p = 5 # p = int(input("Enter the value for 'p': "))
+            
+            
+            
+            break
+        except ValueError:
+            print("Invalid input. Please try again.")
+
+    # Prompt the user to input the folder path
+    # while True:
+        
+    #     if not folder_path_regex.match(folder_path):
+    #         print("Invalid folder path. Folder path must be alphanumeric and can contain underscore, dot, forward slash, backslash, and hyphen.")
+    #         continue
+    #     if os.path.isdir(folder_path):
+    #         break
+    #     else:
+    #         print("Invalid folder path. Please try again.")
+    #folder_path = input("Enter the folder path: ")
+    mode = 'encrypt'
+    folder_path = "E:\enc"
+    
+
+    # Prompt the user to select the mode: 'encrypt' or 'decrypt'
+    # while True:
+    #     mode = 'encrypt' #mode = input("Enter the mode ('encrypt' or 'decrypt'): ")
+    #     if mode in ['encrypt', 'decrypt']:
+    #         breakon
+    #     else:
+    #         print("Invalid mode. Please try again.")
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    ratio = cpu_gpu_ratio()    
+
+    # Call the hkdf_scrypt function with user inputs
+    derived_key = hkdf_scrypt(password, salt, length, n, r, p)
+
+    # Encode the derived key using base64
+    encoded_key = base64.b64encode(derived_key).decode()
+
+    if mode == 'encrypt':
+        # Store the encoded key in a secure location (e.g., a file or database)
+        with open("backup_key.txt", "w") as file:
+            file.write(encoded_key)
+
+    # Encrypt or decrypt the files within the folder using AES-256
+    encrypt_folder(folder_path, derived_key, mode, ratio, device)
