#!/usr/bin/env python2
'''
pCTF 2012 CLI interface

Email bug reports to plaid.parliament.of.pwning@gmail.com.
'''
import sys
import getpass
import json
import cookielib
import urllib
import urllib2
import urlparse
import textwrap
import getpass

try:
    import readline

    def completer(text, state):
        commands = [
            'help',
            'hints',
            'problems',
            'score',
            'show',
            'submit',
            'vote_clear',
            'vote_down',
            'vote_show',
            'vote_summary',
            'vote_up',
            'quit',
        ]

        options = [x for x in commands if x.startswith(text)]
        if state < len(options):
            return options[state]
        else:
            return None

    readline.set_completer(completer)
    readline.parse_and_bind("tab: complete")
except ImportError:
    pass

class PCTFClient(object):
    def __init__(self, username=None, password=None,
                 base_url='http://ctf.plaidctf.com'):
        self._cookiejar = cookielib.CookieJar()
        self._username = username
        self._password = password
        self._base_url = base_url
        self._opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(self._cookiejar))
        self._auth = None
        self._csrftoken = None

    def login(self):
        # Get CSRF token
        url = urlparse.urljoin(self._base_url, '/dropdowns/login')
        self._opener.open(url)

        for c in self._cookiejar:
            if c.name == 'csrftoken':
                self._csrftoken = c.value

        url = urlparse.urljoin(self._base_url, '/teams/login')
        form_data = {
            'username': self._username,
            'password': self._password,
            'csrfmiddlewaretoken': self._csrftoken
        }
        try:
            response = self._opener.open(url, urllib.urlencode(form_data))
        except urllib2.URLError, reason:
            return False

        if 'Please try again.' in response.read():
            return False

        for c in self._cookiejar:
            if c.name == 'auth':
                self._auth = c.value

        return True

    def request(self, path, data):
        url = urlparse.urljoin(self._base_url, path)

        if data is None:
            req = urllib2.Request(url)
        else:
            form_data = {
                'data': json.dumps(data),
                'csrfmiddlewaretoken': self._csrftoken
            }
            req = urllib2.Request(url, urllib.urlencode(form_data))

        req.add_header('X-Auth', self._auth)
        response = self._opener.open(req)
        return response.read()

    def getteam(self):
        data = self.request('/teams/my_team', None)
        my_team = json.loads(data)
        if my_team['score'] is None:
            my_team['score'] = 0
        return my_team

    def clean(self, s):
        s = s.replace('<br />', '\n\n')
        s = s.replace('<br/>', '\n\n')
        s = s.replace('<br>', '\n\n')
        s = s.replace('&lt;', '<')
        s = s.replace('&gt;', '>')
        return '\n'.join(textwrap.wrap(s))

    def getproblems(self, unsolved_only=True):
        if unsolved_only:
            data = self.request('/problems/list?unsolved_only', None)
        else:
            data = self.request('/problems/list', None)
        problems = json.loads(data)

        for p in problems:
            p['description'] = self.clean(p['description'])

        return problems

    def clean_hints(self, s):
        # silly split and join to get rid of <script>
        s = '\n'.join(s.splitlines()[5:])
        s = s.replace('<br />', '')
        # silly hack to *emphasize* <h2>
        s = s.replace('<h2>', '*')
        s = s.replace('</h2>', '*')
        s = s.replace('<strong>', '')
        s = s.replace('</strong>', '')
        return s.strip()

    def gethints(self):
        hints = self.request('/hints', None)
        return self.clean_hints(hints)

    def submit(self, problem_id, key):
        data = {
            'key': key,
            'problem': problem_id,
        }

        try:
            self.request('/problems/submit', data)
        except urllib2.HTTPError:
            return False

        return True

    def vote_summary(self):
        try:
            data = self.request('/problems/vote_summary', None)
            return json.loads(data)
        except urllib2.HTTPError:
            return None

    def vote(self, vote=None):
        try:
            data = self.request('/problems/vote', vote)
            return json.loads(data)
        except urllib2.HTTPError:
            return None

def print_ssh(my_team):
    print '''
SSH username: %(ssh_user)s
SSH password: %(ssh_pass)s
'''.strip() % my_team

def print_problems(problems):
    for p in problems:
        print '%(id)d: %(name)s (%(category)s %(value)d)' % p

def print_problem(problem):
    print '''
id: %(id)d
name: %(name)s
points: %(value)d
category: %(category)s

%(description)s
'''.strip() % problem

def print_vote(vote):
    for k, v in vote.iteritems():
        if v:
            vote = 1
        else:
            vote = -1

        print '%3s %3d' % (k, vote)

def print_vote_summary(vote_summary):
    for p in vote_summary['problems']:
        votes = '+%(upvotes)d/-%(downvotes)d' % p
        out = '%(id)2d' % p
        out += '%8s   ' % votes
        out += '%(name)-24s   (%(category)s %(value)s)' % p
        print out
    print
    print vote_summary['minutes'], 'minutes remaining to vote'

def menu(client):
    while True:
        try:
            command = raw_input('> ')
            command = command.strip().split()

            if not command:
                continue

            if command[0] == 'score':
                my_team = client.getteam()
                print my_team['score']

            elif command[0] == 'ssh':
                my_team = client.getteam()
                print_ssh(my_team)

            elif command[0] == 'problems':
                unsolved_only = True
                if len(command) > 1 and command[1] == 'all':
                    unsolved_only = False
                print_problems(client.getproblems(unsolved_only))

            elif command[0] == 'show':
                if len(command) < 2:
                    print 'Usage: show problem_id'
                    continue

                found = False
                problems = client.getproblems(unsolved_only=False)
                for p in problems:
                    if str(p['id']) == command[1]:
                        print_problem(p)
                        found = True
                        break
                if not found:
                    print 'Problem not found.'

            elif command[0] == 'submit':
                if len(command) < 2:
                    print 'Usage: submit problem_id'
                    print '(you will be prompted for the key)'
                    continue

                try:
                    problem_id = int(command[1])
                except ValueError:
                    print 'Invalid problem id.'
                    continue

                key = raw_input('Key: ')
                if client.submit(problem_id, key):
                    print 'Key accepted!'
                else:
                    print 'Incorrect key.'

            elif command[0] == 'vote_up':
                if len(command) < 2:
                    print 'Usage: vote_up problem_id'
                    print '(you will be prompted for the key)'
                    continue

                problem_id = command[1]
                try:
                    int(problem_id)
                except ValueError:
                    print 'Invalid problem id.'
                    continue

                my_votes = client.vote()
                if problem_id in my_votes:
                    if not my_votes[problem_id]:
                        del my_votes[problem_id]
                    else:
                        print 'You have already upvoted for this problem.'
                else:
                    my_votes[problem_id] = True

                if client.vote(my_votes) is None:
                    print 'Placing vote failed.'
                    continue

                print 'Vote placed.'

            elif command[0] == 'vote_down':
                if len(command) < 2:
                    print 'Usage: vote_down problem_id'
                    print '(you will be prompted for the key)'
                    continue

                problem_id = command[1]
                try:
                    int(problem_id)
                except ValueError:
                    print 'Invalid problem id.'
                    continue

                my_votes = client.vote()
                if problem_id in my_votes:
                    if my_votes[problem_id]:
                        del my_votes[problem_id]
                    else:
                        print 'You have already downvoted for this problem.'
                else:
                    my_votes[problem_id] = False

                if client.vote(my_votes) is None:
                    print 'Placing vote failed.'
                    continue

                print 'Vote placed.'

            elif command[0] == 'vote_clear':
                if client.vote({}) is None:
                    print 'Clearing vote failed.'
                    continue

                print 'Vote cleared.'

            elif command[0] == 'vote_show':
                my_votes = client.vote()
                if my_votes is None:
                    print 'It is not time to vote.'
                else:
                    print_vote(my_votes)

            elif command[0] == 'vote_show':
                my_votes = client.vote()
                print_vote(my_votes)

            elif command[0] == 'vote_summary':
                vote_summary = client.vote_summary()
                if not vote_summary:
                    print 'It is not time to vote.'
                else:
                    print_vote_summary(vote_summary)

            elif command[0] == 'hints':
                hints = client.gethints()
                print hints

            elif command[0] == 'help':
                print '''
help                     Show this help.
hints                    Show the list of hints.
problems                 List all unsolved problems.
problems all             List all problems.
score                    Get your team score.
show [problem_id]        Show problem.
submit [problem_id]      Submit key for problem (will prompt you for the key).
vote_summary             Show status of problem voting.
vote_show                Show your vote.
vote_clear               Clear your vote.
vote_up [problem_id]     Upvote problem.
vote_down [problem_id]   Downvote problem.
quit                     Quit.
'''.strip()
            else:
                print 'Unknown command.'

        except EOFError:
            break

if __name__ == '__main__':
    print '''
pCTF 2012 CLI interface

You can use this interface to play Plaid CTF 2012.
On machines with readline, you can tab-complete commands.
'''.strip()

    print
    username = raw_input('Username: ').strip()
    password = getpass.getpass('Password: ').strip()

    client = PCTFClient(username, password)
    if not client.login():
        print 'Login failed.'
        sys.exit(1)

    print '''
Login succeeded.

For help, type help.
'''.strip()
    menu(client)
    sys.exit(0)