import os import json import sys import subprocess import requests from concurrent.futures import ThreadPoolExecutor, as_completed import time from collections import defaultdict # Configuration MAX_RETRIES = 3 # Maximum retry attempts for downloading files NUM_WORKERS = 5 # Default number of workers TIMEOUT = 30 # Timeout in seconds for each file download PROCESSING_DIR = "/processing" # Directory to store zip, .txt, and .json files # Function to download the JSON file from the specified URL def download_json(bucket_name, download_directory): try: url = f"https://livetimelapse.com.au/ai/chatgpt/internal/b2/servers/www.php?bucket_name={bucket_name}" response = requests.get(url, stream=True) response.raise_for_status() # Raise an exception for HTTP errors # Save JSON to local file with new naming convention json_filename = f"{bucket_name}-files.json" json_path = os.path.join(PROCESSING_DIR, json_filename) with open(json_path, "w") as json_file: json.dump(response.json(), json_file, indent=4) print(f"JSON file downloaded and saved to: {json_path}") return json_path except Exception as e: print(f"Error downloading JSON file: {e}") sys.exit(1) # Function to load the JSON file into memory def load_json_file(json_path): try: with open(json_path, "r") as json_file: return json.load(json_file) except Exception as e: print(f"Error loading JSON file: {e}") sys.exit(1) # Function to compress a folder into a ZIP using 7z def compress_with_7zip(directory, zip_name): """ directory: The local path of the folder you want to zip zip_name: The filename of the resulting .zip (no path) """ try: zip_path = os.path.join(PROCESSING_DIR, zip_name) command = f'7z a "{zip_path}" "{directory}/*"' print(f"[ZIPPING] Compressing folder: {directory} -> {zip_path}") result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) # Save process output to a log file log_file = zip_path.replace(".zip", "-log.txt") with open(log_file, "w") as log: log.write(result.stdout) log.write("\n\nErrors (if any):\n") log.write(result.stderr) print(f"[ZIPPING] Done compressing: {zip_path}") return zip_path, log_file except subprocess.CalledProcessError as e: print(f"Error compressing folder '{directory}': {e}") return None, None # Function to move files to Google Drive in a screen session def move_to_gdrive_with_screen(file_path, rclone_target): """ Moves the given file to Google Drive using rclone move in a detached screen session. """ try: # A small delay to help ensure the file is closed time.sleep(5) log_file = f"/tmp/{os.path.basename(file_path)}.log" session_name = f"rclone_{os.path.basename(file_path)}" command = f'screen -dmS {session_name} sh -c "rclone move {file_path} {rclone_target} -P > {log_file} 2>&1"' print(f"[RCLONE] Moving {file_path} to {rclone_target} in screen session: {session_name}") os.system(command) except Exception as e: print(f"Error starting rclone move command for {file_path}: {e}") # Function to copy files to Google Drive (blocking) def copy_to_gdrive(file_path, rclone_target): """ Copies the given file to Google Drive using rclone copy (blocking). """ try: command = f'rclone copy {file_path} {rclone_target} -P' print(f"[RCLONE] Copying {file_path} to {rclone_target}...") os.system(command) except Exception as e: print(f"Error copying {file_path} to {rclone_target}: {e}") def zip_and_upload_folder(folder_path, folder_name, rclone_target): """ Compresses the given folder (folder_path) into -www.zip, moves the zip to Google Drive in the background, and then removes the folder. """ # Create zip file name zip_filename = f"{folder_name}-www.zip" # Compress zip_path, log_file = compress_with_7zip(folder_path, zip_filename) if zip_path and os.path.exists(zip_path): # Move ZIP to Google Drive in background move_to_gdrive_with_screen(zip_path, rclone_target) # Also upload the compression log, if it exists if log_file and os.path.exists(log_file): copy_to_gdrive(log_file, rclone_target) # Remove folder after successful compression print(f"[CLEANUP] Removing folder: {folder_path}") subprocess.run(f"rm -Rf '{folder_path}'", shell=True, check=False) # Function to download a single file def download_file(file_info, base_dir, stats, folder_counts, folder_completed, zipping_executor, rclone_target): file_path = file_info["path"] cdn_url = file_info["cdn"] # The top-level folder name: everything before the first slash (or the entire path if no slash) top_level_folder = file_path.split("/")[0] # For printing print(f"\nDownloading: {cdn_url}") local_path = os.path.join(base_dir, file_path) os.makedirs(os.path.dirname(local_path), exist_ok=True) # If file already exists, skip if os.path.exists(local_path): print(f"File already exists, skipping: {local_path}") stats["skipped"] += 1 return # Attempt up to MAX_RETRIES for attempt in range(1, MAX_RETRIES + 1): try: response = requests.get(cdn_url, stream=True, timeout=TIMEOUT) response.raise_for_status() with open(local_path, "wb") as file: for chunk in response.iter_content(chunk_size=8192): file.write(chunk) print(f"Saved to: {local_path}") stats["downloaded"] += 1 break except requests.exceptions.RequestException as e: print(f"Attempt {attempt} failed for {cdn_url}: {e}") if attempt < MAX_RETRIES: time.sleep(2 ** attempt) # Exponential backoff else: print(f"Failed to download {cdn_url} after {MAX_RETRIES} attempts.") stats["failed"] += 1 return # Do not proceed if it failed # Once downloaded, increment the folder_completed counter folder_completed[top_level_folder] += 1 # If we've finished all files for this folder, zip & upload it in the background if folder_completed[top_level_folder] == folder_counts[top_level_folder]: folder_path = os.path.join(base_dir, top_level_folder) zipping_executor.submit(zip_and_upload_folder, folder_path, top_level_folder, rclone_target) # Main function def main(num_workers, bucket_name, download_directory, rclone_target): os.makedirs(download_directory, exist_ok=True) os.makedirs(PROCESSING_DIR, exist_ok=True) summary = {"downloaded": 0, "skipped": 0, "failed": 0, "download_time": 0} # Step 1: Download JSON json_path = download_json(bucket_name, PROCESSING_DIR) # Step 2: Load JSON and prepare bucket directory file_list = load_json_file(json_path) bucket_dir = os.path.join(download_directory, bucket_name) os.makedirs(bucket_dir, exist_ok=True) # Build counts of how many files per top-level folder folder_counts = defaultdict(int) for info in file_list: top_folder = info["path"].split("/")[0] folder_counts[top_folder] += 1 # Track completion counts per folder folder_completed = defaultdict(int) print(f"Downloading {len(file_list)} images into: {bucket_dir}") # We'll use one executor for downloads, another for zipping with ThreadPoolExecutor(max_workers=num_workers) as download_executor, \ ThreadPoolExecutor(max_workers=2) as zipping_executor: # limit zipping concurrency if desired # Submit all download tasks futures = [] for file_info in file_list: futures.append( download_executor.submit( download_file, file_info, bucket_dir, summary, folder_counts, folder_completed, zipping_executor, rclone_target ) ) # Wait for all download tasks to complete for _ in as_completed(futures): pass print(f"Finished downloading files for bucket: {bucket_name}") # Wait for all zipping tasks to complete zipping_executor.shutdown(wait=True) # Copy the JSON file to Google Drive copy_to_gdrive(json_path, rclone_target) # Remove the entire bucket directory (whatever might be left) print(f"[CLEANUP] Removing bucket directory: {bucket_dir}") subprocess.run(f"rm -Rf '{bucket_dir}'", shell=True, check=False) print("All tasks completed!") print(f"Summary: {summary}") if __name__ == "__main__": if len(sys.argv) != 5: print("Usage: python3 script.py ") sys.exit(1) num_workers = int(sys.argv[1]) bucket_name = sys.argv[2] download_directory = sys.argv[3] rclone_target = sys.argv[4] main(num_workers, bucket_name, download_directory, rclone_target)