79 lines
2.1 KiB
Python
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['']
|
|
|
|
|
|
|