diff options
Diffstat (limited to 'iv/orodja/napad/submission.py')
| -rwxr-xr-x | iv/orodja/napad/submission.py | 47 |
1 files changed, 30 insertions, 17 deletions
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() |