1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
|
#!/usr/bin/python3
import os
import asyncio
import re
import sqlite3
import aiohttp
db = sqlite3.connect(os.getenv("SUBMISSION_DB", "flags.db"))
db.execute("CREATE TABLE IF NOT EXISTS flags (id INTEGER PRIMARY KEY, flag TEXT NOT NULL UNIQUE, sent INTEGER NOT NULL DEFAULT 0, date TEXT DEFAULT (strftime('%FT%R:%f', 'now')) NOT NULL, status TEXT, msg TEXT) STRICT")
flag_regex = re.compile(os.getenv("FLAG_REGEX", "^[A-Z0-9]{31}=$").encode(), re.ASCII | re.DOTALL | re.VERBOSE)
async def submitter ():
while True:
print("submitter loop")
flags = []
for row in db.execute("SELECT flag FROM flags WHERE sent == 0 ORDER BY date DESC"):
if len(flags) < int(os.getenv("SUBMISSION_MAX_FLAGS", "2560")):
flags.append(row[0])
if len(flags) == 0:
await asyncio.sleep(1)
for i in [1]:
async with aiohttp.ClientSession(headers={"X-Team-Token": os.getenv("SUBMISSION_TEAM_TOKEN")}) as session:
async with session.put(os.getenv("SUBMISSION_URL", 'http://10.10.0.1:8080/flags'), json=flags) as response:
if response.status // 100 != 2:
print("submitter error: " + await response.text())
break
cursor = db.cursor()
for obj in await response.json():
cursor.execute("UPDATE flags SET sent=?, status=?, msg=? WHERE flag=?", [int(obj.get("status") != "RESUBMIT"), obj.get("status"), obj.get("msg"), obj.get("flag")])
db.commit()
await asyncio.sleep(int(os.getenv("SUBMISSION_DELAY", "15")))
async def handle_client (reader, writer):
while True:
incoming = await reader.readuntil(b'\n')
if len(incoming) == 0:
break
buffer = incoming.replace(b'\r', b'').replace(b'\n', b'')
if buffer.startswith(b' '):
for row in db.execute(buffer[1:].decode()):
writer.write(str(row).encode() + b'\n')
continue
if buffer.startswith(b'@'):
writer.write(str(db.execute(buffer[1:].decode()).fetchall()).encode() + b'\n')
continue
if buffer.startswith(b'#'):
writer.write(str(len(db.execute(buffer[1:].decode()).fetchall())).encode() + b'\n')
continue
if re.match(flag_regex, buffer) == None:
writer.write(b'BAD_FLAG\n')
continue
flag = buffer.decode()
try:
db.execute("INSERT INTO flags (flag) VALUES (?)", [flag])
except sqlite3.IntegrityError:
status, msg, date = [x for x in db.execute("SELECT status, msg, date FROM flags WHERE flag=?", [flag])][0]
writer.write(b"OLD_FLAG " + date.encode() + b" " + str(status).encode() + b" " + str(msg).encode() + b"\n")
else:
writer.write(b'NEW_FLAG\n')
writer.close()
async def run_server ():
server = await asyncio.start_server(handle_client, os.getenv("SUBMISSION_BIND", "::"), os.getenv("SUBMISSION_PORT", "21502"))
event_loop = asyncio.get_event_loop()
event_loop.create_task(submitter())
async with server:
await server.serve_forever()
asyncio.run(run_server())
|
