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, team
from nizkctf.repohost import RepoHost
from nizkctf.subrepo import SubRepo
from nizkctf.challenge import Challenge, derive_keypair, random_salt
from nizkctf.six import to_unicode, PY2


def read_opt(msg, opts):
    while True:
        inp = raw_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 = raw_input('Username: ')
            password = getpass.getpass('Password: ')
            RepoHost.login(username=username, password=password)
            return True
        elif opt == '2':
            print('Enter your auth token:')
            token = raw_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 = raw_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_add_challenge(args):
    id = raw_input('Challenge id (digits, letters, underscore): ').strip()
    title = raw_input('Title: ').strip()
    description = raw_input('Description: ').strip()
    points = int(raw_input('Points: ').strip())
    tags = raw_input('Tags (separate tags with space): ').strip().split()
    salt = raw_input('Salt (empty string for random salt): ').strip()
    if salt == '':
        salt = random_salt()
    flag = raw_input('Flag: ').strip()

    pk, sk = derive_keypair(salt, flag)
    chall = Challenge(id=id)
    chall['id'] = id
    chall['title'] = title
    chall['description'] = description
    chall['points'] = points
    chall['tags'] = tags
    chall['salt'] = salt
    chall['pk'] = pk
    chall.save()


def main():
    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,
        'add': cmd_add_challenge,
    }

    parser = argparse.ArgumentParser(description='nizk CTF cli')
    subparsers = parser.add_subparsers(help='command help')

    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_add_challenge = subparsers.add_parser('add', help='add challenge')
    parser_add_challenge.set_defaults(command='add')

    args = parser.parse_args()
    commands[args.command](args)

if __name__ == '__main__':
    main()