flagbot.py

   1 #!/usr/bin/env python
   2 #
   3 # Copyright 2015 by Jan Vales <jan@jvales.net> (Someone <someone@somenet.org>)
   4 # send me your changes. credit author(s). do not publish. share alike.
   5 # to be done: find a suitable licence text.
   6 #
   7
   8 import psycopg2
   9 import psycopg2.extras
  10 import sys
  11 import time
  12 import signal
  13 import socket
  14 import threading
  15
  16 import ictf
  17
  18 team = None
  19 dbconnstring = "host=127.0.0.1 port=5432 dbname=flagbot user=flagbot password=flagbotpw"
  20
  21 def login(force = False):
  22     global team
  23     if force == True or team is None:
  24         team = ictf.iCTF().login(# TODO: insert credentials here
  25         print("login(): logged in.")
  26
  27 def getTargets():
  28     print("getTargets(): starting...")
  29     login(True)
  30     global team
  31
  32     service_list = team.get_service_list()
  33     service_names = [s['service_name'] for s in service_list]
  34     try:
  35         dbconn = psycopg2.connect(dbconnstring)
  36         cur = dbconn.cursor()
  37         for service in service_names:
  38             try:
  39                 targets = team.get_targets(service)
  40                 for target in targets:
  41                     try:
  42 # TODO: wait for ictf folks to finish this one!
  43                         cur.execute("INSERT INTO flag_ids (service, host, flag_id) VALUES (%s, %s, %s)", (service, target['ip'], target['flag_id']))
  44                         dbconn.commit()
  45                     except psycopg2.DatabaseError as e:
  46                         pass
  47             except Exception as e:
  48                 print 'Error %s' % e
  49                 print "Unexpected error:", sys.exc_info()[0]
  50
  51     except Exception as e:
  52         print 'Error %s' % e
  53         print "Unexpected error:", sys.exc_info()[0]
  54
  55     try:
  56         dbconn.close()
  57     except psycopg2.DatabaseError as e:
  58         print 'Error %s' % e
  59     dbconn = None
  60
  61     time.sleep(5.0)
  62     gettargetth = threading.Timer(0.5, getTargets)
  63     gettargetth.start()
  64     print("getTargets(): ended")
  65
  66 def submitFlags():
  67     print("submitFlags(): starting...")
  68     login()
  69     global team
  70
  71     try:
  72         dbconn = psycopg2.connect(dbconnstring)
  73         cur = dbconn.cursor(cursor_factory=psycopg2.extras.DictCursor)
  74         cur.execute("SELECT * from flags where status = 0 or status = 3 limit 500")
  75         print "Fetched " + str(cur.rowcount) + " rows"
  76         for row in cur.fetchall():
  77             try:
  78 # TODO: wait for ictf folks to finish this one!
  79                 print("Would submit:"+row['flag'])
  80                 result = team.submit_flag([row['flag']])
  81
  82 # TODO: wait for ictf folks to finish this one!
  83 #                if "correct" in result:
  84 #                    cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), "
  85 #                                "status = %s, srvresponse = %s WHERE fid = %s",
  86 #                                (1, "correct", row['fid']))
  87 # TODO: wait for ictf folks to finish this one!
  88
  89                 dbconn.commit()
  90             except psycopg2.DatabaseError as e:
  91                 print 'Error %s' % e
  92
  93     except Exception as e:
  94         print 'Error %s' % e
  95         print "Unexpected error:", sys.exc_info()[0]
  96
  97     try:
  98         dbconn.close()
  99     except psycopg2.DatabaseError as e:
 100         print 'Error %s' % e
 101     dbconn = None
 102
 103     time.sleep(5.0)
 104     submitth = threading.Timer(0.5, submitFlags)
 105     submitth.setDaemon(True)
 106     submitth.start()
 107     print("submitFlags(): ended")
 108
 109 def main():
 110     print "*** starting ..."
 111     dbconn = psycopg2.connect("host=127.0.0.1 port=5432 dbname=flagbot user=flagbot password=flagbotpw")
 112     cur = dbconn.cursor()
 113     cur.execute("CREATE TABLE IF NOT EXISTS flag_ids ("
 114         "service character varying(32) NOT NULL,"
 115         "host character varying(32) NOT NULL,"
 116         "flag_id character varying(128) NOT NULL,"
 117         "received timestamp without time zone NOT NULL DEFAULT date_trunc('second'::text, now()),"
 118         "handed_out timestamp without time zone,"
 119         "status integer NOT NULL DEFAULT 0,"
 120         "CONSTRAINT flag_ids_pkey PRIMARY KEY (service, flag_id)"
 121         ")")
 122     cur.execute("CREATE TABLE IF NOT EXISTS flags ("
 123         "fid serial NOT NULL,"
 124         "flag character varying(32) NOT NULL,"
 125         "service character varying(32),"
 126         "flag_id character varying(128),"
 127         "submitter character varying(32),"
 128         "received timestamp without time zone NOT NULL DEFAULT date_trunc('second'::text, now()),"
 129         "submitted timestamp without time zone,"
 130         "status integer NOT NULL DEFAULT 0, "
 131         "srvresponse character varying(128),"
 132         "CONSTRAINT flags_pkey PRIMARY KEY (fid),"
 133         "CONSTRAINT flags_flag_key UNIQUE (flag)"
 134         ")")
 135     dbconn.commit()
 136     cur.close()
 137     cur = None
 138     dbconn = None
 139     print "Connected to DB + tables created"
 140
 141     # start one other thread
 142     submitth = threading.Timer(1.0, submitFlags)
 143     submitth.setDaemon(True)
 144     submitth.start()
 145
 146     getTargets()
 147
 148 if __name__ == "__main__":
 149     def signal_handler(signal, frame):
 150         print 'SIG received. exitting!'
 151         sys.exit(0)
 152     signal.signal(signal.SIGINT, signal_handler)
 153     main()
 154