import asyncio
import base64
import socket

def shift(b, s):
    r = b + s
    if r >= 256:
        r -= 256
    elif r < 0:
        r += 256
    return r

def shift_seq(seq, byte_shift):
    result = bytearray(seq)
    for i in range(min(len(seq),128)):
        result[i] = (result[i] + byte_shift) & 0xff
    return result

def encrypt(data):
    return shift_seq(data, 15)

def decrypt(data):
    return shift_seq(data, -15)


class Duplex:
    def __init__(self, s1, s1_remote, s1_transform, s2, s2_remote, s2_transform, loop, mtu):
        self.mtu = mtu
        self.loop = loop

        #loop.create_task(self.stream(s1,s2, s2_remote, s2_transform))
        loop.create_task(self.stream(s2,s1, s1_remote, s1_transform))

    async def stream(self, rx, tx, tx_addr, transform):
        print("Starting stream", rx.getsockname(), '>>',  tx_addr)
        while True:
            data = await self.loop.sock_recv(rx, self.mtu)
            #print(rx.getsockname(),'>>', tx_addr, len(data), data[:20].hex(), transform(data)[:20].hex())
            await self.loop.sock_sendto(tx, transform(data), tx_addr)



async def main():
    rx = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    rx.setblocking(False)
    rx.bind((RX_IP, RX_PORT))

    connections = {}

    print("App started")
    loop = asyncio.get_running_loop()

    while True:
        data, addr = await loop.sock_recvfrom(rx, RX_MTU) 
        #print("Received "+str(len(data))+" bytes", addr, data[:20])

        if addr in connections.keys():
            tx = connections[addr]
        else:
            connections[addr] = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            tx = connections[addr]
            tx.setblocking(False)
            tx.bind(tx.getsockname())
            Duplex(rx, addr, RX_ALGO, connections[addr], (TX_IP, TX_PORT), TX_ALGO, loop, RX_MTU)

        try:
            tx.sendto(TX_ALGO(data), (TX_IP, TX_PORT))
        except Exception:
            print("reconnect", addr)
            del connections[addr]
        #print("Sent", TX_ALGO(data)[:20], "to", (TX_IP, TX_PORT))

if __name__ == '__main__':
    RX_IP = "127.0.0.1"
    RX_PORT = 5005
    RX_MTU = 1500
    RX_ALGO = encrypt

    TX_IP = "194.135.105.21"
    TX_PORT = 25565
    TX_ALGO = decrypt

    asyncio.run(main())