#!/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, team INTEGER, service BLOB, round INTEGER, context BLOB, 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_balance = dict()
unsent_flags = 0
for flag, team, service in db.execute("SELECT flag, team, service FROM flags WHERE sent == 0 ORDER BY date DESC"):
if (team, service) not in flags_balance.keys():
flags_balance[(team, service)] = []
flags_balance[(team, service)].append(flag)
unsent_flags += 1
flags = []
while len(flags) < int(os.getenv("SUBMISSION_MAX_FLAGS", "2560")) and unsent_flags > 0: # to zna biti počasno, najdi lepši način
for key in flags_balance.keys():
try:
zastava = flags_balance[key].pop(0)
except IndexError:
flags_balance.pop(key)
else:
flags.append(zastava)
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.split(b' ')[0]) == None:
writer.write(b'BAD_FLAG\n')
continue
# SUBMISSION LINE FORMAT: "flag teamnumber roundnumber service any other context"
flag = buffer.split(b' ')[0].decode()
context = b' '.join(buffer.split(b' ')[1:])
try:
team = int(buffer.split(b' ')[1].decode())
except (ValueError, UnicodeDecodeError):
team = -1
try:
runda = int(buffer.split(b' ')[2].decode())
except (ValueError, UnicodeDecodeError):
runda = -1
service = buffer.split(b' ')[3]
try:
db.execute("INSERT INTO flags (flag, team, service, round, context) VALUES (?, ?)", [flag, team, service, runda, context])
except sqlite3.IntegrityError:
status, msg, date, context = [x for x in db.execute("SELECT status, msg, date, context FROM flags WHERE flag=?", [flag])][0]
writer.write(b"OLD_FLAG " + date.encode() + b"\t" + str(status).encode() + b"\t" + str(context).encode() + b"\t" + 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())
