1640 lines
70 KiB
Python
1640 lines
70 KiB
Python
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 <URL> [选项][/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=<KB/s>[/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 <URL> [--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() |