ctf

#! /usr/bin/env python
# -*- encoding: utf-8 -*-

from __future__ import unicode_literals, division, print_function,\
                       absolute_import
import os
import sys
import codecs
import argparse
import getpass
import base64
import pysodium
import traceback

from nizkctf.cli import scoreboard, challenges, news, team, localserver
from nizkctf.repohost import RepoHost
from nizkctf.subrepo import SubRepo
from nizkctf.challenge import Challenge, derive_keypair, random_salt
from nizkctf.six import input, to_unicode, PY2
from nizkctf.settings import Settings


def ensure_unicode_locale():
    if sys.stdout.encoding != 'UTF-8':
        sys.stderr.write('\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')


def read_opt(msg, opts):
    while True:
        inp = input(msg)
        if inp.strip() in opts:
            return inp.strip()


def read_auth(opt):
    while True:
        if opt == '1':
            print('Enter your username and password:')
            username = input('Username: ')
            password = getpass.getpass('Password: ')
            OTP = input('OTP (only if you use two-factor authentication): ')
            if OTP != "":
                RepoHost.login(username=username, password=password, OTP=OTP)
            else:
                RepoHost.login(username=username, password=password)
            return True
        elif opt == '2':
            print('Enter your auth token (with public_repo scope):')
            token = input('Token: ')
            RepoHost.login(token=token)
            return True


def cmd_scoreboard(args):
    if args.pull:
        SubRepo.pull()

    ranking, submissions = scoreboard.rank()
    scoreboard.pprint(ranking, top=args.top, show_names=args.names)

    if args.chart:
        scoreboard.plot(ranking, submissions)


def parse_countries(countries):
    if countries == '':
        return []
    return [c.lower().strip() for c in countries.split(',')]


def cmd_register(args):
    team.register(args.name, parse_countries(args.countries))


def cmd_login(args):
    if args.token:
        RepoHost.login(token=args.token)
    elif args.username and args.password:
        RepoHost.login(username=args.username, password=args.password)
    else:
        read_auth('1')
    print('Credentials stored')
    print('Cloning submissions repository')
    SubRepo.clone()


def cmd_init(args):
    print('NIZKCTF initializing your environment.')
    print('First of all, we need your github/gitlab credentials:')
    print('[1] auth via username / password')
    print('[2] auth with personal access token (with public_repo scope)')
    opt = read_opt('>> ', {'1', '2'})
    print('')

    read_auth(opt)
    print('')

    print('Cloning submissions repository')
    SubRepo.clone()

    print('Do you want to register a new team? [y/n]')
    opt = read_opt('>> ', {'y', 'n'})
    if opt == 'y':
        print('Enter your team name:')
        team_name = input('>> ')
        print("Enter a comma-separated list of your member's countries.")
        print("Use 2-letter ISO country codes: http://flag-icon-css.lip.is")
        print("Example: For a team composed by brazilians and russians: BR,RU")
        print("(leave empty if you want to omit nationality)")
        team_countries = parse_countries(input('>> '))
        team.register(team_name, team_countries)
    print('')

    print('We are all set!')


def cmd_submit(args):
    print('Checking flag: %s' % args.flag)
    ret, msg = challenges.submit_flag(args.flag, args.chall)
    print(msg)
    if not ret:
        sys.exit(1)


def cmd_challenges(args):
    challenges.pprint()


def cmd_serve(args):
    localserver.main(port=args.port)


def cmd_add_challenge(args):
    id = input('Challenge id (digits, letters, underscore): ').strip()
    title = input('Title: ').strip()
    description = input('Description: ').strip()
    if not Settings.dynamic_scoring:
        points = int(input('Points: ').strip())
    tags = input('Tags (separate tags with space): ').strip().split()
    salt = input('Salt (empty string for random salt): ').strip()
    if salt == '':
        salt = random_salt()
    while True:
        level = input('Pwhash level (interactive, moderate, sensitive): ').strip().upper()
        try:
            opslimit = getattr(pysodium, 'crypto_pwhash_argon2id_OPSLIMIT_'+level)
            memlimit = getattr(pysodium, 'crypto_pwhash_argon2id_MEMLIMIT_'+level)
            break
        except:
            traceback.print_exc()
    flag = input('Flag: ').strip()

    pk, sk = derive_keypair(salt, opslimit, memlimit, flag)
    chall = Challenge(id=id)
    chall['id'] = id
    chall['title'] = title
    chall['description'] = description
    if not Settings.dynamic_scoring:
        chall['points'] = points
    chall['tags'] = tags
    chall['salt'] = salt
    chall['opslimit'] = opslimit
    chall['memlimit'] = memlimit
    chall['pk'] = pk
    chall.save()


def cmd_news(args):
    if args.pull:
        SubRepo.pull()
    news.pprint(news.News(), team_only=args.team)


def cmd_add_news(args):
    news.submit(args.msg, args.to)


def main():
    ensure_unicode_locale()
    if PY2:
        sys.argv = map(to_unicode, sys.argv)

    commands = {
        'init': cmd_init,
        'login': cmd_login,
        'score': cmd_scoreboard,
        'register': cmd_register,
        'submit': cmd_submit,
        'challs': cmd_challenges,
        'serve': cmd_serve,

        'add': cmd_add_challenge,

        'news': cmd_news,
        'add_news': cmd_add_news,
    }

    parser = argparse.ArgumentParser(description='nizk CTF cli')
    subparsers = parser.add_subparsers(help='command help',
                                       metavar='{init,login,score,register,'
                                               'submit,challs,serve,news}')

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

    parser_login = subparsers.add_parser('login',
                                         help='authenticate in gitlab/github')
    parser_login.set_defaults(command='login')
    parser_login.add_argument('--username', type=str, default=None,
                              metavar='USERNAME',
                              help='username for logging in')
    parser_login.add_argument('--password', type=str, default=None,
                              metavar='PASSWORD',
                              help='password for logging in')
    parser_login.add_argument('--token', type=str, default=None,
                              metavar='TOKEN',
                              help='use personal access token (with public_repo scope) instead of '
                                   'username/password')

    parser_score = subparsers.add_parser('score', help='scoreboard help')
    parser_score.set_defaults(command='score')
    parser_score.add_argument('--top', type=int,
                              default=0, metavar='N',
                              help='size of ranking to display')
    parser_score.add_argument('--pull', action='store_true',
                              help='pull submissions before displaying scores')
    parser_score.add_argument('--names', action='store_true',
                              help='display team names')
    parser_score.add_argument('--chart', action='store_true',
                              help='display chart')

    parser_register = subparsers.add_parser('register',
                                            help='register a new team')
    parser_register.set_defaults(command='register')
    parser_register.add_argument('name', metavar='NAME', help='team name')
    parser_register.add_argument('countries', metavar='COUNTRIES',
                                 help='Comma-separated list of 2-letter ISO '
                                      'country codes (pass an empty string to '
                                      'omit nationality)')

    parser_submit = subparsers.add_parser('submit', help='submit a flag')
    parser_submit.set_defaults(command='submit')
    parser_submit.add_argument('flag', metavar='FLAG', help='flag')
    parser_submit.add_argument('--chall', type=str, default=None,
                               metavar='CHALL_ID', help='challenge id')

    parser_challenges = subparsers.add_parser('challs', help='list challenges')
    parser_challenges.set_defaults(command='challs')

    parser_serve = subparsers.add_parser('serve', help='start a local server')
    parser_serve.set_defaults(command='serve')
    parser_serve.add_argument('--port', type=int, default=8000,
                              metavar='PORT', help='port')

    parser_add_challenge = subparsers.add_parser('add')
    parser_add_challenge.set_defaults(command='add')

    parser_news = subparsers.add_parser('news', help='pull and show news')
    parser_news.set_defaults(command='news')
    parser_news.add_argument('--pull', action='store_true',
                             help='pull news before displaying news')
    parser_news.add_argument('--team', action='store_true',
                             help='show only team news')

    parser_add_news = subparsers.add_parser('add_news')
    parser_add_news.set_defaults(command='add_news')
    parser_add_news.add_argument('--msg', required=True,
                                 help='msg to be added')
    parser_add_news.add_argument('--to', help='team name')

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


if __name__ == '__main__':
    main()