|
| 1 | +# cipher_tool.py |
| 2 | +""" |
| 3 | +A simple AES-based file encryption and decryption tool for automating code protection. |
| 4 | +""" |
| 5 | +import os |
| 6 | +from Crypto.Cipher import AES |
| 7 | +from Crypto.Random import get_random_bytes |
| 8 | +from Crypto.Protocol.KDF import PBKDF2 |
| 9 | + |
| 10 | +import hashlib |
| 11 | + |
| 12 | +BLOCK_SIZE = 16 # AES block size in bytes |
| 13 | +SALT_SIZE = 16 # Salt size in bytes |
| 14 | +KEY_SIZE = 32 # AES-256 |
| 15 | +ITERATIONS = 100_000 |
| 16 | + |
| 17 | +def pad(data): |
| 18 | + padding_len = BLOCK_SIZE - len(data) % BLOCK_SIZE |
| 19 | + return data + bytes([padding_len]) * padding_len |
| 20 | + |
| 21 | +def unpad(data): |
| 22 | + padding_len = data[-1] |
| 23 | + return data[:-padding_len] |
| 24 | + |
| 25 | +def derive_key(password, salt): |
| 26 | + return PBKDF2(password, salt, dkLen=KEY_SIZE, count=ITERATIONS) |
| 27 | + |
| 28 | +def encrypt_file(input_path, output_path, password): |
| 29 | + salt = get_random_bytes(SALT_SIZE) |
| 30 | + key = derive_key(password.encode(), salt) |
| 31 | + cipher = AES.new(key, AES.MODE_CBC) |
| 32 | + with open(input_path, 'rb') as f: |
| 33 | + plaintext = f.read() |
| 34 | + padded = pad(plaintext) |
| 35 | + ciphertext = cipher.encrypt(padded) |
| 36 | + with open(output_path, 'wb') as f: |
| 37 | + f.write(salt + cipher.iv + ciphertext) |
| 38 | + print(f"Encrypted {input_path} -> {output_path}") |
| 39 | + |
| 40 | +def decrypt_file(input_path, output_path, password): |
| 41 | + with open(input_path, 'rb') as f: |
| 42 | + salt = f.read(SALT_SIZE) |
| 43 | + iv = f.read(BLOCK_SIZE) |
| 44 | + ciphertext = f.read() |
| 45 | + key = derive_key(password.encode(), salt) |
| 46 | + cipher = AES.new(key, AES.MODE_CBC, iv=iv) |
| 47 | + padded = cipher.decrypt(ciphertext) |
| 48 | + plaintext = unpad(padded) |
| 49 | + with open(output_path, 'wb') as f: |
| 50 | + f.write(plaintext) |
| 51 | + print(f"Decrypted {input_path} -> {output_path}") |
| 52 | + |
| 53 | +def generate_checksum(file_path, algo='sha256'): |
| 54 | + """ |
| 55 | + Generate a checksum (SHA-256 by default) for a file. |
| 56 | + """ |
| 57 | + h = hashlib.new(algo) |
| 58 | + with open(file_path, 'rb') as f: |
| 59 | + for chunk in iter(lambda: f.read(4096), b''): |
| 60 | + h.update(chunk) |
| 61 | + return h.hexdigest() |
| 62 | + |
| 63 | +def verify_checksum(file_path, expected_checksum, algo='sha256'): |
| 64 | + """ |
| 65 | + Verify a file's checksum matches the expected value. |
| 66 | + """ |
| 67 | + actual = generate_checksum(file_path, algo) |
| 68 | + return actual == expected_checksum, actual |
| 69 | + |
| 70 | +if __name__ == "__main__": |
| 71 | + import argparse |
| 72 | + parser = argparse.ArgumentParser(description="Encrypt or decrypt files using AES.") |
| 73 | + parser.add_argument('mode', choices=['encrypt', 'decrypt'], help='Mode: encrypt or decrypt') |
| 74 | + parser.add_argument('input', help='Input file path') |
| 75 | + parser.add_argument('output', help='Output file path') |
| 76 | + parser.add_argument('password', help='Password for encryption/decryption') |
| 77 | + args = parser.parse_args() |
| 78 | + if args.mode == 'encrypt': |
| 79 | + encrypt_file(args.input, args.output, args.password) |
| 80 | + else: |
| 81 | + decrypt_file(args.input, args.output, args.password) |
| 82 | + |
| 83 | +# Example usage: |
| 84 | +# checksum = generate_checksum('somefile.py') |
| 85 | +# ok, actual = verify_checksum('somefile.py', checksum) |
0 commit comments