sftp download via aria2c + max concurrency

sftp-aria2c-crawl.py

#!/usr/bin/env python3

import paramiko
import os
import stat
import subprocess
from concurrent.futures import ThreadPoolExecutor

# Define SFTP credentials and server information
SFTP_HOST = "sftp.bitport.io"
SFTP_PORT = 22
SFTP_USER = "<username>"
SFTP_PASS = "<password>"
REMOTE_BASE_DIR = "/"           # Change this to the base directory on the server
LOCAL_BASE_DIR = "./downloads"  # Local directory to save downloaded files

# Function to download a single file using aria2c
def download_file(remote_path, local_dir, item_filename):
    aria2c_command = [
        "aria2c",
        "--file-allocation=none",
        "-x8",  # Number of connections
        f"--ftp-user={SFTP_USER}",
        f"--ftp-passwd={SFTP_PASS}",
        f"sftp://{SFTP_HOST}:{SFTP_PORT}{remote_path}",
        "-d", local_dir,  # Set the download directory
        "-o", item_filename  # Set the output filename
    ]

    print(f"Starting download: {remote_path} -> {os.path.join(local_dir, item_filename)}")
    process = subprocess.Popen(aria2c_command)
    process.wait()  # Wait for the process to complete
    return process.returncode  # Return the exit code

# Function to recursively crawl directories and prepare downloads
def sftp_recursive_download(sftp, remote_dir, local_dir, executor):
    # Ensure the local directory exists
    if not os.path.exists(local_dir):
        os.makedirs(local_dir)

    # List the contents of the remote directory
    for item in sftp.listdir_attr(remote_dir):
        remote_path = os.path.join(remote_dir, item.filename)
        local_path = os.path.join(local_dir, item.filename)

        # Check if it's a directory
        if stat.S_ISDIR(item.st_mode):
            # Recursively crawl into the directory
            print(f"Entering directory: {remote_path}")
            sftp_recursive_download(sftp, remote_path, local_path, executor)
        else:
            # Submit the download task to the executor
            executor.submit(download_file, remote_path, local_dir, item.filename)

def main():
    # Ensure the local base directory exists
    if not os.path.exists(LOCAL_BASE_DIR):
        os.makedirs(LOCAL_BASE_DIR)

    # Create an SSH client
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        # Connect to the SFTP server
        print("Connecting to SFTP server...")
        ssh.connect(SFTP_HOST, port=SFTP_PORT, username=SFTP_USER, password=SFTP_PASS)

        # Open an SFTP session
        sftp = ssh.open_sftp()

        # Start recursive download from the base directory
        print(f"Starting download from {REMOTE_BASE_DIR}...")

        # Use ThreadPoolExecutor to limit concurrent downloads
        with ThreadPoolExecutor(max_workers=4) as executor:
            sftp_recursive_download(sftp, REMOTE_BASE_DIR, LOCAL_BASE_DIR, executor)

        # Close the SFTP session
        sftp.close()
        print("All downloads started. Waiting for them to finish...")

    finally:
        # Close the SSH connection
        ssh.close()

    print("All downloads complete.")

if __name__ == "__main__":
    main()