ctf

#! /usr/bin/env python3
import os
import sys
import glob
import argparse
import pysodium
import traceback
from pathlib import Path


LANGS = ['en', 'pt']


def ensure_unicode_locale():
    if sys.stdout.encoding != 'utf-8':
        print('\033[1mWARNING\033[00m: This CTF accepts '
              'international characters in team names and\nproblem '
              'descriptions. You are currently not using a Unicode '
              'locale,\ntherefore you may experience random '
              'UnicodeEncodeError exceptions.\n\nPlease fix by '
              'changing to a Unicode locale, e.g.\n\n'
              '  export LC_ALL=en_US.UTF-8\n\n\n',
              file=sys.stderr)


def cmd_init(args):
    from nizkctf.repo import contents, trail

    print('Cloning Git repositories', file=sys.stderr)
    contents.clone()
    trail.clone()
    print('We are all set!', file=sys.stderr)


def cmd_audit(args):
    from nizkctf.audit import audit
    from nizkctf.cli.audit import audit_record

    if args.record:
        audit_record()
    else:
        audit()
    print('OK')


def cmd_add_challenge(args):
    from nizkctf.challenge import Challenge, \
        chall_dir, derive_keypair, random_salt

    id = input('Challenge id (digits, letters, underscore): ').strip()
    title = input('Title: ').strip()
    tags = input('Tags (separate tags with space): ').strip().split()
    salt = random_salt()
    while True:
        flag = input('Flag: ').strip()
        if flag.startswith('CTF-BR{') and flag.endswith('}'):
            break
        print('Please respect flag format: CTF-BR{...}')
    while True:
        level = input(
            'Pwhash level (interactive, moderate, sensitive): ').strip().upper()
        if level == 'INTERACTIVE' and len(flag) < 24:
            print('Please use moderate or sensitive level for short flags')
            continue
        try:
            opslimit = getattr(
                pysodium, 'crypto_pwhash_argon2id_OPSLIMIT_'+level)
            memlimit = getattr(
                pysodium, 'crypto_pwhash_argon2id_MEMLIMIT_'+level)
            break
        except:
            traceback.print_exc()

    pk, _ = derive_keypair(salt, opslimit, memlimit, flag)
    chall = Challenge(id=id)
    chall['id'] = id
    chall['title'] = title
    chall['tags'] = tags
    chall['salt'] = salt
    chall['opslimit'] = opslimit
    chall['memlimit'] = memlimit
    chall['pk'] = pk
    chall.save()

    print('\nPlease add challenge description to:')
    for lang in LANGS:
        filename = os.path.join(chall_dir, '{}.{}.md'.format(id, lang))
        Path(filename).touch()
        print(' - {}'.format(filename))
    print('\nto make it public, call:')
    print('./ctf publish_chall {}'.format(id))


def cmd_publish_chall(args):
    from nizkctf.repo import contents
    from nizkctf.challenge import ChallengeIndex, chall_dir

    if args.chall == []:
        for filename in glob.glob(os.path.join(chall_dir, '*.json')):
            chall_id, _ = os.path.splitext(os.path.basename(filename))
            if chall_id == 'index':
                continue
            print(' - {}'.format(chall_id))
            args.chall.append(chall_id)
        ans = input('Going to publish all challenges, are you sure? (y/n) ')
        if ans != 'y':
            print('Aborted')
            return

    index = ChallengeIndex()
    for chall_id in args.chall:
        files = ['{}.json'.format(chall_id)]
        files += ['{}.{}.md'.format(chall_id, lang) for lang in LANGS]
        files = [os.path.join(chall_dir, filename) for filename in files]
        contents.add(files)
        if chall_id not in index:
            index.append(chall_id)
    index.save()
    contents.add([index.path()])

    contents.push(commit_message='Publish challenges')


def cmd_add_news(args):
    from nizkctf.cli import news

    news.submit(args.msg)


def main():
    ensure_unicode_locale()

    commands = {
        'init': cmd_init,
        'audit': cmd_audit,
        'add_chall': cmd_add_challenge,
        'publish_chall': cmd_publish_chall,
        'add_news': cmd_add_news,
    }

    parser = argparse.ArgumentParser(description='NIZKCTF admin toolkit')
    subparsers = parser.add_subparsers(help='command help',
                                       metavar='{init,audit,add_chall,'
                                       'publish_chall,add_news}')

    parser_init = subparsers.add_parser('init', help='init ctf environment')
    parser_init.set_defaults(command='init')

    parser_audit = subparsers.add_parser('audit', help='audit ctf scoreboard')
    parser_audit.set_defaults(command='audit', record=False)
    parser_audit.add_argument(
        '--record', action='store_true', help='record audit trail to git')

    parser_add_challenge = subparsers.add_parser(
        'add_chall', help='add challenge')
    parser_add_challenge.set_defaults(command='add_chall')

    parser_publish_challenge = subparsers.add_parser(
        'publish_chall', help='publish challenge(s)')
    parser_publish_challenge.set_defaults(command='publish_chall')
    parser_publish_challenge.add_argument(
        'chall', nargs='*', help='id of challenges (empty for all)')

    parser_add_news = subparsers.add_parser('add_news', help='add news')
    parser_add_news.set_defaults(command='add_news')
    parser_add_news.add_argument('msg', help='msg to be added')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    args = parser.parse_args()
    commands[args.command](args)


if __name__ == '__main__':
    main()