MoviePilot/app/modules/filemanager/storages/smb.py

import signal
import threading
import time
from pathlib import Path
from typing import List, Optional, Union

import smbclient
from smbclient import ClientConfig, register_session, reset_connection_cache
from smbprotocol.exceptions import SMBException, SMBResponseException, SMBAuthenticationError

from app import schemas
from app.core.config import settings, global_vars
from app.log import logger
from app.modules.filemanager import StorageBase
from app.modules.filemanager.storages import transfer_process
from app.schemas.types import StorageSchema
from app.utils.singleton import WeakSingleton

lock = threading.Lock()

# 全局SMB实例列表，用于信号处理
_smb_instances = []


def _signal_handler(signum, frame):
    """
    信号处理器，用于在进程终止时清理SMB连接
    """
    logger.info(f"【SMB】收到信号 {signum}，开始清理SMB连接...")
    for smb_instance in _smb_instances:
        try:
            if hasattr(smb_instance, '_cleanup_connection'):
                smb_instance._cleanup_connection()
        except Exception as e:
            logger.debug(f"【SMB】清理实例时出错: {e}")


# 注册信号处理器
signal.signal(signal.SIGTERM, _signal_handler)
signal.signal(signal.SIGINT, _signal_handler)


def _safe_smb_operation(operation_func, *args, **kwargs):
    """
    安全地执行SMB操作，处理连接错误和线程问题
    """
    max_retries = 3
    retry_count = 0
    
    while retry_count < max_retries:
        try:
            return operation_func(*args, **kwargs)
        except (SMBResponseException, SMBException) as e:
            retry_count += 1
            logger.warning(f"【SMB】操作失败，尝试重试 ({retry_count}/{max_retries}): {e}")
            if retry_count >= max_retries:
                logger.error(f"【SMB】操作失败，已达到最大重试次数: {e}")
                raise
            # 短暂等待后重试
            import time
            time.sleep(1)
        except Exception as e:
            logger.error(f"【SMB】操作时发生未知错误: {e}")
            raise


class SMBConnectionError(Exception):
    """
    SMB 连接错误
    """
    pass


class SMB(StorageBase, metaclass=WeakSingleton):
    """
    SMB网络挂载存储相关操作 - 使用 smbclient 高级接口
    """

    # 存储类型
    schema = StorageSchema.SMB

    # 支持的整理方式
    transtype = {
        "move": "移动",
        "copy": "复制",
    }

    # 文件块大小，默认10MB
    chunk_size = 10 * 1024 * 1024

    def __enter__(self):
        """
        上下文管理器入口
        """
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        上下文管理器出口，确保连接被正确清理
        """
        try:
            self._cleanup_connection()
        except Exception as e:
            logger.debug(f"【SMB】上下文管理器清理连接失败: {e}")

    def __init__(self):
        super().__init__()
        self._connected = False
        self._server_path = None
        self._host = None
        self._username = None
        self._password = None
        
        # 注册到全局实例列表
        _smb_instances.append(self)
        
        self._init_connection()

    def _init_connection(self):
        """
        初始化SMB连接配置
        """
        try:
            conf = self.get_conf()
            if not conf:
                return

            self._host = conf.get("host")
            self._username = conf.get("username")
            self._password = conf.get("password")
            domain = conf.get("domain", "")
            share = conf.get("share", "")
            port = conf.get("port", 445)

            if not all([self._host, share]):
                logger.error("【SMB】缺少必要的连接参数：host 和 share")
                return

            # 构建服务器路径
            self._server_path = f"\\\\{self._host}\\{share}"

            # 配置全局客户端设置
            ClientConfig(
                username=self._username,
                password=self._password,
                domain=domain if domain else None,
                connection_timeout=60,
                port=port,
                auth_protocol="negotiate",  # 使用协商认证
                require_secure_negotiate=False  # 匿名访问时可能需要关闭安全协商
            )

            # 注册会话以启用连接池
            register_session(
                self._host,
                username=self._username,
                password=self._password,
                port=port,
                encrypt=False,  # 根据需要启用加密
                connection_timeout=60
            )

            # 测试连接
            self._test_connection()

            self._connected = True
            # 判断是否为匿名访问
            if self._is_anonymous_access():
                logger.info(f"【SMB】匿名连接成功：{self._server_path}")
            else:
                logger.info(f"【SMB】认证连接成功：{self._server_path} (用户：{self._username})")

        except Exception as e:
            logger.error(f"【SMB】连接初始化失败：{e}")
            self._connected = False

    def _test_connection(self):
        """
        测试SMB连接
        """
        try:
            # 使用安全操作包装器测试连接
            _safe_smb_operation(smbclient.listdir, self._server_path)
        except SMBAuthenticationError as e:
            raise SMBConnectionError(f"SMB认证失败：{e}")
        except SMBResponseException as e:
            raise SMBConnectionError(f"SMB响应错误：{e}")
        except SMBException as e:
            raise SMBConnectionError(f"SMB连接错误：{e}")
        except Exception as e:
            raise SMBConnectionError(f"连接测试失败：{e}")

    def _is_anonymous_access(self) -> bool:
        """
        检查是否为匿名访问
        """
        return not self._username and not self._password

    def _check_connection(self):
        """
        检查SMB连接状态
        """
        if not self._connected or not self._server_path:
            raise SMBConnectionError("【SMB】连接未建立或已断开，请检查配置！")
        
        # 尝试重新连接如果连接已断开
        try:
            self._test_connection()
        except Exception as e:
            logger.warning(f"【SMB】连接检查失败，尝试重新连接: {e}")
            self._reconnect()

    def _normalize_path(self, path: Union[str, Path]) -> str:
        """
        标准化路径格式为SMB路径
        """
        path_str = str(path)

        # 处理根路径
        if path_str in ["/", "\\"]:
            return self._server_path

        # 去除前导斜杠
        if path_str.startswith("/"):
            path_str = path_str[1:]

        # 构建完整的SMB路径
        if path_str:
            return f"{self._server_path}\\{path_str.replace('/', '\\')}"
        else:
            return self._server_path

    def _create_fileitem(self, stat_result, file_path: str, name: str) -> schemas.FileItem:
        """
        创建文件项
        """
        try:
            # 检查是否为目录
            is_directory = smbclient.path.isdir(file_path)

            # 处理路径
            relative_path = file_path.replace(self._server_path, "").replace("\\", "/")
            if not relative_path.startswith("/"):
                relative_path = "/" + relative_path

            if is_directory and not relative_path.endswith("/"):
                relative_path += "/"

            # 获取时间戳
            try:
                modify_time = int(stat_result.st_mtime)
            except (AttributeError, TypeError):
                modify_time = int(time.time())

            if is_directory:
                return schemas.FileItem(
                    storage=self.schema.value,
                    type="dir",
                    path=relative_path,
                    name=name,
                    basename=name,
                    modify_time=modify_time
                )
            else:
                return schemas.FileItem(
                    storage=self.schema.value,
                    type="file",
                    path=relative_path,
                    name=name,
                    basename=Path(name).stem,
                    extension=Path(name).suffix[1:] if Path(name).suffix else None,
                    size=getattr(stat_result, 'st_size', 0),
                    modify_time=modify_time
                )
        except Exception as e:
            logger.error(f"【SMB】创建文件项失败：{e}")
            # 返回基本的文件项信息
            return schemas.FileItem(
                storage=self.schema.value,
                type="file",
                path=file_path.replace(self._server_path, "").replace("\\", "/"),
                name=name,
                basename=Path(name).stem,
                modify_time=int(time.time())
            )

    def init_storage(self):
        """
        初始化存储
        """
        # 重置连接缓存
        reset_connection_cache()
        self._init_connection()

    def check(self) -> bool:
        """
        检查存储是否可用
        """
        if not self._connected:
            return False

        try:
            self._test_connection()
            return True
        except Exception as e:
            logger.debug(f"【SMB】连接检查失败：{e}")
            self._connected = False
            # 尝试重新连接
            try:
                self._reconnect()
                if self._connected:
                    return True
            except Exception as reconnect_error:
                logger.debug(f"【SMB】重新连接失败：{reconnect_error}")
            return False

    def list(self, fileitem: schemas.FileItem) -> List[schemas.FileItem]:
        """
        浏览文件
        """
        try:
            self._check_connection()

            if fileitem.type == "file":
                item = self.detail(fileitem)
                if item:
                    return [item]
                return []

            # 构建SMB路径
            smb_path = self._normalize_path(fileitem.path.rstrip("/"))

            # 列出目录内容
            try:
                entries = _safe_smb_operation(smbclient.listdir, smb_path)
            except Exception as e:
                logger.error(f"【SMB】列出目录失败: {smb_path} - {e}")
                # 尝试重新连接
                self._reconnect()
                return []

            items = []
            for entry in entries:
                if entry in [".", ".."]:
                    continue

                entry_path = f"{smb_path}\\{entry}"
                try:
                    stat_result = smbclient.stat(entry_path)
                    item = self._create_fileitem(stat_result, entry_path, entry)
                    items.append(item)
                except Exception as e:
                    logger.debug(f"【SMB】获取文件信息失败: {entry_path} - {e}")
                    continue

            return items
        except Exception as e:
            logger.error(f"【SMB】列出文件失败: {e}")
            # 尝试重新连接
            self._reconnect()
            return []

    def create_folder(self, fileitem: schemas.FileItem, name: str) -> Optional[schemas.FileItem]:
        """
        创建目录
        """
        try:
            self._check_connection()

            parent_path = self._normalize_path(fileitem.path.rstrip("/"))
            new_path = f"{parent_path}\\{name}"

            # 创建目录
            smbclient.mkdir(new_path)

            # 返回创建的目录信息
            return schemas.FileItem(
                storage=self.schema.value,
                type="dir",
                path=f"{fileitem.path.rstrip('/')}/{name}/",
                name=name,
                basename=name,
                modify_time=int(time.time())
            )
        except Exception as e:
            logger.error(f"【SMB】创建目录失败: {e}")
            return None

    def get_folder(self, path: Path) -> Optional[schemas.FileItem]:
        """
        获取目录，如目录不存在则创建
        """
        # 检查目录是否存在
        folder = self.get_item(path)
        if folder:
            return folder

        # 逐级创建目录
        parts = path.parts
        current_path = Path("/")

        for part in parts[1:]:  # 跳过根目录
            current_path = current_path / part
            folder = self.get_item(current_path)
            if not folder:
                parent_folder = self.get_item(current_path.parent)
                if not parent_folder:
                    logger.error(f"【SMB】父目录不存在: {current_path.parent}")
                    return None
                folder = self.create_folder(parent_folder, part)
                if not folder:
                    return None

        return folder

    def get_item(self, path: Path) -> Optional[schemas.FileItem]:
        """
        获取文件或目录，不存在返回None
        """
        try:
            self._check_connection()

            # 处理根目录
            if str(path) == "/":
                return schemas.FileItem(
                    storage=self.schema.value,
                    type="dir",
                    path="/",
                    name="",
                    basename="",
                    modify_time=int(time.time())
                )

            smb_path = self._normalize_path(str(path).rstrip("/"))

            # 检查路径是否存在
            if not smbclient.path.exists(smb_path):
                return None

            stat_result = smbclient.stat(smb_path)
            file_name = Path(path).name

            return self._create_fileitem(stat_result, smb_path, file_name)

        except Exception as e:
            logger.debug(f"【SMB】获取文件项失败: {e}")
            return None

    def detail(self, fileitem: schemas.FileItem) -> Optional[schemas.FileItem]:
        """
        获取文件详情
        """
        return self.get_item(Path(fileitem.path))

    def delete(self, fileitem: schemas.FileItem) -> bool:
        """
        删除文件或目录
        """
        try:
            self._check_connection()

            smb_path = self._normalize_path(fileitem.path.rstrip("/"))
            logger.info(f"【SMB】开始删除: {fileitem.path} (类型: {fileitem.type})")

            # 先检查路径是否存在
            if not smbclient.path.exists(smb_path):
                logger.warn(f"【SMB】路径不存在，跳过删除: {fileitem.path}")
                return True

            if fileitem.type == "dir":
                # 递归删除目录及其内容
                logger.debug(f"【SMB】递归删除目录: {smb_path}")
                self._recursive_delete(smb_path)
            else:
                # 删除文件
                logger.debug(f"【SMB】删除文件: {smb_path}")
                smbclient.remove(smb_path)

            logger.info(f"【SMB】删除成功: {fileitem.path}")
            return True
        except SMBConnectionError as e:
            logger.error(f"【SMB】删除失败 - 连接错误: {fileitem.path} - {e}")
            return False
        except SMBResponseException as e:
            logger.error(f"【SMB】删除失败 - SMB响应错误: {fileitem.path} - {e}")
            return False
        except SMBException as e:
            logger.error(f"【SMB】删除失败 - SMB错误: {fileitem.path} - {e}")
            return False
        except Exception as e:
            logger.error(f"【SMB】删除失败 - 未知错误: {fileitem.path} - {e}")
            return False

    def _recursive_delete(self, smb_path: str):
        """
        递归删除目录及其所有内容
        """
        try:
            # 检查路径是否存在
            if not smbclient.path.exists(smb_path):
                logger.debug(f"【SMB】路径不存在，跳过删除: {smb_path}")
                return

            # 如果是文件，直接删除
            if smbclient.path.isfile(smb_path):
                logger.debug(f"【SMB】删除文件: {smb_path}")
                smbclient.remove(smb_path)
                return

            # 如果是目录，先删除其内容
            if smbclient.path.isdir(smb_path):
                logger.debug(f"【SMB】开始删除目录内容: {smb_path}")
                try:
                    # 列出目录内容
                    entries = smbclient.listdir(smb_path)
                    logger.debug(f"【SMB】目录 {smb_path} 包含 {len(entries)} 个项目")

                    for entry in entries:
                        if entry in [".", ".."]:
                            continue
                        entry_path = f"{smb_path}\\{entry}"
                        logger.debug(f"【SMB】递归删除子项: {entry_path}")
                        # 递归删除子项
                        self._recursive_delete(entry_path)

                    # 删除空目录
                    logger.debug(f"【SMB】删除空目录: {smb_path}")
                    smbclient.rmdir(smb_path)
                    logger.debug(f"【SMB】目录删除成功: {smb_path}")

                except SMBResponseException as e:
                    # 如果目录不为空，尝试强制删除
                    logger.warn(f"【SMB】目录不为空，尝试强制删除: {smb_path} - {e}")
                    # 使用remove方法尝试删除（某些SMB服务器支持）
                    try:
                        smbclient.remove(smb_path)
                        logger.info(f"【SMB】强制删除目录成功: {smb_path}")
                    except Exception as remove_error:
                        # 如果还是失败，记录错误并抛出异常
                        logger.error(f"【SMB】无法删除非空目录: {smb_path} - {remove_error}")
                        raise SMBConnectionError(f"无法删除非空目录 {smb_path}: {remove_error}")
                except SMBException as e:
                    logger.error(f"【SMB】SMB操作失败: {smb_path} - {e}")
                    raise SMBConnectionError(f"SMB操作失败 {smb_path}: {e}")

        except SMBConnectionError:
            # 重新抛出SMB连接错误
            raise
        except Exception as e:
            logger.error(f"【SMB】递归删除失败: {smb_path} - {e}")
            raise SMBConnectionError(f"递归删除失败 {smb_path}: {e}")

    def rename(self, fileitem: schemas.FileItem, name: str) -> bool:
        """
        重命名文件
        """
        try:
            self._check_connection()

            old_path = self._normalize_path(fileitem.path.rstrip("/"))
            parent_path = Path(fileitem.path).parent
            new_path = self._normalize_path(str(parent_path / name))

            # 重命名
            smbclient.rename(old_path, new_path)

            logger.info(f"【SMB】重命名成功: {fileitem.path} -> {name}")
            return True
        except Exception as e:
            logger.error(f"【SMB】重命名失败: {e}")
            return False

    def download(self, fileitem: schemas.FileItem, path: Path = None) -> Optional[Path]:
        """
        带实时进度显示的下载
        """
        local_path = path or settings.TEMP_PATH / fileitem.name
        smb_path = self._normalize_path(fileitem.path)
        try:
            self._check_connection()

            # 确保本地目录存在
            local_path.parent.mkdir(parents=True, exist_ok=True)

            # 获取文件大小
            file_size = fileitem.size

            # 初始化进度条
            logger.info(f"【SMB】开始下载: {fileitem.name} -> {local_path}")
            progress_callback = transfer_process(Path(fileitem.path).as_posix())

            # 使用更高效的文件传输方式
            with smbclient.open_file(smb_path, mode="rb") as src_file:
                with open(local_path, "wb") as dst_file:
                    downloaded_size = 0
                    while True:
                        if global_vars.is_transfer_stopped(fileitem.path):
                            logger.info(f"【SMB】{fileitem.path} 下载已取消！")
                            return None
                        chunk = src_file.read(self.chunk_size)
                        if not chunk:
                            break
                        dst_file.write(chunk)
                        downloaded_size += len(chunk)
                        # 更新进度
                        if file_size:
                            progress = (downloaded_size * 100) / file_size
                            progress_callback(progress)

            # 完成下载
            progress_callback(100)
            logger.info(f"【SMB】下载完成: {fileitem.name}")
            return local_path

        except Exception as e:
            logger.error(f"【SMB】下载失败: {fileitem.name} - {e}")
            # 删除可能部分下载的文件
            if local_path.exists():
                local_path.unlink()
            return None

    def upload(self, fileitem: schemas.FileItem, path: Path,
               new_name: Optional[str] = None) -> Optional[schemas.FileItem]:
        """
        带实时进度显示的上传
        """
        target_name = new_name or path.name
        target_path = Path(fileitem.path) / target_name
        smb_path = self._normalize_path(str(target_path))

        try:
            self._check_connection()

            # 获取文件大小
            file_size = path.stat().st_size

            # 初始化进度条
            logger.info(f"【SMB】开始上传: {path} -> {target_path}")
            progress_callback = transfer_process(path.as_posix())

            # 使用更高效的文件传输方式
            with open(path, "rb") as src_file:
                with smbclient.open_file(smb_path, mode="wb") as dst_file:
                    uploaded_size = 0
                    while True:
                        if global_vars.is_transfer_stopped(path.as_posix()):
                            logger.info(f"【SMB】{path} 上传已取消！")
                            return None
                        chunk = src_file.read(self.chunk_size)
                        if not chunk:
                            break
                        dst_file.write(chunk)
                        uploaded_size += len(chunk)
                        # 更新进度
                        if file_size:
                            progress = (uploaded_size * 100) / file_size
                            progress_callback(progress)

            # 完成上传
            progress_callback(100)
            logger.info(f"【SMB】上传完成: {target_name}")

            # 返回上传后的文件信息
            return self.get_item(target_path)

        except Exception as e:
            logger.error(f"【SMB】上传失败: {target_name} - {e}")
            return None

    def copy(self, fileitem: schemas.FileItem, path: Path, new_name: str) -> bool:
        """
        复制文件
        """
        try:
            # 下载到临时文件
            temp_file = self.download(fileitem)
            if not temp_file:
                return False

            # 获取目标目录
            target_folder = self.get_item(path)
            if not target_folder:
                return False

            # 上传到目标位置
            result = self.upload(target_folder, temp_file, new_name)

            # 删除临时文件
            if temp_file.exists():
                temp_file.unlink()

            return result is not None
        except Exception as e:
            logger.error(f"【SMB】复制失败: {e}")
            return False

    def move(self, fileitem: schemas.FileItem, path: Path, new_name: str) -> bool:
        """
        移动文件
        """
        try:
            # 先复制
            if not self.copy(fileitem, path, new_name):
                return False

            # 再删除原文件
            if not self.delete(fileitem):
                logger.warn(f"【SMB】删除原文件失败: {fileitem.path}")
                return False

            return True
        except Exception as e:
            logger.error(f"【SMB】移动失败: {e}")
            return False

    def link(self, fileitem: schemas.FileItem, target_file: Path) -> bool:
        pass

    def softlink(self, fileitem: schemas.FileItem, target_file: Path) -> bool:
        pass

    def usage(self) -> Optional[schemas.StorageUsage]:
        """
        存储使用情况
        """
        try:
            self._check_connection()
            volume_stat = smbclient.stat_volume(self._server_path)
            return schemas.StorageUsage(
                total=volume_stat.total_size,
                available=volume_stat.caller_available_size
            )

        except Exception as e:
            logger.error(f"【SMB】获取存储使用情况失败: {e}")
            return None

    def __del__(self):
        """
        析构函数，清理连接
        """
        try:
            # 从全局实例列表中移除
            if self in _smb_instances:
                _smb_instances.remove(self)
            
            # smbclient 自动管理连接池，但我们可以重置缓存
            if hasattr(self, '_connected') and self._connected:
                self._cleanup_connection()
        except Exception as e:
            logger.debug(f"【SMB】清理连接失败: {e}")

    def _reconnect(self):
        """
        重新建立SMB连接
        """
        try:
            logger.info("【SMB】尝试重新连接...")
            
            # 清理现有连接
            self._cleanup_connection()
            
            # 重新初始化连接
            self._init_connection()
            
            if self._connected:
                logger.info("【SMB】重新连接成功")
            else:
                logger.error("【SMB】重新连接失败")
                
        except Exception as e:
            logger.error(f"【SMB】重新连接过程中出现错误: {e}")
            self._connected = False

    def _cleanup_connection(self):
        """
        安全地清理SMB连接
        """
        try:
            # 标记连接为断开状态
            self._connected = False
            
            # 重置连接缓存，这会清理所有连接
            reset_connection_cache()
            
            logger.debug("【SMB】连接清理完成")
        except Exception as e:
            logger.debug(f"【SMB】连接清理过程中出现错误: {e}")
            # 即使清理失败，也要确保连接状态被重置
            self._connected = False