<commit_before>#! /usr/bin/env python
from __future__ import print_function

import base64
import hashlib
import hmac
import os
import re
import sys
import time
from datetime import datetime, timedelta
from distutils.util import strtobool

try:
    from configparser import ConfigParser
except ImportError:  # python 2
    from ConfigParser import ConfigParser

try:
    from urllib.parse import quote
except ImportError:  # python 2
    from urllib import quote

import pytz

import requests
from requests.structures import CaseInsensitiveDict

PY2 = sys.version_info < (3, 0)

if PY2:
    text_type = unicode  # noqa
    string_type = basestring  # noqa
    integer_types = int, long  # noqa
    binary_type = str
else:
    text_type = str
    string_type = str
    integer_types = int
    binary_type = bytes

if sys.version_info >= (3, 5):
    try:
        from . import AIOCloudStack  # noqa
    except ImportError:
        pass

TIMEOUT = 10
PAGE_SIZE = 500
POLL_INTERVAL = 2.0
EXPIRATION = timedelta(minutes=10)
EXPIRES_FORMAT = "%Y-%m-%dT%H:%M:%S%z"

REQUIRED_CONFIG_KEYS = {"endpoint", "key", "secret", "method", "timeout"}
ALLOWED_CONFIG_KEYS = {"verify", "cert", "retry", "theme", "expiration",
                       "poll_interval", "trace", "dangerous_no_tls_verify"}
DEFAULT_CONFIG = {
    "timeout": 10,
    "method": "get",
    "retry": 0,
    "verify": None,
    "cert": None,
    "name": None,
    "expiration": 600,
    "poll_interval": POLL_INTERVAL,
    "trace": None,
    "dangerous_no_tls_verify": False,
}

PENDING = 0
SUCCESS = 1
FAILURE = 2

def cs_encode(s):
    """Encode URI component like CloudStack would do before signing.

    java.net.URLEncoder.encode(s).replace('+', '%20')
    """
    if PY2 and isinstance(s, text_type):
        s = s.encode("utf-8")
    return quote(s, safe="*")

def transform(params):
    """
    Transforms an heterogeneous map of params into a CloudStack
    ready mapping of parameter to values.

    It handles lists and dicts.

    >>> p = {"a": 1, "b": "foo", "c": ["eggs", "spam"], "d": {"key": "value"}}
    >>> transform(p)
    >>> print(p)
    {'a': '1', 'b': 'foo', 'c': 'eggs,spam', 'd[0].key': 'value'}
    """
    for key, value in list(params.items()):
        if value is None:
            params.pop(key)
            continue

        if isinstance(value, (string_type, binary_type)):
            continue

        if isinstance(value, integer_types):
            params[key] = text_type(value)
        elif isinstance(value, (list, tuple, set, dict)):
            if not value:
                params.pop(key)
            else:
                if isinstance(value, dict):
                    value = [value]
                if isinstance(value, set):
                    value = list(value)
                if not isinstance(value[0], dict):
                    params[key] = ",".join(value)
                else:
                    params.pop(key)
                    for index, val in enumerate(value):
                        for name, v in val.items():
                            k = "%s[%d].%s" % (key, index, name)
                            params[k] = text_type(v)
        else:
            raise ValueError(type(value))

class CloudStackException(Exception):
    """Exception nicely wrapping a request response."""
    def __init__(self, *args, **kwargs):
        self.response = kwargs.pop('response')
        assert self.response, "a response keyword argument is required"
        super(CloudStackException, self).__init__(*args, **kwargs)

class CloudStack(object):
    def __init__(self, endpoint, key, secret, timeout=10, method='get',
                 verify=None, cert=None, name=None, retry=0,
                 job_timeout=None, poll_interval=POLL_INTERVAL,
                 expiration=timedelta(minutes=10), trace=False,
                 dangerous_no_tls_verify=False):
        self.endpoint = endpoint
        self.key = key
        self.secret = secret
        self.timeout = int(timeout)
        self.method = method.lower()
        if verify:
            self.verify = verify
        else:
            self.verify = not dangerous_no_tls_verify

        self.cert = cert
        self.name = name
        self.retry = int(retry)
        self.job_timeout = int(job_timeout) if job_timeout else 0
        self.poll_interval = float(poll_interval)
        if not hasattr(expiration, "seconds"):
            expiration = timedelta(seconds=int(expiration))
        self.expiration = expiration
        self.trace = bool(trace)

    def __repr__(self):
        return '<CloudStack: {0}>'.format(self.name or self.endpoint)

    def __getattr__(self, command):
        def handler(**kwargs):
            return self._request(command, **kwargs)
        return handler

    def _prepare_request(self, command, json=True, opcode_name='command',
                         fetch_list=False, **kwargs):
        params = CaseInsensitiveDict(**kwargs)
        params.update({
            'apiKey': self.key,
            opcode_name: command,
        })
        if json:
            params['response'] = 'json'
        if 'page' in kwargs or fetch_list:
            params.setdefault('pagesize', PAGE_SIZE)
        if 'expires' not in params and self.expiration.total_seconds() >= 0:
            params['signatureVersion'] = '3'
            tz = pytz.utc
            expires = tz.localize(datetime.utcnow() + self.expiration)
            params['expires'] = expires.astimezone(tz).strftime(EXPIRES_FORMAT)

        kind = 'params' if self.method == 'get' else 'data'
        return kind, dict(params.items())

    def _request(self, command, json=True, opcode_name='command',
                 fetch_list=False, headers=None, **params):
        fetch_result = params.pop('fetch_result', False)
        kind, params = self._prepare_request(command, json, opcode_name,
                                             fetch_list, **params)

        done = False
        max_retry = self.retry
        final_data = []
        page = 1
        while not done:
            if fetch_list:
                params['page'] = page

            transform(params)
            params.pop('signature', None)
            params['signature'] = self._sign(params)

            req = requests.Request(self.method,
                                   self.endpoint,
                                   headers=headers,
                                   **{kind: params})
            prepped = req.prepare()
            if self.trace:
                print(prepped.method, prepped.url, file=sys.stderr)
                if prepped.headers:
                    print(prepped.headers, "\n", file=sys.stderr)
                if prepped.body:
                    print(prepped.body, file=sys.stderr)
                else:
                    print(file=sys.stderr)

            try:
                with requests.Session() as session:
                    response = session.send(prepped,
                                            timeout=self.timeout,
                                            verify=self.verify,
                                            cert=self.cert)

            except requests.exceptions.ConnectionError:
                max_retry -= 1
                if (
                    max_retry < 0 or
                    not command.startswith(('list', 'queryAsync'))
                ):
                    raise
                continue
            max_retry = self.retry

            if self.trace:
                print(response.status_code, response.reason, file=sys.stderr)
                headersTrace = "\n".join("{}: {}".format(k, v)
                                         for k, v in response.headers.items())
                print(headersTrace, "\n", file=sys.stderr)
                print(response.text, "\n", file=sys.stderr)

            data = self._response_value(response, json)

            if fetch_list:
                try:
                    [key] = [k for k in data.keys() if k != 'count']
                except ValueError:
                    done = True
                else:
                    final_data.extend(data[key])
                    page += 1
                    if len(final_data) >= data.get('count', PAGE_SIZE):
                        done = True
            elif fetch_result and 'jobid' in data:
                final_data = self._jobresult(jobid=data['jobid'],
                                             headers=headers)
                done = True
            else:
                final_data = data
                done = True
        return final_data

    def _response_value(self, response, json=True):
        """Parses the HTTP response as a the cloudstack value.

        It throws an exception if the server didn't answer with a 200.
        """
        if json:
            contentType = response.headers.get("Content-Type", "")
            if not contentType.startswith(("application/json",
                                           "text/javascript")):
                raise CloudStackException(
                    "JSON (application/json) was expected, got {!r}"
                    .format(contentType),
                    response=response)

            try:
                data = response.json()
            except ValueError as e:
                msg = "Make sure endpoint URL '%s' is correct." % self.endpoint
                raise CloudStackException(
                    "HTTP {0} response from CloudStack"
                    .format(response.status_code),
                    "%s. " % str(e) + msg,
                    response=response)

            [key] = data.keys()
            data = data[key]
        else:
            data = response.text

        if response.status_code != 200:
            raise CloudStackException(
                "HTTP {0} response from CloudStack".format(
                    response.status_code),
                data,
                response=response)

        return data

    def _jobresult(self, jobid, json=True, headers=None):
        """Poll the async job result.

        To be run via in a Thread, the result is put within
        the result list which is a hack.
        """
        failures = 0

        total_time = self.job_timeout or 2**30
        remaining = timedelta(seconds=total_time)
        endtime = datetime.now() + remaining

        while remaining.total_seconds() > 0:
            timeout = max(min(self.timeout, remaining.total_seconds()), 1)
            try:
                kind, params = self._prepare_request('queryAsyncJobResult',
                                                     jobid=jobid)

                transform(params)
                params['signature'] = self._sign(params)

                response = getattr(requests, self.method)(self.endpoint,
                                                          headers=headers,
                                                          timeout=timeout,
                                                          verify=self.verify,
                                                          cert=self.cert,
                                                          **{kind: params})

                j = self._response_value(response, json)

                failures = 0
                if j['jobstatus'] != PENDING:
                    if j['jobresultcode'] or j['jobstatus'] != SUCCESS:
                        raise CloudStackException("Job failure",
                                                  response=response)

                    if 'jobresult' not in j:
                        raise CloudStackException("Unknown job result",
                                                  response=response)

                    return j['jobresult']

            except CloudStackException:
                raise

            except Exception as e:
                failures += 1
                if failures > 10:
                    raise e

            time.sleep(self.poll_interval)
            remaining = endtime - datetime.now()

        if response:
            response.status_code = 408

        raise CloudStackException("Timeout waiting for async job result",
                                  jobid,
                                  response=response)

    def _sign(self, data):
        """
        Compute a signature string according to the CloudStack
        signature method (hmac/sha1).
        """

        # Python2/3 urlencode aren't good enough for this task.
        params = "&".join(
            "=".join((key, cs_encode(value)))
            for key, value in sorted(data.items())
        )

        digest = hmac.new(
            self.secret.encode('utf-8'),
            msg=params.lower().encode('utf-8'),
            digestmod=hashlib.sha1).digest()

        return base64.b64encode(digest).decode('utf-8').strip()

def read_config_from_ini(ini_group=None):
    # Config file: $PWD/cloudstack.ini or $HOME/.cloudstack.ini
    # Last read wins in configparser
    paths = [
        os.path.join(os.path.expanduser('~'), '.cloudstack.ini'),
        os.path.join(os.getcwd(), 'cloudstack.ini'),
    ]
    # Look at CLOUDSTACK_CONFIG first if present
    if 'CLOUDSTACK_CONFIG' in os.environ:
        paths.append(os.path.expanduser(os.environ['CLOUDSTACK_CONFIG']))
    if not any([os.path.exists(c) for c in paths]):
        raise SystemExit("Config file not found. Tried {0}".format(
            ", ".join(paths)))
    conf = ConfigParser()
    conf.read(paths)

    if not ini_group:
        ini_group = os.getenv('CLOUDSTACK_REGION', 'cloudstack')

        if not conf.has_section(ini_group):
            return dict(name=None)

    all_keys = REQUIRED_CONFIG_KEYS.union(ALLOWED_CONFIG_KEYS)
    ini_config = {k: v
                  for k, v in conf.items(ini_group)
                  if v and k in all_keys}
    ini_config["name"] = ini_group
    return ini_config

def read_config(ini_group=None):
    """
    Read the configuration from the environment, or config.

    First it try to go for the environment, then it overrides
    those with the cloudstack.ini file.
    """
    env_conf = dict(DEFAULT_CONFIG)
    for key in REQUIRED_CONFIG_KEYS.union(ALLOWED_CONFIG_KEYS):
        env_key = "CLOUDSTACK_{0}".format(key.upper())
        value = os.getenv(env_key)
        if value:
            env_conf[key] = value

    # overrides means we have a .ini to read
    overrides = os.getenv('CLOUDSTACK_OVERRIDES', '').strip()

    if not overrides and set(env_conf).issuperset(REQUIRED_CONFIG_KEYS):
        return env_conf

    ini_conf = read_config_from_ini(ini_group)

    overrides = {s.lower() for s in re.split(r'\W+', overrides)}
    config = dict(dict(env_conf, **ini_conf),
                  **{k: v for k, v in env_conf.items() if k in overrides})

    missings = REQUIRED_CONFIG_KEYS.difference(config)
    if missings:
        raise ValueError("the configuration is missing the following keys: " +
                         ", ".join(missings))

    # convert booleans values.
    bool_keys = ('dangerous_no_tls_verify',)
    for bool_key in bool_keys:
        if isinstance(config[bool_key], string_type):
            try:
                config[bool_key] = strtobool(config[bool_key])
            except ValueError:
                pass

    return config
<commit_msg>client: assert is redundant<commit_after>#! /usr/bin/env python
from __future__ import print_function

import base64
import hashlib
import hmac
import os
import re
import sys
import time
from datetime import datetime, timedelta
from distutils.util import strtobool

try:
    from configparser import ConfigParser
except ImportError:  # python 2
    from ConfigParser import ConfigParser

try:
    from urllib.parse import quote
except ImportError:  # python 2
    from urllib import quote

import pytz

import requests
from requests.structures import CaseInsensitiveDict

PY2 = sys.version_info < (3, 0)

if PY2:
    text_type = unicode  # noqa
    string_type = basestring  # noqa
    integer_types = int, long  # noqa
    binary_type = str
else:
    text_type = str
    string_type = str
    integer_types = int
    binary_type = bytes

if sys.version_info >= (3, 5):
    try:
        from . import AIOCloudStack  # noqa
    except ImportError:
        pass

TIMEOUT = 10
PAGE_SIZE = 500
POLL_INTERVAL = 2.0
EXPIRATION = timedelta(minutes=10)
EXPIRES_FORMAT = "%Y-%m-%dT%H:%M:%S%z"

REQUIRED_CONFIG_KEYS = {"endpoint", "key", "secret", "method", "timeout"}
ALLOWED_CONFIG_KEYS = {"verify", "cert", "retry", "theme", "expiration",
                       "poll_interval", "trace", "dangerous_no_tls_verify"}
DEFAULT_CONFIG = {
    "timeout": 10,
    "method": "get",
    "retry": 0,
    "verify": None,
    "cert": None,
    "name": None,
    "expiration": 600,
    "poll_interval": POLL_INTERVAL,
    "trace": None,
    "dangerous_no_tls_verify": False,
}

PENDING = 0
SUCCESS = 1
FAILURE = 2

def cs_encode(s):
    """Encode URI component like CloudStack would do before signing.

    java.net.URLEncoder.encode(s).replace('+', '%20')
    """
    if PY2 and isinstance(s, text_type):
        s = s.encode("utf-8")
    return quote(s, safe="*")

def transform(params):
    """
    Transforms an heterogeneous map of params into a CloudStack
    ready mapping of parameter to values.

    It handles lists and dicts.

    >>> p = {"a": 1, "b": "foo", "c": ["eggs", "spam"], "d": {"key": "value"}}
    >>> transform(p)
    >>> print(p)
    {'a': '1', 'b': 'foo', 'c': 'eggs,spam', 'd[0].key': 'value'}
    """
    for key, value in list(params.items()):
        if value is None:
            params.pop(key)
            continue

        if isinstance(value, (string_type, binary_type)):
            continue

        if isinstance(value, integer_types):
            params[key] = text_type(value)
        elif isinstance(value, (list, tuple, set, dict)):
            if not value:
                params.pop(key)
            else:
                if isinstance(value, dict):
                    value = [value]
                if isinstance(value, set):
                    value = list(value)
                if not isinstance(value[0], dict):
                    params[key] = ",".join(value)
                else:
                    params.pop(key)
                    for index, val in enumerate(value):
                        for name, v in val.items():
                            k = "%s[%d].%s" % (key, index, name)
                            params[k] = text_type(v)
        else:
            raise ValueError(type(value))

class CloudStackException(Exception):
    """Exception nicely wrapping a request response."""
    def __init__(self, *args, **kwargs):
        self.response = kwargs.pop('response')
        super(CloudStackException, self).__init__(*args, **kwargs)

class CloudStack(object):
    def __init__(self, endpoint, key, secret, timeout=10, method='get',
                 verify=None, cert=None, name=None, retry=0,
                 job_timeout=None, poll_interval=POLL_INTERVAL,
                 expiration=timedelta(minutes=10), trace=False,
                 dangerous_no_tls_verify=False):
        self.endpoint = endpoint
        self.key = key
        self.secret = secret
        self.timeout = int(timeout)
        self.method = method.lower()
        if verify:
            self.verify = verify
        else:
            self.verify = not dangerous_no_tls_verify

        self.cert = cert
        self.name = name
        self.retry = int(retry)
        self.job_timeout = int(job_timeout) if job_timeout else 0
        self.poll_interval = float(poll_interval)
        if not hasattr(expiration, "seconds"):
            expiration = timedelta(seconds=int(expiration))
        self.expiration = expiration
        self.trace = bool(trace)

    def __repr__(self):
        return '<CloudStack: {0}>'.format(self.name or self.endpoint)

    def __getattr__(self, command):
        def handler(**kwargs):
            return self._request(command, **kwargs)
        return handler

    def _prepare_request(self, command, json=True, opcode_name='command',
                         fetch_list=False, **kwargs):
        params = CaseInsensitiveDict(**kwargs)
        params.update({
            'apiKey': self.key,
            opcode_name: command,
        })
        if json:
            params['response'] = 'json'
        if 'page' in kwargs or fetch_list:
            params.setdefault('pagesize', PAGE_SIZE)
        if 'expires' not in params and self.expiration.total_seconds() >= 0:
            params['signatureVersion'] = '3'
            tz = pytz.utc
            expires = tz.localize(datetime.utcnow() + self.expiration)
            params['expires'] = expires.astimezone(tz).strftime(EXPIRES_FORMAT)

        kind = 'params' if self.method == 'get' else 'data'
        return kind, dict(params.items())

    def _request(self, command, json=True, opcode_name='command',
                 fetch_list=False, headers=None, **params):
        fetch_result = params.pop('fetch_result', False)
        kind, params = self._prepare_request(command, json, opcode_name,
                                             fetch_list, **params)

        done = False
        max_retry = self.retry
        final_data = []
        page = 1
        while not done:
            if fetch_list:
                params['page'] = page

            transform(params)
            params.pop('signature', None)
            params['signature'] = self._sign(params)

            req = requests.Request(self.method,
                                   self.endpoint,
                                   headers=headers,
                                   **{kind: params})
            prepped = req.prepare()
            if self.trace:
                print(prepped.method, prepped.url, file=sys.stderr)
                if prepped.headers:
                    print(prepped.headers, "\n", file=sys.stderr)
                if prepped.body:
                    print(prepped.body, file=sys.stderr)
                else:
                    print(file=sys.stderr)

            try:
                with requests.Session() as session:
                    response = session.send(prepped,
                                            timeout=self.timeout,
                                            verify=self.verify,
                                            cert=self.cert)

            except requests.exceptions.ConnectionError:
                max_retry -= 1
                if (
                    max_retry < 0 or
                    not command.startswith(('list', 'queryAsync'))
                ):
                    raise
                continue
            max_retry = self.retry

            if self.trace:
                print(response.status_code, response.reason, file=sys.stderr)
                headersTrace = "\n".join("{}: {}".format(k, v)
                                         for k, v in response.headers.items())
                print(headersTrace, "\n", file=sys.stderr)
                print(response.text, "\n", file=sys.stderr)

            data = self._response_value(response, json)

            if fetch_list:
                try:
                    [key] = [k for k in data.keys() if k != 'count']
                except ValueError:
                    done = True
                else:
                    final_data.extend(data[key])
                    page += 1
                    if len(final_data) >= data.get('count', PAGE_SIZE):
                        done = True
            elif fetch_result and 'jobid' in data:
                final_data = self._jobresult(jobid=data['jobid'],
                                             headers=headers)
                done = True
            else:
                final_data = data
                done = True
        return final_data

    def _response_value(self, response, json=True):
        """Parses the HTTP response as a the cloudstack value.

        It throws an exception if the server didn't answer with a 200.
        """
        if json:
            contentType = response.headers.get("Content-Type", "")
            if not contentType.startswith(("application/json",
                                           "text/javascript")):
                raise CloudStackException(
                    "JSON (application/json) was expected, got {!r}"
                    .format(contentType),
                    response=response)

            try:
                data = response.json()
            except ValueError as e:
                msg = "Make sure endpoint URL '%s' is correct." % self.endpoint
                raise CloudStackException(
                    "HTTP {0} response from CloudStack"
                    .format(response.status_code),
                    "%s. " % str(e) + msg,
                    response=response)

            [key] = data.keys()
            data = data[key]
        else:
            data = response.text

        if response.status_code != 200:
            raise CloudStackException(
                "HTTP {0} response from CloudStack".format(
                    response.status_code),
                data,
                response=response)

        return data

    def _jobresult(self, jobid, json=True, headers=None):
        """Poll the async job result.

        To be run via in a Thread, the result is put within
        the result list which is a hack.
        """
        failures = 0

        total_time = self.job_timeout or 2**30
        remaining = timedelta(seconds=total_time)
        endtime = datetime.now() + remaining

        while remaining.total_seconds() > 0:
            timeout = max(min(self.timeout, remaining.total_seconds()), 1)
            try:
                kind, params = self._prepare_request('queryAsyncJobResult',
                                                     jobid=jobid)

                transform(params)
                params['signature'] = self._sign(params)

                response = getattr(requests, self.method)(self.endpoint,
                                                          headers=headers,
                                                          timeout=timeout,
                                                          verify=self.verify,
                                                          cert=self.cert,
                                                          **{kind: params})

                j = self._response_value(response, json)

                failures = 0
                if j['jobstatus'] != PENDING:
                    if j['jobresultcode'] or j['jobstatus'] != SUCCESS:
                        raise CloudStackException("Job failure",
                                                  response=response)

                    if 'jobresult' not in j:
                        raise CloudStackException("Unknown job result",
                                                  response=response)

                    return j['jobresult']

            except CloudStackException:
                raise

            except Exception as e:
                failures += 1
                if failures > 10:
                    raise e

            time.sleep(self.poll_interval)
            remaining = endtime - datetime.now()

        if response:
            response.status_code = 408

        raise CloudStackException("Timeout waiting for async job result",
                                  jobid,
                                  response=response)

    def _sign(self, data):
        """
        Compute a signature string according to the CloudStack
        signature method (hmac/sha1).
        """

        # Python2/3 urlencode aren't good enough for this task.
        params = "&".join(
            "=".join((key, cs_encode(value)))
            for key, value in sorted(data.items())
        )

        digest = hmac.new(
            self.secret.encode('utf-8'),
            msg=params.lower().encode('utf-8'),
            digestmod=hashlib.sha1).digest()

        return base64.b64encode(digest).decode('utf-8').strip()

def read_config_from_ini(ini_group=None):
    # Config file: $PWD/cloudstack.ini or $HOME/.cloudstack.ini
    # Last read wins in configparser
    paths = [
        os.path.join(os.path.expanduser('~'), '.cloudstack.ini'),
        os.path.join(os.getcwd(), 'cloudstack.ini'),
    ]
    # Look at CLOUDSTACK_CONFIG first if present
    if 'CLOUDSTACK_CONFIG' in os.environ:
        paths.append(os.path.expanduser(os.environ['CLOUDSTACK_CONFIG']))
    if not any([os.path.exists(c) for c in paths]):
        raise SystemExit("Config file not found. Tried {0}".format(
            ", ".join(paths)))
    conf = ConfigParser()
    conf.read(paths)

    if not ini_group:
        ini_group = os.getenv('CLOUDSTACK_REGION', 'cloudstack')

        if not conf.has_section(ini_group):
            return dict(name=None)

    all_keys = REQUIRED_CONFIG_KEYS.union(ALLOWED_CONFIG_KEYS)
    ini_config = {k: v
                  for k, v in conf.items(ini_group)
                  if v and k in all_keys}
    ini_config["name"] = ini_group
    return ini_config

def read_config(ini_group=None):
    """
    Read the configuration from the environment, or config.

    First it try to go for the environment, then it overrides
    those with the cloudstack.ini file.
    """
    env_conf = dict(DEFAULT_CONFIG)
    for key in REQUIRED_CONFIG_KEYS.union(ALLOWED_CONFIG_KEYS):
        env_key = "CLOUDSTACK_{0}".format(key.upper())
        value = os.getenv(env_key)
        if value:
            env_conf[key] = value

    # overrides means we have a .ini to read
    overrides = os.getenv('CLOUDSTACK_OVERRIDES', '').strip()

    if not overrides and set(env_conf).issuperset(REQUIRED_CONFIG_KEYS):
        return env_conf

    ini_conf = read_config_from_ini(ini_group)

    overrides = {s.lower() for s in re.split(r'\W+', overrides)}
    config = dict(dict(env_conf, **ini_conf),
                  **{k: v for k, v in env_conf.items() if k in overrides})

    missings = REQUIRED_CONFIG_KEYS.difference(config)
    if missings:
        raise ValueError("the configuration is missing the following keys: " +
                         ", ".join(missings))

    # convert booleans values.
    bool_keys = ('dangerous_no_tls_verify',)
    for bool_key in bool_keys:
        if isinstance(config[bool_key], string_type):
            try:
                config[bool_key] = strtobool(config[bool_key])
            except ValueError:
                pass

    return config
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import numpy as np
from enum import Enum as BaseEnum

ENUM_ARRAY_DTYPE = np.int16

class Enum(BaseEnum):
    """
        Enum based on `enum34 <https://pypi.python.org/pypi/enum34/>`_, whose items have an index.
    """

    # Tweak enums to add an index attribute to each enum item
    def __init__(self, name):
        # When the enum item is initialized, self._member_names_ contains the names of the previously initialized items, so its length is the index of this item.
        self.index = len(self._member_names_)

    # Bypass the slow Enum.__eq__
    __eq__ = object.__eq__

    __hash__ = object.__hash__  # In Python 3, __hash__ must be defined if __eq__ is defined to stay hashable

    @classmethod
    def encode(cls, array):
        """
            Encode a string numpy array, or an enum item numpy array, into an :any:`EnumArray`. See :any:`EnumArray.decode` for decoding.

            :param numpy.ndarray array: Numpy array of string identifiers, or of enum items, to encode.

            :returns: An :any:`EnumArray` encoding the input array values.
            :rtype: :any:`EnumArray`

            For instance:

            >>> string_identifier_array = numpy.asarray(['free_lodger', 'owner'])
            >>> encoded_array = HousingOccupancyStatus.encode(string_identifier_array)
            >>> encoded_array[0]
            >>> 2  # Encoded value

            >>> enum_item_array = numpy.asarray([HousingOccupancyStatus.free_lodger, HousingOccupancyStatus.owner])
            >>> encoded_array = HousingOccupancyStatus.encode(enum_item_array)
            >>> encoded_array[0]
            >>> 2  # Encoded value

        """
        if type(array) is EnumArray:
            return array
        if array.dtype.kind in {'U', 'S'}:  # String array
            array = np.select([array == item.name for item in cls], [item.index for item in cls]).astype(ENUM_ARRAY_DTYPE)
        elif array.dtype.kind == 'O':  # Enum items arrays
            # Ensure we are comparing the comparable. The problem this fixes:
            #  On entering "cls" this method will generally come from variable.possible_values,
            #  while the array values may come from directly importing a module containing an Enum class.
            #  However, variables (and hence their possible_values) are loaded by a call to load_module,
            #  which gives them a different identity from the ones imported in the usual way.
            #  So, instead of relying on the "cls" passed in, we use only its name to check that
            #  the values in the array, if non-empty, are of the right type.
            if len(array) > 0 and cls.__name__ is array[0].__class__.__name__:
                cls = array[0].__class__
            array = np.select([array == item for item in cls], [item.index for item in cls]).astype(ENUM_ARRAY_DTYPE)
        return EnumArray(array, cls)

class EnumArray(np.ndarray):
    """
        Numpy array subclass representing an array of enum items.

        EnumArrays are encoded as ``int`` arrays to improve performance
    """

    # Subclassing np.ndarray is a little tricky. To read more about the two following methods, see https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html#slightly-more-realistic-example-attribute-added-to-existing-array.
    def __new__(cls, input_array, possible_values = None):
        obj = np.asarray(input_array).view(cls)
        obj.possible_values = possible_values
        return obj

    # See previous comment
    def __array_finalize__(self, obj):
        if obj is None:
            return
        self.possible_values = getattr(obj, 'possible_values', None)

    def __eq__(self, other):
        # When comparing to an item of self.possible_values, use the item index to speed up the comparison
        if other.__class__.__name__ is self.possible_values.__name__:
            return self.view(np.ndarray) == other.index  # use view(np.ndarray) so that the result is a classic ndarray, not an EnumArray
        return self.view(np.ndarray) == other

    def __ne__(self, other):
        return np.logical_not(self == other)

    def _forbidden_operation(self, other):
        raise TypeError("Forbidden operation. The only operations allowed on EnumArrays are '==' and '!='.")

    __add__ = _forbidden_operation
    __mul__ = _forbidden_operation
    __lt__ = _forbidden_operation
    __le__ = _forbidden_operation
    __gt__ = _forbidden_operation
    __ge__ = _forbidden_operation
    __and__ = _forbidden_operation
    __or__ = _forbidden_operation

    def decode(self):
        """
            Return the array of enum items corresponding to self

            >>> enum_array = household('housing_occupancy_status', period)
            >>> enum_array[0]
            >>> 2  # Encoded value
            >>> enum_array.decode()[0]
            >>> <HousingOccupancyStatus.free_lodger: 'Free lodger'>  # Decoded value : enum item
        """
        return np.select([self == item.index for item in self.possible_values], [item for item in self.possible_values])

    def decode_to_str(self):
        """
            Return the array of string identifiers corresponding to self

            >>> enum_array = household('housing_occupancy_status', period)
            >>> enum_array[0]
            >>> 2  # Encoded value
            >>> enum_array.decode_to_str()[0]
            >>> 'free_lodger' # String identifier
        """
        return np.select([self == item.index for item in self.possible_values], [item.name for item in self.possible_values])

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, str(self.decode()))
<commit_msg>Fix swapped words in comment<commit_after># -*- coding: utf-8 -*-

import numpy as np
from enum import Enum as BaseEnum

ENUM_ARRAY_DTYPE = np.int16

class Enum(BaseEnum):
    """
        Enum based on `enum34 <https://pypi.python.org/pypi/enum34/>`_, whose items have an index.
    """

    # Tweak enums to add an index attribute to each enum item
    def __init__(self, name):
        # When the enum item is initialized, self._member_names_ contains the names of the previously initialized items, so its length is the index of this item.
        self.index = len(self._member_names_)

    # Bypass the slow Enum.__eq__
    __eq__ = object.__eq__

    __hash__ = object.__hash__  # In Python 3, __hash__ must be defined if __eq__ is defined to stay hashable

    @classmethod
    def encode(cls, array):
        """
            Encode a string numpy array, or an enum item numpy array, into an :any:`EnumArray`. See :any:`EnumArray.decode` for decoding.

            :param numpy.ndarray array: Numpy array of string identifiers, or of enum items, to encode.

            :returns: An :any:`EnumArray` encoding the input array values.
            :rtype: :any:`EnumArray`

            For instance:

            >>> string_identifier_array = numpy.asarray(['free_lodger', 'owner'])
            >>> encoded_array = HousingOccupancyStatus.encode(string_identifier_array)
            >>> encoded_array[0]
            >>> 2  # Encoded value

            >>> enum_item_array = numpy.asarray([HousingOccupancyStatus.free_lodger, HousingOccupancyStatus.owner])
            >>> encoded_array = HousingOccupancyStatus.encode(enum_item_array)
            >>> encoded_array[0]
            >>> 2  # Encoded value

        """
        if type(array) is EnumArray:
            return array
        if array.dtype.kind in {'U', 'S'}:  # String array
            array = np.select([array == item.name for item in cls], [item.index for item in cls]).astype(ENUM_ARRAY_DTYPE)
        elif array.dtype.kind == 'O':  # Enum items arrays
            # Ensure we are comparing the comparable. The problem this fixes:
            #  On entering this method "cls" will generally come from variable.possible_values,
            #  while the array values may come from directly importing a module containing an Enum class.
            #  However, variables (and hence their possible_values) are loaded by a call to load_module,
            #  which gives them a different identity from the ones imported in the usual way.
            #  So, instead of relying on the "cls" passed in, we use only its name to check that
            #  the values in the array, if non-empty, are of the right type.
            if len(array) > 0 and cls.__name__ is array[0].__class__.__name__:
                cls = array[0].__class__
            array = np.select([array == item for item in cls], [item.index for item in cls]).astype(ENUM_ARRAY_DTYPE)
        return EnumArray(array, cls)

class EnumArray(np.ndarray):
    """
        Numpy array subclass representing an array of enum items.

        EnumArrays are encoded as ``int`` arrays to improve performance
    """

    # Subclassing np.ndarray is a little tricky. To read more about the two following methods, see https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html#slightly-more-realistic-example-attribute-added-to-existing-array.
    def __new__(cls, input_array, possible_values = None):
        obj = np.asarray(input_array).view(cls)
        obj.possible_values = possible_values
        return obj

    # See previous comment
    def __array_finalize__(self, obj):
        if obj is None:
            return
        self.possible_values = getattr(obj, 'possible_values', None)

    def __eq__(self, other):
        # When comparing to an item of self.possible_values, use the item index to speed up the comparison
        if other.__class__.__name__ is self.possible_values.__name__:
            return self.view(np.ndarray) == other.index  # use view(np.ndarray) so that the result is a classic ndarray, not an EnumArray
        return self.view(np.ndarray) == other

    def __ne__(self, other):
        return np.logical_not(self == other)

    def _forbidden_operation(self, other):
        raise TypeError("Forbidden operation. The only operations allowed on EnumArrays are '==' and '!='.")

    __add__ = _forbidden_operation
    __mul__ = _forbidden_operation
    __lt__ = _forbidden_operation
    __le__ = _forbidden_operation
    __gt__ = _forbidden_operation
    __ge__ = _forbidden_operation
    __and__ = _forbidden_operation
    __or__ = _forbidden_operation

    def decode(self):
        """
            Return the array of enum items corresponding to self

            >>> enum_array = household('housing_occupancy_status', period)
            >>> enum_array[0]
            >>> 2  # Encoded value
            >>> enum_array.decode()[0]
            >>> <HousingOccupancyStatus.free_lodger: 'Free lodger'>  # Decoded value : enum item
        """
        return np.select([self == item.index for item in self.possible_values], [item for item in self.possible_values])

    def decode_to_str(self):
        """
            Return the array of string identifiers corresponding to self

            >>> enum_array = household('housing_occupancy_status', period)
            >>> enum_array[0]
            >>> 2  # Encoded value
            >>> enum_array.decode_to_str()[0]
            >>> 'free_lodger' # String identifier
        """
        return np.select([self == item.index for item in self.possible_values], [item.name for item in self.possible_values])

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, str(self.decode()))
<|endoftext|>


<commit_before>from openid.test import datadriven

import unittest

from openid.message import Message, BARE_NS, OPENID_NS, OPENID2_NS
from openid import association
import time
from openid import cryptutil
import warnings

class AssociationSerializationTest(unittest.TestCase):
    def test_roundTrip(self):
        issued = int(time.time())
        lifetime = 600
        assoc = association.Association(
            'handle', 'secret', issued, lifetime, 'HMAC-SHA1')
        s = assoc.serialize()
        assoc2 = association.Association.deserialize(s)
        self.failUnlessEqual(assoc.handle, assoc2.handle)
        self.failUnlessEqual(assoc.issued, assoc2.issued)
        self.failUnlessEqual(assoc.secret, assoc2.secret)
        self.failUnlessEqual(assoc.lifetime, assoc2.lifetime)
        self.failUnlessEqual(assoc.assoc_type, assoc2.assoc_type)

from openid.server.server import \
     DiffieHellmanSHA1ServerSession, \
     DiffieHellmanSHA256ServerSession, \
     PlainTextServerSession

from openid.consumer.consumer import \
     DiffieHellmanSHA1ConsumerSession, \
     DiffieHellmanSHA256ConsumerSession, \
     PlainTextConsumerSession

from openid.dh import DiffieHellman

def createNonstandardConsumerDH():
    nonstandard_dh = DiffieHellman(1315291, 2)
    return DiffieHellmanSHA1ConsumerSession(nonstandard_dh)

class DiffieHellmanSessionTest(datadriven.DataDrivenTestCase):
    secrets = [
        '\x00' * 20,
        '\xff' * 20,
        ' ' * 20,
        'This is a secret....',
        ]

    session_factories = [
        (DiffieHellmanSHA1ConsumerSession, DiffieHellmanSHA1ServerSession),
        (createNonstandardConsumerDH, DiffieHellmanSHA1ServerSession),
        (PlainTextConsumerSession, PlainTextServerSession),
        ]

    def generateCases(cls):
        return [(c, s, sec)
                for c, s in cls.session_factories
                for sec in cls.secrets]

    generateCases = classmethod(generateCases)

    def __init__(self, csess_fact, ssess_fact, secret):
        datadriven.DataDrivenTestCase.__init__(self, csess_fact.__name__)
        self.secret = secret
        self.csess_fact = csess_fact
        self.ssess_fact = ssess_fact

    def runOneTest(self):
        csess = self.csess_fact()
        msg = Message.fromOpenIDArgs(csess.getRequest())
        ssess = self.ssess_fact.fromMessage(msg)
        check_secret = csess.extractSecret(
            Message.fromOpenIDArgs(ssess.answer(self.secret)))
        self.failUnlessEqual(self.secret, check_secret)

class TestMakePairs(unittest.TestCase):
    """Check the key-value formatting methods of associations.
    """

    def setUp(self):
        self.message = m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {
            'mode': 'id_res',
            'identifier': '=example',
            'signed': 'identifier,mode',
            'sig': 'cephalopod',
            })
        m.updateArgs(BARE_NS, {'xey': 'value'})
        self.assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")

    def testMakePairs(self):
        """Make pairs using the OpenID 1.x type signed list."""
        pairs = self.assoc._makePairs(self.message)
        expected = [
            ('identifier', '=example'),
            ('mode', 'id_res'),
            ]
        self.failUnlessEqual(pairs, expected)

class TestMac(unittest.TestCase):
    def setUp(self):
        self.pairs = [('key1', 'value1'),
                      ('key2', 'value2')]

    def test_sha1(self):
        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")
        expected = ('\xe0\x1bv\x04\xf1G\xc0\xbb\x7f\x9a\x8b'
                    '\xe9\xbc\xee}\\\xe5\xbb7*')
        sig = assoc.sign(self.pairs)
        self.failUnlessEqual(sig, expected)

    if cryptutil.SHA256_AVAILABLE:
        def test_sha256(self):
            assoc = association.Association.fromExpiresIn(
                3600, '{sha256SA}', 'very_secret', "HMAC-SHA256")
            expected = ('\xfd\xaa\xfe;\xac\xfc*\x988\xad\x05d6-\xeaVy'
                        '\xd5\xa5Z.<\xa9\xed\x18\x82\\$\x95x\x1c&')
            sig = assoc.sign(self.pairs)
            self.failUnlessEqual(sig, expected)

class TestMessageSigning(unittest.TestCase):
    def setUp(self):
        self.message = m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {'mode': 'id_res',
                                  'identifier': '=example'})
        m.updateArgs(BARE_NS, {'xey': 'value'})
        self.args = {'openid.mode': 'id_res',
                     'openid.identifier': '=example',
                     'xey': 'value'}

    def test_signSHA1(self):
        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")
        signed = assoc.signMessage(self.message)
        self.failUnless(signed.getArg(OPENID_NS, "sig"))
        self.failUnlessEqual(signed.getArg(OPENID_NS, "signed"),
                             "assoc_handle,identifier,mode,ns,signed")
        self.failUnlessEqual(signed.getArg(BARE_NS, "xey"), "value",
                             signed)

    if cryptutil.SHA256_AVAILABLE:
        def test_signSHA256(self):
            assoc = association.Association.fromExpiresIn(
                3600, '{sha1}', 'very_secret', "HMAC-SHA256")
            signed = assoc.signMessage(self.message)
            self.failUnless(signed.getArg(OPENID_NS, "sig"))
            self.failUnlessEqual(signed.getArg(OPENID_NS, "signed"),
                                 "assoc_handle,identifier,mode,ns,signed")
            self.failUnlessEqual(signed.getArg(BARE_NS, "xey"), "value",
                                 signed)

class TestCheckMessageSignature(unittest.TestCase):
    def test_aintGotSignedList(self):
        m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {'mode': 'id_res',
                                  'identifier': '=example',
                                  'sig': 'coyote',
                                  })
        m.updateArgs(BARE_NS, {'xey': 'value'})
        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")
        self.failUnlessRaises(ValueError, assoc.checkMessageSignature, m)

    def test_checkMethodSignature(self):
        m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {'mode': 'id_res',
                                  'identifier': '=example',
                                  })
        m.updateArgs(BARE_NS, {'xey': 'value'})

        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")

        # check with valid signature
        m = assoc.signMessage(m)
        passed = assoc.checkMessageSignature(m)
        self.failUnless(passed, "valid signature failed")

        # check with invalid signature
        m.updateArgs(OPENID2_NS, {'sig': 'badkey'})
        passed = assoc.checkMessageSignature(m)
        self.failUnless(passed == 0, "invalid signature passed")

def pyUnitTests():
    return datadriven.loadTests(__name__)

if __name__ == '__main__':
    suite = pyUnitTests()
    runner = unittest.TextTestRunner()
    runner.run(suite)
<commit_msg>fixed test method name<commit_after>from openid.test import datadriven

import unittest

from openid.message import Message, BARE_NS, OPENID_NS, OPENID2_NS
from openid import association
import time
from openid import cryptutil
import warnings

class AssociationSerializationTest(unittest.TestCase):
    def test_roundTrip(self):
        issued = int(time.time())
        lifetime = 600
        assoc = association.Association(
            'handle', 'secret', issued, lifetime, 'HMAC-SHA1')
        s = assoc.serialize()
        assoc2 = association.Association.deserialize(s)
        self.failUnlessEqual(assoc.handle, assoc2.handle)
        self.failUnlessEqual(assoc.issued, assoc2.issued)
        self.failUnlessEqual(assoc.secret, assoc2.secret)
        self.failUnlessEqual(assoc.lifetime, assoc2.lifetime)
        self.failUnlessEqual(assoc.assoc_type, assoc2.assoc_type)

from openid.server.server import \
     DiffieHellmanSHA1ServerSession, \
     DiffieHellmanSHA256ServerSession, \
     PlainTextServerSession

from openid.consumer.consumer import \
     DiffieHellmanSHA1ConsumerSession, \
     DiffieHellmanSHA256ConsumerSession, \
     PlainTextConsumerSession

from openid.dh import DiffieHellman

def createNonstandardConsumerDH():
    nonstandard_dh = DiffieHellman(1315291, 2)
    return DiffieHellmanSHA1ConsumerSession(nonstandard_dh)

class DiffieHellmanSessionTest(datadriven.DataDrivenTestCase):
    secrets = [
        '\x00' * 20,
        '\xff' * 20,
        ' ' * 20,
        'This is a secret....',
        ]

    session_factories = [
        (DiffieHellmanSHA1ConsumerSession, DiffieHellmanSHA1ServerSession),
        (createNonstandardConsumerDH, DiffieHellmanSHA1ServerSession),
        (PlainTextConsumerSession, PlainTextServerSession),
        ]

    def generateCases(cls):
        return [(c, s, sec)
                for c, s in cls.session_factories
                for sec in cls.secrets]

    generateCases = classmethod(generateCases)

    def __init__(self, csess_fact, ssess_fact, secret):
        datadriven.DataDrivenTestCase.__init__(self, csess_fact.__name__)
        self.secret = secret
        self.csess_fact = csess_fact
        self.ssess_fact = ssess_fact

    def runOneTest(self):
        csess = self.csess_fact()
        msg = Message.fromOpenIDArgs(csess.getRequest())
        ssess = self.ssess_fact.fromMessage(msg)
        check_secret = csess.extractSecret(
            Message.fromOpenIDArgs(ssess.answer(self.secret)))
        self.failUnlessEqual(self.secret, check_secret)

class TestMakePairs(unittest.TestCase):
    """Check the key-value formatting methods of associations.
    """

    def setUp(self):
        self.message = m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {
            'mode': 'id_res',
            'identifier': '=example',
            'signed': 'identifier,mode',
            'sig': 'cephalopod',
            })
        m.updateArgs(BARE_NS, {'xey': 'value'})
        self.assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")

    def testMakePairs(self):
        """Make pairs using the OpenID 1.x type signed list."""
        pairs = self.assoc._makePairs(self.message)
        expected = [
            ('identifier', '=example'),
            ('mode', 'id_res'),
            ]
        self.failUnlessEqual(pairs, expected)

class TestMac(unittest.TestCase):
    def setUp(self):
        self.pairs = [('key1', 'value1'),
                      ('key2', 'value2')]

    def test_sha1(self):
        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")
        expected = ('\xe0\x1bv\x04\xf1G\xc0\xbb\x7f\x9a\x8b'
                    '\xe9\xbc\xee}\\\xe5\xbb7*')
        sig = assoc.sign(self.pairs)
        self.failUnlessEqual(sig, expected)

    if cryptutil.SHA256_AVAILABLE:
        def test_sha256(self):
            assoc = association.Association.fromExpiresIn(
                3600, '{sha256SA}', 'very_secret', "HMAC-SHA256")
            expected = ('\xfd\xaa\xfe;\xac\xfc*\x988\xad\x05d6-\xeaVy'
                        '\xd5\xa5Z.<\xa9\xed\x18\x82\\$\x95x\x1c&')
            sig = assoc.sign(self.pairs)
            self.failUnlessEqual(sig, expected)

class TestMessageSigning(unittest.TestCase):
    def setUp(self):
        self.message = m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {'mode': 'id_res',
                                  'identifier': '=example'})
        m.updateArgs(BARE_NS, {'xey': 'value'})
        self.args = {'openid.mode': 'id_res',
                     'openid.identifier': '=example',
                     'xey': 'value'}

    def test_signSHA1(self):
        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")
        signed = assoc.signMessage(self.message)
        self.failUnless(signed.getArg(OPENID_NS, "sig"))
        self.failUnlessEqual(signed.getArg(OPENID_NS, "signed"),
                             "assoc_handle,identifier,mode,ns,signed")
        self.failUnlessEqual(signed.getArg(BARE_NS, "xey"), "value",
                             signed)

    if cryptutil.SHA256_AVAILABLE:
        def test_signSHA256(self):
            assoc = association.Association.fromExpiresIn(
                3600, '{sha1}', 'very_secret', "HMAC-SHA256")
            signed = assoc.signMessage(self.message)
            self.failUnless(signed.getArg(OPENID_NS, "sig"))
            self.failUnlessEqual(signed.getArg(OPENID_NS, "signed"),
                                 "assoc_handle,identifier,mode,ns,signed")
            self.failUnlessEqual(signed.getArg(BARE_NS, "xey"), "value",
                                 signed)

class TestCheckMessageSignature(unittest.TestCase):
    def test_aintGotSignedList(self):
        m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {'mode': 'id_res',
                                  'identifier': '=example',
                                  'sig': 'coyote',
                                  })
        m.updateArgs(BARE_NS, {'xey': 'value'})
        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")
        self.failUnlessRaises(ValueError, assoc.checkMessageSignature, m)

    def test_checkMessageSignature(self):
        m = Message(OPENID2_NS)
        m.updateArgs(OPENID2_NS, {'mode': 'id_res',
                                  'identifier': '=example',
                                  })
        m.updateArgs(BARE_NS, {'xey': 'value'})

        assoc = association.Association.fromExpiresIn(
            3600, '{sha1}', 'very_secret', "HMAC-SHA1")

        # check with valid signature
        m = assoc.signMessage(m)
        passed = assoc.checkMessageSignature(m)
        self.failUnless(passed, "valid signature failed")

        # check with invalid signature
        m.updateArgs(OPENID2_NS, {'sig': 'badkey'})
        passed = assoc.checkMessageSignature(m)
        self.failUnless(passed == 0, "invalid signature passed")

def pyUnitTests():
    return datadriven.loadTests(__name__)

if __name__ == '__main__':
    suite = pyUnitTests()
    runner = unittest.TextTestRunner()
    runner.run(suite)
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Created on Thu Sep 11 15:29:27 2014

@author: John Swoboda
"""
import pdb
import numpy as np
from tables import *

def readMad_hdf5 (filename, paramstr): #timelims=None
    """@author: Anna Stuhlmacher
    madgrigal h5 read in function for the python implementation of GeoData for Madrigal Sondrestrom data
    Input:
    filename path to hdf5 file
    list of parameters to look at written as strings
    Returns:
    dictionary with keys are the Madrigal parameter string, the value is an array
    rows are unique data locations (data_loc) = (alt, el1, azm)
    columns are unique times
    """
    #open hdf5 file
    files = openFile(filename, mode = "r", title = 'Sondrestrom1')
    all_data = files.getNode('/Data/Table Layout').read()
    sensor_data = files.getNode('/Metadata/Experiment Parameters').read()
    files.close()

    instrument = sensor_data[0][1] #instrument type string
    if "Sondrestrom" in instrument:
        radar = 1
        print("Sondrestrom data")
    elif "Poker Flat" in instrument:
        radar = 2
        print("PFISR data")
    elif "Resolute Bay" in instrument:
        radar = 3
        print ("RISR data")
    else:
        print("Error: Radar type not supported by program in this version.")
        exit()

    # get the data location (range, el1, azm)
    all_loc = []
    if radar == 1:
        angle1 = 'elm'
        alt = all_data['gdalt']
    elif radar == 2:
        angle1 = 'elm'
        alt = all_data['range']
    elif radar ==3:
        angle1 = 'elm'
        alt = all_data['range']

    try:
        el = all_data[angle1]
    except ValueError:
        el = np.nan * np.ones(len(list(alt)))

    try:
        azm = all_data['azm']
    except ValueError:
        azm = np.nan * np.ones(len(list(alt)))
    # take out nans
    nan_ar = np.isnan(alt)|np.isnan(el)|np.isnan(azm)
    notnan = np.logical_not(nan_ar)
    for i in range(len(alt)):
        if notnan[i]:
            all_loc.append([alt[i], el[i], azm[i]])

    #create list of unique data location lists
    dataloc = [list(y) for y in set([tuple(x) for x in all_loc])]
    all_times = []
    times1 = all_data['ut1_unix'][notnan]
    times2 = all_data['ut2_unix'][notnan]
    for i in range(len(times1)):
        all_times.append([times1[i], times2[i]])
    uniq_times = sorted(set(tuple(x) for x in all_times))

    #initialize and fill data dictionary with parameter arrays
    data = {}
    maxcols = len(uniq_times)
    maxrows = len(dataloc)
    for p in paramstr:
        if not p in all_data.dtype.names:
            print 'Warning: ' + p + ' is not a valid parameter name.'
            continue
        tempdata = all_data[p][notnan] #list of parameter pulled from all_data
        temparray = np.empty([maxrows,maxcols]) #converting the tempdata list into array form

        for t in range(len(tempdata)):
            #row
            row = dataloc.index(all_loc[t])
            #column-time
            col = uniq_times.index(tuple(all_times[t]))
            temparray[row][col] = tempdata[t]

        data[p]=temparray

    #get the sensor location (lat, long, alt)
    lat = sensor_data[7][1]
    lon = sensor_data[8][1]
    sensor_alt = sensor_data[9][1]
    sensorloc = np.array([lat,lon,sensor_alt], dtype='f')
    coordnames = 'Spherical'

    return (data,coordnames,np.array(dataloc, dtype='f'),sensorloc,np.asarray(uniq_times, dtype='f'))

#data, coordnames, dataloc, sensorloc, times = readMad_hdf5('/Users/anna/Research/Ionosphere/2008WorldDaysPDB/son081001g.001.hdf5', ['ti', 'dti', 'nel'])
<commit_msg>updated utilityfuncs to common way of setting range and angles.<commit_after>#!/usr/bin/env python
"""
Created on Thu Sep 11 15:29:27 2014

@author: John Swoboda
"""
import pdb
import numpy as np
from tables import *

def readMad_hdf5 (filename, paramstr): #timelims=None
    """@author: Anna Stuhlmacher
    madgrigal h5 read in function for the python implementation of GeoData for Madrigal Sondrestrom data
    Input:
    filename path to hdf5 file
    list of parameters to look at written as strings
    Returns:
    dictionary with keys are the Madrigal parameter string, the value is an array
    rows are unique data locations (data_loc) = (rng, azm, el1)
    columns are unique times
    """
    #open hdf5 file
    files = openFile(filename, mode = "r", title = 'Sondrestrom1')
    all_data = files.getNode('/Data/Table Layout').read()
    sensor_data = files.getNode('/Metadata/Experiment Parameters').read()
    files.close()

    instrument = sensor_data[0][1] #instrument type string
    if "Sondrestrom" in instrument:
        radar = 1
        print("Sondrestrom data")
    elif "Poker Flat" in instrument:
        radar = 2
        print("PFISR data")
    elif "Resolute Bay" in instrument:
        radar = 3
        print ("RISR data")
    else:
        print("Error: Radar type not supported by program in this version.")
        exit()

    # get the data location (range, el1, azm)
    all_loc = []
    if radar == 1:
        angle1 = 'elm'
        rng = all_data['gdalt']
    elif radar == 2:
        angle1 = 'elm'
        rng = all_data['range']
    elif radar ==3:
        angle1 = 'elm'
        rng = all_data['range']

    try:
        el = all_data[angle1]
    except ValueError:
        el = np.nan * np.ones(len(list(rng)))

    try:
        azm = all_data['azm']
    except ValueError:
        azm = np.nan * np.ones(len(list(rng)))
    # take out nans
    nan_ar = np.isnan(rng)|np.isnan(el)|np.isnan(azm)
    notnan = np.logical_not(nan_ar)
    for i in range(len(rng)):
        if notnan[i]:
            all_loc.append([rng[i],azm[i],el[i]])

    #create list of unique data location lists
    dataloc = [list(y) for y in set([tuple(x) for x in all_loc])]
    all_times = []
    times1 = all_data['ut1_unix'][notnan]
    times2 = all_data['ut2_unix'][notnan]
    for i in range(len(times1)):
        all_times.append([times1[i], times2[i]])
    uniq_times = sorted(set(tuple(x) for x in all_times))

    #initialize and fill data dictionary with parameter arrays
    data = {}
    maxcols = len(uniq_times)
    maxrows = len(dataloc)
    for p in paramstr:
        if not p in all_data.dtype.names:
            print 'Warning: ' + p + ' is not a valid parameter name.'
            continue
        tempdata = all_data[p][notnan] #list of parameter pulled from all_data
        temparray = np.empty([maxrows,maxcols]) #converting the tempdata list into array form

        for t in range(len(tempdata)):
            #row
            row = dataloc.index(all_loc[t])
            #column-time
            col = uniq_times.index(tuple(all_times[t]))
            temparray[row][col] = tempdata[t]

        data[p]=temparray

    #get the sensor location (lat, long, rng)
    lat = sensor_data[7][1]
    lon = sensor_data[8][1]
    sensor_alt = sensor_data[9][1]
    sensorloc = np.array([lat,lon,sensor_alt], dtype='f')
    coordnames = 'Spherical'

    return (data,coordnames,np.array(dataloc, dtype='f'),sensorloc,np.asarray(uniq_times, dtype='f'))

#data, coordnames, dataloc, sensorloc, times = readMad_hdf5('/Users/anna/Research/Ionosphere/2008WorldDaysPDB/son081001g.001.hdf5', ['ti', 'dti', 'nel'])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2017-2022, 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 signal
import shutil
import getpass
import tempfile
import subprocess
import multiprocessing
import psutil
from openquake.baselib import (
    zeromq as z, general, parallel, config, sap, InvalidFile)
try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(title):
        "Do nothing"

def _streamer():
    # streamer for zmq workers running on the master node
    port = int(config.zworkers.ctrl_port)
    task_input_url = 'tcp://0.0.0.0:%d' % (port + 2)
    task_output_url = 'tcp://%s:%s' % (config.dbserver.listen, port + 1)
    try:
        z.zmq.proxy(z.bind(task_input_url, z.zmq.PULL),
                    z.bind(task_output_url, z.zmq.PUSH))
    except (KeyboardInterrupt, z.zmq.ContextTerminated):
        pass  # killed cleanly by SIGINT/SIGTERM

def ssh_args(zworkers):
    """
    :yields: triples (hostIP, num_cores, [ssh remote python command])
    """
    remote_python = zworkers.remote_python or sys.executable
    remote_user = zworkers.remote_user or getpass.getuser()
    if zworkers.host_cores.strip():
        for hostcores in config.zworkers.host_cores.split(','):
            host, cores = hostcores.split()
            if host == '127.0.0.1':  # localhost
                yield host, cores, [sys.executable]
            else:
                yield host, cores, [
                    'ssh', '-f', '-T', remote_user + '@' + host, remote_python]

class WorkerMaster(object):
    """
    :param ctrl_port: port on which the worker pools listen
    :param host_cores: names of the remote hosts and number of cores to use
    :param remote_python: path of the Python executable on the remote hosts
    """
    def __init__(self, zworkers=config.zworkers, receiver_ports=None):
        self.zworkers = zworkers
        # NB: receiver_ports is not used but needed for compliance
        self.ctrl_port = int(zworkers.ctrl_port)
        self.host_cores = ([hc.split() for hc in zworkers.host_cores.split(',')]
                           if zworkers.host_cores else [])
        for host, cores in self.host_cores:
            if int(cores) < -1:
                raise InvalidFile('openquake.cfg: found %s %s' %
                                  (host, cores))
        self.remote_python = zworkers.remote_python or sys.executable
        self.remote_user = zworkers.remote_user or getpass.getuser()
        self.popens = []

    def start(self):
        """
        Start multiple workerpools, possibly on remote servers via ssh,
        assuming there is an active streamer.
        """
        starting = []
        for host, cores, args in ssh_args(self.zworkers):
            if general.socket_ready((host, self.ctrl_port)):
                print('%s:%s already running' % (host, self.ctrl_port))
                continue
            ctrl_url = 'tcp://0.0.0.0:%s' % self.ctrl_port
            args += ['-m', 'openquake.baselib.workerpool', ctrl_url,
                     '-n', cores]
            if host != '127.0.0.1':
                print('%s: if it hangs, check the ssh keys' % ' '.join(args))
            self.popens.append(subprocess.Popen(args))
            starting.append(host)
        return 'starting %s' % starting

    def stop(self):
        """
        Send a "stop" command to all worker pools
        """
        stopped = []
        for host, _ in self.host_cores:
            if not general.socket_ready((host, self.ctrl_port)):
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                sock.send('stop')
                stopped.append(host)
        for popen in self.popens:
            popen.terminate()
            # since we are not consuming any output from the spawned process
            # we must call wait() after terminate() to have Popen()
            # fully deallocate the process file descriptors, otherwise
            # zombies will arise
            popen.wait()
        self.popens = []
        return 'stopped %s' % stopped

    def kill(self):
        """
        Send a "killall" command to all worker pools to cleanup everything
        in case of hard out of memory situations
        """
        killed = []
        for host, cores, args in ssh_args(self.zworkers):
            args = args[:-1] + ['killall', '-u', 'openquake']
            if host != '127.0.0.1':
                print(' '.join(args))
            subprocess.run(args)
            killed.append(host)
        return 'killed %s' % killed

    def status(self):
        """
        :returns: a list [(host, running, total), ...]
        """
        executing = []
        for host, _cores in self.host_cores:
            if not general.socket_ready((host, self.ctrl_port)):
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                running = len(sock.send('get_executing').split())
                total = sock.send('get_num_workers')
                executing.append((host, running, total))
        return executing

    def restart(self):
        """
        Stop and start again
        """
        self.stop()
        self.start()
        return 'restarted'

def worker(sock, executing):
    """
    :param sock: a zeromq.Socket of kind PULL
    :param executing: a path inside /tmp/calc_XXX
    """
    setproctitle('oq-zworker')
    with sock:
        for cmd, args, taskno, mon in sock:
            fname = os.path.join(executing, '%s-%s' % (mon.calc_id, taskno))
            # NB: very hackish way of keeping track of the running tasks,
            # used in get_executing, could litter the file system
            open(fname, 'w').close()
            parallel.safely_call(cmd, args, taskno, mon)
            os.remove(fname)

class WorkerPool(object):
    """
    A pool of workers accepting the command 'stop' and 'kill' and reading
    tasks to perform from the task_server_url.

    :param ctrl_url: zmq address of the control socket
    :param num_workers: the number of workers (or -1)
    """
    def __init__(self, ctrl_url, num_workers=-1):
        self.ctrl_url = ctrl_url
        self.task_server_url = 'tcp://%s:%s' % (
            config.dbserver.host, int(config.zworkers.ctrl_port) + 1)
        if num_workers == -1:
            try:
                self.num_workers = len(psutil.Process().cpu_affinity())
            except AttributeError:  # missing cpu_affinity on macOS
                self.num_workers = psutil.cpu_count()
        else:
            self.num_workers = num_workers
        self.executing = tempfile.mkdtemp()
        self.pid = os.getpid()

    def start(self):
        """
        Start worker processes and a control loop
        """
        title = 'oq-zworkerpool %s' % self.ctrl_url[6:]  # strip tcp://
        print('Starting ' + title, file=sys.stderr)
        setproctitle(title)
        # start workers
        self.workers = []
        for _ in range(self.num_workers):
            sock = z.Socket(self.task_server_url, z.zmq.PULL, 'connect')
            sock.proc = multiprocessing.Process(
                target=worker, args=(sock, self.executing))
            sock.proc.start()
            self.workers.append(sock)

        # start control loop accepting the commands stop and kill
        with z.Socket(self.ctrl_url, z.zmq.REP, 'bind') as ctrlsock:
            for cmd in ctrlsock:
                if cmd in ('stop', 'kill'):
                    msg = getattr(self, cmd)()
                    ctrlsock.send(msg)
                    break
                elif cmd == 'getpid':
                    ctrlsock.send(self.proc.pid)
                elif cmd == 'get_num_workers':
                    ctrlsock.send(self.num_workers)
                elif cmd == 'get_executing':
                    ctrlsock.send(' '.join(sorted(os.listdir(self.executing))))
        shutil.rmtree(self.executing)

    def stop(self):
        """
        Send a SIGTERM to all worker processes
        """
        for sock in self.workers:
            os.kill(sock.proc.pid, signal.SIGTERM)
        for sock in self.workers:
            sock.proc.join()
        return 'WorkerPool %s stopped' % self.ctrl_url

    def kill(self):
        """
        Send a SIGKILL to all worker processes
        """
        for sock in self.workers:
            os.kill(sock.proc.pid, signal.SIGKILL)
        for sock in self.workers:
            sock.proc.join()
        return 'WorkerPool %s killed' % self.ctrl_url

def workerpool(worker_url='tcp://0.0.0.0:1909', *, num_workers: int = -1):
    # start a workerpool without a streamer
    WorkerPool(worker_url, num_workers).start()

workerpool.worker_url = dict(
    help='ZMQ address (tcp:///w.x.y.z:port) of the worker')
workerpool.num_workers = dict(help='number of cores to use')

if __name__ == '__main__':
    sap.run(workerpool)
<commit_msg>Fixed wrong usage of config.zworkers<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2017-2022, 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 signal
import shutil
import getpass
import tempfile
import subprocess
import multiprocessing
import psutil
from openquake.baselib import (
    zeromq as z, general, parallel, config, sap, InvalidFile)
try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(title):
        "Do nothing"

def _streamer():
    # streamer for zmq workers running on the master node
    port = int(config.zworkers.ctrl_port)
    task_input_url = 'tcp://0.0.0.0:%d' % (port + 2)
    task_output_url = 'tcp://%s:%s' % (config.dbserver.listen, port + 1)
    try:
        z.zmq.proxy(z.bind(task_input_url, z.zmq.PULL),
                    z.bind(task_output_url, z.zmq.PUSH))
    except (KeyboardInterrupt, z.zmq.ContextTerminated):
        pass  # killed cleanly by SIGINT/SIGTERM

def ssh_args(zworkers):
    """
    :yields: triples (hostIP, num_cores, [ssh remote python command])
    """
    remote_python = zworkers.remote_python or sys.executable
    remote_user = zworkers.remote_user or getpass.getuser()
    if zworkers.host_cores.strip():
        for hostcores in zworkers.host_cores.split(','):
            host, cores = hostcores.split()
            if host == '127.0.0.1':  # localhost
                yield host, cores, [sys.executable]
            else:
                yield host, cores, [
                    'ssh', '-f', '-T', remote_user + '@' + host, remote_python]

class WorkerMaster(object):
    """
    :param ctrl_port: port on which the worker pools listen
    :param host_cores: names of the remote hosts and number of cores to use
    :param remote_python: path of the Python executable on the remote hosts
    """
    def __init__(self, zworkers=config.zworkers, receiver_ports=None):
        self.zworkers = zworkers
        # NB: receiver_ports is not used but needed for compliance
        self.ctrl_port = int(zworkers.ctrl_port)
        self.host_cores = ([hc.split() for hc in zworkers.host_cores.split(',')]
                           if zworkers.host_cores else [])
        for host, cores in self.host_cores:
            if int(cores) < -1:
                raise InvalidFile('openquake.cfg: found %s %s' %
                                  (host, cores))
        self.remote_python = zworkers.remote_python or sys.executable
        self.remote_user = zworkers.remote_user or getpass.getuser()
        self.popens = []

    def start(self):
        """
        Start multiple workerpools, possibly on remote servers via ssh,
        assuming there is an active streamer.
        """
        starting = []
        for host, cores, args in ssh_args(self.zworkers):
            if general.socket_ready((host, self.ctrl_port)):
                print('%s:%s already running' % (host, self.ctrl_port))
                continue
            ctrl_url = 'tcp://0.0.0.0:%s' % self.ctrl_port
            args += ['-m', 'openquake.baselib.workerpool', ctrl_url,
                     '-n', cores]
            if host != '127.0.0.1':
                print('%s: if it hangs, check the ssh keys' % ' '.join(args))
            self.popens.append(subprocess.Popen(args))
            starting.append(host)
        return 'starting %s' % starting

    def stop(self):
        """
        Send a "stop" command to all worker pools
        """
        stopped = []
        for host, _ in self.host_cores:
            if not general.socket_ready((host, self.ctrl_port)):
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                sock.send('stop')
                stopped.append(host)
        for popen in self.popens:
            popen.terminate()
            # since we are not consuming any output from the spawned process
            # we must call wait() after terminate() to have Popen()
            # fully deallocate the process file descriptors, otherwise
            # zombies will arise
            popen.wait()
        self.popens = []
        return 'stopped %s' % stopped

    def kill(self):
        """
        Send a "killall" command to all worker pools to cleanup everything
        in case of hard out of memory situations
        """
        killed = []
        for host, cores, args in ssh_args(self.zworkers):
            args = args[:-1] + ['killall', '-u', 'openquake']
            if host != '127.0.0.1':
                print(' '.join(args))
            subprocess.run(args)
            killed.append(host)
        return 'killed %s' % killed

    def status(self):
        """
        :returns: a list [(host, running, total), ...]
        """
        executing = []
        for host, _cores in self.host_cores:
            if not general.socket_ready((host, self.ctrl_port)):
                continue
            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
            with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                running = len(sock.send('get_executing').split())
                total = sock.send('get_num_workers')
                executing.append((host, running, total))
        return executing

    def restart(self):
        """
        Stop and start again
        """
        self.stop()
        self.start()
        return 'restarted'

def worker(sock, executing):
    """
    :param sock: a zeromq.Socket of kind PULL
    :param executing: a path inside /tmp/calc_XXX
    """
    setproctitle('oq-zworker')
    with sock:
        for cmd, args, taskno, mon in sock:
            fname = os.path.join(executing, '%s-%s' % (mon.calc_id, taskno))
            # NB: very hackish way of keeping track of the running tasks,
            # used in get_executing, could litter the file system
            open(fname, 'w').close()
            parallel.safely_call(cmd, args, taskno, mon)
            os.remove(fname)

class WorkerPool(object):
    """
    A pool of workers accepting the command 'stop' and 'kill' and reading
    tasks to perform from the task_server_url.

    :param ctrl_url: zmq address of the control socket
    :param num_workers: the number of workers (or -1)
    """
    def __init__(self, ctrl_url, num_workers=-1):
        self.ctrl_url = ctrl_url
        self.task_server_url = 'tcp://%s:%s' % (
            config.dbserver.host, int(config.zworkers.ctrl_port) + 1)
        if num_workers == -1:
            try:
                self.num_workers = len(psutil.Process().cpu_affinity())
            except AttributeError:  # missing cpu_affinity on macOS
                self.num_workers = psutil.cpu_count()
        else:
            self.num_workers = num_workers
        self.executing = tempfile.mkdtemp()
        self.pid = os.getpid()

    def start(self):
        """
        Start worker processes and a control loop
        """
        title = 'oq-zworkerpool %s' % self.ctrl_url[6:]  # strip tcp://
        print('Starting ' + title, file=sys.stderr)
        setproctitle(title)
        # start workers
        self.workers = []
        for _ in range(self.num_workers):
            sock = z.Socket(self.task_server_url, z.zmq.PULL, 'connect')
            sock.proc = multiprocessing.Process(
                target=worker, args=(sock, self.executing))
            sock.proc.start()
            self.workers.append(sock)

        # start control loop accepting the commands stop and kill
        with z.Socket(self.ctrl_url, z.zmq.REP, 'bind') as ctrlsock:
            for cmd in ctrlsock:
                if cmd in ('stop', 'kill'):
                    msg = getattr(self, cmd)()
                    ctrlsock.send(msg)
                    break
                elif cmd == 'getpid':
                    ctrlsock.send(self.proc.pid)
                elif cmd == 'get_num_workers':
                    ctrlsock.send(self.num_workers)
                elif cmd == 'get_executing':
                    ctrlsock.send(' '.join(sorted(os.listdir(self.executing))))
        shutil.rmtree(self.executing)

    def stop(self):
        """
        Send a SIGTERM to all worker processes
        """
        for sock in self.workers:
            os.kill(sock.proc.pid, signal.SIGTERM)
        for sock in self.workers:
            sock.proc.join()
        return 'WorkerPool %s stopped' % self.ctrl_url

    def kill(self):
        """
        Send a SIGKILL to all worker processes
        """
        for sock in self.workers:
            os.kill(sock.proc.pid, signal.SIGKILL)
        for sock in self.workers:
            sock.proc.join()
        return 'WorkerPool %s killed' % self.ctrl_url

def workerpool(worker_url='tcp://0.0.0.0:1909', *, num_workers: int = -1):
    # start a workerpool without a streamer
    WorkerPool(worker_url, num_workers).start()

workerpool.worker_url = dict(
    help='ZMQ address (tcp:///w.x.y.z:port) of the worker')
workerpool.num_workers = dict(help='number of cores to use')

if __name__ == '__main__':
    sap.run(workerpool)
<|endoftext|>


<commit_before># =================================================================
#
# Authors: Tom Kralidis <tomkralidis@gmail.com>
#
# Copyright (c) 2014 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

import models
import util

def list_resources(resource_type=None, query=None):
    """return all resources"""

    reliability_values = []
    first_run = None
    last_run = None

    response = {
        'total': 0,
        'success': {
            'number': 0,
            'percentage': 0
        },
        'fail': {
            'number': 0,
            'percentage': 0
        },
        'first_run': None,
        'last_run': None,
        'reliability': 0
    }

    if resource_type is not None:
        response['resources'] = models.Resource.query.filter_by(
            resource_type=resource_type).all()
    if query is not None:
        field, term = get_query_field_term(query)
        response['resources'] = models.Resource.query.filter(
            field.ilike(term)).all()
    else:
        response['resources'] = models.Resource.query.all()

    response['total'] = len(response['resources'])
    for resource in response['resources']:
        if resource.first_run < first_run or first_run is None:
            first_run = resource.first_run
        if resource.last_run < last_run or last_run is None:
            last_run = resource.last_run
        response['first_run'] = first_run
        response['last_run'] = last_run
        if resource.last_run.success:
            response['success']['number'] += 1
        else:
            response['fail']['number'] += 1
        reliability_values.append(resource.reliability)

    response['success']['percentage'] = util.percentage(
        response['success']['number'], response['total'])
    response['fail']['percentage'] = util.percentage(
        response['fail']['number'], response['total'])
    response['reliability'] = util.average(reliability_values)

    return response

def get_resource_by_id(identifier):
    """return one resource by identifier"""
    return models.Resource.query.filter_by(
        identifier=identifier).first_or_404()

def get_resource_types_counts():
    """return frequency counts of registered resource types"""

    mrt = models.get_resource_types_counts()
    return {
        'counts': mrt[0],
        'total': mrt[1]
    }

def get_query_field_term(query):
    """determine query context from q="""

    field = models.Resource.title  # default

    try:
        facet, term = query.split(':')
        term2 = '%%%s%%' % term  # default like
        if facet == 'url':
            field = models.Resource.url
        elif facet == 'title':
            field = models.Resource.title
        elif facet == 'site':
            field = models.Resource.url
            term2 = '%%%s/%%' % term
        elif facet == 'owner':
            field = models.Resource.owner_identifier
        term = term2
    except ValueError:  # default search
        term = '%%%s%%' % query

    return [field, term]
<commit_msg>fix for issue #58, resource type filtering main view<commit_after># =================================================================
#
# Authors: Tom Kralidis <tomkralidis@gmail.com>
#
# Copyright (c) 2014 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

import models
import util

def list_resources(resource_type=None, query=None):
    """return all resources"""

    reliability_values = []
    first_run = None
    last_run = None

    response = {
        'total': 0,
        'success': {
            'number': 0,
            'percentage': 0
        },
        'fail': {
            'number': 0,
            'percentage': 0
        },
        'first_run': None,
        'last_run': None,
        'reliability': 0
    }

    if resource_type is not None:
        response['resources'] = models.Resource.query.filter_by(
            resource_type=resource_type).all()
        print 'resource_type=%s' % resource_type

    if query is not None:
        field, term = get_query_field_term(query)
        response['resources'] = models.Resource.query.filter(
            field.ilike(term)).all()

    # No query nor resource_type provided: fetch all resources
    if response['resources'] is None:
        response['resources'] = models.Resource.query.all()

    response['total'] = len(response['resources'])
    for resource in response['resources']:
        if resource.first_run < first_run or first_run is None:
            first_run = resource.first_run
        if resource.last_run < last_run or last_run is None:
            last_run = resource.last_run
        response['first_run'] = first_run
        response['last_run'] = last_run
        if resource.last_run.success:
            response['success']['number'] += 1
        else:
            response['fail']['number'] += 1
        reliability_values.append(resource.reliability)

    response['success']['percentage'] = util.percentage(
        response['success']['number'], response['total'])
    response['fail']['percentage'] = util.percentage(
        response['fail']['number'], response['total'])
    response['reliability'] = util.average(reliability_values)

    return response

def get_resource_by_id(identifier):
    """return one resource by identifier"""
    return models.Resource.query.filter_by(
        identifier=identifier).first_or_404()

def get_resource_types_counts():
    """return frequency counts of registered resource types"""

    mrt = models.get_resource_types_counts()
    return {
        'counts': mrt[0],
        'total': mrt[1]
    }

def get_query_field_term(query):
    """determine query context from q="""

    field = models.Resource.title  # default

    try:
        facet, term = query.split(':')
        term2 = '%%%s%%' % term  # default like
        if facet == 'url':
            field = models.Resource.url
        elif facet == 'title':
            field = models.Resource.title
        elif facet == 'site':
            field = models.Resource.url
            term2 = '%%%s/%%' % term
        elif facet == 'owner':
            field = models.Resource.owner_identifier
        term = term2
    except ValueError:  # default search
        term = '%%%s%%' % query

    return [field, term]
<|endoftext|>