343 lines
15 KiB
Python
343 lines
15 KiB
Python
import asyncio
|
|
import json
|
|
import random
|
|
import logging
|
|
import argparse
|
|
import socket
|
|
import ssl
|
|
import base64
|
|
import os
|
|
import inspect
|
|
import importlib
|
|
|
|
POKEMON_NAMES = [
|
|
"Bulbasaur", "Ivysaur", "Venusaur", "Charmander", "Charmeleon", "Charizard",
|
|
"Squirtle", "Wartortle", "Blastoise", "Caterpie", "Metapod", "Butterfree",
|
|
"Weedle", "Kakuna", "Beedrill", "Pidgey", "Pidgeotto", "Pidgeot", "Rattata",
|
|
"Raticate", "Spearow", "Fearow", "Ekans", "Arbok", "Pikachu", "Raichu",
|
|
"Sandshrew", "Sandslash", "Nidoran", "Nidorina", "Nidoqueen", "Nidoran♂",
|
|
"Nidorino", "Nidoking", "Clefairy", "Clefable", "Vulpix", "Ninetales",
|
|
"Jigglypuff", "Wigglytuff", "Zubat", "Golbat", "Oddish", "Gloom", "Vileplume",
|
|
"Paras", "Parasect", "Venonat", "Venomoth", "Diglett", "Dugtrio", "Meowth",
|
|
"Persian", "Psyduck", "Golduck", "Mankey", "Primeape", "Growlithe", "Arcanine",
|
|
"Poliwag", "Poliwhirl", "Poliwrath", "Abra", "Kadabra", "Alakazam", "Machop",
|
|
"Machoke", "Machamp", "Bellsprout", "Weepinbell", "Victreebel", "Tentacool",
|
|
"Tentacruel", "Geodude", "Graveler", "Golem", "Ponyta", "Rapidash", "Slowpoke",
|
|
"Slowbro", "Magnemite", "Magneton", "Farfetchd", "Doduo", "Dodrio", "Seel",
|
|
"Dewgong", "Grimer", "Muk", "Shellder", "Cloyster", "Gastly", "Haunter",
|
|
"Gengar", "Onix", "Drowzee", "Hypno", "Krabby", "Kingler", "Voltorb",
|
|
"Electrode", "Exeggcute", "Exeggutor", "Cubone", "Marowak", "Hitmonlee",
|
|
"Hitmonchan", "Lickitung", "Koffing", "Weezing", "Rhyhorn", "Rhydon", "Chansey",
|
|
"Tangela", "Kangaskhan", "Horsea", "Seadra", "Goldeen", "Seaking", "Staryu",
|
|
"Starmie", "MrMime", "Scyther", "Jynx", "Electabuzz", "Magmar", "Pinsir",
|
|
"Tauros", "Magikarp", "Gyarados", "Lapras", "Ditto", "Eevee", "Vaporeon",
|
|
"Jolteon", "Flareon", "Porygon", "Omanyte", "Omastar", "Kabuto", "Kabutops",
|
|
"Aerodactyl", "Snorlax", "Articuno", "Zapdos", "Moltres", "Dratini",
|
|
"Dragonair", "Dragonite", "Mewtwo", "Mew"
|
|
]
|
|
|
|
def generate_nick():
|
|
return f"{random.choice(POKEMON_NAMES)}{random.randint(100, 999)}"
|
|
|
|
class SimpleEncryption:
|
|
def __init__(self, key):
|
|
self.key = key
|
|
|
|
def encrypt(self, data):
|
|
encrypted = bytearray()
|
|
for i, char in enumerate(data):
|
|
encrypted.append(ord(char) ^ self.key[i % len(self.key)])
|
|
return base64.b64encode(encrypted).decode()
|
|
|
|
def decrypt(self, data):
|
|
encrypted = base64.b64decode(data)
|
|
decrypted = bytearray()
|
|
for i, byte in enumerate(encrypted):
|
|
decrypted.append(byte ^ self.key[i % len(self.key)])
|
|
return decrypted.decode()
|
|
|
|
class BasePlugin:
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
|
|
async def on_command(self, sender, channel, command, args):
|
|
pass
|
|
|
|
@property
|
|
def commands(self):
|
|
return {}
|
|
|
|
class CommandHub:
|
|
def __init__(self, hub_address, hub_port, irc_server, irc_port, irc_channel, use_ssl=False, channel_password=None, server_password=None):
|
|
self.hub_config = {
|
|
"address": hub_address,
|
|
"port": hub_port,
|
|
}
|
|
self.leaf_bots = set()
|
|
self.irc_config = {
|
|
"server": irc_server,
|
|
"port": irc_port,
|
|
"channel": irc_channel,
|
|
"use_ssl": use_ssl,
|
|
"channel_password": channel_password,
|
|
"password": server_password,
|
|
}
|
|
self.encryption_key = os.urandom(32)
|
|
self.encryption = SimpleEncryption(self.encryption_key)
|
|
self.plugins = []
|
|
self.commands = {
|
|
"test": "Send a test message to the channel",
|
|
"join": "Join a new channel",
|
|
"leave": "Leave a channel",
|
|
"nick": "Change nickname",
|
|
"UPDATECONF": "Update hub configuration",
|
|
"UPDATECONF.IRC": "Update IRC configuration"
|
|
}
|
|
self.used_nicks = {}
|
|
self.load_plugins()
|
|
|
|
def load_plugins(self):
|
|
new_plugins = []
|
|
plugins_dir = os.path.join(os.path.dirname(__file__), "plugins")
|
|
for filename in os.listdir(plugins_dir):
|
|
if filename.endswith(".py") and not filename.startswith("__"):
|
|
module_name = f"plugins.{filename[:-3]}"
|
|
module = importlib.import_module(module_name)
|
|
for name, obj in inspect.getmembers(module):
|
|
if inspect.isclass(obj) and issubclass(obj, BasePlugin) and obj != BasePlugin:
|
|
plugin = obj(self)
|
|
self.plugins.append(plugin)
|
|
self.commands.update(plugin.commands)
|
|
new_plugins.append(name)
|
|
logging.info(f"Loaded new plugin: {name}")
|
|
return new_plugins
|
|
|
|
def generate_unique_nick(self, base_nick):
|
|
count = 0
|
|
while f"{base_nick}{count if count else ''}" in self.used_nicks:
|
|
count += 1
|
|
new_nick = f"{base_nick}{count if count else ''}"
|
|
self.used_nicks[new_nick] = new_nick
|
|
return new_nick
|
|
|
|
def release_nick(self, nick):
|
|
if nick in self.used_nicks:
|
|
del self.used_nicks[nick]
|
|
|
|
async def handle_leaf_connection(self, reader, writer):
|
|
addr = writer.get_extra_info('peername')
|
|
logging.info(f"New leaf bot connected: {addr}")
|
|
|
|
writer.write(self.encryption_key)
|
|
await writer.drain()
|
|
|
|
leaf_config = self.irc_config.copy()
|
|
leaf_config["nickname"] = self.generate_unique_nick(generate_nick())
|
|
|
|
encrypted_config = self.encryption.encrypt(json.dumps(leaf_config))
|
|
writer.write(encrypted_config.encode() + b'\n')
|
|
await writer.drain()
|
|
|
|
encrypted_commands = self.encryption.encrypt(json.dumps(self.commands))
|
|
writer.write(encrypted_commands.encode() + b'\n')
|
|
await writer.drain()
|
|
|
|
self.leaf_bots.add(writer)
|
|
try:
|
|
while True:
|
|
data = await reader.readline()
|
|
if not data:
|
|
break
|
|
encrypted_message = data.decode().strip()
|
|
message = self.encryption.decrypt(encrypted_message)
|
|
logging.info(f"Received from leaf bot {addr}: {message}")
|
|
await self.process_leaf_message(writer, message)
|
|
finally:
|
|
self.leaf_bots.remove(writer)
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
logging.info(f"Leaf bot disconnected: {addr}")
|
|
|
|
async def process_leaf_message(self, writer, message):
|
|
try:
|
|
data = json.loads(message)
|
|
if data['type'] == 'command':
|
|
response = await self.execute_command(data['sender'], data['channel'], data['command'], data.get('args', []))
|
|
encrypted_response = self.encryption.encrypt(json.dumps(response))
|
|
writer.write(encrypted_response.encode() + b'\n')
|
|
await writer.drain()
|
|
elif data['type'] == 'nick_update':
|
|
if 'old_nick' in data:
|
|
self.release_nick(data['old_nick'])
|
|
if 'new_nick' in data:
|
|
new_nick = self.generate_unique_nick(data['new_nick'])
|
|
response = {"type": "action", "action": "set_nick", "nickname": new_nick}
|
|
encrypted_response = self.encryption.encrypt(json.dumps(response))
|
|
writer.write(encrypted_response.encode() + b'\n')
|
|
await writer.drain()
|
|
elif data['type'] == 'alert':
|
|
logging.warning(f"Alert from leaf bot: {data['message']}")
|
|
else:
|
|
logging.warning(f"Unknown message type from leaf bot: {data['type']}")
|
|
except json.JSONDecodeError:
|
|
logging.error(f"Received invalid JSON from leaf bot: {message}")
|
|
except Exception as e:
|
|
logging.error(f"Error processing leaf message: {e}")
|
|
|
|
async def execute_command(self, sender, channel, command, args):
|
|
if command == "test":
|
|
return {"type": "action", "action": "send_message", "channel": self.irc_config['channel'], "message": "Test message from hub"}
|
|
elif command == "join":
|
|
return {"type": "action", "action": "join_channel", "channel": args[0] if args else channel}
|
|
elif command == "leave":
|
|
return {"type": "action", "action": "leave_channel", "channel": args[0] if args else channel}
|
|
elif command == "nick":
|
|
new_nick = self.generate_unique_nick(generate_nick())
|
|
return {"type": "action", "action": "change_nick", "nickname": new_nick}
|
|
elif command == "request_nick":
|
|
base_nick = args[0] if args else generate_nick()
|
|
new_nick = self.generate_unique_nick(base_nick)
|
|
return {"type": "action", "action": "set_nick", "nickname": new_nick}
|
|
elif command == "release_nick":
|
|
if args:
|
|
self.release_nick(args[0])
|
|
return {"type": "action", "action": "nick_released"}
|
|
elif command == "UPDATECONF":
|
|
return self.update_hub_config(args)
|
|
elif command == "UPDATECONF.IRC":
|
|
return self.update_irc_config(args)
|
|
else:
|
|
for plugin in self.plugins:
|
|
result = await plugin.on_command(sender, channel, command, args)
|
|
if result:
|
|
return result
|
|
return {"type": "error", "message": "Unknown command"}
|
|
|
|
def update_hub_config(self, params):
|
|
if len(params) >= 2:
|
|
self.hub_config["address"] = params[0]
|
|
self.hub_config["port"] = int(params[1])
|
|
elif len(params) == 1:
|
|
if params[0].isdigit():
|
|
self.hub_config["port"] = int(params[0])
|
|
else:
|
|
self.hub_config["address"] = params[0]
|
|
logging.info(f"Updated hub configuration: {self.hub_config}")
|
|
return {"type": "action", "action": "update_hub_config", "config": self.hub_config}
|
|
|
|
def update_irc_config(self, params):
|
|
if len(params) >= 3:
|
|
self.irc_config["server"] = params[0]
|
|
self.irc_config["port"] = int(params[1])
|
|
self.irc_config["channel"] = params[2]
|
|
self.irc_config["channel_password"] = params[3] if len(params) > 3 else None
|
|
self.irc_config["use_ssl"] = True if len(params) > 4 and params[4] == "-ssl" else False
|
|
logging.info(f"Updated IRC configuration: {self.irc_config}")
|
|
return {"type": "action", "action": "update_irc_config", "config": self.irc_config}
|
|
|
|
async def broadcast_command(self, command, params):
|
|
response = await self.execute_command("Console", "Hub", command, params)
|
|
encrypted_response = self.encryption.encrypt(json.dumps(response))
|
|
for bot in self.leaf_bots:
|
|
try:
|
|
bot.write(encrypted_response.encode() + b'\n')
|
|
await bot.drain()
|
|
except Exception as e:
|
|
logging.error(f"Error broadcasting command to leaf bot: {e}")
|
|
|
|
async def run_hub_server(self):
|
|
servers = []
|
|
|
|
# Try to start IPv6 server
|
|
try:
|
|
server_ipv6 = await asyncio.start_server(
|
|
self.handle_leaf_connection, '::', self.hub_config['port'], family=socket.AF_INET6)
|
|
servers.append(server_ipv6)
|
|
logging.info(f'Serving on IPv6: {server_ipv6.sockets[0].getsockname()}')
|
|
except Exception as e:
|
|
logging.warning(f"Failed to start server on IPv6: {e}")
|
|
|
|
# Try to start IPv4 server
|
|
try:
|
|
server_ipv4 = await asyncio.start_server(
|
|
self.handle_leaf_connection, self.hub_config['address'], self.hub_config['port'], family=socket.AF_INET)
|
|
servers.append(server_ipv4)
|
|
logging.info(f'Serving on IPv4: {server_ipv4.sockets[0].getsockname()}')
|
|
except Exception as e:
|
|
logging.warning(f"Failed to start server on IPv4: {e}")
|
|
|
|
if not servers:
|
|
logging.error("Failed to start any servers. Exiting.")
|
|
return
|
|
|
|
await asyncio.gather(*(server.serve_forever() for server in servers))
|
|
|
|
async def console_input(self):
|
|
while True:
|
|
try:
|
|
command = await asyncio.get_event_loop().run_in_executor(None, input, "Enter command: ")
|
|
parts = command.split()
|
|
if parts:
|
|
await self.broadcast_command(parts[0], parts[1:])
|
|
else:
|
|
logging.warning("Empty command entered")
|
|
except Exception as e:
|
|
logging.error(f"Error processing console input: {e}")
|
|
|
|
async def run(self):
|
|
await asyncio.gather(
|
|
self.run_hub_server(),
|
|
self.console_input()
|
|
)
|
|
|
|
def setup_logger(debug=False):
|
|
level = logging.DEBUG if debug else logging.INFO
|
|
logging.basicConfig(level=level,
|
|
format='%(asctime)s | %(levelname)8s | %(message)s',
|
|
datefmt='%Y-%m-%d %H:%M:%S')
|
|
|
|
async def main(args):
|
|
setup_logger(args.debug)
|
|
while True:
|
|
try:
|
|
hub_bot = CommandHub(
|
|
args.hub_address,
|
|
args.hub_port,
|
|
args.server,
|
|
args.port,
|
|
args.channel,
|
|
use_ssl=args.ssl,
|
|
channel_password=args.key,
|
|
server_password=args.password,
|
|
)
|
|
await hub_bot.run()
|
|
except Exception as e:
|
|
logging.error(f"HubBot crashed: {e}")
|
|
logging.info("Restarting HubBot in 30 seconds...")
|
|
await asyncio.sleep(30)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="IRC Hub Bot")
|
|
parser.add_argument("--hub-address", default="0.0.0.0", help="Address for the hub to listen on")
|
|
parser.add_argument("--hub-port", type=int, default=8888, help="Port for the hub to listen on")
|
|
parser.add_argument("--server", required=True, help="The IRC server address")
|
|
parser.add_argument("--port", type=int, default=6667, help="The IRC server port")
|
|
parser.add_argument("--channel", required=True, help="The IRC channel to join")
|
|
parser.add_argument("--ssl", action="store_true", help="Use SSL for IRC connection")
|
|
parser.add_argument("--key", help="The key (password) for the IRC channel, if required")
|
|
parser.add_argument("--password", help="The password for the IRC server, if required")
|
|
parser.add_argument("--debug", action="store_true", help="Enable debug mode")
|
|
args = parser.parse_args()
|
|
|
|
if args.ssl and args.port == 6667:
|
|
args.port = 6697 # Default SSL port
|
|
|
|
try:
|
|
asyncio.run(main(args))
|
|
except KeyboardInterrupt:
|
|
print("Hub bot stopped by user.")
|
|
except Exception as e:
|
|
logging.error(f"Fatal error in main loop: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|