465 lines
17 KiB
Python
465 lines
17 KiB
Python
VERSION: str = ""
|
|
with open("VERSION", "r") as version_file:
|
|
VERSION = "v%s" % version_file.read().strip()
|
|
SOURCE: str = "https://git.io/bitbot"
|
|
URL: str = "https://bitbot.dev"
|
|
|
|
import enum, queue, os, queue, select, socket, sys, threading, time, traceback
|
|
import typing, uuid
|
|
from src import Config, EventManager, Exports, IRCServer, Logging
|
|
from src import ModuleManager, PollHook, PollSource, Socket, Timers, utils
|
|
|
|
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 ListLambdaPollHook(PollHook.PollHook):
|
|
def __init__(self,
|
|
collection: typing.Callable[[], typing.Iterable[typing.Any]],
|
|
func: typing.Callable[[typing.Any], None]):
|
|
self._collection = collection
|
|
self._func = func
|
|
def next(self):
|
|
timeouts = [self._func(i) for i in self._collection()]
|
|
timeouts = [t for t in timeouts if t is not None]
|
|
return min(timeouts or [None])
|
|
|
|
class Bot(object):
|
|
def __init__(self, directory, data_directory, args, cache, config, database,
|
|
events, exports, log, modules, timers):
|
|
self.directory = directory
|
|
self.data_directory = data_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._writing = False
|
|
self._reading = False
|
|
self.servers = {}
|
|
self.reconnections = {}
|
|
|
|
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._poll_timeouts = [] # typing.List[PollHook.PollHook]
|
|
self._poll_timeouts.append(ListLambdaPollHook(
|
|
lambda: self.servers.values(),
|
|
lambda server: server.until_read_timeout()))
|
|
|
|
self._poll_timeouts.append(ListLambdaPollHook(
|
|
lambda: self.servers.values(),
|
|
lambda server: server.until_next_ping()))
|
|
|
|
self._poll_timeouts.append(ListLambdaPollHook(
|
|
lambda: self.servers.values(), self._throttle_timeout))
|
|
|
|
self._poll_sources = [] # typing.List[PollSource.PollSource]
|
|
|
|
def add_poll_hook(self, hook: PollHook.PollHook):
|
|
self._poll_timeouts.append(hook)
|
|
def add_poll_source(self, source: PollSource.PollSource):
|
|
self._poll_sources.append(source)
|
|
|
|
def _throttle_timeout(self, server: IRCServer.Server):
|
|
if server.socket.waiting_throttled_send():
|
|
return server.socket.send_throttle_timeout()
|
|
return None
|
|
|
|
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[typing.Tuple[TriggerResult, typing.Any]
|
|
] = queue.Queue(1)
|
|
|
|
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 and
|
|
isinstance(returned, Exception)):
|
|
raise returned
|
|
elif type == TriggerResult.Return:
|
|
return returned
|
|
|
|
def panic(self, reason):
|
|
exc_info = False
|
|
if any(sys.exc_info()):
|
|
exc_info = True
|
|
|
|
self.log.critical("panic() called: %s", [reason], exc_info=exc_info)
|
|
sys.exit(utils.consts.Exit.PANIC)
|
|
|
|
def get_config(self, name: str) -> Config.Config:
|
|
path = os.path.join(self.data_directory, "%s.conf" % name)
|
|
config = Config.Config(name, path)
|
|
config.load()
|
|
return config
|
|
|
|
def _module_lists(self):
|
|
module_lists = self.get_config("modules")
|
|
whitelist = module_lists.get_list("whitelist")
|
|
blacklist = module_lists.get_list("blacklist")
|
|
|
|
return whitelist, blacklist
|
|
|
|
def load_modules(self
|
|
) -> typing.Tuple[typing.List[str], typing.List[str]]:
|
|
whitelist, blacklist = self._module_lists()
|
|
return self.modules.load_modules(self, whitelist=whitelist,
|
|
blacklist=blacklist)
|
|
def try_reload_modules(self) -> ModuleManager.TryReloadResult:
|
|
whitelist, blacklist = self._module_lists()
|
|
return self.modules.try_reload_modules(self, whitelist=whitelist,
|
|
blacklist=blacklist)
|
|
|
|
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:
|
|
ip = ""
|
|
if server.socket.connected_ip is not None:
|
|
ip = f" ({server.socket.connected_ip})"
|
|
|
|
self.log.warn("Failed to connect to %s%s: %s",
|
|
[str(server), ip, str(e)])
|
|
self.log.debug("Connection failure reason:", exc_info=True)
|
|
return False
|
|
self.servers[server.fileno()] = server
|
|
self._read_poll.register(server.fileno(), select.POLLIN)
|
|
return True
|
|
|
|
def get_poll_timeout(self) -> float:
|
|
timeouts = []
|
|
for poll_timeout in self._poll_timeouts:
|
|
timeouts.append(poll_timeout.next())
|
|
|
|
min_secs = min([timeout for timeout in timeouts if not timeout == None])
|
|
return max([min_secs, 0])
|
|
|
|
def disconnect(self, server: IRCServer.Server):
|
|
del self.servers[server.fileno()]
|
|
self._trigger_both()
|
|
|
|
def _timed_reconnect(self, timer: Timers.Timer):
|
|
server_id = timer.kwargs["server_id"]
|
|
params = timer.kwargs.get("connection_params", None)
|
|
if not self.reconnect(server_id, params):
|
|
timer.redo()
|
|
else:
|
|
del self.reconnections[server_id]
|
|
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=None, prefix: str=None,
|
|
default: typing.Any=[]) -> typing.List[typing.Any]:
|
|
if not pattern == None:
|
|
return self.database.bot_settings.find(pattern, default)
|
|
elif not prefix == None:
|
|
return self.database.bot_settings.find_prefix(prefix, default)
|
|
else:
|
|
raise ValueError("Please provide 'pattern' or 'prefix'")
|
|
|
|
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._writing = True
|
|
self._reading = 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 stop(self, reason: str="Stopping"):
|
|
self._reading = False # disable read thread
|
|
self.trigger_read()
|
|
for server in self.servers.values():
|
|
line = server.send_quit(reason)
|
|
line.events.on("send").hook(self._shutdown_hook(server))
|
|
def _shutdown_hook(self, server):
|
|
def shutdown(e):
|
|
server.disconnect()
|
|
self.disconnect(server)
|
|
if not self.servers:
|
|
self._writing = False
|
|
return shutdown
|
|
|
|
def _kill(self):
|
|
self._writing = False
|
|
self._reading = False
|
|
self._trigger_both()
|
|
|
|
def _event_loop(self):
|
|
while ((self._writing or self._reading) or
|
|
not self._event_queue.empty()):
|
|
try:
|
|
item = self._event_queue.get(block=True,
|
|
timeout=self.get_poll_timeout())
|
|
except queue.Empty:
|
|
# caused by timeout being hit.
|
|
continue
|
|
finally:
|
|
self._check()
|
|
|
|
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()
|
|
break
|
|
|
|
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 Exception as e:
|
|
self.panic("Exception on '%s' thread" % name)
|
|
|
|
def _write_loop(self):
|
|
while self._writing:
|
|
poll_sources = {}
|
|
with self._write_condition:
|
|
fds = []
|
|
for fd, server in self.servers.items():
|
|
if server.socket.waiting_immediate_send():
|
|
fds.append(fd)
|
|
|
|
for poll_source in self._poll_sources:
|
|
for fileno in poll_source.get_writables():
|
|
poll_sources[fileno] = poll_source
|
|
fds.append(fileno)
|
|
|
|
if not fds:
|
|
self._write_condition.wait()
|
|
continue
|
|
else:
|
|
for fd in fds:
|
|
self._write_poll.register(fd, select.POLLOUT)
|
|
|
|
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)
|
|
elif fd in poll_sources:
|
|
def _trigger(source, fd):
|
|
return lambda: source.is_writeable(fd)
|
|
self.trigger(_trigger(poll_sources[fd], fd))
|
|
|
|
def _read_loop(self):
|
|
poll_sources = {}
|
|
while self._reading:
|
|
new_poll_sources = {}
|
|
for poll_source in self._poll_sources:
|
|
for fileno in poll_source.get_readables():
|
|
new_poll_sources[fileno] = poll_source
|
|
for fileno in new_poll_sources:
|
|
if not fileno in poll_sources:
|
|
poll_sources[fileno] = new_poll_sources[fileno]
|
|
self._read_poll.register(fileno, select.POLLIN)
|
|
for fileno in list(poll_sources.keys()):
|
|
if not fileno in new_poll_sources:
|
|
del poll_sources[fileno]
|
|
self._read_poll.unregister(fileno)
|
|
|
|
events = self._read_poll.poll()
|
|
|
|
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
|
|
elif fd in poll_sources:
|
|
def _trigger(source, fd):
|
|
return lambda: source.is_readable(fd)
|
|
self.trigger(_trigger(poll_sources[fd], fd))
|
|
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
|
|
|
|
event_item = TriggerEvent(TriggerEventType.Action,
|
|
self._post_read_factory(server, lines))
|
|
self._event_queue.put(event_item)
|
|
elif event & select.POLLHUP:
|
|
self.log.warn("Recieved POLLHUP for %s", [str(server)])
|
|
server.disconnect()
|
|
|
|
def _check(self):
|
|
for poll_timeout in self._poll_timeouts:
|
|
if poll_timeout.next() == 0:
|
|
poll_timeout.call()
|
|
|
|
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)
|
|
|
|
timer = self._timers.add("timed-reconnect",
|
|
self._timed_reconnect, reconnect_delay,
|
|
server_id=server.id)
|
|
self.reconnections[server.id] = timer
|
|
|
|
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()
|