<commit_before>""" Module for interacting with the problems """
import imp
import json

import api.common
import api.user
import api.team

from datetime import datetime
from api.common import validate, APIException, check
from voluptuous import Schema, Length, Required, Range
from pymongo.errors import DuplicateKeyError
from bson import json_util
from os.path import join, isfile

grader_base_path = "./graders"
check_graders_exist = True

submission_schema = Schema({
    Required("tid"): check(
        (0, "This does not look like a valid tid.", [str, Length(max=100)])),
    Required("pid"): check(
        (0, "This does not look like a valid pid.", [str, Length(max=100)])),
    Required("key"): check(
        (0, "This does not look like a valid key.", [str, Length(max=100)]))
})

problem_schema = Schema({
    Required("basescore"): check(
        (0, "Basescore must be a positive integer.", [int, Range(min=0)])),
    Required("category"): check(
        (0, "Category must be a string.", [str])),
    Required("desc"): check(
        (0, "The problem description must be a string.", [str])),
    Required("displayname"): check(
        (0, "The problem's display name must be a string.", [str])),
    Required("grader"): check(
        (0, "A grader does not exist at that path.", [
            lambda grader: not check_graders_exist or
            isfile(join(grader_base_path, grader))])),
    Required("threshold"): check(
        (0, "Threshold must be a positive integer.", [int, Range(min=0)])),

    "disabled": check(
        (0, "A problem's disabled state is either True or False.", [
            lambda disabled: type(disabled) == bool])),
    "autogen": check(
        (0, "A problem should either be autogenerated or not, True/False", [
            lambda autogen: type(autogen) == bool])),
    "relatedproblems": check(
        (0, "Related problems should be a list of related problems.", [list])),
    "pid": check(
        (0, "The problem's pid should be a string.", [str])),
    "weightmap": check(
        (0, "Weightmap should be a dict.", [dict])),
    "tags": check(
        (0, "Tags must be described as a list.", [list])),
    "hint": check(
        (0, "A hint must be a string.", [str])),

    "_id": check(
        (0, "Your problems should not already have _ids.", [lambda id: False]))
})

def insert_problem(problem):
    """
    Inserts a problem into the database. Does sane validation.

    Args:
        Problem dict.
        basescore: points awarded for completing the problem.
        category: problem's category
        desc: description of the problem.
        grader: path relative to grader_base_path
        threshold: Amount of points necessary for a team to unlock this problem.

        Optional:
        disabled: True or False. Defaults to False.
        hint: hint for completing the problem.
        pid: if you care to manually specify a pid
        tags: list of problem tags.
        relatedproblems: list of related problems.
        weightmap: problem's unlock weightmap
        autogen: Whether or not the problem will be auto generated.
    Returns:
        The newly created problem id.
        """

    db = api.common.get_conn()
    validate(problem_schema, problem)

    problem["disabled"] = problem.get("disabled", False)

    if problem.get("pid", None) is None:
        problem["pid"] = api.common.token()

    if len(search_problems({"pid": problem["pid"]}, {"displayname": problem["displayname"]})) > 0:
        raise APIException(0, None, "Problem with identical displayname or pid already exists.")

    db.problems.insert(problem)

    return problem["pid"]

def remove_problem(pid):
    """
    Removes a problem from the given database.

    Args:
        pid: the pid of the problem to remove.
    Returns:
        The removed problem object.
    """

    db = api.common.get_conn()
    problem = get_problem(pid=pid)

    db.problems.remove({"pid": pid})

    return problem

def set_problem_disabled(pid, disabled):
    """
    Updates a problem's availability.

    Args:
        pid: the problem's pid
        disabled: whether or not the problem should be disabled.
    Returns:
        The updated problem object.
    """
    return update_problem(pid, {"disabled": disabled})

def update_problem(pid, updated_problem):
    """
    Updates a problem with new properties.

    Args:
        pid: the pid of the problem to update.
        updated_problem: an updated problem object.
    Returns:
        The updated problem object.
    """

    db = api.common.get_conn()

    problem = get_problem(pid=pid)
    problem.update(updated_problem)

    validate(problem_schema, problem)

    db.problems.update({"pid": pid}, problem)

    return problem

def search_problems(*conditions):
    """
    Aggregates all problems that contain all of the given properties from the list specified.

    Args:
        conditions: multiple mongo queries to search.
    Returns:
        The list of matching problems.
    """

    db = api.common.get_conn()

    return list(db.problems.find({"$or": list(conditions)}))

def insert_problem_from_json(blob):
    """
    Converts json blob of problem(s) into dicts. Runs insert_problem on each one.
    See insert_problem for more information.

    Returns:
        A list of the created problem pids if an array of problems is specified.
    """

    result = json.loads(blob, default=json_util.default)

    if type(result) == list:
        return [insert_problem(problem) for problem in result]
    elif type(result) == dict:
        return insert_problem(result)
    else:
        raise APIException(0, "JSON blob does not appear to be a list of problems or a single problem.")

def grade_problem(uid=None, pid, key):
    """
    Grades the problem with its associated grader script.

    Args:
        uid: uid if provided
        pid: problem's pid
        key: user's submission
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    problem = api.problem.get_problem(pid)

    try:
        (correct, message) = imp.load_source(
            problem["grader"][:-3], join(grader_base_path, problem["grader"])
        ).grade(uid, key)
    except FileNotFoundError:
        raise APIException(0, None, "Problem {} grader is offline.".format(problem["pid"]))

    return {
        "correct": correct,
        "points": problem["basescore"],
        "message": message
    }

def submit_key(tid, pid, key, uid=None, ip=None):
    """
    User problem submission. Problem submission is inserted into the database.

    Args:
        tid: user's team id
        pid: problem's pid
        key: answer text
        uid: user's uid
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    db = api.common.get_conn()
    validate(submission_schema, {"tid": tid, "pid": pid, "key": key})

    if pid not in get_unlocked_pids(tid):
        raise APIException(0, None, "You can't submit flags to problems you haven't unlocked.")

    if pid in get_solved_pids(tid):
        raise APIException(0, None, "You have already solved this problem.")

    user = api.user.get_user(uid=uid)
    if user is None:
        raise APIException(0, None, "User submitting flag does not exist.")
    uid = user["uid"]

    result = grade_problem(uid, pid, key)

    problem = get_problem(pid=pid)

    submission = {
        'uid': uid,
        'tid': tid,
        'timestamp': datetime.now(),
        'pid': pid,
        'ip': ip,
        'key': key,
        'category': problem['category'],
        'correct': result['correct']
    }

    if key in [submission["key"] for submission in  get_submissions(tid=tid)]:
        raise APIException(0, None, "You or one of your teammates has already tried this solution.")

    db.submissions.insert(submission)

    return result

def get_submissions(pid=None, uid=None, tid=None, category=None):
    """
    Gets the submissions from a team or user.
    Optional filters of pid or category.

    Args:
        uid: the user id
        tid: the team id

        category: category filter.
        pid: problem filter.
    Returns:
        A list of submissions from the given entity
    """

    db = api.common.get_conn()

    match = {}

    if uid is not None:
      match.update({"uid": uid})
    elif tid is not None:
      match.update({"tid": tid})

    if pid is not None:
      match.update({"pid": pid})

    if category is not None:
      match.update({"category": category})

    return list(db.submissions.find(match))

def clear_all_submissions():
    """
    Removes all submissions from the database.
    """

    db = api.common.get_conn()
    db.submissions.remove()

def clear_submissions(uid=None, tid=None):
    """
    Clear submissions from a given team or user.

    Args:
        uid: the user's uid to clear from.
        tid: the team's tid to clear from.
    """

    db = api.common.get_conn()

    match = {}

    if uid is not None:
        match.update({"uid": uid})
    elif tid is not None:
        match.update({"tid": tid})
    else:
        raise APIException(0, None, "You must supply either a tid or uid")

    return db.submissions.remove(match)

def invalidate_submissions(pid=None, uid=None, tid=None):
    """
    Invalidates the submissions for a given problem. Can be filtered by uid or tid.
    Passing no arguments will invalidate all submissions.

    Args:
        pid: the pid of the problem.
        uid: the user's uid that will his submissions invalidated.
        tid: the team's tid that will have their submissions invalidated.
    """

    db = api.common.get_conn()

    match = {}

    if pid is not None:
        match.update({"pid": pid})

    if uid is not None:
        match.update({"uid": uid})
    elif tid is not None:
        match.update({"tid": tid})

    db.submissions.update(match, {"correct": False})

def reevaluate_submissions_for_problem(pid):
    """
    In the case of the problem or grader being updated, this will reevaluate all submissions.
    This will NOT work for auto generated problems.

    Args:
        pid: the pid of the problem to be reevaluated.
    Returns:
        A list of affected tids.
    """

    db = api.common.get_conn()

    problem = get_problem(pid=pid)

    keys = {}
    for submission in get_submissions(pid=pid):
        key = submission["key"]
        if key not in keys:
            result = grade_problem(pid, key)
            if result["correct"] != submission["correct"]:
                keys[key] = result["correct"]
            else:
                keys[key] = None

    for key, change in keys.items():
        if change is not None:
            db.submissions.update({"key": key}, {"correct": change}, multi=True)

def get_problem(pid=None, name=None, show_disabled=False):
    """
    Gets a single problem.

    Args:
        pid: The problem id
        name: The name of the problem
        show_disabled: Boolean indicating whether or not to show disabled problems.
    Returns:
        The problem dictionary from the database
    """

    db = api.common.get_conn()

    match = {"disabled": show_disabled}
    if pid is not None:
        match.update({'pid': pid})
    elif name is not None:
        match.update({'displayname': name})
    else:
        raise APIException(0, None, "Problem information not given")

    db = api.common.get_conn()
    problem = db.problems.find_one(match)

    if problem is None:
        raise APIException(0, None, "Could not find problem")

    return problem

def get_all_problems(category=None, show_disabled=False):
    """
    Gets all of the problems in the database.

    Args:
        category: Optional parameter to restrict which problems are returned
        show_disabled: Boolean indicating whether or not to show disabled problems.
    Returns:
        List of problems from the database
    """

    db = api.common.get_conn()

    match = {"disabled": show_disabled}
    if category is not None:
      match.update({'category': category})

    return list(db.problems.find(match))

def get_solved_pids(tid, category=None):
    """
    Gets the solved pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of solved problem ids
    """

    return [sub['pid'] for sub in get_submissions(tid=tid, category=category) if sub['correct'] == True]

def get_solved_problems(tid, category=None):
    """
    Gets the solved problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of solved problem dictionaries
    """

    return [get_problem(pid) for pid in get_solved_pids(tid, category)]

def get_unlocked_pids(tid, category=None):
    """
    Gets the unlocked pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of unlocked problem ids
    """

    solved = get_solved_problems(tid, category)

    unlocked = []
    for problem in get_all_problems():
        if 'weightmap' not in problem or 'threshold' not in problem:
            unlocked.append(problem['pid'])
        else:
            weightsum = sum(problem['weightmap'].get(pid, 0) for pid in get_solved_pids(tid, category))
            if weightsum >= problem['threshold']:
                unlocked.append(problem['pid'])

    return unlocked

def get_unlocked_problems(tid, category=None):
    """
    Gets the unlocked problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of unlocked problem dictionaries
    """

    return [get_problem(pid) for pid in get_unlocked_pids(tid, category)]
<commit_msg>Problem code is passing basic tests<commit_after>""" Module for interacting with the problems """
import imp
import json

import api.common
import api.user
import api.team

from datetime import datetime
from api.common import validate, APIException, check
from voluptuous import Schema, Length, Required, Range
from pymongo.errors import DuplicateKeyError
from bson import json_util
from os.path import join, isfile

grader_base_path = "./graders"
check_graders_exist = True

submission_schema = Schema({
    Required("tid"): check(
        (0, "This does not look like a valid tid.", [str, Length(max=100)])),
    Required("pid"): check(
        (0, "This does not look like a valid pid.", [str, Length(max=100)])),
    Required("key"): check(
        (0, "This does not look like a valid key.", [str, Length(max=100)]))
})

problem_schema = Schema({
    Required("basescore"): check(
        (0, "Basescore must be a positive integer.", [int, Range(min=0)])),
    Required("category"): check(
        (0, "Category must be a string.", [str])),
    Required("desc"): check(
        (0, "The problem description must be a string.", [str])),
    Required("displayname"): check(
        (0, "The problem's display name must be a string.", [str])),
    Required("grader"): check(
        (0, "A grader does not exist at that path.", [
            lambda grader: not check_graders_exist or
            isfile(join(grader_base_path, grader))])),
    Required("threshold"): check(
        (0, "Threshold must be a positive integer.", [int, Range(min=0)])),

    "disabled": check(
        (0, "A problem's disabled state is either True or False.", [
            lambda disabled: type(disabled) == bool])),
    "autogen": check(
        (0, "A problem should either be autogenerated or not, True/False", [
            lambda autogen: type(autogen) == bool])),
    "relatedproblems": check(
        (0, "Related problems should be a list of related problems.", [list])),
    "pid": check(
        (0, "The problem's pid should be a string.", [str])),
    "weightmap": check(
        (0, "Weightmap should be a dict.", [dict])),
    "tags": check(
        (0, "Tags must be described as a list.", [list])),
    "hint": check(
        (0, "A hint must be a string.", [str])),

    "_id": check(
        (0, "Your problems should not already have _ids.", [lambda id: False]))
})

def insert_problem(problem):
    """
    Inserts a problem into the database. Does sane validation.

    Args:
        Problem dict.
        basescore: points awarded for completing the problem.
        category: problem's category
        desc: description of the problem.
        grader: path relative to grader_base_path
        threshold: Amount of points necessary for a team to unlock this problem.

        Optional:
        disabled: True or False. Defaults to False.
        hint: hint for completing the problem.
        pid: if you care to manually specify a pid
        tags: list of problem tags.
        relatedproblems: list of related problems.
        weightmap: problem's unlock weightmap
        autogen: Whether or not the problem will be auto generated.
    Returns:
        The newly created problem id.
        """

    db = api.common.get_conn()
    validate(problem_schema, problem)

    problem["disabled"] = problem.get("disabled", False)

    if problem.get("pid", None) is None:
        problem["pid"] = api.common.token()

    if len(search_problems({"pid": problem["pid"]}, {"displayname": problem["displayname"]})) > 0:
        raise APIException(0, None, "Problem with identical displayname or pid already exists.")

    db.problems.insert(problem)

    return problem["pid"]

def remove_problem(pid):
    """
    Removes a problem from the given database.

    Args:
        pid: the pid of the problem to remove.
    Returns:
        The removed problem object.
    """

    db = api.common.get_conn()
    problem = get_problem(pid=pid)

    db.problems.remove({"pid": pid})

    return problem

def set_problem_disabled(pid, disabled):
    """
    Updates a problem's availability.

    Args:
        pid: the problem's pid
        disabled: whether or not the problem should be disabled.
    Returns:
        The updated problem object.
    """
    return update_problem(pid, {"disabled": disabled})

def update_problem(pid, updated_problem):
    """
    Updates a problem with new properties.

    Args:
        pid: the pid of the problem to update.
        updated_problem: an updated problem object.
    Returns:
        The updated problem object.
    """

    db = api.common.get_conn()

    problem = get_problem(pid=pid)
    problem.update(updated_problem)

    validate(problem_schema, problem)

    db.problems.update({"pid": pid}, problem)

    return problem

def search_problems(*conditions):
    """
    Aggregates all problems that contain all of the given properties from the list specified.

    Args:
        conditions: multiple mongo queries to search.
    Returns:
        The list of matching problems.
    """

    db = api.common.get_conn()

    return list(db.problems.find({"$or": list(conditions)}))

def insert_problem_from_json(blob):
    """
    Converts json blob of problem(s) into dicts. Runs insert_problem on each one.
    See insert_problem for more information.

    Returns:
        A list of the created problem pids if an array of problems is specified.
    """

    result = json.loads(blob, default=json_util.default)

    if type(result) == list:
        return [insert_problem(problem) for problem in result]
    elif type(result) == dict:
        return insert_problem(result)
    else:
        raise APIException(0, "JSON blob does not appear to be a list of problems or a single problem.")

def grade_problem(pid, key, uid=None):
    """
    Grades the problem with its associated grader script.

    Args:
        uid: uid if provided
        pid: problem's pid
        key: user's submission
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    problem = get_problem(pid=pid)

    try:
        (correct, message) = imp.load_source(
            problem["grader"][:-3], join(grader_base_path, problem["grader"])
        ).grade(uid, key)
    except FileNotFoundError:
        raise APIException(0, None, "Problem {} grader is offline.".format(problem["pid"]))

    return {
        "correct": correct,
        "points": problem["basescore"],
        "message": message
    }

def submit_key(tid, pid, key, uid=None, ip=None):
    """
    User problem submission. Problem submission is inserted into the database.

    Args:
        tid: user's team id
        pid: problem's pid
        key: answer text
        uid: user's uid
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    db = api.common.get_conn()
    validate(submission_schema, {"tid": tid, "pid": pid, "key": key})

    if pid not in get_unlocked_pids(tid):
        raise APIException(0, None, "You can't submit flags to problems you haven't unlocked.")

    if pid in get_solved_pids(tid):
        raise APIException(0, None, "You have already solved this problem.")

    user = api.user.get_user(uid=uid)
    if user is None:
        raise APIException(0, None, "User submitting flag does not exist.")
    uid = user["uid"]

    result = grade_problem(pid, key, uid)

    problem = get_problem(pid=pid)

    submission = {
        'uid': uid,
        'tid': tid,
        'timestamp': datetime.now(),
        'pid': pid,
        'ip': ip,
        'key': key,
        'category': problem['category'],
        'correct': result['correct']
    }

    if key in [submission["key"] for submission in  get_submissions(tid=tid)]:
        raise APIException(0, None, "You or one of your teammates has already tried this solution.")

    db.submissions.insert(submission)

    return result

def get_submissions(pid=None, uid=None, tid=None, category=None):
    """
    Gets the submissions from a team or user.
    Optional filters of pid or category.

    Args:
        uid: the user id
        tid: the team id

        category: category filter.
        pid: problem filter.
    Returns:
        A list of submissions from the given entity
    """

    db = api.common.get_conn()

    match = {}

    if uid is not None:
      match.update({"uid": uid})
    elif tid is not None:
      match.update({"tid": tid})

    if pid is not None:
      match.update({"pid": pid})

    if category is not None:
      match.update({"category": category})

    return list(db.submissions.find(match))

def clear_all_submissions():
    """
    Removes all submissions from the database.
    """

    db = api.common.get_conn()
    db.submissions.remove()

def clear_submissions(uid=None, tid=None):
    """
    Clear submissions from a given team or user.

    Args:
        uid: the user's uid to clear from.
        tid: the team's tid to clear from.
    """

    db = api.common.get_conn()

    match = {}

    if uid is not None:
        match.update({"uid": uid})
    elif tid is not None:
        match.update({"tid": tid})
    else:
        raise APIException(0, None, "You must supply either a tid or uid")

    return db.submissions.remove(match)

def invalidate_submissions(pid=None, uid=None, tid=None):
    """
    Invalidates the submissions for a given problem. Can be filtered by uid or tid.
    Passing no arguments will invalidate all submissions.

    Args:
        pid: the pid of the problem.
        uid: the user's uid that will his submissions invalidated.
        tid: the team's tid that will have their submissions invalidated.
    """

    db = api.common.get_conn()

    match = {}

    if pid is not None:
        match.update({"pid": pid})

    if uid is not None:
        match.update({"uid": uid})
    elif tid is not None:
        match.update({"tid": tid})

    db.submissions.update(match, {"correct": False})

def reevaluate_submissions_for_problem(pid):
    """
    In the case of the problem or grader being updated, this will reevaluate all submissions.
    This will NOT work for auto generated problems.

    Args:
        pid: the pid of the problem to be reevaluated.
    Returns:
        A list of affected tids.
    """

    db = api.common.get_conn()

    problem = get_problem(pid=pid)

    keys = {}
    for submission in get_submissions(pid=pid):
        key = submission["key"]
        if key not in keys:
            result = grade_problem(pid, key)
            if result["correct"] != submission["correct"]:
                keys[key] = result["correct"]
            else:
                keys[key] = None

    for key, change in keys.items():
        if change is not None:
            db.submissions.update({"key": key}, {"correct": change}, multi=True)

def get_problem(pid=None, name=None, show_disabled=False):
    """
    Gets a single problem.

    Args:
        pid: The problem id
        name: The name of the problem
        show_disabled: Boolean indicating whether or not to show disabled problems.
    Returns:
        The problem dictionary from the database
    """

    db = api.common.get_conn()

    match = {"disabled": show_disabled}
    if pid is not None:
        match.update({'pid': pid})
    elif name is not None:
        match.update({'displayname': name})
    else:
        raise APIException(0, None, "Problem information not given")

    db = api.common.get_conn()
    problem = db.problems.find_one(match)

    if problem is None:
        raise APIException(0, None, "Could not find problem")

    return problem

def get_all_problems(category=None, show_disabled=False):
    """
    Gets all of the problems in the database.

    Args:
        category: Optional parameter to restrict which problems are returned
        show_disabled: Boolean indicating whether or not to show disabled problems.
    Returns:
        List of problems from the database
    """

    db = api.common.get_conn()

    match = {"disabled": show_disabled}
    if category is not None:
      match.update({'category': category})

    return list(db.problems.find(match))

def get_solved_pids(tid, category=None):
    """
    Gets the solved pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of solved problem ids
    """

    return [sub['pid'] for sub in get_submissions(tid=tid, category=category) if sub['correct'] == True]

def get_solved_problems(tid, category=None):
    """
    Gets the solved problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of solved problem dictionaries
    """

    return [get_problem(pid) for pid in get_solved_pids(tid, category)]

def get_unlocked_pids(tid, category=None):
    """
    Gets the unlocked pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of unlocked problem ids
    """

    solved = get_solved_problems(tid, category)

    unlocked = []
    for problem in get_all_problems():
        if 'weightmap' not in problem or 'threshold' not in problem:
            unlocked.append(problem['pid'])
        else:
            weightsum = sum(problem['weightmap'].get(pid, 0) for pid in get_solved_pids(tid, category))
            if weightsum >= problem['threshold']:
                unlocked.append(problem['pid'])

    return unlocked

def get_unlocked_problems(tid, category=None):
    """
    Gets the unlocked problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of unlocked problem dictionaries
    """

    return [get_problem(pid) for pid in get_unlocked_pids(tid, category)]
<|endoftext|>


<commit_before>""" Module for interacting with the problems """
import imp

import api.common
import api.user
import api.team
from datetime import datetime
from api.common import validate, APIException, check
from voluptuous import Schema, Length, Required

submission_schema = Schema({
    Required("tid"): check(
        (0, "This does not look like a valid tid.", [str, Length(max=100)])),
    Required("pid"): check(
        (0, "This does not look like a valid pid.", [str, Length(max=100)])),
    Required("key"): check(
        (0, "This does not look like a valid key.", [str, Length(max=100)]))
})

def grade_problem(uid, pid, key):
    """
    Grades the problem with its associated grader script.

    Args:
        uid: user's uid
        pid: problem's pid
        key: user's submission
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    problem = api.problem.get_problem(pid)

    try:
        (correct, message) = imp.load_source(
            problem["grader"][:-3], "./graders/" + problem["grader"]
        ).grade(uid, key)
    except FileNotFoundError:
       raise APIException(0, None, "Problem {} grader is offline.".format(problem["pid"]))

    return {
        "correct": correct,
        "points": problem["basescore"],
        "message": message
    }

def submit_problem(tid, pid, key, uid=None, ip=None):
    """
    User problem submission. Problem submission is inserted into the database.

    Args:
        tid: user's team id
        pid: problem's pid
        key: answer text
        uid: user's uid
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    db = api.common.get_conn()
    validate(submission_schema, {"tid": tid, "pid": pid, "key": key})

    if pid not in get_unlocked_pids(tid):
        raise APIException(0, None, "You can submit flags to problems you haven't unlocked.")

    if pid in get_solved_pids(tid):
        raise APIException(0, None, "You have already solved this problem.")

    user = api.user.get_user(uid=uid)
    if user is None:
        raise APIException(0, None, "User submitting flag does not exist.")
    uid = user["uid"]

    result = grade_problem(uid, pid, key)

    submission = {
        'uid': uid,
        'tid': tid,
        'timestamp': datetime.now(),
        'pid': pid,
        'ip': ip,
        'key': key,
        'correct': result["correct"]
    }

    try:
        db.submissions.insert(submission)
    except DuplicateKeyError:
        raise APIException(0, None, "You or one of your teammates has already tried this solution.")

    return result

def get_submissions(uid):
    """
    Gets the submissions from a given user.

    Args:
        uid: the user id

    Returns:
        A list of submissions from the given user
    """

    db = api.common.get_conn()
    return db.submissions.find({'uid': uid})

def get_team_submissions(tid):
    """
    Gets the submissions from the entire team

    Args:
        tid: the team id

    Returns:
        A list of submissions from the entire team
    """

    db = api.common.get_conn()
    return db.submissions.find({'tid': tid})

def get_problem(pid=None, name=None):
    """
    Gets a single problem.

    Args:
        pid: The problem id
        name: The name of the problem

    Returns:
        The problem dictionary from the database
    """

    if pid is not None:
        match = {'pid': pid}
    elif name is not none:
        match = {'displayname': name}
    else:
        raise APIException(0, None, "Problem information not given")

    problem = db.problems.find_one(match)

    if problem is None:
        raise APIException(0, None, "Could not find problem")

    return problem

def get_all_problems(category=None):
    """
    Gets all of the problems in the database.

    Args:
        category: Optional parameter to restrict which problems are returned
    """

    db = api.common.get_conn()

    if category:
        return db.problems.find({'category': category})
    return db.problems.find()

def get_solved_pids(tid, category=None):
    """
    Gets the solved pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned

    Returns:
        List of solved problem ids
    """

    correct_pids = [sub['pid'] for sub in get_team_submission(tid) if sub['correct'] == True]

    solved = []
    for pid in correct_pids:
        solved += pid

    return solved

def get_solved_problems(tid, category=None):
    """
    Gets the solved problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned

    Returns:
        List of solved problem dictionaries
    """

    return [get_problem(pid) for pid in get_solved_pids(tid, category)]

def get_unlocked_pids(tid, category=None):
    """
    Gets the unlocked pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned

    Returns:
        List of unlocked problem ids
    """

    solved = get_solved_problems(tid, category)

    unlocked = []
    for problem in get_all_problems():
        if 'weightmap' not in problem or 'threshold' not in problem:
            unlocked.append(problem['pid'])
        else:
            weightsum = sum(problem['weightmap'].get(pid, 0) for pid in get_solved_pids())
            if weightsum >= problem['threshold']:
                unlocked.append(problem['pid'])

    return unlocked

def get_unlocked_problems(tid, category=None):
    """
    Gets the unlocked problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned

    Returns:
        List of unlocked problem dictionaries
    """

    return [get_problem(pid) for pid in get_unlocked_pids(tid, category)]
<commit_msg>Small fixes; docstring formatting<commit_after>""" Module for interacting with the problems """
import imp

import api.common
import api.user
import api.team
from datetime import datetime
from api.common import validate, APIException, check
from voluptuous import Schema, Length, Required

submission_schema = Schema({
    Required("tid"): check(
        (0, "This does not look like a valid tid.", [str, Length(max=100)])),
    Required("pid"): check(
        (0, "This does not look like a valid pid.", [str, Length(max=100)])),
    Required("key"): check(
        (0, "This does not look like a valid key.", [str, Length(max=100)]))
})

def grade_problem(uid, pid, key):
    """
    Grades the problem with its associated grader script.

    Args:
        uid: user's uid
        pid: problem's pid
        key: user's submission
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    problem = api.problem.get_problem(pid)

    try:
        (correct, message) = imp.load_source(
            problem["grader"][:-3], "./graders/" + problem["grader"]
        ).grade(uid, key)
    except FileNotFoundError:
       raise APIException(0, None, "Problem {} grader is offline.".format(problem["pid"]))

    return {
        "correct": correct,
        "points": problem["basescore"],
        "message": message
    }

def submit_problem(tid, pid, key, uid=None, ip=None):
    """
    User problem submission. Problem submission is inserted into the database.

    Args:
        tid: user's team id
        pid: problem's pid
        key: answer text
        uid: user's uid
    Returns:
        A dict.
        correct: boolean
        points: number of points the problem is worth.
        message: message returned from the grader.
    """

    db = api.common.get_conn()
    validate(submission_schema, {"tid": tid, "pid": pid, "key": key})

    if pid not in get_unlocked_pids(tid):
        raise APIException(0, None, "You can submit flags to problems you haven't unlocked.")

    if pid in get_solved_pids(tid):
        raise APIException(0, None, "You have already solved this problem.")

    user = api.user.get_user(uid=uid)
    if user is None:
        raise APIException(0, None, "User submitting flag does not exist.")
    uid = user["uid"]

    result = grade_problem(uid, pid, key)

    submission = {
        'uid': uid,
        'tid': tid,
        'timestamp': datetime.now(),
        'pid': pid,
        'ip': ip,
        'key': key,
        'correct': result["correct"]
    }

    try:
        db.submissions.insert(submission)
    except DuplicateKeyError:
        raise APIException(0, None, "You or one of your teammates has already tried this solution.")

    return result

def get_submissions(uid):
    """
    Gets the submissions from a given user.

    Args:
        uid: the user id
    Returns:
        A list of submissions from the given user
    """

    db = api.common.get_conn()
    return db.submissions.find({'uid': uid})

def get_team_submissions(tid):
    """
    Gets the submissions from the entire team

    Args:
        tid: the team id
    Returns:
        A list of submissions from the entire team
    """

    db = api.common.get_conn()
    return db.submissions.find({'tid': tid})

def get_problem(pid=None, name=None):
    """
    Gets a single problem.

    Args:
        pid: The problem id
        name: The name of the problem
    Returns:
        The problem dictionary from the database
    """

    if pid is not None:
        match = {'pid': pid}
    elif name is not none:
        match = {'displayname': name}
    else:
        raise APIException(0, None, "Problem information not given")

    db = api.common.get_conn()
    problem = db.problems.find_one(match)

    if problem is None:
        raise APIException(0, None, "Could not find problem")

    return problem

def get_all_problems(category=None):
    """
    Gets all of the problems in the database.

    Args:
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of problems from the database
    """

    db = api.common.get_conn()

    if category:
        return list(db.problems.find({'category': category}))
    return list(db.problems.find())

def get_solved_pids(tid, category=None):
    """
    Gets the solved pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of solved problem ids
    """

    correct_pids = [sub['pid'] for sub in get_team_submission(tid) if sub['correct'] == True]

    solved = []
    for pid in correct_pids:
        solved += pid

    return solved

def get_solved_problems(tid, category=None):
    """
    Gets the solved problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of solved problem dictionaries
    """

    return [get_problem(pid) for pid in get_solved_pids(tid, category)]

def get_unlocked_pids(tid, category=None):
    """
    Gets the unlocked pids for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of unlocked problem ids
    """

    solved = get_solved_problems(tid, category)

    unlocked = []
    for problem in get_all_problems():
        if 'weightmap' not in problem or 'threshold' not in problem:
            unlocked.append(problem['pid'])
        else:
            weightsum = sum(problem['weightmap'].get(pid, 0) for pid in get_solved_pids())
            if weightsum >= problem['threshold']:
                unlocked.append(problem['pid'])

    return unlocked

def get_unlocked_problems(tid, category=None):
    """
    Gets the unlocked problems for a given team.

    Args:
        tid: The team id
        category: Optional parameter to restrict which problems are returned
    Returns:
        List of unlocked problem dictionaries
    """

    return [get_problem(pid) for pid in get_unlocked_pids(tid, category)]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# define some variables
from __future__ import print_function
from builtins import str
import numpy as np
import sys
import matplotlib
if matplotlib.get_backend() != "TKAgg":
  matplotlib.use("TKAgg")
import pylab as plt
try:
  from mpl_toolkits.basemap import Basemap
except ImportError:
  Basemap = None
from pylab import meshgrid

import pmagpy.pmag as pmag
import pmagpy.pmagplotlib as pmagplotlib
from matplotlib import cm
def main():
    """
    NAME
        plot_magmap.py

    DESCRIPTION
        makes a color contour map of desired field model

    SYNTAX
        plot_magmap.py [command line options]

    OPTIONS
        -h prints help and quits
        -f FILE  specify field model file with format:  l m g h
        -fmt [pdf,eps,svg,png]  specify format for output figure  (default is png)
        -mod [arch3k,cals3k,pfm9k,hfm10k,cals10k_2,shadif14k,cals10k] specify model for 3ka to 1900 CE, default is  cals10k
        -alt ALT;  specify altitude in km, default is sealevel (0)
        -age specify date in decimal year, default is 2015
        -lon0: 0 longitude for map, default is 0
        -el: [D,I,B,Br]  specify element for plotting

    """
    cmap='RdYlBu'
    date=2016.
    if not Basemap:
      print("-W- You must intstall the Basemap module to run plot_magmap.py")
      sys.exit()
    dir_path='.'
    lincr=1 # level increment for contours
    if '-WD' in sys.argv:
        ind = sys.argv.index('-WD')
        dir_path=sys.argv[ind+1]
    if '-h' in sys.argv:
        print(main.__doc__)
        sys.exit()
    if '-fmt' in sys.argv:
        ind = sys.argv.index('-fmt')
        fmt=sys.argv[ind+1]
        if fmt=='jpg':
            print ('jpg not a supported option')
            print(main.__doc__)
            sys.exit()
    else: fmt='png'
    if '-el' in sys.argv:
        ind = sys.argv.index('-el')
        el=sys.argv[ind+1]
    else:
        el='B'
    if '-alt' in sys.argv:
        ind = sys.argv.index('-alt')
        alt=sys.argv[ind+1]
    else: alt=0
    if '-lon0' in sys.argv:
        ind=sys.argv.index('-lon0')
        lon_0=float(sys.argv[ind+1])
    else: lon_0=0
    if '-mod' in sys.argv:
        ind=sys.argv.index('-mod')
        mod=sys.argv[ind+1]
    elif '-f' in sys.argv:
        ind=sys.argv.index('-f')
        ghfile=sys.argv[ind+1]
        mod='custom'
        date=''
    else: mod,ghfile='cals10k',''
    if '-age' in sys.argv:
        ind=sys.argv.index('-age')
        date=float(sys.argv[ind+1])
    if '-alt' in sys.argv:
        ind=sys.argv.index('-alt')
        alt=float(sys.argv[ind+1])
    else: alt=0
    Ds,Is,Bs,Brs,lons,lats=pmag.do_mag_map(date,mod=mod,lon_0=lon_0,el=el,alt=alt,file=ghfile)
    m = Basemap(projection='hammer',lon_0=lon_0)
    x,y=m(*meshgrid(lons,lats))
    m.drawcoastlines()
    if mod=='custom':
        d='Custom'
    else: d=str(date)
    if el=='B':
        levmax=Bs.max()+lincr
        levmin=round(Bs.min()-lincr)
        levels=np.arange(levmin,levmax,lincr)
        cs=m.contourf(x,y,Bs,levels=levels,cmap=cmap)
        plt.title('Field strength ($\mu$T): '+d);
    if el=='Brs':
        levmax=Brs.max()+lincr
        levmin=round(Brs.min()-lincr)
        cs=m.contourf(x,y,Brs,levels=np.arange(levmin,levmax,lincr),cmap=cmap)
        plt.title('Radial field strength ($\mu$T): '+str(date));
    if el=='I':
        levmax=Is.max()+lincr
        levmin=round(Is.min()-lincr)
        cs=m.contourf(x,y,Is,levels=np.arange(levmin,levmax,lincr),cmap=cmap)
        m.contour(x,y,Is,levels=np.arange(-80,90,10),colors='black')
        plt.title('Field inclination: '+str(date));
    if el=='D':
        #cs=m.contourf(x,y,Ds,levels=np.arange(-180,180,10),cmap=cmap)
        cs=m.contourf(x,y,Ds,levels=np.arange(-180,180,10),cmap=cmap)
        m.contour(x,y,Ds,levels=np.arange(-180,180,10),colors='black')
        plt.title('Field declination: '+str(date));
    cbar=m.colorbar(cs,location='bottom')
    plt.savefig('igrf'+d+'.'+fmt)
    print('Figure saved as: ','igrf'+d+'.'+fmt)

if __name__ == "__main__":
    main()
<commit_msg>	modified:   plot_magmap.py => added color map options<commit_after>#!/usr/bin/env python
# define some variables
from __future__ import print_function
from builtins import str
import numpy as np
import sys
import matplotlib
if matplotlib.get_backend() != "TKAgg":
  matplotlib.use("TKAgg")
import pylab as plt
try:
  from mpl_toolkits.basemap import Basemap
except ImportError:
  Basemap = None
from pylab import meshgrid

import pmagpy.pmag as pmag
import pmagpy.pmagplotlib as pmagplotlib
from matplotlib import cm
import warnings
warnings.filterwarnings("ignore")

def main():
    """
    NAME
        plot_magmap.py

    DESCRIPTION
        makes a color contour map of desired field model

    SYNTAX
        plot_magmap.py [command line options]

    OPTIONS
        -h prints help and quits
        -f FILE  specify field model file with format:  l m g h
        -fmt [pdf,eps,svg,png]  specify format for output figure  (default is png)
        -mod [arch3k,cals3k,pfm9k,hfm10k,cals10k_2,shadif14k,cals10k] specify model for 3ka to 1900 CE, default is  cals10k
        -alt ALT;  specify altitude in km, default is sealevel (0)
        -age specify date in decimal year, default is 2015
        -lon0: 0 longitude for map, default is 0
        -el: [D,I,B,Br]  specify element for plotting
        -cm: [see https://matplotlib.org/users/colormaps.html] specify color map for plotting (default is RdYlBu)

    """
    cmap='RdYlBu'
    date=2016.
    if not Basemap:
      print("-W- You must intstall the Basemap module to run plot_magmap.py")
      sys.exit()
    dir_path='.'
    lincr=1 # level increment for contours
    if '-WD' in sys.argv:
        ind = sys.argv.index('-WD')
        dir_path=sys.argv[ind+1]
    if '-h' in sys.argv:
        print(main.__doc__)
        sys.exit()
    if '-fmt' in sys.argv:
        ind = sys.argv.index('-fmt')
        fmt=sys.argv[ind+1]
        if fmt=='jpg':
            print ('jpg not a supported option')
            print(main.__doc__)
            sys.exit()
    else: fmt='png'
    if '-cm' in sys.argv:
        ind=sys.argv.index('-cm')
        cmap=sys.argv[ind+1]
    if '-el' in sys.argv:
        ind = sys.argv.index('-el')
        el=sys.argv[ind+1]
    else:
        el='B'
    if '-alt' in sys.argv:
        ind = sys.argv.index('-alt')
        alt=sys.argv[ind+1]
    else: alt=0
    if '-lon0' in sys.argv:
        ind=sys.argv.index('-lon0')
        lon_0=float(sys.argv[ind+1])
    else: lon_0=0
    if '-mod' in sys.argv:
        ind=sys.argv.index('-mod')
        mod=sys.argv[ind+1]
    elif '-f' in sys.argv:
        ind=sys.argv.index('-f')
        ghfile=sys.argv[ind+1]
        mod='custom'
        date=''
    else: mod,ghfile='cals10k',''
    if '-age' in sys.argv:
        ind=sys.argv.index('-age')
        date=float(sys.argv[ind+1])
    if '-alt' in sys.argv:
        ind=sys.argv.index('-alt')
        alt=float(sys.argv[ind+1])
    else: alt=0
    Ds,Is,Bs,Brs,lons,lats=pmag.do_mag_map(date,mod=mod,lon_0=lon_0,el=el,alt=alt,file=ghfile)
    m = Basemap(projection='hammer',lon_0=lon_0)
    x,y=m(*meshgrid(lons,lats))
    m.drawcoastlines()
    if mod=='custom':
        d='Custom'
    else: d=str(date)
    if el=='B':
        levmax=Bs.max()+lincr
        levmin=round(Bs.min()-lincr)
        levels=np.arange(levmin,levmax,lincr)
        cs=m.contourf(x,y,Bs,levels=levels,cmap=cmap)
        plt.title('Field strength ($\mu$T): '+d);
    if el=='Brs':
        levmax=Brs.max()+lincr
        levmin=round(Brs.min()-lincr)
        cs=m.contourf(x,y,Brs,levels=np.arange(levmin,levmax,lincr),cmap=cmap)
        plt.title('Radial field strength ($\mu$T): '+str(date));
    if el=='I':
        levmax=Is.max()+lincr
        levmin=round(Is.min()-lincr)
        cs=m.contourf(x,y,Is,levels=np.arange(levmin,levmax,lincr),cmap=cmap)
        m.contour(x,y,Is,levels=np.arange(-80,90,10),colors='black')
        plt.title('Field inclination: '+str(date));
    if el=='D':
        #cs=m.contourf(x,y,Ds,levels=np.arange(-180,180,10),cmap=cmap)
        cs=m.contourf(x,y,Ds,levels=np.arange(-180,180,10),cmap=cmap)
        m.contour(x,y,Ds,levels=np.arange(-180,180,10),colors='black')
        plt.title('Field declination: '+str(date));
    cbar=m.colorbar(cs,location='bottom')
    plt.savefig('igrf'+d+'.'+fmt)
    print('Figure saved as: ','igrf'+d+'.'+fmt)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>"""
Autogenerate certificate-based ssh keys for minions.

Keyword arguments comes from the saltmaster config where the module is
registered, use that to specify root keys and how to partition the certificates.
You probably want to have different certificates for each environment if you
have multiple environments on the same saltmaster, to prevent a compromised host
in one environment from being able to impersonate a host in another.

Installing and configuring the extension:

    $ sudo mkdir -p /var/lib/salt-ssh-keys
    $ sudo chown root:saltmaster /var/lib/salt-ssh-keys
    $ sudo chmod 750 /var/lib/salt-ssh-keys
    $ sudo ssh-keygen -t ed25519 -N '' -f /var/lib/salt-ssh-keys/root

Put this file in your saltmaster's `extension_modules` directory, and add the
following to /etc/salt/master:

ext_pillar:
    - ssh_keys:
        root_keys:
            - path: /var/lib/salt-ssh-keys/root
              minion_globs:
                - "*"

Modify the root keys to your liking, eg by compartmentalizing by environment or
similar.

This extension works in conjunction with the openssh-server state to provision
the keys.
"""

import fnmatch
import os
import re
import subprocess
import tempfile
from datetime import datetime, timedelta
from logging import getLogger

_logger = getLogger(__name__)
VALID_RE = re.compile(r'Valid: .* ([0-9T:-]+)$')

def ext_pillar(
        minion_id,
        pillar,
        root_keys=None,
        validity=timedelta(days=3),
        key_types=('ed25519',),
        keystore='/var/lib/salt-ssh-keys',
        **kwargs):
    if not root_keys:
        _logger.warning('The ssh_keys pillar extension is not configured with any root certificates')
        return {}

    root_key_path = get_root_key_path(minion_id, root_keys)
    if not root_key_path:
        _logger.warning('No root key for minion %s', minion_id)
        return {}

    ret = {}

    for key_type in key_types:
        key_path, cert_path = get_cert_for_minion(
            minion_id,
            root_key_path,
            key_type,
            keystore,
            validity,
        )

        with open(key_path, 'rb') as fh:
            ret['host_%s_key' % key_type] = fh.read()
        with open(cert_path, 'rb') as fh:
            ret['host_%s_certificate' % key_type] = fh.read()

    return {
        'openssh_server': ret,
    }

def get_cert_for_minion(minion_id, root_key_path, key_type, keystore, validity):
    key_path = os.path.join(keystore, '%s-%s' % (minion_id, key_type))
    cert_path = os.path.join(keystore, '%s-%s-cert.pub' % (minion_id, key_type))
    existing_cert_expiry = get_cert_expiry(cert_path, minion_id)
    if existing_cert_expiry and existing_cert_expiry - datetime.utcnow() > validity/3:
        return key_path, cert_path

    # Re-create cert when less than a third of the lifetime left
    temp_key_path = generate_ssh_key(key_type)
    temp_cert_path = sign_ssh_key(temp_key_path, root_key_path, minion_id, validity)
    os.rename(temp_key_path, key_path)
    os.rename(temp_cert_path, cert_path)

    # Remove the unused pubkey
    os.remove(temp_key_path + '.pub')

    return key_path, cert_path

def get_root_key_path(minion_id, root_keys):
    for potential_key in root_keys:
        globs = potential_key.get('minion_globs')
        if not globs:
            return potential_key['path']

        for glob in globs:
            if fnmatch.fnmatch(minion_id, glob):
                return potential_key['path']

    _logger.warning('No root key found matching minion %s', minion_id)
    return None

def generate_ssh_key(key_type):
    key_path = tempfile.mktemp()
    subprocess.check_call([
        'ssh-keygen',
        '-t', key_type,
        '-f', key_path,
        '-N', '', # no passphrase
        '-q',
    ])
    return key_path

def sign_ssh_key(key_path, root_key_path, name, validity):
    subprocess.check_call([
        'ssh-keygen',
        '-s', root_key_path,
        '-h',
        '-I', name,
        '-q',
        '-V', get_ssh_validity(validity),
        key_path,
    ])

    # ssh-keygen doesn't enable specifying an output filename, but naming
    # follows a consistent template
    return key_path + '-cert.pub'

def get_ssh_validity(validity):
    now = datetime.utcnow()
    validity_end = now + validity
    ssh_time_format = '%Y%m%d%H%M%S'
    # The docs says it should be possible to only specify the end time, but
    # that fails for me
    return '%s:%s' % (now.strftime(ssh_time_format), validity_end.strftime(ssh_time_format))

def get_cert_expiry(cert_path, minion_id):
    try:
        output = subprocess.check_output([
            'ssh-keygen',
            '-L',
            '-f', cert_path,
        ], stderr=subprocess.PIPE)
    except Exception as e:
        return None

    for line in output.split('\n'):
        match = VALID_RE.search(line)
        if not match:
            continue

        formatted_time = match.group(1)
        return datetime.strptime(formatted_time, '%Y-%m-%dT%H:%M:%S')

    # TODO: Ignore and create a new one?
    raise ValueError('Unparseable certificate: %s' % output)

<commit_msg>Specify validity in seconds<commit_after>"""
Autogenerate certificate-based ssh keys for minions.

Keyword arguments comes from the saltmaster config where the module is
registered, use that to specify root keys and how to partition the certificates.
You probably want to have different certificates for each environment if you
have multiple environments on the same saltmaster, to prevent a compromised host
in one environment from being able to impersonate a host in another.

Installing and configuring the extension:

    $ sudo mkdir -p /var/lib/salt-ssh-keys
    $ sudo chown root:saltmaster /var/lib/salt-ssh-keys
    $ sudo chmod 750 /var/lib/salt-ssh-keys
    $ sudo ssh-keygen -t ed25519 -N '' -f /var/lib/salt-ssh-keys/root

Put this file in your saltmaster's `extension_modules` directory, and add the
following to /etc/salt/master:

ext_pillar:
    - ssh_keys:
        root_keys:
            - path: /var/lib/salt-ssh-keys/root
              minion_globs:
                - "*"

Modify the root keys to your liking, eg by compartmentalizing by environment or
similar.

This extension works in conjunction with the openssh-server state to provision
the keys.
"""

import fnmatch
import os
import re
import subprocess
import tempfile
from datetime import datetime, timedelta
from logging import getLogger

_logger = getLogger(__name__)
VALID_RE = re.compile(r'Valid: .* ([0-9T:-]+)$')

def ext_pillar(
        minion_id,
        pillar,
        root_keys=None,
        validity_seconds=259200, # 3 days
        key_types=('ed25519',),
        keystore='/var/lib/salt-ssh-keys',
        **kwargs):
    if not root_keys:
        _logger.warning('The ssh_keys pillar extension is not configured with any root certificates')
        return {}

    root_key_path = get_root_key_path(minion_id, root_keys)
    if not root_key_path:
        _logger.warning('No root key for minion %s', minion_id)
        return {}

    validity = timedelta(seconds=validity_seconds)
    ret = {}

    for key_type in key_types:
        key_path, cert_path = get_cert_for_minion(
            minion_id,
            root_key_path,
            key_type,
            keystore,
            validity,
        )

        with open(key_path, 'rb') as fh:
            ret['host_%s_key' % key_type] = fh.read()
        with open(cert_path, 'rb') as fh:
            ret['host_%s_certificate' % key_type] = fh.read()

    return {
        'openssh_server': ret,
    }

def get_cert_for_minion(minion_id, root_key_path, key_type, keystore, validity):
    key_path = os.path.join(keystore, '%s-%s' % (minion_id, key_type))
    cert_path = os.path.join(keystore, '%s-%s-cert.pub' % (minion_id, key_type))
    existing_cert_expiry = get_cert_expiry(cert_path, minion_id)
    if existing_cert_expiry and existing_cert_expiry - datetime.utcnow() > validity/3:
        return key_path, cert_path

    # Re-create cert when less than a third of the lifetime left
    temp_key_path = generate_ssh_key(key_type)
    temp_cert_path = sign_ssh_key(temp_key_path, root_key_path, minion_id, validity)
    os.rename(temp_key_path, key_path)
    os.rename(temp_cert_path, cert_path)

    # Remove the unused pubkey
    os.remove(temp_key_path + '.pub')

    return key_path, cert_path

def get_root_key_path(minion_id, root_keys):
    for potential_key in root_keys:
        globs = potential_key.get('minion_globs')
        if not globs:
            return potential_key['path']

        for glob in globs:
            if fnmatch.fnmatch(minion_id, glob):
                return potential_key['path']

    _logger.warning('No root key found matching minion %s', minion_id)
    return None

def generate_ssh_key(key_type):
    key_path = tempfile.mktemp()
    subprocess.check_call([
        'ssh-keygen',
        '-t', key_type,
        '-f', key_path,
        '-N', '', # no passphrase
        '-q',
    ])
    return key_path

def sign_ssh_key(key_path, root_key_path, name, validity):
    subprocess.check_call([
        'ssh-keygen',
        '-s', root_key_path,
        '-h',
        '-I', name,
        '-q',
        '-V', get_ssh_validity(validity),
        key_path,
    ])

    # ssh-keygen doesn't enable specifying an output filename, but naming
    # follows a consistent template
    return key_path + '-cert.pub'

def get_ssh_validity(validity):
    now = datetime.utcnow()
    validity_end = now + validity
    ssh_time_format = '%Y%m%d%H%M%S'
    # The docs says it should be possible to only specify the end time, but
    # that fails for me
    return '%s:%s' % (now.strftime(ssh_time_format), validity_end.strftime(ssh_time_format))

def get_cert_expiry(cert_path, minion_id):
    try:
        output = subprocess.check_output([
            'ssh-keygen',
            '-L',
            '-f', cert_path,
        ], stderr=subprocess.PIPE)
    except Exception as e:
        return None

    for line in output.split('\n'):
        match = VALID_RE.search(line)
        if not match:
            continue

        formatted_time = match.group(1)
        return datetime.strptime(formatted_time, '%Y-%m-%dT%H:%M:%S')

    # TODO: Ignore and create a new one?
    raise ValueError('Unparseable certificate: %s' % output)
<|endoftext|>


<commit_before># Copyright (c) 2015 Mirantis, Inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from oslo_config import cfg

service_option = [
    cfg.BoolOpt("murano",
                default=True,
                help="Whether or not murano is expected to be available"),
    cfg.BoolOpt("murano_cfapi",
                default=False,
                help="Whether or not murano-cfapi is expected to be "
                     "unavailable by default"),
    cfg.BoolOpt("glare",
                default=False,
                help="Whether or not glare is expected to be unavailable")
]

application_catalog_group = cfg.OptGroup(name="application_catalog",
                                         title="Application Catalog Options")

service_broker_group = cfg.OptGroup(name="service_broker",
                                    title="Service Broker Options")

artifacts_group = cfg.OptGroup(name="artifacts",
                               title="Glance Artifacts Options")

ApplicationCatalogGroup = [
    # Application catalog tempest configuration
    cfg.StrOpt("region",
               default="",
               help="The application_catalog region name to use. If empty, "
                    "the value of identity.region is used instead. "
                    "If no such region is found in the service catalog, "
                    "the first found one is used."),

    cfg.StrOpt("linux_image",
               default="debian-8-m-agent.qcow2",
               help="Image for linux services"),

    cfg.StrOpt("catalog_type",
               default="application-catalog",
               help="Catalog type of Application Catalog."),

    cfg.StrOpt("endpoint_type",
               default="publicURL",
               choices=["publicURL", "adminURL", "internalURL"],
               help="The endpoint type for application catalog service."),

    cfg.IntOpt("build_interval",
               default=3,
               help="Time in seconds between application catalog"
                    " availability checks."),

    cfg.IntOpt("build_timeout",
               default=500,
               help="Timeout in seconds to wait for a application catalog"
                    " to become available."),
    cfg.BoolOpt("glare_backend",
                default=False,
                help="Tells tempest about murano glare backend "
                     "configuration."),
    cfg.BoolOpt("cinder_volume_tests",
                default=False,
                help="Whether or not cinder volumes attachment tests "
                     "are expected to run"),
    cfg.BoolOpt("deployment_tests",
                default=False,
                help="Whether or not deployment tests are expected to run")
]

ServiceBrokerGroup = [
    # Test runs control
    cfg.StrOpt("run_service_broker_tests",
               default=False,
               help="Defines whether run service broker api tests or not"),

    cfg.StrOpt("catalog_type",
               default="service-broker",
               help="Catalog type of Service Broker API"),

    cfg.StrOpt("endpoint_type",
               default="publicURL",
               choices=["publicURL", "adminURL", "internalURL"],
               help="The endpoint type for service broker service"),

    cfg.IntOpt("build_interval",
               default=3,
               help="Time in seconds between service broker"
                    " availability checks."),

    cfg.IntOpt("build_timeout",
               default=500,
               help="Timeout in seconds to wait for a service broker"
                    " to become available.")

]

ArtifactsGroup = [
    # Glance artifacts options
    cfg.StrOpt("catalog_type",
               default="artifact",
               help="Catalog type of Artifacts API"),

    cfg.StrOpt("endpoint_type",
               default="publicURL",
               choices=["publicURL", "adminURL", "internalURL"],
               help="The endpoint type for artifacts service"),

    cfg.IntOpt("build_interval",
               default=3,
               help="Time in seconds between artifacts"
                    " availability checks."),

    cfg.IntOpt("build_timeout",
               default=500,
               help="Timeout in seconds to wait for a artifacts"
                    " to become available.")
]
<commit_msg>Fix config type of run_service_broker_tests to Boolean<commit_after># Copyright (c) 2015 Mirantis, Inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from oslo_config import cfg

service_option = [
    cfg.BoolOpt("murano",
                default=True,
                help="Whether or not murano is expected to be available"),
    cfg.BoolOpt("murano_cfapi",
                default=False,
                help="Whether or not murano-cfapi is expected to be "
                     "unavailable by default"),
    cfg.BoolOpt("glare",
                default=False,
                help="Whether or not glare is expected to be unavailable")
]

application_catalog_group = cfg.OptGroup(name="application_catalog",
                                         title="Application Catalog Options")

service_broker_group = cfg.OptGroup(name="service_broker",
                                    title="Service Broker Options")

artifacts_group = cfg.OptGroup(name="artifacts",
                               title="Glance Artifacts Options")

ApplicationCatalogGroup = [
    # Application catalog tempest configuration
    cfg.StrOpt("region",
               default="",
               help="The application_catalog region name to use. If empty, "
                    "the value of identity.region is used instead. "
                    "If no such region is found in the service catalog, "
                    "the first found one is used."),

    cfg.StrOpt("linux_image",
               default="debian-8-m-agent.qcow2",
               help="Image for linux services"),

    cfg.StrOpt("catalog_type",
               default="application-catalog",
               help="Catalog type of Application Catalog."),

    cfg.StrOpt("endpoint_type",
               default="publicURL",
               choices=["publicURL", "adminURL", "internalURL"],
               help="The endpoint type for application catalog service."),

    cfg.IntOpt("build_interval",
               default=3,
               help="Time in seconds between application catalog"
                    " availability checks."),

    cfg.IntOpt("build_timeout",
               default=500,
               help="Timeout in seconds to wait for a application catalog"
                    " to become available."),
    cfg.BoolOpt("glare_backend",
                default=False,
                help="Tells tempest about murano glare backend "
                     "configuration."),
    cfg.BoolOpt("cinder_volume_tests",
                default=False,
                help="Whether or not cinder volumes attachment tests "
                     "are expected to run"),
    cfg.BoolOpt("deployment_tests",
                default=False,
                help="Whether or not deployment tests are expected to run")
]

ServiceBrokerGroup = [
    # Test runs control
    cfg.BoolOpt("run_service_broker_tests",
                default=False,
                help="Defines whether run service broker api tests or not"),

    cfg.StrOpt("catalog_type",
               default="service-broker",
               help="Catalog type of Service Broker API"),

    cfg.StrOpt("endpoint_type",
               default="publicURL",
               choices=["publicURL", "adminURL", "internalURL"],
               help="The endpoint type for service broker service"),

    cfg.IntOpt("build_interval",
               default=3,
               help="Time in seconds between service broker"
                    " availability checks."),

    cfg.IntOpt("build_timeout",
               default=500,
               help="Timeout in seconds to wait for a service broker"
                    " to become available.")

]

ArtifactsGroup = [
    # Glance artifacts options
    cfg.StrOpt("catalog_type",
               default="artifact",
               help="Catalog type of Artifacts API"),

    cfg.StrOpt("endpoint_type",
               default="publicURL",
               choices=["publicURL", "adminURL", "internalURL"],
               help="The endpoint type for artifacts service"),

    cfg.IntOpt("build_interval",
               default=3,
               help="Time in seconds between artifacts"
                    " availability checks."),

    cfg.IntOpt("build_timeout",
               default=500,
               help="Timeout in seconds to wait for a artifacts"
                    " to become available.")
]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Paul Nilsson, paul.nilsson@cern.ch, 2018

from pilot.user.atlas.common import get_db_info
from pilot.util.auxiliary import get_logger
from pilot.util.jobmetrics import get_job_metrics_entry
from pilot.util.processes import get_core_count

import logging

logger = logging.getLogger(__name__)

def get_job_metrics(job):
    """
    Return a properly formatted job metrics string.
    The format of the job metrics string is defined by the server. It will be reported to the server during updateJob.

    Example of job metrics:
    Number of events read | Number of events written | vmPeak maximum | vmPeak average | RSS average | ..
    Format: nEvents=<int> nEventsW=<int> vmPeakMax=<int> vmPeakMean=<int> RSSMean=<int> hs06=<float> shutdownTime=<int>
            cpuFactor=<float> cpuLimit=<float> diskLimit=<float> jobStart=<int> memLimit=<int> runLimit=<float>

    :param job: job object.
    :return: job metrics (string).
    """

    log = get_logger(job.jobid)

    job_metrics = ""

    # report core count (will also set corecount in job object)
    corecount = get_core_count(job)
    if corecount is not None and corecount != "NULL" and corecount != 'null':
        job_metrics += get_job_metrics_entry("coreCount", corecount)

    # report number of events
    if job.nevents > 0:
        job_metrics += get_job_metrics_entry("nEvents", job.nevents)
    if job.neventsw > 0:
        job_metrics += get_job_metrics_entry("nEventsW", job.neventsw)

    # add metadata from job report
    if job.metadata:
        db_time, db_data = get_db_info(job.metadata)
    if job.dbtime and job.dbtime != "":
        job_metrics += get_job_metrics_entry("dbTime", job.dbtime)
    if job.dbdata and job.dbdata != "":
        job_metrics += get_job_metrics_entry("dbData", job.dbdata)

    # event service
    # if job.external_stageout_time:
    #     job_metrics += get_job_metrics_entry("ExStageoutTime", job.external_stageout_time)

    # eventservice zip file
    # if job.output_zip_name and job.output_zip_bucket_id:
    #     job_metrics += get_job_metrics_entry("outputZipName", os.path.basename(job.output_zip_name))
    #     job_metrics += get_job_metrics_entry("outputZipBucketID", job.output_zip_bucket_id)

    # report on which OS bucket the log was written to, if any
    # if job.log_bucket_id != -1:
    #     job_metrics += get_job_metrics_entry("logBucketID", job.log_bucket_id)

    # yoda
    # if job.yoda_job_metrics:
    #     for key in job.yoda_job_metrics:
    #         if key == 'startTime' or key == 'endTime':
    #             value = strftime("%Y-%m-%d %H:%M:%S", gmtime(job.yoda_mob_metrics[key]))
    #             job_metrics += get_job_metrics_entry(key, value)
    #         elif key.startswith("min") or key.startswith("max"):
    #             pass
    #         else:
    #             job_metrics += get_job_metrics_entry(key, job.yoda_job_metrics[key])

    # get the max disk space used by the payload (at the end of a job)
    if job.state == "finished" or job.state == "failed" or job.state == "holding":
        max_space = job.get_max_workdir_size()
        if max_space > 0L:
            job_metrics += get_job_metrics_entry("workDirSize", max_space)
        else:
            log.info("will not add max space = %d B to job metrics" % (max_space))

    # get analytics data

    # done with job metrics, now verify the string

    # correct for potential initial and trailing space
    job_metrics = job_metrics.lstrip().rstrip()

    if job_metrics != "":
        log.debug('job metrics=\"%s\"' % (job_metrics))
    else:
        log.debug("no job metrics (all values are zero)")

    # is job_metrics within allowed size?
    if len(job_metrics) > 500:
        log.warning("job_metrics out of size (%d)" % (len(job_metrics)))

        # try to reduce the field size and remove the last entry which might be cut
        job_metrics = job_metrics[:500]
        job_metrics = " ".join(job_metrics.split(" ")[:-1])
        log.warning("job_metrics has been reduced to: %s" % (job_metrics))

    return job_metrics
<commit_msg>Added initial code for adding memory leak info to job metrics using the analytics package<commit_after>#!/usr/bin/env python
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Paul Nilsson, paul.nilsson@cern.ch, 2018

from pilot.api import analytics
from pilot.util.auxiliary import get_logger
from pilot.util.jobmetrics import get_job_metrics_entry
from pilot.util.processes import get_core_count

from .common import get_db_info
from .utilities import get_memory_monitor_output_filename

import os
import logging

logger = logging.getLogger(__name__)

def get_job_metrics(job):
    """
    Return a properly formatted job metrics string.
    The format of the job metrics string is defined by the server. It will be reported to the server during updateJob.

    Example of job metrics:
    Number of events read | Number of events written | vmPeak maximum | vmPeak average | RSS average | ..
    Format: nEvents=<int> nEventsW=<int> vmPeakMax=<int> vmPeakMean=<int> RSSMean=<int> hs06=<float> shutdownTime=<int>
            cpuFactor=<float> cpuLimit=<float> diskLimit=<float> jobStart=<int> memLimit=<int> runLimit=<float>

    :param job: job object.
    :return: job metrics (string).
    """

    log = get_logger(job.jobid)

    job_metrics = ""

    # report core count (will also set corecount in job object)
    corecount = get_core_count(job)
    if corecount is not None and corecount != "NULL" and corecount != 'null':
        job_metrics += get_job_metrics_entry("coreCount", corecount)

    # report number of events
    if job.nevents > 0:
        job_metrics += get_job_metrics_entry("nEvents", job.nevents)
    if job.neventsw > 0:
        job_metrics += get_job_metrics_entry("nEventsW", job.neventsw)

    # add metadata from job report
    if job.metadata:
        db_time, db_data = get_db_info(job.metadata)
    if job.dbtime and job.dbtime != "":
        job_metrics += get_job_metrics_entry("dbTime", job.dbtime)
    if job.dbdata and job.dbdata != "":
        job_metrics += get_job_metrics_entry("dbData", job.dbdata)

    # event service
    # if job.external_stageout_time:
    #     job_metrics += get_job_metrics_entry("ExStageoutTime", job.external_stageout_time)

    # eventservice zip file
    # if job.output_zip_name and job.output_zip_bucket_id:
    #     job_metrics += get_job_metrics_entry("outputZipName", os.path.basename(job.output_zip_name))
    #     job_metrics += get_job_metrics_entry("outputZipBucketID", job.output_zip_bucket_id)

    # report on which OS bucket the log was written to, if any
    # if job.log_bucket_id != -1:
    #     job_metrics += get_job_metrics_entry("logBucketID", job.log_bucket_id)

    # yoda
    # if job.yoda_job_metrics:
    #     for key in job.yoda_job_metrics:
    #         if key == 'startTime' or key == 'endTime':
    #             value = strftime("%Y-%m-%d %H:%M:%S", gmtime(job.yoda_mob_metrics[key]))
    #             job_metrics += get_job_metrics_entry(key, value)
    #         elif key.startswith("min") or key.startswith("max"):
    #             pass
    #         else:
    #             job_metrics += get_job_metrics_entry(key, job.yoda_job_metrics[key])

    # get the max disk space used by the payload (at the end of a job)
    if job.state == "finished" or job.state == "failed" or job.state == "holding":
        max_space = job.get_max_workdir_size()
        if max_space > 0L:
            job_metrics += get_job_metrics_entry("workDirSize", max_space)
        else:
            log.info("will not add max space = %d B to job metrics" % (max_space))

    # get analytics data
    path = os.path.join(job.workdir, get_memory_monitor_output_filename())
    if os.path.exists(path):
        client = analytics.Analytics()
        table = client.get_table(path)

        if table:
            x = table['Time']
            y = table['PSS']

            if len(x) >= 2 and len(y) >= 2:
                fit = client.fit(x, y)

                log.info('current memory leak: %f B/s' % fit.slope())

    # done with job metrics, now verify the string

    # correct for potential initial and trailing space
    job_metrics = job_metrics.lstrip().rstrip()

    if job_metrics != "":
        log.debug('job metrics=\"%s\"' % (job_metrics))
    else:
        log.debug("no job metrics (all values are zero)")

    # is job_metrics within allowed size?
    if len(job_metrics) > 500:
        log.warning("job_metrics out of size (%d)" % (len(job_metrics)))

        # try to reduce the field size and remove the last entry which might be cut
        job_metrics = job_metrics[:500]
        job_metrics = " ".join(job_metrics.split(" ")[:-1])
        log.warning("job_metrics has been reduced to: %s" % (job_metrics))

    return job_metrics
<|endoftext|>


<commit_before>"""Manages jobs."""

from __future__ import absolute_import

import bisect

from .env import CPUS

__all__ = ["QueueManager", "queues", "attr_queue", "config_queue",
           "checksum_queue", "fetch_queue", "build_queue", "install_queue"]

class QueueManager(object):
    """Manages jobs and runs them as resources come available."""

    def __init__(self, load=1):
        """Initialise the manager with an indication of load available."""
        self._load = load
        self._sort = False
        self.queue = []
        self.active = []
        self.stalled = []
        self.active_load = 0

    def __len__(self):
        return len(self.queue) + len(self.active) + len(self.stalled)

    @property
    def load(self):
        """Returns the current allowed load."""
        return self._load

    @load.setter
    def load(self, load):
        """Set the load and start jobs as required."""
        run = load > self._load
        self._load = load
        if run:
            self._run()

    def add(self, job):
        """Add a job to be run."""
        assert(job not in self.queue)
        if self._sort:
            self.queue.append(job)
        else:
            bisect.insort(self.queue, job)
        if self.active_load < self._load:
            self._run()

    def done(self, job):
        """Indicates a job has completed."""
        self.active.remove(job)
        self.active_load -= job.load
        if self.active_load < self._load:
            self._run()

    def reorder(self):
        """Reorder the queued jobs as their priority may have changed."""
        self._sort = True

    def remove(self, job):
        """Remove a job from being run."""
        try:
            self.queue.remove(job)
        except ValueError:
            return False
        return True

    def _run(self):
        """Fills up the remaining load with jobs"""
        from .job import StalledJob
        assert(self.active_load < self._load)

        stalled = []
        if self._sort:
            self._sort = False
            self.stalled.sort()
            self.queue.sort()
        for queue in (self.stalled, self.queue):
            while self.active_load < self._load and len(queue):
                job = self._find_job(self._load - self.active_load, queue)
                try:
                    self.active_load += job.load
                    self.active.append(job)
                    job.run(self)
                except StalledJob:
                    self.active_load -= job.load
                    self.active.remove(job)
                    stalled.append(job)
        if len(stalled):
            self.stalled.extend(stalled)
            self.stalled.sort()

    @staticmethod
    def _find_job(load, queue):
        """Find a job from queue that has at most load."""
        best_job = queue[0]
        best_idx = 0
        if best_job.load <= load:
            return queue.pop(0)
        for idx in range(1, len(queue)):
            job = queue[idx]
            if job.load <= load:
                return queue.pop(idx)
            if best_job.load > job.load:
                best_job = job
                best_idx = idx
        return queue.pop(best_idx)

attr_queue = QueueManager(CPUS * 2)
clean_queue = QueueManager(1)

config_queue = QueueManager(1)
checksum_queue = QueueManager(max(1, CPUS // 2))
fetch_queue = QueueManager(1)
build_queue = QueueManager(CPUS * 2)
install_queue = QueueManager(1)
package_queue = QueueManager(1)
queues = (config_queue, checksum_queue, fetch_queue, build_queue,
          install_queue, package_queue, install_queue)
<commit_msg>Use the install queue for REPOINSTALL.<commit_after>"""Manages jobs."""

from __future__ import absolute_import

import bisect

from .env import CPUS

__all__ = ["QueueManager", "queues", "attr_queue", "config_queue",
           "checksum_queue", "fetch_queue", "build_queue", "install_queue"]

class QueueManager(object):
    """Manages jobs and runs them as resources come available."""

    def __init__(self, load=1):
        """Initialise the manager with an indication of load available."""
        self._load = load
        self._sort = False
        self.queue = []
        self.active = []
        self.stalled = []
        self.active_load = 0

    def __len__(self):
        return len(self.queue) + len(self.active) + len(self.stalled)

    @property
    def load(self):
        """Returns the current allowed load."""
        return self._load

    @load.setter
    def load(self, load):
        """Set the load and start jobs as required."""
        run = load > self._load
        self._load = load
        if run:
            self._run()

    def add(self, job):
        """Add a job to be run."""
        assert(job not in self.queue)
        if self._sort:
            self.queue.append(job)
        else:
            bisect.insort(self.queue, job)
        if self.active_load < self._load:
            self._run()

    def done(self, job):
        """Indicates a job has completed."""
        self.active.remove(job)
        self.active_load -= job.load
        if self.active_load < self._load:
            self._run()

    def reorder(self):
        """Reorder the queued jobs as their priority may have changed."""
        self._sort = True

    def remove(self, job):
        """Remove a job from being run."""
        try:
            self.queue.remove(job)
        except ValueError:
            return False
        return True

    def _run(self):
        """Fills up the remaining load with jobs"""
        from .job import StalledJob
        assert(self.active_load < self._load)

        stalled = []
        if self._sort:
            self._sort = False
            self.stalled.sort()
            self.queue.sort()
        for queue in (self.stalled, self.queue):
            while self.active_load < self._load and len(queue):
                job = self._find_job(self._load - self.active_load, queue)
                try:
                    self.active_load += job.load
                    self.active.append(job)
                    job.run(self)
                except StalledJob:
                    self.active_load -= job.load
                    self.active.remove(job)
                    stalled.append(job)
        if len(stalled):
            self.stalled.extend(stalled)
            self.stalled.sort()

    @staticmethod
    def _find_job(load, queue):
        """Find a job from queue that has at most load."""
        best_job = queue[0]
        best_idx = 0
        if best_job.load <= load:
            return queue.pop(0)
        for idx in range(1, len(queue)):
            job = queue[idx]
            if job.load <= load:
                return queue.pop(idx)
            if best_job.load > job.load:
                best_job = job
                best_idx = idx
        return queue.pop(best_idx)

attr_queue = QueueManager(CPUS * 2)
clean_queue = QueueManager(1)

config_queue = QueueManager(1)
checksum_queue = QueueManager(max(1, CPUS // 2))
fetch_queue = QueueManager(1)
build_queue = QueueManager(CPUS * 2)
install_queue = QueueManager(1)
package_queue = QueueManager(1)
queues = (config_queue, checksum_queue, fetch_queue, build_queue,
          install_queue, package_queue, install_queue, install_queue)
<|endoftext|>


<commit_before>import modules
import logging

class ModuleLoadException(Exception):
	def __init__(self, exception):
		super().__init__()
		self.exception = exception
	def __str__(self):
		return 'ModuleLoadException({})'.format(self.exception)

class ModuleManager( object ):
	def __init__(self, bot, blacklist=None):
		self.bot = bot
		self.modules = {}
		self.loaded_modules = {}
		self.blacklist = blacklist or []
		for module_name in modules.get_modules():
			if module_name in blacklist:
				continue
			logging.info( 'Loading module {0}: {1}'.format( module_name, self.add_module( module_name ) ) )

	def unload( self ):
		for module_name in list( self.modules.keys() ):
			logging.info( 'Unloading module {0}: {1}'.format( module_name, self.remove_module( module_name ) ) )

	def reload_modules( self ):
		"""Reload all modules. Warning: this is fairly destructive"""
		# remove modules that no longer exist
		for module_name in [ m for m in self.modules.keys() if m not in modules.get_modules() ]:
			self.remove_module( module_name )
		# reload all modules
		for module_name in self.modules:
			self.reload_module( module_name )
		# add modules that are not added yet
		for module_name in [ m for m in modules.get_modules() if m not in self.modules ]:
			self.add_module( module_name )

	def get_modules( self ):
		"""Get all found modules"""
		return iter( self.modules.keys() )

	def get_loaded_modules( self ):
		"""Get all loaded modules"""
		return list( self.loaded_modules.items() )

	def module_is_loaded( self, module_name ):
		return module_name in self.loaded_modules
	def get_module( self, module_name ):
		try:
			return self.loaded_modules[ module_name ]
		except:
			return False

	def get_available_modules( self ):
		"""Get all available modules: modules that are found but not loaded"""
		modules = [ key for key in self.modules.keys() if key not in self.loaded_modules ]
		modules.sort()
		return modules

	def add_module( self, module_name ):
		"""Load a module"""
		if module_name in self.modules:
			return 'Module already available'
		try:
			module = modules.get_module( module_name )
			if module:
				self.modules[ module_name ] = module
				return 'Module added'
		except Exception as e:
			return 'Error loading module: {0}'.format( e )
		return 'Module not available'

	def remove_module( self, module_name ):
		"""Unload a module"""
		if not module_name in self.modules:
			return 'Module not available'
		if module_name in self.loaded_modules:
			self.disable_module( module_name )
		del self.modules[ module_name ]
		return 'Module removed'

	def enable_module( self, module_name ):
		"""Enable a module"""
		if not module_name in self.modules:
			return 'Module {} not available'.format( module_name )
		if module_name in self.loaded_modules:
			return 'Module {} already enabled'.format( module_name )
		try:
			self.loaded_modules[ module_name ] = self.modules[ module_name ]( self )
		except Exception as e:
			raise ModuleLoadException(e)
			return 'Module {} failed to load: {}'.format( module_name, e )
		return 'Module {} enabled'.format( module_name )

	def disable_module( self, module_name ):
		"""Disable a module"""
		if not module_name in self.loaded_modules:
			return 'Module not enabled'
		try:
			self.loaded_modules[ module_name ].stop()
		except Exception as e:
			logger.warning( 'Module %s failed to stop: %s', module_name, e )
		del self.loaded_modules[ module_name ]
		return 'Module {} disabled'.format( module_name )

	def restart_module( self, module_name ):
		"""Restart a module"""
		if not module_name in self.modules:
			return 'Module {} not available'.format( module_name )
		if module_name in self.loaded_modules:
			self.disable_module( module_name )
		self.enable_module( module_name )
		return 'Module {} restarted'.format( module_name )

	def reload_module( self, module_name ):
		"""Reload a module"""
		start_module = module_name in self.loaded_modules
		self.remove_module( module_name ) # remove to clear references
		modules.reload_module( module_name ) # actually reload class
		self.add_module( module_name ) # re-add module
		if start_module: # enable if it was enabled
			self.enable_module( module_name )
		return 'Module {} reloaded'.format( module_name )

	# methods from Bot

	def notice( self, target, message ):
		self.bot.notice( target, message )

	def privmsg( self, target, message ):
		self.bot.privmsg( target, message )

	def get_config( self, group, key, default = None ):
		return self.bot.get_config( group, key, default )

	def set_config( self, group, key, value ):
		return self.bot.set_config( group, key, value )

	def get_module(self, name):
		if name in self.loaded_modules:
			return self.loaded_modules[name]
<commit_msg>Fixed bug in ModuleManager.<commit_after>import modules
import logging

class ModuleLoadException(Exception):
	def __init__(self, exception):
		super().__init__()
		self.exception = exception
	def __str__(self):
		return 'ModuleLoadException({})'.format(self.exception)

class ModuleManager( object ):
	def __init__(self, bot, blacklist=None):
		self.bot = bot
		self.modules = {}
		self.loaded_modules = {}
		blacklist = blacklist or []
		for module_name in modules.get_modules():
			if module_name in blacklist:
				continue
			logging.info( 'Loading module {0}: {1}'.format( module_name, self.add_module( module_name ) ) )

	def unload( self ):
		for module_name in list( self.modules.keys() ):
			logging.info( 'Unloading module {0}: {1}'.format( module_name, self.remove_module( module_name ) ) )

	def reload_modules( self ):
		"""Reload all modules. Warning: this is fairly destructive"""
		# remove modules that no longer exist
		for module_name in [ m for m in self.modules.keys() if m not in modules.get_modules() ]:
			self.remove_module( module_name )
		# reload all modules
		for module_name in self.modules:
			self.reload_module( module_name )
		# add modules that are not added yet
		for module_name in [ m for m in modules.get_modules() if m not in self.modules ]:
			self.add_module( module_name )

	def get_modules( self ):
		"""Get all found modules"""
		return iter( self.modules.keys() )

	def get_loaded_modules( self ):
		"""Get all loaded modules"""
		return list( self.loaded_modules.items() )

	def module_is_loaded( self, module_name ):
		return module_name in self.loaded_modules
	def get_module( self, module_name ):
		try:
			return self.loaded_modules[ module_name ]
		except:
			return False

	def get_available_modules( self ):
		"""Get all available modules: modules that are found but not loaded"""
		modules = [ key for key in self.modules.keys() if key not in self.loaded_modules ]
		modules.sort()
		return modules

	def add_module( self, module_name ):
		"""Load a module"""
		if module_name in self.modules:
			return 'Module already available'
		try:
			module = modules.get_module( module_name )
			if module:
				self.modules[ module_name ] = module
				return 'Module added'
		except Exception as e:
			return 'Error loading module: {0}'.format( e )
		return 'Module not available'

	def remove_module( self, module_name ):
		"""Unload a module"""
		if not module_name in self.modules:
			return 'Module not available'
		if module_name in self.loaded_modules:
			self.disable_module( module_name )
		del self.modules[ module_name ]
		return 'Module removed'

	def enable_module( self, module_name ):
		"""Enable a module"""
		if not module_name in self.modules:
			return 'Module {} not available'.format( module_name )
		if module_name in self.loaded_modules:
			return 'Module {} already enabled'.format( module_name )
		try:
			self.loaded_modules[ module_name ] = self.modules[ module_name ]( self )
		except Exception as e:
			raise ModuleLoadException(e)
			return 'Module {} failed to load: {}'.format( module_name, e )
		return 'Module {} enabled'.format( module_name )

	def disable_module( self, module_name ):
		"""Disable a module"""
		if not module_name in self.loaded_modules:
			return 'Module not enabled'
		try:
			self.loaded_modules[ module_name ].stop()
		except Exception as e:
			logger.warning( 'Module %s failed to stop: %s', module_name, e )
		del self.loaded_modules[ module_name ]
		return 'Module {} disabled'.format( module_name )

	def restart_module( self, module_name ):
		"""Restart a module"""
		if not module_name in self.modules:
			return 'Module {} not available'.format( module_name )
		if module_name in self.loaded_modules:
			self.disable_module( module_name )
		self.enable_module( module_name )
		return 'Module {} restarted'.format( module_name )

	def reload_module( self, module_name ):
		"""Reload a module"""
		start_module = module_name in self.loaded_modules
		self.remove_module( module_name ) # remove to clear references
		modules.reload_module( module_name ) # actually reload class
		self.add_module( module_name ) # re-add module
		if start_module: # enable if it was enabled
			self.enable_module( module_name )
		return 'Module {} reloaded'.format( module_name )

	# methods from Bot

	def notice( self, target, message ):
		self.bot.notice( target, message )

	def privmsg( self, target, message ):
		self.bot.privmsg( target, message )

	def get_config( self, group, key, default = None ):
		return self.bot.get_config( group, key, default )

	def set_config( self, group, key, value ):
		return self.bot.set_config( group, key, value )

	def get_module(self, name):
		if name in self.loaded_modules:
			return self.loaded_modules[name]
<|endoftext|>