backup.py

import os
import sys
import tarfile
import shutil
import boto3
from botocore.client import Config
from flask import Flask, jsonify
from datetime import datetime

# Cloudflare R2 Config
R2_ACCOUNT_ID = os.getenv('R2_ACCOUNT_ID')
R2_ACCESS_KEY = os.getenv('R2_ACCESS_KEY')
R2_SECRET_KEY = os.getenv('R2_SECRET_KEY')
R2_BUCKET_NAME = os.getenv('R2_BUCKET_NAME')

STORAGE_DIR = 'storage'
PREFIX = 'apexkit_backup_'
MAX_BACKUPS_TO_KEEP = 5

def get_s3_client():
    if not all([R2_ACCOUNT_ID, R2_ACCESS_KEY, R2_SECRET_KEY, R2_BUCKET_NAME]):
        print("⚠️  R2 credentials missing in Hugging Face secrets. Skipping cloud storage operations.")
        return None
    
    endpoint_url = f"https://{R2_ACCOUNT_ID}.r2.cloudflarestorage.com"
    return boto3.client('s3',
        endpoint_url=endpoint_url,
        aws_access_key_id=R2_ACCESS_KEY,
        aws_secret_access_key=R2_SECRET_KEY,
        config=Config(signature_version='s3v4')
    )

def make_tarfile(output_filename, source_dir):
    with tarfile.open(output_filename, "w:gz") as tar:
        tar.add(source_dir, arcname=os.path.basename(source_dir))

def extract_tarfile(input_filename):
    with tarfile.open(input_filename, "r:gz") as tar:
        tar.extractall(path=".")

def clean_junk_folders():
    """Removes tmp and internal backup files before zipping to save massive amounts of space"""
    for folder in ['storage/tmp', 'storage/backups']:
        if os.path.exists(folder):
            shutil.rmtree(folder, ignore_errors=True)
            os.makedirs(folder, exist_ok=True)

def perform_backup():
    s3 = get_s3_client()
    if not s3:
        return False, "S3 client not configured."
    
    clean_junk_folders()

    # Create a unique timestamped filename to bust Cloudflare's cache
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    archive_name = f"{PREFIX}{timestamp}.tar.gz"
    
    print(f"📦 Compressing {STORAGE_DIR} to {archive_name}...")
    make_tarfile(archive_name, STORAGE_DIR)
    
    print(f"☁️  Uploading {archive_name} to R2 bucket '{R2_BUCKET_NAME}'...")
    try:
        s3.upload_file(archive_name, R2_BUCKET_NAME, archive_name)
        print(f"✅ Upload successful! Size: {os.path.getsize(archive_name) / (1024*1024):.2f} MB")
        
        # Clean up the local tarball so it doesn't waste disk space
        os.remove(archive_name)

        # GET RID OF THE OLD ONES
        response = s3.list_objects_v2(Bucket=R2_BUCKET_NAME, Prefix=PREFIX)
        if 'Contents' in response:
            # Sort backups by date (newest first)
            backups = sorted(response['Contents'], key=lambda x: x['LastModified'], reverse=True)
            
            # If we have more than the max allowed, delete the oldest ones
            if len(backups) > MAX_BACKUPS_TO_KEEP:
                for old_backup in backups[MAX_BACKUPS_TO_KEEP:]:
                    old_key = old_backup['Key']
                    print(f"🗑️  Deleting old backup to save space: {old_key}")
                    s3.delete_object(Bucket=R2_BUCKET_NAME, Key=old_key)

        return True, f"Backup successful: {archive_name}"
    except Exception as e:
        print(f"❌ Backup failed: {e}")
        return False, str(e)

def perform_restore():
    s3 = get_s3_client()
    if not s3:
        return False
    
    print(f"☁️  Scanning for latest backup in R2 bucket '{R2_BUCKET_NAME}'...")
    try:
        # 1. Search for NEW timestamped backups
        response = s3.list_objects_v2(Bucket=R2_BUCKET_NAME, Prefix=PREFIX)
        
        latest_backup_key = None
        
        if 'Contents' in response and len(response['Contents']) > 0:
            # Sort by date and grab the absolute newest one
            backups = sorted(response['Contents'], key=lambda x: x['LastModified'], reverse=True)
            latest_backup_key = backups[0]['Key']
        else:
            # 2. LEGACY FALLBACK: If no new backups exist, look for the old one!
            print("⚠️  No timestamped backups found. Attempting to restore legacy 'apexkit_storage_backup.tar.gz'...")
            latest_backup_key = 'apexkit_storage_backup.tar.gz'

        print(f"📥 Downloading backup: {latest_backup_key}...")
        temp_download_name = "restore.tar.gz"
        
        # Download the file
        s3.download_file(R2_BUCKET_NAME, latest_backup_key, temp_download_name)
        
        print(f"📦 Extracting {latest_backup_key}...")
        extract_tarfile(temp_download_name)
        
        # Clean up downloaded tarball
        os.remove(temp_download_name)
        
        print("✅ Restore successful! Files are fully up to date.")
        return True

    except Exception as e:
        print(f"⚠️  Restore failed (Is the bucket empty?). Error: {e}")
        return False

# --- LOCAL API SERVER ---
app = Flask(__name__)

@app.route('/backup', methods=['POST', 'GET'])
def trigger_backup():
    success, msg = perform_backup()
    if success:
        return jsonify({"status": "success", "message": msg}), 200
    else:
        return jsonify({"status": "error", "message": msg}), 500

if __name__ == '__main__':
    if len(sys.argv) > 1:
        if sys.argv[1] == '--restore':
            perform_restore()
        elif sys.argv[1] == '--serve':
            app.run(host='127.0.0.1', port=5000)
    else:
        print("Use --restore or --serve")