import asyncio
import Utils
import websockets
import functools
from copy import deepcopy
from typing import List, Any, Iterable
from NetUtils import decode, encode, JSONtoTextParser, JSONMessagePart, NetworkItem
from MultiServer import Endpoint
from CommonClient import CommonContext, gui_enabled, ClientCommandProcessor, logger, get_base_parser
DEBUG = False
class AHITJSONToTextParser(JSONtoTextParser):
def _handle_color(self, node: JSONMessagePart):
return self._handle_text(node) # No colors for the in-game text
class AHITCommandProcessor(ClientCommandProcessor):
def _cmd_ahit(self):
"""Check AHIT Connection State"""
if isinstance(self.ctx, AHITContext):
logger.info(f"AHIT Status: {self.ctx.get_ahit_status()}")
class AHITContext(CommonContext):
command_processor = AHITCommandProcessor
game = "A Hat in Time"
def __init__(self, server_address, password):
super().__init__(server_address, password)
self.proxy = None
self.proxy_task = None
self.gamejsontotext = AHITJSONToTextParser(self)
self.autoreconnect_task = None
self.endpoint = None
self.items_handling = 0b111
self.room_info = None
self.connected_msg = None
self.game_connected = False
self.awaiting_info = False
self.full_inventory: List[Any] = []
self.server_msgs: List[Any] = []
async def server_auth(self, password_requested: bool = False):
if password_requested and not self.password:
await super(AHITContext, self).server_auth(password_requested)
await self.get_username()
await self.send_connect()
def get_ahit_status(self) -> str:
if not self.is_proxy_connected():
return "Not connected to A Hat in Time"
return "Connected to A Hat in Time"
async def send_msgs_proxy(self, msgs: Iterable[dict]) -> bool:
""" `msgs` JSON serializable """
if not self.endpoint or not self.endpoint.socket.open or self.endpoint.socket.closed:
return False
if DEBUG:
logger.info(f"Outgoing message: {msgs}")
await self.endpoint.socket.send(msgs)
return True
async def disconnect(self, allow_autoreconnect: bool = False):
await super().disconnect(allow_autoreconnect)
async def disconnect_proxy(self):
if self.endpoint and not self.endpoint.socket.closed:
await self.endpoint.socket.close()
if self.proxy_task is not None:
await self.proxy_task
def is_connected(self) -> bool:
return self.server and self.server.socket.open
def is_proxy_connected(self) -> bool:
return self.endpoint and self.endpoint.socket.open
def on_print_json(self, args: dict):
text = self.gamejsontotext(deepcopy(args["data"]))
msg = {"cmd": "PrintJSON", "data": [{"text": text}], "type": "Chat"}
self.server_msgs.append(encode([msg]))
if self.ui:
self.ui.print_json(args["data"])
else:
text = self.jsontotextparser(args["data"])
logger.info(text)
def update_items(self):
# just to be safe - we might still have an inventory from a different room
if not self.is_connected():
return
self.server_msgs.append(encode([{"cmd": "ReceivedItems", "index": 0, "items": self.full_inventory}]))
def on_package(self, cmd: str, args: dict):
if cmd == "Connected":
self.connected_msg = encode([args])
if self.awaiting_info:
self.server_msgs.append(self.room_info)
self.update_items()
self.awaiting_info = False
elif cmd == "ReceivedItems":
if args["index"] == 0:
self.full_inventory.clear()
for item in args["items"]:
self.full_inventory.append(NetworkItem(*item))
self.server_msgs.append(encode([args]))
elif cmd == "RoomInfo":
self.seed_name = args["seed_name"]
self.room_info = encode([args])
else:
if cmd != "PrintJSON":
self.server_msgs.append(encode([args]))
def run_gui(self):
from kvui import GameManager
class AHITManager(GameManager):
logging_pairs = [
("Client", "Archipelago")
]
base_title = "Archipelago A Hat in Time Client"
self.ui = AHITManager(self)
self.ui_task = asyncio.create_task(self.ui.async_run(), name="UI")
async def proxy(websocket, path: str = "/", ctx: AHITContext = None):
ctx.endpoint = Endpoint(websocket)
try:
await on_client_connected(ctx)
if ctx.is_proxy_connected():
async for data in websocket:
if DEBUG:
logger.info(f"Incoming message: {data}")
for msg in decode(data):
if msg["cmd"] == "Connect":
# Proxy is connecting, make sure it is valid
if msg["game"] != "A Hat in Time":
logger.info("Aborting proxy connection: game is not A Hat in Time")
await ctx.disconnect_proxy()
break
if ctx.seed_name:
seed_name = msg.get("seed_name", "")
if seed_name != "" and seed_name != ctx.seed_name:
logger.info("Aborting proxy connection: seed mismatch from save file")
logger.info(f"Expected: {ctx.seed_name}, got: {seed_name}")
text = encode([{"cmd": "PrintJSON",
"data": [{"text": "Connection aborted - save file to seed mismatch"}]}])
await ctx.send_msgs_proxy(text)
await ctx.disconnect_proxy()
break
if ctx.connected_msg and ctx.is_connected():
await ctx.send_msgs_proxy(ctx.connected_msg)
ctx.update_items()
continue
if not ctx.is_proxy_connected():
break
await ctx.send_msgs([msg])
except Exception as e:
if not isinstance(e, websockets.WebSocketException):
logger.exception(e)
finally:
await ctx.disconnect_proxy()
async def on_client_connected(ctx: AHITContext):
if ctx.room_info and ctx.is_connected():
await ctx.send_msgs_proxy(ctx.room_info)
else:
ctx.awaiting_info = True
async def proxy_loop(ctx: AHITContext):
try:
while not ctx.exit_event.is_set():
if len(ctx.server_msgs) > 0:
for msg in ctx.server_msgs:
await ctx.send_msgs_proxy(msg)
ctx.server_msgs.clear()
await asyncio.sleep(0.1)
except Exception as e:
logger.exception(e)
logger.info("Aborting AHIT Proxy Client due to errors")
def launch():
async def main():
parser = get_base_parser()
args = parser.parse_args()
ctx = AHITContext(args.connect, args.password)
logger.info("Starting A Hat in Time proxy server")
ctx.proxy = websockets.serve(functools.partial(proxy, ctx=ctx),
host="localhost", port=11311, ping_timeout=999999, ping_interval=999999)
ctx.proxy_task = asyncio.create_task(proxy_loop(ctx), name="ProxyLoop")
if gui_enabled:
ctx.run_gui()
ctx.run_cli()
await ctx.proxy
await ctx.proxy_task
await ctx.exit_event.wait()
Utils.init_logging("AHITClient")
# options = Utils.get_options()
import colorama
colorama.init()
asyncio.run(main())
colorama.deinit()