diff --git a/zulip/group_membership.py b/zulip/group_membership.py new file mode 100644 index 0000000..d89e73d --- /dev/null +++ b/zulip/group_membership.py @@ -0,0 +1,117 @@ + +import gevent.monkey +gevent.monkey.patch_all() + +import logging + +import requests + +logging.basicConfig(level='INFO') + +class Client(object): + def __init__(self, base_url, api_key): + self.base_url = base_url + self.api_key = api_key + self.session = requests.Session() + + def request(self, method, *path, **params): + # TODO api key + if method == 'GET': + args = {"params": params} + else: + args = {"data": params} + url = "/".join([self.base_url, "api/v1"] + list(path)) + resp = session.request(method, url, **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["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)