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

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: ')
            RepoHost.login(username=username, password=password)
            return True
        elif opt == '2':
            print('Enter your auth token:')
            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 cmd_register(args):
    team.register(args.name)


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')
    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('>> ')
        team.register(team_name)
    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()
    flag = input('Flag: ').strip()

    pk, sk = derive_keypair(salt, 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['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 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_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()