shelfmark/shelfmark/release_sources/irc/dcc.py

"""DCC (Direct Client-to-Client) protocol implementation.

Handles DCC SEND file transfers used by IRC bots to send files.
"""

import re
import socket
import struct
from dataclasses import dataclass
from ipaddress import ip_address
from pathlib import PureWindowsPath
from typing import TYPE_CHECKING

from shelfmark.core.logger import setup_logger

if TYPE_CHECKING:
    from collections.abc import Callable
    from pathlib import Path
    from threading import Event

logger = setup_logger(__name__)

# Regex to parse DCC SEND messages - handles quoted filenames
# Format: DCC SEND "filename.epub" 2760158537 2050 2321788
#                  |                |          |    |
#                  filename         IP(int)    port size
DCC_REGEX = re.compile(r'DCC SEND "?(.+[^"])"?\s(\d+)\s+(\d+)\s+(\d+)\s*')

# Buffer size for DCC transfers - 4096 bytes provides good performance
BUFFER_SIZE = 4096


@dataclass
class DCCOffer:
    """Parsed DCC SEND offer."""

    filename: str
    ip: str
    port: int
    size: int

    @property
    def address(self) -> tuple[str, int]:
        """Return (ip, port) tuple for socket.connect()."""
        return (self.ip, self.port)


class DCCError(Exception):
    """Base exception for DCC operations."""


class DCCParseError(DCCError):
    """Failed to parse DCC SEND string."""


class DCCSizeError(DCCError):
    """Downloaded size doesn't match expected size."""


class DCCConnectionError(DCCError):
    """Failed to connect to DCC sender."""


class DCCSecurityError(DCCError):
    """Rejected unsafe DCC offer metadata."""


def int_to_ip(ip_int: int) -> str:
    """Convert 32-bit integer (DCC format) to dotted IP notation."""
    packed = struct.pack(">I", ip_int)
    return ".".join(str(b) for b in packed)


def parse_dcc_send(text: str) -> DCCOffer:
    """Parse a DCC SEND message into a DCCOffer. Raises DCCParseError on failure."""
    match = DCC_REGEX.search(text)
    if not match:
        msg = f"Invalid DCC SEND format: {text[:100]}"
        raise DCCParseError(msg)

    filename = match.group(1).strip('"')
    ip_int = int(match.group(2))
    port = int(match.group(3))
    size = int(match.group(4))
    try:
        ip = int_to_ip(ip_int)
    except struct.error as e:
        msg = f"Invalid DCC IP integer: {ip_int}"
        raise DCCParseError(msg) from e

    return DCCOffer(
        filename=safe_dcc_filename(filename),
        ip=ip,
        port=port,
        size=size,
    )


def safe_dcc_filename(filename: str) -> str:
    """Return a DCC filename that cannot escape its destination directory."""
    safe_name = filename.strip()
    windows_path = PureWindowsPath(safe_name)
    if (
        not safe_name
        or safe_name in {".", ".."}
        or "/" in safe_name
        or "\\" in safe_name
        or windows_path.drive
    ):
        msg = f"Rejected unsafe DCC filename: {filename!r}"
        raise DCCSecurityError(msg)
    return safe_name


def validate_dcc_endpoint(offer: DCCOffer) -> None:
    """Reject DCC endpoints that can target local/internal network services."""
    if not 1 <= offer.port <= 65535:
        msg = f"Rejected invalid DCC port: {offer.port}"
        raise DCCSecurityError(msg)

    try:
        address = ip_address(offer.ip)
    except ValueError as e:
        msg = f"Rejected invalid DCC IP address: {offer.ip}"
        raise DCCSecurityError(msg) from e

    if not address.is_global:
        msg = f"Rejected non-public DCC endpoint: {offer.ip}"
        raise DCCSecurityError(msg)


def download_dcc(
    offer: DCCOffer,
    dest_path: Path,
    progress_callback: Callable[[float], None] | None = None,
    cancel_flag: Event | None = None,
    timeout: float = 30.0,
) -> None:
    """Download file via DCC protocol to dest_path. Raises DCCError on failure."""
    validate_dcc_endpoint(offer)
    logger.info("DCC connecting to %s:%s for %s", offer.ip, offer.port, offer.filename)

    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        sock.connect(offer.address)
    except OSError as e:
        msg = f"Failed to connect to {offer.ip}:{offer.port}: {e}"
        raise DCCConnectionError(msg) from e

    try:
        received = 0
        last_progress = -1

        with dest_path.open("wb") as f:
            while received < offer.size:
                # Check for cancellation
                if cancel_flag and cancel_flag.is_set():
                    logger.info("DCC download cancelled")
                    return

                # Read chunk
                try:
                    chunk = sock.recv(BUFFER_SIZE)
                except TimeoutError as e:
                    msg = f"Timeout reading from {offer.ip}:{offer.port}"
                    raise DCCError(msg) from e

                if not chunk:
                    # Connection closed prematurely
                    break

                f.write(chunk)
                received += len(chunk)

                # Report progress (every 1%)
                if progress_callback:
                    progress = int((received / offer.size) * 100)
                    if progress != last_progress:
                        progress_callback(progress)
                        last_progress = progress

        # Verify downloaded size matches expected
        if received != offer.size:
            msg = f"Size mismatch: expected {offer.size} bytes, got {received}"
            raise DCCSizeError(msg)

        logger.info("DCC download complete: %s bytes", received)

    finally:
        sock.close()