import asyncio
import getpass
import sys
import logging
import traceback
from typing import Dict, Any
import aiohttp
from snekbot.rpc import RPC
from snekbot.bot import Bot
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger: logging.Logger = logging.getLogger("snek_cli")
def markdown_to_ansi(text: str) -> str:
import re
def heading_sub(match: re.Match) -> str:
hashes = match.group(1)
content = match.group(2).strip()
level = len(hashes)
if level == 1:
return f"\033[1;4m{content}\033[0m" # Bold + Underline
elif level == 2:
return f"\033[1m{content}\033[0m"
elif level == 3:
return f"\033[1;36m{content}\033[0m"
elif level == 4:
return f"\033[36m{content}\033[0m"
elif level == 5:
return f"\033[2;36m{content}\033[0m"
else:
return content
text = re.sub(r"^(#{1,5})\s+(.+)$", heading_sub, text, flags=re.MULTILINE)
text = re.sub(r"^\s*---+\s*$", "\033[38;5;244m" + "─" * 40 + "\033[0m", text, flags=re.MULTILINE)
def ordered_list_sub(match: re.Match) -> str:
idx = match.group(1)
content = match.group(2)
return f"\033[38;5;33m{idx}.\033[0m {content}"
text = re.sub(r"^(\d+)\.\s+(.*)$", ordered_list_sub, text, flags=re.MULTILINE)
def unordered_list_sub(match: re.Match) -> str:
bullet = match.group(1)
content = match.group(2)
return f"\033[38;5;33m•\033[0m {content}"
text = re.sub(r"^(\s*[-*+])\s+(.*)$", unordered_list_sub, text, flags=re.MULTILINE)
text = re.sub(r"`([^`]+)`", r"\033[38;5;244m\033[48;5;236m \1 \033[0m", text)
text = re.sub(r"\*\*(.*?)\*\*", r"\033[1m\1\033[0m", text)
text = re.sub(r"\*(.*?)\*", r"\033[3m\1\033[0m", text)
return text
class CliClient(Bot):
active_channel_uid: Any = None
_channel_cache: Dict[Any, Any] = {}
def _display(self, message: str) -> None:
ansi_message = markdown_to_ansi(message)
sys.stdout.write('\r' + ' ' * 80 + '\r')
sys.stdout.write(ansi_message + '\n')
sys.stdout.write(f'[{self._get_active_channel_name()}]> ')
sys.stdout.flush()
def _get_active_channel_name(self) -> str:
if self.active_channel_uid and self.active_channel_uid in self._channel_cache:
return self._channel_cache[self.active_channel_uid]['name']
return "No Channel"
async def send_message(self, channel_uid: Any, message: str) -> bool:
if self.ws and not self.ws.closed:
payload: Dict[str, Any] = {
"method": "send_message",
"args": [channel_uid, message, True],
"kwargs": {},
"callId": None
}
await self.ws.send_json(payload)
return True
logger.error("Cannot send message, WebSocket is not connected.")
return False
async def _network_loop(self) -> None:
try:
async with aiohttp.ClientSession() as session:
async with session.ws_connect(self.url) as ws:
self.ws = ws
self.rpc = RPC(self.ws)
await self.rpc.login(self.username, self.password)
self.user: Dict[str, Any] = await self.rpc.get_user(None)
await self.on_init()
while not self.ws.closed:
data = await self.rpc.receive()
if not data:
break
event: str = "?"
try:
event = data.event
except AttributeError:
pass
try:
message: str = data.message.strip()
event = "message"
except AttributeError:
pass
if event == "?":
continue
elif event == "message":
if not data.is_final:
continue
if data.username == self.user["username"]:
continue
else:
await self.on_message(data.username, data.user_nick, data.channel_uid, message)
else:
try:
await getattr(self, "on_" + data.event)(**data.data)
except AttributeError:
logger.debug("Not implemented event: " + event)
except Exception:
logger.error("Network loop disconnected with an error.")
traceback.print_exc()
async def run(self) -> None:
network_task = asyncio.create_task(self._network_loop())
while not self.user:
if network_task.done():
network_task.result()
return
await asyncio.sleep(0.1)
input_task = asyncio.create_task(self._input_loop())
done, pending = await asyncio.wait(
[network_task, input_task],
return_when=asyncio.FIRST_COMPLETED
)
for task in pending:
task.cancel()
async def on_init(self) -> None:
try:
logger.info("Successfully logged in as %s.", self.user['username'])
except KeyError:
raise ValueError("\033[91mUsername or password incorrect!\033[0m")
channels = await self.get_channels()
for channel in channels:
self._channel_cache[channel['uid']] = channel
self._display(f"**Welcome, {self.user['nick']}!** Your available commands are:")
self._display("> /join <channel_name>, /channels, /quit")
await self.on_idle()
async def on_message(self, username: str, user_nick: str, channel_uid: Any, message: str) -> None:
sys.stdout.write('\a')
sys.stdout.flush()
channel_name = self._channel_cache.get(channel_uid, {}).get('name', 'unknown')
self._display(f"**[{channel_name}]** *<{user_nick}>*: {message}")
async def on_own_message(self, channel_uid: Any, message: str) -> None:
channel_name = self._channel_cache.get(channel_uid, {}).get('name', 'unknown')
full_message = f"**[{channel_name}]** *<{self.user['nick']}>*: {message}"
self._display(full_message)
async def on_idle(self) -> None:
sys.stdout.write(f'[{self._get_active_channel_name()}]> ')
sys.stdout.flush()
async def _handle_command(self, line: str) -> bool:
command, *args = line.strip().split()
if command == "/join":
if not args:
self._display("> Usage: `/join <channel_name>`")
return True
channel_name_to_join = args[0]
for uid, channel_data in self._channel_cache.items():
if channel_data['name'] == channel_name_to_join:
self.active_channel_uid = uid
self._display(f"*Active channel set to* **'{channel_name_to_join}'**.")
return True
self._display(f"*** Error: Channel '{channel_name_to_join}' not found.")
elif command == "/channels":
names = [ch['name'] for ch in self._channel_cache.values()]
self._display(f"*Available channels*: `{'`, `'.join(names)}`")
elif command == "/quit":
self._display("*Disconnecting. Goodbye!*")
return False
else:
self._display(f"*** Unknown command: {command}")
return True
async def _input_loop(self) -> None:
loop = asyncio.get_event_loop()
try:
while True:
line = await loop.run_in_executor(None, sys.stdin.readline)
line = line.strip()
if not line:
continue
if line.startswith("/"):
if not await self._handle_command(line):
break
else:
if not self.active_channel_uid:
self._display("> No active channel. Use `/join <channel_name>` to start talking.")
continue
await self.send_message(self.active_channel_uid, line)
await self.on_own_message(self.active_channel_uid, line)
finally:
if self.ws and not self.ws.closed:
await self.ws.close()
if __name__ == "__main__":
print("--- Snekbot CLI Client ---")
url = "wss://snek.molodetz.nl/rpc.ws"
username = input("Username: ")
password = getpass.getpass("Password: ")
client = CliClient(username, password, url=url)
try:
asyncio.run(client.run())
except KeyboardInterrupt:
print("\nClient shut down by user.")
except Exception as e:
logger.error(f"An unexpected error occurred: {e}")
Reference in New Issue View Git Blame Copy Permalink