<commit_before># Use modern Python
from __future__ import absolute_import, print_function, unicode_literals

# Standard library imports
from collections import defaultdict
from decimal import Decimal
import logging
from optparse import make_option

# Django imports
from django.apps import apps as default_apps
from django.core.management.base import BaseCommand

from corehq.apps.accounting.exceptions import AccountingError
from corehq.apps.accounting.models import (
    SoftwareProductType, SoftwarePlanEdition, SoftwarePlanVisibility, FeatureType,
    UNLIMITED_FEATURE_USAGE,
)

logger = logging.getLogger(__name__)

BOOTSTRAP_EDITION_TO_ROLE = {
    SoftwarePlanEdition.COMMUNITY: 'community_plan_v0',
    SoftwarePlanEdition.STANDARD: 'standard_plan_v0',
    SoftwarePlanEdition.PRO: 'pro_plan_v0',
    SoftwarePlanEdition.ADVANCED: 'advanced_plan_v0',
    SoftwarePlanEdition.ENTERPRISE: 'enterprise_plan_v0',
}
EDITIONS = [
    SoftwarePlanEdition.COMMUNITY,
    SoftwarePlanEdition.STANDARD,
    SoftwarePlanEdition.PRO,
    SoftwarePlanEdition.ADVANCED,
    SoftwarePlanEdition.ENTERPRISE,
]
FEATURE_TYPES = [f[0] for f in FeatureType.CHOICES]
PRODUCT_TYPES = [p[0] for p in SoftwareProductType.CHOICES]

class Command(BaseCommand):
    help = 'Populate a fresh db with standard set of Software Plans.'

    option_list = BaseCommand.option_list + (
        make_option('--dry-run', action='store_true', default=False,
                    help='Do not actually modify the database, just verbosely log what happen'),
        make_option('--verbose', action='store_true', default=False,
                    help='Enable debug output'),
        make_option('--testing', action='store_true', default=False,
                    help='Run this command for testing purposes.'),
    )

    def handle(self, dry_run=False, verbose=False, testing=False, *args, **options):
        logger.info('Bootstrapping standard plans. Enterprise plans will have to be created via the admin UIs.')

        for_tests = testing
        if for_tests:
            logger.info("Initializing Plans and Roles for Testing")

        ensure_plans(dry_run=dry_run, verbose=verbose, for_tests=for_tests, apps=default_apps)

def ensure_plans(dry_run, verbose, for_tests, apps):
    DefaultProductPlan = apps.get_model('accounting', 'DefaultProductPlan')
    SoftwarePlan = apps.get_model('accounting', 'SoftwarePlan')
    SoftwarePlanVersion = apps.get_model('accounting', 'SoftwarePlanVersion')
    Role = apps.get_model('django_prbac', 'Role')

    edition_to_features = _ensure_features(dry_run=dry_run, verbose=verbose, apps=apps)
    for product_type in PRODUCT_TYPES:
        for edition in EDITIONS:
            role_slug = BOOTSTRAP_EDITION_TO_ROLE[edition]
            try:
                role = Role.objects.get(slug=role_slug)
            except Role.DoesNotExist:
                logger.info("Could not find the role '%s'. Did you forget to run cchq_prbac_bootstrap?")
                logger.info("Aborting. You should figure this out.")
                return
            software_plan_version = SoftwarePlanVersion(role=role)

            product, product_rates = _ensure_product_and_rate(product_type, edition, dry_run=dry_run, verbose=verbose, apps=apps)
            feature_rates = _ensure_feature_rates(edition_to_features[edition], edition, dry_run=dry_run, verbose=verbose, for_tests=for_tests, apps=apps)
            software_plan = SoftwarePlan(
                name='%s Edition' % product.name, edition=edition, visibility=SoftwarePlanVisibility.PUBLIC
            )
            if dry_run:
                logger.info("[DRY RUN] Creating Software Plan: %s" % software_plan.name)
            else:
                try:
                    software_plan = SoftwarePlan.objects.get(name=software_plan.name)
                    if verbose:
                        logger.info("Plan '%s' already exists. Using existing plan to add version."
                                    % software_plan.name)
                except SoftwarePlan.DoesNotExist:
                    software_plan.save()
                    if verbose:
                        logger.info("Creating Software Plan: %s" % software_plan.name)

                    software_plan_version.plan = software_plan

                    # must save before assigning many-to-many relationship
                    if hasattr(SoftwarePlanVersion, 'product_rates'):
                        software_plan_version.save()

                    for product_rate in product_rates:
                        product_rate.save()
                        if hasattr(SoftwarePlanVersion, 'product_rates'):
                            software_plan_version.product_rates.add(product_rate)
                        elif hasattr(SoftwarePlanVersion, 'product_rate'):
                            assert len(product_rates) == 1
                            software_plan_version.product_rate = product_rate
                        else:
                            raise AccountingError('SoftwarePlanVersion does not have product_rate or product_rates field')

                    # must save before assigning many-to-many relationship
                    if hasattr(SoftwarePlanVersion, 'product_rate'):
                        software_plan_version.save()

                    for feature_rate in feature_rates:
                        feature_rate.save()
                        software_plan_version.feature_rates.add(feature_rate)
                    software_plan_version.save()

            if edition == SoftwarePlanEdition.ADVANCED:
                trials = [True, False]
            else:
                trials = [False]
            for is_trial in trials:
                default_product_plan = DefaultProductPlan(product_type=product.product_type, edition=edition, is_trial=is_trial)
                if dry_run:
                    logger.info("[DRY RUN] Setting plan as default for product '%s' and edition '%s'." %
                            (product.product_type, default_product_plan.edition))
                else:
                    try:
                        default_product_plan = DefaultProductPlan.objects.get(product_type=product.product_type,
                                                                              edition=edition, is_trial=is_trial)
                        if verbose:
                            logger.info("Default for product '%s' and edition "
                                        "'%s' already exists." % (
                                            product.product_type, default_product_plan.edition
                                        ))
                    except DefaultProductPlan.DoesNotExist:
                        default_product_plan.plan = software_plan
                        default_product_plan.save()
                        if verbose:
                            logger.info("Setting plan as default for product '%s' and edition '%s'." %
                                        (product.product_type,
                                         default_product_plan.edition))

def _ensure_product_and_rate(product_type, edition, dry_run, verbose, apps):
    """
    Ensures that all the necessary SoftwareProducts and SoftwareProductRates are created for the plan.
    """
    SoftwareProduct = apps.get_model('accounting', 'SoftwareProduct')
    SoftwareProductRate = apps.get_model('accounting', 'SoftwareProductRate')

    if verbose:
        logger.info('Ensuring Products and Product Rates')

    product = SoftwareProduct(name='%s %s' % (product_type, edition), product_type=product_type)
    if edition == SoftwarePlanEdition.ENTERPRISE:
        product.name = "Dimagi Only %s" % product.name

    product_rates = []
    BOOTSTRAP_PRODUCT_RATES = {
        SoftwarePlanEdition.COMMUNITY: [
            SoftwareProductRate(),  # use all the defaults
        ],
        SoftwarePlanEdition.STANDARD: [
            SoftwareProductRate(monthly_fee=Decimal('100.00')),
        ],
        SoftwarePlanEdition.PRO: [
            SoftwareProductRate(monthly_fee=Decimal('500.00')),
        ],
        SoftwarePlanEdition.ADVANCED: [
            SoftwareProductRate(monthly_fee=Decimal('1000.00')),
        ],
        SoftwarePlanEdition.ENTERPRISE: [
            SoftwareProductRate(monthly_fee=Decimal('0.00')),
        ],
    }

    for product_rate in BOOTSTRAP_PRODUCT_RATES[edition]:
        if dry_run:
            logger.info("[DRY RUN] Creating Product: %s" % product)
            logger.info("[DRY RUN] Corresponding product rate of $%d created." % product_rate.monthly_fee)
        else:
            try:
                product = SoftwareProduct.objects.get(name=product.name)
                if verbose:
                    logger.info("Product '%s' already exists. Using "
                                "existing product to add rate."
                                % product.name)
            except SoftwareProduct.DoesNotExist:
                product.save()
                if verbose:
                    logger.info("Creating Product: %s" % product)
            if verbose:
                logger.info("Corresponding product rate of $%d created."
                            % product_rate.monthly_fee)
        product_rate.product = product
        product_rates.append(product_rate)
    return product, product_rates

def _ensure_features(dry_run, verbose, apps):
    """
    Ensures that all the Features necessary for the plans are created.
    """
    Feature = apps.get_model('accounting', 'Feature')

    if verbose:
        logger.info('Ensuring Features')

    edition_to_features = defaultdict(list)
    for edition in EDITIONS:
        for feature_type in FEATURE_TYPES:
            feature = Feature(name='%s %s' % (feature_type, edition), feature_type=feature_type)
            if edition == SoftwarePlanEdition.ENTERPRISE:
                feature.name = "Dimagi Only %s" % feature.name
            if dry_run:
                logger.info("[DRY RUN] Creating Feature: %s" % feature)
            else:
                try:
                    feature = Feature.objects.get(name=feature.name)
                    if verbose:
                        logger.info("Feature '%s' already exists. Using "
                                    "existing feature to add rate."
                                    % feature.name)
                except Feature.DoesNotExist:
                    feature.save()
                    if verbose:
                        logger.info("Creating Feature: %s" % feature)
            edition_to_features[edition].append(feature)
    return edition_to_features

def _ensure_feature_rates(features, edition, dry_run, verbose, for_tests, apps):
    """
    Ensures that all the FeatureRates necessary for the plans are created.
    """
    FeatureRate = apps.get_model('accounting', 'FeatureRate')

    if verbose:
        logger.info('Ensuring Feature Rates')

    feature_rates = []
    BOOTSTRAP_FEATURE_RATES = {
        SoftwarePlanEdition.COMMUNITY: {
            FeatureType.USER: FeatureRate(monthly_limit=2 if for_tests else 50,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=0),  # use defaults here
        },
        SoftwarePlanEdition.STANDARD: {
            FeatureType.USER: FeatureRate(monthly_limit=4 if for_tests else 100,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=3 if for_tests else 100),
        },
        SoftwarePlanEdition.PRO: {
            FeatureType.USER: FeatureRate(monthly_limit=6 if for_tests else 500,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=5 if for_tests else 500),
        },
        SoftwarePlanEdition.ADVANCED: {
            FeatureType.USER: FeatureRate(monthly_limit=8 if for_tests else 1000,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=7 if for_tests else 1000),
        },
        SoftwarePlanEdition.ENTERPRISE: {
            FeatureType.USER: FeatureRate(monthly_limit=UNLIMITED_FEATURE_USAGE, per_excess_fee=Decimal('0.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=UNLIMITED_FEATURE_USAGE),
        },
    }
    for feature in features:
        feature_rate = BOOTSTRAP_FEATURE_RATES[edition][feature.feature_type]
        feature_rate.feature = feature
        if dry_run:
            logger.info("[DRY RUN] Creating rate for feature '%s': %s" % (feature.name, feature_rate))
        elif verbose:
            logger.info("Creating rate for feature '%s': %s" % (feature.name, feature_rate))
        feature_rates.append(feature_rate)
    return feature_rates
<commit_msg>fix logger text<commit_after># Use modern Python
from __future__ import absolute_import, print_function, unicode_literals

# Standard library imports
from collections import defaultdict
from decimal import Decimal
import logging
from optparse import make_option

# Django imports
from django.apps import apps as default_apps
from django.core.management.base import BaseCommand

from corehq.apps.accounting.exceptions import AccountingError
from corehq.apps.accounting.models import (
    SoftwareProductType, SoftwarePlanEdition, SoftwarePlanVisibility, FeatureType,
    UNLIMITED_FEATURE_USAGE,
)

logger = logging.getLogger(__name__)

BOOTSTRAP_EDITION_TO_ROLE = {
    SoftwarePlanEdition.COMMUNITY: 'community_plan_v0',
    SoftwarePlanEdition.STANDARD: 'standard_plan_v0',
    SoftwarePlanEdition.PRO: 'pro_plan_v0',
    SoftwarePlanEdition.ADVANCED: 'advanced_plan_v0',
    SoftwarePlanEdition.ENTERPRISE: 'enterprise_plan_v0',
}
EDITIONS = [
    SoftwarePlanEdition.COMMUNITY,
    SoftwarePlanEdition.STANDARD,
    SoftwarePlanEdition.PRO,
    SoftwarePlanEdition.ADVANCED,
    SoftwarePlanEdition.ENTERPRISE,
]
FEATURE_TYPES = [f[0] for f in FeatureType.CHOICES]
PRODUCT_TYPES = [p[0] for p in SoftwareProductType.CHOICES]

class Command(BaseCommand):
    help = 'Populate a fresh db with standard set of Software Plans.'

    option_list = BaseCommand.option_list + (
        make_option('--dry-run', action='store_true', default=False,
                    help='Do not actually modify the database, just verbosely log what happen'),
        make_option('--verbose', action='store_true', default=False,
                    help='Enable debug output'),
        make_option('--testing', action='store_true', default=False,
                    help='Run this command for testing purposes.'),
    )

    def handle(self, dry_run=False, verbose=False, testing=False, *args, **options):
        logger.info('Bootstrapping standard plans. Custom plans will have to be created via the admin UIs.')

        for_tests = testing
        if for_tests:
            logger.info("Initializing Plans and Roles for Testing")

        ensure_plans(dry_run=dry_run, verbose=verbose, for_tests=for_tests, apps=default_apps)

def ensure_plans(dry_run, verbose, for_tests, apps):
    DefaultProductPlan = apps.get_model('accounting', 'DefaultProductPlan')
    SoftwarePlan = apps.get_model('accounting', 'SoftwarePlan')
    SoftwarePlanVersion = apps.get_model('accounting', 'SoftwarePlanVersion')
    Role = apps.get_model('django_prbac', 'Role')

    edition_to_features = _ensure_features(dry_run=dry_run, verbose=verbose, apps=apps)
    for product_type in PRODUCT_TYPES:
        for edition in EDITIONS:
            role_slug = BOOTSTRAP_EDITION_TO_ROLE[edition]
            try:
                role = Role.objects.get(slug=role_slug)
            except Role.DoesNotExist:
                logger.info("Could not find the role '%s'. Did you forget to run cchq_prbac_bootstrap?")
                logger.info("Aborting. You should figure this out.")
                return
            software_plan_version = SoftwarePlanVersion(role=role)

            product, product_rates = _ensure_product_and_rate(product_type, edition, dry_run=dry_run, verbose=verbose, apps=apps)
            feature_rates = _ensure_feature_rates(edition_to_features[edition], edition, dry_run=dry_run, verbose=verbose, for_tests=for_tests, apps=apps)
            software_plan = SoftwarePlan(
                name='%s Edition' % product.name, edition=edition, visibility=SoftwarePlanVisibility.PUBLIC
            )
            if dry_run:
                logger.info("[DRY RUN] Creating Software Plan: %s" % software_plan.name)
            else:
                try:
                    software_plan = SoftwarePlan.objects.get(name=software_plan.name)
                    if verbose:
                        logger.info("Plan '%s' already exists. Using existing plan to add version."
                                    % software_plan.name)
                except SoftwarePlan.DoesNotExist:
                    software_plan.save()
                    if verbose:
                        logger.info("Creating Software Plan: %s" % software_plan.name)

                    software_plan_version.plan = software_plan

                    # must save before assigning many-to-many relationship
                    if hasattr(SoftwarePlanVersion, 'product_rates'):
                        software_plan_version.save()

                    for product_rate in product_rates:
                        product_rate.save()
                        if hasattr(SoftwarePlanVersion, 'product_rates'):
                            software_plan_version.product_rates.add(product_rate)
                        elif hasattr(SoftwarePlanVersion, 'product_rate'):
                            assert len(product_rates) == 1
                            software_plan_version.product_rate = product_rate
                        else:
                            raise AccountingError('SoftwarePlanVersion does not have product_rate or product_rates field')

                    # must save before assigning many-to-many relationship
                    if hasattr(SoftwarePlanVersion, 'product_rate'):
                        software_plan_version.save()

                    for feature_rate in feature_rates:
                        feature_rate.save()
                        software_plan_version.feature_rates.add(feature_rate)
                    software_plan_version.save()

            if edition == SoftwarePlanEdition.ADVANCED:
                trials = [True, False]
            else:
                trials = [False]
            for is_trial in trials:
                default_product_plan = DefaultProductPlan(product_type=product.product_type, edition=edition, is_trial=is_trial)
                if dry_run:
                    logger.info("[DRY RUN] Setting plan as default for product '%s' and edition '%s'." %
                            (product.product_type, default_product_plan.edition))
                else:
                    try:
                        default_product_plan = DefaultProductPlan.objects.get(product_type=product.product_type,
                                                                              edition=edition, is_trial=is_trial)
                        if verbose:
                            logger.info("Default for product '%s' and edition "
                                        "'%s' already exists." % (
                                            product.product_type, default_product_plan.edition
                                        ))
                    except DefaultProductPlan.DoesNotExist:
                        default_product_plan.plan = software_plan
                        default_product_plan.save()
                        if verbose:
                            logger.info("Setting plan as default for product '%s' and edition '%s'." %
                                        (product.product_type,
                                         default_product_plan.edition))

def _ensure_product_and_rate(product_type, edition, dry_run, verbose, apps):
    """
    Ensures that all the necessary SoftwareProducts and SoftwareProductRates are created for the plan.
    """
    SoftwareProduct = apps.get_model('accounting', 'SoftwareProduct')
    SoftwareProductRate = apps.get_model('accounting', 'SoftwareProductRate')

    if verbose:
        logger.info('Ensuring Products and Product Rates')

    product = SoftwareProduct(name='%s %s' % (product_type, edition), product_type=product_type)
    if edition == SoftwarePlanEdition.ENTERPRISE:
        product.name = "Dimagi Only %s" % product.name

    product_rates = []
    BOOTSTRAP_PRODUCT_RATES = {
        SoftwarePlanEdition.COMMUNITY: [
            SoftwareProductRate(),  # use all the defaults
        ],
        SoftwarePlanEdition.STANDARD: [
            SoftwareProductRate(monthly_fee=Decimal('100.00')),
        ],
        SoftwarePlanEdition.PRO: [
            SoftwareProductRate(monthly_fee=Decimal('500.00')),
        ],
        SoftwarePlanEdition.ADVANCED: [
            SoftwareProductRate(monthly_fee=Decimal('1000.00')),
        ],
        SoftwarePlanEdition.ENTERPRISE: [
            SoftwareProductRate(monthly_fee=Decimal('0.00')),
        ],
    }

    for product_rate in BOOTSTRAP_PRODUCT_RATES[edition]:
        if dry_run:
            logger.info("[DRY RUN] Creating Product: %s" % product)
            logger.info("[DRY RUN] Corresponding product rate of $%d created." % product_rate.monthly_fee)
        else:
            try:
                product = SoftwareProduct.objects.get(name=product.name)
                if verbose:
                    logger.info("Product '%s' already exists. Using "
                                "existing product to add rate."
                                % product.name)
            except SoftwareProduct.DoesNotExist:
                product.save()
                if verbose:
                    logger.info("Creating Product: %s" % product)
            if verbose:
                logger.info("Corresponding product rate of $%d created."
                            % product_rate.monthly_fee)
        product_rate.product = product
        product_rates.append(product_rate)
    return product, product_rates

def _ensure_features(dry_run, verbose, apps):
    """
    Ensures that all the Features necessary for the plans are created.
    """
    Feature = apps.get_model('accounting', 'Feature')

    if verbose:
        logger.info('Ensuring Features')

    edition_to_features = defaultdict(list)
    for edition in EDITIONS:
        for feature_type in FEATURE_TYPES:
            feature = Feature(name='%s %s' % (feature_type, edition), feature_type=feature_type)
            if edition == SoftwarePlanEdition.ENTERPRISE:
                feature.name = "Dimagi Only %s" % feature.name
            if dry_run:
                logger.info("[DRY RUN] Creating Feature: %s" % feature)
            else:
                try:
                    feature = Feature.objects.get(name=feature.name)
                    if verbose:
                        logger.info("Feature '%s' already exists. Using "
                                    "existing feature to add rate."
                                    % feature.name)
                except Feature.DoesNotExist:
                    feature.save()
                    if verbose:
                        logger.info("Creating Feature: %s" % feature)
            edition_to_features[edition].append(feature)
    return edition_to_features

def _ensure_feature_rates(features, edition, dry_run, verbose, for_tests, apps):
    """
    Ensures that all the FeatureRates necessary for the plans are created.
    """
    FeatureRate = apps.get_model('accounting', 'FeatureRate')

    if verbose:
        logger.info('Ensuring Feature Rates')

    feature_rates = []
    BOOTSTRAP_FEATURE_RATES = {
        SoftwarePlanEdition.COMMUNITY: {
            FeatureType.USER: FeatureRate(monthly_limit=2 if for_tests else 50,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=0),  # use defaults here
        },
        SoftwarePlanEdition.STANDARD: {
            FeatureType.USER: FeatureRate(monthly_limit=4 if for_tests else 100,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=3 if for_tests else 100),
        },
        SoftwarePlanEdition.PRO: {
            FeatureType.USER: FeatureRate(monthly_limit=6 if for_tests else 500,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=5 if for_tests else 500),
        },
        SoftwarePlanEdition.ADVANCED: {
            FeatureType.USER: FeatureRate(monthly_limit=8 if for_tests else 1000,
                                          per_excess_fee=Decimal('1.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=7 if for_tests else 1000),
        },
        SoftwarePlanEdition.ENTERPRISE: {
            FeatureType.USER: FeatureRate(monthly_limit=UNLIMITED_FEATURE_USAGE, per_excess_fee=Decimal('0.00')),
            FeatureType.SMS: FeatureRate(monthly_limit=UNLIMITED_FEATURE_USAGE),
        },
    }
    for feature in features:
        feature_rate = BOOTSTRAP_FEATURE_RATES[edition][feature.feature_type]
        feature_rate.feature = feature
        if dry_run:
            logger.info("[DRY RUN] Creating rate for feature '%s': %s" % (feature.name, feature_rate))
        elif verbose:
            logger.info("Creating rate for feature '%s': %s" % (feature.name, feature_rate))
        feature_rates.append(feature_rate)
    return feature_rates
<|endoftext|>


<commit_before>import re
from datetime import datetime
from itertools import chain

from couchdbkit import ResourceNotFound

from corehq.apps.app_manager.dbaccessors import get_app
from corehq.apps.app_manager.models import Application
from corehq.apps.app_manager.util import get_correct_app_class
from corehq.apps.es import FormES
from corehq.apps.es.filters import NOT, doc_type
from corehq.util.couch import IterDB
from corehq.util.quickcache import quickcache
from couchforms.const import ATTACHMENT_NAME
from couchforms.models import XFormInstance
from dimagi.utils.couch.database import iter_docs
from django.core.management.base import BaseCommand
from optparse import make_option

ONE_HOUR = 60 * 60
INFO = "info"
ERROR = "error"
WARNING = "warning"

SET_XMLNS = "set-xmlns-on-xfrom"
ERROR_SAVING = "error-saving-xform"
MULTI_MATCH = "multiple-possible-forms"

class Command(BaseCommand):
    help = ("Running this command will fix xform submissions with 'undefined' xmlns."
     " It will only fix xforms that are submitted against builds that have"
     " already been repaired.")

    args = '<log_path>'

    option_list = BaseCommand.option_list + (
        make_option(
            '--dry-run',
            action='store_true',
            dest='dry_run',
            default=False,
            help="Don't do the actual modifications, but still log what would be affected"
        ),
    )

    def handle(self, *args, **options):

        dry_run = options.get("dry_run", True)
        log_path = args[0].strip()

        with open(log_path, "w") as log_file:
            self.fix_xforms(log_file, dry_run)

    @staticmethod
    def fix_xforms(log_file, dry_run):
        unfixable_builds = set()
        total, submissions = get_submissions_without_xmlns()
        xform_db = IterDB(XFormInstance.get_db())
        with xform_db as xform_db:
            for i, xform_instance in enumerate(submissions):
                Command._print_progress(i, total)
                try:
                    xmlns = get_correct_xmlns(xform_instance)
                except MultiplePreviouslyFixedForms as e:
                    if xform_instance.build_id not in unfixable_builds:
                        unfixable_builds.add(xform_instance.build_id)
                        print e.message
                    _log(log_file, WARNING, MULTI_MATCH, xform_instance)
                    continue

                if xmlns:
                    set_xmlns_on_submission(
                        xform_instance,
                        xmlns,
                        xform_db,
                        log_file,
                        dry_run,
                    )

        for error_id in xform_db.error_ids:
            _log(ERROR, ERROR_SAVING, xform_id=error_id)

    @staticmethod
    def _print_progress(i, total_submissions):
        if i % 200 == 0 and i != 0:
            print "Progress: {} of {} ({})  {}".format(
                i, total_submissions, round(i / float(total_submissions), 2), datetime.now()
            )

def _log(stream, level, event, xform=None, xform_id=None):
    basic_template = "[{level}] {event}, xform_id={xform_id}"
    full_template = basic_template + ", domain={xform.domain}, username={xform.metadata.username}, " \
                                     "app_id={xform.app_id}, build_id={xform.build_id}, " \
                                     "xmlns={xform.xmlns}"
    if xform_id:
        msg = basic_template.format(level=level, event=event, xform_id=xform_id)
    else:
        msg = full_template.format(level=level, event=event, xform_id=xform._id, xform=xform)
    stream.write(msg + "\n")

def parse_log_message(line):
    match = re.match('^\[(.*)\] ([^,]*), (.*)', line)
    level, event, extras = match.groups()
    extras_dict = {}
    for pair in extras.split(", "):
        key, val = pair.split("=", 1)
        extras_dict[key] = val
    return level, event, extras_dict

def get_submissions_without_xmlns():
    submissions = XFormInstance.get_db().view(
        'couchforms/by_xmlns',
        key="undefined",
        include_docs=False,
        reduce=False,
    ).all()
    total_submissions = len(submissions)
    submission_id_generator = (s['id'] for s in submissions)
    submissions_doc_generator = (
        XFormInstance.wrap(i)
        for i in iter_docs(XFormInstance.get_db(), submission_id_generator)
    )

    total_error_submissions, error_submissions_doc_generator = _get_error_submissions_without_xmlns()
    return (
        total_submissions + total_error_submissions,
        chain(submissions_doc_generator, error_submissions_doc_generator)
    )

def _get_error_submissions_without_xmlns():

    query = (FormES()
             .xmlns('undefined')
             .remove_default_filter("is_xform_instance")
             .filter(NOT(doc_type('xforminstance')))
             .source(['_id']))
    result = query.run()
    total_error_submissions = result.total
    error_submissions = (
        XFormInstance.wrap(i)
        for i in iter_docs(XFormInstance.get_db(), (x['_id'] for x in result.hits))
    )
    return total_error_submissions, error_submissions

def set_xmlns_on_submission(xform_instance, xmlns, xform_db, log_file, dry_run):
    """
    Set the xmlns on an XFormInstance, and the save the document.
    """
    old_xml = xform_instance.get_xml()
    assert old_xml.count('xmlns="undefined"') == 1
    new_xml = old_xml.replace('xmlns="undefined"', 'xmlns="{}"'.format(xmlns))
    if not dry_run:
        replace_xml(xform_instance, new_xml)

    xform_instance.xmlns = xmlns
    xform_instance.form['@xmlns'] = xmlns
    xform_instance.form_migrated_from_undefined_xmlns = datetime.utcnow()
    if not dry_run:
        xform_db.save(xform_instance)
    _log(log_file, INFO, SET_XMLNS, xform_instance)

def get_forms_without_xmlns(app):
    return [form for form in app.get_forms() if form.xmlns == "undefined"]

def get_previously_fixed_forms(app):
    ret = []
    for form in app.get_forms():
        try:
            if form.form_migrated_from_undefined_xmlns:
                ret.append(form)
        except AttributeError:
            pass
    return ret

def replace_xml(xform, new_xml):
    if (
        xform._attachments and
        ATTACHMENT_NAME in xform._attachments and
        'data' in xform._attachments[ATTACHMENT_NAME]
    ):
        raise Exception("Unexpected attachment format: _attachments")

    else:
        try:
            xform.put_attachment(new_xml, name=ATTACHMENT_NAME, content_type='text/xml')
        except ResourceNotFound:
            raise Exception("Unexpected attachment format: old attachment scheme")

def get_saved_apps(app):
    saved_apps = Application.get_db().view(
        'app_manager/saved_app',
        startkey=[app.domain, app._id],
        endkey=[app.domain, app._id, {}],
        include_docs=True,
    )
    return [get_correct_app_class(row['doc']).wrap(row['doc']) for row in saved_apps]

class MultiplePreviouslyFixedForms(Exception):
    def __init__(self, build_id, app_id):
        template = "Unable to determine matching form. Multiple forms in app " \
                   "{} were previously repaired. Build submitted against was {}"
        super(MultiplePreviouslyFixedForms, self).__init__(template.format(app_id, build_id))

@quickcache(["xform_instance.build_id"], memoize_timeout=ONE_HOUR)
def get_correct_xmlns(xform_instance):
    if xform_instance.build_id is None:
        return None
    build = get_app(xform_instance.domain, xform_instance.build_id)
    # TODO: What if the app has been deleted?

    previously_fixed_forms_in_build = get_previously_fixed_forms(build)
    if len(previously_fixed_forms_in_build) == 1:
        return previously_fixed_forms_in_build[0].xmlns
    else:
        raise MultiplePreviouslyFixedForms(xform_instance.build_id, xform_instance.app_id)
<commit_msg>Add dry run check<commit_after>import re
from datetime import datetime
from itertools import chain

from couchdbkit import ResourceNotFound

from corehq.apps.app_manager.dbaccessors import get_app
from corehq.apps.app_manager.models import Application
from corehq.apps.app_manager.util import get_correct_app_class
from corehq.apps.es import FormES
from corehq.apps.es.filters import NOT, doc_type
from corehq.util.couch import IterDB
from corehq.util.quickcache import quickcache
from couchforms.const import ATTACHMENT_NAME
from couchforms.models import XFormInstance
from dimagi.utils.couch.database import iter_docs
from django.core.management.base import BaseCommand
from optparse import make_option

ONE_HOUR = 60 * 60
INFO = "info"
ERROR = "error"
WARNING = "warning"

SET_XMLNS = "set-xmlns-on-xfrom"
ERROR_SAVING = "error-saving-xform"
MULTI_MATCH = "multiple-possible-forms"

class Command(BaseCommand):
    help = ("Running this command will fix xform submissions with 'undefined' xmlns."
     " It will only fix xforms that are submitted against builds that have"
     " already been repaired.")

    args = '<log_path>'

    option_list = BaseCommand.option_list + (
        make_option(
            '--dry-run',
            action='store_true',
            dest='dry_run',
            default=False,
            help="Don't do the actual modifications, but still log what would be affected"
        ),
    )

    def handle(self, *args, **options):

        dry_run = options.get("dry_run", True)
        log_path = args[0].strip()

        if not options.get('no_input', False) and not dry_run:
            confirm = raw_input(
                u"""
                Are you sure you want to fix XFormInstances with "undefined" xmlns?
                This is NOT a dry run. y/N?
                """
            )
            if confirm != "y":
                print "\n\t\tSwap duplicates cancelled."
                return

        with open(log_path, "w") as log_file:
            self.fix_xforms(log_file, dry_run)

    @staticmethod
    def fix_xforms(log_file, dry_run):
        unfixable_builds = set()
        total, submissions = get_submissions_without_xmlns()
        xform_db = IterDB(XFormInstance.get_db())
        with xform_db as xform_db:
            for i, xform_instance in enumerate(submissions):
                Command._print_progress(i, total)
                try:
                    xmlns = get_correct_xmlns(xform_instance)
                except MultiplePreviouslyFixedForms as e:
                    if xform_instance.build_id not in unfixable_builds:
                        unfixable_builds.add(xform_instance.build_id)
                        print e.message
                    _log(log_file, WARNING, MULTI_MATCH, xform_instance)
                    continue

                if xmlns:
                    set_xmlns_on_submission(
                        xform_instance,
                        xmlns,
                        xform_db,
                        log_file,
                        dry_run,
                    )

        for error_id in xform_db.error_ids:
            _log(ERROR, ERROR_SAVING, xform_id=error_id)

    @staticmethod
    def _print_progress(i, total_submissions):
        if i % 200 == 0 and i != 0:
            print "Progress: {} of {} ({})  {}".format(
                i, total_submissions, round(i / float(total_submissions), 2), datetime.now()
            )

def _log(stream, level, event, xform=None, xform_id=None):
    basic_template = "[{level}] {event}, xform_id={xform_id}"
    full_template = basic_template + ", domain={xform.domain}, username={xform.metadata.username}, " \
                                     "app_id={xform.app_id}, build_id={xform.build_id}, " \
                                     "xmlns={xform.xmlns}"
    if xform_id:
        msg = basic_template.format(level=level, event=event, xform_id=xform_id)
    else:
        msg = full_template.format(level=level, event=event, xform_id=xform._id, xform=xform)
    stream.write(msg + "\n")

def parse_log_message(line):
    match = re.match('^\[(.*)\] ([^,]*), (.*)', line)
    level, event, extras = match.groups()
    extras_dict = {}
    for pair in extras.split(", "):
        key, val = pair.split("=", 1)
        extras_dict[key] = val
    return level, event, extras_dict

def get_submissions_without_xmlns():
    submissions = XFormInstance.get_db().view(
        'couchforms/by_xmlns',
        key="undefined",
        include_docs=False,
        reduce=False,
    ).all()
    total_submissions = len(submissions)
    submission_id_generator = (s['id'] for s in submissions)
    submissions_doc_generator = (
        XFormInstance.wrap(i)
        for i in iter_docs(XFormInstance.get_db(), submission_id_generator)
    )

    total_error_submissions, error_submissions_doc_generator = _get_error_submissions_without_xmlns()
    return (
        total_submissions + total_error_submissions,
        chain(submissions_doc_generator, error_submissions_doc_generator)
    )

def _get_error_submissions_without_xmlns():

    query = (FormES()
             .xmlns('undefined')
             .remove_default_filter("is_xform_instance")
             .filter(NOT(doc_type('xforminstance')))
             .source(['_id']))
    result = query.run()
    total_error_submissions = result.total
    error_submissions = (
        XFormInstance.wrap(i)
        for i in iter_docs(XFormInstance.get_db(), (x['_id'] for x in result.hits))
    )
    return total_error_submissions, error_submissions

def set_xmlns_on_submission(xform_instance, xmlns, xform_db, log_file, dry_run):
    """
    Set the xmlns on an XFormInstance, and the save the document.
    """
    old_xml = xform_instance.get_xml()
    assert old_xml.count('xmlns="undefined"') == 1
    new_xml = old_xml.replace('xmlns="undefined"', 'xmlns="{}"'.format(xmlns))
    if not dry_run:
        replace_xml(xform_instance, new_xml)

    xform_instance.xmlns = xmlns
    xform_instance.form['@xmlns'] = xmlns
    xform_instance.form_migrated_from_undefined_xmlns = datetime.utcnow()
    if not dry_run:
        xform_db.save(xform_instance)
    _log(log_file, INFO, SET_XMLNS, xform_instance)

def get_forms_without_xmlns(app):
    return [form for form in app.get_forms() if form.xmlns == "undefined"]

def get_previously_fixed_forms(app):
    ret = []
    for form in app.get_forms():
        try:
            if form.form_migrated_from_undefined_xmlns:
                ret.append(form)
        except AttributeError:
            pass
    return ret

def replace_xml(xform, new_xml):
    if (
        xform._attachments and
        ATTACHMENT_NAME in xform._attachments and
        'data' in xform._attachments[ATTACHMENT_NAME]
    ):
        raise Exception("Unexpected attachment format: _attachments")

    else:
        try:
            xform.put_attachment(new_xml, name=ATTACHMENT_NAME, content_type='text/xml')
        except ResourceNotFound:
            raise Exception("Unexpected attachment format: old attachment scheme")

def get_saved_apps(app):
    saved_apps = Application.get_db().view(
        'app_manager/saved_app',
        startkey=[app.domain, app._id],
        endkey=[app.domain, app._id, {}],
        include_docs=True,
    )
    return [get_correct_app_class(row['doc']).wrap(row['doc']) for row in saved_apps]

class MultiplePreviouslyFixedForms(Exception):
    def __init__(self, build_id, app_id):
        template = "Unable to determine matching form. Multiple forms in app " \
                   "{} were previously repaired. Build submitted against was {}"
        super(MultiplePreviouslyFixedForms, self).__init__(template.format(app_id, build_id))

@quickcache(["xform_instance.build_id"], memoize_timeout=ONE_HOUR)
def get_correct_xmlns(xform_instance):
    if xform_instance.build_id is None:
        return None
    build = get_app(xform_instance.domain, xform_instance.build_id)
    # TODO: What if the app has been deleted?

    previously_fixed_forms_in_build = get_previously_fixed_forms(build)
    if len(previously_fixed_forms_in_build) == 1:
        return previously_fixed_forms_in_build[0].xmlns
    else:
        raise MultiplePreviouslyFixedForms(xform_instance.build_id, xform_instance.app_id)
<|endoftext|>


<commit_before>#  ========= CONFIDENTIAL =========
#
#  Copyright (C) 2010-2014 Dell, Inc. - ALL RIGHTS RESERVED
#
#  ======================================================================
#   NOTICE: All information contained herein is, and remains the property
#   of Dell, Inc. The intellectual and technical concepts contained herein
#   are proprietary to Dell, Inc. and may be covered by U.S. and Foreign
#   Patents, patents in process, and are protected by trade secret or
#   copyright law. Dissemination of this information or reproduction of
#   this material is strictly forbidden unless prior written permission
#   is obtained from Dell, Inc.
#  ======================================================================
import logging
import os

from dcm.agent.cloudmetadata import CLOUD_TYPES
import dcm.agent.config as config
import dcm.agent.exceptions as exceptions
import dcm.agent.jobs.direct_pass as direct_pass
import dcm.agent.utils as utils

_cloud_stack_map = {
    "1": "xvdb",
    "2": "xvdc",
    "4": "xvde",
    "5": "xvdf",
    "6": "xvdg",
    "7": "xvdh",
    "8": "xvdi",
    "9": "xvdj"
}

_support_matrix = {
    CLOUD_TYPES.Amazon.lower(): [config.PLATFORM_TYPES.PLATFORM_UBUNTU.lower()]
}

_g_logger = logging.getLogger(__name__)

_g_platform_dep_installer = {
    config.PLATFORM_TYPES.PLATFORM_UBUNTU: ["debInstall",
                                            "--no-install-recommends",
                                            "mdadm"],
}

_g_platform_dep_crypto_installer = {
    config.PLATFORM_TYPES.PLATFORM_UBUNTU: ["debInstall",
                                            "cryptsetup"],
}

def _is_supported(conf):
    try:
        supported_distros = _support_matrix[conf.cloud_type.lower()]
    except KeyError:
        return False
    return conf.platform_name.lower() in supported_distros

class MountVolume(direct_pass.DirectPass):

    protocol_arguments = {
        "formatVolume":
        ("A boolean indicating if the volume should be formatted.",
         True, bool, None),
        "fileSystem":
        ("The file system type to which the volume will be formatted",
         True, str, None),
        "raidLevel": ("The RAID configuration to use", True, str, "NONE"),
        "encryptedFsEncryptionKey":
        ("The encryption key for encrypted volumes", False,
         utils.base64type_binary_convertor, None),
        "mountPoint": ("The directory on which the volume will be mounted",
                       False, str, None),
        "devices": ("The list of devices that will be used.",
                    True, list, None)
    }

    def __init__(self, conf, job_id, items_map, name, arguments):
        super(MountVolume, self).__init__(
            conf, job_id, items_map, name, arguments)

    def _install_deps(self, dep_dict):
        if self.conf.platform_name in dep_dict:
            _g_logger.debug("Installing packaging deps")
            pkg_installer_cmd = \
                dep_dict[self.conf.platform_name]
            if pkg_installer_cmd:
                cmd_path = self.conf.get_script_location(pkg_installer_cmd[0])
                pkg_installer_cmd[0] = cmd_path
                (stdout, stderr, rc) = utils.run_command(
                    self.conf, pkg_installer_cmd)
                _g_logger.debug("Results of install: stdout: %s, stderr: "
                                "%s, rc %d" % (str(stdout), str(stderr), rc))

    def setup_encryption(self, device_id, encrypted_device_id, key_file_path):
        command = [self.conf.get_script_location("setupEncryption"),
                   device_id,
                   encrypted_device_id,
                   key_file_path]
        (stdout, stderr, rc) = utils.run_command(self.conf, command)
        if rc != 0:
            raise exceptions.AgentExecutableException(
                command, rc, stdout, stderr)
        return rc

    def write_key_file(self, block_device):
        if self.args.encryptedFsEncryptionKey is None:
            return None

        if block_device:
            key_file_dir = self.conf.storage_temppath
        else:
            key_file_dir = os.path.join(self.conf.storage_base_dir, "tmp")
        key_file_path = os.path.join(key_file_dir, "fskey.txt")

        with open(key_file_path, "w") as fptr:
            fptr.write(self.args.encryptedFsEncryptionKey)

        return key_file_path

    def format(self, device_id):
        return utils.agent_format(
            self.conf, device_id, self.args.fileSystem,
            self.args.mountPoint, self.args.encryptedFsEncryptionKey)

    def configure_raid(self, device_id):
        _g_logger.info("Configuring RAID on " + str(device_id))
        if self.args.formatVolume:
            exe = self.conf.get_script_location("configureRaid")
        else:
            exe = self.conf.get_script_location("assembleRaid")

        if self.args.raidLevel.upper() == "NONE":
            raise exceptions.AgentPluginBadParameterException(
                "mount_volume",
                "When using multiple volumes you must specify a RAID level")

        try:
            raid_level = str(int(self.args.raidLevel[4:]))
        except:
            raise exceptions.AgentPluginBadParameterException(
                "mount_volume",
                "Invalid RAID level")

        cmd = [exe, raid_level, device_id]

        for d in self.args.devices:
            cmd.append(d)

        _g_logger.debug("Running the raid configuration command %s" % str(cmd))
        (stdout, stderr, rc) = utils.run_command(self.conf, cmd)
        _g_logger.debug("configure raid results: %d stdout=%s stderr=%s" %
                        (rc, str(stdout), str(stderr)))

        if rc != 0:
            _g_logger.error(
                "Failed to run raid configuration: stdout=%s\nstderr=%s"
                % (str(stdout), str(stderr)))
            raise exceptions.AgentExecutableException(
                cmd, rc, stdout, stderr)

    def _normalize_device(self):
        if self.conf.cloud_type != CLOUD_TYPES.CloudStack and\
                self.conf.cloud_type != CLOUD_TYPES.CloudStack3:
            return self.args.devices[:]
        modified_device_list = []
        for target_device in self.args.devices:
            if target_device not in _cloud_stack_map:
                raise exceptions.AgentPluginBadParameterException(
                    "mount_volume",
                    "When using cloud stack the device must be one of: %s" %
                    str(list(_cloud_stack_map.keys())))
            modified_device_list.append(_cloud_stack_map[target_device])
        return modified_device_list

    def mount_block_volume(self):
        if not self.args.devices:
            return 0

        if len(self.args.devices) > 1 and\
                self.args.raidLevel.upper() == "NONE":
            raise exceptions.AgentJobException(
                "Must specify a RAID volume with mounting multiple devices at "
                "once.")

        if len(self.args.devices) > 1:
            target_device = "md0"
        else:
            target_device = self.args.devices[0]

        _g_logger.debug("target device is " + target_device)

        td = target_device
        if self.args.encryptedFsEncryptionKey is not None:
            encrypted_device = "es" + target_device
            td = encrypted_device
            self._install_deps(_g_platform_dep_crypto_installer)

        device_mappings = utils.get_device_mappings(self.conf)
        for mapping in device_mappings:
            if mapping["device_id"] == td:
                return 0

        if len(self.args.devices) > 1:
            self.configure_raid(target_device)

        if self.args.formatVolume:
            if self.args.encryptedFsEncryptionKey:
                key_file_path = None
                try:
                    key_file_path = self.write_key_file(True)
                    self.setup_encryption(
                        target_device, encrypted_device, key_file_path)
                    utils.open_encrypted_device(self.conf,
                                                target_device,
                                                encrypted_device,
                                                key_file_path)
                    target_device = encrypted_device
                finally:
                    utils.safe_delete(key_file_path)
            self.format(target_device)
        elif self.args.encryptedFsEncryptionKey:
            key_file_path = None
            try:
                key_file_path = self.write_key_file(True)
                utils.open_encrypted_device(self.conf,
                                            target_device,
                                            encrypted_device,
                                            key_file_path)
                target_device = encrypted_device
            finally:
                utils.safe_delete(key_file_path)
        utils.mount(self.conf,
                    target_device,
                    self.args.fileSystem,
                    self.args.mountPoint)
        return 0

    def run(self):
        if not _is_supported(self.conf):
            raise exceptions.AgentUnsupportedCloudFeature(
                "mount is not supported on the distro " +
                self.conf.cloud_type.lower() + " and cloud " +
                self.conf.platform_name.lower())

        if self.args.mountPoint is None:
            self.args.mountPoint = self.conf.storage_mountpoint

        # if the mount point exists and is not empty we cannot use it
        if os.path.exists(self.args.mountPoint):
            if os.listdir(self.args.mountPoint):
                raise exceptions.AgentOptionException(
                    "The path %s exists.  For safety the agent only mounts "
                    "volumes on paths that do not exist or are empty.")

        if self.args.fileSystem is None:
            self.args.fileSystem = self.conf.storage_default_file_system

        self.args.devices = self._normalize_device()
        self._install_deps(_g_platform_dep_installer)
        rc = self.mount_block_volume()

        reply = {"return_code": rc, "message": "",
                 "error_message": "", "reply_type": "void"}
        return reply

def load_plugin(conf, job_id, items_map, name, arguments):
    return MountVolume(conf, job_id, items_map, name, arguments)

def get_features(conf):
    if _is_supported(conf):
        return {'mount': True,
                'format': True}
    return {}
<commit_msg>Python 3 encoding cleanup for mount.<commit_after>#  ========= CONFIDENTIAL =========
#
#  Copyright (C) 2010-2014 Dell, Inc. - ALL RIGHTS RESERVED
#
#  ======================================================================
#   NOTICE: All information contained herein is, and remains the property
#   of Dell, Inc. The intellectual and technical concepts contained herein
#   are proprietary to Dell, Inc. and may be covered by U.S. and Foreign
#   Patents, patents in process, and are protected by trade secret or
#   copyright law. Dissemination of this information or reproduction of
#   this material is strictly forbidden unless prior written permission
#   is obtained from Dell, Inc.
#  ======================================================================
import logging
import os

from dcm.agent.cloudmetadata import CLOUD_TYPES
import dcm.agent.config as config
import dcm.agent.exceptions as exceptions
import dcm.agent.jobs.direct_pass as direct_pass
import dcm.agent.utils as utils

_cloud_stack_map = {
    "1": "xvdb",
    "2": "xvdc",
    "4": "xvde",
    "5": "xvdf",
    "6": "xvdg",
    "7": "xvdh",
    "8": "xvdi",
    "9": "xvdj"
}

_support_matrix = {
    CLOUD_TYPES.Amazon.lower(): [config.PLATFORM_TYPES.PLATFORM_UBUNTU.lower()]
}

_g_logger = logging.getLogger(__name__)

_g_platform_dep_installer = {
    config.PLATFORM_TYPES.PLATFORM_UBUNTU: ["debInstall",
                                            "--no-install-recommends",
                                            "mdadm"],
}

_g_platform_dep_crypto_installer = {
    config.PLATFORM_TYPES.PLATFORM_UBUNTU: ["debInstall",
                                            "cryptsetup"],
}

def _is_supported(conf):
    try:
        supported_distros = _support_matrix[conf.cloud_type.lower()]
    except KeyError:
        return False
    return conf.platform_name.lower() in supported_distros

class MountVolume(direct_pass.DirectPass):

    protocol_arguments = {
        "formatVolume":
        ("A boolean indicating if the volume should be formatted.",
         True, bool, None),
        "fileSystem":
        ("The file system type to which the volume will be formatted",
         True, str, None),
        "raidLevel": ("The RAID configuration to use", True, str, "NONE"),
        "encryptedFsEncryptionKey":
        ("The encryption key for encrypted volumes", False,
         utils.base64type_binary_convertor, None),
        "mountPoint": ("The directory on which the volume will be mounted",
                       False, str, None),
        "devices": ("The list of devices that will be used.",
                    True, list, None)
    }

    def __init__(self, conf, job_id, items_map, name, arguments):
        super(MountVolume, self).__init__(
            conf, job_id, items_map, name, arguments)

    def _install_deps(self, dep_dict):
        if self.conf.platform_name in dep_dict:
            _g_logger.debug("Installing packaging deps")
            pkg_installer_cmd = \
                dep_dict[self.conf.platform_name]
            if pkg_installer_cmd:
                cmd_path = self.conf.get_script_location(pkg_installer_cmd[0])
                pkg_installer_cmd[0] = cmd_path
                (stdout, stderr, rc) = utils.run_command(
                    self.conf, pkg_installer_cmd)
                _g_logger.debug("Results of install: stdout: %s, stderr: "
                                "%s, rc %d" % (str(stdout), str(stderr), rc))

    def setup_encryption(self, device_id, encrypted_device_id, key_file_path):
        command = [self.conf.get_script_location("setupEncryption"),
                   device_id,
                   encrypted_device_id,
                   key_file_path]
        (stdout, stderr, rc) = utils.run_command(self.conf, command)
        if rc != 0:
            raise exceptions.AgentExecutableException(
                command, rc, stdout, stderr)
        return rc

    def write_key_file(self, block_device):
        if self.args.encryptedFsEncryptionKey is None:
            return None

        if block_device:
            key_file_dir = self.conf.storage_temppath
        else:
            key_file_dir = os.path.join(self.conf.storage_base_dir, "tmp")
        key_file_path = os.path.join(key_file_dir, "fskey.txt")

        with open(key_file_path, "wb") as fptr:
            fptr.write(self.args.encryptedFsEncryptionKey)

        return key_file_path

    def format(self, device_id):
        return utils.agent_format(
            self.conf, device_id, self.args.fileSystem,
            self.args.mountPoint, self.args.encryptedFsEncryptionKey)

    def configure_raid(self, device_id):
        _g_logger.info("Configuring RAID on " + str(device_id))
        if self.args.formatVolume:
            exe = self.conf.get_script_location("configureRaid")
        else:
            exe = self.conf.get_script_location("assembleRaid")

        if self.args.raidLevel.upper() == "NONE":
            raise exceptions.AgentPluginBadParameterException(
                "mount_volume",
                "When using multiple volumes you must specify a RAID level")

        try:
            raid_level = str(int(self.args.raidLevel[4:]))
        except:
            raise exceptions.AgentPluginBadParameterException(
                "mount_volume",
                "Invalid RAID level")

        cmd = [exe, raid_level, device_id]

        for d in self.args.devices:
            cmd.append(d)

        _g_logger.debug("Running the raid configuration command %s" % str(cmd))
        (stdout, stderr, rc) = utils.run_command(self.conf, cmd)
        _g_logger.debug("configure raid results: %d stdout=%s stderr=%s" %
                        (rc, str(stdout), str(stderr)))

        if rc != 0:
            _g_logger.error(
                "Failed to run raid configuration: stdout=%s\nstderr=%s"
                % (str(stdout), str(stderr)))
            raise exceptions.AgentExecutableException(
                cmd, rc, stdout, stderr)

    def _normalize_device(self):
        if self.conf.cloud_type != CLOUD_TYPES.CloudStack and\
                self.conf.cloud_type != CLOUD_TYPES.CloudStack3:
            return self.args.devices[:]
        modified_device_list = []
        for target_device in self.args.devices:
            if target_device not in _cloud_stack_map:
                raise exceptions.AgentPluginBadParameterException(
                    "mount_volume",
                    "When using cloud stack the device must be one of: %s" %
                    str(list(_cloud_stack_map.keys())))
            modified_device_list.append(_cloud_stack_map[target_device])
        return modified_device_list

    def mount_block_volume(self):
        if not self.args.devices:
            return 0

        if len(self.args.devices) > 1 and\
                self.args.raidLevel.upper() == "NONE":
            raise exceptions.AgentJobException(
                "Must specify a RAID volume with mounting multiple devices at "
                "once.")

        if len(self.args.devices) > 1:
            target_device = "md0"
        else:
            target_device = self.args.devices[0]

        _g_logger.debug("target device is " + target_device)

        td = target_device
        if self.args.encryptedFsEncryptionKey is not None:
            encrypted_device = "es" + target_device
            td = encrypted_device
            self._install_deps(_g_platform_dep_crypto_installer)

        device_mappings = utils.get_device_mappings(self.conf)
        for mapping in device_mappings:
            if mapping["device_id"] == td:
                return 0

        if len(self.args.devices) > 1:
            self.configure_raid(target_device)

        if self.args.formatVolume:
            if self.args.encryptedFsEncryptionKey:
                key_file_path = None
                try:
                    key_file_path = self.write_key_file(True)
                    self.setup_encryption(
                        target_device, encrypted_device, key_file_path)
                    utils.open_encrypted_device(self.conf,
                                                target_device,
                                                encrypted_device,
                                                key_file_path)
                    target_device = encrypted_device
                finally:
                    if key_file_path:
                        utils.safe_delete(key_file_path)
            self.format(target_device)
        elif self.args.encryptedFsEncryptionKey:
            key_file_path = None
            try:
                key_file_path = self.write_key_file(True)
                utils.open_encrypted_device(self.conf,
                                            target_device,
                                            encrypted_device,
                                            key_file_path)
                target_device = encrypted_device
            finally:
                utils.safe_delete(key_file_path)
        utils.mount(self.conf,
                    target_device,
                    self.args.fileSystem,
                    self.args.mountPoint)
        return 0

    def run(self):
        if not _is_supported(self.conf):
            raise exceptions.AgentUnsupportedCloudFeature(
                "mount is not supported on the distro " +
                self.conf.cloud_type.lower() + " and cloud " +
                self.conf.platform_name.lower())

        if self.args.mountPoint is None:
            self.args.mountPoint = self.conf.storage_mountpoint

        # if the mount point exists and is not empty we cannot use it
        if os.path.exists(self.args.mountPoint):
            if os.listdir(self.args.mountPoint):
                raise exceptions.AgentOptionException(
                    "The path %s exists.  For safety the agent only mounts "
                    "volumes on paths that do not exist or are empty.")

        if self.args.fileSystem is None:
            self.args.fileSystem = self.conf.storage_default_file_system

        self.args.devices = self._normalize_device()
        self._install_deps(_g_platform_dep_installer)
        rc = self.mount_block_volume()

        reply = {"return_code": rc, "message": "",
                 "error_message": "", "reply_type": "void"}
        return reply

def load_plugin(conf, job_id, items_map, name, arguments):
    return MountVolume(conf, job_id, items_map, name, arguments)

def get_features(conf):
    if _is_supported(conf):
        return {'mount': True,
                'format': True}
    return {}
<|endoftext|>


<commit_before>#
# Copyright (c) 2009, Georgia Tech Research Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the Georgia Tech Research Corporation nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY GEORGIA TECH RESEARCH CORPORATION ''AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL GEORGIA TECH BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# Author: Advait Jain

import numpy as np, math

import roslib; roslib.load_manifest('hrl_cody_arms')
import PyKDL as kdl

from equilibrium_point_control.hrl_arm import HRLArmKinematics
import create_IK_guess_dict as cgd

import hrl_lib.transforms as tr
import hrl_lib.util as ut
import hrl_lib.kdl_utils as ku

class CodyArmKinematics(HRLArmKinematics):
    # arm - 'r' or 'l'
    def __init__(self, arm):
        HRLArmKinematics.__init__(self, n_jts = 7)

        # create joint limit dicts
        if arm == 'r':
            max_lim = np.radians([ 120.00, 122.15, 77.5, 144., 122.,  45.,  45.])
            min_lim = np.radians([ -47.61,  -20., -77.5,   0., -80., -45., -45.])
        else:
            max_lim = np.radians([ 120.00,   20.,  77.5, 144.,   80.,  45.,  45.])
            min_lim = np.radians([ -47.61, -122.15, -77.5,   0., -122., -45., -45.])

        self.joint_lim_dict = {}
        self.joint_lim_dict['max'] = max_lim
        self.joint_lim_dict['min'] = min_lim

        wrist_stub_length = 0.0135 + 0.04318 # wrist linkange and FT sensor lengths
        self.setup_kdl_chains(arm, wrist_stub_length)

        if arm == 'r':
            pkl_nm = 'q_guess_right_dict.pkl'
        else:
            pkl_nm = 'q_guess_left_dict.pkl'

        pth = roslib.rospack.rospackexec(['find', 'hrl_cody_arms'])
        q_guess_pkl = pth + '/src/hrl_cody_arms/'+pkl_nm
        self.q_guess_dict = ut.load_pickle(q_guess_pkl)

    #--------------- KDL stuff ----------------
    # KDL joint array to meka joint list. (7->7)
    def kdl_angles_to_meka(self, q_jnt_arr):
        if q_jnt_arr == None:
            return None

        q_rad = [0. for i in range(7)]
        q_rad[0] = -q_jnt_arr[0]
        q_rad[1] = -q_jnt_arr[1]
        q_rad[2] = -q_jnt_arr[2]
        q_rad[3] = -q_jnt_arr[3]
        q_rad[4] = -q_jnt_arr[4]
        q_rad[5] = -q_jnt_arr[5]
        q_rad[6] = -q_jnt_arr[6]
        return q_rad

    # meka joint list to KDL joint array (7->7)
    def meka_angles_to_kdl(self, q_list):
        if q_list == None:
            return None

        n_joints = len(q_list)
        q = kdl.JntArray(n_joints)
        q[0] = -q_list[0]
        q[1] = -q_list[1]
        q[2] = -q_list[2]
        q[3] = -q_list[3]
        q[4] = -q_list[4]
        q[5] = -q_list[5]
        q[6] = -q_list[6]
        return q

    def create_right_chain(self, end_effector_length):
        ch = kdl.Chain()
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,-0.18493,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,-0.03175,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.00635,0.,-0.27795))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,-0.27853))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,0.,-end_effector_length))))
        return ch

    def create_left_chain(self, end_effector_length):
        ch = kdl.Chain()
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.18493,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,0.03175,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.00635,0.,-0.27795))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,-0.27853))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,0.,-end_effector_length))))
        return ch

    def create_solvers(self, ch):
         fk = kdl.ChainFkSolverPos_recursive(ch)
         ik_v = kdl.ChainIkSolverVel_pinv(ch)
         ik_p = kdl.ChainIkSolverPos_NR(ch, fk, ik_v)
         jac = kdl.ChainJntToJacSolver(ch)
         return fk, ik_v, ik_p, jac

    def setup_kdl_chains(self, arm, end_effector_length):
        if arm == 'r':
            #right arm
            ch = self.create_right_chain(end_effector_length)
            fk, ik_v, ik_p, jac = self.create_solvers(ch)
        else:
            #left arm
            ch = self.create_left_chain(end_effector_length)
            fk, ik_v, ik_p, jac = self.create_solvers(ch)

        kdl_chains = {}
        kdl_chains['chain'] = ch
        kdl_chains['nJnts'] = ch.getNrOfJoints()
        kdl_chains['fk_p'] = fk
        kdl_chains['ik_v'] = ik_v
        kdl_chains['ik_p'] = ik_p
        kdl_chains['jacobian_solver'] = jac

        #Add both chains to dictionary
        self.kdl_chains = kdl_chains

    def FK_kdl(self, q, link_number):
        fk_solver = self.kdl_chains['fk_p']
        endeffec_frame = kdl.Frame()
        kinematics_status = fk_solver.JntToCart(q, endeffec_frame,
                                                link_number)
        if kinematics_status >= 0:
#            print 'End effector transformation matrix:', endeffec_frame
            return endeffec_frame
        else:
            print 'Could not compute forward kinematics.'
            return None

    def IK_kdl(self,frame, q_init):
        nJnts = self.kdl_chains['nJnts']
        ik_solver = self.kdl_chains['ik_p']
        q = kdl.JntArray(nJnts)
        if ik_solver.CartToJnt(q_init,frame,q)>=0:
            for i in range(nJnts):
                q[i] = tr.angle_within_mod180(q[i])
            return q
        else:
            return None

    #-------------- implementation of HRLArmKinematics -----------

    def FK_vanilla(self, q, link_number = 7):
        q = self.meka_angles_to_kdl(q)
        frame = self.FK_kdl(q, link_number)
        pos = frame.p
        pos = ku.kdl_vec_to_np(pos)
        m = frame.M
        rot = ku.kdl_rot_to_np(m)
        return pos, rot

    ## compute Jacobian at point pos.
    # p is in the ground coord frame.
    def Jacobian(self, q, pos=None):
        if pos == None:
            pos = self.FK(q)[0]

        ch = self.kdl_chains['chain']
        v_list = []
        w_list = []

        for i in xrange(self.n_jts):
            p, rot = self.FK_vanilla(q, i)
            r = pos - p
            z_idx = ch.getSegment(i).getJoint().getType() - 1
            z = rot[:, z_idx]

            # this is a nasty trick. The rotation matrix returned by
            # FK_vanilla describes the orientation of a frame of the
            # KDL chain in the base frame. It just so happens that the
            # way Advait defined the KDL chain, the axis of rotation
            # in KDL is the -ve of the axis of rotation on the real
            # robot.
            # Advait apologizes for this.
            z = -z

            v_list.append(np.matrix(np.cross(z.A1, r.A1)).T)
            w_list.append(z)

        J = np.row_stack((np.column_stack(v_list), np.column_stack(w_list)))
        return J

    def IK_vanilla(self, p, rot, q_guess=None):
        p_kdl = ku.np_vec_to_kdl(p)
        rot_kdl = ku.np_rot_to_kdl(rot)
        fr = kdl.Frame(rot_kdl, p_kdl)

        if q_guess == None:
            q_guess = cgd.find_good_config(p, self.q_guess_dict)

        q_guess = self.meka_angles_to_kdl(q_guess)
        q_res = self.IK_kdl(arm, fr, q_guess)
        q_res = self.kdl_angles_to_meka(arm, q_res)
        return q_res

    #----------- extra functions -----------------

    ## clamp joint angles to their physical limits.
    # @param q - list of 7 joint angles.
    # The joint limits for IK are larger that the physical limits.
    def clamp_to_joint_limits(self, q, delta_list=[0.,0.,0.,0.,0.,0.,0.]):
        max_arr = self.joint_lim_dict['max']
        min_arr = self.joint_lim_dict['min']
        q_arr = np.array(q)
        d_arr = np.array(delta_list)
        return np.clip(q_arr, min_arr-d_arr, max_arr+d_arr)

    def within_joint_limits(self, q, delta_list=[0.,0.,0.,0.,0.,0.,0.]):
        max_arr = self.joint_lim_dict['max']
        min_arr = self.joint_lim_dict['min']
        q_arr = np.array(q)
        d_arr = np.array(delta_list)
        return np.all((q_arr <= max_arr+d_arr, q_arr >= min_arr-d_arr))

<commit_msg>M    cody_arm_kinematics.py bug in Jacobian computation producing better behavior on the robot. going to try to replicate this in simulation. this commit has bug free Jacobian (I think) but with debugging print statements.<commit_after>#
# Copyright (c) 2009, Georgia Tech Research Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the Georgia Tech Research Corporation nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY GEORGIA TECH RESEARCH CORPORATION ''AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL GEORGIA TECH BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# Author: Advait Jain

import numpy as np, math

import roslib; roslib.load_manifest('hrl_cody_arms')
import PyKDL as kdl

from equilibrium_point_control.hrl_arm import HRLArmKinematics
import create_IK_guess_dict as cgd

import hrl_lib.transforms as tr
import hrl_lib.util as ut
import hrl_lib.kdl_utils as ku

class CodyArmKinematics(HRLArmKinematics):
    # arm - 'r' or 'l'
    def __init__(self, arm):
        HRLArmKinematics.__init__(self, n_jts = 7)

        # create joint limit dicts
        if arm == 'r':
            max_lim = np.radians([ 120.00, 122.15, 77.5, 144., 122.,  45.,  45.])
            min_lim = np.radians([ -47.61,  -20., -77.5,   0., -80., -45., -45.])
        else:
            max_lim = np.radians([ 120.00,   20.,  77.5, 144.,   80.,  45.,  45.])
            min_lim = np.radians([ -47.61, -122.15, -77.5,   0., -122., -45., -45.])

        self.joint_lim_dict = {}
        self.joint_lim_dict['max'] = max_lim
        self.joint_lim_dict['min'] = min_lim

        wrist_stub_length = 0.0135 + 0.04318 # wrist linkange and FT sensor lengths
        self.setup_kdl_chains(arm, wrist_stub_length)

        if arm == 'r':
            pkl_nm = 'q_guess_right_dict.pkl'
        else:
            pkl_nm = 'q_guess_left_dict.pkl'

        pth = roslib.rospack.rospackexec(['find', 'hrl_cody_arms'])
        q_guess_pkl = pth + '/src/hrl_cody_arms/'+pkl_nm
        self.q_guess_dict = ut.load_pickle(q_guess_pkl)

    #--------------- KDL stuff ----------------
    # KDL joint array to meka joint list. (7->7)
    def kdl_angles_to_meka(self, q_jnt_arr):
        if q_jnt_arr == None:
            return None

        q_rad = [0. for i in range(7)]
        q_rad[0] = -q_jnt_arr[0]
        q_rad[1] = -q_jnt_arr[1]
        q_rad[2] = -q_jnt_arr[2]
        q_rad[3] = -q_jnt_arr[3]
        q_rad[4] = -q_jnt_arr[4]
        q_rad[5] = -q_jnt_arr[5]
        q_rad[6] = -q_jnt_arr[6]
        return q_rad

    # meka joint list to KDL joint array (7->7)
    def meka_angles_to_kdl(self, q_list):
        if q_list == None:
            return None

        n_joints = len(q_list)
        q = kdl.JntArray(n_joints)
        q[0] = -q_list[0]
        q[1] = -q_list[1]
        q[2] = -q_list[2]
        q[3] = -q_list[3]
        q[4] = -q_list[4]
        q[5] = -q_list[5]
        q[6] = -q_list[6]
        return q

    def create_right_chain(self, end_effector_length):
        ch = kdl.Chain()
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,-0.18493,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,-0.03175,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.00635,0.,-0.27795))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,-0.27853))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,0.,-end_effector_length))))
        return ch

    def create_left_chain(self, end_effector_length):
        ch = kdl.Chain()
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.18493,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,0.03175,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.00635,0.,-0.27795))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,-0.27853))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotZ),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotY),kdl.Frame(kdl.Vector(0.,0.,0.))))
        ch.addSegment(kdl.Segment(kdl.Joint(kdl.Joint.RotX),kdl.Frame(kdl.Vector(0.,0.,-end_effector_length))))
        return ch

    def create_solvers(self, ch):
         fk = kdl.ChainFkSolverPos_recursive(ch)
         ik_v = kdl.ChainIkSolverVel_pinv(ch)
         ik_p = kdl.ChainIkSolverPos_NR(ch, fk, ik_v)
         jac = kdl.ChainJntToJacSolver(ch)
         return fk, ik_v, ik_p, jac

    def setup_kdl_chains(self, arm, end_effector_length):
        if arm == 'r':
            #right arm
            ch = self.create_right_chain(end_effector_length)
            fk, ik_v, ik_p, jac = self.create_solvers(ch)
        else:
            #left arm
            ch = self.create_left_chain(end_effector_length)
            fk, ik_v, ik_p, jac = self.create_solvers(ch)

        kdl_chains = {}
        kdl_chains['chain'] = ch
        kdl_chains['nJnts'] = ch.getNrOfJoints()
        kdl_chains['fk_p'] = fk
        kdl_chains['ik_v'] = ik_v
        kdl_chains['ik_p'] = ik_p
        kdl_chains['jacobian_solver'] = jac

        #Add both chains to dictionary
        self.kdl_chains = kdl_chains

    def FK_kdl(self, q, link_number):
        fk_solver = self.kdl_chains['fk_p']
        endeffec_frame = kdl.Frame()
        kinematics_status = fk_solver.JntToCart(q, endeffec_frame,
                                                link_number)
        if kinematics_status >= 0:
#            print 'End effector transformation matrix:', endeffec_frame
            return endeffec_frame
        else:
            print 'Could not compute forward kinematics.'
            return None

    def IK_kdl(self,frame, q_init):
        nJnts = self.kdl_chains['nJnts']
        ik_solver = self.kdl_chains['ik_p']
        q = kdl.JntArray(nJnts)
        if ik_solver.CartToJnt(q_init,frame,q)>=0:
            for i in range(nJnts):
                q[i] = tr.angle_within_mod180(q[i])
            return q
        else:
            return None

    #-------------- implementation of HRLArmKinematics -----------

    def FK_vanilla(self, q, link_number = 7):
        q = self.meka_angles_to_kdl(q)
        frame = self.FK_kdl(q, link_number)
        pos = frame.p
        pos = ku.kdl_vec_to_np(pos)
        m = frame.M
        rot = ku.kdl_rot_to_np(m)
        return pos, rot

    ## compute Jacobian at point pos.
    # p is in the ground coord frame.
    def Jacobian(self, q, pos=None):
        print '-------------------------'
        print 'pos:', pos.A1
        print 'ee:', self.FK(q)[0].A1
        print '-------------------------'
        #pos = None
        if pos == None:
            pos = self.FK(q)[0]

        ch = self.kdl_chains['chain']
        v_list = []
        w_list = []

        for i in xrange(self.n_jts):
            p, rot = self.FK_vanilla(q, i)
            r = pos - p
            z_idx = ch.getSegment(i).getJoint().getType() - 1
            z = rot[:, z_idx]

            # this is a nasty trick. The rotation matrix returned by
            # FK_vanilla describes the orientation of a frame of the
            # KDL chain in the base frame. It just so happens that the
            # way Advait defined the KDL chain, the axis of rotation
            # in KDL is the -ve of the axis of rotation on the real
            # robot.
            # Advait apologizes for this.
            z = -z

            v_list.append(np.matrix(np.cross(z.A1, r.A1)).T)
            w_list.append(z)

        J = np.row_stack((np.column_stack(v_list), np.column_stack(w_list)))
        return J

    def IK_vanilla(self, p, rot, q_guess=None):
        p_kdl = ku.np_vec_to_kdl(p)
        rot_kdl = ku.np_rot_to_kdl(rot)
        fr = kdl.Frame(rot_kdl, p_kdl)

        if q_guess == None:
            q_guess = cgd.find_good_config(p, self.q_guess_dict)

        q_guess = self.meka_angles_to_kdl(q_guess)
        q_res = self.IK_kdl(arm, fr, q_guess)
        q_res = self.kdl_angles_to_meka(arm, q_res)
        return q_res

    #----------- extra functions -----------------

    ## clamp joint angles to their physical limits.
    # @param q - list of 7 joint angles.
    # The joint limits for IK are larger that the physical limits.
    def clamp_to_joint_limits(self, q, delta_list=[0.,0.,0.,0.,0.,0.,0.]):
        max_arr = self.joint_lim_dict['max']
        min_arr = self.joint_lim_dict['min']
        q_arr = np.array(q)
        d_arr = np.array(delta_list)
        return np.clip(q_arr, min_arr-d_arr, max_arr+d_arr)

    def within_joint_limits(self, q, delta_list=[0.,0.,0.,0.,0.,0.,0.]):
        max_arr = self.joint_lim_dict['max']
        min_arr = self.joint_lim_dict['min']
        q_arr = np.array(q)
        d_arr = np.array(delta_list)
        return np.all((q_arr <= max_arr+d_arr, q_arr >= min_arr-d_arr))

<|endoftext|>


<commit_before>#!/usr/bin/env python
import rospy
import random

from flexbe_core import EventState, Logger

class PickJokeState(EventState):
	'''
	Select a joke or pick a random one.

	-- joke_selection 	int 	Index of the joke or 0 for random.

	#> joke 			string 	Provides the selected joke.

	<= done 			Joke has been selected.

	'''

	RANDOM = 0

	def __init__(self, joke_selection):
		super(PickJokeState, self).__init__(outcomes = ['done'],
											output_keys = ['joke'])

		self._joke_selection = joke_selection
		self._jokes = [	'Boy is your name homework because I am not doing you and I should be',
                                 'Are you David Beckham? Because I would bend for you', 'What"s a nice guy like you doing with a body like that'
		]

	def execute(self, userdata):
		userdata.joke = self._jokes[self._joke_selection - 1]
		return 'done'

	def on_enter(self, userdata):
		if self._joke_selection == 0:
			self._joke_selection = random.randint(1, len(self._jokes))

<commit_msg>jokes updated<commit_after>#!/usr/bin/env python
import rospy
import random

from flexbe_core import EventState, Logger

class PickJokeState(EventState):
	'''
	Select a joke or pick a random one.

	-- joke_selection 	int 	Index of the joke or 0 for random.

	#> joke 			string 	Provides the selected joke.

	<= done 			Joke has been selected.

	'''

	RANDOM = 0

	def __init__(self, joke_selection):
		super(PickJokeState, self).__init__(outcomes = ['done'],
											output_keys = ['joke'])

		self._joke_selection = joke_selection
		self._jokes = [	'Boy is your name homework because I am not doing you and I should be',
                                 'Are you David Beckham? Because I would bend for you', 'What"s a nice guy like you doing with a body like that',
                   'I am not drunk, I am just intoxicated by you'.

		]

	def execute(self, userdata):
		userdata.joke = self._jokes[self._joke_selection - 1]
		return 'done'

	def on_enter(self, userdata):
		if self._joke_selection == 0:
			self._joke_selection = random.randint(1, len(self._jokes))

<|endoftext|>


<commit_before>from django import forms
from django.contrib.auth.models import Permission
from django.utils.translation import ugettext_lazy as _, pgettext_lazy

from oscar.core.loading import get_model
from oscar.core.compat import existing_user_fields, get_user_model
from oscar.apps.customer.forms import EmailUserCreationForm
from oscar.core.validators import password_validators

User = get_user_model()
Partner = get_model('partner', 'Partner')
PartnerAddress = get_model('partner', 'PartnerAddress')

class PartnerSearchForm(forms.Form):
    name = forms.CharField(
        required=False, label=pgettext_lazy(u"Partner's name", u"Name"))

class PartnerCreateForm(forms.ModelForm):

    def __init__(self, *args, **kwargs):
        super(PartnerCreateForm, self).__init__(*args, **kwargs)
        # Partner.name is optional and that is okay. But if creating through
        # the dashboard, it seems sensible to enforce as it's the only field
        # in the form.
        self.fields['name'].required = True

    class Meta:
        model = Partner
        fields = ('name',)

ROLE_CHOICES = (
    ('staff', _('Full dashboard access')),
    ('limited', _('Limited dashboard access')),
)

class NewUserForm(EmailUserCreationForm):
    role = forms.ChoiceField(choices=ROLE_CHOICES, widget=forms.RadioSelect,
                             label=_('User role'), initial='limited')

    def __init__(self, partner, *args, **kwargs):
        self.partner = partner
        super(NewUserForm, self).__init__(host=None, *args, **kwargs)

    def save(self):
        role = self.cleaned_data.get('role', 'limited')
        user = super(NewUserForm, self).save(commit=False)
        user.is_staff = role == 'staff'
        user.save()
        self.partner.users.add(user)
        if role == 'limited':
            dashboard_access_perm = Permission.objects.get(
                codename='dashboard_access', content_type__app_label='partner')
            user.user_permissions.add(dashboard_access_perm)
        return user

    class Meta:
        model = User
        fields = existing_user_fields(
            ['first_name', 'last_name', 'email']) + ['password1', 'password2']

class ExistingUserForm(forms.ModelForm):
    """
    Slightly different form that makes
    * makes saving password optional
    * doesn't regenerate username
    * doesn't allow changing email till #668 is resolved
    """
    role = forms.ChoiceField(choices=ROLE_CHOICES, widget=forms.RadioSelect,
                             label=_('User role'))
    password1 = forms.CharField(
        label=_('Password'),
        widget=forms.PasswordInput,
        required=False,
        validators=password_validators)
    password2 = forms.CharField(
        required=False,
        label=_('Confirm Password'),
        widget=forms.PasswordInput)

    def clean_password2(self):
        password1 = self.cleaned_data.get('password1', '')
        password2 = self.cleaned_data.get('password2', '')

        if password1 != password2:
            raise forms.ValidationError(
                _("The two password fields didn't match."))
        return password2

    def __init__(self, *args, **kwargs):
        user = kwargs['instance']
        role = 'staff' if user.is_staff else 'limited'
        kwargs.get('initial', {}).setdefault('role', role)
        super(ExistingUserForm, self).__init__(*args, **kwargs)

    def save(self):
        role = self.cleaned_data.get('role', 'none')
        user = super(ExistingUserForm, self).save(commit=False)
        user.is_staff = role == 'staff'
        if self.cleaned_data['password1']:
            user.set_password(self.cleaned_data['password1'])
        user.save()

        dashboard_perm = Permission.objects.get(
            codename='dashboard_access', content_type__app_label='partner')
        user_has_perm = user.user_permissions.filter(
            pk=dashboard_perm.pk).exists()
        if role == 'limited' and not user_has_perm:
            user.user_permissions.add(dashboard_perm)
        elif role == 'staff' and user_has_perm:
            user.user_permissions.remove(dashboard_perm)
        return user

    class Meta:
        model = User
        fields = existing_user_fields(
            ['first_name', 'last_name']) + ['password1', 'password2']

class UserEmailForm(forms.Form):
    # We use a CharField so that a partial email address can be entered
    email = forms.CharField(
        label=_("Email address"), max_length=100)

class PartnerAddressForm(forms.ModelForm):
    name = forms.CharField(
        required=False, label=pgettext_lazy(u"Partner's name", u"Name"))

    class Meta:
        fields = ('name', 'line1', 'line2', 'line3', 'line4',
                  'state', 'postcode', 'country')
        model = PartnerAddress
<commit_msg>PartnerAddressForm's name field is a CharField with no max_length. It causes a DataError if users enter a name with more than 128 characters (max_length of Partner's name).<commit_after>from django import forms
from django.contrib.auth.models import Permission
from django.utils.translation import ugettext_lazy as _, pgettext_lazy

from oscar.core.loading import get_model
from oscar.core.compat import existing_user_fields, get_user_model
from oscar.apps.customer.forms import EmailUserCreationForm
from oscar.core.validators import password_validators

User = get_user_model()
Partner = get_model('partner', 'Partner')
PartnerAddress = get_model('partner', 'PartnerAddress')

class PartnerSearchForm(forms.Form):
    name = forms.CharField(
        required=False, label=pgettext_lazy(u"Partner's name", u"Name"))

class PartnerCreateForm(forms.ModelForm):

    def __init__(self, *args, **kwargs):
        super(PartnerCreateForm, self).__init__(*args, **kwargs)
        # Partner.name is optional and that is okay. But if creating through
        # the dashboard, it seems sensible to enforce as it's the only field
        # in the form.
        self.fields['name'].required = True

    class Meta:
        model = Partner
        fields = ('name',)

ROLE_CHOICES = (
    ('staff', _('Full dashboard access')),
    ('limited', _('Limited dashboard access')),
)

class NewUserForm(EmailUserCreationForm):
    role = forms.ChoiceField(choices=ROLE_CHOICES, widget=forms.RadioSelect,
                             label=_('User role'), initial='limited')

    def __init__(self, partner, *args, **kwargs):
        self.partner = partner
        super(NewUserForm, self).__init__(host=None, *args, **kwargs)

    def save(self):
        role = self.cleaned_data.get('role', 'limited')
        user = super(NewUserForm, self).save(commit=False)
        user.is_staff = role == 'staff'
        user.save()
        self.partner.users.add(user)
        if role == 'limited':
            dashboard_access_perm = Permission.objects.get(
                codename='dashboard_access', content_type__app_label='partner')
            user.user_permissions.add(dashboard_access_perm)
        return user

    class Meta:
        model = User
        fields = existing_user_fields(
            ['first_name', 'last_name', 'email']) + ['password1', 'password2']

class ExistingUserForm(forms.ModelForm):
    """
    Slightly different form that makes
    * makes saving password optional
    * doesn't regenerate username
    * doesn't allow changing email till #668 is resolved
    """
    role = forms.ChoiceField(choices=ROLE_CHOICES, widget=forms.RadioSelect,
                             label=_('User role'))
    password1 = forms.CharField(
        label=_('Password'),
        widget=forms.PasswordInput,
        required=False,
        validators=password_validators)
    password2 = forms.CharField(
        required=False,
        label=_('Confirm Password'),
        widget=forms.PasswordInput)

    def clean_password2(self):
        password1 = self.cleaned_data.get('password1', '')
        password2 = self.cleaned_data.get('password2', '')

        if password1 != password2:
            raise forms.ValidationError(
                _("The two password fields didn't match."))
        return password2

    def __init__(self, *args, **kwargs):
        user = kwargs['instance']
        role = 'staff' if user.is_staff else 'limited'
        kwargs.get('initial', {}).setdefault('role', role)
        super(ExistingUserForm, self).__init__(*args, **kwargs)

    def save(self):
        role = self.cleaned_data.get('role', 'none')
        user = super(ExistingUserForm, self).save(commit=False)
        user.is_staff = role == 'staff'
        if self.cleaned_data['password1']:
            user.set_password(self.cleaned_data['password1'])
        user.save()

        dashboard_perm = Permission.objects.get(
            codename='dashboard_access', content_type__app_label='partner')
        user_has_perm = user.user_permissions.filter(
            pk=dashboard_perm.pk).exists()
        if role == 'limited' and not user_has_perm:
            user.user_permissions.add(dashboard_perm)
        elif role == 'staff' and user_has_perm:
            user.user_permissions.remove(dashboard_perm)
        return user

    class Meta:
        model = User
        fields = existing_user_fields(
            ['first_name', 'last_name']) + ['password1', 'password2']

class UserEmailForm(forms.Form):
    # We use a CharField so that a partial email address can be entered
    email = forms.CharField(
        label=_("Email address"), max_length=100)

class PartnerAddressForm(forms.ModelForm):
    name = forms.CharField(
        required=False, max_length=128,
        label=pgettext_lazy(u"Partner's name", u"Name"))

    class Meta:
        fields = ('name', 'line1', 'line2', 'line3', 'line4',
                  'state', 'postcode', 'country')
        model = PartnerAddress
<|endoftext|>