satori

#!/usr/bin/env python2.7
# -*- mode: python -*-
import argparse
import getpass
import json
import os
import io
import requests
import webbrowser
import random
import sys
import time
import subprocess

from pyquery import PyQuery as pq
from lxml.html import etree

__version__ = '0.1'

BASE = 'https://satori.tcs.uj.edu.pl'
LOG_REQUESTS = True

def parse_html(html):
    parser = etree.HTMLParser()
    tree = etree.parse(io.BytesIO(html), parser)
    return tree

class SatoriError(Exception):
    __module__ = 'satori'

def match_code(query, code):
    query = query.lower()
    code = code.lower().encode('utf8')
    return query == code or query + '*' == code

class Cache(object):
    def __init__(self, path):
        self.path = path
        self.data = None

    def load(self):
        if self.data is not None:
            return
        try:
            with open(self.path) as f:
                self.data = json.load(f)
        except IOError:
            self.data = {}

    def save(self):
        dir = os.path.dirname(self.path)
        if not os.path.exists(dir):
            os.mkdir(dir)
        suf = '.tmp%d' % random.randrange(10000)
        with open(self.path + suf, 'w') as f:
            json.dump(self.data, f)
        os.rename(self.path + suf, self.path)

    def __setitem__(self, key, value):
        self.load()
        self.data[key] = value
        self.save()

    def __getitem__(self, key):
        self.load()
        return self.data[key]

    def __contains__(self, key):
        self.load()
        return key in self.data

    def get(self, key, default=None):
        self.load()
        return self.data.get(key, default)

def cached(cache_name):
    def wrapper1(func):
        def wrapper(self, *args):
            key = repr(args)
            cache = getattr(self, cache_name)
            if key not in cache:
                cache[key] = func(self, *args)
            return cache[key]

        return wrapper

    return wrapper1

class Session(object):
    def __init__(self, path='~/.config/satori.json',
                 cache_path='~/.cache/satori'):
        self.path = os.path.expanduser(path)
        self.cache_path = os.path.expanduser(cache_path)
        self.settings = {}
        self.match_contest_cache = Cache(self.cache_path + '/match_contest.json')
        self.match_submit_problem_cache = Cache(self.cache_path + '/match_submit_problem.json')
        self.match_problem_cache = Cache(self.cache_path + '/match_problem.json')

    def load(self):
        try:
            with open(self.path, 'r') as f:
                self.settings = json.load(f)
        except IOError:
            pass

    def save(self):
        with open(self.path + '.tmp', 'w') as f:
            json.dump(self.settings, f)
        os.rename(self.path + '.tmp', self.path)

    def login(self, username, password):
        self.settings['username'] = username
        self.settings['password'] = password.encode('hex')

    # HTTP

    def _do_request(self, method, path, parse=True, cookie=True, data={}, files={}):
        if LOG_REQUESTS:
            log_line = '{} {}'.format(method, path)
            color1 = '\33[0;33m'
            color2 = '\33[1;30m'
            endcolor = '\33[m'
            sys.stderr.write(color1 + log_line + '...'  + endcolor)
            sys.stderr.flush()
        r = requests.request(method, BASE + path,
                             files=files,
                             allow_redirects=False,
                             data=data,
                             headers={'Cookie':
                                      'satori_token='
                                      + self.settings['satori_token']} if cookie else {})

        if LOG_REQUESTS:
            sys.stderr.write(color2 + '\r' + log_line + ' -> {}'.format(r.status_code) + endcolor + '\n')

        return r

    def request(self, method, path, **kwargs):
        parse = kwargs.get('parse', True)

        if not 'satori_token' in self.settings:
            self._login()

        r = self._do_request(method, path, **kwargs)

        if r.status_code == 302 and r.headers['location'].startswith(BASE + '/login'):
            self._login()
            self.save()
            r = self._do_request(method, path)

        if parse:
            return pq(r.content)
        else:
            return r

    def _login(self):
        if 'username' not in self.settings:
            raise SatoriError('Please login with `satori login`.')
        resp = self._do_request('POST', '/login',
                                cookie=False,
                                parse=False,
                                data={
                                    'login': self.settings['username'],
                                    'password': self.settings['password'].decode('hex')})
        if resp.status_code == 302:
            cookie = resp.headers['set-cookie']
            token = cookie.split('=')[1].split(';')[0]
            self.settings['satori_token'] = token
        else:
            raise SatoriError('invalid login')

    # API

    def get_contests(self, other=False, archived=False):
        resp = self.request('GET', '/contest/select' + ('?participating_limit=0&participating_filter_archived=1' if archived else ''))
        tables = resp.find('.results')
        if not other:
            tables = pq(tables[0])
        for row in tables.find('tr'):
            row = pq(row)
            name = row.find('a.stdlink').text()
            link = row.find('a.stdlink').attr('href')
            if link.startswith('/contest/'):
                id = link.split('/')[2]
                yield name, int(id)

    @cached('match_contest_cache')
    def match_contest(self, query):
        try:
            return int(query)
        except ValueError:
            for name, id in self.get_contests(query):
                if query.lower() in name.lower():
                    return id
            raise SatoriError('no contest %r found' % query)

    def print_contests(self, **kwargs):
        for name, id in self.get_contests(**kwargs):
            print u'{}\t{}'.format(id, name)

    def get_problems(self, contest):
        url = '/contest/%d/problems' % self.match_contest(contest)

        for row in self.request('GET', url).find('.results').find('tr'):
            row = pq(row)
            cols = [ pq(c) for c in row.find('td') ]
            if not cols: continue
            code = cols[0].text().strip()
            name = cols[1].text()
            desc = cols[3].text()
            pdf = cols[2].find('a').attr('href')
            if pdf:
                id = int(pdf.split('/')[3])
            else:
                id = None
            url = cols[1].find('.stdlink')
            if url: url = url.attr('href')
            yield id, pdf, url, code, name, desc

    def print_problems(self, contest):
        for id, pdf, url, code, title, desc in self.get_problems(contest):
            print u'{:<9} {: <5} {: <30} {}' \
                .format(id or '', code, title, desc)


    def get_submit_problems(self, contest):
        url = '/contest/%d/submit' % self.match_contest(contest)
        body = self.request('GET', url)
        opts = body.find('[name=problem]').find('option')
        for opt in opts:
            opt = pq(opt)
            if opt.attr('value'):
                code, name = opt.text().split(':', 1)
                yield int(opt.attr('value')), code, name

    def print_submit_problems(self, contest):
        for id, code, desc in self.get_submit_problems(contest):
            print u'{:<9} {: <5} {}' \
                .format(id or '', code, desc)

    @cached('match_submit_problem_cache')
    def match_submit_problem(self, contest, problem):
        for id, code, desc in self.get_submit_problems(contest):
            if match_code(problem, code):
                return (id, code)
        raise SatoriError('unknown problem %r' % problem)

    @cached('match_problem_cache')
    def match_problem(self, contest, problem):
        try:
            return int(problem)
        except ValueError:
            pass

        for id, pdf, url, code, title, desc in self.get_problems(contest):
            if match_code(problem, code):
                return (id, pdf, url)
        raise SatoriError('unknown problem %r' % problem)

    def cache_write(self, name, data):
        if not os.path.exists(self.cache_path):
            os.mkdir(self.cache_path)

        path = self.cache_path + '/' + name
        with open(path, 'w') as f:
            f.write(data)

        return path

    def get_pdf(self, contest, problem):
        id, pdf, url = self.match_problem(contest, problem)
        if not pdf:
            return None
        ret = self.request('GET', pdf, parse=False)
        ret.raise_for_status()
        return self.cache_write('%d.pdf' % id, ret.content)

    def get_status(self, contest, id):
        url = '/contest/%d/results/%d' % (self.match_contest(contest), id)
        body = self.request('GET', url)
        header = body.find('table.results').find('tr')[1]
        header = pq(header).find('td')

        problem = pq(header[2]).text()
        status = pq(header[-1]).text()

        tests = []
        for row in body.find('.mainsphinx table.docutils tr')[1:]:
            row = pq(row)
            cols = map(pq, row.find('td'))
            tests.append((cols[0].text(), cols[1].text()))

        return problem, status, tests

    def print_status(self, contest, id, out=None):
        out = out or sys.stdout
        problem, status, tests = self.get_status(contest, id)
        print >>out, 'Test report for %s' % problem
        print >>out, 'Status:', status
        for name, status in tests:
            print >>out, ' - {: <10} {}'.format(name, status)

    def submit(self, contest, problem, file):
        name = file.split('/')[-1]
        url = '/contest/%d/submit' % self.match_contest(contest)
        problem, problemcode = self.match_submit_problem(contest, problem)

        ret = self.request('POST', url,
                           data={'problem': problem},
                           files={'codefile': (name, open(file), 'text/plain')},
                           parse=False)

        if ret.status_code != 302:
            err_path = self.cache_path + '/error.html'
            with open(err_path, 'w') as f:
                f.write(ret.content)
            raise SatoriError('submit failed (check error: {})'.format(err_path))

        return problemcode

    def get_last_submit(self, contest):
        url = '/contest/%d/results' % self.match_contest(contest)
        body = self.request('GET', url)
        tr = body.find('table.results').find('tr')[1]
        return int(pq(tr).find('a').text())

    def get_submits(self, contest):
        url = '/contest/%d/results?results_limit=2000' % self.match_contest(contest)
        body = self.request('GET', url)
        rows = body.find('table.results').find('tr')[1:]
        for row in rows:
            status = pq(row).find('.submitstatus').text()
            id = int(pq(row).find('a').text())
            name = row[1].text
            yield id, name, status

    def print_submits(self, contest):
        for id, name, status in self.get_submits(contest):
            print u'{}\t{}\t{}'.format(id, name, status)

    def download_submit(self, submit):
        resp = self.request('GET', '/download/Submit/%d/data/content/info.txt' % submit, parse=False)
        return resp.content

def notify_status(problem, status):
    subprocess.check_call(['notify-send', 'New status for {}: {}'.format(problem, status)])

def wait(sess, contest, id, notify=True):
    contest = sess.match_contest(contest)
    last_status = ''
    while True:
        problem, status, tests = sess.get_status(contest, id)
        if status != last_status:
            if notify: notify_status(problem, status)
            last_status = status
        if status and status != 'QUE':
            break
        time.sleep(10)

    if notify:
        with open(sess.cache_path + '/last.txt', 'a') as out:
            sess.print_status(contest, id, out)
            print >>out

    if not notify:
        sess.print_status(contest, id)

def download_submits(sess, contest, dir):
    from magic import Magic
    import mimetypes

    magic = Magic(mime=True, uncompress=False)
    mimetype_to_ext = dict( (v, k) for k, v in mimetypes.types_map.items() )
    mimetype_to_ext['text/x-c'] = '.cpp'

    submits = list(sess.get_submits(contest))

    for id, name, status in submits:
        status = status.split()[0]

        if '/' in name:
            raise ValueError('bad task name')
        if '/' in status:
            raise ValueError('bad status name')

        if not status:
            status = 'UNK'

        data = sess.download_submit(id)
        mime = magic.from_buffer(data).split(';')[0]
        ext = mimetype_to_ext.get(mime, '.txt')

        path = dir + '/%s/%d_%s%s' % (name, id, status, ext)
        dirs = os.path.dirname(path)
        if not os.path.exists(dirs):
            os.makedirs(dirs)
        with open(path, 'wb') as f:
            f.write(data)

def main():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='command')
    subparsers.add_parser(
        'login',
        help='Set Satori credentials.')

    contests_parser = subparsers.add_parser(
        'contests',
        help='List contests')
    contests_parser.add_argument('--show-other', action='store_true')
    contests_parser.add_argument('--show-archived', action='store_true')

    problems_parser = subparsers.add_parser(
        'problems',
        help='List problems.')
    problems_parser.add_argument('--submit',
                                 action='store_true',
                                 help='show submittable problems.')
    problems_parser.add_argument('contest')

    problem_parser = subparsers.add_parser(
        'problem',
        help='Open problem statement.')
    problem_parser.add_argument('contest')
    problem_parser.add_argument('problem')
    problem_parser.add_argument('--pdf', action='store_true',
                                help='Open PDF instead of problem page.')

    downloadproblems_parser = subparsers.add_parser(
        'downloadproblems',
        help='Download all problem PDFs from a given contest.')
    downloadproblems_parser.add_argument('contest')

    status_parser = subparsers.add_parser(
        'status',
        help='Show submit status.')
    status_parser.add_argument('contest')
    status_parser.add_argument('id', type=int, default=0, nargs='?')

    wait_parser = subparsers.add_parser(
        'wait',
        help='Wait for submit status change and show notification.')
    wait_parser.add_argument('contest')
    wait_parser.add_argument('id', type=int, default=0, nargs='?')

    submits_parser = subparsers.add_parser(
        'submits',
        help='List submits.')
    submits_parser.add_argument('contest')

    downloadsubmits_parser = subparsers.add_parser(
        'downloadsubmits',
        help='Downloads your submissions.')
    downloadsubmits_parser.add_argument('contest')
    downloadsubmits_parser.add_argument('dir')

    submit_parser = subparsers.add_parser(
        'submit',
        help='Submit a solution..')
    submit_parser.add_argument('contest')
    submit_parser.add_argument('problem')
    submit_parser.add_argument('file')

    subparsers.add_parser(
        'last',
        help='Show results of recently submitted tasks.')

    subparsers.add_parser(
        'clear-cache',
        help='Clear cache.')

    ns = parser.parse_args()
    sess = Session()
    sess.load()

    if ns.command == 'login':
        sess.login(raw_input('Username: '),
                   getpass.getpass('Password: '))
        sess.print_contests()
        sess.save()

    elif ns.command == 'contests':
        sess.print_contests(other=ns.show_other, archived=ns.show_archived)

    elif ns.command == 'problems':
        if ns.submit:
            sess.print_submit_problems(ns.contest)
        else:
            sess.print_problems(ns.contest)

    elif ns.command == 'problem':
        def open_pdf():
            path = sess.get_pdf(ns.contest, ns.problem)
            if not path:
                print 'Problem statement missing.'
            else:
                webbrowser.open('file://' + path)

        if ns.pdf:
            open_pdf()
        else:
            id, pdf, url = sess.match_problem(ns.contest, ns.problem)
            if not url:
                print 'No HTML for this problem'
                open_pdf()
            else:
                webbrowser.open(BASE + url)

    elif ns.command == 'downloadproblems':
        for id, pdf, url, code, title, desc in sess.get_problems(ns.contest):
            if '/' in code or '\0' in code:
                raise Exception('invalid task code')

            if not pdf:
                print 'skip task', title
                continue

            ret = sess.request('GET', pdf, parse=False)
            ret.raise_for_status()
            with open(code + '.pdf', 'wb') as f:
                f.write(ret.content)

    elif ns.command == 'status':
        id = ns.id
        if id == 0:
            id = sess.get_last_submit(ns.contest)
        sess.print_status(ns.contest, id)

    elif ns.command == 'wait':
        id = ns.id
        if id == 0:
            id = sess.get_last_submit(ns.contest)
        wait(sess, ns.contest, id, notify=False)

    elif ns.command == 'submit':
        sess.submit(ns.contest, ns.problem, ns.file)
        submit_id = sess.get_last_submit(ns.contest)
        print 'Submitted as', submit_id
        check_new_version()

        pid = os.fork()
        if pid == 0:
            sys.stdout = sys.stderr = open(sess.cache_path + '/waiter.log', 'a', 1)
            wait(sess, ns.contest, submit_id, notify=True)
            os._exit(0)
        else:
            print 'Forked waiter with PID', pid

    elif ns.command == 'submits':
        sess.print_submits(ns.contest)

    elif ns.command == 'downloadsubmits':
        download_submits(sess, ns.contest, ns.dir)

    elif ns.command == 'last':
        subprocess.check_call(['tail', '-n', '200', sess.cache_path + '/last.txt'])

    elif ns.command == 'clear-cache':
        path = os.path.expanduser('~/.cache/satori')
        for file in os.listdir(path):
            os.unlink(path + '/' + file)

def check_new_version():
    try:
        data = requests.get('https://raw.githubusercontent.com/zielmicha/satori-cli/master/version').content.strip()
        if data != __version__:
            print 'New version of satori-cli is available. '
            print 'Download it from https://github.com/zielmicha/satori-cli'
    except Exception as ex:
        print 'Failed to check for new version:', ex

if __name__ == '__main__':
    HTTP_DEBUG = False

    if HTTP_DEBUG:
        import httplib as http_client
        http_client.HTTPConnection.debuglevel = 1
        import logging
        logging.basicConfig()
        logging.getLogger().setLevel(logging.DEBUG)
        requests_log = logging.getLogger("requests.packages.urllib3")
        requests_log.setLevel(logging.DEBUG)
        requests_log.propagate = True

    try:
        main()
    except SatoriError as err:
        sys.exit('Error: ' + str(err))
    except KeyboardInterrupt:
        sys.exit('Interrupted')