92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
import socket, struct, os, select, math, hashlib, re#, brotli
|
|
|
|
max_frame_payload=508
|
|
header = "!L16s"
|
|
header_size = struct.calcsize(header)
|
|
max_payload = max_frame_payload - header_size
|
|
|
|
localport = 1337
|
|
localip = '0.0.0.0'
|
|
|
|
|
|
UDPServerSocket = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
|
|
|
|
UDPServerSocket.bind((localip, localport))
|
|
UDPServerSocket.setblocking(0)
|
|
|
|
print("UDP server up and listening")
|
|
|
|
packet_request_format = re.compile("[a-zA-Z0-9.-_ ]+/[0-9]+")
|
|
hash_request_format = re.compile("[a-zA-Z0-9.-_ ]+:")
|
|
|
|
while True:
|
|
|
|
ready = select.select([UDPServerSocket], [], [], 0.1)
|
|
|
|
if ready[0]:
|
|
|
|
bytesAddressPair = UDPServerSocket.recvfrom(max_frame_payload)
|
|
|
|
message = bytesAddressPair[0]
|
|
|
|
if ".." in message.decode():
|
|
continue
|
|
|
|
address = bytesAddressPair[1]
|
|
|
|
if re.match(packet_request_format, message.decode()):
|
|
filename = "./files/"+message.decode().split('/')[0]
|
|
if not os.path.isfile(filename):
|
|
continue
|
|
with open(filename, 'rb') as f:
|
|
part = int(message.decode().split('/')[1])
|
|
f.seek(part*max_payload)
|
|
data = f.read(max_payload)
|
|
checksum = hashlib.sha3_512(data).digest()[:16]
|
|
data = struct.pack(f"{header}{max_payload}s", part, checksum, data)
|
|
UDPServerSocket.sendto(data, address)
|
|
print(f'sent lost packet {message.decode()}')
|
|
continue
|
|
|
|
if re.match(hash_request_format, message.decode()):
|
|
print("hash request")
|
|
filename = "./files/"+message.decode().split(':')[0]
|
|
if not os.path.isfile(filename):
|
|
continue
|
|
sha512 = hashlib.sha1()
|
|
with open(filename, 'rb') as f:
|
|
data = f.read(max_payload)
|
|
sha512.update(data)
|
|
UDPServerSocket.sendto(struct.pack("64s", sha512.digest()), address)
|
|
continue
|
|
|
|
filename = "./files/"+message.decode()
|
|
if not os.path.isfile(filename):
|
|
continue
|
|
file_len = math.ceil(os.path.getsize(filename)/max_payload)
|
|
|
|
#calculate last packet size
|
|
last_packet_size = os.path.getsize(filename) % max_payload
|
|
if last_packet_size == 0:
|
|
last_packet_size = max_payload # if last packet is full, set size to max
|
|
#print(max_payload,last_packet_size)
|
|
|
|
tosend = struct.pack(f'!LH{len(message.decode())}s', file_len, last_packet_size, message)
|
|
|
|
UDPServerSocket.sendto(tosend, address)
|
|
|
|
sha512 = hashlib.sha512()
|
|
|
|
with open(filename, 'rb') as f:
|
|
for i in range(file_len):
|
|
fdata = f.read(max_payload)
|
|
sha512.update(fdata)
|
|
#if i == 5: #simulate lost packet
|
|
# continue
|
|
checksum = hashlib.sha3_512(fdata).digest()[:16]
|
|
data = struct.pack(f"{header}{max_payload}s", i, checksum, fdata)
|
|
#print("sending data",fdata)
|
|
UDPServerSocket.sendto(data, address)
|
|
print(f'{i+1}/{file_len}', end='\r')
|
|
UDPServerSocket.sendto(struct.pack("64s", sha512.digest()), address)
|
|
print(' '*100,end='\r') |