AstrBot/astrbot/core/utils/io.py

import asyncio
import base64
import inspect
import logging
import os
import re
import shutil
import socket
import ssl
import time
import uuid
import zipfile
from ipaddress import IPv4Address, IPv6Address, ip_address
from pathlib import Path

import aiohttp
import anyio
import certifi
import psutil
from PIL import Image

from .astrbot_path import get_astrbot_data_path, get_astrbot_path, get_astrbot_temp_path
from .version_comparator import VersionComparator

logger = logging.getLogger("astrbot")


def _get_aiohttp():
    import aiohttp

    return aiohttp


def on_error(func, path, exc_info) -> None:
    """A callback of the rmtree function."""
    import stat

    if not os.access(path, os.W_OK):
        os.chmod(path, stat.S_IWUSR)
        func(path)
    else:
        raise exc_info[1]


def remove_dir(file_path: str) -> bool:
    if not os.path.lexists(file_path):
        return True
    if os.path.isfile(file_path) or os.path.islink(file_path):
        os.remove(file_path)
    else:
        shutil.rmtree(file_path, onerror=on_error)
    return True


def ensure_dir(dir_path: str | Path) -> None:
    """确保目录存在。如果路径处存在非目录的文件或损坏的符号链接，则先将其删除。"""
    p = Path(dir_path)
    if (p.exists() or p.is_symlink()) and not p.is_dir():
        logger.warning(f"路径 {p} 已存在但不是目录，正在清理以创建目录。")
        try:
            if p.is_dir():
                shutil.rmtree(p, onerror=on_error)
            else:
                p.unlink()
        except Exception as e:
            logger.error(f"清理冲突路径 {p} 失败: {e!s}")
            raise RuntimeError(f"无法清理冲突路径 {p}：{e!s}") from e

    try:
        p.mkdir(parents=True, exist_ok=True)
    except Exception as e:
        logger.error(f"创建目录 {p} 失败: {e!s}")
        raise RuntimeError(f"无法创建目录 {p}：{e!s}") from e


def port_checker(port: int, host: str = "localhost") -> bool:
    sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sk.settimeout(1)
    try:
        sk.connect((host, port))
        sk.close()
        return True
    except Exception:
        sk.close()
        return False


def save_temp_img(img: Image.Image | bytes) -> str:
    temp_dir = get_astrbot_temp_path()
    # 获得时间戳
    timestamp = f"{int(time.time())}_{uuid.uuid4().hex[:8]}"
    p = os.path.join(temp_dir, f"io_temp_img_{timestamp}.jpg")

    if isinstance(img, Image.Image):
        img.save(p)
    else:
        with open(p, "wb") as f:
            f.write(img)
    return p


async def download_image_by_url(
    url: str,
    post: bool = False,
    post_data: dict | None = None,
    path: str | None = None,
) -> str:
    """下载图片, 返回 path"""
    aiohttp = _get_aiohttp()
    try:
        ssl_context = ssl.create_default_context(
            cafile=certifi.where(),
        )  # 使用 certifi 提供的 CA 证书
        connector = aiohttp.TCPConnector(ssl=ssl_context)  # 使用 certifi 的根证书
        async with aiohttp.ClientSession(
            trust_env=True,
            connector=connector,
        ) as session:
            if post:
                async with session.post(url, json=post_data) as resp:
                    if not path:
                        return save_temp_img(await resp.read())
                    async with await anyio.open_file(path, "wb") as f:
                        await f.write(await resp.read())
                    return path
            else:
                async with session.get(url) as resp:
                    if not path:
                        return save_temp_img(await resp.read())
                    async with await anyio.open_file(path, "wb") as f:
                        await f.write(await resp.read())
                    return path
    except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError):
        # 关闭SSL验证(仅在证书验证失败时作为fallback)
        logger.warning(
            f"SSL certificate verification failed for {url}. "
            "Disabling SSL verification (CERT_NONE) as a fallback. "
            "This is insecure and exposes the application to man-in-the-middle attacks. "
            "Please investigate and resolve certificate issues.",
        )
        ssl_context = ssl.create_default_context()
        ssl_context.check_hostname = False
        ssl_context.verify_mode = ssl.CERT_NONE
        async with aiohttp.ClientSession() as session:
            if post:
                async with session.post(url, json=post_data, ssl=ssl_context) as resp:
                    if not path:
                        return save_temp_img(await resp.read())
                    async with await anyio.open_file(path, "wb") as f:
                        await f.write(await resp.read())
                    return path
            else:
                async with session.get(url, ssl=ssl_context) as resp:
                    if not path:
                        return save_temp_img(await resp.read())
                    async with await anyio.open_file(path, "wb") as f:
                        await f.write(await resp.read())
                    return path
    except Exception as e:
        raise e


async def _emit_download_progress(progress_callback, payload: dict) -> None:
    if not progress_callback:
        return
    result = progress_callback(payload)
    if inspect.isawaitable(result):
        await result


class DownloadFileHTTPError(RuntimeError):
    """Raised when a file download returns an unsuccessful HTTP status."""


def _raise_for_download_status(resp, url: str) -> None:
    if resp.status == 200:
        return
    logger.error(
        "Failed to download file from %s. HTTP status code: %s",
        url,
        resp.status,
    )
    raise DownloadFileHTTPError(
        f"Failed to download file from {url}. HTTP status code: {resp.status}"
    )


async def download_file(
    url: str,
    path: str,
    show_progress: bool = False,
    progress_callback=None,
) -> None:
    """从指定 url 下载文件到指定路径 path"""
    aiohttp = _get_aiohttp()
    try:
        ssl_context = ssl.create_default_context(
            cafile=certifi.where(),
        )  # 使用 certifi 提供的 CA 证书
        connector = aiohttp.TCPConnector(ssl=ssl_context)
        async with aiohttp.ClientSession(
            trust_env=True,
            connector=connector,
        ) as session:
            async with session.get(
                url,
                timeout=aiohttp.ClientTimeout(total=1800),
            ) as resp:
                _raise_for_download_status(resp, url)
                total_size = int(resp.headers.get("content-length", 0))
                downloaded_size = 0
                start_time = time.time()
                if show_progress:
                    logger.info(
                        f"Downloading: {url} | Size: {total_size / 1024:.2f} KB",
                    )
                await _emit_download_progress(
                    progress_callback,
                    {
                        "url": url,
                        "downloaded": 0,
                        "total": total_size,
                        "percent": 0,
                        "speed": 0,
                    },
                )
                async with await anyio.open_file(path, "wb") as f:
                    while True:
                        chunk = await resp.content.read(8192)
                        if not chunk:
                            break
                        await f.write(chunk)
                        downloaded_size += len(chunk)
                        elapsed_time = (
                            time.time() - start_time
                            if time.time() - start_time > 0
                            else 1
                        )
                        speed = downloaded_size / 1024 / elapsed_time  # KB/s
                        percent = downloaded_size / total_size if total_size > 0 else 0
                        await _emit_download_progress(
                            progress_callback,
                            {
                                "url": url,
                                "downloaded": downloaded_size,
                                "total": total_size,
                                "percent": percent,
                                "speed": speed,
                            },
                        )
                        if show_progress:
                            logger.info(
                                f"Progress: {percent:.2%} Speed: {speed:.2f} KB/s",
                            )
                await _emit_download_progress(
                    progress_callback,
                    {
                        "url": url,
                        "downloaded": downloaded_size,
                        "total": total_size,
                        "percent": 1,
                        "speed": 0,
                    },
                )
    except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError):
        # 关闭SSL验证(仅在证书验证失败时作为fallback)
        logger.warning(
            f"SSL certificate verification failed for {url}. "
            "Falling back to unverified connection (CERT_NONE). ",
        )
        logger.warning(
            f"SSL certificate verification failed for {url}. "
            "Falling back to unverified connection (CERT_NONE). "
            "This is insecure and exposes the application to man-in-the-middle attacks. "
            "Please investigate certificate issues with the remote server.",
        )
        ssl_context = ssl.create_default_context()
        ssl_context.check_hostname = False
        ssl_context.verify_mode = ssl.CERT_NONE
        async with (
            aiohttp.ClientSession() as session,
            session.get(
                url,
                ssl=ssl_context,
                timeout=aiohttp.ClientTimeout(total=120),
            ) as resp,
        ):
            _raise_for_download_status(resp, url)
            total_size = int(resp.headers.get("content-length", 0))
            downloaded_size = 0
            start_time = time.time()
            if show_progress:
                logger.info(f"Size: {total_size / 1024:.2f} KB | URL: {url}")
            await _emit_download_progress(
                progress_callback,
                {
                    "url": url,
                    "downloaded": 0,
                    "total": total_size,
                    "percent": 0,
                    "speed": 0,
                },
            )
            async with await anyio.open_file(path, "wb") as f:
                while True:
                    chunk = await resp.content.read(8192)
                    if not chunk:
                        break
                    await f.write(chunk)
                    downloaded_size += len(chunk)
                    elapsed_time = (
                        time.time() - start_time if time.time() - start_time > 0 else 1
                    )
                    speed = downloaded_size / 1024 / elapsed_time  # KB/s
                    percent = downloaded_size / total_size if total_size > 0 else 0
                    await _emit_download_progress(
                        progress_callback,
                        {
                            "url": url,
                            "downloaded": downloaded_size,
                            "total": total_size,
                            "percent": percent,
                            "speed": speed,
                        },
                    )
                    if show_progress:
                        logger.info(
                            f"Progress: {percent:.2%} Speed: {speed:.2f} KB/s",
                        )
            await _emit_download_progress(
                progress_callback,
                {
                    "url": url,
                    "downloaded": downloaded_size,
                    "total": total_size,
                    "percent": 1,
                    "speed": 0,
                },
            )
    if show_progress:
        logger.info("下载完成")


def file_to_base64(file_path: str) -> str:
    with open(file_path, "rb") as f:
        data_bytes = f.read()
        base64_str = base64.b64encode(data_bytes).decode()
    return "base64://" + base64_str


def get_local_ip_addresses() -> list[IPv4Address | IPv6Address]:
    net_interfaces = psutil.net_if_addrs()
    network_ips: list[IPv4Address | IPv6Address] = []

    for _, addrs in net_interfaces.items():
        for addr in addrs:
            if addr.family == socket.AF_INET:
                network_ips.append(ip_address(addr.address))
            elif addr.family == socket.AF_INET6:
                # 过滤掉 IPv6 的 link-local 地址(fe80:...)
                ip = ip_address(addr.address.split("%")[0])  # 处理带 zone index 的情况
                if not ip.is_link_local:
                    network_ips.append(ip)

    return network_ips


async def get_public_ip_address() -> list[IPv4Address | IPv6Address]:
    urls = [
        "https://api64.ipify.org",
        "https://ident.me",
        "https://ifconfig.me",
        "https://icanhazip.com",
    ]
    found_ips: dict[int, IPv4Address | IPv6Address] = {}

    async def fetch(session: aiohttp.ClientSession, url: str):
        try:
            async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
                if resp.status == 200:
                    raw_ip = (await resp.text()).strip()
                    ip = ip_address(raw_ip)
                    if ip.version not in found_ips:
                        found_ips[ip.version] = ip
        except Exception as e:
            # Ignore errors from individual services so that a single failing
            # endpoint does not prevent discovering the public IP from others.
            logger.debug("Failed to fetch public IP from %s: %s", url, e)

    async with aiohttp.ClientSession() as session:
        tasks = [fetch(session, url) for url in urls]
        await asyncio.gather(*tasks)

    # 返回找到的所有 IP 对象列表
    return list(found_ips.values())


def _read_dashboard_dist_version(dist_dir: str | Path) -> str | None:
    version_file = Path(dist_dir) / "assets" / "version"
    if version_file.exists():
        return version_file.read_text(encoding="utf-8").strip()
    return None


def get_dashboard_dist_version(dist_dir: str | Path) -> str | None:
    """Read the version marker from a given dashboard dist directory."""
    return _read_dashboard_dist_version(dist_dir)


def get_bundled_dashboard_dist_path() -> Path:
    return Path(get_astrbot_path()) / "astrbot" / "dashboard" / "dist"


def _normalize_dashboard_version(version: str) -> str:
    version = version.strip()
    if version[:1].lower() == "v":
        version = version[1:]
    if not re.match(
        r"^[0-9]+(?:\.[0-9]+)*"
        r"(?:-[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?"
        r"(?:\+.+)?$",
        version,
    ):
        raise ValueError(f"invalid dashboard version: {version!r}")
    return version


def should_use_bundled_dashboard_dist(
    user_dist: str | Path,
    current_version: str,
) -> bool:
    user_version = _read_dashboard_dist_version(user_dist)
    bundled_dist = get_bundled_dashboard_dist_path()
    if user_version is None or not bundled_dist.exists():
        return False
    try:
        return (
            VersionComparator.compare_version(
                _normalize_dashboard_version(current_version),
                _normalize_dashboard_version(user_version),
            )
            > 0
        )
    except (TypeError, ValueError):
        return False


async def get_dashboard_version():
    # First check user data directory (manually updated / downloaded dashboard).
    dist_dir = os.path.join(get_astrbot_data_path(), "dist")
    if await asyncio.to_thread(os.path.exists, dist_dir):
        from astrbot.core.config.default import VERSION

        if should_use_bundled_dashboard_dist(dist_dir, VERSION):
            bundled_version = _read_dashboard_dist_version(
                get_bundled_dashboard_dist_path(),
            )
            if bundled_version is not None:
                return bundled_version
        return _read_dashboard_dist_version(dist_dir)

    bundled = get_bundled_dashboard_dist_path()
    if bundled.exists():
        return _read_dashboard_dist_version(bundled)
    return None


async def download_dashboard(
    path: str | None = None,
    extract_path: str = "data",
    latest: bool = True,
    version: str | None = None,
    proxy: str | None = None,
    progress_callback=None,
    extract: bool = True,
) -> None:
    """下载管理面板文件"""
    if path is None:
        zip_path = anyio.Path(get_astrbot_data_path()) / "dashboard.zip"
    else:
        zip_path = anyio.Path(path)
    await zip_path.parent.mkdir(parents=True, exist_ok=True)

    if latest or len(str(version)) != 40:
        ver_name = "latest" if latest else version
        dashboard_release_url = f"https://astrbot-registry.soulter.top/download/astrbot-dashboard/{ver_name}/dist.zip"
        logger.info(
            f"Downloading AstrBot WebUI from {dashboard_release_url}",
        )
        try:
            await download_file(
                dashboard_release_url,
                str(zip_path),
                show_progress=True,
                progress_callback=progress_callback,
            )
            if not zipfile.is_zipfile(zip_path):
                raise RuntimeError(
                    "Downloaded dashboard package is not a valid ZIP file"
                )
        except BaseException as _:
            if latest:
                # Resolve latest release tag from GitHub API to construct correct asset URL
                ssl_context = ssl.create_default_context(cafile=certifi.where())
                async with (
                    aiohttp.ClientSession(
                        connector=aiohttp.TCPConnector(ssl=ssl_context),
                        trust_env=True,
                    ) as session,
                    session.get(
                        "https://api.github.com/repos/AstrBotDevs/AstrBot/releases/latest",
                        timeout=30,
                        headers={"Accept": "application/vnd.github+json"},
                    ) as api_resp,
                ):
                    api_resp.raise_for_status()
                    release_data = await api_resp.json()
                    tag = release_data["tag_name"]
            else:
                tag = version
            dashboard_release_url = f"https://github.com/AstrBotDevs/AstrBot/releases/download/{tag}/AstrBot-{tag}-dashboard.zip"
            if proxy:
                dashboard_release_url = f"{proxy}/{dashboard_release_url}"
            await download_file(
                dashboard_release_url,
                str(zip_path),
                show_progress=True,
                progress_callback=progress_callback,
            )
    else:
        url = f"https://github.com/AstrBotDevs/astrbot-release-harbour/releases/download/release-{version}/dist.zip"
        logger.info(f"Downloading AstrBot WebUI from {url}")
        if proxy:
            url = f"{proxy}/{url}"
        await download_file(
            url,
            str(zip_path),
            show_progress=True,
            progress_callback=progress_callback,
        )
    if not zipfile.is_zipfile(zip_path):
        raise RuntimeError("Downloaded dashboard package is not a valid ZIP file")
    if extract:
        extract_dashboard(zip_path, extract_path)


def extract_dashboard(zip_path: str | Path, extract_path: str | Path) -> None:
    """Extract a downloaded dashboard archive."""
    extract_root = Path(extract_path).resolve()
    ensure_dir(extract_root)
    with zipfile.ZipFile(zip_path, "r") as z:
        for member in z.infolist():
            target_path = (extract_root / member.filename).resolve()
            if not target_path.is_relative_to(extract_root):
                raise ValueError(
                    f"Unsafe dashboard archive path: {member.filename}",
                )
            z.extract(member, extract_root)