A simple file encryption Python script

Sep 16, 2024#安全140

AI Translation

This post is translated from Chinese into English through AI.View Original

AI-generated summary

The text discusses a simple Python script created to encrypt sensitive files, such as cryptocurrency wallet private keys. The author previously stored wallet recovery phrases in plain text files but became concerned about potential data loss or security risks, especially after experiencing hardware failures. The script uses the PBKDF2 algorithm for encryption, requiring a password for decryption, which enhances security against data leaks when stored in the cloud. The script includes functions for encrypting and decrypting files, and the author provides usage instructions.

To save the private key of the encrypted wallet, I wrote a simple file encryption script, which can also be used to encrypt other things.

Previously, I kept the wallet mnemonic pasted in a txt file on my computer, but I thought about what would happen if the hard drive suddenly failed. Last month, the motherboard of my computer broke down and it couldn't be repaired, but the hard drive was fine and all the data was intact. However, some wallets in browser extensions that didn't back up the mnemonic were directly lost, but they were all wallets with little value.

Then I uploaded the file to the cloud (that's how I used to back it up), but recently I saw a bug reported online about a certain cloud drive, where users could see files stored by others. Although I didn't investigate whether it was a rumor, it made me hesitant to store plaintext directly in the cloud. Below is the code for this Python script:

locker.py

from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import os
import sys
import getpass

def generate_key(password: str, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=default_backend()
    )
    return kdf.derive(password.encode())

def encrypt_file(password: str, input_file: str, output_file: str):
    salt = os.urandom(16)
    key = generate_key(password, salt)
    aesgcm = AESGCM(key)
    nonce = os.urandom(12)
  
    with open(input_file, 'rb') as f:
        data = f.read()
  
    encrypted_data = aesgcm.encrypt(nonce, data, None)
  
    with open(output_file, 'wb') as f:
        f.write(salt + nonce + encrypted_data)

def decrypt_file(password: str, input_file: str, output_file: str):
    with open(input_file, 'rb') as f:
        salt = f.read(16)
        nonce = f.read(12)
        encrypted_data = f.read()
  
    key = generate_key(password, salt)
    aesgcm = AESGCM(key)
  
    decrypted_data = aesgcm.decrypt(nonce, encrypted_data, None)
  
    with open(output_file, 'wb') as f:
        f.write(decrypted_data)

args = sys.argv

if len(args) != 4:
    print("Usage: python locker.py <op_type(encrypt/decrypt)> <input_file> <output_file>")
    sys.exit(1)

op_type = args[1]
input_file = args[2]
output_file = args[3]

password = getpass.getpass("Enter password: ")
if op_type == "encrypt":
    re_password = getpass.getpass("Re-enter password: ")
    if password != re_password:
        print("Passwords do not match")
        sys.exit(1)

if op_type == "encrypt":
    encrypt_file(password, input_file, output_file)
elif op_type == "decrypt":
    try:
        decrypt_file(password, input_file, output_file)
    except:
        print("Invalid password")
        sys.exit(1)
else:
    print("Invalid operation type")
    sys.exit(1)

Using the PBKDF2 algorithm for encryption requires entering your own password to decrypt, which at least improves security. The probability of loss due to cloud leakage is also relatively small.

Usage:

python locker.py <op_type(encrypt/decrypt)> <input_file> <output_file>

from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from cryptography.hazmat.primitives import hashes from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives.ciphers.aead import AESGCM import os import sys import getpass def generate_key(password: str, salt: bytes) -> bytes: kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=default_backend() ) return kdf.derive(password.encode()) def encrypt_file(password: str, input_file: str, output_file: str): salt = os.urandom(16) key = generate_key(password, salt) aesgcm = AESGCM(key) nonce = os.urandom(12) with open(input_file, 'rb') as f: data = f.read() encrypted_data = aesgcm.encrypt(nonce, data, None) with open(output_file, 'wb') as f: f.write(salt + nonce + encrypted_data) def decrypt_file(password: str, input_file: str, output_file: str): with open(input_file, 'rb') as f: salt = f.read(16) nonce = f.read(12) encrypted_data = f.read() key = generate_key(password, salt) aesgcm = AESGCM(key) decrypted_data = aesgcm.decrypt(nonce, encrypted_data, None) with open(output_file, 'wb') as f: f.write(decrypted_data) args = sys.argv if len(args) != 4: print("Usage: python locker.py <op_type(encrypt/decrypt)> <input_file> <output_file>") sys.exit(1) op_type = args[1] input_file = args[2] output_file = args[3] password = getpass.getpass("Enter password: ") if op_type == "encrypt": re_password = getpass.getpass("Re-enter password: ") if password != re_password: print("Passwords do not match") sys.exit(1) if op_type == "encrypt": encrypt_file(password, input_file, output_file) elif op_type == "decrypt": try: decrypt_file(password, input_file, output_file) except: print("Invalid password") sys.exit(1) else: print("Invalid operation type") sys.exit(1)