mirror of https://github.com/ekimekim/wubloader
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
118 lines
3.4 KiB
Python
118 lines
3.4 KiB
Python
|
|
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)
|