ctftool

#!/usr/bin/env python3
from database import *
from datetime import datetime, timedelta
import config
import getpass
import hashlib
import os
import os.path
import shutil
import sys
import redis
import random
import utils
import utils.admin
import yaml

tables = [Team, TeamAccess, Challenge, ChallengeSolve, ChallengeFailure, NewsItem, TroubleTicket, TicketComment, Notification, ScoreAdjustment, AdminUser]

operation = sys.argv[1]
if operation == "create-tables":
    [i.create_table() for i in tables]
    print("Tables created")

elif operation == "drop-tables":
    if input("Are you sure? Type yes to continue: ") == "yes":
        [i.drop_table() for i in tables]
        print("Done")
    else:
        print("Okay, nothing happened.")

elif operation == "add-challenge":
    challengefile = sys.argv[2]
    with open(challengefile) as f:
        chal = Challenge.create(**yaml.load(f))
        print("Challenge added with id {}".format(chal.id))

elif operation == "gen-challenge":
    n = int(sys.argv[2])
    for i in range(n):
        name = str(random.randint(0, 999999999))
        chal = Challenge.create(name="Challenge {}".format(name), category="Generated", description="Lorem ipsum, dolor sit amet. The flag is {}".format(name), points=random.randint(50, 400), flag=name, author="autogen")
        print("Challenge added with id {}".format(chal.id))

elif operation == "gen-team":
    n = int(sys.argv[2])
    chals = list(Challenge.select())
    ctz = datetime.now()
    diff = timedelta(minutes=5)
    for i in range(n):
        name = "Team {}".format(i + 1)
        t = Team.create(name=name, email="none@none.com", affiliation="Autogenerated", eligible=True, key="", email_confirmation_key="autogen", email_confirmed=True)
        t.key = "autogen{}".format(t.id)
        t.save()
        print("Team added with id {}".format(t.id))

elif operation == "add-admin":
    username = input("Username: ")
    password = getpass.getpass().encode()
    pwhash = utils.admin.create_password(password)
    r = random.SystemRandom()
    secret = "".join([r.choice("ABCDEFGHIJKLMNOPQRSTUVWXYZ234567") for i in range(16)])
    AdminUser.create(username=username, password=pwhash, secret=secret)
    print("AdminUser created; Enter the following key into your favorite TOTP application (Google Authenticator Recommended): {}".format(secret))

elif operation == "scan":
    path = sys.argv[2]
    dirs = [j for j in [os.path.join(path, i) for i in os.listdir(path)] if os.path.isdir(j)]
    print(dirs)
    n = 0

    for d in dirs:
        staticpaths = {}
        if os.path.exists(os.path.join(d, "static.yml")):
            with open(os.path.join(d, "static.yml")) as f:
                statics = yaml.load(f)
                for static in statics:
                    h = hashlib.sha256()
                    with open(os.path.join(d, static), "rb") as staticfile:
                        while True:
                            buf = staticfile.read(4096)
                            h.update(buf)
                            if not buf:
                                break

                    if "." in static:
                        name, ext = static.split(".", maxsplit=1)
                        fn = "{}_{}.{}".format(name, h.hexdigest(), ext)
                    else:
                        fn = "{}_{}".format(static, h.hexdigest())
                    staticpaths[static] = fn
                    shutil.copy(os.path.join(d, static), os.path.join(config.static_dir, fn))
                    print(fn)

        if os.path.exists(os.path.join(d, "problem.yml")):
            with open(os.path.join(d, "problem.yml")) as f:
                n += 1
                data = yaml.load(f)
                for i in staticpaths:
                    print("looking for |{}|".format(i))
                    data["description"] = data["description"].replace("|{}|".format(i), "{}{}".format(config.static_prefix, staticpaths[i]))

                query = Challenge.select().where(Challenge.name == data["name"])
                if query.exists():
                    print("Updating " + str(data["name"]) + "...")
                    q = Challenge.update(**data).where(Challenge.name == data["name"])
                    q.execute()
                else:
                    Challenge.create(**data)

    print(n, "challenges loaded")

elif operation == "recache-solves":
    r = redis.StrictRedis()
    for chal in Challenge.select():
        r.hset("solves", chal.id, chal.solves.count())
    print(r.hvals("solves"))

elif operation == "list-challenges":
    for chal in Challenge.select():
        print("{}  {}".format(str(chal.id).rjust(3), chal.name))

elif operation == "del-challenge":
    id = sys.argv[2]
    c = Challenge.get(id=id)
    ChallengeFailure.delete().where(ChallengeFailure.challenge == c).execute()
    ChallengeSolve.delete().where(ChallengeSolve.challenge == c).execute()
    c.delete_instance()

# vim: syntax=python:ft=python