ai_wht_wechat/backend/storage_state_manager.py

"""
小红书Storage State文件管理工具
用于管理和清理storage_state文件
"""
import os
import json
import time
from datetime import datetime, timedelta
from pathlib import Path


STORAGE_DIR = "storage_states"


def get_storage_files():
    """获取所有storage_state文件"""
    if not os.path.exists(STORAGE_DIR):
        return []
    
    files = []
    for filename in os.listdir(STORAGE_DIR):
        if filename.endswith('.json'):
            filepath = os.path.join(STORAGE_DIR, filename)
            stat = os.stat(filepath)
            files.append({
                'filename': filename,
                'filepath': filepath,
                'size': stat.st_size,
                'modified_time': stat.st_mtime,
                'modified_date': datetime.fromtimestamp(stat.st_mtime)
            })
    return files


def cleanup_old_files(days=30):
    """清理超过指定天数未使用的文件"""
    files = get_storage_files()
    cutoff_time = time.time() - (days * 24 * 60 * 60)
    deleted_count = 0
    
    print(f"\n开始清理{days}天前的storage_state文件...")
    for file_info in files:
        if file_info['modified_time'] < cutoff_time:
            try:
                os.remove(file_info['filepath'])
                print(f"  已删除: {file_info['filename']} (最后修改: {file_info['modified_date']})")
                deleted_count += 1
            except Exception as e:
                print(f"  删除失败 {file_info['filename']}: {e}")
    
    print(f"\n清理完成！共删除 {deleted_count} 个文件")
    return deleted_count


def list_storage_files():
    """列出所有storage_state文件"""
    files = get_storage_files()
    
    if not files:
        print("\n未找到任何storage_state文件")
        return
    
    print(f"\n找到 {len(files)} 个storage_state文件:\n")
    print(f"{'文件名':<40} {'大小':<10} {'最后修改时间'}")
    print("-" * 80)
    
    for file_info in sorted(files, key=lambda x: x['modified_time'], reverse=True):
        size_kb = file_info['size'] / 1024
        print(f"{file_info['filename']:<40} {size_kb:>8.1f}KB {file_info['modified_date']}")
    
    total_size = sum(f['size'] for f in files) / 1024 / 1024
    print(f"\n总大小: {total_size:.2f} MB")


def validate_storage_file(phone):
    """验证指定手机号的storage_state文件是否有效"""
    filepath = os.path.join(STORAGE_DIR, f"xhs_{phone}.json")
    
    if not os.path.exists(filepath):
        print(f"\n❌ 文件不存在: {filepath}")
        return False
    
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 检查必要字段
        if 'cookies' not in data:
            print(f"\n❌ 文件格式错误: 缺少cookies字段")
            return False
        
        if 'origins' not in data:
            print(f"\n⚠️ 文件格式不完整: 缺少origins字段")
        
        cookie_count = len(data.get('cookies', []))
        print(f"\n✅ 文件有效")
        print(f"   Cookie数量: {cookie_count}")
        print(f"   文件大小: {os.path.getsize(filepath) / 1024:.1f}KB")
        print(f"   最后修改: {datetime.fromtimestamp(os.path.getmtime(filepath))}")
        
        return True
        
    except json.JSONDecodeError:
        print(f"\n❌ 文件格式错误: 不是有效的JSON")
        return False
    except Exception as e:
        print(f"\n❌ 验证失败: {e}")
        return False


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("用法:")
        print("  python storage_state_manager.py list                    # 列出所有文件")
        print("  python storage_state_manager.py cleanup [days]          # 清理旧文件（默认30天）")
        print("  python storage_state_manager.py validate <phone>        # 验证指定手机号的文件")
        sys.exit(1)
    
    command = sys.argv[1]
    
    if command == "list":
        list_storage_files()
    elif command == "cleanup":
        days = int(sys.argv[2]) if len(sys.argv) > 2 else 30
        cleanup_old_files(days)
    elif command == "validate":
        if len(sys.argv) < 3:
            print("错误: 请提供手机号")
            sys.exit(1)
        phone = sys.argv[2]
        validate_storage_file(phone)
    else:
        print(f"未知命令: {command}")
        sys.exit(1)