#!/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 submit(flag):
    global team

    try:
        result = team.submit_flag([flag])
        result = result[0]

        if "notactive" in result:
            return (2, result)
        elif "incorrect" in result:
            return (3, result)
        elif "correct" in result:
            return (1, result)
        else:
            return (4, result)

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


def main():
    sleeptime = 1
    dbconn = None
    while True:
        try:
            print("*** sleeping "+str(sleeptime)+" sec...")
            time.sleep(sleeptime)
            dbconn = psycopg2.connect("host=s.i port=5432 dbname=flagbot user=flagbot password=flagbotpw")
            print("Connected to DB")

            cur = None
            login(True)
            global team

            while True:
                try:
                    cur = dbconn.cursor(cursor_factory=psycopg2.extras.DictCursor)
                    cur.execute("SELECT * from flags where status = 0 or status = 4 order by random() limit 1")
                    if cur.rowcount == 0:
                        print("*** sleeping another "+str(sleeptime)+" sec...")
                        time.sleep(sleeptime)
                        continue
                    for row in cur.fetchall():
                        if row['flag'] is None or row['flag'].strip() == '':
                            continue
                        (status, resp) = submit(row['flag'])
                        print("submitted: "+row['flag']+" - response: ("+str(status)+") "+resp)
                        if status != 0:
                            cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), status = %s, srvresponse = %s WHERE fid = %s and status = 0 or status = 4",
                                (status, resp, row['fid']))
                            dbconn.commit()
                except psycopg2.DatabaseError as e:
                    print('Error %s' % e)
                try:
                    cur.close()
                    dbconn.rollback()
                except psycopg2.DatabaseError as e:
                    print('Error %s' % e)
                cur = None
                dbconn.rollback()
            sock.shutdown(socket.SHUT_WR)
            sock.close()

        except psycopg2.DatabaseError as e:
            print('Error %s' % e)
        try:
            dbconn.close()
        except psycopg2.DatabaseError as e:
            print('Error %s' % e)
        dbconn = None

    print("should never be reached")


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