<commit_before>import logging
import uuid

import jsonschema
import newrelic.agent
import slugid

from treeherder.etl.common import to_timestamp
from treeherder.etl.exceptions import MissingPushException
from treeherder.etl.jobs import store_job_data
from treeherder.etl.schema import job_json_schema
from treeherder.model.models import (Push,
                                     Repository)

logger = logging.getLogger(__name__)

class JobLoader(object):
    """Validate, transform and load a list of Jobs"""

    TEST_RESULT_MAP = {
        "success": "success",
        "fail": "testfailed",
        "exception": "exception",
        "canceled": "usercancel",
        "superseded": "superseded",
        "unknown": "unknown"
    }
    BUILD_RESULT_MAP = {
        "success": "success",
        "fail": "busted",
        "exception": "exception",
        "canceled": "usercancel",
        "superseded": "superseded",
        "unknown": "unknown"
    }
    TIME_FIELD_MAP = {
        "submit_timestamp": "timeScheduled",
        "start_timestamp": "timeStarted",
        "end_timestamp": "timeCompleted"
    }
    PLATFORM_FIELD_MAP = {
        "build_platform": "buildMachine",
        "machine_platform": "runMachine"
    }

    def process_job(self, pulse_job):
        if self._is_valid_job(pulse_job):
            try:
                project = pulse_job["origin"]["project"]
                newrelic.agent.add_custom_parameter("project", project)

                repository = Repository.objects.get(name=project)

                if pulse_job["state"] != "unscheduled":
                    try:
                        self.validate_revision(repository, pulse_job)
                        store_job_data(repository, [self.transform(pulse_job)])
                    except AttributeError:
                        logger.warning("Skipping job due to bad attribute", exc_info=1)
            except Repository.DoesNotExist:
                logger.info("Job with unsupported project: %s", project)

    def validate_revision(self, repository, pulse_job):
        revision = pulse_job["origin"].get("revision")
        # will raise an exception if repository with name does not
        # exist (which we want, I think, to draw attention to the problem)
        # check the revision for this job has an existing push
        # If it doesn't, then except out so that the celery task will
        # retry till it DOES exist.
        revision_field = 'revision__startswith' if len(revision) < 40 else 'revision'
        filter_kwargs = {'repository': repository, revision_field: revision}

        if revision_field == 'revision__startswith':
            newrelic.agent.record_custom_event(
                'short_revision_job_loader',
                {'error': 'Revision <40 chars', 'revision': revision, 'job': pulse_job}
            )

        if not Push.objects.filter(**filter_kwargs).exists():
            raise MissingPushException(
                "No push found in {} for revision {}".format(
                    pulse_job["origin"]["project"],
                    revision))

    def transform(self, pulse_job):
        """
        Transform a pulse job into a job that can be written to disk.  Log
        References and artifacts will also be transformed and loaded with the
        job.

        We can rely on the structure of ``pulse_job`` because it will
        already have been validated against the JSON Schema at this point.
        """
        job_guid = pulse_job["taskId"]

        x = {
            "job": {
                "job_guid": job_guid,
                "name": pulse_job["display"].get("jobName", "unknown"),
                "job_symbol": self._get_job_symbol(pulse_job),
                "group_name": pulse_job["display"].get("groupName", "unknown"),
                "group_symbol": pulse_job["display"].get("groupSymbol"),
                "product_name": pulse_job.get("productName", "unknown"),
                "state": pulse_job["state"],
                "result": self._get_result(pulse_job),
                "reason": pulse_job.get("reason", "unknown"),
                "who": pulse_job.get("owner", "unknown"),
                "build_system_type": pulse_job["buildSystem"],
                "tier": pulse_job.get("tier", 1),
                "machine": self._get_machine(pulse_job),
                "option_collection": self._get_option_collection(pulse_job),
                "log_references": self._get_log_references(pulse_job),
                "artifacts": self._get_artifacts(pulse_job, job_guid),
            },
            "superseded": pulse_job.get("coalesced", []),
            "revision": pulse_job["origin"]["revision"]
        }

        # some or all the time fields may not be present in some cases
        for k, v in self.TIME_FIELD_MAP.items():
            if v in pulse_job:
                x["job"][k] = to_timestamp(pulse_job[v])

        # if only one platform is given, use it.
        default_platform = pulse_job.get(
            "buildMachine",
            pulse_job.get("runMachine", {}))

        for k, v in self.PLATFORM_FIELD_MAP.items():
            platform_src = pulse_job[v] if v in pulse_job else default_platform
            x["job"][k] = self._get_platform(platform_src)

        # add some taskcluster metadata if it's available
        # currently taskcluster doesn't pass the taskId directly, so we'll
        # derive it from the guid, where it is stored in uncompressed
        # guid form of a slug (see: https://github.com/taskcluster/slugid)
        # FIXME: add support for processing the taskcluster information
        # properly, when it's available:
        # https://bugzilla.mozilla.org/show_bug.cgi?id=1323110#c7
        try:
            (decoded_task_id, retry_id) = job_guid.split('/')
            real_task_id = slugid.encode(uuid.UUID(decoded_task_id))
            x["job"].update({
                "taskcluster_task_id": real_task_id,
                "taskcluster_retry_id": int(retry_id)
            })
        # TODO: Figure out what exception types we actually expect here.
        except Exception:
            pass

        return x

    def _get_job_symbol(self, job):
        return "{}{}".format(
            job["display"].get("jobSymbol", ""),
            job["display"].get("chunkId", "")
        )

    def _get_artifacts(self, job, job_guid):
        artifact_funcs = [self._get_job_info_artifact,
                          self._get_text_log_summary_artifact]
        pulse_artifacts = []
        for artifact_func in artifact_funcs:
            artifact = artifact_func(job, job_guid)
            if artifact:
                pulse_artifacts.append(artifact)

        # add in any arbitrary artifacts included in the "extra" section
        pulse_artifacts.extend(self._get_extra_artifacts(job, job_guid))
        return pulse_artifacts

    def _get_job_info_artifact(self, job, job_guid):
        if "jobInfo" in job:
            ji = job["jobInfo"]
            job_details = []
            if "links" in ji:
                for link in ji["links"]:
                    job_details.append({
                        "url": link["url"],
                        "content_type": "link",
                        "value": link["linkText"],
                        "title": link["label"]
                    })

            artifact = {
                "blob": {
                    "job_details": job_details
                },
                "type": "json",
                "name": "Job Info",
                "job_guid": job_guid
            }
            return artifact

    def _get_text_log_summary_artifact(self, job, job_guid):
        # We can only have one text_log_summary artifact,
        # so pick the first log with steps to create it.

        if "logs" in job:
            for log in job["logs"]:
                if "steps" in log:
                    old_steps = log["steps"]
                    new_steps = []

                    for idx, step in enumerate(old_steps):
                        errors = step.get("errors", [])

                        new_steps.append({
                            "name": step["name"],
                            "result": self._get_step_result(job, step["result"]),
                            "started": step["timeStarted"],
                            "finished": step["timeFinished"],
                            "started_linenumber": step["lineStarted"],
                            "finished_linenumber": step["lineFinished"],
                            "errors": errors,
                            "order": idx
                        })

                    return {
                        "blob": {
                            "step_data": {
                                "steps": new_steps,
                                "errors_truncated": log.get("errorsTruncated")
                            },
                            "logurl": log["url"]
                        },
                        "type": "json",
                        "name": "text_log_summary",
                        "job_guid": job_guid
                    }

    def _get_extra_artifacts(self, job, job_guid):
        artifacts = []
        if "extra" in job and "artifacts" in job["extra"]:
            for extra_artifact in job["extra"]["artifacts"]:
                artifact = extra_artifact
                artifact["job_guid"] = job_guid
                artifacts.append(artifact)

        return artifacts

    def _get_log_references(self, job):
        log_references = []
        for logref in job.get("logs", []):
            log_references.append({
                "name": logref["name"],
                "url": logref["url"],
                "parse_status": "parsed" if "steps" in logref else "pending"
            })
        log_references.extend(self._get_errorsummary_log_references(job))
        return log_references

    def _get_errorsummary_log_references(self, job):
        log_references = []
        try:
            links = job["jobInfo"]["links"]
        except KeyError:
            return []
        for link in links:
            text = link.get("linkText", "")
            if "url" in link and text.endswith("_errorsummary.log"):
                log_references.append({"name": "errorsummary_json",
                                       "url": link["url"],
                                       "parse_status": "pending"})
        return log_references

    def _get_option_collection(self, job):
        option_collection = {"opt": True}
        if "labels" in job:
            option_collection = {}
            for option in job["labels"]:
                option_collection[option] = True
        return option_collection

    def _get_platform(self, platform_src):
        platform = {}
        if platform_src:
            platform = {
                "platform": platform_src["platform"],
                "os_name": platform_src["os"],
                "architecture": platform_src["architecture"]
            }
        return platform

    def _get_machine(self, job):
        machine = "unknown"
        if "buildMachine" in job:
            machine = job["buildMachine"].get("name", machine)
        elif "runMachine" in job:
            machine = job["runMachine"].get("name", machine)
        return machine

    def _get_result(self, job):
        if job["state"] == "completed":
            resmap = self.TEST_RESULT_MAP if job["jobKind"] == "test" else self.BUILD_RESULT_MAP
            result = job.get("result", "unknown")
            if job.get("isRetried", False):
                return "retry"
            return resmap[result]
        return "unknown"

    def _get_step_result(self, job, result):
        resmap = self.TEST_RESULT_MAP if job["jobKind"] == "test" else self.BUILD_RESULT_MAP
        return resmap[result]

    def _is_valid_job(self, pulse_job):
        try:
            jsonschema.validate(pulse_job, job_json_schema)
        except (jsonschema.ValidationError, jsonschema.SchemaError) as e:
            logger.error("JSON Schema validation error during job ingestion: %s", e)
            return False
        return True
<commit_msg>Bug 1526763 - Explicitly convert encoded task IDs back to a string (#4618)<commit_after>import logging
import uuid

import jsonschema
import newrelic.agent
import slugid

from treeherder.etl.common import to_timestamp
from treeherder.etl.exceptions import MissingPushException
from treeherder.etl.jobs import store_job_data
from treeherder.etl.schema import job_json_schema
from treeherder.model.models import (Push,
                                     Repository)

logger = logging.getLogger(__name__)

class JobLoader(object):
    """Validate, transform and load a list of Jobs"""

    TEST_RESULT_MAP = {
        "success": "success",
        "fail": "testfailed",
        "exception": "exception",
        "canceled": "usercancel",
        "superseded": "superseded",
        "unknown": "unknown"
    }
    BUILD_RESULT_MAP = {
        "success": "success",
        "fail": "busted",
        "exception": "exception",
        "canceled": "usercancel",
        "superseded": "superseded",
        "unknown": "unknown"
    }
    TIME_FIELD_MAP = {
        "submit_timestamp": "timeScheduled",
        "start_timestamp": "timeStarted",
        "end_timestamp": "timeCompleted"
    }
    PLATFORM_FIELD_MAP = {
        "build_platform": "buildMachine",
        "machine_platform": "runMachine"
    }

    def process_job(self, pulse_job):
        if self._is_valid_job(pulse_job):
            try:
                project = pulse_job["origin"]["project"]
                newrelic.agent.add_custom_parameter("project", project)

                repository = Repository.objects.get(name=project)

                if pulse_job["state"] != "unscheduled":
                    try:
                        self.validate_revision(repository, pulse_job)
                        store_job_data(repository, [self.transform(pulse_job)])
                    except AttributeError:
                        logger.warning("Skipping job due to bad attribute", exc_info=1)
            except Repository.DoesNotExist:
                logger.info("Job with unsupported project: %s", project)

    def validate_revision(self, repository, pulse_job):
        revision = pulse_job["origin"].get("revision")
        # will raise an exception if repository with name does not
        # exist (which we want, I think, to draw attention to the problem)
        # check the revision for this job has an existing push
        # If it doesn't, then except out so that the celery task will
        # retry till it DOES exist.
        revision_field = 'revision__startswith' if len(revision) < 40 else 'revision'
        filter_kwargs = {'repository': repository, revision_field: revision}

        if revision_field == 'revision__startswith':
            newrelic.agent.record_custom_event(
                'short_revision_job_loader',
                {'error': 'Revision <40 chars', 'revision': revision, 'job': pulse_job}
            )

        if not Push.objects.filter(**filter_kwargs).exists():
            raise MissingPushException(
                "No push found in {} for revision {}".format(
                    pulse_job["origin"]["project"],
                    revision))

    def transform(self, pulse_job):
        """
        Transform a pulse job into a job that can be written to disk.  Log
        References and artifacts will also be transformed and loaded with the
        job.

        We can rely on the structure of ``pulse_job`` because it will
        already have been validated against the JSON Schema at this point.
        """
        job_guid = pulse_job["taskId"]

        x = {
            "job": {
                "job_guid": job_guid,
                "name": pulse_job["display"].get("jobName", "unknown"),
                "job_symbol": self._get_job_symbol(pulse_job),
                "group_name": pulse_job["display"].get("groupName", "unknown"),
                "group_symbol": pulse_job["display"].get("groupSymbol"),
                "product_name": pulse_job.get("productName", "unknown"),
                "state": pulse_job["state"],
                "result": self._get_result(pulse_job),
                "reason": pulse_job.get("reason", "unknown"),
                "who": pulse_job.get("owner", "unknown"),
                "build_system_type": pulse_job["buildSystem"],
                "tier": pulse_job.get("tier", 1),
                "machine": self._get_machine(pulse_job),
                "option_collection": self._get_option_collection(pulse_job),
                "log_references": self._get_log_references(pulse_job),
                "artifacts": self._get_artifacts(pulse_job, job_guid),
            },
            "superseded": pulse_job.get("coalesced", []),
            "revision": pulse_job["origin"]["revision"]
        }

        # some or all the time fields may not be present in some cases
        for k, v in self.TIME_FIELD_MAP.items():
            if v in pulse_job:
                x["job"][k] = to_timestamp(pulse_job[v])

        # if only one platform is given, use it.
        default_platform = pulse_job.get(
            "buildMachine",
            pulse_job.get("runMachine", {}))

        for k, v in self.PLATFORM_FIELD_MAP.items():
            platform_src = pulse_job[v] if v in pulse_job else default_platform
            x["job"][k] = self._get_platform(platform_src)

        # add some taskcluster metadata if it's available
        # currently taskcluster doesn't pass the taskId directly, so we'll
        # derive it from the guid, where it is stored in uncompressed
        # guid form of a slug (see: https://github.com/taskcluster/slugid)
        # FIXME: add support for processing the taskcluster information
        # properly, when it's available:
        # https://bugzilla.mozilla.org/show_bug.cgi?id=1323110#c7
        try:
            (decoded_task_id, retry_id) = job_guid.split('/')
            # Using `.decode('ascii')` to match the `slugid.encode()` implementation:
            # https://github.com/taskcluster/slugid.py/blob/v1.0.7/slugid/slugid.py#L22
            real_task_id = slugid.encode(uuid.UUID(decoded_task_id)).decode('ascii')
            x["job"].update({
                "taskcluster_task_id": real_task_id,
                "taskcluster_retry_id": int(retry_id)
            })
        # TODO: Figure out what exception types we actually expect here.
        except Exception:
            pass

        return x

    def _get_job_symbol(self, job):
        return "{}{}".format(
            job["display"].get("jobSymbol", ""),
            job["display"].get("chunkId", "")
        )

    def _get_artifacts(self, job, job_guid):
        artifact_funcs = [self._get_job_info_artifact,
                          self._get_text_log_summary_artifact]
        pulse_artifacts = []
        for artifact_func in artifact_funcs:
            artifact = artifact_func(job, job_guid)
            if artifact:
                pulse_artifacts.append(artifact)

        # add in any arbitrary artifacts included in the "extra" section
        pulse_artifacts.extend(self._get_extra_artifacts(job, job_guid))
        return pulse_artifacts

    def _get_job_info_artifact(self, job, job_guid):
        if "jobInfo" in job:
            ji = job["jobInfo"]
            job_details = []
            if "links" in ji:
                for link in ji["links"]:
                    job_details.append({
                        "url": link["url"],
                        "content_type": "link",
                        "value": link["linkText"],
                        "title": link["label"]
                    })

            artifact = {
                "blob": {
                    "job_details": job_details
                },
                "type": "json",
                "name": "Job Info",
                "job_guid": job_guid
            }
            return artifact

    def _get_text_log_summary_artifact(self, job, job_guid):
        # We can only have one text_log_summary artifact,
        # so pick the first log with steps to create it.

        if "logs" in job:
            for log in job["logs"]:
                if "steps" in log:
                    old_steps = log["steps"]
                    new_steps = []

                    for idx, step in enumerate(old_steps):
                        errors = step.get("errors", [])

                        new_steps.append({
                            "name": step["name"],
                            "result": self._get_step_result(job, step["result"]),
                            "started": step["timeStarted"],
                            "finished": step["timeFinished"],
                            "started_linenumber": step["lineStarted"],
                            "finished_linenumber": step["lineFinished"],
                            "errors": errors,
                            "order": idx
                        })

                    return {
                        "blob": {
                            "step_data": {
                                "steps": new_steps,
                                "errors_truncated": log.get("errorsTruncated")
                            },
                            "logurl": log["url"]
                        },
                        "type": "json",
                        "name": "text_log_summary",
                        "job_guid": job_guid
                    }

    def _get_extra_artifacts(self, job, job_guid):
        artifacts = []
        if "extra" in job and "artifacts" in job["extra"]:
            for extra_artifact in job["extra"]["artifacts"]:
                artifact = extra_artifact
                artifact["job_guid"] = job_guid
                artifacts.append(artifact)

        return artifacts

    def _get_log_references(self, job):
        log_references = []
        for logref in job.get("logs", []):
            log_references.append({
                "name": logref["name"],
                "url": logref["url"],
                "parse_status": "parsed" if "steps" in logref else "pending"
            })
        log_references.extend(self._get_errorsummary_log_references(job))
        return log_references

    def _get_errorsummary_log_references(self, job):
        log_references = []
        try:
            links = job["jobInfo"]["links"]
        except KeyError:
            return []
        for link in links:
            text = link.get("linkText", "")
            if "url" in link and text.endswith("_errorsummary.log"):
                log_references.append({"name": "errorsummary_json",
                                       "url": link["url"],
                                       "parse_status": "pending"})
        return log_references

    def _get_option_collection(self, job):
        option_collection = {"opt": True}
        if "labels" in job:
            option_collection = {}
            for option in job["labels"]:
                option_collection[option] = True
        return option_collection

    def _get_platform(self, platform_src):
        platform = {}
        if platform_src:
            platform = {
                "platform": platform_src["platform"],
                "os_name": platform_src["os"],
                "architecture": platform_src["architecture"]
            }
        return platform

    def _get_machine(self, job):
        machine = "unknown"
        if "buildMachine" in job:
            machine = job["buildMachine"].get("name", machine)
        elif "runMachine" in job:
            machine = job["runMachine"].get("name", machine)
        return machine

    def _get_result(self, job):
        if job["state"] == "completed":
            resmap = self.TEST_RESULT_MAP if job["jobKind"] == "test" else self.BUILD_RESULT_MAP
            result = job.get("result", "unknown")
            if job.get("isRetried", False):
                return "retry"
            return resmap[result]
        return "unknown"

    def _get_step_result(self, job, result):
        resmap = self.TEST_RESULT_MAP if job["jobKind"] == "test" else self.BUILD_RESULT_MAP
        return resmap[result]

    def _is_valid_job(self, pulse_job):
        try:
            jsonschema.validate(pulse_job, job_json_schema)
        except (jsonschema.ValidationError, jsonschema.SchemaError) as e:
            logger.error("JSON Schema validation error during job ingestion: %s", e)
            return False
        return True
<|endoftext|>


<commit_before>#!/usr/bin/python3
"""
Problem management script.
"""

import argparse
import sys
import imp
import random
import shutil
import os

import api
from api.common import APIException, InternalException

from os import path
from bson import json_util

def insert_problems(files):
    problems = get_problems(files)
    for problem in problems:
        try:
            api.problem.insert_problem(problem)
        except APIException as error:
            raise
            exit(1)
    errors = api.problem.analyze_problems()
    for error in errors:
        print(error)

def get_problems(files):
    problems = []
    for contents in files.values():
        for line in contents:
            problem = json_util.loads(line.strip())
            problems.append(problem)
    return problems

def migrate_problems(files, output_file, debug):

    migration_key = {
        "displayname": "name",
        "basescore": "score",
        "desc": "description",
        "relatedproblems": "related_problems"
    }

    migration_overwrites = {
        "grader": "test.py",
        "autogen": False
    }

    deletion_key = ["_id", "pid", "generator", "submissiontype", "devnotes"]

    problems = get_problems(files)
    output = ""

    def get_display_name_from_pid(problems, pid):
        for problem in problems:
            if problem.get("pid") == pid:
                return problem.get("displayname")

    for problem in problems:
        if problem.get("weightmap"):
            new_map = {}
            for pid, num in problem["weightmap"].items():
                name = get_display_name_from_pid(problems, pid)
                new_map[name] = num
            problem["weightmap"] = new_map

    for problem in problems:
        if "desc" not in problem:
            problem["desc"] = "I'm bad."

        for key in list(problem.keys()):

            if key in migration_key:
                problem[migration_key[key]] = problem[key]

            if key in migration_overwrites:
                problem[key] = migration_overwrites[key]

            if key in migration_key or key in deletion_key:
                problem.pop(key, None)
        output += json_util.dumps(problem) + "\n"
    output_file.write(output)

def build_autogen(instances):
    problems = api.problem.get_all_problems(show_disabled=True)
    for problem in problems:
        if problem.get("autogen", False):
            api.autogen.build_problem_instances(problem["pid"], instances)

def list_problems():
    #TODO: This could be improved
    problems = api.problem.get_all_problems(show_disabled=True)
    for problem in problems:
        print("{} ({}) - {} points".format(problem["name"], "disabled" if problem["disabled"] else "enabled", problem["score"]))

def drop_collections(collections):
    db = api.common.get_conn()
    for collection in collections:
        db[collection].remove()
        print("Dropped -> {}".format(collection))

def get_output_file(output):
    if output == sys.stdout:
        return output
    else:
        try:
            return open(output, "w")
        except IOError as error:
            print(error)
            exit(1)

def main():
    parser = argparse.ArgumentParser(description='picoCTF problem manager')

    parser.add_argument("-l", action="store_true", dest="show_list", help="View problem list")
    #TODO: Implement this?
    #parser.add_argument("-f", action="append", default=[], dest="filters", help="Key:value pairs that are used to search the database for problems. Used in conjuction with -l.")
    parser.add_argument("-m", action="store_true", dest="migrate", help="Migrate old 2013 problems to new format.", default=False)
    parser.add_argument("-b", "--build-autogen", action="store", type=int, help="Generate a specified amount of instances for a given list of problems.", default=0)
    parser.add_argument("-d", action="store_true", dest="debug", help="Debug mode", default=False)
    parser.add_argument("-o", action="store", dest="output_file", help="Output file.", default=sys.stdout)
    parser.add_argument("--no-confirm", action="store_true", dest="no_confirm", help="Remove confirmation and assume default action.")
    parser.add_argument("--drop-collections", dest="collections_string", action="store", help="Remove all entries from the listed collections in the database.")

    parser.add_argument("files", nargs="*", help="Files containing problems to insert.")

    args = parser.parse_args()

    files = {}

    #Check that all listed files exist.
    for file_path in args.files:
        if not path.isfile(file_path):
            print("{}: File does not exist!".format(file_path))
            exit(1)
        with open(file_path, "r") as f:
            files[file_path] = f.readlines()

    output_file = get_output_file(args.output_file)

    if args.collections_string:
        collections = map(lambda s: s.strip(), args.collections_string.split(","))
        drop_collections(collections)
    if args.build_autogen > 0:
        build_autogen(args.build_autogen)
    if args.show_list:
        list_problems()
    if args.migrate:
        migrate_problems(files, output_file, args.debug)
    else:
        insert_problems(files)
main()
<commit_msg>Added short flag for dropping collections, -c<commit_after>#!/usr/bin/python3
"""
Problem management script.
"""

import argparse
import sys
import imp
import random
import shutil
import os

import api
from api.common import APIException, InternalException

from os import path
from bson import json_util

def insert_problems(files):
    problems = get_problems(files)
    for problem in problems:
        try:
            api.problem.insert_problem(problem)
        except APIException as error:
            raise
            exit(1)
    errors = api.problem.analyze_problems()
    for error in errors:
        print(error)

def get_problems(files):
    problems = []
    for contents in files.values():
        for line in contents:
            problem = json_util.loads(line.strip())
            problems.append(problem)
    return problems

def migrate_problems(files, output_file, debug):

    migration_key = {
        "displayname": "name",
        "basescore": "score",
        "desc": "description",
        "relatedproblems": "related_problems"
    }

    migration_overwrites = {
        "grader": "test.py",
        "autogen": False
    }

    deletion_key = ["_id", "pid", "generator", "submissiontype", "devnotes"]

    problems = get_problems(files)
    output = ""

    def get_display_name_from_pid(problems, pid):
        for problem in problems:
            if problem.get("pid") == pid:
                return problem.get("displayname")

    for problem in problems:
        if problem.get("weightmap"):
            new_map = {}
            for pid, num in problem["weightmap"].items():
                name = get_display_name_from_pid(problems, pid)
                new_map[name] = num
            problem["weightmap"] = new_map

    for problem in problems:
        if "desc" not in problem:
            problem["desc"] = "I'm bad."

        for key in list(problem.keys()):

            if key in migration_key:
                problem[migration_key[key]] = problem[key]

            if key in migration_overwrites:
                problem[key] = migration_overwrites[key]

            if key in migration_key or key in deletion_key:
                problem.pop(key, None)
        output += json_util.dumps(problem) + "\n"
    output_file.write(output)

def build_autogen(instances):
    problems = api.problem.get_all_problems(show_disabled=True)
    for problem in problems:
        if problem.get("autogen", False):
            api.autogen.build_problem_instances(problem["pid"], instances)

def list_problems():
    #TODO: This could be improved
    problems = api.problem.get_all_problems(show_disabled=True)
    for problem in problems:
        print("{} ({}) - {} points".format(problem["name"], "disabled" if problem["disabled"] else "enabled", problem["score"]))

def drop_collections(collections):
    db = api.common.get_conn()
    for collection in collections:
        db[collection].remove()
        print("Dropped -> {}".format(collection))

def get_output_file(output):
    if output == sys.stdout:
        return output
    else:
        try:
            return open(output, "w")
        except IOError as error:
            print(error)
            exit(1)

def main():
    parser = argparse.ArgumentParser(description='picoCTF problem manager')

    parser.add_argument("-l", action="store_true", dest="show_list", help="View problem list")
    #TODO: Implement this?
    #parser.add_argument("-f", action="append", default=[], dest="filters", help="Key:value pairs that are used to search the database for problems. Used in conjuction with -l.")
    parser.add_argument("-m", action="store_true", dest="migrate", help="Migrate old 2013 problems to new format.", default=False)
    parser.add_argument("-b", "--build-autogen", action="store", type=int, help="Generate a specified amount of instances for a given list of problems.", default=0)
    parser.add_argument("-d", action="store_true", dest="debug", help="Debug mode", default=False)
    parser.add_argument("-o", action="store", dest="output_file", help="Output file.", default=sys.stdout)
    parser.add_argument("--no-confirm", action="store_true", dest="no_confirm", help="Remove confirmation and assume default action.")
    parser.add_argument("-c", "--drop-collections", dest="collections_string", action="store", help="Remove all entries from the listed collections in the database.")

    parser.add_argument("files", nargs="*", help="Files containing problems to insert.")

    args = parser.parse_args()

    files = {}

    #Check that all listed files exist.
    for file_path in args.files:
        if not path.isfile(file_path):
            print("{}: File does not exist!".format(file_path))
            exit(1)
        with open(file_path, "r") as f:
            files[file_path] = f.readlines()

    output_file = get_output_file(args.output_file)

    if args.collections_string:
        collections = map(lambda s: s.strip(), args.collections_string.split(","))
        drop_collections(collections)
    if args.build_autogen > 0:
        build_autogen(args.build_autogen)
    if args.show_list:
        list_problems()
    if args.migrate:
        migrate_problems(files, output_file, args.debug)
    else:
        insert_problems(files)
main()
<|endoftext|>


<commit_before>import re
import os
import logging
import time
import hashlib
from malware import utils
from malware.snort import SnortRule
from malware import pickle_tool
from malware import apikey
from virus_total_apis import PrivateApi as VirusTotal

logger = logging.getLogger(__name__)

request_times = 0
init_timer = time.time()

def list_pcap(path):
    for dirPath, dirNames, fileNames in os.walk(path):
        for f in fileNames:
            if f.split('.')[1] == 'pcap':
                yield os.path.join(dirPath, f)

def _check_timer_counter():
    global request_times
    global init_timer
    if request_times == 300:
        request_times = 0
        period = time.time() - init_timer
        waiting = 60 - period + 1
        if waiting > 0:
            logger.info("Waiting %s seconds", (str(waiting)))
            time.sleep(waiting)
        init_timer = time.time()

def is_valid_url(url):
    import re
    regex = re.compile(
        # r'^(?:[a-z0-9\.\-]*)://'  # scheme is validated separately
        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}(?<!-)\.?)|'  # domain...
        r'localhost|'  # localhost...
        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}|'  # ...or ipv4
        r'\[?[A-F0-9]*:[A-F0-9:]+\]?)'  # ...or ipv6
        r'(?::\d+)?'  # optional port
        r'(?:/?|[/?]\S+)$', re.IGNORECASE)
    return url is not None and regex.search(url)

def get_url_positive(resource, cache={}):
    m = hashlib.sha1(resource)
    urlkey = m.hexdigest()
    if urlkey in cache.keys():
        # logger.info("%s in cache" % resource)
        positives = cache.get(urlkey)[1]
        # logger.info('%s in cache positives: %s' % (resource, positives))
        return positives
    else:
        global request_times
        global init_timer
        request_times += 1
        logger.info("Search on VirusTotal counter: %s", str(request_times))
        vt = VirusTotal(apikey.APIKEY_0)
        positives = vt.get_url_report(resource).get('results').get('positives')
        _check_timer_counter()

        if positives >= 0:
            cache[urlkey] = [resource, positives]
            return positives
        elif positives is None:
            logger.info("No report. Submmit the URL to VirusTotal")
            vt.scan_url(resource)
            request_times += 1
            _check_timer_counter()
            return None
        else:
            logger.debug("Get reports failed.")
            return None

def iterpayload(path):
    connection = utils.follow_tcp_stream(path)
    for conn, frame in connection.iteritems():
        for seq, content in frame.iteritems():
            if content:
                # Generate the content and 5-tuple
                yield content, conn
            else:
                # Some packets have no payload
                pass

def gen_rule(pcap_path):
    # rule_set = list()
    cache = pickle_tool.check_pickle()
    # sid = 1000000
    for content, conn in iterpayload('%s' % (pcap_path)):
        # print content, utils.connection_key_2_str(conn)
        get_method = re.search('GET (.*) ', content)
        host = re.search('Host: (.*)', content)
        if host and get_method:
            if get_method.group(1) == '/':
                url = is_valid_url(host.group(1).rstrip())
            else:
                url = is_valid_url(host.group(1).rstrip() + get_method.group(1))

            if url is not None:
                # valid_utf8 = True
                try:
                    url.group(0).decode('utf-8')
                except UnicodeDecodeError:
                    with open('invalid_utf8', 'a') as fp:
                        fp.write('{u}\n'.format(u=url.group(0)))
                    url = None
                    # valid_utf8 = False

            if url is not None:
                host_content = host.group(0).rstrip()
                uricontent = get_method.group(1)
                pos = get_url_positive(url.group(0), cache)

                if pos > 0:
                    if uricontent == '/':
                        uricontent = None
                    rule = make_uri_rule(host_content, uricontent, conn[3])
                    with open('uricontent.rules', 'a') as fp:
                        fp.write('{r}\n'.format(r=str(rule)))
                    yield rule
                else:
                    # positives == 0 or positives == None
                    pass
            else:
                # invalid_url
                pass

        else:
            # print "no URL"
            pass
            # print repr(content)
    # for r in rule_set:
    #     print r
    # return rule_set
    pickle_tool.update_pickle(cache)

def make_host_rule(content, dst_port, sid=0):
    rule = SnortRule()
    pattern = dict()
    pattern['msg'] = '"Trojan.Gen"'
    pattern['content'] = ['"{host}"'.format(host=content), 'nocase']
    # pattern['sid'] = sid
    pattern['dst_port'] = dst_port
    rule.set_malicious_pattern(**pattern)
    return rule

def make_uri_rule(content, uricontent, dst_port, sid=0):
    rule = SnortRule()
    pattern = dict()
    pattern['msg'] = '"Trojan.Gen.uricontent"'
    pattern['content'] = ['"{host}"'.format(host=content), 'nocase']
    pattern['uricontent'] = ['"{uri}"'.format(uri=uricontent), 'nocase']
    # pattern['sid'] = sid
    pattern['dst_port'] = dst_port
    rule.set_malicious_pattern(**pattern)
    return rule

def main():
    logging.basicConfig(level=logging.INFO,
                        format='[%(levelname)s] %(message)s',)

    rule_instances = list()
    RULE = list()

    for pcap in list_pcap('./PCAPLog/'):
        for rule in gen_rule(pcap):
            rule_instances.append(rule)

    for ruleobj in rule_instances:
        RULE.append(str(ruleobj))

    RULE = list(set(RULE))

    for r in RULE:
        print r

if __name__ == "__main__":
    main()
<commit_msg>make it OO<commit_after>import re
import sys
import os
import logging
import time
import hashlib
from malware import utils
from malware.snort import SnortRule
from malware import pickle_tool
from malware import apikey
from virus_total_apis import PrivateApi as VirusTotal

logger = logging.getLogger(__name__)

class Validator(object):
    def is_valid_url(self, url):
        regex = re.compile(
            # r'^(?:[a-z0-9\.\-]*)://'  # scheme is validated separately
            r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}(?<!-)\.?)|'  # domain...
            r'localhost|'  # localhost...
            r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}|'  # ...or ipv4
            r'\[?[A-F0-9]*:[A-F0-9:]+\]?)'  # ...or ipv6
            r'(?::\d+)?'  # optional port
            r'(?:/?|[/?]\S+)$', re.IGNORECASE)
        return url is not None and regex.search(url)

    def is_hsot(self, content):
        regex = re.compile('Host: (.*)')
        return content is not None and regex.search(content)

    def is_get_method(self, content):
        regex = re.compile('GET (.*) ')
        return content is not None and regex.search(content)

class RuleEngine(object):
    def __init__(self):
        self.vt_req_counter = 0
        self.vt_req_timer = time.time()
        self.vd = Validator()
        self.vt = VirusTotal(apikey.APIKEY_1)
        self.cache = {}

    def iterpcap(self, path):
        for dirPath, dirNames, fileNames in os.walk(path):
            for f in fileNames:
                if f.split('.')[1] == 'pcap':
                    # check the file is pcap file
                    yield os.path.join(dirPath, f)

    def _iterpayload(self, path):
        connection = utils.follow_tcp_stream(path)
        for conn, frame in connection.iteritems():
            for seq, content in frame.iteritems():
                if content:
                    # Generate the content and 5-tuple
                    yield content, conn
                else:
                    # Some packets have no payload
                    pass

    def _check_timer_counter(self):
        if self.vt_req_counter == 4:
            self.vt_req_counter = 0
            period = time.time() - self.vt_req_timer
            waiting = 60 - period + 1
            if waiting > 0:
                logger.info("Waiting %s seconds", (str(waiting)))
                time.sleep(waiting)
            self.vt_req_timer = time.time()

    def _make_rule(self, content, uricontent, dst_port, sid=0):
        rule = SnortRule()
        pattern = dict()
        pattern['msg'] = '"Trojan.Gen.uricontent"'
        pattern['content'] = ['"{host}"'.format(host=content), 'nocase']
        pattern['uricontent'] = ['"{uri}"'.format(uri=uricontent), 'nocase']
        # pattern['sid'] = sid
        pattern['dst_port'] = dst_port
        rule.set_malicious_pattern(**pattern)
        return rule

    def _get_url_positive(self, resource):
        m = hashlib.sha1(resource)
        urlkey = m.hexdigest()
        if urlkey in self.cache.keys():
            # logger.info("%s in cache" % resource)
            positives = self.cache.get(urlkey)[1]
            return positives
        else:
            self.vt_req_counter += 1
            logger.info("Search on VirusTotal counter: %s",
                        str(self.vt_req_counter))

            response = self.vt.get_url_report(resource)

            if response.get('error') is not None:
                logger.info("Error: {e}".format(e=response.get('error')))
                sys.exit(0)

            results = response.get('results')
            positives = results.get('positives')
            self._check_timer_counter()

            if positives >= 0:
                self.cache[urlkey] = [resource, positives]
                return positives
            elif positives is None:
                logger.info('''No report.
                            Submmit the URL to VirusTotal countert: %s''',
                            str(self.vt_req_counter))
                self.vt.scan_url(resource)
                self._check_timer_counter()
                return None
            else:
                logger.debug("Get reports failed.")
                return None

    def gen_rule(self, pcap_path):
        self.cache = pickle_tool.check_json()
        for content, conn in self._iterpayload('%s' % (pcap_path)):
            # print content, utils.connection_key_2_str(conn)
            get_method = self.vd.is_get_method(content)
            host = self.vd.is_hsot(content)
            if host and get_method:
                if get_method.group(1) == '/':
                    url = self.vd.is_valid_url(host.group(1).rstrip())
                else:
                    url = self.vd.is_valid_url(host.group(1).rstrip() +
                                               get_method.group(1))

                if url is not None:
                    # valid_utf8 = True
                    try:
                        url.group(0).decode('utf-8')
                    except UnicodeDecodeError:
                        with open('invalid_utf8', 'a') as fp:
                            fp.write('{u}\n'.format(u=url.group(0)))
                        url = None
                        # valid_utf8 = False

                if url is not None:
                    host_content = host.group(0).rstrip()
                    uricontent = get_method.group(1)
                    pos = self._get_url_positive(url.group(0))

                    if pos > 0:
                        if uricontent == '/':
                            uricontent = None
                        #print host_content
                        rule = self._make_rule(host_content, uricontent,
                                               conn[3], 0)

                        with open('uricontent.rules', 'a') as fp:
                            fp.write('{r}\n'.format(r=str(rule)))
                        yield rule
                    else:
                        # positives == 0 or positives == None
                        pass
                else:
                    # invalid_url
                    pass
            else:
                pass

        pickle_tool.update_json(self.cache)

class PayloadIterator2(object):
    def __init__(self, path):
        self.index = 0
        self.path = path
        self.pcap_list = list()
        self.content = list()
        self.five_tuple = list()

    def __iter__(self):
        for dirPath, dirNames, fileNames in os.walk(self.path):
            for f in fileNames:
                if f.split('.')[1] == 'pcap':
                    self.pcap_list.append(os.path.join(dirPath, f))
                else:
                    # Not a pcap file
                    pass

        for p in self.pcap_list:
            connection = utils.follow_tcp_stream(p)
            for five_tuple, frame in connection.iteritems():
                for seq, content in frame.iteritems():
                    if content:
                        # Generate the content and 5-tuple
                        self.content.append(content)
                        self.five_tuple.append(five_tuple)
                    else:
                        # Some packets have no payload
                        pass
        return self

    def next(self):
        try:
            five_tuple = self.five_tuple[self.index]
            content = self.content[self.index]
        except IndexError:
            raise StopIteration
        self.index += 1
        return content, five_tuple

def main():
    logging.basicConfig(level=logging.INFO,
                        format='[%(levelname)s] %(message)s',)

    rule_engine = RuleEngine()
    # print dir(rule_engine)
    rule_instances = list()
    RULE = list()

    for pcap in rule_engine.iterpcap('./PCAPLog/'):
        # print pcap
        for rule in rule_engine.gen_rule(pcap):
            rule_instances.append(rule)

    for ruleobj in rule_instances:
        RULE.append(str(ruleobj))

    RULE = list(set(RULE))

    for r in RULE:
        print r

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


<commit_before>#!/usr/bin/env python3

# ----------------------------------------------------------------------------
# Copyright (c) 2013--, Qiyun Zhu and Katharina Dittmar.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

from setuptools import setup, find_packages
from glob import glob

classes = """
    Development Status :: 4 - Beta
    License :: OSI Approved :: BSD License
    Topic :: Scientific/Engineering
    Topic :: Scientific/Engineering :: Bio-Informatics
    Programming Language :: Python :: 3
    Programming Language :: Python :: 3 :: Only
    Operating System :: Unix
    Operating System :: POSIX
    Operating System :: MacOS
    Operating System :: Windows
"""

classifiers = [s.strip() for s in classes.split('\n') if s]

description = (
    'Genome-wide detection of putative horizontal gene transfer (HGT) events '
    'based on sequence homology search hit distribution statistics')

long_description = open('README.md').read()

setup(
    name='hgtector',
    version='2.0b1',
    license='BSD-3-Clause',
    description=description,
    long_description=long_description,
    long_description_content_type='text/markdown',
    author='Qiyun Zhu',
    author_email='qiyunzhu@gmail.com',
    url='https://github.com/DittmarLab/HGTector',
    packages=find_packages(),
    scripts=glob('scripts/hgtector'),
    package_data={'hgtector': ['config.yml']},
    include_package_data=True,
    install_requires=[
        'pyyaml',
        'numpy',
        'scipy',
        'pandas',
        'scikit-learn',
        'matplotlib'
    ],
    classifiers=classifiers,
    python_requires='>=3.6',
    entry_points='[console_scripts]'
)
<commit_msg>testing install 2<commit_after>#!/usr/bin/env python3

# ----------------------------------------------------------------------------
# Copyright (c) 2013--, Qiyun Zhu and Katharina Dittmar.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

from setuptools import setup, find_packages
from glob import glob

classes = """
    Development Status :: 4 - Beta
    License :: OSI Approved :: BSD License
    Topic :: Scientific/Engineering
    Topic :: Scientific/Engineering :: Bio-Informatics
    Programming Language :: Python :: 3
    Programming Language :: Python :: 3 :: Only
    Operating System :: Unix
    Operating System :: POSIX
    Operating System :: MacOS
    Operating System :: Windows
"""

classifiers = [s.strip() for s in classes.split('\n') if s]

description = (
    'Genome-wide detection of putative horizontal gene transfer (HGT) events '
    'based on sequence homology search hit distribution statistics')

long_description = open('README.md').read()

setup(
    name='hgtector',
    version='2.0b1',
    license='BSD-3-Clause',
    description=description,
    long_description=long_description,
    long_description_content_type='text/markdown',
    author='Qiyun Zhu',
    author_email='qiyunzhu@gmail.com',
    url='https://github.com/DittmarLab/HGTector',
    packages=find_packages(),
    scripts=glob('scripts/hgtector'),
    data_files={'config': ['config.yml']},
    include_package_data=True,
    install_requires=[
        'pyyaml',
        'numpy',
        'scipy',
        'pandas',
        'scikit-learn',
        'matplotlib'
    ],
    classifiers=classifiers,
    python_requires='>=3.6',
    entry_points='[console_scripts]'
)
<|endoftext|>


<commit_before>"""
Twisted framework specific code for base plugin functionality. Used by each plugin definition.

"""

import copy
import json
import os
import sys
from twisted import web
from twisted.application import service
from twisted.application.internet import TimerService, StreamServerEndpointService
from twisted.internet.endpoints import TCP4ServerEndpoint, SSL4ServerEndpoint

from twisted.internet import ssl, reactor
from twisted.internet.task import LoopingCall
from twisted.python import log
from twisted.python import usage
from twisted.web.resource import Resource
from twisted.web import wsgi, rewrite
from twisted.web import server

from anchore_engine.apis.ssl import _load_ssl_key, _load_ssl_cert
from anchore_engine.subsys import logger
from anchore_engine.configuration import localconfig
from anchore_engine.service import ApiService
from anchore_engine import utils

class CommonOptions(usage.Options):
    """
    Default Anchroe CLI options for the twistd plugins
    """

    optParameters = [
        ["config", "c", None, "Configuration directory location."],
        ["validate-responses", "r", False, "Enable response validation.", lambda x: x in ['True', 'true', 't', True]]
    ]

# simple twisted resource for health check route
class EmptyResource(Resource):
    """
    A simple resource to return empty if rendered
    """
    isLeaf = True

    def render_GET(self, request):
        return b''

# simple twisted resource for version check route
class VersionResource(Resource):
    """
    A simple resource to return version information if rendered
    """
    isLeaf = True

    def __init__(self, versions):
        super().__init__()
        self.versions = versions

    def render_GET(self, request):
        request.responseHeaders.addRawHeader(b"Content-Type", b"application/json")

        ret = {
            'service': {
                'version': self.versions.get('service_version', None),
            },
            'api': {
            },
            'db': {
                'schema_version': self.versions.get('db_version', None),
            }
        }

        try:
            response = utils.ensure_bytes(json.dumps(ret))
        except:
            response = utils.ensure_bytes(json.dumps({}))

        return response

def _load_config(config_option, validate_params=None):
    try:
        # config and init
        configfile = configdir = None
        if config_option:
            configdir = config_option
            configfile = os.path.join(config_option, 'config.yaml')

        localconfig.load_config(configdir=configdir, configfile=configfile, validate_params=validate_params)
        my_config = localconfig.get_config()
        my_config['myservices'] = []
        logger.spew("localconfig=" + json.dumps(my_config, indent=4, sort_keys=True))
        return my_config
    except Exception as err:
        logger.error("cannot load configuration: exception - " + str(err))
        raise err

def _validate_options(options):
    if 'config' not in options:
        raise Exception('Invalid startup options. "config" must be specified')
    return options

class WsgiApiServiceMaker(object):
    """
    A service maker that builds twistd api handlers as well. Includes the /health resource by default

    """
    service_cls = None
    tapname = None  # e.g. "anchore-api"
    description = None  # e.g. "Anchore Service"
    options = CommonOptions

    # The base child paths served by this service (e.g. /health, /v1/...). A list of web.resource.Resource() objects
    _default_resource_nodes = {
        b'health': EmptyResource(),
        b'version': VersionResource({})
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.global_configuration = None
        self.service_config = None
        self.anchore_service = None
        self.root_resource = None
        self.twistd_service = None
        self.resource_nodes = copy.deepcopy(self._default_resource_nodes)

    def _init_logging(self):
        if self.global_configuration is None:
            log.err('No configuration found to initialize logging for. Expecting other errors, so setting log level to DEBUG')
            log_level = 'DEBUG'
            log_to_db = False
        else:
            try:
                service_config = self.global_configuration['services'][self.service_cls.__service_name__]
                log_level = service_config.get('log_level', self.global_configuration.get('log_level', 'INFO'))
                log_to_db = self.global_configuration.get('log_to_db', False)
            except Exception as err:
                log.err("error checking for enabled services, check config file - exception: " + str(err))
                raise Exception("error checking for enabled services, check config file - exception: " + str(err))

        logger.set_log_level(log_level, log_to_db=log_to_db)

    def _check_enabled(self):
        if not self.global_configuration.get('services', {}).get(self.service_cls.__service_name__, {}).get('enabled', False):
            log.err("Service {} not enabled in configuration file: shutting down".format(self.service_cls.__service_name__))
            sys.exit(0)

    def _init_config(self, options):
        _validate_options(options)
        self.global_configuration = _load_config(options['config'])
        self.service_config = self.global_configuration.get('services', {}).get(self.service_cls.__service_name__, {})

    def _get_non_api_monitor(self, service):
        return service.get_monitor_thread(monitor_thread_wrapper=lambda target, kwargs: TimerService(1, target, **kwargs))

    def _get_api_monitor(self, service):
        return service.get_monitor_thread(monitor_thread_wrapper=LoopingCall)

    def makeService(self, options):

        try:
            logger.info('Initializing configuration')
            self._init_config(options)

            logger.info('Initializing logging')
            self._init_logging()

            self._check_enabled()

            #logger.enable_bootstrap_logging(self.tapname)

            assert (issubclass(self.service_cls, ApiService))
            self.anchore_service = self.service_cls(options=options)
            self.anchore_service.initialize(self.global_configuration)
            self.resource_nodes[b'version'].versions = self.anchore_service.versions

            # application object
            application = service.Application("Service-" + '-'.join(self.anchore_service.name))
            self.twistd_service = service.MultiService()
            self.twistd_service.setServiceParent(application)

            if not os.environ.get('ANCHORE_ENGINE_DISABLE_MONITORS'):
                logger.info('Starting monitor thread')
                lc = self._get_api_monitor(self.anchore_service)
                lc.start(1)
            else:
                logger.warn('Skipped start of monitor threads due to ANCHORE_ENGINE_DISABLE_MONITORS set in environment')

            thread_stats_interval = int(self.service_config.get('debug_thread_stats_dump_interval', 0))
            if thread_stats_interval > 0:
                logger.info('Based on service config, starting the thread stats dumper')
                monitor = LoopingCall(dump_stats)
                monitor.start(thread_stats_interval)

            logger.info('Building api handlers')
            s = self._build_api_service()
            s.setServiceParent(self.twistd_service)

            return self.twistd_service

        except Exception as err:
            logger.exception("cannot create/init/register service: " + self.service_cls.__service_name__ + " - exception: " + str(err))
            raise Exception("cannot start service (see above for information)")
        finally:
            pass
            #logger.disable_bootstrap_logging()

    def _add_resource(self, name, resource):
        """
        Add a resource to this resource as a child
        :param resource: Resource subclass object
        :param name: path name for the resource (as bytes object)
        :return:
        """

        self.resource_nodes[name] = resource

    def _build_api_service(self):
        """
        Once called, the resource is initialized. Any calls to self._add_resource() should be done before calling this fn.
        :return:
        """

        thread_count = int(self.service_config.get('max_request_threads', localconfig.DEFAULT_SERVICE_THREAD_COUNT))

        wsgi_app = self.anchore_service.get_api_application()
        wsgi_site = wsgi.WSGIResource(reactor, reactor.getThreadPool(), application=wsgi_app)
        reactor.getThreadPool().adjustPoolsize(maxthreads=thread_count)
        logger.debug('Thread pool size stats. Min={}, Max={}'.format(reactor.getThreadPool().min, reactor.getThreadPool().max))

        self._add_resource(self.anchore_service.__service_api_version__.encode('utf-8'), wsgi_site)
        self.root_resource = web.resource.Resource()

        # Add nodes
        for name, resource in self.resource_nodes.items():
            self.root_resource.putChild(name, resource)

        # this will rewrite any calls that do not have an explicit version to the base path before being processed by flask
        self._api_version_bytes = self.anchore_service.__service_api_version__.encode('utf-8') # This is optimization

        # Handle the auth vs non-auth child resources to not consume a path element
        root = rewrite.RewriterResource(self.root_resource, self._default_version_rewrite)

        # Build the main site server
        site = server.Site(root)
        listen = self.anchore_service.configuration['listen']

        # Disable the twisted access logging by overriding the log function as it uses a raw 'write' and cannot otherwise be disabled, iff enable_access_logging is set to False in either the service or global config
        try:
            eal = True
            if "enable_access_logging" in self.anchore_service.configuration:
                eal = self.anchore_service.configuration.get("enable_access_logging", True)
            elif "enable_access_logging" in self.configuration:
                eal = self.configuration.get("enable_access_logging", True)

            if not eal:
                def _null_logger(request):
                    pass
                site.log = _null_logger

        except:
            pass

        if str(self.anchore_service.configuration.get('ssl_enable', '')).lower() == 'true':
            try:
                ssl_data = {
                    'ssl_cert': _load_ssl_cert(self.anchore_service.configuration['ssl_cert']) if 'ssl_cert' in self.anchore_service.configuration else None,
                    'ssl_chain': _load_ssl_cert(self.anchore_service.configuration['ssl_chain']) if 'ssl_chain' in self.anchore_service.configuration else None,
                    'ssl_key': _load_ssl_key(self.anchore_service.configuration['ssl_key']) if 'ssl_key' in self.anchore_service.configuration else None
                }

                if ssl_data['ssl_chain']:
                    sfact = ssl.CertificateOptions(privateKey=ssl_data['ssl_key'], certificate=ssl_data['ssl_cert'],
                                                   extraCertChain=[ssl_data['ssl_chain']])
                else:
                    sfact = ssl.CertificateOptions(privateKey=ssl_data['ssl_key'], certificate=ssl_data['ssl_cert'])

                endpoint = SSL4ServerEndpoint(reactor=reactor, port=int(self.anchore_service.configuration['port']), sslContextFactory=sfact, interface=listen)
            except Exception as err:
                raise err
        else:
            endpoint = TCP4ServerEndpoint(reactor=reactor, port=int(self.anchore_service.configuration['port']), interface=listen)

        ret_svc = StreamServerEndpointService(endpoint=endpoint, factory=site)
        ret_svc.setName(self.anchore_service.name)

        return ret_svc

    def _default_version_rewrite(self, request):
        try:
            if request.postpath:
                #if request.postpath[0] != b'health' and request.postpath[0] != self._api_version_bytes:
                if request.postpath[0] not in [b'health', b'version'] and request.postpath[0] != self._api_version_bytes:
                    request.postpath.insert(0, self._api_version_bytes)
                    request.path = b'/' + self._api_version_bytes + request.path
        except Exception as err:
            logger.error("rewrite exception: " + str(err))
            raise err

def dump_stats():
    """
    Dump some basic stats about the reactor pool and threads at info level
    :return:
    """

    logger.info('Reactor queue stats: {}'.format(reactor.getThreadPool()._team.statistics().__dict__))
<commit_msg>Adds some debugging tools to the service stack enabled by environment variables. These are dangerous, not for production use.<commit_after>"""
Twisted framework specific code for base plugin functionality. Used by each plugin definition.

"""

import copy
import json
import os
import sys
import datetime
from twisted import web
from twisted.application import service
from twisted.application.internet import TimerService, StreamServerEndpointService
from twisted.internet.endpoints import TCP4ServerEndpoint, SSL4ServerEndpoint

from twisted.internet import ssl, reactor
from twisted.internet.task import LoopingCall
from twisted.python import log
from twisted.python import usage
from twisted.web.resource import Resource
from twisted.web import wsgi, rewrite
from twisted.web import server

from anchore_engine.apis.ssl import _load_ssl_key, _load_ssl_cert
from anchore_engine.subsys import logger
from anchore_engine.configuration import localconfig
from anchore_engine.service import ApiService
from anchore_engine import utils
import faulthandler

# For the debug CLI, require a code modification to enable it. This allows on-host edits of the script and restart, but no accidental config from env vars or config.
enable_dangerous_debug_cli = False

# Thread dumper is safer since read-only and only on local-host, so allow it to configure from env var
enable_thread_dumper = (os.getenv('ANCHORE_ENABLE_DANGEROUS_THREAD_DUMP_API', 'false').lower() == 'true')

if enable_dangerous_debug_cli or enable_thread_dumper:
    from twisted.application import internet, service
    from twisted.conch.insults import insults
    from twisted.conch.manhole import ColoredManhole
    from twisted.conch.telnet import TelnetTransport, TelnetBootstrapProtocol
    from twisted.internet import protocol

class CommonOptions(usage.Options):
    """
    Default Anchroe CLI options for the twistd plugins
    """

    optParameters = [
        ["config", "c", None, "Configuration directory location."],
        ["validate-responses", "r", False, "Enable response validation.", lambda x: x in ['True', 'true', 't', True]]
    ]

# simple twisted resource for health check route
class EmptyResource(Resource):
    """
    A simple resource to return empty if rendered
    """
    isLeaf = True

    def render_GET(self, request):
        return b''

class ThreadDumperResource(Resource):
    isLeaf = True

    def __init__(self):
        super().__init__()
        logger.info("Initializing thread dumper resource")

    def render_GET(self, request):
        logger.info("Handling thread dump request")

        try:
            with open('/var/log/anchore/pid_{}_thread_dump-{}'.format(os.getpid(), datetime.datetime.now().isoformat()), 'w') as dest:
                faulthandler.dump_traceback(dest, all_threads=True)
        except:
            logger.exception('Error dumping thread frames')
            return b'Failed'

        return b'Sucess'

# simple twisted resource for version check route
class VersionResource(Resource):
    """
    A simple resource to return version information if rendered
    """
    isLeaf = True

    def __init__(self, versions):
        super().__init__()
        self.versions = versions

    def render_GET(self, request):
        request.responseHeaders.addRawHeader(b"Content-Type", b"application/json")

        ret = {
            'service': {
                'version': self.versions.get('service_version', None),
            },
            'api': {
            },
            'db': {
                'schema_version': self.versions.get('db_version', None),
            }
        }

        try:
            response = utils.ensure_bytes(json.dumps(ret))
        except:
            response = utils.ensure_bytes(json.dumps({}))

        return response

def _load_config(config_option, validate_params=None):
    try:
        # config and init
        configfile = configdir = None
        if config_option:
            configdir = config_option
            configfile = os.path.join(config_option, 'config.yaml')

        localconfig.load_config(configdir=configdir, configfile=configfile, validate_params=validate_params)
        my_config = localconfig.get_config()
        my_config['myservices'] = []
        logger.spew("localconfig=" + json.dumps(my_config, indent=4, sort_keys=True))
        return my_config
    except Exception as err:
        logger.error("cannot load configuration: exception - " + str(err))
        raise err

def _validate_options(options):
    if 'config' not in options:
        raise Exception('Invalid startup options. "config" must be specified')
    return options

class WsgiApiServiceMaker(object):
    """
    A service maker that builds twistd api handlers as well. Includes the /health resource by default

    """
    service_cls = None
    tapname = None  # e.g. "anchore-api"
    description = None  # e.g. "Anchore Service"
    options = CommonOptions

    # The base child paths served by this service (e.g. /health, /v1/...). A list of web.resource.Resource() objects
    _default_resource_nodes = {
        b'health': EmptyResource(),
        b'version': VersionResource({})
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.global_configuration = None
        self.service_config = None
        self.anchore_service = None
        self.root_resource = None
        self.twistd_service = None
        self.resource_nodes = copy.deepcopy(self._default_resource_nodes)

    def _init_logging(self):
        if self.global_configuration is None:
            log.err('No configuration found to initialize logging for. Expecting other errors, so setting log level to DEBUG')
            log_level = 'DEBUG'
            log_to_db = False
        else:
            try:
                service_config = self.global_configuration['services'][self.service_cls.__service_name__]
                log_level = service_config.get('log_level', self.global_configuration.get('log_level', 'INFO'))
                log_to_db = self.global_configuration.get('log_to_db', False)
            except Exception as err:
                log.err("error checking for enabled services, check config file - exception: " + str(err))
                raise Exception("error checking for enabled services, check config file - exception: " + str(err))

        logger.set_log_level(log_level, log_to_db=log_to_db)

    def _check_enabled(self):
        if not self.global_configuration.get('services', {}).get(self.service_cls.__service_name__, {}).get('enabled', False):
            log.err("Service {} not enabled in configuration file: shutting down".format(self.service_cls.__service_name__))
            sys.exit(0)

    def _init_config(self, options):
        _validate_options(options)
        self.global_configuration = _load_config(options['config'])
        self.service_config = self.global_configuration.get('services', {}).get(self.service_cls.__service_name__, {})

    def _get_non_api_monitor(self, service):
        return service.get_monitor_thread(monitor_thread_wrapper=lambda target, kwargs: TimerService(1, target, **kwargs))

    def _get_api_monitor(self, service):
        return service.get_monitor_thread(monitor_thread_wrapper=LoopingCall)

    def makeDebugCLIService(self, args):
        """
        This is dangerous, and should only ever be enabled by explicit user config and only for non-production use

        :param args:
        :return:
        """

        f = protocol.ServerFactory()
        f.protocol = lambda: TelnetTransport(TelnetBootstrapProtocol,
                                             insults.ServerProtocol,
                                             args['protocolFactory'],
                                             *args.get('protocolArgs', ()),
                                             **args.get('protocolKwArgs', {}))
        return internet.TCPServer(args['telnet'], f)

    def makeService(self, options):

        try:
            logger.info('Initializing configuration')
            self._init_config(options)

            logger.info('Initializing logging')
            self._init_logging()

            self._check_enabled()

            #logger.enable_bootstrap_logging(self.tapname)

            assert (issubclass(self.service_cls, ApiService))
            self.anchore_service = self.service_cls(options=options)
            self.anchore_service.initialize(self.global_configuration)
            self.resource_nodes[b'version'].versions = self.anchore_service.versions

            # application object
            application = service.Application("Service-" + '-'.join(self.anchore_service.name))
            self.twistd_service = service.MultiService()
            self.twistd_service.setServiceParent(application)

            if not os.environ.get('ANCHORE_ENGINE_DISABLE_MONITORS'):
                logger.info('Starting monitor thread')
                lc = self._get_api_monitor(self.anchore_service)
                lc.start(1)
            else:
                logger.warn('Skipped start of monitor threads due to ANCHORE_ENGINE_DISABLE_MONITORS set in environment')

            thread_stats_interval = int(self.service_config.get('debug_thread_stats_dump_interval', 0))
            if thread_stats_interval > 0:
                logger.info('Based on service config, starting the thread stats dumper')
                monitor = LoopingCall(dump_stats)
                monitor.start(thread_stats_interval)

            logger.info('Building api handlers')
            s = self._build_api_service()
            s.setServiceParent(self.twistd_service)

            if enable_dangerous_debug_cli:
                logger.warn('Loading *dangerous* debug/telnet service as specified by debug config')
                self.makeDebugCLIService({'protocolFactory': ColoredManhole,
                                  'protocolArgs': (None,),
                                  'telnet': 6023,
                                  }).setServiceParent(self.twistd_service)

            return self.twistd_service

        except Exception as err:
            logger.exception("cannot create/init/register service: " + self.service_cls.__service_name__ + " - exception: " + str(err))
            raise Exception("cannot start service (see above for information)")
        finally:
            pass
            #logger.disable_bootstrap_logging()

    def _add_resource(self, name, resource):
        """
        Add a resource to this resource as a child
        :param resource: Resource subclass object
        :param name: path name for the resource (as bytes object)
        :return:
        """

        self.resource_nodes[name] = resource

    def _build_api_service(self):
        """
        Once called, the resource is initialized. Any calls to self._add_resource() should be done before calling this fn.
        :return:
        """

        thread_count = int(self.service_config.get('max_request_threads', localconfig.DEFAULT_SERVICE_THREAD_COUNT))

        wsgi_app = self.anchore_service.get_api_application()
        wsgi_site = wsgi.WSGIResource(reactor, reactor.getThreadPool(), application=wsgi_app)
        reactor.getThreadPool().adjustPoolsize(maxthreads=thread_count)
        logger.debug('Thread pool size stats. Min={}, Max={}'.format(reactor.getThreadPool().min, reactor.getThreadPool().max))

        self._add_resource(self.anchore_service.__service_api_version__.encode('utf-8'), wsgi_site)

        if enable_thread_dumper:
            logger.warn("Adding thread dump route for debugging since debug flag is set. This is dangerous and should not be done in normal production")
            self._add_resource(b'threads', ThreadDumperResource())

        self.root_resource = web.resource.Resource()

        # Add nodes
        for name, resource in self.resource_nodes.items():
            self.root_resource.putChild(name, resource)

        # this will rewrite any calls that do not have an explicit version to the base path before being processed by flask
        self._api_version_bytes = self.anchore_service.__service_api_version__.encode('utf-8') # This is optimization

        # Handle the auth vs non-auth child resources to not consume a path element
        root = rewrite.RewriterResource(self.root_resource, self._default_version_rewrite)

        # Build the main site server
        site = server.Site(root)
        listen = self.anchore_service.configuration['listen']

        # Disable the twisted access logging by overriding the log function as it uses a raw 'write' and cannot otherwise be disabled, iff enable_access_logging is set to False in either the service or global config
        try:
            eal = True
            if "enable_access_logging" in self.anchore_service.configuration:
                eal = self.anchore_service.configuration.get("enable_access_logging", True)
            elif "enable_access_logging" in self.configuration:
                eal = self.configuration.get("enable_access_logging", True)

            if not eal:
                def _null_logger(request):
                    pass
                site.log = _null_logger

        except:
            pass

        if str(self.anchore_service.configuration.get('ssl_enable', '')).lower() == 'true':
            try:
                ssl_data = {
                    'ssl_cert': _load_ssl_cert(self.anchore_service.configuration['ssl_cert']) if 'ssl_cert' in self.anchore_service.configuration else None,
                    'ssl_chain': _load_ssl_cert(self.anchore_service.configuration['ssl_chain']) if 'ssl_chain' in self.anchore_service.configuration else None,
                    'ssl_key': _load_ssl_key(self.anchore_service.configuration['ssl_key']) if 'ssl_key' in self.anchore_service.configuration else None
                }

                if ssl_data['ssl_chain']:
                    sfact = ssl.CertificateOptions(privateKey=ssl_data['ssl_key'], certificate=ssl_data['ssl_cert'],
                                                   extraCertChain=[ssl_data['ssl_chain']])
                else:
                    sfact = ssl.CertificateOptions(privateKey=ssl_data['ssl_key'], certificate=ssl_data['ssl_cert'])

                endpoint = SSL4ServerEndpoint(reactor=reactor, port=int(self.anchore_service.configuration['port']), sslContextFactory=sfact, interface=listen)
            except Exception as err:
                raise err
        else:
            endpoint = TCP4ServerEndpoint(reactor=reactor, port=int(self.anchore_service.configuration['port']), interface=listen)

        ret_svc = StreamServerEndpointService(endpoint=endpoint, factory=site)
        ret_svc.setName(self.anchore_service.name)

        return ret_svc

    def _default_version_rewrite(self, request):
        try:
            if request.postpath:
                #if request.postpath[0] != b'health' and request.postpath[0] != self._api_version_bytes:
                if request.postpath[0] not in self.resource_nodes.keys() and request.postpath[0] != self._api_version_bytes:
                    request.postpath.insert(0, self._api_version_bytes)
                    request.path = b'/' + self._api_version_bytes + request.path
        except Exception as err:
            logger.error("rewrite exception: " + str(err))
            raise err

def dump_stats():
    """
    Dump some basic stats about the reactor pool and threads at info level
    :return:
    """

    logger.info('Reactor queue stats: {}'.format(reactor.getThreadPool()._team.statistics().__dict__))
<|endoftext|>