1+ #!/usr/bin/env python3
2+
3+ import sys
4+ import os
5+ from binascii import unhexlify , hexlify
6+ import string
7+ from Crypto .PublicKey import RSA
8+ from hashlib import sha256
9+
10+ import asyncio
11+ import aiosqlite
12+
13+ from crypto import decrypt , encrypt , sign , verify
14+
15+ # def timeout_handler():
16+ # print("TIMEOUT!")
17+ # os._exit(os.EX_OK)
18+
19+ if os .path .isfile ('key.pem' ):
20+ try :
21+ KEY = RSA .importKey (open ('key.pem' ,'r' ).read ())
22+ except :
23+ KEY = RSA .generate (2048 )
24+ key_file = open ('key.pem' ,'wb' )
25+ key_file .write (KEY .export_key ())
26+ key_file .close ()
27+ else :
28+ KEY = RSA .generate (2048 )
29+ key_file = open ('key.pem' ,'wb' )
30+ key_file .write (KEY .export_key ())
31+ key_file .close ()
32+
33+ PUBLIC_KEY = KEY .public_key ().export_key ()
34+
35+ if os .path .isfile ('checker.pubkey' ):
36+ CHECKER_KEY = RSA .importKey (open ('checker.pubkey' ,'r' ).read ())
37+ else :
38+ raise Exception ('Public Key of Checker not found' )
39+
40+ global DB_CONN
41+
42+ class Store (object ):
43+ def __init__ (self , rx , tx ):
44+ self .rx = rx
45+ self .tx = tx
46+
47+ async def run (self ):
48+ while True :
49+ try :
50+ self .tx .write (b"command: " )
51+ await self .tx .drain ()
52+
53+ line = await self .rx .readline ()
54+ # print(f"Received command: {line}")
55+
56+ input_data = line .strip ()
57+ if not input_data :
58+ break
59+
60+ await self .process_command (input_data )
61+ await self .tx .drain ()
62+
63+ except EOFError :
64+ break
65+ except (UnicodeError , IndexError ) as e :
66+ print (e , file = sys .stderr )
67+ break
68+
69+
70+ async def process_command (self , input_data : bytes ) -> str :
71+ args = input_data .decode ().split (' ' ) # split signature
72+ command = args [0 ]
73+
74+ if command == 'receive' :
75+ #checking signature
76+ if len (args ) != 5 :
77+ self .tx .write (b"Entered line must have format \" receive category data tick signature\" \n " )
78+ return
79+
80+ signature = args [- 1 ]
81+ args = args [1 :- 1 ]
82+ if not verify (' ' .join (args ), signature , CHECKER_KEY ):
83+ self .tx .write (b"invalid signature\n " )
84+ return
85+
86+ await self .receive (* args )
87+
88+ elif command == 'send' :
89+ try :
90+ tick = int (args [1 ])
91+ except ValueError :
92+ self .tx .write (b'First argument must be integer\n ' )
93+ await self .send (args [1 ])
94+
95+ elif command == 'send_pubkey' :
96+ self .tx .write (PUBLIC_KEY + b'\n ' )
97+ else :
98+ self .tx .write (b'Unknown command\n ' )
99+ #print("Entered line must have format \"command [params]* [signature]\" separated by spaces")
100+
101+ async def receive (self , category : str , data : str , tick : str ) -> str :
102+ if category == 'flag' :
103+ data = decrypt (data , privkey = KEY )
104+ else :
105+ data = unhexlify (data ).decode ()
106+ #store data in DB
107+ try :
108+ tick = int (tick )
109+ except ValueError :
110+ self .tx .write (b'tick must be integer\n ' )
111+ return
112+
113+ if all ([char in string .printable for char in data ]):
114+ async with DB_CONN .execute ('insert into store (tick, category, data) values (?,?,?);' , (tick , category , data )) as cursor :
115+ last_id = cursor .lastrowid
116+
117+ if last_id is None :
118+ self .tx .write (b'Failed to add element to db.\n ' )
119+ return
120+
121+ await DB_CONN .commit ()
122+ self .tx .write (f"{ sha256 (data .encode ()).hexdigest ()} { last_id } \n " .encode ())
123+ else :
124+ self .tx .write (f'Data not correctly decrypted: { data .encode ()} \n ' .encode ())
125+
126+ async def send (self , flag_id : int ) -> str :
127+ async with DB_CONN .execute ('select data,category from store where id = ' + str (flag_id ) + ';' ) as cursor :
128+ try :
129+ content , category = await cursor .fetchone ()
130+ except TypeError :
131+ self .tx .write (b'Key not in Database\n ' )
132+ return
133+
134+ # print(content, category, file=sys.stderr)
135+ if category == 'flag' :
136+ self .tx .write (f"{ encrypt (content , CHECKER_KEY )} \n " .encode ())
137+ else :
138+ self .tx .write (hexlify (content .encode ()) + b'\n ' )
139+
140+ async def handle_connection (reader , writer ):
141+ s = Store (reader , writer )
142+ await s .run ()
143+ writer .close ()
144+ await writer .wait_closed ()
145+
146+ async def main ():
147+ global DB_CONN
148+ DB_CONN = await aiosqlite .connect ('data/store.db' )
149+
150+ server = await asyncio .start_server (handle_connection , host = "0.0.0.0" , port = "9122" )
151+
152+ addr = server .sockets [0 ].getsockname ()
153+ print (f'Serving on { addr } )
154+
155+ async with server :
156+ await server .serve_forever ()
157+
158+ if __name__ == "__main__" :
159+ asyncio .run (main ())
0 commit comments