From fda4eb4f383ffeaf038a5734d6a10a5a2e310517 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anton=20Luka=20=C5=A0ijanec?= <anton@sijanec.eu>
Date: Tue, 3 Sep 2024 01:51:44 +0200
Subject: napad now works

---
 iv/orodja/napad/submission.py | 47 +++++++++++++++++++++++++++----------------
 1 file changed, 30 insertions(+), 17 deletions(-)

(limited to 'iv/orodja/napad/submission.py')

diff --git a/iv/orodja/napad/submission.py b/iv/orodja/napad/submission.py
index 30e8257..fce0a27 100755
--- a/iv/orodja/napad/submission.py
+++ b/iv/orodja/napad/submission.py
@@ -4,12 +4,12 @@ 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", "^[A-Z0-9]{31}=$").encode(), re.ASCII | re.DOTALL | re.VERBOSE)
+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:
-		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"):
@@ -19,29 +19,39 @@ async def submitter ():
 			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():
+			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:
-				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()
+				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:
-		incoming = await reader.readuntil(b'\n')
+		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'')
@@ -63,18 +73,21 @@ async def handle_client (reader, writer):
 		context = b' '.join(buffer.split(b' ')[1:])
 		try:
 			team = int(buffer.split(b' ')[1].decode())
-		except (ValueError, UnicodeDecodeError):
+		except (ValueError, UnicodeDecodeError, IndexError):
 			team = -1
 		try:
 			runda = int(buffer.split(b' ')[2].decode())
-		except (ValueError, UnicodeDecodeError):
+		except (ValueError, UnicodeDecodeError, IndexError):
 			runda = -1
-		service = buffer.split(b' ')[3]
 		try:
-			db.execute("INSERT INTO flags (flag, team, service, round, context) VALUES (?, ?)", [flag, team, service, runda, context])
+			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 " + date.encode() + b"\t" + str(status).encode() + b"\t" + str(context).encode() + b"\t" + str(msg).encode() + b"\n")
+			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()
-- 
cgit v1.2.3