#!/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) STRICT")
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=? 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:
			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 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, 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())