<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from bs4 import BeautifulSoup
from django.conf import settings
from django.contrib.auth.models import AnonymousUser
from django.core.mail import EmailMessage, EmailMultiAlternatives
from django.db import models
from django.db.models import Q
from django.http.request import HttpRequest
from django.template import Context, engines
from django.utils.translation import ugettext_lazy as _, override as translation_override
from django.utils.six.moves.urllib.parse import urlparse

from post_office import mail
from post_office.models import Email as OriginalEmail, EmailTemplate

from filer.fields.file import FilerFileField

from shop import app_settings
from .customer import CustomerModel

class Email(OriginalEmail):
    """
    Patch class `post_office.models.Email` which fixes https://github.com/ui/django-post_office/issues/116
    and additionally converts HTML messages into plain text messages.
    """
    class Meta:
        proxy = True

    def email_message(self, connection=None):
        if self.template is not None:
            render_language = self.context.get('render_language', settings.LANGUAGE_CODE)
            context = Context(self.context)
            with translation_override(render_language):
                subject = engines['django'].from_string(self.template.subject).render(context)
                message = engines['django'].from_string(self.template.content).render(context)
                html_message = engines['django'].from_string(self.template.html_content).render(context)
        else:
            subject = self.subject
            message = self.message
            html_message = self.html_message

        if html_message:
            if not message:
                message = BeautifulSoup(html_message).text
            mailmsg = EmailMultiAlternatives(
                subject=subject, body=message, from_email=self.from_email,
                to=self.to, bcc=self.bcc, cc=self.cc,
                connection=connection, headers=self.headers)
            mailmsg.attach_alternative(html_message, 'text/html')
        else:
            mailmsg = EmailMessage(
                subject=subject, body=message, from_email=self.from_email,
                to=self.to, bcc=self.bcc, cc=self.cc,
                connection=connection, headers=self.headers)

        for attachment in self.attachments.all():
            mailmsg.attach(attachment.name, attachment.file.read())
        return mailmsg

class EmailManager(models.Manager):
    def get_queryset(self):
        return Email.objects.get_queryset()

OriginalEmail.add_to_class('objects', EmailManager())

class Notification(models.Model):
    """
    A task executed on receiving a signal.
    """
    name = models.CharField(max_length=50, verbose_name=_("Name"))
    transition_target = models.CharField(max_length=50, verbose_name=_("Event"))
    mail_to = models.PositiveIntegerField(verbose_name=_("Mail to"), null=True,
                                          blank=True, default=None)
    mail_template = models.ForeignKey(EmailTemplate, verbose_name=_("Template"),
                                      limit_choices_to=Q(language__isnull=True) | Q(language=''))

    class Meta:
        app_label = 'shop'
        verbose_name = _("Notification")
        verbose_name_plural = _("Notifications")
        ordering = ('transition_target', 'mail_to')

    def get_recipient(self, order):
        if self.mail_to is None:
            return
        if self.mail_to == 0:
            return order.customer.email
        return CustomerModel.objects.get(pk=self.mail_to).email

class NotificationAttachment(models.Model):
    notification = models.ForeignKey(Notification)
    attachment = FilerFileField(null=True, blank=True, related_name='email_attachment')

    class Meta:
        app_label = 'shop'

class EmulateHttpRequest(HttpRequest):
    """
    Use this class to emulate a HttpRequest object, when templates must be rendered
    asynchronously, for instance when an email must be generated out of an Order object.
    """
    def __init__(self, customer, stored_request):
        super(EmulateHttpRequest, self).__init__()
        parsedurl = urlparse(stored_request.get('absolute_base_uri'))
        self.path = self.path_info = parsedurl.path
        self.environ = {}
        self.META['PATH_INFO'] = parsedurl.path
        self.META['SCRIPT_NAME'] = ''
        self.META['HTTP_HOST'] = parsedurl.netloc
        self.META['HTTP_X_FORWARDED_PROTO'] = parsedurl.scheme
        self.META['QUERY_STRING'] = parsedurl.query
        self.META['HTTP_USER_AGENT'] = stored_request.get('user_agent')
        self.META['REMOTE_ADDR'] = stored_request.get('remote_ip')
        self.method = 'GET'
        self.LANGUAGE_CODE = self.COOKIES['django_language'] = stored_request.get('language')
        self.customer = customer
        self.user = customer.is_anonymous() and AnonymousUser or customer.user
        self.current_page = None

def order_event_notification(sender, instance=None, target=None, **kwargs):
    from shop.models.order import OrderModel
    from shop.serializers.order import OrderDetailSerializer

    if not isinstance(instance, OrderModel._materialized_model):
        return
    for notification in Notification.objects.filter(transition_target=target):
        recipient = notification.get_recipient(instance)
        if recipient is None:
            continue

        # emulate a request object which behaves similar to that one, when the customer submitted its order
        emulated_request = EmulateHttpRequest(instance.customer, instance.stored_request)
        customer_serializer = app_settings.CUSTOMER_SERIALIZER(instance.customer)
        order_serializer = OrderDetailSerializer(instance, context={'request': emulated_request})
        language = instance.stored_request.get('language')
        context = {
            'customer': customer_serializer.data,
            'data': order_serializer.data,
            'ABSOLUTE_BASE_URI': emulated_request.build_absolute_uri().rstrip('/'),
            'render_language': language,
        }
        try:
            template = notification.mail_template.translated_templates.get(language=language)
        except EmailTemplate.DoesNotExist:
            template = notification.mail_template
        attachments = {}
        for notiatt in notification.notificationattachment_set.all():
            attachments[notiatt.attachment.original_filename] = notiatt.attachment.file.file
        mail.send(recipient, template=notification.mail_template, context=context,
                  attachments=attachments, render_on_delivery=True)
<commit_msg>Improve isinstance() check in order notification<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from bs4 import BeautifulSoup
from django.conf import settings
from django.contrib.auth.models import AnonymousUser
from django.core.mail import EmailMessage, EmailMultiAlternatives
from django.db import models
from django.db.models import Q
from django.http.request import HttpRequest
from django.template import Context, engines
from django.utils.translation import ugettext_lazy as _, override as translation_override
from django.utils.six.moves.urllib.parse import urlparse
from shop.models.order import BaseOrder

from post_office import mail
from post_office.models import Email as OriginalEmail, EmailTemplate

from filer.fields.file import FilerFileField

from shop import app_settings
from .customer import CustomerModel

class Email(OriginalEmail):
    """
    Patch class `post_office.models.Email` which fixes https://github.com/ui/django-post_office/issues/116
    and additionally converts HTML messages into plain text messages.
    """
    class Meta:
        proxy = True

    def email_message(self, connection=None):
        if self.template is not None:
            render_language = self.context.get('render_language', settings.LANGUAGE_CODE)
            context = Context(self.context)
            with translation_override(render_language):
                subject = engines['django'].from_string(self.template.subject).render(context)
                message = engines['django'].from_string(self.template.content).render(context)
                html_message = engines['django'].from_string(self.template.html_content).render(context)
        else:
            subject = self.subject
            message = self.message
            html_message = self.html_message

        if html_message:
            if not message:
                message = BeautifulSoup(html_message).text
            mailmsg = EmailMultiAlternatives(
                subject=subject, body=message, from_email=self.from_email,
                to=self.to, bcc=self.bcc, cc=self.cc,
                connection=connection, headers=self.headers)
            mailmsg.attach_alternative(html_message, 'text/html')
        else:
            mailmsg = EmailMessage(
                subject=subject, body=message, from_email=self.from_email,
                to=self.to, bcc=self.bcc, cc=self.cc,
                connection=connection, headers=self.headers)

        for attachment in self.attachments.all():
            mailmsg.attach(attachment.name, attachment.file.read())
        return mailmsg

class EmailManager(models.Manager):
    def get_queryset(self):
        return Email.objects.get_queryset()

OriginalEmail.add_to_class('objects', EmailManager())

class Notification(models.Model):
    """
    A task executed on receiving a signal.
    """
    name = models.CharField(max_length=50, verbose_name=_("Name"))
    transition_target = models.CharField(max_length=50, verbose_name=_("Event"))
    mail_to = models.PositiveIntegerField(verbose_name=_("Mail to"), null=True,
                                          blank=True, default=None)
    mail_template = models.ForeignKey(EmailTemplate, verbose_name=_("Template"),
                                      limit_choices_to=Q(language__isnull=True) | Q(language=''))

    class Meta:
        app_label = 'shop'
        verbose_name = _("Notification")
        verbose_name_plural = _("Notifications")
        ordering = ('transition_target', 'mail_to')

    def get_recipient(self, order):
        if self.mail_to is None:
            return
        if self.mail_to == 0:
            return order.customer.email
        return CustomerModel.objects.get(pk=self.mail_to).email

class NotificationAttachment(models.Model):
    notification = models.ForeignKey(Notification)
    attachment = FilerFileField(null=True, blank=True, related_name='email_attachment')

    class Meta:
        app_label = 'shop'

class EmulateHttpRequest(HttpRequest):
    """
    Use this class to emulate a HttpRequest object, when templates must be rendered
    asynchronously, for instance when an email must be generated out of an Order object.
    """
    def __init__(self, customer, stored_request):
        super(EmulateHttpRequest, self).__init__()
        parsedurl = urlparse(stored_request.get('absolute_base_uri'))
        self.path = self.path_info = parsedurl.path
        self.environ = {}
        self.META['PATH_INFO'] = parsedurl.path
        self.META['SCRIPT_NAME'] = ''
        self.META['HTTP_HOST'] = parsedurl.netloc
        self.META['HTTP_X_FORWARDED_PROTO'] = parsedurl.scheme
        self.META['QUERY_STRING'] = parsedurl.query
        self.META['HTTP_USER_AGENT'] = stored_request.get('user_agent')
        self.META['REMOTE_ADDR'] = stored_request.get('remote_ip')
        self.method = 'GET'
        self.LANGUAGE_CODE = self.COOKIES['django_language'] = stored_request.get('language')
        self.customer = customer
        self.user = customer.is_anonymous() and AnonymousUser or customer.user
        self.current_page = None

def order_event_notification(sender, instance=None, target=None, **kwargs):
    from shop.serializers.order import OrderDetailSerializer

    if not isinstance(instance, BaseOrder):
        return
    for notification in Notification.objects.filter(transition_target=target):
        recipient = notification.get_recipient(instance)
        if recipient is None:
            continue

        # emulate a request object which behaves similar to that one, when the customer submitted its order
        emulated_request = EmulateHttpRequest(instance.customer, instance.stored_request)
        customer_serializer = app_settings.CUSTOMER_SERIALIZER(instance.customer)
        order_serializer = OrderDetailSerializer(instance, context={'request': emulated_request})
        language = instance.stored_request.get('language')
        context = {
            'customer': customer_serializer.data,
            'data': order_serializer.data,
            'ABSOLUTE_BASE_URI': emulated_request.build_absolute_uri().rstrip('/'),
            'render_language': language,
        }
        try:
            template = notification.mail_template.translated_templates.get(language=language)
        except EmailTemplate.DoesNotExist:
            template = notification.mail_template
        attachments = {}
        for notiatt in notification.notificationattachment_set.all():
            attachments[notiatt.attachment.original_filename] = notiatt.attachment.file.file
        mail.send(recipient, template=notification.mail_template, context=context,
                  attachments=attachments, render_on_delivery=True)
<|endoftext|>


<commit_before># Copyright (C) 2004 Anthony Baxter
"""Tests for shtoom.app.
"""

# XXX Todo: Pull in examples from
# http://www.watersprings.org/pub/id/draft-smith-sipping-auth-examples-00.txt

from twisted.trial import unittest

class AuthTest(unittest.TestCase):

    def test_sipSimpleAuth(self):
        from shtoom.sip import Registration
        reg = Registration(None, None)
        resp = reg.calcAuth('REGISTER', 'sip:divmod.com',
                     'Digest realm="asterisk", nonce="24a52b3a"',
                     ('anthony', 'foo'))
        correct = 'Digest username="anthony",realm="asterisk",nonce="24a52b3a",response="925bd204843bf802a1ed8a2494557c01",uri="sip:divmod.com"'
        correct = correct.split(',')
        correct = [ x.strip() for x in correct ]
        correct.sort()
        resp = resp.split(',')
        resp = [ x.strip() for x in resp ]
        resp.sort()
        self.assertEqual(resp, correct)
<commit_msg>registration was changed<commit_after># Copyright (C) 2004 Anthony Baxter
"""Tests for shtoom.app.
"""

# XXX Todo: Pull in examples from
# http://www.watersprings.org/pub/id/draft-smith-sipping-auth-examples-00.txt

from twisted.trial import unittest

class AuthTest(unittest.TestCase):

    def test_sipSimpleAuth(self):
        from shtoom.sip import Registration
        reg = Registration(None)
        resp = reg.calcAuth('REGISTER', 'sip:divmod.com',
                     'Digest realm="asterisk", nonce="24a52b3a"',
                     ('anthony', 'foo'))
        correct = 'Digest username="anthony",realm="asterisk",nonce="24a52b3a",response="925bd204843bf802a1ed8a2494557c01",uri="sip:divmod.com"'
        correct = correct.split(',')
        correct = [ x.strip() for x in correct ]
        correct.sort()
        resp = resp.split(',')
        resp = [ x.strip() for x in resp ]
        resp.sort()
        self.assertEqual(resp, correct)
<|endoftext|>


<commit_before>"""LAMMPS SimPhoNy Wrapper.

This module provides a wrapper for LAMMPS-md
"""
import contextlib
import os
import shutil
import tempfile

from io.lammps_fileio_data_manager import LammpsFileIoDataManager
from io.lammps_process import LammpsProcess

from common.atom_style import AtomStyle

from config.script_writer import ScriptWriter

from cuba_extension import CUBAExtension

from internal.lammps_internal_data_manager import LammpsInternalDataManager

from simphony.core import CUBA
from simphony.core.data_container import DataContainer
from simphony.cuds import CUDS
from simphony.cuds.abc_modeling_engine import ABCModelingEngine
from simphony.cuds.abc_particles import ABCParticles
from simphony.cuds.meta import api
import simphony.cuds.particles as scp

def _temp_directory():
    """Provide a temp directory.

    The name of the created temp directory is returned when context is entered
    and this directory is deleted when context is exited
    """
    temp_dir = tempfile.mkdtemp()
    yield temp_dir
    shutil.rmtree(temp_dir)

_temp_directory = contextlib.contextmanager(_temp_directory)

class LammpsWrapper(ABCModelingEngine):
    """Wrapper to LAMMPS-md."""

    def __init__(self, use_internal_interface=False, **kwargs):
        """Constructor.

        Parameters
        ----------
        engine_type : EngineType
            type of engine

        use_internal_interface : bool, optional
            If true, then the internal interface (library) is used when
            communicating with LAMMPS, if false, then file-io interface is
            used where input/output files are used to communicate with LAMMPS
        """
        self.BC = DataContainer()
        self.CM = DataContainer()
        self.SP = DataContainer()
        self.CM_extension = {}
        self.SP_extension = {}
        self.BC_extension = {}
        self.SD = kwargs.get('cuds', CUDS())

        self._use_internal_interface = use_internal_interface
        atom_style = AtomStyle.ATOMIC
        self._executable_name = 'lammps'
        self._script_writer = ScriptWriter(atom_style)

        if self._use_internal_interface:
            import lammps
            # lammps.lammps(cmdargs=["-screen", "none", "-log", "none"])
            self._lammps = lammps.lammps()
            self._data_manager = LammpsInternalDataManager(self._lammps,
                                                           self.SD,
                                                           atom_style)
        else:
            self._data_manager = LammpsFileIoDataManager(self.SD, atom_style)

        # Number of runs
        self._run_count = 0

        # Dataset uids which are added.
        self._dataset_uids = []

        # Call the base class in order to load CUDS
        super(LammpsWrapper, self).__init__(**kwargs)

    def _count_of(self, cuds, item_type):
        """Workaround for broken CUDS counter."""
        count = 0
        for c in cuds.iter():
            if isinstance(c, item_type):
                count += 1
        return count

    def _load_cuds(self):
        '''Load CUDS data into lammps engine.'''
        cuds = self.get_cuds()
        if not cuds:
            return

        # FIXME: `count_of` is broken for non-meta classes,
        # e.g. Particles, Particle, Mesh, etc.
        # if cuds.count_of(CUBA.PARTICLES) != 1:
        if self._count_of(cuds, scp.Particles) != 1:
            raise Exception('This version of simlammps'
                            ' needs only one particle dataset, not %s' %
                            self._count_of(cuds, scp.Particles))
        if cuds.count_of(CUBA.MATERIAL) != 1:
            raise Exception('This version of simlammps needs'
                            ' only one material not %s' %
                            cuds.count_of(CUBA.MATERIAL))

        material_to_atom = {}
        number_atom_types = 0
        for mat in cuds.iter(item_type=CUBA.MATERIAL):
            if CUBA.MASS not in mat.data:
                raise Exception('Material needs mass property!')
            number_atom_types += 1
            material_to_atom[mat.uid] = number_atom_types

        for particles in cuds.iter(item_type=CUBA.PARTICLES):
            update_list = []
            for single_particle in particles.iter():
                single_particle.data[CUBA.MATERIAL_TYPE] = mat.uid
                update_list.append(single_particle)
            particles.update(update_list)

        for b in cuds.iter(item_type=CUBA.BOX):
            pass

        particle_sum = 0
        for ds in cuds.iter(item_type=CUBA.PARTICLES):
            particle_sum += len(ds)
            ds.data = mat.data
            ds.data_extension = {
                CUBAExtension.BOX_VECTORS: b.vector}
            self.add_dataset(ds)

        if particle_sum == 0:
            raise Exception('simlammps needs some particles')

        if cuds.count_of(CUBA.MOLECULAR_DYNAMICS) != 1:
            raise Exception('simlammps supports only MD')

        if cuds.count_of(CUBA.BOX) != 1:
            raise Exception('simlammps needs one box')

        for termo in cuds.iter(item_type=CUBA.THERMOSTAT):
            if isinstance(termo, api.TemperatureRescaling):
                self.CM_extension[CUBA.THERMODYNAMIC_ENSEMBLE] = 'NVE'
                continue

        if not cuds.count_of(CUBA.INTEGRATION_TIME):
            raise Exception('Only one integration time setup is accepted')

        for i in cuds.iter(item_type=CUBA.INTEGRATION_TIME):
            self.CM[CUBA.TIME_STEP] = i.step
            self.CM[CUBA.NUMBER_OF_TIME_STEPS] = int(i.final / i.step)

        if cuds.count_of(CUBA.CONDITION) != 1:
            raise Exception('Sorry only one condition is accepted, not %s' %
                            cuds.count_of(CUBA.CONDITION))

        for c in cuds.iter(item_type=CUBA.CONDITION):
            if isinstance(c, api.Periodic):
                self.BC_extension[CUBAExtension.BOX_FACES] = [
                    'periodic',
                    'periodic',
                    'periodic']
                continue
            raise Exception('Sorry, I am confused!')

        if cuds.count_of(CUBA.INTERATOMIC_POTENTIAL) != 1:
            raise Exception('Lammps needs one interatomic potention')

        for ip in cuds.iter(item_type=CUBA.INTERATOMIC_POTENTIAL):
            pass

        self.SP_extension[CUBAExtension.PAIR_POTENTIALS] =\
            'lj:\n  global_cutoff: {global_cutoff}\n'\
            '  parameters:\n  - pair: [{mat1}, {mat2}]\n'\
            '    epsilon: {energy_well_depth}\n'\
            '    sigma: {vanderwaals_radious}\n'\
            '    cutoff: {cutoff}\n'\
            .format(global_cutoff=1.12246,
                    mat1=material_to_atom[ip.material[0].uid],
                    mat2=material_to_atom[ip.material[1].uid],
                    energy_well_depth=ip.energy_well_depth,
                    vanderwaals_radious=ip.van_der_waals_radius,
                    cutoff=ip.cutoff_distance)

    def add_dataset(self, container):
        """Add a CUDS container.

        Parameters
        ----------
        container : {ABCParticles}
            The CUDS container to add to the engine.

        Raises
        ------
        TypeError:
            If the container type is not supported (i.e. ABCLattice, ABCMesh).
        ValueError:
            If there is already a dataset with the given name.

        """
        if not isinstance(container, ABCParticles):
            raise TypeError(
                'The type of the dataset container is not supported')
        if container.name in self._data_manager:
            raise ValueError("Particle container '{}' already exists"
                             .format(container.name))
        self._data_manager.new_particles(container)

    def get_dataset(self, name):
        """Get the dataset.

        The returned particle container can be used to query
        and change the related data inside LAMMPS.

        Parameters
        ----------
        name: str
            name of CUDS container to be retrieved.

        Returns
        -------
        container :
            A proxy of the dataset named ``name`` that is stored
            internally in the Engine.

        Raises
        ------
        ValueError:
            If there is no dataset with the given name

        """
        if name in self._data_manager:
            return self._data_manager[name]

        raise ValueError("Particle container '{}` does not exist"
                         .format(name))

    def get_dataset_names(self):
        """Return the names of all the datasets."""
        return [name for name in self._data_manager]

    def remove_dataset(self, name):
        """Remove a dataset.

        Parameters
        ----------
        name: str
            name of CUDS container to be deleted

        Raises
        ------
        ValueError:
            If there is no dataset with the given name

        """
        if name in self._data_manager:
            del self._data_manager[name]
        else:
            raise ValueError("Particles '{}\\' does not exist".format(name))

    def iter_datasets(self, names=None):
        """Return an iterator over a subset or all of the containers.

        Parameters
        ----------
        names : sequence of str, optional
            names of specific containers to be iterated over. If names is not
            given, then all containers will be iterated over.

        """
        if not names:
            for name in self._data_manager:
                yield self._data_manager[name]
        else:
            for name in names:
                if name in self._data_manager:
                    yield self._data_manager[name]
                    continue
                raise ValueError("Particle container '{}\\` does not exist"
                                 .format(name))

    def run(self):
        """Run lammps-engine based on configuration and data."""
        if self._run_count > 0:
            self._load_cuds()

        if self._use_internal_interface:
            self._data_manager.flush()
            commands = ''
            commands += ScriptWriter.get_pair_style(
                _combine(self.SP, self.SP_extension))
            commands += ScriptWriter.get_fix(CM=_combine(self.CM,
                                                         self.CM_extension))
            commands += ScriptWriter.get_pair_coeff(_combine(
                self.SP, self.SP_extension))
            commands += ScriptWriter.get_boundary(_combine(
                self.BC, self.BC_extension), change_existing_boundary=True)
            commands += ScriptWriter.get_run(CM=_combine(
                self.CM, self.CM_extension))
            for command in commands.splitlines():
                self._lammps.command(command)
            # after running, we read any changes from lammps
            self._data_manager.read()
        else:
            with _temp_directory() as temp_dir:
                input_data_filename = os.path.join(temp_dir, 'data_in.lammps')
                output_data_filename = os.path.join(temp_dir,
                                                    'data_out.lammps')
                self._data_manager.flush(input_data_filename)

                commands = self._script_writer.get_configuration(
                    input_data_file=input_data_filename,
                    output_data_file=output_data_filename,
                    BC=_combine(self.BC, self.BC_extension),
                    CM=_combine(self.CM, self.CM_extension),
                    SP=_combine(self.SP, self.SP_extension),
                    materials=[
                        material for material in
                        self.SD.iter(item_type=CUBA.MATERIAL)])

                process = LammpsProcess(lammps_name=self._executable_name,
                                        log_directory=temp_dir)

                process.run(commands)
                self._data_manager.read(output_data_filename)
        # A naive flag for the next run.
        self._run_count += 1

        # Replace datasets in CUDS with proxy ones.
        for ds_name in self._data_manager:
            proxy_dataset = self.get_dataset(ds_name)
            cuds_dataset = self._cuds.get_by_name(ds_name)
            proxy_dataset._uid = cuds_dataset.uid
            self._cuds.update([self.get_dataset(ds_name)])

def _combine(data_container, data_container_extension):
    """Combine a the approved CUBA with non-approved CUBA key-values.

    Parameters
    ----------
    data_container : DataContainer
        data container with CUBA attributes
    data_container_extension : dict
        data container with non-approved CUBA attributes

    Returns
    ----------
    dict
        dictionary containing the approved adn non-approved
        CUBA key-values

    """
    result = dict(data_container_extension)
    result.update(data_container)
    return result
<commit_msg>Refactor. Extract CUDS checks.<commit_after>"""LAMMPS SimPhoNy Wrapper.

This module provides a wrapper for LAMMPS-md
"""
import contextlib
import os
import shutil
import tempfile

from io.lammps_fileio_data_manager import LammpsFileIoDataManager
from io.lammps_process import LammpsProcess

from common.atom_style import AtomStyle

from config.script_writer import ScriptWriter

from cuba_extension import CUBAExtension

from internal.lammps_internal_data_manager import LammpsInternalDataManager

from simphony.core import CUBA
from simphony.core.data_container import DataContainer
from simphony.cuds import CUDS
from simphony.cuds.abc_modeling_engine import ABCModelingEngine
from simphony.cuds.abc_particles import ABCParticles
from simphony.cuds.meta import api
import simphony.cuds.particles as scp

def _temp_directory():
    """Provide a temp directory.

    The name of the created temp directory is returned when context is entered
    and this directory is deleted when context is exited
    """
    temp_dir = tempfile.mkdtemp()
    yield temp_dir
    shutil.rmtree(temp_dir)

_temp_directory = contextlib.contextmanager(_temp_directory)

class LammpsWrapper(ABCModelingEngine):
    """Wrapper to LAMMPS-md."""

    def __init__(self, use_internal_interface=False, **kwargs):
        """Constructor.

        Parameters
        ----------
        engine_type : EngineType
            type of engine

        use_internal_interface : bool, optional
            If true, then the internal interface (library) is used when
            communicating with LAMMPS, if false, then file-io interface is
            used where input/output files are used to communicate with LAMMPS
        """
        self.BC = DataContainer()
        self.CM = DataContainer()
        self.SP = DataContainer()
        self.CM_extension = {}
        self.SP_extension = {}
        self.BC_extension = {}
        self.SD = kwargs.get('cuds', CUDS())

        self._use_internal_interface = use_internal_interface
        atom_style = AtomStyle.ATOMIC
        self._executable_name = 'lammps'
        self._script_writer = ScriptWriter(atom_style)

        if self._use_internal_interface:
            import lammps
            # lammps.lammps(cmdargs=["-screen", "none", "-log", "none"])
            self._lammps = lammps.lammps()
            self._data_manager = LammpsInternalDataManager(self._lammps,
                                                           self.SD,
                                                           atom_style)
        else:
            self._data_manager = LammpsFileIoDataManager(self.SD, atom_style)

        # Number of runs
        self._run_count = 0

        # Dataset uids which are added.
        self._dataset_uids = []

        # Call the base class in order to load CUDS
        super(LammpsWrapper, self).__init__(**kwargs)

    def _count_of(self, cuds, item_type):
        """Workaround for broken CUDS counter."""
        count = 0
        for c in cuds.iter():
            if isinstance(c, item_type):
                count += 1
        return count

    def _check_cuds(self, cuds):
        """Check the given cuds for consistency."""
        # FIXME: `count_of` is broken for non-meta classes,
        # e.g. Particles, Particle, Mesh, etc.
        # if cuds.count_of(CUBA.PARTICLES) != 1:
        if self._count_of(cuds, scp.Particles) != 1:
            raise Exception('This version of simlammps'
                            ' needs only one particle dataset, not %s' %
                            self._count_of(cuds, scp.Particles))
        if cuds.count_of(CUBA.MATERIAL) != 1:
            raise Exception('This version of simlammps needs'
                            ' only one material not %s' %
                            cuds.count_of(CUBA.MATERIAL))

        if cuds.count_of(CUBA.MOLECULAR_DYNAMICS) != 1:
            raise Exception('simlammps supports only MD')

        if cuds.count_of(CUBA.BOX) != 1:
            raise Exception('simlammps needs one box')

        if cuds.count_of(CUBA.INTERATOMIC_POTENTIAL) != 1:
            raise Exception('Lammps needs one interatomic potention')

        if cuds.count_of(CUBA.CONDITION) != 1:
            raise Exception('Sorry only one condition is accepted, not %s' %
                            cuds.count_of(CUBA.CONDITION))

        if not cuds.count_of(CUBA.INTEGRATION_TIME):
            raise Exception('Only one integration time setup is accepted')

    def _load_cuds(self):
        """Load CUDS data into lammps engine."""
        cuds = self.get_cuds()
        if not cuds:
            return

        # Move checks to a separate method
        self._check_cuds(cuds)

        material_to_atom = {}
        number_atom_types = 0
        for mat in cuds.iter(item_type=CUBA.MATERIAL):
            if CUBA.MASS not in mat.data:
                raise Exception('Material needs mass property!')
            number_atom_types += 1
            material_to_atom[mat.uid] = number_atom_types

        for particle_container in cuds.iter(item_type=CUBA.PARTICLES):
            update_list = []
            for single_particle in particle_container.iter():
                single_particle.data[CUBA.MATERIAL_TYPE] = mat.uid
                update_list.append(single_particle)
            particle_container.update(update_list)

        for b in cuds.iter(item_type=CUBA.BOX):
            pass
        #b = cuds.get_one(CUBA.BOX)

        particle_sum = 0
        for ds in cuds.iter(item_type=CUBA.PARTICLES):
            particle_sum += len(ds)
            ds.data = mat.data
            ds.data_extension = {
                CUBAExtension.BOX_VECTORS: b.vector}
            # Add dataset when it is not already there. Rely on uid.
            if ds.uid not in self._dataset_uids:
                self.add_dataset(ds)
                self._dataset_uids.append(ds.uid)
                # Replace dataset in CUDS with proxy one.
                # proxy_dataset = self.get_dataset(ds.name)
                # proxy_dataset._uid = ds.uid
                # cuds.update([self.get_dataset(ds.name)])

        if particle_sum == 0:
            raise Exception('simlammps needs some particles')

        # TODO: add thermo check.
        # TODO: else?
        for thermo in cuds.iter(item_type=CUBA.THERMOSTAT):
            if isinstance(thermo, api.TemperatureRescaling):
                self.CM_extension[CUBA.THERMODYNAMIC_ENSEMBLE] = 'NVE'

        for i in cuds.iter(item_type=CUBA.INTEGRATION_TIME):
            self.CM[CUBA.TIME_STEP] = i.step
            self.CM[CUBA.NUMBER_OF_TIME_STEPS] = int(i.final / i.step)

        for c in cuds.iter(item_type=CUBA.CONDITION):
            if isinstance(c, api.Periodic):
                self.BC_extension[CUBAExtension.BOX_FACES] = [
                    'periodic',
                    'periodic',
                    'periodic']
                continue
            raise Exception('Sorry, I am confused!')

        for ip in cuds.iter(item_type=CUBA.INTERATOMIC_POTENTIAL):
            pass

        self.SP_extension[CUBAExtension.PAIR_POTENTIALS] =\
            'lj:\n  global_cutoff: {global_cutoff}\n'\
            '  parameters:\n  - pair: [{mat1}, {mat2}]\n'\
            '    epsilon: {energy_well_depth}\n'\
            '    sigma: {vanderwaals_radious}\n'\
            '    cutoff: {cutoff}\n'\
            .format(global_cutoff=1.12246,
                    mat1=material_to_atom[ip.material[0].uid],
                    mat2=material_to_atom[ip.material[1].uid],
                    energy_well_depth=ip.energy_well_depth,
                    vanderwaals_radious=ip.van_der_waals_radius,
                    cutoff=ip.cutoff_distance)

    def add_dataset(self, container):
        """Add a CUDS container.

        Parameters
        ----------
        container : {ABCParticles}
            The CUDS container to add to the engine.

        Raises
        ------
        TypeError:
            If the container type is not supported (i.e. ABCLattice, ABCMesh).
        ValueError:
            If there is already a dataset with the given name.

        """
        if not isinstance(container, ABCParticles):
            raise TypeError(
                'The type of the dataset container is not supported')
        if container.name in self._data_manager:
            raise ValueError("Particle container '{}' already exists"
                             .format(container.name))
        self._data_manager.new_particles(container)

    def get_dataset(self, name):
        """Get the dataset.

        The returned particle container can be used to query
        and change the related data inside LAMMPS.

        Parameters
        ----------
        name: str
            name of CUDS container to be retrieved.

        Returns
        -------
        container :
            A proxy of the dataset named ``name`` that is stored
            internally in the Engine.

        Raises
        ------
        ValueError:
            If there is no dataset with the given name

        """
        if name in self._data_manager:
            return self._data_manager[name]

        raise ValueError("Particle container '{}` does not exist"
                         .format(name))

    def get_dataset_names(self):
        """Return the names of all the datasets."""
        return [name for name in self._data_manager]

    def remove_dataset(self, name):
        """Remove a dataset.

        Parameters
        ----------
        name: str
            name of CUDS container to be deleted

        Raises
        ------
        ValueError:
            If there is no dataset with the given name

        """
        if name in self._data_manager:
            del self._data_manager[name]
        else:
            raise ValueError("Particles '{}\\' does not exist".format(name))

    def iter_datasets(self, names=None):
        """Return an iterator over a subset or all of the containers.

        Parameters
        ----------
        names : sequence of str, optional
            names of specific containers to be iterated over. If names is not
            given, then all containers will be iterated over.

        """
        if not names:
            for name in self._data_manager:
                yield self._data_manager[name]
        else:
            for name in names:
                if name in self._data_manager:
                    yield self._data_manager[name]
                    continue
                raise ValueError("Particle container '{}\\` does not exist"
                                 .format(name))

    def run(self):
        """Run lammps-engine based on configuration and data."""
        if self._run_count > 0:
            self._load_cuds()

        if self._use_internal_interface:
            self._data_manager.flush()
            commands = ''
            commands += ScriptWriter.get_pair_style(
                _combine(self.SP, self.SP_extension))
            commands += ScriptWriter.get_fix(CM=_combine(self.CM,
                                                         self.CM_extension))
            commands += ScriptWriter.get_pair_coeff(_combine(
                self.SP, self.SP_extension))
            commands += ScriptWriter.get_boundary(_combine(
                self.BC, self.BC_extension), change_existing_boundary=True)
            commands += ScriptWriter.get_run(CM=_combine(
                self.CM, self.CM_extension))
            for command in commands.splitlines():
                self._lammps.command(command)
            # after running, we read any changes from lammps
            self._data_manager.read()
        else:
            with _temp_directory() as temp_dir:
                input_data_filename = os.path.join(temp_dir, 'data_in.lammps')
                output_data_filename = os.path.join(temp_dir,
                                                    'data_out.lammps')
                self._data_manager.flush(input_data_filename)

                commands = self._script_writer.get_configuration(
                    input_data_file=input_data_filename,
                    output_data_file=output_data_filename,
                    BC=_combine(self.BC, self.BC_extension),
                    CM=_combine(self.CM, self.CM_extension),
                    SP=_combine(self.SP, self.SP_extension),
                    materials=[
                        material for material in
                        self.SD.iter(item_type=CUBA.MATERIAL)])

                process = LammpsProcess(lammps_name=self._executable_name,
                                        log_directory=temp_dir)

                process.run(commands)
                self._data_manager.read(output_data_filename)
        # A naive flag for the next run.
        self._run_count += 1

        # Replace datasets in CUDS with proxy ones.
        for ds_name in self._data_manager:
            proxy_dataset = self.get_dataset(ds_name)
            cuds_dataset = self._cuds.get_by_name(ds_name)
            proxy_dataset._uid = cuds_dataset.uid
            self._cuds.update([self.get_dataset(ds_name)])

def _combine(data_container, data_container_extension):
    """Combine a the approved CUBA with non-approved CUBA key-values.

    Parameters
    ----------
    data_container : DataContainer
        data container with CUBA attributes
    data_container_extension : dict
        data container with non-approved CUBA attributes

    Returns
    ----------
    dict
        dictionary containing the approved adn non-approved
        CUBA key-values

    """
    result = dict(data_container_extension)
    result.update(data_container)
    return result
<|endoftext|>


<commit_before>from SimPEG import Survey, Problem, Utils, np, sp
from simpegEM.Utils import SrcUtils
from simpegEM.Utils.EMUtils import omega

####################################################
# Receivers
####################################################

class RxFDEM(Survey.BaseRx):

    knownRxTypes = {
                    'exr':['e', 'Ex', 'real'],
                    'eyr':['e', 'Ey', 'real'],
                    'ezr':['e', 'Ez', 'real'],
                    'exi':['e', 'Ex', 'imag'],
                    'eyi':['e', 'Ey', 'imag'],
                    'ezi':['e', 'Ez', 'imag'],

                    'bxr':['b', 'Fx', 'real'],
                    'byr':['b', 'Fy', 'real'],
                    'bzr':['b', 'Fz', 'real'],
                    'bxi':['b', 'Fx', 'imag'],
                    'byi':['b', 'Fy', 'imag'],
                    'bzi':['b', 'Fz', 'imag'],

                    'jxr':['j', 'Fx', 'real'],
                    'jyr':['j', 'Fy', 'real'],
                    'jzr':['j', 'Fz', 'real'],
                    'jxi':['j', 'Fx', 'imag'],
                    'jyi':['j', 'Fy', 'imag'],
                    'jzi':['j', 'Fz', 'imag'],

                    'hxr':['h', 'Ex', 'real'],
                    'hyr':['h', 'Ey', 'real'],
                    'hzr':['h', 'Ez', 'real'],
                    'hxi':['h', 'Ex', 'imag'],
                    'hyi':['h', 'Ey', 'imag'],
                    'hzi':['h', 'Ez', 'imag'],
                   }
    radius = None

    def __init__(self, locs, rxType):
        Survey.BaseRx.__init__(self, locs, rxType)

    @property
    def projField(self):
        """Field Type projection (e.g. e b ...)"""
        return self.knownRxTypes[self.rxType][0]

    @property
    def projGLoc(self):
        """Grid Location projection (e.g. Ex Fy ...)"""
        return self.knownRxTypes[self.rxType][1]

    @property
    def projComp(self):
        """Component projection (real/imag)"""
        return self.knownRxTypes[self.rxType][2]

    def projectFields(self, src, mesh, u):
        P = self.getP(mesh)
        u_part_complex = u[src, self.projField]
        # get the real or imag component
        real_or_imag = self.projComp
        u_part = getattr(u_part_complex, real_or_imag)
        return P*u_part

    def projectFieldsDeriv(self, src, mesh, u, v, adjoint=False):
        P = self.getP(mesh)

        if not adjoint:
            Pv_complex = P * v
            real_or_imag = self.projComp
            Pv = getattr(Pv_complex, real_or_imag)
        elif adjoint:
            Pv_real = P.T * v

            real_or_imag = self.projComp
            if real_or_imag == 'imag':
                Pv = 1j*Pv_real
            elif real_or_imag == 'real':
                Pv = Pv_real.astype(complex)
            else:
                raise NotImplementedError('must be real or imag')

        return Pv

####################################################
# Sources
####################################################

class SrcFDEM(Survey.BaseSrc):
    freq = None
    rxPair = RxFDEM

    def eval(self, prob):
        S_m = self._getS_m(prob)
        S_e = self._getS_e(prob)

        if S_m is not None:
            if len(S_m.shape) == 1: S_m = Utils.mkvc(S_m,2)
        if S_e is not None:
            if len(S_e.shape) == 1: S_e = Utils.mkvc(S_e,2)
        return S_m, S_e

    def evalDeriv(self, prob, v, adjoint=None):
        return self._getS_mDeriv(prob,v,adjoint), self._getS_eDeriv(prob,v,adjoint)

class SrcFDEM_RawVec_e(SrcFDEM):
    """
        RawVec electric source. It is defined by the user provided vector S_e

        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, rxList, freq, S_e):
        self.S_e = np.array(S_e,dtype=float)
        self.freq = float(freq)
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self, prob):
        return None

    def _getS_e(self, prob):
        return self.S_e

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_RawVec_m(SrcFDEM):
    """
        RawVec magnetic source. It is defined by the user provided vector S_m

        :param numpy.array S_m: magnetic source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, rxList, freq, S_m):
        self.S_m = np.array(S_m,dtype=float)
        self.freq = float(freq)
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self, prob):
        return self.S_m

    def _getS_e(self, prob):
        return None

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_RawVec(SrcFDEM):
    """
        RawVec source. It is defined by the user provided vectors S_m, S_e

        :param numpy.array S_m: magnetic source term
        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """
    def __init__(self, rxList, freq, S_m, S_e):
        self.S_m = np.array(S_m,dtype=float)
        self.S_e = np.array(S_e,dtype=float)
        self.freq = float(freq)
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self,prob):
        return self.S_m

    def _getS_e(self,prob):
        return self.S_e

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_MagDipole(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, rxList, freq, loc, orientation='Z', moment=1.):
        self.freq = float(freq)
        self.loc = loc
        self.orientation = orientation
        self.moment = moment
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self,prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(self.loc, gridY, 'y')

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x')
            ay = srcfct(self.loc, gridY, 'y')
            az = srcfct(self.loc, gridZ, 'z')
            a = np.concatenate((ax, ay, az))

        S_m = -1j*omega(self.freq)*C*a

        return S_m

    def _getS_e(self,prob):
        return None

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_MagDipole_Bfield(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    #TODO: neither does moment
    def __init__(self, rxList, freq, loc, orientation='Z', moment=1.):
        self.freq = float(freq)
        self.orientation = orientation
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self,prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl.T

        srcfct = SrcUtils.MagneticDipoleFields
        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            bx = srcfct(self.loc, gridX, 'x')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,bz))
        else:
            bx = srcfct(self.loc, gridX, 'x')
            by = srcfct(self.loc, gridY, 'y')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,by,bz))

        return -1j*omega(self.freq)*b

    def _getS_e(self,prob):
        return None

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_CircularLoop(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, rxList, freq, loc, orientation='Z', radius = 1.):
        self.freq = float(freq)
        self.orientation = orientation
        self.radius = radius
        SrcFDEM.__init__(self, rxList)

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(src.loc, gridY, 'y', self.radius)

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x', self.radius)
            ay = srcfct(self.loc, gridY, 'y', self.radius)
            az = srcfct(self.loc, gridZ, 'z', self.radius)
            a = np.concatenate((ax, ay, az))

        return -1j*omega(self.freq)*C*a

####################################################
# Survey
####################################################

class SurveyFDEM(Survey.BaseSurvey):
    """
        docstring for SurveyFDEM
    """

    srcPair = SrcFDEM

    def __init__(self, srcList, **kwargs):
        # Sort these by frequency
        self.srcList = srcList
        Survey.BaseSurvey.__init__(self, **kwargs)

        _freqDict = {}
        for src in srcList:
            if src.freq not in _freqDict:
                _freqDict[src.freq] = []
            _freqDict[src.freq] += [src]

        self._freqDict = _freqDict
        self._freqs = sorted([f for f in self._freqDict])

    @property
    def freqs(self):
        """Frequencies"""
        return self._freqs

    @property
    def nFreq(self):
        """Number of frequencies"""
        return len(self._freqDict)

    @property
    def nSrcByFreq(self):
        if getattr(self, '_nSrcByFreq', None) is None:
            self._nSrcByFreq = {}
            for freq in self.freqs:
                self._nSrcByFreq[freq] = len(self.getSrcByFreq(freq))
        return self._nSrcByFreq

    def getSrcByFreq(self, freq):
        """Returns the sources associated with a specific frequency."""
        assert freq in self._freqDict, "The requested frequency is not in this survey."
        return self._freqDict[freq]

    def projectFields(self, u):
        data = Survey.Data(self)
        for src in self.srcList:
            for rx in src.rxList:
                data[src, rx] = rx.projectFields(src, self.mesh, u)
        return data

    def projectFieldsDeriv(self, u):
        raise Exception('Use Sources to project fields deriv.')
<commit_msg>set loc and moment inside of MagDipoleB source<commit_after>from SimPEG import Survey, Problem, Utils, np, sp
from simpegEM.Utils import SrcUtils
from simpegEM.Utils.EMUtils import omega

####################################################
# Receivers
####################################################

class RxFDEM(Survey.BaseRx):

    knownRxTypes = {
                    'exr':['e', 'Ex', 'real'],
                    'eyr':['e', 'Ey', 'real'],
                    'ezr':['e', 'Ez', 'real'],
                    'exi':['e', 'Ex', 'imag'],
                    'eyi':['e', 'Ey', 'imag'],
                    'ezi':['e', 'Ez', 'imag'],

                    'bxr':['b', 'Fx', 'real'],
                    'byr':['b', 'Fy', 'real'],
                    'bzr':['b', 'Fz', 'real'],
                    'bxi':['b', 'Fx', 'imag'],
                    'byi':['b', 'Fy', 'imag'],
                    'bzi':['b', 'Fz', 'imag'],

                    'jxr':['j', 'Fx', 'real'],
                    'jyr':['j', 'Fy', 'real'],
                    'jzr':['j', 'Fz', 'real'],
                    'jxi':['j', 'Fx', 'imag'],
                    'jyi':['j', 'Fy', 'imag'],
                    'jzi':['j', 'Fz', 'imag'],

                    'hxr':['h', 'Ex', 'real'],
                    'hyr':['h', 'Ey', 'real'],
                    'hzr':['h', 'Ez', 'real'],
                    'hxi':['h', 'Ex', 'imag'],
                    'hyi':['h', 'Ey', 'imag'],
                    'hzi':['h', 'Ez', 'imag'],
                   }
    radius = None

    def __init__(self, locs, rxType):
        Survey.BaseRx.__init__(self, locs, rxType)

    @property
    def projField(self):
        """Field Type projection (e.g. e b ...)"""
        return self.knownRxTypes[self.rxType][0]

    @property
    def projGLoc(self):
        """Grid Location projection (e.g. Ex Fy ...)"""
        return self.knownRxTypes[self.rxType][1]

    @property
    def projComp(self):
        """Component projection (real/imag)"""
        return self.knownRxTypes[self.rxType][2]

    def projectFields(self, src, mesh, u):
        P = self.getP(mesh)
        u_part_complex = u[src, self.projField]
        # get the real or imag component
        real_or_imag = self.projComp
        u_part = getattr(u_part_complex, real_or_imag)
        return P*u_part

    def projectFieldsDeriv(self, src, mesh, u, v, adjoint=False):
        P = self.getP(mesh)

        if not adjoint:
            Pv_complex = P * v
            real_or_imag = self.projComp
            Pv = getattr(Pv_complex, real_or_imag)
        elif adjoint:
            Pv_real = P.T * v

            real_or_imag = self.projComp
            if real_or_imag == 'imag':
                Pv = 1j*Pv_real
            elif real_or_imag == 'real':
                Pv = Pv_real.astype(complex)
            else:
                raise NotImplementedError('must be real or imag')

        return Pv

####################################################
# Sources
####################################################

class SrcFDEM(Survey.BaseSrc):
    freq = None
    rxPair = RxFDEM

    def eval(self, prob):
        S_m = self._getS_m(prob)
        S_e = self._getS_e(prob)

        if S_m is not None:
            if len(S_m.shape) == 1: S_m = Utils.mkvc(S_m,2)
        if S_e is not None:
            if len(S_e.shape) == 1: S_e = Utils.mkvc(S_e,2)
        return S_m, S_e

    def evalDeriv(self, prob, v, adjoint=None):
        return self._getS_mDeriv(prob,v,adjoint), self._getS_eDeriv(prob,v,adjoint)

class SrcFDEM_RawVec_e(SrcFDEM):
    """
        RawVec electric source. It is defined by the user provided vector S_e

        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, rxList, freq, S_e):
        self.S_e = np.array(S_e,dtype=float)
        self.freq = float(freq)
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self, prob):
        return None

    def _getS_e(self, prob):
        return self.S_e

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_RawVec_m(SrcFDEM):
    """
        RawVec magnetic source. It is defined by the user provided vector S_m

        :param numpy.array S_m: magnetic source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, rxList, freq, S_m):
        self.S_m = np.array(S_m,dtype=float)
        self.freq = float(freq)
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self, prob):
        return self.S_m

    def _getS_e(self, prob):
        return None

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_RawVec(SrcFDEM):
    """
        RawVec source. It is defined by the user provided vectors S_m, S_e

        :param numpy.array S_m: magnetic source term
        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """
    def __init__(self, rxList, freq, S_m, S_e):
        self.S_m = np.array(S_m,dtype=float)
        self.S_e = np.array(S_e,dtype=float)
        self.freq = float(freq)
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self,prob):
        return self.S_m

    def _getS_e(self,prob):
        return self.S_e

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_MagDipole(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, rxList, freq, loc, orientation='Z', moment=1.):
        self.freq = float(freq)
        self.loc = loc
        self.orientation = orientation
        self.moment = moment
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self,prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(self.loc, gridY, 'y')

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x')
            ay = srcfct(self.loc, gridY, 'y')
            az = srcfct(self.loc, gridZ, 'z')
            a = np.concatenate((ax, ay, az))

        S_m = -1j*omega(self.freq)*C*a

        return S_m

    def _getS_e(self,prob):
        return None

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_MagDipole_Bfield(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    #TODO: neither does moment
    def __init__(self, rxList, freq, loc, orientation='Z', moment=1.):
        self.freq = float(freq)
        self.loc = loc
        self.orientation = orientation
        self.moment = moment
        SrcFDEM.__init__(self, rxList)

    def _getS_m(self,prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl.T

        srcfct = SrcUtils.MagneticDipoleFields
        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            bx = srcfct(self.loc, gridX, 'x')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,bz))
        else:
            bx = srcfct(self.loc, gridX, 'x')
            by = srcfct(self.loc, gridY, 'y')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,by,bz))

        return -1j*omega(self.freq)*b

    def _getS_e(self,prob):
        return None

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

class SrcFDEM_CircularLoop(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, rxList, freq, loc, orientation='Z', radius = 1.):
        self.freq = float(freq)
        self.orientation = orientation
        self.radius = radius
        SrcFDEM.__init__(self, rxList)

    def _getS_mDeriv(self, prob, v, adjoint = False):
        return None

    def _getS_eDeriv(self, prob, v, adjoint = False):
        return None

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(src.loc, gridY, 'y', self.radius)

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x', self.radius)
            ay = srcfct(self.loc, gridY, 'y', self.radius)
            az = srcfct(self.loc, gridZ, 'z', self.radius)
            a = np.concatenate((ax, ay, az))

        return -1j*omega(self.freq)*C*a

####################################################
# Survey
####################################################

class SurveyFDEM(Survey.BaseSurvey):
    """
        docstring for SurveyFDEM
    """

    srcPair = SrcFDEM

    def __init__(self, srcList, **kwargs):
        # Sort these by frequency
        self.srcList = srcList
        Survey.BaseSurvey.__init__(self, **kwargs)

        _freqDict = {}
        for src in srcList:
            if src.freq not in _freqDict:
                _freqDict[src.freq] = []
            _freqDict[src.freq] += [src]

        self._freqDict = _freqDict
        self._freqs = sorted([f for f in self._freqDict])

    @property
    def freqs(self):
        """Frequencies"""
        return self._freqs

    @property
    def nFreq(self):
        """Number of frequencies"""
        return len(self._freqDict)

    @property
    def nSrcByFreq(self):
        if getattr(self, '_nSrcByFreq', None) is None:
            self._nSrcByFreq = {}
            for freq in self.freqs:
                self._nSrcByFreq[freq] = len(self.getSrcByFreq(freq))
        return self._nSrcByFreq

    def getSrcByFreq(self, freq):
        """Returns the sources associated with a specific frequency."""
        assert freq in self._freqDict, "The requested frequency is not in this survey."
        return self._freqDict[freq]

    def projectFields(self, u):
        data = Survey.Data(self)
        for src in self.srcList:
            for rx in src.rxList:
                data[src, rx] = rx.projectFields(src, self.mesh, u)
        return data

    def projectFieldsDeriv(self, u):
        raise Exception('Use Sources to project fields deriv.')
<|endoftext|>


<commit_before># This file is part of victims-web.
#
# Copyright (C) 2013 The Victims Project
#
# 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 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/>.
"""
Version 2 of the webservice. Remember service versions are not the same as
application versions.
"""
import datetime
import json

from flask import Blueprint, Response, request, current_app

from victims_web.cache import cache
from victims_web.handlers.security import apiauth, api_request_user
from victims_web.handlers.sslify import ssl_exclude
from victims_web.models import Hash, Removal, JsonifyMixin, Coordinates
from victims_web.submissions import submit, upload
from victims_web.util import groups

v2 = Blueprint('service_v2', __name__)

# Module globals
EOL = None
MIME_TYPE = 'application/json'

def make_response(data, code=200):
    return Response(
        response=data,
        status=code,
        mimetype=MIME_TYPE
    )

def error(msg='Could not understand request.', code=400):
    """
    Returns an error json response.

    :Parameters:
        - `msg`: Error message to be returned in json string.
        - `code`: The code to return as status code for the response.
    """
    return make_response(json.dumps([{'error': msg}]), code)

def success(msg='Request successful.', code=201):
    """
    Returns a success json resposne.

    :Paramenters:
        - `msg`: Error message to be returned in json string.
        - `code`: The code to return as status code for the response.
    """
    return make_response(json.dumps([{'success': msg}]), code)

class StreamedSerialResponseValue(object):
    """
    A thin wrapper class around the cleaned/filtered results to enable
    streaming and caching simultaneously.
    """

    def __init__(self, result, fields=None):
        """
        Creates the streamed iterator.

        :Parameters:
           - `result`: The result to iterate over.
        """
        self.result = result.clone()
        self.fields = fields
        # NOTE: We must do the count else the cursor will stop at 100
        self.result_count = self.result.count()

    def _json(self, item):
        if isinstance(item, JsonifyMixin):
            return item.jsonify(self.fields)
        elif isinstance(item, str) or isinstance(item, unicode):
            return str(item)
        else:
            return json.dumps(item)

    def __getstate__(self):
        """
        The state returned is just the json string of the object
        """
        dump = [self._json(o) for o in self.result]
        return json.dumps((dump, self.fields, self.result_count))

    def __setstate__(self, state):
        """
        When unpickling, convert the json string into an py-object
        """
        (self.result, self.fields, self.result_count) = json.loads(state)

    def __iter__(self):
        """
        The iterator implementing result to json string generator and
        splitting the results by newlines.
        """
        yield "[\n"
        count = 0
        for item in self.result:
            count += 1
            jsons = self._json(item)
            if jsons == '{}':
                continue
            data = '{"fields": ' + jsons + '}'
            if count != self.result_count:
                yield data + ",\n"
            else:
                yield data
        yield "]"

def stream_items(items, fields=None):
    return make_response(StreamedSerialResponseValue(items, fields))

@v2.route('/status.json')
@cache.cached()
def status():
    """
    Return the status of the service.
    """
    data = json.dumps({
        'eol': EOL,
        'supported': True,
        'version': '2',
        'recommended': True,
        'endpoint': '/service/v2/'
    })

    return make_response(data)

@v2.route('/update/<group>/<since>/', methods=['GET'])
def update_for_group(group, since):
    """
    Returns all items updated  past a specific date in utc.

    :Parameters:
       - `since`: a specific date in utc
       - `group`: group to limit items to
    """
    try:
        items = Hash.objects(
            date__gt=datetime.datetime.strptime(since, "%Y-%m-%dT%H:%M:%S"),
            group=group
        )

        fields = current_app.config['API_UPDATES_DEFAULT_FIELDS']

        fields_arg = request.args.get('fields', None)
        if fields_arg is not None:
            fields = [
                Hash.modelname(field)
                for field in fields_arg.replace(' ', '').split(',')
            ]

        items = items.only(*fields)
        return stream_items(items, fields)
    except Exception as e:
        current_app.logger.debug(e)
        return error()

@v2.route('/update/<group>/all', methods=['GET'])
def update_all(group):
    """
    A convinience call to get all updates from the begining of time.

    :Parameters:
        - `group`: group to limit items to
    """
    return update_for_group(group, '1970-01-01T00:00:00')

@v2.route('/update/<since>/', methods=['GET'])
def update(since):
    """
    Default update service.

    :Parameters:
       - `since`: a specific date in utc
    """
    return update_for_group(current_app.config['DEFAULT_GROUP'], since)

@v2.route('/remove/<group>/<since>/')
@cache.memoize()
def remove_for_group(group, since):
    """
    Returns all items to remove past a specific date in utc.

    :Parameters:
       - `since`: a specific date in utc
       - `group`: group to limit items to
    """
    try:
        timestamp = datetime.datetime.strptime(since, "%Y-%m-%dT%H:%M:%S")
        items = Removal.objects(date__gt=timestamp, group=group)
        return stream_items(items)
    except:
        return error()

@v2.route('/remove/<since>/')
@cache.memoize()
def remove(since):
    """
    Default remove service.

    :Parameters:
       - `since`: a specific date in utc
    """
    return remove_for_group(current_app.config['DEFAULT_GROUP'], since)

@v2.route('/cves/<algorithm>/<arg>/', methods=['GET'])
def cves_algorithm(algorithm, arg):
    """
    Returns any cves that match the given the request.

    If GET, we check only the combined hashes for the given algorithm for
    matches.

    :Parameters:
       - `algorithm`: Fingerprinting algorithm.
       - `arg`: The fingerprint.
    """
    try:
        algorithms = ['sha512', 'sha1', 'md5']
        if algorithm not in algorithms:
            return error('Invalid alogrithm. Use any of %s.' % (
                ', '.join(algorithms)))
        elif len(arg) not in [32, 40, 128]:
            return error('Invalid checksum length for %s' % (algorithm))

        kwargs = {("hashes__%s__combined" % (algorithm)): arg}
        cves = Hash.objects.only('cves').filter(**kwargs)
        return stream_items(cves, ['cves'])
    except Exception:
        return error()

@v2.route('/cves/<group>/', methods=['GET'])
def cves(group):
    """
    Get cves that match the given coordinates for the specified group.

    Expectes coordinates as arguments.

    :Parameters:
        - `group`: The group for which to search in
    """
    try:
        kwargs = {
            'coordinates__%s' % (coord): request.args.get(coord).strip()
            for coord in current_app.config['SUBMISSION_GROUPS'].get(group)
            if coord in request.args and coord in Coordinates._fields
        }

        if len(kwargs) == 0:
            raise ValueError('No coordinates given')

        kwargs['group'] = group
        fields = ['cves', 'coordinates']
        cves = Hash.objects.only(*fields).filter(**kwargs)
        return stream_items(cves, fields)
    except ValueError as ve:
        return error(ve.message)
    except Exception as e:
        current_app.logger.debug(e.message)
        return error()

@v2.route('/submit/hash/<group>/', methods=['PUT'])
@apiauth
def submit_hash(group):
    """
    Allows for authenticated users to submit hashes via json.
    """
    user = '%s' % api_request_user()
    try:
        if group not in groups():
            raise ValueError('Invalid group specified')
        json_data = request.get_json()
        if 'cves' not in json_data:
            raise ValueError('No CVE provided')
        entry = Hash()
        entry.mongify(json_data)
        entry.submitter = user
        submit(
            user, 'json-api-hash', group, suffix='Hash', entry=entry,
            approval='PENDING_APPROVAL')
        return success()
    except ValueError as ve:
        return error(ve.message)
    except Exception as e:
        current_app.logger.info('Invalid submission by %s' % (user))
        current_app.logger.debug(e)
        return error()

@v2.route('/submit/archive/<group>', methods=['PUT'])
@apiauth
def submit_archive(group):
    """
    Allows for authenticated users to submit archives
    """
    user = '%s' % api_request_user()
    try:
        if group not in groups():
            raise ValueError('Invalid group specified')

        if 'cves' not in request.args:
            raise ValueError('CVE(s) required')

        cves = [cve.strip() for cve in request.args['cves'].split(',')]

        coordinates = Coordinates(**{
            coord: request.args.get(coord).strip()
            for coord in current_app.config['SUBMISSION_GROUPS'].get(group)
            if coord in request.args
        })
        files = upload(group, request.files.get('archive', None), coordinates)

        for (ondisk, filename, suffix) in files:
            submit(
                user, ondisk, group, filename, suffix, cves,
                coordinates=coordinates
            )

        return success()
    except ValueError as ve:
        current_app.logger.info('Invalid submission by %s: %s' %
                                (user, ve.message))
        return error(ve.message)
    except Exception as e:
        current_app.logger.info(e.message)
        return error()

SUBMISSION_ROUTES = [submit_hash, submit_archive]

for v in [update, remove, cves]:
    ssl_exclude(update)
<commit_msg>Fix submission route<commit_after># This file is part of victims-web.
#
# Copyright (C) 2013 The Victims Project
#
# 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 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/>.
"""
Version 2 of the webservice. Remember service versions are not the same as
application versions.
"""
import datetime
import json

from flask import Blueprint, Response, request, current_app

from victims_web.cache import cache
from victims_web.handlers.security import apiauth, api_request_user
from victims_web.handlers.sslify import ssl_exclude
from victims_web.models import Hash, Removal, JsonifyMixin, Coordinates
from victims_web.submissions import submit, upload
from victims_web.util import groups

v2 = Blueprint('service_v2', __name__)

# Module globals
EOL = None
MIME_TYPE = 'application/json'

def make_response(data, code=200):
    return Response(
        response=data,
        status=code,
        mimetype=MIME_TYPE
    )

def error(msg='Could not understand request.', code=400):
    """
    Returns an error json response.

    :Parameters:
        - `msg`: Error message to be returned in json string.
        - `code`: The code to return as status code for the response.
    """
    return make_response(json.dumps([{'error': msg}]), code)

def success(msg='Request successful.', code=201):
    """
    Returns a success json resposne.

    :Paramenters:
        - `msg`: Error message to be returned in json string.
        - `code`: The code to return as status code for the response.
    """
    return make_response(json.dumps([{'success': msg}]), code)

class StreamedSerialResponseValue(object):
    """
    A thin wrapper class around the cleaned/filtered results to enable
    streaming and caching simultaneously.
    """

    def __init__(self, result, fields=None):
        """
        Creates the streamed iterator.

        :Parameters:
           - `result`: The result to iterate over.
        """
        self.result = result.clone()
        self.fields = fields
        # NOTE: We must do the count else the cursor will stop at 100
        self.result_count = self.result.count()

    def _json(self, item):
        if isinstance(item, JsonifyMixin):
            return item.jsonify(self.fields)
        elif isinstance(item, str) or isinstance(item, unicode):
            return str(item)
        else:
            return json.dumps(item)

    def __getstate__(self):
        """
        The state returned is just the json string of the object
        """
        dump = [self._json(o) for o in self.result]
        return json.dumps((dump, self.fields, self.result_count))

    def __setstate__(self, state):
        """
        When unpickling, convert the json string into an py-object
        """
        (self.result, self.fields, self.result_count) = json.loads(state)

    def __iter__(self):
        """
        The iterator implementing result to json string generator and
        splitting the results by newlines.
        """
        yield "[\n"
        count = 0
        for item in self.result:
            count += 1
            jsons = self._json(item)
            if jsons == '{}':
                continue
            data = '{"fields": ' + jsons + '}'
            if count != self.result_count:
                yield data + ",\n"
            else:
                yield data
        yield "]"

def stream_items(items, fields=None):
    return make_response(StreamedSerialResponseValue(items, fields))

@v2.route('/status.json')
@cache.cached()
def status():
    """
    Return the status of the service.
    """
    data = json.dumps({
        'eol': EOL,
        'supported': True,
        'version': '2',
        'recommended': True,
        'endpoint': '/service/v2/'
    })

    return make_response(data)

@v2.route('/update/<group>/<since>/', methods=['GET'])
def update_for_group(group, since):
    """
    Returns all items updated  past a specific date in utc.

    :Parameters:
       - `since`: a specific date in utc
       - `group`: group to limit items to
    """
    try:
        items = Hash.objects(
            date__gt=datetime.datetime.strptime(since, "%Y-%m-%dT%H:%M:%S"),
            group=group
        )

        fields = current_app.config['API_UPDATES_DEFAULT_FIELDS']

        fields_arg = request.args.get('fields', None)
        if fields_arg is not None:
            fields = [
                Hash.modelname(field)
                for field in fields_arg.replace(' ', '').split(',')
            ]

        items = items.only(*fields)
        return stream_items(items, fields)
    except Exception as e:
        current_app.logger.debug(e)
        return error()

@v2.route('/update/<group>/all', methods=['GET'])
def update_all(group):
    """
    A convinience call to get all updates from the begining of time.

    :Parameters:
        - `group`: group to limit items to
    """
    return update_for_group(group, '1970-01-01T00:00:00')

@v2.route('/update/<since>/', methods=['GET'])
def update(since):
    """
    Default update service.

    :Parameters:
       - `since`: a specific date in utc
    """
    return update_for_group(current_app.config['DEFAULT_GROUP'], since)

@v2.route('/remove/<group>/<since>/')
@cache.memoize()
def remove_for_group(group, since):
    """
    Returns all items to remove past a specific date in utc.

    :Parameters:
       - `since`: a specific date in utc
       - `group`: group to limit items to
    """
    try:
        timestamp = datetime.datetime.strptime(since, "%Y-%m-%dT%H:%M:%S")
        items = Removal.objects(date__gt=timestamp, group=group)
        return stream_items(items)
    except:
        return error()

@v2.route('/remove/<since>/')
@cache.memoize()
def remove(since):
    """
    Default remove service.

    :Parameters:
       - `since`: a specific date in utc
    """
    return remove_for_group(current_app.config['DEFAULT_GROUP'], since)

@v2.route('/cves/<algorithm>/<arg>/', methods=['GET'])
def cves_algorithm(algorithm, arg):
    """
    Returns any cves that match the given the request.

    If GET, we check only the combined hashes for the given algorithm for
    matches.

    :Parameters:
       - `algorithm`: Fingerprinting algorithm.
       - `arg`: The fingerprint.
    """
    try:
        algorithms = ['sha512', 'sha1', 'md5']
        if algorithm not in algorithms:
            return error('Invalid alogrithm. Use any of %s.' % (
                ', '.join(algorithms)))
        elif len(arg) not in [32, 40, 128]:
            return error('Invalid checksum length for %s' % (algorithm))

        kwargs = {("hashes__%s__combined" % (algorithm)): arg}
        cves = Hash.objects.only('cves').filter(**kwargs)
        return stream_items(cves, ['cves'])
    except Exception:
        return error()

@v2.route('/cves/<group>/', methods=['GET'])
def cves(group):
    """
    Get cves that match the given coordinates for the specified group.

    Expectes coordinates as arguments.

    :Parameters:
        - `group`: The group for which to search in
    """
    try:
        kwargs = {
            'coordinates__%s' % (coord): request.args.get(coord).strip()
            for coord in current_app.config['SUBMISSION_GROUPS'].get(group)
            if coord in request.args and coord in Coordinates._fields
        }

        if len(kwargs) == 0:
            raise ValueError('No coordinates given')

        kwargs['group'] = group
        fields = ['cves', 'coordinates']
        cves = Hash.objects.only(*fields).filter(**kwargs)
        return stream_items(cves, fields)
    except ValueError as ve:
        return error(ve.message)
    except Exception as e:
        current_app.logger.debug(e.message)
        return error()

@v2.route('/submit/hash/<group>/', methods=['PUT'])
@apiauth
def submit_hash(group):
    """
    Allows for authenticated users to submit hashes via json.
    """
    user = '%s' % api_request_user()
    try:
        if group not in groups():
            raise ValueError('Invalid group specified')
        json_data = request.get_json()
        if 'cves' not in json_data:
            raise ValueError('No CVE provided')
        entry = Hash()
        entry.mongify(json_data)
        entry.submitter = user
        submit(
            user, 'json-api-hash', group, suffix='Hash', entry=entry,
            approval='PENDING_APPROVAL')
        return success()
    except ValueError as ve:
        return error(ve.message)
    except Exception as e:
        current_app.logger.info('Invalid submission by %s' % (user))
        current_app.logger.debug(e)
        return error()

@v2.route('/submit/archive/<group>/', methods=['PUT'])
@apiauth
def submit_archive(group):
    """
    Allows for authenticated users to submit archives
    """
    user = '%s' % api_request_user()
    try:
        if group not in groups():
            raise ValueError('Invalid group specified')

        if 'cves' not in request.args:
            raise ValueError('CVE(s) required')

        cves = [cve.strip() for cve in request.args['cves'].split(',')]

        coordinates = Coordinates(**{
            coord: request.args.get(coord).strip()
            for coord in current_app.config['SUBMISSION_GROUPS'].get(group)
            if coord in request.args
        })
        files = upload(group, request.files.get('archive', None), coordinates)

        for (ondisk, filename, suffix) in files:
            submit(
                user, ondisk, group, filename, suffix, cves,
                coordinates=coordinates
            )

        return success()
    except ValueError as ve:
        current_app.logger.info('Invalid submission by %s: %s' %
                                (user, ve.message))
        return error(ve.message)
    except Exception as e:
        current_app.logger.info(e.message)
        return error()

SUBMISSION_ROUTES = [submit_hash, submit_archive]

for v in [update, remove, cves]:
    ssl_exclude(update)
<|endoftext|>


<commit_before>from distutils.core import setup

with open('requirements.txt') as f:
    requirements = f.read().splitlines()

setup(name='DataHandling',
      version='1.11',
      description='Python package with functions for data analysis',
      author='Ashley Setter',
      auther_email='A.Setter@soton.ac.uk',
      url=None,
      packages=['DataHandling', 'DataHandling.LeCroy'],
      package_dir={'DataHandling': 'DataHandling',
      },
      install_requires=requirements,
)
<commit_msg>update minor version number<commit_after>from distutils.core import setup

with open('requirements.txt') as f:
    requirements = f.read().splitlines()

setup(name='DataHandling',
      version='1.1.1',
      description='Python package with functions for data analysis',
      author='Ashley Setter',
      auther_email='A.Setter@soton.ac.uk',
      url=None,
      packages=['DataHandling', 'DataHandling.LeCroy'],
      package_dir={'DataHandling': 'DataHandling',
      },
      install_requires=requirements,
)
<|endoftext|>


<commit_before># NFS RPC client -- RFC 1094

# XXX This is not yet complete.
# XXX Only GETATTR, SETTTR, LOOKUP and READDIR are supported.

# (See mountclient.py for some hints on how to write RPC clients in
# Python in general)

import rpc
from rpc import UDPClient, TCPClient
from mountclient import FHSIZE, MountPacker, MountUnpacker

NFS_PROGRAM = 100003
NFS_VERSION = 2

# enum stat
NFS_OK = 0
# (...many error values...)

# enum ftype
NFNON = 0
NFREG = 1
NFDIR = 2
NFBLK = 3
NFCHR = 4
NFLNK = 5

class NFSPacker(MountPacker):

	def pack_sattrargs(self, sa):
		file, attributes = sa
		self.pack_fhandle(file)
		self.pack_sattr(attributes)

	def pack_sattr(self, sa):
		mode, uid, gid, size, atime, mtime = sa
		self.pack_uint(mode)
		self.pack_uint(uid)
		self.pack_uint(gid)
		self.pack_uint(size)
		self.pack_timeval(atime)
		self.pack_timeval(mtime)

	def pack_diropargs(self, da):
		dir, name = da
		self.pack_fhandle(dir)
		self.pack_string(name)

	def pack_readdirargs(self, ra):
		dir, cookie, count = ra
		self.pack_fhandle(dir)
		self.pack_uint(cookie)
		self.pack_uint(count)

	def pack_timeval(self, tv):
		secs, usecs = tv
		self.pack_uint(secs)
		self.pack_uint(usecs)

class NFSUnpacker(MountUnpacker):

	def unpack_readdirres(self):
		status = self.unpack_enum()
		if status == NFS_OK:
			entries = self.unpack_list(self.unpack_entry)
			eof = self.unpack_bool()
			rest = (entries, eof)
		else:
			rest = None
		return (status, rest)

	def unpack_entry(self):
		fileid = self.unpack_uint()
		name = self.unpack_string()
		cookie = self.unpack_uint()
		return (fileid, name, cookie)

	def unpack_diropres(self):
		status = self.unpack_enum()
		if status == NFS_OK:
			fh = self.unpack_fhandle()
			fa = self.unpack_fattr()
			rest = (fh, fa)
		else:
			rest = None
		return (status, rest)

	def unpack_attrstat(self):
		status = self.unpack_enum()
		if status == NFS_OK:
			attributes = self.unpack_fattr()
		else:
			attributes = None
		return status, attributes

	def unpack_fattr(self):
              type = self.unpack_enum()
              mode = self.unpack_uint()
              nlink = self.unpack_uint()
              uid = self.unpack_uint()
              gid = self.unpack_uint()
              size = self.unpack_uint()
              blocksize = self.unpack_uint()
              rdev = self.unpack_uint()
              blocks = self.unpack_uint()
              fsid = self.unpack_uint()
              fileid = self.unpack_uint()
              atime = self.unpack_timeval()
              mtime = self.unpack_timeval()
              ctime = self.unpack_timeval()
	      return (type, mode, nlink, uid, gid, size, blocksize, \
		      rdev, blocks, fsid, fileid, atime, mtime, ctime)

	def unpack_timeval(self):
		secs = self.unpack_uint()
		usecs = self.unpack_uint()
		return (secs, usecs)

class NFSClient(UDPClient):

	def init(self, host):
		return UDPClient.init(self, host, NFS_PROGRAM, NFS_VERSION)

	def addpackers(self):
		self.packer = NFSPacker().init()
		self.unpacker = NFSUnpacker().init('')

	def mkcred(self, proc):
		if self.cred == None:
			self.cred = rpc.AUTH_UNIX, rpc.make_auth_unix_default()
		return self.cred

	def Getattr(self, fh):
		self.start_call(1)
		self.packer.pack_fhandle(fh)
		self.do_call()
		as = self.unpacker.unpack_attrstat()
		self.end_call()
		return as

	def Setattr(self, sa):
		self.start_call(2)
		self.packer.pack_sattrargs(sa)
		self.do_call()
		as = self.unpacker.unpack_attrstat()
		self.end_call()
		return as

	# Root() is obsolete

	def Lookup(self, da):
		self.start_call(4)
		self.packer.pack_diropargs(da)
		self.do_call()
		dr = self.unpacker.unpack_diropres()
		self.end_call()
		return dr

	# ...

	def Readdir(self, ra):
		self.start_call(16)
		self.packer.pack_readdirargs(ra)
		self.do_call()
		rr = self.unpacker.unpack_readdirres()
		self.end_call()
		return rr

	# Shorthand to get the entire contents of a directory
	def Listdir(self, dir):
		list = []
		ra = (dir, 0, 2000)
		while 1:
			(status, rest) = self.Readdir(ra)
			if status <> NFS_OK:
				break
			entries, eof = rest
			last_cookie = None
			for fileid, name, cookie in entries:
				list.append(fileid, name)
				last_cookie = cookie
			if eof or last_cookie == None:
				break
			ra = (ra[0], last_cookie, ra[2])
		return list

def test():
	import sys
	if sys.argv[1:]: host = sys.argv[1]
	else: host = ''
	if sys.argv[2:]: filesys = sys.argv[2]
	else: filesys = None
	from mountclient import UDPMountClient, TCPMountClient
	mcl = TCPMountClient().init(host)
	if filesys == None:
		list = mcl.Export()
		for item in list:
			print item
		return
	sf = mcl.Mnt(filesys)
	print sf
	fh = sf[1]
	if fh:
		ncl = NFSClient().init(host)
		as = ncl.Getattr(fh)
		print as
		list = ncl.Listdir(fh)
		for item in list: print item
		mcl.Umnt(filesys)
<commit_msg>Changed mkcred interface<commit_after># NFS RPC client -- RFC 1094

# XXX This is not yet complete.
# XXX Only GETATTR, SETTTR, LOOKUP and READDIR are supported.

# (See mountclient.py for some hints on how to write RPC clients in
# Python in general)

import rpc
from rpc import UDPClient, TCPClient
from mountclient import FHSIZE, MountPacker, MountUnpacker

NFS_PROGRAM = 100003
NFS_VERSION = 2

# enum stat
NFS_OK = 0
# (...many error values...)

# enum ftype
NFNON = 0
NFREG = 1
NFDIR = 2
NFBLK = 3
NFCHR = 4
NFLNK = 5

class NFSPacker(MountPacker):

	def pack_sattrargs(self, sa):
		file, attributes = sa
		self.pack_fhandle(file)
		self.pack_sattr(attributes)

	def pack_sattr(self, sa):
		mode, uid, gid, size, atime, mtime = sa
		self.pack_uint(mode)
		self.pack_uint(uid)
		self.pack_uint(gid)
		self.pack_uint(size)
		self.pack_timeval(atime)
		self.pack_timeval(mtime)

	def pack_diropargs(self, da):
		dir, name = da
		self.pack_fhandle(dir)
		self.pack_string(name)

	def pack_readdirargs(self, ra):
		dir, cookie, count = ra
		self.pack_fhandle(dir)
		self.pack_uint(cookie)
		self.pack_uint(count)

	def pack_timeval(self, tv):
		secs, usecs = tv
		self.pack_uint(secs)
		self.pack_uint(usecs)

class NFSUnpacker(MountUnpacker):

	def unpack_readdirres(self):
		status = self.unpack_enum()
		if status == NFS_OK:
			entries = self.unpack_list(self.unpack_entry)
			eof = self.unpack_bool()
			rest = (entries, eof)
		else:
			rest = None
		return (status, rest)

	def unpack_entry(self):
		fileid = self.unpack_uint()
		name = self.unpack_string()
		cookie = self.unpack_uint()
		return (fileid, name, cookie)

	def unpack_diropres(self):
		status = self.unpack_enum()
		if status == NFS_OK:
			fh = self.unpack_fhandle()
			fa = self.unpack_fattr()
			rest = (fh, fa)
		else:
			rest = None
		return (status, rest)

	def unpack_attrstat(self):
		status = self.unpack_enum()
		if status == NFS_OK:
			attributes = self.unpack_fattr()
		else:
			attributes = None
		return status, attributes

	def unpack_fattr(self):
              type = self.unpack_enum()
              mode = self.unpack_uint()
              nlink = self.unpack_uint()
              uid = self.unpack_uint()
              gid = self.unpack_uint()
              size = self.unpack_uint()
              blocksize = self.unpack_uint()
              rdev = self.unpack_uint()
              blocks = self.unpack_uint()
              fsid = self.unpack_uint()
              fileid = self.unpack_uint()
              atime = self.unpack_timeval()
              mtime = self.unpack_timeval()
              ctime = self.unpack_timeval()
	      return (type, mode, nlink, uid, gid, size, blocksize, \
		      rdev, blocks, fsid, fileid, atime, mtime, ctime)

	def unpack_timeval(self):
		secs = self.unpack_uint()
		usecs = self.unpack_uint()
		return (secs, usecs)

class NFSClient(UDPClient):

	def init(self, host):
		return UDPClient.init(self, host, NFS_PROGRAM, NFS_VERSION)

	def addpackers(self):
		self.packer = NFSPacker().init()
		self.unpacker = NFSUnpacker().init('')

	def mkcred(self):
		if self.cred == None:
			self.cred = rpc.AUTH_UNIX, rpc.make_auth_unix_default()
		return self.cred

	def Getattr(self, fh):
		self.start_call(1)
		self.packer.pack_fhandle(fh)
		self.do_call()
		as = self.unpacker.unpack_attrstat()
		self.end_call()
		return as

	def Setattr(self, sa):
		self.start_call(2)
		self.packer.pack_sattrargs(sa)
		self.do_call()
		as = self.unpacker.unpack_attrstat()
		self.end_call()
		return as

	# Root() is obsolete

	def Lookup(self, da):
		self.start_call(4)
		self.packer.pack_diropargs(da)
		self.do_call()
		dr = self.unpacker.unpack_diropres()
		self.end_call()
		return dr

	# ...

	def Readdir(self, ra):
		self.start_call(16)
		self.packer.pack_readdirargs(ra)
		self.do_call()
		rr = self.unpacker.unpack_readdirres()
		self.end_call()
		return rr

	# Shorthand to get the entire contents of a directory
	def Listdir(self, dir):
		list = []
		ra = (dir, 0, 2000)
		while 1:
			(status, rest) = self.Readdir(ra)
			if status <> NFS_OK:
				break
			entries, eof = rest
			last_cookie = None
			for fileid, name, cookie in entries:
				list.append(fileid, name)
				last_cookie = cookie
			if eof or last_cookie == None:
				break
			ra = (ra[0], last_cookie, ra[2])
		return list

def test():
	import sys
	if sys.argv[1:]: host = sys.argv[1]
	else: host = ''
	if sys.argv[2:]: filesys = sys.argv[2]
	else: filesys = None
	from mountclient import UDPMountClient, TCPMountClient
	mcl = TCPMountClient().init(host)
	if filesys == None:
		list = mcl.Export()
		for item in list:
			print item
		return
	sf = mcl.Mnt(filesys)
	print sf
	fh = sf[1]
	if fh:
		ncl = NFSClient().init(host)
		as = ncl.Getattr(fh)
		print as
		list = ncl.Listdir(fh)
		for item in list: print item
		mcl.Umnt(filesys)
<|endoftext|>