iShotaBot/defs/bilibili_download.py

import contextlib
import os
import time
from asyncio import create_subprocess_shell, subprocess, Lock
from io import BytesIO
from typing import Tuple, Dict, Union, Optional

import aiofiles
from bilibili_api import HEADERS
from bilibili_api.video import Video, VideoDownloadURLDataDetecter, VideoQuality
from httpx import AsyncClient, Response
from pyrogram.types import Message

from defs.request import cache_dir
from init import bot, logger, request

FFMPEG_PATH = "ffmpeg"
FFPROBE_PATH = "ffprobe"
LOCK = Lock()
EDIT_TEMP_SECONDS = 10.0
MESSAGE_MAP: Dict[int, float] = {}
UPLOAD_MESSAGE_MAP: Dict[int, int] = {}


class BilibiliDownloaderError(Exception):
    """Bilibili 下载器错误"""

    MSG = "Bilibili 下载器错误"

    def __init__(self, msg: str = None):
        self.MSG = msg or self.MSG


class FileTooBig(BilibiliDownloaderError):
    """文件过大，超过2GB"""

    MSG = "文件过大，超过2GB"


class FileNoSize(BilibiliDownloaderError):
    """文件大小未知"""

    MSG = "文件大小未知"


class FFmpegError(BilibiliDownloaderError):
    """FFmpeg 转换失败"""

    MSG = "FFmpeg 转换失败"


def should_edit(m: Message) -> bool:
    if m.id in MESSAGE_MAP:
        last_time = MESSAGE_MAP[m.id]
        if last_time + EDIT_TEMP_SECONDS < time.time():
            return True
    else:
        return True
    return False


def format_bytes(size: Union[int, float]) -> str:
    """格式化文件大小"""
    power = 1024
    n = 0
    power_labels = {0: "", 1: "K", 2: "M", 3: "G", 4: "T"}
    while size > power:
        size /= power
        n += 1
    if n > 4:
        n = 4
    return f"{round(size, 2)} {power_labels[n]}B"


def format_seconds(seconds: Union[int, float]) -> str:
    """格式化秒数"""
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    s = round(s, 2)
    text = ""
    if h > 0:
        text += f" {h} 小时"
    if m > 0:
        text += f" {m} 分钟"
    if s > 0:
        text += f" {s} 秒"
    return text.strip()


async def safe_edit(m: Message, text: str):
    try:
        await m.edit_text(text)
    except Exception:
        pass


async def fail_edit(m: Message, text: str):
    try:
        await m.edit_text(text)
    except Exception:
        with contextlib.suppress(Exception):
            await m.reply(text, quote=True)


async def execute(command: str) -> Tuple[str, int]:
    process = await create_subprocess_shell(
        command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE
    )
    stdout, stderr = await process.communicate()
    try:
        result = str(stdout.decode().strip()) + str(stderr.decode().strip())
    except UnicodeDecodeError:
        result = str(stdout.decode("gbk").strip()) + str(stderr.decode("gbk").strip())
    return result, process.returncode


def safe_remove(path: str):
    if os.path.exists(path):
        os.remove(path)


async def message_edit(
    length: int,
    total_downloaded: int,
    temp_downloaded: int,
    chunk_time: float,
    m: Message,
    t: str,
):
    speed = temp_downloaded / (chunk_time if chunk_time > 0 else 1)
    text = (
        f"{t}进度\n\n"
        f"{format_bytes(total_downloaded)} / {format_bytes(length)} "
        f"({round(total_downloaded / length * 100.0, 2)}%)\n\n"
        f"传输区间速度：{format_bytes(speed)}/s\n"
        f"预计剩余时间：{format_seconds((length - total_downloaded) / speed)}"
    )
    await safe_edit(m, text)


async def download_url(url: str, out: str, m: Message, start: str):
    async with AsyncClient(headers=HEADERS, timeout=60) as sess:
        async with sess.stream("GET", url) as resp:
            logger.info(f"Downloading {start}")
            resp: Response
            length = resp.headers.get("content-length")
            if not length:
                raise FileNoSize
            length = int(length)
            if length > 1.9 * 1024 * 1024 * 1024:
                raise FileTooBig
            total_downloaded = 0
            temp_downloaded = 0
            MESSAGE_MAP[m.id] = time.time() - EDIT_TEMP_SECONDS
            async with aiofiles.open(out, "wb") as f:
                async for chunk in resp.aiter_bytes(1024):
                    if not chunk:
                        break
                    chunk_len = len(chunk)
                    total_downloaded += chunk_len
                    temp_downloaded += chunk_len
                    async with LOCK:
                        _should_edit = should_edit(m)
                        if _should_edit:
                            now = time.time()
                            chunk_time = now - MESSAGE_MAP[m.id]
                            MESSAGE_MAP[m.id] = now
                    if _should_edit:
                        bot.loop.create_task(
                            message_edit(
                                length,
                                total_downloaded,
                                temp_downloaded,
                                chunk_time,
                                m,
                                f"{start}下载",
                            )
                        )
                        temp_downloaded = 0
                    await f.write(chunk)


async def get_video_duration(path: str) -> float:
    """获取视频时长"""
    video_duration, code = await execute(
        f"{FFPROBE_PATH} -v error -select_streams v:0 -show_entries format=duration "
        f"-of default=noprint_wrappers=1:nokey=1 {path}"
    )
    if code != 0:
        raise FFmpegError("视频时长获取失败")
    return round(float(video_duration.split("[")[0].strip()), 2)


async def get_video_height_width(path: str) -> Tuple[int, int]:
    """获取视频高度和宽度"""
    result, code = await execute(
        f"{FFPROBE_PATH} -v error -select_streams v:0 -show_entries stream=width,height "
        f"-of csv=s=x:p=0 {path}"
    )
    if code != 0:
        raise FFmpegError("视频宽高度获取失败")
    video_width, video_height = result.split("[")[0].split("x")
    return int(video_height), int(video_width)


async def take_screenshot(v: Video) -> Optional[BytesIO]:
    """获取视频封面"""
    try:
        info = await v.get_info()
        pic_get = (await request.get(info["pic"])).content
        pic = BytesIO(pic_get)
        pic.name = "screenshot.jpg"
        return pic
    except Exception:
        logger.exception("获取视频封面失败")
        return None


async def go_download(v: Video, p_num: int, m: Message):
    video_path = cache_dir / f"{v.get_aid()}_{p_num}.mp4"
    safe_remove(video_path)
    flv_temp_path = cache_dir / f"{v.get_aid()}_{p_num}_temp.flv"
    video_temp_path = cache_dir / f"{v.get_aid()}_{p_num}_video.m4s"
    audio_temp_path = cache_dir / f"{v.get_aid()}_{p_num}_audio.m4s"
    # 有 MP4 流 / FLV 流两种可能
    try:
        # 获取视频下载链接
        download_url_data = await v.get_download_url(p_num)
        # 解析视频下载信息
        detector = VideoDownloadURLDataDetecter(data=download_url_data)
        streams = detector.detect_best_streams(
            video_max_quality=VideoQuality._1080P_60,  # noqa
        )
        if not streams:
            raise BilibiliDownloaderError("无法获取下载链接")
        if detector.check_flv_stream():
            # FLV 流下载
            await download_url(streams[0].url, flv_temp_path, m, "视频 FLV ")
            # 转换文件格式
            _, result = await execute(f"{FFMPEG_PATH} -i {flv_temp_path} {video_path}")
        else:
            if len(streams) < 2:
                raise BilibiliDownloaderError("获取下载链接异常")
            # MP4 流下载
            await download_url(streams[0].url, video_temp_path, m, "视频 m4s ")
            await download_url(streams[1].url, audio_temp_path, m, "音频 m4s ")
            # 混流
            logger.info("Merging video and audio")
            _, result = await execute(
                f'{FFMPEG_PATH} -i "{video_temp_path}" -i "{audio_temp_path}" '
                f"-c:v copy -c:a copy -movflags +faststart "
                f'-y "{video_path}"'
            )
        if result != 0:
            raise FFmpegError
        bot.loop.create_task(go_upload(v, p_num, m))
    except BilibiliDownloaderError as e:
        await fail_edit(m, e.MSG)
    except Exception as e:
        logger.exception("Downloading video failed")
        await fail_edit(m, f"下载失败：{e}")
    finally:
        # 删除临时文件
        safe_remove(flv_temp_path)
        safe_remove(video_temp_path)
        safe_remove(audio_temp_path)


async def go_upload_progress(current: int, total: int, m: Message):
    if current == 0:
        return
    async with LOCK:
        _should_edit = should_edit(m)
        if _should_edit:
            now = time.time()
            chunk_time = now - MESSAGE_MAP[m.id]
            MESSAGE_MAP[m.id] = now
    if _should_edit:
        t = UPLOAD_MESSAGE_MAP[m.id] if m.id in UPLOAD_MESSAGE_MAP else 0
        UPLOAD_MESSAGE_MAP[m.id] = current
        chunk = current - t
        chunk = chunk if chunk > 0 else 0
        await message_edit(total, current, chunk, chunk_time, m, "上传")


async def go_upload(v: Video, p_num: int, m: Message):
    video_path = cache_dir / f"{v.get_aid()}_{p_num}.mp4"
    if not video_path.exists():
        await fail_edit(m, "视频文件不存在")
        return
    try:
        video_duration = await get_video_duration(video_path)
        video_height, video_width = await get_video_height_width(video_path)
        video_jpg = await take_screenshot(v)
        logger.info(f"Uploading {video_path}")
        await bot.send_video(
            chat_id=m.chat.id,
            video=str(video_path),
            duration=int(video_duration),
            width=video_width,
            height=video_height,
            thumb=video_jpg,
            supports_streaming=True,
            progress=go_upload_progress,
            progress_args=(m,),
        )
        logger.info(f"Upload {video_path} success")
    except BilibiliDownloaderError as e:
        await fail_edit(m, e.MSG)
        return
    except Exception as e:
        logger.exception("Uploading video failed")
        await fail_edit(m, f"上传失败：{e}")
        return
    finally:
        safe_remove(video_path)
        if m.id in MESSAGE_MAP:
            del MESSAGE_MAP[m.id]
        if m.id in UPLOAD_MESSAGE_MAP:
            del UPLOAD_MESSAGE_MAP[m.id]
    with contextlib.suppress(Exception):
        await m.delete()