asyncronous service and checker fixes

checker/checker.py

@@ -2,7 +2,9 @@
 #!/usr/bin/env python3
 from enochecker import *
 
+import binascii
 from binascii import unhexlify, hexlify
+
 from Crypto.PublicKey import RSA
 from Crypto.Cipher import DES
 from Crypto.Util.number import long_to_bytes, bytes_to_long
@@ -53,13 +55,15 @@ def putflag(self):  # type: () -> None
 				self.get_pubkey()
 				key = RSA.import_key(self.team_db['pubkey'])
 
+				conn = self.connect()
+				expect_command_prompt(conn)
+
 				content = 'flag %s %d' % (encrypt(self.flag, key), self.flag_round)
 				signature = sign(content, private_key)
 
 				input_data = ('receive %s %s' % (content, signature)).encode()
-				conn = self.connect()
-				expect_command_prompt(conn)
 				conn.write(input_data + b"\n")
+				self.debug(f"Sent msg to client: {input_data}")
 
 				try:
 					ret = expect_command_prompt(conn).decode().strip().split(":")
@@ -210,7 +214,12 @@ def getnoise(self):  # type: () -> None
 				joke_id = self.team_db[self.noise_key() + "joke_id"]
 				conn.write(f"send {joke_id}\n".encode() )
 				joke_hex = expect_command_prompt(conn).decode().strip()
-				joke = unhexlify(joke_hex).decode()
+				self.debug(f"joke recieved: {joke_hex}, len {len(joke_hex)}")
+				try:
+					joke = unhexlify(joke_hex).decode()
+				except binascii.Error:
+					self.debug("failed to decode joke-hex")
+					raise BrokenServiceException("Retrieved invalid joke")
 
 				joke_orig = self.team_db[self.noise_key() + "joke"]
 				self.debug(f"{joke_orig}, {joke}")
@@ -251,7 +260,7 @@ def exploit(self):
 		pass
 
 def expect_command_prompt(conn):
-	return conn.readline_expect(b'command: ',b'command: ').split(b'command')[0] # need colon and space in split?
+	return conn.readline_expect(b'command: ',b'command: ').split(b'command: ')[0] # need colon and space in split?
 
 app = CryStoreChecker.service  # This can be used for uswgi.
 if __name__ == "__main__":

service/Dockerfile

@@ -1,23 +1,24 @@
-FROM ubuntu:18.04
+FROM ubuntu:20.04
 
 RUN apt-get update && apt-get -y upgrade
 RUN apt-get install -y python3-pip
-RUN apt-get install -y socat
+# RUN apt-get install -y socat
 RUN apt-get install -y sqlite
 
 RUN python3 -m pip install pycryptodome
+RUN python3 -m pip install aiosqlite
 
 RUN mkdir /service
 WORKDIR /service
 RUN mkdir data
 
-COPY cry.py crypto.py init_client.sql checker.pubkey run.sh ./
+COPY cry_async.py crypto.py init_client.sql checker.pubkey run.sh ./
 
 RUN groupadd -g 2000 cryptogroup \
 && useradd -m -u 2001 -g cryptogroup cryptodude
 
 RUN chown -R cryptodude /service
-RUN chmod +x cry.py run.sh
+RUN chmod +x cry_async.py run.sh
 USER cryptodude
 
 CMD ["/service/run.sh"]

service/cry_async.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+from binascii import unhexlify, hexlify
+import string
+from Crypto.PublicKey import RSA
+from hashlib import sha256
+
+import asyncio
+import aiosqlite
+
+from crypto import decrypt, encrypt, sign, verify
+
+# def timeout_handler():
+#     print("TIMEOUT!")
+#     os._exit(os.EX_OK)
+
+if os.path.isfile('key.pem'):
+    try:
+        KEY = RSA.importKey(open('key.pem','r').read())
+    except:
+        KEY = RSA.generate(2048)
+        key_file = open('key.pem','wb')
+        key_file.write(KEY.export_key())
+        key_file.close()
+else:
+    KEY = RSA.generate(2048)
+    key_file = open('key.pem','wb')
+    key_file.write(KEY.export_key())
+    key_file.close()
+
+PUBLIC_KEY = KEY.public_key().export_key()
+
+if os.path.isfile('checker.pubkey'):
+    CHECKER_KEY = RSA.importKey(open('checker.pubkey','r').read())
+else:
+    raise Exception('Public Key of Checker not found')
+
+global DB_CONN
+
+class Store(object):
+    def __init__(self, rx, tx):
+        self.rx = rx
+        self.tx = tx
+
+    async def run(self):
+        while True:
+            try:
+                self.tx.write(b"command: ")
+                await self.tx.drain()
+
+                line = await self.rx.readline()
+                # print(f"Received command: {line}")
+
+                input_data = line.strip()
+                if not input_data:
+                    break
+
+                await self.process_command(input_data)
+                await self.tx.drain()
+
+            except EOFError:
+                break
+            except (UnicodeError, IndexError) as e:
+                print(e, file=sys.stderr)
+                break
+
+
+    async def process_command(self, input_data : bytes) -> str:
+        args = input_data.decode().split(' ') # split signature
+        command = args[0]
+
+        if command == 'receive':
+            #checking signature
+            if len(args) != 5:
+                self.tx.write(b"Entered line must have format \"receive category data tick signature\"\n")
+                return
+
+            signature = args[-1]
+            args = args[1:-1]
+            if not verify(' '.join(args), signature, CHECKER_KEY):
+                self.tx.write(b"invalid signature\n")
+                return
+
+            await self.receive(*args)
+
+        elif command == 'send':
+            try:
+                tick = int(args[1])
+            except ValueError:
+                self.tx.write(b'First argument must be integer\n')
+            await self.send(args[1])
+
+        elif command == 'send_pubkey':
+            self.tx.write(PUBLIC_KEY + b'\n')
+        else:
+            self.tx.write(b'Unknown command\n')
+            #print("Entered line must have format \"command [params]* [signature]\" separated by spaces")
+
+    async def receive(self, category : str, data : str, tick : str) -> str:
+        if category == 'flag':
+            data = decrypt(data, privkey = KEY)
+        else:
+            data = unhexlify(data).decode()
+        #store data in DB
+        try:
+            tick = int(tick)
+        except ValueError:
+            self.tx.write(b'tick must be integer\n')
+            return
+
+        if all([char in string.printable for char in data]):
+            async with DB_CONN.execute('insert into store (tick, category, data) values (?,?,?);', (tick, category, data)) as cursor:
+                last_id = cursor.lastrowid
+
+            if last_id is None:
+                self.tx.write(b'Failed to add element to db.\n')
+                return 
+
+            await DB_CONN.commit()
+            self.tx.write(f"{sha256(data.encode()).hexdigest()}:{last_id}\n".encode())
+        else:
+            self.tx.write(f'Data not correctly decrypted: {data.encode()}\n'.encode())
+
+    async def send(self, flag_id : int) -> str:
+        async with DB_CONN.execute('select data,category from store where id = ' + str(flag_id) + ';') as cursor:
+            try:
+                content, category = await cursor.fetchone() 
+            except TypeError:
+                self.tx.write(b'Key not in Database\n')
+                return
+
+        # print(content, category, file=sys.stderr)
+        if category == 'flag':
+            self.tx.write(f"{encrypt(content, CHECKER_KEY)}\n".encode())
+        else:
+            self.tx.write(hexlify(content.encode()) + b'\n')
+
+async def handle_connection(reader, writer):
+    s = Store(reader, writer)
+    await s.run()
+    writer.close()
+    await writer.wait_closed()
+
+async def main():
+    global DB_CONN
+    DB_CONN = await aiosqlite.connect('data/store.db')
+
+    server = await asyncio.start_server(handle_connection, host="0.0.0.0", port="9122")
+
+    addr = server.sockets[0].getsockname()
+    print(f'Serving on {addr}')
+
+    async with server:
+        await server.serve_forever()
+
+if __name__ == "__main__":
+    asyncio.run(main())

service/run.sh

@@ -1,4 +1,5 @@
 #!/bin/sh
 touch data/store.db
 sqlite3 data/store.db < init_client.sql
-socat TCP4-LISTEN:9122,fork,reuseaddr EXEC:'python3 /service/cry.py'
+#socat TCP4-LISTEN:9122,fork,reuseaddr EXEC:'python3 /service/cry.py'
+python3 /service/cry_async.py