<commit_before>#!/usr/bin/python3

import api

"""
teams = api.teams.filter(score > 0)
for server in api.shell_servers.get_online_servers():
  shell = server.get_ssh()
  for team in teams:
    current_head = db.submission_heads.find(tid)
    most_recent_sub = db.submissions.get_most_recent(tid)

    if current_head is None or current_head < most_recent_sub:
      unlocked_problems = api.problems.get_unlocked(team)

      for user in team.get_enabled_users():
        user_sym_links = shell.get_sym_links(user)
        additional_sym_links = unlocked_problems / user_sym_links
        removable_sym_links = user_sym_links / unlocked_problems

        shell.sym_links.add(user, additional_sym_links)
        shell.sym_links.remove(user, removable_sym_links)
"""

def run():
    db = api.common.get_conn()

    teams = api.team.get_all_teams(show_ineligible=True)

    for server in api.shell_servers.get_servers():
        shell = api.shell_servers.get_connection(server["sid"])

        submission_heads = db.submission_heads.find_one({"sid": server["sid"]})
        if submission_heads is None:
            db.submission_heads.insert({"sid": server["sid"], "heads": {}})
            submission_heads = db.submission_heads.find_one({"sid": server["sid"]})

        submission_heads = submission_heads["heads"]
        for team in teams:
            submission = api.problem.get_most_recent_submission(tid=team["tid"], correctness=True)
            team_head_pid = submission_heads.get(team["tid"], None)
            if submission is not None and (team_head_pid is None or submission["pid"] != team_head_pid):
                print("Team '%s' is outdated. %s -> %s" % (team["team_name"], team_head_pid, submission["pid"]))

                unlocked_problems = api.problem.get_unlocked_problems(tid=team["tid"])

                db.submission_heads.update({"sid": server["sid"]}, {"$set": {"heads."+team["tid"]: submission["pid"]}})
            else:
                print("Team '%s' up-to-date." % team["team_name"])


<commit_msg>Added full implementation based on unlocked problems<commit_after>#!/usr/bin/python3

import api
import spur
import time
from os.path import join

def cache(f, *args, **kwargs):
    result = f(cache=False, *args, **kwargs)
    key = api.cache.get_mongo_key(f, *args, **kwargs)
    api.cache.set(key, result)
    return result

class CantConnectException(Exception):
    pass

# open shell server connections
connections = {}

def get_connection(sid):
    global connections
    if connections.get(sid, None) == None:
        try:
            connections[sid] = api.shell_servers.get_connection(sid)
        except api.common.WebException as e:
            raise CantConnectException
    return connections[sid]

def get_home_dir(shell, user):
    # first make sure the user has an account
    if shell.run(["id", "-u", user], allow_error=True).return_code != 0:
        return None

    home_dir = shell.run(["bash", "-c", "echo ~%s" % user]).output.decode("utf-8").strip()

    return home_dir

@api.cache.memoize()
def get_symlinks(sid, user):
    shell = get_connection(sid)

    home_dir = get_home_dir(shell, user)
    if home_dir == None:
        return None

    problems_path = join(home_dir, "problems")
    try:
        result = shell.run(["sudo", "ls", problems_path]).output.decode("utf-8")
        return result.split("\n")[:-1]
    except spur.results.RunProcessError as e:
        # directory must not exist
        result = shell.run(["sudo", "mkdir", problems_path], allow_error=True)
        if result.return_code != 0:
            # something else is wrong, so give up
            return None

        # try again
        get_symlinks(shell, user)

def rm_symlinks(sid, user, to_remove):
    shell = get_connection(sid)

    home_dir = get_home_dir(shell, user)
    if home_dir == None:
        return None

    problems_path = join(home_dir, "problems")
    for name in to_remove:
        symlink = join(problems_path, name)
        try:
            shell.run(["sudo", "rm", symlink])
            print("Removed symlink %s" % symlink)
        except spur.results.RunProcessError as e:
            print("Failed to remove symlink %s" % symlink)

def add_symlinks(sid, user, to_add):
    shell = get_connection(sid)

    home_dir = get_home_dir(shell, user)
    if home_dir == None:
        return None

    problems_path = join(home_dir, "problems")
    for name, path in to_add.items():
        symlink = join(problems_path, name)
        try:
            shell.run(["sudo", "ln", "-s", path, symlink])
            print("Added symlink %s --> %s" % (symlink, path))
        except spur.results.RunProcessError as e:
            print("Failed to add symlink %s->%s" % (symlink, path))

def run():
    global connections

    teams = api.team.get_all_teams(show_ineligible=True)

    for server in api.shell_servers.get_servers():

        try:
            for team in teams:
                unlocked_problems = api.problem.get_unlocked_problems(tid=team["tid"])
                correct_symlinks = {p["name"]:p["deployment_directory"] for p in unlocked_problems if p["should_symlink"]}
                correct_names = set(correct_symlinks.keys())

                for user in api.team.get_team_members(tid=team["tid"]):
                    start = time.time()
                    current_symlinks = get_symlinks(server["sid"], user["username"])
                    if current_symlinks == None:
                        continue

                    symlinked_names = set(current_symlinks)
                    to_add = correct_names - symlinked_names
                    to_remove = symlinked_names - correct_names

                    if len(to_add) > 0:
                        add_symlinks(server["sid"], user["username"], {name:correct_symlinks[name] for name in to_add})
                    if len(to_remove) > 0:
                        rm_symlinks(server["sid"], user["username"], to_remove)

                    if len(to_add) > 0 or len(to_remove) > 0:
                        cache(get_symlinks, server["sid"], user["username"])

                    print("Took %.2f seconds." % (time.time() - start))

        except CantConnectException as e:
            print("Couldn't connect to server \"%s\"" % server["name"])

        connections[server["sid"]] = None

{'fraction_non_alphanumeric': 0.1229624940655167, 'fraction_numerical': 0.0020572875454977053, 'mean_word_length': 2.640552995391705, 'pattern_counts': {'":': 6, '<': 4, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3209496', 'n_tokens_mistral': 1989, 'n_tokens_neox': 1914, 'n_words': 435}