upload_to_s3.py

#!/usr/bin/python3

# this file is used to upload files onto s3 and notify admins of successful
# uploads. the uploading includes grabbing files from the upload list,
# compressing them, and encrypting them.
# USAGE: python3 upload_to_s3.py BACKUP_DIR

import os
import sys
import boto3
import gzip
import shutil
from discord_webhook import DiscordWebhook
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad

BUCKET_NAME = "duke2025-esa03-backup-storage"
WEBHOOK_URL = 'https://discord.com/api/webhooks/1313315767424520262/X6OqgMGop-'
WEBHOOK_URL += 'lKTDl4Q20sTxT1Hkg7nvneYJwuzozAHjCxlvgkoW5rZp-LfHtG2HSTfvxd'
UPLOAD_FILE_LIST = '/mnt/btrfs_project_drive/upload_files.txt'
ROOT_DIR = '/mnt/btrfs_project_drive'
successful_upload = True

# notify_discord(content)
#  parameters: content      the notification to send
# description: this function pings a discord server with the given message
def notify_discord(content):
    webhook = DiscordWebhook(url=WEBHOOK_URL, content=content)
    response = webhook.execute()

# encrypt_file(in_filename, out_filename, key)
#  parameters: in_filename  the name of the file to encrypt
#              out_filename the location to write the encrypted file to
#              key          the encryption key
# description: this funciton encrypts the given file and writes the output to
#              a new location
def encrypt_file(in_filename, out_filename, key):
    iv = get_random_bytes(16)  # 128-bit IV
    cipher = AES.new(key, AES.MODE_CBC, iv)
    
    with open(in_filename, 'rb') as f_in:
        plaintext = f_in.read()
        ciphertext = cipher.encrypt(pad(plaintext, AES.block_size))
    
    with open(out_filename, 'wb') as f_out:
        f_out.write(iv + ciphertext)

# upload_file_to_s3(local_path, s3_key)
#  parameters: local_path   the file to be uploaded
#              s3_key       the name for the file on s3
# description: this function compresses, encrypts, and uploads a single file
def upload_file_to_s3(local_path, s3_key):
    s3 = boto3.client('s3')
    try:
        compressed_path = local_path + '.gz'
        encrypted_path = compressed_path + '.enc'
        with open('aes_key.bin', 'rb') as key_file:
            key = key_file.read()

        # compress
        with open(local_path, 'rb') as f_in, gzip.open(compressed_path, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)

        # encrypt
        encrypt_file(compressed_path, encrypted_path, key)

        # upload
        s3.upload_file(encrypted_path, BUCKET_NAME, s3_key)

        # cleanup by deleting local compressed and encrypted files
        os.remove(compressed_path)
        os.remove(encrypted_path)
    
    # on any failures, notify admins
    except Exception as e:
        notify_discord(f'Backup failed while uploading metadata: {e}')

# upload_new_files()
# description: this function traverses files found in the upload list and calls
#              to upload each one
def upload_new_files():
    with open(UPLOAD_FILE_LIST, 'r') as f:
        for line in f:
            line = line.strip()
            file_hash, file_path = line.split('$')
            file_path = os.path.join(ROOT_DIR, "snapshots", file_path)
            if os.path.isfile(file_path):
                s3_key = f"files/{file_hash}"
                upload_file_to_s3(file_path, s3_key)

# main()
# description: this is the driver function of the script, calling functions to
#              begin parsing and uploading
def main():
    if len(sys.argv) != 2:
        sys.exit(1)
    
    snapshot_date = sys.argv[1]
    snapshot_dir = os.path.join(ROOT_DIR, "snapshots", snapshot_date)
    
    snapshot_metadata_path = os.path.join(snapshot_dir, "metadata.txt")
    snapshot_metadata_s3_key = f"snapshots/{snapshot_date}/metadata.txt"
    
    successful_upload = True
    upload_file_to_s3(snapshot_metadata_path, snapshot_metadata_s3_key)

    upload_new_files()
    if successful_upload:
        notify_discord(f'Backup completed uploading successfully')

if __name__ == "__main__":
    main()