generate_test_data.py

#!/usr/bin/env python3
"""
测试数据生成脚本
为CTF平台生成测试用的解题记录和积分数据，用于测试排行榜功能
"""

import random
from datetime import datetime, timedelta
from app import create_app
from app.models import User, Challenge, Solve, Category, UserRole, Competition, db

def generate_test_data():
    """生成测试数据"""
    app = create_app()
    
    with app.app_context():
        # 获取所有用户和题目
        users = User.query.all()
        challenges = Challenge.query.all()
        
        print(f"找到 {len(users)} 个用户和 {len(challenges)} 个题目")
        
        if len(users) == 0 or len(challenges) == 0:
            print("没有足够的用户或题目数据，请先创建用户和题目")
            return
        
        # 清除现有的解题记录
        existing_solves = Solve.query.all()
        for solve in existing_solves:
            db.session.delete(solve)
        
        print(f"清除了 {len(existing_solves)} 条现有解题记录")
        
        # 为每个用户生成随机的解题记录
        total_solves = 0
        
        for user in users:
            # 跳过管理员账户，让其保持较高排名
            if user.role.name == 'ADMIN':
                # 管理员解决所有题目
                user_challenges = challenges.copy()
                solve_count = len(challenges)
            else:
                # 其他用户随机解决一些题目
                solve_count = random.randint(0, len(challenges))
                user_challenges = random.sample(challenges, solve_count)
            
            print(f"为用户 {user.username} 生成 {solve_count} 条解题记录")
            
            for i, challenge in enumerate(user_challenges):
                # 生成随机的解题时间（过去30天内）
                days_ago = random.randint(0, 30)
                hours_ago = random.randint(0, 23)
                minutes_ago = random.randint(0, 59)
                
                solve_time = datetime.now() - timedelta(
                    days=days_ago, 
                    hours=hours_ago, 
                    minutes=minutes_ago
                )
                
                # 计算积分（基础分数 + 随机奖励分数）
                base_points = challenge.points
                bonus_points = random.randint(0, 20)  # 0-20的随机奖励分
                total_points = base_points + bonus_points
                
                # 创建解题记录
                solve = Solve(
                    user_id=user.id,
                    challenge_id=challenge.id,
                    flag_submitted=challenge.flag,
                    points_earned=challenge.points,
                    created_at=solve_time
                )
                
                db.session.add(solve)
                total_solves += 1
        
        # 提交所有更改
        try:
            db.session.commit()
            print(f"成功生成 {total_solves} 条解题记录")
            
            # 显示生成后的统计信息
            print("\n=== 生成后的数据统计 ===")
            for user in users:
                user_solves = Solve.query.filter_by(user_id=user.id).all()
                total_score = sum(solve.points_earned for solve in user_solves)
                print(f"{user.username}: {len(user_solves)} 题, {total_score} 分")
                
        except Exception as e:
            db.session.rollback()
            print(f"生成测试数据时出错: {e}")

def add_more_users():
    """添加更多测试用户"""
    app = create_app()
    
    with app.app_context():
        # 生成一些有趣的用户名
        usernames = [
            "hacker_alice", "crypto_bob", "web_charlie", "pwn_diana",
            "reverse_eve", "forensics_frank", "misc_grace", "osint_henry",
            "steganography_ivy", "binary_jack", "network_kate", "mobile_leo",
            "blockchain_mary", "ai_nick", "quantum_olivia", "social_paul"
        ]
        
        existing_users = {user.username for user in User.query.all()}
        new_users_count = 0
        
        for username in usernames:
            if username not in existing_users:
                # 随机分配角色
                role_choice = random.choice([UserRole.STUDENT, UserRole.TEACHER])
                
                user = User(
                    username=username,
                    email=f"{username}@test.com",
                    password_hash="pbkdf2:sha256:600000$test$test",  # 测试密码
                    role=role_choice,
                    created_at=datetime.now() - timedelta(days=random.randint(1, 90))
                )
                
                db.session.add(user)
                new_users_count += 1
                print(f"添加用户: {username} ({role_choice})")
        
        try:
            db.session.commit()
            print(f"成功添加 {new_users_count} 个新用户")
        except Exception as e:
            db.session.rollback()
            print(f"添加用户时出错: {e}")

def add_more_challenges():
    """添加更多测试题目"""
    app = create_app()
    
    with app.app_context():
        challenges_data = [
            {"title": "Web进阶 - XSS攻击", "category": "Web", "points": 120, "difficulty": "medium"},
            {"title": "Web高级 - CSRF防护绕过", "category": "Web", "points": 200, "difficulty": "hard"},
            {"title": "Crypto进阶 - RSA解密", "category": "Crypto", "points": 180, "difficulty": "medium"},
            {"title": "Crypto高级 - 椭圆曲线", "category": "Crypto", "points": 250, "difficulty": "hard"},
            {"title": "PWN入门 - 栈溢出", "category": "PWN", "points": 150, "difficulty": "medium"},
            {"title": "PWN进阶 - ROP链构造", "category": "PWN", "points": 300, "difficulty": "hard"},
            {"title": "Reverse入门 - 简单逆向", "category": "Reverse", "points": 100, "difficulty": "easy"},
            {"title": "Reverse进阶 - 反调试", "category": "Reverse", "points": 220, "difficulty": "hard"},
            {"title": "Misc进阶 - 流量分析", "category": "Misc", "points": 160, "difficulty": "medium"},
            {"title": "Misc高级 - 内存取证", "category": "Misc", "points": 280, "difficulty": "hard"},
        ]
        
        existing_titles = {challenge.title for challenge in Challenge.query.all()}
        new_challenges_count = 0
        
        for challenge_data in challenges_data:
            if challenge_data["title"] not in existing_titles:
                # 获取第一个分类作为默认分类
                default_category = Category.query.first()
                if not default_category:
                    # 如果没有分类，创建一个默认分类
                    default_category = Category(name="默认分类", description="默认分类")
                    db.session.add(default_category)
                    db.session.flush()  # 获取ID
                
                # 获取第一个管理员作为作者
                admin_user = User.query.filter_by(role=UserRole.ADMIN).first()
                if not admin_user:
                    admin_user = User.query.first()  # 如果没有管理员，使用第一个用户
                
                challenge = Challenge(
                    title=challenge_data["title"],
                    description=f"这是一道{challenge_data['difficulty']}难度的{challenge_data['category']}题目",
                    category_id=default_category.id,
                    author_id=admin_user.id,
                    points=challenge_data["points"],
                    difficulty=challenge_data["difficulty"],
                    flag=f"flag{{{challenge_data['title'].replace(' ', '_').lower()}}}",
                    is_active=True,
                    created_at=datetime.now() - timedelta(days=random.randint(1, 60))
                )
                
                db.session.add(challenge)
                new_challenges_count += 1
                print(f"添加题目: {challenge_data['title']} ({challenge_data['points']}分)")
        
        try:
            db.session.commit()
            print(f"成功添加 {new_challenges_count} 个新题目")
        except Exception as e:
            db.session.rollback()
            print(f"添加题目时出错: {e}")

def add_test_competitions():
    """添加测试比赛数据"""
    app = create_app()
    
    with app.app_context():
        # 获取管理员用户作为比赛创建者
        admin_user = User.query.filter_by(role=UserRole.ADMIN).first()
        teacher_user = User.query.filter_by(role=UserRole.TEACHER).first()
        
        if not admin_user:
            admin_user = User.query.first()
        
        creator = admin_user if admin_user else teacher_user
        
        if not creator:
            print("没有找到管理员或教师用户，无法创建比赛")
            return
        
        # 定义测试比赛数据
        competitions_data = [
            {
                "title": "新生CTF入门赛",
                "description": "专为CTF新手设计的入门级比赛，涵盖Web、Crypto、Misc等基础题目。",
                "start_time": datetime.now() + timedelta(days=7),  # 7天后开始
                "end_time": datetime.now() + timedelta(days=9),    # 9天后结束
                "freeze_time": datetime.now() + timedelta(days=8, hours=20),  # 封榜时间
                "is_public": True
            },
            {
                "title": "春季CTF挑战赛",
                "description": "中等难度的综合性CTF比赛，适合有一定基础的参赛者。包含Web安全、密码学、逆向工程等多个方向。",
                "start_time": datetime.now() - timedelta(days=2),  # 2天前开始（进行中）
                "end_time": datetime.now() + timedelta(days=5),    # 5天后结束
                "freeze_time": datetime.now() + timedelta(days=4, hours=18),
                "is_public": True
            },
            {
                "title": "高级渗透测试竞赛",
                "description": "高难度的专业级CTF比赛，主要面向有经验的安全研究人员和渗透测试工程师。",
                "start_time": datetime.now() + timedelta(days=14), # 14天后开始
                "end_time": datetime.now() + timedelta(days=17),   # 17天后结束
                "freeze_time": datetime.now() + timedelta(days=16, hours=12),
                "is_public": True
            },
            {
                "title": "内部训练赛",
                "description": "仅限内部人员参加的训练比赛，用于技能提升和团队协作训练。",
                "start_time": datetime.now() + timedelta(days=21), # 21天后开始
                "end_time": datetime.now() + timedelta(days=23),   # 23天后结束
                "freeze_time": None,  # 无封榜时间
                "is_public": False
            },
            {
                "title": "历史经典CTF回顾赛",
                "description": "回顾经典CTF题目的比赛，已经结束。包含了历年来的经典题目和解题思路分享。",
                "start_time": datetime.now() - timedelta(days=30), # 30天前开始（已结束）
                "end_time": datetime.now() - timedelta(days=27),   # 27天前结束
                "freeze_time": datetime.now() - timedelta(days=28),
                "is_public": True
            },
            {
                "title": "Web安全专项赛",
                "description": "专注于Web安全的专项比赛，包含SQL注入、XSS、CSRF、文件上传等各类Web漏洞。",
                "start_time": datetime.now() + timedelta(hours=2), # 2小时后开始
                "end_time": datetime.now() + timedelta(days=3),    # 3天后结束
                "freeze_time": datetime.now() + timedelta(days=2, hours=20),
                "is_public": True
            }
        ]
        
        existing_titles = {comp.title for comp in Competition.query.all()}
        new_competitions_count = 0
        
        for comp_data in competitions_data:
            if comp_data["title"] not in existing_titles:
                competition = Competition(
                    title=comp_data["title"],
                    description=comp_data["description"],
                    creator_id=creator.id,
                    start_time=comp_data["start_time"],
                    end_time=comp_data["end_time"],
                    freeze_time=comp_data["freeze_time"],
                    is_public=comp_data["is_public"],
                    created_at=datetime.now() - timedelta(days=random.randint(1, 10))
                )
                
                db.session.add(competition)
                new_competitions_count += 1
                print(f"添加比赛: {comp_data['title']}")
        
        try:
            db.session.commit()
            print(f"成功添加 {new_competitions_count} 个测试比赛")
            
            # 显示比赛状态统计
            print("\n=== 比赛状态统计 ===")
            competitions = Competition.query.all()
            for comp in competitions:
                status = comp.get_status()
                print(f"{comp.title}: {status.value}")
                
        except Exception as e:
            db.session.rollback()
            print(f"添加比赛时出错: {e}")

def add_competition_participants():
    """为比赛添加参赛者"""
    app = create_app()
    
    with app.app_context():
        competitions = Competition.query.all()
        students = User.query.filter_by(role=UserRole.STUDENT).all()
        
        if not competitions or not students:
            print("没有找到比赛或学生用户")
            return
        
        total_participants = 0
        
        for competition in competitions:
            # 为每个比赛随机添加参赛者
            participant_count = random.randint(3, min(len(students), 15))
            participants = random.sample(students, participant_count)
            
            for participant in participants:
                if participant not in competition.participants:
                    competition.participants.append(participant)
                    total_participants += 1
            
            print(f"为比赛 '{competition.title}' 添加了 {len(participants)} 名参赛者")
        
        try:
            db.session.commit()
            print(f"成功添加 {total_participants} 个参赛记录")
        except Exception as e:
            db.session.rollback()
            print(f"添加参赛者时出错: {e}")

if __name__ == "__main__":
    print("=== CTF平台测试数据生成器 ===")
    print("1. 添加更多用户")
    print("2. 添加更多题目") 
    print("3. 生成解题记录")
    print("4. 添加测试比赛")
    print("5. 添加比赛参赛者")
    print("6. 全部执行")
    
    choice = input("请选择操作 (1-6): ").strip()
    
    if choice == "1":
        add_more_users()
    elif choice == "2":
        add_more_challenges()
    elif choice == "3":
        generate_test_data()
    elif choice == "4":
        add_test_competitions()
    elif choice == "5":
        add_competition_participants()
    elif choice == "6":
        print("执行全部操作...")
        add_more_users()
        add_more_challenges()
        generate_test_data()
        add_test_competitions()
        add_competition_participants()
    else:
        print("无效选择")