<commit_before># -*- coding: utf-8 -*-
# Copyright (C) 2014-present Taiga Agile LLC
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import uuid

from django.utils.translation import ugettext as _
from django.conf import settings

from taiga.base.api import viewsets
from taiga.base import response
from taiga.base import exceptions as exc
from taiga.base.decorators import list_route
from taiga.users.models import AuthData, User
from taiga.users.services import get_user_photo_url
from taiga.users.gravatar import get_user_gravatar_id

from taiga.importers import permissions
from taiga.importers import exceptions
from taiga.importers.services import resolve_users_bindings
from .normal import JiraNormalImporter
from .agile import JiraAgileImporter
from . import tasks

class JiraImporterViewSet(viewsets.ViewSet):
    permission_classes = (permissions.ImporterPermission,)

    def _get_token(self, request):
        token_data = request.DATA.get('token', "").split(".")

        token = {
            "access_token": token_data[0],
            "access_token_secret": token_data[1],
            "key_cert": settings.IMPORTERS.get('jira', {}).get('cert', None),
            "consumer_key": settings.IMPORTERS.get('jira', {}).get('consumer_key', None)
        }
        return token

    @list_route(methods=["POST"])
    def list_users(self, request, *args, **kwargs):
        self.check_permissions(request, "list_users", None)

        url = request.DATA.get('url', None)
        token = self._get_token(request)
        project_id = request.DATA.get('project', None)

        if not project_id:
            raise exc.WrongArguments(_("The project param is needed"))
        if not url:
            raise exc.WrongArguments(_("The url param is needed"))

        importer = JiraNormalImporter(request.user, url, token)
        users = importer.list_users()
        for user in users:
            user['user'] = None
            if not user['email']:
                continue

            try:
                taiga_user = User.objects.get(email=user['email'])
            except User.DoesNotExist:
                continue

            user['user'] = {
                'id': taiga_user.id,
                'full_name': taiga_user.get_full_name(),
                'gravatar_id': get_user_gravatar_id(taiga_user),
                'photo': get_user_photo_url(taiga_user),
            }
        return response.Ok(users)

    @list_route(methods=["POST"])
    def list_projects(self, request, *args, **kwargs):
        self.check_permissions(request, "list_projects", None)
        url = request.DATA.get('url', None)
        if not url:
            raise exc.WrongArguments(_("The url param is needed"))

        token = self._get_token(request)
        importer = JiraNormalImporter(request.user, url, token)
        agile_importer = JiraAgileImporter(request.user, url, token)
        projects = importer.list_projects()
        boards = agile_importer.list_projects()
        return response.Ok(sorted(projects + boards, key=lambda x: x['name']))

    @list_route(methods=["POST"])
    def import_project(self, request, *args, **kwargs):
        self.check_permissions(request, "import_project", None)

        url = request.DATA.get('url', None)
        token = self._get_token(request)
        project_id = request.DATA.get('project', None)
        if not project_id:
            raise exc.WrongArguments(_("The project param is needed"))

        if not url:
            raise exc.WrongArguments(_("The url param is needed"))

        options = {
            "name": request.DATA.get('name', None),
            "description": request.DATA.get('description', None),
            "users_bindings": resolve_users_bindings(request.DATA.get("users_bindings", {})),
            "keep_external_reference": request.DATA.get("keep_external_reference", False),
            "is_private": request.DATA.get("is_private", False),
        }

        importer_type = request.DATA.get('importer_type', "normal")
        if importer_type == "agile":
            importer = JiraAgileImporter(request.user, url, token)
        else:
            project_type = request.DATA.get("project_type", "scrum")
            if project_type == "kanban":
                options['template'] = "kanban"
            else:
                options['template'] = "scrum"

            importer = JiraNormalImporter(request.user, url, token)

            types_bindings = {
                "epic": [],
                "us": [],
                "task": [],
                "issue": [],
            }
            for issue_type in importer.list_issue_types(project_id):
                if project_type in ['scrum', 'kanban']:
                    # Set the type bindings
                    if issue_type['subtask']:
                        types_bindings['task'].append(issue_type)
                    elif issue_type['name'].upper() == "EPIC":
                        types_bindings["epic"].append(issue_type)
                    elif issue_type['name'].upper() in ["US", "USERSTORY", "USER STORY"]:
                        types_bindings["us"].append(issue_type)
                    elif issue_type['name'].upper() in ["ISSUE", "BUG", "ENHANCEMENT"]:
                        types_bindings["issue"].append(issue_type)
                    else:
                        types_bindings["us"].append(issue_type)
                elif project_type == "issues":
                    # Set the type bindings
                    if issue_type['subtask']:
                        continue
                    types_bindings["issue"].append(issue_type)
                elif project_type == "issues-with-subissues":
                    types_bindings["issue"].append(issue_type)
                else:
                    raise exc.WrongArguments(_("Invalid project_type {}").format(project_type))

            options["types_bindings"] = types_bindings

        if settings.CELERY_ENABLED:
            task = tasks.import_project.delay(request.user.id, url, token, project_id, options, importer_type)
            return response.Accepted({"task_id": task.id})

        project = importer.import_project(project_id, options)
        project_data = {
            "slug": project.slug,
            "my_permissions": ["view_us"],
            "is_backlog_activated": project.is_backlog_activated,
            "is_kanban_activated": project.is_kanban_activated,
        }

        return response.Ok(project_data)

    @list_route(methods=["GET"])
    def auth_url(self, request, *args, **kwargs):
        self.check_permissions(request, "auth_url", None)
        jira_url = request.QUERY_PARAMS.get('url', None)

        if not jira_url:
            raise exc.WrongArguments(_("The url param is needed"))

        try:
            (oauth_token, oauth_secret, url) = JiraNormalImporter.get_auth_url(
                jira_url,
                settings.IMPORTERS.get('jira', {}).get('consumer_key', None),
                settings.IMPORTERS.get('jira', {}).get('cert', None),
                True
            )
        except exceptions.InvalidServiceConfiguration:
            raise exc.BadRequest(_("Invalid Jira server configuration."))

        (auth_data, created) = AuthData.objects.get_or_create(
            user=request.user,
            key="jira-oauth",
            defaults={
                "value": uuid.uuid4().hex,
                "extra": {},
            }
        )
        auth_data.extra = {
            "oauth_token": oauth_token,
            "oauth_secret": oauth_secret,
            "url": jira_url,
        }
        auth_data.save()

        return response.Ok({"url": url})

    @list_route(methods=["POST"])
    def authorize(self, request, *args, **kwargs):
        self.check_permissions(request, "authorize", None)

        try:
            oauth_data = request.user.auth_data.get(key="jira-oauth")
            oauth_verifier = request.DATA.get("oauth_verifier", None)
            oauth_token = oauth_data.extra['oauth_token']
            oauth_secret = oauth_data.extra['oauth_secret']
            server_url = oauth_data.extra['url']
            oauth_data.delete()

            jira_token = JiraNormalImporter.get_access_token(
                server_url,
                settings.IMPORTERS.get('jira', {}).get('consumer_key', None),
                settings.IMPORTERS.get('jira', {}).get('cert', None),
                oauth_token,
                oauth_secret,
                oauth_verifier,
                False
            )
        except Exception as e:
            raise exc.WrongArguments(_("Invalid or expired auth token"))

        return response.Ok({
            "token": jira_token['access_token'] + "." + jira_token['access_token_secret'],
            "url": server_url
        })
<commit_msg>(fix) capture exception when Jira fails due to unsupported version<commit_after># -*- coding: utf-8 -*-
# Copyright (C) 2014-present Taiga Agile LLC
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import uuid

from django.utils.translation import ugettext as _
from django.conf import settings

from taiga.base.api import viewsets
from taiga.base import response
from taiga.base import exceptions as exc
from taiga.base.decorators import list_route
from taiga.users.models import AuthData, User
from taiga.users.services import get_user_photo_url
from taiga.users.gravatar import get_user_gravatar_id

from taiga.importers import permissions
from taiga.importers import exceptions
from taiga.importers.services import resolve_users_bindings
from .normal import JiraNormalImporter
from .agile import JiraAgileImporter
from . import tasks

class JiraImporterViewSet(viewsets.ViewSet):
    permission_classes = (permissions.ImporterPermission,)

    def _get_token(self, request):
        token_data = request.DATA.get('token', "").split(".")

        token = {
            "access_token": token_data[0],
            "access_token_secret": token_data[1],
            "key_cert": settings.IMPORTERS.get('jira', {}).get('cert', None),
            "consumer_key": settings.IMPORTERS.get('jira', {}).get('consumer_key', None)
        }
        return token

    @list_route(methods=["POST"])
    def list_users(self, request, *args, **kwargs):
        self.check_permissions(request, "list_users", None)

        url = request.DATA.get('url', None)
        token = self._get_token(request)
        project_id = request.DATA.get('project', None)

        if not project_id:
            raise exc.WrongArguments(_("The project param is needed"))
        if not url:
            raise exc.WrongArguments(_("The url param is needed"))

        importer = JiraNormalImporter(request.user, url, token)
        try:
            users = importer.list_users()
        except Exception as e:
            # common error due to modern Jira versions which are unsupported by Taiga
            raise exc.BadRequest(_("""
                There was an error; probably due to an unsupported Jira version.
                Taiga does not support Jira releases from 8.6."""
            ))

        for user in users:
            user['user'] = None
            if not user['email']:
                continue

            try:
                taiga_user = User.objects.get(email=user['email'])
            except User.DoesNotExist:
                continue

            user['user'] = {
                'id': taiga_user.id,
                'full_name': taiga_user.get_full_name(),
                'gravatar_id': get_user_gravatar_id(taiga_user),
                'photo': get_user_photo_url(taiga_user),
            }
        return response.Ok(users)

    @list_route(methods=["POST"])
    def list_projects(self, request, *args, **kwargs):
        self.check_permissions(request, "list_projects", None)
        url = request.DATA.get('url', None)
        if not url:
            raise exc.WrongArguments(_("The url param is needed"))

        token = self._get_token(request)
        importer = JiraNormalImporter(request.user, url, token)
        agile_importer = JiraAgileImporter(request.user, url, token)
        projects = importer.list_projects()
        boards = agile_importer.list_projects()
        return response.Ok(sorted(projects + boards, key=lambda x: x['name']))

    @list_route(methods=["POST"])
    def import_project(self, request, *args, **kwargs):
        self.check_permissions(request, "import_project", None)

        url = request.DATA.get('url', None)
        token = self._get_token(request)
        project_id = request.DATA.get('project', None)
        if not project_id:
            raise exc.WrongArguments(_("The project param is needed"))

        if not url:
            raise exc.WrongArguments(_("The url param is needed"))

        options = {
            "name": request.DATA.get('name', None),
            "description": request.DATA.get('description', None),
            "users_bindings": resolve_users_bindings(request.DATA.get("users_bindings", {})),
            "keep_external_reference": request.DATA.get("keep_external_reference", False),
            "is_private": request.DATA.get("is_private", False),
        }

        importer_type = request.DATA.get('importer_type', "normal")
        if importer_type == "agile":
            importer = JiraAgileImporter(request.user, url, token)
        else:
            project_type = request.DATA.get("project_type", "scrum")
            if project_type == "kanban":
                options['template'] = "kanban"
            else:
                options['template'] = "scrum"

            importer = JiraNormalImporter(request.user, url, token)

            types_bindings = {
                "epic": [],
                "us": [],
                "task": [],
                "issue": [],
            }
            for issue_type in importer.list_issue_types(project_id):
                if project_type in ['scrum', 'kanban']:
                    # Set the type bindings
                    if issue_type['subtask']:
                        types_bindings['task'].append(issue_type)
                    elif issue_type['name'].upper() == "EPIC":
                        types_bindings["epic"].append(issue_type)
                    elif issue_type['name'].upper() in ["US", "USERSTORY", "USER STORY"]:
                        types_bindings["us"].append(issue_type)
                    elif issue_type['name'].upper() in ["ISSUE", "BUG", "ENHANCEMENT"]:
                        types_bindings["issue"].append(issue_type)
                    else:
                        types_bindings["us"].append(issue_type)
                elif project_type == "issues":
                    # Set the type bindings
                    if issue_type['subtask']:
                        continue
                    types_bindings["issue"].append(issue_type)
                elif project_type == "issues-with-subissues":
                    types_bindings["issue"].append(issue_type)
                else:
                    raise exc.WrongArguments(_("Invalid project_type {}").format(project_type))

            options["types_bindings"] = types_bindings

        if settings.CELERY_ENABLED:
            task = tasks.import_project.delay(request.user.id, url, token, project_id, options, importer_type)
            return response.Accepted({"task_id": task.id})

        project = importer.import_project(project_id, options)
        project_data = {
            "slug": project.slug,
            "my_permissions": ["view_us"],
            "is_backlog_activated": project.is_backlog_activated,
            "is_kanban_activated": project.is_kanban_activated,
        }

        return response.Ok(project_data)

    @list_route(methods=["GET"])
    def auth_url(self, request, *args, **kwargs):
        self.check_permissions(request, "auth_url", None)
        jira_url = request.QUERY_PARAMS.get('url', None)

        if not jira_url:
            raise exc.WrongArguments(_("The url param is needed"))

        try:
            (oauth_token, oauth_secret, url) = JiraNormalImporter.get_auth_url(
                jira_url,
                settings.IMPORTERS.get('jira', {}).get('consumer_key', None),
                settings.IMPORTERS.get('jira', {}).get('cert', None),
                True
            )
        except exceptions.InvalidServiceConfiguration:
            raise exc.BadRequest(_("Invalid Jira server configuration."))

        (auth_data, created) = AuthData.objects.get_or_create(
            user=request.user,
            key="jira-oauth",
            defaults={
                "value": uuid.uuid4().hex,
                "extra": {},
            }
        )
        auth_data.extra = {
            "oauth_token": oauth_token,
            "oauth_secret": oauth_secret,
            "url": jira_url,
        }
        auth_data.save()

        return response.Ok({"url": url})

    @list_route(methods=["POST"])
    def authorize(self, request, *args, **kwargs):
        self.check_permissions(request, "authorize", None)

        try:
            oauth_data = request.user.auth_data.get(key="jira-oauth")
            oauth_verifier = request.DATA.get("oauth_verifier", None)
            oauth_token = oauth_data.extra['oauth_token']
            oauth_secret = oauth_data.extra['oauth_secret']
            server_url = oauth_data.extra['url']
            oauth_data.delete()

            jira_token = JiraNormalImporter.get_access_token(
                server_url,
                settings.IMPORTERS.get('jira', {}).get('consumer_key', None),
                settings.IMPORTERS.get('jira', {}).get('cert', None),
                oauth_token,
                oauth_secret,
                oauth_verifier,
                False
            )
        except Exception as e:
            raise exc.WrongArguments(_("Invalid or expired auth token"))

        return response.Ok({
            "token": jira_token['access_token'] + "." + jira_token['access_token_secret'],
            "url": server_url
        })
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
##############################################################################
#
#    Base Location Geonames Import module for OpenERP
#    Copyright (C) 2014 Akretion (http://www.akretion.com)
#    @author Alexis de Lattre <alexis.delattre@akretion.com>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from openerp.osv import orm, fields
from openerp.tools.translate import _
import requests
import tempfile
import StringIO
import unicodecsv
import zipfile
import os
import logging

logger = logging.getLogger(__name__)

class better_zip_geonames_import(orm.TransientModel):
    _name = 'better.zip.geonames.import'
    _description = 'Import Better Zip from Geonames'

    _columns = {
        'country_id': fields.many2one('res.country', 'Country', required=True),
    }

    def _prepare_better_zip(
            self, cr, uid, row, country_id, states, context=None):
        '''This function is designed to be inherited'''
        state_id = False
        if states and row[4] and row[4] in states:
            state_id = states[row[4].upper()]
        if row[0] == 'FR' and 'CEDEX' in row[1]:
            return False
        vals = {
            'name': row[1],
            'city': row[2],
            'state_id': state_id,
            'country_id': country_id,
            }
        return vals

    def create_better_zip(
            self, cr, uid, row, country_id, country_code, states,
            context=None):
        bzip_id = False
        if row[0] != country_code:
            raise orm.except_orm(
                _('Error:'),
                _("The country code inside the file (%s) doesn't "
                    "correspond to the selected country (%s).")
                % (row[0], country_code))
        logger.debug('ZIP = %s - City = %s' % (row[1], row[2]))
        if row[1] and row[2]:
            vals = self._prepare_better_zip(
                cr, uid, row, country_id, states, context=context)
            if vals:
                bzip_id = self.pool['res.better.zip'].create(
                    cr, uid, vals, context=context)
        return bzip_id

    def run_import(self, cr, uid, ids, context=None):
        assert len(ids) == 1, 'Only one ID for the better zip import wizard'
        bzip_obj = self.pool['res.better.zip']
        wizard = self.browse(cr, uid, ids[0], context=context)
        country_id = wizard.country_id.id
        country_code = wizard.country_id.code.upper()
        url = 'http://download.geonames.org/export/zip/%s.zip' % country_code
        logger.info('Starting to download %s' % url)
        res_request = requests.get(url)
        if res_request.status_code != requests.codes.ok:
            raise orm.except_orm(
                _('Error:'),
                _('Got an error %d when trying to download the file %s.')
                % (res_request.status_code, url))
        bzip_ids_to_delete = bzip_obj.search(
            cr, uid, [('country_id', '=', country_id)], context=context)
        if bzip_ids_to_delete:
            cr.execute('SELECT id FROM res_better_zip WHERE id in %s '
                'FOR UPDATE NOWAIT', (tuple(bzip_ids_to_delete), ))
            bzip_obj.unlink(cr, uid, bzip_ids_to_delete, context=context)
            logger.info(
                '%d better zip entries deleted for country %s'
                % (len(bzip_ids_to_delete), wizard.country_id.name))
        state_ids = self.pool['res.country.state'].search(
            cr, uid, [('country_id', '=', country_id)], context=context)
        states = {}
        # key = code of the state ; value = ID of the state in OpenERP
        if state_ids:
            states_r = self.pool['res.country.state'].read(
                cr, uid, state_ids, ['code', 'country_id'], context=context)
            for state in states_r:
                states[state['code'].upper()] = state['id']
        f_geonames = zipfile.ZipFile(StringIO.StringIO(res_request.content))
        tempdir = tempfile.mkdtemp(prefix='openerp')
        f_geonames.extract('%s.txt' % country_code, tempdir)
        logger.info('The geonames zipfile has been decompressed')
        data_file = open(os.path.join(tempdir, '%s.txt' % country_code), 'r')
        data_file.seek(0)
        logger.info(
            'Starting to create the better zip entries %s state information'
            % (states and 'with' or 'without'))
        for row in unicodecsv.reader(
                data_file, encoding='utf-8', delimiter='	'):
            self.create_better_zip(
                cr, uid, row, country_id, country_code, states,
                context=context)
        data_file.close()
        logger.info(
            'The wizard to create better zip entries from geonames '
            'has been successfully completed.')
        return True
<commit_msg>Remove France-specific code ; moved to new module l10n_fr_base_location_geonames_import on lp:openerp-french-localization<commit_after># -*- encoding: utf-8 -*-
##############################################################################
#
#    Base Location Geonames Import module for OpenERP
#    Copyright (C) 2014 Akretion (http://www.akretion.com)
#    @author Alexis de Lattre <alexis.delattre@akretion.com>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from openerp.osv import orm, fields
from openerp.tools.translate import _
import requests
import tempfile
import StringIO
import unicodecsv
import zipfile
import os
import logging

logger = logging.getLogger(__name__)

class better_zip_geonames_import(orm.TransientModel):
    _name = 'better.zip.geonames.import'
    _description = 'Import Better Zip from Geonames'

    _columns = {
        'country_id': fields.many2one('res.country', 'Country', required=True),
    }

    def _prepare_better_zip(
            self, cr, uid, row, country_id, states, context=None):
        '''This function is designed to be inherited'''
        state_id = False
        if states and row[4] and row[4] in states:
            state_id = states[row[4].upper()]
        vals = {
            'name': row[1],
            'city': row[2],
            'state_id': state_id,
            'country_id': country_id,
            }
        return vals

    def create_better_zip(
            self, cr, uid, row, country_id, country_code, states,
            context=None):
        bzip_id = False
        if row[0] != country_code:
            raise orm.except_orm(
                _('Error:'),
                _("The country code inside the file (%s) doesn't "
                    "correspond to the selected country (%s).")
                % (row[0], country_code))
        logger.debug('ZIP = %s - City = %s' % (row[1], row[2]))
        if row[1] and row[2]:
            vals = self._prepare_better_zip(
                cr, uid, row, country_id, states, context=context)
            if vals:
                bzip_id = self.pool['res.better.zip'].create(
                    cr, uid, vals, context=context)
        return bzip_id

    def run_import(self, cr, uid, ids, context=None):
        assert len(ids) == 1, 'Only one ID for the better zip import wizard'
        bzip_obj = self.pool['res.better.zip']
        wizard = self.browse(cr, uid, ids[0], context=context)
        country_id = wizard.country_id.id
        country_code = wizard.country_id.code.upper()
        url = 'http://download.geonames.org/export/zip/%s.zip' % country_code
        logger.info('Starting to download %s' % url)
        res_request = requests.get(url)
        if res_request.status_code != requests.codes.ok:
            raise orm.except_orm(
                _('Error:'),
                _('Got an error %d when trying to download the file %s.')
                % (res_request.status_code, url))
        bzip_ids_to_delete = bzip_obj.search(
            cr, uid, [('country_id', '=', country_id)], context=context)
        if bzip_ids_to_delete:
            cr.execute('SELECT id FROM res_better_zip WHERE id in %s '
                'FOR UPDATE NOWAIT', (tuple(bzip_ids_to_delete), ))
            bzip_obj.unlink(cr, uid, bzip_ids_to_delete, context=context)
            logger.info(
                '%d better zip entries deleted for country %s'
                % (len(bzip_ids_to_delete), wizard.country_id.name))
        state_ids = self.pool['res.country.state'].search(
            cr, uid, [('country_id', '=', country_id)], context=context)
        states = {}
        # key = code of the state ; value = ID of the state in OpenERP
        if state_ids:
            states_r = self.pool['res.country.state'].read(
                cr, uid, state_ids, ['code', 'country_id'], context=context)
            for state in states_r:
                states[state['code'].upper()] = state['id']
        f_geonames = zipfile.ZipFile(StringIO.StringIO(res_request.content))
        tempdir = tempfile.mkdtemp(prefix='openerp')
        f_geonames.extract('%s.txt' % country_code, tempdir)
        logger.info('The geonames zipfile has been decompressed')
        data_file = open(os.path.join(tempdir, '%s.txt' % country_code), 'r')
        data_file.seek(0)
        logger.info(
            'Starting to create the better zip entries %s state information'
            % (states and 'with' or 'without'))
        for row in unicodecsv.reader(
                data_file, encoding='utf-8', delimiter='	'):
            self.create_better_zip(
                cr, uid, row, country_id, country_code, states,
                context=context)
        data_file.close()
        logger.info(
            'The wizard to create better zip entries from geonames '
            'has been successfully completed.')
        return True
<|endoftext|>


<commit_before># Copyright 2013-2020 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class Delphes(CMakePackage):
    """A high energy physics framework for fast simulation
    of a generic collider experiment.
    """

    homepage = "https://cp3.irmp.ucl.ac.be/projects/delphes"
    git = "https://github.com/delphes/delphes.git"
    url = "http://cp3.irmp.ucl.ac.be/downloads/Delphes-3.4.2.tar.gz"

    maintainers = ['drbenmorgan', 'vvolkl', 'selvaggi']

    version('master', branch='master')
    version('3.4.2', sha256='d46a7c5474de650befdb89377115feee31f1743107ceb3d8da699be9d48c097b')
    version('3.4.1', sha256='4b5a2aeac326643f45b6d45c39ba2302e323eeb86d8cb58843c6e73949b1208a')
    version('3.4.0', sha256='c0f9500663a0c3a5c1eddcee598a67b5bcfc9318303195c6cacc0590b4023fa1')
    version('3.3.3', sha256='404de818a6b7852b01187ccf598d8ac19d308b9361f128751ef003cde248ff00')
    version('3.3.2', sha256='b8dc066e480678bb50ea0b68d157c391d47f66c084bda602d3d498538e682622')
    version('3.3.1', sha256='d8fcaa9711b5892ba24b2c7be38158dedbe552b159961f9d29887b2cc7eb2e83')
    version('3.3.0', sha256='3fcdcd31827227ff3d0d56df908b12289c67aa6d01c5725a2a9441c200f3966f')
    version('3.2.0', sha256='3510b0852c750120425f9b014cada25d48b90b29c512b974a9ffbd7aa80ccde4')
    version('3.1.2', sha256='edfccc47f7666d3607e86db82c6c79cfb10716423b496f0c0bdd4060b717ea1d')
    version('3.1.1', sha256='c4128575b6103239ca13de392f47da2eaedfd93c3789b1ecb32eea09da3408e4')
    version('3.1.0', sha256='c37b07aea3e57b39d34bf07f8afd209e36b278cf3792cd6e970d96a2c3b114eb')
    version('3.0.12', sha256='55b4cf25f681c75457e33ad4ee615b9ab66317216675ca5f466ab256aa85cd16')
    version('3.0.11', sha256='870921c8070762dc56aa0b8e0e07d1756584399e8740c848d330fc0fcb5e0604')
    version('3.0.10', sha256='872a386baf298cab14e42aac198dbf7184a2ab7c28ee712448060e1dec078d34')
    version('3.0.9', sha256='d12592fe66062a51e513a8d070fe1f49b672a4328bad2aa5cdb682937391a639')
    version('3.0.8', sha256='8ab146ca3c163932ab21df9168d8ca86dbb1c3494b7bdc3e143743d769803c23')
    version('3.0.7', sha256='7f43c84bca38fb8a41d7840dd2d7fab52456182babaa1e528791d0f4e517aba8')
    version('3.0.6', sha256='9e225731d57d2a76d35886841f8eff121bb3a45560b16077bd8c351151581d88')
    version('3.0.5', sha256='ab64ec6d2476fbfa40562e7edb510a8ab4c4fe5be77a4353ebf315c2af181a80')

    depends_on('cmake', type='build')
    depends_on('root')

    variant('build_type', default='Release',
            description='The build type to build',
            values=('Debug', 'Release'))

    variant('cxxstd',
            default='17',
            values=('14', '17'),
            multi=False,
            description='Use the specified C++ standard when building.')

    def cmake_args(self):
        args = []
        # C++ Standard
        args.append('-DCMAKE_CXX_STANDARD=%s'
                    % self.spec.variants['cxxstd'].value)
        return args

    def setup_run_environment(self, env):
        # make the cards distributed with delphes more easily accessible
        env.set('DELPHESCARDS', self.prefix.cards)
<commit_msg>delphes: better specify the dependency on root (#17175)<commit_after># Copyright 2013-2020 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class Delphes(CMakePackage):
    """A high energy physics framework for fast simulation
    of a generic collider experiment.
    """

    homepage = "https://cp3.irmp.ucl.ac.be/projects/delphes"
    git = "https://github.com/delphes/delphes.git"
    url = "http://cp3.irmp.ucl.ac.be/downloads/Delphes-3.4.2.tar.gz"

    maintainers = ['drbenmorgan', 'vvolkl', 'selvaggi']

    version('master', branch='master')
    version('3.4.2', sha256='d46a7c5474de650befdb89377115feee31f1743107ceb3d8da699be9d48c097b')
    version('3.4.1', sha256='4b5a2aeac326643f45b6d45c39ba2302e323eeb86d8cb58843c6e73949b1208a')
    version('3.4.0', sha256='c0f9500663a0c3a5c1eddcee598a67b5bcfc9318303195c6cacc0590b4023fa1')
    version('3.3.3', sha256='404de818a6b7852b01187ccf598d8ac19d308b9361f128751ef003cde248ff00')
    version('3.3.2', sha256='b8dc066e480678bb50ea0b68d157c391d47f66c084bda602d3d498538e682622')
    version('3.3.1', sha256='d8fcaa9711b5892ba24b2c7be38158dedbe552b159961f9d29887b2cc7eb2e83')
    version('3.3.0', sha256='3fcdcd31827227ff3d0d56df908b12289c67aa6d01c5725a2a9441c200f3966f')
    version('3.2.0', sha256='3510b0852c750120425f9b014cada25d48b90b29c512b974a9ffbd7aa80ccde4')
    version('3.1.2', sha256='edfccc47f7666d3607e86db82c6c79cfb10716423b496f0c0bdd4060b717ea1d')
    version('3.1.1', sha256='c4128575b6103239ca13de392f47da2eaedfd93c3789b1ecb32eea09da3408e4')
    version('3.1.0', sha256='c37b07aea3e57b39d34bf07f8afd209e36b278cf3792cd6e970d96a2c3b114eb')
    version('3.0.12', sha256='55b4cf25f681c75457e33ad4ee615b9ab66317216675ca5f466ab256aa85cd16')
    version('3.0.11', sha256='870921c8070762dc56aa0b8e0e07d1756584399e8740c848d330fc0fcb5e0604')
    version('3.0.10', sha256='872a386baf298cab14e42aac198dbf7184a2ab7c28ee712448060e1dec078d34')
    version('3.0.9', sha256='d12592fe66062a51e513a8d070fe1f49b672a4328bad2aa5cdb682937391a639')
    version('3.0.8', sha256='8ab146ca3c163932ab21df9168d8ca86dbb1c3494b7bdc3e143743d769803c23')
    version('3.0.7', sha256='7f43c84bca38fb8a41d7840dd2d7fab52456182babaa1e528791d0f4e517aba8')
    version('3.0.6', sha256='9e225731d57d2a76d35886841f8eff121bb3a45560b16077bd8c351151581d88')
    version('3.0.5', sha256='ab64ec6d2476fbfa40562e7edb510a8ab4c4fe5be77a4353ebf315c2af181a80')

    depends_on('cmake', type='build')
    depends_on('root cxxstd=14', when='cxxstd=14')
    depends_on('root cxxstd=17', when='cxxstd=17')

    variant('build_type', default='Release',
            description='The build type to build',
            values=('Debug', 'Release'))

    variant('cxxstd',
            default='17',
            values=('14', '17'),
            multi=False,
            description='Use the specified C++ standard when building.')

    def cmake_args(self):
        args = []
        # C++ Standard
        args.append('-DCMAKE_CXX_STANDARD=%s'
                    % self.spec.variants['cxxstd'].value)
        return args

    def setup_run_environment(self, env):
        # make the cards distributed with delphes more easily accessible
        env.set('DELPHESCARDS', self.prefix.cards)
<|endoftext|>


<commit_before>from common import *

def cached_property(func):
    pname = "_" + func.__name__
    @property
    @wraps(func)
    def caching(self, *args, **kwargs):
        if not hasattr(self, pname):
            setattr(self, pname, func(self, *args, **kwargs))
        return getattr(self, pname)
    return caching

def show_queries(func):
    @wraps(func)
    def queries(self, *args, **kwargs):
        connection.queries[:] = []
        stripped = [arg for arg in args if arg != "querylog"]
        try:
            return func(self, *stripped, **kwargs)
        finally:
            if "querylog" in args:
                cherrypy.response.headers["Content-type"] = "text/plain"
                return pformat(connection.queries)
    return queries

def csrf_protected(func):
    @wraps(func)
    def protected(*args, csrf_token, **kwargs):
        check_csrf(csrf_token)
        return func(*args, **kwargs)
    return protected

def ajax(func):
    @wraps(func)
    def returns_json(*args, **kwargs):
        cherrypy.response.headers["Content-Type"] = "application/json"
        assert cherrypy.request.method == "POST", "POST required"
        check_csrf(kwargs.pop("csrf_token", None))
        return json.dumps(func(*args, **kwargs)).encode("utf-8")
    return returns_json

def csv_file(func):
    @wraps(func)
    def csvout(self):
        cherrypy.response.headers["Content-Type"] = "application/csv"
        cherrypy.response.headers["Content-Disposition"] = "attachment; filename=" + func.__name__ + ".csv"
        writer = StringIO()
        func(self, csv.writer(writer))
        return writer.getvalue()
    return csvout

def credit_card(func):
    @wraps(func)
    def charge(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except HTTPRedirect:
            raise
        except:
            # TODO: email eli a stacktrace and return a template
            return traceback.format_exc()
    return charge

def render(template, data = None):
    import constants
    from models import Account, all_models
    data = data or {}
    data.update({m.__name__: m for m in all_models()})
    data.update({k: v for k,v in constants.__dict__.items() if re.match("^[_A-Z0-9]*$", k)})
    data.update({k: getattr(state, k) for k in dir(state) if re.match("^[_A-Z0-9]*$", k)})
    data.update({
        "now":   datetime.now(),
        "PAGE":  cherrypy.request.path_info.split("/")[-1],
        "CSRF_TOKEN":  getattr(cherrypy, "session", {}).get("csrf_token")
    })

    access = Account.access_set()
    for acctype in ["ACCOUNTS","PEOPLE","STUFF","MONEY","CHALLENGES","CHECKINS"]:
        data["HAS_" + acctype + "_ACCESS"] = getattr(constants, acctype) in access

    rendered = loader.get_template(template).render( Context(data) )
    if not state.AT_THE_CON and Account.admin_name() == "Nick Marinelli":
        rendered = rendered.replace("Fest", "Con")
    return rendered

def renderable(func):
    @wraps(func)
    def with_rendering(*args, **kwargs):
        res = func(*args, **kwargs)
        if isinstance(res, dict):
            mod_name = func.__module__.split(".")[1]
            template = "{}/{}.html".format(mod_name, func.__name__)
            return render(template, res)
        else:
            return res
    return with_rendering

def unrestricted(func):
    func.restricted = False
    return func

def restricted(func):
    @wraps(func)
    def with_restrictions(*args, **kwargs):
        if func.restricted:
            if func.restricted == (SIGNUPS,):
                if not cherrypy.session.get("staffer_id"):
                    raise HTTPRedirect("../signups/login?message=You+are+not+logged+in")

            elif cherrypy.session.get("account_id") is None:
                raise HTTPRedirect("../accounts/login?message=You+are+not+logged+in")

            else:
                from models import Account
                if not set(func.restricted).intersection( Account.access_set() ):
                    if len(func.restricted) == 1:
                        return "You need {} access for this page".format(dict(ACCESS_OPTS)[func.restricted[0]])
                    else:
                        return ("You need at least one of the following access levels to view this page: "
                            + ", ".join(dict(ACCESS_OPTS)[r] for r in func.restricted))

        return func(*args, **kwargs)
    return with_restrictions

class all_renderable:
    def __init__(self, *needs_access):
        self.needs_access = needs_access

    def __call__(self, klass):
        for name,func in klass.__dict__.items():
            if hasattr(func, "__call__"):
                func.restricted = getattr(func, "restricted", self.needs_access)
                new_func = show_queries(restricted(renderable(func)))
                new_func.exposed = True
                setattr(klass, name, new_func)
        return klass

register = template.Library()
def tag(klass):
    @register.tag(klass.__name__)
    def tagged(parser, token):
        return klass(*token.split_contents()[1:])
    return klass
<commit_msg>fixed csrf bug in email templates<commit_after>from common import *

def cached_property(func):
    pname = "_" + func.__name__
    @property
    @wraps(func)
    def caching(self, *args, **kwargs):
        if not hasattr(self, pname):
            setattr(self, pname, func(self, *args, **kwargs))
        return getattr(self, pname)
    return caching

def show_queries(func):
    @wraps(func)
    def queries(self, *args, **kwargs):
        connection.queries[:] = []
        stripped = [arg for arg in args if arg != "querylog"]
        try:
            return func(self, *stripped, **kwargs)
        finally:
            if "querylog" in args:
                cherrypy.response.headers["Content-type"] = "text/plain"
                return pformat(connection.queries)
    return queries

def csrf_protected(func):
    @wraps(func)
    def protected(*args, csrf_token, **kwargs):
        check_csrf(csrf_token)
        return func(*args, **kwargs)
    return protected

def ajax(func):
    @wraps(func)
    def returns_json(*args, **kwargs):
        cherrypy.response.headers["Content-Type"] = "application/json"
        assert cherrypy.request.method == "POST", "POST required"
        check_csrf(kwargs.pop("csrf_token", None))
        return json.dumps(func(*args, **kwargs)).encode("utf-8")
    return returns_json

def csv_file(func):
    @wraps(func)
    def csvout(self):
        cherrypy.response.headers["Content-Type"] = "application/csv"
        cherrypy.response.headers["Content-Disposition"] = "attachment; filename=" + func.__name__ + ".csv"
        writer = StringIO()
        func(self, csv.writer(writer))
        return writer.getvalue()
    return csvout

def credit_card(func):
    @wraps(func)
    def charge(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except HTTPRedirect:
            raise
        except:
            # TODO: email eli a stacktrace and return a template
            return traceback.format_exc()
    return charge

def render(template, data = None):
    import constants
    from models import Account, all_models
    data = data or {}
    data.update({m.__name__: m for m in all_models()})
    data.update({k: v for k,v in constants.__dict__.items() if re.match("^[_A-Z0-9]*$", k)})
    data.update({k: getattr(state, k) for k in dir(state) if re.match("^[_A-Z0-9]*$", k)})
    data.update({
        "now":   datetime.now(),
        "PAGE":  cherrypy.request.path_info.split("/")[-1]
    })
    try:
        data["CSRF_TOKEN"] = cherrypy.session["csrf_token"]
    except:
        pass

    access = Account.access_set()
    for acctype in ["ACCOUNTS","PEOPLE","STUFF","MONEY","CHALLENGES","CHECKINS"]:
        data["HAS_" + acctype + "_ACCESS"] = getattr(constants, acctype) in access

    rendered = loader.get_template(template).render( Context(data) )
    if not state.AT_THE_CON and Account.admin_name() == "Nick Marinelli":
        rendered = rendered.replace("Fest", "Con")
    return rendered

def renderable(func):
    @wraps(func)
    def with_rendering(*args, **kwargs):
        res = func(*args, **kwargs)
        if isinstance(res, dict):
            mod_name = func.__module__.split(".")[1]
            template = "{}/{}.html".format(mod_name, func.__name__)
            return render(template, res)
        else:
            return res
    return with_rendering

def unrestricted(func):
    func.restricted = False
    return func

def restricted(func):
    @wraps(func)
    def with_restrictions(*args, **kwargs):
        if func.restricted:
            if func.restricted == (SIGNUPS,):
                if not cherrypy.session.get("staffer_id"):
                    raise HTTPRedirect("../signups/login?message=You+are+not+logged+in")

            elif cherrypy.session.get("account_id") is None:
                raise HTTPRedirect("../accounts/login?message=You+are+not+logged+in")

            else:
                from models import Account
                if not set(func.restricted).intersection( Account.access_set() ):
                    if len(func.restricted) == 1:
                        return "You need {} access for this page".format(dict(ACCESS_OPTS)[func.restricted[0]])
                    else:
                        return ("You need at least one of the following access levels to view this page: "
                            + ", ".join(dict(ACCESS_OPTS)[r] for r in func.restricted))

        return func(*args, **kwargs)
    return with_restrictions

class all_renderable:
    def __init__(self, *needs_access):
        self.needs_access = needs_access

    def __call__(self, klass):
        for name,func in klass.__dict__.items():
            if hasattr(func, "__call__"):
                func.restricted = getattr(func, "restricted", self.needs_access)
                new_func = show_queries(restricted(renderable(func)))
                new_func.exposed = True
                setattr(klass, name, new_func)
        return klass

register = template.Library()
def tag(klass):
    @register.tag(klass.__name__)
    def tagged(parser, token):
        return klass(*token.split_contents()[1:])
    return klass
<|endoftext|>


<commit_before># coding: utf-8
from __future__ import absolute_import
from __future__ import unicode_literals
from codecs import BOM_UTF8
from contextlib import closing
import io
import os

from django.test import SimpleTestCase
from lxml import html, etree
from mock import patch, Mock

from couchexport.export import export_from_tables
from couchexport.models import Format
from couchexport.writers import (
    MAX_XLS_COLUMNS,
    CsvFileWriter,
    PythonDictWriter,
    XlsLengthException,
    ZippedExportWriter,
)
from six.moves import range

class ZippedExportWriterTests(SimpleTestCase):

    def setUp(self):
        self.zip_file_patch = patch('zipfile.ZipFile')
        self.MockZipFile = self.zip_file_patch.start()

        self.path_mock = Mock()
        self.path_mock.get_path.return_value = 'tmp'

        self.writer = ZippedExportWriter()
        self.writer.archive_basepath = '✓path'.encode('utf-8')
        self.writer.tables = [self.path_mock]
        self.writer.file = Mock()

    def tearDown(self):
        self.zip_file_patch.stop()
        del self.writer

    def test_zipped_export_writer_unicode(self):
        mock_zip_file = self.MockZipFile.return_value
        self.writer.table_names = {0: 'ひらがな'}
        self.writer._write_final_result()
        mock_zip_file.write.assert_called_with(
            'tmp',
            os.path.join(self.writer.archive_basepath, 'ひらがな.csv'.encode('utf-8'))
        )

    def test_zipped_export_writer_utf8(self):
        mock_zip_file = self.MockZipFile.return_value
        self.writer.table_names = {0: b'\xe3\x81\xb2\xe3\x82\x89\xe3\x81\x8c\xe3\x81\xaa'}
        self.writer._write_final_result()
        mock_zip_file.write.assert_called_with(
            'tmp',
            os.path.join(self.writer.archive_basepath, 'ひらがな.csv'.encode('utf-8'))
        )

class CsvFileWriterTests(SimpleTestCase):

    def test_csv_file_writer_bom(self):
        """
        CsvFileWriter should prepend a byte-order mark to the start of the CSV file for Excel
        """
        writer = CsvFileWriter()
        headers = ['ham', 'spam', 'eggs']
        writer.open('Spam')
        writer.write_row(headers)
        writer.finish()
        file_start = writer.get_file().read(6)
        self.assertEqual(file_start, BOM_UTF8 + b'ham')

    def test_csv_file_writer_utf8(self):
        writer = CsvFileWriter()
        headers = ['hám', 'spam', 'eggs']
        writer.open('Spam')
        writer.write_row(headers)
        writer.finish()
        file_start = writer.get_file().read(7)
        self.assertEqual(file_start, BOM_UTF8 + 'hám'.encode('utf-8'))

    def test_csv_file_writer_int(self):
        writer = CsvFileWriter()
        headers = [100, 'spam', 'eggs']
        writer.open('Spam')
        writer.write_row(headers)
        writer.finish()
        file_start = writer.get_file().read(6)
        self.assertEqual(file_start, BOM_UTF8 + b'100')

class HtmlExportWriterTests(SimpleTestCase):

    def test_nones_transformed(self):
        headers = ('Breakfast', 'Breakfast', 'Amuse-Bouche', 'Breakfast')
        row = ('spam', 'spam', None, 'spam')
        table = (headers, row, row, row)
        export_tables = (('Spam', table),)

        with closing(io.BytesIO()) as file_:
            export_from_tables(export_tables, file_, Format.HTML)
            html_string = file_.getvalue()

        root = html.fromstring(html_string)
        html_rows = [
            [etree.tostring(td).strip() for td in tr.xpath('./td')]
            for tr in root.xpath('./body/table/tbody/tr')
        ]
        self.assertEqual(html_rows,
                         [['spam	', 'spam	', '', 'spam	'],
                          ['spam	', 'spam	', '', 'spam	'],
                          ['spam	', 'spam	', '', 'spam	']])

class Excel2007ExportWriterTests(SimpleTestCase):

    def test_bytestrings(self):
        format_ = Format.XLS_2007
        file_ = io.BytesIO()
        table = [
            [b'heading\xe2\x80\x931', b'heading\xe2\x80\x932', b'heading\xe2\x80\x933'],
            [b'row1\xe2\x80\x931', b'row1\xe2\x80\x932', b'row1\xe2\x80\x933'],
            [b'row2\xe2\x80\x931', b'row2\xe2\x80\x932', b'row2\xe2\x80\x933'],
        ]
        tables = [[b'table\xe2\x80\x93title', table]]
        export_from_tables(tables, file_, format_)

class Excel2003ExportWriterTests(SimpleTestCase):

    def test_data_length(self):
        format_ = Format.XLS
        file_ = io.BytesIO()
        table = [
            ['header{}'.format(i) for i in range(MAX_XLS_COLUMNS)],
            ['row{}'.format(i) for i in range(MAX_XLS_COLUMNS)],
        ]
        tables = [['title', table]]

        with self.assertRaises(XlsLengthException):
            export_from_tables(tables, file_, format_)

        table = [
            ['header{}'.format(i) for i in range(MAX_XLS_COLUMNS - 1)],
            ['row{}'.format(i) for i in range(MAX_XLS_COLUMNS - 1)],
        ]
        tables = [['title', table]]
        export_from_tables(tables, file_, format_)

class HeaderNameTest(SimpleTestCase):

    def test_names_matching_case(self):
        writer = PythonDictWriter()
        stringio = io.StringIO()
        table_index_1 = "case_Sensitive"
        table_index_2 = "case_sensitive"
        table_headers = [[]]
        writer.open(
            [
                (table_index_1, table_headers),
                (table_index_2, table_headers)
            ],
            stringio
        )
        writer.close()
        preview = writer.get_preview()

        first_sheet_name = preview[0]['table_name']
        second_sheet_name = preview[1]['table_name']
        self.assertNotEqual(
            first_sheet_name.lower(),
            second_sheet_name.lower(),
            "Sheet names must not be equal. Comparison is NOT case sensitive. Names were '{}' and '{}'".format(
                first_sheet_name, second_sheet_name
            )
        )

    def test_max_header_length(self):
        writer = PythonDictWriter()
        writer.max_table_name_size = 10
        stringio = io.StringIO()
        table_index = "my_table_index"
        table_headers = [("first header", "second header")]
        writer.open(
            [(table_index, table_headers)],
            stringio
        )
        writer.close()
        preview = writer.get_preview()
        self.assertGreater(len(table_index), writer.max_table_name_size)
        self.assertLessEqual(len(preview[0]['table_name']), writer.max_table_name_size)

    def test_max_header_length_duplicates(self):
        writer = PythonDictWriter()
        writer.max_table_name_size = 7
        stringio = io.StringIO()
        table_headers = [("header1", "header2")]
        writer.open(
            [
                ("prefix1: index", table_headers),
                ("prefix2- index", table_headers),
            ],
            stringio
        )
        writer.close()
        preview = writer.get_preview()
        table_names = {table['table_name'] for table in preview}
        self.assertEqual(len(table_names), 2)
<commit_msg>decode<commit_after># coding: utf-8
from __future__ import absolute_import
from __future__ import unicode_literals
from codecs import BOM_UTF8
from contextlib import closing
import io
import os

from django.test import SimpleTestCase
from lxml import html, etree
from mock import patch, Mock

from couchexport.export import export_from_tables
from couchexport.models import Format
from couchexport.writers import (
    MAX_XLS_COLUMNS,
    CsvFileWriter,
    PythonDictWriter,
    XlsLengthException,
    ZippedExportWriter,
)
from six.moves import range

class ZippedExportWriterTests(SimpleTestCase):

    def setUp(self):
        self.zip_file_patch = patch('zipfile.ZipFile')
        self.MockZipFile = self.zip_file_patch.start()

        self.path_mock = Mock()
        self.path_mock.get_path.return_value = 'tmp'

        self.writer = ZippedExportWriter()
        self.writer.archive_basepath = '✓path'.encode('utf-8')
        self.writer.tables = [self.path_mock]
        self.writer.file = Mock()

    def tearDown(self):
        self.zip_file_patch.stop()
        del self.writer

    def test_zipped_export_writer_unicode(self):
        mock_zip_file = self.MockZipFile.return_value
        self.writer.table_names = {0: 'ひらがな'}
        self.writer._write_final_result()
        mock_zip_file.write.assert_called_with(
            'tmp',
            os.path.join(self.writer.archive_basepath, 'ひらがな.csv'.encode('utf-8'))
        )

    def test_zipped_export_writer_utf8(self):
        mock_zip_file = self.MockZipFile.return_value
        self.writer.table_names = {0: b'\xe3\x81\xb2\xe3\x82\x89\xe3\x81\x8c\xe3\x81\xaa'}
        self.writer._write_final_result()
        mock_zip_file.write.assert_called_with(
            'tmp',
            os.path.join(self.writer.archive_basepath, 'ひらがな.csv'.encode('utf-8'))
        )

class CsvFileWriterTests(SimpleTestCase):

    def test_csv_file_writer_bom(self):
        """
        CsvFileWriter should prepend a byte-order mark to the start of the CSV file for Excel
        """
        writer = CsvFileWriter()
        headers = ['ham', 'spam', 'eggs']
        writer.open('Spam')
        writer.write_row(headers)
        writer.finish()
        file_start = writer.get_file().read(6)
        self.assertEqual(file_start, BOM_UTF8 + b'ham')

    def test_csv_file_writer_utf8(self):
        writer = CsvFileWriter()
        headers = ['hám', 'spam', 'eggs']
        writer.open('Spam')
        writer.write_row(headers)
        writer.finish()
        file_start = writer.get_file().read(7)
        self.assertEqual(file_start, BOM_UTF8 + 'hám'.encode('utf-8'))

    def test_csv_file_writer_int(self):
        writer = CsvFileWriter()
        headers = [100, 'spam', 'eggs']
        writer.open('Spam')
        writer.write_row(headers)
        writer.finish()
        file_start = writer.get_file().read(6)
        self.assertEqual(file_start, BOM_UTF8 + b'100')

class HtmlExportWriterTests(SimpleTestCase):

    def test_nones_transformed(self):
        headers = ('Breakfast', 'Breakfast', 'Amuse-Bouche', 'Breakfast')
        row = ('spam', 'spam', None, 'spam')
        table = (headers, row, row, row)
        export_tables = (('Spam', table),)

        with closing(io.BytesIO()) as file_:
            export_from_tables(export_tables, file_, Format.HTML)
            html_string = file_.getvalue()

        root = html.fromstring(html_string)
        html_rows = [
            [etree.tostring(td).strip().decode('utf-8') for td in tr.xpath('./td')]
            for tr in root.xpath('./body/table/tbody/tr')
        ]
        self.assertEqual(html_rows,
                         [['spam	', 'spam	', '', 'spam	'],
                          ['spam	', 'spam	', '', 'spam	'],
                          ['spam	', 'spam	', '', 'spam	']])

class Excel2007ExportWriterTests(SimpleTestCase):

    def test_bytestrings(self):
        format_ = Format.XLS_2007
        file_ = io.BytesIO()
        table = [
            [b'heading\xe2\x80\x931', b'heading\xe2\x80\x932', b'heading\xe2\x80\x933'],
            [b'row1\xe2\x80\x931', b'row1\xe2\x80\x932', b'row1\xe2\x80\x933'],
            [b'row2\xe2\x80\x931', b'row2\xe2\x80\x932', b'row2\xe2\x80\x933'],
        ]
        tables = [[b'table\xe2\x80\x93title', table]]
        export_from_tables(tables, file_, format_)

class Excel2003ExportWriterTests(SimpleTestCase):

    def test_data_length(self):
        format_ = Format.XLS
        file_ = io.BytesIO()
        table = [
            ['header{}'.format(i) for i in range(MAX_XLS_COLUMNS)],
            ['row{}'.format(i) for i in range(MAX_XLS_COLUMNS)],
        ]
        tables = [['title', table]]

        with self.assertRaises(XlsLengthException):
            export_from_tables(tables, file_, format_)

        table = [
            ['header{}'.format(i) for i in range(MAX_XLS_COLUMNS - 1)],
            ['row{}'.format(i) for i in range(MAX_XLS_COLUMNS - 1)],
        ]
        tables = [['title', table]]
        export_from_tables(tables, file_, format_)

class HeaderNameTest(SimpleTestCase):

    def test_names_matching_case(self):
        writer = PythonDictWriter()
        stringio = io.StringIO()
        table_index_1 = "case_Sensitive"
        table_index_2 = "case_sensitive"
        table_headers = [[]]
        writer.open(
            [
                (table_index_1, table_headers),
                (table_index_2, table_headers)
            ],
            stringio
        )
        writer.close()
        preview = writer.get_preview()

        first_sheet_name = preview[0]['table_name']
        second_sheet_name = preview[1]['table_name']
        self.assertNotEqual(
            first_sheet_name.lower(),
            second_sheet_name.lower(),
            "Sheet names must not be equal. Comparison is NOT case sensitive. Names were '{}' and '{}'".format(
                first_sheet_name, second_sheet_name
            )
        )

    def test_max_header_length(self):
        writer = PythonDictWriter()
        writer.max_table_name_size = 10
        stringio = io.StringIO()
        table_index = "my_table_index"
        table_headers = [("first header", "second header")]
        writer.open(
            [(table_index, table_headers)],
            stringio
        )
        writer.close()
        preview = writer.get_preview()
        self.assertGreater(len(table_index), writer.max_table_name_size)
        self.assertLessEqual(len(preview[0]['table_name']), writer.max_table_name_size)

    def test_max_header_length_duplicates(self):
        writer = PythonDictWriter()
        writer.max_table_name_size = 7
        stringio = io.StringIO()
        table_headers = [("header1", "header2")]
        writer.open(
            [
                ("prefix1: index", table_headers),
                ("prefix2- index", table_headers),
            ],
            stringio
        )
        writer.close()
        preview = writer.get_preview()
        table_names = {table['table_name'] for table in preview}
        self.assertEqual(len(table_names), 2)
<|endoftext|>


<commit_before>import ConfigParser
import os
import logging
import logging.config
import platform
import string
import subprocess
import sys
from optparse import OptionParser, Values
from cStringIO import StringIO
from util import getOS

# CONSTANTS
DATADOG_CONF = "datadog.conf"
DEFAULT_CHECK_FREQUENCY = 15 # seconds
DEFAULT_STATSD_FREQUENCY = 10 # seconds
PUP_STATSD_FREQUENCY = 2 # seconds

class DDConfigNotFound(Exception): pass

def get_parsed_args():
    parser = OptionParser()
    parser.add_option('-d', '--dd_url', action='store', default=None,
                        dest='dd_url')
    parser.add_option('-c', '--clean', action='store_true', default=False,
                        dest='clean')
    parser.add_option('-u', '--use-local-forwarder', action='store_true',
                        default=False,dest='use_forwarder')
    parser.add_option('-n', '--disable-dd', action='store_true', default=False,
                        dest="disable_dd")
    try:
        options, args = parser.parse_args()
    except SystemExit:
        options, args = Values({'dd_url': None,
                                'clean': False,
                                'use_forwarder':False,
                                'disable_dd':False}), [] # Ignore parse errors
    return options, args

def get_version():
    return "3.1.3"

def skip_leading_wsp(f):
    "Works on a file, returns a file-like object"
    return StringIO("\n".join(map(string.strip, f.readlines())))

def initialize_logging(config_path, os_name=None):
    try:
        logging.config.fileConfig(config_path)
    except Exception, e:
        sys.stderr.write("Couldn't initialize logging: %s" % str(e))

def _windows_config_path():
    # Find the 'common appdata' path
    import ctypes
    from ctypes import wintypes, windll

    CSIDL_COMMON_APPDATA = 35

    _SHGetFolderPath = windll.shell32.SHGetFolderPathW
    _SHGetFolderPath.argtypes = [wintypes.HWND,
                                ctypes.c_int,
                                wintypes.HANDLE,
                                wintypes.DWORD, wintypes.LPCWSTR]

    path_buf = wintypes.create_unicode_buffer(wintypes.MAX_PATH)
    result = _SHGetFolderPath(0, CSIDL_COMMON_APPDATA, 0, 0, path_buf)
    common_data = path_buf.value

    path = os.path.join(common_data, 'Datadog Agent', DATADOG_CONF)
    if os.path.exists(path):
        return path
    raise DDConfigNotFound(path)

def _unix_config_path():
    path = os.path.join('/etc/dd-agent', DATADOG_CONF)
    if os.path.exists(path):
        return path
    raise DDConfigNotFound(path)

def get_config_path(cfg_path=None, os_name=None):
    # Check if there's an override and if it exists
    if cfg_path is not None and os.path.exists(cfg_path):
        return cfg_path

    # Check for an OS-specific path, continue on not-found exceptions
    exc = None
    if os_name == 'windows':
        try:
            return _windows_config_path()
        except DDConfigNotFound, e:
            exc = e
    else:
        try:
            return _unix_config_path()
        except DDConfigNotFound, e:
            exc = e

    # Check if there's a config stored in the current agent directory
    path = os.path.realpath(__file__)
    path = os.path.dirname(path)
    if os.path.exists(os.path.join(path, DATADOG_CONF)):
        return os.path.join(path, DATADOG_CONF)

    # If all searches fail, exit the agent with an error
    sys.stderr.write("Please supply a configuration file at %s or in the directory where the agent is currently deployed.\n" % exc.message)
    sys.exit(3)

def get_config(parse_args = True, cfg_path=None, init_logging=False, options=None):
    if parse_args:
        options, args = get_parsed_args()
    elif not options:
        args = None

    # General config
    agentConfig = {
        'check_freq': DEFAULT_CHECK_FREQUENCY,
        'debug_mode': False,
        'dogstatsd_interval': DEFAULT_STATSD_FREQUENCY,
        'dogstatsd_port': 8125,
        'dogstatsd_target': 'http://localhost:17123',
        'graphite_listen_port': None,
        'hostname': None,
        'listen_port': None,
        'tags': None,
        'use_ec2_instance_id': False,
        'version': get_version(),
        'watchdog': True,
    }

    dogstatsd_interval = DEFAULT_STATSD_FREQUENCY

    # Config handling
    try:
        # Find the right config file
        path = os.path.realpath(__file__)
        path = os.path.dirname(path)

        config_path = get_config_path(cfg_path, os_name=getOS())
        config = ConfigParser.ConfigParser()
        config.readfp(skip_leading_wsp(open(config_path)))

        if init_logging:
            initialize_logging(config_path, os_name=getOS())

        # bulk import
        for option in config.options('Main'):
            agentConfig[option] = config.get('Main', option)

        #
        # Core config
        #

        if config.has_option('Main', 'use_dd'):
            agentConfig['use_dd'] = config.get('Main', 'use_dd').lower() in ("yes", "true")
        else:
            agentConfig['use_dd'] = True

        if options is not None and options.use_forwarder:
            listen_port = 17123
            if config.has_option('Main','listen_port'):
                listen_port = config.get('Main','listen_port')
            agentConfig['dd_url'] = "http://localhost:" + str(listen_port)
        elif options is not None and not options.disable_dd and options.dd_url:
            agentConfig['dd_url'] = options.dd_url
        else:
            agentConfig['dd_url'] = config.get('Main', 'dd_url')
        if agentConfig['dd_url'].endswith('/'):
            agentConfig['dd_url'] = agentConfig['dd_url'][:-1]

        # Whether also to send to Pup
        if config.has_option('Main', 'use_pup'):
            agentConfig['use_pup'] = config.get('Main', 'use_pup').lower() in ("yes", "true")
        else:
            agentConfig['use_pup'] = True

        if agentConfig['use_pup']:
            if config.has_option('Main', 'pup_url'):
                agentConfig['pup_url'] = config.get('Main', 'pup_url')
            else:
                agentConfig['pup_url'] = 'http://localhost:17125'

            pup_port = 17125
            if config.has_option('Main', 'pup_port'):
                agentConfig['pup_port'] = int(config.get('Main', 'pup_port'))

        # Increases the frequency of statsd metrics when only sending to Pup
        if not agentConfig['use_dd'] and agentConfig['use_pup']:
            dogstatsd_interval = PUP_STATSD_FREQUENCY

        if not agentConfig['use_dd'] and not agentConfig['use_pup']:
            sys.stderr.write("Please specify at least one endpoint to send metrics to. This can be done in datadog.conf.")
            exit(2)

        # Which API key to use
        agentConfig['api_key'] = config.get('Main', 'api_key')

        # Debug mode
        agentConfig['debug_mode'] = config.get('Main', 'debug_mode').lower() in ("yes", "true")

        if config.has_option('Main', 'use_ec2_instance_id'):
            use_ec2_instance_id = config.get('Main', 'use_ec2_instance_id')
            # translate yes into True, the rest into False
            agentConfig['use_ec2_instance_id'] = (use_ec2_instance_id.lower() == 'yes')

        if config.has_option('Main', 'check_freq'):
            try:
                agentConfig['check_freq'] = int(config.get('Main', 'check_freq'))
            except:
                pass

        # Disable Watchdog (optionally)
        if config.has_option('Main', 'watchdog'):
            if config.get('Main', 'watchdog').lower() in ('no', 'false'):
                agentConfig['watchdog'] = False

        # Optional graphite listener
        if config.has_option('Main','graphite_listen_port'):
            agentConfig['graphite_listen_port'] = int(config.get('Main','graphite_listen_port'))
        else:
            agentConfig['graphite_listen_port'] = None

        # Dogstatsd config
        dogstatsd_defaults = {
            'dogstatsd_port' : 8125,
            'dogstatsd_target' : 'http://localhost:17123',
            'dogstatsd_interval' : dogstatsd_interval
        }
        for key, value in dogstatsd_defaults.iteritems():
            if config.has_option('Main', key):
                agentConfig[key] = config.get('Main', key)
            else:
                agentConfig[key] = value

        # Optional config
        # FIXME not the prettiest code ever...
        if config.has_option('Main', 'use_mount'):
            agentConfig['use_mount'] = config.get('Main', 'use_mount').lower() in ("yes", "true", "1")

        if config.has_option('datadog', 'ddforwarder_log'):
            agentConfig['has_datadog'] = True

        # Dogstream config
        if config.has_option("Main", "dogstream_log"):
            # Older version, single log support
            log_path = config.get("Main", "dogstream_log")
            if config.has_option("Main", "dogstream_line_parser"):
                agentConfig["dogstreams"] = ':'.join([log_path, config.get("Main", "dogstream_line_parser")])
            else:
                agentConfig["dogstreams"] = log_path

        elif config.has_option("Main", "dogstreams"):
            agentConfig["dogstreams"] = config.get("Main", "dogstreams")

        if config.has_option("Main", "nagios_perf_cfg"):
            agentConfig["nagios_perf_cfg"] = config.get("Main", "nagios_perf_cfg")

        if config.has_section('WMI'):
            agentConfig['WMI'] = {}
            for key, value in config.items('WMI'):
                agentConfig['WMI'][key] = value

    except ConfigParser.NoSectionError, e:
        sys.stderr.write('Config file not found or incorrectly formatted.\n')
        sys.exit(2)

    except ConfigParser.ParsingError, e:
        sys.stderr.write('Config file not found or incorrectly formatted.\n')
        sys.exit(2)

    except ConfigParser.NoOptionError, e:
        sys.stderr.write('There are some items missing from your config file, but nothing fatal [%s]' % e)

    if 'apache_status_url' in agentConfig and agentConfig['apache_status_url'] == None:
        sys.stderr.write('You must provide a config value for apache_status_url. If you do not wish to use Apache monitoring, leave it as its default value - http://www.example.com/server-status/?auto.\n')
        sys.exit(2)

    if 'nginx_status_url' in agentConfig and agentConfig['nginx_status_url'] == None:
        sys.stderr.write('You must provide a config value for nginx_status_url. If you do not wish to use Nginx monitoring, leave it as its default value - http://www.example.com/nginx_status.\n')
        sys.exit(2)

    if 'mysql_server' in agentConfig and agentConfig['mysql_server'] != '' and 'mysql_user' in agentConfig and agentConfig['mysql_user'] != '' and 'mysql_pass' in agentConfig:
        try:
            import MySQLdb
        except ImportError:
            sys.stderr.write('You have configured MySQL for monitoring, but the MySQLdb module is not installed. For more info, see: http://help.datadoghq.com.\n')
            sys.exit(2)

    if 'mongodb_server' in agentConfig and agentConfig['mongodb_server'] != '':
        try:
            import pymongo
        except ImportError:
            sys.stderr.write('You have configured MongoDB for monitoring, but the pymongo module is not installed.\n')
            sys.exit(2)

    return agentConfig

def get_system_stats():
    systemStats = {
        'machine': platform.machine(),
        'platform': sys.platform,
        'processor': platform.processor(),
        'pythonV': platform.python_version()
    }

    if sys.platform == 'linux2':
        grep = subprocess.Popen(['grep', 'model name', '/proc/cpuinfo'], stdout=subprocess.PIPE, close_fds=True)
        wc = subprocess.Popen(['wc', '-l'], stdin=grep.stdout, stdout=subprocess.PIPE, close_fds=True)
        systemStats['cpuCores'] = int(wc.communicate()[0])

    if sys.platform == 'darwin':
        systemStats['cpuCores'] = int(subprocess.Popen(['sysctl', 'hw.ncpu'], stdout=subprocess.PIPE, close_fds=True).communicate()[0].split(': ')[1])

    if sys.platform == 'linux2':
        systemStats['nixV'] = platform.dist()

    elif sys.platform == 'darwin':
        systemStats['macV'] = platform.mac_ver()

    elif sys.platform.find('freebsd') != -1:
        version = platform.uname()[2]
        systemStats['fbsdV'] = ('freebsd', version, '') # no codename for FreeBSD

    return systemStats

def set_win32_cert_path():
    ''' In order to use tornado.httpclient with the packaged .exe on Windows we
    need to override the default ceritifcate location which is based on the path
    to tornado and will give something like "C:\path\to\program.exe\tornado/cert-file".

    If pull request #379 is accepted (https://github.com/facebook/tornado/pull/379) we
    will be able to override this in a clean way. For now, we have to monkey patch
    tornado.httpclient._DEFAULT_CA_CERTS
    '''
    crt_path = os.path.join(os.environ['PROGRAMFILES'], 'Datadog Agent',
        'ca-certificates.crt')
    import tornado.simple_httpclient
    tornado.simple_httpclient._DEFAULT_CA_CERTS = crt_path
<commit_msg>Change the config and program location once more<commit_after>import ConfigParser
import os
import logging
import logging.config
import platform
import string
import subprocess
import sys
from optparse import OptionParser, Values
from cStringIO import StringIO
from util import getOS

# CONSTANTS
DATADOG_CONF = "datadog.conf"
DEFAULT_CHECK_FREQUENCY = 15 # seconds
DEFAULT_STATSD_FREQUENCY = 10 # seconds
PUP_STATSD_FREQUENCY = 2 # seconds

class DDConfigNotFound(Exception): pass

def get_parsed_args():
    parser = OptionParser()
    parser.add_option('-d', '--dd_url', action='store', default=None,
                        dest='dd_url')
    parser.add_option('-c', '--clean', action='store_true', default=False,
                        dest='clean')
    parser.add_option('-u', '--use-local-forwarder', action='store_true',
                        default=False,dest='use_forwarder')
    parser.add_option('-n', '--disable-dd', action='store_true', default=False,
                        dest="disable_dd")
    try:
        options, args = parser.parse_args()
    except SystemExit:
        options, args = Values({'dd_url': None,
                                'clean': False,
                                'use_forwarder':False,
                                'disable_dd':False}), [] # Ignore parse errors
    return options, args

def get_version():
    return "3.1.3"

def skip_leading_wsp(f):
    "Works on a file, returns a file-like object"
    return StringIO("\n".join(map(string.strip, f.readlines())))

def initialize_logging(config_path, os_name=None):
    try:
        logging.config.fileConfig(config_path)
    except Exception, e:
        sys.stderr.write("Couldn't initialize logging: %s" % str(e))

def _windows_config_path():
    # Find the 'common appdata' path
    import ctypes
    from ctypes import wintypes, windll

    CSIDL_COMMON_APPDATA = 35

    _SHGetFolderPath = windll.shell32.SHGetFolderPathW
    _SHGetFolderPath.argtypes = [wintypes.HWND,
                                ctypes.c_int,
                                wintypes.HANDLE,
                                wintypes.DWORD, wintypes.LPCWSTR]

    path_buf = wintypes.create_unicode_buffer(wintypes.MAX_PATH)
    result = _SHGetFolderPath(0, CSIDL_COMMON_APPDATA, 0, 0, path_buf)
    common_data = path_buf.value

    path = os.path.join(common_data, 'Datadog', DATADOG_CONF)
    if os.path.exists(path):
        return path
    raise DDConfigNotFound(path)

def _unix_config_path():
    path = os.path.join('/etc/dd-agent', DATADOG_CONF)
    if os.path.exists(path):
        return path
    raise DDConfigNotFound(path)

def get_config_path(cfg_path=None, os_name=None):
    # Check if there's an override and if it exists
    if cfg_path is not None and os.path.exists(cfg_path):
        return cfg_path

    # Check for an OS-specific path, continue on not-found exceptions
    exc = None
    if os_name == 'windows':
        try:
            return _windows_config_path()
        except DDConfigNotFound, e:
            exc = e
    else:
        try:
            return _unix_config_path()
        except DDConfigNotFound, e:
            exc = e

    # Check if there's a config stored in the current agent directory
    path = os.path.realpath(__file__)
    path = os.path.dirname(path)
    if os.path.exists(os.path.join(path, DATADOG_CONF)):
        return os.path.join(path, DATADOG_CONF)

    # If all searches fail, exit the agent with an error
    sys.stderr.write("Please supply a configuration file at %s or in the directory where the agent is currently deployed.\n" % exc.message)
    sys.exit(3)

def get_config(parse_args = True, cfg_path=None, init_logging=False, options=None):
    if parse_args:
        options, args = get_parsed_args()
    elif not options:
        args = None

    # General config
    agentConfig = {
        'check_freq': DEFAULT_CHECK_FREQUENCY,
        'debug_mode': False,
        'dogstatsd_interval': DEFAULT_STATSD_FREQUENCY,
        'dogstatsd_port': 8125,
        'dogstatsd_target': 'http://localhost:17123',
        'graphite_listen_port': None,
        'hostname': None,
        'listen_port': None,
        'tags': None,
        'use_ec2_instance_id': False,
        'version': get_version(),
        'watchdog': True,
    }

    dogstatsd_interval = DEFAULT_STATSD_FREQUENCY

    # Config handling
    try:
        # Find the right config file
        path = os.path.realpath(__file__)
        path = os.path.dirname(path)

        config_path = get_config_path(cfg_path, os_name=getOS())
        config = ConfigParser.ConfigParser()
        config.readfp(skip_leading_wsp(open(config_path)))

        if init_logging:
            initialize_logging(config_path, os_name=getOS())

        # bulk import
        for option in config.options('Main'):
            agentConfig[option] = config.get('Main', option)

        #
        # Core config
        #

        if config.has_option('Main', 'use_dd'):
            agentConfig['use_dd'] = config.get('Main', 'use_dd').lower() in ("yes", "true")
        else:
            agentConfig['use_dd'] = True

        if options is not None and options.use_forwarder:
            listen_port = 17123
            if config.has_option('Main','listen_port'):
                listen_port = config.get('Main','listen_port')
            agentConfig['dd_url'] = "http://localhost:" + str(listen_port)
        elif options is not None and not options.disable_dd and options.dd_url:
            agentConfig['dd_url'] = options.dd_url
        else:
            agentConfig['dd_url'] = config.get('Main', 'dd_url')
        if agentConfig['dd_url'].endswith('/'):
            agentConfig['dd_url'] = agentConfig['dd_url'][:-1]

        # Whether also to send to Pup
        if config.has_option('Main', 'use_pup'):
            agentConfig['use_pup'] = config.get('Main', 'use_pup').lower() in ("yes", "true")
        else:
            agentConfig['use_pup'] = True

        if agentConfig['use_pup']:
            if config.has_option('Main', 'pup_url'):
                agentConfig['pup_url'] = config.get('Main', 'pup_url')
            else:
                agentConfig['pup_url'] = 'http://localhost:17125'

            pup_port = 17125
            if config.has_option('Main', 'pup_port'):
                agentConfig['pup_port'] = int(config.get('Main', 'pup_port'))

        # Increases the frequency of statsd metrics when only sending to Pup
        if not agentConfig['use_dd'] and agentConfig['use_pup']:
            dogstatsd_interval = PUP_STATSD_FREQUENCY

        if not agentConfig['use_dd'] and not agentConfig['use_pup']:
            sys.stderr.write("Please specify at least one endpoint to send metrics to. This can be done in datadog.conf.")
            exit(2)

        # Which API key to use
        agentConfig['api_key'] = config.get('Main', 'api_key')

        # Debug mode
        agentConfig['debug_mode'] = config.get('Main', 'debug_mode').lower() in ("yes", "true")

        if config.has_option('Main', 'use_ec2_instance_id'):
            use_ec2_instance_id = config.get('Main', 'use_ec2_instance_id')
            # translate yes into True, the rest into False
            agentConfig['use_ec2_instance_id'] = (use_ec2_instance_id.lower() == 'yes')

        if config.has_option('Main', 'check_freq'):
            try:
                agentConfig['check_freq'] = int(config.get('Main', 'check_freq'))
            except:
                pass

        # Disable Watchdog (optionally)
        if config.has_option('Main', 'watchdog'):
            if config.get('Main', 'watchdog').lower() in ('no', 'false'):
                agentConfig['watchdog'] = False

        # Optional graphite listener
        if config.has_option('Main','graphite_listen_port'):
            agentConfig['graphite_listen_port'] = int(config.get('Main','graphite_listen_port'))
        else:
            agentConfig['graphite_listen_port'] = None

        # Dogstatsd config
        dogstatsd_defaults = {
            'dogstatsd_port' : 8125,
            'dogstatsd_target' : 'http://localhost:17123',
            'dogstatsd_interval' : dogstatsd_interval
        }
        for key, value in dogstatsd_defaults.iteritems():
            if config.has_option('Main', key):
                agentConfig[key] = config.get('Main', key)
            else:
                agentConfig[key] = value

        # Optional config
        # FIXME not the prettiest code ever...
        if config.has_option('Main', 'use_mount'):
            agentConfig['use_mount'] = config.get('Main', 'use_mount').lower() in ("yes", "true", "1")

        if config.has_option('datadog', 'ddforwarder_log'):
            agentConfig['has_datadog'] = True

        # Dogstream config
        if config.has_option("Main", "dogstream_log"):
            # Older version, single log support
            log_path = config.get("Main", "dogstream_log")
            if config.has_option("Main", "dogstream_line_parser"):
                agentConfig["dogstreams"] = ':'.join([log_path, config.get("Main", "dogstream_line_parser")])
            else:
                agentConfig["dogstreams"] = log_path

        elif config.has_option("Main", "dogstreams"):
            agentConfig["dogstreams"] = config.get("Main", "dogstreams")

        if config.has_option("Main", "nagios_perf_cfg"):
            agentConfig["nagios_perf_cfg"] = config.get("Main", "nagios_perf_cfg")

        if config.has_section('WMI'):
            agentConfig['WMI'] = {}
            for key, value in config.items('WMI'):
                agentConfig['WMI'][key] = value

    except ConfigParser.NoSectionError, e:
        sys.stderr.write('Config file not found or incorrectly formatted.\n')
        sys.exit(2)

    except ConfigParser.ParsingError, e:
        sys.stderr.write('Config file not found or incorrectly formatted.\n')
        sys.exit(2)

    except ConfigParser.NoOptionError, e:
        sys.stderr.write('There are some items missing from your config file, but nothing fatal [%s]' % e)

    if 'apache_status_url' in agentConfig and agentConfig['apache_status_url'] == None:
        sys.stderr.write('You must provide a config value for apache_status_url. If you do not wish to use Apache monitoring, leave it as its default value - http://www.example.com/server-status/?auto.\n')
        sys.exit(2)

    if 'nginx_status_url' in agentConfig and agentConfig['nginx_status_url'] == None:
        sys.stderr.write('You must provide a config value for nginx_status_url. If you do not wish to use Nginx monitoring, leave it as its default value - http://www.example.com/nginx_status.\n')
        sys.exit(2)

    if 'mysql_server' in agentConfig and agentConfig['mysql_server'] != '' and 'mysql_user' in agentConfig and agentConfig['mysql_user'] != '' and 'mysql_pass' in agentConfig:
        try:
            import MySQLdb
        except ImportError:
            sys.stderr.write('You have configured MySQL for monitoring, but the MySQLdb module is not installed. For more info, see: http://help.datadoghq.com.\n')
            sys.exit(2)

    if 'mongodb_server' in agentConfig and agentConfig['mongodb_server'] != '':
        try:
            import pymongo
        except ImportError:
            sys.stderr.write('You have configured MongoDB for monitoring, but the pymongo module is not installed.\n')
            sys.exit(2)

    return agentConfig

def get_system_stats():
    systemStats = {
        'machine': platform.machine(),
        'platform': sys.platform,
        'processor': platform.processor(),
        'pythonV': platform.python_version()
    }

    if sys.platform == 'linux2':
        grep = subprocess.Popen(['grep', 'model name', '/proc/cpuinfo'], stdout=subprocess.PIPE, close_fds=True)
        wc = subprocess.Popen(['wc', '-l'], stdin=grep.stdout, stdout=subprocess.PIPE, close_fds=True)
        systemStats['cpuCores'] = int(wc.communicate()[0])

    if sys.platform == 'darwin':
        systemStats['cpuCores'] = int(subprocess.Popen(['sysctl', 'hw.ncpu'], stdout=subprocess.PIPE, close_fds=True).communicate()[0].split(': ')[1])

    if sys.platform == 'linux2':
        systemStats['nixV'] = platform.dist()

    elif sys.platform == 'darwin':
        systemStats['macV'] = platform.mac_ver()

    elif sys.platform.find('freebsd') != -1:
        version = platform.uname()[2]
        systemStats['fbsdV'] = ('freebsd', version, '') # no codename for FreeBSD

    return systemStats

def set_win32_cert_path():
    ''' In order to use tornado.httpclient with the packaged .exe on Windows we
    need to override the default ceritifcate location which is based on the path
    to tornado and will give something like "C:\path\to\program.exe\tornado/cert-file".

    If pull request #379 is accepted (https://github.com/facebook/tornado/pull/379) we
    will be able to override this in a clean way. For now, we have to monkey patch
    tornado.httpclient._DEFAULT_CA_CERTS
    '''
    crt_path = os.path.join(os.environ['PROGRAMFILES'], 'Datadog', 'Datadog Agent',
        'ca-certificates.crt')
    import tornado.simple_httpclient
    tornado.simple_httpclient._DEFAULT_CA_CERTS = crt_path
<|endoftext|>