itte/itte.py

186 lines
7.4 KiB
Python
Raw Normal View History

import argparse
2018-09-12 20:01:24 +00:00
import socket
2018-09-15 23:18:44 +00:00
import yaml
from time import sleep
2018-09-17 03:51:09 +00:00
from random import randint
from sys import exit
2018-09-12 20:01:24 +00:00
2018-09-16 04:48:50 +00:00
class Util:
"""Utility functions."""
def yml(self, yml_file):
"""Open a YAML file and return a dictionary of values."""
try:
fh = open(yml_file, "r")
data = yaml.safe_load(fh)
fh.close()
except TypeError:
exit("[debug][err] Cannot load YML file. Please check it exists.")
2018-09-16 04:48:50 +00:00
return data
2018-09-17 03:51:09 +00:00
def rand(self, lst):
"""Return a random item from a given list."""
return lst[randint(0, len(lst)-1)]
2018-09-16 04:48:50 +00:00
def cli_flags(self):
"""Parse command line flags."""
self.argp = argparse.ArgumentParser()
self.argp.add_argument("-c", "--config", help="Config file")
return self.argp.parse_args()
2018-09-12 20:01:24 +00:00
class IRC:
"""Methods for basic IRC communication."""
2018-09-12 20:01:24 +00:00
def config(self, def_conf):
2018-09-15 23:18:44 +00:00
"""Load runtime settings from a YAML config file, and returns a
dictionary of config values. Looks for the file in a runtime path or in
the default location."""
2018-09-16 04:48:50 +00:00
self.util = Util()
# Check for runtime config locatiion
flags = self.util.cli_flags()
if flags.config != "":
cfg = self.util.yml(flags.config)
else:
cfg = self.util.yml(def_conf)
2018-09-15 23:18:44 +00:00
self.server = (cfg["server"]["host"], cfg["server"]["port"])
self.channels = cfg["channels"]
self.bot_nick = cfg["bot_nick"]
self.req_prefix = cfg["req_prefix"]
self.admin_user = cfg["admin"]["user"]
self.admin_code = cfg["admin"]["code"]
self.debug = cfg["debug"]
return cfg
2018-09-12 20:01:24 +00:00
2018-09-15 23:18:44 +00:00
def run(self, listen_hook):
"""A routine that connects to a server, joins channels, and attaches
the request listener hook to a loop."""
self.connect(self.server, self.bot_nick)
# Wait for server to reply before joining channels
svr_greet = self.receive()
while ("001 " + self.bot_nick) not in svr_greet:
sleep(1)
svr_greet = self.receive()
2018-09-15 23:18:44 +00:00
self.join_channels(self.channels)
2018-09-14 05:54:21 +00:00
while 1:
2018-09-14 23:37:47 +00:00
data = self.receive()
2018-10-13 17:18:47 +00:00
self.keepalive(data, self.bot_nick)
2018-09-15 23:18:44 +00:00
self.msg = self.parse(data, self.req_prefix)
for c in self.channels:
2018-09-14 23:37:47 +00:00
# Pass in a context dict for handlers
listen_hook({"msg": self.msg, "listen_chan": c})
2018-09-12 20:01:24 +00:00
def connect(self, server, bot_nick):
2018-09-12 20:01:24 +00:00
"""Connect to the server and sends user/nick information."""
try:
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.sock.connect(server)
except ConnectionError as err:
2018-10-13 17:18:47 +00:00
exit("[debug][err] " + str(err))
self.send("USER", bot_nick + " " + bot_nick + " " + \
bot_nick + " " + bot_nick)
self.send("NICK", bot_nick)
2018-09-14 23:37:47 +00:00
def disconnect(self, resp_msg, quit_msg):
"""Notify the admin user and disconnect from the server."""
2018-09-14 23:37:47 +00:00
self.send("PRIVMSG", resp_msg, recvr=self.admin_user)
self.send("QUIT", ":" + quit_msg)
2018-09-14 23:37:47 +00:00
# Currently only one server per app instance is supported, so
# disconnect also exits the app
exit("Shutting down ...")
2018-09-17 03:51:09 +00:00
2018-10-13 17:18:47 +00:00
def keepalive(self, line, bot_nick):
2018-09-14 23:37:47 +00:00
"""Stay connected to a server by responding to server pings."""
2018-10-13 17:18:47 +00:00
if line.split(" ")[0] == "PING":
resp = line.replace("PING", "PONG", 1)
if self.debug:
print("[debug][send] " + resp)
self.sock.sendall(bytes(resp + "\r\n", "utf-8"))
2018-10-13 17:18:47 +00:00
# Fallback to ensure process exits if timeout occurs
elif (bot_nick in line.split(" ")[0]) and \
("QUIT :Ping timeout:" in line):
exit("[debug][err] Ping timeout, exited.")
2018-09-14 23:37:47 +00:00
def join_channels(self, channels):
"""Join channels given a list of channel names."""
for c in channels:
if c.strip() != "" or c.strip() != "#":
self.send("JOIN", c)
2018-09-14 05:54:21 +00:00
def send(self, command, text, *args, **kwargs):
2018-09-17 03:51:09 +00:00
"""Send messages given the IRC command and text. Optionally specify a
2018-09-14 05:54:21 +00:00
message recipient with `recvr=user`."""
recvr = kwargs.get("recvr", "")
if recvr != "":
recvr += " :"
if self.debug:
print("[debug][send] " + command + " " + recvr + text)
bs = bytes(command + " " + recvr + text + "\r\n", "utf-8")
try:
self.sock.sendall(bs)
except BrokenPipeError as err:
if self.debug:
print("[debug] " + str(err) + " at `" + \
bs.decode("utf-8").strip() + "`")
pass
2018-09-14 23:37:47 +00:00
def receive(self):
"""Get messages from the connected socket."""
2019-01-15 17:26:21 +00:00
data = self.sock.recv(4096).decode("utf-8").strip("\r\n")
if self.debug:
print("[debug][recv] " + data)
return data
2018-09-14 23:37:47 +00:00
def parse(self, line, req_prefix):
2018-09-14 05:54:21 +00:00
"""Using received data from a socket, extract the request, the nick and
username of the requester, the channel where the request originated and
return a dictionary of values."""
2018-09-12 20:01:24 +00:00
data = {"req": "", "req_chan": "", "nick": "", "user": ""}
if (":" + req_prefix) in line:
2018-09-12 20:01:24 +00:00
data["req"] = line.split("PRIVMSG", 1)[1].split(":" + \
req_prefix, 1)[1].strip()
data["req_chan"] = line.split("PRIVMSG ", \
1)[1].split(" :", 1)[0]
2018-09-12 20:01:24 +00:00
data["nick"] = line.split("!~", 1)[0][1:]
data["user"] = line.split("!~", 2)[0][0:].split("@", 1)[0]
2018-09-12 20:01:24 +00:00
return data
2018-09-14 23:37:47 +00:00
def listen(self, context, trigger, handler, *args, **kwargs):
"""Listen for a trigger and call the handler. It takes a context
dictionary (to determine the channel and recipient), trigger string and
corresponding handler function. Optional flags (chan, query, admin) can
be used to specify whether a trigger is active in channels or by
/msg. e.g. `chan=False, query=True` to make a trigger query-only. By
default, it will listen in both channels and private messages."""
in_chan = kwargs.get("chan", True)
in_query = kwargs.get("query", True)
in_admin = kwargs.get("admin", False)
# Admin requests are query/pm only
if in_admin:
in_chan = False
msg = context["msg"]
channel = context["listen_chan"]
# Responses are sent via pm to the user by default, while requests made
# in a channel are sent to the channel. While it's possible to override
# the recvr, it usually easier to enable a trigger in the same location
# where the response will be sent.
context["recvr"] = msg["user"]
if msg["req"] == trigger:
# Respond only in the channel where the request was made
if in_chan and channel == msg["req_chan"]:
context["recvr"] = msg["req_chan"]
handler(context)
# Respond to query/pm
elif in_query and msg["req_chan"] == self.bot_nick:
handler(context)
# Respond only to the admin user
elif in_admin and msg["user"].lower() == \
self.admin_user.lower() and self.admin_code in msg["req"]:
handler(context)
def reply(self, cxt, text):
"""Alias of send() with PRIVMSG command preset."""
self.send("PRIVMSG", text, recvr=cxt["recvr"])