<commit_before>import socket
from geoip import geolite2

class BearStorage():
    def __init__(self, ip, rawrequest, timestamp, parsed_request, isDetected,
                 hostname):
        self.ip = ip
        self.rawrequest = rawrequest
        self.timestamp = timestamp
        self.path = ""
        self.command = ""
        self.version = ""
        if hasattr(parsed_request, 'path'):
            self.path = parsed_request.path
        if parsed_request.command is not None:
            self.command = parsed_request.command
        if hasattr(parsed_request, 'request_version'):
            self.version = parsed_request.request_version
        if hasattr(parsed_request, 'headers'):
            self.headers = parsed_request.headers
            if 'User-Agent' in parsed_request.headers.keys():
                self.ua = parsed_request.headers['User-Agent']
            else:
                self.ua = ""
        else:
            self.headers = ""
            self.ua = ""
            self.version = ""
        self.isDetected = isDetected
        self.hostname = hostname
        location = geolite2.lookup(ip)
        if location is not None:
            self.country = location.country
            self.continent = location.continent
            self.timezone = location.timezone
        else:
            self.country = ''
            self.continent = ''
            self.timezone = ''
        self.tracert = ""  # TODO
        try:
            self.dnsname = socket.gethostbyaddr(ip)[0]
        except:
            self.dnsname = ""

    def __str__(self):
        if self.path is not "":
            output = "hostname: " + self.hostname + "\r\n"
            output += "IP: " + self.ip + "\r\n"
            output += "timestamp: " + self.timestamp + "\r\n"
            output += "User-Agent: " + self.ua + "\r\n"
            output += "path: " + self.path + "\r\n"
            output += "command: " + self.command + "\r\n"
            output += "version: " + self.version + "\r\n"
            output += "country: " + self.country + "\r\n"
            if self.isDetected:
                output += "Detected: Yes" + "\r\n"
            else:
                output += "Detected: No" + "\r\n"
        else:
            output = "hostname: " + self.hostname + "\r\n"
            output += "IP: " + self.ip + "\r\n"
            output += "timestamp: " + self.timestamp + "\r\n"
            output += "rawrequest: " + self.rawrequest
            output += "country: " + self.country + "\r\n"
        return output

    def __repr__(self):
        return self.__str__()
<commit_msg>Fix useragent parsing in bearstorage<commit_after>import socket
from geoip import geolite2

class BearStorage():
    def __init__(self, ip, rawrequest, timestamp, parsed_request, isDetected,
                 hostname):
        self.ip = ip
        self.rawrequest = rawrequest
        self.timestamp = timestamp
        self.path = ""
        self.command = ""
        self.version = ""
        self.ua = ""
        self.headers = ""
        self.country = ""
        self.continent = ""
        self.timezone = ""
        self.dnsname = ""
        self.tracert = ""  # TODO
        if hasattr(parsed_request, 'path'):
            self.path = parsed_request.path
        if parsed_request.command is not None:
            self.command = parsed_request.command
        if hasattr(parsed_request, 'request_version'):
            self.version = parsed_request.request_version
        if hasattr(parsed_request, 'headers'):
            self.headers = parsed_request.headers
            if 'user-agent' in parsed_request.headers:
                self.ua = parsed_request.headers['user-agent']
        self.isDetected = isDetected
        self.hostname = hostname
        location = geolite2.lookup(ip)
        if location is not None:
            self.country = location.country
            self.continent = location.continent
            self.timezone = location.timezone
        try:
            self.dnsname = socket.gethostbyaddr(ip)[0]
        except:
            pass

    def __str__(self):
        if self.path is not "":
            output = "hostname: " + self.hostname + "\r\n"
            output += "IP: " + self.ip + "\r\n"
            output += "timestamp: " + self.timestamp + "\r\n"
            output += "User-Agent: " + self.ua + "\r\n"
            output += "path: " + self.path + "\r\n"
            output += "command: " + self.command + "\r\n"
            output += "version: " + self.version + "\r\n"
            output += "country: " + self.country + "\r\n"
            if self.isDetected:
                output += "Detected: Yes" + "\r\n"
            else:
                output += "Detected: No" + "\r\n"
        else:
            output = "hostname: " + self.hostname + "\r\n"
            output += "IP: " + self.ip + "\r\n"
            output += "timestamp: " + self.timestamp + "\r\n"
            output += "rawrequest: " + self.rawrequest
            output += "country: " + self.country + "\r\n"
        return output

    def __repr__(self):
        return self.__str__()
<|endoftext|>


<commit_before>import os
import yaml

from twisted.internet import reactor, threads, defer

from ooni import otime
from ooni.utils import Storage

reports = Storage()
scapyFactory = None
stateDict = None
state = Storage()

# XXX refactor this to use a database
resume_lock = defer.DeferredLock()

basic = None
cmd_line_options = None
resume_filename = None

# XXX-Twisted this is used to check if we have started the reactor or not. It
# is necessary because if the tests are already concluded because we have
# resumed a test session then it will call reactor.run() even though there is
# no condition that will ever stop it.
# There should be a more twisted way of doing this.
start_reactor = True

tor_state = None
tor_control = None

config_file = None
sample_config_file = None

# This is used to store the probes IP address obtained via Tor
probe_ip = None

def get_root_path():
    this_directory = os.path.dirname(__file__)
    root = os.path.join(this_directory, '..')
    root = os.path.abspath(root)
    return root

def createConfigFile():
    """
    XXX implement me
    """
    sample_config_file = os.path.join(get_root_path(), 'ooniprobe.conf.sample')

def loadConfigFile():
    """
    This is a helper function that makes sure that the configuration attributes
    are singletons.
    """
    config_file = os.path.join(get_root_path(), 'ooniprobe.conf')
    try:
        f = open(config_file)
    except IOError:
        createConfigFile()
        raise Exception("Unable to open config file. "\
                    "Copy ooniprobe.conf.sample to ooniprobe.conf")

    config_file_contents = '\n'.join(f.readlines())
    configuration = yaml.safe_load(config_file_contents)

    # Process the basic configuration options
    basic = Storage()
    for k, v in configuration['basic'].items():
        basic[k] = v

    # Process the privacy configuration options
    privacy = Storage()
    for k, v in configuration['privacy'].items():
        privacy[k] = v

    # Process the advanced configuration options
    advanced = Storage()
    for k, v in configuration['advanced'].items():
        advanced[k] = v

    # Process the tor configuration options
    tor = Storage()
    for k, v in configuration['tor'].items():
        tor[k] = v

    return basic, privacy, advanced, tor

class TestFilenameNotSet(Exception):
    pass

def generateReportFilenames():
    try:
        test_filename = os.path.basename(cmd_line_options['test'])
    except IndexError:
        raise TestFilenameNotSet

    test_name = '.'.join(test_filename.split(".")[:-1])
    frm_str = "report_%s_"+otime.timestamp()+".%s"
    reports.yamloo = frm_str % (test_name, "yamloo")
    reports.pcap = frm_str % (test_name, "pcap")

if not basic:
    # Here we make sure that we instance the config file attributes only once
    basic, privacy, advanced, tor = loadConfigFile()

if not resume_filename:
    resume_filename = os.path.join(get_root_path(), 'ooniprobe.resume')
    try:
        with open(resume_filename) as f: pass
    except IOError as e:
        with open(resume_filename, 'w+') as f: pass

# This is used to keep track of the state of the sniffer
sniffer_running = None
<commit_msg>Properly handle empty tor configuration parameters<commit_after>import os
import yaml

from twisted.internet import reactor, threads, defer

from ooni import otime
from ooni.utils import Storage

reports = Storage()
scapyFactory = None
stateDict = None
state = Storage()

# XXX refactor this to use a database
resume_lock = defer.DeferredLock()

basic = None
cmd_line_options = None
resume_filename = None

# XXX-Twisted this is used to check if we have started the reactor or not. It
# is necessary because if the tests are already concluded because we have
# resumed a test session then it will call reactor.run() even though there is
# no condition that will ever stop it.
# There should be a more twisted way of doing this.
start_reactor = True

tor_state = None
tor_control = None

config_file = None
sample_config_file = None

# This is used to store the probes IP address obtained via Tor
probe_ip = None

def get_root_path():
    this_directory = os.path.dirname(__file__)
    root = os.path.join(this_directory, '..')
    root = os.path.abspath(root)
    return root

def createConfigFile():
    """
    XXX implement me
    """
    sample_config_file = os.path.join(get_root_path(), 'ooniprobe.conf.sample')

def loadConfigFile():
    """
    This is a helper function that makes sure that the configuration attributes
    are singletons.
    """
    config_file = os.path.join(get_root_path(), 'ooniprobe.conf')
    try:
        f = open(config_file)
    except IOError:
        createConfigFile()
        raise Exception("Unable to open config file. "\
                    "Copy ooniprobe.conf.sample to ooniprobe.conf")

    config_file_contents = '\n'.join(f.readlines())
    configuration = yaml.safe_load(config_file_contents)

    # Process the basic configuration options
    basic = Storage()
    for k, v in configuration['basic'].items():
        basic[k] = v

    # Process the privacy configuration options
    privacy = Storage()
    for k, v in configuration['privacy'].items():
        privacy[k] = v

    # Process the advanced configuration options
    advanced = Storage()
    for k, v in configuration['advanced'].items():
        advanced[k] = v

    # Process the tor configuration options
    tor = Storage()
    try:
        for k, v in configuration['tor'].items():
            tor[k] = v
    except AttributeError:
        pass

    return basic, privacy, advanced, tor

class TestFilenameNotSet(Exception):
    pass

def generateReportFilenames():
    try:
        test_filename = os.path.basename(cmd_line_options['test'])
    except IndexError:
        raise TestFilenameNotSet

    test_name = '.'.join(test_filename.split(".")[:-1])
    frm_str = "report_%s_"+otime.timestamp()+".%s"
    reports.yamloo = frm_str % (test_name, "yamloo")
    reports.pcap = frm_str % (test_name, "pcap")

if not basic:
    # Here we make sure that we instance the config file attributes only once
    basic, privacy, advanced, tor = loadConfigFile()

if not resume_filename:
    resume_filename = os.path.join(get_root_path(), 'ooniprobe.resume')
    try:
        with open(resume_filename) as f: pass
    except IOError as e:
        with open(resume_filename, 'w+') as f: pass

# This is used to keep track of the state of the sniffer
sniffer_running = None
<|endoftext|>


<commit_before>__author__ = 'Cedric Da Costa Faro'

from datetime import datetime
import hashlib
from werkzeug.security import generate_password_hash, check_password_hash
from flask import request
from flask.ext.login import UserMixin
from . import db, login_manager

# We define here user table with all required fields,
# we also retrieve user's avatar from Gravatar if any,
# We make sure that user's password is encrypted
class User(UserMixin, db.Model):
    __tablename__ = 'users'
    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(64), nullable=False, unique=True, index=True)
    username = db.Column(db.String(64), nullable=False, unique=True, index=True)
    is_admin = db.Column(db.Boolean)
    password_hash = db.Column(db.String(128))
    name = db.Column(db.String(64))
    location = db.Column(db.String(64))
    bio = db.Column(db.Text())
    avatar_hash = db.Column(db.String(32))
    projects = db.relationship('Project', lazy='dynamic', backref='author')

    def __init__(self, **kwargs):
        super(User, self).__init__(**kwargs)
        if self.email is not None and self.avatar_hash is None:
            self.avatar_hash = hashlib.md5(self.email.encode('utf-8')).hexdigest()

    @property
    def password(self):
        raise AttributeError('password is not a readable attribute')

    @password.setter
    def password(self, password):
        self.password_hash = generate_password_hash(password)

    def verify_password(self, password):
        return check_password_hash(self.password_hash, password)

    def gravatar(self, size=100, default='identicon', rating='g'):
        if request.is_secure:
            url = 'htpps://secure.gravatar.com/avatar'
        else:
            url = 'http://www.gravatar.com/avatar'
        hash = self.avatar_hash or hashlib.md5(self.email.encode('utf-8')).hexdigest()
        return '{url}/{hash}?s={size}&d={default}&r={rating}'.format(
            url=url, hash=hash, size=size, default=default, rating=rating)

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

# We define here the structure of the project table and link it to a user via a foreign key
class Project(db.Model):
    __tablename__ = 'projects'
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(128), nullable=False)
    description = db.Column(db.Text)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'))
    date = db.Column(db.DateTime())
<commit_msg>typo on https for Gravatar (v0.24b)<commit_after>__author__ = 'Cedric Da Costa Faro'

from datetime import datetime
import hashlib
from werkzeug.security import generate_password_hash, check_password_hash
from flask import request
from flask.ext.login import UserMixin
from . import db, login_manager

# We define here user table with all required fields,
# we also retrieve user's avatar from Gravatar if any,
# We make sure that user's password is encrypted
class User(UserMixin, db.Model):
    __tablename__ = 'users'
    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(64), nullable=False, unique=True, index=True)
    username = db.Column(db.String(64), nullable=False, unique=True, index=True)
    is_admin = db.Column(db.Boolean)
    password_hash = db.Column(db.String(128))
    name = db.Column(db.String(64))
    location = db.Column(db.String(64))
    bio = db.Column(db.Text())
    avatar_hash = db.Column(db.String(32))
    projects = db.relationship('Project', lazy='dynamic', backref='author')

    def __init__(self, **kwargs):
        super(User, self).__init__(**kwargs)
        if self.email is not None and self.avatar_hash is None:
            self.avatar_hash = hashlib.md5(self.email.encode('utf-8')).hexdigest()

    @property
    def password(self):
        raise AttributeError('password is not a readable attribute')

    @password.setter
    def password(self, password):
        self.password_hash = generate_password_hash(password)

    def verify_password(self, password):
        return check_password_hash(self.password_hash, password)

    def gravatar(self, size=100, default='identicon', rating='g'):
        if request.is_secure:
            url = 'https://secure.gravatar.com/avatar'
        else:
            url = 'http://www.gravatar.com/avatar'
        hash = self.avatar_hash or hashlib.md5(self.email.encode('utf-8')).hexdigest()
        return '{url}/{hash}?s={size}&d={default}&r={rating}'.format(
            url=url, hash=hash, size=size, default=default, rating=rating)

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

# We define here the structure of the project table and link it to a user via a foreign key
class Project(db.Model):
    __tablename__ = 'projects'
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(128), nullable=False)
    description = db.Column(db.Text)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'))
    date = db.Column(db.DateTime())
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import os
import sys
import getpass
import logging
from openquake.baselib import sap, config, datastore
from openquake.baselib.general import safeprint
from openquake.hazardlib import valid
from openquake.commonlib import util, logs, readinput
from openquake.engine import engine as eng
from openquake.engine.export import core
from openquake.engine.utils import confirm
from openquake.engine.tools.make_html_report import make_report
from openquake.server import dbserver
from openquake.calculators.export import export
from openquake.commands.abort import abort

HAZARD_CALCULATION_ARG = "--hazard-calculation-id"
MISSING_HAZARD_MSG = "Please specify '%s=<id>'" % HAZARD_CALCULATION_ARG

def get_job_id(job_id, username=None):
    job = logs.dbcmd('get_job', job_id, username)
    if not job:
        sys.exit('Job %s not found' % job_id)
    return job.id

def run_job(job_ini, log_level='info', log_file=None, exports='',
            username=getpass.getuser(), **kw):
    """
    Run a job using the specified config file and other options.

    :param str job_ini:
        Path to calculation config (INI-style) files.
    :param str log_level:
        'debug', 'info', 'warn', 'error', or 'critical'
    :param str log_file:
        Path to log file.
    :param exports:
        A comma-separated string of export types requested by the user.
    :param username:
        Name of the user running the job
    :param kw:
        Extra parameters like hazard_calculation_id and calculation_mode
    """
    job_id = logs.init('job', getattr(logging, log_level.upper()))
    with logs.handle(job_id, log_level, log_file):
        job_ini = os.path.abspath(job_ini)
        oqparam = eng.job_from_file(job_ini, job_id, username, **kw)
        kw['username'] = username
        eng.run_calc(job_id, oqparam, exports, **kw)
        for line in logs.dbcmd('list_outputs', job_id, False):
            safeprint(line)
    return job_id

def run_tile(job_ini, sites_slice):
    """
    Used in tiling calculations
    """
    return run_job(job_ini, sites_slice=(sites_slice.start, sites_slice.stop))

def del_calculation(job_id, confirmed=False):
    """
    Delete a calculation and all associated outputs.
    """
    if logs.dbcmd('get_job', job_id) is None:
        print('There is no job %d' % job_id)
        return

    if confirmed or confirm(
            'Are you sure you want to (abort and) delete this calculation and '
            'all associated outputs?\nThis action cannot be undone. (y/n): '):
        try:
            abort.func(job_id)
            resp = logs.dbcmd('del_calc', job_id, getpass.getuser())
        except RuntimeError as err:
            safeprint(err)
        else:
            if 'success' in resp:
                print('Removed %d' % job_id)
            else:
                print(resp['error'])

def smart_run(job_ini, oqparam, log_level, log_file, exports, reuse_hazard):
    """
    Run calculations by storing their hazard checksum and reusing previous
    calculations if requested.
    """
    haz_checksum = readinput.get_checksum32(oqparam, hazard=True)
    # retrieve an old calculation with the right checksum, if any
    job = logs.dbcmd('get_job_from_checksum', haz_checksum)
    reuse = reuse_hazard and job and os.path.exists(job.ds_calc_dir + '.hdf5')
    # recompute the hazard and store the checksum
    ebr = (oqparam.calculation_mode == 'event_based_risk' and
           'gmfs' not in oqparam.inputs)
    if ebr:
        kw = dict(calculation_mode='event_based')
        if (oqparam.sites or 'sites' in oqparam.inputs or
                'site_model' in oqparam.inputs):
            # remove exposure from the hazard
            kw['exposure_file'] = ''
    else:
        kw = {}
    if not reuse:
        hc_id = run_job(job_ini, log_level, log_file, exports, **kw)
        if job is None:
            logs.dbcmd('add_checksum', hc_id, haz_checksum)
        elif not reuse_hazard or not os.path.exists(job.ds_calc_dir + '.hdf5'):
            logs.dbcmd('update_job_checksum', hc_id, haz_checksum)
        if ebr:
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
        else:
            job_id = hc_id
    else:
        hc_id = job.id
        logging.info('Reusing job #%d', job.id)
        job_id = run_job(job_ini, log_level, log_file,
                         exports, hazard_calculation_id=hc_id)
    if ebr and oqparam.aggregate_by:
        logging.info('Exporting aggregated data')
        dstore = util.read(job_id)
        aggby = 'aggregate_by/%s/' % ','.join(oqparam.aggregate_by)
        fnames = []
        fnames.extend(export((aggby + 'avg_losses', 'csv'), dstore))
        fnames.extend(export((aggby + 'curves', 'csv'), dstore))
        for fname in fnames:
            logging.info('Exported %s', fname)

@sap.script
def engine(log_file, no_distribute, yes, config_file, make_html_report,
           upgrade_db, db_version, what_if_I_upgrade, run,
           list_hazard_calculations, list_risk_calculations,
           delete_calculation, delete_uncompleted_calculations,
           hazard_calculation_id, list_outputs, show_log,
           export_output, export_outputs, exports='',
           log_level='info', reuse_hazard=False):
    """
    Run a calculation using the traditional command line API
    """
    if not run:
        # configure a basic logging
        logs.init()

    if config_file:
        config.read(os.path.abspath(os.path.expanduser(config_file)),
                    soft_mem_limit=int, hard_mem_limit=int, port=int,
                    multi_user=valid.boolean, multi_node=valid.boolean)

    if no_distribute:
        os.environ['OQ_DISTRIBUTE'] = 'no'

    # check if the datadir exists
    datadir = datastore.get_datadir()
    if not os.path.exists(datadir):
        os.makedirs(datadir)

    dbserver.ensure_on()
    # check if we are talking to the right server
    err = dbserver.check_foreign()
    if err:
        sys.exit(err)

    if upgrade_db:
        msg = logs.dbcmd('what_if_I_upgrade', 'read_scripts')
        if msg.startswith('Your database is already updated'):
            pass
        elif yes or confirm('Proceed? (y/n) '):
            logs.dbcmd('upgrade_db')
        sys.exit(0)

    if db_version:
        safeprint(logs.dbcmd('db_version'))
        sys.exit(0)

    if what_if_I_upgrade:
        safeprint(logs.dbcmd('what_if_I_upgrade', 'extract_upgrade_scripts'))
        sys.exit(0)

    # check if the db is outdated
    outdated = logs.dbcmd('check_outdated')
    if outdated:
        sys.exit(outdated)

    # hazard or hazard+risk
    if hazard_calculation_id == -1:
        # get the latest calculation of the current user
        hc_id = get_job_id(hazard_calculation_id, getpass.getuser())
    elif hazard_calculation_id:
        # make it possible to use calculations made by another user
        hc_id = get_job_id(hazard_calculation_id)
    else:
        hc_id = None
    if run:
        log_file = os.path.expanduser(log_file) \
            if log_file is not None else None
        job_inis = [os.path.expanduser(f) for f in run]
        if len(job_inis) == 1 and not hc_id:
            # init logs before calling get_oqparam
            logs.init('nojob', getattr(logging, log_level.upper()))
            # not using logs.handle that logs on the db
            oq = readinput.get_oqparam(job_inis[0])
            smart_run(job_inis[0], oq, log_level, log_file,
                      exports, reuse_hazard)
            return
        for i, job_ini in enumerate(job_inis):
            open(job_ini, 'rb').read()  # IOError if the file does not exist
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
            if not hc_id:  # use the first calculation as base for the others
                hc_id = job_id
    # hazard
    elif list_hazard_calculations:
        for line in logs.dbcmd(
                'list_calculations', 'hazard', getpass.getuser()):
            safeprint(line)
    elif delete_calculation is not None:
        del_calculation(delete_calculation, yes)
    # risk
    elif list_risk_calculations:
        for line in logs.dbcmd('list_calculations', 'risk', getpass.getuser()):
            safeprint(line)

    # export
    elif make_html_report:
        safeprint('Written %s' % make_report(make_html_report))
        sys.exit(0)

    elif list_outputs is not None:
        hc_id = get_job_id(list_outputs)
        for line in logs.dbcmd('list_outputs', hc_id):
            safeprint(line)
    elif show_log is not None:
        hc_id = get_job_id(show_log)
        for line in logs.dbcmd('get_log', hc_id):
            safeprint(line)

    elif export_output is not None:
        output_id, target_dir = export_output
        dskey, calc_id, datadir = logs.dbcmd('get_output', int(output_id))
        for line in core.export_output(
                dskey, calc_id, datadir, os.path.expanduser(target_dir),
                exports or 'csv,xml'):
            safeprint(line)

    elif export_outputs is not None:
        job_id, target_dir = export_outputs
        hc_id = get_job_id(job_id)
        for line in core.export_outputs(
                hc_id, os.path.expanduser(target_dir), exports or 'csv,xml'):
            safeprint(line)

    elif delete_uncompleted_calculations:
        logs.dbcmd('delete_uncompleted_calculations', getpass.getuser())

    else:
        engine.parentparser.prog = 'oq engine'
        engine.parentparser.print_usage()

engine._add('log_file', '--log-file', '-L', help='''\
Location where to store log messages; if not specified, log messages
will be printed to the console (to stderr)''')
engine._add('no_distribute', '--no-distribute', '--nd', help='''\
Disable calculation task distribution and run the
computation in a single process. This is intended for
use in debugging and profiling.''', action='store_true')
engine.flg('yes', 'Automatically answer "yes" when asked to confirm an action')
engine.opt('config_file', 'Custom openquake.cfg file, to override default '
           'configurations')
engine._add('make_html_report', '--make-html-report', '--r',
            help='Build an HTML report of the computation at the given date',
            metavar='YYYY-MM-DD|today')
engine.flg('upgrade_db', 'Upgrade the openquake database')
engine.flg('db_version', 'Show the current version of the openquake database')
engine.flg('what_if_I_upgrade', 'Show what will happen to the openquake '
           'database if you upgrade')
engine._add('run', '--run', help='Run a job with the specified config file',
            metavar='JOB_INI', nargs='+')
engine._add('list_hazard_calculations', '--list-hazard-calculations', '--lhc',
            help='List hazard calculation information', action='store_true')
engine._add('list_risk_calculations', '--list-risk-calculations', '--lrc',
            help='List risk calculation information', action='store_true')
engine._add('delete_calculation', '--delete-calculation', '--dc',
            help='Delete a calculation and all associated outputs',
            metavar='CALCULATION_ID', type=int)
engine._add('delete_uncompleted_calculations',
            '--delete-uncompleted-calculations', '--duc',
            help='Delete all the uncompleted calculations',
            action='store_true')
engine._add('hazard_calculation_id', '--hazard-calculation-id', '--hc',
            help='Use the given job as input for the next job')
engine._add('list_outputs', '--list-outputs', '--lo',
            help='List outputs for the specified calculation',
            metavar='CALCULATION_ID')
engine._add('show_log', '--show-log', '--sl',
            help='Show the log of the specified calculation',
            metavar='CALCULATION_ID')
engine._add('export_output', '--export-output', '--eo',
            nargs=2, metavar=('OUTPUT_ID', 'TARGET_DIR'),
            help='Export the desired output to the specified directory')
engine._add('export_outputs', '--export-outputs', '--eos',
            nargs=2, metavar=('CALCULATION_ID', 'TARGET_DIR'),
            help='Export all of the calculation outputs to the '
            'specified directory')
engine.opt('exports', 'Comma-separated string specifing the export formats, '
           'in order of priority')
engine.opt('log_level', 'Defaults to "info"',
           choices=['debug', 'info', 'warn', 'error', 'critical'])
engine.flg('reuse_hazard', 'Reuse the event based hazard if available')
<commit_msg>[demos]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import os
import sys
import getpass
import logging
from openquake.baselib import sap, config, datastore
from openquake.baselib.general import safeprint
from openquake.hazardlib import valid
from openquake.commonlib import util, logs, readinput
from openquake.engine import engine as eng
from openquake.engine.export import core
from openquake.engine.utils import confirm
from openquake.engine.tools.make_html_report import make_report
from openquake.server import dbserver
from openquake.calculators.export import export
from openquake.commands.abort import abort

HAZARD_CALCULATION_ARG = "--hazard-calculation-id"
MISSING_HAZARD_MSG = "Please specify '%s=<id>'" % HAZARD_CALCULATION_ARG

def get_job_id(job_id, username=None):
    job = logs.dbcmd('get_job', job_id, username)
    if not job:
        sys.exit('Job %s not found' % job_id)
    return job.id

def run_job(job_ini, log_level='info', log_file=None, exports='',
            username=getpass.getuser(), **kw):
    """
    Run a job using the specified config file and other options.

    :param str job_ini:
        Path to calculation config (INI-style) files.
    :param str log_level:
        'debug', 'info', 'warn', 'error', or 'critical'
    :param str log_file:
        Path to log file.
    :param exports:
        A comma-separated string of export types requested by the user.
    :param username:
        Name of the user running the job
    :param kw:
        Extra parameters like hazard_calculation_id and calculation_mode
    """
    job_id = logs.init('job', getattr(logging, log_level.upper()))
    with logs.handle(job_id, log_level, log_file):
        job_ini = os.path.abspath(job_ini)
        oqparam = eng.job_from_file(job_ini, job_id, username, **kw)
        kw['username'] = username
        eng.run_calc(job_id, oqparam, exports, **kw)
        for line in logs.dbcmd('list_outputs', job_id, False):
            safeprint(line)
    return job_id

def run_tile(job_ini, sites_slice):
    """
    Used in tiling calculations
    """
    return run_job(job_ini, sites_slice=(sites_slice.start, sites_slice.stop))

def del_calculation(job_id, confirmed=False):
    """
    Delete a calculation and all associated outputs.
    """
    if logs.dbcmd('get_job', job_id) is None:
        print('There is no job %d' % job_id)
        return

    if confirmed or confirm(
            'Are you sure you want to (abort and) delete this calculation and '
            'all associated outputs?\nThis action cannot be undone. (y/n): '):
        try:
            abort(job_id)
            resp = logs.dbcmd('del_calc', job_id, getpass.getuser())
        except RuntimeError as err:
            safeprint(err)
        else:
            if 'success' in resp:
                print('Removed %d' % job_id)
            else:
                print(resp['error'])

def smart_run(job_ini, oqparam, log_level, log_file, exports, reuse_hazard):
    """
    Run calculations by storing their hazard checksum and reusing previous
    calculations if requested.
    """
    haz_checksum = readinput.get_checksum32(oqparam, hazard=True)
    # retrieve an old calculation with the right checksum, if any
    job = logs.dbcmd('get_job_from_checksum', haz_checksum)
    reuse = reuse_hazard and job and os.path.exists(job.ds_calc_dir + '.hdf5')
    # recompute the hazard and store the checksum
    ebr = (oqparam.calculation_mode == 'event_based_risk' and
           'gmfs' not in oqparam.inputs)
    if ebr:
        kw = dict(calculation_mode='event_based')
        if (oqparam.sites or 'sites' in oqparam.inputs or
                'site_model' in oqparam.inputs):
            # remove exposure from the hazard
            kw['exposure_file'] = ''
    else:
        kw = {}
    if not reuse:
        hc_id = run_job(job_ini, log_level, log_file, exports, **kw)
        if job is None:
            logs.dbcmd('add_checksum', hc_id, haz_checksum)
        elif not reuse_hazard or not os.path.exists(job.ds_calc_dir + '.hdf5'):
            logs.dbcmd('update_job_checksum', hc_id, haz_checksum)
        if ebr:
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
        else:
            job_id = hc_id
    else:
        hc_id = job.id
        logging.info('Reusing job #%d', job.id)
        job_id = run_job(job_ini, log_level, log_file,
                         exports, hazard_calculation_id=hc_id)
    if ebr and oqparam.aggregate_by:
        logging.info('Exporting aggregated data')
        dstore = util.read(job_id)
        aggby = 'aggregate_by/%s/' % ','.join(oqparam.aggregate_by)
        fnames = []
        fnames.extend(export((aggby + 'avg_losses', 'csv'), dstore))
        fnames.extend(export((aggby + 'curves', 'csv'), dstore))
        for fname in fnames:
            logging.info('Exported %s', fname)

@sap.script
def engine(log_file, no_distribute, yes, config_file, make_html_report,
           upgrade_db, db_version, what_if_I_upgrade, run,
           list_hazard_calculations, list_risk_calculations,
           delete_calculation, delete_uncompleted_calculations,
           hazard_calculation_id, list_outputs, show_log,
           export_output, export_outputs, exports='',
           log_level='info', reuse_hazard=False):
    """
    Run a calculation using the traditional command line API
    """
    if not run:
        # configure a basic logging
        logs.init()

    if config_file:
        config.read(os.path.abspath(os.path.expanduser(config_file)),
                    soft_mem_limit=int, hard_mem_limit=int, port=int,
                    multi_user=valid.boolean, multi_node=valid.boolean)

    if no_distribute:
        os.environ['OQ_DISTRIBUTE'] = 'no'

    # check if the datadir exists
    datadir = datastore.get_datadir()
    if not os.path.exists(datadir):
        os.makedirs(datadir)

    dbserver.ensure_on()
    # check if we are talking to the right server
    err = dbserver.check_foreign()
    if err:
        sys.exit(err)

    if upgrade_db:
        msg = logs.dbcmd('what_if_I_upgrade', 'read_scripts')
        if msg.startswith('Your database is already updated'):
            pass
        elif yes or confirm('Proceed? (y/n) '):
            logs.dbcmd('upgrade_db')
        sys.exit(0)

    if db_version:
        safeprint(logs.dbcmd('db_version'))
        sys.exit(0)

    if what_if_I_upgrade:
        safeprint(logs.dbcmd('what_if_I_upgrade', 'extract_upgrade_scripts'))
        sys.exit(0)

    # check if the db is outdated
    outdated = logs.dbcmd('check_outdated')
    if outdated:
        sys.exit(outdated)

    # hazard or hazard+risk
    if hazard_calculation_id == -1:
        # get the latest calculation of the current user
        hc_id = get_job_id(hazard_calculation_id, getpass.getuser())
    elif hazard_calculation_id:
        # make it possible to use calculations made by another user
        hc_id = get_job_id(hazard_calculation_id)
    else:
        hc_id = None
    if run:
        log_file = os.path.expanduser(log_file) \
            if log_file is not None else None
        job_inis = [os.path.expanduser(f) for f in run]
        if len(job_inis) == 1 and not hc_id:
            # init logs before calling get_oqparam
            logs.init('nojob', getattr(logging, log_level.upper()))
            # not using logs.handle that logs on the db
            oq = readinput.get_oqparam(job_inis[0])
            smart_run(job_inis[0], oq, log_level, log_file,
                      exports, reuse_hazard)
            return
        for i, job_ini in enumerate(job_inis):
            open(job_ini, 'rb').read()  # IOError if the file does not exist
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
            if not hc_id:  # use the first calculation as base for the others
                hc_id = job_id
    # hazard
    elif list_hazard_calculations:
        for line in logs.dbcmd(
                'list_calculations', 'hazard', getpass.getuser()):
            safeprint(line)
    elif delete_calculation is not None:
        del_calculation(delete_calculation, yes)
    # risk
    elif list_risk_calculations:
        for line in logs.dbcmd('list_calculations', 'risk', getpass.getuser()):
            safeprint(line)

    # export
    elif make_html_report:
        safeprint('Written %s' % make_report(make_html_report))
        sys.exit(0)

    elif list_outputs is not None:
        hc_id = get_job_id(list_outputs)
        for line in logs.dbcmd('list_outputs', hc_id):
            safeprint(line)
    elif show_log is not None:
        hc_id = get_job_id(show_log)
        for line in logs.dbcmd('get_log', hc_id):
            safeprint(line)

    elif export_output is not None:
        output_id, target_dir = export_output
        dskey, calc_id, datadir = logs.dbcmd('get_output', int(output_id))
        for line in core.export_output(
                dskey, calc_id, datadir, os.path.expanduser(target_dir),
                exports or 'csv,xml'):
            safeprint(line)

    elif export_outputs is not None:
        job_id, target_dir = export_outputs
        hc_id = get_job_id(job_id)
        for line in core.export_outputs(
                hc_id, os.path.expanduser(target_dir), exports or 'csv,xml'):
            safeprint(line)

    elif delete_uncompleted_calculations:
        logs.dbcmd('delete_uncompleted_calculations', getpass.getuser())

    else:
        engine.parentparser.prog = 'oq engine'
        engine.parentparser.print_usage()

engine._add('log_file', '--log-file', '-L', help='''\
Location where to store log messages; if not specified, log messages
will be printed to the console (to stderr)''')
engine._add('no_distribute', '--no-distribute', '--nd', help='''\
Disable calculation task distribution and run the
computation in a single process. This is intended for
use in debugging and profiling.''', action='store_true')
engine.flg('yes', 'Automatically answer "yes" when asked to confirm an action')
engine.opt('config_file', 'Custom openquake.cfg file, to override default '
           'configurations')
engine._add('make_html_report', '--make-html-report', '--r',
            help='Build an HTML report of the computation at the given date',
            metavar='YYYY-MM-DD|today')
engine.flg('upgrade_db', 'Upgrade the openquake database')
engine.flg('db_version', 'Show the current version of the openquake database')
engine.flg('what_if_I_upgrade', 'Show what will happen to the openquake '
           'database if you upgrade')
engine._add('run', '--run', help='Run a job with the specified config file',
            metavar='JOB_INI', nargs='+')
engine._add('list_hazard_calculations', '--list-hazard-calculations', '--lhc',
            help='List hazard calculation information', action='store_true')
engine._add('list_risk_calculations', '--list-risk-calculations', '--lrc',
            help='List risk calculation information', action='store_true')
engine._add('delete_calculation', '--delete-calculation', '--dc',
            help='Delete a calculation and all associated outputs',
            metavar='CALCULATION_ID', type=int)
engine._add('delete_uncompleted_calculations',
            '--delete-uncompleted-calculations', '--duc',
            help='Delete all the uncompleted calculations',
            action='store_true')
engine._add('hazard_calculation_id', '--hazard-calculation-id', '--hc',
            help='Use the given job as input for the next job')
engine._add('list_outputs', '--list-outputs', '--lo',
            help='List outputs for the specified calculation',
            metavar='CALCULATION_ID')
engine._add('show_log', '--show-log', '--sl',
            help='Show the log of the specified calculation',
            metavar='CALCULATION_ID')
engine._add('export_output', '--export-output', '--eo',
            nargs=2, metavar=('OUTPUT_ID', 'TARGET_DIR'),
            help='Export the desired output to the specified directory')
engine._add('export_outputs', '--export-outputs', '--eos',
            nargs=2, metavar=('CALCULATION_ID', 'TARGET_DIR'),
            help='Export all of the calculation outputs to the '
            'specified directory')
engine.opt('exports', 'Comma-separated string specifing the export formats, '
           'in order of priority')
engine.opt('log_level', 'Defaults to "info"',
           choices=['debug', 'info', 'warn', 'error', 'critical'])
engine.flg('reuse_hazard', 'Reuse the event based hazard if available')
<|endoftext|>


<commit_before>import os

import dj_database_url

CONF_ROOT = os.path.dirname(__file__)

DATABASES = {
    'default': dj_database_url.config(),
}

# CACHES = {
#     'default': {
#         'BACKEND': 'django.core.cache.backends.memcached.MemcachedCache',
#         'LOCATION': ['127.0.0.1:11211'],
#     }
# }

# SENTRY_USE_QUEUE = True
# BROKER_URL = 'redis://localhost:6379'

# SENTRY_BUFFER = 'sentry.buffer.redis.RedisBuffer'
# SENTRY_BUFFER_OPTIONS = {
#     'hosts': {
#         0: {
#             'host': '127.0.0.1',
#             'port': 6379,
#         }
#     }
# }

SENTRY_KEY = os.getenv('SENTRY_KEY')

# Set this to false to require authentication
SENTRY_PUBLIC = False

# SENTRY_URL_PREFIX = 'http://sentry.example.com'  # No trailing slash!

SENTRY_WEB_HOST = '0.0.0.0'
SENTRY_WEB_PORT = os.getenv('PORT')
SENTRY_WEB_OPTIONS = {
    'workers': 3,
}

# Mail server configuration
EMAIL_BACKEND = 'database_email_backend.backend.DatabaseEmailBackend'

# http://twitter.com/apps/new
TWITTER_CONSUMER_KEY = ''
TWITTER_CONSUMER_SECRET = ''

# http://developers.facebook.com/setup/
FACEBOOK_APP_ID = ''
FACEBOOK_API_SECRET = ''

# http://code.google.com/apis/accounts/docs/OAuth2.html#Registering
GOOGLE_OAUTH2_CLIENT_ID = ''
GOOGLE_OAUTH2_CLIENT_SECRET = ''

# https://github.com/settings/applications/new
GITHUB_APP_ID = ''
GITHUB_API_SECRET = ''

# https://trello.com/1/appKey/generate
TRELLO_API_KEY = ''
TRELLO_API_SECRET = ''
<commit_msg>Add ALLOWED_HOSTS to sentry.conf.py<commit_after>import os

import dj_database_url

CONF_ROOT = os.path.dirname(__file__)

DATABASES = {
    'default': dj_database_url.config(),
}

# CACHES = {
#     'default': {
#         'BACKEND': 'django.core.cache.backends.memcached.MemcachedCache',
#         'LOCATION': ['127.0.0.1:11211'],
#     }
# }

# SENTRY_USE_QUEUE = True
# BROKER_URL = 'redis://localhost:6379'

# SENTRY_BUFFER = 'sentry.buffer.redis.RedisBuffer'
# SENTRY_BUFFER_OPTIONS = {
#     'hosts': {
#         0: {
#             'host': '127.0.0.1',
#             'port': 6379,
#         }
#     }
# }

SENTRY_KEY = os.getenv('SENTRY_KEY')

# Set this to false to require authentication
SENTRY_PUBLIC = False

# SENTRY_URL_PREFIX = 'http://sentry.example.com'  # No trailing slash!

SENTRY_WEB_HOST = '0.0.0.0'
SENTRY_WEB_PORT = os.getenv('PORT')
SENTRY_WEB_OPTIONS = {
    'workers': 3,
}

ALLOWED_HOSTS = [
    '.herokuapp.com',
]

# Mail server configuration
EMAIL_BACKEND = 'database_email_backend.backend.DatabaseEmailBackend'

# http://twitter.com/apps/new
TWITTER_CONSUMER_KEY = ''
TWITTER_CONSUMER_SECRET = ''

# http://developers.facebook.com/setup/
FACEBOOK_APP_ID = ''
FACEBOOK_API_SECRET = ''

# http://code.google.com/apis/accounts/docs/OAuth2.html#Registering
GOOGLE_OAUTH2_CLIENT_ID = ''
GOOGLE_OAUTH2_CLIENT_SECRET = ''

# https://github.com/settings/applications/new
GITHUB_APP_ID = ''
GITHUB_API_SECRET = ''

# https://trello.com/1/appKey/generate
TRELLO_API_KEY = ''
TRELLO_API_SECRET = ''
<|endoftext|>


<commit_before># Copyright (c) 2013, ReMake Electric ehf
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
"""
This is a module for a single message consumption process
It listens to a topic and expects to see a known sequence of messages.
"""

from __future__ import division

import collections
import errno
import logging
import math
import os
import stat
import tempfile
import time

import fuse
import mosquitto

from beem.trackers import ObservedMessage as MsgStatus

class TrackingListener():
    """
    An MQTT message subscriber that tracks an expected message sequence
    and generates timing, duplicate/missing and monitors for drops
    """

    msg_statuses = []

    def __init__(self, host, port, opts):
        self.options = opts
        self.cid = opts.clientid
        self.log = logging.getLogger(__name__ + ":" + self.cid)
        self.mqttc = mosquitto.Mosquitto(self.cid)
        self.mqttc.on_message = self.msg_handler
        self.listen_topic = opts.topic
        self.time_start = None
        # TODO - you _probably_ want to tweak this
        self.mqttc.max_inflight_messages_set(200)
        rc = self.mqttc.connect(host, port, 60)
        if rc:
            raise Exception("Couldn't even connect! ouch! rc=%d" % rc)
            # umm, how?
        self.mqttc.subscribe('$SYS/broker/publish/messages/dropped', 0)
        self.drop_count = None
        self.dropping = False
        self.mqttc.loop_start()

    def msg_handler(self, mosq, userdata, msg):
        # WARNING: this _must_ release as quickly as possible!
        # get the sequence id from the topic
        #self.log.debug("heard a message on topic: %s", msg.topic)
        if msg.topic == '$SYS/broker/publish/messages/dropped':
            if self.drop_count:
                self.log.warn("Drop count has increased by %d",
                              int(msg.payload) - self.drop_count)
                self.dropping = True
            else:
                self.drop_count = int(msg.payload)
                self.log.debug("Initial drop count: %d", self.drop_count)
            return
        if not self.time_start:
            self.time_start = time.time()

        try:
            ms = MsgStatus(msg)
            self.msg_statuses.append(ms)
        except Exception:
            self.log.exception("Failed to parse a received message. (Is the publisher sending time-tracking information with -t?)")

    def run(self, qos=1):
        """
        Start a (long lived) process waiting for messages to arrive.
        The number of clients and messages per client that are expected
        are set at creation time

        """
        self.expected = self.options.msg_count * self.options.client_count
        self.log.info(
            "Listening for %d messages on topic %s (q%d)",
            self.expected, self.listen_topic, qos)
        rc = self.mqttc.subscribe(self.listen_topic, qos)
        #assert rc == 0, "Failed to subscribe?! this isn't handled!", rc
        print ("subscribe rc was ", rc)
        while len(self.msg_statuses) < self.expected:
            # let the mosquitto thread fill us up
            time.sleep(1)
            self.log.info("Still waiting for %d messages",
                          self.expected - len(self.msg_statuses))
            if self.dropping:
                self.log.error("Detected drops are occuring, aborting test!")
                break
        self.time_end = time.time()
        self.mqttc.disconnect()

    def stats(self):
        msg_count = len(self.msg_statuses)
        flight_times = [x.time_flight() for x in self.msg_statuses]
        mean = sum(flight_times) / len(flight_times)
        squares = [x * x for x in [q - mean for q in flight_times]]
        stddev = math.sqrt(sum(squares) / len(flight_times))

        actual_clients = set([x.cid for x in self.msg_statuses])
        per_client_expected = range(1, self.options.msg_count + 1)
        per_client_real = {}
        per_client_missing = {}
        for cid in actual_clients:
            per_client_real[cid] = [x.mid for x in self.msg_statuses if x.cid == cid]
            per_client_missing[cid] = list(set(per_client_expected).difference(set(per_client_real[cid])))

        return {
            "clientid": self.cid,
            "client_count": len(actual_clients),
            "test_complete": not self.dropping,
            "msg_duplicates": [x for x, y in collections.Counter(self.msg_statuses).items() if y > 1],
            "msg_missing": per_client_missing,
            "msg_count": msg_count,
            "ms_per_msg": (self.time_end - self.time_start) / msg_count * 1000,
            "msg_per_sec": msg_count / (self.time_end - self.time_start),
            "time_total": self.time_end - self.time_start,
            "flight_time_mean": mean,
            "flight_time_stddev": stddev,
            "flight_time_max": max(flight_times),
            "flight_time_min": min(flight_times)
        }

def static_file_attrs(content=None):
    now = time.time()
    if content:
        size = len(content)
    else:
        size = 20
    return {
            "file": dict(st_mode=(stat.S_IFREG | 0444), st_nlink=1,
                            st_size=size,
                            st_ctime=now, st_mtime=now,
                            st_atime=now),
            "content": content
            }

class MalariaWatcherStatsFS(fuse.LoggingMixIn, fuse.Operations):

    file_attrs = dict(st_mode=(stat.S_IFREG | 0444), st_nlink=1,
                            st_size=20000,
                            st_ctime=time.time(), st_mtime=time.time(),
                            st_atime=time.time())

    dir_attrs = dict(st_mode=(stat.S_IFDIR | 0755),  st_nlink=2,
                            st_ctime=time.time(), st_mtime=time.time(),
                            st_atime=time.time())
    README_STATFS = """
This is a FUSE filesystem that contains a set of files representing various
statistics we have gathered about the MQTT broker we are watching and the
topics we are subscribed to.
"""
    msgs_total = 0
    msgs_stored = 0
    drop_count = 0
    drop_count_initial = None

    def handle_msgs_total(self):
        """Total number of messages seen since we started"""
        return self.msgs_total

    def handle_msgs_stored(self):
        """Total number of stored ($sys/broker/messages/stored)"""
        return self.msgs_stored

    def handle_uptime(self):
        """Time in seconds this watcher has been running"""
        return time.time() - self.time_start

    def handle_drop_count(self):
        """Total drops since this watcher has been running"""
        return self.drop_count

    def handle_topic(self):
        """The topic this watcher is subscribing too
        (Yes, this should be a list)"""
        return self.listen_topic

    def handle_readme(self):
        """Returns 'this' readme ;)"""
        rval = self.README_STATFS
        useful = [x for x in self.handlers if x != '/']
        file_field = "File                   "
        rval += "\n" + file_field + "Description\n\n"
        for h in useful:
            func = self.handlers[h].get("handler", None)
            desc = None
            if not func:
                desc = "Raw file, no further description"
            if func:
                desc = func.__doc__
            if not desc:
                desc = "No description in handler method! (Fix pydoc!)"
            # pad file line to line up with the description
            line = "%s%s\n" % (str.ljust(h[1:], len(file_field)), desc)
            rval += line
        return rval

    handlers = {
            "/": {"file": dir_attrs, "handler": None},
            "/msgs_total": {"file": file_attrs, "handler": handle_msgs_total},
            "/msgs_stored": {"file": file_attrs, "handler": handle_msgs_stored},
            "/uptime": {"file": file_attrs, "handler": handle_uptime},
            "/topic": {"file": file_attrs, "handler": handle_topic},
            "/drop_count": {"file": file_attrs, "handler": handle_drop_count},
            "/README": static_file_attrs(README_STATFS),
            "/README.detailed": {"file": file_attrs, "handler": handle_readme}
        }

    def __init__(self, options):
        print("listener operations __init__")
        self.options = options
        self.time_start = time.time()

    def msg_handler(self, mosq, userdata, msg):
        # WARNING: this _must_ release as quickly as possible!
        # get the sequence id from the topic
        #self.log.debug("heard a message on topic: %s", msg.topic)
        if msg.topic == '$SYS/broker/publish/messages/dropped':
            if self.drop_count_initial:
                self.log.warn("Drop count has increased by %d",
                              int(msg.payload) - self.drop_count_initial)
                self.drop_count = int(msg.payload) - self.drop_count_initial
            else:
                self.drop_count_initial = int(msg.payload)
                self.log.debug("Initial drops: %d", self.drop_count_initial)
            return
        if "messages/stored" in msg.topic:
            self.msgs_stored = int(msg.payload)
            return
        self.msgs_total += 1

    def init(self, path):
        """
        Fuse calls this when it's ready, so we can start our actual mqtt
        processes here.
        """
        print("listener post init init(), path=", path)
        self.cid = self.options.clientid
        self.log = logging.getLogger(__name__ + ":" + self.cid)
        self.mqttc = mosquitto.Mosquitto(self.cid)
        self.mqttc.on_message = self.msg_handler
        self.listen_topic = self.options.topic
        # TODO - you _probably_ want to tweak this
        self.mqttc.max_inflight_messages_set(200)
        rc = self.mqttc.connect(self.options.host, self.options.port, 60)
        if rc:
            raise Exception("Couldn't even connect! ouch! rc=%d" % rc)
            # umm, how?
        # b/p/m for >= 1.2, b/m for 1.1.x
        self.mqttc.subscribe('$SYS/broker/publish/messages/dropped', 0)
        self.mqttc.subscribe('$SYS/broker/messages/dropped', 0)
        self.mqttc.subscribe('$SYS/broker/messages/stored', 0)
        self.mqttc.loop_start()
        self.mqttc.subscribe(self.options.topic, self.options.qos)

    def getattr(self, path, fh=None):
        if path not in self.handlers:
            raise fuse.FuseOSError(errno.ENOENT)

        return self.handlers[path]["file"]

    def read(self, path, size, offset, fh):
        if self.handlers[path].get("content", False):
            return self.handlers[path]["content"]
        funcy = self.handlers[path]["handler"]
        return str(funcy(self)) + "\n"

    def readdir(self, path, fh):
        return ['.', '..'] + [x[1:] for x in self.handlers if x != '/']

class CensusListener():
    """
    Create a listener that just watches all the messages go past.
    It doesn't care about time in flight or expected vs actual, it just cares
    about what it has seen, and maintains long term stats on whatever
    it does see.
    """
    def __init__(self, options):
        self.log = logging.getLogger(__name__)
        path_provided = True
        if not options.directory:
            path_provided = False
            options.directory = tempfile.mkdtemp()
        self.log.info("Statistics files will be available in %s", options.directory)
        fuse.FUSE(MalariaWatcherStatsFS(options),
                  options.directory, foreground=True)
        if not path_provided:
            self.log.info("Automatically removing statsfs: %s", options.directory)
            os.rmdir(options.directory)
<commit_msg>mqttfs: Support drops from mosquitto 1.1.x and 1.2.x<commit_after># Copyright (c) 2013, ReMake Electric ehf
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
"""
This is a module for a single message consumption process
It listens to a topic and expects to see a known sequence of messages.
"""

from __future__ import division

import collections
import errno
import logging
import math
import os
import stat
import tempfile
import time

import fuse
import mosquitto

from beem.trackers import ObservedMessage as MsgStatus

class TrackingListener():
    """
    An MQTT message subscriber that tracks an expected message sequence
    and generates timing, duplicate/missing and monitors for drops
    """

    msg_statuses = []

    def __init__(self, host, port, opts):
        self.options = opts
        self.cid = opts.clientid
        self.log = logging.getLogger(__name__ + ":" + self.cid)
        self.mqttc = mosquitto.Mosquitto(self.cid)
        self.mqttc.on_message = self.msg_handler
        self.listen_topic = opts.topic
        self.time_start = None
        # TODO - you _probably_ want to tweak this
        self.mqttc.max_inflight_messages_set(200)
        rc = self.mqttc.connect(host, port, 60)
        if rc:
            raise Exception("Couldn't even connect! ouch! rc=%d" % rc)
            # umm, how?
        self.mqttc.subscribe('$SYS/broker/publish/messages/dropped', 0)
        self.drop_count = None
        self.dropping = False
        self.mqttc.loop_start()

    def msg_handler(self, mosq, userdata, msg):
        # WARNING: this _must_ release as quickly as possible!
        # get the sequence id from the topic
        #self.log.debug("heard a message on topic: %s", msg.topic)
        if msg.topic == '$SYS/broker/publish/messages/dropped':
            if self.drop_count:
                self.log.warn("Drop count has increased by %d",
                              int(msg.payload) - self.drop_count)
                self.dropping = True
            else:
                self.drop_count = int(msg.payload)
                self.log.debug("Initial drop count: %d", self.drop_count)
            return
        if not self.time_start:
            self.time_start = time.time()

        try:
            ms = MsgStatus(msg)
            self.msg_statuses.append(ms)
        except Exception:
            self.log.exception("Failed to parse a received message. (Is the publisher sending time-tracking information with -t?)")

    def run(self, qos=1):
        """
        Start a (long lived) process waiting for messages to arrive.
        The number of clients and messages per client that are expected
        are set at creation time

        """
        self.expected = self.options.msg_count * self.options.client_count
        self.log.info(
            "Listening for %d messages on topic %s (q%d)",
            self.expected, self.listen_topic, qos)
        rc = self.mqttc.subscribe(self.listen_topic, qos)
        #assert rc == 0, "Failed to subscribe?! this isn't handled!", rc
        print ("subscribe rc was ", rc)
        while len(self.msg_statuses) < self.expected:
            # let the mosquitto thread fill us up
            time.sleep(1)
            self.log.info("Still waiting for %d messages",
                          self.expected - len(self.msg_statuses))
            if self.dropping:
                self.log.error("Detected drops are occuring, aborting test!")
                break
        self.time_end = time.time()
        self.mqttc.disconnect()

    def stats(self):
        msg_count = len(self.msg_statuses)
        flight_times = [x.time_flight() for x in self.msg_statuses]
        mean = sum(flight_times) / len(flight_times)
        squares = [x * x for x in [q - mean for q in flight_times]]
        stddev = math.sqrt(sum(squares) / len(flight_times))

        actual_clients = set([x.cid for x in self.msg_statuses])
        per_client_expected = range(1, self.options.msg_count + 1)
        per_client_real = {}
        per_client_missing = {}
        for cid in actual_clients:
            per_client_real[cid] = [x.mid for x in self.msg_statuses if x.cid == cid]
            per_client_missing[cid] = list(set(per_client_expected).difference(set(per_client_real[cid])))

        return {
            "clientid": self.cid,
            "client_count": len(actual_clients),
            "test_complete": not self.dropping,
            "msg_duplicates": [x for x, y in collections.Counter(self.msg_statuses).items() if y > 1],
            "msg_missing": per_client_missing,
            "msg_count": msg_count,
            "ms_per_msg": (self.time_end - self.time_start) / msg_count * 1000,
            "msg_per_sec": msg_count / (self.time_end - self.time_start),
            "time_total": self.time_end - self.time_start,
            "flight_time_mean": mean,
            "flight_time_stddev": stddev,
            "flight_time_max": max(flight_times),
            "flight_time_min": min(flight_times)
        }

def static_file_attrs(content=None):
    now = time.time()
    if content:
        size = len(content)
    else:
        size = 20
    return {
            "file": dict(st_mode=(stat.S_IFREG | 0444), st_nlink=1,
                            st_size=size,
                            st_ctime=now, st_mtime=now,
                            st_atime=now),
            "content": content
            }

class MalariaWatcherStatsFS(fuse.LoggingMixIn, fuse.Operations):

    file_attrs = dict(st_mode=(stat.S_IFREG | 0444), st_nlink=1,
                            st_size=20000,
                            st_ctime=time.time(), st_mtime=time.time(),
                            st_atime=time.time())

    dir_attrs = dict(st_mode=(stat.S_IFDIR | 0755),  st_nlink=2,
                            st_ctime=time.time(), st_mtime=time.time(),
                            st_atime=time.time())
    README_STATFS = """
This is a FUSE filesystem that contains a set of files representing various
statistics we have gathered about the MQTT broker we are watching and the
topics we are subscribed to.
"""
    msgs_total = 0
    msgs_stored = 0
    drop_count = 0
    drop_count_initial = None

    def handle_msgs_total(self):
        """Total number of messages seen since we started"""
        return self.msgs_total

    def handle_msgs_stored(self):
        """Total number of stored ($sys/broker/messages/stored)"""
        return self.msgs_stored

    def handle_uptime(self):
        """Time in seconds this watcher has been running"""
        return time.time() - self.time_start

    def handle_drop_count(self):
        """Total drops since this watcher has been running"""
        return self.drop_count

    def handle_topic(self):
        """The topic this watcher is subscribing too
        (Yes, this should be a list)"""
        return self.listen_topic

    def handle_readme(self):
        """Returns 'this' readme ;)"""
        rval = self.README_STATFS
        useful = [x for x in self.handlers if x != '/']
        file_field = "File                   "
        rval += "\n" + file_field + "Description\n\n"
        for h in useful:
            func = self.handlers[h].get("handler", None)
            desc = None
            if not func:
                desc = "Raw file, no further description"
            if func:
                desc = func.__doc__
            if not desc:
                desc = "No description in handler method! (Fix pydoc!)"
            # pad file line to line up with the description
            line = "%s%s\n" % (str.ljust(h[1:], len(file_field)), desc)
            rval += line
        return rval

    handlers = {
            "/": {"file": dir_attrs, "handler": None},
            "/msgs_total": {"file": file_attrs, "handler": handle_msgs_total},
            "/msgs_stored": {"file": file_attrs, "handler": handle_msgs_stored},
            "/uptime": {"file": file_attrs, "handler": handle_uptime},
            "/topic": {"file": file_attrs, "handler": handle_topic},
            "/drop_count": {"file": file_attrs, "handler": handle_drop_count},
            "/README": static_file_attrs(README_STATFS),
            "/README.detailed": {"file": file_attrs, "handler": handle_readme}
        }

    def __init__(self, options):
        print("listener operations __init__")
        self.options = options
        self.time_start = time.time()

    def msg_handler(self, mosq, userdata, msg):
        # WARNING: this _must_ release as quickly as possible!
        # get the sequence id from the topic
        #self.log.debug("heard a message on topic: %s", msg.topic)
        if "/messages/dropped" in msg.topic:
            if self.drop_count_initial:
                self.log.warn("Drop count has increased by %d",
                              (int(msg.payload) - self.drop_count_initial))
                self.drop_count = int(msg.payload) - self.drop_count_initial
            else:
                self.drop_count_initial = int(msg.payload)
                self.log.debug("Initial drops: %d", self.drop_count_initial)
            return
        if "messages/stored" in msg.topic:
            self.msgs_stored = int(msg.payload)
            return
        self.msgs_total += 1

    def init(self, path):
        """
        Fuse calls this when it's ready, so we can start our actual mqtt
        processes here.
        """
        print("listener post init init(), path=", path)
        self.cid = self.options.clientid
        self.log = logging.getLogger(__name__ + ":" + self.cid)
        self.mqttc = mosquitto.Mosquitto(self.cid)
        self.mqttc.on_message = self.msg_handler
        self.listen_topic = self.options.topic
        # TODO - you _probably_ want to tweak this
        self.mqttc.max_inflight_messages_set(200)
        rc = self.mqttc.connect(self.options.host, self.options.port, 60)
        if rc:
            raise Exception("Couldn't even connect! ouch! rc=%d" % rc)
            # umm, how?
        # b/p/m for >= 1.2, b/m for 1.1.x
        self.mqttc.subscribe('$SYS/broker/publish/messages/dropped', 0)
        self.mqttc.subscribe('$SYS/broker/messages/dropped', 0)
        self.mqttc.subscribe('$SYS/broker/messages/stored', 0)
        self.mqttc.loop_start()
        self.mqttc.subscribe(self.options.topic, self.options.qos)

    def getattr(self, path, fh=None):
        if path not in self.handlers:
            raise fuse.FuseOSError(errno.ENOENT)

        return self.handlers[path]["file"]

    def read(self, path, size, offset, fh):
        if self.handlers[path].get("content", False):
            return self.handlers[path]["content"]
        funcy = self.handlers[path]["handler"]
        return str(funcy(self)) + "\n"

    def readdir(self, path, fh):
        return ['.', '..'] + [x[1:] for x in self.handlers if x != '/']

class CensusListener():
    """
    Create a listener that just watches all the messages go past.
    It doesn't care about time in flight or expected vs actual, it just cares
    about what it has seen, and maintains long term stats on whatever
    it does see.
    """
    def __init__(self, options):
        self.log = logging.getLogger(__name__)
        path_provided = True
        if not options.directory:
            path_provided = False
            options.directory = tempfile.mkdtemp()
        self.log.info("Statistics files will be available in %s", options.directory)
        fuse.FUSE(MalariaWatcherStatsFS(options),
                  options.directory, foreground=True)
        if not path_provided:
            self.log.info("Automatically removing statsfs: %s", options.directory)
            os.rmdir(options.directory)
<|endoftext|>


<commit_before>import os
import sys
import types
import time
import inspect

from twisted.internet import defer, reactor
from twisted.python import reflect, failure, usage

from twisted.python import log as tlog

from twisted.trial import unittest
from twisted.trial.runner import TrialRunner, TestLoader
from twisted.trial.runner import isPackage, isTestCase, ErrorHolder
from twisted.trial.runner import filenameToModule, _importFromFile

from ooni.reporter import ReporterFactory
from ooni.inputunit import InputUnitFactory
from ooni.nettest import InputTestSuite
from ooni import nettest
from ooni.utils import log, geodata, date
from ooni.plugoo import tests as oonitests

def isTestCase(thing):
    try:
        return issubclass(thing, unittest.TestCase)
    except TypeError:
        return False

def isLegacyTest(obj):
    """
    Returns True if the test in question is written using the OONITest legacy
    class.
    We do this for backward compatibility of the OONIProbe API.
    """
    try:
        if issubclass(obj, oonitests.OONITest) and not obj == oonitests.OONITest:
                return True
        else:
            return False
    except TypeError:
        return False

def adaptLegacyTest(obj, config):
    """
    We take a legacy OONITest class and convert it into a nettest.TestCase.
    This allows backward compatibility of old OONI tests.

    XXX perhaps we could implement another extra layer that makes the even
    older test cases compatible with the new OONI.
    """
    class legacy_reporter(object):
        def __init__(self, report_target):
            self.report_target = report_target

        def __call__(self, what):
            self.report_target.append(what)

    class LegacyOONITest(nettest.TestCase):
        try:
            name = obj.shortName
        except:
            name = "LegacyOONITest"

        originalTest = obj

        subOptions = obj.options()
        subOptions.parseOptions(config['subArgs'])

        test_class = obj(None, None, None, None)
        test_class.local_options = subOptions
        assets = test_class.load_assets()

        inputs = [None]
        # XXX here we are only taking assets that are set to one item only.
        for key, inputs in assets.items():
            pass

        inputs = inputs
        local_options = subOptions

        @defer.inlineCallbacks
        def test_start_legacy_test(self):

            self.legacy_report = []

            my_test = self.originalTest(None, None, None)
            my_test.report = legacy_reporter(self.legacy_report)
            args = {}
            args[self.key] = self.input
            result = yield my_test.startTest(args)
            self.report['result'] = result

    return LegacyOONITest

def processTest(obj, config):
    if obj.optParameters or obj.inputFile:
        if not obj.optParameters:
            obj.optParameters = []

        class Options(usage.Options):
            optParameters = obj.optParameters

        inputFile = obj.inputFile
        if inputFile:
            Options.optParameters.append(inputFile)

        print Options.optParameters
        options = Options()
        options.parseOptions(config['subArgs'])

        obj.localOptions = options

        if inputFile:
            obj.inputFile = options[inputFile[0]]
        try:
            tmp_obj = obj()
            tmp_obj.getOptions()
        except usage.UsageError:
            options.opt_help()

    return obj

def findTestClassesFromConfig(config):
    """
    Takes as input the command line config parameters and returns the test
    case classes.
    If it detects that a certain test class is using the old OONIProbe format,
    then it will adapt it to the new testing system.
    """
    filename = config['test']

    classes = []

    module = filenameToModule(filename)
    for name, val in inspect.getmembers(module):
        if isTestCase(val):
            classes.append(processTest(val, config))
        elif isLegacyTest(val):
            classes.append(adaptLegacyTest(val, config))
    return classes

def makeTestCases(klass, tests, methodPrefix):
    """
    Takes a class some tests and returns the test cases. methodPrefix is how
    the test case functions should be prefixed with.
    """
    cases = []
    for test in tests:
        cases.append(klass(methodPrefix+test))
    return cases

def loadTestsAndOptions(classes):
    """
    Takes a list of classes and returnes their testcases and options.
    Legacy tests will be adapted.
    """
    methodPrefix = 'test'
    suiteFactory = InputTestSuite
    options = []
    testCases = []
    names = []

    for klass in classes:
        try:
            k = klass()
            opts = k.getOptions()
            options.append(opts)
        except AttributeError:
            options.append([])
        tests = reflect.prefixedMethodNames(klass, methodPrefix)
        if tests:
            cases = makeTestCases(klass, tests, methodPrefix)
            testCases.append(cases)
        else:
            options.pop()

    return testCases, options

class ORunner(object):
    """
    This is a specialized runner used by the ooniprobe command line tool.
    I am responsible for reading the inputs from the test files and splitting
    them in input units. I also create all the report instances required to run
    the tests.
    """
    def __init__(self, cases, options=None, config=None, *arg, **kw):
        self.baseSuite = InputTestSuite
        self.cases = cases
        self.options = options
        self.inputs = options['inputs']
        try:
            reportFile = open(config['reportfile'], 'a+')
        except:
            filename = 'report_'+date.timestamp()+'.yaml'
            reportFile = open(filename, 'a+')
        self.reporterFactory = ReporterFactory(reportFile,
                                    testSuite=self.baseSuite(self.cases))

    def runWithInputUnit(self, inputUnit):
        idx = 0
        result = self.reporterFactory.create()

        for input in inputUnit:
            result.reporterFactory = self.reporterFactory

            suite = self.baseSuite(self.cases)
            suite.input = input
            suite(result, idx)

            # XXX refactor all of this index bullshit to avoid having to pass
            # this index around. Probably what I want to do is go and make
            # changes to report to support the concept of having multiple runs
            # of the same test.
            # We currently need to do this addition in order to get the number
            # of times the test cases that have run inside of the test suite.
            idx += (suite._idx - idx)

        result.done()

    def run(self):
        log.start()

        self.reporterFactory.options = self.options

        for inputUnit in InputUnitFactory(self.inputs):
            self.runWithInputUnit(inputUnit)

<commit_msg>Work on input file handling<commit_after>import os
import sys
import types
import time
import inspect

from twisted.internet import defer, reactor
from twisted.python import reflect, failure, usage

from twisted.python import log as tlog

from twisted.trial import unittest
from twisted.trial.runner import TrialRunner, TestLoader
from twisted.trial.runner import isPackage, isTestCase, ErrorHolder
from twisted.trial.runner import filenameToModule, _importFromFile

from ooni.reporter import ReporterFactory
from ooni.inputunit import InputUnitFactory
from ooni.nettest import InputTestSuite
from ooni import nettest
from ooni.utils import log, geodata, date
from ooni.plugoo import tests as oonitests

def isTestCase(thing):
    try:
        return issubclass(thing, unittest.TestCase)
    except TypeError:
        return False

def isLegacyTest(obj):
    """
    Returns True if the test in question is written using the OONITest legacy
    class.
    We do this for backward compatibility of the OONIProbe API.
    """
    try:
        if issubclass(obj, oonitests.OONITest) and not obj == oonitests.OONITest:
                return True
        else:
            return False
    except TypeError:
        return False

def adaptLegacyTest(obj, config):
    """
    We take a legacy OONITest class and convert it into a nettest.TestCase.
    This allows backward compatibility of old OONI tests.

    XXX perhaps we could implement another extra layer that makes the even
    older test cases compatible with the new OONI.
    """
    class legacy_reporter(object):
        def __init__(self, report_target):
            self.report_target = report_target

        def __call__(self, what):
            self.report_target.append(what)

    class LegacyOONITest(nettest.TestCase):
        try:
            name = obj.shortName
        except:
            name = "LegacyOONITest"

        originalTest = obj

        subOptions = obj.options()
        subOptions.parseOptions(config['subArgs'])

        test_class = obj(None, None, None, None)
        test_class.local_options = subOptions
        assets = test_class.load_assets()

        inputs = [None]
        # XXX here we are only taking assets that are set to one item only.
        for key, inputs in assets.items():
            pass

        inputs = inputs
        local_options = subOptions

        @defer.inlineCallbacks
        def test_start_legacy_test(self):

            self.legacy_report = []

            my_test = self.originalTest(None, None, None)
            my_test.report = legacy_reporter(self.legacy_report)
            args = {}
            args[self.key] = self.input
            result = yield my_test.startTest(args)
            self.report['result'] = result

    return LegacyOONITest

def processTest(obj, config):
    inputFile = obj.inputFile

    if obj.optParameters or inputFile:
        if not obj.optParameters:
            obj.optParameters = []

        if inputFile:
            obj.optParameters.append(inputFile)

        class Options(usage.Options):
            optParameters = obj.optParameters

        options = Options()
        options.parseOptions(config['subArgs'])

        obj.localOptions = options

        if inputFile:
            obj.inputFile = options[inputFile[0]]
        try:
            tmp_obj = obj()
            tmp_obj.getOptions()
        except usage.UsageError:
            options.opt_help()

    return obj

def findTestClassesFromConfig(config):
    """
    Takes as input the command line config parameters and returns the test
    case classes.
    If it detects that a certain test class is using the old OONIProbe format,
    then it will adapt it to the new testing system.
    """
    filename = config['test']

    classes = []

    module = filenameToModule(filename)
    for name, val in inspect.getmembers(module):
        if isTestCase(val):
            classes.append(processTest(val, config))
        elif isLegacyTest(val):
            classes.append(adaptLegacyTest(val, config))
    return classes

def makeTestCases(klass, tests, methodPrefix):
    """
    Takes a class some tests and returns the test cases. methodPrefix is how
    the test case functions should be prefixed with.
    """
    cases = []
    for test in tests:
        cases.append(klass(methodPrefix+test))
    return cases

def loadTestsAndOptions(classes):
    """
    Takes a list of classes and returnes their testcases and options.
    Legacy tests will be adapted.
    """
    methodPrefix = 'test'
    suiteFactory = InputTestSuite
    options = []
    testCases = []
    names = []

    for klass in classes:
        try:
            k = klass()
            opts = k.getOptions()
            options.append(opts)
        except AttributeError:
            options.append([])
        tests = reflect.prefixedMethodNames(klass, methodPrefix)
        if tests:
            cases = makeTestCases(klass, tests, methodPrefix)
            testCases.append(cases)
        else:
            options.pop()

    return testCases, options

class ORunner(object):
    """
    This is a specialized runner used by the ooniprobe command line tool.
    I am responsible for reading the inputs from the test files and splitting
    them in input units. I also create all the report instances required to run
    the tests.
    """
    def __init__(self, cases, options=None, config=None, *arg, **kw):
        self.baseSuite = InputTestSuite
        self.cases = cases
        self.options = options
        self.inputs = options['inputs']
        try:
            reportFile = open(config['reportfile'], 'a+')
        except:
            filename = 'report_'+date.timestamp()+'.yaml'
            reportFile = open(filename, 'a+')
        self.reporterFactory = ReporterFactory(reportFile,
                                    testSuite=self.baseSuite(self.cases))

    def runWithInputUnit(self, inputUnit):
        idx = 0
        result = self.reporterFactory.create()

        for input in inputUnit:
            result.reporterFactory = self.reporterFactory

            suite = self.baseSuite(self.cases)
            suite.input = input
            suite(result, idx)

            # XXX refactor all of this index bullshit to avoid having to pass
            # this index around. Probably what I want to do is go and make
            # changes to report to support the concept of having multiple runs
            # of the same test.
            # We currently need to do this addition in order to get the number
            # of times the test cases that have run inside of the test suite.
            idx += (suite._idx - idx)

        result.done()

    def run(self):
        log.start()

        self.reporterFactory.options = self.options

        for inputUnit in InputUnitFactory(self.inputs):
            self.runWithInputUnit(inputUnit)

<|endoftext|>


<commit_before>try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict
from django.db.models import Q
from avocado.formatters import registry as formatters
from serrano.formatters import HTMLFormatter
from django.conf import settings

class AlamutFormatter(HTMLFormatter):
    href = settings.ALAMUT_URL + '/show?request={0}'
    request = 'chr{chr}:g.{pos}{ref}>{alt}'

    def to_html(self, values, **context):
        request = self.request.format(**values)
        href = self.href.format(request)
        return '{label}'.format(
            href=href, label=request)
    to_html.process_multiple = True

    def to_excel(self, values, **context):
        request = self.request.format(**values)
        href = self.href.format(request)

        return '=HYPERLINK("{href}", "{label}")'.format(
            href=href, label=request)
    to_excel.process_multiple = True

class SampleFormatter(HTMLFormatter):
    def to_html(self, values, **kwargs):
        return values['sample__label']

    to_html.process_multiple = True

class ReadDepthFormatter(HTMLFormatter):
    template = '{read_depth}x '\
        '({coverage_ref}/{coverage_alt})'

    def to_html(self, values, **kwargs):
        label_class = ''
        if values['read_depth'] < 10:
            label_class = 'warning'
        elif values['read_depth'] >= 30:
            label_class = 'success'
        return self.template.format(label_class=label_class, **values)

    to_html.process_multiple = True

class CohortsFormatter(HTMLFormatter):
    # Cache cohort data since cannot and should not change during this
    # formatter's usage
    def _get_cohorts(self, **context):
        if not hasattr(self, '_cohorts'):
            from vdw.samples.models import Cohort
            # Augment resource with cohort-related details
            # (e.g. allele frequencies).
            perms = Q(user=None, published=True)

            if 'request' in context:
                perms |= Q(user=context['request'].user)

            # All cohorts this user has permission to
            self._cohorts = list(Cohort.objects.filter(perms, batch=None)
                                 .only('name', 'count')
                                 .order_by('order', 'name'))
        return self._cohorts

    def _get_cohort_variants(self, value, **context):
        from vdw.samples.models import CohortVariant
        cohorts = self._get_cohorts(**context)

        # All cohort variants for the above cohorts
        variants = dict(
            CohortVariant.objects.filter(
                variant__pk=value, cohort__in=cohorts)
            .values_list('cohort__pk', 'af').distinct())
        output = []
        for c in cohorts:
            output.append((c.name, variants.get(c.pk, 0), c.count))
        return output

    def to_html(self, value, **context):
        variants = self._get_cohort_variants(value, **context)

        html = ['
']
        for name, af, count in variants:
            html.append('
{0} {1}% ({2})'
                        '
'.format(name, str(af * 100), count))
        html.append('
')

        return ''.join(html)

    def to_excel(self, value, **context):
        variants = self._get_cohort_variants(value, **context)
        output = OrderedDict()
        for name, af, count in variants:
            key = '{0} Cohort Allele Freq ({1} samples)'.format(name, count)
            output[key] = af
        return output

    to_csv = to_excel

formatters.register(AlamutFormatter, 'Alamut Query Link')
formatters.register(SampleFormatter, 'Sample')
formatters.register(ReadDepthFormatter, 'Read Depth')
formatters.register(CohortsFormatter, 'Cohorts')
<commit_msg>Limit cohort formatter to cohorts from projects user has permision for<commit_after>try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict
from django.conf import settings
from django.db.models import Q
from guardian.shortcuts import get_objects_for_user
from avocado.formatters import registry as formatters
from serrano.formatters import HTMLFormatter

class AlamutFormatter(HTMLFormatter):
    href = settings.ALAMUT_URL + '/show?request={0}'
    request = 'chr{chr}:g.{pos}{ref}>{alt}'

    def to_html(self, values, **context):
        request = self.request.format(**values)
        href = self.href.format(request)
        return '{label}'.format(
            href=href, label=request)
    to_html.process_multiple = True

    def to_excel(self, values, **context):
        request = self.request.format(**values)
        href = self.href.format(request)

        return '=HYPERLINK("{href}", "{label}")'.format(
            href=href, label=request)
    to_excel.process_multiple = True

class SampleFormatter(HTMLFormatter):
    def to_html(self, values, **kwargs):
        return values['sample__label']

    to_html.process_multiple = True

class ReadDepthFormatter(HTMLFormatter):
    template = '{read_depth}x '\
        '({coverage_ref}/{coverage_alt})'

    def to_html(self, values, **kwargs):
        label_class = ''
        if values['read_depth'] < 10:
            label_class = 'warning'
        elif values['read_depth'] >= 30:
            label_class = 'success'
        return self.template.format(label_class=label_class, **values)

    to_html.process_multiple = True

class CohortsFormatter(HTMLFormatter):
    # Cache cohort data since cannot and should not change during this
    # formatter's usage
    def _get_cohorts(self, **context):
        if not hasattr(self, '_cohorts'):
            from vdw.samples.models import Cohort
            # Augment resource with cohort-related details
            # (e.g. allele frequencies).
            perms = Q(user=None, published=True)

            # All cohorts this user has permission to.
            projects = None
            if 'request' in context:
                perms |= Q(user=context['request'].user)
                projects = get_objects_for_user(
                    context['request'].user, 'samples.view_project')

            cohorts = Cohort.objects.filter(perms, batch=None)

            if projects is not None:
                cohorts = cohorts.filter(project__in=projects)

            self._cohorts = list(
                cohorts.only('name', 'count').order_by('order', 'name'))

        return self._cohorts

    def _get_cohort_variants(self, value, **context):
        from vdw.samples.models import CohortVariant
        cohorts = self._get_cohorts(**context)

        # All cohort variants for the above cohorts
        variants = dict(
            CohortVariant.objects.filter(
                variant__pk=value, cohort__in=cohorts)
            .values_list('cohort__pk', 'af').distinct())
        output = []
        for c in cohorts:
            output.append((c.name, variants.get(c.pk, 0), c.count))
        return output

    def to_html(self, value, **context):
        variants = self._get_cohort_variants(value, **context)

        html = ['
']
        for name, af, count in variants:
            html.append('
{0} {1}% ({2})'
                        '
'.format(name, str(af * 100), count))
        html.append('
')

        return ''.join(html)

    def to_excel(self, value, **context):
        variants = self._get_cohort_variants(value, **context)
        output = OrderedDict()
        for name, af, count in variants:
            key = '{0} Cohort Allele Freq ({1} samples)'.format(name, count)
            output[key] = af
        return output

    to_csv = to_excel

formatters.register(AlamutFormatter, 'Alamut Query Link')
formatters.register(SampleFormatter, 'Sample')
formatters.register(ReadDepthFormatter, 'Read Depth')
formatters.register(CohortsFormatter, 'Cohorts')
<|endoftext|>