1 files changed, 89 insertions, 0 deletions
diff --git a/iv/orodja/napad/submission.py b/iv/orodja/napad/submission.py
new file mode 100755
index 0000000..c345bdb
--- /dev/null
+++ b/iv/orodja/napad/submission.py
@@ -0,0 +1,89 @@
+#!/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())