<commit_before>from operator import attrgetter

import graphene
from django.utils.translation import npgettext_lazy, pgettext_lazy
from django_prices.templatetags import prices_i18n
from graphene.types import InputObjectType
from graphql_jwt.decorators import permission_required
from payments import PaymentError, PaymentStatus

from ...account.models import Address
from ...core.utils.taxes import ZERO_TAXED_MONEY
from ...order import (
    CustomPaymentChoices, FulfillmentStatus, OrderStatus, models)
from ...order.utils import cancel_order
from ...shipping.models import ANY_COUNTRY
from ..core.mutations import BaseMutation, ModelDeleteMutation, ModelMutation
from ..core.types import Decimal, Error
from ..utils import get_node
from .types import Order

def try_payment_action(action, money, errors):
    try:
        action(money.amount)
    except (PaymentError, ValueError) as e:
        message_dict = e.message_dict
        for field in message_dict:
            for message in message_dict[field]:
                errors.append(Error(field=field, message=message))

class AddressInput(graphene.InputObjectType):
    first_name = graphene.String(description='Given name.')
    last_name = graphene.String(description='Family name.')
    company_name = graphene.String(description='Company or organization.')
    street_address_1 = graphene.String(description='Address.')
    street_address_2 = graphene.String(description='Address.')
    city = graphene.String(description='City.')
    city_area = graphene.String(description='District.')
    postal_code = graphene.String(description='Postal code.')
    country = graphene.String(description='Country.')
    country_area = graphene.String(description='State or province.')
    phone = graphene.String(description='Phone number.')
    email = graphene.String(description='Email address.')

class LineInput(graphene.InputObjectType):
    variant_id = graphene.ID(description='Product variant ID.')
    quantity = graphene.Int(
        description='The number of products in the order.')

    class Meta:
        description = 'Represents a permission object in a friendly form.'

class DraftOrderInput(InputObjectType):
    billing_address = AddressInput(
        description='Address associated with the payment.')
    user = graphene.ID(
        descripton='Customer associated with the draft order.')
    email = graphene.String(description='Email address of the customer.')
    discount = Decimal(description='Discount amount for the order.')
    lines = graphene.List(
        LineInput,
        description="""Variant line input consisting of variant ID
        and quantity of products.""")
    shipping_address = AddressInput(
        description='Address to where the order will be shipped.')
    shipping_method = graphene.ID(
        description='ID of a selected shipping method.')
    voucher = graphene.ID(
        description='ID of the voucher associated with the order')

class OrderUpdateInput(graphene.InputObjectType):
    billing_address = AddressInput(
        description='Address associated with the payment.')
    user_email = graphene.String(description='Email address of the customer.')
    shipping_address = AddressInput(
        description='Address to where the order will be shipped.')

class DraftOrderCreate(ModelMutation):
    class Arguments:
        input = DraftOrderInput(
            required=True,
            description='Fields required to create an order.')

    class Meta:
        description = 'Creates a new variant for a product.'
        model = models.Order

    @classmethod
    def clean_input(cls, info, instance, input, errors):
        shipping_address = input.pop('shipping_address', None)
        billing_address = input.pop('billing_address', None)
        cleaned_input = super().clean_input(info, instance, input, errors)
        cleaned_input['status'] = OrderStatus.DRAFT
        display_gross_prices = info.context.site.settings.display_gross_prices
        cleaned_input['display_gross_prices'] = display_gross_prices

        # Set up default addresses if possible
        user = cleaned_input.get('user')
        if user and not shipping_address:
            cleaned_input[
                'shipping_address'] = user.default_shipping_address
        if user and not billing_address:
            cleaned_input[
                'billing_address'] = user.default_billing_address
        if shipping_address:
            shipping_address = Address(**shipping_address)
            if not cls.clean_instance(shipping_address, errors):
                shipping_address.save()
            cleaned_input['shipping_address'] = shipping_address

        if billing_address:
            billing_address = Address(**billing_address)
            if not cls.clean_instance(billing_address, errors):
                billing_address.save()
            cleaned_input['billing_address'] = billing_address

        return cleaned_input

    @classmethod
    def construct_instance(cls, instance, cleaned_data):
        instance = super().construct_instance(instance, cleaned_data)
        instance.shipping_address = cleaned_data.get('shipping_address')
        instance.billing_address = cleaned_data.get('billing_address')
        return instance

    @classmethod
    def user_is_allowed(cls, user, input):
        return user.has_perm('order.edit_order')

class DraftOrderUpdate(DraftOrderCreate):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of an order to update.')
        input = DraftOrderInput(
            required=True,
            description='Fields required to update an order.')

    class Meta:
        description = 'Updates a draft order.'
        model = models.Order

class DraftOrderDelete(ModelDeleteMutation):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of a draft order to delete.')

    class Meta:
        description = 'Deletes a draft order.'
        model = models.Order

    @classmethod
    def user_is_allowed(cls, user, input):
        return user.has_perm('order.edit_order')

def check_for_draft_order_errors(order):
    """Return a list of errors associated with the order.

    Checks, if given order has a proper shipping address and method
    set up and return list of errors if not.
    """
    errors = []
    if order.get_total_quantity() == 0:
        errors.append(
            Error(
                field='lines',
                message='Could not create order without any products.'))
    method = order.shipping_method
    shipping_address = order.shipping_address
    shipping_not_valid = (
        method and shipping_address and
        method.country_code != ANY_COUNTRY and
        shipping_address.country.code != method.country_code)
    if shipping_not_valid:
        errors.append(
            Error(
                field='shipping',
                message='Shipping method is not valid for chosen shipping '
                        'address'))
    return errors

class DraftOrderComplete(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True,
            description='ID of the order that will be completed.')

    class Meta:
        description = 'Completes creating an order.'

    order = graphene.Field(
        Order, description='Completed order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id):
        order = get_node(info, order_id, only_type=Order)
        errors = check_for_draft_order_errors(order)
        if errors:
            return cls(errors=errors)

        order.status = OrderStatus.UNFULFILLED
        if order.user:
            order.user_email = order.user.email
        remove_shipping_address = False
        if not order.is_shipping_required():
            order.shipping_method_name = None
            order.shipping_price = ZERO_TAXED_MONEY
            if order.shipping_address:
                remove_shipping_address = True
        order.save()
        if remove_shipping_address:
            order.shipping_address.delete()
        msg = 'Order created from draft order'
        order.history.create(content=msg, user=info.context.user)
        return DraftOrderComplete(order=order)

class OrderUpdate(DraftOrderUpdate):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of an order to update.')
        input = OrderUpdateInput(
            required=True,
            description='Fields required to update an order.')

    class Meta:
        description = 'Updates an order.'
        model = models.Order

class OrderAddNoteInput(graphene.InputObjectType):
    order = graphene.ID(description='ID of the order.')
    user = graphene.ID(description='ID of the user who added note.')
    content = graphene.String(description='Note content.')
    is_public = graphene.String(
        description='Determine if note is visible by customer or not.')

class OrderAddNote(ModelMutation):
    class Arguments:
        input = OrderAddNoteInput(
            required=True,
            description='Fields required to add note to order.')

    class Meta:
        description = 'Adds note to order.'
        model = models.OrderNote

    @classmethod
    def user_is_allowed(cls, user, input):
        return user.has_perm('order.edit_order')

    @classmethod
    def save(cls, info, instance, cleaned_input):
        super.save(info, instance, cleaned_input)
        msg = 'Added note'
        instance.order.history.create(content=msg, user=info.context.user)

class OrderCancel(BaseMutation):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of the order to cancel.')
        restock = graphene.Boolean(
            required=True,
            description='Determine if lines will be restocked or not.')

    order = graphene.Field(
        Order, description='Canceled order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, id, restock):
        order = get_node(info, id, only_type=Order)
        cancel_order(order=order, restock=restock)
        if restock:
            restock_msg = npgettext_lazy(
                'Dashboard message',
                'Restocked %(quantity)d item',
                'Restocked %(quantity)d items',
                'quantity') % {'quantity': order.get_total_quantity()}
            order.history.create(content=restock_msg, user=info.context.user)
        else:
            msg = pgettext_lazy('Dashboard message', 'Order canceled')
            order.history.create(content=msg, user=info.context.user)
        return OrderCancel(order=order)

class OrderMarkAsPaid(BaseMutation):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of the order to mark paid.')

    order = graphene.Field(
        Order, description='Mark order as manually paid.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, id):
        order = get_node(info, id, only_type=Order)
        if order.payments.exists():
            field = 'payment'
            msg = 'Orders with payments can not be manually marked as paid.'
            return cls(errors=[Error(field=field, message=msg)])
        defaults = {
            'total': order.total.gross.amount,
            'tax': order.total.tax.amount,
            'currency': order.total.currency,
            'delivery': order.shipping_price.net.amount,
            'description': pgettext_lazy(
                'Payment description', 'Order %(order)s') % {'order': order},
            'captured_amount': order.total.gross.amount}
        models.Payment.objects.get_or_create(
            variant=CustomPaymentChoices.MANUAL,
            status=PaymentStatus.CONFIRMED, order=order,
            defaults=defaults)
        msg = 'Order manually marked as paid'
        order.history.create(content=msg, user=info.context.user)
        return OrderMarkAsPaid(order=order)

class OrderCapture(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True, description='ID of the order to capture.')

    order = graphene.Field(
        Order, description='Captured order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id):
        order = get_node(info, order_id, only_type=Order)
        payment = order.get_last_payment()
        amount = order.total.quantize('0.01').gross
        errors = []
        try_payment_action(payment.capture, amount, errors)
        if errors:
            return cls(errors=errors)

        msg = 'Captured %(amount)s' % {
            'amount': prices_i18n.amount(amount)}
        order.history.create(content=msg, user=info.context.user)
        return OrderCapture(order=order)

class OrderRelease(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True, description='ID of the order to release.')

    order = graphene.Field(
        Order, description='released order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id):
        order = get_node(info, order_id, only_type=Order)
        payment = order.get_last_payment()
        errors = []
        if payment.status != PaymentStatus.PREAUTH:
            errors.append(
                Error(field='payment',
                      message='Only pre-authorized payments can be released'))
        try:
            payment.release()
        except (PaymentError, ValueError) as e:
            message_dict = e.message_dict
            for field in message_dict:
                for message in message_dict[field]:
                    errors.append(Error(field=field, message=message))
        if errors:
            return cls(errors=errors)

        msg = 'Released payment'
        order.history.create(content=msg, user=info.context.user)
        return OrderRelease(order=order)

class OrderRefund(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True, description='ID of the order to refund.')

    order = graphene.Field(
        Order, description='released order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id):
        order = get_node(info, order_id, only_type=Order)
        payment = order.get_last_payment()
        amount = payment.captured_amount
        errors = []
        if payment.variant == CustomPaymentChoices.MANUAL:
            errors.append(
                Error(field='payment',
                      message='Manual payments can not be refunded.'))
        try_payment_action(payment.refund, amount, errors)
        if errors:
            return cls(errors=errors)

        msg = 'Refunded %(amount)s' % {
            'amount': prices_i18n.amount(amount)}
        order.history.create(content=msg, user=info.context.user)
        return OrderRefund(order=order)
<commit_msg>Add amount variable to proper mutations<commit_after>from operator import attrgetter

import graphene
from django.utils.translation import npgettext_lazy, pgettext_lazy
from django_prices.templatetags import prices_i18n
from graphene.types import InputObjectType
from graphql_jwt.decorators import permission_required
from payments import PaymentError, PaymentStatus

from ...account.models import Address
from ...core.utils.taxes import ZERO_TAXED_MONEY
from ...order import (
    CustomPaymentChoices, FulfillmentStatus, OrderStatus, models)
from ...order.utils import cancel_order
from ...shipping.models import ANY_COUNTRY
from ..core.mutations import BaseMutation, ModelDeleteMutation, ModelMutation
from ..core.types import Decimal, Error
from ..utils import get_node
from .types import Order

def try_payment_action(action, money, errors):
    try:
        action(money)
    except (PaymentError, ValueError) as e:
        errors.append(Error(field='payment', message=str(e)))

class AddressInput(graphene.InputObjectType):
    first_name = graphene.String(description='Given name.')
    last_name = graphene.String(description='Family name.')
    company_name = graphene.String(description='Company or organization.')
    street_address_1 = graphene.String(description='Address.')
    street_address_2 = graphene.String(description='Address.')
    city = graphene.String(description='City.')
    city_area = graphene.String(description='District.')
    postal_code = graphene.String(description='Postal code.')
    country = graphene.String(description='Country.')
    country_area = graphene.String(description='State or province.')
    phone = graphene.String(description='Phone number.')
    email = graphene.String(description='Email address.')

class LineInput(graphene.InputObjectType):
    variant_id = graphene.ID(description='Product variant ID.')
    quantity = graphene.Int(
        description='The number of products in the order.')

    class Meta:
        description = 'Represents a permission object in a friendly form.'

class DraftOrderInput(InputObjectType):
    billing_address = AddressInput(
        description='Address associated with the payment.')
    user = graphene.ID(
        descripton='Customer associated with the draft order.')
    email = graphene.String(description='Email address of the customer.')
    discount = Decimal(description='Discount amount for the order.')
    lines = graphene.List(
        LineInput,
        description="""Variant line input consisting of variant ID
        and quantity of products.""")
    shipping_address = AddressInput(
        description='Address to where the order will be shipped.')
    shipping_method = graphene.ID(
        description='ID of a selected shipping method.')
    voucher = graphene.ID(
        description='ID of the voucher associated with the order')

class OrderUpdateInput(graphene.InputObjectType):
    billing_address = AddressInput(
        description='Address associated with the payment.')
    user_email = graphene.String(description='Email address of the customer.')
    shipping_address = AddressInput(
        description='Address to where the order will be shipped.')

class DraftOrderCreate(ModelMutation):
    class Arguments:
        input = DraftOrderInput(
            required=True,
            description='Fields required to create an order.')

    class Meta:
        description = 'Creates a new variant for a product.'
        model = models.Order

    @classmethod
    def clean_input(cls, info, instance, input, errors):
        shipping_address = input.pop('shipping_address', None)
        billing_address = input.pop('billing_address', None)
        cleaned_input = super().clean_input(info, instance, input, errors)
        cleaned_input['status'] = OrderStatus.DRAFT
        display_gross_prices = info.context.site.settings.display_gross_prices
        cleaned_input['display_gross_prices'] = display_gross_prices

        # Set up default addresses if possible
        user = cleaned_input.get('user')
        if user and not shipping_address:
            cleaned_input[
                'shipping_address'] = user.default_shipping_address
        if user and not billing_address:
            cleaned_input[
                'billing_address'] = user.default_billing_address
        if shipping_address:
            shipping_address = Address(**shipping_address)
            if not cls.clean_instance(shipping_address, errors):
                shipping_address.save()
            cleaned_input['shipping_address'] = shipping_address

        if billing_address:
            billing_address = Address(**billing_address)
            if not cls.clean_instance(billing_address, errors):
                billing_address.save()
            cleaned_input['billing_address'] = billing_address

        return cleaned_input

    @classmethod
    def construct_instance(cls, instance, cleaned_data):
        instance = super().construct_instance(instance, cleaned_data)
        instance.shipping_address = cleaned_data.get('shipping_address')
        instance.billing_address = cleaned_data.get('billing_address')
        return instance

    @classmethod
    def user_is_allowed(cls, user, input):
        return user.has_perm('order.edit_order')

class DraftOrderUpdate(DraftOrderCreate):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of an order to update.')
        input = DraftOrderInput(
            required=True,
            description='Fields required to update an order.')

    class Meta:
        description = 'Updates a draft order.'
        model = models.Order

class DraftOrderDelete(ModelDeleteMutation):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of a draft order to delete.')

    class Meta:
        description = 'Deletes a draft order.'
        model = models.Order

    @classmethod
    def user_is_allowed(cls, user, input):
        return user.has_perm('order.edit_order')

def check_for_draft_order_errors(order):
    """Return a list of errors associated with the order.

    Checks, if given order has a proper shipping address and method
    set up and return list of errors if not.
    """
    errors = []
    if order.get_total_quantity() == 0:
        errors.append(
            Error(
                field='lines',
                message='Could not create order without any products.'))
    method = order.shipping_method
    shipping_address = order.shipping_address
    shipping_not_valid = (
        method and shipping_address and
        method.country_code != ANY_COUNTRY and
        shipping_address.country.code != method.country_code)
    if shipping_not_valid:
        errors.append(
            Error(
                field='shipping',
                message='Shipping method is not valid for chosen shipping '
                        'address'))
    return errors

class DraftOrderComplete(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True,
            description='ID of the order that will be completed.')

    class Meta:
        description = 'Completes creating an order.'

    order = graphene.Field(
        Order, description='Completed order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id):
        order = get_node(info, order_id, only_type=Order)
        errors = check_for_draft_order_errors(order)
        if errors:
            return cls(errors=errors)

        order.status = OrderStatus.UNFULFILLED
        if order.user:
            order.user_email = order.user.email
        remove_shipping_address = False
        if not order.is_shipping_required():
            order.shipping_method_name = None
            order.shipping_price = ZERO_TAXED_MONEY
            if order.shipping_address:
                remove_shipping_address = True
        order.save()
        if remove_shipping_address:
            order.shipping_address.delete()
        msg = 'Order created from draft order'
        order.history.create(content=msg, user=info.context.user)
        return DraftOrderComplete(order=order)

class OrderUpdate(DraftOrderUpdate):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of an order to update.')
        input = OrderUpdateInput(
            required=True,
            description='Fields required to update an order.')

    class Meta:
        description = 'Updates an order.'
        model = models.Order

class OrderAddNoteInput(graphene.InputObjectType):
    order = graphene.ID(description='ID of the order.')
    user = graphene.ID(description='ID of the user who added note.')
    content = graphene.String(description='Note content.')
    is_public = graphene.String(
        description='Determine if note is visible by customer or not.')

class OrderAddNote(ModelMutation):
    class Arguments:
        input = OrderAddNoteInput(
            required=True,
            description='Fields required to add note to order.')

    class Meta:
        description = 'Adds note to order.'
        model = models.OrderNote

    @classmethod
    def user_is_allowed(cls, user, input):
        return user.has_perm('order.edit_order')

    @classmethod
    def save(cls, info, instance, cleaned_input):
        super.save(info, instance, cleaned_input)
        msg = 'Added note'
        instance.order.history.create(content=msg, user=info.context.user)

class OrderCancel(BaseMutation):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of the order to cancel.')
        restock = graphene.Boolean(
            required=True,
            description='Determine if lines will be restocked or not.')

    order = graphene.Field(
        Order, description='Canceled order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, id, restock):
        order = get_node(info, id, only_type=Order)
        cancel_order(order=order, restock=restock)
        if restock:
            restock_msg = npgettext_lazy(
                'Dashboard message',
                'Restocked %(quantity)d item',
                'Restocked %(quantity)d items',
                'quantity') % {'quantity': order.get_total_quantity()}
            order.history.create(content=restock_msg, user=info.context.user)
        else:
            msg = pgettext_lazy('Dashboard message', 'Order canceled')
            order.history.create(content=msg, user=info.context.user)
        return OrderCancel(order=order)

class OrderMarkAsPaid(BaseMutation):
    class Arguments:
        id = graphene.ID(
            required=True, description='ID of the order to mark paid.')

    order = graphene.Field(
        Order, description='Mark order as manually paid.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, id):
        order = get_node(info, id, only_type=Order)
        if order.payments.exists():
            field = 'payment'
            msg = 'Orders with payments can not be manually marked as paid.'
            return cls(errors=[Error(field=field, message=msg)])
        defaults = {
            'total': order.total.gross.amount,
            'tax': order.total.tax.amount,
            'currency': order.total.currency,
            'delivery': order.shipping_price.net.amount,
            'description': pgettext_lazy(
                'Payment description', 'Order %(order)s') % {'order': order},
            'captured_amount': order.total.gross.amount}
        models.Payment.objects.get_or_create(
            variant=CustomPaymentChoices.MANUAL,
            status=PaymentStatus.CONFIRMED, order=order,
            defaults=defaults)
        msg = 'Order manually marked as paid.'
        order.history.create(content=msg, user=info.context.user)
        return OrderMarkAsPaid(order=order)

class OrderCapture(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True, description='ID of the order to capture.')
        amount = Decimal(
            required=True, description='Amount of money to capture.')

    order = graphene.Field(
        Order, description='Captured order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id, amount):
        order = get_node(info, order_id, only_type=Order)
        payment = order.get_last_payment()
        errors = []
        try_payment_action(payment.capture, amount, errors)
        if errors:
            return cls(errors=errors)

        msg = 'Captured %(amount)s' % {'amount': amount}
        order.history.create(content=msg, user=info.context.user)
        return OrderCapture(order=order)

class OrderRelease(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True, description='ID of the order to release.')

    order = graphene.Field(
        Order, description='released order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id):
        order = get_node(info, order_id, only_type=Order)
        payment = order.get_last_payment()
        errors = []
        if payment.status != PaymentStatus.PREAUTH:
            errors.append(
                Error(field='payment',
                      message='Only pre-authorized payments can be released'))
        try:
            payment.release()
        except (PaymentError, ValueError) as e:
            message_dict = e.message_dict
            for field in message_dict:
                for message in message_dict[field]:
                    errors.append(Error(field=field, message=message))
        if errors:
            return cls(errors=errors)

        msg = 'Released payment'
        order.history.create(content=msg, user=info.context.user)
        return OrderRelease(order=order)

class OrderRefund(BaseMutation):
    class Arguments:
        order_id = graphene.ID(
            required=True, description='ID of the order to refund.')
        amount = Decimal(
            required=True, description='Amount of money to refund.')

    order = graphene.Field(
        Order, description='released order.')

    @classmethod
    @permission_required('order.edit_order')
    def mutate(cls, root, info, order_id, amount):
        order = get_node(info, order_id, only_type=Order)
        payment = order.get_last_payment()
        errors = []
        if payment.variant == CustomPaymentChoices.MANUAL:
            errors.append(
                Error(field='payment',
                      message='Manual payments can not be refunded.'))
        import ipdb; ipdb.set_trace()
        try_payment_action(payment.refund, amount, errors)
        if errors:
            return cls(errors=errors)

        msg = 'Refunded %(amount)s' % {'amount': amount}
        order.history.create(content=msg, user=info.context.user)
        return OrderRefund(order=order)
<|endoftext|>


<commit_before>#  Copyright (c) 2021, Manfred Moitzi
#  License: MIT License
from typing import Callable, Optional, Dict
import abc
import sys
import os
import glob
import signal
import logging
from pathlib import Path

from ezdxf import recover
from ezdxf.lldxf import const
from ezdxf.lldxf.validator import is_dxf_file

__all__ = ["get", "add_parsers"]

logger = logging.getLogger("ezdxf")

def get(cmd: str) -> Optional[Callable]:
    cls = _commands.get(cmd)
    if cls:
        return cls.run
    return None

def add_parsers(subparsers) -> None:
    for cmd in _commands.values():  # in order of registration
        cmd.add_parser(subparsers)

class Command:
    """abstract base class for launcher commands"""

    NAME = "command"

    @staticmethod
    @abc.abstractmethod
    def add_parser(subparsers) -> None:
        pass

    @staticmethod
    @abc.abstractmethod
    def run(args) -> None:
        pass

_commands: Dict[str, Command] = dict()

def register(cls: Command):
    """Register a launcher sub-command."""
    _commands[cls.NAME] = cls
    return cls

@register
class PrettyPrint(Command):
    """Launcher sub-command: pp"""

    NAME = "pp"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            PrettyPrint.NAME, help="pretty print DXF files as HTML file"
        )
        parser.add_argument(
            "files",
            metavar="FILE",
            nargs="+",
            help="DXF files pretty print",
        )
        parser.add_argument(
            "-o",
            "--open",
            action="store_true",
            help="open generated HTML file by the default web browser",
        )
        parser.add_argument(
            "-r",
            "--raw",
            action="store_true",
            help="raw mode, no DXF structure interpretation",
        )
        parser.add_argument(
            "-x",
            "--nocompile",
            action="store_true",
            help="don't compile points coordinates into single tags "
                 "(only in raw mode)",
        )
        parser.add_argument(
            "-l",
            "--legacy",
            action="store_true",
            help="legacy mode, reorder DXF point coordinates",
        )
        parser.add_argument(
            "-s",
            "--sections",
            action="store",
            default="hctbeo",
            help="choose sections to include and their order, h=HEADER, c=CLASSES, "
                 "t=TABLES, b=BLOCKS, e=ENTITIES, o=OBJECTS",
        )

    @staticmethod
    def run(args):
        from ezdxf.pp import run

        run(args)

@register
class Audit(Command):
    """Launcher sub-command: audit"""

    NAME = "audit"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Audit.NAME, help="audit and repair DXF files"
        )
        parser.add_argument(
            "files",
            metavar="FILE",
            nargs="+",
            help="audit DXF files",
        )
        parser.add_argument(
            "-s",
            "--save",
            action="store_true",
            help='save recovered files with extension ".rec.dxf" ',
        )

    @staticmethod
    def run(args):
        def build_outname(name: str) -> str:
            p = Path(name)
            return p.parent / (p.stem + ".rec.dxf")

        def log_fixes(auditor):
            for error in auditor.fixes:
                logger.info("fixed:" + error.message)

        def log_errors(auditor):
            for error in auditor.errors:
                logger.error(error.message)

        def _audit(filename: str) -> None:
            msg = f"auditing file: {filename}"
            print(msg)
            logger.info(msg)
            try:
                doc, auditor = recover.readfile(filename)
            except IOError:
                msg = "Not a DXF file or a generic I/O error."
                print(msg)
                logger.error(msg)
                return  # keep on processing additional files
            except const.DXFStructureError:
                msg = "Invalid or corrupted DXF file."
                print(msg)
                logger.error(msg)
                return  # keep on processing additional files

            if auditor.has_errors:
                auditor.print_error_report()
                log_errors(auditor)
            if auditor.has_fixes:
                auditor.print_fixed_errors()
                log_fixes(auditor)

            if auditor.has_errors is False and auditor.has_fixes is False:
                print("No errors found.")
            else:
                print(
                    f"Found {len(auditor.errors)} errors, "
                    f"applied {len(auditor.fixes)} fixes"
                )

            if args.save:
                outname = build_outname(filename)
                doc.saveas(outname)
                print(f"Saved recovered file as: {outname}")

        for pattern in args.files:
            names = list(glob.glob(pattern))
            if len(names) == 0:
                msg = f"File(s) '{pattern}' not found."
                print(msg)
                logger.error(msg)
                continue
            for filename in names:
                if not os.path.exists(filename):
                    msg = f"File '{filename}' not found."
                    print(msg)
                    logger.error(msg)
                    continue
                if not is_dxf_file(filename):
                    msg = f"File '{filename}' is not a DXF file."
                    print(msg)
                    logger.error(msg)
                    continue
                _audit(filename)

def load_document(filename: str):
    try:
        doc, auditor = recover.readfile(filename)
    except IOError:
        msg = f"Not a DXF file or a generic I/O error: {filename}"
        print(msg)
        logger.error(msg)
        sys.exit(2)
    except const.DXFStructureError:
        msg = f"Invalid or corrupted DXF file: {filename}"
        print(msg)
        logger.error(msg)
        sys.exit(3)

    if auditor.has_errors:
        # But is most likely good enough for rendering.
        msg = f"Found {len(auditor.errors)} unrecoverable errors."
        print(msg)
        logger.error(msg)
    if auditor.has_fixes:
        msg = f"Fixed {len(auditor.fixes)} errors."
        print(msg)
        logger.info(msg)
    return doc, auditor

@register
class Draw(Command):
    """Launcher sub-command: draw"""

    NAME = "draw"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Draw.NAME, help="draw and convert DXF files by Matplotlib"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="?",
            help="DXF file to view or convert",
        )
        parser.add_argument(
            "--formats",
            action="store_true",
            help="show all supported export formats and exit",
        )
        parser.add_argument(
            "--layout",
            default="Model",
            help="select the layout to draw",
        )
        parser.add_argument(
            "--all-layers-visible",
            help="draw all layers including the ones marked as invisible",
        )
        parser.add_argument(
            "--all-entities-visible",
            help="draw all entities including the ones marked as invisible "
                 "(some entities are individually marked as invisible even "
                 "if the layer is visible)",
        )
        parser.add_argument(
            "-o", "--out", required=False, help="output filename for export"
        )
        parser.add_argument(
            "--dpi",
            type=int,
            default=300,
            help="target render resolution, default is 300",
        )
        parser.add_argument(
            "--ltype",
            default="internal",
            choices=["internal", "ezdxf"],
            help="select the line type rendering engine, default is internal",
        )

    @staticmethod
    def run(args):
        # Import on demand for a quicker startup:
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            print("Matplotlib package not found.")
            sys.exit(2)

        from ezdxf.addons.drawing import RenderContext, Frontend
        from ezdxf.addons.drawing.matplotlib import MatplotlibBackend

        # Print all supported export formats:
        if args.formats:
            fig = plt.figure()
            for (
                extension,
                description,
            ) in fig.canvas.get_supported_filetypes().items():
                print(f"{extension}: {description}")
            sys.exit(0)

        if args.file:
            filename = args.file
        else:
            print("argument FILE is required")
            sys.exit(1)

        doc, _ = load_document(filename)

        try:
            layout = doc.layouts.get(args.layout)
        except KeyError:
            print(f'Could not find layout "{args.layout}". '
                  f'Valid layouts: {[l.name for l in doc.layouts]}')
            sys.exit(1)

        fig = plt.figure()
        ax = fig.add_axes([0, 0, 1, 1])
        ctx = RenderContext(doc)
        out = MatplotlibBackend(ax, params={"linetype_renderer": args.ltype})

        if args.all_layers_visible:
            for layer_properties in ctx.layers.values():
                layer_properties.is_visible = True

        if args.all_entities_visible:
            class AllVisibleFrontend(Frontend):
                def override_properties(self, entity: "DXFGraphic",
                                        properties: "Properties") -> None:
                    properties.is_visible = True

            frontend = AllVisibleFrontend(ctx, out)
        else:
            frontend = Frontend(ctx, out)

        frontend.draw_layout(layout, finalize=True)

        if args.out is not None:
            print(f'exporting to "{args.out}"')
            fig.savefig(args.out, dpi=args.dpi)
            plt.close(fig)
        else:
            plt.show()

@register
class View(Command):
    """Launcher sub-command: view"""

    NAME = "view"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            View.NAME, help="view DXF files by the PyQt viewer"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="?",
            help="DXF file to view",
        )
        parser.add_argument(
            "--ltype",
            default="internal",
            choices=["internal", "ezdxf"],
            help="select the line type rendering engine, default is internal",
        )
        # disable lineweight at all by default:
        parser.add_argument(
            "--lwscale",
            type=float,
            default=0,
            help="set custom line weight scaling, default is 0 to disable "
                 "line weights at all",
        )

    @staticmethod
    def run(args):
        # Import on demand for a quicker startup:
        try:
            from PyQt5 import QtWidgets
        except ImportError:
            print("PyQt5 package not found.")
            sys.exit(1)
        from ezdxf.addons.drawing.qtviewer import CadViewer

        signal.signal(signal.SIGINT, signal.SIG_DFL)  # handle Ctrl+C properly
        app = QtWidgets.QApplication(sys.argv)
        set_app_icon(app)
        viewer = CadViewer(
            params={
                "linetype_renderer": args.ltype,
                "lineweight_scaling": args.lwscale,
            }
        )
        filename = args.file
        if filename:
            doc, auditor = load_document(filename)
            viewer.set_document(doc, auditor)
            viewer.draw_layout("Model")
        sys.exit(app.exec_())

@register
class Browse(Command):
    """Launcher sub-command: browse"""

    NAME = "browse"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Browse.NAME, help="browse DXF file structure"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="?",
            help="DXF file to browse",
        )
        parser.add_argument(
            "-l", "--line", type=int, required=False, help="go to line number"
        )
        parser.add_argument(
            "-g",
            "--handle",
            required=False,
            help="go to entity by HANDLE, HANDLE has to be a hex value without "
                 "any prefix like 'fefe'",
        )

    @staticmethod
    def run(args):
        try:
            from PyQt5.QtWidgets import QApplication
        except ImportError:
            print("PyQt5 package not found.")
            sys.exit(1)
        from ezdxf.addons import browser

        signal.signal(signal.SIGINT, signal.SIG_DFL)  # handle Ctrl+C properly
        app = QApplication(sys.argv)
        set_app_icon(app)
        main_window = browser.DXFStructureBrowser(
            args.file, line=args.line, handle=args.handle
        )
        main_window.show()
        sys.exit(app.exec_())

@register
class Strip(Command):
    """Launcher sub-command: strip"""

    NAME = "strip"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Strip.NAME, help="strip comments from DXF files"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="+",
            help='DXF file to process, wildcards "*" and "?" are supported',
        )
        parser.add_argument(
            "-b",
            "--backup",
            action="store_true",
            required=False,
            help='make a backup copy with extension ".bak" from the '
                 "DXF file, overwrites existing backup files",
        )
        parser.add_argument(
            "-t",
            "--thumbnail",
            action="store_true",
            required=False,
            help="strip THUMBNAILIMAGE section",
        )
        parser.add_argument(
            "-v",
            "--verbose",
            action="store_true",
            required=False,
            help="give more output",
        )

    @staticmethod
    def run(args):
        from ezdxf.tools.strip import strip

        for pattern in args.file:
            for filename in glob.glob(pattern):
                strip(
                    filename,
                    backup=args.backup,
                    thumbnail=args.thumbnail,
                    verbose=args.verbose,
                )

def set_app_icon(app):
    from PyQt5 import QtGui, QtCore

    app_icon = QtGui.QIcon()
    p = resources_path()
    app_icon.addFile(str(p / "16x16.png"), QtCore.QSize(16, 16))
    app_icon.addFile(str(p / "24x24.png"), QtCore.QSize(24, 24))
    app_icon.addFile(str(p / "32x32.png"), QtCore.QSize(32, 32))
    app_icon.addFile(str(p / "48x48.png"), QtCore.QSize(48, 48))
    app_icon.addFile(str(p / "64x64.png"), QtCore.QSize(64, 64))
    app_icon.addFile(str(p / "256x256.png"), QtCore.QSize(256, 256))
    app.setWindowIcon(app_icon)

def resources_path():
    from pathlib import Path

    return Path(__file__).parent / "resources"
<commit_msg>added layout option to Qt viewer<commit_after>#  Copyright (c) 2021, Manfred Moitzi
#  License: MIT License
from typing import Callable, Optional, Dict
import abc
import sys
import os
import glob
import signal
import logging
from pathlib import Path

from ezdxf import recover
from ezdxf.lldxf import const
from ezdxf.lldxf.validator import is_dxf_file

__all__ = ["get", "add_parsers"]

logger = logging.getLogger("ezdxf")

def get(cmd: str) -> Optional[Callable]:
    cls = _commands.get(cmd)
    if cls:
        return cls.run
    return None

def add_parsers(subparsers) -> None:
    for cmd in _commands.values():  # in order of registration
        cmd.add_parser(subparsers)

class Command:
    """abstract base class for launcher commands"""

    NAME = "command"

    @staticmethod
    @abc.abstractmethod
    def add_parser(subparsers) -> None:
        pass

    @staticmethod
    @abc.abstractmethod
    def run(args) -> None:
        pass

_commands: Dict[str, Command] = dict()

def register(cls: Command):
    """Register a launcher sub-command."""
    _commands[cls.NAME] = cls
    return cls

@register
class PrettyPrint(Command):
    """Launcher sub-command: pp"""

    NAME = "pp"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            PrettyPrint.NAME, help="pretty print DXF files as HTML file"
        )
        parser.add_argument(
            "files",
            metavar="FILE",
            nargs="+",
            help="DXF files pretty print",
        )
        parser.add_argument(
            "-o",
            "--open",
            action="store_true",
            help="open generated HTML file by the default web browser",
        )
        parser.add_argument(
            "-r",
            "--raw",
            action="store_true",
            help="raw mode, no DXF structure interpretation",
        )
        parser.add_argument(
            "-x",
            "--nocompile",
            action="store_true",
            help="don't compile points coordinates into single tags "
                 "(only in raw mode)",
        )
        parser.add_argument(
            "-l",
            "--legacy",
            action="store_true",
            help="legacy mode, reorder DXF point coordinates",
        )
        parser.add_argument(
            "-s",
            "--sections",
            action="store",
            default="hctbeo",
            help="choose sections to include and their order, h=HEADER, c=CLASSES, "
                 "t=TABLES, b=BLOCKS, e=ENTITIES, o=OBJECTS",
        )

    @staticmethod
    def run(args):
        from ezdxf.pp import run

        run(args)

@register
class Audit(Command):
    """Launcher sub-command: audit"""

    NAME = "audit"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Audit.NAME, help="audit and repair DXF files"
        )
        parser.add_argument(
            "files",
            metavar="FILE",
            nargs="+",
            help="audit DXF files",
        )
        parser.add_argument(
            "-s",
            "--save",
            action="store_true",
            help='save recovered files with extension ".rec.dxf" ',
        )

    @staticmethod
    def run(args):
        def build_outname(name: str) -> str:
            p = Path(name)
            return p.parent / (p.stem + ".rec.dxf")

        def log_fixes(auditor):
            for error in auditor.fixes:
                logger.info("fixed:" + error.message)

        def log_errors(auditor):
            for error in auditor.errors:
                logger.error(error.message)

        def _audit(filename: str) -> None:
            msg = f"auditing file: {filename}"
            print(msg)
            logger.info(msg)
            try:
                doc, auditor = recover.readfile(filename)
            except IOError:
                msg = "Not a DXF file or a generic I/O error."
                print(msg)
                logger.error(msg)
                return  # keep on processing additional files
            except const.DXFStructureError:
                msg = "Invalid or corrupted DXF file."
                print(msg)
                logger.error(msg)
                return  # keep on processing additional files

            if auditor.has_errors:
                auditor.print_error_report()
                log_errors(auditor)
            if auditor.has_fixes:
                auditor.print_fixed_errors()
                log_fixes(auditor)

            if auditor.has_errors is False and auditor.has_fixes is False:
                print("No errors found.")
            else:
                print(
                    f"Found {len(auditor.errors)} errors, "
                    f"applied {len(auditor.fixes)} fixes"
                )

            if args.save:
                outname = build_outname(filename)
                doc.saveas(outname)
                print(f"Saved recovered file as: {outname}")

        for pattern in args.files:
            names = list(glob.glob(pattern))
            if len(names) == 0:
                msg = f"File(s) '{pattern}' not found."
                print(msg)
                logger.error(msg)
                continue
            for filename in names:
                if not os.path.exists(filename):
                    msg = f"File '{filename}' not found."
                    print(msg)
                    logger.error(msg)
                    continue
                if not is_dxf_file(filename):
                    msg = f"File '{filename}' is not a DXF file."
                    print(msg)
                    logger.error(msg)
                    continue
                _audit(filename)

def load_document(filename: str):
    try:
        doc, auditor = recover.readfile(filename)
    except IOError:
        msg = f"Not a DXF file or a generic I/O error: {filename}"
        print(msg)
        logger.error(msg)
        sys.exit(2)
    except const.DXFStructureError:
        msg = f"Invalid or corrupted DXF file: {filename}"
        print(msg)
        logger.error(msg)
        sys.exit(3)

    if auditor.has_errors:
        # But is most likely good enough for rendering.
        msg = f"Found {len(auditor.errors)} unrecoverable errors."
        print(msg)
        logger.error(msg)
    if auditor.has_fixes:
        msg = f"Fixed {len(auditor.fixes)} errors."
        print(msg)
        logger.info(msg)
    return doc, auditor

@register
class Draw(Command):
    """Launcher sub-command: draw"""

    NAME = "draw"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Draw.NAME, help="draw and convert DXF files by Matplotlib"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="?",
            help="DXF file to view or convert",
        )
        parser.add_argument(
            "--formats",
            action="store_true",
            help="show all supported export formats and exit",
        )
        parser.add_argument(
            "--layout",
            default="Model",
            help="select the layout to draw",
        )
        parser.add_argument(
            "--all-layers-visible",
            help="draw all layers including the ones marked as invisible",
        )
        parser.add_argument(
            "--all-entities-visible",
            help="draw all entities including the ones marked as invisible "
                 "(some entities are individually marked as invisible even "
                 "if the layer is visible)",
        )
        parser.add_argument(
            "-o", "--out", required=False, help="output filename for export"
        )
        parser.add_argument(
            "--dpi",
            type=int,
            default=300,
            help="target render resolution, default is 300",
        )
        parser.add_argument(
            "--ltype",
            default="internal",
            choices=["internal", "ezdxf"],
            help="select the line type rendering engine, default is internal",
        )

    @staticmethod
    def run(args):
        # Import on demand for a quicker startup:
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            print("Matplotlib package not found.")
            sys.exit(2)

        from ezdxf.addons.drawing import RenderContext, Frontend
        from ezdxf.addons.drawing.matplotlib import MatplotlibBackend

        # Print all supported export formats:
        if args.formats:
            fig = plt.figure()
            for (
                extension,
                description,
            ) in fig.canvas.get_supported_filetypes().items():
                print(f"{extension}: {description}")
            sys.exit(0)

        if args.file:
            filename = args.file
        else:
            print("argument FILE is required")
            sys.exit(1)

        doc, _ = load_document(filename)

        try:
            layout = doc.layouts.get(args.layout)
        except KeyError:
            print(f'Could not find layout "{args.layout}". '
                  f'Valid layouts: {[l.name for l in doc.layouts]}')
            sys.exit(1)

        fig = plt.figure()
        ax = fig.add_axes([0, 0, 1, 1])
        ctx = RenderContext(doc)
        out = MatplotlibBackend(ax, params={"linetype_renderer": args.ltype})

        if args.all_layers_visible:
            for layer_properties in ctx.layers.values():
                layer_properties.is_visible = True

        if args.all_entities_visible:
            class AllVisibleFrontend(Frontend):
                def override_properties(self, entity: "DXFGraphic",
                                        properties: "Properties") -> None:
                    properties.is_visible = True

            frontend = AllVisibleFrontend(ctx, out)
        else:
            frontend = Frontend(ctx, out)

        frontend.draw_layout(layout, finalize=True)

        if args.out is not None:
            print(f'exporting to "{args.out}"')
            fig.savefig(args.out, dpi=args.dpi)
            plt.close(fig)
        else:
            plt.show()

@register
class View(Command):
    """Launcher sub-command: view"""

    NAME = "view"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            View.NAME, help="view DXF files by the PyQt viewer"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="?",
            help="DXF file to view",
        )
        parser.add_argument(
            "--layout",
            default="Model",
            help="select the layout to draw",
        )
        parser.add_argument(
            "--ltype",
            default="internal",
            choices=["internal", "ezdxf"],
            help="select the line type rendering engine, default is internal",
        )
        # disable lineweight at all by default:
        parser.add_argument(
            "--lwscale",
            type=float,
            default=0,
            help="set custom line weight scaling, default is 0 to disable "
                 "line weights at all",
        )

    @staticmethod
    def run(args):
        # Import on demand for a quicker startup:
        try:
            from PyQt5 import QtWidgets
        except ImportError:
            print("PyQt5 package not found.")
            sys.exit(1)
        from ezdxf.addons.drawing.qtviewer import CadViewer

        signal.signal(signal.SIGINT, signal.SIG_DFL)  # handle Ctrl+C properly
        app = QtWidgets.QApplication(sys.argv)
        set_app_icon(app)
        viewer = CadViewer(
            params={
                "linetype_renderer": args.ltype,
                "lineweight_scaling": args.lwscale,
            }
        )
        filename = args.file
        if filename:
            doc, auditor = load_document(filename)
            viewer.set_document(doc, auditor)
            viewer.draw_layout(args.layout)
        sys.exit(app.exec_())

@register
class Browse(Command):
    """Launcher sub-command: browse"""

    NAME = "browse"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Browse.NAME, help="browse DXF file structure"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="?",
            help="DXF file to browse",
        )
        parser.add_argument(
            "-l", "--line", type=int, required=False, help="go to line number"
        )
        parser.add_argument(
            "-g",
            "--handle",
            required=False,
            help="go to entity by HANDLE, HANDLE has to be a hex value without "
                 "any prefix like 'fefe'",
        )

    @staticmethod
    def run(args):
        try:
            from PyQt5.QtWidgets import QApplication
        except ImportError:
            print("PyQt5 package not found.")
            sys.exit(1)
        from ezdxf.addons import browser

        signal.signal(signal.SIGINT, signal.SIG_DFL)  # handle Ctrl+C properly
        app = QApplication(sys.argv)
        set_app_icon(app)
        main_window = browser.DXFStructureBrowser(
            args.file, line=args.line, handle=args.handle
        )
        main_window.show()
        sys.exit(app.exec_())

@register
class Strip(Command):
    """Launcher sub-command: strip"""

    NAME = "strip"

    @staticmethod
    def add_parser(subparsers):
        parser = subparsers.add_parser(
            Strip.NAME, help="strip comments from DXF files"
        )
        parser.add_argument(
            "file",
            metavar="FILE",
            nargs="+",
            help='DXF file to process, wildcards "*" and "?" are supported',
        )
        parser.add_argument(
            "-b",
            "--backup",
            action="store_true",
            required=False,
            help='make a backup copy with extension ".bak" from the '
                 "DXF file, overwrites existing backup files",
        )
        parser.add_argument(
            "-t",
            "--thumbnail",
            action="store_true",
            required=False,
            help="strip THUMBNAILIMAGE section",
        )
        parser.add_argument(
            "-v",
            "--verbose",
            action="store_true",
            required=False,
            help="give more output",
        )

    @staticmethod
    def run(args):
        from ezdxf.tools.strip import strip

        for pattern in args.file:
            for filename in glob.glob(pattern):
                strip(
                    filename,
                    backup=args.backup,
                    thumbnail=args.thumbnail,
                    verbose=args.verbose,
                )

def set_app_icon(app):
    from PyQt5 import QtGui, QtCore

    app_icon = QtGui.QIcon()
    p = resources_path()
    app_icon.addFile(str(p / "16x16.png"), QtCore.QSize(16, 16))
    app_icon.addFile(str(p / "24x24.png"), QtCore.QSize(24, 24))
    app_icon.addFile(str(p / "32x32.png"), QtCore.QSize(32, 32))
    app_icon.addFile(str(p / "48x48.png"), QtCore.QSize(48, 48))
    app_icon.addFile(str(p / "64x64.png"), QtCore.QSize(64, 64))
    app_icon.addFile(str(p / "256x256.png"), QtCore.QSize(256, 256))
    app.setWindowIcon(app_icon)

def resources_path():
    from pathlib import Path

    return Path(__file__).parent / "resources"
<|endoftext|>


<commit_before>from time import time, sleep
from copy import deepcopy
import threading
import sqlite3
import json
import os
import decimal
from bitcoin.core import COIN, b2lx
import feemodel.config
from feemodel.config import config, historyFile
from feemodel.util import proxy, logWrite, StoppableThread

pollPeriod = config['pollPeriod']
keepHistory = config['keepHistory']
historyLock = threading.Lock()

class TxMempool(StoppableThread):
    # Have to handle RPC errors
    # Writehistory means write to db the mempool state at each block.
    # We keep <keepHistory> number of past blocks.

    def update(self):
        currHeight, mapTxNew = proxy.pollMempool()
        if currHeight > self.bestSeenBlock:
            threading.Thread(target=self.processBlocks,
                args=(range(self.bestSeenBlock+1,currHeight+1),
                    self.mapTx, deepcopy(mapTxNew)),
                name='mempool-processBlocks').start()
            self.mapTx = mapTxNew
            self.bestSeenBlock = currHeight
            return True
            # Be careful here: may have to pass deepcopy of self.mapTx to processBlocks,
            # if we are going to return self.mapTx for other functions to use.
        else:
            self.mapTx = mapTxNew
            return False

    def run(self):
        feemodel.config.apprun = True
        logWrite("Starting mempool")
        self.bestSeenBlock, self.mapTx = proxy.pollMempool()
        while not self.isStopped():
            self.update()
            self.sleep(pollPeriod)
        logWrite("Closing up mempool...")
        for thread in threading.enumerate():
            if thread.name.startswith('mempool'):
                thread.join()
        logWrite("Finished everything.")

    @staticmethod
    def processBlocks(blockHeightRange, currPool, newPool, blockTime=None):
        with historyLock:
            if not blockTime:
                blockTime = int(time())
            blocks = []
            for blockHeight in blockHeightRange:
                blockData = proxy.getblock(proxy.getblockhash(blockHeight))
                blockSize = len(blockData.serialize())
                blockTxList = [b2lx(tx.GetHash()) for tx in blockData.vtx]
                entries = deepcopy(currPool)
                numMempoolTxsInBlock = 0

                for txid, entry in entries.iteritems():
                    if txid in blockTxList:
                        entry['inBlock'] = True
                        numMempoolTxsInBlock += 1
                        del currPool[txid]
                    else:
                        entry['inBlock'] = False
                    entry['leadTime'] = blockTime - entry['time']
                    entry['feeRate'] = int(entry['fee']*COIN) * 1000 // entry['size']
                    entry['isConflict'] = False

                blocks.append(Block(entries,blockHeight,blockSize,blockTime))
                logWrite(str(numMempoolTxsInBlock) + ' of ' +
                    str(len(blockTxList)-1) + ' in block ' + str(blockHeight))

            # To-do: insert warnings if block inclusion ratio is too low, or conflicts are too high
            conflicts = set(currPool) - set(newPool)

            numConflicts = 0
            for block in blocks:
                numConflicts += block.removeConflicts(conflicts)

            logWrite("%d conflicts removed." % numConflicts)

            if feemodel.config.apprun:
                for block in blocks:
                    block.writeHistory()

            return blocks

class Block(object):
    def __init__(self, entries, blockHeight, blockSize, blockTime):
        # To-do: add a 'tx-coverage' field
        self.entries = entries
        self.height = blockHeight
        self.size = blockSize
        self.time = blockTime

    def removeConflicts(self, conflicts):
        numConflicts = 0
        for txid in self.entries.keys():
            if txid in conflicts:
                self.entries[txid]['isConflict'] = True
                numConflicts += 1

        return numConflicts

    def writeHistory(self, dbFile=historyFile):
        db = None
        dbExists = os.path.exists(dbFile)
        try:
            db = sqlite3.connect(dbFile)
            if not dbExists:
                with db:
                    db.execute('CREATE TABLE blocks (height INTEGER UNIQUE, size INTEGER, time REAL)')
                    db.execute('CREATE TABLE txs (blockheight INTEGER, txid TEXT, data TEXT)')
            with db:
                db.execute('INSERT INTO blocks VALUES (?,?,?)', (self.height, self.size, self.time))
                db.executemany('INSERT INTO txs VALUES (?,?,?)',
                    [(self.height, txid, json.dumps(entry,default=decimalDefault))
                    for txid,entry in self.entries.iteritems()])
                historyLimit = self.height - keepHistory
                if keepHistory > 0:
                    db.execute('DELETE FROM blocks WHERE height<=?', (historyLimit,))
                    db.execute('DELETE FROM txs WHERE blockheight<=?', (historyLimit,))
        except Exception as e:
            logWrite(repr(e))
            logWrite("Exception in writing/cleaning history.")
        finally:
            if db:
                db.close()

    @classmethod
    def blockFromHistory(cls, blockHeight, dbFile=historyFile):
        db = None
        try:
            db = sqlite3.connect(dbFile)
            block = db.execute('SELECT size,time FROM blocks WHERE height=?',
                (blockHeight,)).fetchall()
            if block:
                blockSize,blockTime = block[0]
            else:
                return None
            txlist = db.execute('SELECT txid,data FROM txs WHERE blockheight=?', (blockHeight,))
            entries = {txid: json.loads(str(data)) for txid,data in txlist}
            for entry in entries.itervalues():
                entry['fee'] = decimal.Decimal(entry['fee'])
                entry['startingpriority'] = decimal.Decimal(entry['startingpriority'])
                entry['currentpriority'] = decimal.Decimal(entry['currentpriority'])
            return cls(entries,blockHeight,blockSize,blockTime)
        except Exception as e:
            logWrite(repr(e))
            return None
        finally:
            if db:
                db.close()

    def __eq__(self,other):
        if not isinstance(other,Block):
            return False
        return all([
            self.entries == other.entries,
            self.height == other.height,
            self.size == other.size,
            self.time == other.time,
        ])

class LoadHistory(object):
    def __init__(self, dbFile=historyFile):
        self.fns = []
        self.dbFile = dbFile

    def registerFn(self, fn, blockHeightRange):
        # blockHeightRange tuple (start,end) includes start but not end, to adhere to range() convention
        self.fns.append((fn, blockHeightRange))

    def loadBlocks(self):
        startHeight = min([f[1][0] for f in self.fns])
        endHeight = max([f[1][1] for f in self.fns])

        for height in range(startHeight, endHeight):
            block = Block.blockFromHistory(height, self.dbFile)
            for fn, blockHeightRange in self.fns:
                if height >= blockHeightRange[0] and height < blockHeightRange[1]:
                    fn([block])

def decimalDefault(obj):
    if isinstance(obj, decimal.Decimal):
        return str(obj)
    raise TypeError

<commit_msg>sql index into history; separate insert from delete<commit_after>from time import time, sleep
from copy import deepcopy
import threading
import sqlite3
import json
import os
import decimal
from bitcoin.core import COIN, b2lx
import feemodel.config
from feemodel.config import config, historyFile
from feemodel.util import proxy, logWrite, StoppableThread

pollPeriod = config['pollPeriod']
keepHistory = config['keepHistory']
historyLock = threading.Lock()

class TxMempool(StoppableThread):
    # Have to handle RPC errors
    # Writehistory means write to db the mempool state at each block.
    # We keep <keepHistory> number of past blocks.

    def update(self):
        currHeight, mapTxNew = proxy.pollMempool()
        if currHeight > self.bestSeenBlock:
            threading.Thread(target=self.processBlocks,
                args=(range(self.bestSeenBlock+1,currHeight+1),
                    self.mapTx, deepcopy(mapTxNew)),
                name='mempool-processBlocks').start()
            self.mapTx = mapTxNew
            self.bestSeenBlock = currHeight
            return True
            # Be careful here: may have to pass deepcopy of self.mapTx to processBlocks,
            # if we are going to return self.mapTx for other functions to use.
        else:
            self.mapTx = mapTxNew
            return False

    def run(self):
        feemodel.config.apprun = True
        logWrite("Starting mempool")
        self.bestSeenBlock, self.mapTx = proxy.pollMempool()
        while not self.isStopped():
            self.update()
            self.sleep(pollPeriod)
        logWrite("Closing up mempool...")
        for thread in threading.enumerate():
            if thread.name.startswith('mempool'):
                thread.join()
        logWrite("Finished everything.")

    @staticmethod
    def processBlocks(blockHeightRange, currPool, newPool, blockTime=None):
        with historyLock:
            if not blockTime:
                blockTime = int(time())
            blocks = []
            for blockHeight in blockHeightRange:
                blockData = proxy.getblock(proxy.getblockhash(blockHeight))
                blockSize = len(blockData.serialize())
                blockTxList = [b2lx(tx.GetHash()) for tx in blockData.vtx]
                entries = deepcopy(currPool)
                numMempoolTxsInBlock = 0

                for txid, entry in entries.iteritems():
                    if txid in blockTxList:
                        entry['inBlock'] = True
                        numMempoolTxsInBlock += 1
                        del currPool[txid]
                    else:
                        entry['inBlock'] = False
                    entry['leadTime'] = blockTime - entry['time']
                    entry['feeRate'] = int(entry['fee']*COIN) * 1000 // entry['size']
                    entry['isConflict'] = False

                blocks.append(Block(entries,blockHeight,blockSize,blockTime))
                logWrite('%d of %d in block %d' % (
                    numMempoolTxsInBlock, len(blockTxList)-1, blockHeight))

            # To-do: insert warnings if block inclusion ratio is too low, or conflicts are too high
            conflicts = set(currPool) - set(newPool)

            numConflicts = 0
            for block in blocks:
                numConflicts += block.removeConflicts(conflicts)

            logWrite("%d conflicts removed." % numConflicts)

            if feemodel.config.apprun:
                for block in blocks:
                    block.writeHistory()

            return blocks

class Block(object):
    def __init__(self, entries, blockHeight, blockSize, blockTime):
        # To-do: add a 'tx-coverage' field
        self.entries = entries
        self.height = blockHeight
        self.size = blockSize
        self.time = blockTime

    def removeConflicts(self, conflicts):
        numConflicts = 0
        for txid in self.entries.keys():
            if txid in conflicts:
                self.entries[txid]['isConflict'] = True
                numConflicts += 1

        return numConflicts

    def writeHistory(self, dbFile=historyFile):
        db = None
        dbExists = os.path.exists(dbFile)
        try:
            db = sqlite3.connect(dbFile)
            if not dbExists:
                with db:
                    db.execute('CREATE TABLE blocks (height INTEGER UNIQUE, size INTEGER, time REAL)')
                    db.execute('CREATE TABLE txs (blockheight INTEGER, txid TEXT, data TEXT)')
            db.execute('CREATE INDEX IF NOT EXISTS heightidx ON txs (blockheight)')
            with db:
                db.execute('INSERT INTO blocks VALUES (?,?,?)', (self.height, self.size, self.time))
                db.executemany('INSERT INTO txs VALUES (?,?,?)',
                    [(self.height, txid, json.dumps(entry,default=decimalDefault))
                    for txid,entry in self.entries.iteritems()])
            historyLimit = self.height - keepHistory
            if keepHistory > 0:
                with db:
                    db.execute('DELETE FROM blocks WHERE height<=?', (historyLimit,))
                    db.execute('DELETE FROM txs WHERE blockheight<=?', (historyLimit,))
        except Exception as e:
            logWrite(repr(e))
            logWrite("Exception in writing/cleaning history.")
        finally:
            if db:
                db.close()

    @classmethod
    def blockFromHistory(cls, blockHeight, dbFile=historyFile):
        db = None
        try:
            db = sqlite3.connect(dbFile)
            block = db.execute('SELECT size,time FROM blocks WHERE height=?',
                (blockHeight,)).fetchall()
            if block:
                blockSize,blockTime = block[0]
            else:
                return None
            txlist = db.execute('SELECT txid,data FROM txs WHERE blockheight=?', (blockHeight,))
            entries = {txid: json.loads(str(data)) for txid,data in txlist}
            for entry in entries.itervalues():
                entry['fee'] = decimal.Decimal(entry['fee'])
                entry['startingpriority'] = decimal.Decimal(entry['startingpriority'])
                entry['currentpriority'] = decimal.Decimal(entry['currentpriority'])
            return cls(entries,blockHeight,blockSize,blockTime)
        except Exception as e:
            logWrite(repr(e))
            return None
        finally:
            if db:
                db.close()

    def __eq__(self,other):
        if not isinstance(other,Block):
            return False
        return all([
            self.entries == other.entries,
            self.height == other.height,
            self.size == other.size,
            self.time == other.time,
        ])

class LoadHistory(object):
    def __init__(self, dbFile=historyFile):
        self.fns = []
        self.dbFile = dbFile

    def registerFn(self, fn, blockHeightRange):
        # blockHeightRange tuple (start,end) includes start but not end, to adhere to range() convention
        self.fns.append((fn, blockHeightRange))

    def loadBlocks(self):
        startHeight = min([f[1][0] for f in self.fns])
        endHeight = max([f[1][1] for f in self.fns])

        for height in range(startHeight, endHeight):
            block = Block.blockFromHistory(height, self.dbFile)
            for fn, blockHeightRange in self.fns:
                if height >= blockHeightRange[0] and height < blockHeightRange[1]:
                    fn([block])

def decimalDefault(obj):
    if isinstance(obj, decimal.Decimal):
        return str(obj)
    raise TypeError

<|endoftext|>


<commit_before>### $Id: slices.py 15842 2009-11-22 09:56:13Z anil $
### $URL: https://svn.planet-lab.org/svn/sfa/trunk/sfa/plc/slices.py $

import datetime
import time
import traceback
import sys
from copy import deepcopy
from lxml import etree
from StringIO import StringIO
from types import StringTypes
from sfa.util.rspecHelper import merge_rspecs
from sfa.util.namespace import *
from sfa.util.rspec import *
from sfa.util.specdict import *
from sfa.util.faults import *
from sfa.util.record import SfaRecord
from sfa.util.policy import Policy
from sfa.util.prefixTree import prefixTree
from sfa.util.sfaticket import *
from sfa.trust.credential import Credential
from sfa.util.threadmanager import ThreadManager
import sfa.util.xmlrpcprotocol as xmlrpcprotocol
from sfa.util.debug import log
import sfa.plc.peers as peers
from copy import copy

def get_version():
    version = {}
    version['geni_api'] = 1
    version['sfa'] = 1
    return version

def slice_status(api, slice_xrn, creds ):
    result = {}
    result['geni_urn'] = slice_xrn
    result['geni_status'] = 'unknown'
    result['geni_resources'] = {}
    return result

def create_slice(api, xrn, creds, rspec, users):
    hrn, type = urn_to_hrn(xrn)

    # Validate the RSpec against PlanetLab's schema --disabled for now
    # The schema used here needs to aggregate the PL and VINI schemas
    # schema = "/var/www/html/schemas/pl.rng"
    schema = None
    if schema:
        try:
            tree = etree.parse(StringIO(rspec))
        except etree.XMLSyntaxError:
            message = str(sys.exc_info()[1])
            raise InvalidRSpec(message)

        relaxng_doc = etree.parse(schema)
        relaxng = etree.RelaxNG(relaxng_doc)

        if not relaxng(tree):
            error = relaxng.error_log.last_error
            message = "%s (line %s)" % (error.message, error.line)
            raise InvalidRSpec(message)

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        # Just send entire RSpec to each aggregate
        server = api.aggregates[aggregate]
        threads.run(server.CreateSliver, xrn, cred, rspec, users)

    results = threads.get_results()
    merged_rspec = merge_rspecs(results)
    return merged_rspec

def renew_slice(api, xrn, creds, expiration_time):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.RenewSliver, xrn, credential, expiration_time)
    threads.get_results()
    return 1

def get_ticket(api, xrn, creds, rspec, users):
    slice_hrn, type = urn_to_hrn(xrn)
    # get the netspecs contained within the clients rspec
    aggregate_rspecs = {}
    tree= etree.parse(StringIO(rspec))
    elements = tree.findall('./network')
    for element in elements:
        aggregate_hrn = element.values()[0]
        aggregate_rspecs[aggregate_hrn] = rspec

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate, aggregate_rspec in aggregate_rspecs.items():
        server = None
        if aggregate in api.aggregates:
            server = api.aggregates[aggregate]
        else:
            net_urn = hrn_to_urn(aggregate, 'authority')
            # we may have a peer that knows about this aggregate
            for agg in api.aggregates:
                target_aggs = api.aggregates[agg].get_aggregates(credential, net_urn)
                if not target_aggs or not 'hrn' in target_aggs[0]:
                    continue
                # send the request to this address
                url = target_aggs[0]['url']
                server = xmlrpcprotocol.get_server(url, api.key_file, api.cert_file)
                # aggregate found, no need to keep looping
                break
        if server is None:
            continue
        threads.run(server.GetTicket, xrn, credential, aggregate_rspec, users)

    results = threads.get_results()

    # gather information from each ticket
    rspecs = []
    initscripts = []
    slivers = []
    object_gid = None
    for result in results:
        agg_ticket = SfaTicket(string=result)
        attrs = agg_ticket.get_attributes()
        if not object_gid:
            object_gid = agg_ticket.get_gid_object()
        rspecs.append(agg_ticket.get_rspec())
        initscripts.extend(attrs.get('initscripts', []))
        slivers.extend(attrs.get('slivers', []))

    # merge info
    attributes = {'initscripts': initscripts,
                 'slivers': slivers}
    merged_rspec = merge_rspecs(rspecs)

    # create a new ticket
    ticket = SfaTicket(subject = slice_hrn)
    ticket.set_gid_caller(api.auth.client_gid)
    ticket.set_issuer(key=api.key, subject=api.hrn)
    ticket.set_gid_object(object_gid)
    ticket.set_pubkey(object_gid.get_pubkey())
    #new_ticket.set_parent(api.auth.hierarchy.get_auth_ticket(auth_hrn))
    ticket.set_attributes(attributes)
    ticket.set_rspec(merged_rspec)
    ticket.encode()
    ticket.sign()
    return ticket.save_to_string(save_parents=True)

def delete_slice(api, xrn, origin_hrn=None):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.DeleteSliver, xrn, credential)
    threads.get_results()
    return 1

def start_slice(api, xrn, creds):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.Start, xrn, credential)
    threads.get_results()
    return 1

def stop_slice(api, xrn, creds):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.Stop, xrn, credential)
    threads.get_results()
    return 1

def reset_slice(api, xrn):
    """
    Not implemented
    """
    return 1

def shutdown(api, xrn, creds):
    """
    Not implemented
    """
    return 1

def status(api, xrn, creds):
    """
    Not implemented
    """
    return 1

def get_slices(api, creds):
    # look in cache first
    if api.cache:
        slices = api.cache.get('slices')
        if slices:
            return slices

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    # fetch from aggregates
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.ListSlices, credential)

    # combime results
    results = threads.get_results()
    slices = []
    for result in results:
        slices.extend(result)

    # cache the result
    if api.cache:
        api.cache.add('slices', slices)

    return slices

def get_rspec(api, creds, options):
    # get slice's hrn from options
    xrn = options.get('geni_slice_urn', None)
    hrn, type = urn_to_hrn(xrn)

    # get hrn of the original caller
    origin_hrn = options.get('origin_hrn', None)
    if not origin_hrn:
        origin_hrn = Credential(string=creds[0]).get_gid_caller().get_hrn()

    # look in cache first
    if api.cache and not xrn:
        rspec =  api.cache.get('nodes')
        if rspec:
            return rspec

    hrn, type = urn_to_hrn(xrn)
    rspec = None

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        # get the rspec from the aggregate
        server = api.aggregates[aggregate]
        my_opts = copy(options)
        my_opts['geni_compressed'] = False
        threads.run(server.ListResources, cred, my_opts)
        #threads.run(server.get_resources, cred, xrn, origin_hrn)

    results = threads.get_results()
    # combine the rspecs into a single rspec
    for agg_rspec in results:
        try:
            tree = etree.parse(StringIO(agg_rspec))
        except etree.XMLSyntaxError:
            message = str(agg_rspec) + ": " + str(sys.exc_info()[1])
            raise InvalidRSpec(message)

        root = tree.getroot()
        if root.get("type") in ["SFA"]:
            if rspec == None:
                rspec = root
            else:
                for network in root.iterfind("./network"):
                    rspec.append(deepcopy(network))
                for request in root.iterfind("./request"):
                    rspec.append(deepcopy(request))

    rspec =  etree.tostring(rspec, xml_declaration=True, pretty_print=True)
    # cache the result
    if api.cache and not xrn:
        api.cache.add('nodes', rspec)

    return rspec

def main():
    r = RSpec()
    r.parseFile(sys.argv[1])
    rspec = r.toDict()
    create_slice(None,'plc.princeton.tmacktestslice',rspec)

if __name__ == "__main__":
    main()

<commit_msg>fix NameError (cred should be credential)<commit_after>### $Id: slices.py 15842 2009-11-22 09:56:13Z anil $
### $URL: https://svn.planet-lab.org/svn/sfa/trunk/sfa/plc/slices.py $

import datetime
import time
import traceback
import sys
from copy import deepcopy
from lxml import etree
from StringIO import StringIO
from types import StringTypes
from sfa.util.rspecHelper import merge_rspecs
from sfa.util.namespace import *
from sfa.util.rspec import *
from sfa.util.specdict import *
from sfa.util.faults import *
from sfa.util.record import SfaRecord
from sfa.util.policy import Policy
from sfa.util.prefixTree import prefixTree
from sfa.util.sfaticket import *
from sfa.trust.credential import Credential
from sfa.util.threadmanager import ThreadManager
import sfa.util.xmlrpcprotocol as xmlrpcprotocol
from sfa.util.debug import log
import sfa.plc.peers as peers
from copy import copy

def get_version():
    version = {}
    version['geni_api'] = 1
    version['sfa'] = 1
    return version

def slice_status(api, slice_xrn, creds ):
    result = {}
    result['geni_urn'] = slice_xrn
    result['geni_status'] = 'unknown'
    result['geni_resources'] = {}
    return result

def create_slice(api, xrn, creds, rspec, users):
    hrn, type = urn_to_hrn(xrn)

    # Validate the RSpec against PlanetLab's schema --disabled for now
    # The schema used here needs to aggregate the PL and VINI schemas
    # schema = "/var/www/html/schemas/pl.rng"
    schema = None
    if schema:
        try:
            tree = etree.parse(StringIO(rspec))
        except etree.XMLSyntaxError:
            message = str(sys.exc_info()[1])
            raise InvalidRSpec(message)

        relaxng_doc = etree.parse(schema)
        relaxng = etree.RelaxNG(relaxng_doc)

        if not relaxng(tree):
            error = relaxng.error_log.last_error
            message = "%s (line %s)" % (error.message, error.line)
            raise InvalidRSpec(message)

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        # Just send entire RSpec to each aggregate
        server = api.aggregates[aggregate]
        threads.run(server.CreateSliver, xrn, cred, rspec, users)

    results = threads.get_results()
    merged_rspec = merge_rspecs(results)
    return merged_rspec

def renew_slice(api, xrn, creds, expiration_time):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.RenewSliver, xrn, credential, expiration_time)
    threads.get_results()
    return 1

def get_ticket(api, xrn, creds, rspec, users):
    slice_hrn, type = urn_to_hrn(xrn)
    # get the netspecs contained within the clients rspec
    aggregate_rspecs = {}
    tree= etree.parse(StringIO(rspec))
    elements = tree.findall('./network')
    for element in elements:
        aggregate_hrn = element.values()[0]
        aggregate_rspecs[aggregate_hrn] = rspec

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate, aggregate_rspec in aggregate_rspecs.items():
        server = None
        if aggregate in api.aggregates:
            server = api.aggregates[aggregate]
        else:
            net_urn = hrn_to_urn(aggregate, 'authority')
            # we may have a peer that knows about this aggregate
            for agg in api.aggregates:
                target_aggs = api.aggregates[agg].get_aggregates(credential, net_urn)
                if not target_aggs or not 'hrn' in target_aggs[0]:
                    continue
                # send the request to this address
                url = target_aggs[0]['url']
                server = xmlrpcprotocol.get_server(url, api.key_file, api.cert_file)
                # aggregate found, no need to keep looping
                break
        if server is None:
            continue
        threads.run(server.GetTicket, xrn, credential, aggregate_rspec, users)

    results = threads.get_results()

    # gather information from each ticket
    rspecs = []
    initscripts = []
    slivers = []
    object_gid = None
    for result in results:
        agg_ticket = SfaTicket(string=result)
        attrs = agg_ticket.get_attributes()
        if not object_gid:
            object_gid = agg_ticket.get_gid_object()
        rspecs.append(agg_ticket.get_rspec())
        initscripts.extend(attrs.get('initscripts', []))
        slivers.extend(attrs.get('slivers', []))

    # merge info
    attributes = {'initscripts': initscripts,
                 'slivers': slivers}
    merged_rspec = merge_rspecs(rspecs)

    # create a new ticket
    ticket = SfaTicket(subject = slice_hrn)
    ticket.set_gid_caller(api.auth.client_gid)
    ticket.set_issuer(key=api.key, subject=api.hrn)
    ticket.set_gid_object(object_gid)
    ticket.set_pubkey(object_gid.get_pubkey())
    #new_ticket.set_parent(api.auth.hierarchy.get_auth_ticket(auth_hrn))
    ticket.set_attributes(attributes)
    ticket.set_rspec(merged_rspec)
    ticket.encode()
    ticket.sign()
    return ticket.save_to_string(save_parents=True)

def delete_slice(api, xrn, origin_hrn=None):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.DeleteSliver, xrn, credential)
    threads.get_results()
    return 1

def start_slice(api, xrn, creds):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.Start, xrn, credential)
    threads.get_results()
    return 1

def stop_slice(api, xrn, creds):
    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.Stop, xrn, credential)
    threads.get_results()
    return 1

def reset_slice(api, xrn):
    """
    Not implemented
    """
    return 1

def shutdown(api, xrn, creds):
    """
    Not implemented
    """
    return 1

def status(api, xrn, creds):
    """
    Not implemented
    """
    return 1

def get_slices(api, creds):
    # look in cache first
    if api.cache:
        slices = api.cache.get('slices')
        if slices:
            return slices

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    # fetch from aggregates
    for aggregate in api.aggregates:
        server = api.aggregates[aggregate]
        threads.run(server.ListSlices, credential)

    # combime results
    results = threads.get_results()
    slices = []
    for result in results:
        slices.extend(result)

    # cache the result
    if api.cache:
        api.cache.add('slices', slices)

    return slices

def get_rspec(api, creds, options):
    # get slice's hrn from options
    xrn = options.get('geni_slice_urn', None)
    hrn, type = urn_to_hrn(xrn)

    # get hrn of the original caller
    origin_hrn = options.get('origin_hrn', None)
    if not origin_hrn:
        origin_hrn = Credential(string=creds[0]).get_gid_caller().get_hrn()

    # look in cache first
    if api.cache and not xrn:
        rspec =  api.cache.get('nodes')
        if rspec:
            return rspec

    hrn, type = urn_to_hrn(xrn)
    rspec = None

    # attempt to use delegated credential first
    credential = api.getDelegatedCredential(creds)
    if not credential:
        credential = api.getCredential()
    threads = ThreadManager()
    for aggregate in api.aggregates:
        # get the rspec from the aggregate
        server = api.aggregates[aggregate]
        my_opts = copy(options)
        my_opts['geni_compressed'] = False
        threads.run(server.ListResources, credential, my_opts)
        #threads.run(server.get_resources, cred, xrn, origin_hrn)

    results = threads.get_results()
    # combine the rspecs into a single rspec
    for agg_rspec in results:
        try:
            tree = etree.parse(StringIO(agg_rspec))
        except etree.XMLSyntaxError:
            message = str(agg_rspec) + ": " + str(sys.exc_info()[1])
            raise InvalidRSpec(message)

        root = tree.getroot()
        if root.get("type") in ["SFA"]:
            if rspec == None:
                rspec = root
            else:
                for network in root.iterfind("./network"):
                    rspec.append(deepcopy(network))
                for request in root.iterfind("./request"):
                    rspec.append(deepcopy(request))

    rspec =  etree.tostring(rspec, xml_declaration=True, pretty_print=True)
    # cache the result
    if api.cache and not xrn:
        api.cache.add('nodes', rspec)

    return rspec

def main():
    r = RSpec()
    r.parseFile(sys.argv[1])
    rspec = r.toDict()
    create_slice(None,'plc.princeton.tmacktestslice',rspec)

if __name__ == "__main__":
    main()

<|endoftext|>