path: root/iv/orodja/napad/submission.py

#!/usr/bin/python3
import os
import asyncio
import re
import sqlite3
import aiohttp
import traceback
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, submitted TEXT) STRICT") # submitted is date
flag_regex = re.compile(os.getenv("FLAG_REGEX_MATCH", "^[A-Z0-9]{31}=$").encode(), re.ASCII | re.DOTALL | re.VERBOSE)
async def submitter ():
	while True:
		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 [x for x in flags_balance.keys()]:
				try:
					zastava = flags_balance[key].pop(0)
				except IndexError:
					flags_balance.pop(key)
				else:
					flags.append(zastava)
					unsent_flags -= 1
		if len(flags) == 0:
			await asyncio.sleep(1)
			continue
		for i in [1]:
			async with aiohttp.ClientSession(headers={"X-Team-Token": os.getenv("SUBMISSION_TEAM_TOKEN")}) as session:
				try:
					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=?, submitted=strftime('%FT%R:%f', 'now') WHERE flag=?", [int(obj.get("status") != "RESUBMIT"), obj.get("status"), obj.get("msg"), obj.get("flag")])
						db.commit()
				except Exception as e:
					traceback.print_exc()
		await asyncio.sleep(int(os.getenv("SUBMISSION_DELAY", "15")))
async def handle_client (reader, writer):
	while True:
		try: # SUBMISSION LINE FORMAT: "flag teamnumber roundnumber service any other context"
			incoming = await reader.readuntil(b'\n')
		except asyncio.exceptions.IncompleteReadError as e:
			if int(str(e).split(" ")[0]) == 0:
				break
			raise e
		if len(incoming) == 0:
			break
		buffer = incoming.replace(b'\r', b'').replace(b'\n', b'')
		if re.match(flag_regex, buffer.split(b' ')[0]) == None:
			writer.write(b'BAD_FLAG\n')
			continue
		flag = buffer.split(b' ')[0].decode()
		context = b' '.join(buffer.split(b' ')[1:])
		try:
			team = int(buffer.split(b' ')[1].decode())
		except (ValueError, UnicodeDecodeError, IndexError):
			team = -1
		try:
			runda = int(buffer.split(b' ')[2].decode())
		except (ValueError, UnicodeDecodeError, IndexError):
			runda = -1
		try:
			service = buffer.split(b' ')[3]
		except IndexError:
			service = None
		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 " + str(status).encode() + b" " + date.encode() + b" " + context + 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())