CTFd/tests/helpers.py

import datetime
import gc
import random
import string
import uuid
from collections import namedtuple
from unittest.mock import Mock, patch

import requests
from flask.testing import FlaskClient
from sqlalchemy.engine.url import make_url
from sqlalchemy_utils import drop_database
from werkzeug.datastructures import Headers

from CTFd import create_app
from CTFd.cache import cache, clear_standings
from CTFd.config import TestingConfig
from CTFd.models import (
    Awards,
    ChallengeComments,
    ChallengeFiles,
    Challenges,
    Comments,
    Fails,
    Fields,
    Files,
    Flags,
    Hints,
    Notifications,
    PageComments,
    PageFiles,
    Pages,
    Solves,
    Tags,
    TeamComments,
    Teams,
    Tokens,
    Tracking,
    Unlocks,
    UserComments,
    Users,
)

text_type = str
binary_type = bytes


FakeRequest = namedtuple("FakeRequest", ["form"])


class CTFdTestClient(FlaskClient):
    def open(self, *args, **kwargs):
        if kwargs.get("json") is not None:
            with self.session_transaction() as sess:
                api_key_headers = Headers({"CSRF-Token": sess.get("nonce")})
                headers = kwargs.pop("headers", Headers())
                if isinstance(headers, dict):
                    headers = Headers(headers)
                headers.extend(api_key_headers)
                kwargs["headers"] = headers
        return super(CTFdTestClient, self).open(*args, **kwargs)


def create_ctfd(
    ctf_name="CTFd",
    ctf_description="CTF description",
    name="admin",
    email="admin@ctfd.io",
    password="password",
    user_mode="users",
    setup=True,
    enable_plugins=False,
    application_root="/",
    config=TestingConfig,
):
    if enable_plugins:
        config.SAFE_MODE = False
    else:
        config.SAFE_MODE = True

    config.APPLICATION_ROOT = application_root
    url = make_url(config.SQLALCHEMY_DATABASE_URI)
    if url.database:
        url.database = str(uuid.uuid4())
    config.SQLALCHEMY_DATABASE_URI = str(url)

    app = create_app(config)
    app.test_client_class = CTFdTestClient

    if setup:
        app = setup_ctfd(
            app,
            ctf_name=ctf_name,
            ctf_description=ctf_description,
            name=name,
            email=email,
            password=password,
            user_mode=user_mode,
        )
    return app


def setup_ctfd(
    app,
    ctf_name="CTFd",
    ctf_description="CTF description",
    name="admin",
    email="admin@ctfd.io",
    password="password",
    user_mode="users",
):
    with app.app_context():
        with app.test_client() as client:
            client.get("/setup")  # Populate session with nonce
            with client.session_transaction() as sess:
                data = {
                    "ctf_name": ctf_name,
                    "ctf_description": ctf_description,
                    "name": name,
                    "email": email,
                    "password": password,
                    "user_mode": user_mode,
                    "nonce": sess.get("nonce"),
                }
            client.post("/setup", data=data)
    return app


def destroy_ctfd(app):
    with app.app_context():
        gc.collect()  # Garbage collect (necessary in the case of dataset freezes to clean database connections)
        cache.clear()
        drop_database(app.config["SQLALCHEMY_DATABASE_URI"])


def register_user(
    app, name="user", email="user@ctfd.io", password="password", raise_for_error=True
):
    with app.app_context():
        with app.test_client() as client:
            client.get("/register")
            with client.session_transaction() as sess:
                data = {
                    "name": name,
                    "email": email,
                    "password": password,
                    "nonce": sess.get("nonce"),
                }
            client.post("/register", data=data)
            if raise_for_error:
                with client.session_transaction() as sess:
                    assert sess["id"]
                    assert sess["nonce"]
                    assert sess["hash"]


def register_team(app, name="team", password="password", raise_for_error=True):
    with app.app_context():
        with app.test_client() as client:
            client.get("/team")
            with client.session_transaction() as sess:
                data = {"name": name, "password": password, "nonce": sess.get("nonce")}
            r = client.post("/teams/new", data=data)
            if raise_for_error:
                assert r.status_code == 302
            return client


def login_as_user(app, name="user", password="password", raise_for_error=True):
    with app.app_context():
        with app.test_client() as client:
            client.get("/login")
            with client.session_transaction() as sess:
                data = {"name": name, "password": password, "nonce": sess.get("nonce")}
            client.post("/login", data=data)
            if raise_for_error:
                with client.session_transaction() as sess:
                    assert sess["id"]
                    assert sess["nonce"]
                    assert sess["hash"]
            return client


def login_with_mlc(
    app,
    name="user",
    scope="profile%20team",
    email="user@ctfd.io",
    oauth_id=1337,
    team_name="TestTeam",
    team_oauth_id=1234,
    raise_for_error=True,
):
    with app.test_client() as client, patch.object(
        requests, "get"
    ) as fake_get_request, patch.object(requests, "post") as fake_post_request:
        client.get("/login")
        with client.session_transaction() as sess:
            nonce = sess["nonce"]

            redirect_url = "{endpoint}?response_type=code&client_id={client_id}&scope={scope}&state={state}".format(
                endpoint=app.config["OAUTH_AUTHORIZATION_ENDPOINT"],
                client_id=app.config["OAUTH_CLIENT_ID"],
                scope=scope,
                state=nonce,
            )

        r = client.get("/oauth", follow_redirects=False)
        assert r.location == redirect_url

        fake_post_response = Mock()
        fake_post_request.return_value = fake_post_response
        fake_post_response.status_code = 200
        fake_post_response.json = lambda: {"access_token": "fake_mlc_access_token"}

        fake_get_response = Mock()
        fake_get_request.return_value = fake_get_response
        fake_get_response.status_code = 200
        fake_get_response.json = lambda: {
            "id": oauth_id,
            "name": name,
            "email": email,
            "team": {"id": team_oauth_id, "name": team_name},
        }

        client.get(
            "/redirect?code={code}&state={state}".format(
                code="mlc_test_code", state=nonce
            ),
            follow_redirects=False,
        )

        if raise_for_error:
            with client.session_transaction() as sess:
                assert sess["id"]
                assert sess["nonce"]
                assert sess["hash"]
        return client


def get_scores(user):
    r = user.get("/api/v1/scoreboard")
    scores = r.get_json()
    return scores["data"]


def random_string(n=5):
    return "".join(
        random.choice(string.ascii_letters + string.digits) for _ in range(n)
    )


def random_int(start=2147483647, stop=None, step=1):
    return random.randrange(start, stop, step)


def gen_challenge(
    db,
    name="chal_name",
    description="chal_description",
    value=100,
    category="chal_category",
    type="standard",
    state="visible",
    **kwargs
):
    chal = Challenges(
        name=name,
        description=description,
        value=value,
        category=category,
        type=type,
        state=state,
        **kwargs
    )
    db.session.add(chal)
    db.session.commit()
    return chal


def gen_award(db, user_id, team_id=None, name="award_name", value=100):
    award = Awards(user_id=user_id, team_id=team_id, name=name, value=value)
    award.date = datetime.datetime.utcnow()
    db.session.add(award)
    db.session.commit()
    clear_standings()
    return award


def gen_tag(db, challenge_id, value="tag_tag", **kwargs):
    tag = Tags(challenge_id=challenge_id, value=value, **kwargs)
    db.session.add(tag)
    db.session.commit()
    return tag


def gen_file(db, location, challenge_id=None, page_id=None):
    if challenge_id:
        f = ChallengeFiles(challenge_id=challenge_id, location=location)
    elif page_id:
        f = PageFiles(page_id=page_id, location=location)
    else:
        f = Files(location=location)
    db.session.add(f)
    db.session.commit()
    return f


def gen_flag(db, challenge_id, content="flag", type="static", data=None, **kwargs):
    flag = Flags(challenge_id=challenge_id, content=content, type=type, **kwargs)
    if data:
        flag.data = data
    db.session.add(flag)
    db.session.commit()
    return flag


def gen_user(db, name="user_name", email="user@ctfd.io", password="password", **kwargs):
    user = Users(name=name, email=email, password=password, **kwargs)
    db.session.add(user)
    db.session.commit()
    return user


def gen_team(
    db,
    name="team_name",
    email="team@ctfd.io",
    password="password",
    member_count=4,
    **kwargs
):
    team = Teams(name=name, email=email, password=password, **kwargs)
    for i in range(member_count):
        name = "user-{}-{}".format(random_string(), str(i))
        user = gen_user(db, name=name, email=name + "@ctfd.io", team_id=team.id)
        if i == 0:
            team.captain_id = user.id
        team.members.append(user)
    db.session.add(team)
    db.session.commit()
    return team


def gen_hint(
    db, challenge_id, content="This is a hint", cost=0, type="standard", **kwargs
):
    hint = Hints(
        challenge_id=challenge_id, content=content, cost=cost, type=type, **kwargs
    )
    db.session.add(hint)
    db.session.commit()
    return hint


def gen_unlock(db, user_id, team_id=None, target=None, type="hints"):
    unlock = Unlocks(user_id=user_id, team_id=team_id, target=target, type=type)
    db.session.add(unlock)
    db.session.commit()
    return unlock


def gen_solve(
    db,
    user_id,
    team_id=None,
    challenge_id=None,
    ip="127.0.0.1",
    provided="rightkey",
    **kwargs
):
    solve = Solves(
        user_id=user_id,
        team_id=team_id,
        challenge_id=challenge_id,
        ip=ip,
        provided=provided,
        **kwargs
    )
    solve.date = datetime.datetime.utcnow()
    db.session.add(solve)
    db.session.commit()
    clear_standings()
    return solve


def gen_fail(
    db,
    user_id,
    team_id=None,
    challenge_id=None,
    ip="127.0.0.1",
    provided="wrongkey",
    **kwargs
):
    fail = Fails(
        user_id=user_id,
        team_id=team_id,
        challenge_id=challenge_id,
        ip=ip,
        provided=provided,
        **kwargs
    )
    fail.date = datetime.datetime.utcnow()
    db.session.add(fail)
    db.session.commit()
    return fail


def gen_tracking(db, user_id=None, ip="127.0.0.1", **kwargs):
    tracking = Tracking(ip=ip, user_id=user_id, **kwargs)
    db.session.add(tracking)
    db.session.commit()
    return tracking


def gen_page(db, title, route, content, draft=False, auth_required=False, **kwargs):
    page = Pages(
        title=title,
        route=route,
        content=content,
        draft=draft,
        auth_required=auth_required,
        **kwargs
    )
    db.session.add(page)
    db.session.commit()
    return page


def gen_notification(db, title="title", content="content"):
    notif = Notifications(title=title, content=content)
    db.session.add(notif)
    db.session.commit()


def gen_token(db, type="user", user_id=None, expiration=None):
    token = Tokens(type=type, user_id=user_id, expiration=expiration)
    db.session.add(token)
    db.session.commit()
    return token


def gen_comment(db, content="comment", author_id=None, type="challenge", **kwargs):
    if type == "challenge":
        model = ChallengeComments
    elif type == "user":
        model = UserComments
    elif type == "team":
        model = TeamComments
    elif type == "page":
        model = PageComments
    else:
        model = Comments

    comment = model(content=content, author_id=author_id, type=type, **kwargs)
    db.session.add(comment)
    db.session.commit()
    return comment


def gen_field(
    db,
    name="CustomField",
    type="user",
    field_type="text",
    description="CustomFieldDescription",
    required=True,
    public=True,
    editable=True,
):
    field = Fields(
        name=name,
        type=type,
        field_type=field_type,
        description=description,
        required=required,
        public=public,
        editable=editable,
    )
    db.session.add(field)
    db.session.commit()
    return field


def simulate_user_activity(db, user):
    gen_tracking(db, user_id=user.id)
    gen_award(db, user_id=user.id)
    challenge = gen_challenge(db)
    flag = gen_flag(db, challenge_id=challenge.id)
    hint = gen_hint(db, challenge_id=challenge.id)

    for _ in range(5):
        gen_fail(db, user_id=user.id, challenge_id=challenge.id)

    gen_unlock(db, user_id=user.id, target=hint.id, type="hints")
    gen_solve(db, user_id=user.id, challenge_id=challenge.id, provided=flag.content)