#!/usr/bin/env python3 """NAS 文件整理主程序""" import os import sys import yaml import logging import signal from pathlib import Path from datetime import datetime from typing import Optional from tqdm import tqdm from nas_client import create_nas_client from classifier import FileClassifier # 全局标志 running = True def signal_handler(signum, frame): """处理中断信号""" global running logger.info("收到中断信号，正在停止...") running = False # 注册信号处理 signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) # 配置日志 def setup_logging(config: dict) -> logging.Logger: """设置日志""" log_config = config.get('logging', {}) log_level = getattr(logging, log_config.get('level', 'INFO')) log_file = log_config.get('file', 'logs/nas-organizer.log') # 确保日志目录存在 log_dir = os.path.dirname(log_file) if log_dir and not os.path.exists(log_dir): os.makedirs(log_dir, exist_ok=True) # 配置日志处理器 handlers = [ logging.FileHandler(log_file, encoding='utf-8'), logging.StreamHandler() ] logging.basicConfig( level=log_level, format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S', handlers=handlers ) return logging.getLogger(__name__) class NASOrganizer: """NAS 文件整理器""" def __init__(self, config_path: str = 'config.yaml'): self.config_path = config_path with open(config_path, 'r', encoding='utf-8') as f: self.config = yaml.safe_load(f) # 设置日志 global logger logger = setup_logging(self.config) # 初始化组件 self.nas = create_nas_client(self.config['nas']) self.classifier = FileClassifier(self.config['classification']) self.source = self.config['paths']['source'].rstrip('/') self.destination = self.config['paths']['destination'].rstrip('/') self.options = self.config.get('options', {}) self.mode = self.config.get('mode', {'type': 'once'}) # 统计信息 self.stats = { 'processed': 0, 'success': 0, 'failed': 0, 'skipped': 0, 'total_size': 0 } def run_once(self) -> bool: """运行一次文件整理""" logger.info("=" * 60) logger.info("开始 NAS 文件整理") logger.info("=" * 60) # 连接 NAS logger.info("正在连接 NAS...") if not self.nas.connect(): logger.error("NAS 连接失败") return False logger.info("✓ NAS 连接成功") logger.info(f"源目录：{self.source}") logger.info(f"目标目录：{self.destination}") try: # 确保目标目录存在 self.nas.create_dir(self.destination) # 获取文件列表 logger.info("正在扫描文件...") files = self.nas.list_files(self.source) files = [f for f in files if not f['is_dir']] logger.info(f"发现 {len(files)} 个文件") if not files: logger.info("没有需要整理的文件") return True # 显示分类统计 self._show_preview(files) if self.options.get('dry_run', False): logger.info("[试运行模式] 不会实际移动文件") return True # 整理文件 self._process_files(files) # 清理空目录 if self.options.get('delete_empty_dirs', False): self._cleanup_empty_dirs() # 显示统计 self._show_stats() return self.stats['failed'] == 0 except Exception as e: logger.error(f"整理过程出错：{e}") return False finally: self.nas.disconnect() def run_watch(self) -> bool: """持续监控模式""" logger.info("启动监控模式...") interval = self.mode.get('watch', {}).get('interval', 60) # 连接 NAS if not self.nas.connect(): logger.error("NAS 连接失败") return False last_check = {} while running: try: logger.info(f"检查新文件 (间隔：{interval}秒)...") files = self.nas.list_files(self.source) files = [f for f in files if not f['is_dir']] # 过滤新文件 new_files = [] for f in files: file_key = f['path'] mtime = f['modified'].timestamp() if file_key not in last_check or last_check[file_key] != mtime: new_files.append(f) last_check[file_key] = mtime if new_files: logger.info(f"发现 {len(new_files)} 个新文件") self._process_files(new_files) else: logger.debug("没有新文件") # 等待下一次检查 import time for _ in range(interval): if not running: break time.sleep(1) except Exception as e: logger.error(f"监控过程出错：{e}") import time time.sleep(10) # 出错后等待 10 秒 self.nas.disconnect() logger.info("监控模式已停止") return True def _show_preview(self, files: list): """显示分类预览""" stats = self.classifier.get_category_stats(files) logger.info("\n分类预览:") logger.info("-" * 40) for category, data in sorted(stats.items()): logger.info(f" {category}: {data['count']} 个文件 ({data['size_human']})") logger.info("-" * 40) def _process_files(self, files: list): """处理文件列表""" for file_info in tqdm(files, desc="整理文件", unit="file"): if not running: logger.info("用户中断，停止处理") break self.stats['processed'] += 1 try: # 获取分类 category = self.classifier.classify(file_info) # 构建目标路径 dst_path = f"{self.destination}/{category}/{file_info['name']}" src_path = file_info['path'] # 检查目标文件是否已存在 if self.options.get('skip_existing', True): if self.nas.file_exists(dst_path): logger.debug(f"跳过已存在：{file_info['name']}") self.stats['skipped'] += 1 continue # 移动文件 if self.nas.move_file(src_path, dst_path): logger.debug(f"✓ {file_info['name']} -> {category}/") self.stats['success'] += 1 self.stats['total_size'] += file_info.get('size', 0) else: logger.warning(f"✗ 移动失败：{file_info['name']}") self.stats['failed'] += 1 except Exception as e: logger.error(f"处理文件 {file_info['name']} 失败：{e}") self.stats['failed'] += 1 def _cleanup_empty_dirs(self): """清理源目录中的空文件夹""" # 注意：SMB 协议删除目录需要特殊处理 # 这里仅记录，实际删除需要 NAS 客户端支持 logger.info("空目录清理功能需要 NAS 客户端支持，已跳过") def _show_stats(self): """显示统计信息""" logger.info("\n" + "=" * 60) logger.info("整理完成") logger.info("=" * 60) logger.info(f"处理文件：{self.stats['processed']} 个") logger.info(f"成功：{self.stats['success']} 个") logger.info(f"失败：{self.stats['failed']} 个") logger.info(f"跳过：{self.stats['skipped']} 个") logger.info(f"总大小：{FileClassifier._format_size(self.stats['total_size'])}") logger.info("=" * 60) def main(): """主函数""" # 获取配置文件路径 config_file = sys.argv[1] if len(sys.argv) > 1 else 'config.yaml' if not os.path.exists(config_file): print(f"❌ 配置文件不存在：{config_file}") print("\n请先创建配置文件:") print(" cp config.yaml.example config.yaml") print(" vim config.yaml") sys.exit(1) # 创建整理器 organizer = NASOrganizer(config_file) # 根据模式运行 mode_type = organizer.mode.get('type', 'once') if mode_type == 'watch': success = organizer.run_watch() else: success = organizer.run_once() sys.exit(0 if success else 1) if __name__ == '__main__': main()