diff options
Diffstat (limited to 'iv/orodja/napad/submission.py')
-rwxr-xr-x | iv/orodja/napad/submission.py | 64 |
1 files changed, 64 insertions, 0 deletions
diff --git a/iv/orodja/napad/submission.py b/iv/orodja/napad/submission.py new file mode 100755 index 0000000..dcf75ca --- /dev/null +++ b/iv/orodja/napad/submission.py @@ -0,0 +1,64 @@ +#!/usr/bin/python3 +import os +import asyncio +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") +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 = [] + 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]) + if len(flags) == 0: + await asyncio.sleep(1) + 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() + await asyncio.sleep(int(os.getenv("SUBMISSION_DELAY", "15"))) +async def handle_client (reader, writer): + while True: + incoming = await reader.readuntil(b'\n') + 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) == None: + writer.write(b'BAD_FLAG\n') + continue + flag = buffer.decode() + try: + db.execute("INSERT INTO flags (flag) VALUES (?)", [flag]) + 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") + 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()) |