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

import re
import sys
import time
import threading
from math import ceil

import six
if six.PY2:
    from backports.shutil_get_terminal_size import get_terminal_size
else:
    from shutil import get_terminal_size

last_output_lines = 0
overflow_flag = False
is_atty = sys.stdout.isatty()

widths = [
    (126,    1), (159,    0), (687,     1), (710,   0), (711,   1),
    (727,    0), (733,    1), (879,     0), (1154,  1), (1161,  0),
    (4347,   1), (4447,   2), (7467,    1), (7521,  0), (8369,  1),
    (8426,   0), (9000,   1), (9002,    2), (11021, 1), (12350, 2),
    (12351,  1), (12438,  2), (12442,   0), (19893, 2), (19967, 1),
    (55203,  2), (63743,  1), (64106,   2), (65039, 1), (65059, 0),
    (65131,  2), (65279,  1), (65376,   2), (65500, 1), (65510, 2),
    (120831, 1), (262141, 2), (1114109, 1),
]

def get_char_width(char):
    global widths
    o = ord(char)
    if o == 0xe or o == 0xf:
        return 0
    for num, wid in widths:
        if o <= num:
            return wid
    return 1

def preprocess(content):
    """
    对输出内容进行预处理，转为str类型，并替换行内\r\t\n等字符为空格
    do pre-process to the content, turn it into str, and replace \r\t\n with space
    """

    if six.PY2:
        _content = unicode(content, encoding=sys.stdin.encoding) if isinstance(content, str) else content
        assert isinstance(_content, unicode)
    elif six.PY3:
        _content = str(content)

    _content = re.sub(r'\r|\t|\n', ' ', _content)

    return _content

def cut_off_at(content, width):
    if line_width(content) > width:
        now = content[:width]
        while line_width(now) > width:
            now = now[:-1]
        now += "$" * (width - line_width(now))
        return now
    else:
        return content

def print_line(content, columns, force_single_line):

    padding = " " * ((columns - line_width(content)) % columns)
    output = "{content}{padding}".format(content=content, padding=padding)
    if force_single_line:
        output = cut_off_at(output, columns)
    print(output, end='')
    sys.stdout.flush()

def line_width(line):
    """
    计算本行在输出到命令行后所占的宽度
    calculate the width of output in terminal
    """
    if six.PY2:
        assert isinstance(line, unicode)
    result = sum(map(get_char_width, line))
    return result

def lines_of_content(content, width):
    """
    计算内容在特定输出宽度下实际显示的行数
    calculate the actual rows with specific terminal width
    """
    result = 0
    if isinstance(content, list):
        for line in content:
            _line = preprocess(line)
            result += ceil(line_width(_line) / width)
    elif isinstance(content, dict):
        for k, v in content.items():
            # 加2是算上行内冒号和空格的宽度
            # adding 2 for the for the colon and space ": "
            _k, _v = map(preprocess, (k, v))
            result += ceil((line_width(_k) + line_width(_v) + 2) / width)
    return int(result)

def print_multi_line(content, force_single_line):

    global last_output_lines
    global overflow_flag
    global is_atty

    if not is_atty:
        if isinstance(content, list):
            for line in content:
                print(line)
        elif isinstance(content, dict):
            for k, v in sorted(content.items(), key=lambda x: x[0]):
                print("{}: {}".format(k, v))
        else:
            raise TypeError("Excepting types: list, dict. Got: {}".format(type(content)))
        return

    columns, rows = get_terminal_size()
    lines = lines_of_content(content, columns)
    if force_single_line is False and lines > rows:
        overflow_flag = True
    elif force_single_line is True and len(content) > rows:
        overflow_flag = True

    # 确保初始输出位置是位于最左处的
    # to make sure the cursor is at the left most
    print("\b" * columns, end="")

    if isinstance(content, list):
        for line in content:
            _line = preprocess(line)
            print_line(_line, columns, force_single_line)
    elif isinstance(content, dict):
        for k, v in sorted(content.items(), key=lambda x: x[0]):
            _k, _v = map(preprocess, (k, v))
            print_line("{}: {}".format(_k, _v), columns, force_single_line)
    else:
        raise TypeError("Excepting types: list, dict. Got: {}".format(type(content)))

    # 输出额外的空行来清除上一次输出的剩余内容
    # do extra blank lines to wipe the remaining of last output
    print(" " * columns * (last_output_lines - lines), end="")

    # 回到初始输出位置
    # back to the origin pos
    print("\r\b\r" * (max(last_output_lines, lines) + 1), end="")
    sys.stdout.flush()
    last_output_lines = lines

class output:

    class SignalList(list):

        def __init__(self, parent, obj):
            super(output.SignalList, self).__init__(obj)
            self.parent = parent
            self.lock = threading.Lock()

        def change(self, newlist):
            with self.lock:
                if six.PY2:
                    self[:] = []
                else:
                    self.clear()
                self.extend(newlist)
                self.parent.refresh(int(time.time()*1000), forced=False)

        def __setitem__(self, key, value):
            global is_atty
            with self.lock:
                super(output.SignalList, self).__setitem__(key, value)
                if not is_atty:
                    print("{}".format(value))
                else:
                    self.parent.refresh(int(time.time()*1000), forced=False)

    class SignalDict(dict):

        def __init__(self, parent, obj):
            super(output.SignalDict, self).__init__(obj)
            self.parent = parent
            self.lock = threading.Lock()

        def change(self, newlist):
            with self.lock:
                self.clear()
                super(output.SignalDict, self).update(newlist)
                self.parent.refresh(int(time.time()*1000), forced=False)

        def __setitem__(self, key, value):
            global is_atty

            with self.lock:
                super(output.SignalDict, self).__setitem__(key, value)
                if not is_atty:
                    print("{}: {}".format(key, value))
                else:
                    self.parent.refresh(int(time.time()*1000), forced=False)

    def __init__(self, output_type="list", initial_len=1, interval=0, force_single_line=False):

        if output_type == "list":
            self.warped_obj = output.SignalList(self, [''] * initial_len)
        elif output_type == "dict":
            self.warped_obj = output.SignalDict(self, {})

        self.interval = interval
        self.force_single_line = force_single_line
        self._last_update = int(time.time()*1000)

    def refresh(self, new_time=0, forced=True):
        if new_time - self._last_update >= self.interval or forced:
            print_multi_line(self.warped_obj, self.force_single_line)
            self._last_update = new_time

    def __enter__(self):
        global is_atty
        if not is_atty:
            print("Not in terminal, reprint now using normal build-in print function.")

        return self.warped_obj

    def __exit__(self, exc_type, exc_val, exc_tb):
        global is_atty

        self.refresh(forced=True)
        if is_atty:
            columns, _ = get_terminal_size()
            if self.force_single_line:
                print('\n' * len(self.warped_obj), end="")
            else:
                print('\n' * lines_of_content(self.warped_obj, columns), end="")
            global last_output_lines
            global overflow_flag
            last_output_lines = 0
            if overflow_flag:
                print("Detected that the lines of output has been exceeded the height of terminal windows, which \
                caused the former output remained and keep adding new lines.")
                print("检测到输出过程中, 输出行数曾大于命令行窗口行数, 这会导致输出清除不完整, 而使输出不停增长。请注意控制输出行数。")
<commit_msg>Fix: add  warning when fail to get terminal size #4<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function, division, unicode_literals

import re
import sys
import time
import threading
from math import ceil

import six
if six.PY2:
    from backports.shutil_get_terminal_size import get_terminal_size
    input = raw_input
else:
    from shutil import get_terminal_size
    from builtins import input

last_output_lines = 0
overflow_flag = False
is_atty = sys.stdout.isatty()

widths = [
    (126,    1), (159,    0), (687,     1), (710,   0), (711,   1),
    (727,    0), (733,    1), (879,     0), (1154,  1), (1161,  0),
    (4347,   1), (4447,   2), (7467,    1), (7521,  0), (8369,  1),
    (8426,   0), (9000,   1), (9002,    2), (11021, 1), (12350, 2),
    (12351,  1), (12438,  2), (12442,   0), (19893, 2), (19967, 1),
    (55203,  2), (63743,  1), (64106,   2), (65039, 1), (65059, 0),
    (65131,  2), (65279,  1), (65376,   2), (65500, 1), (65510, 2),
    (120831, 1), (262141, 2), (1114109, 1),
]

def get_char_width(char):
    global widths
    o = ord(char)
    if o == 0xe or o == 0xf:
        return 0
    for num, wid in widths:
        if o <= num:
            return wid
    return 1

def preprocess(content):
    """
    对输出内容进行预处理，转为str类型，并替换行内\r\t\n等字符为空格
    do pre-process to the content, turn it into str, and replace \r\t\n with space
    """

    if six.PY2:
        _content = unicode(content, encoding=sys.stdin.encoding) if isinstance(content, str) else content
        assert isinstance(_content, unicode)
    elif six.PY3:
        _content = str(content)

    _content = re.sub(r'\r|\t|\n', ' ', _content)

    return _content

def cut_off_at(content, width):
    if line_width(content) > width:
        now = content[:width]
        while line_width(now) > width:
            now = now[:-1]
        now += "$" * (width - line_width(now))
        return now
    else:
        return content

def print_line(content, columns, force_single_line):

    padding = " " * ((columns - line_width(content)) % columns)
    output = "{content}{padding}".format(content=content, padding=padding)
    if force_single_line:
        output = cut_off_at(output, columns)
    print(output, end='')
    sys.stdout.flush()

def line_width(line):
    """
    计算本行在输出到命令行后所占的宽度
    calculate the width of output in terminal
    """
    if six.PY2:
        assert isinstance(line, unicode)
    result = sum(map(get_char_width, line))
    return result

def lines_of_content(content, width):
    """
    计算内容在特定输出宽度下实际显示的行数
    calculate the actual rows with specific terminal width
    """
    result = 0
    if isinstance(content, list):
        for line in content:
            _line = preprocess(line)
            result += ceil(line_width(_line) / width)
    elif isinstance(content, dict):
        for k, v in content.items():
            # 加2是算上行内冒号和空格的宽度
            # adding 2 for the for the colon and space ": "
            _k, _v = map(preprocess, (k, v))
            result += ceil((line_width(_k) + line_width(_v) + 2) / width)
    return int(result)

def print_multi_line(content, force_single_line):

    global last_output_lines
    global overflow_flag
    global is_atty

    if not is_atty:
        if isinstance(content, list):
            for line in content:
                print(line)
        elif isinstance(content, dict):
            for k, v in sorted(content.items(), key=lambda x: x[0]):
                print("{}: {}".format(k, v))
        else:
            raise TypeError("Excepting types: list, dict. Got: {}".format(type(content)))
        return

    columns, rows = get_terminal_size()
    lines = lines_of_content(content, columns)
    if force_single_line is False and lines > rows:
        overflow_flag = True
    elif force_single_line is True and len(content) > rows:
        overflow_flag = True

    # 确保初始输出位置是位于最左处的
    # to make sure the cursor is at the left most
    print("\b" * columns, end="")

    if isinstance(content, list):
        for line in content:
            _line = preprocess(line)
            print_line(_line, columns, force_single_line)
    elif isinstance(content, dict):
        for k, v in sorted(content.items(), key=lambda x: x[0]):
            _k, _v = map(preprocess, (k, v))
            print_line("{}: {}".format(_k, _v), columns, force_single_line)
    else:
        raise TypeError("Excepting types: list, dict. Got: {}".format(type(content)))

    # 输出额外的空行来清除上一次输出的剩余内容
    # do extra blank lines to wipe the remaining of last output
    print(" " * columns * (last_output_lines - lines), end="")

    # 回到初始输出位置
    # back to the origin pos
    print("\r\b\r" * (max(last_output_lines, lines) + 1), end="")
    sys.stdout.flush()
    last_output_lines = lines

class output:

    class SignalList(list):

        def __init__(self, parent, obj):
            super(output.SignalList, self).__init__(obj)
            self.parent = parent
            self.lock = threading.Lock()

        def change(self, newlist):
            with self.lock:
                if six.PY2:
                    self[:] = []
                else:
                    self.clear()
                self.extend(newlist)
                self.parent.refresh(int(time.time()*1000), forced=False)

        def __setitem__(self, key, value):
            global is_atty
            with self.lock:
                super(output.SignalList, self).__setitem__(key, value)
                if not is_atty:
                    print("{}".format(value))
                else:
                    self.parent.refresh(int(time.time()*1000), forced=False)

    class SignalDict(dict):

        def __init__(self, parent, obj):
            super(output.SignalDict, self).__init__(obj)
            self.parent = parent
            self.lock = threading.Lock()

        def change(self, newlist):
            with self.lock:
                self.clear()
                super(output.SignalDict, self).update(newlist)
                self.parent.refresh(int(time.time()*1000), forced=False)

        def __setitem__(self, key, value):
            global is_atty

            with self.lock:
                super(output.SignalDict, self).__setitem__(key, value)
                if not is_atty:
                    print("{}: {}".format(key, value))
                else:
                    self.parent.refresh(int(time.time()*1000), forced=False)

    def __init__(self, output_type="list", initial_len=1, interval=0, force_single_line=False, no_warning=False):
        self.no_warning = no_warning
        no_warning and print("All reprint warning diabled.")

        global is_atty
        # reprint does not work in the IDLE terminal, and any other environment that can't get terminal_size
        if is_atty and not all(get_terminal_size()):
            if not no_warning:
                r = input("Fail to get terminal size, we got {}, continue anyway? (y/N)".format(get_terminal_size()))
                if not (r and isinstance(r, str) and r.lower()[0] in ['y','t','1']):
                    sys.exit(0)

            is_atty = False

        if output_type == "list":
            self.warped_obj = output.SignalList(self, [''] * initial_len)
        elif output_type == "dict":
            self.warped_obj = output.SignalDict(self, {})

        self.interval = interval
        self.force_single_line = force_single_line
        self._last_update = int(time.time()*1000)

    def refresh(self, new_time=0, forced=True):
        if new_time - self._last_update >= self.interval or forced:
            print_multi_line(self.warped_obj, self.force_single_line)
            self._last_update = new_time

    def __enter__(self):
        global is_atty
        if not is_atty:
            if not self.no_warning:
                print("Not in terminal, reprint now using normal build-in print function.")

        return self.warped_obj

    def __exit__(self, exc_type, exc_val, exc_tb):
        global is_atty

        self.refresh(forced=True)
        if is_atty:
            columns, _ = get_terminal_size()
            if self.force_single_line:
                print('\n' * len(self.warped_obj), end="")
            else:
                print('\n' * lines_of_content(self.warped_obj, columns), end="")
            global last_output_lines
            global overflow_flag
            last_output_lines = 0
            if overflow_flag:
                if not self.no_warning:
                    print("Detected that the lines of output has been exceeded the height of terminal windows, which \
                    caused the former output remained and keep adding new lines.")
                    print("检测到输出过程中, 输出行数曾大于命令行窗口行数, 这会导致输出清除不完整, 而使输出不停增长。请注意控制输出行数。")
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from datetime import timedelta
import logging
import os
import re
import time

from django.conf import settings
from django.db import models
from django.db.models import Q
from django.utils.crypto import get_random_string
from django.utils.encoding import python_2_unicode_compatible
from django.utils.http import same_origin
from django.utils.timezone import now
from django.utils.translation import ugettext_lazy as _

import requests

from mama_cas.compat import gevent
from mama_cas.compat import user_model
from mama_cas.exceptions import BadPgt
from mama_cas.exceptions import InvalidProxyCallback
from mama_cas.exceptions import InvalidRequest
from mama_cas.exceptions import InvalidService
from mama_cas.exceptions import InvalidTicket
from mama_cas.request import SingleSignOutRequest
from mama_cas.utils import add_query_params
from mama_cas.utils import is_scheme_https
from mama_cas.utils import clean_service_url
from mama_cas.utils import is_valid_service_url

if gevent:
    from gevent import monkey
    monkey.patch_all(thread=False, select=False)

logger = logging.getLogger(__name__)

class TicketManager(models.Manager):
    def create_ticket(self, ticket=None, **kwargs):
        """
        Create a new ``Ticket``. Additional arguments are passed to the
        ``create()`` function. Return the newly created ``Ticket``.
        """
        if not ticket:
            ticket = self.create_ticket_str()
        if 'service' in kwargs:
            kwargs['service'] = clean_service_url(kwargs['service'])
        if 'expires' not in kwargs:
            expires = now() + timedelta(seconds=self.model.TICKET_EXPIRE)
            kwargs['expires'] = expires
        t = self.create(ticket=ticket, **kwargs)
        logger.debug("Created %s %s" % (t.name, t.ticket))
        return t

    def create_ticket_str(self, prefix=None):
        """
        Generate a sufficiently opaque ticket string to ensure the ticket is
        not guessable. If a prefix is provided, prepend it to the string.
        """
        if not prefix:
            prefix = self.model.TICKET_PREFIX
        return "%s-%d-%s" % (prefix, int(time.time()),
                             get_random_string(length=self.model.TICKET_RAND_LEN))

    def validate_ticket(self, ticket, service, renew=False):
        """
        Given a ticket string and service identifier, validate the
        corresponding ``Ticket``. If validation succeeds, return the
        ``Ticket``. If validation fails, raise an appropriate error.

        If ``renew`` is provided, the validation will only succeed if the
        ticket was issued from the presentation of the user's primary
        credentials.
        """
        if not ticket:
            raise InvalidRequest("No ticket string provided")

        if not self.model.TICKET_RE.match(ticket):
            raise InvalidTicket("Ticket string %s is invalid" % ticket)

        try:
            t = self.get(ticket=ticket)
        except self.model.DoesNotExist:
            raise InvalidTicket("Ticket %s does not exist" % ticket)

        if t.is_consumed():
            raise InvalidTicket("%s %s has already been used" %
                                (t.name, ticket))
        t.consume()

        if t.is_expired():
            raise InvalidTicket("%s %s has expired" % (t.name, ticket))

        if not service:
            raise InvalidRequest("No service identifier provided")

        if not is_valid_service_url(service):
            raise InvalidService("Service %s is not a valid %s URL" %
                                 (service, t.name))

        if not same_origin(t.service, service):
            raise InvalidService("%s %s for service %s is invalid for service "
                                 "%s" % (t.name, ticket, t.service, service))

        if renew and not t.is_primary():
            raise InvalidTicket("%s %s was not issued via primary "
                                "credentials" % (t.name, ticket))

        logger.debug("Validated %s %s" % (t.name, ticket))
        return t

    def delete_invalid_tickets(self):
        """
        Delete consumed or expired ``Ticket``s that are not referenced
        by other ``Ticket``s. Invalid tickets are no longer valid for
        authentication and can be safely deleted.

        A custom management command is provided that executes this method
        on all applicable models by running ``manage.py cleanupcas``.
        """
        for ticket in self.filter(Q(consumed__isnull=False) |
                                  Q(expires__lte=now())).order_by('-expires'):
            try:
                ticket.delete()
            except models.ProtectedError:
                pass

    def consume_tickets(self, user):
        """
        Consume all valid ``Ticket``s for a specified user. This is run
        when the user logs out to ensure all issued tickets are no longer
        valid for future authentication attempts.
        """
        for ticket in self.filter(user=user, consumed__isnull=True,
                                  expires__gt=now()):
            ticket.consume()

@python_2_unicode_compatible
class Ticket(models.Model):
    """
    ``Ticket`` is an abstract base class implementing common methods
    and fields for CAS tickets.
    """
    TICKET_EXPIRE = getattr(settings, 'MAMA_CAS_TICKET_EXPIRE', 90)
    TICKET_RAND_LEN = getattr(settings, 'MAMA_CAS_TICKET_RAND_LEN', 32)
    TICKET_RE = re.compile("^[A-Z]{2,3}-[0-9]{10,}-[a-zA-Z0-9]{%d}$" % TICKET_RAND_LEN)

    ticket = models.CharField(_('ticket'), max_length=255, unique=True)
    user = models.ForeignKey(user_model, verbose_name=_('user'))
    expires = models.DateTimeField(_('expires'))
    consumed = models.DateTimeField(_('consumed'), null=True)

    objects = TicketManager()

    class Meta:
        abstract = True

    def __str__(self):
        return self.ticket

    @property
    def name(self):
        return self._meta.verbose_name

    def consume(self):
        """
        Consume a ``Ticket`` by populating the ``consumed`` field with
        the current datetime. A consumed ``Ticket`` is invalid for future
        authentication attempts.
        """
        self.consumed = now()
        self.save()

    def is_consumed(self):
        """
        Check a ``Ticket``s consumed state. Return ``True`` if the ticket is
        consumed, and ``False`` otherwise.
        """
        if self.consumed:
            return True
        return False

    def is_expired(self):
        """
        Check a ``Ticket``s expired state. Return ``True`` if the ticket is
        expired, and ``False`` otherwise.
        """
        return self.expires <= now()

class ServiceTicketManager(TicketManager):
    def request_sign_out(self, user):
        """
        Send a single sign-out request to each service accessed by a
        specified user. This is called at logout when single sign-out
        is enabled.

        If gevent is installed, asynchronous requests will be sent.
        Otherwise, synchronous requests will be sent.
        """
        tickets = list(self.filter(user=user, consumed__gte=user.last_login))

        if gevent:
            from gevent.pool import Pool
            pool = Pool(4)
            requests = [pool.spawn(t.request_sign_out) for t in tickets]
            gevent.joinall(requests)
        else:
            for ticket in tickets:
                ticket.request_sign_out()

class ServiceTicket(Ticket):
    """
    (3.1) A ``ServiceTicket`` is used by the client as a credential to
    obtain access to a service. It is obtained upon a client's presentation
    of credentials and a service identifier to /login.
    """
    TICKET_PREFIX = 'ST'

    service = models.CharField(_('service'), max_length=255)
    primary = models.BooleanField(_('primary'), default=False)

    objects = ServiceTicketManager()

    class Meta:
        verbose_name = _('service ticket')
        verbose_name_plural = _('service tickets')

    def is_primary(self):
        """
        Check the credential origin for a ``ServiceTicket``. If the ticket was
        issued from the presentation of the user's primary credentials,
        return ``True``, otherwise return ``False``.
        """
        if self.primary:
            return True
        return False

    def request_sign_out(self):
        """
        Send a POST request to the ``ServiceTicket``s service URL to
        request sign-out. The remote session is identified by the
        service ticket string that instantiated the session. Clients
        that do not support single sign-out will notice these spurious
        requests in their logs.
        """
        request = SingleSignOutRequest(context={'ticket': self})
        try:
            resp = requests.post(self.service, data=request.render_content(),
                                 headers=request.headers())
            resp.raise_for_status()
        except requests.exceptions.RequestException as e:
            logger.warning("Single sign-out request to %s returned %s" %
                           (self.service, e))
        else:
            logger.debug("Single sign-out request sent to %s" % self.service)

class ProxyTicket(Ticket):
    """
    (3.2) A ``ProxyTicket`` is used by a service as a credential to obtain
    access to a back-end service on behalf of a client. It is obtained upon
    a service's presentation of a ``ProxyGrantingTicket`` and a service
    identifier.
    """
    TICKET_PREFIX = 'PT'

    service = models.CharField(_('service'), max_length=255)
    granted_by_pgt = models.ForeignKey('ProxyGrantingTicket',
                                       verbose_name=_('granted by proxy-granting ticket'))

    class Meta:
        verbose_name = _('proxy ticket')
        verbose_name_plural = _('proxy tickets')

class ProxyGrantingTicketManager(TicketManager):
    def create_ticket(self, pgturl, validate=True, **kwargs):
        """
        When a ``pgtUrl`` parameter is provided to ``/serviceValidate`` or
        ``/proxyValidate``, attempt to create a new ``ProxyGrantingTicket``.
        Start by creating the necessary ticket strings and then validate the
        callback URL. If validation succeeds, create and return the
        ``ProxyGrantingTicket``. If validation fails, return ``None``.

        If ``validate`` is set to False, ``pgtUrl`` validation is skipped.
        This is intended only for testing purposes, so a PGT can be created
        without a valid callback URL present.
        """
        pgtid = self.create_ticket_str()
        pgtiou = self.create_ticket_str(prefix=self.model.IOU_PREFIX)
        try:
            if validate:
                self.validate_callback(pgturl, pgtid, pgtiou)
        except InvalidProxyCallback as e:
            # pgtUrl validation failed, so nothing has been created
            logger.warning("%s %s" % (e.code, e))
            return None
        else:
            # pgtUrl validation succeeded, so create a new PGT with the
            # already created ticket strings
            return super(ProxyGrantingTicketManager, self).create_ticket(ticket=pgtid,
                                                                         iou=pgtiou,
                                                                         **kwargs)

    def validate_callback(self, url, pgtid, pgtiou):
        """
        Verify the provided proxy callback URL. This verification process
        requires three steps:

        1. The URL scheme must be HTTPS
        2. The SSL certificate must be valid and its name must match that
           of the service
        3. The callback URL must respond with a 200 or 3xx response code

        It is not required for validation that 3xx redirects be followed.
        """
        # Ensure the scheme is HTTPS before proceeding
        if not is_scheme_https(url):
            raise InvalidProxyCallback("Proxy callback %s is not HTTPS" % url)

        # Connect to proxy callback URL, checking the SSL certificate
        url_params = add_query_params(url, {'pgtId': pgtid, 'pgtIou': pgtiou})
        verify = os.environ.get('REQUESTS_CA_BUNDLE', True)
        try:
            r = requests.get(url_params, verify=verify, timeout=3.0)
        except requests.exceptions.SSLError:
            msg = "SSL cert validation failed for proxy callback %s" % url
            raise InvalidProxyCallback(msg)
        except requests.exceptions.ConnectionError:
            msg = "Error connecting to proxy callback %s" % url
            raise InvalidProxyCallback(msg)
        except requests.exceptions.Timeout:
            msg = "Timeout connecting to proxy callback %s" % url
            raise InvalidProxyCallback(msg)

        # Check the returned HTTP status code
        try:
            r.raise_for_status()
        except requests.exceptions.HTTPError as e:
            msg = "Proxy callback %s returned %s" % (url, e)
            raise InvalidProxyCallback(msg)

    def validate_ticket(self, ticket, service):
        """
        Given a ticket string and service identifier, validate the
        corresponding ``Ticket``. If validation succeeds, return the
        ``Ticket``. If validation fails, raise an appropriate error.
        """
        if not ticket:
            raise InvalidRequest("No ticket string provided")

        if not service:
            raise InvalidRequest("No service identifier provided")

        if not self.model.TICKET_RE.match(ticket):
            raise InvalidTicket("Ticket string %s is invalid" % ticket)

        try:
            t = self.get(ticket=ticket)
        except self.model.DoesNotExist:
            raise BadPgt("Ticket %s does not exist" % ticket)

        if t.is_consumed():
            raise InvalidTicket("%s %s has already been used" %
                                (t.name, ticket))

        if t.is_expired():
            raise InvalidTicket("%s %s has expired" % (t.name, ticket))

        if not is_valid_service_url(service):
            raise InvalidService("Service %s is not a valid %s URL" %
                                 (service, t.name))

        logger.debug("Validated %s %s" % (t.name, ticket))
        return t

class ProxyGrantingTicket(Ticket):
    """
    (3.3) A ``ProxyGrantingTicket`` is used by a service to obtain proxy
    tickets for obtaining access to a back-end service on behalf of a
    client. It is obtained upon validation of a ``ServiceTicket`` or a
    ``ProxyTicket``.
    """
    TICKET_PREFIX = 'PGT'
    IOU_PREFIX = 'PGTIOU'
    TICKET_EXPIRE = getattr(settings, 'SESSION_COOKIE_AGE')

    iou = models.CharField(_('iou'), max_length=255, unique=True)
    granted_by_st = models.ForeignKey(ServiceTicket, null=True, blank=True,
                                      on_delete=models.PROTECT,
                                      verbose_name=_('granted by service ticket'))
    granted_by_pt = models.ForeignKey(ProxyTicket, null=True, blank=True,
                                      on_delete=models.PROTECT,
                                      verbose_name=_('granted by proxy ticket'))

    objects = ProxyGrantingTicketManager()

    class Meta:
        verbose_name = _('proxy-granting ticket')
        verbose_name_plural = _('proxy-granting tickets')
<commit_msg>Allow gevent pool size to be configured<commit_after>from __future__ import unicode_literals

from datetime import timedelta
import logging
import os
import re
import time

from django.conf import settings
from django.db import models
from django.db.models import Q
from django.utils.crypto import get_random_string
from django.utils.encoding import python_2_unicode_compatible
from django.utils.http import same_origin
from django.utils.timezone import now
from django.utils.translation import ugettext_lazy as _

import requests

from mama_cas.compat import gevent
from mama_cas.compat import user_model
from mama_cas.exceptions import BadPgt
from mama_cas.exceptions import InvalidProxyCallback
from mama_cas.exceptions import InvalidRequest
from mama_cas.exceptions import InvalidService
from mama_cas.exceptions import InvalidTicket
from mama_cas.request import SingleSignOutRequest
from mama_cas.utils import add_query_params
from mama_cas.utils import is_scheme_https
from mama_cas.utils import clean_service_url
from mama_cas.utils import is_valid_service_url

if gevent:
    from gevent.pool import Pool
    from gevent import monkey
    monkey.patch_all(thread=False, select=False)

logger = logging.getLogger(__name__)

class TicketManager(models.Manager):
    def create_ticket(self, ticket=None, **kwargs):
        """
        Create a new ``Ticket``. Additional arguments are passed to the
        ``create()`` function. Return the newly created ``Ticket``.
        """
        if not ticket:
            ticket = self.create_ticket_str()
        if 'service' in kwargs:
            kwargs['service'] = clean_service_url(kwargs['service'])
        if 'expires' not in kwargs:
            expires = now() + timedelta(seconds=self.model.TICKET_EXPIRE)
            kwargs['expires'] = expires
        t = self.create(ticket=ticket, **kwargs)
        logger.debug("Created %s %s" % (t.name, t.ticket))
        return t

    def create_ticket_str(self, prefix=None):
        """
        Generate a sufficiently opaque ticket string to ensure the ticket is
        not guessable. If a prefix is provided, prepend it to the string.
        """
        if not prefix:
            prefix = self.model.TICKET_PREFIX
        return "%s-%d-%s" % (prefix, int(time.time()),
                             get_random_string(length=self.model.TICKET_RAND_LEN))

    def validate_ticket(self, ticket, service, renew=False):
        """
        Given a ticket string and service identifier, validate the
        corresponding ``Ticket``. If validation succeeds, return the
        ``Ticket``. If validation fails, raise an appropriate error.

        If ``renew`` is provided, the validation will only succeed if the
        ticket was issued from the presentation of the user's primary
        credentials.
        """
        if not ticket:
            raise InvalidRequest("No ticket string provided")

        if not self.model.TICKET_RE.match(ticket):
            raise InvalidTicket("Ticket string %s is invalid" % ticket)

        try:
            t = self.get(ticket=ticket)
        except self.model.DoesNotExist:
            raise InvalidTicket("Ticket %s does not exist" % ticket)

        if t.is_consumed():
            raise InvalidTicket("%s %s has already been used" %
                                (t.name, ticket))
        t.consume()

        if t.is_expired():
            raise InvalidTicket("%s %s has expired" % (t.name, ticket))

        if not service:
            raise InvalidRequest("No service identifier provided")

        if not is_valid_service_url(service):
            raise InvalidService("Service %s is not a valid %s URL" %
                                 (service, t.name))

        if not same_origin(t.service, service):
            raise InvalidService("%s %s for service %s is invalid for service "
                                 "%s" % (t.name, ticket, t.service, service))

        if renew and not t.is_primary():
            raise InvalidTicket("%s %s was not issued via primary "
                                "credentials" % (t.name, ticket))

        logger.debug("Validated %s %s" % (t.name, ticket))
        return t

    def delete_invalid_tickets(self):
        """
        Delete consumed or expired ``Ticket``s that are not referenced
        by other ``Ticket``s. Invalid tickets are no longer valid for
        authentication and can be safely deleted.

        A custom management command is provided that executes this method
        on all applicable models by running ``manage.py cleanupcas``.
        """
        for ticket in self.filter(Q(consumed__isnull=False) |
                                  Q(expires__lte=now())).order_by('-expires'):
            try:
                ticket.delete()
            except models.ProtectedError:
                pass

    def consume_tickets(self, user):
        """
        Consume all valid ``Ticket``s for a specified user. This is run
        when the user logs out to ensure all issued tickets are no longer
        valid for future authentication attempts.
        """
        for ticket in self.filter(user=user, consumed__isnull=True,
                                  expires__gt=now()):
            ticket.consume()

@python_2_unicode_compatible
class Ticket(models.Model):
    """
    ``Ticket`` is an abstract base class implementing common methods
    and fields for CAS tickets.
    """
    TICKET_EXPIRE = getattr(settings, 'MAMA_CAS_TICKET_EXPIRE', 90)
    TICKET_RAND_LEN = getattr(settings, 'MAMA_CAS_TICKET_RAND_LEN', 32)
    TICKET_RE = re.compile("^[A-Z]{2,3}-[0-9]{10,}-[a-zA-Z0-9]{%d}$" % TICKET_RAND_LEN)

    ticket = models.CharField(_('ticket'), max_length=255, unique=True)
    user = models.ForeignKey(user_model, verbose_name=_('user'))
    expires = models.DateTimeField(_('expires'))
    consumed = models.DateTimeField(_('consumed'), null=True)

    objects = TicketManager()

    class Meta:
        abstract = True

    def __str__(self):
        return self.ticket

    @property
    def name(self):
        return self._meta.verbose_name

    def consume(self):
        """
        Consume a ``Ticket`` by populating the ``consumed`` field with
        the current datetime. A consumed ``Ticket`` is invalid for future
        authentication attempts.
        """
        self.consumed = now()
        self.save()

    def is_consumed(self):
        """
        Check a ``Ticket``s consumed state. Return ``True`` if the ticket is
        consumed, and ``False`` otherwise.
        """
        if self.consumed:
            return True
        return False

    def is_expired(self):
        """
        Check a ``Ticket``s expired state. Return ``True`` if the ticket is
        expired, and ``False`` otherwise.
        """
        return self.expires <= now()

class ServiceTicketManager(TicketManager):
    def request_sign_out(self, user):
        """
        Send a single sign-out request to each service accessed by a
        specified user. This is called at logout when single sign-out
        is enabled.

        If gevent is installed, asynchronous requests will be sent.
        Otherwise, synchronous requests will be sent.
        """
        def spawn(ticket, pool=None):
            if pool != None:
                return pool.spawn(ticket.request_sign_out)
            return gevent.spawn(ticket.request_sign_out)

        tickets = list(self.filter(user=user, consumed__gte=user.last_login))

        if gevent:
            size = getattr(settings, 'MAMA_CAS_ASYNC_CONCURRENCY', None)
            pool = Pool(size) if size else None
            requests = [spawn(t) for t in tickets]
            gevent.joinall(requests)
        else:
            for ticket in tickets:
                ticket.request_sign_out()

class ServiceTicket(Ticket):
    """
    (3.1) A ``ServiceTicket`` is used by the client as a credential to
    obtain access to a service. It is obtained upon a client's presentation
    of credentials and a service identifier to /login.
    """
    TICKET_PREFIX = 'ST'

    service = models.CharField(_('service'), max_length=255)
    primary = models.BooleanField(_('primary'), default=False)

    objects = ServiceTicketManager()

    class Meta:
        verbose_name = _('service ticket')
        verbose_name_plural = _('service tickets')

    def is_primary(self):
        """
        Check the credential origin for a ``ServiceTicket``. If the ticket was
        issued from the presentation of the user's primary credentials,
        return ``True``, otherwise return ``False``.
        """
        if self.primary:
            return True
        return False

    def request_sign_out(self):
        """
        Send a POST request to the ``ServiceTicket``s service URL to
        request sign-out. The remote session is identified by the
        service ticket string that instantiated the session. Clients
        that do not support single sign-out will notice these spurious
        requests in their logs.
        """
        request = SingleSignOutRequest(context={'ticket': self})
        try:
            resp = requests.post(self.service, data=request.render_content(),
                                 headers=request.headers())
            resp.raise_for_status()
        except requests.exceptions.RequestException as e:
            logger.warning("Single sign-out request to %s returned %s" %
                           (self.service, e))
        else:
            logger.debug("Single sign-out request sent to %s" % self.service)

class ProxyTicket(Ticket):
    """
    (3.2) A ``ProxyTicket`` is used by a service as a credential to obtain
    access to a back-end service on behalf of a client. It is obtained upon
    a service's presentation of a ``ProxyGrantingTicket`` and a service
    identifier.
    """
    TICKET_PREFIX = 'PT'

    service = models.CharField(_('service'), max_length=255)
    granted_by_pgt = models.ForeignKey('ProxyGrantingTicket',
                                       verbose_name=_('granted by proxy-granting ticket'))

    class Meta:
        verbose_name = _('proxy ticket')
        verbose_name_plural = _('proxy tickets')

class ProxyGrantingTicketManager(TicketManager):
    def create_ticket(self, pgturl, validate=True, **kwargs):
        """
        When a ``pgtUrl`` parameter is provided to ``/serviceValidate`` or
        ``/proxyValidate``, attempt to create a new ``ProxyGrantingTicket``.
        Start by creating the necessary ticket strings and then validate the
        callback URL. If validation succeeds, create and return the
        ``ProxyGrantingTicket``. If validation fails, return ``None``.

        If ``validate`` is set to False, ``pgtUrl`` validation is skipped.
        This is intended only for testing purposes, so a PGT can be created
        without a valid callback URL present.
        """
        pgtid = self.create_ticket_str()
        pgtiou = self.create_ticket_str(prefix=self.model.IOU_PREFIX)
        try:
            if validate:
                self.validate_callback(pgturl, pgtid, pgtiou)
        except InvalidProxyCallback as e:
            # pgtUrl validation failed, so nothing has been created
            logger.warning("%s %s" % (e.code, e))
            return None
        else:
            # pgtUrl validation succeeded, so create a new PGT with the
            # already created ticket strings
            return super(ProxyGrantingTicketManager, self).create_ticket(ticket=pgtid,
                                                                         iou=pgtiou,
                                                                         **kwargs)

    def validate_callback(self, url, pgtid, pgtiou):
        """
        Verify the provided proxy callback URL. This verification process
        requires three steps:

        1. The URL scheme must be HTTPS
        2. The SSL certificate must be valid and its name must match that
           of the service
        3. The callback URL must respond with a 200 or 3xx response code

        It is not required for validation that 3xx redirects be followed.
        """
        # Ensure the scheme is HTTPS before proceeding
        if not is_scheme_https(url):
            raise InvalidProxyCallback("Proxy callback %s is not HTTPS" % url)

        # Connect to proxy callback URL, checking the SSL certificate
        url_params = add_query_params(url, {'pgtId': pgtid, 'pgtIou': pgtiou})
        verify = os.environ.get('REQUESTS_CA_BUNDLE', True)
        try:
            r = requests.get(url_params, verify=verify, timeout=3.0)
        except requests.exceptions.SSLError:
            msg = "SSL cert validation failed for proxy callback %s" % url
            raise InvalidProxyCallback(msg)
        except requests.exceptions.ConnectionError:
            msg = "Error connecting to proxy callback %s" % url
            raise InvalidProxyCallback(msg)
        except requests.exceptions.Timeout:
            msg = "Timeout connecting to proxy callback %s" % url
            raise InvalidProxyCallback(msg)

        # Check the returned HTTP status code
        try:
            r.raise_for_status()
        except requests.exceptions.HTTPError as e:
            msg = "Proxy callback %s returned %s" % (url, e)
            raise InvalidProxyCallback(msg)

    def validate_ticket(self, ticket, service):
        """
        Given a ticket string and service identifier, validate the
        corresponding ``Ticket``. If validation succeeds, return the
        ``Ticket``. If validation fails, raise an appropriate error.
        """
        if not ticket:
            raise InvalidRequest("No ticket string provided")

        if not service:
            raise InvalidRequest("No service identifier provided")

        if not self.model.TICKET_RE.match(ticket):
            raise InvalidTicket("Ticket string %s is invalid" % ticket)

        try:
            t = self.get(ticket=ticket)
        except self.model.DoesNotExist:
            raise BadPgt("Ticket %s does not exist" % ticket)

        if t.is_consumed():
            raise InvalidTicket("%s %s has already been used" %
                                (t.name, ticket))

        if t.is_expired():
            raise InvalidTicket("%s %s has expired" % (t.name, ticket))

        if not is_valid_service_url(service):
            raise InvalidService("Service %s is not a valid %s URL" %
                                 (service, t.name))

        logger.debug("Validated %s %s" % (t.name, ticket))
        return t

class ProxyGrantingTicket(Ticket):
    """
    (3.3) A ``ProxyGrantingTicket`` is used by a service to obtain proxy
    tickets for obtaining access to a back-end service on behalf of a
    client. It is obtained upon validation of a ``ServiceTicket`` or a
    ``ProxyTicket``.
    """
    TICKET_PREFIX = 'PGT'
    IOU_PREFIX = 'PGTIOU'
    TICKET_EXPIRE = getattr(settings, 'SESSION_COOKIE_AGE')

    iou = models.CharField(_('iou'), max_length=255, unique=True)
    granted_by_st = models.ForeignKey(ServiceTicket, null=True, blank=True,
                                      on_delete=models.PROTECT,
                                      verbose_name=_('granted by service ticket'))
    granted_by_pt = models.ForeignKey(ProxyTicket, null=True, blank=True,
                                      on_delete=models.PROTECT,
                                      verbose_name=_('granted by proxy ticket'))

    objects = ProxyGrantingTicketManager()

    class Meta:
        verbose_name = _('proxy-granting ticket')
        verbose_name_plural = _('proxy-granting tickets')
<|endoftext|>


<commit_before>"""Easy access to common pipeline environment variables."""
# -----------------------------------------------------------------------------
# Module: dpa.env.vars
# Contact: Josh Tomlinson (jtomlin)
# -----------------------------------------------------------------------------

# -----------------------------------------------------------------------------
# Imports:
# -----------------------------------------------------------------------------
from dpa.env import EnvVar, ListVar, PathVar

# -----------------------------------------------------------------------------
# Classes:
# -----------------------------------------------------------------------------
class DpaVars(object):
    """A convenience class for accessing common pipeline environment variables.

    The point of this class is to document and provide easy access to common
    environment variables, either specific to or used by the pipeline, via the
    :py:mod:`dpa.env` APIs. Ideally, all pipeline related environment variables
    would be defined here to prevent hardcoding their names throughout the
    codebase. This also provides a single location where reliance on
    environment variables in the pipeline can be vetted.

    Each method accepts an optional ``default`` keyword for the default value
    of the variable. All methods return instances of :py:obj:`dpa.env.EnvVar`
    (or a subclass like :py:obj:`dpa.env.PathVar`).

    Usage::

        >>> from dpa.env.vars import DpaVars
        >>> python_path = DpaVars.python_path()

    """

    # -------------------------------------------------------------------------
    # Class attributes:
    # -------------------------------------------------------------------------

    ld_library_path = staticmethod(
        lambda default="": PathVar('LD_LIBRARY_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$LD_LIBRARY_PATH``

    The search path for binary library loading.

    """

    ld_library_path_base = staticmethod(
        lambda default="": PathVar('DPA_BASE_LD_LIRARY_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$DPA_BASE_LD_LIBRARY_PATH``

    A baseline for ``$LD_LIBRARY_PATH`` used when setting/unsetting ptasks.

    """

    path = staticmethod(
        lambda default="": PathVar('PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$PATH``

    The search path for executables.

    """

    path_base = staticmethod(
        lambda default="": PathVar('DPA_BASE_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$DPA_BASE_PATH``

    A baseline for ``$PATH`` used when setting/unsetting ptasks.

    """

    python_path = staticmethod(
        lambda default="": PathVar('PYTHONPATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$PYTHONPATH``

    The search path for python packages.

    """

    python_path_base = staticmethod(
        lambda default="": PathVar('DPA_BASE_PYTHONPATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$DPA_BASE_PYTHONPATH``

    A baseline for ``$PYTHONPATH`` used when setting/unsetting ptasks.

    """

    filesystem_root = staticmethod(
        lambda default="": EnvVar('DPA_FILESYSTEM_ROOT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_FILESYSTEM_ROOT``

    The root filesystem directory for the current pipeline location.

    """
    projects_root = staticmethod(
        lambda default="": EnvVar('DPA_PROJECTS_ROOT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PROJECTS_ROOT``

    The root directory for dpa projects.

    """

    ptask_custom_vars = staticmethod(
        lambda default="": ListVar('DPA_PTASK_CUSTOM_VARS', default)
    )
    """Returns an instance of :py:obj:`dpa.env.ListVar` for ``$DPA_PTASK_CUSTOM_VARS``

    A list of env variable names specific to the current ptask. This list is
    used to keep track of what custom variables have been set in the ptask so
    that they can be unset when the ptask context is changed.

    """

    ptask_history_file = staticmethod(
        lambda default="~/.ptask_history": \
            EnvVar('DPA_PTASK_HISTORY_FILE', default)
    )
    """Returns an instance of a :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_HISTORY_FILE``

    The path to a file storing the history of recently set ptasks. Set
    ``$DPA_PTASK_HISTORY_SIZE`` to set the number of recent ptasks to store in
    the history file.

    """

    ptask_history_size= staticmethod(
        lambda default=10: \
            EnvVar('DPA_PTASK_HISTORY_SIZE', default)
    )
    """Returns an instance of a :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_HISTORY_SIZE``

    The number of ptasks to store in the ptask history file in the current
    user's ``$DPA_PTASK_HISTORY_FILE``.

    """

    ptask_spec = staticmethod(
        lambda default="": EnvVar('DPA_PTASK_SPEC', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_SPEC``

    Used to identify the ptask currently set in the environment.

    """

    ptask_version= staticmethod(
        lambda default="": EnvVar('DPA_PTASK_VERSION', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_VERSION``

    Used to identify the version of the currently set ptask.

    """

    ptask_path = staticmethod(
        lambda default="": EnvVar('DPA_PTASK_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_PATH``

    Used to identify the path of the currently set ptask.

    """

    python_importer_disable = staticmethod(
        lambda default="": EnvVar('DPA_PYTHON_FINDER_DISABLE', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PYTHON_FINDER_DISABLE``

    Setting this env variable to any value will disable the specialized DPA
    python importer which allows portions of packages to exist in multiple
    locations along the python path. For more information, see the
    documentation for ``dpa_import_hook``.

    """

    share_logs = staticmethod(
        lambda default="": EnvVar('DPA_SHARE_LOGS', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_SHARE_LOGS``

    The shared logging directory for dpa pipeline log files.

    """

    ptask_prompt = staticmethod(
        lambda default="": EnvVar('DPA_PTASK_PROMPT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_PROMPT``

    This defines the shell prompt when a ptask is set.

    """

    no_ptask_prompt = staticmethod(
        lambda default="": EnvVar('DPA_NO_PTASK_PROMPT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_NO_PTASK_PROMPT``

    This defines the shell prompt when a ptask is not set.

    """

    data_server = staticmethod(
        lambda default="": EnvVar('DPA_DATA_SERVER', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_DATA_SERVER``

    This defines the server to connect to in order to query pipeline data.

    """

    cheesyq_data_server = staticmethod(
        lambda default="": EnvVar('CHEESYQ_DATA_SERVER', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$CHEESYQ_DATA_SERVER``

    This defines the server to connect to in order to query pipeline data.

    """

    location_code = staticmethod(
        lambda default="": EnvVar('DPA_LOCATION_CODE', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_LOCATION_CODE``

    This defines the physical location of the current environment.

    """

<commit_msg>Fixed typo in ld library path variable<commit_after>"""Easy access to common pipeline environment variables."""
# -----------------------------------------------------------------------------
# Module: dpa.env.vars
# Contact: Josh Tomlinson (jtomlin)
# -----------------------------------------------------------------------------

# -----------------------------------------------------------------------------
# Imports:
# -----------------------------------------------------------------------------
from dpa.env import EnvVar, ListVar, PathVar

# -----------------------------------------------------------------------------
# Classes:
# -----------------------------------------------------------------------------
class DpaVars(object):
    """A convenience class for accessing common pipeline environment variables.

    The point of this class is to document and provide easy access to common
    environment variables, either specific to or used by the pipeline, via the
    :py:mod:`dpa.env` APIs. Ideally, all pipeline related environment variables
    would be defined here to prevent hardcoding their names throughout the
    codebase. This also provides a single location where reliance on
    environment variables in the pipeline can be vetted.

    Each method accepts an optional ``default`` keyword for the default value
    of the variable. All methods return instances of :py:obj:`dpa.env.EnvVar`
    (or a subclass like :py:obj:`dpa.env.PathVar`).

    Usage::

        >>> from dpa.env.vars import DpaVars
        >>> python_path = DpaVars.python_path()

    """

    # -------------------------------------------------------------------------
    # Class attributes:
    # -------------------------------------------------------------------------

    ld_library_path = staticmethod(
        lambda default="": PathVar('LD_LIBRARY_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$LD_LIBRARY_PATH``

    The search path for binary library loading.

    """

    ld_library_path_base = staticmethod(
        lambda default="": PathVar('DPA_BASE_LD_LIBRARY_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$DPA_BASE_LD_LIBRARY_PATH``

    A baseline for ``$LD_LIBRARY_PATH`` used when setting/unsetting ptasks.

    """

    path = staticmethod(
        lambda default="": PathVar('PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$PATH``

    The search path for executables.

    """

    path_base = staticmethod(
        lambda default="": PathVar('DPA_BASE_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$DPA_BASE_PATH``

    A baseline for ``$PATH`` used when setting/unsetting ptasks.

    """

    python_path = staticmethod(
        lambda default="": PathVar('PYTHONPATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$PYTHONPATH``

    The search path for python packages.

    """

    python_path_base = staticmethod(
        lambda default="": PathVar('DPA_BASE_PYTHONPATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.PathVar` for ``$DPA_BASE_PYTHONPATH``

    A baseline for ``$PYTHONPATH`` used when setting/unsetting ptasks.

    """

    filesystem_root = staticmethod(
        lambda default="": EnvVar('DPA_FILESYSTEM_ROOT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_FILESYSTEM_ROOT``

    The root filesystem directory for the current pipeline location.

    """
    projects_root = staticmethod(
        lambda default="": EnvVar('DPA_PROJECTS_ROOT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PROJECTS_ROOT``

    The root directory for dpa projects.

    """

    ptask_custom_vars = staticmethod(
        lambda default="": ListVar('DPA_PTASK_CUSTOM_VARS', default)
    )
    """Returns an instance of :py:obj:`dpa.env.ListVar` for ``$DPA_PTASK_CUSTOM_VARS``

    A list of env variable names specific to the current ptask. This list is
    used to keep track of what custom variables have been set in the ptask so
    that they can be unset when the ptask context is changed.

    """

    ptask_history_file = staticmethod(
        lambda default="~/.ptask_history": \
            EnvVar('DPA_PTASK_HISTORY_FILE', default)
    )
    """Returns an instance of a :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_HISTORY_FILE``

    The path to a file storing the history of recently set ptasks. Set
    ``$DPA_PTASK_HISTORY_SIZE`` to set the number of recent ptasks to store in
    the history file.

    """

    ptask_history_size= staticmethod(
        lambda default=10: \
            EnvVar('DPA_PTASK_HISTORY_SIZE', default)
    )
    """Returns an instance of a :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_HISTORY_SIZE``

    The number of ptasks to store in the ptask history file in the current
    user's ``$DPA_PTASK_HISTORY_FILE``.

    """

    ptask_spec = staticmethod(
        lambda default="": EnvVar('DPA_PTASK_SPEC', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_SPEC``

    Used to identify the ptask currently set in the environment.

    """

    ptask_version= staticmethod(
        lambda default="": EnvVar('DPA_PTASK_VERSION', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_VERSION``

    Used to identify the version of the currently set ptask.

    """

    ptask_path = staticmethod(
        lambda default="": EnvVar('DPA_PTASK_PATH', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_PATH``

    Used to identify the path of the currently set ptask.

    """

    python_importer_disable = staticmethod(
        lambda default="": EnvVar('DPA_PYTHON_FINDER_DISABLE', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PYTHON_FINDER_DISABLE``

    Setting this env variable to any value will disable the specialized DPA
    python importer which allows portions of packages to exist in multiple
    locations along the python path. For more information, see the
    documentation for ``dpa_import_hook``.

    """

    share_logs = staticmethod(
        lambda default="": EnvVar('DPA_SHARE_LOGS', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_SHARE_LOGS``

    The shared logging directory for dpa pipeline log files.

    """

    ptask_prompt = staticmethod(
        lambda default="": EnvVar('DPA_PTASK_PROMPT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_PTASK_PROMPT``

    This defines the shell prompt when a ptask is set.

    """

    no_ptask_prompt = staticmethod(
        lambda default="": EnvVar('DPA_NO_PTASK_PROMPT', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_NO_PTASK_PROMPT``

    This defines the shell prompt when a ptask is not set.

    """

    data_server = staticmethod(
        lambda default="": EnvVar('DPA_DATA_SERVER', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_DATA_SERVER``

    This defines the server to connect to in order to query pipeline data.

    """

    cheesyq_data_server = staticmethod(
        lambda default="": EnvVar('CHEESYQ_DATA_SERVER', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$CHEESYQ_DATA_SERVER``

    This defines the server to connect to in order to query pipeline data.

    """

    location_code = staticmethod(
        lambda default="": EnvVar('DPA_LOCATION_CODE', default)
    )
    """Returns an instance of :py:obj:`dpa.env.EnvVar` for ``$DPA_LOCATION_CODE``

    This defines the physical location of the current environment.

    """

<|endoftext|>


<commit_before>import numpy as np

from astropy.io import fits
from astropy.io.fits import PrimaryHDU, ImageHDU, CompImageHDU, Header, HDUList
from astropy.wcs import WCS
from astropy.extern import six

def parse_input_data(input_data, hdu_in=None):
    """
    Parse input data to return a Numpy array and WCS object.
    """

    if isinstance(input_data, six.string_types):
        return parse_input_data(fits.open(input_data), hdu_in=hdu_in)
    elif isinstance(input_data, HDUList):
        if hdu_in is None:
            if len(input_data) > 1:
                raise ValueError("More than one HDU is present, please specify HDU to use with ``hdu_in=`` option")
            else:
                hdu_in = 0
        return parse_input_data(input_data[hdu_in])
    elif isinstance(input_data, (PrimaryHDU, ImageHDU, CompImageHDU)):
        return input_data.data, WCS(input_data.header)
    elif isinstance(input_data, tuple) and isinstance(input_data[0], np.ndarray):
        if isinstance(input_data[1], Header):
            return input_data[0], WCS(input_data[1])
        else:
            return input_data
    else:
        raise TypeError("input_data should either be an HDU object or a tuple of (array, WCS) or (array, Header)")

def parse_output_projection(output_projection, shape_out=None):

    if isinstance(output_projection, Header):
        wcs_out = WCS(output_projection)
        try:
            shape_out = [output_projection['NAXIS{0}'.format(i + 1)] for i in range(output_projection['NAXIS'])][::-1]
        except KeyError:
            if shape_out is None:
                raise ValueError("Need to specify shape since output header does not contain complete shape information")
		if output_projection['NAXIS'] == 0:
			raise ValueError("The shape of the output header should not be an empty tuple")
    elif isinstance(output_projection, WCS):
        wcs_out = output_projection
        if shape_out is None:
            raise ValueError("Need to specify shape when specifying output_projection as WCS object")
    else:
        raise TypeError('output_projection should either be a Header or a WCS object')

    return wcs_out, shape_out
<commit_msg>Fix inconsistent indention<commit_after>import numpy as np

from astropy.io import fits
from astropy.io.fits import PrimaryHDU, ImageHDU, CompImageHDU, Header, HDUList
from astropy.wcs import WCS
from astropy.extern import six

def parse_input_data(input_data, hdu_in=None):
    """
    Parse input data to return a Numpy array and WCS object.
    """

    if isinstance(input_data, six.string_types):
        return parse_input_data(fits.open(input_data), hdu_in=hdu_in)
    elif isinstance(input_data, HDUList):
        if hdu_in is None:
            if len(input_data) > 1:
                raise ValueError("More than one HDU is present, please specify HDU to use with ``hdu_in=`` option")
            else:
                hdu_in = 0
        return parse_input_data(input_data[hdu_in])
    elif isinstance(input_data, (PrimaryHDU, ImageHDU, CompImageHDU)):
        return input_data.data, WCS(input_data.header)
    elif isinstance(input_data, tuple) and isinstance(input_data[0], np.ndarray):
        if isinstance(input_data[1], Header):
            return input_data[0], WCS(input_data[1])
        else:
            return input_data
    else:
        raise TypeError("input_data should either be an HDU object or a tuple of (array, WCS) or (array, Header)")

def parse_output_projection(output_projection, shape_out=None):

    if isinstance(output_projection, Header):
        wcs_out = WCS(output_projection)
        try:
            shape_out = [output_projection['NAXIS{0}'.format(i + 1)] for i in range(output_projection['NAXIS'])][::-1]
        except KeyError:
            if shape_out is None:
                raise ValueError("Need to specify shape since output header does not contain complete shape information")
        if output_projection['NAXIS'] == 0:
            raise ValueError("The shape of the output header should not be an empty tuple")
    elif isinstance(output_projection, WCS):
        wcs_out = output_projection
        if shape_out is None:
            raise ValueError("Need to specify shape when specifying output_projection as WCS object")
    else:
        raise TypeError('output_projection should either be a Header or a WCS object')

    return wcs_out, shape_out
<|endoftext|>


<commit_before>##############################################################################
#
#    Copyright (C) 2020 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Soulaymane Lamrani <soulaymane.lamrani@protonmail.com>
#
#    The licence is in the file __manifest__.py
#
##############################################################################
import datetime

from addons import crm
from odoo import models, fields, api

class MailActivity(models.Model):
    _inherit = "mail.activity"

    phonecall_id = fields.Many2one("crm.phonecall", "Phonecall")

    @api.model
    def create(self, vals):
        if "activity_type_id" in vals:
            if vals["activity_type_id"] == \
                    self.env.ref("mail.mail_activity_data_call").id:
                if vals["res_model_id"] == self.env["ir.model"].search([(
                    "model", "=", "res.partner"
                )]).id:
                    vals["phonecall_id"] = self.env["crm.phonecall"].create({
                        "date": vals["date_deadline"],
                        "name": vals["summary"],
                        "partner_id": vals["res_id"],
                        "user_id": vals["user_id"],
                        "direction": "outbound",
                        "state": "open"
                    }).id

                if vals["res_model_id"] == self.env["ir.model"].search([(
                    "model", "=", "crm.lead"
                )]).id:
                    vals["phonecall_ids"] = self.env["crm.phonecall"].create({
                        "date": vals["date_deadline"],
                        "name": vals["summary"],
                        "opportunity_id": vals["res_id"],
                        "partner_id":
                            self.env["crm.lead"].browse(vals["res_id"]).partner_id.id,
                        "user_id": vals["user_id"],
                        "direction": "outbound",
                        "state": "open"
                    }).id
        return super(MailActivity, self.sudo()).create(vals)

    def action_feedback(self, feedback=False):
        self.mapped("phonecall_id").write({
            "state": "done"
        })
        return super(MailActivity, self).action_feedback(feedback)
<commit_msg>CI-752 Code refactor<commit_after>##############################################################################
#
#    Copyright (C) 2020 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Soulaymane Lamrani <soulaymane.lamrani@protonmail.com>
#
#    The licence is in the file __manifest__.py
#
##############################################################################
from odoo import models, fields, api

class MailActivity(models.Model):
    _inherit = "mail.activity"

    phonecall_id = fields.Many2one("crm.phonecall", "Phonecall")

    @api.model
    def create(self, vals_list):
        """
        Links mail activities on crm.phonecall
        """
        if isinstance(vals_list, dict):
            vals_list = [vals_list]
        call_activity_id = self.env.ref("mail.mail_activity_data_call").id
        partner_model_id = self.env.ref("base.model_res_partner").id
        lead_model_id = self.env.ref("crm.model_crm_lead").id
        for vals in vals_list:
            if vals.get("activity_type_id") == call_activity_id:
                phonecall_vals = {
                    "date": vals["date_deadline"],
                    "name": vals.get("summary", "Phonecall"),
                    "user_id": vals["user_id"],
                    "direction": "outbound",
                    "state": "open"
                }
                if vals["res_model_id"] == partner_model_id:
                    phonecall_vals["partner_id"] = vals["res_id"]

                elif vals["res_model_id"] == lead_model_id:
                    phonecall_vals.update({
                        "opportunity_id": vals["res_id"],
                        "partner_id": self.env["crm.lead"].browse(
                            vals["res_id"]).partner_id.id,
                    })
                else:
                    model = self.env["ir.model"].browse(vals["res_model_id"]).model
                    if hasattr(self.env[model], "partner_id"):
                        phonecall_vals["partner_id"] = self.env[model].browse(
                            vals["res_id"]).partner_id.id
                vals["phonecall_id"] = self.env["crm.phonecall"].create(
                    phonecall_vals).id
        return super().create(vals_list)

    def action_feedback(self, feedback=False):
        self.mapped("phonecall_id").write({
            "state": "done"
        })
        return super().action_feedback(feedback)
<|endoftext|>


<commit_before>#!/usr/bin/env python
'''
This is a quick demo hack of a wavenet generator.

Carl Quillen
'''

from __future__ import print_function
import optparse
import sys
import time
from operator import mul

import numpy as np
import tensorflow as tf
import tensorflow.contrib.layers as layers
import librosa
from tensorflow.contrib.framework import arg_scope
from ops import mu_law_encode, mu_law_decode
from wavenet import wavenet

parser = optparse.OptionParser()
parser.add_option('-p', '--param_file', dest='param_file',
                  default=None, help='File to set parameters')
parser.add_option('-l', '--logdir', dest='logdir',
                  default=None, help='Tensorflow event logdir')
parser.add_option('-i', '--input_file', dest='input_file',
                  default=None, help='Input checkpoint file')
parser.add_option('-o', '--output_file', dest='output_file',
                  default='wn_sample.wav', help='Output generated wav file')
parser.add_option('-Z', '--audio_chunk_size', dest='audio_chunk_size',
                  type=int, default=100000, help='Audio chunk size per batch.')
parser.add_option('-H', '--histogram_summaries', dest='histogram_summaries',
                  action='store_true', default=False,
                  help='Do histogram summaries')
parser.add_option('-b', '--batch_norm', dest='batch_norm',
                  action='store_true', default=False,
                  help='Do batch normalization')
parser.add_option('-n', '--num_samples', default=32000, dest='num_samples',
                  type=int, help='Samples to generate')

opts, cmdline_args = parser.parse_args()

# Further options *must* come from a parameter file.
# TODO: add checks that everything is defined.
if opts.param_file is None:
    print("You must provide a parameter file (-p).", file=sys.stderr)
    exit(1)

with open(opts.param_file) as f:
    exec(f)

if opts.input_file is None:
    print("You must provide an input model (-i).", file=sys.stderr)
    exit(1)

input_dim = opts.quantization_channels if opts.one_hot_input else 1
prev_out = np.zeros((1, 1, input_dim), dtype=np.float32)
last_sample = tf.placeholder(tf.float32, shape=(1, 1, input_dim),
                             name='last_sample')

with tf.name_scope("Generate"):
    out = tf.nn.softmax(wavenet(last_sample, opts, is_training=False))

    max_likeli_sample = tf.reshape(
        mu_law_decode(tf.argmax(out, axis=2), opts.quantization_channels), ())

    # Sample from the output distribution to feed back into the input:
    pick = tf.cumsum(out, axis=2)
    select = tf.random_uniform(shape=())
    x = tf.reduce_sum(tf.cast(pick < select, tf.int32), axis=2)
    if opts.one_hot_input:
        out = tf.one_hot(x, depth=opts.quantization_channels)
    else:
        gen_sample = tf.reshape(
            mu_law_decode(x, opts.quantization_channels), ())
        out = tf.reshape(gen_sample, (1, 1, 1))

saver = tf.train.Saver(tf.trainable_variables())
init = tf.global_variables_initializer()

# Finalize the graph, so that any new ops cannot be created.
# this is good for avoiding memory leaks.
tf.get_default_graph().finalize()

# Run on cpu only.
sess = tf.Session(config=tf.ConfigProto(device_count={'GPU': 0}))
sess.run(init)

print("Restoring from", opts.input_file)
saver.restore(sess, opts.input_file)

output = np.zeros((opts.num_samples), dtype=np.float32)

last_time = time.time()
for sample in xrange(opts.num_samples):
    output[sample], prev_out = sess.run(
        fetches=[max_likeli_sample, out], feed_dict={last_sample: prev_out})
    if sample % 1000 == 999:
        new_time = time.time()
        print("{} samples generated dt={:.02f}".format(sample + 1,
                                                       new_time-last_time))
        last_time = new_time
sess.close()

print("Writing to ", opts.output_file)
librosa.output.write_wav(opts.output_file, output, opts.sample_rate)
<commit_msg>Zeroize state in generation.<commit_after>#!/usr/bin/env python
'''
This is a quick demo hack of a wavenet generator.

Carl Quillen
'''

from __future__ import print_function
import optparse
import sys
import time
from operator import mul

import numpy as np
import tensorflow as tf
import tensorflow.contrib.layers as layers
import librosa
from tensorflow.contrib.framework import arg_scope
from ops import mu_law_encode, mu_law_decode
from wavenet import wavenet

parser = optparse.OptionParser()
parser.add_option('-p', '--param_file', dest='param_file',
                  default=None, help='File to set parameters')
parser.add_option('-l', '--logdir', dest='logdir',
                  default=None, help='Tensorflow event logdir')
parser.add_option('-i', '--input_file', dest='input_file',
                  default=None, help='Input checkpoint file')
parser.add_option('-o', '--output_file', dest='output_file',
                  default='wn_sample.wav', help='Output generated wav file')
parser.add_option('-b', '--batch_norm', dest='batch_norm',
                  action='store_true', default=False,
                  help='Do batch normalization')
parser.add_option("-z", '--zeros', default=16000, dest='initial_zeros',
                  type=int, help='Initial warm-up zero-buffer samples')
parser.add_option('-n', '--num_samples', default=32000, dest='num_samples',
                  type=int, help='Samples to generate')

opts, cmdline_args = parser.parse_args()

opts.histogram_summaries = False

# Further options *must* come from a parameter file.
# TODO: add checks that everything is defined.
if opts.param_file is None:
    print("You must provide a parameter file (-p).", file=sys.stderr)
    exit(1)

with open(opts.param_file) as f:
    exec(f)

if opts.input_file is None:
    print("You must provide an input model (-i).", file=sys.stderr)
    exit(1)

input_dim = opts.quantization_channels if opts.one_hot_input else 1
prev_out = np.zeros((1, 1, input_dim), dtype=np.float32)
last_sample = tf.placeholder(tf.float32, shape=(1, 1, input_dim),
                             name='last_sample')

with tf.name_scope("Generate"):
    out = tf.nn.softmax(wavenet(last_sample, opts, is_training=False))

    max_likeli_sample = tf.reshape(
        mu_law_decode(tf.argmax(out, axis=2), opts.quantization_channels), ())

    # Sample from the output distribution to feed back into the input:
    pick = tf.cumsum(out, axis=2)
    select = tf.random_uniform(shape=())
    x = tf.reduce_sum(tf.cast(pick < select, tf.int32), axis=2)
    if opts.one_hot_input:
        out = tf.one_hot(x, depth=opts.quantization_channels)
    else:
        gen_sample = tf.reshape(
            mu_law_decode(x, opts.quantization_channels), ())
        out = tf.reshape(gen_sample, (1, 1, 1))

saver = tf.train.Saver(tf.trainable_variables())
init = tf.global_variables_initializer()

if opts.initial_zeros > 0:
    with tf.name_scope("Zeroize_state"):
        if opts.one_hot_input:
            zero = tf.constant(value=opts.quantization_channels/2,
                               shape=(1, opts.initial_zeros))
            zero = tf.one_hot(zero, depth=opts.quantization_channels)
        else:
            zero = tf.constant(value=0.0, shape=(1, opts.initial_zeros, 1))
        zeroize = wavenet(zero, opts, reuse=True, pad_reuse=True,
                          is_training=False)

# Finalize the graph, so that any new ops cannot be created.
# this is good for avoiding memory leaks.
tf.get_default_graph().finalize()

# Run on cpu only.
sess = tf.Session(config=tf.ConfigProto(device_count={'GPU': 0}))
sess.run(init)

print("Restoring from", opts.input_file)
saver.restore(sess, opts.input_file)

if opts.initial_zeros > 0:
    print("Zeroing state")
    sess.run(zeroize)
    print("Starting generation")

output = np.zeros((opts.num_samples), dtype=np.float32)

last_time = time.time()
for sample in xrange(opts.num_samples):
    output[sample], prev_out = sess.run(
        fetches=[max_likeli_sample, out], feed_dict={last_sample: prev_out})
    if sample % 1000 == 999:
        new_time = time.time()
        print("{} samples generated dt={:.02f}".format(sample + 1,
                                                       new_time-last_time))
        last_time = new_time
sess.close()

print("Writing to ", opts.output_file)
librosa.output.write_wav(opts.output_file, output, opts.sample_rate)
<|endoftext|>


<commit_before>from __future__ import absolute_import
import subprocess
import tarfile
import os
import zipfile
import six
import sys
import traceback
from threading import Thread

from django.utils.text import get_valid_filename
from django.core.files import File
from django.conf import settings

from celery import Task
from celery import app
from celery.schedules import crontab
from celery.signals import worker_process_init

from . import settings as wooey_settings

try:
    from Queue import Empty, Queue
except ImportError:
    from queue import Empty, Queue  # python 3.x

ON_POSIX = 'posix' in sys.builtin_module_names

celery_app = app.app_or_default()

def enqueue_output(out, q):
    for line in iter(out.readline, b''):
        q.put(line)
    out.close()

def output_monitor_queue(out):
    q = Queue()
    p = Thread(target=enqueue_output, args=(out, q))
    p.start()
    return q, p

def update_from_output_queue(q, out):
    try:
        line = q.get_nowait()
    except Empty:
        return out

    out += str(line)
    return out

@worker_process_init.connect
def configure_workers(*args, **kwargs):
    # this sets up Django on nodes started by the worker daemon.
    import django
    django.setup()

class WooeyTask(Task):
    pass

    # def after_return(self, status, retval, task_id, args, kwargs, einfo):
    #     job, created = WooeyJob.objects.get_or_create(wooey_celery_id=task_id)
    #     job.content_type.wooey_celery_state = status
    #     job.save()

@celery_app.task(base=WooeyTask)
def submit_script(**kwargs):
    job_id = kwargs.pop('wooey_job')
    resubmit = kwargs.pop('wooey_resubmit', False)
    from .backend import utils
    from .models import WooeyJob, UserFile
    job = WooeyJob.objects.get(pk=job_id)

    command = utils.get_job_commands(job=job)
    if resubmit:
        # clone ourselves, setting pk=None seems hackish but it works
        job.pk = None

    # This is where the script works from -- it is what is after the media_root since that may change between
    # setups/where our user uploads are stored.
    cwd = job.get_output_path()

    abscwd = os.path.abspath(os.path.join(settings.MEDIA_ROOT, cwd))
    job.command = ' '.join(command)
    job.save_path = cwd

    utils.mkdirs(abscwd)
    # make sure we have the script, otherwise download it. This can happen if we have an ephemeral file system or are
    # executing jobs on a worker node.
    script_path = job.script_version.script_path
    if not utils.get_storage(local=True).exists(script_path.path):
        utils.get_storage(local=True).save(script_path.path, script_path.file)

    job.status = WooeyJob.RUNNING
    job.save()

    stdout, stderr = '', ''
    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=abscwd, universal_newlines=True)

    # We need to use subprocesses to capture the IO, otherwise they will block one another
    # i.e. a check against stderr will sit waiting on stderr before returning
    # we use Queues to communicate
    qout, pout = output_monitor_queue(proc.stdout)
    qerr, perr = output_monitor_queue(proc.stderr)

    prev_std = None

    # Loop until the process is complete + both stdout/stderr have EOFd
    while proc.poll() is None or pout.is_alive() or perr.is_alive():

        # Check for updates from either (non-blocking)
        stdout = update_from_output_queue(qout, stdout)
        stderr = update_from_output_queue(qerr, stderr)

        # If there are changes, update the db
        if (stdout, stderr) != prev_std:
            job.update_realtime(stdout=stdout, stderr=stderr)
            prev_std = (stdout, stderr)

    # tar/zip up the generated content for bulk downloads
    def get_valid_file(cwd, name, ext):
        out = os.path.join(cwd, name)
        index = 0
        while os.path.exists(six.u('{}.{}').format(out, ext)):
            index += 1
            out = os.path.join(cwd, six.u('{}_{}').format(name, index))
        return six.u('{}.{}').format(out, ext)

    # fetch the job again in case the database connection was lost during the job or something else changed.
    job = WooeyJob.objects.get(pk=job_id)
    # if there are files generated, make zip/tar files for download
    if len(os.listdir(abscwd)):
        tar_out = get_valid_file(abscwd, get_valid_filename(job.job_name), 'tar.gz')
        tar = tarfile.open(tar_out, "w:gz")
        tar_name = os.path.splitext(os.path.splitext(os.path.split(tar_out)[1])[0])[0]
        tar.add(abscwd, arcname=tar_name)
        tar.close()

        zip_out = get_valid_file(abscwd, get_valid_filename(job.job_name), 'zip')
        zip = zipfile.ZipFile(zip_out, "w")
        arcname = os.path.splitext(os.path.split(zip_out)[1])[0]
        zip.write(abscwd, arcname=arcname)
        for root, folders, filenames in os.walk(os.path.split(zip_out)[0]):
            for filename in filenames:
                path = os.path.join(root, filename)
                if path == tar_out:
                    continue
                if path == zip_out:
                    continue
                try:
                    zip.write(path, arcname=os.path.join(arcname, filename))
                except:
                    stderr = '{}\n{}'.format(stderr, traceback.format_exc())
        try:
            zip.close()
        except:
            stderr = '{}\n{}'.format(stderr, traceback.format_exc())

        # save all the files generated as well to our default storage for ephemeral storage setups
        if wooey_settings.WOOEY_EPHEMERAL_FILES:
            for root, folders, files in os.walk(abscwd):
                for filename in files:
                    filepath = os.path.join(root, filename)
                    s3path = os.path.join(root[root.find(cwd):], filename)
                    remote = utils.get_storage(local=False)
                    exists = remote.exists(s3path)
                    filesize = remote.size(s3path) if exists else 0
                    if not exists or (exists and filesize == 0):
                        if exists:
                            remote.delete(s3path)
                        remote.save(s3path, File(open(filepath, 'rb')))
    utils.create_job_fileinfo(job)

    job.stdout = stdout
    job.stderr = stderr
    job.status = WooeyJob.COMPLETED
    job.update_realtime(delete=True)
    job.save()

    return (stdout, stderr)

@celery_app.task(base=WooeyTask)
def cleanup_wooey_jobs(**kwargs):
    from django.utils import timezone
    from .models import WooeyJob

    cleanup_settings = wooey_settings.WOOEY_JOB_EXPIRATION
    anon_settings = cleanup_settings.get('anonymous')
    now = timezone.now()
    if anon_settings:
        WooeyJob.objects.filter(user=None, created_date__lte=now-anon_settings).delete()
    user_settings = cleanup_settings.get('user')
    if user_settings:
        WooeyJob.objects.filter(user__isnull=False, created_date__lte=now-user_settings).delete()

celery_app.conf.update(
    CELERYBEAT_SCHEDULE = {
        'cleanup-old-jobs': {
            'task': 'wooey.tasks.cleanup_wooey_jobs',
            'schedule': crontab(hour=0, minute=0), # cleanup at midnight each day
        },
    }
)
<commit_msg>Only use universal_newlines on windows machines<commit_after>from __future__ import absolute_import
import subprocess
import tarfile
import os
import zipfile
import six
import sys
import traceback
from threading import Thread

from django.utils.text import get_valid_filename
from django.core.files import File
from django.conf import settings

from celery import Task
from celery import app
from celery.schedules import crontab
from celery.signals import worker_process_init

from . import settings as wooey_settings

try:
    from Queue import Empty, Queue
except ImportError:
    from queue import Empty, Queue  # python 3.x

ON_POSIX = 'posix' in sys.builtin_module_names
ON_WINDOWS = os.name == 'nt'

celery_app = app.app_or_default()

def enqueue_output(out, q):
    for line in iter(out.readline, b''):
        q.put(line)
    out.close()

def output_monitor_queue(out):
    q = Queue()
    p = Thread(target=enqueue_output, args=(out, q))
    p.start()
    return q, p

def update_from_output_queue(q, out):
    try:
        line = q.get_nowait()
    except Empty:
        return out

    out += str(line)
    return out

@worker_process_init.connect
def configure_workers(*args, **kwargs):
    # this sets up Django on nodes started by the worker daemon.
    import django
    django.setup()

class WooeyTask(Task):
    pass

    # def after_return(self, status, retval, task_id, args, kwargs, einfo):
    #     job, created = WooeyJob.objects.get_or_create(wooey_celery_id=task_id)
    #     job.content_type.wooey_celery_state = status
    #     job.save()

@celery_app.task(base=WooeyTask)
def submit_script(**kwargs):
    job_id = kwargs.pop('wooey_job')
    resubmit = kwargs.pop('wooey_resubmit', False)
    from .backend import utils
    from .models import WooeyJob, UserFile
    job = WooeyJob.objects.get(pk=job_id)

    command = utils.get_job_commands(job=job)
    if resubmit:
        # clone ourselves, setting pk=None seems hackish but it works
        job.pk = None

    # This is where the script works from -- it is what is after the media_root since that may change between
    # setups/where our user uploads are stored.
    cwd = job.get_output_path()

    abscwd = os.path.abspath(os.path.join(settings.MEDIA_ROOT, cwd))
    job.command = ' '.join(command)
    job.save_path = cwd

    utils.mkdirs(abscwd)
    # make sure we have the script, otherwise download it. This can happen if we have an ephemeral file system or are
    # executing jobs on a worker node.
    script_path = job.script_version.script_path
    if not utils.get_storage(local=True).exists(script_path.path):
        utils.get_storage(local=True).save(script_path.path, script_path.file)

    job.status = WooeyJob.RUNNING
    job.save()

    stdout, stderr = '', ''
    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=abscwd, universal_newlines=True if ON_WINDOWS else False)

    # We need to use subprocesses to capture the IO, otherwise they will block one another
    # i.e. a check against stderr will sit waiting on stderr before returning
    # we use Queues to communicate
    qout, pout = output_monitor_queue(proc.stdout)
    qerr, perr = output_monitor_queue(proc.stderr)

    prev_std = None

    # Loop until the process is complete + both stdout/stderr have EOFd
    while proc.poll() is None or pout.is_alive() or perr.is_alive():

        # Check for updates from either (non-blocking)
        stdout = update_from_output_queue(qout, stdout)
        stderr = update_from_output_queue(qerr, stderr)

        # If there are changes, update the db
        if (stdout, stderr) != prev_std:
            job.update_realtime(stdout=stdout, stderr=stderr)
            prev_std = (stdout, stderr)

    # tar/zip up the generated content for bulk downloads
    def get_valid_file(cwd, name, ext):
        out = os.path.join(cwd, name)
        index = 0
        while os.path.exists(six.u('{}.{}').format(out, ext)):
            index += 1
            out = os.path.join(cwd, six.u('{}_{}').format(name, index))
        return six.u('{}.{}').format(out, ext)

    # fetch the job again in case the database connection was lost during the job or something else changed.
    job = WooeyJob.objects.get(pk=job_id)
    # if there are files generated, make zip/tar files for download
    if len(os.listdir(abscwd)):
        tar_out = get_valid_file(abscwd, get_valid_filename(job.job_name), 'tar.gz')
        tar = tarfile.open(tar_out, "w:gz")
        tar_name = os.path.splitext(os.path.splitext(os.path.split(tar_out)[1])[0])[0]
        tar.add(abscwd, arcname=tar_name)
        tar.close()

        zip_out = get_valid_file(abscwd, get_valid_filename(job.job_name), 'zip')
        zip = zipfile.ZipFile(zip_out, "w")
        arcname = os.path.splitext(os.path.split(zip_out)[1])[0]
        zip.write(abscwd, arcname=arcname)
        for root, folders, filenames in os.walk(os.path.split(zip_out)[0]):
            for filename in filenames:
                path = os.path.join(root, filename)
                if path == tar_out:
                    continue
                if path == zip_out:
                    continue
                try:
                    zip.write(path, arcname=os.path.join(arcname, filename))
                except:
                    stderr = '{}\n{}'.format(stderr, traceback.format_exc())
        try:
            zip.close()
        except:
            stderr = '{}\n{}'.format(stderr, traceback.format_exc())

        # save all the files generated as well to our default storage for ephemeral storage setups
        if wooey_settings.WOOEY_EPHEMERAL_FILES:
            for root, folders, files in os.walk(abscwd):
                for filename in files:
                    filepath = os.path.join(root, filename)
                    s3path = os.path.join(root[root.find(cwd):], filename)
                    remote = utils.get_storage(local=False)
                    exists = remote.exists(s3path)
                    filesize = remote.size(s3path) if exists else 0
                    if not exists or (exists and filesize == 0):
                        if exists:
                            remote.delete(s3path)
                        remote.save(s3path, File(open(filepath, 'rb')))
    utils.create_job_fileinfo(job)

    job.stdout = stdout
    job.stderr = stderr
    job.status = WooeyJob.COMPLETED
    job.update_realtime(delete=True)
    job.save()

    return (stdout, stderr)

@celery_app.task(base=WooeyTask)
def cleanup_wooey_jobs(**kwargs):
    from django.utils import timezone
    from .models import WooeyJob

    cleanup_settings = wooey_settings.WOOEY_JOB_EXPIRATION
    anon_settings = cleanup_settings.get('anonymous')
    now = timezone.now()
    if anon_settings:
        WooeyJob.objects.filter(user=None, created_date__lte=now-anon_settings).delete()
    user_settings = cleanup_settings.get('user')
    if user_settings:
        WooeyJob.objects.filter(user__isnull=False, created_date__lte=now-user_settings).delete()

celery_app.conf.update(
    CELERYBEAT_SCHEDULE = {
        'cleanup-old-jobs': {
            'task': 'wooey.tasks.cleanup_wooey_jobs',
            'schedule': crontab(hour=0, minute=0), # cleanup at midnight each day
        },
    }
)
<|endoftext|>