<commit_before>import os
import logging
import datetime

from unittest2 import TestCase, SkipTest

from smashrun import Smashrun

client_id = os.environ.get("SMASHRUN_CLIENT_ID")
client_secret = os.environ.get("SMASHRUN_CLIENT_SECRET")
refresh_token = os.environ.get("SMASHRUN_REFRESH_TOKEN")

if os.environ.get("DEBUG"):
    logging.basicConfig(level=logging.DEBUG)

class TestSmashrun(TestCase):

    """Tests for the smashrun client.

    Note that these tests are *not* unit tests and actually run against the
    live Smashrun API, so be nice and don't repeatedly run them a rediculous
    amount.

    """

    @classmethod
    def setUpClass(cls):
        if not all((client_id, client_secret, refresh_token)):
            raise SkipTest("Missing SMASHRUN_REFRESH_TOKEN, "
                           "SMASHRUN_CLIENT_ID, and SMASHRUN_CLIENT_SECRET "
                           "environment variables")

    def setUp(self):
        self.client = Smashrun(client_id=client_id,
                               client_secret=client_secret,
                               redirect_uri='http://localhost')
        self.client.refresh_token(refresh_token=refresh_token)

    def test_get_activities(self):
        activities = list(self.client.get_activities())
        self.assertNotEqual(activities, [])  # small chance this will fail...

    def test_get_activities_style(self):
        activities = self.client.get_activities(count=2, style='ids')
        activity = next(activities)
        self.assertIsInstance(activity, int)

    def test_get_activities_invalid_style(self):
        with self.assertRaises(ValueError):
            self.client.get_activities(style='foobar')

    def test_get_activities_since(self):
        one_month_ago = (datetime.datetime.now() -
                         datetime.timedelta(days=30))
        # TODO: it would be easier to have a test for this if each date field
        # was converted to a datetime...
        list(self.client.get_activities(since=one_month_ago))

    def test_get_activity(self):
        activities = self.client.get_activities()
        activity = next(activities)
        id_num = activity['activityId']
        fetched_activity = self.client.get_activity(id_num)
        self.assertEqual(activity['activityId'],
                         fetched_activity['activityId'])
        self.assertIn('recordingValues', fetched_activity)

    def test_get_badges(self):
        # if it doesn't raise an exception, that's good enough for me
        self.client.get_badges()

    def test_get_current_weight(self):
        self.client.get_current_weight()

    def test_get_weight_history(self):
        self.client.get_weight_history()

    def test_get_stats(self):
        self.client.get_stats()

    def test_get_stats_year(self):
        self.client.get_stats(year=2014)

    def test_get_stats_year_month(self):
        self.client.get_stats(year=2015, month=2)

    def test_get_stats_month_and_no_year(self):
        with self.assertRaises(ValueError):
            self.client.get_stats(month=1)

    def test_get_auth_url(self):
        url = self.client.get_auth_url()[0]
        self.assertIn('client_id', url)
        self.assertIn('client_secret', url)
        self.assertIn('redirect_uri', url)

    def test_get_userinfo(self):
        r = self.client.get_userinfo()
        self.assertIn('id', r)
<commit_msg>Spelling fix in comment<commit_after>import os
import logging
import datetime

from unittest2 import TestCase, SkipTest

from smashrun import Smashrun

client_id = os.environ.get("SMASHRUN_CLIENT_ID")
client_secret = os.environ.get("SMASHRUN_CLIENT_SECRET")
refresh_token = os.environ.get("SMASHRUN_REFRESH_TOKEN")

if os.environ.get("DEBUG"):
    logging.basicConfig(level=logging.DEBUG)

class TestSmashrun(TestCase):

    """Tests for the smashrun client.

    Note that these tests are *not* unit tests and actually run against the
    live Smashrun API, so be nice and don't repeatedly run them a ridiculous
    amount.

    """

    @classmethod
    def setUpClass(cls):
        if not all((client_id, client_secret, refresh_token)):
            raise SkipTest("Missing SMASHRUN_REFRESH_TOKEN, "
                           "SMASHRUN_CLIENT_ID, and SMASHRUN_CLIENT_SECRET "
                           "environment variables")

    def setUp(self):
        self.client = Smashrun(client_id=client_id,
                               client_secret=client_secret,
                               redirect_uri='http://localhost')
        self.client.refresh_token(refresh_token=refresh_token)

    def test_get_activities(self):
        activities = list(self.client.get_activities())
        self.assertNotEqual(activities, [])  # small chance this will fail...

    def test_get_activities_style(self):
        activities = self.client.get_activities(count=2, style='ids')
        activity = next(activities)
        self.assertIsInstance(activity, int)

    def test_get_activities_invalid_style(self):
        with self.assertRaises(ValueError):
            self.client.get_activities(style='foobar')

    def test_get_activities_since(self):
        one_month_ago = (datetime.datetime.now() -
                         datetime.timedelta(days=30))
        # TODO: it would be easier to have a test for this if each date field
        # was converted to a datetime...
        list(self.client.get_activities(since=one_month_ago))

    def test_get_activity(self):
        activities = self.client.get_activities()
        activity = next(activities)
        id_num = activity['activityId']
        fetched_activity = self.client.get_activity(id_num)
        self.assertEqual(activity['activityId'],
                         fetched_activity['activityId'])
        self.assertIn('recordingValues', fetched_activity)

    def test_get_badges(self):
        # if it doesn't raise an exception, that's good enough for me
        self.client.get_badges()

    def test_get_current_weight(self):
        self.client.get_current_weight()

    def test_get_weight_history(self):
        self.client.get_weight_history()

    def test_get_stats(self):
        self.client.get_stats()

    def test_get_stats_year(self):
        self.client.get_stats(year=2014)

    def test_get_stats_year_month(self):
        self.client.get_stats(year=2015, month=2)

    def test_get_stats_month_and_no_year(self):
        with self.assertRaises(ValueError):
            self.client.get_stats(month=1)

    def test_get_auth_url(self):
        url = self.client.get_auth_url()[0]
        self.assertIn('client_id', url)
        self.assertIn('client_secret', url)
        self.assertIn('redirect_uri', url)

    def test_get_userinfo(self):
        r = self.client.get_userinfo()
        self.assertIn('id', r)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
#  obliterate_tests.py:  testing Obliterate
#
#  Subversion is a tool for revision control.
#  See http://subversion.apache.org for more information.
#
# ====================================================================
#    Licensed to the Apache Software Foundation (ASF) under one
#    or more contributor license agreements.  See the NOTICE file
#    distributed with this work for additional information
#    regarding copyright ownership.  The ASF licenses this file
#    to you under the Apache License, Version 2.0 (the
#    "License"); you may not use this file except in compliance
#    with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing,
#    software distributed under the License is distributed on an
#    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#    KIND, either express or implied.  See the License for the
#    specific language governing permissions and limitations
#    under the License.
######################################################################

# General modules
import shutil, sys, re, os

# Our testing module
import svntest
from svntest import main, actions, wc, objects

# (abbreviation)
Item = wc.StateItem
XFail = svntest.testcase.XFail
Skip = svntest.testcase.Skip
SkipUnless = svntest.testcase.SkipUnless

######################################################################
# Test utilities
#

#obliteration_dirs = ['f-mod', 'f-add', 'f-del', 'f-rpl', 'f-mov']
obliteration_dirs = ['f-mod']

def create_dd1_scenarios(wc, repo):
  """Create, in the initially empty repository of the SvnWC WC, the
     obliteration test scenarios depicted in each "Example 1" in
     <notes/obliterate/fspec-dd1/dd1-file-ops.svg>."""

  # r1: base directories
  for dir in obliteration_dirs:
    wc.svn_mkdir(dir)
  wc.svn_commit()

  # r2 to r8 inclusive, just so that the obliteration rev is a round and
  # consistent number (10), no matter what complexity of history we have.
  while repo.head_rev < 8:
    repo.svn_mkdirs('tmp/' + str(repo.head_rev + 1))
  wc.svn_update()

  # r9: add the files used in the scenarios
  if 'f-mod' in obliteration_dirs:
    wc.svn_file_create_add('f-mod/F', "Pear\n")
  if 'f-add' in obliteration_dirs:
    pass  # nothing needed
  if 'f-del' in obliteration_dirs:
    wc.svn_file_create_add('f-del/F', "Pear\n")
  if 'f-rpl' in obliteration_dirs:
    wc.svn_file_create_add('f-rpl/F', "Pear\n")
  if 'f-mov' in obliteration_dirs:
    wc.svn_file_create_add('f-mov/E', "Pear\n")  # 'E' will be moved to 'F'
  wc.svn_commit()

  # r10: the rev in which files named 'F' are to be obliterated
  if 'f-mod' in obliteration_dirs:
    wc.file_modify('f-mod/F', 'Apple\n')
  if 'f-add' in obliteration_dirs:
    wc.svn_file_create_add('f-add/F', 'Apple\n')
  if 'f-del' in obliteration_dirs:
    wc.svn_delete('f-del/F')
  if 'f-rpl' in obliteration_dirs:
    wc.svn_delete('f-rpl/F')
    wc.svn_file_create_add('f-rpl/F', 'Apple\n')
  if 'f-mov' in obliteration_dirs:
    wc.svn_move('f-mov/E', 'f-mov/F')
    wc.file_modify('f-mov/F', 'Apple\n')
  wc.svn_commit(log='Rev to be obliterated')

  # r11: some more recent history that refers to the revision we changed
  # (We are not ready to test this yet.)
  #for dir in obliteration_dirs:
  #  if not dir == 'f-del':
  #    wc.file_modify(dir + '/F', 'Orange\n')
  #wc.svn_commit()

  return 10

def hook_enable(repo):
  """Make a pre-obliterate hook in REPO (an svntest.objects.SvnRepository)
  that allows the user 'jrandom' to do any obliteration."""

  hook_path = os.path.join(repo.repo_absdir, 'hooks', 'pre-obliterate')
  # Embed the text carefully: it might include characters like "%" and "'".
  main.create_python_hook_script(hook_path, 'import sys\n'
    'repos = sys.argv[1]\n'
    'user = sys.argv[2]\n'
    'if user == "jrandom":\n'
    '  sys.exit(0)\n'
    'sys.exit(1)\n')

######################################################################
# Tests
#
#   Each test must return on success or raise on failure.

#----------------------------------------------------------------------

def obliterate_1(sbox):
  "test svn obliterate"

  # Create empty repos and WC
  actions.guarantee_empty_repository(sbox.repo_dir)
  expected_out = svntest.wc.State(sbox.wc_dir, {})
  expected_disk = svntest.wc.State(sbox.wc_dir, {})
  actions.run_and_verify_checkout(sbox.repo_url, sbox.wc_dir, expected_out,
                                  expected_disk)

  # Create test utility objects
  repo = objects.SvnRepository(sbox.repo_url, sbox.repo_dir)
  wc = objects.SvnWC(sbox.wc_dir, repo)

  os.chdir(sbox.wc_dir)

  # Create scenarios ready for obliteration
  apple_rev = create_dd1_scenarios(wc, repo)

  # Dump the repository state, if possible, for debugging
  try:
    repo.dump('before.dump')
  except:
    pass

  hook_enable(repo)

  # Obliterate a file in the revision where the file content was 'Apple'
  for dir in obliteration_dirs:
    repo.obliterate_node_rev(dir + '/F', apple_rev)

  # Dump the repository state, if possible, for debugging
  try:
    repo.dump('after.dump')
  except:
    pass

def pre_obliterate_hook(sbox):
  "test the pre-obliterate hook"

  # Create empty repos and WC
  actions.guarantee_empty_repository(sbox.repo_dir)
  expected_out = svntest.wc.State(sbox.wc_dir, {})
  expected_disk = svntest.wc.State(sbox.wc_dir, {})
  actions.run_and_verify_checkout(sbox.repo_url, sbox.wc_dir, expected_out,
                                  expected_disk)

  # Create test utility objects
  repo = objects.SvnRepository(sbox.repo_url, sbox.repo_dir)
  wc = objects.SvnWC(sbox.wc_dir, repo)

  os.chdir(sbox.wc_dir)

  # Create scenarios ready for obliteration
  apple_rev = create_dd1_scenarios(wc, repo)
  dir = obliteration_dirs[0]

  hook_path = os.path.join(repo.repo_absdir, 'hooks', 'pre-obliterate')

  # Try an obliterate that should be forbidden as no hook is installed
  exp_err = 'svn: Repository has not been enabled to accept obliteration\n'
  repo.obliterate_node_rev(dir + '/F', apple_rev,
                           exp_out=[], exp_err=exp_err, exp_exit=1)

  # Try an obliterate that should be forbidden by the hook
  main.create_python_hook_script(hook_path, 'import sys\n'
    'sys.stderr.write("Pre-oblit, %s, %s" %\n'
    '                 (sys.argv[1], sys.argv[2]))\n'
    'sys.exit(1)\n')
  exp_err = 'svn: Obliteration blocked by pre-obliterate hook (exit code 1) with output:|' + \
            'Pre-oblit, /.*, jrandom'
  repo.obliterate_node_rev(dir + '/F', apple_rev,
                           exp_out=[], exp_err=exp_err, exp_exit=1)

  # Try an obliterate that should be allowed by the hook
  hook_enable(repo)
  repo.obliterate_node_rev(dir + '/F', apple_rev)

########################################################################
# Run the tests

def supports_obliterate():
  return svntest.main.is_ra_type_file() and not svntest.main.is_fs_type_fsfs()

# list all tests here, starting with None:
test_list = [ None,
              SkipUnless(obliterate_1, supports_obliterate),
              SkipUnless(pre_obliterate_hook, supports_obliterate),
             ]

if __name__ == '__main__':
  svntest.main.run_tests(test_list)
  # NOTREACHED
<commit_msg>Avoid spurious obliterate XFAIL with BDB.<commit_after>#!/usr/bin/env python
#
#  obliterate_tests.py:  testing Obliterate
#
#  Subversion is a tool for revision control.
#  See http://subversion.apache.org for more information.
#
# ====================================================================
#    Licensed to the Apache Software Foundation (ASF) under one
#    or more contributor license agreements.  See the NOTICE file
#    distributed with this work for additional information
#    regarding copyright ownership.  The ASF licenses this file
#    to you under the Apache License, Version 2.0 (the
#    "License"); you may not use this file except in compliance
#    with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing,
#    software distributed under the License is distributed on an
#    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#    KIND, either express or implied.  See the License for the
#    specific language governing permissions and limitations
#    under the License.
######################################################################

# General modules
import shutil, sys, re, os

# Our testing module
import svntest
from svntest import main, actions, wc, objects

# (abbreviation)
Item = wc.StateItem
XFail = svntest.testcase.XFail
Skip = svntest.testcase.Skip
SkipUnless = svntest.testcase.SkipUnless

######################################################################
# Test utilities
#

#obliteration_dirs = ['f-mod', 'f-add', 'f-del', 'f-rpl', 'f-mov']
obliteration_dirs = ['f-mod']

def create_dd1_scenarios(wc, repo):
  """Create, in the initially empty repository of the SvnWC WC, the
     obliteration test scenarios depicted in each "Example 1" in
     <notes/obliterate/fspec-dd1/dd1-file-ops.svg>."""

  # r1: base directories
  for dir in obliteration_dirs:
    wc.svn_mkdir(dir)
  wc.svn_commit()

  # r2 to r8 inclusive, just so that the obliteration rev is a round and
  # consistent number (10), no matter what complexity of history we have.
  while repo.head_rev < 8:
    repo.svn_mkdirs('tmp/' + str(repo.head_rev + 1))
  wc.svn_update()

  # r9: add the files used in the scenarios
  if 'f-mod' in obliteration_dirs:
    wc.svn_file_create_add('f-mod/F', "Pear\n")
  if 'f-add' in obliteration_dirs:
    pass  # nothing needed
  if 'f-del' in obliteration_dirs:
    wc.svn_file_create_add('f-del/F', "Pear\n")
  if 'f-rpl' in obliteration_dirs:
    wc.svn_file_create_add('f-rpl/F', "Pear\n")
  if 'f-mov' in obliteration_dirs:
    wc.svn_file_create_add('f-mov/E', "Pear\n")  # 'E' will be moved to 'F'
  wc.svn_commit()

  # r10: the rev in which files named 'F' are to be obliterated
  if 'f-mod' in obliteration_dirs:
    wc.file_modify('f-mod/F', 'Apple\n')
  if 'f-add' in obliteration_dirs:
    wc.svn_file_create_add('f-add/F', 'Apple\n')
  if 'f-del' in obliteration_dirs:
    wc.svn_delete('f-del/F')
  if 'f-rpl' in obliteration_dirs:
    wc.svn_delete('f-rpl/F')
    wc.svn_file_create_add('f-rpl/F', 'Apple\n')
  if 'f-mov' in obliteration_dirs:
    wc.svn_move('f-mov/E', 'f-mov/F')
    wc.file_modify('f-mov/F', 'Apple\n')
  wc.svn_commit(log='Rev to be obliterated')

  # r11: some more recent history that refers to the revision we changed
  # (We are not ready to test this yet.)
  #for dir in obliteration_dirs:
  #  if not dir == 'f-del':
  #    wc.file_modify(dir + '/F', 'Orange\n')
  #wc.svn_commit()

  return 10

def hook_enable(repo):
  """Make a pre-obliterate hook in REPO (an svntest.objects.SvnRepository)
  that allows the user 'jrandom' to do any obliteration."""

  hook_path = os.path.join(repo.repo_absdir, 'hooks', 'pre-obliterate')
  # Embed the text carefully: it might include characters like "%" and "'".
  main.create_python_hook_script(hook_path, 'import sys\n'
    'repos = sys.argv[1]\n'
    'user = sys.argv[2]\n'
    'if user == "jrandom":\n'
    '  sys.exit(0)\n'
    'sys.exit(1)\n')

######################################################################
# Tests
#
#   Each test must return on success or raise on failure.

#----------------------------------------------------------------------

def obliterate_1(sbox):
  "test svn obliterate"

  # Create empty repos and WC
  actions.guarantee_empty_repository(sbox.repo_dir)
  expected_out = svntest.wc.State(sbox.wc_dir, {})
  expected_disk = svntest.wc.State(sbox.wc_dir, {})
  actions.run_and_verify_checkout(sbox.repo_url, sbox.wc_dir, expected_out,
                                  expected_disk)

  # Create test utility objects
  repo = objects.SvnRepository(sbox.repo_url, sbox.repo_dir)
  wc = objects.SvnWC(sbox.wc_dir, repo)

  os.chdir(sbox.wc_dir)

  # Create scenarios ready for obliteration
  apple_rev = create_dd1_scenarios(wc, repo)

  # Dump the repository state, if possible, for debugging
  try:
    repo.dump('before.dump')
  except:
    pass

  hook_enable(repo)

  # Obliterate a file in the revision where the file content was 'Apple'
  for dir in obliteration_dirs:
    repo.obliterate_node_rev(dir + '/F', apple_rev)

  # Dump the repository state, if possible, for debugging
  try:
    repo.dump('after.dump')
  except:
    pass

def pre_obliterate_hook(sbox):
  "test the pre-obliterate hook"

  # Create empty repos and WC
  actions.guarantee_empty_repository(sbox.repo_dir)
  expected_out = svntest.wc.State(sbox.wc_dir, {})
  expected_disk = svntest.wc.State(sbox.wc_dir, {})
  actions.run_and_verify_checkout(sbox.repo_url, sbox.wc_dir, expected_out,
                                  expected_disk)

  # Create test utility objects
  repo = objects.SvnRepository(sbox.repo_url, sbox.repo_dir)
  wc = objects.SvnWC(sbox.wc_dir, repo)

  os.chdir(sbox.wc_dir)

  # Create scenarios ready for obliteration
  apple_rev = create_dd1_scenarios(wc, repo)
  dir = obliteration_dirs[0]

  hook_path = os.path.join(repo.repo_absdir, 'hooks', 'pre-obliterate')

  # Try an obliterate that should be forbidden as no hook is installed
  exp_err = 'svn: Repository has not been enabled to accept obliteration\n'
  repo.obliterate_node_rev(dir + '/F', apple_rev,
                           exp_out=[], exp_err=exp_err, exp_exit=1)

  # Try an obliterate that should be forbidden by the hook
  main.create_python_hook_script(hook_path, 'import sys\n'
    'sys.stderr.write("Pre-oblit, %s, %s" %\n'
    '                 (sys.argv[1], sys.argv[2]))\n'
    'sys.exit(1)\n')
  exp_err = 'svn: Obliteration blocked by pre-obliterate hook (exit code 1) with output:|' + \
            'Pre-oblit, /.*, jrandom'
  repo.obliterate_node_rev(dir + '/F', apple_rev,
                           exp_out=[], exp_err=exp_err, exp_exit=1)

  # Try an obliterate that should be allowed by the hook
  hook_enable(repo)
  repo.obliterate_node_rev(dir + '/F', apple_rev)

########################################################################
# Run the tests

def supports_obliterate():
  if svntest.main.is_ra_type_file() and not svntest.main.is_fs_type_fsfs():
    code, output, error = svntest.main.run_svn(None, "help")
    for line in output:
      if line.find("obliterate") != -1:
        return True
  return False

# list all tests here, starting with None:
test_list = [ None,
              SkipUnless(obliterate_1, supports_obliterate),
              SkipUnless(pre_obliterate_hook, supports_obliterate),
             ]

if __name__ == '__main__':
  svntest.main.run_tests(test_list)
  # NOTREACHED
<|endoftext|>


<commit_before>from __future__ import unicode_literals, division, absolute_import
from builtins import *  # pylint: disable=unused-import, redefined-builtin

import json
import logging
import os
import re
from collections import deque
from functools import wraps

from flask import Flask, request, jsonify, make_response
from flask_compress import Compress
from flask_cors import CORS
from flask_restplus import Model, Api as RestPlusAPI
from flask_restplus import Resource
from flexget import manager
from flexget.config_schema import process_config
from flexget.utils.database import with_session
from flexget.webserver import User
from jsonschema import RefResolutionError
from werkzeug.http import generate_etag

from . import __path__

__version__ = '1.0-RC1'

log = logging.getLogger('api')

class APIClient(object):
    """
    This is an client which can be used as a more pythonic interface to the rest api.

    It skips http, and is only usable from within the running flexget process.
    """

    def __init__(self):
        self.app = api_app.test_client()

    def __getattr__(self, item):
        return APIEndpoint('/api/' + item, self.get_endpoint)

    def get_endpoint(self, url, data=None, method=None):
        if method is None:
            method = 'POST' if data is not None else 'GET'
        auth_header = dict(Authorization='Token %s' % api_key())
        response = self.app.open(url, data=data, follow_redirects=True, method=method, headers=auth_header)
        result = json.loads(response.get_data(as_text=True))
        # TODO: Proper exceptions
        if 200 > response.status_code >= 300:
            raise Exception(result['error'])
        return result

class APIEndpoint(object):
    def __init__(self, endpoint, caller):
        self.endpoint = endpoint
        self.caller = caller

    def __getattr__(self, item):
        return self.__class__(self.endpoint + '/' + item, self.caller)

    __getitem__ = __getattr__

    def __call__(self, data=None, method=None):
        return self.caller(self.endpoint, data=data, method=method)

def api_version(f):
    """ Add the 'X-API-Version' header to all responses """

    @wraps(f)
    def wrapped(*args, **kwargs):
        rv = f(*args, **kwargs)
        rv.headers['API-Version'] = __version__
        return rv

    return wrapped

class APIResource(Resource):
    """All api resources should subclass this class."""
    method_decorators = [with_session, api_version]

    def __init__(self, api, *args, **kwargs):
        self.manager = manager.manager
        super(APIResource, self).__init__(api, *args, **kwargs)

class APISchemaModel(Model):
    """A flask restplus :class:`flask_restplus.models.ApiModel` which can take a json schema directly."""

    def __init__(self, name, schema, *args, **kwargs):
        super(APISchemaModel, self).__init__(name, *args, **kwargs)
        self._schema = schema

    @property
    def __schema__(self):
        if self.__parent__:
            return {
                'allOf': [
                    {'$ref': '#/definitions/{0}'.format(self.__parent__.name)},
                    self._schema
                ]
            }
        else:
            return self._schema

    def __nonzero__(self):
        return bool(self._schema)

    def __bool__(self):
        return self._schema is not None

    def __repr__(self):
        return '<ApiSchemaModel(%r)>' % self._schema

class API(RestPlusAPI):
    """
    Extends a flask restplus :class:`flask_restplus.Api` with:
      - methods to make using json schemas easier
      - methods to auto document and handle :class:`ApiError` responses
    """

    def _rewrite_refs(self, schema):
        if isinstance(schema, list):
            for value in schema:
                self._rewrite_refs(value)

        if isinstance(schema, dict):
            for key, value in schema.items():
                if isinstance(value, (list, dict)):
                    self._rewrite_refs(value)

                if key == '$ref' and value.startswith('/'):
                    schema[key] = '#definitions%s' % value

    def schema(self, name, schema, **kwargs):
        """
        Register a json schema.

        Usable like :meth:`flask_restplus.Api.model`, except takes a json schema as its argument.

        :returns: An :class:`ApiSchemaModel` instance registered to this api.
        """
        model = APISchemaModel(name, schema, **kwargs)
        model.__apidoc__.update(kwargs)
        self.models[name] = model
        return model

    def inherit(self, name, parent, fields):
        """
        Extends :meth:`flask_restplus.Api.inherit` to allow `fields` to be a json schema, if `parent` is a
        :class:`ApiSchemaModel`.
        """
        if isinstance(parent, APISchemaModel):
            model = APISchemaModel(name, fields)
            model.__apidoc__['name'] = name
            model.__parent__ = parent
            self.models[name] = model
            return model
        return super(API, self).inherit(name, parent, fields)

    def validate(self, model, schema_override=None, description=None):
        """
        When a method is decorated with this, json data submitted to the endpoint will be validated with the given
        `model`. This also auto-documents the expected model, as well as the possible :class:`ValidationError` response.
        """

        def decorator(func):
            @api.expect((model, description))
            @api.response(ValidationError)
            @wraps(func)
            def wrapper(*args, **kwargs):
                payload = request.json
                try:
                    schema = schema_override if schema_override else model.__schema__
                    errors = process_config(config=payload, schema=schema, set_defaults=False)

                    if errors:
                        raise ValidationError(errors)
                except RefResolutionError as e:
                    raise APIError(str(e))
                return func(*args, **kwargs)

            return wrapper

        return decorator

    def response(self, code_or_apierror, description='Success', model=None, **kwargs):
        """
        Extends :meth:`flask_restplus.Api.response` to allow passing an :class:`ApiError` class instead of
        response code. If an `ApiError` is used, the response code, and expected response model, is automatically
        documented.
        """
        try:
            if issubclass(code_or_apierror, APIError):
                description = code_or_apierror.description or description
                return self.doc(
                    responses={code_or_apierror.status_code: (description, code_or_apierror.response_model)}, **kwargs)
        except TypeError:
            # If first argument isn't a class this happens
            pass
        return super(API, self).response(code_or_apierror, description, model=model, **kwargs)

    def pagination_parser(self, parser=None, sort_choices=None, default=None, add_sort=None):
        """
        Return a standardized pagination parser, to be used for any endpoint that has pagination.

        :param parser: Can extend a given parser or create a new one
        :param sort_choices: A tuple of strings, to be used as server side attribute searches
        :param default: The default sort string, used `sort_choices[0]` if not given
        :return: An api.parser() instance with pagination and sorting arguments.
        """
        pagination = parser.copy() if parser else self.parser()
        pagination.add_argument('page', type=int, default=1, help='Page number')
        pagination.add_argument('per_page', type=int, default=50, help='Results per page')
        if sort_choices or add_sort:
            pagination.add_argument('order', choices=('desc', 'asc'), default='desc', help='Sorting order')
        if sort_choices:
            pagination.add_argument('sort_by', choices=sort_choices, default=default or sort_choices[0],
                                    help='Sort by attribute')

        return pagination

api_app = Flask(__name__, template_folder=os.path.join(os.path.dirname(__path__[0]), 'templates'))
api_app.config['REMEMBER_COOKIE_NAME'] = 'flexget.token'
api_app.config['DEBUG'] = True
api_app.config['ERROR_404_HELP'] = False

CORS(api_app)
Compress(api_app)

api = API(
    api_app,
    title='Flexget API v{}'.format(__version__),
    version=__version__,
    description='View and manage flexget core operations and plugins. Open each endpoint for detailed usage or'
                ' navigate to http://flexget.com/API for more details.'
)

base_message = {
    'type': 'object',
    'properties': {
        'status_code': {'type': 'integer'},
        'message': {'type': 'string'},
        'status': {'type': 'string'}
    },
    'required': ['status_code', 'message', 'status']

}

base_message_schema = api.schema('base_message', base_message)

class APIError(Exception):
    description = 'Server error'
    status_code = 500
    status = 'Error'
    response_model = base_message_schema

    def __init__(self, message=None, payload=None):
        self.message = message
        self.payload = payload

    def to_dict(self):
        rv = self.payload or {}
        rv.update(status_code=self.status_code, message=self.message, status=self.status)
        return rv

    @classmethod
    def schema(cls):
        return cls.response_model.__schema__

class NotFoundError(APIError):
    status_code = 404
    description = 'Not found'

class Unauthorized(APIError):
    status_code = 401
    description = 'Unauthorized'

class BadRequest(APIError):
    status_code = 400
    description = 'Bad request'

class Conflict(APIError):
    status_code = 409
    description = 'Conflict'

class PreconditionFailed(APIError):
    status_code = 412
    description = 'Precondition failed'

class NotModified(APIError):
    status_code = 304
    description = 'not modified'

class ValidationError(APIError):
    status_code = 422
    description = 'Validation error'

    response_model = api.inherit('validation_error', APIError.response_model, {
        'type': 'object',
        'properties': {
            'validation_errors': {
                'type': 'array',
                'items': {
                    'type': 'object',
                    'properties': {
                        'message': {'type': 'string', 'description': 'A human readable message explaining the error.'},
                        'validator': {'type': 'string', 'description': 'The name of the failed validator.'},
                        'validator_value': {
                            'type': 'string', 'description': 'The value for the failed validator in the schema.'
                        },
                        'path': {'type': 'string'},
                        'schema_path': {'type': 'string'},
                    }
                }
            }
        },
        'required': ['validation_errors']
    })

    verror_attrs = (
        'message', 'cause', 'validator', 'validator_value',
        'path', 'schema_path', 'parent'
    )

    def __init__(self, validation_errors, message='validation error'):
        payload = {'validation_errors': [self._verror_to_dict(error) for error in validation_errors]}
        super(ValidationError, self).__init__(message, payload=payload)

    def _verror_to_dict(self, error):
        error_dict = {}
        for attr in self.verror_attrs:
            if isinstance(getattr(error, attr), deque):
                error_dict[attr] = list(getattr(error, attr))
            else:
                error_dict[attr] = str(getattr(error, attr))
        return error_dict

empty_response = api.schema('empty', {'type': 'object'})

def success_response(message, status_code=200, status='success'):
    rsp_dict = {
        'message': message,
        'status_code': status_code,
        'status': status
    }
    rsp = jsonify(rsp_dict)
    rsp.status_code = status_code
    return rsp

@api.errorhandler(APIError)
@api.errorhandler(NotFoundError)
@api.errorhandler(ValidationError)
@api.errorhandler(BadRequest)
@api.errorhandler(Unauthorized)
@api.errorhandler(Conflict)
@api.errorhandler(NotModified)
@api.errorhandler(PreconditionFailed)
def api_errors(error):
    return error.to_dict(), error.status_code

@with_session
def api_key(session=None):
    log.debug('fetching token for internal lookup')
    return session.query(User).first().token

def etag(f):
    """
    A decorator that add an ETag header to the response and checks for the "If-Match" and "If-Not-Match" headers to
     return an appropriate response.

    :param f: A GET or HEAD flask method to wrap
    :return: The method's response with the ETag and Cache-Control headers, raises a 412 error or returns a 304 response
    """

    @wraps(f)
    def wrapped(*args, **kwargs):
        # Identify if this is a GET or HEAD in order to proceed
        assert request.method in ['HEAD', 'GET'], '@etag is only supported for GET requests'
        rv = f(*args, **kwargs)
        rv = make_response(rv)
        etag = generate_etag(rv.get_data())
        rv.headers['Cache-Control'] = 'max-age=86400'
        rv.headers['ETag'] = etag
        if_match = request.headers.get('If-Match')
        if_none_match = request.headers.get('If-None-Match')

        if if_match:
            etag_list = [tag.strip() for tag in if_match.split(',')]
            if etag not in etag_list and '*' not in etag_list:
                raise PreconditionFailed('etag does not match')
        elif if_none_match:
            etag_list = [tag.strip() for tag in if_none_match.split(',')]
            if etag in etag_list or '*' in etag_list:
                raise NotModified

        return rv

    return wrapped

def pagination_headers(total_pages, total_items, page_count, request):
    """
    Creates the `Link`. 'Count' and  'Total-Count' headers, to be used for pagination traversing

    :param total_pages: Total number of pages
    :param total_items: Total number of items in all the pages
    :param page_count: Item count for page (may differ from page size request)
    :param request: The flask request used, required to build other reoccurring vars like url and such.
    :return:
    """

    # Build constant variables from request data
    url = request.url_root + request.path.lstrip('/')
    per_page = request.args.get('per_page', 50)
    page = int(request.args.get('page', 1))

    # Build the base template
    LINKTEMPLATE = '<{}?per_page={}&'.format(url, per_page)

    # Removed page and per_page from query string
    query_string = re.sub(b'&?per_page=\d+&?', b'', request.query_string)
    query_string = re.sub(b'&?page=\d+&?', b'', query_string)

    # Add all original query params
    LINKTEMPLATE = LINKTEMPLATE.rstrip('&') + query_string.decode() + '&page={}>; rel="{}"'

    link_string = ''

    if page > 1:
        link_string += LINKTEMPLATE.format(page - 1, 'prev') + ', '
    if page < total_pages:
        link_string += LINKTEMPLATE.format(page + 1, 'next') + ', '
    link_string += LINKTEMPLATE.format(total_pages, 'last')

    return {
        'Link': link_string,
        'Total-Count': total_items,
        'Count': page_count
    }
<commit_msg>minor doc change<commit_after>from __future__ import unicode_literals, division, absolute_import
from builtins import *  # pylint: disable=unused-import, redefined-builtin

import json
import logging
import os
import re
from collections import deque
from functools import wraps

from flask import Flask, request, jsonify, make_response
from flask_compress import Compress
from flask_cors import CORS
from flask_restplus import Model, Api as RestPlusAPI
from flask_restplus import Resource
from flexget import manager
from flexget.config_schema import process_config
from flexget.utils.database import with_session
from flexget.webserver import User
from jsonschema import RefResolutionError
from werkzeug.http import generate_etag

from . import __path__

__version__ = '1.0-RC1'

log = logging.getLogger('api')

class APIClient(object):
    """
    This is an client which can be used as a more pythonic interface to the rest api.

    It skips http, and is only usable from within the running flexget process.
    """

    def __init__(self):
        self.app = api_app.test_client()

    def __getattr__(self, item):
        return APIEndpoint('/api/' + item, self.get_endpoint)

    def get_endpoint(self, url, data=None, method=None):
        if method is None:
            method = 'POST' if data is not None else 'GET'
        auth_header = dict(Authorization='Token %s' % api_key())
        response = self.app.open(url, data=data, follow_redirects=True, method=method, headers=auth_header)
        result = json.loads(response.get_data(as_text=True))
        # TODO: Proper exceptions
        if 200 > response.status_code >= 300:
            raise Exception(result['error'])
        return result

class APIEndpoint(object):
    def __init__(self, endpoint, caller):
        self.endpoint = endpoint
        self.caller = caller

    def __getattr__(self, item):
        return self.__class__(self.endpoint + '/' + item, self.caller)

    __getitem__ = __getattr__

    def __call__(self, data=None, method=None):
        return self.caller(self.endpoint, data=data, method=method)

def api_version(f):
    """ Add the 'X-API-Version' header to all responses """

    @wraps(f)
    def wrapped(*args, **kwargs):
        rv = f(*args, **kwargs)
        rv.headers['API-Version'] = __version__
        return rv

    return wrapped

class APIResource(Resource):
    """All api resources should subclass this class."""
    method_decorators = [with_session, api_version]

    def __init__(self, api, *args, **kwargs):
        self.manager = manager.manager
        super(APIResource, self).__init__(api, *args, **kwargs)

class APISchemaModel(Model):
    """A flask restplus :class:`flask_restplus.models.ApiModel` which can take a json schema directly."""

    def __init__(self, name, schema, *args, **kwargs):
        super(APISchemaModel, self).__init__(name, *args, **kwargs)
        self._schema = schema

    @property
    def __schema__(self):
        if self.__parent__:
            return {
                'allOf': [
                    {'$ref': '#/definitions/{0}'.format(self.__parent__.name)},
                    self._schema
                ]
            }
        else:
            return self._schema

    def __nonzero__(self):
        return bool(self._schema)

    def __bool__(self):
        return self._schema is not None

    def __repr__(self):
        return '<ApiSchemaModel(%r)>' % self._schema

class API(RestPlusAPI):
    """
    Extends a flask restplus :class:`flask_restplus.Api` with:
      - methods to make using json schemas easier
      - methods to auto document and handle :class:`ApiError` responses
    """

    def _rewrite_refs(self, schema):
        if isinstance(schema, list):
            for value in schema:
                self._rewrite_refs(value)

        if isinstance(schema, dict):
            for key, value in schema.items():
                if isinstance(value, (list, dict)):
                    self._rewrite_refs(value)

                if key == '$ref' and value.startswith('/'):
                    schema[key] = '#definitions%s' % value

    def schema(self, name, schema, **kwargs):
        """
        Register a json schema.

        Usable like :meth:`flask_restplus.Api.model`, except takes a json schema as its argument.

        :returns: An :class:`ApiSchemaModel` instance registered to this api.
        """
        model = APISchemaModel(name, schema, **kwargs)
        model.__apidoc__.update(kwargs)
        self.models[name] = model
        return model

    def inherit(self, name, parent, fields):
        """
        Extends :meth:`flask_restplus.Api.inherit` to allow `fields` to be a json schema, if `parent` is a
        :class:`ApiSchemaModel`.
        """
        if isinstance(parent, APISchemaModel):
            model = APISchemaModel(name, fields)
            model.__apidoc__['name'] = name
            model.__parent__ = parent
            self.models[name] = model
            return model
        return super(API, self).inherit(name, parent, fields)

    def validate(self, model, schema_override=None, description=None):
        """
        When a method is decorated with this, json data submitted to the endpoint will be validated with the given
        `model`. This also auto-documents the expected model, as well as the possible :class:`ValidationError` response.
        """

        def decorator(func):
            @api.expect((model, description))
            @api.response(ValidationError)
            @wraps(func)
            def wrapper(*args, **kwargs):
                payload = request.json
                try:
                    schema = schema_override if schema_override else model.__schema__
                    errors = process_config(config=payload, schema=schema, set_defaults=False)

                    if errors:
                        raise ValidationError(errors)
                except RefResolutionError as e:
                    raise APIError(str(e))
                return func(*args, **kwargs)

            return wrapper

        return decorator

    def response(self, code_or_apierror, description='Success', model=None, **kwargs):
        """
        Extends :meth:`flask_restplus.Api.response` to allow passing an :class:`ApiError` class instead of
        response code. If an `ApiError` is used, the response code, and expected response model, is automatically
        documented.
        """
        try:
            if issubclass(code_or_apierror, APIError):
                description = code_or_apierror.description or description
                return self.doc(
                    responses={code_or_apierror.status_code: (description, code_or_apierror.response_model)}, **kwargs)
        except TypeError:
            # If first argument isn't a class this happens
            pass
        return super(API, self).response(code_or_apierror, description, model=model, **kwargs)

    def pagination_parser(self, parser=None, sort_choices=None, default=None, add_sort=None):
        """
        Return a standardized pagination parser, to be used for any endpoint that has pagination.

        :param parser: Can extend a given parser or create a new one
        :param sort_choices: A tuple of strings, to be used as server side attribute searches
        :param default: The default sort string, used `sort_choices[0]` if not given
        :return: An api.parser() instance with pagination and sorting arguments.
        """
        pagination = parser.copy() if parser else self.parser()
        pagination.add_argument('page', type=int, default=1, help='Page number')
        pagination.add_argument('per_page', type=int, default=50, help='Results per page')
        if sort_choices or add_sort:
            pagination.add_argument('order', choices=('desc', 'asc'), default='desc', help='Sorting order')
        if sort_choices:
            pagination.add_argument('sort_by', choices=sort_choices, default=default or sort_choices[0],
                                    help='Sort by attribute')

        return pagination

api_app = Flask(__name__, template_folder=os.path.join(os.path.dirname(__path__[0]), 'templates'))
api_app.config['REMEMBER_COOKIE_NAME'] = 'flexget.token'
api_app.config['DEBUG'] = True
api_app.config['ERROR_404_HELP'] = False

CORS(api_app)
Compress(api_app)

api = API(
    api_app,
    title='Flexget API v{}'.format(__version__),
    version=__version__,
    description='View and manage flexget core operations and plugins. Open each endpoint view for usage information.'
                ' Navigate to http://flexget.com/API for more details.'
)

base_message = {
    'type': 'object',
    'properties': {
        'status_code': {'type': 'integer'},
        'message': {'type': 'string'},
        'status': {'type': 'string'}
    },
    'required': ['status_code', 'message', 'status']

}

base_message_schema = api.schema('base_message', base_message)

class APIError(Exception):
    description = 'Server error'
    status_code = 500
    status = 'Error'
    response_model = base_message_schema

    def __init__(self, message=None, payload=None):
        self.message = message
        self.payload = payload

    def to_dict(self):
        rv = self.payload or {}
        rv.update(status_code=self.status_code, message=self.message, status=self.status)
        return rv

    @classmethod
    def schema(cls):
        return cls.response_model.__schema__

class NotFoundError(APIError):
    status_code = 404
    description = 'Not found'

class Unauthorized(APIError):
    status_code = 401
    description = 'Unauthorized'

class BadRequest(APIError):
    status_code = 400
    description = 'Bad request'

class Conflict(APIError):
    status_code = 409
    description = 'Conflict'

class PreconditionFailed(APIError):
    status_code = 412
    description = 'Precondition failed'

class NotModified(APIError):
    status_code = 304
    description = 'not modified'

class ValidationError(APIError):
    status_code = 422
    description = 'Validation error'

    response_model = api.inherit('validation_error', APIError.response_model, {
        'type': 'object',
        'properties': {
            'validation_errors': {
                'type': 'array',
                'items': {
                    'type': 'object',
                    'properties': {
                        'message': {'type': 'string', 'description': 'A human readable message explaining the error.'},
                        'validator': {'type': 'string', 'description': 'The name of the failed validator.'},
                        'validator_value': {
                            'type': 'string', 'description': 'The value for the failed validator in the schema.'
                        },
                        'path': {'type': 'string'},
                        'schema_path': {'type': 'string'},
                    }
                }
            }
        },
        'required': ['validation_errors']
    })

    verror_attrs = (
        'message', 'cause', 'validator', 'validator_value',
        'path', 'schema_path', 'parent'
    )

    def __init__(self, validation_errors, message='validation error'):
        payload = {'validation_errors': [self._verror_to_dict(error) for error in validation_errors]}
        super(ValidationError, self).__init__(message, payload=payload)

    def _verror_to_dict(self, error):
        error_dict = {}
        for attr in self.verror_attrs:
            if isinstance(getattr(error, attr), deque):
                error_dict[attr] = list(getattr(error, attr))
            else:
                error_dict[attr] = str(getattr(error, attr))
        return error_dict

empty_response = api.schema('empty', {'type': 'object'})

def success_response(message, status_code=200, status='success'):
    rsp_dict = {
        'message': message,
        'status_code': status_code,
        'status': status
    }
    rsp = jsonify(rsp_dict)
    rsp.status_code = status_code
    return rsp

@api.errorhandler(APIError)
@api.errorhandler(NotFoundError)
@api.errorhandler(ValidationError)
@api.errorhandler(BadRequest)
@api.errorhandler(Unauthorized)
@api.errorhandler(Conflict)
@api.errorhandler(NotModified)
@api.errorhandler(PreconditionFailed)
def api_errors(error):
    return error.to_dict(), error.status_code

@with_session
def api_key(session=None):
    log.debug('fetching token for internal lookup')
    return session.query(User).first().token

def etag(f):
    """
    A decorator that add an ETag header to the response and checks for the "If-Match" and "If-Not-Match" headers to
     return an appropriate response.

    :param f: A GET or HEAD flask method to wrap
    :return: The method's response with the ETag and Cache-Control headers, raises a 412 error or returns a 304 response
    """

    @wraps(f)
    def wrapped(*args, **kwargs):
        # Identify if this is a GET or HEAD in order to proceed
        assert request.method in ['HEAD', 'GET'], '@etag is only supported for GET requests'
        rv = f(*args, **kwargs)
        rv = make_response(rv)
        etag = generate_etag(rv.get_data())
        rv.headers['Cache-Control'] = 'max-age=86400'
        rv.headers['ETag'] = etag
        if_match = request.headers.get('If-Match')
        if_none_match = request.headers.get('If-None-Match')

        if if_match:
            etag_list = [tag.strip() for tag in if_match.split(',')]
            if etag not in etag_list and '*' not in etag_list:
                raise PreconditionFailed('etag does not match')
        elif if_none_match:
            etag_list = [tag.strip() for tag in if_none_match.split(',')]
            if etag in etag_list or '*' in etag_list:
                raise NotModified

        return rv

    return wrapped

def pagination_headers(total_pages, total_items, page_count, request):
    """
    Creates the `Link`. 'Count' and  'Total-Count' headers, to be used for pagination traversing

    :param total_pages: Total number of pages
    :param total_items: Total number of items in all the pages
    :param page_count: Item count for page (may differ from page size request)
    :param request: The flask request used, required to build other reoccurring vars like url and such.
    :return:
    """

    # Build constant variables from request data
    url = request.url_root + request.path.lstrip('/')
    per_page = request.args.get('per_page', 50)
    page = int(request.args.get('page', 1))

    # Build the base template
    LINKTEMPLATE = '<{}?per_page={}&'.format(url, per_page)

    # Removed page and per_page from query string
    query_string = re.sub(b'&?per_page=\d+&?', b'', request.query_string)
    query_string = re.sub(b'&?page=\d+&?', b'', query_string)

    # Add all original query params
    LINKTEMPLATE = LINKTEMPLATE.rstrip('&') + query_string.decode() + '&page={}>; rel="{}"'

    link_string = ''

    if page > 1:
        link_string += LINKTEMPLATE.format(page - 1, 'prev') + ', '
    if page < total_pages:
        link_string += LINKTEMPLATE.format(page + 1, 'next') + ', '
    link_string += LINKTEMPLATE.format(total_pages, 'last')

    return {
        'Link': link_string,
        'Total-Count': total_items,
        'Count': page_count
    }
<|endoftext|>


<commit_before>import json

from elasticsearch import Elasticsearch
from elasticsearch_dsl import Search
from elasticsearch.exceptions import NotFoundError
from elasticsearch.helpers import bulk, BulkIndexError

from query_models import SearchQuery, TermMatch, AggregatedResults, SimpleResults
from bulk_queue import BulkQueue

from utilities.logger import logger

from event import Event

class ElasticsearchBadServer(Exception):
    def __str__(self):
        return "Bad ES Server defined"

class ElasticsearchException(Exception):
    def __str__(self):
        return "Exception in ES encountered"

class ElasticsearchInvalidIndex(Exception):
    def __init__(self, index_name):
        self.index_name = index_name

    def __str__(self):
        return "Invalid index: " + str(self.index_name)

class ElasticsearchClient():

    def __init__(self, servers, bulk_amount=100, bulk_refresh_time=30):
        self.es_connection = Elasticsearch(servers)
        self.es_connection.ping()
        self.bulk_queue = BulkQueue(self, threshold=bulk_amount, flush_time=bulk_refresh_time)

    def delete_index(self, index_name, ignore_fail=False):
        ignore_codes = []
        if ignore_fail is True:
            ignore_codes = [400, 404]

        self.es_connection.indices.delete(index=index_name, ignore=ignore_codes)

    def get_indices(self):
        return self.es_connection.indices.stats()['indices'].keys()

    def index_exists(self, index_name):
        return self.es_connection.indices.exists(index_name)

    def create_index(self, index_name, index_config=None):
        if not index_config:
            index_config = '''
            {
              "mappings":{}
            }'''
        self.es_connection.indices.create(index=index_name, update_all_types='true', body=index_config)

    def create_alias(self, alias, index):
        actions = []
        if self.es_connection.indices.exists_alias('*', alias):
            actions.append({
                'remove': {'index': '*', 'alias': alias}
            })

        actions.append({
            'add': {'index': index, 'alias': alias}
        })
        self.es_connection.indices.update_aliases(dict(actions=actions))

    def create_alias_multiple_indices(self, alias_name, indices):
        actions = []
        if self.es_connection.indices.exists_alias('*', alias_name):
            actions.append({
                'remove': {'index': '*', 'alias': alias_name}
            })

        for index in indices:
            actions.append({'add': {'index': index, 'alias': alias_name}})
        self.es_connection.indices.update_aliases(dict(actions=actions))

    def get_alias(self, alias_name):
        return self.es_connection.indices.get_alias(index='*', name=alias_name).keys()

    def flush(self, index_name):
        self.es_connection.indices.flush(index=index_name)

    def search(self, search_query, indices, size, request_timeout):
        results = []
        try:
            results = Search(using=self.es_connection, index=indices).params(size=size, request_timeout=request_timeout).filter(search_query).execute()
        except NotFoundError:
            raise ElasticsearchInvalidIndex(indices)

        result_set = SimpleResults(results)
        return result_set

    def aggregated_search(self, search_query, indices, aggregations, size, request_timeout):
        search_obj = Search(using=self.es_connection, index=indices).params(size=size, request_timeout=request_timeout)
        query_obj = search_obj.filter(search_query)
        for aggregation in aggregations:
            query_obj.aggs.bucket(name=aggregation.to_dict()['terms']['field'], agg_type=aggregation)
        results = query_obj.execute()

        result_set = AggregatedResults(results)
        return result_set

    def save_documents(self, documents):
        try:
            bulk(self.es_connection, documents)
        except BulkIndexError as e:
            logger.error("Error bulk indexing: " + str(e))

    def start_bulk_timer(self):
        if not self.bulk_queue.started():
            self.bulk_queue.start_timer()

    def finish_bulk(self):
        self.bulk_queue.stop_timer()

    def __bulk_save_document(self, index, doc_type, body, doc_id=None):
        self.start_bulk_timer()
        self.bulk_queue.add(index=index, doc_type=doc_type, body=body, doc_id=doc_id)

    def __save_document(self, index, doc_type, body, doc_id=None, bulk=False):
        if bulk:
            self.__bulk_save_document(index=index, doc_type=doc_type, body=body, doc_id=doc_id)
        else:
            return self.es_connection.index(index=index, doc_type=doc_type, id=doc_id, body=body)

    def __parse_document(self, body, doc_type):
        if type(body) is str:
            body = json.loads(body)

        if '_type' in body:
            doc_type = body['_type']

        doc_body = body
        if '_source' in body:
            doc_body = body['_source']
        return doc_body, doc_type

    def save_object(self, body, index, doc_type, doc_id=None, bulk=False):
        doc_body, doc_type = self.__parse_document(body, doc_type)
        return self.__save_document(index=index, doc_type=doc_type, body=doc_body, doc_id=doc_id, bulk=bulk)

    def save_alert(self, body, index='alerts', doc_type='alert', doc_id=None, bulk=False):
        doc_body, doc_type = self.__parse_document(body, doc_type)
        return self.__save_document(index=index, doc_type=doc_type, body=doc_body, doc_id=doc_id, bulk=bulk)

    def save_event(self, body, index='events', doc_type='event', doc_id=None, bulk=False):
        doc_body, doc_type = self.__parse_document(body, doc_type)
        event = Event(doc_body)
        event.add_required_fields()
        return self.__save_document(index=index, doc_type=doc_type, body=event, doc_id=doc_id, bulk=bulk)

    def get_object_by_id(self, object_id, indices):
        id_match = TermMatch('_id', object_id)
        search_query = SearchQuery()
        search_query.add_must(id_match)
        results = search_query.execute(self, indices=indices)
        if len(results['hits']) == 0:
            return None
        else:
            return results['hits'][0]

    def get_alert_by_id(self, alert_id):
        return self.get_object_by_id(alert_id, ['alerts'])

    def get_event_by_id(self, event_id):
        return self.get_object_by_id(event_id, ['events'])

    def save_dashboard(self, dash_file, dash_name):
        f = open(dash_file)
        dashboardjson = json.load(f)
        f.close()
        title = dashboardjson['title']
        dashid = dash_name.replace(' ', '-')
        if dash_name:
            title = dash_name
        dashboarddata = {
            "user": "guest",
            "group": "guest",
            "title": title,
            "dashboard": json.dumps(dashboardjson)
        }

        return self.es_connection.index(index='.kibana', doc_type='dashboard', body=dashboarddata, id=dashid)

    def get_cluster_health(self):
        health_dict = self.es_connection.cluster.health()
        # To line up with the health stats from ES1, we're
        # removing certain keys
        health_dict.pop('active_shards_percent_as_number', None)
        health_dict.pop('delayed_unassigned_shards', None)
        health_dict.pop('number_of_in_flight_fetch', None)
        health_dict.pop('number_of_pending_tasks', None)
        health_dict.pop('task_max_waiting_in_queue_millis', None)

        return health_dict
<commit_msg>Adding open and close abstraction for managing indexes<commit_after>import json

from elasticsearch import Elasticsearch
from elasticsearch_dsl import Search
from elasticsearch.exceptions import NotFoundError
from elasticsearch.helpers import bulk, BulkIndexError

from query_models import SearchQuery, TermMatch, AggregatedResults, SimpleResults
from bulk_queue import BulkQueue

from utilities.logger import logger

from event import Event

class ElasticsearchBadServer(Exception):
    def __str__(self):
        return "Bad ES Server defined"

class ElasticsearchException(Exception):
    def __str__(self):
        return "Exception in ES encountered"

class ElasticsearchInvalidIndex(Exception):
    def __init__(self, index_name):
        self.index_name = index_name

    def __str__(self):
        return "Invalid index: " + str(self.index_name)

class ElasticsearchClient():

    def __init__(self, servers, bulk_amount=100, bulk_refresh_time=30):
        self.es_connection = Elasticsearch(servers)
        self.es_connection.ping()
        self.bulk_queue = BulkQueue(self, threshold=bulk_amount, flush_time=bulk_refresh_time)

    def close_index(self, index_name)
        self.es_connection.indices.close(index=index_name)

    def open_index(self, index_name)
        self.es_connection.indices.open(index=index_name)

    def delete_index(self, index_name, ignore_fail=False):
        ignore_codes = []
        if ignore_fail is True:
            ignore_codes = [400, 404]

        self.es_connection.indices.delete(index=index_name, ignore=ignore_codes)

    def get_indices(self):
        return self.es_connection.indices.stats()['indices'].keys()

    def index_exists(self, index_name):
        return self.es_connection.indices.exists(index_name)

    def create_index(self, index_name, index_config=None):
        if not index_config:
            index_config = '''
            {
              "mappings":{}
            }'''
        self.es_connection.indices.create(index=index_name, update_all_types='true', body=index_config)

    def create_alias(self, alias, index):
        actions = []
        if self.es_connection.indices.exists_alias('*', alias):
            actions.append({
                'remove': {'index': '*', 'alias': alias}
            })

        actions.append({
            'add': {'index': index, 'alias': alias}
        })
        self.es_connection.indices.update_aliases(dict(actions=actions))

    def create_alias_multiple_indices(self, alias_name, indices):
        actions = []
        if self.es_connection.indices.exists_alias('*', alias_name):
            actions.append({
                'remove': {'index': '*', 'alias': alias_name}
            })

        for index in indices:
            actions.append({'add': {'index': index, 'alias': alias_name}})
        self.es_connection.indices.update_aliases(dict(actions=actions))

    def get_alias(self, alias_name):
        return self.es_connection.indices.get_alias(index='*', name=alias_name).keys()

    def flush(self, index_name):
        self.es_connection.indices.flush(index=index_name)

    def search(self, search_query, indices, size, request_timeout):
        results = []
        try:
            results = Search(using=self.es_connection, index=indices).params(size=size, request_timeout=request_timeout).filter(search_query).execute()
        except NotFoundError:
            raise ElasticsearchInvalidIndex(indices)

        result_set = SimpleResults(results)
        return result_set

    def aggregated_search(self, search_query, indices, aggregations, size, request_timeout):
        search_obj = Search(using=self.es_connection, index=indices).params(size=size, request_timeout=request_timeout)
        query_obj = search_obj.filter(search_query)
        for aggregation in aggregations:
            query_obj.aggs.bucket(name=aggregation.to_dict()['terms']['field'], agg_type=aggregation)
        results = query_obj.execute()

        result_set = AggregatedResults(results)
        return result_set

    def save_documents(self, documents):
        try:
            bulk(self.es_connection, documents)
        except BulkIndexError as e:
            logger.error("Error bulk indexing: " + str(e))

    def start_bulk_timer(self):
        if not self.bulk_queue.started():
            self.bulk_queue.start_timer()

    def finish_bulk(self):
        self.bulk_queue.stop_timer()

    def __bulk_save_document(self, index, doc_type, body, doc_id=None):
        self.start_bulk_timer()
        self.bulk_queue.add(index=index, doc_type=doc_type, body=body, doc_id=doc_id)

    def __save_document(self, index, doc_type, body, doc_id=None, bulk=False):
        if bulk:
            self.__bulk_save_document(index=index, doc_type=doc_type, body=body, doc_id=doc_id)
        else:
            return self.es_connection.index(index=index, doc_type=doc_type, id=doc_id, body=body)

    def __parse_document(self, body, doc_type):
        if type(body) is str:
            body = json.loads(body)

        if '_type' in body:
            doc_type = body['_type']

        doc_body = body
        if '_source' in body:
            doc_body = body['_source']
        return doc_body, doc_type

    def save_object(self, body, index, doc_type, doc_id=None, bulk=False):
        doc_body, doc_type = self.__parse_document(body, doc_type)
        return self.__save_document(index=index, doc_type=doc_type, body=doc_body, doc_id=doc_id, bulk=bulk)

    def save_alert(self, body, index='alerts', doc_type='alert', doc_id=None, bulk=False):
        doc_body, doc_type = self.__parse_document(body, doc_type)
        return self.__save_document(index=index, doc_type=doc_type, body=doc_body, doc_id=doc_id, bulk=bulk)

    def save_event(self, body, index='events', doc_type='event', doc_id=None, bulk=False):
        doc_body, doc_type = self.__parse_document(body, doc_type)
        event = Event(doc_body)
        event.add_required_fields()
        return self.__save_document(index=index, doc_type=doc_type, body=event, doc_id=doc_id, bulk=bulk)

    def get_object_by_id(self, object_id, indices):
        id_match = TermMatch('_id', object_id)
        search_query = SearchQuery()
        search_query.add_must(id_match)
        results = search_query.execute(self, indices=indices)
        if len(results['hits']) == 0:
            return None
        else:
            return results['hits'][0]

    def get_alert_by_id(self, alert_id):
        return self.get_object_by_id(alert_id, ['alerts'])

    def get_event_by_id(self, event_id):
        return self.get_object_by_id(event_id, ['events'])

    def save_dashboard(self, dash_file, dash_name):
        f = open(dash_file)
        dashboardjson = json.load(f)
        f.close()
        title = dashboardjson['title']
        dashid = dash_name.replace(' ', '-')
        if dash_name:
            title = dash_name
        dashboarddata = {
            "user": "guest",
            "group": "guest",
            "title": title,
            "dashboard": json.dumps(dashboardjson)
        }

        return self.es_connection.index(index='.kibana', doc_type='dashboard', body=dashboarddata, id=dashid)

    def get_cluster_health(self):
        health_dict = self.es_connection.cluster.health()
        # To line up with the health stats from ES1, we're
        # removing certain keys
        health_dict.pop('active_shards_percent_as_number', None)
        health_dict.pop('delayed_unassigned_shards', None)
        health_dict.pop('number_of_in_flight_fetch', None)
        health_dict.pop('number_of_pending_tasks', None)
        health_dict.pop('task_max_waiting_in_queue_millis', None)

        return health_dict
<|endoftext|>


<commit_before>###################################################################
# Copyright 2013-2015 All Rights Reserved
# Authors: The Paradrop Team
###################################################################

from paradrop.lib.utils.output import out, logPrefix
import docker
from docker import Client
import json

def startChute(update):
    print update.__dict__
    out.warn('** %s TODO implement me\n' % logPrefix())

    repo = update.name + ":latest"
    dockerfile = update.dockerfile
    name = update.name

    host_config=docker.utils.create_host_config(
        port_bindings={ 80: 9000 },
        restart_policy={"MaximumRetryCount": 0, "Name": "always"}
    )

    c = Client(base_url="unix://var/run/docker.sock")

    for line in c.build(rm=True, tag=repo, fileobj=dockerfile):
        for key, value in json.loads(line).iteritems():
            if (key == "stream"):
                #out.info("-- %s %s\n" % (logPrefix(), str(value)))
                update.pkg.request.write(str(value))
            elif (isinstance(value, dict)):
                continue
            else:
                #out.info("-- %s %s\n" % (logPrefix(), str(value)))
                update.pkg.request.write(str(value))
    container = c.create_container(
        image=repo, name=name, host_config=host_config
    )

    out.info("-- %s %s\n" % (logPrefix(), str(container.get('Id'))))
    c.start(container.get('Id'))

def stopChute(update):
    out.warn('** %s TODO implement me\n' % logPrefix())
    c = Client(base_url='unix://var/run/docker.sock')
    repo = update.name + ":latest"
    name = update.name
    c.stop(container=name)
    c.remove_container(container=name)
    c.remove_image(image=repo)
<commit_msg>fixed import issue. Please use fully-qualified package names as default and only use "from" when it makes sense. Takes care of these issures."<commit_after>###################################################################
# Copyright 2013-2015 All Rights Reserved
# Authors: The Paradrop Team
###################################################################

from paradrop.lib.utils.output import out, logPrefix
import docker
import json

def startChute(update):
    print update.__dict__
    out.warn('** %s TODO implement me\n' % logPrefix())

    repo = update.name + ":latest"
    dockerfile = update.dockerfile
    name = update.name

    host_config = docker.utils.create_host_config(
        port_bindings={80: 9000},
        restart_policy={"MaximumRetryCount": 0, "Name": "always"}
    )

    c = docker.Client(base_url="unix://var/run/docker.sock")

    for line in c.build(rm=True, tag=repo, fileobj=dockerfile):
        for key, value in json.loads(line).iteritems():
            if (key == "stream"):
                #out.info("-- %s %s\n" % (logPrefix(), str(value)))
                update.pkg.request.write(str(value))
            elif (isinstance(value, dict)):
                continue
            else:
                #out.info("-- %s %s\n" % (logPrefix(), str(value)))
                update.pkg.request.write(str(value))
    container = c.create_container(
        image=repo, name=name, host_config=host_config
    )

    out.info("-- %s %s\n" % (logPrefix(), str(container.get('Id'))))
    c.start(container.get('Id'))

def stopChute(update):
    out.warn('** %s TODO implement me\n' % logPrefix())
    c = docker.Client(base_url='unix://var/run/docker.sock')
    repo = update.name + ":latest"
    name = update.name
    c.stop(container=name)
    c.remove_container(container=name)
    c.remove_image(image=repo)
<|endoftext|>


<commit_before>import datetime
import os
import pytest
import json
from owslib import crs

from birdy.client import converters
from birdy import WPSClient

# These tests assume Emu is running on the localhost
url = "http://localhost:5000/wps"

def data_path(*args):
    return os.path.join(os.path.dirname(__file__), "resources", *args)

@pytest.fixture(scope="module")
def wps():
    return WPSClient(url=url)

@pytest.mark.online
def test_wps_client_backward_compability():
    from birdy import BirdyClient
    BirdyClient(url=url)
    from birdy import import_wps
    import_wps(url=url)

@pytest.mark.online
def test_wps_client_single_output(wps):
    msg, = wps.hello("david")
    assert msg == "Hello david"
    ans, = wps.binaryoperatorfornumbers(inputa=1, inputb=2, operator="add")
    assert ans == 3.0

@pytest.mark.online
def test_wps_client_multiple_output(wps):
    # For multiple outputs, the output is a namedtuple
    out = wps.dummyprocess(10, 20)
    x1, x2 = out
    assert x1 == "11"
    assert x2 == "19"
    assert out.output1 == "11"
    assert out.output2 == "19"

@pytest.mark.online
def test_interactive(capsys):
    m = WPSClient(url=url, interactive=True)
    assert m.hello("david").output == "Hello david"
    captured = capsys.readouterr()
    assert captured.out.startswith(str(datetime.date.today()))
    assert m.binaryoperatorfornumbers()[0] == 5

@pytest.mark.skip(reason="Complex Output is not working.")
@pytest.mark.online
def test_wps_client_complex_output(wps):
    # As reference
    wps._convert_objects = False
    out_r, ref_r = wps.multiple_outputs(2)
    assert out_r.startswith("http")
    assert out_r.endswith(".txt")
    # TODO: fix ComplexDataInput
    assert ref_r.value.startswith("http")
    assert ref_r.value.endswith(".json")

    # As objects
    wps._convert_objects = True
    out_o, ref_o = wps.multiple_outputs(2)
    assert out_o == "my output file number 0"
    assert isinstance(ref_o, dict)
    wps._convert_objects = False

@pytest.mark.online
def test_process_subset_only_one():
    m = WPSClient(url=url, processes=["nap"])
    assert count_class_methods(m) == 1

    m = WPSClient(url=url, processes="nap")
    assert count_class_methods(m) == 1

@pytest.mark.online
def test_process_subset_names():
    with pytest.raises(ValueError, match="missing"):
        WPSClient(url=url, processes=["missing"])
    with pytest.raises(ValueError, match="wrong, process, names"):
        WPSClient(url=url, processes=["wrong", "process", "names"])

@pytest.mark.online
def test_inputs(wps):
    import netCDF4 as nc
    wps._convert_objects = True
    time_ = datetime.datetime.now().time()
    date_ = datetime.datetime.now().date()
    datetime_ = datetime.datetime.now()
    result = wps.inout(
        string="test string",
        int=3,
        float=3.5,
        boolean=True,
        time=time_.isoformat(),
        date=date_.isoformat(),
        datetime=datetime_.isoformat(sep=" "),
        string_choice="rock",
        string_multiple_choice="sitting duck",
        text="some text",
        dataset="file://" + data_path("dummy.nc"),
    )
    expected = (
        "test string",
        3,
        3.5,
        True,
        time_,
        date_,
        datetime_,
        "rock",
        "sitting duck",
        "some text",
    )
    assert expected == result[:-2]

    expected_netcdf = nc.Dataset(data_path("dummy.nc"))
    netcdf = result[-2]
    assert list(expected_netcdf.dimensions) == list(netcdf.dimensions)
    assert list(expected_netcdf.variables) == list(netcdf.variables)
    assert expected_netcdf.title == netcdf.title

    bbox = result[-1]
    assert bbox.crs == crs.Crs("epsg:4326")
    assert bbox.dimensions == 2
    wps._convert_objects = False

@pytest.mark.online
def test_netcdf(wps):
    import netCDF4 as nc

    wps._convert_objects = True
    if nc.getlibversion() > "4.5":
        m = WPSClient(url=url, processes=["output_formats"], convert_objects=True)
        ncdata, jsondata = m.output_formats()
        assert isinstance(ncdata, nc.Dataset)
        ncdata.close()
        assert isinstance(jsondata, dict)
    wps._convert_objects = False

def count_class_methods(class_):
    import types

    return len(
        [
            f
            for f in class_.__dict__.values()
            if isinstance(f, types.MethodType) and not f.__name__.startswith("_")
        ]
    )

def test_converter():
    j = converters.JSONConverter()
    assert isinstance(j, converters.default_converters["application/json"])

def test_jsonconverter():
    d = {"a": 1}
    s = json.dumps(d)

    j = converters.JSONConverter()
    assert j.convert_data(s) == d
<commit_msg>add angle to inout test<commit_after>import datetime
import os
import pytest
import json
from owslib import crs

from birdy.client import converters
from birdy import WPSClient

# These tests assume Emu is running on the localhost
url = "http://localhost:5000/wps"

def data_path(*args):
    return os.path.join(os.path.dirname(__file__), "resources", *args)

@pytest.fixture(scope="module")
def wps():
    return WPSClient(url=url)

@pytest.mark.online
def test_wps_client_backward_compability():
    from birdy import BirdyClient
    BirdyClient(url=url)
    from birdy import import_wps
    import_wps(url=url)

@pytest.mark.online
def test_wps_client_single_output(wps):
    msg, = wps.hello("david")
    assert msg == "Hello david"
    ans, = wps.binaryoperatorfornumbers(inputa=1, inputb=2, operator="add")
    assert ans == 3.0

@pytest.mark.online
def test_wps_client_multiple_output(wps):
    # For multiple outputs, the output is a namedtuple
    out = wps.dummyprocess(10, 20)
    x1, x2 = out
    assert x1 == "11"
    assert x2 == "19"
    assert out.output1 == "11"
    assert out.output2 == "19"

@pytest.mark.online
def test_interactive(capsys):
    m = WPSClient(url=url, interactive=True)
    assert m.hello("david").output == "Hello david"
    captured = capsys.readouterr()
    assert captured.out.startswith(str(datetime.date.today()))
    assert m.binaryoperatorfornumbers()[0] == 5

@pytest.mark.skip(reason="Complex Output is not working.")
@pytest.mark.online
def test_wps_client_complex_output(wps):
    # As reference
    wps._convert_objects = False
    out_r, ref_r = wps.multiple_outputs(2)
    assert out_r.startswith("http")
    assert out_r.endswith(".txt")
    # TODO: fix ComplexDataInput
    assert ref_r.value.startswith("http")
    assert ref_r.value.endswith(".json")

    # As objects
    wps._convert_objects = True
    out_o, ref_o = wps.multiple_outputs(2)
    assert out_o == "my output file number 0"
    assert isinstance(ref_o, dict)
    wps._convert_objects = False

@pytest.mark.online
def test_process_subset_only_one():
    m = WPSClient(url=url, processes=["nap"])
    assert count_class_methods(m) == 1

    m = WPSClient(url=url, processes="nap")
    assert count_class_methods(m) == 1

@pytest.mark.online
def test_process_subset_names():
    with pytest.raises(ValueError, match="missing"):
        WPSClient(url=url, processes=["missing"])
    with pytest.raises(ValueError, match="wrong, process, names"):
        WPSClient(url=url, processes=["wrong", "process", "names"])

@pytest.mark.online
def test_inputs(wps):
    import netCDF4 as nc
    wps._convert_objects = True
    time_ = datetime.datetime.now().time()
    date_ = datetime.datetime.now().date()
    datetime_ = datetime.datetime.now()
    result = wps.inout(
        string="test string",
        int=3,
        float=3.5,
        boolean=True,
        angle=67.,
        time=time_.isoformat(),
        date=date_.isoformat(),
        datetime=datetime_.isoformat(sep=" "),
        string_choice="rock",
        string_multiple_choice="sitting duck",
        text="some text",
        dataset="file://" + data_path("dummy.nc"),
    )
    expected = (
        "test string",
        3,
        3.5,
        True,
        67.,
        time_,
        date_,
        datetime_,
        "rock",
        "sitting duck",
        "some text",
    )
    assert expected == result[:-2]

    expected_netcdf = nc.Dataset(data_path("dummy.nc"))
    netcdf = result[-2]
    assert list(expected_netcdf.dimensions) == list(netcdf.dimensions)
    assert list(expected_netcdf.variables) == list(netcdf.variables)
    assert expected_netcdf.title == netcdf.title

    bbox = result[-1]
    assert bbox.crs == crs.Crs("epsg:4326")
    assert bbox.dimensions == 2
    wps._convert_objects = False

@pytest.mark.online
def test_netcdf(wps):
    import netCDF4 as nc

    wps._convert_objects = True
    if nc.getlibversion() > "4.5":
        m = WPSClient(url=url, processes=["output_formats"], convert_objects=True)
        ncdata, jsondata = m.output_formats()
        assert isinstance(ncdata, nc.Dataset)
        ncdata.close()
        assert isinstance(jsondata, dict)
    wps._convert_objects = False

def count_class_methods(class_):
    import types

    return len(
        [
            f
            for f in class_.__dict__.values()
            if isinstance(f, types.MethodType) and not f.__name__.startswith("_")
        ]
    )

def test_converter():
    j = converters.JSONConverter()
    assert isinstance(j, converters.default_converters["application/json"])

def test_jsonconverter():
    d = {"a": 1}
    s = json.dumps(d)

    j = converters.JSONConverter()
    assert j.convert_data(s) == d
<|endoftext|>


<commit_before>import logging
import os
import re

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError

from ngw.core.models import FIELD_LOGIN, ContactFieldValue

def ip_int_to_str(val):
    """
    That function takes a 0..2**32 interger and converts it into a string IP
    address.
    For example: 16909060 aka (1<<24)+(2<<16)+(3<<8)+4 will return 1.2.3.4
    """

    if not isinstance(val, int):
        raise TypeError('ip_int_to_str expects a number')
    if val < 0 or val >= (1 << 32):
        raise ValueError('Out of range')

    return '{}.{}.{}.{}'.format(
        (val >> 24) & 255,
        (val >> 16) & 255,
        (val >> 8) & 255,
        val & 255)

def ip_str_to_int(val):
    """
    That function takes a string with a IP address and converts it to a integer
    For example: 1.2.3.4 returns 16909060 aka (1<<24)+(2<<16)+(3<<8)+4
    """

    split = re.fullmatch(r'(\d+)\.(\d+)\.(\d+)\.(\d+)', val)
    if split is None:
        raise ValueError('Not an IP address')
    result = 0
    for idx in range(1, 5):
        frag = int(split.group(idx))
        if frag < 0 or frag > 255:
            raise ValueError('Out of range')
        result = (result << 8) + frag
    return result

class Command(BaseCommand):
    help = 'Plugin for openvpn --client-connect'

    def handle(self, *args, **options):
        logger = logging.getLogger('vpn')
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter(
            '{asctime} {name} {levelname!s:8} {message}', style='{'))
        logger.addHandler(handler)
        verbosity = int(options['verbosity'])
        if verbosity == 0:
            logger.setLevel(logging.ERROR)
        elif verbosity == 1:
            logger.setLevel(logging.WARNING)
        elif verbosity == 2:
            logger.setLevel(logging.INFO)
        elif verbosity == 3:
            logger.setLevel(logging.DEBUG)

        logger.debug('vpnconfig.handle starts.')

        # logger.debug('vpnconfig args: {}'.format(args))
        # logger.debug('vpnconfig environment: {}'.format(os.environ))

        if len(args) != 1:
            msg = ("This script takes exactly one argument. Are you really "
                   "using openvpn --client-connect?")
            logger.critical(msg)
            raise CommandError(msg)

        filename = args[0]
        logger.debug('Appending configuration to file {}'.format(filename))
        try:
            f = open(filename, 'a')
        except FileNotFoundError:
            msg = "Error opening temporary file {}".format(filename)
            logger.critical(msg)
            raise CommandError(msg)

        try:
            login = os.environ['common_name']
        except KeyError:
            msg = "Environment doesn't contains common_name."
            logger.critical(msg)
            raise CommandError(msg)

        logger.debug('Setting up network for {}'.format(login))

        try:
            login_value = ContactFieldValue.objects.get(
                contact_field_id=FIELD_LOGIN, value=login)
        except ContactFieldValue.DoesNotExist:
            msg = 'Contact with login "{}" not found.'.format(login)
            logger.info(msg)
            raise CommandError(msg)

        cid = login_value.contact_id
        # contact = login_value.contact

        try:
            baseip = settings.VPN_BASEIP
        except AttributeError:
            msg = 'Settings must define VPN_BASEIP'
            logger.critical(msg)
            raise CommandError(msg)
        baseip = ip_str_to_int(baseip)

        try:
            maxip = settings.VPN_MAXIP
        except AttributeError:
            msg = 'Settings must define VPN_MAXIP'
            logger.critical(msg)
            raise CommandError(msg)
        maxip = ip_str_to_int(maxip)

        addr = baseip+cid*4+1
        gate_addr = baseip+cid*4+2

        if addr > maxip:
            msg = 'IP range is too small to have client #{}'.format(cid)
            logger.critical(msg)
            raise CommandError(msg)

        output = 'ifconfig-push {} {}'.format(addr, gate_addr)
        logger.debug('{} ({}): {}'.format(login, cid, output))
        f.write(output+'\n')

        f.close()
<commit_msg>vpn: Use string format to push routes<commit_after>import logging
import os
import re

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError

from ngw.core.models import FIELD_LOGIN, ContactFieldValue

def ip_int_to_str(val):
    """
    That function takes a 0..2**32 interger and converts it into a string IP
    address.
    For example: 16909060 aka (1<<24)+(2<<16)+(3<<8)+4 will return 1.2.3.4
    """

    if not isinstance(val, int):
        raise TypeError('ip_int_to_str expects a number')
    if val < 0 or val >= (1 << 32):
        raise ValueError('Out of range')

    return '{}.{}.{}.{}'.format(
        (val >> 24) & 255,
        (val >> 16) & 255,
        (val >> 8) & 255,
        val & 255)

def ip_str_to_int(val):
    """
    That function takes a string with a IP address and converts it to a integer
    For example: 1.2.3.4 returns 16909060 aka (1<<24)+(2<<16)+(3<<8)+4
    """

    split = re.fullmatch(r'(\d+)\.(\d+)\.(\d+)\.(\d+)', val)
    if split is None:
        raise ValueError('Not an IP address')
    result = 0
    for idx in range(1, 5):
        frag = int(split.group(idx))
        if frag < 0 or frag > 255:
            raise ValueError('Out of range')
        result = (result << 8) + frag
    return result

class Command(BaseCommand):
    help = 'Plugin for openvpn --client-connect'

    def handle(self, *args, **options):
        logger = logging.getLogger('vpn')
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter(
            '{asctime} {name} {levelname!s:8} {message}', style='{'))
        logger.addHandler(handler)
        verbosity = int(options['verbosity'])
        if verbosity == 0:
            logger.setLevel(logging.ERROR)
        elif verbosity == 1:
            logger.setLevel(logging.WARNING)
        elif verbosity == 2:
            logger.setLevel(logging.INFO)
        elif verbosity == 3:
            logger.setLevel(logging.DEBUG)

        logger.debug('vpnconfig.handle starts.')

        # logger.debug('vpnconfig args: {}'.format(args))
        # logger.debug('vpnconfig environment: {}'.format(os.environ))

        if len(args) != 1:
            msg = ("This script takes exactly one argument. Are you really "
                   "using openvpn --client-connect?")
            logger.critical(msg)
            raise CommandError(msg)

        filename = args[0]
        logger.debug('Appending configuration to file {}'.format(filename))
        try:
            f = open(filename, 'a')
        except FileNotFoundError:
            msg = "Error opening temporary file {}".format(filename)
            logger.critical(msg)
            raise CommandError(msg)

        try:
            login = os.environ['common_name']
        except KeyError:
            msg = "Environment doesn't contains common_name."
            logger.critical(msg)
            raise CommandError(msg)

        logger.debug('Setting up network for {}'.format(login))

        try:
            login_value = ContactFieldValue.objects.get(
                contact_field_id=FIELD_LOGIN, value=login)
        except ContactFieldValue.DoesNotExist:
            msg = 'Contact with login "{}" not found.'.format(login)
            logger.info(msg)
            raise CommandError(msg)

        cid = login_value.contact_id
        # contact = login_value.contact

        try:
            baseip = settings.VPN_BASEIP
        except AttributeError:
            msg = 'Settings must define VPN_BASEIP'
            logger.critical(msg)
            raise CommandError(msg)
        baseip = ip_str_to_int(baseip)

        try:
            maxip = settings.VPN_MAXIP
        except AttributeError:
            msg = 'Settings must define VPN_MAXIP'
            logger.critical(msg)
            raise CommandError(msg)
        maxip = ip_str_to_int(maxip)

        addr = baseip+cid*4+1
        gate_addr = baseip+cid*4+2

        if addr > maxip:
            msg = 'IP range is too small to have client #{}'.format(cid)
            logger.critical(msg)
            raise CommandError(msg)

        addr = ip_int_to_str(addr)
        gate_addr = ip_int_to_str(gate_addr)

        output = 'ifconfig-push {} {}'.format(addr, gate_addr)
        logger.debug('{} ({}): {}'.format(login, cid, output))
        f.write(output+'\n')

        f.close()
<|endoftext|>