wubloader/zulip/group_membership.py

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)