102 lines
3.8 KiB
Python
102 lines
3.8 KiB
Python
|
import irc3
|
||
|
|
from irc3.plugins.command import command
|
||
|
|
from irc3.compat import Queue
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
def strip_nick_prefix(nick):
|
||
|
|
"""Remove IRC mode prefixes from a nickname"""
|
||
|
|
return nick.lstrip('@+%&~!') if nick else ''
|
||
|
|
|
||
|
|
@irc3.plugin
|
||
|
|
class MassMessagePlugin:
|
||
|
|
"""Mass messaging plugin using async queue system"""
|
||
|
|
|
||
|
|
def __init__(self, bot):
|
||
|
|
self.bot = bot
|
||
|
|
self.delay = 0.0001 # Delay between messages in seconds
|
||
|
|
self.queue = Queue() # Using irc3's compatibility Queue
|
||
|
|
self.count = 0 # Counter for successfully sent messages
|
||
|
|
|
||
|
|
async def _worker(self, message, total):
|
||
|
|
"""Worker task to process messages from the queue"""
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
nick = await self.queue.get()
|
||
|
|
try:
|
||
|
|
if nick == "FUCKYOU":
|
||
|
|
nick = "Zodiac"
|
||
|
|
self.bot.privmsg(nick, message)
|
||
|
|
self.count += 1
|
||
|
|
self.bot.log.info(f"Sent to {nick} ({self.count}/{total})")
|
||
|
|
except Exception as e:
|
||
|
|
self.bot.log.error(f"Failed to message {nick}: {e}")
|
||
|
|
finally:
|
||
|
|
self.queue.task_done() # Mark the task as done after processing
|
||
|
|
|
||
|
|
if self.delay > 0:
|
||
|
|
await asyncio.sleep(self.delay)
|
||
|
|
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
# Exit silently on cancellation
|
||
|
|
break
|
||
|
|
except Exception as e:
|
||
|
|
self.bot.log.error(f"Worker error: {e}")
|
||
|
|
break
|
||
|
|
|
||
|
|
@command(permission="admin", options_first=True)
|
||
|
|
async def msgall(self, mask, target, args):
|
||
|
|
"""Send a message to all users in the channel using async queue
|
||
|
|
|
||
|
|
%%msgall <message>...
|
||
|
|
"""
|
||
|
|
if not target.is_channel:
|
||
|
|
return "This command can only be used in a channel"
|
||
|
|
|
||
|
|
message = ' '.join(args['<message>'])
|
||
|
|
workers = [] # Ensure workers is defined in the scope of the try block
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Fetch the list of users in the channel
|
||
|
|
result = await self.bot.async_cmds.names(target)
|
||
|
|
nicknames = [strip_nick_prefix(n) for n in result['names']]
|
||
|
|
recipients = [n for n in nicknames if n != self.bot.nick]
|
||
|
|
|
||
|
|
if not recipients:
|
||
|
|
return "No valid recipients found"
|
||
|
|
|
||
|
|
total = len(recipients)
|
||
|
|
self.count = 0 # Reset the counter for this run
|
||
|
|
|
||
|
|
# Add all recipients to the queue
|
||
|
|
for nick in recipients:
|
||
|
|
await self.queue.put(nick)
|
||
|
|
|
||
|
|
# Create worker tasks
|
||
|
|
workers = [asyncio.create_task(self._worker(message, total)) for _ in range(1)]
|
||
|
|
|
||
|
|
# Send initial confirmation
|
||
|
|
self.bot.privmsg(target, f"Starting mass message to {total} users...")
|
||
|
|
|
||
|
|
# Wait for the queue to be fully processed
|
||
|
|
await self.queue.join()
|
||
|
|
|
||
|
|
# Cancel worker tasks after the queue is empty
|
||
|
|
for task in workers:
|
||
|
|
task.cancel()
|
||
|
|
# Wait for workers to finish cancellation
|
||
|
|
await asyncio.gather(*workers, return_exceptions=True)
|
||
|
|
|
||
|
|
return f"Mass message completed. Sent to {self.count}/{total} users."
|
||
|
|
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
self.bot.log.info("Mass message command cancelled. Cleaning up workers.")
|
||
|
|
# Cancel any existing worker tasks
|
||
|
|
for task in workers:
|
||
|
|
task.cancel()
|
||
|
|
# Allow workers to handle cancellation
|
||
|
|
await asyncio.gather(*workers, return_exceptions=True)
|
||
|
|
# Re-raise the cancellation to inform the bot framework
|
||
|
|
raise
|
||
|
|
except Exception as e:
|
||
|
|
self.bot.log.error(f"Error in msgall: {e}")
|
||
|
|
return f"Mass message failed: {str(e)}"
|