688 lines
33 KiB
Python
688 lines
33 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Tool to send a download to an aria2c daemon via RPC.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import sys
|
|
import os
|
|
import glob
|
|
import shutil
|
|
import re
|
|
import shlex
|
|
import time
|
|
from urllib.parse import urljoin
|
|
|
|
try:
|
|
import aria2p
|
|
from aria2p.utils import human_readable_bytes
|
|
except ImportError:
|
|
print("aria2p is not installed. Please install it with: pip install aria2p", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
logger = logging.getLogger('download_aria_tool')
|
|
|
|
class TimeoutError(Exception):
|
|
pass
|
|
|
|
|
|
def add_download_aria_parser(subparsers):
|
|
"""Add the parser for the 'download aria-rpc' command."""
|
|
parser = subparsers.add_parser(
|
|
'aria-rpc',
|
|
description='Send a download to an aria2c daemon via RPC, using an info.json from stdin or a file.',
|
|
formatter_class=argparse.RawTextHelpFormatter,
|
|
help='Download a specific format using aria2c RPC.',
|
|
epilog="""
|
|
Usage Notes for Fragmented Downloads (e.g., DASH):
|
|
|
|
To download and automatically merge fragmented formats, you must:
|
|
1. Use '--wait' to make the operation synchronous.
|
|
2. Use '--auto-merge-fragments' to enable the merge logic.
|
|
3. Ensure this script has access to the directory where aria2c saves files.
|
|
|
|
Example for a remote aria2c daemon:
|
|
- The remote daemon saves files to '/srv/downloads' on its machine.
|
|
- This directory is mounted locally at '/mnt/remote_aria2_downloads'.
|
|
|
|
cat latest-info.json | yt-ops-client download aria-rpc -f "299/137" \\
|
|
--wait --auto-merge-fragments \\
|
|
--remote-dir /srv/downloads \\
|
|
--fragments-dir /mnt/remote_aria2_downloads
|
|
"""
|
|
)
|
|
parser.add_argument('--load-info-json', type=argparse.FileType('r', encoding='utf-8'), help="Path to the info.json file. If not provided, reads from stdin.")
|
|
parser.add_argument('-f', '--format', required=True, help='The format ID to download. Supports yt-dlp style format selectors (e.g., "137/136,140").')
|
|
parser.add_argument('--output-dir', help='Local directory to save the final merged file. Defaults to the current directory.')
|
|
parser.add_argument('--fragments-dir', help='The local path where this script should look for downloaded fragments. If the aria2c daemon is remote, this should be a local mount point corresponding to --remote-dir. Defaults to --output-dir.')
|
|
parser.add_argument('--remote-dir', help='The absolute path to the download directory on the remote aria2c host. This is passed via RPC.')
|
|
parser.add_argument('--aria-host', default='localhost', help='The host of the aria2c RPC server. Default: localhost.')
|
|
parser.add_argument('--aria-port', type=int, default=6800, help='The port of the aria2c RPC server. Default: 6800.')
|
|
parser.add_argument('--aria-secret', help='The secret token for the aria2c RPC server (often required, e.g., "SQGCQPLVFQIASMPNPOJYLVGJYLMIDIXDXAIXOTX").')
|
|
parser.add_argument('--proxy', help='Proxy to use for the download, e.g., "socks5://127.0.0.1:1080".')
|
|
parser.add_argument('--downloader-args', help='Arguments for aria2c, in yt-dlp format (e.g., "aria2c:[-x 8, -k 1M]").')
|
|
parser.add_argument('--wait', action='store_true', help='Wait for the download to complete and report its status. Note: This makes the operation synchronous and will block until the download finishes.')
|
|
parser.add_argument('--wait-timeout', help='Timeout in seconds for waiting on downloads. Use "auto" to calculate based on a minimum speed of 200KiB/s. Requires --wait. Default: no timeout.')
|
|
parser.add_argument('--auto-merge-fragments', action='store_true', help='Automatically merge fragments after download. Requires --wait and assumes the script has filesystem access to the aria2c host.')
|
|
parser.add_argument('--remove-fragments-after-merge', action='store_true', help='Delete individual fragment files after a successful merge. Requires --auto-merge-fragments.')
|
|
parser.add_argument('--cleanup', action='store_true', help='After a successful download, remove the final file(s) from the filesystem. For fragmented downloads, this implies --remove-fragments-after-merge.')
|
|
parser.add_argument('--remove-on-complete', action=argparse.BooleanOptionalAction, default=True, help='Remove the download from aria2c history on successful completion. Use --no-remove-on-complete to disable. May fail on older aria2c daemons.')
|
|
parser.add_argument('--purge-on-complete', action='store_true', help='Use aria2.purgeDownloadResult to clear ALL completed/failed downloads from history on success. Use as a workaround for older daemons.')
|
|
parser.add_argument('--verbose', action='store_true', help='Enable verbose output for this script.')
|
|
return parser
|
|
|
|
def cleanup_aria_download(api, downloads):
|
|
"""Pause and remove downloads from aria2c."""
|
|
if not downloads:
|
|
return
|
|
try:
|
|
logger.info(f"Attempting to clean up {len(downloads)} download(s) from aria2c...")
|
|
# Filter out downloads that might already be gone
|
|
valid_downloads = [d for d in downloads if hasattr(d, 'gid')]
|
|
if not valid_downloads:
|
|
logger.info("No valid downloads to clean up.")
|
|
return
|
|
api.pause(valid_downloads)
|
|
# Give aria2c a moment to process the pause command before removing
|
|
time.sleep(0.5)
|
|
api.remove(valid_downloads)
|
|
logger.info("Cleanup successful.")
|
|
except Exception as e:
|
|
logger.warning(f"An error occurred during aria2c cleanup: {e}")
|
|
|
|
|
|
def parse_aria_error(download):
|
|
"""Parses an aria2p Download object to get a detailed error message."""
|
|
error_code = download.error_code
|
|
error_message = download.error_message
|
|
|
|
if not error_message:
|
|
return f"Unknown aria2c error (Code: {error_code})"
|
|
|
|
# Check for common HTTP errors in the message
|
|
http_status_match = re.search(r'HTTP status (\d+)', error_message)
|
|
if http_status_match:
|
|
status_code = int(http_status_match.group(1))
|
|
if status_code == 403:
|
|
return f"HTTP Error 403: Forbidden. The URL may have expired or requires valid cookies/headers."
|
|
elif status_code == 404:
|
|
return f"HTTP Error 404: Not Found. The resource is unavailable."
|
|
else:
|
|
return f"HTTP Error {status_code}."
|
|
|
|
if "Timeout" in error_message or "timed out" in error_message.lower():
|
|
return "Download timed out."
|
|
|
|
# Fallback to the raw error message
|
|
return f"Aria2c error (Code: {error_code}): {error_message}"
|
|
|
|
|
|
def parse_aria_args_to_options(args_str):
|
|
"""
|
|
Parses yt-dlp style downloader args for aria2c.
|
|
Example: "aria2c:[-x 8, -k 1M]" or just "-x 8 -k 1M"
|
|
Returns a dictionary of options for aria2p.
|
|
"""
|
|
if not args_str or not args_str.strip():
|
|
return {}
|
|
|
|
inner_args_str = args_str.strip()
|
|
match = re.match(r'aria2c:\s*\[(.*)\]', inner_args_str)
|
|
if match:
|
|
# Handle yt-dlp's format
|
|
inner_args_str = match.group(1).replace(',', ' ')
|
|
else:
|
|
# If it doesn't match, assume the whole string is a set of arguments.
|
|
logger.debug(f"Downloader args '{args_str}' does not match 'aria2c:[...]' format. Parsing as a raw argument string.")
|
|
|
|
arg_list = shlex.split(inner_args_str)
|
|
|
|
# Use a mini-parser to handle CLI-style args
|
|
parser = argparse.ArgumentParser(add_help=False, prog="aria2c_args_parser")
|
|
parser.add_argument('-x', '--max-connection-per-server')
|
|
parser.add_argument('-k', '--min-split-size')
|
|
parser.add_argument('-s', '--split')
|
|
parser.add_argument('--all-proxy')
|
|
|
|
try:
|
|
# We only care about known arguments
|
|
known_args, unknown_args = parser.parse_known_args(arg_list)
|
|
if unknown_args:
|
|
logger.warning(f"Ignoring unknown arguments in --downloader-args: {unknown_args}")
|
|
# Convert to dict, removing None values
|
|
return {k: v for k, v in vars(known_args).items() if v is not None}
|
|
except Exception:
|
|
logger.warning(f"Failed to parse arguments inside --downloader-args: '{inner_args_str}'")
|
|
return {}
|
|
|
|
|
|
def main_download_aria(args):
|
|
"""Main logic for the 'download-aria' command."""
|
|
log_level = logging.DEBUG if args.verbose else logging.INFO
|
|
logging.basicConfig(level=log_level, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', stream=sys.stderr)
|
|
|
|
if args.remove_fragments_after_merge and not args.auto_merge_fragments:
|
|
logger.error("--remove-fragments-after-merge requires --auto-merge-fragments.")
|
|
return 1
|
|
if args.auto_merge_fragments and not args.wait:
|
|
logger.error("--auto-merge-fragments requires --wait.")
|
|
return 1
|
|
if args.wait_timeout and not args.wait:
|
|
logger.error("--wait-timeout requires --wait.")
|
|
return 1
|
|
|
|
if args.wait:
|
|
logger.info("Will wait for download to complete and report status. This is a synchronous operation.")
|
|
else:
|
|
logger.info("Will submit download and exit immediately (asynchronous).")
|
|
|
|
info_json_content = ""
|
|
input_source_name = ""
|
|
if args.load_info_json:
|
|
info_json_content = args.load_info_json.read()
|
|
input_source_name = args.load_info_json.name
|
|
else:
|
|
info_json_content = sys.stdin.read()
|
|
input_source_name = "stdin"
|
|
|
|
if not info_json_content.strip():
|
|
logger.error(f"Failed to read info.json from {input_source_name}. Input is empty.")
|
|
return 1
|
|
|
|
try:
|
|
info_data = json.loads(info_json_content)
|
|
logger.info(f"Successfully loaded info.json from {input_source_name}.")
|
|
except json.JSONDecodeError:
|
|
logger.error(f"Failed to parse info.json from {input_source_name}. Is the input valid JSON?")
|
|
return 1
|
|
|
|
# Find the requested format, supporting yt-dlp style selectors
|
|
target_format = None
|
|
# A format selector can be a comma-separated list of preferences,
|
|
# where each preference can be a slash-separated list of format_ids.
|
|
# e.g., "299/137/136,140" means try 299, then 137, then 136, then 140.
|
|
format_preferences = [item.strip() for sublist in (i.split('/') for i in args.format.split(',')) for item in sublist if item.strip()]
|
|
|
|
available_formats_map = {f['format_id']: f for f in info_data.get('formats', []) if 'format_id' in f}
|
|
|
|
for format_id in format_preferences:
|
|
if format_id in available_formats_map:
|
|
target_format = available_formats_map[format_id]
|
|
logger.info(f"Selected format ID '{format_id}' from selector '{args.format}'.")
|
|
break
|
|
|
|
if not target_format:
|
|
logger.error(f"No suitable format found for selector '{args.format}' in info.json.")
|
|
return 1
|
|
|
|
# Get file size for auto-timeout and dynamic options
|
|
total_filesize = target_format.get('filesize') or target_format.get('filesize_approx')
|
|
|
|
# Construct filename
|
|
video_id = info_data.get('id', 'unknown_video_id')
|
|
title = info_data.get('title', 'unknown_title')
|
|
ext = target_format.get('ext', 'mp4')
|
|
# Sanitize title for filename
|
|
safe_title = "".join([c for c in title if c.isalpha() or c.isdigit() or c in (' ', '-', '_')]).rstrip()
|
|
filename = f"{safe_title} [{video_id}].f{target_format['format_id']}.{ext}"
|
|
|
|
# Prepare options for aria2
|
|
aria_options = {
|
|
# Options from yt-dlp's aria2c integration for performance and reliability
|
|
'max-connection-per-server': 16,
|
|
'split': 16,
|
|
'min-split-size': '1M',
|
|
'http-accept-gzip': 'true',
|
|
'file-allocation': 'none',
|
|
}
|
|
|
|
if args.proxy:
|
|
aria_options['all-proxy'] = args.proxy
|
|
|
|
custom_options = parse_aria_args_to_options(args.downloader_args)
|
|
|
|
# Dynamically set min-split-size if not overridden by user
|
|
if 'min_split_size' not in custom_options and total_filesize:
|
|
if total_filesize > 100 * 1024 * 1024: # 100 MiB
|
|
aria_options['min-split-size'] = '5M'
|
|
logger.info("File is > 100MiB, dynamically setting min-split-size to 5M.")
|
|
|
|
if custom_options:
|
|
aria_options.update(custom_options)
|
|
logger.info(f"Applied custom aria2c options from --downloader-args: {custom_options}")
|
|
|
|
aria_options['out'] = filename
|
|
|
|
# Add headers from info.json, mimicking yt-dlp's behavior for aria2c
|
|
headers = target_format.get('http_headers')
|
|
if headers:
|
|
header_list = [f'{key}: {value}' for key, value in headers.items()]
|
|
aria_options['header'] = header_list
|
|
logger.info(f"Adding {len(header_list)} HTTP headers to the download.")
|
|
if args.verbose:
|
|
for h in header_list:
|
|
if h.lower().startswith('cookie:'):
|
|
logger.debug(f" Header: Cookie: [REDACTED]")
|
|
else:
|
|
logger.debug(f" Header: {h}")
|
|
|
|
is_fragmented = 'fragments' in target_format
|
|
if not is_fragmented:
|
|
url = target_format.get('url')
|
|
if not url:
|
|
logger.error(f"Format ID '{args.format}' has neither a URL nor fragments.")
|
|
return 1
|
|
|
|
try:
|
|
logger.info(f"Connecting to aria2c RPC at http://{args.aria_host}:{args.aria_port}")
|
|
client = aria2p.Client(
|
|
host=f"http://{args.aria_host}",
|
|
port=args.aria_port,
|
|
secret=args.aria_secret or ""
|
|
)
|
|
api = aria2p.API(client)
|
|
|
|
timeout_seconds = None
|
|
if args.wait_timeout:
|
|
if args.wait_timeout.lower() == 'auto':
|
|
if total_filesize:
|
|
# Min speed: 200 KiB/s. Min timeout: 30s.
|
|
min_speed = 200 * 1024
|
|
calculated_timeout = int(total_filesize / min_speed)
|
|
timeout_seconds = max(30, calculated_timeout)
|
|
total_filesize_hr, _ = human_readable_bytes(total_filesize)
|
|
logger.info(f"Auto-calculated timeout: {timeout_seconds}s (based on {total_filesize_hr} at 200KiB/s).")
|
|
else:
|
|
logger.warning("Cannot use 'auto' timeout: file size not available in info.json. Timeout disabled.")
|
|
else:
|
|
try:
|
|
timeout_seconds = int(args.wait_timeout)
|
|
if timeout_seconds <= 0:
|
|
raise ValueError
|
|
except ValueError:
|
|
logger.error(f"Invalid --wait-timeout value: '{args.wait_timeout}'. Must be a positive integer or 'auto'.")
|
|
return 1
|
|
|
|
if is_fragmented:
|
|
return download_fragments_aria(args, api, target_format, filename, aria_options, timeout_seconds, remote_dir=args.remote_dir)
|
|
else:
|
|
return download_url_aria(args, api, url, filename, aria_options, timeout_seconds, remote_dir=args.remote_dir)
|
|
|
|
except Exception as e:
|
|
logger.error(f"An error occurred while communicating with aria2c: {e}", exc_info=args.verbose)
|
|
return 1
|
|
|
|
def download_url_aria(args, api, url, filename, aria_options, timeout_seconds, remote_dir=None):
|
|
"""Handle downloading a single URL with aria2c."""
|
|
if remote_dir:
|
|
aria_options['dir'] = remote_dir
|
|
logger.info(f"Adding download for format '{args.format}' with URL: {url[:70]}...")
|
|
downloads = api.add_uris([url], options=aria_options)
|
|
|
|
if not downloads:
|
|
logger.error("Failed to add download to aria2c. The API returned an empty result.")
|
|
return 1
|
|
|
|
# Handle older aria2p versions that return a single Download object instead of a list
|
|
download = downloads[0] if isinstance(downloads, list) else downloads
|
|
logger.info(f"Successfully added download to aria2c. GID: {download.gid}")
|
|
|
|
if args.wait:
|
|
logger.info(f"Waiting for download {download.gid} to complete...")
|
|
start_time = time.time()
|
|
try:
|
|
while True:
|
|
if timeout_seconds and (time.time() - start_time > timeout_seconds):
|
|
raise TimeoutError(f"Download did not complete within {timeout_seconds}s timeout.")
|
|
|
|
# Re-fetch the download object to get the latest status
|
|
download.update()
|
|
# A download is no longer active if it's complete, errored, paused, or removed.
|
|
if download.status not in ('active', 'waiting'):
|
|
break
|
|
|
|
progress_info = (
|
|
f"\rGID {download.gid}: {download.status} "
|
|
f"{download.progress_string()} "
|
|
f"({download.download_speed_string()}) "
|
|
f"ETA: {download.eta_string()}"
|
|
)
|
|
sys.stdout.write(progress_info)
|
|
sys.stdout.flush()
|
|
time.sleep(0.5)
|
|
except (KeyboardInterrupt, TimeoutError) as e:
|
|
sys.stdout.write('\n')
|
|
if isinstance(e, KeyboardInterrupt):
|
|
logger.warning("Wait interrupted by user. Cleaning up download...")
|
|
cleanup_aria_download(api, [download])
|
|
return 130
|
|
else: # TimeoutError
|
|
logger.error(f"Download timed out. Cleaning up... Error: {e}")
|
|
cleanup_aria_download(api, [download])
|
|
return 1
|
|
except aria2p.ClientException as e:
|
|
# This can happen if the download completes and is removed by aria2c
|
|
# before we can check its final status. Assume success in this case.
|
|
logger.warning(f"Could not get final status for GID {download.gid} (maybe removed on completion?): {e}. Assuming success.")
|
|
print(f"Download for GID {download.gid} presumed successful.")
|
|
return 0
|
|
|
|
sys.stdout.write('\n') # Newline after progress bar
|
|
|
|
# Final status check (no need to update again, we have the latest status)
|
|
if download.status == 'complete':
|
|
logger.info(f"Download {download.gid} completed successfully.")
|
|
|
|
downloaded_filepath_remote = None
|
|
if download.files:
|
|
downloaded_filepath_remote = download.files[0].path
|
|
print(f"Download successful: {downloaded_filepath_remote}")
|
|
else:
|
|
print("Download successful, but no file path reported by aria2c.")
|
|
|
|
if args.cleanup and downloaded_filepath_remote:
|
|
local_filepath = None
|
|
# To map remote path to local, we need remote_dir and a local equivalent.
|
|
# We'll use fragments_dir as the local equivalent, which defaults to output_dir.
|
|
local_base_dir = args.fragments_dir or args.output_dir or '.'
|
|
if remote_dir:
|
|
if downloaded_filepath_remote.startswith(remote_dir):
|
|
relative_path = os.path.relpath(downloaded_filepath_remote, remote_dir)
|
|
local_filepath = os.path.join(local_base_dir, relative_path)
|
|
else:
|
|
logger.warning(f"Cleanup: Downloaded file path '{downloaded_filepath_remote}' does not start with remote-dir '{remote_dir}'. Cannot map to local path.")
|
|
else:
|
|
logger.warning(f"Cleanup: --remote-dir not specified. Assuming download path is accessible locally as '{downloaded_filepath_remote}'.")
|
|
local_filepath = downloaded_filepath_remote
|
|
|
|
if local_filepath:
|
|
try:
|
|
if os.path.exists(local_filepath):
|
|
os.remove(local_filepath)
|
|
logger.info(f"Cleanup: Removed downloaded file '{local_filepath}'")
|
|
else:
|
|
logger.warning(f"Cleanup: File not found at expected local path '{local_filepath}'. Skipping removal.")
|
|
except OSError as e:
|
|
logger.error(f"Cleanup failed: Could not remove file '{local_filepath}': {e}")
|
|
elif args.cleanup:
|
|
logger.warning("Cleanup requested, but no downloaded file path was reported by aria2c.")
|
|
|
|
if args.purge_on_complete:
|
|
try:
|
|
api.purge_download_result()
|
|
logger.info("Purged all completed/failed downloads from aria2c history.")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to purge download history: {e}")
|
|
elif args.remove_on_complete:
|
|
try:
|
|
api.remove_download_result(download)
|
|
logger.info(f"Removed download {download.gid} from aria2c history.")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to remove download {download.gid} from history: {e}")
|
|
|
|
return 0
|
|
else:
|
|
detailed_error = parse_aria_error(download)
|
|
logger.error(f"Download {download.gid} failed. Error: {detailed_error}")
|
|
return 1
|
|
else:
|
|
print(f"Successfully added download. GID: {download.gid}")
|
|
return 0
|
|
|
|
def download_fragments_aria(args, api, target_format, filename, aria_options, timeout_seconds, remote_dir=None):
|
|
"""Handle downloading fragmented formats with aria2c."""
|
|
logger.info(f"Format '{args.format}' is fragmented. Adding all fragments to download queue.")
|
|
fragment_base_url = target_format.get('fragment_base_url')
|
|
fragments = target_format['fragments']
|
|
|
|
MAX_FRAGMENTS = 50000
|
|
if len(fragments) > MAX_FRAGMENTS:
|
|
logger.error(
|
|
f"The number of fragments ({len(fragments)}) exceeds the safety limit of {MAX_FRAGMENTS}. "
|
|
f"This is to prevent overwhelming the aria2c server. Aborting."
|
|
)
|
|
return 1
|
|
|
|
# We need to set the 'dir' option for all fragments if specified.
|
|
# The 'out' option will be set per-fragment.
|
|
frag_aria_options = aria_options.copy()
|
|
frag_aria_options.pop('out', None) # Remove the main 'out' option
|
|
|
|
if remote_dir:
|
|
frag_aria_options['dir'] = remote_dir
|
|
logger.info(f"Instructing remote aria2c to save fragments to: {remote_dir}")
|
|
|
|
base_filename, file_ext = os.path.splitext(filename)
|
|
|
|
calls = []
|
|
for i, fragment in enumerate(fragments):
|
|
frag_url = fragment.get('url')
|
|
if not frag_url:
|
|
if not fragment_base_url:
|
|
logger.error(f"Fragment {i} has no URL and no fragment_base_url is available. Aborting.")
|
|
return 1
|
|
frag_url = urljoin(fragment_base_url, fragment['path'])
|
|
|
|
# Use the base filename from the main file, but add fragment identifier
|
|
fragment_filename = f"{base_filename}-Frag{i}{file_ext}"
|
|
|
|
current_frag_options = frag_aria_options.copy()
|
|
current_frag_options['out'] = os.path.basename(fragment_filename)
|
|
|
|
# Prepare parameters for multicall in the format:
|
|
# {"methodName": "aria2.addUri", "params": [["url"], {"out": "file.mp4"}]}
|
|
# The secret token is automatically added by aria2p.
|
|
params = [[frag_url], current_frag_options]
|
|
call_struct = {
|
|
"methodName": api.client.ADD_URI,
|
|
"params": params
|
|
}
|
|
calls.append(call_struct)
|
|
|
|
results = api.client.multicall(calls)
|
|
if not results:
|
|
logger.error("Failed to add fragments to aria2c. The API returned an empty result.")
|
|
return 1
|
|
|
|
# The result of a multicall of addUri is a list of lists, where each inner list
|
|
# contains the GID of one download, e.g., [['gid1'], ['gid2']].
|
|
# A failed call for a fragment may result in a fault struct dict instead of a list.
|
|
# We extract GIDs from successful calls.
|
|
gids = [result[0] for result in results if isinstance(result, list) and result]
|
|
|
|
if len(gids) != len(fragments):
|
|
failed_count = len(fragments) - len(gids)
|
|
logger.warning(f"{failed_count} out of {len(fragments)} fragments failed to be added to aria2c.")
|
|
|
|
if not gids:
|
|
logger.error("Failed to add any fragments to aria2c. All submissions failed.")
|
|
return 1
|
|
|
|
logger.info(f"Successfully added {len(gids)} fragments to aria2c.")
|
|
if args.verbose:
|
|
logger.debug(f"GIDs: {gids}")
|
|
|
|
if args.wait:
|
|
logger.info(f"Waiting for {len(gids)} fragments to complete...")
|
|
start_time = time.time()
|
|
downloads_to_cleanup = []
|
|
try:
|
|
while True:
|
|
if timeout_seconds and (time.time() - start_time > timeout_seconds):
|
|
raise TimeoutError(f"Fragment downloads did not complete within {timeout_seconds}s timeout.")
|
|
|
|
downloads = api.get_downloads(gids)
|
|
downloads_to_cleanup = downloads # Store for potential cleanup
|
|
# A download is considered "active" if it's currently downloading or waiting in the queue.
|
|
# It is "not active" if it is complete, errored, paused, or removed.
|
|
active_downloads = [d for d in downloads if d.status in ('active', 'waiting')]
|
|
if not active_downloads:
|
|
break # All downloads are complete or have stopped for other reasons
|
|
|
|
for d in active_downloads:
|
|
d.update()
|
|
|
|
completed_count = len(downloads) - len(active_downloads)
|
|
total_bytes = sum(d.total_length for d in downloads)
|
|
downloaded_bytes = sum(d.completed_length for d in downloads)
|
|
total_speed = sum(d.download_speed for d in downloads)
|
|
progress_percent = (downloaded_bytes / total_bytes * 100) if total_bytes > 0 else 0
|
|
|
|
progress_info = (
|
|
f"\rProgress: {completed_count}/{len(downloads)} fragments | "
|
|
f"{progress_percent:.1f}% "
|
|
f"({human_readable_bytes(downloaded_bytes)}/{human_readable_bytes(total_bytes)}) "
|
|
f"Speed: {human_readable_bytes(total_speed)}/s"
|
|
)
|
|
sys.stdout.write(progress_info)
|
|
sys.stdout.flush()
|
|
time.sleep(0.5)
|
|
except (KeyboardInterrupt, TimeoutError) as e:
|
|
sys.stdout.write('\n')
|
|
if isinstance(e, KeyboardInterrupt):
|
|
logger.warning("Wait interrupted by user. Cleaning up fragments...")
|
|
cleanup_aria_download(api, downloads_to_cleanup)
|
|
return 130
|
|
else: # TimeoutError
|
|
logger.error(f"Download timed out. Cleaning up fragments... Error: {e}")
|
|
cleanup_aria_download(api, downloads_to_cleanup)
|
|
return 1
|
|
except aria2p.ClientException as e:
|
|
# This can happen if downloads complete and are removed by aria2c
|
|
# before we can check their final status. Assume success in this case.
|
|
logger.warning(f"Could not get final status for some fragments (maybe removed on completion?): {e}. Assuming success.")
|
|
|
|
sys.stdout.write('\n')
|
|
|
|
# Final status check
|
|
failed_downloads = []
|
|
try:
|
|
downloads = api.get_downloads(gids)
|
|
failed_downloads = [d for d in downloads if d.status != 'complete']
|
|
except aria2p.ClientException as e:
|
|
logger.warning(f"Could not perform final status check for fragments (maybe removed on completion?): {e}. Assuming success.")
|
|
# If we can't check, we assume success based on the earlier wait loop not failing catastrophically.
|
|
failed_downloads = []
|
|
|
|
if failed_downloads:
|
|
logger.error(f"{len(failed_downloads)} fragments failed to download.")
|
|
for d in failed_downloads:
|
|
detailed_error = parse_aria_error(d)
|
|
logger.error(f" GID {d.gid}: {detailed_error}")
|
|
return 1
|
|
else:
|
|
logger.info("All fragments downloaded successfully.")
|
|
output_dir = args.output_dir or '.'
|
|
final_filepath = os.path.join(output_dir, filename)
|
|
fragments_lookup_dir = args.fragments_dir or output_dir
|
|
|
|
if args.auto_merge_fragments:
|
|
logger.info(f"Attempting to merge fragments into: {final_filepath}")
|
|
logger.info(f"Searching for fragments in local directory: {os.path.abspath(fragments_lookup_dir)}")
|
|
try:
|
|
# base_filename and file_ext are available from earlier in the function
|
|
# We must escape the base filename in case it contains glob special characters like [ or ].
|
|
escaped_base = glob.escape(base_filename)
|
|
search_path = os.path.join(fragments_lookup_dir, f"{escaped_base}-Frag*{file_ext}")
|
|
fragment_files = glob.glob(search_path)
|
|
|
|
if not fragment_files:
|
|
logger.error(f"No fragment files found with pattern: {search_path}")
|
|
return 1
|
|
|
|
def fragment_sort_key(f):
|
|
match = re.search(r'Frag(\d+)', os.path.basename(f))
|
|
return int(match.group(1)) if match else -1
|
|
fragment_files.sort(key=fragment_sort_key)
|
|
|
|
with open(final_filepath, 'wb') as dest_file:
|
|
for frag_path in fragment_files:
|
|
with open(frag_path, 'rb') as src_file:
|
|
shutil.copyfileobj(src_file, dest_file)
|
|
|
|
logger.info(f"Successfully merged {len(fragment_files)} fragments into {final_filepath}")
|
|
|
|
if args.remove_fragments_after_merge or args.cleanup:
|
|
logger.info("Removing fragment files...")
|
|
for frag_path in fragment_files:
|
|
os.remove(frag_path)
|
|
logger.info("Fragment files removed.")
|
|
|
|
if args.cleanup:
|
|
try:
|
|
os.remove(final_filepath)
|
|
logger.info(f"Cleanup: Removed merged file '{final_filepath}'")
|
|
except OSError as e:
|
|
logger.error(f"Cleanup failed: Could not remove merged file '{final_filepath}': {e}")
|
|
|
|
print(f"Download and merge successful: {final_filepath}")
|
|
|
|
if args.purge_on_complete:
|
|
try:
|
|
api.purge_download_result()
|
|
logger.info("Purged all completed/failed downloads from aria2c history.")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to purge download history: {e}")
|
|
elif args.remove_on_complete:
|
|
try:
|
|
# The `downloads` variable from the last status check should be valid here.
|
|
api.remove_download_result(downloads)
|
|
logger.info(f"Removed {len(downloads)} fragment downloads from aria2c history.")
|
|
except aria2p.ClientException as e:
|
|
logger.warning(f"Could not remove fragment downloads from history (maybe already gone?): {e}")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to remove fragment downloads from history: {e}")
|
|
|
|
return 0
|
|
|
|
except Exception as e:
|
|
logger.error(f"An error occurred during merging: {e}", exc_info=args.verbose)
|
|
logger.error("Fragments were downloaded but not merged.")
|
|
return 1
|
|
else:
|
|
print("Download successful. Fragments now need to be merged manually.")
|
|
print(f"The final merged file should be named: {final_filepath}")
|
|
print("You can merge them with a command like:")
|
|
print(f" cat `ls -v '{os.path.join(fragments_lookup_dir, base_filename)}'-Frag*'{file_ext}'` > '{final_filepath}'")
|
|
|
|
if args.cleanup:
|
|
logger.info("Cleanup requested. Removing downloaded fragments...")
|
|
try:
|
|
# base_filename and file_ext are available from earlier in the function
|
|
escaped_base = glob.escape(base_filename)
|
|
search_path = os.path.join(fragments_lookup_dir, f"{escaped_base}-Frag*{file_ext}")
|
|
fragment_files = glob.glob(search_path)
|
|
|
|
if not fragment_files:
|
|
logger.warning(f"Cleanup: No fragment files found with pattern: {search_path}")
|
|
else:
|
|
for frag_path in fragment_files:
|
|
os.remove(frag_path)
|
|
logger.info(f"Removed {len(fragment_files)} fragment files.")
|
|
except Exception as e:
|
|
logger.error(f"An error occurred during fragment cleanup: {e}", exc_info=args.verbose)
|
|
|
|
if args.purge_on_complete:
|
|
try:
|
|
api.purge_download_result()
|
|
logger.info("Purged all completed/failed downloads from aria2c history.")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to purge download history: {e}")
|
|
elif args.remove_on_complete:
|
|
try:
|
|
# The `downloads` variable from the last status check should be valid here.
|
|
api.remove_download_result(downloads)
|
|
logger.info(f"Removed {len(downloads)} fragment downloads from aria2c history.")
|
|
except aria2p.ClientException as e:
|
|
logger.warning(f"Could not remove fragment downloads from history (maybe already gone?): {e}")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to remove fragment downloads from history: {e}")
|
|
|
|
return 0
|
|
else:
|
|
print(f"Successfully added {len(gids)} fragments. GIDs: {gids}")
|
|
print("These fragments will need to be merged manually after download.")
|
|
return 0
|