pylons/pylon.py
2024-07-15 20:58:54 -06:00

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()