diff --git a/relay.py b/relay.py new file mode 100644 index 0000000..2820c75 --- /dev/null +++ b/relay.py @@ -0,0 +1,98 @@ +import asyncio + +import telnetlib3 +from telnetlib3 import TelnetReaderUnicode, TelnetWriterUnicode + +from tngame.utils import setup_logger + +log = setup_logger(True) + + +async def shell(reader: TelnetReaderUnicode, writer: TelnetWriterUnicode): + # Get the size of the terminal + async def get_size() -> tuple[int, int]: + writer.write('\x1b[18t') + await writer.drain() + size = await reader.read(100) + + # Parse the size, it's in the format of \x1b[8;{height};{width}t + height, width = size.split(';')[1:3] + height, width = int(height), int(width[:-1]) + + return height, width + + # Run tngame-rs + h, w = await get_size() + proc = await asyncio.create_subprocess_exec( + './tngame-rs/target/release/tngame-rs', + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + env={'TN_TERM_SIZE': f'{w}x{h}'} + ) + + async def relay_stdout(): + # Listen input + while not writer.is_closing(): + # Close if subprocess is closed + if proc.returncode is not None: + break + + # Read output + try: + out = await proc.stdout.readuntil(b'\00\00\00') + if out: + # print("Sending output from tngame-rs to telnet:", repr(out)) + writer.write(out.decode()) + await writer.drain() + except asyncio.IncompleteReadError: + break + + print("Closing reader.") + raise asyncio.CancelledError("Closing reader.") + + async def relay_stdin(): + while not reader.connection_closed and not reader.at_eof(): + # Close if subprocess is closed + if proc.returncode is not None: + break + + inp: str = await reader.read(3) + if inp: + # print("Sending input from telnet to tngame-rs:", repr(inp)) + proc.stdin.write(inp.encode()) + await proc.stdin.drain() + + print("Closing writer.") + raise asyncio.CancelledError("Closing writer.") + + tasks = [asyncio.create_task(relay_stdout()), asyncio.create_task(relay_stdin())] + try: + await asyncio.gather(*tasks) + except asyncio.CancelledError: + for task in tasks: + task.cancel() + + if proc.returncode is None: + print("Killing process.") + proc.kill() + + # Print exit message + writer.write("\r\nThanks for visiting <3\r\n") + + # Make cursor visible again + writer.write("\x1b[?25h") + + await writer.drain() + writer.close() + reader.close() + print("Connection closed.") + + +if __name__ == '__main__': + # Create a new event loop, start the server and wait for it to close + print("Starting server on port 2323.") + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + coro = telnetlib3.create_server(port=2323, shell=shell) + server = loop.run_until_complete(coro) + loop.run_until_complete(server.wait_closed())