#!/usr/bin/env python2
#
# Copyright 2015-2017 by Jan Vales <jan@jvales.net> (Someone <someone@somenet.org>)
# send me your changes. credit author(s). do not publish. share alike.
# to be done: find a suitable licence text.
#

import psycopg2
import psycopg2.extras
import sys
import time
import signal
import socket
import threading
import traceback

import ictf

team = None
dbconnstring = "host=127.0.0.1 port=5432 dbname=flagbot user=flagbot password=flagbotpw"

def login(force = False):
    global team
    if force == True or team is None:
        team = ictf.iCTF().login('ctf@w0y.at','ZmtphWHFUDwRWk6m')
        print("login(): logged in.")


def submitFlags():
    print("submitFlag(): starting...")
    dbconn = None

    try:
        login()
        global team

        dbconn = psycopg2.connect(dbconnstring)
        cur = dbconn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        cur.execute("SELECT * from flags where status = 0 or status = 4 order by random() limit 400")
        print("submitFlag(): submitting " + str(cur.rowcount) + " flags...")
        flist = [x['flag'] for x in cur.fetchall()]
        result = team.submit_flag(flist)
        resultlist = list(zip(flist,result))

        for res in resultlist:
            try:
                result = res[1]

                if "notactive" in result:
                    cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), "
                                "status = %s, srvresponse = %s WHERE flag = %s",
                                (2, result, res[0]))

                elif "incorrect" in result:
                    cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), "
                                "status = %s, srvresponse = %s WHERE flag = %s",
                                (2, result, res[0]))

                elif "correct" in result:
                    cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), "
                                "status = %s, srvresponse = %s WHERE flag = %s",
                                (1, result, res[0]))

                elif "already" in result:
                    cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), "
                                "status = %s, srvresponse = %s WHERE flag = %s",
                                (3, result, res[0]))


                else:
                    import pprint
                    pprint.pprint(result)
                    cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), "
                                "status = %s, srvresponse = %s WHERE flag = %s",
                                (4, result, res[0]))

                dbconn.commit()
            except psycopg2.DatabaseError as e:
                print('Error %s' % e)
                traceback.print_exc(file=sys.stdout)

    except Exception as e:
        print('Error %s' % e)
        traceback.print_exc(file=sys.stdout)

    try:
        dbconn.close()
    except AttributeError as e:
        pass
#        print('Error %s' % e)
#        traceback.print_exc(file=sys.stdout)

    except psycopg2.DatabaseError as e:
        print('Error %s' % e)
        traceback.print_exc(file=sys.stdout)

    print("submitFlag(): ended.")
    time.sleep(5.0)


def main():
    print("*** starting ...")
    while True:
        submitFlags()

if __name__ == "__main__":
    def signal_handler(signal, frame):
        print('SIG received. exitting!')
        sys.exit(0)
    signal.signal(signal.SIGINT, signal_handler)
    main()