ForcePilot/backend/package/yuxi/channel/extensions/irc/client.py

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)