ColdCore/ctftool

130 lines
4.9 KiB
Python
Executable File

#!/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