#!/usr/bin/env python

# 2014 by Jan "Someone" Vales <someone@somenet.org>
# do not publish!

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

def readlines(sock, recv_buffer=4096, delim='\n'):
	buffer = ''
	data = True
	while data:
		data = sock.recv(recv_buffer)
		buffer += data

		while buffer.find(delim) != -1:
			line, buffer = buffer.split('\n', 1)
			yield line
	return

def submit(sock,flag):
    submission_success = False
    fs=sock.makefile()

    print "*** Submitting flag: "+flag

    sock.sendall(flag+"\n")

    resp = fs.readline()+""
    print resp
    if 'Accepted' in resp:
        return (1, resp.replace(flag,''))

    if 'Denied: no such flag' in resp:
        return (2, 'Denied: no such flag')

    if 'Denied: invalid flag' in resp:
        return (2, 'Denied: invalid flag')

    if 'Denied: flag is too old' in resp:
        return (2, 'Denied: flag is too old')

    if 'Denied: you already submitted this flag' in resp:
        return (3, 'Denied: you already submitted this flag')

    if 'Denied: flag is your own' in resp:
        return (2, 'Denied: flag is your own')

    if 'Denied: your appropriate service' in resp:
        return (4, 'Denied: your appropriate service')

    # RETURN (success?, srvresponse)
    print resp
    return (4, 'WTF?!')

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
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect(("flags.ructfe.org", 31337))
            eating = True
            fs=sock.makefile()
            while eating:
                 resp = fs.readline()+""
                 if resp.startswith('Enter your flags, finished with newline'):
                     eating = False

            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")
                    #print "Fetched " + str(cur.rowcount) + " rows"
                    if cur.rowcount == 0:
                        print "*** sleeping another "+str(sleeptime)+" sec..."
                        time.sleep(sleeptime)
                        continue
                    for row in cur.fetchall():
                    #    print "here" + str(row['fid'])
                        if row['flag'] is None or row['flag'].strip() == '':
                            continue						
                        (success, resp) = submit(sock,row['flag'])
                        if success != 0:
                            cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), status = %s, srvresponse = %s WHERE fid = %s and status = 0 or status = 4",
                                (success, 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()