<commit_before># -*- coding: utf-8 -*-
"""File wrappers for use as exchange data between views and responses."""
from __future__ import absolute_import
from io import BytesIO
from six.moves.urllib.parse import urlparse

from django.core.files.base import File
from django.utils.encoding import force_bytes

import requests

class StorageFile(File):
    """A file in a Django storage.

    This class looks like :py:class:`django.db.models.fields.files.FieldFile`,
    but unrelated to model instance.

    """
    def __init__(self, storage, name, file=None):
        """Constructor.

        storage:
          Some :py:class:`django.core.files.storage.Storage` instance.

        name:
          File identifier in storage, usually a filename as a string.

        """
        self.storage = storage
        self.name = name
        self.file = file

    def _get_file(self):
        """Getter for :py:attr:``file`` property."""
        if not hasattr(self, '_file') or self._file is None:
            self._file = self.storage.open(self.name, 'rb')
        return self._file

    def _set_file(self, file):
        """Setter for :py:attr:``file`` property."""
        self._file = file

    def _del_file(self):
        """Deleter for :py:attr:``file`` property."""
        del self._file

    #: Required by django.core.files.utils.FileProxy.
    file = property(_get_file, _set_file, _del_file)

    def open(self, mode='rb'):
        """Retrieves the specified file from storage and return open() result.

        Proxy to self.storage.open(self.name, mode).

        """
        return self.storage.open(self.name, mode)

    def save(self, content):
        """Saves new content to the file.

        Proxy to self.storage.save(self.name).

        The content should be a proper File object, ready to be read from the
        beginning.

        """
        return self.storage.save(self.name, content)

    @property
    def path(self):
        """Return a local filesystem path which is suitable for open().

        Proxy to self.storage.path(self.name).

        May raise NotImplementedError if storage doesn't support file access
        with Python's built-in open() function

        """
        return self.storage.path(self.name)

    def delete(self):
        """Delete the specified file from the storage system.

        Proxy to self.storage.delete(self.name).

        """
        return self.storage.delete(self.name)

    def exists(self):
        """Return True if file already exists in the storage system.

        If False, then the name is available for a new file.

        """
        return self.storage.exists(self.name)

    @property
    def size(self):
        """Return the total size, in bytes, of the file.

        Proxy to self.storage.size(self.name).

        """
        return self.storage.size(self.name)

    @property
    def url(self):
        """Return an absolute URL where the file's contents can be accessed.

        Proxy to self.storage.url(self.name).

        """
        return self.storage.url(self.name)

    @property
    def accessed_time(self):
        """Return the last accessed time (as datetime object) of the file.

        Proxy to self.storage.accessed_time(self.name).

        """
        return self.storage.accessed(self.name)

    @property
    def created_time(self):
        """Return the creation time (as datetime object) of the file.

        Proxy to self.storage.created_time(self.name).

        """
        return self.storage.created_time(self.name)

    @property
    def modified_time(self):
        """Return the last modification time (as datetime object) of the file.

        Proxy to self.storage.modified_time(self.name).

        """
        return self.storage.modified_time(self.name)

class VirtualFile(File):
    """Wrapper for files that live in memory."""
    def __init__(self, file=None, name=u'', url='', size=None):
        """Constructor.

        file:
          File object. Typically an io.StringIO.

        name:
          File basename.

        url:
          File URL.

        """
        super(VirtualFile, self).__init__(file, name)
        self.url = url
        if size is not None:
            self._size = size

    def _get_size(self):
        try:
            return self._size
        except AttributeError:
            try:
                self._size = self.file.size
            except AttributeError:
                self._size = len(self.file.getvalue())
        return self._size

    def _set_size(self, value):
        return super(VirtualFile, self)._set_size(value)

    size = property(_get_size, _set_size)

    def __iter__(self):
        """Same as ``File.__iter__()`` but using ``force_bytes()``.

        See https://code.djangoproject.com/ticket/21321

        """
        # Iterate over this file-like object by newlines
        buffer_ = None
        for chunk in self.chunks():
            chunk_buffer = BytesIO(force_bytes(chunk))

            for line in chunk_buffer:
                if buffer_:
                    line = buffer_ + line
                    buffer_ = None

                # If this is the end of a line, yield
                # otherwise, wait for the next round
                if line[-1] in ('\n', '\r'):
                    yield line
                else:
                    buffer_ = line

        if buffer_ is not None:
            yield buffer_

class HTTPFile(File):
    """Wrapper for files that live on remote HTTP servers.

    Acts as a proxy.

    Uses https://pypi.python.org/pypi/requests.

    Always sets "stream=True" in requests kwargs.

    """
    def __init__(self, request_factory=requests.get, url='', name=u'',
                 **kwargs):
        self.request_factory = request_factory
        self.url = url
        if name is None:
            parts = urlparse(url)
            if parts.path:  # Name from path.
                self.name = parts.path.strip('/').rsplit('/', 1)[-1]
            else:  # Name from domain.
                self.name = parts.netloc
        else:
            self.name = name
        kwargs['stream'] = True
        self.request_kwargs = kwargs

    @property
    def request(self):
        try:
            return self._request
        except AttributeError:
            self._request = self.request_factory(self.url,
                                                 **self.request_kwargs)
            return self._request

    @property
    def file(self):
        return self.request.raw

    @property
    def size(self):
        """Return the total size, in bytes, of the file.

        Reads response's "content-length" header.

        """
        return self.request.headers['Content-Length']
<commit_msg>Refs #97 - HTTPFile proxies decoded file wrapper (was undecoded urllib3 file wrapper).<commit_after># -*- coding: utf-8 -*-
"""File wrappers for use as exchange data between views and responses."""
from __future__ import absolute_import
from io import BytesIO
from six.moves.urllib.parse import urlparse

from django.core.files.base import File
from django.utils.encoding import force_bytes

import requests

from django_downloadview.io import StringIteratorIO

class StorageFile(File):
    """A file in a Django storage.

    This class looks like :py:class:`django.db.models.fields.files.FieldFile`,
    but unrelated to model instance.

    """
    def __init__(self, storage, name, file=None):
        """Constructor.

        storage:
          Some :py:class:`django.core.files.storage.Storage` instance.

        name:
          File identifier in storage, usually a filename as a string.

        """
        self.storage = storage
        self.name = name
        self.file = file

    def _get_file(self):
        """Getter for :py:attr:``file`` property."""
        if not hasattr(self, '_file') or self._file is None:
            self._file = self.storage.open(self.name, 'rb')
        return self._file

    def _set_file(self, file):
        """Setter for :py:attr:``file`` property."""
        self._file = file

    def _del_file(self):
        """Deleter for :py:attr:``file`` property."""
        del self._file

    #: Required by django.core.files.utils.FileProxy.
    file = property(_get_file, _set_file, _del_file)

    def open(self, mode='rb'):
        """Retrieves the specified file from storage and return open() result.

        Proxy to self.storage.open(self.name, mode).

        """
        return self.storage.open(self.name, mode)

    def save(self, content):
        """Saves new content to the file.

        Proxy to self.storage.save(self.name).

        The content should be a proper File object, ready to be read from the
        beginning.

        """
        return self.storage.save(self.name, content)

    @property
    def path(self):
        """Return a local filesystem path which is suitable for open().

        Proxy to self.storage.path(self.name).

        May raise NotImplementedError if storage doesn't support file access
        with Python's built-in open() function

        """
        return self.storage.path(self.name)

    def delete(self):
        """Delete the specified file from the storage system.

        Proxy to self.storage.delete(self.name).

        """
        return self.storage.delete(self.name)

    def exists(self):
        """Return True if file already exists in the storage system.

        If False, then the name is available for a new file.

        """
        return self.storage.exists(self.name)

    @property
    def size(self):
        """Return the total size, in bytes, of the file.

        Proxy to self.storage.size(self.name).

        """
        return self.storage.size(self.name)

    @property
    def url(self):
        """Return an absolute URL where the file's contents can be accessed.

        Proxy to self.storage.url(self.name).

        """
        return self.storage.url(self.name)

    @property
    def accessed_time(self):
        """Return the last accessed time (as datetime object) of the file.

        Proxy to self.storage.accessed_time(self.name).

        """
        return self.storage.accessed(self.name)

    @property
    def created_time(self):
        """Return the creation time (as datetime object) of the file.

        Proxy to self.storage.created_time(self.name).

        """
        return self.storage.created_time(self.name)

    @property
    def modified_time(self):
        """Return the last modification time (as datetime object) of the file.

        Proxy to self.storage.modified_time(self.name).

        """
        return self.storage.modified_time(self.name)

class VirtualFile(File):
    """Wrapper for files that live in memory."""
    def __init__(self, file=None, name=u'', url='', size=None):
        """Constructor.

        file:
          File object. Typically an io.StringIO.

        name:
          File basename.

        url:
          File URL.

        """
        super(VirtualFile, self).__init__(file, name)
        self.url = url
        if size is not None:
            self._size = size

    def _get_size(self):
        try:
            return self._size
        except AttributeError:
            try:
                self._size = self.file.size
            except AttributeError:
                self._size = len(self.file.getvalue())
        return self._size

    def _set_size(self, value):
        return super(VirtualFile, self)._set_size(value)

    size = property(_get_size, _set_size)

    def __iter__(self):
        """Same as ``File.__iter__()`` but using ``force_bytes()``.

        See https://code.djangoproject.com/ticket/21321

        """
        # Iterate over this file-like object by newlines
        buffer_ = None
        for chunk in self.chunks():
            chunk_buffer = BytesIO(force_bytes(chunk))

            for line in chunk_buffer:
                if buffer_:
                    line = buffer_ + line
                    buffer_ = None

                # If this is the end of a line, yield
                # otherwise, wait for the next round
                if line[-1] in ('\n', '\r'):
                    yield line
                else:
                    buffer_ = line

        if buffer_ is not None:
            yield buffer_

class HTTPFile(File):
    """Wrapper for files that live on remote HTTP servers.

    Acts as a proxy.

    Uses https://pypi.python.org/pypi/requests.

    Always sets "stream=True" in requests kwargs.

    """
    def __init__(self, request_factory=requests.get, url='', name=u'',
                 **kwargs):
        self.request_factory = request_factory
        self.url = url
        if name is None:
            parts = urlparse(url)
            if parts.path:  # Name from path.
                self.name = parts.path.strip('/').rsplit('/', 1)[-1]
            else:  # Name from domain.
                self.name = parts.netloc
        else:
            self.name = name
        kwargs['stream'] = True
        self.request_kwargs = kwargs

    @property
    def request(self):
        try:
            return self._request
        except AttributeError:
            self._request = self.request_factory(self.url,
                                                 **self.request_kwargs)
            return self._request

    @property
    def file(self):
        try:
            return self._file
        except AttributeError:
            content = self.request.iter_content()
            self._file = StringIteratorIO(content)
            return self._file

    @property
    def size(self):
        """Return the total size, in bytes, of the file.

        Reads response's "content-length" header.

        """
        return self.request.headers['Content-Length']
<|endoftext|>


<commit_before>import logging
import re

from django.core.exceptions import ValidationError
from django.db.models import signals

from django.db.models.fields import CharField, Field, FloatField, NOT_PROVIDED
from django import forms

from django_measurement import measure, utils

logger = logging.getLogger(__name__)

class OriginalUnitField(CharField):
    pass

class MeasureNameField(CharField):
    pass

class MeasurementValueField(FloatField):
    pass

class MeasurementFormField(forms.CharField):
    MEASUREMENT_RE = re.compile(r'(?P<value>[0-9.]+) ?(?P<unit>[A-Za-z_]+)')

    def to_python(self, value):
        if not value:
            return None
        match = self.MEASUREMENT_RE.match(value)
        if not match:
            raise ValidationError('%s could not be parsed into a known measurement' % value)
        value, unit = match.groups()
        unit = unit.lower()
        measurement = utils.guess_measurement(value, unit)
        if not match:
            raise ValidationError('%s is not a known unit' % unit)
        return measurement

def get_measurement_parts(value):
    if isinstance(value, measure.UnknownMeasure):
        return value.get_measurement_parts()
    measure_name = '%s(%s)' % (
        value.__class__.__name__,
        value.STANDARD_UNIT
    )
    original_unit = value.unit
    standard_value = value.standard
    return measure_name, original_unit, standard_value

class MeasurementFieldDescriptor(object):
    MEASUREMENT_NAME_RE = re.compile(r'([a-zA-Z0-9]+)(?:\(([a-zA-Z0-9_]+)\))?')

    def __init__(self, field, measurement_field_name, original_unit_field_name, measure_field_name):
        self.field = field
        self.measurement_field_name = measurement_field_name
        self.original_unit_field_name = original_unit_field_name
        self.measure_field_name = measure_field_name

    def _get_measure_by_name(self, measure_name):
        measures = utils.build_measure_list()
        return measures[measure_name]

    def _get_measure_name_and_std_unit(self, measure_string):
        return self.MEASUREMENT_NAME_RE.search(measure_string).groups()

    def __get__(self, instance, instance_type=None):
        if instance is None:
            return self

        measure_packed = getattr(
            instance,
            self.measure_field_name
        )
        if measure_packed:
            measure_name, std_unit = self._get_measure_name_and_std_unit(
                    measure_packed
            )
        else:
            measure_name = ''
            std_unit = ''
        measurement_value = getattr(
            instance,
            self.measurement_field_name,
        )
        original_unit = getattr(
            instance,
            self.original_unit_field_name,
        )
        if not measure_name or not original_unit:
            return ''

        try:
            instance_measure = self._get_measure_by_name(measure_name)
            if std_unit and instance_measure.STANDARD_UNIT != std_unit:
                raise ValueError(
                    'Measurement %s base unit %s does not match stored %s' %
                    (
                        measure_name,
                        instance_measure.STANDARD_UNIT,
                        std_unit
                    )
                )
            measurement = utils.get_measurement(
                instance_measure,
                measurement_value,
                instance_measure.STANDARD_UNIT,
                original_unit=original_unit
            )
        except (AttributeError, KeyError, ImportError):
            measurement = measure.UnknownMeasure(
                measure=measure_name,
                original_unit=original_unit,
                value=measurement_value,
            )

        return measurement

    def __set__(self, instance, measurement):
        if instance is None:
            raise AttributeError("Must be accessed via instance")

        if measurement:
            measure, original_unit, standard_value = get_measurement_parts(
                measurement
            )
        else:
            measure = ''
            original_unit = ''
            standard_value = 0

        setattr(instance, self.measure_field_name, measure)
        setattr(instance, self.original_unit_field_name, original_unit)
        setattr(instance, self.measurement_field_name, standard_value)

class MeasurementField(Field):
    def __init__(self,
            *args,
            **kwargs
        ):
        super(MeasurementField, self).__init__(*args, **kwargs)
        self.original_unit_field_name = '%s_unit'
        self.measure_field_name = '%s_measure'
        self.measurement_field_name = '%s_value'

    def get_original_unit_field_name(self):
        return self.original_unit_field_name % self.name

    def get_measure_field_name(self):
        return self.measure_field_name % self.name

    def get_measurement_field_name(self):
        return self.measurement_field_name % self.name

    def contribute_to_class(self, cls, name):
        self.name = name

        if self.default is not NOT_PROVIDED:
            parts = get_measurement_parts(self.default)
            default_name, default_unit, default_value = parts
        else:
            default_name, default_unit, default_value = '', '', 0.0

        original_unit_field_name = self.get_original_unit_field_name()
        self.original_unit_field = OriginalUnitField(
            editable=False,
            blank=True,
            default=default_unit,
            max_length=50,
        )
        cls.add_to_class(original_unit_field_name, self.original_unit_field)

        measure_field_name = self.get_measure_field_name()
        self.measure_field = MeasureNameField(
            editable=False,
            blank=True,
            default=default_name,
            max_length=255,
        )
        cls.add_to_class(measure_field_name, self.measure_field)

        measurement_field_name = self.get_measurement_field_name()
        self.measurement_field = MeasurementValueField(
            editable=False,
            blank=True,
            default=default_value,
            max_length=50,
        )
        cls.add_to_class(measurement_field_name, self.measurement_field)

        field = MeasurementFieldDescriptor(
            self,
            measurement_field_name=measurement_field_name,
            original_unit_field_name=original_unit_field_name,
            measure_field_name=measure_field_name,
        )

        signals.pre_init.connect(self.instance_pre_init, sender=cls, weak=False)

        setattr(cls, self.name, field)

    def instance_pre_init(self, signal, sender, args, kwargs, **_kwargs):
        if self.name in kwargs:
            value = kwargs.pop(self.name)
            measure_name, original_unit, standard_value = get_measurement_parts(
                value
            )
            kwargs[self.get_measure_field_name()] = measure_name
            kwargs[self.get_original_unit_field_name()] = original_unit
            kwargs[self.get_measurement_field_name()] = standard_value

try:
    from south.modelsinspector import add_introspection_rules
    add_introspection_rules([], ["^django_measurement\.fields\.MeasurementField"])
    add_introspection_rules([], ["^django_measurement\.fields\.OriginalUnitField"])
    add_introspection_rules([], ["^django_measurement\.fields\.MeasureNameField"])
    add_introspection_rules([], ["^django_measurement\.fields\.MeasurementValueField"])
except ImportError:
    pass
<commit_msg>make custom QuerySet and Manager to allow queries for measurements<commit_after>import logging
import re

from django.core.exceptions import ValidationError
from django.db.models import signals, Manager
from django.db.models.query import QuerySet
from django.db.models.fields import CharField, Field, FloatField, NOT_PROVIDED
from django import forms

from measurement.base import MeasureBase
from django_measurement import measure, utils

logger = logging.getLogger(__name__)

class OriginalUnitField(CharField):
    pass

class MeasureNameField(CharField):
    pass

class MeasurementValueField(FloatField):
    pass

class MeasurementFormField(forms.CharField):
    MEASUREMENT_RE = re.compile(r'(?P<value>[0-9.]+) ?(?P<unit>[A-Za-z_]+)')

    def to_python(self, value):
        if not value:
            return None
        match = self.MEASUREMENT_RE.match(value)
        if not match:
            raise ValidationError('%s could not be parsed into a known measurement' % value)
        value, unit = match.groups()
        unit = unit.lower()
        measurement = utils.guess_measurement(value, unit)
        if not match:
            raise ValidationError('%s is not a known unit' % unit)
        return measurement

def get_measurement_parts(value):
    if isinstance(value, measure.UnknownMeasure):
        return value.get_measurement_parts()
    measure_name = '%s(%s)' % (
        value.__class__.__name__,
        value.STANDARD_UNIT
    )
    original_unit = value.unit
    standard_value = value.standard
    return measure_name, original_unit, standard_value

class MeasurementFieldDescriptor(object):
    MEASUREMENT_NAME_RE = re.compile(r'([a-zA-Z0-9]+)(?:\(([a-zA-Z0-9_]+)\))?')

    def __init__(self, field, measurement_field_name, original_unit_field_name, measure_field_name):
        self.field = field
        self.measurement_field_name = measurement_field_name
        self.original_unit_field_name = original_unit_field_name
        self.measure_field_name = measure_field_name

    def _get_measure_by_name(self, measure_name):
        measures = utils.build_measure_list()
        return measures[measure_name]

    def _get_measure_name_and_std_unit(self, measure_string):
        return self.MEASUREMENT_NAME_RE.search(measure_string).groups()

    def __get__(self, instance, instance_type=None):
        if instance is None:
            return self

        measure_packed = getattr(
            instance,
            self.measure_field_name
        )
        if measure_packed:
            measure_name, std_unit = self._get_measure_name_and_std_unit(
                    measure_packed
            )
        else:
            measure_name = ''
            std_unit = ''
        measurement_value = getattr(
            instance,
            self.measurement_field_name,
        )
        original_unit = getattr(
            instance,
            self.original_unit_field_name,
        )
        if not measure_name or not original_unit:
            return ''

        try:
            instance_measure = self._get_measure_by_name(measure_name)
            if std_unit and instance_measure.STANDARD_UNIT != std_unit:
                raise ValueError(
                    'Measurement %s base unit %s does not match stored %s' %
                    (
                        measure_name,
                        instance_measure.STANDARD_UNIT,
                        std_unit
                    )
                )
            measurement = utils.get_measurement(
                instance_measure,
                measurement_value,
                instance_measure.STANDARD_UNIT,
                original_unit=original_unit
            )
        except (AttributeError, KeyError, ImportError):
            measurement = measure.UnknownMeasure(
                measure=measure_name,
                original_unit=original_unit,
                value=measurement_value,
            )

        return measurement

    def __set__(self, instance, measurement):
        if instance is None:
            raise AttributeError("Must be accessed via instance")

        if measurement:
            measure, original_unit, standard_value = get_measurement_parts(
                measurement
            )
        else:
            measure = ''
            original_unit = ''
            standard_value = 0

        setattr(instance, self.measure_field_name, measure)
        setattr(instance, self.original_unit_field_name, original_unit)
        setattr(instance, self.measurement_field_name, standard_value)

class MeasurementField(Field):
    def __init__(self,
            *args,
            **kwargs
        ):
        super(MeasurementField, self).__init__(*args, **kwargs)
        self.original_unit_field_name = '%s_unit'
        self.measure_field_name = '%s_measure'
        self.measurement_field_name = '%s_value'

    def get_original_unit_field_name(self):
        return self.original_unit_field_name % self.name

    def get_measure_field_name(self):
        return self.measure_field_name % self.name

    def get_measurement_field_name(self):
        return self.measurement_field_name % self.name

    def contribute_to_class(self, cls, name):
        self.name = name

        if self.default is not NOT_PROVIDED:
            parts = get_measurement_parts(self.default)
            default_name, default_unit, default_value = parts
        else:
            default_name, default_unit, default_value = '', '', 0.0

        original_unit_field_name = self.get_original_unit_field_name()
        self.original_unit_field = OriginalUnitField(
            editable=False,
            blank=True,
            default=default_unit,
            max_length=50,
        )
        cls.add_to_class(original_unit_field_name, self.original_unit_field)

        measure_field_name = self.get_measure_field_name()
        self.measure_field = MeasureNameField(
            editable=False,
            blank=True,
            default=default_name,
            max_length=255,
        )
        cls.add_to_class(measure_field_name, self.measure_field)

        measurement_field_name = self.get_measurement_field_name()
        self.measurement_field = MeasurementValueField(
            editable=False,
            blank=True,
            default=default_value,
            max_length=50,
        )
        cls.add_to_class(measurement_field_name, self.measurement_field)

        field = MeasurementFieldDescriptor(
            self,
            measurement_field_name=measurement_field_name,
            original_unit_field_name=original_unit_field_name,
            measure_field_name=measure_field_name,
        )

        signals.pre_init.connect(self.instance_pre_init, sender=cls, weak=False)

        setattr(cls, self.name, field)

        class MeasurementManager(Manager):
            use_for_related_fields = True

            def get_query_set(self):
                return MeasurementQuerySet(model=cls)
        cls.add_to_class('objects', MeasurementManager())

    def instance_pre_init(self, signal, sender, args, kwargs, **_kwargs):
        if self.name in kwargs:
            value = kwargs.pop(self.name)
            measure_name, original_unit, standard_value = get_measurement_parts(
                value
            )
            kwargs[self.get_measure_field_name()] = measure_name
            kwargs[self.get_original_unit_field_name()] = original_unit
            kwargs[self.get_measurement_field_name()] = standard_value

class MeasurementQuerySet(QuerySet):
    def filter(self, *args, **kwargs):
        kwargs = self._fix_fields(kwargs)
        return super(MeasurementQuerySet, self).filter(*args, **kwargs)

    def exclude(self, *args, **kwargs):
        kwargs = self._fix_fields(kwargs)
        return super(MeasurementQuerySet, self).exclude(*args, **kwargs)

    def get(self, *args, **kwargs):
        kwargs = self._fix_fields(kwargs)
        return super(MeasurementQuerySet, self).get(*args, **kwargs)

    def _fix_fields(self, kwargs):
        copy = kwargs.copy()
        for key, val in kwargs.iteritems():
            if issubclass(val.__class__, MeasureBase):
                name, unit, value = get_measurement_parts(val)
                copy[key + '_measure'] = name
                copy[key + '_unit'] = unit
                copy[key + '_value'] = value
                del copy[key]
        return copy

try:
    from south.modelsinspector import add_introspection_rules
    add_introspection_rules([], ["^django_measurement\.fields\.MeasurementField"])
    add_introspection_rules([], ["^django_measurement\.fields\.OriginalUnitField"])
    add_introspection_rules([], ["^django_measurement\.fields\.MeasureNameField"])
    add_introspection_rules([], ["^django_measurement\.fields\.MeasurementValueField"])
except ImportError:
    pass
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Author: Konstantinos Livieratos <livieratos.konstantinos@gmail.com>
# https://github.com/koslibpro/django-traffic/

import logging

import django
from django.conf import settings
from django.utils import timezone
from django.utils.deprecation import MiddlewareMixin
from django.contrib.gis.geoip2 import GeoIP2

from elasticsearch import Elasticsearch

class ESTrafficInfoMiddleware(MiddlewareMixin):
    def __init__(self, get_response=None):
        self.get_response = get_response

        if settings.TRAFFIC_INDEX_NAME:
            self.index_name = settings.TRAFFIC_INDEX_NAME
        else:
            self.index_name = "django-traffic"

        self.geo_db_path = settings.GEO_DB_PATH

        if settings.ES_CLIENT:
            self.es = settings.ES_CLIENT
        else:
            assert settings.ES_HOST, 'ES_HOST definition in settings.py is required'
            self.es = Elasticsearch(
                hosts=[settings.ES_HOST]
            )

        super(ESTrafficInfoMiddleware, self).__init__(get_response=get_response)

    def process_request(self, request, *args, **kwargs):
        self.es_upstream(request)
        return None

    def ip_to_cordinates(self, device_ip):
        if self.geo_db_path is None:
            # assume the user has set GEOIP_PATH in settings.py
            g = GeoIP2()
        else:
            g = GeoIP2(path=self.geo_db_path)
        try:
            lat, lng = g.lat_lon(device_ip)
        except Exception as e:
            logging.error("[django-traffic] Error while getting lan/lon from GEOIP2: %s " % e)
            return None, None
        else:
            return lat, lng

    def es_upstream(self, request):
        device_ip = request.META.get('HTTP_X_FORWARDED_FOR') or request.META.get('REMOTE_ADDR')
        device_ip = str(device_ip).split(',')[0]
        logging.info("[django-traffic] Device IP: %s" % device_ip)

        if not self.es.indices.exists(index=self.index_name):
            mapping = {
                "mappings": {
                    "request-info": {
                        "properties": {
                            "timestamp": {
                                "type": "date"
                            },
                            "text": {
                                "type": "string"
                            },
                            "location": {
                                "type": "geo_point"
                            },
                            "method": {
                                "type": "string"
                            },
                            "body": {
                                "type": "string"
                            },
                            "path": {
                                "type": "string"
                            },
                            "path_info": {
                                "type": "string"
                            },
                            "scheme": {
                                "type": "string"
                            },
                            "encoding": {
                                "type": "string"
                            },
                            "encoding_type": {
                                "type": "string"
                            },
                            "ip_addr": {
                                "type": "ip"
                            }
                        }
                    }
                }
            }
            logging.info("[django-traffic] Creating new elasticsearch indice...")
            self.es.indices.create(index=self.index_name, body=mapping)

        doc = {
            "timestamp": timezone.now(),
            "text": "django-traffic geo-point object",
            "method": request.method,
            "body": request.body,
            "path": request.path,
            "path_info": request.path_info,
            "scheme": request.scheme,
            "encoding": request.encoding,
            "encoding_type": getattr(request, 'encoding_type', ''),
            "ip_addr": device_ip
        }

        lat, lng = self.ip_to_cordinates(device_ip)
        if lat and lng:
            doc["location"] = {
                "lat": lat,
                "lon": lng
            }
        else:
            if not getattr(settings, 'LOG_WITHOUT_LOCATION', False):
                return

        res = self.es.index(index=self.index_name, doc_type='request-info', body=doc)
        if res['created']:
            logging.info("[django-traffic] request indexed")
<commit_msg>removed odd import<commit_after># -*- coding: utf-8 -*-
# Author: Konstantinos Livieratos <livieratos.konstantinos@gmail.com>
# https://github.com/koslibpro/django-traffic/

import logging

from django.conf import settings
from django.utils import timezone
from django.utils.deprecation import MiddlewareMixin
from django.contrib.gis.geoip2 import GeoIP2

from elasticsearch import Elasticsearch

class ESTrafficInfoMiddleware(MiddlewareMixin):
    def __init__(self, get_response=None):
        self.get_response = get_response

        if settings.TRAFFIC_INDEX_NAME:
            self.index_name = settings.TRAFFIC_INDEX_NAME
        else:
            self.index_name = "django-traffic"

        self.geo_db_path = settings.GEO_DB_PATH

        if settings.ES_CLIENT:
            self.es = settings.ES_CLIENT
        else:
            assert settings.ES_HOST, 'ES_HOST definition in settings.py is required'
            self.es = Elasticsearch(
                hosts=[settings.ES_HOST]
            )

        super(ESTrafficInfoMiddleware, self).__init__(get_response=get_response)

    def process_request(self, request, *args, **kwargs):
        self.es_upstream(request)
        return None

    def ip_to_cordinates(self, device_ip):
        if self.geo_db_path is None:
            # assume the user has set GEOIP_PATH in settings.py
            g = GeoIP2()
        else:
            g = GeoIP2(path=self.geo_db_path)
        try:
            lat, lng = g.lat_lon(device_ip)
        except Exception as e:
            logging.error("[django-traffic] Error while getting lan/lon from GEOIP2: %s " % e)
            return None, None
        else:
            return lat, lng

    def es_upstream(self, request):
        device_ip = request.META.get('HTTP_X_FORWARDED_FOR') or request.META.get('REMOTE_ADDR')
        device_ip = str(device_ip).split(',')[0]
        logging.info("[django-traffic] Device IP: %s" % device_ip)

        if not self.es.indices.exists(index=self.index_name):
            mapping = {
                "mappings": {
                    "request-info": {
                        "properties": {
                            "timestamp": {
                                "type": "date"
                            },
                            "text": {
                                "type": "string"
                            },
                            "location": {
                                "type": "geo_point"
                            },
                            "method": {
                                "type": "string"
                            },
                            "body": {
                                "type": "string"
                            },
                            "path": {
                                "type": "string"
                            },
                            "path_info": {
                                "type": "string"
                            },
                            "scheme": {
                                "type": "string"
                            },
                            "encoding": {
                                "type": "string"
                            },
                            "encoding_type": {
                                "type": "string"
                            },
                            "ip_addr": {
                                "type": "ip"
                            }
                        }
                    }
                }
            }
            logging.info("[django-traffic] Creating new elasticsearch indice...")
            self.es.indices.create(index=self.index_name, body=mapping)

        doc = {
            "timestamp": timezone.now(),
            "text": "django-traffic geo-point object",
            "method": request.method,
            "body": request.body,
            "path": request.path,
            "path_info": request.path_info,
            "scheme": request.scheme,
            "encoding": request.encoding,
            "encoding_type": getattr(request, 'encoding_type', ''),
            "ip_addr": device_ip
        }

        lat, lng = self.ip_to_cordinates(device_ip)
        if lat and lng:
            doc["location"] = {
                "lat": lat,
                "lon": lng
            }
        else:
            if not getattr(settings, 'LOG_WITHOUT_LOCATION', False):
                return

        res = self.es.index(index=self.index_name, doc_type='request-info', body=doc)
        if res['created']:
            logging.info("[django-traffic] request indexed")
<|endoftext|>


<commit_before># coding=utf-8
from django.db.models.query import QuerySet
from django.conf import settings
from django.contrib.auth.models import UserManager, AbstractBaseUser, PermissionsMixin
from django.core.mail import send_mail
from django.db import models
from django.utils import timezone
from django.utils.translation import ugettext_lazy as _
from django.core import validators

# basic model stuff
##########################################
class DataQuerySet(QuerySet):
    def delete(self):
        self.update(active=False, deleted=timezone.now())

class AllObjectsQuerySet(QuerySet):
    def undelete(self):
        self.update(active=True, deleted=None)

class DataManager(models.Manager):
    use_for_related_fields = True

    def get_queryset(self):
        qs = DataQuerySet(self.model, using=self._db)
        return qs.filter(active=True)

    def delete(self):
        self.update(active=False, deleted=timezone.now())

class AllObjectsManager(models.Manager):

    def get_queryset(self):
        return DataQuerySet(self.model, using=self._db)

    def deleted(self):
        return self.filter(active=False)

# base model with useful stuff
class BaseModel(models.Model):

    created = models.DateTimeField(auto_now_add=True, editable=False, db_index=True)
    modified = models.DateTimeField(auto_now=True, editable=False, db_index=True)
    deleted = models.DateTimeField(editable=False, null=True)

    active = models.BooleanField(default=True, editable=False, db_index=True)

    # access only active data objects
    data = DataManager()
    # access all (including deleted) data
    objects = AllObjectsManager()

    class Meta:
        abstract = True
        ordering = ['-created']

    # deleted data is bad - doing it you shouldn't!
    def delete(self, using=None):
        self.active = False
        self.save()

    def undelete(self, using=None):
        self.active = True
        self.save()

    def save(self, force_insert=False, force_update=False, using=None, update_fields=None):

        if self.active:
            self.deleted = None
        elif not self.deleted:
            self.deleted = timezone.now()

        super(BaseModel, self).save(force_insert=force_insert, using=using,
                                    update_fields=update_fields)

    def reload_me(self):
        if self.id:
            return self.__class__.objects.get(id=self.id)

    def pprint(self):
        pprint(self.__dict__)

class NamedModel(BaseModel):
    name = models.CharField('Name', max_length=150, db_index=True)

    class Meta:
        ordering = ['name']
        abstract = True

    def __unicode__(self):
        return self.name

class UserDataManager(UserManager, DataManager):
    pass

class User(AbstractBaseUser, BaseModel, PermissionsMixin):
    username = models.CharField(_('username'), max_length=30, unique=True,
                                help_text=_('user.login.username_help'),
                                validators=[
                                    validators.RegexValidator(
                                        r'^[\w.@+-]+$', _('forms.errors.enter_valid_username.'), 'username_invalid')
                                ])
    first_name = models.CharField(_('first name'), max_length=30, blank=True)
    last_name = models.CharField(_('last name'), max_length=30, blank=True)
    email = models.EmailField(_('email address'), blank=False)
    is_staff = models.BooleanField(_('staff status'), default=False,
                                   help_text=_('Designates whether the user can log into this admin site.'))
    is_active = models.BooleanField(_('active'), default=True,
                                    help_text=_('Designates whether this user should be treated as '
                                                'active. Unselect this instead of deleting accounts.'))
    date_joined = models.DateTimeField(_('date joined'), default=timezone.now)

    USERNAME_FIELD = 'username'
    REQUIRED_FIELDS = ['email']

    data = UserDataManager()
    objects = UserManager()

    class Meta:
        abstract = True
        verbose_name = _('user')
        verbose_name_plural = _('users')

    def get_full_name(self):
        """ Returns the first_name plus the last_name, with a space in between. """
        full_name = '%s %s' % (self.first_name, self.last_name)
        return full_name.strip()

    def get_short_name(self):
        """ Returns the short name for the user. """
        return self.first_name

    def email_user(self, subject, message, from_email=settings.DEFAULT_FROM_EMAIL, **kwargs):
        """
         Sends an email to this User.
         If settings.EMAIL_OVERRIDE_ADDRESS is set, this mail will be redirected to the alternate mail address.

        """
        receiver = self.email
        if settings.EMAIL_OVERRIDE_ADDRESS:
            receiver = settings.EMAIL_OVERRIDE_ADDRESS

        send_mail(subject, message, from_email, [receiver], **kwargs)
<commit_msg>me more compatible with python3 added ability to concael models instead of just being able to delete em and also the ability to revael em again changed active to is_deleted to be more compatible to other modules added a foce option to DataManager.delete() and Base)mode.delete() to actually delete stuff<commit_after># coding=utf-8
from __future__ import absolute_import, unicode_literals

from pprint import pprint

from django.db.models.query import QuerySet
from django.conf import settings
from django.contrib.auth.models import UserManager, AbstractBaseUser, PermissionsMixin
from django.core.mail import send_mail
from django.db import models
from django.utils import timezone
from django.utils.translation import ugettext_lazy as _
from django.core import validators

# basic model managers
##########################################
class DataQuerySet(QuerySet):
    def delete(self):
        self.update(is_deleted=True, deleted=timezone.now())

    def conceal(self):
        """
        Actually deleting data (without force=True) will cause DoesNotExist errors to
        appear more frequently in your project, because foreign keys are not deleted
        automatically for you anymore. But because the data manager is the default manager
        of your models, it cant load those models anymore and thus its saver to just hide
        models which are referenced in other models that might still be accessed afterwards.
        """
        self.update(concealed=True)

    def reveal(self):
        self.update(concealed=False)

class AllObjectsQuerySet(QuerySet):
    def undelete(self):
        self.update(is_deleted=False, deleted=None)

class DataManager(object, models.Manager):
    use_for_related_fields = True

    def get_queryset(self):
        qs = DataQuerySet(self.model, using=self._db)
        return qs.filter(is_deleted=False)

    def delete(self, force=False):
        if force:
            return super(DataManager, self).delete()
        else:
            return self.update(is_deleted=True, deleted=timezone.now())

    def visible(self):
        return self.filter(concealed=False)

class AllObjectsManager(object, models.Manager):
    def get_queryset(self):
        return DataQuerySet(self.model, using=self._db)

    def deleted(self):
        return self.filter(is_deleted=True)

# base model with useful stuff
##########################################
class BaseModel(models.Model):
    created = models.DateTimeField(_('created'), auto_now_add=True, editable=False, db_index=True)
    modified = models.DateTimeField(auto_now=True, editable=False)
    deleted = models.DateTimeField(editable=False, null=True)

    is_deleted = models.BooleanField(default=True, editable=False, db_index=True)
    concealed = models.BooleanField(default=False)

    # access non deleted data only
    data = DataManager()
    # access all (including deleted) data
    objects = AllObjectsManager()

    class Meta:
        abstract = True
        ordering = ['-created']

    # deleted data is bad - doing it you shouldn't! (but if u really want, u can)
    def delete(self, using=None, force=False):
        if force:
            super(BaseModel, self).delete(using=using)
        else:
            self.is_deleted = True
            self.save()

    def undelete(self, using=None):
        self.__class__.objects.filter(id=self.id).update(is_deleted=False, deleted=None)
        return self.reload_me()

    def save(self, force_insert=False, force_update=False, using=None, update_fields=None):

        if self.is_deleted:
            self.deleted = None
        elif not self.deleted:
            self.deleted = timezone.now()

        super(BaseModel, self).save(force_insert=force_insert, using=using,
                                    update_fields=update_fields)

    def reload_me(self):
        """ Django 1.8 introduced .refresh_from_db() - plz use this instead if possible :) """
        if self.id:
            return self.__class__.objects.get(id=self.id)

    def pprint(self):
        pprint(self.__dict__)

class NamedModel(BaseModel):
    name = models.CharField('Name', max_length=150, db_index=True)

    class Meta:
        ordering = ['name']
        abstract = True

    def __unicode__(self):
        return self.name

class UserDataManager(UserManager, DataManager):
    pass

# abstract base user with data manager
##########################################
class User(AbstractBaseUser, BaseModel, PermissionsMixin):
    username = models.CharField(
        _('username'), max_length=30, unique=True, help_text=_('user.login.username_help'),
        validators=[validators.RegexValidator(r'^[\w.@+-]+$', _('forms.errors.enter_valid_username.'), 'username_invalid')])
    first_name = models.CharField(_('first name'), max_length=30, blank=True)
    last_name = models.CharField(_('last name'), max_length=30, blank=True)
    email = models.EmailField(_('email address'), blank=False)
    is_staff = models.BooleanField(_('staff status'), default=False,
                                   help_text=_('Designates whether the user can log into this admin site.'))
    is_active = models.BooleanField(_('active'), default=True,
                                    help_text=_('Designates whether this user should be treated as '
                                                'active. Unselect this instead of deleting accounts.'))
    date_joined = models.DateTimeField(_('date joined'), default=timezone.now)

    USERNAME_FIELD = 'username'
    REQUIRED_FIELDS = ['email']

    data = UserDataManager()
    objects = UserManager()

    class Meta:
        abstract = True
        verbose_name = _('user')
        verbose_name_plural = _('users')

    def get_full_name(self):
        """ Returns the first_name plus the last_name, with a space in between. """
        full_name = '%s %s' % (self.first_name, self.last_name)
        return full_name.strip()

    def get_short_name(self):
        """ Returns the short name for the user. """
        return self.first_name

    def email_user(self, subject, message, from_email=settings.DEFAULT_FROM_EMAIL, **kwargs):
        """
         Sends an email to this User.
         If settings.EMAIL_OVERRIDE_ADDRESS is set, this mail will be redirected to the alternate mail address.

        """
        receiver = self.email
        if settings.EMAIL_OVERRIDE_ADDRESS:
            receiver = settings.EMAIL_OVERRIDE_ADDRESS

        send_mail(subject, message, from_email, [receiver], **kwargs)
<|endoftext|>


<commit_before>#!/usr/bin/env python
__author__ = 'chris'
description = """
Convert things to django
"""
import sys
import traceback
import os
import imp
import subprocess
import shutil
from argparse import ArgumentParser
from django.template import Context, Engine

from djangui.backend.nodes import ArgParseNodeBuilder

def main():
    parser = ArgumentParser(description=description)
    parser.add_argument('-s', '--scripts', help='A file or folder to look for scripts.', type=str, required=True)
    parser.add_argument('-r', '--recursive', help='Recursively search for scripts.', action='store_true')
    parser.add_argument('-a', '--app', help='The name of the app to create.', type=str, required=True)
    parser.add_argument('-p', '--project', help='The name of the django project to create.', type=str)
    parser.add_argument('--fake', help='Dont do anything.', action='store_true')
    args = parser.parse_args()
    scripts_path = args.scripts
    scripts = []
    if os.path.isdir(scripts_path):
        if args.recursive:
            for root, folders, files in os.path.walk(scripts_path):
                scripts += [os.path.join(root, filename) for filename in files if filename.endswith('.py')]
        else:
            scripts += [os.path.join(scripts_path, filename) for filename in os.listdir(scripts_path) if filename.endswith('.py')]
    else:
        scripts.append(scripts_path)

    project_name = args.project
    app_name = args.app
    app_base_dir = ''
    if project_name:
        if not args.fake:
            subprocess.call(['django-admin.py', 'startproject', project_name])
        project_root = os.path.join(app_base_dir, project_name)
        project_base_dir = os.path.join(project_root, project_name)
        app_base_dir = os.path.join(app_base_dir, project_name)
    app_base_dir = os.path.join(app_base_dir, app_name)

    if not args.fake:
        com = ['django-admin.py', 'startapp', app_name]
        if project_name:
            app_path = os.path.join(os.path.abspath(project_name), app_name)
            os.mkdir(app_path)
            com.append(app_path)
        print com
        subprocess.call(com)

    djanguify_folder = os.path.split(os.path.realpath(__file__))[0]
    app_template_dir = os.path.join(djanguify_folder, '..', 'conf', 'script_template')
    project_template_dir = os.path.join(djanguify_folder, '..', 'conf', 'project_template')

    app_models = []

    for script in scripts:
        arg_nodes = {}
        basename, extension = os.path.splitext(script)
        basename = basename
        filename = os.path.split(basename)[1]
        try:
            module = imp.load_source(basename, script)
        except:
            sys.stderr.write('Error while loading %s:\n'.format(script))
            sys.stderr.write('{0}\n'.format(traceback.format_exc()))
            continue
        module_parser = module.parser
        parser = ArgParseNodeBuilder(filename, module_parser, script)
        app_models.append(parser.getModelDict())

    context = Context(
        dict({
            'app_name': app_name,
            'project_name': project_name,
        },
        **{
            'models': app_models,
        }),
        autoescape=False
    )

    template_files = []
    def walk_dir(templates, dest):
        l = []
        for root, folders, files in os.walk(templates):
            for filename in files:
                if filename.endswith('.pyc'):
                    continue
                relative_dir = '.{0}'.format(os.path.split(os.path.join(root, filename).replace(templates, ''))[0])
                l.append((os.path.join(root, filename), os.path.join(dest, relative_dir)))
        return l

    template_files += walk_dir(app_template_dir, app_base_dir)
    template_files += walk_dir(project_template_dir, project_base_dir)

    # remove files of directories we are overriding
    os.remove(os.path.join(app_base_dir, 'models.py'))

    # move the django settings to the settings path so we don't have to chase Django changes.
    shutil.move(os.path.join(project_base_dir, 'settings.py'), os.path.join(project_base_dir, 'settings', 'django_settings.py'))

    for template_file, dest_dir in template_files:
        template_file = open(template_file)
        content = template_file.read()
        content = content.decode('utf-8')
        template = Engine().from_string(content)
        content = template.render(context)
        content = content.encode('utf-8')
        # import pdb; pdb.set_trace();
        to_name = os.path.join(dest_dir, os.path.split(template_file.name)[1])
        # print template_file, dest_dir, to_name
        try:
            os.mkdir(dest_dir)
        except:
            pass
        with open(to_name, 'wb') as new_file:
            new_file.write(content)

    if project_name:
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'makemigrations'])
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'migrate'])
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'collectstatic', '--noinput'])
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'runserver'])

if __name__ == "__main__":
    sys.exit(main())<commit_msg>bug fix in ordering of django settings<commit_after>#!/usr/bin/env python
__author__ = 'chris'
description = """
Convert things to django
"""
import sys
import traceback
import os
import imp
import subprocess
import shutil
from argparse import ArgumentParser
from django.template import Context, Engine

from djangui.backend.nodes import ArgParseNodeBuilder

def main():
    parser = ArgumentParser(description=description)
    parser.add_argument('-s', '--scripts', help='A file or folder to look for scripts.', type=str, required=True)
    parser.add_argument('-r', '--recursive', help='Recursively search for scripts.', action='store_true')
    parser.add_argument('-a', '--app', help='The name of the app to create.', type=str, required=True)
    parser.add_argument('-p', '--project', help='The name of the django project to create.', type=str)
    parser.add_argument('--fake', help='Dont do anything.', action='store_true')
    args = parser.parse_args()
    scripts_path = args.scripts
    scripts = []
    if os.path.isdir(scripts_path):
        if args.recursive:
            for root, folders, files in os.path.walk(scripts_path):
                scripts += [os.path.join(root, filename) for filename in files if filename.endswith('.py')]
        else:
            scripts += [os.path.join(scripts_path, filename) for filename in os.listdir(scripts_path) if filename.endswith('.py')]
    else:
        scripts.append(scripts_path)

    project_name = args.project
    app_name = args.app
    app_base_dir = ''
    if project_name:
        if not args.fake:
            subprocess.call(['django-admin.py', 'startproject', project_name])
        project_root = os.path.join(app_base_dir, project_name)
        project_base_dir = os.path.join(project_root, project_name)
        app_base_dir = os.path.join(app_base_dir, project_name)
    app_base_dir = os.path.join(app_base_dir, app_name)

    if not args.fake:
        com = ['django-admin.py', 'startapp', app_name]
        if project_name:
            app_path = os.path.join(os.path.abspath(project_name), app_name)
            os.mkdir(app_path)
            com.append(app_path)
        print com
        subprocess.call(com)

    djanguify_folder = os.path.split(os.path.realpath(__file__))[0]
    app_template_dir = os.path.join(djanguify_folder, '..', 'conf', 'script_template')
    project_template_dir = os.path.join(djanguify_folder, '..', 'conf', 'project_template')

    app_models = []

    for script in scripts:
        basename, extension = os.path.splitext(script)
        basename = basename
        filename = os.path.split(basename)[1]
        try:
            module = imp.load_source(basename, script)
        except:
            sys.stderr.write('Error while loading %s:\n'.format(script))
            sys.stderr.write('{0}\n'.format(traceback.format_exc()))
            continue
        module_parser = module.parser
        parser = ArgParseNodeBuilder(filename, module_parser, script)
        app_models.append(parser.getModelDict())

    context = Context(
        dict({
            'app_name': app_name,
            'project_name': project_name,
        },
        **{
            'models': app_models,
        }),
        autoescape=False
    )

    template_files = []
    def walk_dir(templates, dest):
        l = []
        for root, folders, files in os.walk(templates):
            for filename in files:
                if filename.endswith('.pyc'):
                    continue
                relative_dir = '.{0}'.format(os.path.split(os.path.join(root, filename).replace(templates, ''))[0])
                l.append((os.path.join(root, filename), os.path.join(dest, relative_dir)))
        return l

    template_files += walk_dir(app_template_dir, app_base_dir)
    template_files += walk_dir(project_template_dir, project_base_dir)

    # remove files of directories we are overriding
    os.remove(os.path.join(app_base_dir, 'models.py'))

    for template_file, dest_dir in template_files:
        template_file = open(template_file)
        content = template_file.read()
        content = content.decode('utf-8')
        template = Engine().from_string(content)
        content = template.render(context)
        content = content.encode('utf-8')
        # import pdb; pdb.set_trace();
        to_name = os.path.join(dest_dir, os.path.split(template_file.name)[1])
        # print template_file, dest_dir, to_name
        try:
            os.mkdir(dest_dir)
        except:
            pass
        with open(to_name, 'wb') as new_file:
            new_file.write(content)

    # move the django settings to the settings path so we don't have to chase Django changes.
    shutil.move(os.path.join(project_base_dir, 'settings.py'), os.path.join(project_base_dir, 'settings', 'django_settings.py'))

    if project_name:
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'makemigrations'])
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'migrate'])
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'collectstatic', '--noinput'])
        subprocess.call(['python', os.path.join(project_root, 'manage.py'), 'runserver'])

if __name__ == "__main__":
    sys.exit(main())<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.core.exceptions import ValidationError
import re
from django.core import validators
from django.utils.translation import ugettext_lazy as _
from dk.collections import pset
from collections import defaultdict
from django.db.models.fields import Field
from django.db import models

class StatusValue(object):
    def __init__(self, name=None, verbose=None, categories=()):
        self.name = name.strip()
        self.verbose = verbose.strip()
        if isinstance(categories, basestring):
            self.categories = re.split(r'[,\s]+', categories)
        else:
            self.categories = categories

    def __unicode__(self):
        return self.verbose

    def __str__(self):
        return self.name

    def __repr__(self):
        return 'StatusValue(name=%r, verbose=%r, categories=%r)' % (
            self.name, self.verbose, self.categories)

class StatusDef(object):
    """Compact way to define status fields, that takes care of most of
       the manual drudgery involved.

       Usage::

           class MyModel(models.Model):
               ..fields..
               STATUSDEF = statusdef.StatusDef(u'''

                    @start-tagname
                    =========== =================================== ==========
                    status      verbose explanation                 category
                    =========== =================================== ==========
                    ok          Ny Bestilling                       # [init]
                    locked      Last (legges til ny jobb)           # [init]
                    produseres  Klar for QA (lagt til ny jobb)      # [ready]
                    hentet      Kort trykkes                        # [done]
                    sendt       Kort er trykket og sendt med posten # [done]
                    dblbest     Fjernet fordi det var flere         # [err]
                    fnr-feil    Feil i fnr                          # [err]
                    =========== =================================== ==========
                    @end-tagname
               ''')
               status = models.CharField(max_length=10,  # max length of status
                                         choices=STATUSDEF.options)

       @something and @something-else can be used with the sphinx
       `include <http://docutils.sourceforge.net/docs/ref/rst/directives.html#include>`_
       directive's :start-after: and :end-before: options.

       Empty lines, and lines starting with `@` are disregarded.

       You are not required to format the string as a rst table, but it seems
       foolish not to do so.

    """

    # TODO:  create a custom model field that maps to <select>

    defre = re.compile(ur'''
    \s*(?P<name>[a-z][-a-z0-9]*)\s*(?P<verbose>[^\#]*)\#\s*\[(?P<categories>[^]]*)\]
    ''', re.VERBOSE)

    # noinspection PyMethodMayBeStatic
    def _parse(self, txt):  # pylint:disable=R0201
        lines = [line for line in txt.split('\n') if line.strip()]
        defs = pset()

        in_header = False  # a state-machine...

        for line in lines:
            line = line.strip()

            # skip header lines
            if in_header and not line.startswith('='):
                continue
            if line.startswith('@'):
                continue
            if line.startswith('=') and not in_header:
                in_header = True
                continue
            if line.startswith('=') and in_header:
                in_header = False
                continue

            m = StatusDef.defre.match(line)
            if m:
                gdict = pset((str(k), v)
                             for (k, v) in m.groupdict().items())
                sval = StatusValue(name=gdict.name,
                                   verbose=gdict.verbose,
                                   categories=gdict.categories)
                defs[sval.name] = sval
            else:
                print 'error parsing:', repr(line)

        return defs

    def __init__(self, txt):
        self.status = self._parse(txt)
        self._defs = self.status.values()
        self._categories = set()
        for d in self._defs:
            for cat in d.categories:
                self._categories.add(cat)
        self._cat2status = defaultdict(set)
        for d in self._defs:
            for cat in d.categories:
                self._cat2status[cat].add(d)

    @property
    def namelength(self):
        if not self._defs:
            return 0
        return max(len(d.name) for d in self._defs)

    def is_category(self, txt):
        return txt in self._categories

    def category(self, status):
        """Return the category that status belongs to.
        """
        return self.status[status].category

    def category2status(self, category):
        """Return all statuses belonging to `category`.
        """
        return self._cat2status[category]

    def valid_status(self, s):
        """Is `s` a well-defined status value?
        """
        return s in self.status.keys()

    @property
    def options(self):
        """Return list of pairs of (status, description), useful for
           selects boxes etc.
        """
        return [(name, gdict.verbose) for name, gdict in self.status]

class StatusField(Field):
    """Character status field.
    """
    __metaclass__ = models.SubfieldBase

    description = _("Status field")

    def __init__(self, *args, **kw):
        self.txt = args[0] if args else ""
        self.statusdef = StatusDef(self.txt)
        self.max_length = kw['max_length'] = kw.get('max_length', self.statusdef.namelength)
        kw['choices'] = self.statusdef.options
        super(StatusField, self).__init__(**kw)
        self.validators.append(validators.MaxLengthValidator(self.max_length))

    def to_python(self, value):
        """Converts the input ``value`` into a StatusValue instance,
           raising ValueError if the data can't be converted.
        """
        if not value:
            return None

        if isinstance(value, StatusValue):
            return value

        if isinstance(value, basestring):
            if value in self.statusdef.status:
                return self.statusdef.status[value]
            raise ValueError("Unknown status: %r", value)

        return value

    def get_internal_type(self):
        return "StatusField"

    def db_type(self, connection):
        return 'VARCHAR(%s)' % self.max_length

    def get_prep_lookup(self, lookup_type, value):
        """Return a value prepared for database lookup.
        """
        res = set()
        if lookup_type == 'in':
            if isinstance(value, basestring):
                if self.statusdef.is_category(value):
                    res |= self.statusdef.category2status(value)
                else:
                    res.add(value)
            else:
                for v in value:
                    if self.statusdef.is_category(v):
                        res |= self.statusdef.category2status(v)
                    else:
                        res.add(v)
            res = [self.get_prep_value(v) for v in res]
        else:
            res.add(value)
        return res

    def get_prep_value(self, value):
        """Convert to a value useable as a parameter in a query.
        """
        return self.to_python(value).name

    def formfield(self, **kwargs):
        # Passing max_length to forms.CharField means that the value's length
        # will be validated twice. This is considered acceptable since we want
        # the value in the form field (to pass into widget for example).
        defaults = {'max_length': self.max_length}
        defaults.update(kwargs)
        return super(StatusField, self).formfield(**defaults)

# from south.modelsinspector import add_introspection_rules
# add_introspection_rules([],["^dkmodelfields\.statusfield\.StatusField"])
<commit_msg>statusfield needs deconstruct method for django 1.7<commit_after># -*- coding: utf-8 -*-
from django.core.exceptions import ValidationError
import re
from django.core import validators
from django.utils.translation import ugettext_lazy as _
from dk.collections import pset
from collections import defaultdict
from django.db.models.fields import Field
from django.db import models

class StatusValue(object):
    def __init__(self, name=None, verbose=None, categories=()):
        self.name = name.strip()
        self.verbose = verbose.strip()
        if isinstance(categories, basestring):
            self.categories = re.split(r'[,\s]+', categories)
        else:
            self.categories = categories

    def __unicode__(self):
        return self.verbose

    def __str__(self):
        return self.name

    def __repr__(self):
        return 'StatusValue(name=%r, verbose=%r, categories=%r)' % (
            self.name, self.verbose, self.categories)

class StatusDef(object):
    """Compact way to define status fields, that takes care of most of
       the manual drudgery involved.

       Usage::

           class MyModel(models.Model):
               ..fields..
               STATUSDEF = statusdef.StatusDef(u'''

                    @start-tagname
                    =========== =================================== ==========
                    status      verbose explanation                 category
                    =========== =================================== ==========
                    ok          Ny Bestilling                       # [init]
                    locked      Last (legges til ny jobb)           # [init]
                    produseres  Klar for QA (lagt til ny jobb)      # [ready]
                    hentet      Kort trykkes                        # [done]
                    sendt       Kort er trykket og sendt med posten # [done]
                    dblbest     Fjernet fordi det var flere         # [err]
                    fnr-feil    Feil i fnr                          # [err]
                    =========== =================================== ==========
                    @end-tagname
               ''')
               status = models.CharField(max_length=10,  # max length of status
                                         choices=STATUSDEF.options)

       @something and @something-else can be used with the sphinx
       `include <http://docutils.sourceforge.net/docs/ref/rst/directives.html#include>`_
       directive's :start-after: and :end-before: options.

       Empty lines, and lines starting with `@` are disregarded.

       You are not required to format the string as a rst table, but it seems
       foolish not to do so.

    """

    # TODO:  create a custom model field that maps to <select>

    defre = re.compile(ur'''
    \s*(?P<name>[a-z][-a-z0-9]*)\s*(?P<verbose>[^\#]*)\#\s*\[(?P<categories>[^]]*)\]
    ''', re.VERBOSE)

    # noinspection PyMethodMayBeStatic
    def _parse(self, txt):  # pylint:disable=R0201
        lines = [line for line in txt.split('\n') if line.strip()]
        defs = pset()

        in_header = False  # a state-machine...

        for line in lines:
            line = line.strip()

            # skip header lines
            if in_header and not line.startswith('='):
                continue
            if line.startswith('@'):
                continue
            if line.startswith('=') and not in_header:
                in_header = True
                continue
            if line.startswith('=') and in_header:
                in_header = False
                continue

            m = StatusDef.defre.match(line)
            if m:
                gdict = pset((str(k), v)
                             for (k, v) in m.groupdict().items())
                sval = StatusValue(name=gdict.name,
                                   verbose=gdict.verbose,
                                   categories=gdict.categories)
                defs[sval.name] = sval
            else:
                print 'error parsing:', repr(line)

        return defs

    def __init__(self, txt):
        self.status = self._parse(txt)
        self._defs = self.status.values()
        self._categories = set()
        for d in self._defs:
            for cat in d.categories:
                self._categories.add(cat)
        self._cat2status = defaultdict(set)
        for d in self._defs:
            for cat in d.categories:
                self._cat2status[cat].add(d)

    @property
    def namelength(self):
        if not self._defs:
            return 0
        return max(len(d.name) for d in self._defs)

    def is_category(self, txt):
        return txt in self._categories

    def category(self, status):
        """Return the category that status belongs to.
        """
        return self.status[status].category

    def category2status(self, category):
        """Return all statuses belonging to `category`.
        """
        return self._cat2status[category]

    def valid_status(self, s):
        """Is `s` a well-defined status value?
        """
        return s in self.status.keys()

    @property
    def options(self):
        """Return list of pairs of (status, description), useful for
           selects boxes etc.
        """
        return [(name, gdict.verbose) for name, gdict in self.status]

class StatusField(Field):
    """Character status field.
    """
    __metaclass__ = models.SubfieldBase

    description = _("Status field")

    def __init__(self, *args, **kw):
        self.txt = args[0] if args else ""
        self.statusdef = StatusDef(self.txt)
        self.max_length = kw['max_length'] = kw.get('max_length', self.statusdef.namelength)
        kw['choices'] = self.statusdef.options
        super(StatusField, self).__init__(**kw)
        self.validators.append(validators.MaxLengthValidator(self.max_length))

    def deconstruct(self):
        name, path, args, kwargs = super(StatusField, self).deconstruct()
        del kwargs['max_length']
        kwargs['choices'] = self.statusdef.options
        return name, path, args, kwargs

    def to_python(self, value):
        """Converts the input ``value`` into a StatusValue instance,
           raising ValueError if the data can't be converted.
        """
        if not value:
            return None

        if isinstance(value, StatusValue):
            return value

        if isinstance(value, basestring):
            if value in self.statusdef.status:
                return self.statusdef.status[value]
            raise ValueError("Unknown status: %r", value)

        return value

    def get_internal_type(self):
        return "StatusField"

    def db_type(self, connection):
        return 'VARCHAR(%s)' % self.max_length

    def get_prep_lookup(self, lookup_type, value):
        """Return a value prepared for database lookup.
        """
        res = set()
        if lookup_type == 'in':
            if isinstance(value, basestring):
                if self.statusdef.is_category(value):
                    res |= self.statusdef.category2status(value)
                else:
                    res.add(value)
            else:
                for v in value:
                    if self.statusdef.is_category(v):
                        res |= self.statusdef.category2status(v)
                    else:
                        res.add(v)
            res = [self.get_prep_value(v) for v in res]
        else:
            res.add(value)
        return res

    def get_prep_value(self, value):
        """Convert to a value useable as a parameter in a query.
        """
        if value is None:
            return value
        return self.to_python(value).name

    def formfield(self, **kwargs):
        # Passing max_length to forms.CharField means that the value's length
        # will be validated twice. This is considered acceptable since we want
        # the value in the form field (to pass into widget for example).
        defaults = {'max_length': self.max_length}
        defaults.update(kwargs)
        return super(StatusField, self).formfield(**defaults)

# from south.modelsinspector import add_introspection_rules
# add_introspection_rules([],["^dkmodelfields\.statusfield\.StatusField"])
<|endoftext|>


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

from __future__ import unicode_literals

import django

from django.db import IntegrityError, models
from django.template.defaultfilters import slugify as default_slugify
from django.utils.encoding import python_2_unicode_compatible
from django.utils.translation import ugettext_lazy as _

from aldryn_translation_tools.models import (
    TranslatedAutoSlugifyMixin, TranslationHelperMixin)
from parler import appsettings
from parler.managers import TranslatableManager, TranslatableQuerySet
from parler.models import TranslatableModel, TranslatedFields
from treebeard.ns_tree import NS_Node, NS_NodeManager, NS_NodeQuerySet

LANGUAGE_CODES = appsettings.PARLER_LANGUAGES.get_active_choices()

class CategoryQuerySet(TranslatableQuerySet, NS_NodeQuerySet):
    pass

class CategoryManager(TranslatableManager, NS_NodeManager):
    queryset_class = CategoryQuerySet

    def get_queryset(self):
        return self.queryset_class(
            self.model,
            using=self._db
        ).order_by('tree_id', 'lft')

    if django.VERSION < (1, 8):
        get_query_set = get_queryset

#
# TODO: I would have preferred to make this a base class "CategoryBase" which
# is Abstract, then subclass it as a concrete Category class. But, Parler
# cannot be applied to an Abstract class.
#
# TODO: At some point, consider an approach like this:
#     https://gist.github.com/GaretJax/7c7a9acc055c05c65041
#
@python_2_unicode_compatible
class Category(TranslatedAutoSlugifyMixin, TranslationHelperMixin,
               TranslatableModel, NS_Node):
    """
    A category is hierarchical. The structure is implemented with django-
    treebeard's Nested Sets trees, which has the performance characteristics
    we're after, namely: fast reads at the expense of write-speed.
    """
    slug_source_field_name = 'name'

    translations = TranslatedFields(
        name=models.CharField(
            _('name'),
            blank=False,
            default='',
            max_length=255,
        ),
        slug=models.SlugField(
            _('slug'),
            blank=True,
            default='',
            help_text=_('Provide a “slug” or leave blank for an automatically '
                        'generated one.'),
            max_length=255,
        ),
        meta={'unique_together': (('language_code', 'slug', ), )}
    )

    class Meta:
        verbose_name = _('category')
        verbose_name_plural = _('categories')

    objects = CategoryManager()

    def delete(self, using=None):
        #
        # We're simply managing how the two superclasses perform deletion
        # together here.
        #
        self.__class__.objects.filter(pk=self.pk).delete(using)
        super(TranslatableModel, self).delete()

    def __str__(self):
        return self.safe_translation_getter('name', any_language=True)
<commit_msg>Remove unused imports from models.py<commit_after># -*- coding: utf-8 -*-

from __future__ import unicode_literals

import django

from django.db import models
from django.utils.encoding import python_2_unicode_compatible
from django.utils.translation import ugettext_lazy as _

from aldryn_translation_tools.models import (
    TranslatedAutoSlugifyMixin, TranslationHelperMixin)
from parler import appsettings
from parler.managers import TranslatableManager, TranslatableQuerySet
from parler.models import TranslatableModel, TranslatedFields
from treebeard.ns_tree import NS_Node, NS_NodeManager, NS_NodeQuerySet

LANGUAGE_CODES = appsettings.PARLER_LANGUAGES.get_active_choices()

class CategoryQuerySet(TranslatableQuerySet, NS_NodeQuerySet):
    pass

class CategoryManager(TranslatableManager, NS_NodeManager):
    queryset_class = CategoryQuerySet

    def get_queryset(self):
        return self.queryset_class(
            self.model,
            using=self._db
        ).order_by('tree_id', 'lft')

    if django.VERSION < (1, 8):
        get_query_set = get_queryset

#
# TODO: I would have preferred to make this a base class "CategoryBase" which
# is Abstract, then subclass it as a concrete Category class. But, Parler
# cannot be applied to an Abstract class.
#
# TODO: At some point, consider an approach like this:
#     https://gist.github.com/GaretJax/7c7a9acc055c05c65041
#
@python_2_unicode_compatible
class Category(TranslatedAutoSlugifyMixin, TranslationHelperMixin,
               TranslatableModel, NS_Node):
    """
    A category is hierarchical. The structure is implemented with django-
    treebeard's Nested Sets trees, which has the performance characteristics
    we're after, namely: fast reads at the expense of write-speed.
    """
    slug_source_field_name = 'name'

    translations = TranslatedFields(
        name=models.CharField(
            _('name'),
            blank=False,
            default='',
            max_length=255,
        ),
        slug=models.SlugField(
            _('slug'),
            blank=True,
            default='',
            help_text=_('Provide a “slug” or leave blank for an automatically '
                        'generated one.'),
            max_length=255,
        ),
        meta={'unique_together': (('language_code', 'slug', ), )}
    )

    class Meta:
        verbose_name = _('category')
        verbose_name_plural = _('categories')

    objects = CategoryManager()

    def delete(self, using=None):
        #
        # We're simply managing how the two superclasses perform deletion
        # together here.
        #
        self.__class__.objects.filter(pk=self.pk).delete(using)
        super(TranslatableModel, self).delete()

    def __str__(self):
        return self.safe_translation_getter('name', any_language=True)
<|endoftext|>


<commit_before>from datetime import datetime, timedelta
import pytz
from typing import Any, Callable, Dict, List, NamedTuple, Optional

from mozdef_util.elasticsearch_client import ElasticsearchClient as ESClient
from mozdef_util.utilities.toUTC import toUTC
from mozdef_util.query_models import SearchQuery, TermMatch

# Default radius (in Kilometres) that a locality should have.
_DEFAULT_RADIUS_KM = 50.0

# TODO: Switch to dataclasses when we move to Python3.7+

class Locality(NamedTuple):
    '''Represents a specific locality.
    '''

    sourceipaddress: str
    city: str
    country: str
    lastaction: datetime
    latitude: float
    longitude: float
    radius: int

class State(NamedTuple):
    '''Represents the state tracked for each user regarding their localities.
    '''

    type_: str  # Not used here, but we want it in ES.
    username: str
    localities: List[Locality]

class Entry(NamedTuple):
    '''A top-level container for locality state that will be inserted into
    ElasticSearch.
    The `identifier` field here is the `_id` field of the ES document.  When
    this id is `None`, a new document is inserted whereas when the id is known,
    the existing document is updated.
    '''

    identifier: Optional[str]
    state: State

    def new(state: State) -> 'Entry':
        '''Construct a new `Entry` that, when journaled, will result in a new
        state entry being recorded rather than replacing an existing one.
        '''

        return Entry('', state)

class Update(NamedTuple):
    '''Produced by calls to functions operating on lists of `State`s to
    indicate when an update was applied without having to maintain distinct
    lists.
    '''

    state: State
    did_update: bool

    def flat_map(fn: Callable[[State], 'Update'], u: 'Update') -> 'Update':
        '''Apply a function to a `State` that produces an `Update` against the
        state contained within an established `Update`.  The resulting `Update`
        will have its `did_update` field set to `True` if either the original
        or the new `Update` are `True`.
        '''

        new = fn(u.state)

        return Update(new.state, u.did_update or new.did_update)

JournalInterface = Callable[[Entry, str], None]
QueryInterface = Callable[[SearchQuery, str], Optional[Entry]]

def _dict_take(dictionary, keys):
    return {key: dictionary[key] for key in keys}

def wrap_journal(client: ESClient) -> JournalInterface:
    '''Wrap an `ElasticsearchClient` in a closure of type `JournalInterface`.
    '''

    def wrapper(entry: Entry, esindex: str):
        document = dict(entry.state._asdict())

        client.save_object(
            index=esindex,
            body=document,
            doc_id=entry.identifier)

    return wrapper

def wrap_query(client: ESClient) -> QueryInterface:
    '''Wrap an `ElasticsearchClient` in a closure of type `QueryInterface`.
    '''

    def wrapper(query: SearchQuery, esindex: str) -> Optional[Entry]:
        results = query.execute(client, indices=[esindex]).get('hits', [])

        if len(results) == 0:
            return None

        state_dict = results[0].get('_source', {})
        try:
            state_dict['localities'] = [
                # Convert dictionary localities into `Locality`s after
                # parsing the `datetime` from `lastaction`.
                Locality(**_dict_take({
                    k: v if k != 'lastaction' else _parse_datetime(v)
                    for k, v in loc.items()
                }, Locality._fields))
                for loc in state_dict['localities']
            ]

            eid = results[0]['_id']
            state = State(**_dict_take(state_dict, State._fields))

            return Entry(eid, state)
        except TypeError:
            return None
        except KeyError:
            return None

    return wrapper

def from_event(
        event: Dict[str, Any],
        radius=_DEFAULT_RADIUS_KM
) -> Optional[Locality]:
    '''Extract locality information from an event if it is present in order
    to produce a `Locality` for which an authenticated action was taken.
    '''

    _source = event.get('_source', {})

    source_ip = _source.get('sourceipaddress')
    geo_data = _source.get('sourceipgeolocation')

    if source_ip is None or geo_data is None:
        return None

    now = toUTC(datetime.now()).isoformat()
    active_time_str = _source.get('utctimestamp', now)
    active_time = _parse_datetime(active_time_str)

    (city, country, lat, lon) = (
        geo_data.get('city'),
        geo_data.get('country_code'),
        geo_data.get('latitude'),
        geo_data.get('longitude')
    )

    if any([v is None for v in [city, country, lat, lon]]):
        return None

    return Locality(source_ip, city, country, active_time, lat, lon, radius)

def find(qes: QueryInterface, username: str, index: str) -> Optional[Entry]:
    '''Retrieve the locality state for one user from ElasticSearch.
    '''

    search = SearchQuery()
    search.add_must([
        TermMatch('type_', 'locality'),
        TermMatch('username', username)
    ])

    return qes(search, index)

def update(state: State, from_evt: State) -> Update:
    '''Update the localities stored under an existing `State` against those
    contained in a new `State` constructed from events.
    '''

    did_update = False

    for loc1 in from_evt.localities:
        did_find = False

        for index, loc2 in enumerate(state.localities):
            # If we find that the new state's locality has been recorded
            # for the user in question, we only want to update it if either
            # their IP changed or the new time of activity is more recent.
            if loc1.city == loc2.city and loc1.country == loc2.country:
                did_find = True

                new_more_recent = loc1.lastaction > loc2.lastaction
                new_ip = loc1.sourceipaddress != loc2.sourceipaddress

                if new_more_recent or new_ip:
                    state.localities[index] = loc1
                    did_update = True

                # Stop looking for the locality in the records pulled from ES.
                break

        if not did_find:
            state.localities.append(loc1)
            did_update = True

    return Update(state, did_update)

def remove_outdated(state: State, days_valid: int) -> Update:
    '''Update a state by removing localities that are outdated, determined
    by checking if the last activity within a given locality was at least
    some number of days ago.
    '''

    now = toUTC(datetime.now())
    last_valid_date = now - timedelta(days=days_valid)

    new_localities = [
        loc
        for loc in state.localities
        if loc.lastaction >= last_valid_date
    ]

    return Update(
        state=State(state.type_, state.username, new_localities),
        did_update=len(new_localities) != len(state.localities))

def _parse_datetime(datetime_str):
    parsers = [
        lambda s: datetime.strptime(
            ''.join(s.rsplit(':', 1)),
            '%Y-%m-%dT%H:%M:%S%z'),
        lambda s: datetime.strptime(
            ''.join(s.rsplit(':', 1)),
            '%Y-%m-%dT%H:%M:%S.%f%z'),
        lambda s: datetime.strptime(s, '%Y-%m-%dT%H:%M:%S%z'),
        lambda s: datetime.strptime(s, '%Y-%m-%dT%H:%M:%S.%f%z'),
        lambda s: datetime.strptime(s, '%Y-%m-%dT%H:%M:%S'),
        lambda s: datetime.strptime(s, '%Y-%m-%dT%H:%M:%S.%f')
    ]

    for parser in parsers:
        try:
            return parser(datetime_str).replace(tzinfo=pytz.UTC)
        except ValueError:
            continue

    raise ValueError(
        'Unexpected datetime string format: {}'.format(datetime_str))
<commit_msg>Use toUTC instead of our own datetime parser<commit_after>from datetime import datetime, timedelta
import pytz
from typing import Any, Callable, Dict, List, NamedTuple, Optional

from mozdef_util.elasticsearch_client import ElasticsearchClient as ESClient
from mozdef_util.utilities.toUTC import toUTC
from mozdef_util.query_models import SearchQuery, TermMatch

# Default radius (in Kilometres) that a locality should have.
_DEFAULT_RADIUS_KM = 50.0

# TODO: Switch to dataclasses when we move to Python3.7+

class Locality(NamedTuple):
    '''Represents a specific locality.
    '''

    sourceipaddress: str
    city: str
    country: str
    lastaction: datetime
    latitude: float
    longitude: float
    radius: int

class State(NamedTuple):
    '''Represents the state tracked for each user regarding their localities.
    '''

    type_: str  # Not used here, but we want it in ES.
    username: str
    localities: List[Locality]

class Entry(NamedTuple):
    '''A top-level container for locality state that will be inserted into
    ElasticSearch.
    The `identifier` field here is the `_id` field of the ES document.  When
    this id is `None`, a new document is inserted whereas when the id is known,
    the existing document is updated.
    '''

    identifier: Optional[str]
    state: State

    def new(state: State) -> 'Entry':
        '''Construct a new `Entry` that, when journaled, will result in a new
        state entry being recorded rather than replacing an existing one.
        '''

        return Entry('', state)

class Update(NamedTuple):
    '''Produced by calls to functions operating on lists of `State`s to
    indicate when an update was applied without having to maintain distinct
    lists.
    '''

    state: State
    did_update: bool

    def flat_map(fn: Callable[[State], 'Update'], u: 'Update') -> 'Update':
        '''Apply a function to a `State` that produces an `Update` against the
        state contained within an established `Update`.  The resulting `Update`
        will have its `did_update` field set to `True` if either the original
        or the new `Update` are `True`.
        '''

        new = fn(u.state)

        return Update(new.state, u.did_update or new.did_update)

JournalInterface = Callable[[Entry, str], None]
QueryInterface = Callable[[SearchQuery, str], Optional[Entry]]

def _dict_take(dictionary, keys):
    return {key: dictionary[key] for key in keys}

def wrap_journal(client: ESClient) -> JournalInterface:
    '''Wrap an `ElasticsearchClient` in a closure of type `JournalInterface`.
    '''

    def wrapper(entry: Entry, esindex: str):
        document = dict(entry.state._asdict())

        client.save_object(
            index=esindex,
            body=document,
            doc_id=entry.identifier)

    return wrapper

def wrap_query(client: ESClient) -> QueryInterface:
    '''Wrap an `ElasticsearchClient` in a closure of type `QueryInterface`.
    '''

    def wrapper(query: SearchQuery, esindex: str) -> Optional[Entry]:
        results = query.execute(client, indices=[esindex]).get('hits', [])

        if len(results) == 0:
            return None

        state_dict = results[0].get('_source', {})
        try:
            state_dict['localities'] = [
                # Convert dictionary localities into `Locality`s after
                # parsing the `datetime` from `lastaction`.
                Locality(**_dict_take({
                    k: v if k != 'lastaction' else toUTC(v)
                    for k, v in loc.items()
                }, Locality._fields))
                for loc in state_dict['localities']
            ]

            eid = results[0]['_id']
            state = State(**_dict_take(state_dict, State._fields))

            return Entry(eid, state)
        except TypeError:
            return None
        except KeyError:
            return None

    return wrapper

def from_event(
        event: Dict[str, Any],
        radius=_DEFAULT_RADIUS_KM
) -> Optional[Locality]:
    '''Extract locality information from an event if it is present in order
    to produce a `Locality` for which an authenticated action was taken.
    '''

    _source = event.get('_source', {})

    source_ip = _source.get('sourceipaddress')
    geo_data = _source.get('sourceipgeolocation')

    if source_ip is None or geo_data is None:
        return None

    now = toUTC(datetime.now()).isoformat()
    active_time_str = _source.get('utctimestamp', now)
    active_time = toUTC(active_time_str)

    (city, country, lat, lon) = (
        geo_data.get('city'),
        geo_data.get('country_code'),
        geo_data.get('latitude'),
        geo_data.get('longitude')
    )

    if any([v is None for v in [city, country, lat, lon]]):
        return None

    return Locality(source_ip, city, country, active_time, lat, lon, radius)

def find(qes: QueryInterface, username: str, index: str) -> Optional[Entry]:
    '''Retrieve the locality state for one user from ElasticSearch.
    '''

    search = SearchQuery()
    search.add_must([
        TermMatch('type_', 'locality'),
        TermMatch('username', username)
    ])

    return qes(search, index)

def update(state: State, from_evt: State) -> Update:
    '''Update the localities stored under an existing `State` against those
    contained in a new `State` constructed from events.
    '''

    did_update = False

    for loc1 in from_evt.localities:
        did_find = False

        for index, loc2 in enumerate(state.localities):
            # If we find that the new state's locality has been recorded
            # for the user in question, we only want to update it if either
            # their IP changed or the new time of activity is more recent.
            if loc1.city == loc2.city and loc1.country == loc2.country:
                did_find = True

                new_more_recent = loc1.lastaction > loc2.lastaction
                new_ip = loc1.sourceipaddress != loc2.sourceipaddress

                if new_more_recent or new_ip:
                    state.localities[index] = loc1
                    did_update = True

                # Stop looking for the locality in the records pulled from ES.
                break

        if not did_find:
            state.localities.append(loc1)
            did_update = True

    return Update(state, did_update)

def remove_outdated(state: State, days_valid: int) -> Update:
    '''Update a state by removing localities that are outdated, determined
    by checking if the last activity within a given locality was at least
    some number of days ago.
    '''

    now = toUTC(datetime.now())
    last_valid_date = now - timedelta(days=days_valid)

    new_localities = [
        loc
        for loc in state.localities
        if loc.lastaction >= last_valid_date
    ]

    return Update(
        state=State(state.type_, state.username, new_localities),
        did_update=len(new_localities) != len(state.localities))
<|endoftext|>