import asyncio import logging import ssl import time from dataclasses import dataclass from yuxi.channel.extensions.irc.protocol import parse_irc_line from yuxi.channel.extensions.irc.types import ResolvedIrcAccount logger = logging.getLogger(__name__) class IRCFatalError(Exception): pass @dataclass class IrcClient: account: ResolvedIrcAccount reader: asyncio.StreamReader writer: asyncio.StreamWriter current_nick: str connected_at: float = 0.0 async def _negotiate_caps( reader: asyncio.StreamReader, writer: asyncio.StreamWriter, requested: list[str], timeout: float, ) -> set[str]: acked: set[str] = set() send_line(writer, f"CAP REQ :{' '.join(requested)}") await writer.drain() try: while True: line_bytes = await asyncio.wait_for(reader.readline(), timeout=timeout) if not line_bytes: break line = line_bytes.decode("utf-8", errors="replace").rstrip("\r\n") parsed = parse_irc_line(line) if parsed.command == "CAP": subcmd = parsed.params[1] if len(parsed.params) > 1 else "" if subcmd in ("ACK", "NAK"): cap_list = parsed.trailing or (parsed.params[2] if len(parsed.params) > 2 else "") caps = {c.strip() for c in cap_list.split()} if subcmd == "ACK": acked |= caps return acked except asyncio.TimeoutError: logger.warning("CAP negotiation timed out") return acked async def connect_irc_client(account: ResolvedIrcAccount, timeout: float = 15.0) -> IrcClient: ssl_context = ssl.create_default_context() if account.tls else None reader, writer = await asyncio.wait_for( asyncio.open_connection(account.host, account.port, ssl=ssl_context), timeout=timeout, ) caps_to_request = [] if account.config.sasl_enabled and account.nickserv_password: caps_to_request.append("sasl") caps_ack: set[str] = set() if caps_to_request: send_line(writer, "CAP LS 302") await writer.drain() caps_ack = await _negotiate_caps(reader, writer, caps_to_request, timeout) if account.password: send_line(writer, f"PASS {account.password}") if "sasl" in caps_ack and account.nickserv_password: send_line(writer, "AUTHENTICATE PLAIN") await writer.drain() import base64 auth_str = f"\x00{account.username or account.nick}\x00{account.nickserv_password}" auth_b64 = base64.b64encode(auth_str.encode()).decode() send_line(writer, f"AUTHENTICATE {auth_b64}") await writer.drain() if caps_to_request: send_line(writer, "CAP END") await writer.drain() send_line(writer, f"NICK {account.nick}") send_line(writer, f"USER {account.username} 0 * :{account.realname}") await writer.drain() connected_at = time.monotonic() logger.info( "IRC connected to %s:%d (tls=%s) as %s", account.host, account.port, account.tls, account.nick, ) return IrcClient( account=account, reader=reader, writer=writer, current_nick=account.nick, connected_at=connected_at, ) async def read_lines(reader: asyncio.StreamReader, cancel_event: asyncio.Event): while not cancel_event.is_set(): try: line_bytes = await asyncio.wait_for(reader.readline(), timeout=300.0) except asyncio.TimeoutError: continue if not line_bytes: break line = line_bytes.decode("utf-8", errors="replace").rstrip("\r\n") yield line def send_line(writer: asyncio.StreamWriter, line: str) -> None: writer.write(f"{line}\r\n".encode()) async def send_privmsg(writer: asyncio.StreamWriter, target: str, text: str) -> None: send_line(writer, f"PRIVMSG {target} :{text}") await writer.drain() async def handle_ping(writer: asyncio.StreamWriter, payload: str | None) -> None: if payload: send_line(writer, f"PONG :{payload}") else: send_line(writer, "PONG") await writer.drain() async def join_channel(writer: asyncio.StreamWriter, channel: str) -> None: send_line(writer, f"JOIN {channel}") await writer.drain() async def part_channel(writer: asyncio.StreamWriter, channel: str, reason: str = "") -> None: if reason: send_line(writer, f"PART {channel} :{reason}") else: send_line(writer, f"PART {channel}") await writer.drain() async def identify_nickserv(writer: asyncio.StreamWriter, password: str) -> None: send_line(writer, f"PRIVMSG NickServ :IDENTIFY {password}") await writer.drain() async def ghost_nick(writer: asyncio.StreamWriter, nick: str, password: str) -> None: send_line(writer, f"PRIVMSG NickServ :GHOST {nick} {password}") await writer.drain() async def handle_nick_collision( writer: asyncio.StreamWriter, account: ResolvedIrcAccount, current_nick: str, ) -> str: base_nick = account.nick if account.nickserv_password: try: await ghost_nick(writer, base_nick, account.nickserv_password) await asyncio.sleep(1) send_line(writer, f"NICK {base_nick}") await writer.drain() logger.info("Nick collision: GHOST sent for %s, reclaiming %s", base_nick, base_nick) return base_nick except Exception as e: logger.warning("Nick collision: GHOST failed for %s: %s", base_nick, e) fallback = f"{base_nick}_"[:30] send_line(writer, f"NICK {fallback}") await writer.drain() logger.info("Nick collision: fallback nick %s", fallback) return fallback async def graceful_disconnect(client: IrcClient) -> None: try: send_line(client.writer, "QUIT :shutdown") await client.writer.drain() except Exception: pass try: client.writer.close() await client.writer.wait_closed() except Exception: pass logger.info("IRC disconnected from %s:%d", client.account.host, client.account.port)