Upload files to "/"

PCS-A-1-i-Proto.py

@@ -0,0 +1,130 @@
+import os
+import getpass
+import sys
+from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+from pathlib import Path
+import base64
+
+# TODO finish functions
+# TODO test this on some unimportant data
+# TODO Integrate file choosing in both encrypt and decrypt!
+
+# Encrypting data
+def encrypt_data():
+
+    # Pesterish file asker
+    while True:
+        file_name = input("\nEnter the name of a zip file, without the extension that is in the same directory as this script. E.g. for encrypt.zip you would type\"encrypt\".\nName > ") + ".zip"
+        if Path(file_name).is_file():
+            break
+        else:
+            print("This zip file does not exist!")
+
+    file_chosen_name = input("\nMake a name for your encrypted file, e.g. encrypted (press enter for default)\nName > ")
+    if file_chosen_name == "":
+        file_chosen_name = "encrypted"
+        
+    # Pestering password asker until you confirm it
+    while True:
+        password = getpass.getpass("Creating new password > ")
+        confirm = input(f"Did you type the correct password, and remember it? [Y]/[n]")
+        if confirm == "Y" or confirm == "y":
+            # Confirms with user one last time and reminders
+            password = password.encode()
+            break
+                
+    # 1. Setup Salt
+    salt = os.urandom(16)
+
+    # 2. Derive Key
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        iterations=480000,
+    )
+    key = base64.urlsafe_b64encode(kdf.derive(password))
+    f = Fernet(key)
+
+    print("Encrypting...", end = " ")
+    try:
+        # 3. Encrypt and save SALT + DATA together
+        with open(file_name, 'rb') as file:
+            original_data = file.read()
+
+        encrypted_data = f.encrypt(original_data)
+
+        with open(f"{file_chosen_name}.p7c_enc", 'wb') as file:
+            # Write the 16-byte salt first, then the encrypted data
+            file.write(salt)
+            file.write(encrypted_data)
+
+        print("Encrypted! ")
+        delete_original = input("Delete the original zip file for security? [Y]/[n]")
+        if delete_original == "Y" or delete_original == "y":
+            os.remove(f"{file_name}")
+            print("Original file removed.")
+    except Exception as e:
+        print(f"Error! {e}")
+
+# Decrypting data
+def decrypt_data():
+    # 1. Get the password from the user
+    while True:
+        choose_file = input("\nName of the p7c_enc file without extension, e.g. \"encrypted\"\nName > ") + ".p7c_enc"
+        if Path(choose_file).is_file():
+            break
+        else:
+            print("This p7c_enc file does not exist!")
+    
+    password = getpass.getpass("Password > ").encode()
+
+    # 2. Open the encrypted file and extract the salt + data
+    with open(choose_file, 'rb') as file:
+        # Read exactly 16 bytes for the salt
+        file_salt = file.read(16)
+        # Read everything else as the encrypted data
+        encrypted_data = file.read()
+
+    # 3. Re-derive the EXACT same key using that salt
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=file_salt,
+        iterations=480000,
+    )
+    key = base64.urlsafe_b64encode(kdf.derive(password))
+    f = Fernet(key)
+
+    # 4. Decrypt and save the original file
+    print("Decrypting...", end = " ")
+    try:
+        decrypted_data = f.decrypt(encrypted_data)
+        
+        with open('unencrypted.zip', 'wb') as file:
+            file.write(decrypted_data)
+        print("Success.")
+        delete_encrypted = input("Delete the encrypted file? [Y]/[n]")
+        if delete_encrypted == "Y" or delete_encrypted == "y":
+            os.remove(choose_file)
+            print("Encrypted file removed.")
+        
+        
+    except Exception as e:
+        print(f"Could not decrypt!.\nThis might be due to a wrong file or corrupted data?\nDetailed info:\n{e}.")
+
+# Main function
+
+while True:
+    print("\nPCS - P7MJ's enCryption System | V A-1-i Proto | P7MJ")
+    choice = input("[1] Encrypt [2] Decrypt [x] Exit > ")
+    if choice == "1":
+        encrypt_data()
+    elif choice == "2":
+        decrypt_data()
+    elif choice == "x":
+        sys.exit(0)
+    else:
+        print("Not an option!!!")