]> git.somenet.org - ctf/pub/submit_bot.git/blob - submitbot_tcp.py
Prepared for ictf2015
[ctf/pub/submit_bot.git] / 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()