<commit_before># Copyright 2009 the Melange authors.
#
# Licensed 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.

"""Module containing the GCI Callback.
"""

class Callback(object):
  """Callback object that handles interaction between the core.
  """

  API_VERSION = 1

  def __init__(self, core):
    """Initializes a new Callback object for the specified core.
    """

    self.core = core
    self.views = []

  def registerViews(self):
    """Instantiates all view objects.
    """
    from soc.modules.gci.views import accepted_orgs
    from soc.modules.gci.views import admin
    from soc.modules.gci.views import age_check
    from soc.modules.gci.views import all_tasks
    from soc.modules.gci.views import bulk_create
    from soc.modules.gci.views import dashboard
    from soc.modules.gci.views import delete_account
    from soc.modules.gci.views import document
    from soc.modules.gci.views import homepage
    from soc.modules.gci.views import invite
    from soc.modules.gci.views import leaderboard
    from soc.modules.gci.views import org_app
    from soc.modules.gci.views import org_home
    from soc.modules.gci.views import org_profile
    from soc.modules.gci.views import org_score
    from soc.modules.gci.views import participants
    from soc.modules.gci.views import profile
    from soc.modules.gci.views import profile_show
    from soc.modules.gci.views import program
    from soc.modules.gci.views import request
    from soc.modules.gci.views import student_forms
    from soc.modules.gci.views import students_info
    from soc.modules.gci.views import task
    from soc.modules.gci.views import task_list
    from soc.modules.gci.views import task_create

    self.views.append(accepted_orgs.AcceptedOrgsPage())
    self.views.append(accepted_orgs.AcceptedOrgsAdminPage())
    self.views.append(admin.DashboardPage())
    self.views.append(admin.LookupLinkIdPage())
    self.views.append(all_tasks.TaskListPage())
    self.views.append(age_check.AgeCheck())
    self.views.append(bulk_create.BulkCreate())
    self.views.append(dashboard.DashboardPage())
    self.views.append(delete_account.DeleteAccountPage())
    self.views.append(document.DocumentPage())
    self.views.append(document.EditDocumentPage())
    self.views.append(document.EventsPage())
    self.views.append(document.DocumentListPage())
    self.views.append(homepage.Homepage())
    self.views.append(invite.InvitePage())
    self.views.append(invite.ManageInvite())
    self.views.append(invite.RespondInvite())
    self.views.append(invite.ListUserInvitesPage())
    self.views.append(leaderboard.LeaderboardPage())
    self.views.append(leaderboard.StudentTasksPage())
    self.views.append(org_app.GCIOrgAppEditPage())
    self.views.append(org_profile.OrgProfilePage())
    self.views.append(org_app.GCIOrgAppPreviewPage())
    self.views.append(org_app.GCIOrgAppRecordsList())
    self.views.append(org_app.GCIOrgAppShowPage())
    self.views.append(org_app.GCIOrgAppTakePage())
    self.views.append(org_home.GCIBanOrgPost())
    self.views.append(org_home.OrgHomepage())
    self.views.append(org_score.ChooseOrganizationForOrgScorePage())
    self.views.append(org_score.OrgScoresForOrgzanizationPage())
    self.views.append(participants.MentorsListAdminPage())
    self.views.append(profile.GCIProfilePage())
    self.views.append(profile_show.GCIProfileShowPage())
    self.views.append(profile_show.GCIProfileShowAdminPage())
    self.views.append(program.ProgramPage())
    self.views.append(program.GCIProgramMessagesPage())
    self.views.append(program.TimelinePage())
    self.views.append(request.ListUserRequestsPage())
    self.views.append(request.SendRequestPage())
    self.views.append(request.ManageRequestPage())
    self.views.append(request.RespondRequestPage())
    self.views.append(student_forms.StudentFormUpload())
    self.views.append(student_forms.StudentFormDownload())
    self.views.append(students_info.StudentsInfoPage())
    self.views.append(task.TaskViewPage())
    self.views.append(task.WorkSubmissionDownload())
    self.views.append(task_list.AllOrganizationTasksPage())
    self.views.append(task_list.ChooseOrganizationPage())
    self.views.append(task_list.StudentTasksForOrganizationPage())
    self.views.append(task_list.TaskListPage())
    self.views.append(task_create.TaskCreatePage())

    # Google Appengine Tasks
    from soc.modules.gci.tasks.bulk_create import BulkCreateTask
    from soc.modules.gci.tasks.ranking_update import RankingUpdater
    from soc.modules.gci.tasks.score_update import ScoreUpdate
    from soc.modules.gci.tasks.task_update import TaskUpdate

    self.views.append(BulkCreateTask())
    self.views.append(RankingUpdater())
    self.views.append(TaskUpdate())
    self.views.append(ScoreUpdate())

  def registerWithSitemap(self):
    """Called by the server when sitemap entries should be registered.
    """

    self.core.requireUniqueService('registerWithSitemap')

    # Redesigned view registration
    for view in self.views:
      self.core.registerSitemapEntry(view.djangoURLPatterns())

  def registerWithProgramMap(self):
    """Called by the server when program_map entries should be registered.
    """

    self.core.requireUniqueService('registerWithProgramMap')

    from soc.modules.gci.models.program import GCIProgram
    program_entities = GCIProgram.all().fetch(1000)
    map = ('GCI Programs', [
        (str(e.key()), e.name) for e in program_entities])

    self.core.registerProgramEntry(map)
<commit_msg>Register the account deletion moderation view.<commit_after># Copyright 2009 the Melange authors.
#
# Licensed 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.

"""Module containing the GCI Callback.
"""

class Callback(object):
  """Callback object that handles interaction between the core.
  """

  API_VERSION = 1

  def __init__(self, core):
    """Initializes a new Callback object for the specified core.
    """

    self.core = core
    self.views = []

  def registerViews(self):
    """Instantiates all view objects.
    """
    from soc.modules.gci.views import accepted_orgs
    from soc.modules.gci.views import admin
    from soc.modules.gci.views import age_check
    from soc.modules.gci.views import all_tasks
    from soc.modules.gci.views import bulk_create
    from soc.modules.gci.views import dashboard
    from soc.modules.gci.views import delete_account
    from soc.modules.gci.views import document
    from soc.modules.gci.views import homepage
    from soc.modules.gci.views import invite
    from soc.modules.gci.views import leaderboard
    from soc.modules.gci.views import moderate_delete_account
    from soc.modules.gci.views import org_app
    from soc.modules.gci.views import org_home
    from soc.modules.gci.views import org_profile
    from soc.modules.gci.views import org_score
    from soc.modules.gci.views import participants
    from soc.modules.gci.views import profile
    from soc.modules.gci.views import profile_show
    from soc.modules.gci.views import program
    from soc.modules.gci.views import request
    from soc.modules.gci.views import student_forms
    from soc.modules.gci.views import students_info
    from soc.modules.gci.views import task
    from soc.modules.gci.views import task_list
    from soc.modules.gci.views import task_create

    self.views.append(accepted_orgs.AcceptedOrgsPage())
    self.views.append(accepted_orgs.AcceptedOrgsAdminPage())
    self.views.append(admin.DashboardPage())
    self.views.append(admin.LookupLinkIdPage())
    self.views.append(all_tasks.TaskListPage())
    self.views.append(age_check.AgeCheck())
    self.views.append(bulk_create.BulkCreate())
    self.views.append(dashboard.DashboardPage())
    self.views.append(delete_account.DeleteAccountPage())
    self.views.append(document.DocumentPage())
    self.views.append(document.EditDocumentPage())
    self.views.append(document.EventsPage())
    self.views.append(document.DocumentListPage())
    self.views.append(homepage.Homepage())
    self.views.append(invite.InvitePage())
    self.views.append(invite.ManageInvite())
    self.views.append(invite.RespondInvite())
    self.views.append(invite.ListUserInvitesPage())
    self.views.append(leaderboard.LeaderboardPage())
    self.views.append(leaderboard.StudentTasksPage())
    self.views.append(moderate_delete_account.ModerateDeleteAccountPage())
    self.views.append(org_app.GCIOrgAppEditPage())
    self.views.append(org_profile.OrgProfilePage())
    self.views.append(org_app.GCIOrgAppPreviewPage())
    self.views.append(org_app.GCIOrgAppRecordsList())
    self.views.append(org_app.GCIOrgAppShowPage())
    self.views.append(org_app.GCIOrgAppTakePage())
    self.views.append(org_home.GCIBanOrgPost())
    self.views.append(org_home.OrgHomepage())
    self.views.append(org_score.ChooseOrganizationForOrgScorePage())
    self.views.append(org_score.OrgScoresForOrgzanizationPage())
    self.views.append(participants.MentorsListAdminPage())
    self.views.append(profile.GCIProfilePage())
    self.views.append(profile_show.GCIProfileShowPage())
    self.views.append(profile_show.GCIProfileShowAdminPage())
    self.views.append(program.ProgramPage())
    self.views.append(program.GCIProgramMessagesPage())
    self.views.append(program.TimelinePage())
    self.views.append(request.ListUserRequestsPage())
    self.views.append(request.SendRequestPage())
    self.views.append(request.ManageRequestPage())
    self.views.append(request.RespondRequestPage())
    self.views.append(student_forms.StudentFormUpload())
    self.views.append(student_forms.StudentFormDownload())
    self.views.append(students_info.StudentsInfoPage())
    self.views.append(task.TaskViewPage())
    self.views.append(task.WorkSubmissionDownload())
    self.views.append(task_list.AllOrganizationTasksPage())
    self.views.append(task_list.ChooseOrganizationPage())
    self.views.append(task_list.StudentTasksForOrganizationPage())
    self.views.append(task_list.TaskListPage())
    self.views.append(task_create.TaskCreatePage())

    # Google Appengine Tasks
    from soc.modules.gci.tasks.bulk_create import BulkCreateTask
    from soc.modules.gci.tasks.ranking_update import RankingUpdater
    from soc.modules.gci.tasks.score_update import ScoreUpdate
    from soc.modules.gci.tasks.task_update import TaskUpdate

    self.views.append(BulkCreateTask())
    self.views.append(RankingUpdater())
    self.views.append(TaskUpdate())
    self.views.append(ScoreUpdate())

  def registerWithSitemap(self):
    """Called by the server when sitemap entries should be registered.
    """

    self.core.requireUniqueService('registerWithSitemap')

    # Redesigned view registration
    for view in self.views:
      self.core.registerSitemapEntry(view.djangoURLPatterns())

  def registerWithProgramMap(self):
    """Called by the server when program_map entries should be registered.
    """

    self.core.requireUniqueService('registerWithProgramMap')

    from soc.modules.gci.models.program import GCIProgram
    program_entities = GCIProgram.all().fetch(1000)
    map = ('GCI Programs', [
        (str(e.key()), e.name) for e in program_entities])

    self.core.registerProgramEntry(map)
<|endoftext|>


<commit_before>from django.conf.urls import url

from . import views

urlpatterns = [
    url(r'^$', views.index, name='index'),
    url(r'^train/$', views.train, name='train'),
    url(r'^getError/$', views.getError, name='getError'),
    url(r'^getInput/$', views.getInput, name='getInput'),
    url(r'^getReceptiveField/$', views.getReceptiveField, name='getReceptiveField'),
    url(r'^getHistogram/$', views.getHistogram, name='getHistogram'),
]
<commit_msg>Minor update urls.py<commit_after>from django.conf.urls import url

from . import views

urlpatterns = [
    url(r'^$', views.index, name='index'),
    url(r'^train/$', views.train),
    url(r'^getError/$', views.getError),
    url(r'^getInput/$', views.getInput),
    url(r'^getReceptiveField/$', views.getReceptiveField),
    url(r'^getHistogram/$', views.getHistogram),
]
<|endoftext|>


<commit_before>import threading
import logging
import Queue
import sys
import os
import time
from gpiozero import MCP3008

class Battery(object):

	def __init__(self,piManager):
		# create the object yo
		logging.debug("Create the Battery object")
		self.name = "Battery"
		self.piManager = piManager
		# Create the local queue
		self.Queue = Queue.PriorityQueue()
		# The voltage mathing
		self.vref = 2.96
		self.r1 = 51000
		self.r2 = 6800
		self.correctionFactor = 1
		# Create and start the threads
		self.listenerThread = threading.Thread(target=self.__listener, name=self.name+"-listener")
		self.listenerThread.daemon = True
		self.listenerThread.start()
		self.consumerThread = threading.Thread(target=self.__queueConsumer, name=self.name+"-consumer")
		self.consumerThread.daemon = True
		self.consumerThread.start()
		# Register for messages
		self.piManager.register(self,"SystemTest")

	def __listener(self):
		name = threading.current_thread().getName()
		logging.debug('Running the %s thread', name)
		batV = MCP3008(0).value
		while True:
			voltage = self.__calculateBatteryVoltage()
			logging.debug('Battery voltage is %s', voltage )
			time.sleep(2)

	def __calculateBatteryVoltage(self):
		# The ratio the divider is dropping our voltage by
		divider = ( (self.r1 + self.r2) / self.r2 )
		logging.debug('Voltage divider modification %s', divider)
		# expand the value read from 0 - 1 back to its real value based on vref
		temp = MCP3008(0).value
		logging.debug('Read %s from MCP3008' , temp)
		value = MCP3008(0).value * self.vref

		# the output voltage is the divider * the value
		rawVoltage = divider * value
		# our return voltage is the rawVoltage x our correction factor
		returnVoltage = rawVoltage * self.correctionFactor
		return(returnVoltage)

	def addToQueue(self,event,priority=99):
		self.Queue.put( (priority,event) )

	def __queueConsumer(self):
		name = threading.current_thread().getName()
		logging.debug("Running the "+name+" thread")
		# process queue objects as the come in run the thread forever
		while 1:
			item = self.Queue.get(True)
			task = item[1].getTask()
			logging.debug("Process Queue task "+task )<commit_msg>battery testing<commit_after>import threading
import logging
import Queue
import sys
import os
import time
from gpiozero import MCP3008

class Battery(object):

	def __init__(self,piManager):
		# create the object yo
		logging.debug("Create the Battery object")
		self.name = "Battery"
		self.piManager = piManager
		# Create the local queue
		self.Queue = Queue.PriorityQueue()
		# The voltage mathing
		self.vref = 2.96
		self.r1 = 51000
		self.r2 = 6800
		self.correctionFactor = 1
		# Create and start the threads
		self.listenerThread = threading.Thread(target=self.__listener, name=self.name+"-listener")
		self.listenerThread.daemon = True
		self.listenerThread.start()
		self.consumerThread = threading.Thread(target=self.__queueConsumer, name=self.name+"-consumer")
		self.consumerThread.daemon = True
		self.consumerThread.start()
		# Register for messages
		self.piManager.register(self,"SystemTest")

	def __listener(self):
		name = threading.current_thread().getName()
		logging.debug('Running the %s thread', name)
		batV = MCP3008(0).value
		while True:
			voltage = self.__calculateBatteryVoltage()
			logging.debug('Battery voltage is %s', voltage )
			time.sleep(2)

	def __calculateBatteryVoltage(self):
		# The ratio the divider is dropping our voltage by
		divider = float( (self.r1 + self.r2) / self.r2 )
		logging.debug('Voltage divider modification %s', divider)
		# expand the value read from 0 - 1 back to its real value based on vref
		temp = MCP3008(0).value
		logging.debug('Read %s from MCP3008' , temp)
		value = MCP3008(0).value * self.vref

		# the output voltage is the divider * the value
		rawVoltage = divider * value
		# our return voltage is the rawVoltage x our correction factor
		returnVoltage = rawVoltage * self.correctionFactor
		return(returnVoltage)

	def addToQueue(self,event,priority=99):
		self.Queue.put( (priority,event) )

	def __queueConsumer(self):
		name = threading.current_thread().getName()
		logging.debug("Running the "+name+" thread")
		# process queue objects as the come in run the thread forever
		while 1:
			item = self.Queue.get(True)
			task = item[1].getTask()
			logging.debug("Process Queue task "+task )<|endoftext|>


<commit_before>from __future__ import unicode_literals

import contextlib
import distutils.spawn
import os
import sys

import virtualenv

from pre_commit.languages import helpers
from pre_commit.util import clean_path_on_failure
from pre_commit.util import shell_escape

ENVIRONMENT_DIR = 'py_env'

class PythonEnv(helpers.Environment):
    @property
    def env_prefix(self):
        return ". '{{prefix}}{0}activate' &&".format(
            virtualenv.path_locations(
                helpers.environment_dir(ENVIRONMENT_DIR, self.language_version)
            )[-1].rstrip(os.sep) + os.sep,
            'activate',
        )

@contextlib.contextmanager
def in_env(repo_cmd_runner, language_version):
    yield PythonEnv(repo_cmd_runner, language_version)

def norm_version(version):
    version = os.path.expanduser(version)
    if os.name == 'nt':  # pragma: no cover (windows)
        if not distutils.spawn.find_executable(version):
            # expanduser introduces a leading slash
            version = version.strip('\\')
            # The default place for python on windows is:
            # C:\PythonXX\python.exe
            version = r'C:\{0}\python.exe'.format(version.replace('.', ''))
    return version

def install_environment(
        repo_cmd_runner,
        version='default',
        additional_dependencies=None,
):
    assert repo_cmd_runner.exists('setup.py')
    directory = helpers.environment_dir(ENVIRONMENT_DIR, version)

    # Install a virtualenv
    with clean_path_on_failure(repo_cmd_runner.path(directory)):
        venv_cmd = [
            sys.executable, '-m', 'virtualenv',
            '{{prefix}}{0}'.format(directory)
        ]
        if version != 'default':
            venv_cmd.extend(['-p', norm_version(version)])
        repo_cmd_runner.run(venv_cmd)
        with in_env(repo_cmd_runner, version) as env:
            env.run("cd '{prefix}' && pip install .")
            if additional_dependencies:
                env.run(
                    "cd '{prefix}' && pip install " +
                    ' '.join(
                        shell_escape(dep) for dep in additional_dependencies
                    )
                )

def run_hook(repo_cmd_runner, hook, file_args):
    with in_env(repo_cmd_runner, hook['language_version']) as env:
        return helpers.run_hook(env, hook, file_args)
<commit_msg>Remove unneeded format chunk<commit_after>from __future__ import unicode_literals

import contextlib
import distutils.spawn
import os
import sys

import virtualenv

from pre_commit.languages import helpers
from pre_commit.util import clean_path_on_failure
from pre_commit.util import shell_escape

ENVIRONMENT_DIR = 'py_env'

class PythonEnv(helpers.Environment):
    @property
    def env_prefix(self):
        return ". '{{prefix}}{0}activate' &&".format(
            virtualenv.path_locations(
                helpers.environment_dir(ENVIRONMENT_DIR, self.language_version)
            )[-1].rstrip(os.sep) + os.sep,
        )

@contextlib.contextmanager
def in_env(repo_cmd_runner, language_version):
    yield PythonEnv(repo_cmd_runner, language_version)

def norm_version(version):
    version = os.path.expanduser(version)
    if os.name == 'nt':  # pragma: no cover (windows)
        if not distutils.spawn.find_executable(version):
            # expanduser introduces a leading slash
            version = version.strip('\\')
            # The default place for python on windows is:
            # C:\PythonXX\python.exe
            version = r'C:\{0}\python.exe'.format(version.replace('.', ''))
    return version

def install_environment(
        repo_cmd_runner,
        version='default',
        additional_dependencies=None,
):
    assert repo_cmd_runner.exists('setup.py')
    directory = helpers.environment_dir(ENVIRONMENT_DIR, version)

    # Install a virtualenv
    with clean_path_on_failure(repo_cmd_runner.path(directory)):
        venv_cmd = [
            sys.executable, '-m', 'virtualenv',
            '{{prefix}}{0}'.format(directory)
        ]
        if version != 'default':
            venv_cmd.extend(['-p', norm_version(version)])
        repo_cmd_runner.run(venv_cmd)
        with in_env(repo_cmd_runner, version) as env:
            env.run("cd '{prefix}' && pip install .")
            if additional_dependencies:
                env.run(
                    "cd '{prefix}' && pip install " +
                    ' '.join(
                        shell_escape(dep) for dep in additional_dependencies
                    )
                )

def run_hook(repo_cmd_runner, hook, file_args):
    with in_env(repo_cmd_runner, hook['language_version']) as env:
        return helpers.run_hook(env, hook, file_args)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
#  prestans, a standards based WSGI compliant REST framework for Python
#  http://prestans.org
#
#  Copyright (c) 2013, Eternity Technologies Pty Ltd.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of Eternity Technologies nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL ETERNITY TECHNOLOGIES BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
__all__ = ['closure', 'templates']

import inspect
import jinja2
import os
import sys

import prestans.devel.gen.closure
import prestans.types

class Inspector(object):

    def __init__(self, model_file):
        self._model_file = model_file

    def inspect(self):

        blueprints = list()

        try:
            sys.path.append(os.path.dirname(self._model_file))
            models = __import__("models", globals(), locals(), [])
        except ImportError, import_error:
            print "\npreplate requires "
            sys.exit(1)

        for name, type_instance in models.__dict__.iteritems():
            if name.startswith('__') or inspect.ismethod(type_instance) or not inspect.isclass(type_instance) or not issubclass(type_instance, prestans.types.Model):
                continue

            blueprint = type_instance().blueprint()
            for field_name, field_blueprint in blueprint['fields'].iteritems():
                field_blueprint['map_name'] = type_instance().attribute_rewrite_map()[field_name]
            blueprints.append(blueprint)

        return blueprints

class Preplate(object):

    def __init__(self, template_type, model_file, namespace, output_directory):

        self._template_type = template_type
        self._model_file = model_file
        self._namespace = namespace
        self._output_directory = output_directory
        loader = jinja2.PackageLoader('prestans', 'devel/gen/templates')
        self._template_engine = jinja2.Environment(trim_blocks=True, loader=loader)
        self._template_engine.globals["isinstance"] = isinstance
        self._template_engine.globals["list"] = list

    def run(self):

        template = None
        if self._template_type == "closure.model":
            template = prestans.devel.gen.closure.Model(template_engine=self._template_engine, model_file=self._model_file, namespace=self._namespace, output_directory=self._output_directory)
        elif self._template_type == "closure.filter":
            template = prestans.devel.gen.closure.Filter(template_engine=self._template_engine, model_file=self._model_file, namespace=self._namespace, output_directory=self._output_directory)

        if template is None:
            return 1

        return template.run()<commit_msg>Modified line length<commit_after># -*- coding: utf-8 -*-
#
#  prestans, a standards based WSGI compliant REST framework for Python
#  http://prestans.org
#
#  Copyright (c) 2013, Eternity Technologies Pty Ltd.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of Eternity Technologies nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL ETERNITY TECHNOLOGIES BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
__all__ = ['closure', 'templates']

import inspect
import jinja2
import os
import sys

import prestans.devel.gen.closure
import prestans.types

class Inspector(object):

    def __init__(self, model_file):
        self._model_file = model_file

    def inspect(self):

        blueprints = list()

        try:
            sys.path.append(os.path.dirname(self._model_file))
            models = __import__("models", globals(), locals(), [])
        except ImportError, import_error:
            print "\npreplate requires "
            sys.exit(1)

        for name, type_instance in models.__dict__.iteritems():
            if name.startswith('__') or inspect.ismethod(type_instance) or not inspect.isclass(type_instance) or\
             not issubclass(type_instance, prestans.types.Model):
                continue

            blueprint = type_instance().blueprint()
            for field_name, field_blueprint in blueprint['fields'].iteritems():
                field_blueprint['map_name'] = type_instance().attribute_rewrite_map()[field_name]
            blueprints.append(blueprint)

        return blueprints

class Preplate(object):

    def __init__(self, template_type, model_file, namespace, output_directory):

        self._template_type = template_type
        self._model_file = model_file
        self._namespace = namespace
        self._output_directory = output_directory
        loader = jinja2.PackageLoader('prestans', 'devel/gen/templates')
        self._template_engine = jinja2.Environment(trim_blocks=True, loader=loader)
        self._template_engine.globals["isinstance"] = isinstance
        self._template_engine.globals["list"] = list

    def run(self):

        template = None
        if self._template_type == "closure.model":
            template = prestans.devel.gen.closure.Model(template_engine=self._template_engine,
                model_file=self._model_file, namespace=self._namespace, output_directory=self._output_directory)
        elif self._template_type == "closure.filter":
            template = prestans.devel.gen.closure.Filter(template_engine=self._template_engine,
                model_file=self._model_file, namespace=self._namespace, output_directory=self._output_directory)

        if template is None:
            return 1

        return template.run()<|endoftext|>


<commit_before>import plpy
import presto_client
from collections import namedtuple
from copy import copy
import time
import json
import re

# Maximum length for identifiers (e.g. table names, column names, function names)
# defined in pg_config_manual.h
PG_NAMEDATALEN = 64

JSON_TYPE_PATTERN = re.compile("^(?:row|array|map)(?![a-zA-Z])", re.IGNORECASE)

# See the document about system column names: http://www.postgresql.org/docs/9.3/static/ddl-system-columns.html
SYSTEM_COLUMN_NAMES = set(["oid", "tableoid", "xmin", "cmin", "xmax", "cmax", "ctid"])

# convert Presto query result field types to PostgreSQL types
def _pg_result_type(presto_type):
    if presto_type == "varchar":  # for old Presto
        return "varchar(255)"
    elif presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif JSON_TYPE_PATTERN.match(presto_type):
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# convert Presto table column types to PostgreSQL types
def _pg_table_type(presto_type):
    if presto_type == "varchar":  # for old Presto
        return "varchar(255)"
    elif presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif JSON_TYPE_PATTERN.match(presto_type):
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# queries can include same column name twice but tables can't.
def _rename_duplicated_column_names(column_names, rename_system_columns, where):
    renamed = []
    if rename_system_columns:
        used_names = copy(SYSTEM_COLUMN_NAMES)
    else:
        used_names = set()
    for original_name in column_names:
        name = original_name
        while name in used_names:
            name += "_"
        if name != original_name:
            if rename_system_columns and name in SYSTEM_COLUMN_NAMES:
                plpy.warning("Column %s is renamed to %s because the name in %s conflicts with PostgreSQL system column names" % \
                        (plpy.quote_ident(original_name), plpy.quote_ident(name), where))
            else:
                plpy.warning("Column %s is renamed to %s because the name appears twice in %s" % \
                        (plpy.quote_ident(original_name), plpy.quote_ident(name), where))
        used_names.add(name)
        renamed.append(name)
    return renamed

# build CREATE TEMPORARY TABLE statement
def _build_create_temp_table_sql(table_name, column_names, column_types):
    create_sql = ["create temporary table %s (\n  " % plpy.quote_ident(table_name)]

    first = True
    for column_name, column_type in zip(column_names, column_types):
        if first:
            first = False
        else:
            create_sql.append(",\n  ")

        create_sql.append(plpy.quote_ident(column_name))
        create_sql.append(" ")
        create_sql.append(column_type)

    create_sql.append("\n)")
    return ''.join(create_sql)

# build CREATE TABLE statement
def _build_create_table(schema_name, table_name, column_names, column_types, not_nulls):
    alter_sql = ["create table %s.%s (\n  " % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name))]

    first = True
    for column_name, column_type, not_null in zip(column_names, column_types, not_nulls):
        if first:
            first = False
        else:
            alter_sql.append(",\n  ")

        alter_sql.append("%s %s" % (plpy.quote_ident(column_name), column_type))

        if not_null:
            alter_sql.append(" not null")

    alter_sql.append("\n)")
    return ''.join(alter_sql)

def _get_session_time_zone():
    rows = plpy.execute("show timezone")
    return rows[0].values()[0]

def _get_session_search_path_array():
    rows = plpy.execute("select ('{' || current_setting('search_path') || '}')::text[]")
    return rows[0].values()[0]

class QueryAutoClose(object):
    def __init__(self, query):
        self.query = query
        self.column_names = None
        self.column_types = None

    def __del__(self):
        self.query.close()

class QueryAutoCloseIterator(object):
    def __init__(self, gen, query_auto_close):
        self.gen = gen
        self.query_auto_close = query_auto_close

    def __iter__(self):
        return self

    def next(self):
        return next(self.gen)

class QueryAutoCloseIteratorWithJsonConvert(QueryAutoCloseIterator):
    def __init__(self, gen, query_auto_close, json_columns):
        QueryAutoCloseIterator.__init__(self, gen, query_auto_close)
        self.json_columns = json_columns

    def next(self):
        row = next(self.gen)
        for i in self.json_columns:
            row[i] = json.dumps(row[i])
        return row

class SessionData(object):
    def __init__(self):
        self.query_auto_close = None

session = SessionData()

def start_presto_query(presto_server, presto_user, presto_catalog, presto_schema, function_name, query):
    try:
        # preserve search_path if explicitly set
        search_path = _get_session_search_path_array()
        if search_path != ['$user', 'public'] and len(search_path) > 0:
            # search_path is changed explicitly. use the first schema
            presto_schema = search_path[0]

        # start query
        client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema=presto_schema, time_zone=_get_session_time_zone())

        query = client.query(query)
        session.query_auto_close = QueryAutoClose(query)
        try:
            # result schema
            column_names = []
            column_types = []
            for column in query.columns():
                column_names.append(column.name)
                column_types.append(_pg_result_type(column.type))

            column_names = _rename_duplicated_column_names(column_names, False, "a query result")
            session.query_auto_close.column_names = column_names
            session.query_auto_close.column_types = column_types

            # CREATE TABLE for return type of the function
            type_name = function_name + "_type"
            create_type_sql = _build_create_temp_table_sql(type_name, column_names, column_types)

            # CREATE FUNCTION
            create_function_sql = \
                """
                create or replace function pg_temp.%s()
                returns setof pg_temp.%s as $$
                    import prestogres
                    return prestogres.fetch_presto_query_results()
                $$ language plpythonu
                """ % \
                (plpy.quote_ident(function_name), plpy.quote_ident(type_name))

            # run statements
            plpy.execute("drop table if exists pg_temp.%s cascade" % \
                    (plpy.quote_ident(type_name)))
            plpy.execute(create_type_sql)
            plpy.execute(create_function_sql)

            query = None

        finally:
            if query is not None:
                # close query
                session.query_auto_close = None

    except (plpy.SPIError, presto_client.PrestoException) as e:
        # PL/Python converts an exception object in Python to an error message in PostgreSQL
        # using exception class name if exc.__module__ is either of "builtins", "exceptions",
        # or "__main__". Otherwise using "module.name" format. Set __module__ = "__module__"
        # to generate pretty messages.
        e.__class__.__module__ = "__main__"
        raise

def fetch_presto_query_results():
    try:
        # TODO should throw an exception?
        #if session.query_auto_close is None:

        query_auto_close = session.query_auto_close
        session.query_auto_close = None  # close of the iterator closes query

        results = query_auto_close.query.results()
        json_columns = []
        for i, t in enumerate(query_auto_close.column_types):
            if t == "json":
                json_columns.append(i)

        if json_columns:
            return QueryAutoCloseIteratorWithJsonConvert(results, query_auto_close, json_columns)
        else:
            return QueryAutoCloseIterator(results, query_auto_close)

    except (plpy.SPIError, presto_client.PrestoException) as e:
        e.__class__.__module__ = "__main__"
        raise

Column = namedtuple("Column", ("name", "type", "nullable"))

def setup_system_catalog(presto_server, presto_user, presto_catalog, presto_schema, access_role):
    search_path = _get_session_search_path_array()
    if search_path == ['$user', 'public']:
        # search_path is default value.
        plpy.execute("set search_path to %s" % plpy.quote_ident(presto_schema))

    client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema='default')

    # get table list
    sql = "select table_schema, table_name, column_name, is_nullable, data_type" \
          " from information_schema.columns"
    columns, rows = client.run(sql)
    if rows is None:
        rows = []

    schemas = {}

    for row in rows:
        schema_name = row[0]
        table_name = row[1]
        column_name = row[2]
        is_nullable = row[3]
        column_type = row[4]

        if schema_name == "sys" or schema_name == "information_schema":
            # skip system schemas
            continue

        if len(schema_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Schema %s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), PG_NAMEDATALEN - 1))
            continue

        tables = schemas.setdefault(schema_name, {})

        if len(table_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Table %s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), PG_NAMEDATALEN - 1))
            continue

        columns = tables.setdefault(table_name, [])

        if len(column_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Column %s.%s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), \
                     plpy.quote_ident(column_name), PG_NAMEDATALEN - 1))
            continue

        columns.append(Column(column_name, column_type, is_nullable))

    # drop all schemas excepting prestogres_catalog, information_schema and pg_%
    sql = "select n.nspname as schema_name from pg_catalog.pg_namespace n" \
          " where n.nspname not in ('prestogres_catalog', 'information_schema')" \
          " and n.nspname not like 'pg_%'"
    for row in plpy.cursor(sql):
        plpy.execute("drop schema %s cascade" % plpy.quote_ident(row["schema_name"]))

    # create schema and tables
    for schema_name, tables in sorted(schemas.items(), key=lambda (k,v): k):
        try:
            plpy.execute("create schema %s" % (plpy.quote_ident(schema_name)))
        except:
            # ignore error?
            pass

        for table_name, columns in sorted(tables.items(), key=lambda (k,v): k):
            column_names = []
            column_types = []
            not_nulls = []
            for column in columns:
                column_names.append(column.name)
                column_types.append(_pg_table_type(column.type))
                not_nulls.append(not column.nullable)

            # change columns
            column_names = _rename_duplicated_column_names(column_names, True,
                    "%s.%s table" % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name)))
            create_sql = _build_create_table(schema_name, table_name, column_names, column_types, not_nulls)
            plpy.execute(create_sql)

        # grant access on the schema to the restricted user so that
        # pg_table_is_visible(reloid) used by \d of psql command returns true
        plpy.execute("grant usage on schema %s to %s" % \
                (plpy.quote_ident(schema_name), plpy.quote_ident(access_role)))
        # this SELECT privilege is unnecessary because queries against those tables
        # won't run on PostgreSQL. causing an exception is good if Prestogres has
        # a bug sending a presto query to PostgreSQL without rewriting.
        # TODO however, it's granted for now because some BI tools might check
        #      has_table_privilege. the best solution is to grant privilege but
        #      actually selecting from those tables causes an exception.
        plpy.execute("grant select on all tables in schema %s to %s" % \
                (plpy.quote_ident(schema_name), plpy.quote_ident(access_role)))

    # fake current_database() to return Presto's catalog name to be compatible with some
    # applications that use db.schema.table syntax to identify a table
    if plpy.execute("select pg_catalog.current_database()")[0].values()[0] != presto_catalog:
        plpy.execute("delete from pg_catalog.pg_proc where proname='current_database'")
        plpy.execute("create function pg_catalog.current_database() returns name as $$begin return %s::name; end$$ language plpgsql stable strict" % \
                plpy.quote_literal(presto_catalog))

<commit_msg>system column names are reserved in temporary tables as well as regular tables<commit_after>import plpy
import presto_client
from collections import namedtuple
from copy import copy
import time
import json
import re

# Maximum length for identifiers (e.g. table names, column names, function names)
# defined in pg_config_manual.h
PG_NAMEDATALEN = 64

JSON_TYPE_PATTERN = re.compile("^(?:row|array|map)(?![a-zA-Z])", re.IGNORECASE)

# See the document about system column names: http://www.postgresql.org/docs/9.3/static/ddl-system-columns.html
SYSTEM_COLUMN_NAMES = set(["oid", "tableoid", "xmin", "cmin", "xmax", "cmax", "ctid"])

# convert Presto query result field types to PostgreSQL types
def _pg_result_type(presto_type):
    if presto_type == "varchar":  # for old Presto
        return "varchar(255)"
    elif presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif JSON_TYPE_PATTERN.match(presto_type):
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# convert Presto table column types to PostgreSQL types
def _pg_table_type(presto_type):
    if presto_type == "varchar":  # for old Presto
        return "varchar(255)"
    elif presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif JSON_TYPE_PATTERN.match(presto_type):
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# queries can include same column name twice but tables can't.
def _rename_duplicated_column_names(column_names, where):
    renamed = []
    used_names = copy(SYSTEM_COLUMN_NAMES)
    for original_name in column_names:
        name = original_name
        while name in used_names:
            name += "_"
        if name != original_name:
            if name in SYSTEM_COLUMN_NAMES:
                plpy.warning("Column %s is renamed to %s because the name in %s conflicts with PostgreSQL system column names" % \
                        (plpy.quote_ident(original_name), plpy.quote_ident(name), where))
            else:
                plpy.warning("Column %s is renamed to %s because the name appears twice in %s" % \
                        (plpy.quote_ident(original_name), plpy.quote_ident(name), where))
        used_names.add(name)
        renamed.append(name)
    return renamed

# build CREATE TEMPORARY TABLE statement
def _build_create_temp_table_sql(table_name, column_names, column_types):
    create_sql = ["create temporary table %s (\n  " % plpy.quote_ident(table_name)]

    first = True
    for column_name, column_type in zip(column_names, column_types):
        if first:
            first = False
        else:
            create_sql.append(",\n  ")

        create_sql.append(plpy.quote_ident(column_name))
        create_sql.append(" ")
        create_sql.append(column_type)

    create_sql.append("\n)")
    return ''.join(create_sql)

# build CREATE TABLE statement
def _build_create_table(schema_name, table_name, column_names, column_types, not_nulls):
    alter_sql = ["create table %s.%s (\n  " % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name))]

    first = True
    for column_name, column_type, not_null in zip(column_names, column_types, not_nulls):
        if first:
            first = False
        else:
            alter_sql.append(",\n  ")

        alter_sql.append("%s %s" % (plpy.quote_ident(column_name), column_type))

        if not_null:
            alter_sql.append(" not null")

    alter_sql.append("\n)")
    return ''.join(alter_sql)

def _get_session_time_zone():
    rows = plpy.execute("show timezone")
    return rows[0].values()[0]

def _get_session_search_path_array():
    rows = plpy.execute("select ('{' || current_setting('search_path') || '}')::text[]")
    return rows[0].values()[0]

class QueryAutoClose(object):
    def __init__(self, query):
        self.query = query
        self.column_names = None
        self.column_types = None

    def __del__(self):
        self.query.close()

class QueryAutoCloseIterator(object):
    def __init__(self, gen, query_auto_close):
        self.gen = gen
        self.query_auto_close = query_auto_close

    def __iter__(self):
        return self

    def next(self):
        return next(self.gen)

class QueryAutoCloseIteratorWithJsonConvert(QueryAutoCloseIterator):
    def __init__(self, gen, query_auto_close, json_columns):
        QueryAutoCloseIterator.__init__(self, gen, query_auto_close)
        self.json_columns = json_columns

    def next(self):
        row = next(self.gen)
        for i in self.json_columns:
            row[i] = json.dumps(row[i])
        return row

class SessionData(object):
    def __init__(self):
        self.query_auto_close = None

session = SessionData()

def start_presto_query(presto_server, presto_user, presto_catalog, presto_schema, function_name, query):
    try:
        # preserve search_path if explicitly set
        search_path = _get_session_search_path_array()
        if search_path != ['$user', 'public'] and len(search_path) > 0:
            # search_path is changed explicitly. use the first schema
            presto_schema = search_path[0]

        # start query
        client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema=presto_schema, time_zone=_get_session_time_zone())

        query = client.query(query)
        session.query_auto_close = QueryAutoClose(query)
        try:
            # result schema
            column_names = []
            column_types = []
            for column in query.columns():
                column_names.append(column.name)
                column_types.append(_pg_result_type(column.type))

            column_names = _rename_duplicated_column_names(column_names, "a query result")
            session.query_auto_close.column_names = column_names
            session.query_auto_close.column_types = column_types

            # CREATE TABLE for return type of the function
            type_name = function_name + "_type"
            create_type_sql = _build_create_temp_table_sql(type_name, column_names, column_types)

            # CREATE FUNCTION
            create_function_sql = \
                """
                create or replace function pg_temp.%s()
                returns setof pg_temp.%s as $$
                    import prestogres
                    return prestogres.fetch_presto_query_results()
                $$ language plpythonu
                """ % \
                (plpy.quote_ident(function_name), plpy.quote_ident(type_name))

            # run statements
            plpy.execute("drop table if exists pg_temp.%s cascade" % \
                    (plpy.quote_ident(type_name)))
            plpy.execute(create_type_sql)
            plpy.execute(create_function_sql)

            query = None

        finally:
            if query is not None:
                # close query
                session.query_auto_close = None

    except (plpy.SPIError, presto_client.PrestoException) as e:
        # PL/Python converts an exception object in Python to an error message in PostgreSQL
        # using exception class name if exc.__module__ is either of "builtins", "exceptions",
        # or "__main__". Otherwise using "module.name" format. Set __module__ = "__module__"
        # to generate pretty messages.
        e.__class__.__module__ = "__main__"
        raise

def fetch_presto_query_results():
    try:
        # TODO should throw an exception?
        #if session.query_auto_close is None:

        query_auto_close = session.query_auto_close
        session.query_auto_close = None  # close of the iterator closes query

        results = query_auto_close.query.results()
        json_columns = []
        for i, t in enumerate(query_auto_close.column_types):
            if t == "json":
                json_columns.append(i)

        if json_columns:
            return QueryAutoCloseIteratorWithJsonConvert(results, query_auto_close, json_columns)
        else:
            return QueryAutoCloseIterator(results, query_auto_close)

    except (plpy.SPIError, presto_client.PrestoException) as e:
        e.__class__.__module__ = "__main__"
        raise

Column = namedtuple("Column", ("name", "type", "nullable"))

def setup_system_catalog(presto_server, presto_user, presto_catalog, presto_schema, access_role):
    search_path = _get_session_search_path_array()
    if search_path == ['$user', 'public']:
        # search_path is default value.
        plpy.execute("set search_path to %s" % plpy.quote_ident(presto_schema))

    client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema='default')

    # get table list
    sql = "select table_schema, table_name, column_name, is_nullable, data_type" \
          " from information_schema.columns"
    columns, rows = client.run(sql)
    if rows is None:
        rows = []

    schemas = {}

    for row in rows:
        schema_name = row[0]
        table_name = row[1]
        column_name = row[2]
        is_nullable = row[3]
        column_type = row[4]

        if schema_name == "sys" or schema_name == "information_schema":
            # skip system schemas
            continue

        if len(schema_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Schema %s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), PG_NAMEDATALEN - 1))
            continue

        tables = schemas.setdefault(schema_name, {})

        if len(table_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Table %s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), PG_NAMEDATALEN - 1))
            continue

        columns = tables.setdefault(table_name, [])

        if len(column_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Column %s.%s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), \
                     plpy.quote_ident(column_name), PG_NAMEDATALEN - 1))
            continue

        columns.append(Column(column_name, column_type, is_nullable))

    # drop all schemas excepting prestogres_catalog, information_schema and pg_%
    sql = "select n.nspname as schema_name from pg_catalog.pg_namespace n" \
          " where n.nspname not in ('prestogres_catalog', 'information_schema')" \
          " and n.nspname not like 'pg_%'"
    for row in plpy.cursor(sql):
        plpy.execute("drop schema %s cascade" % plpy.quote_ident(row["schema_name"]))

    # create schema and tables
    for schema_name, tables in sorted(schemas.items(), key=lambda (k,v): k):
        try:
            plpy.execute("create schema %s" % (plpy.quote_ident(schema_name)))
        except:
            # ignore error?
            pass

        for table_name, columns in sorted(tables.items(), key=lambda (k,v): k):
            column_names = []
            column_types = []
            not_nulls = []
            for column in columns:
                column_names.append(column.name)
                column_types.append(_pg_table_type(column.type))
                not_nulls.append(not column.nullable)

            # change columns
            column_names = _rename_duplicated_column_names(column_names,
                    "%s.%s table" % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name)))
            create_sql = _build_create_table(schema_name, table_name, column_names, column_types, not_nulls)
            plpy.execute(create_sql)

        # grant access on the schema to the restricted user so that
        # pg_table_is_visible(reloid) used by \d of psql command returns true
        plpy.execute("grant usage on schema %s to %s" % \
                (plpy.quote_ident(schema_name), plpy.quote_ident(access_role)))
        # this SELECT privilege is unnecessary because queries against those tables
        # won't run on PostgreSQL. causing an exception is good if Prestogres has
        # a bug sending a presto query to PostgreSQL without rewriting.
        # TODO however, it's granted for now because some BI tools might check
        #      has_table_privilege. the best solution is to grant privilege but
        #      actually selecting from those tables causes an exception.
        plpy.execute("grant select on all tables in schema %s to %s" % \
                (plpy.quote_ident(schema_name), plpy.quote_ident(access_role)))

    # fake current_database() to return Presto's catalog name to be compatible with some
    # applications that use db.schema.table syntax to identify a table
    if plpy.execute("select pg_catalog.current_database()")[0].values()[0] != presto_catalog:
        plpy.execute("delete from pg_catalog.pg_proc where proname='current_database'")
        plpy.execute("create function pg_catalog.current_database() returns name as $$begin return %s::name; end$$ language plpgsql stable strict" % \
                plpy.quote_literal(presto_catalog))

<|endoftext|>


<commit_before>from procasl import datasets
import numpy as np
from nose import with_setup
from nose.tools import assert_equal, assert_true, assert_not_equal
from numpy.testing import assert_array_equal

from nilearn.datasets import utils
from nilearn.datasets.tests import test_utils as tst

def setup_mock():
    return tst.setup_mock(utils, mri)

def teardown_mock():
    return tst.teardown_mock(utils, mri)

@with_setup(setup_mock, teardown_mock)
@with_setup(tst.setup_tmpdata, tst.teardown_tmpdata)
def test_fetch_kirby():
    # First session, all subjects
    kirby = datasets.fetch_kirby(subjects=range(21), data_dir=tst.tmpdir,
                                 verbose=1)
    assert_equal(len(tst.mock_url_request.urls), 21)
    assert_equal(len(kirby.anat), 21)
    assert_equal(len(kirby.asl), 21)
    assert_equal(len(kirby.m0), 21)
    assert_true(np.all(np.asarray(kirby.session) == 1))

    # Both sessions, 12 subjects
    tst.mock_url_request.reset()
    kirby = mri.fetch_csd(data_dir=tst.tmpdir, sessions=[1, 2],
                        subjects=range(12), verbose=0)
    # Session 1 has already been downloaded
    assert_equal(len(tst.mock_url_request.urls), 24)
    assert_equal(len(kirby.anat), 24)
    assert_equal(len(kirby.asl), 24)
    assert_equal(len(kirby.m0), 24)
    s = np.asarray(kirby.session)
    assert_true(np.all(s[:12] == 1))
    assert_true(np.all(s[12:24] == 2))
    assert_not_equal(kirby.description, '')
<commit_msg>fix url number and remove description test<commit_after>from procasl import datasets
import numpy as np
from nose import with_setup
from nose.tools import assert_equal, assert_true, assert_not_equal
from numpy.testing import assert_array_equal

from nilearn.datasets import utils
from nilearn.datasets.tests import test_utils as tst

def setup_mock():
    return tst.setup_mock(utils, datasets)

def teardown_mock():
    return tst.teardown_mock(utils, datasets)

@with_setup(setup_mock, teardown_mock)
@with_setup(tst.setup_tmpdata, tst.teardown_tmpdata)
def test_fetch_kirby():
    # First session, all subjects
    kirby = datasets.fetch_kirby(subjects=range(21), data_dir=tst.tmpdir,
                                 verbose=1)
    assert_equal(len(tst.mock_url_request.urls), 21)
    assert_equal(len(kirby.anat), 21)
    assert_equal(len(kirby.asl), 21)
    assert_equal(len(kirby.m0), 21)
    assert_true(np.all(np.asarray(kirby.session) == 1))

    # Both sessions, 12 subjects
    tst.mock_url_request.reset()
    kirby = datasets.fetch_kirby(data_dir=tst.tmpdir, sessions=[1, 2],
                                 subjects=range(12), verbose=0)
    # Session 1 has already been downloaded
    assert_equal(len(tst.mock_url_request.urls), 12)
    assert_equal(len(kirby.anat), 24)
    assert_equal(len(kirby.asl), 24)
    assert_equal(len(kirby.m0), 24)
    s = np.asarray(kirby.session)
    assert_true(np.all(s[:12] == 1))
    assert_true(np.all(s[12:24] == 2))
<|endoftext|>


<commit_before>#!/usr/bin/env python3.3
# -*- coding: utf8 -*-
#
# Toolkit script for all types of functions, needed throughout the project
#
# Copyright (c) 2015	NorthernSec
# Copyright (c)	2015	Pieter-Jan Moreels
# This software is licensed under the Original BSD License

# Imports
import os
import re
from urllib.request import urlopen
import urllib
import json
from dateutil import tz
import dateutil.parser

from lib.Config import Configuration

# string to dict
def make_dict(s):
  # break into list of keys and values
  chunks = re.split('\s*(\w+\:)\s*',s)
  res={}
  # work backwards in value, key pairs
  args=[reversed(chunks)]*2
  for value,key in zip(*args):
    key=key.rstrip(':')
    if value:
      #add to current result-dict
      res[key]=value
    else:
      #start a higher-level result-dict
      res={key:res}
  return res

def writeJson(file, data):
  if os.path.exists(file):
    os.remove(file)
  with open(file, 'w') as dump:
    json.dump(data, dump, indent=2)

def toLocalTime(utc):
  timezone = tz.tzlocal()
  utc = dateutil.parser.parse(utc)
  output = utc.astimezone(timezone)
  output = output.strftime('%d-%m-%Y - %H:%M')
  return output

def toHuman(cpe):
  cpe = cpe[7:]
  result = cpe.split(':')[0] + " - "
  for c in cpe.split(':')[1:]:
    c = c.replace(':', ' ')
    c = c.replace('_', ' ')
    result += (" %s" %(c))
  result = result.title()
  return result

def filter(vulns, exploitsOnly=False, filters={}):
  if 'access' in filters and filters['access']!=0:
    vulns=[x for x in vulns if x['access']['vector'] in filters['access']]
  if 'impact' in filters and filters['impact']!=0:
    for fil in filters['impact']: vulns=[x for x in vulns if x['impact'][fil] !='NONE']
  # exploits only
  if exploitsOnly: vulns=[x for x in vulns if ('map_cve_exploitdb' in x or 'map_cve_msf' in x)]
  return vulns

def enhance(systems, exploitsOnly=False, filters={}):
  host,port=Configuration.getCVESearch()
  # deal with filters
  for system in systems:
    cpe=system['cpes'] if 'cpes' in system else None
    if cpe:
      cpes=[]
      for c in cpe:
        cEnc=urllib.parse.quote_plus(c).lower()
        data = (urlopen('http://%s:%s/api/cvefor/%s'%(host,port,cEnc)).read()).decode('utf8')
        vulns=json.loads(str(data))
        # apply filter
        filter(vulns, exploitsOnly, filters)
        # done
        cpes.append({'cpe':c, 'cves':vulns})
      system['cpes']=cpes
      #TODO get possible dpe info and store in dpe
    for service in system['services']:
      if 'cpe' in service:
        c=urllib.parse.quote_plus(service['cpe']).lower()
        data = (urlopen('http://%s:%s/api/cvefor/%s'%(host,port,c)).read()).decode('utf8')
        vulns=json.loads(str(data))
        # apply filter
        filter(vulns, exploitsOnly, filters)
        service['cves']=vulns

      #TODO get dpe info for service
  return systems

def splitByLength(string, size=50):
  return [string[i:i+size] for i in range(0,len(string),size)]
<commit_msg>bugfix<commit_after>#!/usr/bin/env python3.3
# -*- coding: utf8 -*-
#
# Toolkit script for all types of functions, needed throughout the project
#
# Copyright (c) 2015	NorthernSec
# Copyright (c)	2015	Pieter-Jan Moreels
# This software is licensed under the Original BSD License

# Imports
import os
import re
from urllib.request import urlopen
import urllib
import json
from dateutil import tz
import dateutil.parser

from lib.Config import Configuration

# string to dict
def make_dict(s):
  # break into list of keys and values
  chunks = re.split('\s*(\w+\:)\s*',s)
  res={}
  # work backwards in value, key pairs
  args=[reversed(chunks)]*2
  for value,key in zip(*args):
    key=key.rstrip(':')
    if value:
      #add to current result-dict
      res[key]=value
    else:
      #start a higher-level result-dict
      res={key:res}
  return res

def writeJson(file, data):
  if os.path.exists(file):
    os.remove(file)
  with open(file, 'w') as dump:
    json.dump(data, dump, indent=2)

def toLocalTime(utc):
  timezone = tz.tzlocal()
  utc = dateutil.parser.parse(utc)
  output = utc.astimezone(timezone)
  output = output.strftime('%d-%m-%Y - %H:%M')
  return output

def toHuman(cpe):
  cpe = cpe[7:]
  result = cpe.split(':')[0] + " - "
  for c in cpe.split(':')[1:]:
    c = c.replace(':', ' ')
    c = c.replace('_', ' ')
    result += (" %s" %(c))
  result = result.title()
  return result

def filter(vulns, exploitsOnly=False, filters={}):
  if 'access' in filters and len(filters['access'])!=0:
    vulns=[x for x in vulns if x['access']['vector'] in filters['access']]
  if 'impact' in filters and len(filters['impact'])!=0:
    for fil in filters['impact']: vulns=[x for x in vulns if x['impact'][fil] !='NONE']
  # exploits only
  if exploitsOnly: vulns=[x for x in vulns if ('map_cve_exploitdb' in x or 'map_cve_msf' in x)]
  return vulns

def enhance(systems, exploitsOnly=False, filters={}):
  host,port=Configuration.getCVESearch()
  # deal with filters
  for system in systems:
    cpe=system['cpes'] if 'cpes' in system else None
    if cpe:
      cpes=[]
      for c in cpe:
        cEnc=urllib.parse.quote_plus(c).lower()
        data = (urlopen('http://%s:%s/api/cvefor/%s'%(host,port,cEnc)).read()).decode('utf8')
        vulns=json.loads(str(data))
        # apply filter
        vulns=filter(vulns, exploitsOnly, filters)
        # done
        cpes.append({'cpe':c, 'cves':vulns})
      system['cpes']=cpes
      #TODO get possible dpe info and store in dpe
    for service in system['services']:
      if 'cpe' in service:
        c=urllib.parse.quote_plus(service['cpe']).lower()
        data = (urlopen('http://%s:%s/api/cvefor/%s'%(host,port,c)).read()).decode('utf8')
        vulns=json.loads(str(data))
        # apply filter
        vulns=filter(vulns, exploitsOnly, filters)
        service['cves']=vulns

      #TODO get dpe info for service
  return systems

def splitByLength(string, size=50):
  return [string[i:i+size] for i in range(0,len(string),size)]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2015 ARM Limited
#
# Licensed under the Apache License, Version 2.0
# See LICENSE file for details.

from .exporter import Exporter

class gdb_definitions():

    SUPPORTED_MCUS = {
        'K64F': {
        }
    }

class GDB(Exporter):
    def __init__(self, workspace, env_settings):
        self.workspace = workspace
        self.env_settings = env_settings

    def export_project(self, data, env_settings):
        # for native debugging, no command files are necessary
        return None, []

    def supports_target(self, target):
        # !!! TODO: should be yes for native targets
        return False

    def is_supported_by_default(self, target):
        # does not require additional information
        return True

class ARMNoneEABIGDB(GDB):
    SUPPORTED = gdb_definitions.SUPPORTED_MCUS

    generated_project = {
        'path': '',
        'files': {
            'startupfile': '',
        }
    }

    def __init__(self, workspace, env_settings):
        super(ARMNoneEABIGDB, self).__init__(workspace, env_settings)

    def export_project(self):
        generated_projects = copy.deepcopy(self.generated_project)
        expanded_dic = self.workspace.copy()

        # !!! TODO: store and read settings from gdb_definitions
        expanded_dic['gdb_server_port'] = 3333

        project_path, startupfile = self.gen_file_jinja(
            'gdb.tmpl', expanded_dic, '%s.gdbstartup' % expanded_dic['name'], expanded_dic['output_dir']['path'])
        generated_projects['path'] = project_path
        generated_projects['files']['startupfile'] = startupfile
        return generated_projects

    def get_generated_project_files(self):
        return {'path': self.workspace['path'], 'files': [self.workspace['files']['startupfile']]}

    def supports_target(self, target):
        return target in self.SUPPORTED

    def is_supported_by_default(self, target):
        # does not require additional information
        return True
<commit_msg>gdb - fix import copy<commit_after>#!/usr/bin/env python
# Copyright 2015 ARM Limited
#
# Licensed under the Apache License, Version 2.0
# See LICENSE file for details.

import copy

from .exporter import Exporter

class gdb_definitions():

    SUPPORTED_MCUS = {
        'K64F': {
        }
    }

class GDB(Exporter):
    def __init__(self, workspace, env_settings):
        self.workspace = workspace
        self.env_settings = env_settings

    def export_project(self, data, env_settings):
        # for native debugging, no command files are necessary
        return None, []

    def supports_target(self, target):
        # !!! TODO: should be yes for native targets
        return False

    def is_supported_by_default(self, target):
        # does not require additional information
        return True

class ARMNoneEABIGDB(GDB):
    SUPPORTED = gdb_definitions.SUPPORTED_MCUS

    generated_project = {
        'path': '',
        'files': {
            'startupfile': '',
        }
    }

    def __init__(self, workspace, env_settings):
        super(ARMNoneEABIGDB, self).__init__(workspace, env_settings)

    def export_project(self):
        generated_projects = copy.deepcopy(self.generated_project)
        expanded_dic = self.workspace.copy()

        # !!! TODO: store and read settings from gdb_definitions
        expanded_dic['gdb_server_port'] = 3333

        project_path, startupfile = self.gen_file_jinja(
            'gdb.tmpl', expanded_dic, '%s.gdbstartup' % expanded_dic['name'], expanded_dic['output_dir']['path'])
        generated_projects['path'] = project_path
        generated_projects['files']['startupfile'] = startupfile
        return generated_projects

    def get_generated_project_files(self):
        return {'path': self.workspace['path'], 'files': [self.workspace['files']['startupfile']]}

    def supports_target(self, target):
        return target in self.SUPPORTED

    def is_supported_by_default(self, target):
        # does not require additional information
        return True
<|endoftext|>


<commit_before>import json
import re
import base64
from socket import MSG_PEEK

import gevent.pywsgi
from gevent.queue import Queue
from gevent.pool import Group
from gevent.socket import create_connection

from gservice.config import Option
from gservice.core import Service

from ws4py.server.geventserver import UpgradableWSGIHandler
from ws4py.client.geventclient import WebSocketClient
from ws4py.server.wsgi.middleware import WebSocketUpgradeMiddleware

def encode_data_packet(conn_id, data):
    return ''.join([chr(conn_id), data])

def decode_data_packet(data):
    return data[0], str(data[1:])

class CodependentGroup(Group):
    """
    A greenlet group that will kill all greenlets if a single one dies.
    """
    def discard(self, greenlet):
        super(CodependentGroup, self).discard(greenlet)
        if not hasattr(self, '_killing'):
            self._killing = True
            gevent.spawn(self.kill)

class TunnelBroker(Service):
    port = Option('port', default=80)
    address = Option('address', default='127.0.0.1')

    def __init__(self):
        self.frontend = BrokerFrontend(self)
        self.add_service(self.frontend)

        self.tunnels = {}

    #def do_start(self):
    #    gevent.spawn(self.visual_heartbeat)

    def visual_heartbeat(self):
        while True:
            print "."
            gevent.sleep(1)

    def open(self, name):
        tunnel = Tunnel()
        self.tunnels[name] = tunnel
        return tunnel

    def close(self, name):
        tunnel = self.tunnels.pop(name)
        tunnel.close()

    def lookup(self, name):
        return self.tunnels[name]

class BrokerFrontend(gevent.pywsgi.WSGIServer):
    hostname = Option('hostname', default="localtunnel.com")

    def __init__(self, broker):
        gevent.pywsgi.WSGIServer.__init__(self, (broker.address, broker.port))
        self.broker = broker

    def handle(self, socket, address):
        hostname = ''
        hostheader = re.compile('host: ([^\(\);:,<>]+)', re.I)
        for n in [128, 256, 512]:
            bytes = socket.recv(n, MSG_PEEK)
            if not bytes:
                break
            for line in bytes.split('\r\n'):
                match = hostheader.match(line)
                if match:
                    hostname = match.group(1)
            if hostname:
                break
        hostname = hostname.split(':')[0]
        if hostname.endswith('.%s' % self.hostname):
            handler = ProxyHandler(socket, hostname, self.broker)
            handler.handle()
        else:
            handler = TunnelHandler(socket, address, self.broker)
            handler.handle()

class ProxyHandler(object):
    def __init__(self, socket, hostname, broker):
        self.socket = socket
        self.hostname = hostname
        self.broker = broker

    def handle(self):
        tunnel = self.broker.lookup(self.hostname)
        conn = tunnel.create_connection()
        group = CodependentGroup([
            gevent.spawn(self._proxy_in, self.socket, conn),
            gevent.spawn(self._proxy_out, conn, self.socket),
        ])
        gevent.joinall(group.greenlets)
        try:
            self.socket.shutdown(0)
            self.socket.close()
        except:
            pass

    def _proxy_in(self, socket, conn):
        while True:
            data = socket.recv(2048)
            if not data:
                return
            conn.send(data)

    def _proxy_out(self, conn, socket):
        while True:
            data = conn.recv()
            if data is None:
                return
            socket.send(data)

class TunnelHandler(UpgradableWSGIHandler):
    def __init__(self, socket, address, broker):
        UpgradableWSGIHandler.__init__(self, socket, address, broker.frontend)
        self.server.application = WebSocketUpgradeMiddleware(
            self.handle_websocket, self.handle_http)
        self.broker = broker

    def handle_http(self, environ, start_response):
        start_response("200 ok", [])
        return ['hello']

    def handle_websocket(self, websocket, environ):
        name = environ.get('PATH_INFO', '').split('/')[-1]
        tunnel = self.broker.open(name)
        group = CodependentGroup([
            gevent.spawn(self._tunnel_in, tunnel, websocket),
            gevent.spawn(self._tunnel_out, websocket, tunnel),
        ])
        gevent.joinall(group.greenlets)
        self.broker.close(name)
        websocket.close()

    def _tunnel_in(self, tunnel, websocket):
        for type, msg in tunnel:
            if type == 'binary':
                websocket.send(msg, binary=True)
            elif type == 'text':
                websocket.send(msg)

    def _tunnel_out(self, websocket, tunnel):
        while True:
            msg = websocket.receive(msg_obj=True)
            if msg is None:
                return
            if msg.is_text:
                tunnel.dispatch(message=str(msg))
            elif msg.is_binary:
                tunnel.dispatch(data=msg.data)

class Tunnel(object):
    def __init__(self):
        self.connections = {}
        self.tunnelq = Queue()

    def create_connection(self):
        id = 0
        while id in self.connections.keys():
            id += 1
            id %= 2**31
        conn = ConnectionProxy(id, self)
        self.connections[id] = conn
        return conn

    def close(self):
        for conn_id in self.connections:
            conn = self.connections.pop(conn_id)
            conn.close()

    def __iter__(self):
        return self

    def next(self):
        return self.tunnelq.get()

    def dispatch(self, message=None, data=None):
        """ From the tunnel (server) to the proxy (client) """
        if message:
            try:
                parsed = json.loads(message)
            except ValueError:
                raise
            conn_id, event = parsed[0:2]
            if conn_id not in self.connections:
                return
            if event == 'closed':
                conn = self.connections.pop(conn_id)
                conn.close()
        elif data:
            conn_id, data = decode_data_packet(data)
            self.connections[conn_id].recvq.put(data)

class ConnectionProxy(object):
    def __init__(self, id, tunnel):
        self.tunnel = tunnel
        self.id = id
        self.recvq = Queue()
        self.send(open=True)

    def recv(self):
        return self.recvq.get()

    def send(self, data=None, open=None):
        """ From the proxy (client) to the tunnel (server) """
        if open is True:
            msg = [self.id, 'open']
            self.tunnel.tunnelq.put(('text', json.dumps(msg)))
        elif open is False:
            msg = [self.id, 'closed']
            self.tunnel.tunnelq.put(('text', json.dumps(msg)))
        else:
            data = encode_data_packet(self.id, data)
            self.tunnel.tunnelq.put(('binary', data))

    def close(self):
        self.recvq.put(None)
        self.send(open=False)

# XXX: Yeah, obviously this should not live in server.py
class TunnelClient(Service):
    client_port = Option('client_port')
    server_port = Option('port')
    hostname = Option('hostname', default="localtunnel.com")

    def __init__(self):
        self.ws = WebSocketClient('http://%s:%s/t/test.localtunnel.local' %
                    (self.hostname, self.server_port))
        self.connections = {}

    def do_start(self):
        self.ws.connect()
        gevent.spawn(self.listen)
        #gevent.spawn(self.visual_heartbeat)

    def visual_heartbeat(self):
        while True:
            print "."
            gevent.sleep(1)

    def listen(self):
        while True:
            msg = self.ws.receive(msg_obj=True)
            if msg is None:
                print "Trying to stop"
                self.stop()
            if msg.is_text:
                parsed = json.loads(str(msg))
                conn_id, event = parsed[0:2]
                if event == 'open':
                    self.local_open(conn_id)
                elif event == 'closed':
                    self.local_close(conn_id)
            elif msg.is_binary:
                conn_id, data = decode_data_packet(msg.data)
                self.local_send(conn_id, data)

    def local_open(self, conn_id):
        socket = create_connection(('127.0.0.1', self.client_port))
        self.connections[conn_id] = socket
        gevent.spawn(self.local_recv, conn_id)

    def local_close(self, conn_id):
        socket = self.connections.pop(conn_id)
        try:
            socket.shutdown(0)
            socket.close()
        except:
            pass

    def local_send(self, conn_id, data):
        self.connections[conn_id].send(data)

    def local_recv(self, conn_id):
        while True:
            data = self.connections[conn_id].recv(4096)
            if not data:
                break
            self.tunnel_send(conn_id, data)
        self.tunnel_send(conn_id, open=False)

    def tunnel_send(self, conn_id, data=None, open=None):
        if open is False:
            msg = [conn_id, 'closed']
            self.ws.send(json.dumps(msg))
        elif data:
            msg = encode_data_packet(conn_id, data)
            self.ws.send(msg, binary=True)
        else:
            return<commit_msg>bind to 0.0.0.0<commit_after>import json
import re
import base64
from socket import MSG_PEEK

import gevent.pywsgi
from gevent.queue import Queue
from gevent.pool import Group
from gevent.socket import create_connection

from gservice.config import Option
from gservice.core import Service

from ws4py.server.geventserver import UpgradableWSGIHandler
from ws4py.client.geventclient import WebSocketClient
from ws4py.server.wsgi.middleware import WebSocketUpgradeMiddleware

def encode_data_packet(conn_id, data):
    return ''.join([chr(conn_id), data])

def decode_data_packet(data):
    return data[0], str(data[1:])

class CodependentGroup(Group):
    """
    A greenlet group that will kill all greenlets if a single one dies.
    """
    def discard(self, greenlet):
        super(CodependentGroup, self).discard(greenlet)
        if not hasattr(self, '_killing'):
            self._killing = True
            gevent.spawn(self.kill)

class TunnelBroker(Service):
    port = Option('port', default=80)
    address = Option('address', default='0.0.0.0')

    def __init__(self):
        self.frontend = BrokerFrontend(self)
        self.add_service(self.frontend)

        self.tunnels = {}

    #def do_start(self):
    #    gevent.spawn(self.visual_heartbeat)

    def visual_heartbeat(self):
        while True:
            print "."
            gevent.sleep(1)

    def open(self, name):
        tunnel = Tunnel()
        self.tunnels[name] = tunnel
        return tunnel

    def close(self, name):
        tunnel = self.tunnels.pop(name)
        tunnel.close()

    def lookup(self, name):
        return self.tunnels[name]

class BrokerFrontend(gevent.pywsgi.WSGIServer):
    hostname = Option('hostname', default="localtunnel.com")

    def __init__(self, broker):
        gevent.pywsgi.WSGIServer.__init__(self, (broker.address, broker.port))
        self.broker = broker

    def handle(self, socket, address):
        hostname = ''
        hostheader = re.compile('host: ([^\(\);:,<>]+)', re.I)
        for n in [128, 256, 512]:
            bytes = socket.recv(n, MSG_PEEK)
            if not bytes:
                break
            for line in bytes.split('\r\n'):
                match = hostheader.match(line)
                if match:
                    hostname = match.group(1)
            if hostname:
                break
        hostname = hostname.split(':')[0]
        if hostname.endswith('.%s' % self.hostname):
            handler = ProxyHandler(socket, hostname, self.broker)
            handler.handle()
        else:
            handler = TunnelHandler(socket, address, self.broker)
            handler.handle()

class ProxyHandler(object):
    def __init__(self, socket, hostname, broker):
        self.socket = socket
        self.hostname = hostname
        self.broker = broker

    def handle(self):
        tunnel = self.broker.lookup(self.hostname)
        conn = tunnel.create_connection()
        group = CodependentGroup([
            gevent.spawn(self._proxy_in, self.socket, conn),
            gevent.spawn(self._proxy_out, conn, self.socket),
        ])
        gevent.joinall(group.greenlets)
        try:
            self.socket.shutdown(0)
            self.socket.close()
        except:
            pass

    def _proxy_in(self, socket, conn):
        while True:
            data = socket.recv(2048)
            if not data:
                return
            conn.send(data)

    def _proxy_out(self, conn, socket):
        while True:
            data = conn.recv()
            if data is None:
                return
            socket.send(data)

class TunnelHandler(UpgradableWSGIHandler):
    def __init__(self, socket, address, broker):
        UpgradableWSGIHandler.__init__(self, socket, address, broker.frontend)
        self.server.application = WebSocketUpgradeMiddleware(
            self.handle_websocket, self.handle_http)
        self.broker = broker

    def handle_http(self, environ, start_response):
        start_response("200 ok", [])
        return ['hello']

    def handle_websocket(self, websocket, environ):
        name = environ.get('PATH_INFO', '').split('/')[-1]
        tunnel = self.broker.open(name)
        group = CodependentGroup([
            gevent.spawn(self._tunnel_in, tunnel, websocket),
            gevent.spawn(self._tunnel_out, websocket, tunnel),
        ])
        gevent.joinall(group.greenlets)
        self.broker.close(name)
        websocket.close()

    def _tunnel_in(self, tunnel, websocket):
        for type, msg in tunnel:
            if type == 'binary':
                websocket.send(msg, binary=True)
            elif type == 'text':
                websocket.send(msg)

    def _tunnel_out(self, websocket, tunnel):
        while True:
            msg = websocket.receive(msg_obj=True)
            if msg is None:
                return
            if msg.is_text:
                tunnel.dispatch(message=str(msg))
            elif msg.is_binary:
                tunnel.dispatch(data=msg.data)

class Tunnel(object):
    def __init__(self):
        self.connections = {}
        self.tunnelq = Queue()

    def create_connection(self):
        id = 0
        while id in self.connections.keys():
            id += 1
            id %= 2**31
        conn = ConnectionProxy(id, self)
        self.connections[id] = conn
        return conn

    def close(self):
        for conn_id in self.connections:
            conn = self.connections.pop(conn_id)
            conn.close()

    def __iter__(self):
        return self

    def next(self):
        return self.tunnelq.get()

    def dispatch(self, message=None, data=None):
        """ From the tunnel (server) to the proxy (client) """
        if message:
            try:
                parsed = json.loads(message)
            except ValueError:
                raise
            conn_id, event = parsed[0:2]
            if conn_id not in self.connections:
                return
            if event == 'closed':
                conn = self.connections.pop(conn_id)
                conn.close()
        elif data:
            conn_id, data = decode_data_packet(data)
            self.connections[conn_id].recvq.put(data)

class ConnectionProxy(object):
    def __init__(self, id, tunnel):
        self.tunnel = tunnel
        self.id = id
        self.recvq = Queue()
        self.send(open=True)

    def recv(self):
        return self.recvq.get()

    def send(self, data=None, open=None):
        """ From the proxy (client) to the tunnel (server) """
        if open is True:
            msg = [self.id, 'open']
            self.tunnel.tunnelq.put(('text', json.dumps(msg)))
        elif open is False:
            msg = [self.id, 'closed']
            self.tunnel.tunnelq.put(('text', json.dumps(msg)))
        else:
            data = encode_data_packet(self.id, data)
            self.tunnel.tunnelq.put(('binary', data))

    def close(self):
        self.recvq.put(None)
        self.send(open=False)

# XXX: Yeah, obviously this should not live in server.py
class TunnelClient(Service):
    client_port = Option('client_port')
    server_port = Option('port')
    hostname = Option('hostname', default="localtunnel.com")

    def __init__(self):
        self.ws = WebSocketClient('http://%s:%s/t/test.localtunnel.local' %
                    (self.hostname, self.server_port))
        self.connections = {}

    def do_start(self):
        self.ws.connect()
        gevent.spawn(self.listen)
        #gevent.spawn(self.visual_heartbeat)

    def visual_heartbeat(self):
        while True:
            print "."
            gevent.sleep(1)

    def listen(self):
        while True:
            msg = self.ws.receive(msg_obj=True)
            if msg is None:
                print "Trying to stop"
                self.stop()
            if msg.is_text:
                parsed = json.loads(str(msg))
                conn_id, event = parsed[0:2]
                if event == 'open':
                    self.local_open(conn_id)
                elif event == 'closed':
                    self.local_close(conn_id)
            elif msg.is_binary:
                conn_id, data = decode_data_packet(msg.data)
                self.local_send(conn_id, data)

    def local_open(self, conn_id):
        socket = create_connection(('127.0.0.1', self.client_port))
        self.connections[conn_id] = socket
        gevent.spawn(self.local_recv, conn_id)

    def local_close(self, conn_id):
        socket = self.connections.pop(conn_id)
        try:
            socket.shutdown(0)
            socket.close()
        except:
            pass

    def local_send(self, conn_id, data):
        self.connections[conn_id].send(data)

    def local_recv(self, conn_id):
        while True:
            data = self.connections[conn_id].recv(4096)
            if not data:
                break
            self.tunnel_send(conn_id, data)
        self.tunnel_send(conn_id, open=False)

    def tunnel_send(self, conn_id, data=None, open=None):
        if open is False:
            msg = [conn_id, 'closed']
            self.ws.send(json.dumps(msg))
        elif data:
            msg = encode_data_packet(conn_id, data)
            self.ws.send(msg, binary=True)
        else:
            return<|endoftext|>