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

from collections import defaultdict, OrderedDict

from dateutil.parser import  parse as parse_datetime
from lxml import etree

class LithiumType(object):
    def __init__(self, xml_tree, api):
        self.id = xml_tree.get('id')
        self.href = xml_tree.get('href')
        self._tree = xml_tree
        self._populated = False
        self._api = api

        for child in xml_tree:
            if hasattr(self, '_handle_%s' % child.tag):
                getattr(self, '_handle_%s' % child.tag)(child)
            else:
                setattr(self, child.tag, xml_to_type(child, api))

    def __getattr__(self, item):
        if item.startswith('_'):
            raise AttributeError(item)
        if self._api and self.href and not self._populated:
            obj = self._api.raw(self.href)
            self.__dict__.update(obj.__dict__)
        self._populated = True
        if item in self.__dict__:
            return getattr(self, item)
        else:
            raise AttributeError(item)

    def __repr__(self):
        return '<%s>: %s' % (self.__class__.__name__, pprint.pformat({k:v for k, v in self.__dict__.iteritems() if not k.startswith('_')}))

class Board(LithiumType):
    pass

class Blog(LithiumType):
    pass

class Category(LithiumType):
    pass

class User(LithiumType):
    pass

class Message(LithiumType):
    pass

class MessageStatus(LithiumType):
    pass

class Label(LithiumType):
    pass

class Thread(LithiumType):
    def _handle_messages(self, child):
        """
        ``<messages>`` has a somewhat nonstandard structure::

            <messages>
                <topic type="message" href="/messages/id/123" /> <!-- first message in thread -->
                <read> <!-- read count (?) -->
                    <count type="int">5</count>
                </read>
                <linear> <!-- list of messages -->
                    <!-- note: first message is equal to topic -->
                    <message type="message" href="/messages/id/123" />
                    <message type="message" href="/messages/id/124" />
                    <message type="message" href="/messages/id/125" />
                </linear>
            </messages>
        """

        messages = OrderedDict()
        for elem in child:
            if elem.tag == 'topic':
                self.topic = xml_to_type(elem, self._api)
                messages[self.topic.href] = self.topic
            elif elem.tag == 'linear':
                for message in elem:
                    m = xml_to_type(message, self._api)
                    if m.href not in messages:
                        messages[m.href] = m
                    if m.parent and m.parent.href in messages:
                        m.parent = messages[m.parent.href]
            elif elem.tag == 'read':
                self.read_count = elem.text
        self.messages = messages.values()

class EventSubscriptionManager(LithiumType):
    pass

class EventSubscription(LithiumType):

    def __unicode__(self):
        return u"""Type:\t\t%s
Callback:\t%s
Token:\t\t%s""" % (self.event_type, self.callback_url, self.token)

# this maps lithium primitive types to python types. For the most part,
# it leaves them as-is -- unicode strings.

primitive_type_map = defaultdict(lambda: lambda value: value.strip() if value else value, {
    'boolean': lambda value: value == 'true',
    'double': float,
    'float': float,
    'int': int,
    'long': long,
    'date_time': parse_datetime,
})

def primitive_type(xml_tree):
    if xml_tree.getchildren():
        raise ValueError('this is not a primitive type:\n%s' % etree.tostring(xml_tree))
    return primitive_type_map[xml_tree.get('type')](xml_tree.text)

def xml_to_type(xml_tree, api=None):
    """
    recursive function to turn an XML response from the REST API into
    python objects
    """
    lithium_type = xml_tree.get('type')
    if xml_tree.attrib.get('null') == 'true':
        return None
    if lithium_type in types_map:
        # this is a complex type
        lithium_type = types_map.get(xml_tree.get('type'))
        return lithium_type(xml_tree, api)
    elif lithium_type is None:
        # this is a list or dictionary
        children = xml_tree.getchildren()
        if not children:
            # empty, can't say if list or dictionary
            return []
        if 'name' in children[0].attrib:
            # it's a dictionary
            result = {}
            for child in xml_tree.getchildren():
                result[child.get('name')] = xml_to_type(child, api)
        else:
            # it's a list
            result = []
            for child in xml_tree.getchildren():
                result.append(xml_to_type(child, api))
        return result
    else:
        # gotta be a primitive type, then!
        return primitive_type(xml_tree)

types_map = {
    'category': Category,
    'board': Board,
    'blog': Blog,
    'user': User,
    'thread': Thread,
    'message': Message,
    'message_status': MessageStatus,
    'label': Label,

    'event_subscription_manager': EventSubscriptionManager,
    'event_subscription': EventSubscription,

    'value': primitive_type,
}
<commit_msg>python 2.6 compat<commit_after># -*- coding: utf-8 -*-
import pprint

from collections import defaultdict

from dateutil.parser import  parse as parse_datetime
from lxml import etree

from django.utils.datastructures import SortedDict

class LithiumType(object):
    def __init__(self, xml_tree, api):
        self.id = xml_tree.get('id')
        self.href = xml_tree.get('href')
        self._tree = xml_tree
        self._populated = False
        self._api = api

        for child in xml_tree:
            if hasattr(self, '_handle_%s' % child.tag):
                getattr(self, '_handle_%s' % child.tag)(child)
            else:
                setattr(self, child.tag, xml_to_type(child, api))

    def __getattr__(self, item):
        if item.startswith('_'):
            raise AttributeError(item)
        if self._api and self.href and not self._populated:
            obj = self._api.raw(self.href)
            self.__dict__.update(obj.__dict__)
        self._populated = True
        if item in self.__dict__:
            return getattr(self, item)
        else:
            raise AttributeError(item)

    def __repr__(self):
        return '<%s>: %s' % (self.__class__.__name__, pprint.pformat(dict((k,v) for k, v in self.__dict__.iteritems() if not k.startswith('_'))))

class Board(LithiumType):
    pass

class Blog(LithiumType):
    pass

class Category(LithiumType):
    pass

class User(LithiumType):
    pass

class Message(LithiumType):
    pass

class MessageStatus(LithiumType):
    pass

class Label(LithiumType):
    pass

class Thread(LithiumType):
    def _handle_messages(self, child):
        """
        ``<messages>`` has a somewhat nonstandard structure::

            <messages>
                <topic type="message" href="/messages/id/123" /> <!-- first message in thread -->
                <read> <!-- read count (?) -->
                    <count type="int">5</count>
                </read>
                <linear> <!-- list of messages -->
                    <!-- note: first message is equal to topic -->
                    <message type="message" href="/messages/id/123" />
                    <message type="message" href="/messages/id/124" />
                    <message type="message" href="/messages/id/125" />
                </linear>
            </messages>
        """

        messages = SortedDict()
        for elem in child:
            if elem.tag == 'topic':
                self.topic = xml_to_type(elem, self._api)
                messages[self.topic.href] = self.topic
            elif elem.tag == 'linear':
                for message in elem:
                    m = xml_to_type(message, self._api)
                    if m.href not in messages:
                        messages[m.href] = m
                    if m.parent and m.parent.href in messages:
                        m.parent = messages[m.parent.href]
            elif elem.tag == 'read':
                self.read_count = elem.text
        self.messages = messages.values()

class EventSubscriptionManager(LithiumType):
    pass

class EventSubscription(LithiumType):

    def __unicode__(self):
        return u"""Type:\t\t%s
Callback:\t%s
Token:\t\t%s""" % (self.event_type, self.callback_url, self.token)

# this maps lithium primitive types to python types. For the most part,
# it leaves them as-is -- unicode strings.

primitive_type_map = defaultdict(lambda: lambda value: value.strip() if value else value, {
    'boolean': lambda value: value == 'true',
    'double': float,
    'float': float,
    'int': int,
    'long': long,
    'date_time': parse_datetime,
})

def primitive_type(xml_tree):
    if xml_tree.getchildren():
        raise ValueError('this is not a primitive type:\n%s' % etree.tostring(xml_tree))
    return primitive_type_map[xml_tree.get('type')](xml_tree.text)

def xml_to_type(xml_tree, api=None):
    """
    recursive function to turn an XML response from the REST API into
    python objects
    """
    lithium_type = xml_tree.get('type')
    if xml_tree.attrib.get('null') == 'true':
        return None
    if lithium_type in types_map:
        # this is a complex type
        lithium_type = types_map.get(xml_tree.get('type'))
        return lithium_type(xml_tree, api)
    elif lithium_type is None:
        # this is a list or dictionary
        children = xml_tree.getchildren()
        if not children:
            # empty, can't say if list or dictionary
            return []
        if 'name' in children[0].attrib:
            # it's a dictionary
            result = {}
            for child in xml_tree.getchildren():
                result[child.get('name')] = xml_to_type(child, api)
        else:
            # it's a list
            result = []
            for child in xml_tree.getchildren():
                result.append(xml_to_type(child, api))
        return result
    else:
        # gotta be a primitive type, then!
        return primitive_type(xml_tree)

types_map = {
    'category': Category,
    'board': Board,
    'blog': Blog,
    'user': User,
    'thread': Thread,
    'message': Message,
    'message_status': MessageStatus,
    'label': Label,

    'event_subscription_manager': EventSubscriptionManager,
    'event_subscription': EventSubscription,

    'value': primitive_type,
}
<|endoftext|>


<commit_before># This file should have NO imports and be entirely standalone.
# This allows it to import into the runtime DMP as well as
# setup.py during installation.

__version__ = '3.3.18'

# Reminder on uploading to pypi and removing the build folders:
'''

python setup.py sdist upload
rm -rf dist/ django_mako_plus.egg-info/

'''
<commit_msg>update to 1.10<commit_after># This file should have NO imports and be entirely standalone.
# This allows it to import into the runtime DMP as well as
# setup.py during installation.

__version__ = '3.4.1'

# Reminder on uploading to pypi and removing the build folders:
'''

python setup.py sdist upload
rm -rf dist/ django_mako_plus.egg-info/

'''
<|endoftext|>


<commit_before>"""
Instaclone configuration handling.
"""

__author__ = 'jlevy'

import logging as log
import os
import re
import sys
from collections import namedtuple, OrderedDict

from enum import Enum  # enum34
import yaml  # PyYAML

from log_calls import log_calls

import utils

_other_fields = "name version version_hashable"
_required_fields = "local_path remote_path remote_prefix copy_type upload_command download_command"

ConfigBase = namedtuple("ConfigBase", _other_fields + " " + _required_fields)

CONFIGS_REQUIRED = _required_fields.split()

def _stringify_config_field(value):
  return value.name if isinstance(value, Enum) else str(value)

class Config(ConfigBase):
  """Configuration for a single item."""

  def as_string_dict(self):
    d = dict(self._asdict())
    del d["name"]
    return {k: _stringify_config_field(v) for (k, v) in d.iteritems() if v is not None}

CONFIG_NAME = "instaclone"
CONFIG_FILE_SEARCH_ORDER = [".", "$INSTACLONE_DIR", "$HOME/.instaclone"]
CONFIG_DIR_SEARCH_ORDER = ["$INSTACLONE_DIR", "$HOME/.instaclone"]
CONFIG_DIR_DEFAULT = "$HOME/.instaclone"

DEFAULT_ITEM_NAME = "default"

class ConfigError(RuntimeError):
  pass

CopyType = Enum("CopyType", "copy symlink hardlink")

def _locate_config_dir(search_dirs):
  """Look in common locations for config directory."""
  tried = []
  for path in search_dirs:
    log.debug("searching for config dir: %s", path)
    try:
      path = utils.shell_expand_variables(path, os.environ)
    except ValueError:
      continue
    tried.append("no config dir found in: %s" % ", ".join(tried))
    if os.path.isdir(path):
      log.debug("using config dir: %s", path)
      return path
  raise ConfigError("no config file found in: %s" % ", ".join(tried))

def _locate_config_file(search_dirs):
  """Look in common locations for config file."""
  tried = []
  for base in search_dirs:
    try:
      base_expanded = utils.shell_expand_variables(base, os.environ)
    except ValueError:
      continue
    for path in [os.path.join(base_expanded, CONFIG_NAME + suffix) for suffix in ".yml", ".json"]:
      log.debug("searching for config file: %s", path)
      tried.append(path)
      if os.path.isfile(path):
        log.debug("using config file: %s", path)
        return path
  raise ConfigError("no config file found in: %s" % ", ".join(tried))

@log_calls
def _load_raw_configs(override_path, search_dirs):
  if override_path:
    path = override_path
  else:
    path = _locate_config_file(search_dirs)
    if not path:
      raise ConfigError("could not locate config file (be sure a config is in working dir or home folder)")

  with open(path) as f:
    parsed_configs = yaml.safe_load(f)

  out = OrderedDict()
  try:
    items = parsed_configs["items"]
    for (key, config_dict) in items.iteritems():
      # TODO: Could overlay global configs and local ones here.
      combined = {key: None for key in Config._fields}
      combined.update(config_dict)
      combined["name"] = key

      try:
        out[key] = combined
      except TypeError as e:
        raise ConfigError("error in config value: %s: %s" % (e, config_dict))
  except ValueError as e:
    raise ConfigError("error reading config file: %s" % e)

  return out

def _parse_validate(raw_config_list):
  for (name, raw) in raw_config_list.items():
    # Validation.
    for key in CONFIGS_REQUIRED:
      if key not in raw or raw[key] is None:
        raise ConfigError("must specify '%s' in item config: %s" % (key, raw))

    if not re.compile("^[\\w.-]+$").match(raw["name"]):
      raise ConfigError("invalid name for item: '%s'" % raw["name"])

    if "version" in raw and not re.compile("^[\\w.-]+$").match(str(raw["version"])):
      raise ConfigError("invalid version string: '%s'" % raw["version"])
    if "version" not in raw and "version_hashable" not in raw:
      raise ConfigError("must specify 'version' or 'version_hashable' in item config: %s" % raw)

    # Validate shell templates.
    for key in "upload_command", "download_command":
      try:
        utils.shell_expand_to_popen(raw[key], {"REMOTE": "dummy", "LOCAL": "dummy"})
      except ValueError as e:
        raise ConfigError("invalid command in config: %s: %s" % (e, raw[key]))

    # Normalize.
    for path in "local_path", "remote_path":
      if path.startswith("/"):
        raise ConfigError("currently only support relative paths for local_path and remote_path: %s" % path)
      raw[path] = raw[path].rstrip("/")

    # Parse enums.
    try:
      raw["copy_type"] = CopyType[raw["copy_type"]]
    except KeyError:
      raise ConfigError("invalid copy type: %s" % raw["copy_type"])

    yield Config(**raw)

@log_calls
def cache_dir():
  try:
    config_dir = _locate_config_dir(CONFIG_DIR_SEARCH_ORDER)
  except ConfigError:
    config_dir = utils.shell_expand_variables(CONFIG_DIR_DEFAULT, os.environ)
    log.info("config dir not found, so creating: %s", config_dir)
  cache_dir = os.path.join(config_dir, "cache")
  utils.make_all_dirs(cache_dir)
  return cache_dir

def load(override_path, search_dirs=CONFIG_FILE_SEARCH_ORDER):
  """Load all configs from a single file. Use override_path or the first one found in search_dirs."""
  if not search_dirs:
    search_dirs = []
  return list(_parse_validate(_load_raw_configs(override_path, search_dirs)))

def print_configs(configs, stream=sys.stdout):
  yaml.dump({"items": OrderedDict([config.name, config.as_string_dict()] for config in configs)},
            stream=stream, default_flow_style=False)

def _yaml_ordering_support():
  """
  Get yaml lib to handle OrderedDicts.
  See http://stackoverflow.com/questions/5121931/in-python-how-can-you-load-yaml-mappings-as-ordereddicts
  """
  _mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG

  def dict_representer(dumper, data):
    return dumper.represent_dict(data.iteritems())

  def dict_constructor(loader, node):
    return OrderedDict(loader.construct_pairs(node))

  yaml.add_representer(OrderedDict, dict_representer)
  yaml.add_constructor(_mapping_tag, dict_constructor)

_yaml_ordering_support()
<commit_msg>Adjust config format to remove useless "name" key.<commit_after>"""
Instaclone configuration handling.
"""

__author__ = 'jlevy'

import logging as log
import os
import re
import sys
from collections import namedtuple, OrderedDict

from enum import Enum  # enum34
import yaml  # PyYAML

from log_calls import log_calls

import utils

_required_fields = "local_path remote_path remote_prefix copy_type upload_command download_command"
_other_fields = "version version_hashable"

ConfigBase = namedtuple("ConfigBase", _other_fields + " " + _required_fields)

CONFIGS_REQUIRED = _required_fields.split()

def _stringify_config_field(value):
  return value.name if isinstance(value, Enum) else str(value)

class Config(ConfigBase):
  """Configuration for a single item."""

  def as_string_dict(self):
    d = dict(self._asdict())
    return {k: _stringify_config_field(v) for (k, v) in d.iteritems() if v is not None}

CONFIG_NAME = "instaclone"
CONFIG_FILE_SEARCH_ORDER = [".", "$INSTACLONE_DIR", "$HOME/.instaclone"]
CONFIG_DIR_SEARCH_ORDER = ["$INSTACLONE_DIR", "$HOME/.instaclone"]
CONFIG_DIR_DEFAULT = "$HOME/.instaclone"

DEFAULT_ITEM_NAME = "default"

class ConfigError(RuntimeError):
  pass

CopyType = Enum("CopyType", "copy symlink hardlink")

def _locate_config_dir(search_dirs):
  """Look in common locations for config directory."""
  tried = []
  for path in search_dirs:
    log.debug("searching for config dir: %s", path)
    try:
      path = utils.shell_expand_variables(path, os.environ)
    except ValueError:
      continue
    tried.append("no config dir found in: %s" % ", ".join(tried))
    if os.path.isdir(path):
      log.debug("using config dir: %s", path)
      return path
  raise ConfigError("no config file found in: %s" % ", ".join(tried))

def _locate_config_file(search_dirs):
  """Look in common locations for config file."""
  tried = []
  for base in search_dirs:
    try:
      base_expanded = utils.shell_expand_variables(base, os.environ)
    except ValueError:
      continue
    for path in [os.path.join(base_expanded, CONFIG_NAME + suffix) for suffix in ".yml", ".json"]:
      log.debug("searching for config file: %s", path)
      tried.append(path)
      if os.path.isfile(path):
        log.debug("using config file: %s", path)
        return path
  raise ConfigError("no config file found in: %s" % ", ".join(tried))

@log_calls
def _load_raw_configs(override_path, search_dirs):
  if override_path:
    path = override_path
  else:
    path = _locate_config_file(search_dirs)
    if not path:
      raise ConfigError("could not locate config file (be sure a config is in working dir or home folder)")

  with open(path) as f:
    parsed_configs = yaml.safe_load(f)

  out = []
  try:
    items = parsed_configs["items"]
    for config_dict in items:
      # TODO: Could overlay global configs and local ones here.
      combined = {key: None for key in Config._fields}
      combined.update(config_dict)

      try:
        out.append(combined)
      except TypeError as e:
        raise ConfigError("error in config value: %s: %s" % (e, config_dict))
  except ValueError as e:
    raise ConfigError("error reading config file: %s" % e)

  return out

def _parse_validate(raw_config_list):
  for raw in raw_config_list:
    # Validation.
    for key in CONFIGS_REQUIRED:
      if key not in raw or raw[key] is None:
        raise ConfigError("must specify '%s' in item config: %s" % (key, raw))

    if "version" in raw and not re.compile("^[\\w.-]+$").match(str(raw["version"])):
      raise ConfigError("invalid version string: '%s'" % raw["version"])
    if "version" not in raw and "version_hashable" not in raw:
      raise ConfigError("must specify 'version' or 'version_hashable' in item config: %s" % raw)

    # Validate shell templates.
    for key in "upload_command", "download_command":
      try:
        utils.shell_expand_to_popen(raw[key], {"REMOTE": "dummy", "LOCAL": "dummy"})
      except ValueError as e:
        raise ConfigError("invalid command in config: %s: %s" % (e, raw[key]))

    # Normalize.
    for path in "local_path", "remote_path":
      if path.startswith("/"):
        raise ConfigError("currently only support relative paths for local_path and remote_path: %s" % path)
      raw[path] = raw[path].rstrip("/")

    # Parse enums.
    try:
      raw["copy_type"] = CopyType[raw["copy_type"]]
    except KeyError:
      raise ConfigError("invalid copy type: %s" % raw["copy_type"])

    yield Config(**raw)

@log_calls
def cache_dir():
  try:
    config_dir = _locate_config_dir(CONFIG_DIR_SEARCH_ORDER)
  except ConfigError:
    config_dir = utils.shell_expand_variables(CONFIG_DIR_DEFAULT, os.environ)
    log.info("config dir not found, so creating: %s", config_dir)
  cache_dir = os.path.join(config_dir, "cache")
  utils.make_all_dirs(cache_dir)
  return cache_dir

def load(override_path, search_dirs=CONFIG_FILE_SEARCH_ORDER):
  """Load all configs from a single file. Use override_path or the first one found in search_dirs."""
  if not search_dirs:
    search_dirs = []
  return list(_parse_validate(_load_raw_configs(override_path, search_dirs)))

def print_configs(configs, stream=sys.stdout):
  yaml.dump({"items": [config.as_string_dict() for config in configs]},
            stream=stream, default_flow_style=False)

def _yaml_ordering_support():
  """
  Get yaml lib to handle OrderedDicts.
  See http://stackoverflow.com/questions/5121931/in-python-how-can-you-load-yaml-mappings-as-ordereddicts
  """
  _mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG

  def dict_representer(dumper, data):
    return dumper.represent_dict(data.iteritems())

  def dict_constructor(loader, node):
    return OrderedDict(loader.construct_pairs(node))

  yaml.add_representer(OrderedDict, dict_representer)
  yaml.add_constructor(_mapping_tag, dict_constructor)

_yaml_ordering_support()
<|endoftext|>


<commit_before># This file should have NO imports and be entirely standalone.
# This allows it to import into the runtime DMP as well as
# setup.py during installation.

__version__ = '3.7.17'

# Reminder on uploading to pypi and removing the build folders:
'''

python3 setup.py sdist
twine upload dist/*
rm -rf dist/ django_mako_plus.egg-info/

'''
<commit_msg>version bump<commit_after># This file should have NO imports and be entirely standalone.
# This allows it to import into the runtime DMP as well as
# setup.py during installation.

__version__ = '3.7.18'

# Reminder on uploading to pypi and removing the build folders:
'''

python3 setup.py sdist
twine upload dist/*
rm -rf dist/ django_mako_plus.egg-info/

'''
<|endoftext|>


<commit_before># coding=utf-8
"""URI Routing configuration for this apps."""
from django.conf.urls import patterns, url, include

# Needed by haystack views
from feti.forms.search import DefaultSearchForm
from haystack.query import SearchQuerySet
from haystack.views import search_view_factory, SearchView

sqs = SearchQuerySet()

urlpatterns = patterns(
    '',
    url(
        r'^$',
        'feti.views.landing_page.landing_page'
    ),
    url(
        r'^project-team/$',
        'feti.views.project_team.project_team'
    ),
    url(
        r'^project-team/add-campus/$',
        'feti.views.add_campus.add_campus'
    ),
    url(
        r'^search/',
        include('haystack.urls')
    ),
    url(
        r'^customsearch/',
        search_view_factory(
            view_class=SearchView,
            template='search/custom_search.html',
            searchqueryset=sqs,
            form_class=DefaultSearchForm),
        name='haystack_search'),
)
<commit_msg>Fix django admin page.<commit_after># coding=utf-8
"""URI Routing configuration for this apps."""
from django.conf.urls import patterns, url, include

# Needed by haystack views
from feti.forms.search import DefaultSearchForm
from haystack.query import SearchQuerySet
from haystack.views import search_view_factory, SearchView

sqs = SearchQuerySet()

urlpatterns = patterns(
    '',
    url(
        r'^$',
        'feti.views.landing_page.landing_page'
    ),
    url(
        r'^project-team/$',
        'feti.views.project_team.project_team'
    ),
    # url(
    #     r'^project-team/add-campus/$',
    #     'feti.views.add_campus.add_campus'
    ),
    url(
        r'^search/',
        include('haystack.urls')
    ),
    url(
        r'^customsearch/',
        search_view_factory(
            view_class=SearchView,
            template='search/custom_search.html',
            searchqueryset=sqs,
            form_class=DefaultSearchForm),
        name='haystack_search'),
)
<|endoftext|>


<commit_before>
from django_socketio.utils import NoSocket, send, broadcast, broadcast_channel

__version__ = "0.3.4"

<commit_msg>Version bump to 0.3.5<commit_after>
from django_socketio.utils import NoSocket, send, broadcast, broadcast_channel

__version__ = "0.3.5"

<|endoftext|>


<commit_before>"""Records stats about the system."""
import os
import time

import requests

from dnstwister import dnstwist
from dnstwister.storage import redis_stats_store

def get_delta_domains(url=os.getenv('DELTAS_URL')):
    """Return a list of all the domains in all the delta reports.

    If this stops scaling I'll switch to an iterator off a DB query.
    """
    if url is None:
        raise Exception('Delta report URL configuration not set!')

    json = requests.get(os.getenv('DELTAS_URL'), timeout=10).json()
    return [domain
            for (domain,)
            in json['values']
            if dnstwist.validate_domain(domain)]

def main():
    """Main code for worker.

    Run as often as like on a schedule.
    """
    start = time.time()
    store = redis_stats_store.RedisStatsStore()

    delta_domains = get_delta_domains()
    for domain in delta_domains:
        store.note(domain)

    print 'Processed {} in {} seconds'.format(
        len(delta_domains),
        round(time.time() - start, 2)
    )

if __name__ == '__main__':
    main()
<commit_msg>Improved logging for the stats worker<commit_after>"""Records stats about the system."""
import os
import time

import requests

from dnstwister import dnstwist
from dnstwister.storage import redis_stats_store

def get_delta_domains(url=os.getenv('DELTAS_URL')):
    """Return a list of all the domains in all the delta reports.

    If this stops scaling I'll switch to an iterator off a DB query.
    """
    if url is None:
        raise Exception('Delta report URL configuration not set!')

    json = requests.get(os.getenv('DELTAS_URL'), timeout=10).json()
    return [domain
            for (domain,)
            in json['values']
            if dnstwist.validate_domain(domain)]

def main():
    """Main code for worker.

    Run as often as like on a schedule.

    I'm doing it via a Heroku scheduler calling:

        python -m dnstwister.workers.stats

    """
    start = time.time()
    store = redis_stats_store.RedisStatsStore()

    delta_domains = get_delta_domains()
    for domain in delta_domains:
        store.note(domain)

    print 'Processed stats for {} domains in {} seconds'.format(
        len(delta_domains),
        round(time.time() - start, 2)
    )

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


<commit_before>import matplotlib.pyplot as plt
from axes_template import axes_template
from plot_manager import Plot_Manager

def extended_figure(*args,**kwargs):
    kwargs["FigureClass"]=Extended_Figure
    return plt.figure(*args,**kwargs)

class Extended_Figure(plt.Figure):
    '''
    Creates a figure to plot data in multiple axes
    comming from files or from objects with the correct format
    '''
    def __init__(self,*args,**kwargs):
        super(Extended_Figure,self).__init__(*args,**kwargs)
        self.ncol=1
        self.nrow=1
        self.manager=Plot_Manager(self)

    def blank_figure(self):
        # Creates a blank figure
        self.define_template(matrix=[[1]])
        ax=self.axes_dict[1]
        ax.set_xticks([])
        ax.set_yticks([])
        self.tight_layout()

    def define_template(self,**kwargs):
        '''
        Defines the number of rows and colums
        for the subplots. It is possible also to define a custom grid
        that will include the columspans
        '''
        self.template=axes_template(**kwargs)
        self.ncol=self.template.ncol
        self.nrow=self.template.nrow
        self.axes_dict=self.template.apply_template(self)
        self.manager=Plot_Manager(self)
        self.index_dict=dict()
        self.legend_status=dict()
        for index in self.axes_dict:
            self.index_dict[self.axes_dict[index]]=index
            self.legend_status[index]=True

    def view_template(self):
        """
        It generates a figure with the template and numbers
        that identify each axis
        """
        self.template.generate_preview()

    def plot(self,index,*args,**kwargs):
        xcol=None
        ycol=None
        string=None
        if "xcol" in kwargs:
            xcol=kwargs.pop("xcol")
        if "ycol" in kwargs:
            ycol=kwargs.pop("ycol")
        if "string" in kwargs:
            string=kwargs.pop("string")
        plot_label=""
        if "plot_label" in kwargs:
            plot_label=kwargs.pop("plot_label")
        if index in self.axes_dict:
            plot=self.axes_dict[index].plot(*args,**kwargs)[0]
            self.manager.add_plot(index,plot,kwargs,path=plot_label,xcol=xcol,ycol=ycol,string=string)
        else:
            print("This value for the axis is not defined")
            return

    def hist(self,index,*args,**kwargs):
        xcol=None
        ycol=None
        if "xcol" in kwargs:
            xcol=kwargs.pop("xcol")
        if "ycol" in kwargs:
            ycol=kwargs.pop("ycol")
        if "plot_label" in kwargs:
            plot_label=kwargs.pop("plot_label")
        if index in self.axes_dict:
            if "width" in kwargs:
                if kwargs["width"]=="auto":
                    lens=[args[0][i]-args[0][i-1] for i in range(1,len(args[0]))]
                    kwargs["width"]=0.8*min(lens)
            self.axes_dict[index].bar(*args,**kwargs)
        else:
            print("This value for the axis is not defined")
            return

    def hist_file(self,filename,index,xcol=0,ycol=1,**kwargs):
        if "plot_label" not in kwargs:
            kwargs["plot_label"]=filename
        kwargs["xcol"]=xcol
        kwargs["ycol"]=ycol
        x,y=read_file(filename,xcol,ycol)
        self.hist(index,x,y,**kwargs)

    def plot_file(self,filename,index,xcol=0,ycol=1,**kwargs):
        if "plot_label" not in kwargs:
            kwargs["plot_label"]=filename
        kwargs["xcol"]=int(xcol)
        kwargs["ycol"]=int(ycol)
        x,y=read_file(filename,int(xcol),int(ycol))
        # Make a copy to avoid nasty changes in plots
        self.plot(index,x,y,**kwargs)
        return

    def plot_class(self,objeto,index,**kwargs):
        x=objeto.data[0]
        y=objeto.data[1]
        self.plot(index,x,y,**kwargs)

    def update_legend(self,index,state):
        legend=self.axes_dict[index]
        legend.set_visible(state)

    def join_xy(self):
        if not (self.template.xjoinable and self.template.yjoinable):
            print("The tempalte does not allow this kind of joining"+
                  " try with only one kind of axis")
            return
        else:
            self.join_x()
            self.join_y()

    def join_x(self):
        """
        Joins and share the x axis of the subplots
        of each column
        """
        if not self.template.xjoinable:
            print("This configuration does not support x axis joinability")
            return
        matrix=self.template.matrix

        # Join all the axis in each column
        joined_index=[]
        maximuns = []
        minimuns = []
        for column in range(self.ncol):
            joined_index.append(matrix[-1][column])
            ref_axis=self.axes_dict[matrix[-1][column]]
            maximuns.append(ref_axis.get_xlim()[1])
            minimuns.append(ref_axis.get_xlim()[0])
            for row in range(self.nrow):
                index=matrix[row][column]
                if index not in joined_index:
                    joined_index.append(index)
                    ref_axis.get_shared_x_axes().join(self.axes_dict[index],ref_axis)
                mini,maxi=self.axes_dict[matrix[row][column]].get_xlim()
                maximuns[-1] = max(maximuns[-1],maxi)
                minimuns[-1] = min(minimuns[-1],mini)

        # Change the range of all the subplots
        for column in range(self.ncol):
            self.axes_dict[matrix[-1][column]].set_xlim((minimuns[column],maximuns[column]))
        self.fix_x_ticks()

    def fix_x_ticks(self):
        # Elimnate the xticks of all the subplots that are not
        # at the bottom. It also pops the top element of each
        # set of ticks if they are not in the top subplot
        matrix=self.template.matrix
        adjusted_ticks=[]

        for column in range(self.ncol):
            bottom_index=matrix[-1][column]
            top_index=matrix[-1][column]
            removed_ticks=[]
            for row in range(self.nrow):
                if matrix[row][column]!=bottom_index and matrix[row][column] not in removed_ticks:
                    for tick in self.axes_dict[matrix[row][column]].xaxis.get_major_ticks():
                        tick.label1.set_visible(False)
                    removed_ticks.append(matrix[row][column])
                if matrix[row][column]!=top_index and matrix[row][column] not in adjusted_ticks:
                    ticks=self.axes_dict[matrix[row][column]].yaxis.get_major_ticks()
                    ticks[0].set_visible(False)
                    adjusted_ticks.append(matrix[row][column])
        self.subplots_adjust(hspace=0)

    def join_y(self):
        """
        Joins and share the y axis of the subplots
        of each column
        """
        if not self.template.yjoinable:
            print("This configuration does not support y axis joinability")
            return
        matrix=self.template.matrix

        # Join all the axis in each row
        joined_index=[]
        maximuns = []
        minimuns = []
        for row in range(self.nrow):
            joined_index.append(matrix[row][0])
            ref_axis=self.axes_dict[matrix[row][0]]
            maximuns.append(ref_axis.get_ylim()[1])
            minimuns.append(ref_axis.get_ylim()[0])
            for column in range(self.ncol):
                index=matrix[row][column]
                if index not in joined_index:
                    joined_index.append(index)
                    ref_axis.get_shared_y_axes().join(self.axes_dict[index],ref_axis)
                mini,maxi=self.axes_dict[matrix[row][column]].get_ylim()
                maximuns[-1] = max(maximuns[-1],maxi)
                minimuns[-1] = min(minimuns[-1],mini)

        # Change the range of all the subplots
        for column in range(self.nrow):
            self.axes_dict[matrix[column][-1]].set_ylim((minimuns[column],maximuns[column]))
        self.fix_y_ticks()

    def fix_y_ticks(self):
        # Elimnate the xticks of all the subplots that are not
        # at the bottom. It also pops the top element of each
        # set of ticks if they are not in the top subplot
        adjusted_ticks=[]
        matrix=self.template.matrix

        for row in range(self.nrow):
            bottom_index=matrix[row][0]
            top_index=matrix[row][-1]
            removed_ticks=[]
            for column in range(self.ncol):
                if matrix[row][column]!=bottom_index and matrix[row][column] not in removed_ticks:
                    for tick in self.axes_dict[matrix[row][column]].yaxis.get_major_ticks():
                        tick.label1.set_visible(False)
                    removed_ticks.append(matrix[row][column])
                if matrix[row][column]!=top_index and matrix[row][column] not in adjusted_ticks:
                    ticks=self.axes_dict[matrix[row][column]].xaxis.get_major_ticks()
                    ticks[-1].set_visible(False)
                    adjusted_ticks.append(matrix[row][column])
        self.subplots_adjust(wspace=0)

    def undo(self):
        self.manager.undo()

    def redo(self):
        self.manager.redo()

    def canUndo(self):
        return self.manager.canUndo()

    def canRedo(self):
        return self.manager.canRedo()

    def reset_ticks(self):
        """
        Makes visisible all the ticks, this can be used if there are
        hidden ticks that should ve visible
        """
        matrix = self.template.matrix
        for column in range(self.ncol):
            for row in range(self.nrow):
                axes = self.axes_dict[matrix[row][column]]
                x_ticks = axes.xaxis.get_major_ticks()
                for tick in x_ticks:
                    tick.set_visible(True)
                    tick.label1.set_visible(True)
                y_ticks = axes.yaxis.get_major_ticks()
                for tick in y_ticks:
                    tick.set_visible(True)
                    tick.label1.set_visible(True)

    ###########################
    # Lading saving functions #
    ###########################

    def save(self,path):
        self.manager.save(path)

    def load(self,path):
        return self.manager.load(path)

def read_file(filename,xcol=0,ycol=1):
    f=open(filename,"ur")
    x=[]
    y=[]
    for l in f:
        if l.startswith("#"):
            pass
        elif l.startswith(r"@"):
            pass
        else:
            try:
                xval=float(l.split()[xcol])
                yval=float(l.split()[ycol])
                x.append(xval)
                y.append(yval)
            except:
                pass
    return x,y
<commit_msg>corrected u with U<commit_after>import matplotlib.pyplot as plt
from axes_template import axes_template
from plot_manager import Plot_Manager

def extended_figure(*args,**kwargs):
    kwargs["FigureClass"]=Extended_Figure
    return plt.figure(*args,**kwargs)

class Extended_Figure(plt.Figure):
    '''
    Creates a figure to plot data in multiple axes
    comming from files or from objects with the correct format
    '''
    def __init__(self,*args,**kwargs):
        super(Extended_Figure,self).__init__(*args,**kwargs)
        self.ncol=1
        self.nrow=1
        self.manager=Plot_Manager(self)

    def blank_figure(self):
        # Creates a blank figure
        self.define_template(matrix=[[1]])
        ax=self.axes_dict[1]
        ax.set_xticks([])
        ax.set_yticks([])
        self.tight_layout()

    def define_template(self,**kwargs):
        '''
        Defines the number of rows and colums
        for the subplots. It is possible also to define a custom grid
        that will include the columspans
        '''
        self.template=axes_template(**kwargs)
        self.ncol=self.template.ncol
        self.nrow=self.template.nrow
        self.axes_dict=self.template.apply_template(self)
        self.manager=Plot_Manager(self)
        self.index_dict=dict()
        self.legend_status=dict()
        for index in self.axes_dict:
            self.index_dict[self.axes_dict[index]]=index
            self.legend_status[index]=True

    def view_template(self):
        """
        It generates a figure with the template and numbers
        that identify each axis
        """
        self.template.generate_preview()

    def plot(self,index,*args,**kwargs):
        xcol=None
        ycol=None
        string=None
        if "xcol" in kwargs:
            xcol=kwargs.pop("xcol")
        if "ycol" in kwargs:
            ycol=kwargs.pop("ycol")
        if "string" in kwargs:
            string=kwargs.pop("string")
        plot_label=""
        if "plot_label" in kwargs:
            plot_label=kwargs.pop("plot_label")
        if index in self.axes_dict:
            plot=self.axes_dict[index].plot(*args,**kwargs)[0]
            self.manager.add_plot(index,plot,kwargs,path=plot_label,xcol=xcol,ycol=ycol,string=string)
        else:
            print("This value for the axis is not defined")
            return

    def hist(self,index,*args,**kwargs):
        xcol=None
        ycol=None
        if "xcol" in kwargs:
            xcol=kwargs.pop("xcol")
        if "ycol" in kwargs:
            ycol=kwargs.pop("ycol")
        if "plot_label" in kwargs:
            plot_label=kwargs.pop("plot_label")
        if index in self.axes_dict:
            if "width" in kwargs:
                if kwargs["width"]=="auto":
                    lens=[args[0][i]-args[0][i-1] for i in range(1,len(args[0]))]
                    kwargs["width"]=0.8*min(lens)
            self.axes_dict[index].bar(*args,**kwargs)
        else:
            print("This value for the axis is not defined")
            return

    def hist_file(self,filename,index,xcol=0,ycol=1,**kwargs):
        if "plot_label" not in kwargs:
            kwargs["plot_label"]=filename
        kwargs["xcol"]=xcol
        kwargs["ycol"]=ycol
        x,y=read_file(filename,xcol,ycol)
        self.hist(index,x,y,**kwargs)

    def plot_file(self,filename,index,xcol=0,ycol=1,**kwargs):
        if "plot_label" not in kwargs:
            kwargs["plot_label"]=filename
        kwargs["xcol"]=int(xcol)
        kwargs["ycol"]=int(ycol)
        x,y=read_file(filename,int(xcol),int(ycol))
        # Make a copy to avoid nasty changes in plots
        self.plot(index,x,y,**kwargs)
        return

    def plot_class(self,objeto,index,**kwargs):
        x=objeto.data[0]
        y=objeto.data[1]
        self.plot(index,x,y,**kwargs)

    def update_legend(self,index,state):
        legend=self.axes_dict[index]
        legend.set_visible(state)

    def join_xy(self):
        if not (self.template.xjoinable and self.template.yjoinable):
            print("The tempalte does not allow this kind of joining"+
                  " try with only one kind of axis")
            return
        else:
            self.join_x()
            self.join_y()

    def join_x(self):
        """
        Joins and share the x axis of the subplots
        of each column
        """
        if not self.template.xjoinable:
            print("This configuration does not support x axis joinability")
            return
        matrix=self.template.matrix

        # Join all the axis in each column
        joined_index=[]
        maximuns = []
        minimuns = []
        for column in range(self.ncol):
            joined_index.append(matrix[-1][column])
            ref_axis=self.axes_dict[matrix[-1][column]]
            maximuns.append(ref_axis.get_xlim()[1])
            minimuns.append(ref_axis.get_xlim()[0])
            for row in range(self.nrow):
                index=matrix[row][column]
                if index not in joined_index:
                    joined_index.append(index)
                    ref_axis.get_shared_x_axes().join(self.axes_dict[index],ref_axis)
                mini,maxi=self.axes_dict[matrix[row][column]].get_xlim()
                maximuns[-1] = max(maximuns[-1],maxi)
                minimuns[-1] = min(minimuns[-1],mini)

        # Change the range of all the subplots
        for column in range(self.ncol):
            self.axes_dict[matrix[-1][column]].set_xlim((minimuns[column],maximuns[column]))
        self.fix_x_ticks()

    def fix_x_ticks(self):
        # Elimnate the xticks of all the subplots that are not
        # at the bottom. It also pops the top element of each
        # set of ticks if they are not in the top subplot
        matrix=self.template.matrix
        adjusted_ticks=[]

        for column in range(self.ncol):
            bottom_index=matrix[-1][column]
            top_index=matrix[-1][column]
            removed_ticks=[]
            for row in range(self.nrow):
                if matrix[row][column]!=bottom_index and matrix[row][column] not in removed_ticks:
                    for tick in self.axes_dict[matrix[row][column]].xaxis.get_major_ticks():
                        tick.label1.set_visible(False)
                    removed_ticks.append(matrix[row][column])
                if matrix[row][column]!=top_index and matrix[row][column] not in adjusted_ticks:
                    ticks=self.axes_dict[matrix[row][column]].yaxis.get_major_ticks()
                    ticks[0].set_visible(False)
                    adjusted_ticks.append(matrix[row][column])
        self.subplots_adjust(hspace=0)

    def join_y(self):
        """
        Joins and share the y axis of the subplots
        of each column
        """
        if not self.template.yjoinable:
            print("This configuration does not support y axis joinability")
            return
        matrix=self.template.matrix

        # Join all the axis in each row
        joined_index=[]
        maximuns = []
        minimuns = []
        for row in range(self.nrow):
            joined_index.append(matrix[row][0])
            ref_axis=self.axes_dict[matrix[row][0]]
            maximuns.append(ref_axis.get_ylim()[1])
            minimuns.append(ref_axis.get_ylim()[0])
            for column in range(self.ncol):
                index=matrix[row][column]
                if index not in joined_index:
                    joined_index.append(index)
                    ref_axis.get_shared_y_axes().join(self.axes_dict[index],ref_axis)
                mini,maxi=self.axes_dict[matrix[row][column]].get_ylim()
                maximuns[-1] = max(maximuns[-1],maxi)
                minimuns[-1] = min(minimuns[-1],mini)

        # Change the range of all the subplots
        for column in range(self.nrow):
            self.axes_dict[matrix[column][-1]].set_ylim((minimuns[column],maximuns[column]))
        self.fix_y_ticks()

    def fix_y_ticks(self):
        # Elimnate the xticks of all the subplots that are not
        # at the bottom. It also pops the top element of each
        # set of ticks if they are not in the top subplot
        adjusted_ticks=[]
        matrix=self.template.matrix

        for row in range(self.nrow):
            bottom_index=matrix[row][0]
            top_index=matrix[row][-1]
            removed_ticks=[]
            for column in range(self.ncol):
                if matrix[row][column]!=bottom_index and matrix[row][column] not in removed_ticks:
                    for tick in self.axes_dict[matrix[row][column]].yaxis.get_major_ticks():
                        tick.label1.set_visible(False)
                    removed_ticks.append(matrix[row][column])
                if matrix[row][column]!=top_index and matrix[row][column] not in adjusted_ticks:
                    ticks=self.axes_dict[matrix[row][column]].xaxis.get_major_ticks()
                    ticks[-1].set_visible(False)
                    adjusted_ticks.append(matrix[row][column])
        self.subplots_adjust(wspace=0)

    def undo(self):
        self.manager.undo()

    def redo(self):
        self.manager.redo()

    def canUndo(self):
        return self.manager.canUndo()

    def canRedo(self):
        return self.manager.canRedo()

    def reset_ticks(self):
        """
        Makes visisible all the ticks, this can be used if there are
        hidden ticks that should ve visible
        """
        matrix = self.template.matrix
        for column in range(self.ncol):
            for row in range(self.nrow):
                axes = self.axes_dict[matrix[row][column]]
                x_ticks = axes.xaxis.get_major_ticks()
                for tick in x_ticks:
                    tick.set_visible(True)
                    tick.label1.set_visible(True)
                y_ticks = axes.yaxis.get_major_ticks()
                for tick in y_ticks:
                    tick.set_visible(True)
                    tick.label1.set_visible(True)

    ###########################
    # Lading saving functions #
    ###########################

    def save(self,path):
        self.manager.save(path)

    def load(self,path):
        return self.manager.load(path)

def read_file(filename,xcol=0,ycol=1):
    f=open(filename,"Ur")
    x=[]
    y=[]
    for l in f:
        if l.startswith("#"):
            pass
        elif l.startswith(r"@"):
            pass
        else:
            try:
                xval=float(l.split()[xcol])
                yval=float(l.split()[ycol])
                x.append(xval)
                y.append(yval)
            except:
                pass
    return x,y
<|endoftext|>


<commit_before>"""
N-Triples RDF graph serializer for RDFLib.
See <http://www.w3.org/TR/rdf-testcases/#ntriples> for details about the
format.
"""
from rdflib.serializer import Serializer
import warnings

class NTSerializer(Serializer):
    """
    Serializes RDF graphs to NTriples format.
    """

    def serialize(self, stream, base=None, encoding=None, **args):
        if base is not None:
            warnings.warn("NTSerializer does not support base.")
        if encoding is not None:
            warnings.warn("NTSerializer does not use custom encoding.")
        encoding = self.encoding
        for triple in self.store:
            stream.write(_nt_row(triple).encode(encoding, "replace"))
        stream.write("\n")

def _nt_row(triple):
    return u"%s %s %s .\n" % (triple[0].n3(),
            triple[1].n3(),
            _xmlcharref_encode(triple[2].n3()))

# from <http://code.activestate.com/recipes/303668/>
def _xmlcharref_encode(unicode_data, encoding="ascii"):
    """Emulate Python 2.3's 'xmlcharrefreplace' encoding error handler."""
    chars = []

    # nothing to do about xmlchars, but replace newlines with escapes:
    unicode_data=unicode_data.replace("\n","\\n")
    if unicode_data.startswith('"""'):
        unicode_data = unicode_data.replace('"""', '"')

    # Step through the unicode_data string one character at a time in
    # order to catch unencodable characters:
    for char in unicode_data:
        try:
            chars.append(char.encode(encoding, 'strict'))
        except UnicodeError:
            chars.append('\u%04X' % ord(char) if ord(char) <= 0xFFFF else '\U%08X' % ord(char))

    return ''.join(chars)

<commit_msg>Fixes issue 177.<commit_after>"""
N-Triples RDF graph serializer for RDFLib.
See <http://www.w3.org/TR/rdf-testcases/#ntriples> for details about the
format.
"""
from rdflib.serializer import Serializer
import warnings

class NTSerializer(Serializer):
    """
    Serializes RDF graphs to NTriples format.
    """

    def serialize(self, stream, base=None, encoding=None, **args):
        if base is not None:
            warnings.warn("NTSerializer does not support base.")
        if encoding is not None:
            warnings.warn("NTSerializer does not use custom encoding.")
        encoding = self.encoding
        for triple in self.store:
            stream.write(_nt_row(triple).encode(encoding, "replace"))
        stream.write("\n")

def _nt_row(triple):
    return u"%s %s %s .\n" % (triple[0].n3(),
            triple[1].n3(),
            _xmlcharref_encode(triple[2].n3()))

# from <http://code.activestate.com/recipes/303668/>
def _xmlcharref_encode(unicode_data, encoding="ascii"):
    """Emulate Python 2.3's 'xmlcharrefreplace' encoding error handler."""
    chars = []

    # nothing to do about xmlchars, but replace newlines with escapes:
    unicode_data=unicode_data.replace("\n","\\n")
    if unicode_data.startswith('"""'):
        unicode_data = unicode_data.replace('"""', '"')

    # Step through the unicode_data string one character at a time in
    # order to catch unencodable characters:
    for char in unicode_data:
        try:
            chars.append(char.encode(encoding, 'strict'))
        except UnicodeError:
            if ord(char) <= 0xFFFF:
                chars.append('\u%04X' % ord(char))
            else:
                chars.append('\U%08X' % ord(char))

    return ''.join(chars)

<|endoftext|>


<commit_before>import logging

from defusedxml import ElementTree

from dojo.models import Endpoint, Finding

logger = logging.getLogger(__name__)

class Outpost24Parser(object):

    def get_scan_types(self):
        return ["Outpost24 Scan"]

    def get_label_for_scan_types(self, scan_type):
        return "Outpost24 Scan"

    def get_description_for_scan_types(self, scan_type):
        return "Import Outpost24 endpoint vulnerability scan in XML format."

    def get_findings(self, file, test):
        tree = ElementTree.parse(file)
        items = list()
        for detail in tree.iterfind('//detaillist/detail'):
            # finding details
            title = detail.findtext('name')
            # date = detail.findtext('date') # can be used for Finding.date?
            cve = detail.findtext('./cve/id')
            url = detail.findtext('./referencelist/reference/[type=\'solution\']/../url')
            description = detail.findtext('description')
            mitigation = detail.findtext('solution')
            impact = detail.findtext('information')
            cvss_score = detail.findtext('cvss_v3_score') or detail.findtext('cvss_score')
            if not cvss_score:
                cvss_score = 0
            if cvss_score:
                score = float(cvss_score)
                if score < 4:
                    severity = 'Low'
                elif score < 7:
                    severity = 'Medium'
                elif score < 9:
                    severity = 'High'
                else:
                    severity = 'Critical'
            else:
                risk = int(detail.findtext('risk'))
                if risk == 0:
                    severity = 'Low'
                elif risk == 1:
                    severity = 'Medium'
                elif risk == 2:
                    severity = 'High'
                else:
                    severity = 'Critical'
            cvss_description = detail.findtext('cvss_vector_description')
            cvss_vector = detail.findtext('cvss_v3_vector') or detail.findtext('cvss_vector')
            severity_justification = "{}\n{}".format(cvss_score, cvss_description)
            finding = Finding(title=title, test=test, cve=cve, url=url, description=description, mitigation=mitigation,
                              impact=impact, severity=severity,
                              severity_justification=severity_justification)
            # endpoint details
            host = detail.findtext('ip')
            if host:
                protocol = detail.findtext('./portinfo/service')
                try:
                    port = int(detail.findtext('./portinfo/portnumber'))
                except ValueError as ve:
                    logger.debug("General port given. Assigning 0 as default.")
                    port = 0
                finding.unsaved_endpoints.append(Endpoint(protocol=protocol, host=host, port=port))
            items.append(finding)
        return items
<commit_msg>Update parser.py (#4509)<commit_after>import logging

from defusedxml import ElementTree

from dojo.models import Endpoint, Finding

logger = logging.getLogger(__name__)

class Outpost24Parser(object):

    def get_scan_types(self):
        return ["Outpost24 Scan"]

    def get_label_for_scan_types(self, scan_type):
        return "Outpost24 Scan"

    def get_description_for_scan_types(self, scan_type):
        return "Import Outpost24 endpoint vulnerability scan in XML format."

    def get_findings(self, file, test):
        tree = ElementTree.parse(file)
        items = list()
        for detail in tree.iterfind('.//detaillist/detail'):
            # finding details
            title = detail.findtext('name')
            # date = detail.findtext('date') # can be used for Finding.date?
            cve = detail.findtext('./cve/id')
            url = detail.findtext('./referencelist/reference/[type=\'solution\']/../url')
            description = detail.findtext('description')
            mitigation = detail.findtext('solution')
            impact = detail.findtext('information')
            cvss_score = detail.findtext('cvss_v3_score') or detail.findtext('cvss_score')
            if not cvss_score:
                cvss_score = 0
            if cvss_score:
                score = float(cvss_score)
                if score < 4:
                    severity = 'Low'
                elif score < 7:
                    severity = 'Medium'
                elif score < 9:
                    severity = 'High'
                else:
                    severity = 'Critical'
            else:
                risk = int(detail.findtext('risk'))
                if risk == 0:
                    severity = 'Low'
                elif risk == 1:
                    severity = 'Medium'
                elif risk == 2:
                    severity = 'High'
                else:
                    severity = 'Critical'
            cvss_description = detail.findtext('cvss_vector_description')
            cvss_vector = detail.findtext('cvss_v3_vector') or detail.findtext('cvss_vector')
            severity_justification = "{}\n{}".format(cvss_score, cvss_description)
            finding = Finding(title=title, test=test, cve=cve, url=url, description=description, mitigation=mitigation,
                              impact=impact, severity=severity,
                              severity_justification=severity_justification)
            # endpoint details
            host = detail.findtext('ip')
            if host:
                protocol = detail.findtext('./portinfo/service')
                try:
                    port = int(detail.findtext('./portinfo/portnumber'))
                except ValueError as ve:
                    logger.debug("General port given. Assigning 0 as default.")
                    port = 0
                finding.unsaved_endpoints.append(Endpoint(protocol=protocol, host=host, port=port))
            items.append(finding)
        return items
<|endoftext|>


<commit_before>from base64 import b64decode
import datetime
import json
import httplib2
import logging
import sys
import time

import requests

from redash import settings
from redash.query_runner import *
from redash.utils import JSONEncoder

logger = logging.getLogger(__name__)

try:
    import apiclient.errors
    from apiclient.discovery import build
    from apiclient.errors import HttpError
    from oauth2client.client import SignedJwtAssertionCredentials
    from oauth2client import gce

    enabled = True
except ImportError:
    enabled = False

types_map = {
    'INTEGER': TYPE_INTEGER,
    'FLOAT': TYPE_FLOAT,
    'BOOLEAN': TYPE_BOOLEAN,
    'STRING': TYPE_STRING,
    'TIMESTAMP': TYPE_DATETIME,
}

def transform_row(row, fields):
    column_index = 0
    row_data = {}

    for cell in row["f"]:
        field = fields[column_index]
        cell_value = cell['v']

        if cell_value is None:
            pass
        # Otherwise just cast the value
        elif field['type'] == 'INTEGER':
            cell_value = int(cell_value)
        elif field['type'] == 'FLOAT':
            cell_value = float(cell_value)
        elif field['type'] == 'BOOLEAN':
            cell_value = cell_value.lower() == "true"
        elif field['type'] == 'TIMESTAMP':
            cell_value = datetime.datetime.fromtimestamp(float(cell_value))

        row_data[field["name"]] = cell_value
        column_index += 1

    return row_data

def _load_key(filename):
    f = file(filename, "rb")
    try:
        return f.read()
    finally:
        f.close()

def _get_query_results(jobs, project_id, job_id, start_index):
    query_reply = jobs.getQueryResults(projectId=project_id, jobId=job_id, startIndex=start_index).execute()
    logging.debug('query_reply %s', query_reply)
    if not query_reply['jobComplete']:
        time.sleep(10)
        return _get_query_results(jobs, project_id, job_id, start_index)

    return query_reply

class BigQuery(BaseQueryRunner):
    @classmethod
    def enabled(cls):
        return enabled

    @classmethod
    def configuration_schema(cls):
        return {
            'type': 'object',
            'properties': {
                'projectId': {
                    'type': 'string',
                    'title': 'Project ID'
                },
                'jsonKeyFile': {
                    "type": "string",
                    'title': 'JSON Key File'
                },
                'maximumTotalMBytesProcessed': {
                    "type": "number",
                    'title': 'Maximum Total MByte Processed'
                }
            },
            'required': ['jsonKeyFile', 'projectId'],
            'secret': ['jsonKeyFile']
        }

    def __init__(self, configuration_json):
        super(BigQuery, self).__init__(configuration_json)

    def _get_bigquery_service(self):
        scope = [
            "https://www.googleapis.com/auth/bigquery",
            ]

        key = json.loads(b64decode(self.configuration['jsonKeyFile']))

        credentials = SignedJwtAssertionCredentials(key['client_email'], key['private_key'], scope=scope)
        http = httplib2.Http(timeout=settings.BIGQUERY_HTTP_TIMEOUT)
        http = credentials.authorize(http)

        return build("bigquery", "v2", http=http)

    def _get_project_id(self):
        return self.configuration["projectId"]

    def _get_total_bytes_processed(self, jobs, query):
        job_data = {
            "query": query,
            "dryRun": True,
        }
        response = jobs.query(projectId=self._get_project_id(), body=job_data).execute()
        return int(response["totalBytesProcessed"])

    def _get_query_result(self, jobs, query):
        project_id = self._get_project_id()
        job_data = {
            "configuration": {
                "query": {
                    "query": query,
                }
            }
        }

        insert_response = jobs.insert(projectId=project_id, body=job_data).execute()
        current_row = 0
        query_reply = _get_query_results(jobs, project_id=project_id,
                                         job_id=insert_response['jobReference']['jobId'], start_index=current_row)

        logger.debug("bigquery replied: %s", query_reply)

        rows = []

        while ("rows" in query_reply) and current_row < query_reply['totalRows']:
            for row in query_reply["rows"]:
                rows.append(transform_row(row, query_reply["schema"]["fields"]))

            current_row += len(query_reply['rows'])
            query_reply = jobs.getQueryResults(projectId=project_id, jobId=query_reply['jobReference']['jobId'],
                                               startIndex=current_row).execute()

        columns = [{'name': f["name"],
                    'friendly_name': f["name"],
                    'type': types_map.get(f['type'], "string")} for f in query_reply["schema"]["fields"]]

        data = {
            "columns": columns,
            "rows": rows
        }

        return data

    def run_query(self, query):
        logger.debug("BigQuery got query: %s", query)

        bigquery_service = self._get_bigquery_service()
        jobs = bigquery_service.jobs()

        try:
            if "maximumTotalMBytesProcessed" in self.configuration:
                maximumMB = self.configuration["maximumTotalMBytesProcessed"]
                processedMB = self._get_total_bytes_processed(jobs, query) / 1000.0 / 1000.0
                if maximumMB < processedMB:
                    return None, "Too large data will be processed (%f MBytes; maximum: %d MBytes)" % (processedMB, maximumMB)

            data = self._get_query_result(jobs, query)
            error = None

            json_data = json.dumps(data, cls=JSONEncoder)
        except apiclient.errors.HttpError, e:
            json_data = None
            error = e.content
        except KeyboardInterrupt:
            error = "Query cancelled by user."
            json_data = None
        except Exception:
            raise sys.exc_info()[1], None, sys.exc_info()[2]

        return json_data, error

class BigQueryGCE(BigQuery):
    @classmethod
    def type(cls):
        return "bigquery_gce"

    @classmethod
    def configuration_schema(cls):
        return {}

    def _get_project_id(self):
        return requests.get('http://metadata/computeMetadata/v1/project/project-id', headers={'Metadata-Flavor': 'Google'}).content

    def _get_bigquery_service(self):
        credentials = gce.AppAssertionCredentials(scope='https://www.googleapis.com/auth/bigquery')
        http = httplib2.Http()
        http = credentials.authorize(http)

        return build("bigquery", "v2", http=http)

register(BigQuery)
register(BigQueryGCE)
<commit_msg>Rename maximumTotalMBytesProcessed to totalMBytesProcessedLimit.<commit_after>from base64 import b64decode
import datetime
import json
import httplib2
import logging
import sys
import time

import requests

from redash import settings
from redash.query_runner import *
from redash.utils import JSONEncoder

logger = logging.getLogger(__name__)

try:
    import apiclient.errors
    from apiclient.discovery import build
    from apiclient.errors import HttpError
    from oauth2client.client import SignedJwtAssertionCredentials
    from oauth2client import gce

    enabled = True
except ImportError:
    enabled = False

types_map = {
    'INTEGER': TYPE_INTEGER,
    'FLOAT': TYPE_FLOAT,
    'BOOLEAN': TYPE_BOOLEAN,
    'STRING': TYPE_STRING,
    'TIMESTAMP': TYPE_DATETIME,
}

def transform_row(row, fields):
    column_index = 0
    row_data = {}

    for cell in row["f"]:
        field = fields[column_index]
        cell_value = cell['v']

        if cell_value is None:
            pass
        # Otherwise just cast the value
        elif field['type'] == 'INTEGER':
            cell_value = int(cell_value)
        elif field['type'] == 'FLOAT':
            cell_value = float(cell_value)
        elif field['type'] == 'BOOLEAN':
            cell_value = cell_value.lower() == "true"
        elif field['type'] == 'TIMESTAMP':
            cell_value = datetime.datetime.fromtimestamp(float(cell_value))

        row_data[field["name"]] = cell_value
        column_index += 1

    return row_data

def _load_key(filename):
    f = file(filename, "rb")
    try:
        return f.read()
    finally:
        f.close()

def _get_query_results(jobs, project_id, job_id, start_index):
    query_reply = jobs.getQueryResults(projectId=project_id, jobId=job_id, startIndex=start_index).execute()
    logging.debug('query_reply %s', query_reply)
    if not query_reply['jobComplete']:
        time.sleep(10)
        return _get_query_results(jobs, project_id, job_id, start_index)

    return query_reply

class BigQuery(BaseQueryRunner):
    @classmethod
    def enabled(cls):
        return enabled

    @classmethod
    def configuration_schema(cls):
        return {
            'type': 'object',
            'properties': {
                'projectId': {
                    'type': 'string',
                    'title': 'Project ID'
                },
                'jsonKeyFile': {
                    "type": "string",
                    'title': 'JSON Key File'
                },
                'totalMBytesProcessedLimit': {
                    "type": "number",
                    'title': 'Total MByte Processed Limit'
                }
            },
            'required': ['jsonKeyFile', 'projectId'],
            'secret': ['jsonKeyFile']
        }

    def __init__(self, configuration_json):
        super(BigQuery, self).__init__(configuration_json)

    def _get_bigquery_service(self):
        scope = [
            "https://www.googleapis.com/auth/bigquery",
            ]

        key = json.loads(b64decode(self.configuration['jsonKeyFile']))

        credentials = SignedJwtAssertionCredentials(key['client_email'], key['private_key'], scope=scope)
        http = httplib2.Http(timeout=settings.BIGQUERY_HTTP_TIMEOUT)
        http = credentials.authorize(http)

        return build("bigquery", "v2", http=http)

    def _get_project_id(self):
        return self.configuration["projectId"]

    def _get_total_bytes_processed(self, jobs, query):
        job_data = {
            "query": query,
            "dryRun": True,
        }
        response = jobs.query(projectId=self._get_project_id(), body=job_data).execute()
        return int(response["totalBytesProcessed"])

    def _get_query_result(self, jobs, query):
        project_id = self._get_project_id()
        job_data = {
            "configuration": {
                "query": {
                    "query": query,
                }
            }
        }

        insert_response = jobs.insert(projectId=project_id, body=job_data).execute()
        current_row = 0
        query_reply = _get_query_results(jobs, project_id=project_id,
                                         job_id=insert_response['jobReference']['jobId'], start_index=current_row)

        logger.debug("bigquery replied: %s", query_reply)

        rows = []

        while ("rows" in query_reply) and current_row < query_reply['totalRows']:
            for row in query_reply["rows"]:
                rows.append(transform_row(row, query_reply["schema"]["fields"]))

            current_row += len(query_reply['rows'])
            query_reply = jobs.getQueryResults(projectId=project_id, jobId=query_reply['jobReference']['jobId'],
                                               startIndex=current_row).execute()

        columns = [{'name': f["name"],
                    'friendly_name': f["name"],
                    'type': types_map.get(f['type'], "string")} for f in query_reply["schema"]["fields"]]

        data = {
            "columns": columns,
            "rows": rows
        }

        return data

    def run_query(self, query):
        logger.debug("BigQuery got query: %s", query)

        bigquery_service = self._get_bigquery_service()
        jobs = bigquery_service.jobs()

        try:
            if "totalMBytesProcessedLimit" in self.configuration:
                limitMB = self.configuration["totalMBytesProcessedLimit"]
                processedMB = self._get_total_bytes_processed(jobs, query) / 1000.0 / 1000.0
                if limitMB < processedMB:
                    return None, "Larger than %d MBytes will be processed (%f MBytes)" % (limitMB, processedMB)

            data = self._get_query_result(jobs, query)
            error = None

            json_data = json.dumps(data, cls=JSONEncoder)
        except apiclient.errors.HttpError, e:
            json_data = None
            error = e.content
        except KeyboardInterrupt:
            error = "Query cancelled by user."
            json_data = None
        except Exception:
            raise sys.exc_info()[1], None, sys.exc_info()[2]

        return json_data, error

class BigQueryGCE(BigQuery):
    @classmethod
    def type(cls):
        return "bigquery_gce"

    @classmethod
    def configuration_schema(cls):
        return {}

    def _get_project_id(self):
        return requests.get('http://metadata/computeMetadata/v1/project/project-id', headers={'Metadata-Flavor': 'Google'}).content

    def _get_bigquery_service(self):
        credentials = gce.AppAssertionCredentials(scope='https://www.googleapis.com/auth/bigquery')
        http = httplib2.Http()
        http = credentials.authorize(http)

        return build("bigquery", "v2", http=http)

register(BigQuery)
register(BigQueryGCE)
<|endoftext|>


<commit_before>from contextlib import contextmanager
import datetime

from sqlalchemy import create_engine
from sqlalchemy import (
    Column,
    Date,
    DateTime,
    Index,
    LargeBinary,
    SmallInteger,
    String,
    Unicode,
    UniqueConstraint,
)
from sqlalchemy.dialects.mysql import INTEGER as Integer
from sqlalchemy.dialects.mysql import BIGINT as BigInteger
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

_Model = declarative_base()

RADIO_TYPE = {
    '': -1,
    'gsm': 0,
    'cdma': 1,
    'umts': 2,
    'lte': 3,
}
RADIO_TYPE_KEYS = list(RADIO_TYPE.keys())

STAT_TYPE = {
    '': -1,
    'location': 0,
}
STAT_TYPE_INVERSE = dict((v, k) for k, v in STAT_TYPE.items())

# TODO add signal to list of reserved words
# reported upstream at http://www.sqlalchemy.org/trac/ticket/2791
from sqlalchemy.dialects.mysql import base
base.RESERVED_WORDS.add('signal')
del base

class Cell(_Model):
    __tablename__ = 'cell'
    __table_args__ = (
        Index('cell_idx', 'radio', 'mcc', 'mnc', 'lac', 'cid'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    # mapped via RADIO_TYPE
    radio = Column(SmallInteger)
    # int in the range 0-1000
    mcc = Column(SmallInteger)
    # int in the range 0-1000 for gsm
    # int in the range 0-32767 for cdma (system id)
    mnc = Column(Integer)
    lac = Column(Integer)
    cid = Column(Integer)
    psc = Column(Integer)
    range = Column(Integer)

cell_table = Cell.__table__

class CellMeasure(_Model):
    __tablename__ = 'cell_measure'
    __table_args__ = (
        Index('cell_measure_lat_idx', 'lat'),
        Index('cell_measure_lon_idx', 'lon'),
        Index('cell_measure_key_idx', 'radio', 'mcc', 'mnc', 'lac', 'cid'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
            'mysql_row_format': 'compressed',
            'mysql_key_block_size': '4',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    measure_id = Column(BigInteger(unsigned=True))
    created = Column(DateTime)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    time = Column(DateTime)
    accuracy = Column(Integer)
    altitude = Column(Integer)
    altitude_accuracy = Column(Integer)
    # mapped via RADIO_TYPE
    radio = Column(SmallInteger)
    mcc = Column(SmallInteger)
    mnc = Column(Integer)
    lac = Column(Integer)
    cid = Column(Integer)
    psc = Column(Integer)
    asu = Column(SmallInteger)
    signal = Column(SmallInteger)
    ta = Column(SmallInteger)

    def __init__(self, *args, **kw):
        if 'measure_id' not in kw:
            kw['measure_id'] = 0
        if 'created' not in kw:
            kw['created'] = datetime.datetime.utcnow()
        super(CellMeasure, self).__init__(*args, **kw)

cell_measure_table = CellMeasure.__table__

class Wifi(_Model):
    __tablename__ = 'wifi'
    __table_args__ = (
        Index('wifi_key_idx', 'key'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    key = Column(String(40))
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    range = Column(Integer)

wifi_table = Wifi.__table__

class WifiMeasure(_Model):
    __tablename__ = 'wifi_measure'
    __table_args__ = (
        Index('wifi_measure_lat_idx', 'lat'),
        Index('wifi_measure_lon_idx', 'lon'),
        Index('wifi_measure_key_idx', 'key'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
            'mysql_row_format': 'compressed',
            'mysql_key_block_size': '4',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    measure_id = Column(BigInteger(unsigned=True))
    created = Column(DateTime)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    time = Column(DateTime)
    accuracy = Column(Integer)
    altitude = Column(Integer)
    altitude_accuracy = Column(Integer)
    key = Column(String(40))
    channel = Column(SmallInteger)
    signal = Column(SmallInteger)

    def __init__(self, *args, **kw):
        if 'measure_id' not in kw:
            kw['measure_id'] = 0
        if 'created' not in kw:
            kw['created'] = datetime.datetime.utcnow()
        super(WifiMeasure, self).__init__(*args, **kw)

wifi_measure_table = WifiMeasure.__table__

class Measure(_Model):
    __tablename__ = 'measure'
    __table_args__ = (
        Index('measure_created_idx', 'created'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
            'mysql_row_format': 'compressed',
            'mysql_key_block_size': '4',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    created = Column(DateTime)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    time = Column(DateTime)
    accuracy = Column(Integer)
    altitude = Column(Integer)
    altitude_accuracy = Column(Integer)
    radio = Column(SmallInteger)  # mapped via RADIO_TYPE
    # json blobs
    cell = Column(LargeBinary)
    wifi = Column(LargeBinary)

    def __init__(self, *args, **kw):
        if 'created' not in kw:
            kw['created'] = datetime.datetime.utcnow()
        super(Measure, self).__init__(*args, **kw)

measure_table = Measure.__table__

class Score(_Model):
    __tablename__ = 'score'
    __table_args__ = {
        'mysql_engine': 'InnoDB',
        'mysql_charset': 'utf8',
    }

    id = Column(Integer(unsigned=True),
                primary_key=True, autoincrement=True)
    userid = Column(Integer(unsigned=True), index=True, unique=True)
    value = Column(Integer)

score_table = Score.__table__

class Stat(_Model):
    __tablename__ = 'stat'
    __table_args__ = (
        UniqueConstraint('key', 'time', name='stat_key_time_unique'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(Integer(unsigned=True),
                primary_key=True, autoincrement=True)
    # mapped via STAT_TYPE
    key = Column(SmallInteger)
    time = Column(Date)
    value = Column(Integer(unsigned=True))

    @property
    def name(self):
        return STAT_TYPE_INVERSE.get(self.key, '')

    @name.setter
    def name(self, value):
        self.key = STAT_TYPE[value]

stat_table = Stat.__table__

class User(_Model):
    __tablename__ = 'user'
    __table_args__ = (
        Index('user_token_idx', 'token'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(Integer(unsigned=True),
                primary_key=True, autoincrement=True)
    token = Column(String(36))
    nickname = Column(Unicode(128))

user_table = User.__table__

# the request db_sessions and db_tween_factory are inspired by pyramid_tm
# to provide lazy session creation, session closure and automatic
# rollback in case of errors

def db_master_session(request):
    session = getattr(request, '_db_master_session', None)
    if session is None:
        db = request.registry.db_master
        request._db_master_session = session = db.session()
    return session

def db_slave_session(request):
    session = getattr(request, '_db_slave_session', None)
    if session is None:
        db = request.registry.db_slave
        request._db_slave_session = session = db.session()
    return session

@contextmanager
def db_worker_session(database):
    try:
        session = database.session()
        yield session
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

def db_tween_factory(handler, registry):

    def db_tween(request):
        response = handler(request)
        master_session = getattr(request, '_db_master_session', None)
        if master_session is not None:
            # only deal with requests with a session
            if response.status.startswith(('4', '5')):  # pragma: no cover
                # never commit on error
                master_session.rollback()
            master_session.close()
        slave_session = getattr(request, '_db_slave_session', None)
        if slave_session is not None:
            # always rollback/close the `read-only` slave sessions
            try:
                slave_session.rollback()
            finally:
                slave_session.close()
        return response

    return db_tween

class Database(object):

    def __init__(self, uri, socket=None, create=True, echo=False,
                 isolation_level='REPEATABLE READ'):
        options = {
            'pool_recycle': 3600,
            'pool_size': 10,
            'pool_timeout': 10,
            'echo': echo,
            # READ COMMITTED
            'isolation_level': isolation_level,
        }
        options['connect_args'] = {'charset': 'utf8'}
        if socket:  # pragma: no cover
            options['connect_args'] = {'unix_socket': socket}
        options['execution_options'] = {'autocommit': False}
        self.engine = create_engine(uri, **options)
        self.session_factory = sessionmaker(
            bind=self.engine, autocommit=False, autoflush=False)

        # create tables
        if create:
            with self.engine.connect() as conn:
                trans = conn.begin()
                _Model.metadata.create_all(self.engine)
                trans.commit()

    def session(self):
        return self.session_factory()
<commit_msg>add more stat keys<commit_after>from contextlib import contextmanager
import datetime

from sqlalchemy import create_engine
from sqlalchemy import (
    Column,
    Date,
    DateTime,
    Index,
    LargeBinary,
    SmallInteger,
    String,
    Unicode,
    UniqueConstraint,
)
from sqlalchemy.dialects.mysql import INTEGER as Integer
from sqlalchemy.dialects.mysql import BIGINT as BigInteger
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

_Model = declarative_base()

RADIO_TYPE = {
    '': -1,
    'gsm': 0,
    'cdma': 1,
    'umts': 2,
    'lte': 3,
}
RADIO_TYPE_KEYS = list(RADIO_TYPE.keys())

STAT_TYPE = {
    '': -1,
    'location': 0,
    'cell': 1,
    'unique_cell': 2,
    'wifi': 3,
    'unique_wifi': 4,
}
STAT_TYPE_INVERSE = dict((v, k) for k, v in STAT_TYPE.items())

# TODO add signal to list of reserved words
# reported upstream at http://www.sqlalchemy.org/trac/ticket/2791
from sqlalchemy.dialects.mysql import base
base.RESERVED_WORDS.add('signal')
del base

class Cell(_Model):
    __tablename__ = 'cell'
    __table_args__ = (
        Index('cell_idx', 'radio', 'mcc', 'mnc', 'lac', 'cid'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    # mapped via RADIO_TYPE
    radio = Column(SmallInteger)
    # int in the range 0-1000
    mcc = Column(SmallInteger)
    # int in the range 0-1000 for gsm
    # int in the range 0-32767 for cdma (system id)
    mnc = Column(Integer)
    lac = Column(Integer)
    cid = Column(Integer)
    psc = Column(Integer)
    range = Column(Integer)

cell_table = Cell.__table__

class CellMeasure(_Model):
    __tablename__ = 'cell_measure'
    __table_args__ = (
        Index('cell_measure_lat_idx', 'lat'),
        Index('cell_measure_lon_idx', 'lon'),
        Index('cell_measure_key_idx', 'radio', 'mcc', 'mnc', 'lac', 'cid'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
            'mysql_row_format': 'compressed',
            'mysql_key_block_size': '4',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    measure_id = Column(BigInteger(unsigned=True))
    created = Column(DateTime)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    time = Column(DateTime)
    accuracy = Column(Integer)
    altitude = Column(Integer)
    altitude_accuracy = Column(Integer)
    # mapped via RADIO_TYPE
    radio = Column(SmallInteger)
    mcc = Column(SmallInteger)
    mnc = Column(Integer)
    lac = Column(Integer)
    cid = Column(Integer)
    psc = Column(Integer)
    asu = Column(SmallInteger)
    signal = Column(SmallInteger)
    ta = Column(SmallInteger)

    def __init__(self, *args, **kw):
        if 'measure_id' not in kw:
            kw['measure_id'] = 0
        if 'created' not in kw:
            kw['created'] = datetime.datetime.utcnow()
        super(CellMeasure, self).__init__(*args, **kw)

cell_measure_table = CellMeasure.__table__

class Wifi(_Model):
    __tablename__ = 'wifi'
    __table_args__ = (
        Index('wifi_key_idx', 'key'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    key = Column(String(40))
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    range = Column(Integer)

wifi_table = Wifi.__table__

class WifiMeasure(_Model):
    __tablename__ = 'wifi_measure'
    __table_args__ = (
        Index('wifi_measure_lat_idx', 'lat'),
        Index('wifi_measure_lon_idx', 'lon'),
        Index('wifi_measure_key_idx', 'key'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
            'mysql_row_format': 'compressed',
            'mysql_key_block_size': '4',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    measure_id = Column(BigInteger(unsigned=True))
    created = Column(DateTime)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    time = Column(DateTime)
    accuracy = Column(Integer)
    altitude = Column(Integer)
    altitude_accuracy = Column(Integer)
    key = Column(String(40))
    channel = Column(SmallInteger)
    signal = Column(SmallInteger)

    def __init__(self, *args, **kw):
        if 'measure_id' not in kw:
            kw['measure_id'] = 0
        if 'created' not in kw:
            kw['created'] = datetime.datetime.utcnow()
        super(WifiMeasure, self).__init__(*args, **kw)

wifi_measure_table = WifiMeasure.__table__

class Measure(_Model):
    __tablename__ = 'measure'
    __table_args__ = (
        Index('measure_created_idx', 'created'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
            'mysql_row_format': 'compressed',
            'mysql_key_block_size': '4',
        }
    )

    id = Column(BigInteger(unsigned=True),
                primary_key=True, autoincrement=True)
    created = Column(DateTime)
    # lat/lon * decimaljson.FACTOR
    lat = Column(Integer)
    lon = Column(Integer)
    time = Column(DateTime)
    accuracy = Column(Integer)
    altitude = Column(Integer)
    altitude_accuracy = Column(Integer)
    radio = Column(SmallInteger)  # mapped via RADIO_TYPE
    # json blobs
    cell = Column(LargeBinary)
    wifi = Column(LargeBinary)

    def __init__(self, *args, **kw):
        if 'created' not in kw:
            kw['created'] = datetime.datetime.utcnow()
        super(Measure, self).__init__(*args, **kw)

measure_table = Measure.__table__

class Score(_Model):
    __tablename__ = 'score'
    __table_args__ = {
        'mysql_engine': 'InnoDB',
        'mysql_charset': 'utf8',
    }

    id = Column(Integer(unsigned=True),
                primary_key=True, autoincrement=True)
    userid = Column(Integer(unsigned=True), index=True, unique=True)
    value = Column(Integer)

score_table = Score.__table__

class Stat(_Model):
    __tablename__ = 'stat'
    __table_args__ = (
        UniqueConstraint('key', 'time', name='stat_key_time_unique'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(Integer(unsigned=True),
                primary_key=True, autoincrement=True)
    # mapped via STAT_TYPE
    key = Column(SmallInteger)
    time = Column(Date)
    value = Column(Integer(unsigned=True))

    @property
    def name(self):
        return STAT_TYPE_INVERSE.get(self.key, '')

    @name.setter
    def name(self, value):
        self.key = STAT_TYPE[value]

stat_table = Stat.__table__

class User(_Model):
    __tablename__ = 'user'
    __table_args__ = (
        Index('user_token_idx', 'token'),
        {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8',
        }
    )

    id = Column(Integer(unsigned=True),
                primary_key=True, autoincrement=True)
    token = Column(String(36))
    nickname = Column(Unicode(128))

user_table = User.__table__

# the request db_sessions and db_tween_factory are inspired by pyramid_tm
# to provide lazy session creation, session closure and automatic
# rollback in case of errors

def db_master_session(request):
    session = getattr(request, '_db_master_session', None)
    if session is None:
        db = request.registry.db_master
        request._db_master_session = session = db.session()
    return session

def db_slave_session(request):
    session = getattr(request, '_db_slave_session', None)
    if session is None:
        db = request.registry.db_slave
        request._db_slave_session = session = db.session()
    return session

@contextmanager
def db_worker_session(database):
    try:
        session = database.session()
        yield session
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

def db_tween_factory(handler, registry):

    def db_tween(request):
        response = handler(request)
        master_session = getattr(request, '_db_master_session', None)
        if master_session is not None:
            # only deal with requests with a session
            if response.status.startswith(('4', '5')):  # pragma: no cover
                # never commit on error
                master_session.rollback()
            master_session.close()
        slave_session = getattr(request, '_db_slave_session', None)
        if slave_session is not None:
            # always rollback/close the `read-only` slave sessions
            try:
                slave_session.rollback()
            finally:
                slave_session.close()
        return response

    return db_tween

class Database(object):

    def __init__(self, uri, socket=None, create=True, echo=False,
                 isolation_level='REPEATABLE READ'):
        options = {
            'pool_recycle': 3600,
            'pool_size': 10,
            'pool_timeout': 10,
            'echo': echo,
            # READ COMMITTED
            'isolation_level': isolation_level,
        }
        options['connect_args'] = {'charset': 'utf8'}
        if socket:  # pragma: no cover
            options['connect_args'] = {'unix_socket': socket}
        options['execution_options'] = {'autocommit': False}
        self.engine = create_engine(uri, **options)
        self.session_factory = sessionmaker(
            bind=self.engine, autocommit=False, autoflush=False)

        # create tables
        if create:
            with self.engine.connect() as conn:
                trans = conn.begin()
                _Model.metadata.create_all(self.engine)
                trans.commit()

    def session(self):
        return self.session_factory()
<|endoftext|>