CTFd/populate.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import datetime
import hashlib
import random
import argparse

from CTFd import create_app
from CTFd.cache import clear_config, clear_standings, clear_pages
from CTFd.models import (
    Users,
    Teams,
    Challenges,
    Flags,
    Awards,
    ChallengeFiles,
    Fails,
    Solves,
)
from faker import Faker

fake = Faker()

parser = argparse.ArgumentParser()

parser.add_argument("--mode", help="Set user mode", default="teams")
parser.add_argument("--users", help="Amount of users to generate", default=50, type=int)
parser.add_argument("--teams", help="Amount of teams to generate", default=10, type=int)
parser.add_argument(
    "--challenges", help="Amount of challenges to generate", default=20, type=int
)
parser.add_argument(
    "--awards", help="Amount of awards to generate", default=5, type=int
)

args = parser.parse_args()

app = create_app()

mode = args.mode
USER_AMOUNT = args.users
TEAM_AMOUNT = args.teams if args.mode == "teams" else 0
CHAL_AMOUNT = args.challenges
AWARDS_AMOUNT = args.awards

categories = [
    "Exploitation",
    "Reversing",
    "Web",
    "Forensics",
    "Scripting",
    "Cryptography",
    "Networking",
]
companies = ["Corp", "Inc.", "Squad", "Team"]
icons = [
    None,
    "shield",
    "bug",
    "crown",
    "crosshairs",
    "ban",
    "lightning",
    "code",
    "cowboy",
    "angry",
]


def gen_sentence():
    return fake.text()


def gen_name():
    return fake.first_name()


def gen_team_name():
    return fake.word().capitalize() + str(random.randint(1, 1000))


def gen_email():
    return fake.email()


def gen_category():
    return random.choice(categories)


def gen_affiliation():
    return (fake.word() + " " + random.choice(companies)).title()


def gen_value():
    return random.choice(range(100, 500, 50))


def gen_word():
    return fake.word()


def gen_icon():
    return random.choice(icons)


def gen_file():
    return fake.file_name()


def random_date(start, end):
    return start + datetime.timedelta(
        seconds=random.randint(0, int((end - start).total_seconds()))
    )


def random_chance():
    return random.random() > 0.5


if __name__ == "__main__":
    with app.app_context():
        db = app.db

        # Generating Challenges
        print("GENERATING CHALLENGES")
        for x in range(CHAL_AMOUNT):
            word = gen_word()
            chal = Challenges(
                name=word,
                description=gen_sentence(),
                value=gen_value(),
                category=gen_category(),
            )
            db.session.add(chal)
            db.session.commit()
            f = Flags(challenge_id=x + 1, content=word, type="static")
            db.session.add(f)
            db.session.commit()

        # Generating Files
        print("GENERATING FILES")
        AMT_CHALS_WITH_FILES = int(CHAL_AMOUNT * (3.0 / 4.0))
        for x in range(AMT_CHALS_WITH_FILES):
            chal = random.randint(1, CHAL_AMOUNT)
            filename = gen_file()
            md5hash = hashlib.md5(filename.encode("utf-8")).hexdigest()
            chal_file = ChallengeFiles(
                challenge_id=chal, location=md5hash + "/" + filename
            )
            db.session.add(chal_file)

        db.session.commit()

        # Generating Teams
        print("GENERATING TEAMS")
        used = []
        used_oauth_ids = []
        count = 0
        while count < TEAM_AMOUNT:
            name = gen_team_name()
            if name not in used:
                used.append(name)
                team = Teams(name=name, password="password")
                if random_chance():
                    team.affiliation = gen_affiliation()
                if random_chance():
                    oauth_id = random.randint(1, 1000)
                    while oauth_id in used_oauth_ids:
                        oauth_id = random.randint(1, 1000)
                    used_oauth_ids.append(oauth_id)
                    team.oauth_id = oauth_id
                db.session.add(team)
                count += 1

        db.session.commit()

        # Generating Users
        print("GENERATING USERS")
        used = []
        used_oauth_ids = []
        count = 0
        while count < USER_AMOUNT:
            name = gen_name()
            if name not in used:
                used.append(name)
                try:
                    user = Users(name=name, email=gen_email(), password="password")
                    user.verified = True
                    if random_chance():
                        user.affiliation = gen_affiliation()
                    if random_chance():
                        oauth_id = random.randint(1, 1000)
                        while oauth_id in used_oauth_ids:
                            oauth_id = random.randint(1, 1000)
                        used_oauth_ids.append(oauth_id)
                        user.oauth_id = oauth_id
                    if mode == "teams":
                        user.team_id = random.randint(1, TEAM_AMOUNT)
                    db.session.add(user)
                    count += 1
                except Exception:
                    pass

        db.session.commit()

        if mode == "teams":
            # Assign Team Captains
            print("GENERATING TEAM CAPTAINS")
            teams = Teams.query.all()
            for team in teams:
                captain = (
                    Users.query.filter_by(team_id=team.id)
                    .order_by(Users.id)
                    .limit(1)
                    .first()
                )
                if captain:
                    team.captain_id = captain.id
            db.session.commit()

        # Generating Solves
        print("GENERATING SOLVES")
        if mode == "users":
            for x in range(USER_AMOUNT):
                used = []
                base_time = datetime.datetime.utcnow() + datetime.timedelta(
                    minutes=-10000
                )
                for y in range(random.randint(1, CHAL_AMOUNT)):
                    chalid = random.randint(1, CHAL_AMOUNT)
                    if chalid not in used:
                        used.append(chalid)
                        user = Users.query.filter_by(id=x + 1).first()
                        solve = Solves(
                            user_id=user.id,
                            team_id=user.team_id,
                            challenge_id=chalid,
                            ip="127.0.0.1",
                            provided=gen_word(),
                        )

                        new_base = random_date(
                            base_time,
                            base_time
                            + datetime.timedelta(minutes=random.randint(30, 60)),
                        )
                        solve.date = new_base
                        base_time = new_base

                        db.session.add(solve)
                        db.session.commit()
        elif mode == "teams":
            for x in range(1, TEAM_AMOUNT):
                used_teams = []
                used_users = []
                base_time = datetime.datetime.utcnow() + datetime.timedelta(
                    minutes=-10000
                )
                team = Teams.query.filter_by(id=x).first()
                members_ids = [member.id for member in team.members]
                for y in range(random.randint(1, CHAL_AMOUNT)):
                    chalid = random.randint(1, CHAL_AMOUNT)
                    user_id = random.choice(members_ids)
                    if (chalid, team.id) not in used_teams:
                        if (chalid, user_id) not in used_users:
                            solve = Solves(
                                user_id=user_id,
                                team_id=team.id,
                                challenge_id=chalid,
                                ip="127.0.0.1",
                                provided=gen_word(),
                            )
                            new_base = random_date(
                                base_time,
                                base_time
                                + datetime.timedelta(minutes=random.randint(30, 60)),
                            )
                            solve.date = new_base
                            base_time = new_base
                            db.session.add(solve)
                            db.session.commit()
                            used_teams.append((chalid, team.id))
                            used_users.append((chalid, user_id))

        db.session.commit()

        # Generating Awards
        print("GENERATING AWARDS")
        for x in range(USER_AMOUNT):
            base_time = datetime.datetime.utcnow() + datetime.timedelta(minutes=-10000)
            for _ in range(random.randint(0, AWARDS_AMOUNT)):
                user = Users.query.filter_by(id=x + 1).first()
                award = Awards(
                    user_id=user.id,
                    team_id=user.team_id,
                    name=gen_word(),
                    value=random.randint(-10, 10),
                    icon=gen_icon(),
                )
                new_base = random_date(
                    base_time,
                    base_time + datetime.timedelta(minutes=random.randint(30, 60)),
                )
                award.date = new_base
                base_time = new_base

                db.session.add(award)

        db.session.commit()

        # Generating Wrong Flags
        print("GENERATING WRONG FLAGS")
        for x in range(USER_AMOUNT):
            used = []
            base_time = datetime.datetime.utcnow() + datetime.timedelta(minutes=-10000)
            for y in range(random.randint(1, CHAL_AMOUNT * 20)):
                chalid = random.randint(1, CHAL_AMOUNT)
                if chalid not in used:
                    used.append(chalid)
                    user = Users.query.filter_by(id=x + 1).first()
                    wrong = Fails(
                        user_id=user.id,
                        team_id=user.team_id,
                        challenge_id=chalid,
                        ip="127.0.0.1",
                        provided=gen_word(),
                    )

                    new_base = random_date(
                        base_time,
                        base_time + datetime.timedelta(minutes=random.randint(30, 60)),
                    )
                    wrong.date = new_base
                    base_time = new_base

                    db.session.add(wrong)
                    db.session.commit()

        db.session.commit()
        db.session.close()

        clear_config()
        clear_standings()
        clear_pages()