bitbot-3.11-fork/src/IRCBot.py

395 lines
14 KiB
Python

import enum, queue, os, queue, select, socket, sys, threading, time, traceback
import typing, uuid
from src import EventManager, Exports, IRCServer, Logging, ModuleManager
from src import Socket, utils
VERSION = "v1.9.2"
SOURCE = "https://git.io/bitbot"
class TriggerResult(enum.Enum):
Return = 1
Exception = 2
class TriggerEventType(enum.Enum):
Action = 1
Kill = 2
class TriggerEvent(object):
def __init__(self, type: TriggerEventType,
callback: typing.Callable[[], None]=None):
self.type = type
self.callback = callback
class BitBotPanic(Exception):
pass
class Bot(object):
def __init__(self, directory, args, cache, config, database, events,
exports, log, modules, timers):
self.directory = directory
self.args = args
self.cache = cache
self.config = config
self.database = database
self._events = events
self._exports = exports
self.log = log
self.modules = modules
self._timers = timers
self.start_time = time.time()
self.running = False
self.servers = {}
self._event_queue = queue.Queue() # type: typing.Queue[TriggerEvent]
self._read_poll = select.poll()
self._write_poll = select.poll()
self._rtrigger_server, self._rtrigger_client = socket.socketpair()
self._read_poll.register(self._rtrigger_server.fileno(), select.POLLIN)
self._rtrigger_lock = threading.Lock()
self._rtriggered = False
self._write_condition = threading.Condition()
self._read_thread = None
self._write_thread = None
self._events.on("timer.reconnect").hook(self._timed_reconnect)
def _trigger_both(self):
self.trigger_read()
self.trigger_write()
def trigger_read(self):
with self._rtrigger_lock:
if not self._rtriggered:
self._rtriggered = True
self._rtrigger_client.send(b"TRIGGER")
def trigger_write(self):
with self._write_condition:
self._write_condition.notify()
def trigger(self,
func: typing.Optional[typing.Callable[[], typing.Any]]=None,
trigger_threads=True) -> typing.Any:
func = func or (lambda: None)
if utils.is_main_thread():
returned = func()
if trigger_threads:
self._trigger_both()
return returned
func_queue = queue.Queue(1) # type: queue.Queue[str]
def _action():
try:
returned = func()
type = TriggerResult.Return
except Exception as e:
returned = e
type = TriggerResult.Exception
func_queue.put([type, returned])
event_item = TriggerEvent(TriggerEventType.Action, _action)
self._event_queue.put(event_item)
type, returned = func_queue.get(block=True)
if trigger_threads:
self._trigger_both()
if type == TriggerResult.Exception:
raise returned
elif type == TriggerResult.Return:
return returned
def panic(self, reason=None):
callback = None
if not reason == None:
self.log.critical("panic() called: %s", [reason])
exception = sys.exc_info()[1]
if exception:
def _raise():
raise exception
callback = _raise
self._event_queue.put(TriggerEvent(TriggerEventType.Kill, callback))
raise BitBotPanic()
def load_modules(self, safe: bool=False
) -> typing.Tuple[typing.List[str], typing.List[str]]:
db_blacklist = set(self.get_setting("module-blacklist", []))
db_whitelist = set(self.get_setting("module-whitelist", []))
conf_blacklist = self.config.get("module-blacklist", "").split(",")
conf_whitelist = self.config.get("module-whitelist", "").split(",")
conf_blacklist = set(filter(None, conf_blacklist))
conf_whitelist = set(filter(None, conf_whitelist))
blacklist = db_blacklist|conf_blacklist
whitelist = db_whitelist|conf_whitelist
return self.modules.load_modules(self, whitelist=whitelist,
blacklist=blacklist, safe=safe)
def add_server(self, server_id: int, connect: bool = True,
connection_param_args: typing.Dict[str, str]={}
) -> IRCServer.Server:
connection_params = utils.irc.IRCConnectionParameters(
*self.database.servers.get(server_id))
connection_params.args = connection_param_args
new_server = IRCServer.Server(self, self._events,
connection_params.id, connection_params.alias, connection_params)
self._events.on("new.server").call(server=new_server)
if not connect:
return new_server
self.connect(new_server)
return new_server
def get_server_by_id(self, id: int) -> typing.Optional[IRCServer.Server]:
for server in self.servers.values():
if server.id == id:
return server
return None
def get_server_by_alias(self, alias: str) -> typing.Optional[IRCServer.Server]:
alias_lower = alias.lower()
for server in self.servers.values():
if server.alias.lower() == alias_lower:
return server
return None
def connect(self, server: IRCServer.Server) -> bool:
try:
server.connect()
except Exception as e:
self.log.warn("Failed to connect to %s: %s",
[str(server), str(e)])
return False
self.servers[server.fileno()] = server
self._read_poll.register(server.fileno(), select.POLLIN)
return True
def next_send(self) -> typing.Optional[float]:
next = None
for server in self.servers.values():
timeout = server.socket.send_throttle_timeout()
if (server.socket.waiting_throttled_send() and
(next == None or timeout < next)):
next = timeout
return next
def next_ping(self) -> typing.Optional[float]:
timeouts = []
for server in self.servers.values():
timeout = server.until_next_ping()
if not timeout == None:
timeouts.append(timeout)
if not timeouts:
return None
return min(timeouts)
def next_read_timeout(self) -> typing.Optional[float]:
timeouts = []
for server in self.servers.values():
timeouts.append(server.until_read_timeout())
if not timeouts:
return None
return min(timeouts)
def get_poll_timeout(self) -> float:
timeouts = []
timeouts.append(self._timers.next())
timeouts.append(self.next_send())
timeouts.append(self.next_ping())
timeouts.append(self.next_read_timeout())
timeouts.append(self.cache.next_expiration())
min_secs = min([timeout for timeout in timeouts if not timeout == None])
return min_secs*1000 # return milliseconds
def disconnect(self, server: IRCServer.Server):
del self.servers[server.fileno()]
self._trigger_both()
def _timed_reconnect(self, event: EventManager.Event):
if not self.reconnect(event["server_id"],
event.get("connection_params", None)):
event["timer"].redo()
def reconnect(self, server_id: int, connection_params: typing.Optional[
utils.irc.IRCConnectionParameters]=None) -> bool:
args = {} # type: typing.Dict[str, str]
if not connection_params == None:
args = typing.cast(utils.irc.IRCConnectionParameters,
connection_params).args
server = self.add_server(server_id, False, args)
server.reconnected = True
if self.connect(server):
self.servers[server.fileno()] = server
return True
return False
def set_setting(self, setting: str, value: typing.Any):
self.database.bot_settings.set(setting, value)
def get_setting(self, setting: str, default: typing.Any=None) -> typing.Any:
return self.database.bot_settings.get(setting, default)
def find_settings(self, pattern: str, default: typing.Any=[]
) -> typing.List[typing.Any]:
return self.database.bot_settings.find(pattern, default)
def find_settings_prefix(self, prefix: str, default: typing.Any=[]
) -> typing.List[typing.Any]:
return self.database.bot_settings.find_prefix(
prefix, default)
def del_setting(self, setting: str):
self.database.bot_settings.delete(setting)
def _daemon_thread(self, target: typing.Callable[[], None]):
thread = threading.Thread(target=target)
thread.daemon = True
thread.start()
return thread
def run(self):
self.running = True
self._read_thread = self._daemon_thread(
lambda: self._loop_catch("read", self._read_loop))
self._write_thread = self._daemon_thread(
lambda: self._loop_catch("write", self._write_loop))
self._event_loop()
def _kill(self):
self.running = False
self._trigger_both()
def _event_loop(self):
while self.running or not self._event_queue.empty():
if not self.servers:
self._kill()
self.log.warn("No servers, exiting")
break
kill = False
item = self._event_queue.get(block=True, timeout=None)
if item.type == TriggerEventType.Action:
try:
item.callback()
except:
self._kill()
raise
elif item.type == TriggerEventType.Kill:
self._kill()
if not item.callback == None:
item.callback()
def _post_send_factory(self, server, lines):
return lambda: server._post_send(lines)
def _post_read_factory(self, server, lines):
return lambda: server._post_read(lines)
def _loop_catch(self, name: str, loop: typing.Callable[[], None]):
try:
loop()
except BitBotPanic:
return
except Exception as e:
self.panic("Exception on '%s' thread" % name)
def _write_loop(self):
while self.running:
with self._write_condition:
writeable = False
for fd, server in self.servers.items():
if server.socket.waiting_immediate_send():
self._write_poll.register(fd, select.POLLOUT)
writeable = True
if not writeable:
self._write_condition.wait()
continue
events = self._write_poll.poll()
for fd, event in events:
if event & select.POLLOUT:
self._write_poll.unregister(fd)
if fd in self.servers:
server = self.servers[fd]
try:
lines = server._send()
except:
self.log.error("Failed to write to %s",
[str(server)])
raise
event_item = TriggerEvent(TriggerEventType.Action,
self._post_send_factory(server, lines))
self._event_queue.put(event_item)
def _read_loop(self):
while self.running:
events = self._read_poll.poll(self.get_poll_timeout())
self.trigger(self._check, False)
for fd, event in events:
if fd == self._rtrigger_server.fileno():
# throw away data from trigger socket
with self._rtrigger_lock:
self._rtrigger_server.recv(1024)
self._rtriggered = False
else:
if not fd in self.servers:
self._read_poll.unregister(fd)
continue
server = self.servers[fd]
if event & select.POLLIN:
lines = server.read()
if lines == None:
server.disconnect()
continue
self.trigger(self._post_read_factory(server, lines),
False)
elif event & select.POLLHUP:
self.log.warn("Recieved POLLHUP for %s", [str(server)])
server.disconnect()
def _check(self):
self._timers.call()
self.cache.expire()
throttle_filled = False
for server in list(self.servers.values()):
if server.read_timed_out():
self.log.warn("Pinged out from %s", [str(server)])
server.disconnect()
elif server.ping_due() and not server.ping_sent:
server.send_ping()
server.ping_sent = True
if not server.socket.connected:
self._events.on("server.disconnect").call(server=server)
self.disconnect(server)
if not self.get_server_by_id(server.id):
reconnect_delay = self.config.get("reconnect-delay", 10)
self._timers.add("reconnect", reconnect_delay,
server_id=server.id)
self.log.warn(
"Disconnected from %s, reconnecting in %d seconds",
[str(server), reconnect_delay])
elif (server.socket.waiting_throttled_send() and
server.socket.throttle_done()):
server.socket._fill_throttle()
throttle_filled = True
if throttle_filled:
self.trigger_write()