black-mamba/mods/udp_srv.py

79 lines
2.1 KiB
Python

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['']