From ebeb9f3f9e50c776ae66a4fd0523a9d43e6896c5 Mon Sep 17 00:00:00 2001 From: Mike Lang Date: Wed, 26 Oct 2022 04:35:32 -0400 Subject: [PATCH] Combine old schedule bot with group membership bot --- zulip/group_membership.py | 117 -------------------- zulip/schedulebot.py | 224 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 224 insertions(+), 117 deletions(-) delete mode 100644 zulip/group_membership.py create mode 100644 zulip/schedulebot.py diff --git a/zulip/group_membership.py b/zulip/group_membership.py deleted file mode 100644 index f94b0c3..0000000 --- a/zulip/group_membership.py +++ /dev/null @@ -1,117 +0,0 @@ - -import gevent.monkey -gevent.monkey.patch_all() - -import logging - -import requests - -logging.basicConfig(level='INFO') - -class Client(object): - def __init__(self, base_url, email, api_key): - self.base_url = base_url - self.email = email - self.api_key = api_key - self.session = requests.Session() - - def request(self, method, *path, **params): - if method == 'GET': - args = {"params": params} - else: - args = {"data": params} - url = "/".join([self.base_url, "api/v1"] + list(path)) - resp = session.request(method, url, auth=(self.email, self.api_key), **args) - resp.raise_for_status() - return resp.json() - -def get_membership(client): - """Returns {group id: member set}""" - return { - group["id"]: set(group["members"]) - for group in client.request("GET", "user_groups")["user_groups"] - } - -def update_members(client, group_id, old_members, new_members): - added = new_members - old_members - removed = old_members - new_members - client.request("POST", "user_groups", group_id, "members", add=list(added), delete=list(removed)) - -def determine_members(schedules, hour): - return set( - user for user, schedule in schedules.items() - if hour in schedule - ) - -def run_hour(client, user_map, groups, hour): - logging.info("Setting groups for hour {}".format(hour)) - members = get_membership(client) - def run_group(group_id, schedules): - new_members = determine_members(schedules, hour) - new_members = set(user_map[id] for id in new_members) - assert group_id in members, "group {} doesn't exist".format(group_id) - update_members(client, group_id, members[group_id], new_members) - gevent.pool.Group().map(run_group, groups.items()) - -def parse_config(conf_file): - MAX_DAYS=8 - with open(conf_file) as f: - config = yaml.safe_load(f) - assert "url" in config - assert "api_key" in config - assert "start_time" in config - all_users = set() - for group_id, schedules in config["groups"]: - all_users |= set(schedules.keys()) - for user_id, schedule in schedules: - hours = set() - for part in scuedule.split(","): - part = part.strip() - if part.endswith("/24"): - hour = int(part.split('/')[0]) - hours |= set(24 * day + hour for day in range(MAX_DAYS)) - elif part in "ZDAN": - shift = "ZDAN".index(part) - hours |= set(hour for hour in range(24 * MAX_DAYS) if (hour % 24) // 6 == shift) - else: - hours.add(int(part)) - schedules[user_id] = hours - missing = all_users - set(config["members"]) - assert not missing, "missing: {}".format(", ".join(missing)) - - -def main(conf_file, hour=-1): - """ - config: - url: the base url of the instance - api_key: authentication - start_time: Time of the first hour, as epoch int - members: - NAME: USER_ID - groups: - GROUP_ID: - NAME: SCHEDULE - Where: - schedule: Comma-seperated list of hours - hour: One of: - Integer hour, representing that hour of the run - N/24, expanding to hour N of each day (eg. 0/24 is midnight each day) - Z | D | A | N, expanding to all hours of Zeta, Dawn Guard, Alpha Flight, Night Watch respectively. - """ - config = parse_config(conf_file) - client = Client(config["url"], config["email"], config["api_key"]) - user_map = config["members"] - groups = config["groups"] - if hour >= 0: - run_hour(client, user_map, groups, hour) - return - while True: - hour = (time.time() - start_time) // 3600 - run_hour(client, user_map, groups, hour) - next_hour = start_time + 3600 * (hour + 1) - remaining = next_hour - time.time() - if remaining > 0: - time.sleep(remaining) - -if __name__ == '__main__': - argh.dispatch_command(main) diff --git a/zulip/schedulebot.py b/zulip/schedulebot.py new file mode 100644 index 0000000..0cc37d2 --- /dev/null +++ b/zulip/schedulebot.py @@ -0,0 +1,224 @@ + +import gevent.monkey +gevent.monkey.patch_all() + +import csv +import json +import logging +import time +from datetime import datetime, timedelta + +import gevent.pool +import argh + +import requests +session = requests.Session() + +logging.basicConfig(level='INFO') + +class Client(object): + def __init__(self, base_url, email, api_key): + self.base_url = base_url + self.email = email + self.api_key = api_key + + def request(self, method, *path, **params): + if method == 'GET': + args = {"params": params} + else: + args = {"data": { + k: v if isinstance(v, str) else json.dumps(v) + for k, v in params.items() + }} + url = "/".join([self.base_url, "api/v1"] + list(map(str, path))) + resp = session.request(method, url, auth=(self.email, self.api_key), **args) + if not resp.ok: + logging.info(repr(params)) + logging.info(f"Got {resp.status_code} for {url}: {resp.text}") + resp.raise_for_status() + return resp.json() + +def get_membership(client): + """Returns {group id: member set}""" + return { + group["id"]: set(group["members"]) + for group in client.request("GET", "user_groups")["user_groups"] + } + +def update_members(client, group_id, old_members, new_members): + logging.info(f"Updating group {group_id}: {old_members} -> {new_members}") + added = new_members - old_members + removed = old_members - new_members + if added or removed: + client.request("POST", "user_groups", group_id, "members", add=list(added), delete=list(removed)) + +def get_role_at_hour(hours, hour): + if 0 <= hour < len(hours): + return hours[hour] + return "" + +def determine_members(schedule, role, hour): + return set( + user_id for user_id, (_, hours) in schedule.items() + if get_role_at_hour(hours, hour) == role + ) + +def get_display_name(client, user_id): + return client.request("GET", "users", user_id)["user"]["full_name"] + + +def update_groups(client, group_ids, schedule, hour): + logging.info("Setting groups for hour {}".format(hour)) + members = get_membership(client) + def run_group(item): + group_name, group_id = item + new_members = determine_members(schedule, group_name, hour) + assert group_id in members, "group {} doesn't exist".format(group_id) + update_members(client, group_id, members[group_id], new_members) + gevent.pool.Group().map(run_group, group_ids.items()) + + +def post_schedule(client, send_client, start_time, schedule, stream, hour, no_mentions): + going_offline = [] + coming_online = [] + display_names = {} + supervisor = None + for user_id, (_, hours) in schedule.items(): + prev = get_role_at_hour(hours, hour - 1) + now = get_role_at_hour(hours, hour) + if now == "Supervisor": + supervisor = user_id + if user_id not in display_names: + display_names[user_id] = gevent.spawn(get_display_name, client, user_id) + if prev != now: + if prev != "": + going_offline.append((prev, user_id)) + if now != "": + coming_online.append((now, user_id)) + if user_id not in display_names: + display_names[user_id] = gevent.spawn(get_display_name, client, user_id) + # sort by role + going_offline.sort() + coming_online.sort() + logging.info(f"Going offline: {going_offline}") + logging.info(f"Coming online: {coming_online}") + + start_time_pst = datetime.utcfromtimestamp(start_time) - timedelta(hours=8) + current_time = (start_time_pst + timedelta(hours=hour)).replace(minute=0, second=0, microsecond=0) + hour_pst = current_time.hour + shift = hour_pst // 6 + shift = ["night-watch", "zeta", "dawn-guard", "alpha-flight"][shift] + shift_hour = hour_pst % 6 + 1 + + def render_name(user_id, mention=True): + if no_mentions: + mention = False + fallback, _ = schedule[user_id] + try: + result = display_names[user_id].get() + except Exception: + logging.warning(f"Failed to fetch user {user_id}", exc_info=True) + return f"**{fallback}**" + if mention: + return f"@**{result}|{user_id}**" + else: + return f"**{result}**" + + lines = [ + f"**Shift changes for :{shift}: Hour {shift_hour} | Bustime {hour:02d}:00 - {hour+1:02d}:00 | :**", + "Make sure to *mute/unmute* #**current-shift** as needed!", + ] + if supervisor is None: + logging.warning("No supervisor found") + else: + name = render_name(supervisor, mention=False) + lines.append(f"Your shift supervisor is {name}") + if coming_online: + lines += [ + "", + "---", + "Coming online:" + ] + [ + f"- {render_name(user_id)} - {role}" + for role, user_id in coming_online + ] + if going_offline: + lines += [ + "", + "---", + "Going offline:" + ] + [ + f"- {render_name(user_id)} - {role}" + for role, user_id in going_offline + ] + lines += [ + "", + "---", + ] + + send_client.request("POST", "messages", + type="stream", + to=stream, + topic="Schedule", + content="\n".join(lines), + ) + + +def parse_schedule(user_ids, schedule_file): + schedule = {} + with open(schedule_file) as f: + for row in csv.reader(f): + name = row[0] + if name in ["", "Chat Member", "[Counts]", "Hour of the Run"] or name.startswith("-"): + continue + if name not in user_ids: + logging.warning(f"No user id known for user {name}") + continue + user_id = user_ids[name] + schedule[user_id] = name, row[1:] + return schedule + + +def main(conf_file, hour=-1, no_groups=False, stream="General", no_mentions=False): + """ + config: + url: the base url of the instance + api_user: auth used for general api calls + send_user: + auth used for sending messages + defaults to api_user, but you may want a vanity name / avatar + start_time: Time of the first hour, as epoch int + schedule: Path to the schedule CSV file + members: + NAME: USER_ID + groups: + NAME: GROUP_ID + authentication is an object {email, api_key} + """ + with open(conf_file) as f: + config = json.load(f) + client = Client(config["url"], config["api_user"]["email"], config["api_user"]["api_key"]) + send_auth = config.get("send_user", config["api_user"]) + send_client = Client(config["url"], send_auth["email"], send_auth["api_key"]) + group_ids = config["groups"] + start_time = config["start_time"] + schedule = parse_schedule(config["members"], config["schedule"]) + if hour >= 0: + if not no_groups: + update_groups(client, group_ids, schedule, hour) + if stream: + post_schedule(client, send_client, start_time, schedule, stream, hour, no_mentions) + return + while True: + hour = (time.time() - start_time) // 3600 + if not no_groups: + update_groups(client, group_ids, schedule, hour) + if stream: + post_schedule(client, send_client, start_time, schedule, stream, hour, no_mentions) + next_hour = start_time + 3600 * (hour + 1) + remaining = next_hour - time.time() + if remaining > 0: + time.sleep(remaining) + +if __name__ == '__main__': + argh.dispatch_command(main)