From f47d82733cac8c0a7d6d27bab3eea724e4da55ef Mon Sep 17 00:00:00 2001
From: alemidev <me@alemi.dev>
Date: Mon, 22 Nov 2021 12:57:10 +0100
Subject: [PATCH] improved parse_chat

---
 aiocraft/__main__.py     |  2 +-
 aiocraft/client.py       |  7 ++++---
 aiocraft/util/helpers.py | 35 +++++++++++++++++++++++------------
 3 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/aiocraft/__main__.py b/aiocraft/__main__.py
index 184eb77..1b67eb6 100644
--- a/aiocraft/__main__.py
+++ b/aiocraft/__main__.py
@@ -37,7 +37,7 @@ if __name__ == "__main__":
 
 		@client.on_packet(PacketChat, ConnectionState.PLAY)
 		async def print_chat(packet: PacketChat):
-			msg = parse_chat(json.loads(packet.message), color=color)
+			msg = parse_chat(packet.message, ansi_color=color)
 			print(f"[{packet.position}] {msg}")
 
 		client.run() # will block and start asyncio event loop
diff --git a/aiocraft/client.py b/aiocraft/client.py
index f66d40b..83dfaea 100644
--- a/aiocraft/client.py
+++ b/aiocraft/client.py
@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import json
 import uuid
 
 from dataclasses import dataclass
@@ -20,7 +21,7 @@ from .mc.proto.login.serverbound import PacketLoginStart, PacketEncryptionBegin
 from .mc.proto.login.clientbound import (
 	PacketCompress, PacketDisconnect, PacketEncryptionBegin, PacketLoginPluginRequest, PacketSuccess
 )
-from .util import encryption
+from .util import encryption, helpers
 
 LOGGER = logging.getLogger(__name__)
 
@@ -246,7 +247,7 @@ class MinecraftClient(CallbacksHolder, Runnable):
 				self._logger.info("Login success, joining world...")
 				return True
 			elif isinstance(packet, PacketDisconnect):
-				self._logger.error("Kicked while logging in")
+				self._logger.error("Kicked while logging in : %s", helpers.parse_chat(packet.reason))
 				break
 		return False
 
@@ -263,7 +264,7 @@ class MinecraftClient(CallbacksHolder, Runnable):
 					keep_alive_packet = PacketKeepAliveResponse(340, keepAliveId=packet.keepAliveId)
 					await self.dispatcher.write(keep_alive_packet)
 			elif isinstance(packet, PacketKickDisconnect):
-				self._logger.error("Kicked while in game")
+				self._logger.error("Kicked while in game : %s", helpers.parse_chat(packet.reason))
 				break
 			self.run_callbacks(Packet, packet)
 			self.run_callbacks(type(packet), packet)
diff --git a/aiocraft/util/helpers.py b/aiocraft/util/helpers.py
index a60cab7..4cc75be 100644
--- a/aiocraft/util/helpers.py
+++ b/aiocraft/util/helpers.py
@@ -1,4 +1,8 @@
-from termcolor import colored
+import json
+
+from typing import Union
+
+from termcolor import colored # TODO don't use a lib and put ANSI escaped by hand maybe?
 
 _EQUIVALENTS = {
 	"dark_red" : "red",
@@ -32,18 +36,25 @@ def _parse_formatted_block(msg:dict) -> str:
 	else:
 		return colored(msg["text"], "white", attrs=attr)
 
-def parse_chat(msg:dict, color:bool=True) -> str:
+def parse_chat(msg:Union[dict,str], ansi_color:bool=False) -> str:
 	"""Recursive function to parse minecraft chat json, with optional colors"""
+	if isinstance(msg, str):
+		try:
+			data = json.loads(msg)
+		except ValueError:
+			return str(msg) # It's not json, it's already plaintext
+	else:
+		data = msg
 	out = ""
-	if "text" in msg:
-		if color:
-			out += _parse_formatted_block(msg)
+	if "text" in data:
+		if ansi_color:
+			out += _parse_formatted_block(data)
 		else:
-			out += msg["text"]
-	if "with" in msg:
-		for elem in msg["with"]:
-			out += parse_chat(elem, color)
-	if "extra" in msg:
-		for elem in msg["extra"]:
-			out += parse_chat(elem, color)
+			out += data["text"]
+	if "with" in data:
+		for elem in data["with"]:
+			out += parse_chat(elem, ansi_color)
+	if "extra" in data:
+		for elem in data["extra"]:
+			out += parse_chat(elem, ansi_color)
 	return out