A simple file encryption Python script

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>