diff --git a/air2.py b/air2.py new file mode 100644 index 0000000..d079540 --- /dev/null +++ b/air2.py @@ -0,0 +1,1640 @@ +import os +import sys +import time +import json +import signal +import threading +import hashlib +import concurrent.futures +import zipfile +import argparse +from pathlib import Path +from typing import Optional, Dict, List, Tuple, Callable, Any +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from urllib.parse import urlparse +from enum import Enum +import queue +import math +import logging + +# 第三方库 +import requests +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry +from rich.console import Console +from rich.progress import ( + Progress, + TextColumn, + BarColumn, + DownloadColumn, + TransferSpeedColumn, + TimeRemainingColumn, + TaskID +) +from rich.table import Table +from rich.panel import Panel +from rich.live import Live +from rich.text import Text +from rich.prompt import Prompt, Confirm +from rich.style import Style +from rich.theme import Theme + +# 新的现代化CLI依赖 +try: + from prompt_toolkit import PromptSession + from prompt_toolkit.history import FileHistory + from prompt_toolkit.auto_suggest import AutoSuggestFromHistory + from prompt_toolkit.completion import Completer, Completion, WordCompleter + from prompt_toolkit.styles import Style as PtStyle + from prompt_toolkit.key_binding import KeyBindings + PROMPT_TOOLKIT_AVAILABLE = True +except ImportError: + PROMPT_TOOLKIT_AVAILABLE = False + print("提示: 安装 prompt_toolkit 可获得命令补全和历史记录功能") + print("建议执行: pip install prompt_toolkit") + +import urllib3 +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + +# 配置日志 +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler("air2.log"), + logging.StreamHandler(sys.stdout) + ] +) +logger = logging.getLogger("air2") + + +class DownloadStatus(Enum): + WAITING = "等待" + DOWNLOADING = "下载中" + PAUSED = "已暂停" + COMPLETED = "已完成" + ERROR = "错误" + CANCELED = "已取消" + + +@dataclass +class DownloadTask: + """下载任务数据结构""" + url: str + filepath: Path + filesize: int = 0 + downloaded: int = 0 + status: DownloadStatus = DownloadStatus.WAITING + chunks: Dict[int, Tuple[int, int]] = field(default_factory=dict) + threads: int = 8 + max_connections: int = 16 + chunk_size: int = 1024 * 1024 # 动态调整,初始化时设置 + timeout: int = 30 + retry_count: int = 5 + speed_limit: int = 0 + checksum: str = "" + metadata: Dict = field(default_factory=dict) + created_at: datetime = field(default_factory=datetime.now) + error_msg: str = "" + start_time: Optional[datetime] = None + last_update: datetime = field(default_factory=datetime.now) + avg_speed: float = 0.0 + current_speed: float = 0.0 + eta: timedelta = field(default_factory=timedelta) + + @property + def progress(self) -> float: + if self.filesize == 0: + return 0.0 + return (self.downloaded / self.filesize) * 100 + + def update_eta(self): + if self.status != DownloadStatus.DOWNLOADING or self.current_speed <= 0: + self.eta = timedelta(seconds=0) + return + remaining_bytes = self.filesize - self.downloaded + if remaining_bytes <= 0: + self.eta = timedelta(seconds=0) + return + remaining_kb = remaining_bytes / 1024 + seconds = remaining_kb / self.current_speed + self.eta = timedelta(seconds=int(seconds)) + + +class ChunkDownloader: + """分块下载器 - 优化版""" + + def __init__(self, task: DownloadTask, chunk_id: int, start: int, end: int): + self.task = task + self.chunk_id = chunk_id + self.start = start + self.end = end + self.position = start + self.downloaded = 0 + self.running = False + self.session = None + + def download(self, session: requests.Session, progress_queue: queue.Queue): + """下载指定分块(大缓冲区优化)""" + self.running = True + self.session = session + headers = {'Range': f'bytes={self.position}-{self.end}'} + retry_count = 0 + + last_progress_time = time.time() + bytes_since_last_progress = 0 + PROGRESS_INTERVAL_SEC = 1.0 # 每秒最多一次 + PROGRESS_INTERVAL_BYTES = 5 * 1024 * 1024 # 或每5MB + + while retry_count <= self.task.retry_count and self.running: + try: + logger.debug(f"开始下载分块 {self.chunk_id}, 范围: {self.position}-{self.end}") + response = session.get( + self.task.url, + headers=headers, + stream=True, + timeout=(self.task.timeout, self.task.timeout), + verify=False + ) + response.raise_for_status() + + content_range = response.headers.get('Content-Range', '') + if not content_range and self.position > 0: + logger.warning(f"服务器不支持范围请求,重新请求整个文件") + headers = {'Range': f'bytes=0-{self.end}'} + self.position = 0 + continue + + mode = 'rb+' if self.position > 0 else 'wb' + with open(self.task.filepath, mode) as f: + f.seek(self.position) + # 使用 1MB 缓冲区减少系统调用 + for chunk in response.iter_content(chunk_size=1024 * 1024): + if not self.running: + logger.info(f"分块 {self.chunk_id} 被暂停") + break + if chunk: + f.write(chunk) + chunk_len = len(chunk) + self.position += chunk_len + self.downloaded += chunk_len + bytes_since_last_progress += chunk_len + + # 限速控制 + if self.task.speed_limit > 0: + time.sleep(chunk_len / (self.task.speed_limit * 1024)) + + # 智能进度报告 + now = time.time() + if (now - last_progress_time >= PROGRESS_INTERVAL_SEC or + bytes_since_last_progress >= PROGRESS_INTERVAL_BYTES): + progress_queue.put((self.task.url, self.chunk_id, self.downloaded)) + last_progress_time = now + bytes_since_last_progress = 0 + + progress_queue.put((self.task.url, self.chunk_id, self.downloaded)) + + if self.running and self.position >= self.end + 1: + logger.info(f"分块 {self.chunk_id} 下载完成") + return True + elif not self.running: + logger.info(f"分块 {self.chunk_id} 暂停在位置 {self.position}") + return False + else: + logger.warning(f"分块 {self.chunk_id} 未完整下载: {self.position}/{self.end + 1}") + retry_count += 1 + time.sleep(2 ** min(retry_count, 5)) + + except requests.exceptions.RequestException as e: + logger.error(f"分块 {self.chunk_id} 请求失败: {str(e)}") + retry_count += 1 + if retry_count > self.task.retry_count: + progress_queue.put(('error', self.task.url, f"分块{self.chunk_id}: {str(e)}")) + return False + time.sleep(2 ** min(retry_count, 5)) + except IOError as e: + logger.error(f"分块 {self.chunk_id} 文件IO错误: {str(e)}") + progress_queue.put(('error', self.task.url, f"分块{self.chunk_id}: {str(e)}")) + return False + except Exception as e: + logger.exception(f"分块 {self.chunk_id} 未知错误") + progress_queue.put(('error', self.task.url, f"分块{self.chunk_id}: {str(e)}")) + return False + return False + + def pause(self): + self.running = False + if self.session: + try: + self.session.close() + except: + pass + + +class ProtocolHandler: + @staticmethod + def supports(url: str) -> bool: + raise NotImplementedError + + @staticmethod + def get_info(url: str, session: requests.Session) -> Tuple[int, str, Dict]: + raise NotImplementedError + + +class HTTPHandler(ProtocolHandler): + @staticmethod + def supports(url: str) -> bool: + parsed = urlparse(url) + return parsed.scheme in ['http', 'https'] + + @staticmethod + def get_info(url: str, session: requests.Session) -> Tuple[int, str, Dict]: + try: + logger.info(f"获取文件信息: {url}") + response = session.get( + url, + headers={'Range': 'bytes=0-1'}, + timeout=15, + stream=True, + verify=False + ) + response.raise_for_status() + + content_range = response.headers.get('Content-Range', '') + if content_range: + filesize = int(content_range.split('/')[-1]) + else: + filesize = int(response.headers.get('Content-Length', 0)) + + if filesize == 0: + logger.warning("无法获取文件大小,将完整下载文件") + + filename = None + if 'content-disposition' in response.headers: + content_disposition = response.headers['content-disposition'] + if 'filename=' in content_disposition: + filename = content_disposition.split('filename=')[1].strip("\"'") + + if not filename: + parsed = urlparse(url) + filename = os.path.basename(parsed.path) or 'download.bin' + + filename = "".join(c for c in filename if c.isalnum() or c in ('.', '_', '-')).rstrip() + if not filename: + filename = 'download.bin' + + accept_ranges = response.headers.get('Accept-Ranges', '').lower() == 'bytes' + response.close() + + logger.info(f"获取文件信息成功: {filename}, 大小: {filesize} bytes, 支持断点续传: {accept_ranges}") + return filesize, filename, { + 'accept_ranges': accept_ranges, + 'content_type': response.headers.get('Content-Type', ''), + 'last_modified': response.headers.get('Last-Modified', '') + } + except Exception as e: + logger.exception(f"获取文件信息失败: {str(e)}") + raise Exception(f"获取文件信息失败: {str(e)}") + + +class DownloadEngine: + def __init__(self, max_concurrent_tasks: int = 5, max_workers: int = 200): + self.tasks: Dict[str, DownloadTask] = {} + self.active_downloaders: Dict[str, List[ChunkDownloader]] = {} + self.progress_queue = queue.Queue() + self.completed_queue = queue.Queue() + self.task_lock = threading.Lock() + self.running = False + self.console = Console() + self.shutdown_flag = False + + # 增大线程池,提高并发 + self.executor = concurrent.futures.ThreadPoolExecutor( + max_workers=max_workers, + thread_name_prefix='DownloadWorker' + ) + + self.session_pool = {} + self.protocol_handlers = [HTTPHandler()] + self.speed_tracker = {} + self.last_update_time = {} + self.active_task_count = 0 + self.live_active = False + + # 清新配色主题 + custom_theme = Theme({ + "info": "#87CEEB", # 浅蓝 + "success": "#98FB98", # 薄荷绿 + "error": "#F08080", # 淡珊瑚 + "warning": "#FFD700", # 淡金 + "highlight": "#DDA0DD", # 淡紫 + "bar.complete": "#98FB98", + "bar.finished": "#98FB98", + "bar.pulse": "#87CEEB", + "progress.percentage": "#DDA0DD", + }) + self.console.push_theme(custom_theme) + + # 圆角进度条,去除Spinner + self.progress_display = Progress( + TextColumn("[bold #87CEEB]{task.description}"), + BarColumn( + bar_width=None, + complete_style="#98FB98", + finished_style="#98FB98", + pulse_style="#87CEEB" + ), + "[progress.percentage]{task.percentage:>3.1f}%", + "•", + DownloadColumn(), + "•", + TransferSpeedColumn(), + "•", + TimeRemainingColumn(), + console=self.console, + refresh_per_second=4 + ) + self.progress_tasks: Dict[str, TaskID] = {} + self.live = Live(self.progress_display, refresh_per_second=4, console=self.console, auto_refresh=True) + + self.monitor_thread = threading.Thread(target=self._monitor_progress, daemon=True) + self.speed_calc_thread = threading.Thread(target=self._calculate_speeds, daemon=True) + + def _create_session(self, task_id: str) -> requests.Session: + if task_id not in self.session_pool: + session = requests.Session() + + # 优化重试策略 + retry_strategy = Retry( + total=5, + backoff_factor=1, + status_forcelist=[429, 500, 502, 503, 504, 522, 524], + allowed_methods=["GET", "HEAD"], + raise_on_status=False + ) + adapter = HTTPAdapter( + max_retries=retry_strategy, + pool_connections=20, # 增大连接池 + pool_maxsize=100, + pool_block=False + ) + session.mount("http://", adapter) + session.mount("https://", adapter) + + # 优化请求头(支持压缩、长连接) + session.headers.update({ + 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36', + 'Accept': '*/*', + 'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8', + 'Accept-Encoding': 'gzip, deflate, br', + 'Connection': 'keep-alive', + 'Keep-Alive': 'timeout=30, max=1000', + }) + + self.session_pool[task_id] = session + return self.session_pool[task_id] + + def _get_handler(self, url: str) -> ProtocolHandler: + for handler in self.protocol_handlers: + if handler.supports(url): + return handler + raise Exception(f"不支持的协议: {url}") + + def add_task(self, url: str, output_dir: str = ".", **kwargs) -> str: + try: + handler = self._get_handler(url) + session = self._create_session("temp") + filesize, filename, metadata = handler.get_info(url, session) + session.close() + + task_id = hashlib.md5(f"{url}_{datetime.now().timestamp()}".encode()).hexdigest()[:8] + + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + filepath = output_path / filename + counter = 1 + original_filepath = filepath + while filepath.exists(): + stem = original_filepath.stem + suffix = original_filepath.suffix + filepath = output_path / f"{stem}_{counter}{suffix}" + counter += 1 + + # 动态调整分块大小(优化下载速度) + user_chunk_size = kwargs.get('chunk_size', 1024 * 1024) + if filesize > 1024 * 1024 * 1024: # >1GB + dynamic_chunk_size = 16 * 1024 * 1024 + elif filesize > 100 * 1024 * 1024: # >100MB + dynamic_chunk_size = 4 * 1024 * 1024 + else: + dynamic_chunk_size = user_chunk_size + + # 创建任务,确保线程参数正确传递 + task = DownloadTask( + url=url, + filepath=filepath, + filesize=filesize, + metadata=metadata, + threads=kwargs.get('threads', 8), + max_connections=kwargs.get('max_connections', 16), + chunk_size=dynamic_chunk_size, + timeout=kwargs.get('timeout', 30), + retry_count=kwargs.get('retry_count', 5), + speed_limit=kwargs.get('speed_limit', 0) + ) + + with self.task_lock: + self.tasks[task_id] = task + self.speed_tracker[task_id] = {'bytes': 0, 'time': time.time(), 'last_bytes': 0} + + if filepath.exists() and filepath.stat().st_size > 0: + self._resume_task(task_id) + else: + self._prepare_file(task) + + logger.info(f"任务添加成功: {task_id}, 保存到: {filepath}") + return task_id + + except Exception as e: + logger.exception(f"添加任务失败: {str(e)}") + raise + + def _prepare_file(self, task: DownloadTask): + if task.filesize <= 0: + return + try: + # 使用 truncate 快速预分配 + with open(task.filepath, 'wb') as f: + f.truncate(task.filesize) + logger.info(f"预分配文件空间: {task.filepath}, 大小: {task.filesize}") + except Exception as e: + logger.warning(f"预分配文件空间失败,将使用动态写入: {str(e)}") + + def _resume_task(self, task_id: str): + task = self.tasks[task_id] + if not task.filepath.exists(): + return + downloaded = task.filepath.stat().st_size + task.downloaded = downloaded + logger.info(f"恢复下载: {task.filepath}, 已下载: {downloaded}/{task.filesize}") + if task.filesize > 0 and downloaded >= task.filesize: + task.status = DownloadStatus.COMPLETED + + def start_task(self, task_id: str): + with self.task_lock: + if task_id not in self.tasks: + raise Exception("任务不存在") + task = self.tasks[task_id] + if task.status in [DownloadStatus.DOWNLOADING, DownloadStatus.COMPLETED]: + return + if task.status == DownloadStatus.ERROR: + logger.info(f"重新开始失败的任务: {task_id}") + + task.status = DownloadStatus.DOWNLOADING + task.start_time = datetime.now() + task.last_update = datetime.now() + task.current_speed = 0.0 + task.avg_speed = 0.0 + task.eta = timedelta(seconds=0) + + if task_id not in self.progress_tasks: + if task.filesize > 0: + self.progress_tasks[task_id] = self.progress_display.add_task( + description=f"{task.filepath.name[:30]}", + total=task.filesize, + completed=task.downloaded + ) + else: + self.progress_tasks[task_id] = self.progress_display.add_task( + description=f"{task.filepath.name[:30]}", + total=0, + completed=task.downloaded, + start=False + ) + + if not self.live_active: + try: + self.live.start() + self.live_active = True + except Exception as e: + logger.error(f"启动Live显示失败: {str(e)}") + + if not task.chunks and task.filesize > 0: + self._create_chunks(task) + + self._start_downloaders(task_id) + self.active_task_count += 1 + logger.info(f"开始下载任务: {task_id}, 活动任务数: {self.active_task_count}") + + if not self.running: + self.running = True + if not self.monitor_thread.is_alive(): + self.monitor_thread.start() + if not self.speed_calc_thread.is_alive(): + self.speed_calc_thread.start() + + def _create_chunks(self, task: DownloadTask): + if task.filesize <= 0: + return + # 使用 task.threads 作为分块数量 + chunk_count = min(task.threads, max(1, task.filesize // task.chunk_size)) + if chunk_count == 0: + chunk_count = 1 + chunk_size = task.filesize // chunk_count + remainder = task.filesize % chunk_count + + logger.info(f"创建分块: 大小={task.filesize}, 线程={task.threads}, 分块数={chunk_count}, 块大小={chunk_size}") + + for i in range(chunk_count): + start = i * chunk_size + min(i, remainder) + end = start + chunk_size - 1 + if i < remainder: + end += 1 + if i == chunk_count - 1: + end = task.filesize - 1 + task.chunks[i] = (start, end) + + def _start_downloaders(self, task_id: str): + task = self.tasks[task_id] + if task_id not in self.active_downloaders: + self.active_downloaders[task_id] = [] + + session = self._create_session(task_id) + + started_chunks = 0 + for chunk_id, (start, end) in task.chunks.items(): + if task.downloaded >= end + 1: + continue + chunk_start = max(start, task.downloaded) + if chunk_start > end: + continue + downloader = ChunkDownloader(task, chunk_id, chunk_start, end) + self.active_downloaders[task_id].append(downloader) + self.executor.submit(downloader.download, session, self.progress_queue) + started_chunks += 1 + + logger.info(f"启动 {started_chunks} 个分块下载器") + if started_chunks == 0 and task.filesize > 0: + task.status = DownloadStatus.COMPLETED + self.completed_queue.put((task_id, "completed")) + + def _monitor_progress(self): + logger.info("进度监控线程启动") + while self.running and not self.shutdown_flag: + try: + item = self.progress_queue.get(timeout=1.0) + if not isinstance(item, tuple): + continue + + if item[0] == 'error': + _, task_url, error_msg = item + task_id = None + for tid, t in self.tasks.items(): + if t.url == task_url: + task_id = tid + break + if task_id: + logger.error(f"任务 {task_id} 出错: {error_msg}") + with self.task_lock: + if task_id in self.tasks: + task = self.tasks[task_id] + task.error_msg = error_msg + task.status = DownloadStatus.ERROR + if task_id in self.progress_tasks: + self.progress_display.update(self.progress_tasks[task_id], visible=False) + self.completed_queue.put((task_id, "error")) + else: + task_url, chunk_id, downloaded_bytes = item + task_id = None + task = None + with self.task_lock: + for tid, t in self.tasks.items(): + if t.url == task_url: + task_id = tid + task = t + break + + if task_id and task and task.status == DownloadStatus.DOWNLOADING: + total_downloaded = 0 + for cid, downloader in enumerate(self.active_downloaders.get(task_id, [])): + if cid == chunk_id: + total_downloaded += downloaded_bytes + else: + total_downloaded += downloader.downloaded + + for cid, (start, end) in task.chunks.items(): + if cid not in [d.chunk_id for d in self.active_downloaders.get(task_id, [])]: + total_downloaded += (end - start + 1) + + old_downloaded = task.downloaded + task.downloaded = min(total_downloaded, task.filesize) + + if task.downloaded > old_downloaded: + if task_id in self.speed_tracker: + self.speed_tracker[task_id]['bytes'] += (task.downloaded - old_downloaded) + + if task_id in self.progress_tasks: + if task.filesize > 0: + self.progress_display.update( + self.progress_tasks[task_id], + completed=task.downloaded, + visible=True + ) + else: + self.progress_display.update( + self.progress_tasks[task_id], + total=max(task.downloaded * 2, 1024 * 1024), + completed=task.downloaded, + visible=True + ) + task.last_update = datetime.now() + + if task.filesize > 0 and task.downloaded >= task.filesize: + task.status = DownloadStatus.COMPLETED + logger.info(f"任务 {task_id} 下载完成") + if task_id in self.progress_tasks: + self.progress_display.update(self.progress_tasks[task_id], visible=False) + self.completed_queue.put((task_id, "completed")) + self._on_task_complete(task_id) + + except queue.Empty: + continue + except Exception as e: + logger.exception(f"监控进度时出错: {str(e)}") + logger.info("进度监控线程退出") + + def _calculate_speeds(self): + logger.info("速度计算线程启动") + while self.running and not self.shutdown_flag: + time.sleep(1.0) + try: + current_time = time.time() + with self.task_lock: + for task_id, task in self.tasks.items(): + if task.status != DownloadStatus.DOWNLOADING: + continue + if task_id in self.speed_tracker: + tracker = self.speed_tracker[task_id] + time_diff = current_time - tracker['time'] + if time_diff > 0.5: + bytes_diff = tracker['bytes'] + speed = bytes_diff / time_diff / 1024 + task.current_speed = speed + if task.avg_speed == 0: + task.avg_speed = speed + else: + task.avg_speed = task.avg_speed * 0.8 + speed * 0.2 + task.update_eta() + if task_id in self.progress_tasks: + self.progress_display.update( + self.progress_tasks[task_id], + speed=task.avg_speed * 1024 + ) + tracker['bytes'] = 0 + tracker['time'] = current_time + except Exception as e: + logger.exception(f"计算速度时出错: {str(e)}") + logger.info("速度计算线程退出") + + def _on_task_complete(self, task_id: str): + with self.task_lock: + if task_id in self.tasks: + task = self.tasks[task_id] + if task.filesize > 0 and task.filepath.exists(): + actual_size = task.filepath.stat().st_size + if actual_size != task.filesize: + logger.warning(f"文件大小不匹配: 期望 {task.filesize}, 实际 {actual_size}") + if actual_size < task.filesize: + task.status = DownloadStatus.ERROR + task.error_msg = f"文件不完整: 期望 {task.filesize} 字节, 实际 {actual_size} 字节" + return + + if task_id in self.active_downloaders: + for downloader in self.active_downloaders[task_id]: + downloader.pause() + del self.active_downloaders[task_id] + + self.active_task_count = max(0, self.active_task_count - 1) + if self.active_task_count == 0 and self.live_active: + try: + self.live.stop() + self.live_active = False + except Exception as e: + logger.error(f"停止Live显示失败: {str(e)}") + + def pause_task(self, task_id: str): + with self.task_lock: + if task_id not in self.tasks: + return + task = self.tasks[task_id] + if task.status != DownloadStatus.DOWNLOADING: + return + task.status = DownloadStatus.PAUSED + if task_id in self.active_downloaders: + for downloader in self.active_downloaders[task_id]: + downloader.pause() + if task_id in self.progress_tasks: + self.progress_display.update(self.progress_tasks[task_id], visible=False) + self.active_task_count = max(0, self.active_task_count - 1) + if self.active_task_count == 0 and self.live_active: + try: + self.live.stop() + self.live_active = False + except Exception as e: + logger.error(f"停止Live显示失败: {str(e)}") + + def resume_task(self, task_id: str): + with self.task_lock: + if task_id not in self.tasks: + return + task = self.tasks[task_id] + if task.status != DownloadStatus.PAUSED: + return + if task_id in self.progress_tasks: + self.progress_display.update(self.progress_tasks[task_id], visible=True) + self._start_downloaders(task_id) + task.status = DownloadStatus.DOWNLOADING + task.last_update = datetime.now() + self.active_task_count += 1 + if not self.live_active: + try: + self.live.start() + self.live_active = True + except Exception as e: + logger.error(f"启动Live显示失败: {str(e)}") + if not self.running: + self.running = True + if not self.monitor_thread.is_alive(): + self.monitor_thread = threading.Thread(target=self._monitor_progress, daemon=True) + self.monitor_thread.start() + if not self.speed_calc_thread.is_alive(): + self.speed_calc_thread = threading.Thread(target=self._calculate_speeds, daemon=True) + self.speed_calc_thread.start() + + def cancel_task(self, task_id: str): + with self.task_lock: + if task_id not in self.tasks: + return + task = self.tasks[task_id] + if task.status in [DownloadStatus.COMPLETED, DownloadStatus.CANCELED]: + return + old_status = task.status + task.status = DownloadStatus.CANCELED + if task_id in self.active_downloaders: + for downloader in self.active_downloaders[task_id]: + downloader.pause() + if task_id in self.progress_tasks: + self.progress_display.update(self.progress_tasks[task_id], visible=False) + self.completed_queue.put((task_id, "canceled")) + if old_status == DownloadStatus.DOWNLOADING: + self.active_task_count = max(0, self.active_task_count - 1) + if self.active_task_count == 0 and self.live_active: + try: + self.live.stop() + self.live_active = False + except Exception as e: + logger.error(f"停止Live显示失败: {str(e)}") + + def remove_task(self, task_id: str, delete_file: bool = False): + self.cancel_task(task_id) + with self.task_lock: + if task_id not in self.tasks: + return + task = self.tasks[task_id] + if delete_file and task.filepath.exists(): + try: + os.remove(task.filepath) + except Exception as e: + logger.error(f"删除文件失败: {str(e)}") + if task_id in self.session_pool: + try: + self.session_pool[task_id].close() + except: + pass + del self.session_pool[task_id] + if task_id in self.progress_tasks: + try: + self.progress_display.remove_task(self.progress_tasks[task_id]) + except: + pass + del self.progress_tasks[task_id] + del self.tasks[task_id] + if task_id in self.active_downloaders: + del self.active_downloaders[task_id] + if task_id in self.speed_tracker: + del self.speed_tracker[task_id] + + def get_task_info(self, task_id: str) -> Optional[DownloadTask]: + with self.task_lock: + return self.tasks.get(task_id) + + def list_tasks(self) -> List[Tuple[str, DownloadTask]]: + with self.task_lock: + return list(self.tasks.items()) + + def shutdown(self): + logger.info("开始关闭下载引擎") + self.shutdown_flag = True + self.running = False + with self.task_lock: + task_ids = list(self.tasks.keys()) + for task_id in task_ids: + try: + self.pause_task(task_id) + except Exception as e: + logger.error(f"暂停任务 {task_id} 时出错: {str(e)}") + self.executor.shutdown(wait=False) + for session_id, session in self.session_pool.items(): + try: + session.close() + except: + pass + self.session_pool.clear() + if self.live_active: + try: + self.live.stop() + self.live_active = False + except Exception as e: + logger.error(f"停止Live显示失败: {str(e)}") + logger.info("下载引擎已关闭") + + +# ========== 现代化CLI实现(配色同步更新)========== + +class TaskIdCompleter(Completer): + def __init__(self, get_task_ids: Callable[[], List[str]]): + self.get_task_ids = get_task_ids + + def get_completions(self, document, complete_event): + text = document.text_before_cursor + words = text.split() + if len(words) >= 2: + last_word = words[-1] + for tid in self.get_task_ids(): + if tid.startswith(last_word): + yield Completion(tid, start_position=-len(last_word)) + + +class ModernCLI: + def __init__(self): + self.engine = DownloadEngine(max_workers=200) + # 清新配色主题 + cli_theme = Theme({ + "info": "#87CEEB", + "success": "#98FB98", + "error": "#F08080", + "warning": "#FFD700", + "highlight": "#DDA0DD", + }) + self.console = Console(theme=cli_theme) + self.running = True + self.history_file = Path.home() / ".air2_history" + self.commands: Dict[str, Callable] = { + 'help': self.show_help, 'add': self.add_task, 'start': self.start_task, + 'pause': self.pause_task, 'resume': self.resume_task, 'cancel': self.cancel_task, + 'remove': self.remove_task, 'list': self.list_tasks, 'ls': self.list_tasks, + 'status': self.show_status, 'exit': self.exit, 'quit': self.exit, + 'clear': self.clear, 'history': self.show_history, 'unzip': self.unzip_file, + } + self.base_commands = list(self.commands.keys()) + if PROMPT_TOOLKIT_AVAILABLE: + self._setup_prompt_toolkit() + else: + self._setup_fallback() + signal.signal(signal.SIGINT, self._signal_handler) + signal.signal(signal.SIGTERM, self._signal_handler) + + def _setup_prompt_toolkit(self): + self.pt_style = PtStyle.from_dict({ + 'prompt': 'bold #87CEEB', + 'command': 'bold #98FB98', + 'error': 'bold #F08080', + }) + bindings = KeyBindings() + @bindings.add('c-c') + def _(event): + event.app.exit() + command_completer = WordCompleter(self.base_commands, ignore_case=True) + def get_task_ids(): + return [tid for tid, _ in self.engine.list_tasks()] + task_completer = TaskIdCompleter(get_task_ids) + class DynamicCompleter(Completer): + def get_completions(self, document, complete_event): + text = document.text_before_cursor + words = text.split() + if len(words) == 0: + yield from command_completer.get_completions(document, complete_event) + elif len(words) == 1: + if not text.endswith(' '): + yield from command_completer.get_completions(document, complete_event) + else: + cmd = words[0] + if cmd in ['start', 'pause', 'resume', 'cancel', 'remove', 'unzip']: + yield from task_completer.get_completions(document, complete_event) + else: + cmd = words[0] + if cmd in ['start', 'pause', 'resume', 'cancel', 'remove', 'unzip']: + yield from task_completer.get_completions(document, complete_event) + self.session = PromptSession( + history=FileHistory(str(self.history_file)), + auto_suggest=AutoSuggestFromHistory(), + completer=DynamicCompleter(), + style=self.pt_style, + key_bindings=bindings, + complete_while_typing=True, + ) + self.use_prompt_toolkit = True + + def _setup_fallback(self): + self.use_prompt_toolkit = False + self.console.print("[warning]安装 prompt_toolkit 可获得命令补全和历史记录[/warning]") + + def _signal_handler(self, signum, frame): + signal_name = "SIGINT" if signum == signal.SIGINT else "SIGTERM" + self.console.print(f"\n[warning]收到 {signal_name} 信号,正在关闭下载器...[/warning]") + self.running = False + try: + self.engine.shutdown() + except Exception as e: + logger.exception("关闭引擎时出错") + sys.exit(0) + + def _format_size(self, size: int) -> str: + if size == 0: + return "0 B" + for unit in ['B', 'KB', 'MB', 'GB', 'TB']: + if size < 1024.0: + return f"{size:.2f} {unit}" + size /= 1024.0 + return f"{size:.2f} PB" + + def _format_speed(self, speed: float) -> str: + if speed <= 0: + return "-" + elif speed < 1024: + return f"{speed:.1f} KB/s" + elif speed < 1024 * 1024: + return f"{speed / 1024:.1f} MB/s" + else: + return f"{speed / (1024 * 1024):.1f} GB/s" + + def _format_time(self, td: timedelta) -> str: + total_seconds = int(td.total_seconds()) + if total_seconds <= 0: + return "-" + if total_seconds < 60: + return f"{total_seconds}秒" + elif total_seconds < 3600: + return f"{total_seconds // 60}分钟" + elif total_seconds < 86400: + hours = total_seconds // 3600 + minutes = (total_seconds % 3600) // 60 + return f"{hours}小时{minutes}分" + else: + days = total_seconds // 86400 + hours = (total_seconds % 86400) // 3600 + return f"{days}天{hours}小时" + + def show_help(self, args=None): + help_text = """ +[bold #87CEEB]Air2 现代化命令行帮助[/bold #87CEEB] + +[bold]基本命令:[/bold] + [green]add [选项][/green] 添加下载任务 + [green]start <任务ID>[/green] 开始下载 + [green]pause <任务ID>[/green] 暂停任务 + [green]resume <任务ID>[/green] 恢复任务 + [green]cancel <任务ID>[/green] 取消任务 + [green]remove <任务ID> [--delete][/green] 移除任务(可选删除文件) + [green]unzip <任务ID> [输出目录][/green] 解压已下载的 ZIP 文件 + [green]list / ls[/green] 列出所有任务 + [green]status[/green] 显示系统状态 + [green]clear[/green] 清屏 + [green]history[/green] 显示命令历史 + [green]exit / quit[/green] 退出程序 + +[bold]添加任务选项:[/bold] + [yellow]--threads=<数量>[/yellow] 下载线程数 (默认: 8) + [yellow]--output=<目录>[/yellow] 保存目录 (默认: 当前目录) + [yellow]--chunk-size=<大小>[/yellow] 分块大小,支持 K/M/G (默认: 动态) + [yellow]--speed-limit=[/yellow] 下载限速 (0=不限速) + +[bold]快捷键:[/bold] + [bold]Tab[/bold] 自动补全命令/任务ID + [bold]↑/↓[/bold] 浏览历史命令 + [bold]Ctrl+C[/bold] 中止当前输入/退出程序 + +[bold]示例:[/bold] + add https://example.com/file.zip --threads=16 --output=/downloads + add https://example.com/large.iso --chunk-size=10M --speed-limit=1024 + start a1b2c3d4 + pause a1b2c3d4 + remove a1b2c3d4 --delete + unzip a1b2c3d4 /path/to/extract + """ + self.console.print(Panel(help_text, title="帮助", border_style="#87CEEB", padding=(1, 2))) + + def add_task(self, args): + if not args: + self.console.print("[error]错误: 请提供下载URL[/error]") + return + url = args[0] + kwargs = {'threads': 8, 'output': '.', 'chunk_size': 1024 * 1024, 'speed_limit': 0} + for arg in args[1:]: + if arg.startswith('--'): + if '=' in arg: + key, value = arg[2:].split('=', 1) + key = key.replace('-', '_') + if key == 'chunk_size': + value = value.upper() + if value.endswith('K'): + value = int(float(value[:-1]) * 1024) + elif value.endswith('M'): + value = int(float(value[:-1]) * 1024 * 1024) + elif value.endswith('G'): + value = int(float(value[:-1]) * 1024 * 1024 * 1024) + else: + value = int(value) + elif key in ['threads', 'speed_limit', 'timeout', 'retry_count']: + value = int(value) + kwargs[key] = value + try: + output_dir = kwargs.pop('output', '.') + task_id = self.engine.add_task(url, output_dir, **kwargs) + self.console.print(f"[success]任务添加成功! 任务ID: {task_id}[/success]") + self.console.print(f" 使用 [yellow]start {task_id}[/yellow] 开始下载") + except Exception as e: + self.console.print(f"[error]添加任务失败: {str(e)}[/error]") + + def start_task(self, args): + if not args: + self.console.print("[error]错误: 请提供任务ID[/error]") + return + task_id = args[0] + try: + self.engine.start_task(task_id) + except Exception as e: + self.console.print(f"[error]启动任务失败: {str(e)}[/error]") + + def pause_task(self, args): + if not args: + self.console.print("[error]错误: 请提供任务ID[/error]") + return + task_id = args[0] + try: + self.engine.pause_task(task_id) + self.console.print(f"[warning]任务 {task_id} 已暂停[/warning]") + except Exception as e: + self.console.print(f"[error]暂停任务失败: {str(e)}[/error]") + + def resume_task(self, args): + if not args: + self.console.print("[error]错误: 请提供任务ID[/error]") + return + task_id = args[0] + try: + self.engine.resume_task(task_id) + self.console.print(f"[success]任务 {task_id} 已恢复[/success]") + except Exception as e: + self.console.print(f"[error]恢复任务失败: {str(e)}[/error]") + + def cancel_task(self, args): + if not args: + self.console.print("[error]错误: 请提供任务ID[/error]") + return + task_id = args[0] + try: + self.engine.cancel_task(task_id) + self.console.print(f"[warning]任务 {task_id} 已取消[/warning]") + except Exception as e: + self.console.print(f"[error]取消任务失败: {str(e)}[/error]") + + def remove_task(self, args): + if not args: + self.console.print("[error]错误: 请提供任务ID[/error]") + return + task_id = args[0] + delete_file = '--delete' in args + try: + self.engine.remove_task(task_id, delete_file) + self.console.print(f"[success]任务 {task_id} 已移除" + (" (文件已删除)" if delete_file else "") + "[/success]") + except Exception as e: + self.console.print(f"[error]移除任务失败: {str(e)}[/error]") + + def unzip_file(self, args): + """解压已下载的 ZIP 文件""" + if not args: + self.console.print("[error]错误: 请提供任务ID或ZIP文件路径[/error]") + return + + target = args[0] + output_dir = args[1] if len(args) > 1 else None + + # 判断是任务ID还是直接的文件路径 + zip_path = None + task = None + if target in [tid for tid, _ in self.engine.list_tasks()]: + task = self.engine.get_task_info(target) + if not task: + self.console.print(f"[error]任务 {target} 不存在[/error]") + return + if task.status != DownloadStatus.COMPLETED: + self.console.print(f"[error]任务 {target} 尚未完成,无法解压[/error]") + return + zip_path = task.filepath + else: + # 当作文件路径处理 + zip_path = Path(target) + if not zip_path.exists(): + self.console.print(f"[error]文件不存在: {target}[/error]") + return + + if not zipfile.is_zipfile(zip_path): + self.console.print(f"[error]不是有效的 ZIP 文件: {zip_path}[/error]") + return + + # 确定输出目录 + if output_dir: + extract_to = Path(output_dir) + else: + # 默认解压到与 ZIP 同名的文件夹(去掉 .zip 后缀) + extract_to = zip_path.parent / zip_path.stem + + try: + extract_to.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(zip_path, 'r') as zf: + # 显示文件列表并解压 + file_list = zf.namelist() + self.console.print(f"[info]正在解压 {len(file_list)} 个文件到 {extract_to} ...[/info]") + zf.extractall(extract_to) + self.console.print(f"[success]解压完成! 文件已保存至: {extract_to}[/success]") + except Exception as e: + self.console.print(f"[error]解压失败: {str(e)}[/error]") + + def list_tasks(self, args=None): + tasks = self.engine.list_tasks() + if not tasks: + self.console.print("[warning]没有任务[/warning]") + return + table = Table(title="下载任务", show_lines=True, expand=True, border_style="#87CEEB") + table.add_column("ID", style="#87CEEB", no_wrap=True, width=8) + table.add_column("文件名", style="white", width=30) + table.add_column("大小", style="#98FB98", width=12) + table.add_column("进度", justify="right", width=15) + table.add_column("速度", justify="right", width=12) + table.add_column("ETA", justify="right", width=12) + table.add_column("状态", style="bold", width=10) + table.add_column("选项", justify="left", width=15) + + for task_id, task in tasks: + size_str = self._format_size(task.filesize) if task.filesize > 0 else "未知" + progress_str = f"{task.progress:.1f}%" if task.filesize > 0 else "流式" + speed_str = self._format_speed(task.avg_speed) if task.avg_speed > 0 else "-" + eta_str = self._format_time(task.eta) if task.eta.total_seconds() > 0 else "-" + if task.status == DownloadStatus.COMPLETED: + eta_str = "已完成" + elif task.status in [DownloadStatus.PAUSED, DownloadStatus.CANCELED, DownloadStatus.ERROR]: + eta_str = "-" + + status_color = { + DownloadStatus.WAITING: "dim", + DownloadStatus.DOWNLOADING: "#98FB98", + DownloadStatus.PAUSED: "#FFD700", + DownloadStatus.COMPLETED: "#87CEEB", + DownloadStatus.ERROR: "#F08080", + DownloadStatus.CANCELED: "dark_grey" + }.get(task.status, "white") + status_str = f"[{status_color}]{task.status.value}[/{status_color}]" + + options_str = f"线程:{task.threads}" + if task.speed_limit > 0: + options_str += f", 限速:{task.speed_limit}KB/s" + + table.add_row(task_id, str(task.filepath.name)[:30], size_str, progress_str, speed_str, eta_str, status_str, options_str) + self.console.print(table) + + def show_status(self, args=None): + tasks = self.engine.list_tasks() + downloading = sum(1 for _, t in tasks if t.status == DownloadStatus.DOWNLOADING) + paused = sum(1 for _, t in tasks if t.status == DownloadStatus.PAUSED) + completed = sum(1 for _, t in tasks if t.status == DownloadStatus.COMPLETED) + error = sum(1 for _, t in tasks if t.status == DownloadStatus.ERROR) + canceled = sum(1 for _, t in tasks if t.status == DownloadStatus.CANCELED) + total_size = sum(t.filesize for _, t in tasks) + total_downloaded = sum(t.downloaded for _, t in tasks) + total_progress = (total_downloaded / total_size * 100) if total_size > 0 else 0 + + status_text = f""" +[bold #87CEEB]系统状态:[/bold #87CEEB] + +[bold]任务统计:[/bold] + 总计: {len(tasks)} 个任务 + 下载中: [#98FB98]{downloading}[/#98FB98] + 已暂停: [#FFD700]{paused}[/#FFD700] + 已完成: [#87CEEB]{completed}[/#87CEEB] + 错误: [#F08080]{error}[/#F08080] + 已取消: [dark_grey]{canceled}[/dark_grey] + +[bold]数据统计:[/bold] + 总大小: {self._format_size(total_size) if total_size > 0 else "未知"} + 已下载: {self._format_size(total_downloaded)} + 总进度: {total_progress:.1f}% if total_size > 0 else "计算中" + +[bold]性能信息:[/bold] + 活动任务: {self.engine.active_task_count} + 最大线程: 200 + 工作线程: {threading.active_count()} + """ + self.console.print(Panel(status_text, title="系统状态", border_style="#87CEEB", padding=(1, 2))) + + def show_history(self, args=None): + if self.use_prompt_toolkit and hasattr(self.session, 'history'): + history_entries = list(self.session.history.get_strings()) + if history_entries: + self.console.print("[bold]最近命令历史:[/bold]") + for i, entry in enumerate(history_entries[-20:], 1): + self.console.print(f" {i:3d}. {entry}") + else: + self.console.print("[warning]暂无命令历史[/warning]") + else: + self.console.print("[warning]命令历史功能需要 prompt_toolkit[/warning]") + + def clear(self, args=None): + self.console.clear() + + def exit(self, args=None): + confirm = Confirm.ask("确定要退出吗? 活动下载将会暂停") + if confirm: + self.console.print("[warning]正在关闭下载引擎...[/warning]") + try: + self.engine.shutdown() + except Exception as e: + logger.exception("关闭引擎时出错") + self.running = False + self.console.print("[success]再见![/success]") + sys.exit(0) + + def process_completed_tasks(self): + while not self.engine.completed_queue.empty(): + try: + task_id, status = self.engine.completed_queue.get_nowait() + task = self.engine.get_task_info(task_id) + if task: + if status == "completed": + self.console.print(f"\n[success]任务 {task_id} 完成! 文件: {task.filepath}[/success]") + elif status == "error": + self.console.print(f"\n[error]任务 {task_id} 失败: {task.error_msg}[/error]") + elif status == "canceled": + self.console.print(f"\n[warning]任务 {task_id} 已取消[/warning]") + except queue.Empty: + break + + def run(self): + self.console.print(Panel.fit("[bold #87CEEB]Air2[/bold #87CEEB] - 高性能下载引擎", subtitle="输入 'help' 查看帮助", border_style="#87CEEB")) + if not PROMPT_TOOLKIT_AVAILABLE: + self.console.print("[warning]安装 prompt_toolkit 可获得更好的交互体验[/warning]") + self.console.print("pip install prompt_toolkit\n") + while self.running: + try: + self.process_completed_tasks() + if self.use_prompt_toolkit: + try: + user_input = self.session.prompt(">>> ") + except KeyboardInterrupt: + continue + except EOFError: + self.exit() + break + else: + try: + user_input = input("\n>>> ").strip() + except EOFError: + self.exit() + break + except KeyboardInterrupt: + self.console.print("\n[warning]输入 'exit' 退出程序[/warning]") + continue + if not user_input: + continue + parts = user_input.strip().split() + command = parts[0].lower() + args = parts[1:] if len(parts) > 1 else [] + if command in self.commands: + self.commands[command](args) + else: + self.console.print(f"[error]未知命令: {command}[/error]") + self.console.print("输入 'help' 查看可用命令") + except KeyboardInterrupt: + self.console.print("\n[warning]输入 'exit' 退出程序[/warning]") + except Exception as e: + self.console.print(f"[error]发生错误: {str(e)}[/error]") + logger.exception("主循环错误") + try: + self.engine.shutdown() + except: + pass + + +def execute_command_line(): + """命令行直接执行模式,支持 air2 add URL 等用法""" + parser = argparse.ArgumentParser(description="Air2 - 高性能下载引擎", add_help=False) + parser.add_argument("command", nargs="?", help="命令: add, start, pause, resume, cancel, remove, list, status, unzip") + parser.add_argument("args", nargs="*", help="命令参数") + parser.add_argument("--help", action="store_true", help="显示帮助") + + # 为了兼容原有风格,我们手动解析简单的命令 + if len(sys.argv) == 1: + # 无参数,进入交互模式 + return None + + # 处理 help + if sys.argv[1] in ["help", "--help", "-h"]: + print("Air2 命令行用法:") + print(" air2 add [--threads=N] [--output=DIR] [--chunk-size=SIZE] [--speed-limit=KB]") + print(" air2 start <任务ID>") + print(" air2 pause <任务ID>") + print(" air2 resume <任务ID>") + print(" air2 cancel <任务ID>") + print(" air2 remove <任务ID> [--delete]") + print(" air2 unzip <任务ID或ZIP文件> [输出目录]") + print(" air2 list") + print(" air2 status") + print(" air2 help") + sys.exit(0) + + command = sys.argv[1].lower() + args = sys.argv[2:] + + # 创建引擎 + engine = DownloadEngine(max_workers=200) + console = Console() + + # 辅助函数 + def format_size(size): + if size == 0: return "0 B" + for unit in ['B', 'KB', 'MB', 'GB', 'TB']: + if size < 1024.0: + return f"{size:.2f} {unit}" + size /= 1024.0 + return f"{size:.2f} PB" + + def format_speed(speed): + if speed <= 0: return "-" + elif speed < 1024: return f"{speed:.1f} KB/s" + elif speed < 1024*1024: return f"{speed/1024:.1f} MB/s" + else: return f"{speed/(1024*1024):.1f} GB/s" + + def wait_for_task(task_id, console): + """等待任务完成(用于命令行模式)""" + from rich.live import Live + from rich.progress import Progress, BarColumn, DownloadColumn, TransferSpeedColumn, TimeRemainingColumn, TextColumn + + progress = Progress( + TextColumn("[bold #87CEEB]{task.description}"), + BarColumn(complete_style="#98FB98", finished_style="#98FB98"), + "[progress.percentage]{task.percentage:>3.1f}%", + "•", DownloadColumn(), "•", TransferSpeedColumn(), "•", TimeRemainingColumn(), + console=console, refresh_per_second=4 + ) + live = Live(progress, console=console, auto_refresh=True) + task_obj = engine.get_task_info(task_id) + if not task_obj: + console.print(f"[error]任务 {task_id} 不存在") + return False + if task_obj.filesize > 0: + task_id_progress = progress.add_task(f"{task_obj.filepath.name[:30]}", total=task_obj.filesize, completed=task_obj.downloaded) + else: + task_id_progress = progress.add_task(f"{task_obj.filepath.name[:30]}", total=0, completed=task_obj.downloaded) + live.start() + try: + while True: + task = engine.get_task_info(task_id) + if not task: + console.print(f"[error]任务 {task_id} 已消失") + live.stop() + return False + if task.status == DownloadStatus.COMPLETED: + progress.update(task_id_progress, completed=task.filesize) + live.stop() + console.print(f"[success]下载完成: {task.filepath}[/success]") + return True + elif task.status == DownloadStatus.ERROR: + live.stop() + console.print(f"[error]下载失败: {task.error_msg}[/error]") + return False + elif task.status == DownloadStatus.CANCELED: + live.stop() + console.print(f"[warning]下载已取消[/warning]") + return False + else: + if task.filesize > 0: + progress.update(task_id_progress, completed=task.downloaded, total=task.filesize) + else: + progress.update(task_id_progress, completed=task.downloaded, total=max(task.downloaded*2, 1024*1024)) + # 更新速度 + if task.avg_speed > 0: + progress.update(task_id_progress, speed=task.avg_speed*1024) + time.sleep(0.5) + except KeyboardInterrupt: + live.stop() + console.print("\n[warning]用户中断,正在取消任务...[/warning]") + engine.cancel_task(task_id) + return False + finally: + live.stop() + + # 执行命令 + try: + if command == "add": + if not args: + console.print("[error]请提供 URL") + sys.exit(1) + url = args[0] + kwargs = {'threads': 8, 'output': '.', 'chunk_size': 1024*1024, 'speed_limit': 0} + # 解析选项 + for arg in args[1:]: + if arg.startswith('--'): + if '=' in arg: + key, value = arg[2:].split('=', 1) + key = key.replace('-', '_') + if key == 'chunk_size': + value = value.upper() + if value.endswith('K'): value = int(float(value[:-1]) * 1024) + elif value.endswith('M'): value = int(float(value[:-1]) * 1024*1024) + elif value.endswith('G'): value = int(float(value[:-1]) * 1024*1024*1024) + else: value = int(value) + elif key in ['threads', 'speed_limit']: + value = int(value) + kwargs[key] = value + output_dir = kwargs.pop('output', '.') + task_id = engine.add_task(url, output_dir, **kwargs) + console.print(f"[success]任务添加成功: {task_id}[/success]") + # 可选:自动开始下载(默认自动开始) + console.print(f"开始下载 {task_id} ...") + engine.start_task(task_id) + # 等待下载完成 + success = wait_for_task(task_id, console) + sys.exit(0 if success else 1) + + elif command == "start": + if not args: + console.print("[error]请提供任务ID") + sys.exit(1) + task_id = args[0] + engine.start_task(task_id) + console.print(f"等待任务 {task_id} 完成...") + success = wait_for_task(task_id, console) + sys.exit(0 if success else 1) + + elif command == "pause": + if not args: + console.print("[error]请提供任务ID") + sys.exit(1) + task_id = args[0] + engine.pause_task(task_id) + console.print(f"[warning]任务 {task_id} 已暂停[/warning]") + sys.exit(0) + + elif command == "resume": + if not args: + console.print("[error]请提供任务ID") + sys.exit(1) + task_id = args[0] + engine.resume_task(task_id) + console.print(f"[success]任务 {task_id} 已恢复,等待完成...[/success]") + success = wait_for_task(task_id, console) + sys.exit(0 if success else 1) + + elif command == "cancel": + if not args: + console.print("[error]请提供任务ID") + sys.exit(1) + task_id = args[0] + engine.cancel_task(task_id) + console.print(f"[warning]任务 {task_id} 已取消[/warning]") + sys.exit(0) + + elif command == "remove": + if not args: + console.print("[error]请提供任务ID") + sys.exit(1) + task_id = args[0] + delete_file = '--delete' in args + engine.remove_task(task_id, delete_file) + console.print(f"[success]任务 {task_id} 已移除" + (" (文件已删除)" if delete_file else "")) + sys.exit(0) + + elif command == "unzip": + if not args: + console.print("[error]请提供任务ID或ZIP文件路径") + sys.exit(1) + target = args[0] + output_dir = args[1] if len(args) > 1 else None + zip_path = None + task = None + # 判断是否为任务ID + if target in [tid for tid, _ in engine.list_tasks()]: + task = engine.get_task_info(target) + if not task: + console.print(f"[error]任务 {target} 不存在") + sys.exit(1) + if task.status != DownloadStatus.COMPLETED: + console.print(f"[error]任务 {target} 尚未完成,无法解压") + sys.exit(1) + zip_path = task.filepath + else: + zip_path = Path(target) + if not zip_path.exists(): + console.print(f"[error]文件不存在: {target}") + sys.exit(1) + if not zipfile.is_zipfile(zip_path): + console.print(f"[error]不是有效的 ZIP 文件: {zip_path}") + sys.exit(1) + if output_dir: + extract_to = Path(output_dir) + else: + extract_to = zip_path.parent / zip_path.stem + try: + extract_to.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(zip_path, 'r') as zf: + console.print(f"正在解压 {len(zf.namelist())} 个文件到 {extract_to} ...") + zf.extractall(extract_to) + console.print(f"[success]解压完成! 保存至: {extract_to}") + sys.exit(0) + except Exception as e: + console.print(f"[error]解压失败: {str(e)}") + sys.exit(1) + + elif command == "list" or command == "ls": + tasks = engine.list_tasks() + if not tasks: + console.print("[warning]没有任务") + else: + table = Table(title="下载任务", show_lines=True, border_style="#87CEEB") + table.add_column("ID", style="#87CEEB", no_wrap=True) + table.add_column("文件名", width=30) + table.add_column("大小", style="#98FB98") + table.add_column("进度") + table.add_column("状态") + for tid, t in tasks: + size_str = format_size(t.filesize) if t.filesize > 0 else "未知" + prog_str = f"{t.progress:.1f}%" if t.filesize > 0 else "流式" + status_color = { + DownloadStatus.WAITING: "dim", DownloadStatus.DOWNLOADING: "#98FB98", + DownloadStatus.PAUSED: "#FFD700", DownloadStatus.COMPLETED: "#87CEEB", + DownloadStatus.ERROR: "#F08080", DownloadStatus.CANCELED: "dark_grey" + }.get(t.status, "white") + status_str = f"[{status_color}]{t.status.value}[/{status_color}]" + table.add_row(tid, str(t.filepath.name)[:30], size_str, prog_str, status_str) + console.print(table) + sys.exit(0) + + elif command == "status": + tasks = engine.list_tasks() + downloading = sum(1 for _, t in tasks if t.status == DownloadStatus.DOWNLOADING) + paused = sum(1 for _, t in tasks if t.status == DownloadStatus.PAUSED) + completed = sum(1 for _, t in tasks if t.status == DownloadStatus.COMPLETED) + error = sum(1 for _, t in tasks if t.status == DownloadStatus.ERROR) + canceled = sum(1 for _, t in tasks if t.status == DownloadStatus.CANCELED) + total_size = sum(t.filesize for _, t in tasks) + total_downloaded = sum(t.downloaded for _, t in tasks) + console.print(f"任务总数: {len(tasks)} | 下载中: {downloading} | 暂停: {paused} | 完成: {completed} | 错误: {error} | 取消: {canceled}") + if total_size > 0: + console.print(f"总大小: {format_size(total_size)} | 已下载: {format_size(total_downloaded)} | 进度: {total_downloaded/total_size*100:.1f}%") + else: + console.print("总大小: 未知") + sys.exit(0) + + else: + console.print(f"[error]未知命令: {command}") + console.print("输入 air2 help 查看帮助") + sys.exit(1) + + except Exception as e: + console.print(f"[error]执行命令时出错: {str(e)}") + logger.exception("命令行模式错误") + sys.exit(1) + finally: + engine.shutdown() + + +def main(): + # 判断是否进入命令行模式 + if len(sys.argv) > 1 and sys.argv[1] not in ['help', '--help', '-h']: + execute_command_line() + else: + # 交互模式 + try: + cli = ModernCLI() + cli.run() + except Exception as e: + console = Console() + console.print(f"[error]程序启动失败: {str(e)}[/error]") + logger.exception("程序启动失败") + sys.exit(1) + + +if __name__ == "__main__": + main() \ No newline at end of file