FTP同步目录减负，动态路径+使用布隆过滤器 + 跳过已复制文件

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
r"""
按月分目录增量同步工具
- 源根: D:\linux_sync\{biz}\files\YYYY\M
- 目标根: D:\ftpfile\{biz}\files\YYYY\M
- 只扫当月目录，保留原布隆过滤器逻辑
- 新增 --dry-run / -d 参数：模拟运行，不实际复制
"""

import os
import sys
import pickle
import hashlib
import shutil
import time
from datetime import datetime

# ========== 配置 ==========
SOURCE_ROOT = r"D:\linux_sync"
FTP_ROOT = r"D:\ftpfile"
FILES_SUBDIR = "files"
BUSINESS_DIRS = ['lsqy', 'lssj', 'lsyh', 'tzjj', 'tzwl', 'lsjn']

CONFIG_DIR = r"D:\sync_config"
BLOOM_FILE = os.path.join(CONFIG_DIR, "bloom_filter.pkl")
LOG_FILE = os.path.join(CONFIG_DIR, "sync.log")

INITIAL_CAPACITY = 100000
ERROR_RATE = 0.001

# ========== 布隆过滤器（完全保留原实现）==========
class SimpleBloomFilter:
    def __init__(self, capacity=100000, error_rate=0.001):
        self.capacity = capacity
        self.error_rate = error_rate
        self.count = 0
        import math
        self.size = int(-capacity * math.log(error_rate) / (math.log(2) ** 2))
        self.num_hashes = int(self.size / capacity * math.log(2))
        self.size = max(self.size, 1000)
        self.num_hashes = max(self.num_hashes, 3)
        self.num_hashes = min(self.num_hashes, 10)
        self.bits = bytearray((self.size + 7) // 8)
        print(f"[INIT] 布隆过滤器: 容量={capacity}, 误判率={error_rate*100}%, "
              f"位数组={self.size}位, 哈希函数={self.num_hashes}个")

    def _hashes(self, item):
        h1 = int(hashlib.md5(item.encode()).hexdigest(), 16)
        h2 = int(hashlib.sha256(item.encode()).hexdigest(), 16)
        for i in range(self.num_hashes):
            yield (h1 + i * h2) % self.size

    def add(self, item):
        for pos in self._hashes(item):
            byte_idx = pos // 8
            bit_idx = pos % 8
            self.bits[byte_idx] |= (1 << bit_idx)
        self.count += 1

    def __contains__(self, item):
        for pos in self._hashes(item):
            byte_idx = pos // 8
            bit_idx = pos % 8
            if not (self.bits[byte_idx] & (1 << bit_idx)):
                return False
        return True

    def __getstate__(self):
        return {
            'capacity': self.capacity, 'error_rate': self.error_rate,
            'count': self.count, 'size': self.size,
            'num_hashes': self.num_hashes, 'bits': self.bits
        }

    def __setstate__(self, state):
        self.capacity = state['capacity']
        self.error_rate = state['error_rate']
        self.count = state['count']
        self.size = state['size']
        self.num_hashes = state['num_hashes']
        self.bits = state['bits']

# ========== 日志（保留原实现）==========
def log(msg):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    if os.path.exists(LOG_FILE) and os.path.getsize(LOG_FILE) > 1024*1024:
        backup = LOG_FILE + ".old"
        if os.path.exists(backup):
            os.remove(backup)
        shutil.move(LOG_FILE, backup)
        with open(LOG_FILE, 'a', encoding='utf-8') as f:
            f.write(f"[{timestamp}] [ROTATE] 日志超过1MB，已备份\n")
    line = f"[{timestamp}] {msg}"
    print(line)
    with open(LOG_FILE, 'a', encoding='utf-8') as f:
        f.write(line + '\n')

# ========== 自动定位当月目录 ==========
def get_current_month_dirs():
    now = datetime.now()
    year = str(now.year)
    month_variants = [str(now.month), f"{now.month:02d}"]
    
    dirs = []
    for biz in BUSINESS_DIRS:
        source_year = os.path.join(SOURCE_ROOT, biz, FILES_SUBDIR, year)
        if not os.path.exists(source_year):
            continue
        
        month_dir = None
        for m in month_variants:
            mp = os.path.join(source_year, m)
            if os.path.exists(mp):
                month_dir = m
                break
        
        if month_dir:
            source_month = os.path.join(source_year, month_dir)
            target_month = os.path.join(FTP_ROOT, biz, FILES_SUBDIR, year, month_dir)
            dirs.append((source_month, target_month, biz))
    
    return dirs

# ========== 文件签名（基准改为 SOURCE_ROOT，兼容旧 .pkl）==========
def file_signature(filepath):
    try:
        stat = os.stat(filepath)
        rel_path = os.path.relpath(filepath, SOURCE_ROOT)
        key = f"{rel_path}|{stat.st_size}|{int(stat.st_mtime)}"
        return hashlib.md5(key.encode()).hexdigest()
    except Exception as e:
        log(f"[ERROR] 无法获取文件信息: {filepath}, 错误: {e}")
        return None

# ========== 遍历工具（os.scandir 替代 os.walk）==========
def scandir_recursive(path):
    try:
        with os.scandir(path) as it:
            for entry in it:
                if entry.is_dir(follow_symlinks=False):
                    yield from scandir_recursive(entry.path)
                elif entry.is_file(follow_symlinks=False):
                    yield entry
    except (PermissionError, OSError):
        pass

def ensure_dir(path, dry_run=False):
    if dry_run:
        return  # 模拟模式不创建目录
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)

# ========== 复制函数（支持模拟模式）==========
def copy_file(src, dst, dry_run=False):
    if dry_run:
        print(f"  [DRY-RUN] 会复制: {os.path.relpath(src, SOURCE_ROOT)}")
        return True
    try:
        shutil.copy2(src, dst)
        return True
    except Exception as e:
        log(f"[ERROR] 复制失败: {src} -> {dst}, 错误: {e}")
        return False

# ========== 布隆过滤器持久化（保留原实现）==========
def load_bloom():
    if os.path.exists(BLOOM_FILE):
        try:
            with open(BLOOM_FILE, 'rb') as f:
                bloom = pickle.load(f)
            log(f"[LOAD] 加载已有过滤器，已记录约 {bloom.count} 个文件")
            return bloom
        except Exception as e:
            log(f"[WARN] 加载失败，创建新的: {e}")
    bloom = SimpleBloomFilter(INITIAL_CAPACITY, ERROR_RATE)
    log("[INIT] 创建新布隆过滤器")
    return bloom

def save_bloom(bloom):
    try:
        with open(BLOOM_FILE, 'wb') as f:
            pickle.dump(bloom, f)
        size = os.path.getsize(BLOOM_FILE)
        log(f"[SAVE] 过滤器已保存，文件大小: {size} bytes")
        return True
    except Exception as e:
        log(f"[ERROR] 保存失败: {e}")
        return False

# ========== 同步核心（支持模拟模式）==========
def sync_monthly_directories(bloom, dirs, dry_run=False):
    total_copied = total_skipped = total_errors = 0
    
    for source, target, biz in dirs:
        copied = skipped = errors = 0
        ensure_dir(target, dry_run)
        
        mode_tag = "[DRY-RUN]" if dry_run else "[START]"
        log(f"{mode_tag} 扫描 [{biz}]: {source}")
        start_time = time.time()
        
        for entry in scandir_recursive(source):
            sig = file_signature(entry.path)
            if sig is None:
                errors += 1
                continue
            
            if sig in bloom:
                skipped += 1
                if skipped % 1000 == 0:
                    print(f"  [{biz}] 已跳过 {skipped} 个文件...", end='\r')
            else:
                rel_path = os.path.relpath(entry.path, source)
                dst_file = os.path.join(target, rel_path)
                ensure_dir(os.path.dirname(dst_file), dry_run)
                
                if copy_file(entry.path, dst_file, dry_run):
                    if not dry_run:
                        bloom.add(sig)
                    copied += 1
                    stat = os.stat(entry.path)
                    if not dry_run:
                        log(f"[COPY] [{biz}] #{copied}: {rel_path} (大小: {stat.st_size} 字节)")
                    if copied % 100 == 0:
                        print(f"  [{biz}] 已复制 {copied} 个文件...", end='\r')
                else:
                    errors += 1
        
        elapsed = time.time() - start_time
        log(f"[DONE] [{biz}] 完成，耗时: {elapsed:.2f}秒 | 复制:{copied} 跳过:{skipped} 错误:{errors}")
        total_copied += copied
        total_skipped += skipped
        total_errors += errors
    
    return total_copied, total_skipped, total_errors

# ========== 主程序 ==========
def main():
    import argparse
    parser = argparse.ArgumentParser(description='按月分目录增量同步工具')
    parser.add_argument('--reset', action='store_true', help='重置过滤器')
    parser.add_argument('--status', action='store_true', help='查看过滤器状态')
    parser.add_argument('-d', '--dry-run', action='store_true', help='模拟运行，不实际复制')
    args = parser.parse_args()
    
    ensure_dir(CONFIG_DIR)
    
    if args.reset:
        if os.path.exists(BLOOM_FILE):
            backup = BLOOM_FILE + ".backup." + datetime.now().strftime("%Y%m%d%H%M%S")
            shutil.move(BLOOM_FILE, backup)
            log(f"[RESET] 旧过滤器已备份: {backup}")
        log("[RESET] 过滤器已重置")
        return
    
    if args.status:
        if os.path.exists(BLOOM_FILE):
            bloom = load_bloom()
            print(f"\n过滤器状态:")
            print(f"  文件: {BLOOM_FILE}")
            print(f"  大小: {os.path.getsize(BLOOM_FILE)} bytes")
            print(f"  已记录: ~{bloom.count}")
        else:
            print("过滤器不存在")
        return
    
    dirs = get_current_month_dirs()
    if not dirs:
        log(f"[SKIP] 当前年月({datetime.now().strftime('%Y-%m')})无目录需要扫描")
        return
    
    log("=" * 50)
    mode = "模拟运行" if args.dry_run else "按月同步启动"
    log(f"{mode} | 当前年月: {datetime.now().strftime('%Y-%m')}")
    log(f"本次扫描 {len(dirs)} 个业务目录: {[d[2] for d in dirs]}")
    
    bloom = load_bloom()
    copied, skipped, errors = sync_monthly_directories(bloom, dirs, dry_run=args.dry_run)
    
    if args.dry_run:
        print(f"\n{'='*50}")
        print("[DRY-RUN] 模拟完成，未执行任何复制")
        print(f"  本会复制: {copied}")
        print(f"  本会跳过: {skipped}")
        print(f"  错误: {errors}")
        print("=" * 50)
        log(f"[DRY-RUN SUMMARY] 本会复制{copied}, 跳过{skipped}, 错误{errors}")
    else:
        if save_bloom(bloom):
            print(f"\n{'='*50}")
            print("同步完成统计")
            print(f"  复制文件: {copied}")
            print(f"  跳过文件: {skipped}")
            print(f"  错误: {errors}")
            print(f"  过滤器已记录: ~{bloom.count}")
            print("=" * 50)
            log(f"[SUMMARY] 复制{copied}, 跳过{skipped}, 错误{errors}")
        else:
            log("[ERROR] 同步完成但保存过滤器失败！")
            sys.exit(1)

if __name__ == "__main__":
    main()

{/collapse-item}