1 files changed, 32 insertions, 9 deletions

diff --git a/iv/orodja/napad/submission.py b/iv/orodja/napad/submission.py
index dcf75ca..30e8257 100755
--- a/iv/orodja/napad/submission.py
+++ b/iv/orodja/napad/submission.py
@@ -5,15 +5,27 @@ 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")
+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 = []
-		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])
+		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]:
@@ -43,15 +55,26 @@ async def handle_client (reader, writer):
 		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:
+		if re.match(flag_regex, buffer.split(b' ')[0]) == None:
 			writer.write(b'BAD_FLAG\n')
 			continue
-		flag = buffer.decode()
+		# SUBMISSION LINE FORMAT: "flag teamnumber roundnumber service any other context"
+		flag = buffer.split(b' ')[0].decode()
+		context = b' '.join(buffer.split(b' ')[1:])
 		try:
-			db.execute("INSERT INTO flags (flag) VALUES (?)", [flag])
+			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 = [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")
+			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()