submitbot_tcp.py

   1 #!/usr/bin/env python
   2
   3 # 2014 by Jan "Someone" Vales <someone@somenet.org>
   4 # do not publish!
   5
   6 import psycopg2
   7 import psycopg2.extras
   8 import sys
   9 import time
  10 import signal
  11 import socket
  12
  13 def readlines(sock, recv_buffer=4096, delim='\n'):
  14         buffer = ''
  15         data = True
  16         while data:
  17                 data = sock.recv(recv_buffer)
  18                 buffer += data
  19
  20                 while buffer.find(delim) != -1:
  21                         line, buffer = buffer.split('\n', 1)
  22                         yield line
  23         return
  24
  25 def submit(sock,flag):
  26     submission_success = False
  27     fs=sock.makefile()
  28
  29     print "submitting flag: "+flag
  30     sock.sendall(flag+"\n")
  31
  32     resp = fs.readline()+""
  33
  34     if 'Accepted' in resp:
  35         return (1, 'Accepted')
  36
  37     if 'Denied: no such flag' in resp:
  38         return (2, 'Denied: no such flag')
  39
  40     if 'Denied: flag is too old' in resp:
  41         return (2, 'Denied: flag is too old')
  42
  43     if 'Denied: you already submitted this flag' in resp:
  44         return (2, 'Denied: you already submitted this flag')
  45
  46     if 'Denied: flag is your own' in resp:
  47         return (2, 'Denied: flag is your own')
  48
  49     if 'Denied: your appropriate service' in resp:
  50         return (3, 'Denied: your appropriate service')
  51
  52     print(resp)
  53
  54
  55     if 'Status:error' in resp:
  56         wantnext = False
  57         servresponse = ""
  58         for line in resp.splitlines():
  59             if wantnext == True:
  60                 wantnext = False
  61                 servresponse = line
  62             if 'Status:error' in line:
  63                 wantnext = True
  64         return (2, 'Status:error::'+servresponse)
  65
  66     # RETURN (success?, srvresponse)
  67     print resp
  68     return (0, '')
  69
  70 def main():
  71     sleeptime = 5
  72     dbconn = None
  73     while True:
  74         try:
  75             print "*** sleeping "+str(sleeptime)+" sec..."
  76             time.sleep(sleeptime)
  77             dbconn = psycopg2.connect("host=127.0.0.1 port=5432 dbname=flagbot user=flagbot password=flagbotpw")
  78             cur = dbconn.cursor()
  79             cur.execute("CREATE TABLE IF NOT EXISTS flags ("
  80                 "fid serial NOT NULL PRIMARY KEY,"
  81                 "flag character varying(32) NOT NULL UNIQUE,"
  82                 "service character varying(32),"
  83                 "received timestamp without time zone NOT NULL DEFAULT date_trunc('second', NOW()),"
  84                 "submitted timestamp without time zone,"
  85                 "status integer NOT NULL DEFAULT 0,"
  86                 "srvresponse character varying(128)"
  87                 ")")
  88             dbconn.commit()
  89             cur.close()
  90             print "Connected to DB + table created"
  91
  92             cur = None
  93             while True:
  94                 print "*** sleeping another "+str(sleeptime)+" sec..."
  95                 time.sleep(sleeptime)
  96                 try:
  97                     cur = dbconn.cursor(cursor_factory=psycopg2.extras.DictCursor)
  98                     cur.execute("SELECT * from flags where status = 0 or status = 3 limit 500")
  99                     print "Fetched " + str(cur.rowcount) + " rows"
 100                     if cur.rowcount == 0:
 101                         continue
 102                     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 103                     sock.connect(("10.10.10.2", 31337))
 104                     eating = True
 105                     fs=sock.makefile()
 106                     while eating:
 107                         resp = fs.readline()+""
 108                         if resp.startswith('Enter your flags, finished with newline'):
 109                             eating = False
 110
 111                     for row in cur.fetchall():
 112                         (success, resp) = submit(sock,row['flag'])
 113                         if success != 0:
 114                             cur.execute("UPDATE flags SET submitted = date_trunc('second', NOW()), "
 115                                 "status = %s, srvresponse = %s WHERE fid = %s",
 116                                 (success, resp, row['fid']))
 117                             dbconn.commit()
 118                     sock.shutdown(socket.SHUT_WR)
 119                     sock.close()
 120                 except psycopg2.DatabaseError as e:
 121                     print 'Error %s' % e
 122                 try:
 123                     cur.close()
 124                     dbconn.rollback()
 125                 except psycopg2.DatabaseError as e:
 126                     print 'Error %s' % e
 127                 cur = None
 128                 dbconn.rollback()
 129
 130         except psycopg2.DatabaseError as e:
 131             print 'Error %s' % e
 132         try:
 133             dbconn.close()
 134         except psycopg2.DatabaseError as e:
 135             print 'Error %s' % e
 136         dbconn = None
 137     print "should never be reached"
 138
 139 if __name__ == "__main__":
 140     def signal_handler(signal, frame):
 141         print 'SIG received. exitting!'
 142         sys.exit(0)
 143     signal.signal(signal.SIGINT, signal_handler)
 144     main()