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]:
        try:
            bytesAddressPair = UDPServerSocket.recvfrom(max_frame_payload)
        except:
            continue

        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)
        
        filename_checksum = sum(message)

        tosend = struct.pack(f'!IHQ', file_len, last_packet_size, filename_checksum)

        UDPServerSocket.sendto(tosend, address)

        sha512 = hashlib.sha512()

        try:
            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)
        except:
            pass
        print(' '*100,end='\r')