<commit_before>from celery import shared_task
from datetime import datetime
from zeit.cms.i18n import MessageFactory as _
from zeit.cms.workflow.interfaces import CAN_PUBLISH_ERROR
import logging
import os.path
import pytz
import subprocess
import tempfile
import threading
import time
import z3c.celery.celery
import zeit.cms.checkout.interfaces
import zeit.cms.interfaces
import zeit.cms.repository.interfaces
import zeit.cms.workflow.interfaces
import zeit.connector.interfaces
import zeit.objectlog.interfaces
import zope.app.appsetup.product
import zope.app.security.interfaces
import zope.component
import zope.event
import zope.i18n
import zope.interface
import zope.publisher.interfaces
import zope.security.management

logger = logging.getLogger(__name__)
timer_logger = logging.getLogger('zeit.workflow.timer')

class Timer(threading.local):

    def start(self, message):
        self.times = []
        self.mark(message)

    def mark(self, message):
        self.times.append((time.time(), message))

    def get_timings(self):
        result = []
        last = None
        total = 0
        for when, message in self.times:
            if last is None:
                diff = 0
            else:
                diff = when - last
            total += diff
            last = when
            result.append((diff, total, message))
        return result

    def __unicode__(self):
        result = []
        for diff, total, message in self.get_timings():
            result.append(u'%2.4f %2.4f %s' % (diff, total, message))
        return u'\n'.join(result)

timer = Timer()

class Publish(object):

    zope.interface.implements(zeit.cms.workflow.interfaces.IPublish)
    zope.component.adapts(zeit.cms.repository.interfaces.IRepositoryContent)

    def __init__(self, context):
        self.context = context

    def publish(self, priority=None, async=True):
        """Publish object."""
        info = zeit.cms.workflow.interfaces.IPublishInfo(self.context)
        if info.can_publish() == CAN_PUBLISH_ERROR:
            raise zeit.cms.workflow.interfaces.PublishingError(
                "Publish pre-conditions not satisifed.")

        return self._execute_task(
            PUBLISH_TASK, _('Publication scheduled'), priority, async)

    def retract(self, priority=None, async=True):
        """Retract object."""
        return self._execute_task(
            RETRACT_TASK, _('Retracting scheduled'), priority, async)

    def _execute_task(self, task, message, priority, async):
        if async:
            self.log(self.context, message)
            return task.apply_async(
                (self.context.uniqueId,), urgency=self.get_urgency(priority))
        else:
            task(self.context.uniqueId)

    def log(self, obj, msg):
        log = zope.component.getUtility(zeit.objectlog.interfaces.IObjectLog)
        log.log(obj, msg)

    def get_urgency(self, priority):
        if priority is None:
            priority = zeit.cms.workflow.interfaces.IPublishPriority(
                self.context)
        return priority

class PublishRetractTask(object):

    def __init__(self, jobid):
        self.jobid = jobid

    def run(self, uniqueId):
        info = (type(self).__name__, uniqueId, self.jobid)
        __traceback_info__ = info
        timer.start(u'Job %s started: %s (%s)' % info)
        logger.info("Running job %s for %s", self.jobid, uniqueId)

        obj = self.repository.getContent(uniqueId)
        info = zeit.cms.workflow.interfaces.IPublishInfo(obj)
        timer.mark('Looked up object')

        try:
            result = self._run(obj, info)
        except Exception as e:
            logger.error("Error during publish/retract", exc_info=True)
            error_message = _(
                "Error during publish/retract: ${exc}: ${message}",
                mapping=dict(exc=e.__class__.__name__, message=str(e)))
            self._log_timer(uniqueId)
            raise z3c.celery.celery.HandleAfterAbort(
                self.log, obj, error_message,
                message=zope.i18n.translate(
                    error_message, target_language='de'))
        self._log_timer(uniqueId)
        return result

    def _log_timer(self, uniqueId):
        timer.mark('Done %s' % uniqueId)
        timer_logger.debug('Timings:\n%s' % (unicode(timer).encode('utf8'),))
        dummy, total, timer_message = timer.get_timings()[-1]
        logger.info('%s (%2.4fs)' % (timer_message, total))

    def cycle(self, obj):
        """checkout/checkin obj to sync data as necessary.

        The basic idea is that there are some event handlers which sync
        properties to xml on checkout/checkin.

        """
        manager = zeit.cms.checkout.interfaces.ICheckoutManager(obj)
        try:
            # We do not use the user's workingcopy but a "fresh" one which we
            # just throw away afterwards. This has two effects: 1. The users'
            # workingcopy istn't cluttered with ghosts and 2. we can publish in
            # parallel.
            checked_out = manager.checkout(temporary=True, publishing=True)
        except zeit.cms.checkout.interfaces.CheckinCheckoutError:
            logger.warning("Could not checkout %s" % obj.uniqueId)
            return obj
        manager = zeit.cms.checkout.interfaces.ICheckinManager(checked_out)
        try:
            obj = manager.checkin(publishing=True)
        except zeit.cms.checkout.interfaces.CheckinCheckoutError:
            # XXX this codepath is not tested!
            logger.warning("Could not checkin %s" % obj.uniqueId)
            del checked_out.__parent__[checked_out.__name__]
            return obj
        timer.mark('Cycled %s' % obj.uniqueId)
        return obj

    def recurse(self, method, dependencies, obj, *args):
        """Apply method recursively on obj."""
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        DEPENDENCY_PUBLISH_LIMIT = config['dependency-publish-limit']
        stack = [obj]
        seen = set()
        result_obj = None
        while stack:
            current_obj = stack.pop(0)
            if current_obj.uniqueId in seen:
                continue
            seen.add(current_obj.uniqueId)
            logger.debug('%s %s' % (method, current_obj.uniqueId))
            new_obj = method(current_obj, *args)
            timer.mark('Called %s on %s' % (method.__name__,
                                            current_obj.uniqueId))
            if len(seen) > DEPENDENCY_PUBLISH_LIMIT:
                # "strictly greater" comparison since the starting object
                # should not count towards the limit
                break

            # Dive into folders
            if zeit.cms.repository.interfaces.ICollection.providedBy(new_obj):
                stack.extend(new_obj.values())
                timer.mark('Recursed into %s' % (new_obj.uniqueId,))

            if dependencies:
                # Dive into dependent objects
                deps = zeit.workflow.interfaces.IPublicationDependencies(
                    new_obj)
                stack.extend(deps.get_dependencies())
                timer.mark('Got dependencies for %s' % (new_obj.uniqueId,))

            if result_obj is None:
                result_obj = new_obj

        return result_obj

    def get_all_paths(self, obj, dependencies):
        unique_ids = []
        self.recurse(self.get_unique_id, dependencies, obj, unique_ids)
        # The publish/retract scripts doesn't want URLs but local paths, so
        # munge them.
        paths = [self.convert_uid_to_path(uid) for uid in unique_ids]
        return paths

    def get_unique_id(self, obj, unique_ids):
        unique_ids.append(obj.uniqueId)
        return obj

    def convert_uid_to_path(self, uid):
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        path_prefix = config['path-prefix']
        return os.path.join(
            path_prefix,
            uid.replace(zeit.cms.interfaces.ID_NAMESPACE, '', 1))

    @property
    def log(self):
        return zope.component.getUtility(
            zeit.objectlog.interfaces.IObjectLog).log

    @property
    def repository(self):
        return zope.component.getUtility(
            zeit.cms.repository.interfaces.IRepository)

    @staticmethod
    def lock(obj, master=None):
        zope.event.notify(
            zeit.connector.interfaces.ResourceInvaliatedEvent(obj.uniqueId))
        lockable = zope.app.locking.interfaces.ILockable(obj, None)
        if (lockable is not None and
                not lockable.locked() and
                not lockable.ownLock()):
            lockable.lock(timeout=240)
        timer.mark('Locked %s' % obj.uniqueId)
        return obj

    @staticmethod
    def unlock(obj, master=None):
        lockable = zope.app.locking.interfaces.ILockable(obj, None)
        if (lockable is not None and
                lockable.locked() and
                lockable.ownLock()):
            lockable.unlock()
        timer.mark('Unlocked %s' % obj.uniqueId)
        return obj

    @staticmethod
    def call_script(filename, input_data):
        if isinstance(input_data, unicode):
            input_data = input_data.encode('UTF-8')
        with tempfile.NamedTemporaryFile() as f:
            f.write(input_data)
            f.flush()

            out = tempfile.NamedTemporaryFile()
            err = tempfile.NamedTemporaryFile()
            proc = subprocess.Popen([filename, f.name], stdout=out, stderr=err)
            proc.communicate()
            out.seek(0)
            err.seek(0)
            stdout = out.read()
            stderr = err.read()
            out.close()
            err.close()

            if proc.returncode:
                logger.error("%s exited with %s" % (filename, proc.returncode))
            if stdout:
                logger.info("%s:\n%s" % (filename, stdout))
            if stderr:
                logger.error("%s:\n%s" % (filename, stderr))
            if proc.returncode:
                raise zeit.workflow.interfaces.ScriptError(
                    stderr, proc.returncode)

class PublishTask(PublishRetractTask):
    """Publish object."""

    def _run(self, obj, info):
        logger.info('Publishing %s' % obj.uniqueId)
        if info.can_publish() == CAN_PUBLISH_ERROR:
            logger.error("Could not publish %s" % obj.uniqueId)
            self.log(
                obj, _("Could not publish because conditions not satisifed."))
            return "Could not publish because conditions not satisfied."

        obj = self.recurse(self.lock, True, obj, obj)
        obj = self.recurse(self.before_publish, True, obj, obj)
        self.call_publish_script(obj)
        self.recurse(self.after_publish, True, obj, obj)
        obj = self.recurse(self.unlock, True, obj, obj)
        return "Published."

    def before_publish(self, obj, master):
        """Do everything necessary before the actual publish."""

        zope.event.notify(
            zeit.cms.workflow.interfaces.BeforePublishEvent(obj, master))
        timer.mark('Sent BeforePublishEvent for %s' % obj.uniqueId)

        info = zeit.cms.workflow.interfaces.IPublishInfo(obj)
        info.published = True
        info.date_last_published = datetime.now(pytz.UTC)
        timer.mark('Set date_last_published')
        if not info.date_first_released:
            info.date_first_released = info.date_last_published
            timer.mark('Set date_first_releaesd')

        new_obj = self.cycle(obj)
        return new_obj

    def call_publish_script(self, obj):
        """Actually do the publication."""
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        publish_script = config['publish-script']
        paths = self.get_all_paths(obj, True)
        self.call_script(publish_script, '\n'.join(paths))
        timer.mark('Called publish script')

    def after_publish(self, obj, master):
        self.log(obj, _('Published'))
        zope.event.notify(zeit.cms.workflow.interfaces.PublishedEvent(
            obj, master))
        timer.mark('Sent PublishedEvent for %s' % obj.uniqueId)
        return obj

@shared_task(bind=True)
def PUBLISH_TASK(self, uniqueId):
    return PublishTask(self.task_id).run(uniqueId)

class RetractTask(PublishRetractTask):
    """Retract an object."""

    def _run(self, obj, info):
        logger.info('Retracting %s' % obj.uniqueId)
        if not info.published:
            logger.warning(
                "Retracting object %s which is not published." % obj.uniqueId)

        obj = self.recurse(self.before_retract, False, obj, obj)
        self.call_retract_script(obj)
        self.recurse(self.after_retract, False, obj, obj)
        return "Retracted."

    def before_retract(self, obj, master):
        """Do things before the actual retract."""
        self.lock(obj)
        zope.event.notify(
            zeit.cms.workflow.interfaces.BeforeRetractEvent(obj, master))
        info = zeit.cms.workflow.interfaces.IPublishInfo(obj)
        info.published = False
        self.log(obj, _('Retracted'))
        return obj

    def call_retract_script(self, obj):
        """Call the script. This does the actual retract."""
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        retract_script = config['retract-script']
        paths = reversed(self.get_all_paths(obj, False))
        self.call_script(retract_script, '\n'.join(paths))

    def after_retract(self, obj, master):
        """Do things after retract."""
        zope.event.notify(zeit.cms.workflow.interfaces.RetractedEvent(
            obj, master))
        obj = self.cycle(obj)
        self.unlock(obj)
        return obj

    @property
    def repository(self):
        return zope.component.getUtility(
            zeit.cms.repository.interfaces.IRepository)

@shared_task(bind=True)
def RETRACT_TASK(self, uniqueId):
    return RetractTask(self.task_id).run(uniqueId)
<commit_msg>zeit.retresco runs publish inline, do not break in this case.<commit_after>from celery import shared_task
from datetime import datetime
from zeit.cms.i18n import MessageFactory as _
from zeit.cms.workflow.interfaces import CAN_PUBLISH_ERROR
import logging
import os.path
import pytz
import subprocess
import tempfile
import threading
import time
import z3c.celery.celery
import zeit.cms.checkout.interfaces
import zeit.cms.interfaces
import zeit.cms.repository.interfaces
import zeit.cms.workflow.interfaces
import zeit.connector.interfaces
import zeit.objectlog.interfaces
import zope.app.appsetup.product
import zope.app.security.interfaces
import zope.component
import zope.event
import zope.i18n
import zope.interface
import zope.publisher.interfaces
import zope.security.management

logger = logging.getLogger(__name__)
timer_logger = logging.getLogger('zeit.workflow.timer')

class Timer(threading.local):

    def start(self, message):
        self.times = []
        self.mark(message)

    def mark(self, message):
        self.times.append((time.time(), message))

    def get_timings(self):
        result = []
        last = None
        total = 0
        for when, message in self.times:
            if last is None:
                diff = 0
            else:
                diff = when - last
            total += diff
            last = when
            result.append((diff, total, message))
        return result

    def __unicode__(self):
        result = []
        for diff, total, message in self.get_timings():
            result.append(u'%2.4f %2.4f %s' % (diff, total, message))
        return u'\n'.join(result)

timer = Timer()

class Publish(object):

    zope.interface.implements(zeit.cms.workflow.interfaces.IPublish)
    zope.component.adapts(zeit.cms.repository.interfaces.IRepositoryContent)

    def __init__(self, context):
        self.context = context

    def publish(self, priority=None, async=True):
        """Publish object."""
        info = zeit.cms.workflow.interfaces.IPublishInfo(self.context)
        if info.can_publish() == CAN_PUBLISH_ERROR:
            raise zeit.cms.workflow.interfaces.PublishingError(
                "Publish pre-conditions not satisifed.")

        return self._execute_task(
            PUBLISH_TASK, _('Publication scheduled'), priority, async)

    def retract(self, priority=None, async=True):
        """Retract object."""
        return self._execute_task(
            RETRACT_TASK, _('Retracting scheduled'), priority, async)

    def _execute_task(self, task, message, priority, async):
        if async:
            self.log(self.context, message)
            return task.apply_async(
                (self.context.uniqueId,), urgency=self.get_urgency(priority))
        else:
            task(self.context.uniqueId)

    def log(self, obj, msg):
        log = zope.component.getUtility(zeit.objectlog.interfaces.IObjectLog)
        log.log(obj, msg)

    def get_urgency(self, priority):
        if priority is None:
            priority = zeit.cms.workflow.interfaces.IPublishPriority(
                self.context)
        return priority

class PublishRetractTask(object):

    def __init__(self, jobid):
        self.jobid = jobid

    def run(self, uniqueId):
        info = (type(self).__name__, uniqueId, self.jobid)
        __traceback_info__ = info
        timer.start(u'Job %s started: %s (%s)' % info)
        logger.info("Running job %s for %s", self.jobid, uniqueId)

        obj = self.repository.getContent(uniqueId)
        info = zeit.cms.workflow.interfaces.IPublishInfo(obj)
        timer.mark('Looked up object')

        try:
            result = self._run(obj, info)
        except Exception as e:
            logger.error("Error during publish/retract", exc_info=True)
            error_message = _(
                "Error during publish/retract: ${exc}: ${message}",
                mapping=dict(exc=e.__class__.__name__, message=str(e)))
            self._log_timer(uniqueId)
            raise z3c.celery.celery.HandleAfterAbort(
                self.log, obj, error_message,
                message=zope.i18n.translate(
                    error_message, target_language='de'))
        self._log_timer(uniqueId)
        return result

    def _log_timer(self, uniqueId):
        timer.mark('Done %s' % uniqueId)
        timer_logger.debug('Timings:\n%s' % (unicode(timer).encode('utf8'),))
        dummy, total, timer_message = timer.get_timings()[-1]
        logger.info('%s (%2.4fs)' % (timer_message, total))

    def cycle(self, obj):
        """checkout/checkin obj to sync data as necessary.

        The basic idea is that there are some event handlers which sync
        properties to xml on checkout/checkin.

        """
        manager = zeit.cms.checkout.interfaces.ICheckoutManager(obj)
        try:
            # We do not use the user's workingcopy but a "fresh" one which we
            # just throw away afterwards. This has two effects: 1. The users'
            # workingcopy istn't cluttered with ghosts and 2. we can publish in
            # parallel.
            checked_out = manager.checkout(temporary=True, publishing=True)
        except zeit.cms.checkout.interfaces.CheckinCheckoutError:
            logger.warning("Could not checkout %s" % obj.uniqueId)
            return obj
        manager = zeit.cms.checkout.interfaces.ICheckinManager(checked_out)
        try:
            obj = manager.checkin(publishing=True)
        except zeit.cms.checkout.interfaces.CheckinCheckoutError:
            # XXX this codepath is not tested!
            logger.warning("Could not checkin %s" % obj.uniqueId)
            del checked_out.__parent__[checked_out.__name__]
            return obj
        timer.mark('Cycled %s' % obj.uniqueId)
        return obj

    def recurse(self, method, dependencies, obj, *args):
        """Apply method recursively on obj."""
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        DEPENDENCY_PUBLISH_LIMIT = config['dependency-publish-limit']
        stack = [obj]
        seen = set()
        result_obj = None
        while stack:
            current_obj = stack.pop(0)
            if current_obj.uniqueId in seen:
                continue
            seen.add(current_obj.uniqueId)
            logger.debug('%s %s' % (method, current_obj.uniqueId))
            new_obj = method(current_obj, *args)
            timer.mark('Called %s on %s' % (method.__name__,
                                            current_obj.uniqueId))
            if len(seen) > DEPENDENCY_PUBLISH_LIMIT:
                # "strictly greater" comparison since the starting object
                # should not count towards the limit
                break

            # Dive into folders
            if zeit.cms.repository.interfaces.ICollection.providedBy(new_obj):
                stack.extend(new_obj.values())
                timer.mark('Recursed into %s' % (new_obj.uniqueId,))

            if dependencies:
                # Dive into dependent objects
                deps = zeit.workflow.interfaces.IPublicationDependencies(
                    new_obj)
                stack.extend(deps.get_dependencies())
                timer.mark('Got dependencies for %s' % (new_obj.uniqueId,))

            if result_obj is None:
                result_obj = new_obj

        return result_obj

    def get_all_paths(self, obj, dependencies):
        unique_ids = []
        self.recurse(self.get_unique_id, dependencies, obj, unique_ids)
        # The publish/retract scripts doesn't want URLs but local paths, so
        # munge them.
        paths = [self.convert_uid_to_path(uid) for uid in unique_ids]
        return paths

    def get_unique_id(self, obj, unique_ids):
        unique_ids.append(obj.uniqueId)
        return obj

    def convert_uid_to_path(self, uid):
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        path_prefix = config['path-prefix']
        return os.path.join(
            path_prefix,
            uid.replace(zeit.cms.interfaces.ID_NAMESPACE, '', 1))

    @property
    def log(self):
        return zope.component.getUtility(
            zeit.objectlog.interfaces.IObjectLog).log

    @property
    def repository(self):
        return zope.component.getUtility(
            zeit.cms.repository.interfaces.IRepository)

    @staticmethod
    def lock(obj, master=None):
        zope.event.notify(
            zeit.connector.interfaces.ResourceInvaliatedEvent(obj.uniqueId))
        lockable = zope.app.locking.interfaces.ILockable(obj, None)
        if (lockable is not None and
                not lockable.locked() and
                not lockable.ownLock()):
            lockable.lock(timeout=240)
        timer.mark('Locked %s' % obj.uniqueId)
        return obj

    @staticmethod
    def unlock(obj, master=None):
        lockable = zope.app.locking.interfaces.ILockable(obj, None)
        if (lockable is not None and
                lockable.locked() and
                lockable.ownLock()):
            lockable.unlock()
        timer.mark('Unlocked %s' % obj.uniqueId)
        return obj

    @staticmethod
    def call_script(filename, input_data):
        if isinstance(input_data, unicode):
            input_data = input_data.encode('UTF-8')
        with tempfile.NamedTemporaryFile() as f:
            f.write(input_data)
            f.flush()

            out = tempfile.NamedTemporaryFile()
            err = tempfile.NamedTemporaryFile()
            proc = subprocess.Popen([filename, f.name], stdout=out, stderr=err)
            proc.communicate()
            out.seek(0)
            err.seek(0)
            stdout = out.read()
            stderr = err.read()
            out.close()
            err.close()

            if proc.returncode:
                logger.error("%s exited with %s" % (filename, proc.returncode))
            if stdout:
                logger.info("%s:\n%s" % (filename, stdout))
            if stderr:
                logger.error("%s:\n%s" % (filename, stderr))
            if proc.returncode:
                raise zeit.workflow.interfaces.ScriptError(
                    stderr, proc.returncode)

class PublishTask(PublishRetractTask):
    """Publish object."""

    def _run(self, obj, info):
        logger.info('Publishing %s' % obj.uniqueId)
        if info.can_publish() == CAN_PUBLISH_ERROR:
            logger.error("Could not publish %s" % obj.uniqueId)
            self.log(
                obj, _("Could not publish because conditions not satisifed."))
            return "Could not publish because conditions not satisfied."

        obj = self.recurse(self.lock, True, obj, obj)
        obj = self.recurse(self.before_publish, True, obj, obj)
        self.call_publish_script(obj)
        self.recurse(self.after_publish, True, obj, obj)
        obj = self.recurse(self.unlock, True, obj, obj)
        return "Published."

    def before_publish(self, obj, master):
        """Do everything necessary before the actual publish."""

        zope.event.notify(
            zeit.cms.workflow.interfaces.BeforePublishEvent(obj, master))
        timer.mark('Sent BeforePublishEvent for %s' % obj.uniqueId)

        info = zeit.cms.workflow.interfaces.IPublishInfo(obj)
        info.published = True
        info.date_last_published = datetime.now(pytz.UTC)
        timer.mark('Set date_last_published')
        if not info.date_first_released:
            info.date_first_released = info.date_last_published
            timer.mark('Set date_first_releaesd')

        new_obj = self.cycle(obj)
        return new_obj

    def call_publish_script(self, obj):
        """Actually do the publication."""
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        publish_script = config['publish-script']
        paths = self.get_all_paths(obj, True)
        self.call_script(publish_script, '\n'.join(paths))
        timer.mark('Called publish script')

    def after_publish(self, obj, master):
        self.log(obj, _('Published'))
        zope.event.notify(zeit.cms.workflow.interfaces.PublishedEvent(
            obj, master))
        timer.mark('Sent PublishedEvent for %s' % obj.uniqueId)
        return obj

@shared_task(bind=True)
def PUBLISH_TASK(self, uniqueId):
    return PublishTask(getattr(self, 'task_id', 'sync-task')).run(uniqueId)

class RetractTask(PublishRetractTask):
    """Retract an object."""

    def _run(self, obj, info):
        logger.info('Retracting %s' % obj.uniqueId)
        if not info.published:
            logger.warning(
                "Retracting object %s which is not published." % obj.uniqueId)

        obj = self.recurse(self.before_retract, False, obj, obj)
        self.call_retract_script(obj)
        self.recurse(self.after_retract, False, obj, obj)
        return "Retracted."

    def before_retract(self, obj, master):
        """Do things before the actual retract."""
        self.lock(obj)
        zope.event.notify(
            zeit.cms.workflow.interfaces.BeforeRetractEvent(obj, master))
        info = zeit.cms.workflow.interfaces.IPublishInfo(obj)
        info.published = False
        self.log(obj, _('Retracted'))
        return obj

    def call_retract_script(self, obj):
        """Call the script. This does the actual retract."""
        config = zope.app.appsetup.product.getProductConfiguration(
            'zeit.workflow')
        retract_script = config['retract-script']
        paths = reversed(self.get_all_paths(obj, False))
        self.call_script(retract_script, '\n'.join(paths))

    def after_retract(self, obj, master):
        """Do things after retract."""
        zope.event.notify(zeit.cms.workflow.interfaces.RetractedEvent(
            obj, master))
        obj = self.cycle(obj)
        self.unlock(obj)
        return obj

    @property
    def repository(self):
        return zope.component.getUtility(
            zeit.cms.repository.interfaces.IRepository)

@shared_task(bind=True)
def RETRACT_TASK(self, uniqueId):
    return RetractTask(self.task_id).run(uniqueId)
<|endoftext|>


<commit_before># Copyright 2016 Red Hat, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Test DestroyPool.
"""

import time
import unittest

from stratisd_client_dbus import Manager
from stratisd_client_dbus import Pool
from stratisd_client_dbus import StratisdErrorsGen
from stratisd_client_dbus import get_object

from stratisd_client_dbus._implementation import ManagerSpec

from stratisd_client_dbus._constants import TOP_OBJECT

from .._misc import checked_call
from .._misc import _device_list
from .._misc import Service

_MN = ManagerSpec.MethodNames

_DEVICE_STRATEGY = _device_list(0)

class Destroy1TestCase(unittest.TestCase):
    """
    Test 'destroy' on empty database.

    'destroy' should always succeed on an empty database.
    """
    _POOLNAME = 'deadpool'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        """
        self._service = Service()
        self._service.setUp()
        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        Destroy should succeed.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )
        self.assertEqual(rc, self._errors.OK)
        self.assertFalse(result)

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath]
        )

        self.assertEqual(rc1, self._errors.POOL_NOTFOUND)

class Destroy2TestCase(unittest.TestCase):
    """
    Test 'destroy' on database which contains the given pool and an unknown
    number of devices.
    """
    _POOLNAME = 'deadpool'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        """
        self._service = Service()
        self._service.setUp()
        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        Manager.CreatePool(
           self._proxy,
           name=self._POOLNAME,
           redundancy=0,
           force=False,
           devices=_DEVICE_STRATEGY.example()
        )
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        The pool was just created, and may or may not have devices.
        So, it may be possible to destroy it, or it may not be.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath],
        )

        if rc is self._errors.OK:
            expected_rc = self._errors.POOL_NOTFOUND
            expected_result = True
        else:
            self.assertEqual(rc, self._errors.BUSY)
            expected_rc = self._errors.OK
            expected_result = False

        self.assertEqual(rc1, expected_rc)
        self.assertEqual(result, expected_result)

class Destroy3TestCase(unittest.TestCase):
    """
    Test 'destroy' on database which contains the given pool and a volume.
    """
    _POOLNAME = 'deadpool'
    _VOLNAME = 'vol'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        Create a pool and a filesystem.
        """
        self._service = Service()
        self._service.setUp()

        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        (poolpath, _, _) = Manager.CreatePool(
           self._proxy,
           name=self._POOLNAME,
           redundancy=0,
           force=False,
           devices=_DEVICE_STRATEGY.example()
        )
        Pool.CreateFilesystems(
           get_object(poolpath),
           specs=[(self._VOLNAME, '', None)]
        )
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        This should fail since the pool has a filesystem on it.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )
        self.assertEqual(rc, self._errors.BUSY)
        self.assertEqual(result, False)

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath],
        )
        self.assertEqual(rc1, self._errors.OK)

class Destroy4TestCase(unittest.TestCase):
    """
    Test 'destroy' on database which contains the given pool with no devices.
    """
    _POOLNAME = 'deadpool'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        """
        self._service = Service()
        self._service.setUp()
        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        Manager.CreatePool(
           self._proxy,
           name=self._POOLNAME,
           redundancy=0,
           force=False,
           devices=[]
        )
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        The pool was just created and has no devices. It should always be
        possible to destroy it.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath],
        )

        self.assertEqual(rc, self._errors.OK)
        self.assertEqual(result, True)
        self.assertEqual(rc1, self._errors.POOL_NOTFOUND)
<commit_msg>Use == not "is" for return codes. They are only numbers.<commit_after># Copyright 2016 Red Hat, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Test DestroyPool.
"""

import time
import unittest

from stratisd_client_dbus import Manager
from stratisd_client_dbus import Pool
from stratisd_client_dbus import StratisdErrorsGen
from stratisd_client_dbus import get_object

from stratisd_client_dbus._implementation import ManagerSpec

from stratisd_client_dbus._constants import TOP_OBJECT

from .._misc import checked_call
from .._misc import _device_list
from .._misc import Service

_MN = ManagerSpec.MethodNames

_DEVICE_STRATEGY = _device_list(0)

class Destroy1TestCase(unittest.TestCase):
    """
    Test 'destroy' on empty database.

    'destroy' should always succeed on an empty database.
    """
    _POOLNAME = 'deadpool'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        """
        self._service = Service()
        self._service.setUp()
        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        Destroy should succeed.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )
        self.assertEqual(rc, self._errors.OK)
        self.assertFalse(result)

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath]
        )

        self.assertEqual(rc1, self._errors.POOL_NOTFOUND)

class Destroy2TestCase(unittest.TestCase):
    """
    Test 'destroy' on database which contains the given pool and an unknown
    number of devices.
    """
    _POOLNAME = 'deadpool'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        """
        self._service = Service()
        self._service.setUp()
        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        Manager.CreatePool(
           self._proxy,
           name=self._POOLNAME,
           redundancy=0,
           force=False,
           devices=_DEVICE_STRATEGY.example()
        )
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        The pool was just created, and may or may not have devices.
        So, it may be possible to destroy it, or it may not be.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath],
        )

        if rc == self._errors.OK:
            expected_rc = self._errors.POOL_NOTFOUND
            expected_result = True
        elif rc == self._errors.BUSY:
            expected_rc = self._errors.OK
            expected_result = False
        else:
            self.fail("rc must be OK or BUSY, is %d" % rc)

        self.assertEqual(rc1, expected_rc)
        self.assertEqual(result, expected_result)

class Destroy3TestCase(unittest.TestCase):
    """
    Test 'destroy' on database which contains the given pool and a volume.
    """
    _POOLNAME = 'deadpool'
    _VOLNAME = 'vol'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        Create a pool and a filesystem.
        """
        self._service = Service()
        self._service.setUp()

        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        (poolpath, _, _) = Manager.CreatePool(
           self._proxy,
           name=self._POOLNAME,
           redundancy=0,
           force=False,
           devices=_DEVICE_STRATEGY.example()
        )
        Pool.CreateFilesystems(
           get_object(poolpath),
           specs=[(self._VOLNAME, '', None)]
        )
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        This should fail since the pool has a filesystem on it.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )
        self.assertEqual(rc, self._errors.BUSY)
        self.assertEqual(result, False)

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath],
        )
        self.assertEqual(rc1, self._errors.OK)

class Destroy4TestCase(unittest.TestCase):
    """
    Test 'destroy' on database which contains the given pool with no devices.
    """
    _POOLNAME = 'deadpool'

    def setUp(self):
        """
        Start the stratisd daemon with the simulator.
        """
        self._service = Service()
        self._service.setUp()
        time.sleep(1)
        self._proxy = get_object(TOP_OBJECT)
        self._errors = StratisdErrorsGen.get_object()
        Manager.CreatePool(
           self._proxy,
           name=self._POOLNAME,
           redundancy=0,
           force=False,
           devices=[]
        )
        Manager.ConfigureSimulator(self._proxy, denominator=8)

    def tearDown(self):
        """
        Stop the stratisd simulator and daemon.
        """
        self._service.tearDown()

    def testExecution(self):
        """
        The pool was just created and has no devices. It should always be
        possible to destroy it.
        """
        (result, rc, _) = checked_call(
           Manager.DestroyPool(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.DestroyPool]
        )

        (_, rc1, _) = checked_call(
           Manager.GetPoolObjectPath(self._proxy, name=self._POOLNAME),
           ManagerSpec.OUTPUT_SIGS[_MN.GetPoolObjectPath],
        )

        self.assertEqual(rc, self._errors.OK)
        self.assertEqual(result, True)
        self.assertEqual(rc1, self._errors.POOL_NOTFOUND)
<|endoftext|>


<commit_before># ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
import time
import random

from string import ascii_letters

from azure.mgmt.keyvault import KeyVaultManagementClient
from azure.mgmt.keyvault.models import (
    SecretPermissions,
    KeyPermissions,
    CertificatePermissions,
    StoragePermissions,
    Permissions,
    Sku,
    SkuName,
    AccessPolicyEntry,
    VaultProperties,
    VaultCreateOrUpdateParameters,
)

from azure_devtools.scenario_tests.exceptions import(
    AzureTestError, NameInUseError, ReservedResourceNameError
)

from . import AzureMgmtPreparer, ResourceGroupPreparer
from .resource_testcase import RESOURCE_GROUP_PARAM

DEFAULT_PERMISSIONS = Permissions(
    keys=[perm.value for perm in KeyPermissions],
    secrets=[perm.value for perm in SecretPermissions],
    certificates=[perm.value for perm in CertificatePermissions],
    storage=[perm.value for perm in StoragePermissions],
)
DEFAULT_SKU = SkuName.premium.value
CLIENT_OID = "00000000-0000-0000-0000-000000000000"

class KeyVaultPreparer(AzureMgmtPreparer):
    def __init__(
        self,
        name_prefix=random.choice(ascii_letters).lower(),
        sku=DEFAULT_SKU,
        permissions=DEFAULT_PERMISSIONS,
        enabled_for_deployment=True,
        enabled_for_disk_encryption=True,
        enabled_for_template_deployment=True,
        enable_soft_delete=None,
        location='westus',
        parameter_name='vault_uri',
        resource_group_parameter_name=RESOURCE_GROUP_PARAM,
        disable_recording=True,
        playback_fake_resource=None,
        client_kwargs=None,
        random_name_enabled=True
    ):
        super(KeyVaultPreparer, self).__init__(name_prefix, 24,
                                                     disable_recording=disable_recording,
                                                     playback_fake_resource=playback_fake_resource,
                                                     client_kwargs=client_kwargs,
                                                     random_name_enabled=random_name_enabled)
        self.location = location
        self.sku = sku
        self.permissions = permissions
        self.enabled_for_deployment = enabled_for_deployment
        self.enabled_for_disk_encryption = enabled_for_disk_encryption
        self.enabled_for_template_deployment = enabled_for_template_deployment
        self.enable_soft_delete = enable_soft_delete
        self.resource_group_parameter_name = resource_group_parameter_name
        self.parameter_name = parameter_name
        if random_name_enabled:
            self.resource_moniker = "vaultname"
        self.client_oid = None

    def create_resource(self, name, **kwargs):
        self.client_oid = self.test_class_instance.set_value_to_scrub("CLIENT_OID", CLIENT_OID)
        if self.is_live:
            group = self._get_resource_group(**kwargs).name
            access_policies = [
                AccessPolicyEntry(
                    tenant_id=self.test_class_instance.get_settings_value("TENANT_ID"),
                    object_id=self.client_oid,
                    permissions=self.permissions,
                )
            ]
            properties = VaultProperties(
                tenant_id=self.test_class_instance.get_settings_value("TENANT_ID"),
                sku=Sku(name=self.sku),
                access_policies=access_policies,
                vault_uri=None,
                enabled_for_deployment=self.enabled_for_deployment,
                enabled_for_disk_encryption=self.enabled_for_disk_encryption,
                enabled_for_template_deployment=self.enabled_for_template_deployment,
                enable_soft_delete=self.enable_soft_delete,
                enable_purge_protection=None,
            )
            parameters = VaultCreateOrUpdateParameters(location=self.location, properties=properties)
            self.client = self.create_mgmt_client(KeyVaultManagementClient)

            # ARM may return not found at first even though the resource group has been created
            retries = 4
            for i in range(retries):
                try:
                    vault = self.client.vaults.create_or_update(group, name, parameters).result()
                    break
                except Exception as ex:
                    if "VaultAlreadyExists" in str(ex):
                        raise NameInUseError(name)
                    if "ReservedResourceName" in str(ex):
                        raise ReservedResourceNameError(name)
                    if "ResourceGroupNotFound" not in str(ex) or i == retries - 1:
                        raise
                    time.sleep(3)
            self.test_class_instance.scrubber.register_name_pair(
                name,
                self.resource_moniker
            )
            vault_uri = vault.properties.vault_uri
        else:
            # playback => we need only the uri used in the recording
            vault_uri = "https://{}.vault.azure.net/".format(name)
        return {self.parameter_name: vault_uri}

    def remove_resource(self, name, **kwargs):
        if self.is_live:
            group = self._get_resource_group(**kwargs).name
            self.client.vaults.delete(group, name)
            if self.enable_soft_delete:
                self.client.vaults.purge_deleted(name, self.location).wait()

    def _get_resource_group(self, **kwargs):
        try:
            return kwargs[self.resource_group_parameter_name]
        except KeyError:
            template = (
                "To create a key vault a resource group is required. Please add "
                "decorator @{} in front of this storage account preparer."
            )
            raise AzureTestError(template.format(ResourceGroupPreparer.__name__))
<commit_msg>Update KeyVaultPreparer with track 2 mgmt changes (#12060)<commit_after># ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
import time
import random

from string import ascii_letters

from azure.mgmt.keyvault import KeyVaultManagementClient
from azure.mgmt.keyvault.models import (
    SecretPermissions,
    KeyPermissions,
    CertificatePermissions,
    StoragePermissions,
    Permissions,
    Sku,
    SkuName,
    AccessPolicyEntry,
    VaultProperties,
    VaultCreateOrUpdateParameters,
)

from azure_devtools.scenario_tests.exceptions import(
    AzureTestError, NameInUseError, ReservedResourceNameError
)

from . import AzureMgmtPreparer, ResourceGroupPreparer
from .resource_testcase import RESOURCE_GROUP_PARAM

DEFAULT_PERMISSIONS = Permissions(
    keys=[perm.value for perm in KeyPermissions],
    secrets=[perm.value for perm in SecretPermissions],
    certificates=[perm.value for perm in CertificatePermissions],
    storage=[perm.value for perm in StoragePermissions],
)
DEFAULT_SKU = SkuName.premium.value
CLIENT_OID = "00000000-0000-0000-0000-000000000000"

class KeyVaultPreparer(AzureMgmtPreparer):
    def __init__(
        self,
        name_prefix=random.choice(ascii_letters).lower(),
        sku=DEFAULT_SKU,
        permissions=DEFAULT_PERMISSIONS,
        enabled_for_deployment=True,
        enabled_for_disk_encryption=True,
        enabled_for_template_deployment=True,
        enable_soft_delete=None,
        location='westus',
        parameter_name='vault_uri',
        resource_group_parameter_name=RESOURCE_GROUP_PARAM,
        disable_recording=True,
        playback_fake_resource=None,
        client_kwargs=None,
        random_name_enabled=True
    ):
        super(KeyVaultPreparer, self).__init__(name_prefix, 24,
                                                     disable_recording=disable_recording,
                                                     playback_fake_resource=playback_fake_resource,
                                                     client_kwargs=client_kwargs,
                                                     random_name_enabled=random_name_enabled)
        self.location = location
        self.sku = sku
        self.permissions = permissions
        self.enabled_for_deployment = enabled_for_deployment
        self.enabled_for_disk_encryption = enabled_for_disk_encryption
        self.enabled_for_template_deployment = enabled_for_template_deployment
        self.enable_soft_delete = enable_soft_delete
        self.resource_group_parameter_name = resource_group_parameter_name
        self.parameter_name = parameter_name
        if random_name_enabled:
            self.resource_moniker = "vaultname"
        self.client_oid = None

    def create_resource(self, name, **kwargs):
        self.client_oid = self.test_class_instance.set_value_to_scrub("CLIENT_OID", CLIENT_OID)
        if self.is_live:
            group = self._get_resource_group(**kwargs).name
            access_policies = [
                AccessPolicyEntry(
                    tenant_id=self.test_class_instance.get_settings_value("TENANT_ID"),
                    object_id=self.client_oid,
                    permissions=self.permissions,
                )
            ]
            properties = VaultProperties(
                tenant_id=self.test_class_instance.get_settings_value("TENANT_ID"),
                sku=Sku(name=self.sku),
                access_policies=access_policies,
                vault_uri=None,
                enabled_for_deployment=self.enabled_for_deployment,
                enabled_for_disk_encryption=self.enabled_for_disk_encryption,
                enabled_for_template_deployment=self.enabled_for_template_deployment,
                enable_soft_delete=self.enable_soft_delete,
                enable_purge_protection=None,
            )
            parameters = VaultCreateOrUpdateParameters(location=self.location, properties=properties)
            self.client = self.create_mgmt_client(KeyVaultManagementClient)

            # ARM may return not found at first even though the resource group has been created
            retries = 4
            for i in range(retries):
                try:
                    vault = self.client.vaults.begin_create_or_update(group, name, parameters).result()
                    break
                except Exception as ex:
                    if "VaultAlreadyExists" in str(ex):
                        raise NameInUseError(name)
                    if "ReservedResourceName" in str(ex):
                        raise ReservedResourceNameError(name)
                    if "ResourceGroupNotFound" not in str(ex) or i == retries - 1:
                        raise
                    time.sleep(3)
            self.test_class_instance.scrubber.register_name_pair(
                name,
                self.resource_moniker
            )
            vault_uri = vault.properties.vault_uri
        else:
            # playback => we need only the uri used in the recording
            vault_uri = "https://{}.vault.azure.net/".format(name)
        return {self.parameter_name: vault_uri}

    def remove_resource(self, name, **kwargs):
        if self.is_live:
            group = self._get_resource_group(**kwargs).name
            self.client.vaults.delete(group, name)
            if self.enable_soft_delete:
                self.client.vaults.purge_deleted(name, self.location).wait()

    def _get_resource_group(self, **kwargs):
        try:
            return kwargs[self.resource_group_parameter_name]
        except KeyError:
            template = (
                "To create a key vault a resource group is required. Please add "
                "decorator @{} in front of this storage account preparer."
            )
            raise AzureTestError(template.format(ResourceGroupPreparer.__name__))
<|endoftext|>


<commit_before>#!/usr/bin/env python

r"""
This module provides command execution functions such as cmd_fnc and cmd_fnc_u.
"""

import sys
import subprocess

robot_env = 1
try:
    import gen_robot_print as grp
    from robot.libraries.BuiltIn import BuiltIn
except ImportError:
    robot_env = 0
import gen_print as gp
import gen_valid as gv
import gen_misc as gm

###############################################################################
def cmd_fnc(cmd_buf,
            quiet=None,
            test_mode=None,
            debug=None,
            print_output=1,
            show_err=1):

    r"""
    Run the given command in a shell and return the shell return code.

    Description of arguments:
    cmd_buf                         The command string to be run in a shell.
    quiet                           Indicates whether this function should run
                                    the pissuing()
                  function prints an "Issuing: <cmd string>" to stdout.
    test_mode                       If test_mode is set, this function will
                                    not actually run
                  the command.
    debug                           If debug is set, this function will print
                                    extra debug info.
    print_output                    If this is set, this function will print
                                    the stdout/stderr
                  generated by the shell command.
    show_err                        If show_err is set, this function will
                                    print a standardized
                  error report if the shell command returns non-zero.
    """

    quiet = int(gm.global_default(quiet, 0))
    test_mode = int(gm.global_default(test_mode, 0))
    debug = int(gm.global_default(debug, 0))

    if debug:
        if robot_env:
            grp.rprint_var(cmd_buf)
            grp.rprint_var(quiet)
            grp.rprint_var(test_mode)
            grp.rprint_var(debug)
        else:
            gp.print_vars(cmd_buf, quiet, test_mode, debug)

    err_msg = gv.svalid_value(cmd_buf)
    if err_msg != "":
        raise ValueError(err_msg)

    if not quiet:
        if robot_env:
            grp.rpissuing(cmd_buf, test_mode)
        else:
            gp.pissuing(cmd_buf, test_mode)

    if test_mode:
        return 0, ""

    sub_proc = subprocess.Popen(cmd_buf,
                                bufsize=1,
                                shell=True,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT)
    out_buf = ""
    for line in sub_proc.stdout:
        out_buf += line
        if not print_output:
            continue
        if robot_env:
            grp.rprint(line)
        else:
            sys.stdout.write(line)
    if print_output and not robot_env:
        sys.stdout.flush()
    sub_proc.communicate()
    shell_rc = sub_proc.returncode
    if shell_rc != 0 and show_err:
        if robot_env:
            grp.rprint_error_report("The prior command failed.\n" +
                                    gp.sprint_var(shell_rc, 1))
        else:
            gp.print_error_report("The prior command failed.\n" +
                                  gp.sprint_var(shell_rc, 1))

    return shell_rc, out_buf

###############################################################################

###############################################################################
def cmd_fnc_u(cmd_buf,
              quiet=None,
              debug=None,
              print_output=1,
              show_err=1):

    r"""
    Call cmd_fnc with test_mode=0.  See cmd_fnc (above) for details.

    Note the "u" in "cmd_fnc_u" stands for "unconditional".
    """

    return cmd_fnc(cmd_buf, test_mode=0, quiet=quiet, debug=debug,
                   print_output=print_output, show_err=show_err)

###############################################################################
<commit_msg>Changed many grp references in gen_cmd.py to gp equivalents.<commit_after>#!/usr/bin/env python

r"""
This module provides command execution functions such as cmd_fnc and cmd_fnc_u.
"""

import sys
import subprocess

robot_env = 1
try:
    from robot.libraries.BuiltIn import BuiltIn
except ImportError:
    robot_env = 0
import gen_print as gp
import gen_valid as gv
import gen_misc as gm
if robot_env:
    import gen_robot_print as grp

###############################################################################
def cmd_fnc(cmd_buf,
            quiet=None,
            test_mode=None,
            debug=0,
            print_output=1,
            show_err=1):

    r"""
    Run the given command in a shell and return the shell return code.

    Description of arguments:
    cmd_buf                         The command string to be run in a shell.
    quiet                           Indicates whether this function should run
                                    the pissuing()
                  function prints an "Issuing: <cmd string>" to stdout.
    test_mode                       If test_mode is set, this function will
                                    not actually run
                  the command.
    debug                           If debug is set, this function will print
                                    extra debug info.
    print_output                    If this is set, this function will print
                                    the stdout/stderr
                  generated by the shell command.
    show_err                        If show_err is set, this function will
                                    print a standardized
                  error report if the shell command returns non-zero.
    """

    quiet = int(gm.global_default(quiet, 0))
    test_mode = int(gm.global_default(test_mode, 0))

    if debug:
        gp.print_vars(cmd_buf, quiet, test_mode, debug)

    err_msg = gv.svalid_value(cmd_buf)
    if err_msg != "":
        raise ValueError(err_msg)

    if not quiet:
        gp.pissuing(cmd_buf, test_mode)

    if test_mode:
        return 0, ""

    sub_proc = subprocess.Popen(cmd_buf,
                                bufsize=1,
                                shell=True,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT)
    out_buf = ""
    for line in sub_proc.stdout:
        out_buf += line
        if not print_output:
            continue
        if robot_env:
            grp.rprint(line)
        else:
            sys.stdout.write(line)
    if print_output and not robot_env:
        sys.stdout.flush()
    sub_proc.communicate()
    shell_rc = sub_proc.returncode
    if shell_rc != 0 and show_err:
        if robot_env:
            grp.rprint_error_report("The prior command failed.\n" +
                                    gp.sprint_var(shell_rc, 1))
        else:
            gp.print_error_report("The prior command failed.\n" +
                                  gp.sprint_var(shell_rc, 1))

    return shell_rc, out_buf

###############################################################################

###############################################################################
def cmd_fnc_u(cmd_buf,
              quiet=None,
              debug=None,
              print_output=1,
              show_err=1):

    r"""
    Call cmd_fnc with test_mode=0.  See cmd_fnc (above) for details.

    Note the "u" in "cmd_fnc_u" stands for "unconditional".
    """

    return cmd_fnc(cmd_buf, test_mode=0, quiet=quiet, debug=debug,
                   print_output=print_output, show_err=show_err)

###############################################################################
<|endoftext|>


<commit_before>import json

from medallion.filters.basic_filter import BasicFilter
import pytest
from requests.models import Response
from taxii2client import Collection, _filter_kwargs_to_query_params

from stix2 import (Bundle, TAXIICollectionSink, TAXIICollectionSource,
                   TAXIICollectionStore, ThreatActor)
from stix2.datastore import DataSourceError
from stix2.datastore.filters import Filter

COLLECTION_URL = 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'

class MockTAXIICollectionEndpoint(Collection):
    """Mock for taxii2_client.TAXIIClient"""

    def __init__(self, url, **kwargs):
        super(MockTAXIICollectionEndpoint, self).__init__(url, **kwargs)
        self.objects = []

    def add_objects(self, bundle):
        self._verify_can_write()
        if isinstance(bundle, str):
            bundle = json.loads(bundle)
        for object in bundle.get("objects", []):
            self.objects.append(object)

    def get_objects(self, **filter_kwargs):
        self._verify_can_read()
        query_params = _filter_kwargs_to_query_params(filter_kwargs)
        if not isinstance(query_params, dict):
            query_params = json.loads(query_params)
        full_filter = BasicFilter(query_params or {})
        objs = full_filter.process_filter(
            self.objects,
            ("id", "type", "version"),
            []
        )
        if objs:
            return Bundle(objects=objs)
        else:
            resp = Response()
            resp.status_code = 404
            resp.raise_for_status()

    def get_object(self, id_, version=None):
        self._verify_can_read()
        kwargs = {"id": id_}
        if version:
            kwargs.update({"version": version})

        query_params = _filter_kwargs_to_query_params(kwargs)
        full_filter = BasicFilter(query_params or {})
        objs = full_filter.process_filter(
            self.objects,
            ("version", "id"),
            []
        )
        if objs:
            return Bundle(objects=objs)
        else:
            resp = Response()
            resp.status_code = 404
            resp.raise_for_status()

@pytest.fixture
def collection(stix_objs1):
    mock = MockTAXIICollectionEndpoint(COLLECTION_URL, **{
        "id": "91a7b528-80eb-42ed-a74d-c6fbd5a26116",
        "title": "Writable Collection",
        "description": "This collection is a dropbox for submitting indicators",
        "can_read": True,
        "can_write": True,
        "media_types": [
            "application/vnd.oasis.stix+json; version=2.0"
        ]
    })

    mock.objects.extend(stix_objs1)
    return mock

@pytest.fixture
def collection_no_rw_access(stix_objs1):
    mock = MockTAXIICollectionEndpoint(COLLECTION_URL, **{
        "id": "91a7b528-80eb-42ed-a74d-c6fbd5a26116",
        "title": "Not writeable or readable Collection",
        "description": "This collection is a dropbox for submitting indicators",
        "can_read": False,
        "can_write": False,
        "media_types": [
            "application/vnd.oasis.stix+json; version=2.0"
        ]
    })

    mock.objects.extend(stix_objs1)
    return mock

def test_ds_taxii(collection):
    ds = TAXIICollectionSource(collection)
    assert ds.collection is not None

def test_add_stix2_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ])

    tc_sink.add(ta)

def test_add_stix2_with_custom_object(collection):
    tc_sink = TAXIICollectionStore(collection, allow_custom=True)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ],
                     foo="bar",
                     allow_custom=True)

    tc_sink.add(ta)

def test_add_list_object(collection, indicator):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ])

    tc_sink.add([ta, indicator])

def test_add_stix2_bundle_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ])

    tc_sink.add(Bundle(objects=[ta]))

def test_add_str_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = """{
        "type": "threat-actor",
        "id": "threat-actor--eddff64f-feb1-4469-b07c-499a73c96415",
        "created": "2018-04-23T16:40:50.847Z",
        "modified": "2018-04-23T16:40:50.847Z",
        "name": "Teddy Bear",
        "goals": [
            "compromising environment NGOs",
            "water-hole attacks geared towards energy sector"
        ],
        "sophistication": "innovator",
        "resource_level": "government",
        "labels": [
            "nation-state"
        ]
    }"""

    tc_sink.add(ta)

def test_add_dict_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    ta = {
        "type": "threat-actor",
        "id": "threat-actor--eddff64f-feb1-4469-b07c-499a73c96415",
        "created": "2018-04-23T16:40:50.847Z",
        "modified": "2018-04-23T16:40:50.847Z",
        "name": "Teddy Bear",
        "goals": [
            "compromising environment NGOs",
            "water-hole attacks geared towards energy sector"
        ],
        "sophistication": "innovator",
        "resource_level": "government",
        "labels": [
            "nation-state"
        ]
    }

    tc_sink.add(ta)

def test_add_dict_bundle_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    ta = {
        "type": "bundle",
        "id": "bundle--860ccc8d-56c9-4fda-9384-84276fb52fb1",
        "spec_version": "2.0",
        "objects": [
            {
                "type": "threat-actor",
                "id": "threat-actor--dc5a2f41-f76e-425a-81fe-33afc7aabd75",
                "created": "2018-04-23T18:45:11.390Z",
                "modified": "2018-04-23T18:45:11.390Z",
                "name": "Teddy Bear",
                "goals": [
                    "compromising environment NGOs",
                    "water-hole attacks geared towards energy sector"
                ],
                "sophistication": "innovator",
                "resource_level": "government",
                "labels": [
                    "nation-state"
                ]
            }
        ]
    }

    tc_sink.add(ta)

def test_get_stix2_object(collection):
    tc_sink = TAXIICollectionSource(collection)

    objects = tc_sink.get("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")

    assert objects

def test_parse_taxii_filters(collection):
    query = [
        Filter("added_after", "=", "2016-02-01T00:00:01.000Z"),
        Filter("id", "=", "taxii stix object ID"),
        Filter("type", "=", "taxii stix object ID"),
        Filter("version", "=", "first"),
        Filter("created_by_ref", "=", "Bane"),
    ]

    taxii_filters_expected = [
        Filter("added_after", "=", "2016-02-01T00:00:01.000Z"),
        Filter("id", "=", "taxii stix object ID"),
        Filter("type", "=", "taxii stix object ID"),
        Filter("version", "=", "first")
    ]

    ds = TAXIICollectionSource(collection)

    taxii_filters = ds._parse_taxii_filters(query)

    assert taxii_filters == taxii_filters_expected

def test_add_get_remove_filter(collection):
    ds = TAXIICollectionSource(collection)

    # First 3 filters are valid, remaining properties are erroneous in some way
    valid_filters = [
        Filter('type', '=', 'malware'),
        Filter('id', '!=', 'stix object id'),
        Filter('labels', 'in', ["heartbleed", "malicious-activity"]),
    ]

    assert len(ds.filters) == 0

    ds.filters.add(valid_filters[0])
    assert len(ds.filters) == 1

    # Addin the same filter again will have no effect since `filters` acts
    # like a set
    ds.filters.add(valid_filters[0])
    assert len(ds.filters) == 1

    ds.filters.add(valid_filters[1])
    assert len(ds.filters) == 2

    ds.filters.add(valid_filters[2])
    assert len(ds.filters) == 3

    assert valid_filters == [f for f in ds.filters]

    # remove
    ds.filters.remove(valid_filters[0])

    assert len(ds.filters) == 2

    ds.filters.add(valid_filters)

def test_get_all_versions(collection):
    ds = TAXIICollectionStore(collection)

    indicators = ds.all_versions('indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f')
    # There are 3 indicators but 2 share the same 'modified' timestamp
    assert len(indicators) == 2

def test_can_read_error(collection_no_rw_access):
    """create a TAXIICOllectionSource with a taxii2client.Collection
    instance that does not have read access, check ValueError exception is raised"""
    with pytest.raises(DataSourceError) as excinfo:
        TAXIICollectionSource(collection_no_rw_access)
    assert "Collection object provided does not have read access" in str(excinfo.value)

def test_can_write_error(collection_no_rw_access):
    """create a TAXIICOllectionSink with a taxii2client.Collection
    instance that does not have write access, check ValueError exception is raised"""
    with pytest.raises(DataSourceError) as excinfo:
        TAXIICollectionSink(collection_no_rw_access)
    assert "Collection object provided does not have write access" in str(excinfo.value)

def test_get_404(collection):
    """a TAXIICollectionSource.get() call that receives an HTTP 404 response
    code from the taxii2client should be be returned as None.

    TAXII spec states that a TAXII server can return a 404 for
    nonexistent resources or lack of access. Decided that None is acceptable
    reponse to imply that state of the TAXII endpoint.
    """
    ds = TAXIICollectionStore(collection)

    # this will raise 404 from mock TAXII Client but TAXIICollectionStore
    # should handle gracefully and return None
    stix_obj = ds.get("indicator--1")  # this will raise 404 from
    assert stix_obj is None

def test_all_versions_404(collection):
    """ a TAXIICollectionSource.all_version() call that recieves an HTTP 404
    response code from the taxii2client should be returned as an exception"""
    ds = TAXIICollectionStore(collection)
    with pytest.raises(DataSourceError) as excinfo:
        ds.all_versions("indicator--1")
    assert "are either not found or access is denied" in str(excinfo.value)
    assert "404" in str(excinfo.value)

def test_query_404(collection):
    """ a TAXIICollectionSource.query() call that recieves an HTTP 404
    response code from the taxii2client should be returned as an exception"""
    ds = TAXIICollectionStore(collection)
    query = [Filter("type", "=", "malware")]

    with pytest.raises(DataSourceError) as excinfo:
        ds.query(query=query)
    assert "are either not found or access is denied" in str(excinfo.value)
    assert "404" in str(excinfo.value)
<commit_msg>solution to get 404 mocking problem<commit_after>import json

from medallion.filters.basic_filter import BasicFilter
import pytest
from requests.models import Response
from taxii2client import Collection, _filter_kwargs_to_query_params

from stix2 import (Bundle, TAXIICollectionSink, TAXIICollectionSource,
                   TAXIICollectionStore, ThreatActor)
from stix2.datastore import DataSourceError
from stix2.datastore.filters import Filter

COLLECTION_URL = 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'

class MockTAXIICollectionEndpoint(Collection):
    """Mock for taxii2_client.TAXIIClient"""

    def __init__(self, url, **kwargs):
        super(MockTAXIICollectionEndpoint, self).__init__(url, **kwargs)
        self.objects = []

    def add_objects(self, bundle):
        self._verify_can_write()
        if isinstance(bundle, str):
            bundle = json.loads(bundle)
        for object in bundle.get("objects", []):
            self.objects.append(object)

    def get_objects(self, **filter_kwargs):
        self._verify_can_read()
        query_params = _filter_kwargs_to_query_params(filter_kwargs)
        if not isinstance(query_params, dict):
            query_params = json.loads(query_params)
        full_filter = BasicFilter(query_params or {})
        objs = full_filter.process_filter(
            self.objects,
            ("id", "type", "version"),
            []
        )
        if objs:
            return Bundle(objects=objs)
        else:
            resp = Response()
            resp.status_code = 404
            resp.raise_for_status()

    def get_object(self, id, version=None):
        self._verify_can_read()
        query_params = None
        if version:
            query_params = _filter_kwargs_to_query_params({"version": version})
        if query_params:
            query_params = json.loads(query_params)
        full_filter = BasicFilter(query_params or {})
        objs = full_filter.process_filter(
            self.objects,
            ("version",),
            []
        )
        if objs:
            return Bundle(objects=objs)
        else:
            resp = Response()
            resp.status_code = 404
            resp.raise_for_status()

@pytest.fixture
def collection(stix_objs1):
    mock = MockTAXIICollectionEndpoint(COLLECTION_URL, **{
        "id": "91a7b528-80eb-42ed-a74d-c6fbd5a26116",
        "title": "Writable Collection",
        "description": "This collection is a dropbox for submitting indicators",
        "can_read": True,
        "can_write": True,
        "media_types": [
            "application/vnd.oasis.stix+json; version=2.0"
        ]
    })

    mock.objects.extend(stix_objs1)
    return mock

@pytest.fixture
def collection_no_rw_access(stix_objs1):
    mock = MockTAXIICollectionEndpoint(COLLECTION_URL, **{
        "id": "91a7b528-80eb-42ed-a74d-c6fbd5a26116",
        "title": "Not writeable or readable Collection",
        "description": "This collection is a dropbox for submitting indicators",
        "can_read": False,
        "can_write": False,
        "media_types": [
            "application/vnd.oasis.stix+json; version=2.0"
        ]
    })

    mock.objects.extend(stix_objs1)
    return mock

def test_ds_taxii(collection):
    ds = TAXIICollectionSource(collection)
    assert ds.collection is not None

def test_add_stix2_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ])

    tc_sink.add(ta)

def test_add_stix2_with_custom_object(collection):
    tc_sink = TAXIICollectionStore(collection, allow_custom=True)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ],
                     foo="bar",
                     allow_custom=True)

    tc_sink.add(ta)

def test_add_list_object(collection, indicator):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ])

    tc_sink.add([ta, indicator])

def test_add_stix2_bundle_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = ThreatActor(name="Teddy Bear",
                     labels=["nation-state"],
                     sophistication="innovator",
                     resource_level="government",
                     goals=[
                         "compromising environment NGOs",
                         "water-hole attacks geared towards energy sector",
                     ])

    tc_sink.add(Bundle(objects=[ta]))

def test_add_str_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    # create new STIX threat-actor
    ta = """{
        "type": "threat-actor",
        "id": "threat-actor--eddff64f-feb1-4469-b07c-499a73c96415",
        "created": "2018-04-23T16:40:50.847Z",
        "modified": "2018-04-23T16:40:50.847Z",
        "name": "Teddy Bear",
        "goals": [
            "compromising environment NGOs",
            "water-hole attacks geared towards energy sector"
        ],
        "sophistication": "innovator",
        "resource_level": "government",
        "labels": [
            "nation-state"
        ]
    }"""

    tc_sink.add(ta)

def test_add_dict_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    ta = {
        "type": "threat-actor",
        "id": "threat-actor--eddff64f-feb1-4469-b07c-499a73c96415",
        "created": "2018-04-23T16:40:50.847Z",
        "modified": "2018-04-23T16:40:50.847Z",
        "name": "Teddy Bear",
        "goals": [
            "compromising environment NGOs",
            "water-hole attacks geared towards energy sector"
        ],
        "sophistication": "innovator",
        "resource_level": "government",
        "labels": [
            "nation-state"
        ]
    }

    tc_sink.add(ta)

def test_add_dict_bundle_object(collection):
    tc_sink = TAXIICollectionSink(collection)

    ta = {
        "type": "bundle",
        "id": "bundle--860ccc8d-56c9-4fda-9384-84276fb52fb1",
        "spec_version": "2.0",
        "objects": [
            {
                "type": "threat-actor",
                "id": "threat-actor--dc5a2f41-f76e-425a-81fe-33afc7aabd75",
                "created": "2018-04-23T18:45:11.390Z",
                "modified": "2018-04-23T18:45:11.390Z",
                "name": "Teddy Bear",
                "goals": [
                    "compromising environment NGOs",
                    "water-hole attacks geared towards energy sector"
                ],
                "sophistication": "innovator",
                "resource_level": "government",
                "labels": [
                    "nation-state"
                ]
            }
        ]
    }

    tc_sink.add(ta)

def test_get_stix2_object(collection):
    tc_sink = TAXIICollectionSource(collection)

    objects = tc_sink.get("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")

    assert objects

def test_parse_taxii_filters(collection):
    query = [
        Filter("added_after", "=", "2016-02-01T00:00:01.000Z"),
        Filter("id", "=", "taxii stix object ID"),
        Filter("type", "=", "taxii stix object ID"),
        Filter("version", "=", "first"),
        Filter("created_by_ref", "=", "Bane"),
    ]

    taxii_filters_expected = [
        Filter("added_after", "=", "2016-02-01T00:00:01.000Z"),
        Filter("id", "=", "taxii stix object ID"),
        Filter("type", "=", "taxii stix object ID"),
        Filter("version", "=", "first")
    ]

    ds = TAXIICollectionSource(collection)

    taxii_filters = ds._parse_taxii_filters(query)

    assert taxii_filters == taxii_filters_expected

def test_add_get_remove_filter(collection):
    ds = TAXIICollectionSource(collection)

    # First 3 filters are valid, remaining properties are erroneous in some way
    valid_filters = [
        Filter('type', '=', 'malware'),
        Filter('id', '!=', 'stix object id'),
        Filter('labels', 'in', ["heartbleed", "malicious-activity"]),
    ]

    assert len(ds.filters) == 0

    ds.filters.add(valid_filters[0])
    assert len(ds.filters) == 1

    # Addin the same filter again will have no effect since `filters` acts
    # like a set
    ds.filters.add(valid_filters[0])
    assert len(ds.filters) == 1

    ds.filters.add(valid_filters[1])
    assert len(ds.filters) == 2

    ds.filters.add(valid_filters[2])
    assert len(ds.filters) == 3

    assert valid_filters == [f for f in ds.filters]

    # remove
    ds.filters.remove(valid_filters[0])

    assert len(ds.filters) == 2

    ds.filters.add(valid_filters)

def test_get_all_versions(collection):
    ds = TAXIICollectionStore(collection)

    indicators = ds.all_versions('indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f')
    # There are 3 indicators but 2 share the same 'modified' timestamp
    assert len(indicators) == 2

def test_can_read_error(collection_no_rw_access):
    """create a TAXIICOllectionSource with a taxii2client.Collection
    instance that does not have read access, check ValueError exception is raised"""
    with pytest.raises(DataSourceError) as excinfo:
        TAXIICollectionSource(collection_no_rw_access)
    assert "Collection object provided does not have read access" in str(excinfo.value)

def test_can_write_error(collection_no_rw_access):
    """create a TAXIICOllectionSink with a taxii2client.Collection
    instance that does not have write access, check ValueError exception is raised"""
    with pytest.raises(DataSourceError) as excinfo:
        TAXIICollectionSink(collection_no_rw_access)
    assert "Collection object provided does not have write access" in str(excinfo.value)

def test_get_404():
    """a TAXIICollectionSource.get() call that receives an HTTP 404 response
    code from the taxii2client should be be returned as None.

    TAXII spec states that a TAXII server can return a 404 for
    nonexistent resources or lack of access. Decided that None is acceptable
    reponse to imply that state of the TAXII endpoint.
    """

    class TAXIICollection404():
        can_read = True

        def get_object(self, id, version=None):
            resp = Response()
            resp.status_code = 404
            resp.raise_for_status()

    ds = TAXIICollectionSource(TAXIICollection404())

    # this will raise 404 from mock TAXII Client but TAXIICollectionStore
    # should handle gracefully and return None
    stix_obj = ds.get("indicator--1")  # this will raise 404 from
    assert stix_obj is None

def test_all_versions_404(collection):
    """ a TAXIICollectionSource.all_version() call that recieves an HTTP 404
    response code from the taxii2client should be returned as an exception"""
    ds = TAXIICollectionStore(collection)
    with pytest.raises(DataSourceError) as excinfo:
        ds.all_versions("indicator--1")
    assert "are either not found or access is denied" in str(excinfo.value)
    assert "404" in str(excinfo.value)

def test_query_404(collection):
    """ a TAXIICollectionSource.query() call that recieves an HTTP 404
    response code from the taxii2client should be returned as an exception"""
    ds = TAXIICollectionStore(collection)
    query = [Filter("type", "=", "malware")]

    with pytest.raises(DataSourceError) as excinfo:
        ds.query(query=query)
    assert "are either not found or access is denied" in str(excinfo.value)
    assert "404" in str(excinfo.value)
<|endoftext|>