diff --git a/irc/client_aio.py b/irc/client_aio.py index f86962d..46a9835 100644 --- a/irc/client_aio.py +++ b/irc/client_aio.py @@ -27,8 +27,6 @@ that dispatches events to instance methods is included. Current limitations: - * DCC chat has not yet been implemented - * DCC file transfers are not suppored * RFCs 2810, 2811, 2812, and 2813 have not been considered. Notes: @@ -39,15 +37,20 @@ import asyncio import logging +import socket import threading +import warnings +from jaraco.stream import buffer from . import connection from .client import ( + DCCConnectionError, Event, Reactor, ServerConnection, ServerNotConnectedError, SimpleIRCClient, + DCCConnection, _ping_ponger, ) @@ -212,6 +215,239 @@ def disconnect(self, message=""): self._handle_event(Event("disconnect", self.server, "", [message])) +class DCCProtocol(IrcProtocol): + """ + A protocol for handling DCC connections. + + Currently, DCCProtocol uses the same methods as `IrcProtocol` + for handling incoming data. This should be fine for most use + cases, but in the unlikely event that a DCC connection needs to + handle incoming data in a different way than an IRC connection, + this class will need to be overridden. + """ + + def connection_made(self, transport): + if self.connection.passive and not self.connection.connected: + self.connection.transport = transport + self.connection.connected = True + self.connection.peeraddress, self.connection.peerport = transport.get_extra_info('peername') + log.debug("DCC connection from %s:%d", self.connection.peeraddress, self.connection.peerport) + self.connection.reactor._handle_event( + self.connection, Event("dcc_connect", self.connection.peeraddress, None, None) + ) + if hasattr(self.connection, 'server') and self.connection.server: + self.connection.server.close() + return + + # For active connections, ensure transport is set if not already + if not getattr(self.connection, 'transport', None): + self.connection.transport = transport + + +class AioDCCConnection(DCCConnection): + """ + An asyncio-based DCCConnection. + + This class overrides select-based methods with asyncio-based ones. + """ + + reactor: "AioReactor" + buffer_class = buffer.LineBuffer + + protocol_class = DCCProtocol + protocol: DCCProtocol + socket: None + connected: bool + passive: bool + peeraddress: str + peerport: int + + async def connect( + self, address: str, port: int, connect_factory: connection.AioFactory = connection.AioFactory() + ) -> "AioDCCConnection": + """Connect/reconnect to a DCC peer. + + Arguments: + address -- Host/IP address of the peer. + port -- The port number to connect to. + connect_factory -- A callable that takes the event loop and the + server address, and returns a connection (with a socket interface) + + Returns the DCCConnection object. + """ + self.peeraddress = address + self.peerport = port + self.handlers = {} + self.buffer = self.buffer_class() + + self.connect_factory = connect_factory + protocol_instance = self.protocol_class(self, self.reactor.loop) + try: + connection = self.connect_factory(protocol_instance, (self.peeraddress, self.peerport)) + transport, protocol = await connection + except OSError as ex: + raise DCCConnectionError(f"Couldn't connect to socket: {ex}") from ex + + self.transport = transport + self.protocol = protocol + + self.connected = True + self.reactor._on_connect(self.protocol, self.transport) + return self + + async def listen(self, addr=None, port=None, ipv6=False) -> "AioDCCConnection": + """Wait for a connection/reconnection from a DCC peer. + + Returns the DCCConnection object. + + The local IP address and port are available as + self.localaddress and self.localport. After connection from a + peer, the peer address and port are available as + self.peeraddress and self.peerport. + + Arguments: + addr -- Host string or (host, port) tuple to bind to. + If a tuple, the port is only used if `port` is None. + port -- Port to listen on. Can be an int, a (min, max) tuple + to try a range, or a list of ports to try in order. + Overrides the port in `addr` if both are provided. + ipv6 -- Use IPv6 if True. + """ + self.passive = True + self.handlers = {} + self.buffer = self.buffer_class() + + # Resolve host and default port from addr + if addr is None: + host = socket.gethostbyname(socket.gethostname()) + addr_port = 0 + elif isinstance(addr, str): + host = addr + addr_port = 0 + else: + host, addr_port = addr + + # port parameter overrides addr port if specified + if port is None: + port = addr_port + + def factory(): + return self.protocol_class(self, self.reactor.loop) + + family = socket.AF_INET6 if ipv6 else socket.AF_INET + + # Build iterable of ports to try + if isinstance(port, int): + ports = [port] + elif isinstance(port, tuple): + ports = range(port[0], port[1] + 1) + else: + ports = port # assume list/iterable + + last_error = None + for try_port in ports: + try: + self.server = await self.reactor.loop.create_server( + factory, host, try_port, family=family + ) + break + except OSError as ex: + last_error = ex + continue + else: + raise DCCConnectionError(f"Couldn't bind socket: {last_error}") from last_error + + # Get the actual bound address and port + socket_obj = self.server.sockets[0] + self.localaddress, self.localport = socket_obj.getsockname() + + return self + + def disconnect(self, message: str = "") -> None: + """Hang up the connection and close the object. + + Arguments: + + message -- Quit message. + """ + try: + del self.connected + except AttributeError: + return + + try: + if hasattr(self, 'server') and self.server: + self.server.close() + except AttributeError: + pass + + try: + self.transport.close() + except AttributeError: + pass + + self.reactor._handle_event( + self, Event("dcc_disconnect", self.peeraddress, "", [message]) + ) + self.reactor._remove_connection(self) + + def process_data(self, new_data: bytes) -> None: + """ + handles incoming data from the `DCCProtocol` connection. + """ + + if self.dcctype == "chat": + self.buffer.feed(new_data) + + chunks = list(self.buffer) + + if len(self.buffer) > 2**14: + # Bad peer! Naughty peer! + log.info( + "Received >16k from a peer without a newline; " "disconnecting." + ) + self.disconnect() + return + else: + chunks = [new_data] + + command = "dccmsg" + prefix = self.peeraddress + target = None + for chunk in chunks: + log.debug("FROM PEER: %s", chunk) + arguments = [chunk] + log.debug( + "command: %s, source: %s, target: %s, arguments: %s", + command, + prefix, + target, + arguments, + ) + event = Event(command, prefix, target, arguments) + self.reactor._handle_event(self, event) + + def privmsg(self, text: str) -> None: + """ + Send text to DCC peer. + + The text will be padded with a newline if it's a DCC CHAT session. + """ + if self.dcctype == 'chat': + text += '\n' + return self.send_bytes(self.encode(text)) + + def send_bytes(self, data: bytes) -> None: + """ + Send data to DCC peer. + """ + try: + self.transport.write(data) + log.debug("TO PEER: %r\n", data) + except (OSError, AttributeError): + self.disconnect("Connection reset by peer.") + + class AioReactor(Reactor): """ Processes message from on or more asyncio-based IRC server connections. @@ -248,6 +484,7 @@ async def my_repeating_message(connection): """ connection_class = AioConnection + dcc_connection_class = AioDCCConnection def __do_nothing(*args, **kwargs): pass @@ -275,6 +512,21 @@ class definied above. """ self.loop.run_forever() + def dcc(self, dcctype="chat"): + """Creates and returns a DCCConnection object. + + Arguments: + + dcctype -- "chat" for DCC CHAT connections or "raw" for + DCC SEND (or other DCC types). If "chat", + incoming data will be split in newline-separated + chunks. If "raw", incoming data is not touched. + """ + with self.mutex: + conn = self.dcc_connection_class(self, dcctype) + self.connections.append(conn) + return conn + class AioSimpleIRCClient(SimpleIRCClient): """A simple single-server IRC client class. @@ -288,6 +540,25 @@ class AioSimpleIRCClient(SimpleIRCClient): """ reactor_class = AioReactor + reactor: AioReactor def connect(self, *args, **kwargs): self.reactor.loop.run_until_complete(self.connection.connect(*args, **kwargs)) + + def dcc_connect(self, address, port, dcctype="chat"): + """Connect to a DCC peer. + + Returns an AioDCCConnection instance. + """ + warnings.warn("Use self.dcc(type).connect()", DeprecationWarning, stacklevel=2) + dcc = self.dcc(dcctype) + return self.reactor.loop.run_until_complete(dcc.connect(address, port)) + + def dcc_listen(self, dcctype="chat"): + """Listen for connections from a DCC peer. + + Returns an AioDCCConnection instance. + """ + warnings.warn("Use self.dcc(type).listen()", DeprecationWarning, stacklevel=2) + dcc = self.dcc(dcctype) + return self.reactor.loop.run_until_complete(dcc.listen()) diff --git a/irc/tests/test_client_aio.py b/irc/tests/test_client_aio.py index 962cbde..e591311 100644 --- a/irc/tests/test_client_aio.py +++ b/irc/tests/test_client_aio.py @@ -3,6 +3,8 @@ import warnings from unittest.mock import MagicMock +import pytest + from irc import client_aio @@ -38,3 +40,395 @@ def test_privmsg_sends_msg(): mock_transport.write.assert_called_with(b'PRIVMSG #best-channel :You are great\r\n') loop.close() + + +def test_dcc_connection_connect(): + mock_transport = MagicMock() + mock_protocol = MagicMock() + + loop = asyncio.new_event_loop() + loop.create_connection = make_mocked_create_connection( + mock_transport, mock_protocol + ) + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + + assert dcc.connected + assert dcc.transport is mock_transport + assert dcc.protocol is mock_protocol + assert dcc.peeraddress == '127.0.0.1' + assert dcc.peerport == 12345 + + loop.close() + + +def test_dcc_connection_disconnect(): + mock_transport = MagicMock() + mock_protocol = MagicMock() + + loop = asyncio.new_event_loop() + loop.create_connection = make_mocked_create_connection( + mock_transport, mock_protocol + ) + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + dcc.disconnect('test message') + + mock_transport.close.assert_called_once() + assert dcc not in reactor.connections + # Calling disconnect again should be a no-op (idempotent) + dcc.disconnect('test message again') + mock_transport.close.assert_called_once() + + loop.close() + + +def test_dcc_connection_send_bytes(): + mock_transport = MagicMock() + mock_protocol = MagicMock() + + loop = asyncio.new_event_loop() + loop.create_connection = make_mocked_create_connection( + mock_transport, mock_protocol + ) + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + dcc.send_bytes(b'hello') + + mock_transport.write.assert_called_with(b'hello') + + loop.close() + + +def test_dcc_connection_send_bytes_not_connected(): + loop = asyncio.new_event_loop() + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + # Should not raise when not connected + dcc.send_bytes(b'hello') + loop.close() + + +def test_dcc_connection_process_data_chat(): + mock_transport = MagicMock() + mock_protocol = MagicMock() + + loop = asyncio.new_event_loop() + loop.create_connection = make_mocked_create_connection( + mock_transport, mock_protocol + ) + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc('chat') + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + + events = [] + reactor.add_global_handler('dccmsg', lambda c, e: events.append(e)) + dcc.process_data(b'hello\nworld\n') + + assert len(events) == 2 + assert events[0].arguments == [b'hello'] + assert events[1].arguments == [b'world'] + + loop.close() + + +def test_dcc_connection_process_data_raw(): + mock_transport = MagicMock() + mock_protocol = MagicMock() + + loop = asyncio.new_event_loop() + loop.create_connection = make_mocked_create_connection( + mock_transport, mock_protocol + ) + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc('raw') + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + + events = [] + reactor.add_global_handler('dccmsg', lambda c, e: events.append(e)) + dcc.process_data(b'rawdata') + + assert len(events) == 1 + assert events[0].arguments == [b'rawdata'] + + loop.close() + + +def test_dcc_connection_listen(): + loop = asyncio.new_event_loop() + + async def mock_create_server(factory, *args, **kwargs): + mock_server = MagicMock() + mock_socket = MagicMock() + mock_socket.getsockname.return_value = ('127.0.0.1', 54321) + mock_server.sockets = [mock_socket] + return mock_server + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.listen()) + + assert dcc.passive + assert dcc.localaddress == '127.0.0.1' + assert dcc.localport == 54321 + assert hasattr(dcc, 'server') + + loop.close() + + +def test_dcc_connection_listen_ipv6(): + import socket + + loop = asyncio.new_event_loop() + calls = [] + + async def mock_create_server(factory, *args, **kwargs): + calls.append(kwargs) + mock_server = MagicMock() + mock_socket = MagicMock() + mock_socket.getsockname.return_value = ('::1', 54321) + mock_server.sockets = [mock_socket] + return mock_server + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.listen(ipv6=True)) + + assert calls[0]['family'] == socket.AF_INET6 + assert dcc.localaddress == '::1' + + loop.close() + + +def test_dcc_connection_listen_specific_port(): + import socket + + loop = asyncio.new_event_loop() + calls = [] + + async def mock_create_server(factory, *args, **kwargs): + calls.append(args) + mock_server = MagicMock() + mock_socket = MagicMock() + mock_socket.getsockname.return_value = ('127.0.0.1', 12345) + mock_server.sockets = [mock_socket] + return mock_server + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.listen(port=12345)) + + expected_host = socket.gethostbyname(socket.gethostname()) + assert (expected_host, 12345) in calls + assert dcc.localport == 12345 + + loop.close() + + +def test_dcc_connection_listen_port_range(): + loop = asyncio.new_event_loop() + calls = [] + + async def mock_create_server(factory, *args, **kwargs): + calls.append(args) + if args[1] == 1024: + raise OSError('Address in use') + mock_server = MagicMock() + mock_socket = MagicMock() + mock_socket.getsockname.return_value = ('127.0.0.1', args[1]) + mock_server.sockets = [mock_socket] + return mock_server + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.listen(port=(1024, 1026))) + + assert len(calls) == 2 # 1024 fails, 1025 succeeds + assert calls[0][1] == 1024 + assert calls[1][1] == 1025 + assert dcc.localport == 1025 + + loop.close() + + +def test_dcc_connection_listen_port_list(): + loop = asyncio.new_event_loop() + calls = [] + + async def mock_create_server(factory, *args, **kwargs): + calls.append(args) + if args[1] == 8080: + raise OSError('Address in use') + mock_server = MagicMock() + mock_socket = MagicMock() + mock_socket.getsockname.return_value = ('127.0.0.1', args[1]) + mock_server.sockets = [mock_socket] + return mock_server + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.listen(port=[8080, 9090])) + + assert len(calls) == 2 + assert calls[0][1] == 8080 + assert calls[1][1] == 9090 + assert dcc.localport == 9090 + + loop.close() + + +def test_dcc_connection_listen_port_range_all_fail(): + loop = asyncio.new_event_loop() + + async def mock_create_server(factory, *args, **kwargs): + raise OSError('Address in use') + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + + import irc.client + + with pytest.raises(irc.client.DCCConnectionError): + loop.run_until_complete(dcc.listen(port=(1024, 1025))) + + loop.close() + + +def test_dcc_connection_listen_and_accept(): + loop = asyncio.new_event_loop() + + async def mock_create_server(factory, *args, **kwargs): + mock_server = MagicMock() + mock_socket = MagicMock() + mock_socket.getsockname.return_value = ('127.0.0.1', 54321) + mock_server.sockets = [mock_socket] + return mock_server + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.listen()) + + # Simulate a client connecting by calling connection_made on the protocol + mock_transport = MagicMock() + mock_transport.get_extra_info.return_value = ('192.168.1.1', 12345) + + events = [] + reactor.add_global_handler('dcc_connect', lambda c, e: events.append(e)) + + # Create the protocol and simulate connection_made + protocol = dcc.protocol_class(dcc, loop) + protocol.connection_made(mock_transport) + + assert dcc.connected + assert dcc.transport is mock_transport + assert dcc.peeraddress == '192.168.1.1' + assert dcc.peerport == 12345 + assert len(events) == 1 + assert events[0].type == 'dcc_connect' + dcc.server.close.assert_called_once() + + loop.close() + + +def test_dcc_connection_disconnect_with_server(): + loop = asyncio.new_event_loop() + + async def mock_create_server(factory, *args, **kwargs): + mock_server = MagicMock() + mock_socket = MagicMock() + mock_socket.getsockname.return_value = ('127.0.0.1', 54321) + mock_server.sockets = [mock_socket] + return mock_server + + loop.create_server = mock_create_server + + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + loop.run_until_complete(dcc.listen()) + + mock_transport = MagicMock() + mock_transport.get_extra_info.return_value = ('192.168.1.1', 12345) + protocol = dcc.protocol_class(dcc, loop) + protocol.connection_made(mock_transport) + + dcc.disconnect('test message') + + assert dcc.server.close.call_count == 2 # once in connection_made, once in disconnect + mock_transport.close.assert_called_once() + assert dcc not in reactor.connections + # Calling disconnect again should be a no-op (idempotent) + dcc.disconnect('test message again') + assert dcc.server.close.call_count == 2 + mock_transport.close.assert_called_once() + + loop.close() + + +def test_dcc_connection_connect_error(): + async def failing_create_connection(*args, **kwargs): + raise OSError('connection refused') + + loop = asyncio.new_event_loop() + loop.create_connection = failing_create_connection + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc() + + import irc.client + + with pytest.raises(irc.client.DCCConnectionError): + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + + loop.close() + + +def test_dcc_connection_privmsg_chat(): + mock_transport = MagicMock() + mock_protocol = MagicMock() + + loop = asyncio.new_event_loop() + loop.create_connection = make_mocked_create_connection( + mock_transport, mock_protocol + ) + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc('chat') + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + dcc.privmsg('hello') + + mock_transport.write.assert_called_with(b'hello\n') + + loop.close() + + +def test_dcc_connection_privmsg_raw(): + mock_transport = MagicMock() + mock_protocol = MagicMock() + + loop = asyncio.new_event_loop() + loop.create_connection = make_mocked_create_connection( + mock_transport, mock_protocol + ) + reactor = client_aio.AioReactor(loop=loop) + dcc = reactor.dcc('raw') + loop.run_until_complete(dcc.connect('127.0.0.1', 12345)) + dcc.privmsg('hello') + + mock_transport.write.assert_called_with(b'hello') + + loop.close() diff --git a/scripts/dccchat-aio.py b/scripts/dccchat-aio.py new file mode 100644 index 0000000..18d77bf --- /dev/null +++ b/scripts/dccchat-aio.py @@ -0,0 +1,96 @@ +#! /usr/bin/env python +# +# Example program using irc.client_aio. +# +# This program demonstrates DCC CHAT using the asyncio-based client. +# +# This program is free without restrictions; do anything you like with +# it. + +import argparse +import shlex + +import jaraco.logging + +import irc.client +import irc.client_aio + + +class DCCChat(irc.client_aio.AioSimpleIRCClient): + def __init__(self, target=None): + super().__init__() + self.target = target + self.dcc = None + + def on_welcome(self, connection, event): + if self.target: + self.dcc = self.dcc_listen('chat') + msg_parts = [ + 'CHAT', + 'chat', + irc.client.ip_quad_to_numstr(self.dcc.localaddress), + self.dcc.localport, + ] + msg = ' '.join(map(str, msg_parts)) + connection.ctcp('DCC', self.target, msg) + print(f"Listening for DCC CHAT from {self.target} on port {self.dcc.localport}") + else: + print("Waiting for incoming DCC CHAT requests...") + + def on_ctcp(self, connection, event): + payload = event.arguments[1] + parts = shlex.split(payload) + if len(parts) < 4 or parts[0] != 'CHAT': + return + command, dcctype, peer_address, peer_port = parts[:4] + peer_address = irc.client.ip_numstr_to_quad(peer_address) + peer_port = int(peer_port) + print(f"Received DCC CHAT request from {event.source}: {peer_address}:{peer_port}") + self.dcc = self.dcc_connect(peer_address, peer_port, 'chat') + + def on_dcc_connect(self, connection, event): + print("DCC CHAT connected.") + if self.target: + connection.privmsg("Hello! This is an asyncio DCC CHAT echo bot.") + + def on_dcc_disconnect(self, connection, event): + print("DCC CHAT disconnected.") + self.connection.quit() + + def on_dccmsg(self, connection, event): + text = event.arguments[0].decode('utf-8', errors='replace').rstrip('\n') + print(f"< {text}") + response = f"Echo: {text}" + print(f"> {response}") + connection.privmsg(response) + + def on_disconnect(self, connection, event): + raise SystemExit() + + +def get_args(): + parser = argparse.ArgumentParser( + description="Connect to an IRC server and use DCC CHAT." + ) + parser.add_argument('server') + parser.add_argument('nickname') + parser.add_argument( + '-t', '--target', default=None, + help="Nickname to send a DCC CHAT request to (omit to wait for requests)" + ) + parser.add_argument('-p', '--port', default=6667, type=int) + jaraco.logging.add_arguments(parser) + return parser.parse_args() + + +def main(): + args = get_args() + jaraco.logging.setup(args) + + client = DCCChat(target=args.target) + client.connect(args.server, args.port, args.nickname) + client.start() + + +if __name__ == '__main__': + main() diff --git a/scripts/dccreceive-aio.py b/scripts/dccreceive-aio.py new file mode 100644 index 0000000..fb83c1b --- /dev/null +++ b/scripts/dccreceive-aio.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python +# +# Example program using irc.client_aio. +# +# Active DCC file receiver: connects to an IRC server, joins a channel, +# sends "!send " to a target nick, receives the CTCP DCC SEND +# reply, actively connects to the sender, and saves the file. + +import argparse +import asyncio +import os +import shlex +import struct +import sys + +import jaraco.logging + +import irc.client +import irc.client_aio + + +class DCCReceiveAio: + def __init__(self, loop, server, port, nickname, channel, target, filename): + self.loop = loop + self.server = server + self.port = port + self.nickname = nickname + self.channel = channel + self.target = target + self.filename = filename + + self.reactor = irc.client_aio.AioReactor(loop=loop) + self.conn = self.reactor.server() + self.dcc = None + self.file = None + self.received_bytes = 0 + + async def run(self): + await self.conn.connect(self.server, self.port, self.nickname) + self.conn.add_global_handler("welcome", self.on_welcome) + self.conn.add_global_handler("ctcp", self.on_ctcp) + self.conn.add_global_handler("dccmsg", self.on_dccmsg) + self.conn.add_global_handler("dcc_disconnect", self.on_dcc_disconnect) + self.conn.add_global_handler("disconnect", self.on_disconnect) + + self.reactor.process_forever() + + def on_welcome(self, connection, event): + connection.join(self.channel) + connection.privmsg(self.target, f"!send {self.filename}") + + def on_ctcp(self, connection, event): + payload = event.arguments[1] + parts = shlex.split(payload) + if len(parts) < 5: + return + command, filename, peer_address, peer_port, size = parts + if command != "SEND": + return + + save_name = os.path.basename(filename) + if os.path.exists(save_name): + print("A file named", save_name, "already exists. Refusing to save it.") + self.conn.quit() + return + + self.file = open(save_name, "wb") + self.received_bytes = 0 + + peer_address = irc.client.ip_numstr_to_quad(peer_address) + peer_port = int(peer_port) + + self.reactor.loop.create_task(self._do_connect(peer_address, peer_port)) + + async def _do_connect(self, peer_address, peer_port): + self.dcc = self.reactor.dcc("raw") + await self.dcc.connect(peer_address, peer_port) + print(f"Connected to {peer_address}:{peer_port} for DCC receive") + + def on_dccmsg(self, connection, event): + data = event.arguments[0] + self.file.write(data) + self.received_bytes += len(data) + self.dcc.send_bytes(struct.pack("!I", self.received_bytes)) + + def on_dcc_disconnect(self, connection, event): + self.file.close() + print(f"Received file {os.path.basename(self.filename)} ({self.received_bytes} bytes)") + self.conn.quit() + + def on_disconnect(self, connection, event): + raise SystemExit() + + +def get_args(): + parser = argparse.ArgumentParser( + description="Asyncio DCC file receiver. Joins a channel, requests " + "a file from a nick, and saves it." + ) + parser.add_argument('server') + parser.add_argument('nickname') + parser.add_argument('channel') + parser.add_argument('target', help="Nickname of the sender") + parser.add_argument('filename', help="File to request") + parser.add_argument( + '-p', '--port', default=6667, type=int, help="IRC server port" + ) + jaraco.logging.add_arguments(parser) + return parser.parse_args() + + +def main(): + args = get_args() + jaraco.logging.setup(args) + + loop = asyncio.new_event_loop() + client = DCCReceiveAio( + loop, args.server, args.port, args.nickname, + args.channel, args.target, args.filename + ) + + try: + loop.run_until_complete(client.run()) + except irc.client.ServerConnectionError: + print(sys.exc_info()[1]) + raise SystemExit(1) from None + finally: + loop.close() + + +if __name__ == '__main__': + main() diff --git a/scripts/dccsend-aio.py b/scripts/dccsend-aio.py new file mode 100644 index 0000000..8899a61 --- /dev/null +++ b/scripts/dccsend-aio.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python +# +# Example program using irc.client_aio. +# +# Passive DCC file sender: connects to an IRC server, joins a channel, +# and waits for someone to type "!send ". It then listens +# for an incoming DCC connection and sends the file. +# +# Supports a specific port (--port 8080) or a port range fallback +# (--port 1024-65535, the default when --port is omitted). + +import argparse +import asyncio +import os +import struct +import sys + +import jaraco.logging + +import irc.client +import irc.client_aio + + +class DCCSendAio: + def __init__(self, loop, server, port, nickname, channel, port_arg=None): + self.loop = loop + self.server = server + self.port = port + self.nickname = nickname + self.channel = channel + self.port_arg = port_arg + + self.reactor = irc.client_aio.AioReactor(loop=loop) + self.conn = self.reactor.server() + self.dcc = None + self.file = None + self.filesize = 0 + self.sent_bytes = 0 + self.filename = None + + async def run(self): + await self.conn.connect(self.server, self.port, self.nickname) + self.conn.add_global_handler("welcome", self.on_welcome) + self.conn.add_global_handler("privmsg", self.on_privmsg) + self.conn.add_global_handler("dcc_connect", self.on_dcc_connect) + self.conn.add_global_handler("dccmsg", self.on_dccmsg) + self.conn.add_global_handler("dcc_disconnect", self.on_dcc_disconnect) + self.conn.add_global_handler("disconnect", self.on_disconnect) + + self.reactor.process_forever() + + def on_welcome(self, connection, event): + connection.join(self.channel) + + def on_privmsg(self, connection, event): + nick = irc.client.NickMask(event.source).nick + text = event.arguments[0] + + if not text.startswith("!send "): + return + + filename = text[6:].strip() + if not os.path.exists(filename): + connection.privmsg(nick, f"File not found: {filename}") + return + + self.filename = filename + self.reactor.loop.create_task(self._start_send(nick)) + + async def _start_send(self, receiver): + self.file = open(self.filename, 'rb') + self.filesize = os.path.getsize(self.filename) + self.sent_bytes = 0 + + self.dcc = self.reactor.dcc("raw") + + # Determine port: specific, range, or default range + if self.port_arg is None: + port = (1024, 65535) + elif '-' in self.port_arg: + min_port, max_port = map(int, self.port_arg.split('-')) + port = (min_port, max_port) + else: + port = int(self.port_arg) + + await self.dcc.listen(port=port) + + ip_str = irc.client.ip_quad_to_numstr(self.dcc.localaddress) + msg = f"SEND {os.path.basename(self.filename)} {ip_str} {self.dcc.localport} {self.filesize}" + self.conn.ctcp("DCC", receiver, msg) + + print(f"Listening for DCC connection on {self.dcc.localaddress}:{self.dcc.localport}") + + def on_dcc_connect(self, connection, event): + if self.filesize == 0: + self.dcc.disconnect() + return + self._send_chunk() + + def _send_chunk(self): + data = self.file.read(1024) + self.dcc.send_bytes(data) + self.sent_bytes += len(data) + + def on_dccmsg(self, connection, event): + acked = struct.unpack("!I", event.arguments[0])[0] + if acked == self.filesize: + self.dcc.disconnect() + self.conn.quit() + elif acked == self.sent_bytes: + self._send_chunk() + + def on_dcc_disconnect(self, connection, event): + self.file.close() + print(f"Sent file {self.filename} ({self.sent_bytes} bytes)") + self.conn.quit() + + def on_disconnect(self, connection, event): + raise SystemExit() + + +def get_args(): + parser = argparse.ArgumentParser( + description="Asyncio DCC file sender. Waits for '!send ' " + "in a channel, then sends the file via DCC." + ) + parser.add_argument('server') + parser.add_argument('nickname') + parser.add_argument('channel') + parser.add_argument( + '--port', + default=None, + help="Port to listen on. Can be a single port or a range " + "like '1024-65535'. Default is to try 1024-65535.", + ) + parser.add_argument( + '-p', '--irc-port', default=6667, type=int, help="IRC server port" + ) + jaraco.logging.add_arguments(parser) + return parser.parse_args() + + +def main(): + args = get_args() + jaraco.logging.setup(args) + + loop = asyncio.new_event_loop() + client = DCCSendAio( + loop, args.server, args.irc_port, args.nickname, args.channel, args.port + ) + + try: + loop.run_until_complete(client.run()) + except irc.client.ServerConnectionError: + print(sys.exc_info()[1]) + raise SystemExit(1) from None + finally: + loop.close() + + +if __name__ == '__main__': + main()