diff --git a/mods/udp_srv.py b/mods/udp_srv.py new file mode 100644 index 0000000..edf7db4 --- /dev/null +++ b/mods/udp_srv.py @@ -0,0 +1,78 @@ +import socketserver +import pickle +import simplecrypto +from uuid import uuid4 + +BUFSIZE = 8192 +BLOCKSIZE = 4096 + +PACKET_TYPE_RECV_RESULT = 'recv_result' +PACKET_TYPE_DATA_FRAGMENT = 'data_fragment' +PACKET_TYPE_SVC_MESSAGE = 'svc_message' +PACKET_TYPE_NEW_CONNECTION = 'new_connection' + +DATA_TYPE_FILE_CHUNK = 'file_chunk' +DATA_TYPE_CMD = 'cmd' +DATA_TYPE_SEARCH_QUERY = 'search_query' +DATA_TYPE_SEARCH_RESULT = 'search_result' +DATA_TYPE_PEER_PUBKEY = 'peer_pubkey' + +SVC_MESSAGE_BAD_PACKET = 'bad_packet' +SVC_MESSAGE_DECRYPT_ERROR = 'decrypt_error' +SVC_MESSAGE_YOU_ARE_STRANGER = 'you_are_stranger' + +RECV_OK_CODE = 0 +RECV_ERROR_CODE = 1 + +HEADER_SVC_YOU_ARE_STRANGER = {'type': PACKET_TYPE_SVC_MESSAGE, 'msg': SVC_MESSAGE_YOU_ARE_STRANGER} + +HEADER_RECV_OK = {'type': PACKET_TYPE_RECV_RESULT, 'code': RECV_OK_CODE} +HEADER_RECV_ERROR = {'type': PACKET_TYPE_RECV_RESULT, 'code': RECV_ERROR_CODE} + +key = None + +peers = {} + +class Peer: + def __init__(self, peer_addr, pubkey): + self.addr = peer_addr + self.pubkey = pubkey + +def get_key(): + global key + if not key: + key = simplecrypto.RsaKeypair() + return key + +def pickle_data(data): + return pickle.dumps(data, protocol=4) + +def write_svc_msg(fobj, svc_msg, extra_data={}): + MSG = {'type': PACKET_TYPE_SVC_MESSAGE, 'msg': svc_msg} + if extra_data: + MSG.update(extra_data) + fobj.write(pickle_data(MSG)) + +class UDPRequestHandler(socketserver.DatagramRequestHandler): + def handle(self): + datagram = self.rfile.read(BUFSIZE) + if self.client_address not in peers: + write_svc_msg(self.wfile, SVC_MESSAGE_YOU_ARE_STRANGER) + return + try: + unpickled_datagram = pickle.dumps(datagram) + except pickle.UnpicklingError: + write_svc_msg(self.wfile, SVC_MESSAGE_BAD_PACKET) + return + pk = get_key() + try: + data = pickle.loads(pk.decrypt_raw(unpickled_datagram['packet'])) + except: + write_svc_msg(self.wfile, SVC_MESSAGE_DECRYPT_ERROR, {'packet_id': unpickled_datagram['packet_id']}) + return + block_id = data['block_id'] + fragment_num = data['num'] + cmd = data[''] + + +