<commit_before>from django import forms
from django.forms.models import modelformset_factory
from django.utils.safestring import mark_safe
from organizations.models import Organization
from organizations.utils import create_organization

from threebot.utils import get_possible_parameters, get_possible_worker, get_curr_org
from threebot.utils import get_preset_param, get_preset_worker, get_possible_owners, get_possible_lists, get_preset_list
from threebot.models import Worker
from threebot.models import Workflow
from threebot.models import Task
from threebot.models import UserParameter, OrganizationParameter, ParameterList
from threebot.utils import order_workflow_tasks

class UserParameterCreateForm(forms.ModelForm):
    next = forms.CharField(required=False, widget=forms.HiddenInput())

    class Meta:
        model = UserParameter
        fields = ['data_type', 'name', 'value', 'owner', ]
        widgets = {
            'data_type': forms.Select(attrs={'class': 'form-control', }),
            'name': forms.TextInput(attrs={'class': 'form-control', }),
            'value': forms.TextInput(attrs={'class': 'form-control', }),
            'owner': forms.HiddenInput(),
        }

    def __init__(self, *args, **kwargs):
        user = kwargs.pop('user')
        super(UserParameterCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].initial = user

class UserParameterChangeForm(UserParameterCreateForm):
    class Meta(UserParameterCreateForm.Meta):
        exclude = ('owner',)

class OrganizationParameterCreateForm(forms.ModelForm):
    class Meta:
        model = OrganizationParameter
        fields = ['data_type', 'name', 'value', 'owner', ]
        widgets = {
            'data_type': forms.Select(attrs={'class': 'form-control', }),
            'name': forms.TextInput(attrs={'class': 'form-control', }),
            'value': forms.TextInput(attrs={'class': 'form-control', }),
            'owner': forms.HiddenInput(),
        }

    def __init__(self, *args, **kwargs):
        org = kwargs.pop('org')
        super(OrganizationParameterCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].initial = org

class OrganizationParameterChangeForm(OrganizationParameterCreateForm):
    class Meta(OrganizationParameterCreateForm.Meta):
        exclude = ('owner',)

def make_organization_parameter_formset(org, extra=3):
    """
    This is a workaround for passing custom parameters (the parameter owner in our case)
    to a ModelFormset and is based on this stackoverflow answer:
    http://stackoverflow.com/a/1376616
    """
    class _OrganizationParameterCreateForm(forms.ModelForm):
        remove_from_list = forms.BooleanField(required=False)

        class Meta:
            model = OrganizationParameter
            fields = ['data_type', 'name', 'value', 'owner', ]
            widgets = {
                'data_type': forms.Select(attrs={'class': 'form-control', }),
                'name': forms.TextInput(attrs={'class': 'form-control', }),
                'value': forms.TextInput(attrs={'class': 'form-control', }),
                'owner': forms.HiddenInput(),
            }

        def __init__(self, *args, **kwargs):
            # self.org = org
            super(_OrganizationParameterCreateForm, self).__init__(*args, **kwargs)
            self.fields['owner'].initial = org

    return modelformset_factory(OrganizationParameter, form=_OrganizationParameterCreateForm, extra=extra, can_delete=True)

def make_user_parameter_formset(user, extra=3):
    """
    This is a workaround for passing custom parameters (the parameter owner in our case)
    to a ModelFormset and is based on this stackoverflow answer:
    http://stackoverflow.com/a/1376616
    """
    class _UserParameterCreateForm(forms.ModelForm):
        class Meta:
            model = UserParameter
            fields = ['data_type', 'name', 'value', 'owner', ]
            widgets = {
                'data_type': forms.Select(attrs={'class': 'form-control', }),
                'name': forms.TextInput(attrs={'class': 'form-control', }),
                'value': forms.TextInput(attrs={'class': 'form-control', }),
                'owner': forms.HiddenInput(),
            }

        def __init__(self, *args, **kwargs):
            # user = kwargs.pop('user')
            super(_UserParameterCreateForm, self).__init__(*args, **kwargs)
            self.fields['owner'].initial = user

    return modelformset_factory(UserParameter, form=_UserParameterCreateForm, extra=extra, can_delete=True)

class ParameterListSelectForm(forms.Form):
    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        workflow = kwargs.pop('workflow')
        self.workflow = workflow

        super(ParameterListSelectForm, self).__init__(*args, **kwargs)

        possible_lists = [('', '--Choose a Parameter List--')]
        possible_lists += get_possible_lists(request, workflow)
        preset_list_id = get_preset_list(request, workflow, id=True)

        self.fields['parameter_list'] = forms.ChoiceField(
            label="Parameter List",
            choices=possible_lists,
            initial=preset_list_id,
            widget=forms.Select(attrs={'class': 'form-control', }),
        )
        self.fields['parameter_list'].empty_label = None

    def clean_parameter_list(self, *args, **kwargs):
        list_id = self.cleaned_data['parameter_list']
        parameter_list = ParameterList.objects.get(id=list_id)
        workflow_tasks = order_workflow_tasks(self.workflow)

        # make a list of ::all_required_inputs, with name and data type
        all_required_inputs = []
        for wf_task in workflow_tasks:
            for name, data_type in wf_task.task.required_inputs.iteritems():
                if name not in all_required_inputs:
                    all_required_inputs.append((name, data_type))

        # make a list of ::all_parameters from the list, with name and data type
        all_parameters = []
        for parameter in parameter_list.parameters.all():
            all_parameters.append((parameter.name, parameter.data_type))

        # if not every element of ::all_parameters is also in ::all_required_inputs
        # the list is not valid for this workflow, raise an form validation error
        if not False not in [e in all_parameters for e in all_required_inputs]:
            error_message = "Not all required inputs found in ParameterList."
            error_message += "

Required inputs by the Tasks are:"
            for param in all_required_inputs:
                error_message += "
" + str(param[0]) + ":" + str(param[1])

            error_message += "

List contains:"
            for param in all_parameters:
                error_message += "
" + str(param[0]) + ":" + str(param[1])
            raise forms.ValidationError(mark_safe(error_message))
        return list_id

class WorkerSelectForm(forms.Form):
    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        workflow = kwargs.pop('workflow')

        super(WorkerSelectForm, self).__init__(*args, **kwargs)

        possible_workers = get_possible_worker(request, as_list=True)
        preset_worker = get_preset_worker(request, workflow, flat=True)

        self.fields['worker'] = forms.MultipleChoiceField(
            required=True,
            label="Worker",
            choices=possible_workers,
            initial=preset_worker,
            widget=forms.SelectMultiple(attrs={'class': 'form-control', }),
        )
        self.fields['worker'].empty_label = None

        # workaround for displaying a message to the user
        # if no worker is acessible while initializing the form
        if len(possible_workers) <= 0:
            self.cleaned_data = {}
            msg = "No Worker found. Please configure a Worker first."
            self.add_error('worker', msg)

    def clean(self):
        cleaned_data = super(WorkerSelectForm, self).clean()
        worker_ids = cleaned_data.get('worker')
        workers = Worker.objects.filter(id__in=worker_ids)
        for worker in workers:
            if not worker.is_accessible:
                raise forms.ValidationError("worker not aceccible")
        # Always return the full collection of cleaned data.
        # This has changed in django 1.7 but raises errors in django <1.7
        # We want to support both, so we return cleanded_data
        return cleaned_data

class WorkerForm(forms.ModelForm):
    """Base Worker Form"""
    class Meta:
        model = Worker
        fields = ['title', 'ip', 'port', 'muted', 'secret_key', 'pre_task', 'post_task', ]
        widgets = {
            'title': forms.TextInput(attrs={'class': 'form-control', }),
            'ip': forms.TextInput(attrs={'class': 'form-control', }),
            'addr': forms.TextInput(attrs={'class': 'form-control', 'rows': '5', }),
            'port': forms.TextInput(attrs={'class': 'form-control', }),
            'secret_key': forms.TextInput(attrs={'class': 'form-control', }),
            'pre_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'post_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
        }

    def __init__(self, *args, **kwargs):
        super(WorkerForm, self).__init__(*args, **kwargs)

class WorkerCreateForm(WorkerForm):
    """Form to create a Worker"""
    class Meta(WorkerForm.Meta):
        fields = ['owner', ] + WorkerForm.Meta.fields
        exclude = ['muted']
        WorkerForm.Meta.widgets['owner'] = forms.Select(attrs={'class': 'form-control'})

    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        super(WorkerCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].queryset = get_possible_owners(request)
        self.fields['owner'].initial = get_curr_org(request)

class WorkerChangeForm(WorkerForm):
    """Form to edit a Worker"""
    pass

class TaskForm(forms.ModelForm):
    """Base Task Form"""
    class Meta:
        model = Task
        fields = ['title', 'desc', 'template', 'changelog', 'is_readonly', ]
        widgets = {
            'title': forms.TextInput(attrs={'class': 'form-control', }),
            'desc': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'template': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', 'style': 'font-family:monospace;'}),
            'changelog': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'is_readonly': forms.CheckboxInput(attrs={}),
        }

    def __init__(self, *args, **kwargs):
        super(TaskForm, self).__init__(*args, **kwargs)

class TaskCreateForm(TaskForm):
    """Form to create a Task"""
    class Meta(TaskForm.Meta):
        fields = ['owner', ] + TaskForm.Meta.fields
        TaskForm.Meta.widgets['owner'] = forms.Select(attrs={'class': 'form-control'})

    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        super(TaskCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].queryset = get_possible_owners(request)
        self.fields['owner'].initial = get_curr_org(request)

class TaskChangeForm(TaskForm):
    """Form to edit a Task"""
    pass

class TaskImportForm(forms.Form):
    """Form to import a Task"""
    task_json = forms.FileField()

class WorkflowForm(forms.ModelForm):
    """Base Workflow Form"""
    class Meta:
        model = Workflow
        fields = ['title', 'desc', 'pre_task', 'post_task', ]
        widgets = {
            'title': forms.TextInput(attrs={'class': 'form-control', }),
            'desc': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'pre_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'post_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
        }

    def __init__(self, *args, **kwargs):
        super(WorkflowForm, self).__init__(*args, **kwargs)

class WorkflowCreateForm(WorkflowForm):
    """Form to create a Workflow"""
    class Meta(WorkflowForm.Meta):
        fields = ['owner', ] + WorkflowForm.Meta.fields
        WorkflowForm.Meta.widgets['owner'] = forms.Select(attrs={'class': 'form-control'})

    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        super(WorkflowCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].queryset = get_possible_owners(request)
        self.fields['owner'].initial = get_curr_org(request)

class WorkflowChangeForm(WorkflowForm):
    """Form to edit a Workflow"""
    pass

class WorkflowReorderForm(forms.Form):
    workflow_id = forms.CharField(max_length=30, widget=forms.HiddenInput())
    order = forms.CharField(widget=forms.HiddenInput())

class TaskParameterForm(forms.Form):
    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        extra = kwargs.pop('extra')
        workflow_task = kwargs.pop('workflow_task')
        super(TaskParameterForm, self).__init__(*args, **kwargs)

        for name, data_type in extra.iteritems():
            possible_parameters = get_possible_parameters(request, workflow_task, data_type)
            preset_parameter = get_preset_param(request, workflow_task, name, data_type)
            self.fields['wt_task_%i.%s.%s' % (workflow_task.id, data_type, name)] = forms.ChoiceField(
                label="%s (%s)" % (name, data_type),
                choices=possible_parameters,
                initial=preset_parameter,
                widget=forms.Select(attrs={'class': 'form-control', }),
                )

            # each param gets a hidden field for saving prompted data
            self.fields['prompt_wt_task_%i.%s.%s' % (workflow_task.id, data_type, name)] = forms.CharField(widget=forms.HiddenInput(), required=False, )

class OrganizationCreateForm(forms.ModelForm):
    """
    Form class for creating a new organization, complete with new owner, including a
    User instance, OrganizationUser instance, and OrganizationOwner instance.
    """

    def __init__(self, *args, **kwargs):
        self.request = kwargs.pop('request')
        super(OrganizationCreateForm, self).__init__(*args, **kwargs)

    class Meta:
        model = Organization
        exclude = ('users', 'is_active',)

    def save(self, **kwargs):
        """
        Create the organization, then get the user, then make the owner.
        """
        user = self.request.user
        return create_organization(user, self.cleaned_data['name'], self.cleaned_data['slug'], is_active=True)
<commit_msg>Better validation for the WorkerSelectForm<commit_after>from django import forms
from django.forms.models import modelformset_factory
from django.utils.safestring import mark_safe
from organizations.models import Organization
from organizations.utils import create_organization

from threebot.utils import get_possible_parameters, get_possible_worker, get_curr_org
from threebot.utils import get_preset_param, get_preset_worker, get_possible_owners, get_possible_lists, get_preset_list
from threebot.models import Worker
from threebot.models import Workflow
from threebot.models import Task
from threebot.models import UserParameter, OrganizationParameter, ParameterList
from threebot.utils import order_workflow_tasks

class UserParameterCreateForm(forms.ModelForm):
    next = forms.CharField(required=False, widget=forms.HiddenInput())

    class Meta:
        model = UserParameter
        fields = ['data_type', 'name', 'value', 'owner', ]
        widgets = {
            'data_type': forms.Select(attrs={'class': 'form-control', }),
            'name': forms.TextInput(attrs={'class': 'form-control', }),
            'value': forms.TextInput(attrs={'class': 'form-control', }),
            'owner': forms.HiddenInput(),
        }

    def __init__(self, *args, **kwargs):
        user = kwargs.pop('user')
        super(UserParameterCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].initial = user

class UserParameterChangeForm(UserParameterCreateForm):
    class Meta(UserParameterCreateForm.Meta):
        exclude = ('owner',)

class OrganizationParameterCreateForm(forms.ModelForm):
    class Meta:
        model = OrganizationParameter
        fields = ['data_type', 'name', 'value', 'owner', ]
        widgets = {
            'data_type': forms.Select(attrs={'class': 'form-control', }),
            'name': forms.TextInput(attrs={'class': 'form-control', }),
            'value': forms.TextInput(attrs={'class': 'form-control', }),
            'owner': forms.HiddenInput(),
        }

    def __init__(self, *args, **kwargs):
        org = kwargs.pop('org')
        super(OrganizationParameterCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].initial = org

class OrganizationParameterChangeForm(OrganizationParameterCreateForm):
    class Meta(OrganizationParameterCreateForm.Meta):
        exclude = ('owner',)

def make_organization_parameter_formset(org, extra=3):
    """
    This is a workaround for passing custom parameters (the parameter owner in our case)
    to a ModelFormset and is based on this stackoverflow answer:
    http://stackoverflow.com/a/1376616
    """
    class _OrganizationParameterCreateForm(forms.ModelForm):
        remove_from_list = forms.BooleanField(required=False)

        class Meta:
            model = OrganizationParameter
            fields = ['data_type', 'name', 'value', 'owner', ]
            widgets = {
                'data_type': forms.Select(attrs={'class': 'form-control', }),
                'name': forms.TextInput(attrs={'class': 'form-control', }),
                'value': forms.TextInput(attrs={'class': 'form-control', }),
                'owner': forms.HiddenInput(),
            }

        def __init__(self, *args, **kwargs):
            # self.org = org
            super(_OrganizationParameterCreateForm, self).__init__(*args, **kwargs)
            self.fields['owner'].initial = org

    return modelformset_factory(OrganizationParameter, form=_OrganizationParameterCreateForm, extra=extra, can_delete=True)

def make_user_parameter_formset(user, extra=3):
    """
    This is a workaround for passing custom parameters (the parameter owner in our case)
    to a ModelFormset and is based on this stackoverflow answer:
    http://stackoverflow.com/a/1376616
    """
    class _UserParameterCreateForm(forms.ModelForm):
        class Meta:
            model = UserParameter
            fields = ['data_type', 'name', 'value', 'owner', ]
            widgets = {
                'data_type': forms.Select(attrs={'class': 'form-control', }),
                'name': forms.TextInput(attrs={'class': 'form-control', }),
                'value': forms.TextInput(attrs={'class': 'form-control', }),
                'owner': forms.HiddenInput(),
            }

        def __init__(self, *args, **kwargs):
            # user = kwargs.pop('user')
            super(_UserParameterCreateForm, self).__init__(*args, **kwargs)
            self.fields['owner'].initial = user

    return modelformset_factory(UserParameter, form=_UserParameterCreateForm, extra=extra, can_delete=True)

class ParameterListSelectForm(forms.Form):
    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        workflow = kwargs.pop('workflow')
        self.workflow = workflow

        super(ParameterListSelectForm, self).__init__(*args, **kwargs)

        possible_lists = [('', '--Choose a Parameter List--')]
        possible_lists += get_possible_lists(request, workflow)
        preset_list_id = get_preset_list(request, workflow, id=True)

        self.fields['parameter_list'] = forms.ChoiceField(
            label="Parameter List",
            choices=possible_lists,
            initial=preset_list_id,
            widget=forms.Select(attrs={'class': 'form-control', }),
        )
        self.fields['parameter_list'].empty_label = None

    def clean_parameter_list(self, *args, **kwargs):
        list_id = self.cleaned_data['parameter_list']
        parameter_list = ParameterList.objects.get(id=list_id)
        workflow_tasks = order_workflow_tasks(self.workflow)

        # make a list of ::all_required_inputs, with name and data type
        all_required_inputs = []
        for wf_task in workflow_tasks:
            for name, data_type in wf_task.task.required_inputs.iteritems():
                if name not in all_required_inputs:
                    all_required_inputs.append((name, data_type))

        # make a list of ::all_parameters from the list, with name and data type
        all_parameters = []
        for parameter in parameter_list.parameters.all():
            all_parameters.append((parameter.name, parameter.data_type))

        # if not every element of ::all_parameters is also in ::all_required_inputs
        # the list is not valid for this workflow, raise an form validation error
        if not False not in [e in all_parameters for e in all_required_inputs]:
            error_message = "Not all required inputs found in ParameterList."
            error_message += "

Required inputs by the Tasks are:"
            for param in all_required_inputs:
                error_message += "
" + str(param[0]) + ":" + str(param[1])

            error_message += "

List contains:"
            for param in all_parameters:
                error_message += "
" + str(param[0]) + ":" + str(param[1])
            raise forms.ValidationError(mark_safe(error_message))
        return list_id

class WorkerSelectForm(forms.Form):
    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        workflow = kwargs.pop('workflow')

        super(WorkerSelectForm, self).__init__(*args, **kwargs)

        possible_workers = get_possible_worker(request, as_list=True)
        preset_worker = get_preset_worker(request, workflow, flat=True)

        self.fields['worker'] = forms.MultipleChoiceField(
            required=True,
            label="Worker",
            choices=possible_workers,
            initial=preset_worker,
            widget=forms.SelectMultiple(attrs={'class': 'form-control', }),
        )
        self.fields['worker'].empty_label = None

        # workaround for displaying a message to the user
        # if no worker is acessible while initializing the form
        if len(possible_workers) <= 0:
            self.cleaned_data = {}
            msg = "No Worker found. Please configure a Worker first."
            self.add_error('worker', msg)

    def clean(self):
        cleaned_data = super(WorkerSelectForm, self).clean()
        worker_ids = cleaned_data.get('worker', [])
        workers = Worker.objects.filter(id__in=worker_ids)
        for worker in workers:
            if not worker.is_accessible:
                raise forms.ValidationError("worker not aceccible")
        # Always return the full collection of cleaned data.
        # This has changed in django 1.7 but raises errors in django <1.7
        # We want to support both, so we return cleanded_data
        return cleaned_data

class WorkerForm(forms.ModelForm):
    """Base Worker Form"""
    class Meta:
        model = Worker
        fields = ['title', 'ip', 'port', 'muted', 'secret_key', 'pre_task', 'post_task', ]
        widgets = {
            'title': forms.TextInput(attrs={'class': 'form-control', }),
            'ip': forms.TextInput(attrs={'class': 'form-control', }),
            'addr': forms.TextInput(attrs={'class': 'form-control', 'rows': '5', }),
            'port': forms.TextInput(attrs={'class': 'form-control', }),
            'secret_key': forms.TextInput(attrs={'class': 'form-control', }),
            'pre_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'post_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
        }

    def __init__(self, *args, **kwargs):
        super(WorkerForm, self).__init__(*args, **kwargs)

class WorkerCreateForm(WorkerForm):
    """Form to create a Worker"""
    class Meta(WorkerForm.Meta):
        fields = ['owner', ] + WorkerForm.Meta.fields
        exclude = ['muted']
        WorkerForm.Meta.widgets['owner'] = forms.Select(attrs={'class': 'form-control'})

    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        super(WorkerCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].queryset = get_possible_owners(request)
        self.fields['owner'].initial = get_curr_org(request)

class WorkerChangeForm(WorkerForm):
    """Form to edit a Worker"""
    pass

class TaskForm(forms.ModelForm):
    """Base Task Form"""
    class Meta:
        model = Task
        fields = ['title', 'desc', 'template', 'changelog', 'is_readonly', ]
        widgets = {
            'title': forms.TextInput(attrs={'class': 'form-control', }),
            'desc': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'template': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', 'style': 'font-family:monospace;'}),
            'changelog': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'is_readonly': forms.CheckboxInput(attrs={}),
        }

    def __init__(self, *args, **kwargs):
        super(TaskForm, self).__init__(*args, **kwargs)

class TaskCreateForm(TaskForm):
    """Form to create a Task"""
    class Meta(TaskForm.Meta):
        fields = ['owner', ] + TaskForm.Meta.fields
        TaskForm.Meta.widgets['owner'] = forms.Select(attrs={'class': 'form-control'})

    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        super(TaskCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].queryset = get_possible_owners(request)
        self.fields['owner'].initial = get_curr_org(request)

class TaskChangeForm(TaskForm):
    """Form to edit a Task"""
    pass

class TaskImportForm(forms.Form):
    """Form to import a Task"""
    task_json = forms.FileField()

class WorkflowForm(forms.ModelForm):
    """Base Workflow Form"""
    class Meta:
        model = Workflow
        fields = ['title', 'desc', 'pre_task', 'post_task', ]
        widgets = {
            'title': forms.TextInput(attrs={'class': 'form-control', }),
            'desc': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'pre_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
            'post_task': forms.Textarea(attrs={'class': 'form-control', 'rows': '5', }),
        }

    def __init__(self, *args, **kwargs):
        super(WorkflowForm, self).__init__(*args, **kwargs)

class WorkflowCreateForm(WorkflowForm):
    """Form to create a Workflow"""
    class Meta(WorkflowForm.Meta):
        fields = ['owner', ] + WorkflowForm.Meta.fields
        WorkflowForm.Meta.widgets['owner'] = forms.Select(attrs={'class': 'form-control'})

    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        super(WorkflowCreateForm, self).__init__(*args, **kwargs)
        self.fields['owner'].queryset = get_possible_owners(request)
        self.fields['owner'].initial = get_curr_org(request)

class WorkflowChangeForm(WorkflowForm):
    """Form to edit a Workflow"""
    pass

class WorkflowReorderForm(forms.Form):
    workflow_id = forms.CharField(max_length=30, widget=forms.HiddenInput())
    order = forms.CharField(widget=forms.HiddenInput())

class TaskParameterForm(forms.Form):
    def __init__(self, *args, **kwargs):
        request = kwargs.pop('request')
        extra = kwargs.pop('extra')
        workflow_task = kwargs.pop('workflow_task')
        super(TaskParameterForm, self).__init__(*args, **kwargs)

        for name, data_type in extra.iteritems():
            possible_parameters = get_possible_parameters(request, workflow_task, data_type)
            preset_parameter = get_preset_param(request, workflow_task, name, data_type)
            self.fields['wt_task_%i.%s.%s' % (workflow_task.id, data_type, name)] = forms.ChoiceField(
                label="%s (%s)" % (name, data_type),
                choices=possible_parameters,
                initial=preset_parameter,
                widget=forms.Select(attrs={'class': 'form-control', }),
                )

            # each param gets a hidden field for saving prompted data
            self.fields['prompt_wt_task_%i.%s.%s' % (workflow_task.id, data_type, name)] = forms.CharField(widget=forms.HiddenInput(), required=False, )

class OrganizationCreateForm(forms.ModelForm):
    """
    Form class for creating a new organization, complete with new owner, including a
    User instance, OrganizationUser instance, and OrganizationOwner instance.
    """

    def __init__(self, *args, **kwargs):
        self.request = kwargs.pop('request')
        super(OrganizationCreateForm, self).__init__(*args, **kwargs)

    class Meta:
        model = Organization
        exclude = ('users', 'is_active',)

    def save(self, **kwargs):
        """
        Create the organization, then get the user, then make the owner.
        """
        user = self.request.user
        return create_organization(user, self.cleaned_data['name'], self.cleaned_data['slug'], is_active=True)
<|endoftext|>


<commit_before><commit_msg>temp: add special handling for gdrive<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (c) J. A. Zanardo Jr. <zanardo@gmail.com>
#
# https://github.com/zanardo/ticket
#
# Para detalhes do licenciamento, ver COPYING na distribuição
#

import ticket.db
import ticket.user

import config

def tagsdesc():
    # Retorna as descrições de tags
    tagdesc = {}
    c = ticket.db.getcursor()
    c.execute("select tag, description, bgcolor, fgcolor from tagsdesc")
    for r in c:
        tagdesc[r['tag']] = {
            'description': r['description'] or '',
            'bgcolor': r['bgcolor'] or '#00D6D6',
            'fgcolor': r['fgcolor'] or '#4D4D4D'
        }
    return tagdesc

def ticketblocks(ticket_id):
    # Retorna quais ticket são bloqueados por um ticket
    deps = {}
    c = ticket.db.getcursor()
    c.execute("select d.blocks, t.title, t.status, t.admin_only "
        "from dependencies as d inner join tickets as t on t.id = d.blocks "
        "where d.ticket_id = :ticket_id", locals())
    for r in c:
        deps[r[0]] = { 'title': r[1], 'status': r[2], 'admin_only': r[3] }
    return deps

def ticketdepends(ticket_id):
    # Retorna quais ticket dependem de um ticket
    deps = {}
    c = ticket.db.getcursor()
    c.execute("select d.ticket_id, t.title, t.status, t.admin_only "
        "from dependencies as d inner join tickets as t on t.id = d.ticket_id "
        "where d.blocks = :ticket_id", locals())
    for r in c:
        deps[r[0]] = { 'title': r[1], 'status': r[2], 'admin_only': r[3] }
    return deps

def tickettags(ticket_id):
    # Retorna tags de um ticket
    c = ticket.db.getcursor()
    c.execute("select tag from tags where ticket_id = :ticket_id", locals())
    return [r['tag'] for r in c]

def tickettitle(ticket_id):
    # Retorna o título de um ticket
    c = ticket.db.getcursor()
    c.execute("select title from tickets where id = :ticket_id", locals())
    return c.fetchone()['title']

def sanitizecomment(comment):
    # Sanitiza o texto do comentário (quebras de linhas, links, etc)
    subs = [ (r'\r', ''), (r'&', '&'), (r'<', '<'), (r'>', '>'),
             (r'\r?\n', '
\r\n'), (r'\t', '   '),
             (r'  ', '  '), (r'#(\d+)', r'#\1') ]
    for f, t in subs:
        comment = re.sub(f, t, comment)
    return comment
<commit_msg>carrega módulo<commit_after># -*- coding: utf-8 -*-
#
# Copyright (c) J. A. Zanardo Jr. <zanardo@gmail.com>
#
# https://github.com/zanardo/ticket
#
# Para detalhes do licenciamento, ver COPYING na distribuição
#

import re

import ticket.db
import ticket.user

import config

def tagsdesc():
    # Retorna as descrições de tags
    tagdesc = {}
    c = ticket.db.getcursor()
    c.execute("select tag, description, bgcolor, fgcolor from tagsdesc")
    for r in c:
        tagdesc[r['tag']] = {
            'description': r['description'] or '',
            'bgcolor': r['bgcolor'] or '#00D6D6',
            'fgcolor': r['fgcolor'] or '#4D4D4D'
        }
    return tagdesc

def ticketblocks(ticket_id):
    # Retorna quais ticket são bloqueados por um ticket
    deps = {}
    c = ticket.db.getcursor()
    c.execute("select d.blocks, t.title, t.status, t.admin_only "
        "from dependencies as d inner join tickets as t on t.id = d.blocks "
        "where d.ticket_id = :ticket_id", locals())
    for r in c:
        deps[r[0]] = { 'title': r[1], 'status': r[2], 'admin_only': r[3] }
    return deps

def ticketdepends(ticket_id):
    # Retorna quais ticket dependem de um ticket
    deps = {}
    c = ticket.db.getcursor()
    c.execute("select d.ticket_id, t.title, t.status, t.admin_only "
        "from dependencies as d inner join tickets as t on t.id = d.ticket_id "
        "where d.blocks = :ticket_id", locals())
    for r in c:
        deps[r[0]] = { 'title': r[1], 'status': r[2], 'admin_only': r[3] }
    return deps

def tickettags(ticket_id):
    # Retorna tags de um ticket
    c = ticket.db.getcursor()
    c.execute("select tag from tags where ticket_id = :ticket_id", locals())
    return [r['tag'] for r in c]

def tickettitle(ticket_id):
    # Retorna o título de um ticket
    c = ticket.db.getcursor()
    c.execute("select title from tickets where id = :ticket_id", locals())
    return c.fetchone()['title']

def sanitizecomment(comment):
    # Sanitiza o texto do comentário (quebras de linhas, links, etc)
    subs = [ (r'\r', ''), (r'&', '&'), (r'<', '<'), (r'>', '>'),
             (r'\r?\n', '
\r\n'), (r'\t', '   '),
             (r'  ', '  '), (r'#(\d+)', r'#\1') ]
    for f, t in subs:
        comment = re.sub(f, t, comment)
    return comment
<|endoftext|>


<commit_before>import re
import inspect
import logging
import glob

from importlib.util import find_spec, module_from_spec

from zope.interface import implementer, Interface, Attribute
from zope.interface.verify import verifyObject
from zope.interface.exceptions import BrokenImplementation

logger = logging.getLogger('detectem')

def get_plugin_by_name(name, plugins):
    try:
        return [p for p in plugins if p.name == name][0]
    except IndexError:
        return None

class _PluginLoader(object):

    def __init__(self):
        self.plugins = []

    def _full_class_name(self, ins):
        return '{}.{}'.format(ins.__class__.__module__, ins.__class__.__name__)

    def _get_plugin_module_paths(self, plugin_dir):
        filepaths = [
            fp for fp in glob.glob('{}/**/*.py'.format(plugin_dir), recursive=True)
            if not fp.endswith('__init__.py')
        ]
        rel_paths = [re.sub(plugin_dir.rstrip('/') + '/', '', fp) for fp in filepaths]
        module_paths = [rp.replace('/', '.').replace('.py', '') for rp in rel_paths]
        return module_paths

    def _load_plugin(self, klass):
        ins = klass()
        try:
            if verifyObject(IPlugin, ins):
                self.plugins.append(ins)
        except BrokenImplementation:
            logger.warning(
                "Plugin '%(name)s' doesn't provide the plugin interface",
                {'name': self._full_class_name(ins)}
            )

    def load_plugins(self, plugins_package):
        # Resolve directory in the filesystem
        plugin_dir = find_spec(plugins_package).submodule_search_locations[0]

        for module_path in self._get_plugin_module_paths(plugin_dir):
            # Load the module dynamically
            spec = find_spec('{}.{}'.format(plugins_package, module_path))
            m = module_from_spec(spec)
            spec.loader.exec_module(m)

            # Get classes from module and extract the plugin classes
            classes = inspect.getmembers(m, predicate=inspect.isclass)
            for _, klass in classes:
                if klass.__module__ != spec.name:
                    continue
                if not klass.__name__.endswith('Plugin'):
                    continue
                self._load_plugin(klass)

def load_plugins():
    """ Return the list of plugin instances """
    loader = _PluginLoader()
    loader.load_plugins('detectem.plugins')
    return loader.plugins

class IPlugin(Interface):
    name = Attribute(""" Name to identify the plugin. """)
    homepage = Attribute(""" Plugin homepage. """)

@implementer(IPlugin)
class Plugin():
    def _get_matchers(self, value, source='matchers'):
        matchers_dict = getattr(self, source, [])
        return [m[value] for m in matchers_dict if value in m]

    def get_grouped_matchers(self, source='matchers'):
        """ Return dictionary of matchers (not empty ones) """
        data = {}
        for k in ['url', 'body', 'header']:
            m = self._get_matchers(k, source)
            if m:
                data[k] = m

        return data

    @property
    def is_modular(self):
        return bool(hasattr(self, 'modular_matchers'))

    @property
    def is_indicator(self):
        return bool(hasattr(self, 'indicators'))
<commit_msg>Check for duplicated plugins<commit_after>import re
import inspect
import logging
import glob

from importlib.util import find_spec, module_from_spec

from zope.interface import implementer, Interface, Attribute
from zope.interface.verify import verifyObject
from zope.interface.exceptions import BrokenImplementation

logger = logging.getLogger('detectem')

def get_plugin_by_name(name, plugins):
    try:
        return [p for p in plugins if p.name == name][0]
    except IndexError:
        return None

class _PluginLoader(object):

    def __init__(self):
        self.plugins = []

    def _full_class_name(self, ins):
        return '{}.{}'.format(ins.__class__.__module__, ins.__class__.__name__)

    def _get_plugin_module_paths(self, plugin_dir):
        filepaths = [
            fp for fp in glob.glob('{}/**/*.py'.format(plugin_dir), recursive=True)
            if not fp.endswith('__init__.py')
        ]
        rel_paths = [re.sub(plugin_dir.rstrip('/') + '/', '', fp) for fp in filepaths]
        module_paths = [rp.replace('/', '.').replace('.py', '') for rp in rel_paths]
        return module_paths

    def _load_plugin(self, klass):
        ins = klass()
        try:
            if verifyObject(IPlugin, ins):
                reg = get_plugin_by_name(ins.name, self.plugins)
                if not reg:
                    self.plugins.append(ins)
                else:
                    logger.warning(
                        "Plugin '%(name)s' by '%(ins)s' is already provided by '%(reg)s'",
                        {'name': ins.name,
                         'ins': self._full_class_name(ins),
                         'reg': self._full_class_name(reg)}
                    )
        except BrokenImplementation:
            logger.warning(
                "Plugin '%(name)s' doesn't provide the plugin interface",
                {'name': self._full_class_name(ins)}
            )

    def load_plugins(self, plugins_package):
        # Resolve directory in the filesystem
        plugin_dir = find_spec(plugins_package).submodule_search_locations[0]

        for module_path in self._get_plugin_module_paths(plugin_dir):
            # Load the module dynamically
            spec = find_spec('{}.{}'.format(plugins_package, module_path))
            m = module_from_spec(spec)
            spec.loader.exec_module(m)

            # Get classes from module and extract the plugin classes
            classes = inspect.getmembers(m, predicate=inspect.isclass)
            for _, klass in classes:
                if klass.__module__ != spec.name:
                    continue
                if not klass.__name__.endswith('Plugin'):
                    continue
                self._load_plugin(klass)

def load_plugins():
    """ Return the list of plugin instances """
    loader = _PluginLoader()
    loader.load_plugins('detectem.plugins')
    return loader.plugins

class IPlugin(Interface):
    name = Attribute(""" Name to identify the plugin. """)
    homepage = Attribute(""" Plugin homepage. """)

@implementer(IPlugin)
class Plugin():
    def _get_matchers(self, value, source='matchers'):
        matchers_dict = getattr(self, source, [])
        return [m[value] for m in matchers_dict if value in m]

    def get_grouped_matchers(self, source='matchers'):
        """ Return dictionary of matchers (not empty ones) """
        data = {}
        for k in ['url', 'body', 'header']:
            m = self._get_matchers(k, source)
            if m:
                data[k] = m

        return data

    @property
    def is_modular(self):
        return bool(hasattr(self, 'modular_matchers'))

    @property
    def is_indicator(self):
        return bool(hasattr(self, 'indicators'))
<|endoftext|>


<commit_before>from django.contrib.auth.models import User
from django.db import models
from django.utils.translation import ugettext_lazy as _

from tickets import settings

class Ticket(models.Model):
    creator = models.ForeignKey(User, verbose_name=_("Creator"), related_name='tickets')
    date = models.DateTimeField(_("Date"), auto_now_add=True)
    last_update = models.DateTimeField(_("Date"), auto_now=True)
    subject = models.CharField(_("Subject"), max_length=255)
    description = models.TextField(_("Description"), help_text=_("A detailed description of your problem."))
    assignee = models.ForeignKey(User, verbose_name=_("Assignee"), related_name="assigned_tickets", blank=True, null=True)
    status = models.SmallIntegerField(_("Status"), choices=settings.STATUS_CHOICES, default=0)

    class meta:
        verbose_name = _("Ticket")
        verbose_name_plural = _("Tickets")
        ordering = ['date']

    def get_comments_count(self):
        return self.comments.count()

    def __unicode__(self):
        return "%s# %s" % (self.id, self.subject)

    def is_closed(self):
        return self.status in settings.CLOSED_STATUSES

class TicketComment(models.Model):
    ticket = models.ForeignKey(Ticket, verbose_name=_("Ticket"), related_name='comments')
    date = models.DateTimeField(auto_now_add=True, verbose_name=_("Date"), editable=False)
    author = models.ForeignKey(User, verbose_name=_("Author"))
    comment = models.TextField(_("Comment"))

    class Meta:
       verbose_name = _("Ticket comment")
       verbose_name_plural = _("Ticket comments")
       ordering = ['date']

    def __unicode__(self):
        return "Comment on " + unicode(self.ticket)

<commit_msg>remove unnecessary newline at the end of models file<commit_after>from django.contrib.auth.models import User
from django.db import models
from django.utils.translation import ugettext_lazy as _

from tickets import settings

class Ticket(models.Model):
    creator = models.ForeignKey(User, verbose_name=_("Creator"), related_name='tickets')
    date = models.DateTimeField(_("Date"), auto_now_add=True)
    last_update = models.DateTimeField(_("Date"), auto_now=True)
    subject = models.CharField(_("Subject"), max_length=255)
    description = models.TextField(_("Description"), help_text=_("A detailed description of your problem."))
    assignee = models.ForeignKey(User, verbose_name=_("Assignee"), related_name="assigned_tickets", blank=True, null=True)
    status = models.SmallIntegerField(_("Status"), choices=settings.STATUS_CHOICES, default=0)

    class meta:
        verbose_name = _("Ticket")
        verbose_name_plural = _("Tickets")
        ordering = ['date']

    def get_comments_count(self):
        return self.comments.count()

    def __unicode__(self):
        return "%s# %s" % (self.id, self.subject)

    def is_closed(self):
        return self.status in settings.CLOSED_STATUSES

class TicketComment(models.Model):
    ticket = models.ForeignKey(Ticket, verbose_name=_("Ticket"), related_name='comments')
    date = models.DateTimeField(auto_now_add=True, verbose_name=_("Date"), editable=False)
    author = models.ForeignKey(User, verbose_name=_("Author"))
    comment = models.TextField(_("Comment"))

    class Meta:
       verbose_name = _("Ticket comment")
       verbose_name_plural = _("Ticket comments")
       ordering = ['date']

    def __unicode__(self):
        return "Comment on " + unicode(self.ticket)
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
import random
import pdb
import sklearn
import os

from sklearn.metrics import roc_curve, auc
from .feature_squeezing import FeatureSqueezingDetector
from .magnet_mnist import MagNetDetector as MagNetDetectorMNIST
from .magnet_cifar import MagNetDetector as MagNetDetectorCIFAR

from tensorflow.python.platform import flags
FLAGS = flags.FLAGS
from utils.output import write_to_csv

def get_tpr_fpr(true_labels, pred_labels):
    TP = np.sum(np.logical_and(pred_labels == 1, true_labels == 1))
    FP = np.sum(np.logical_and(pred_labels == 1, true_labels == 0))

    AP = np.sum(true_labels)
    AN = np.sum(1-true_labels)

    tpr = TP/AP if AP>0 else np.nan
    fpr = FP/AN if AN>0 else np.nan

    return tpr, fpr, TP, AP

def evalulate_detection_test(Y_detect_test, Y_detect_pred):
    accuracy = sklearn.metrics.accuracy_score(Y_detect_test, Y_detect_pred, normalize=True, sample_weight=None)
    tpr, fpr, tp, ap = get_tpr_fpr(Y_detect_test, Y_detect_pred)
    return accuracy, tpr, fpr, tp, ap

from tinydb import TinyDB, Query

class DetectionEvaluator:
    """
    Get a dataset;
        Failed adv as benign / Failed adv as adversarial.
    For each detector:
        Train
        Test
        Report performance
            Detection rate on each attack.
            Detection on SAEs / FAEs.
            ROC-AUC.

    A detector should have this simplified interface:
        Y_pred = detector(X)
    """
    def __init__(self, model, result_folder, csv_fname, dataset_name):
        pass
        # set_base_model()
        self.model = model
        self.task_dir = result_folder
        self.csv_fpath = os.path.join(result_folder, csv_fname)
        self.dataset_name = dataset_name

        if not os.path.isdir(self.task_dir):
            os.makedirs(self.task_dir)

    def get_attack_id(self, attack_name):
        return self.attack_name_id[attack_name]

    def build_detection_dataset(self, X, Y_label, Y_pred, X_adv_list, Y_adv_pred_list, attack_names, attack_string_hash, clip, Y_test_target_next, Y_test_target_ll):
        # X_train, Y_train, X_test, Y_test, test_idx, failed_adv_idx = \
        #     get_detection_train_test_set(X, Y_label, X_adv_list, Y_adv_pred_list, attack_names)

        """
        Data Model:
            index, attack_id, misclassified, train
            14,     0,           0,             1
        """

        self.attack_names = attack_names
        self.attack_name_id = {}
        self.attack_name_id['legitimate'] = 0
        for i,attack_name in enumerate(attack_names):
            self.attack_name_id[attack_name] = i+1

        X_adv_all = np.concatenate(X_adv_list)
        X_leg_all = X[:len(X_adv_all)]

        self.X_detect = X_detect = np.concatenate([X_leg_all, X_adv_all])
        Y_label_adv = Y_label[:len(X_adv_list[0])]

        detection_db_path = os.path.join(self.task_dir, "detection_db_%s_clip_%s.json" % (attack_string_hash, clip))

        if os.path.isfile(detection_db_path):
            self.db = TinyDB(detection_db_path)
            self.query = Query()
            print ("Loaded an existing detection dataset.")
            return
        else:
            print ("Preparing the detection dataset...")

        # 1. Split Train and Test
        random.seed(1234)
        length = len(X_detect)
        train_ratio = 0.5
        train_idx = random.sample(range(length), int(train_ratio*length))
        train_test_seq = [1 if idx in train_idx else 0 for idx in range(length) ]

        # 2. Tag the misclassified examples, both legitimate and adversarial.
        # TODO: Differentiate the successful examples between targeted and non-targeted.
        misclassified_seq = list(np.argmax(Y_label[:len(X_leg_all)], axis=1) != np.argmax(Y_pred[:len(X_leg_all)], axis=1))
        for Y_adv_pred in Y_adv_pred_list:
            misclassified_seq_adv = list(np.argmax(Y_adv_pred, axis=1) != np.argmax(Y_label_adv, axis=1))
            misclassified_seq += misclassified_seq_adv

        success_adv_seq = [False] * len(X_leg_all)
        for i, Y_adv_pred in enumerate(Y_adv_pred_list):
            attack_name = attack_names[i]
            if 'targeted=ll' in attack_name:
                success_adv_seq_attack = list(np.argmax(Y_adv_pred, axis=1) == np.argmax(Y_test_target_ll, axis=1))
            elif 'targeted=next' in attack_name:
                success_adv_seq_attack = list(np.argmax(Y_adv_pred, axis=1) == np.argmax(Y_test_target_next, axis=1))
            else:
                # The same as misclassified.
                success_adv_seq_attack = list(np.argmax(Y_adv_pred, axis=1) != np.argmax(Y_label_adv, axis=1))
            success_adv_seq += success_adv_seq_attack

        # 3. Tag the attack ID, 0 as legitimate.
        attack_id_seq = [0]*len(X_leg_all)
        for i,attack_name in enumerate(attack_names):
            attack_id_seq += [i+1]*len(X_adv_list[0])

        assert len(X_detect) == len(train_test_seq) == len(misclassified_seq) == len(attack_id_seq)

        self.db = TinyDB(detection_db_path)
        self.query = Query()

        for i in range(len(X_detect)):
            attack_id = attack_id_seq[i]
            misclassified = 1 if misclassified_seq[i] == True else 0
            success = 1 if success_adv_seq[i] == True else 0
            train = train_test_seq[i]
            rec = {'index': i, 'attack_id': attack_id, 'misclassified': misclassified, 'success': success, 'train': train}
            self.db.insert(rec)

    def get_data_from_db_records(self, recs):
        if len(recs) == 0:
            return None, None
        X_idx = [rec['index'] for rec in recs]
        X = self.X_detect[np.array(X_idx)]
        Y = np.array([1 if rec['attack_id']>0 else 0 for rec in recs])
        return X, Y

    def get_training_testing_data(self, train = True):
        db = self.db
        query = self.query

        recs = db.search(query.train == 1)
        X_train, Y_train = self.get_data_from_db_records(recs)

        recs = db.search(query.train == 0)
        X_test, Y_test = self.get_data_from_db_records(recs)

        return X_train, Y_train, X_test, Y_test

    def get_adversarial_data(self, only_testing, success, attack_name=None, include_legitimate=False):
        db = self.db
        query = self.query

        conditions_and = []
        if only_testing:
            conditions_and.append(query.train == 0)

        if attack_name is None:
            conditions_and.append(query.attack_id > 0)
        else:
            attack_id = self.get_attack_id(attack_name)
            conditions_and.append(query.attack_id == attack_id)

        if success:
            conditions_and.append(query.success == 1)
        else:
            conditions_and.append(query.success == 0)

        conditions = reduce(lambda a,b:a&b, conditions_and)
        # print ("conditions: %s " % conditions)

        recs = db.search(conditions)

        if include_legitimate:
            if only_testing:
                conditions = (query.attack_id == 0) & (query.train == 0)
            else:
                conditions = query.attack_id == 0
            # print ("additional conditions: %s " % conditions)
            recs += db.search(conditions)

        return self.get_data_from_db_records(recs)

    def get_sae_testing_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=True, success=True, attack_name=attack_name)

    def get_sae_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=False, success=True, attack_name=attack_name)

    def get_fae_testing_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=True, success=False, attack_name=attack_name)

    def get_fae_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=False, success=False, attack_name=attack_name)

    def get_all_non_fae_testing_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=True, success=True, attack_name=attack_name, include_legitimate=True)

    def get_all_non_fae_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=False, success=True, attack_name=attack_name, include_legitimate=True)

    def get_detector_by_name(self, detector_name):
        model = self.model
        detector = None

        if detector_name.startswith('FeatureSqueezing'):
            detector = FeatureSqueezingDetector(model, detector_name)
        elif detector_name.startswith('MagNet'):
            if self.dataset_name == 'MNIST':
                detector = MagNetDetectorMNIST(model, detector_name)
            elif self.dataset_name == "CIFAR-10":
                detector = MagNetDetectorCIFAR(model, detector_name)

        return detector

    def evaluate_detections(self, params_str):
        X_train, Y_train, X_test, Y_test = self.get_training_testing_data()

        # Example: --detection "FeatureSqueezing?distance_measure=l1&squeezers=median_smoothing_2,bit_depth_4;"
        detector_names = [ele.strip() for ele in params_str.split(';') if ele.strip()!= '']

        dataset_name = self.dataset_name
        csv_fpath = "./detection_%s_saes.csv" % dataset_name
        fieldnames = ['detector', 'threshold', 'fpr'] + self.attack_names + ['overall']
        to_csv = []

        for detector_name in detector_names:
            detector = self.get_detector_by_name(detector_name)
            if detector is None:
                print ("Skipped an unknown detector [%s]" % detector_name.split('?')[0])
                continue
            detector.train(X_train, Y_train)
            Y_test_pred, Y_test_pred_score = detector.test(X_test)

            accuracy, tpr, fpr, tp, ap = evalulate_detection_test(Y_test, Y_test_pred)
            fprs, tprs, thresholds = roc_curve(Y_test, Y_test_pred_score)
            roc_auc = auc(fprs, tprs)

            print ("Detector: %s" % detector_name)
            print ("Accuracy: %f\tTPR: %f\tFPR: %f\tROC-AUC: %f" % (accuracy, tpr, fpr, roc_auc))

            rec = {}
            rec['detector'] = detector_name
            if hasattr(detector, 'threshold'):
                rec['threshold'] = detector.threshold
            else:
                rec['threshold'] = None
            rec['fpr'] = fpr
            overall_detection_rate_saes = 0
            nb_saes = 0
            for attack_name in self.attack_names:
                # No adversarial examples for training for the current detection methods.
                # X_sae, Y_sae = self.get_sae_testing_data(attack_name)
                if FLAGS.detection_train_test_mode:
                    X_sae, Y_sae = self.get_sae_testing_data(attack_name)
                else:
                    X_sae, Y_sae = self.get_sae_data(attack_name)
                Y_test_pred, Y_test_pred_score = detector.test(X_sae)
                _, tpr, _, tp, ap = evalulate_detection_test(Y_sae, Y_test_pred)
                print ("Detection rate on SAEs: %.4f \t %3d/%3d \t %s" % (tpr, tp, ap, attack_name))
                overall_detection_rate_saes += tpr * len(Y_sae)
                nb_saes += len(Y_sae)
                rec[attack_name] = tpr
                # print ("overall_detection_rate_saes/nb_saes: %d/%d" % (overall_detection_rate_saes, nb_saes))

            print ("Overall detection rate on SAEs: %f (%d/%d)" % (overall_detection_rate_saes/nb_saes, overall_detection_rate_saes, nb_saes))
            rec['overall'] = float(overall_detection_rate_saes/nb_saes)
            to_csv.append(rec)

            # No adversarial examples for training for the current detection methods.
            # X_sae_all, Y_sae_all = self.get_sae_testing_data()
            print ("### Excluding FAEs:")
            if FLAGS.detection_train_test_mode:
                X_nfae_all, Y_nfae_all = self.get_all_non_fae_testing_data()
            else:
                X_nfae_all, Y_nfae_all = self.get_all_non_fae_data()
            Y_pred, Y_pred_score = detector.test(X_nfae_all)
            _, tpr, _, tp, ap = evalulate_detection_test(Y_nfae_all, Y_pred)
            fprs, tprs, thresholds = roc_curve(Y_nfae_all, Y_pred_score)

            print ("threshold\tfpr\ttpr")
            for i, threshold  in enumerate(thresholds):
                print ("%.4f\t%.4f\t%.4f" % (threshold, fprs[i], tprs[i]))

            roc_auc = auc(fprs, tprs)
            print ("Overall TPR: %f\tROC-AUC: %f" % (tpr, roc_auc))

            # FAEs
            if FLAGS.detection_train_test_mode:
                X_fae, Y_fae = self.get_fae_testing_data()
            else:
                X_fae, Y_fae = self.get_fae_data()
            Y_test_pred, Y_test_pred_score = detector.test(X_fae)
            _, tpr, _, tp, ap = evalulate_detection_test(Y_fae, Y_test_pred)
            print ("Overall detection rate on FAEs: %.4f \t %3d/%3d" % (tpr, tp, ap))

        write_to_csv(to_csv, csv_fpath, fieldnames)<commit_msg>Reduced the verbose output.<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
import random
import pdb
import sklearn
import os

from sklearn.metrics import roc_curve, auc
from .feature_squeezing import FeatureSqueezingDetector
from .magnet_mnist import MagNetDetector as MagNetDetectorMNIST
from .magnet_cifar import MagNetDetector as MagNetDetectorCIFAR

from tensorflow.python.platform import flags
FLAGS = flags.FLAGS
from utils.output import write_to_csv

def get_tpr_fpr(true_labels, pred_labels):
    TP = np.sum(np.logical_and(pred_labels == 1, true_labels == 1))
    FP = np.sum(np.logical_and(pred_labels == 1, true_labels == 0))

    AP = np.sum(true_labels)
    AN = np.sum(1-true_labels)

    tpr = TP/AP if AP>0 else np.nan
    fpr = FP/AN if AN>0 else np.nan

    return tpr, fpr, TP, AP

def evalulate_detection_test(Y_detect_test, Y_detect_pred):
    accuracy = sklearn.metrics.accuracy_score(Y_detect_test, Y_detect_pred, normalize=True, sample_weight=None)
    tpr, fpr, tp, ap = get_tpr_fpr(Y_detect_test, Y_detect_pred)
    return accuracy, tpr, fpr, tp, ap

from tinydb import TinyDB, Query

class DetectionEvaluator:
    """
    Get a dataset;
        Failed adv as benign / Failed adv as adversarial.
    For each detector:
        Train
        Test
        Report performance
            Detection rate on each attack.
            Detection on SAEs / FAEs.
            ROC-AUC.

    A detector should have this simplified interface:
        Y_pred = detector(X)
    """
    def __init__(self, model, result_folder, csv_fname, dataset_name):
        pass
        # set_base_model()
        self.model = model
        self.task_dir = result_folder
        self.csv_fpath = os.path.join(result_folder, csv_fname)
        self.dataset_name = dataset_name

        if not os.path.isdir(self.task_dir):
            os.makedirs(self.task_dir)

    def get_attack_id(self, attack_name):
        return self.attack_name_id[attack_name]

    def build_detection_dataset(self, X, Y_label, Y_pred, X_adv_list, Y_adv_pred_list, attack_names, attack_string_hash, clip, Y_test_target_next, Y_test_target_ll):
        # X_train, Y_train, X_test, Y_test, test_idx, failed_adv_idx = \
        #     get_detection_train_test_set(X, Y_label, X_adv_list, Y_adv_pred_list, attack_names)

        """
        Data Model:
            index, attack_id, misclassified, train
            14,     0,           0,             1
        """

        self.attack_names = attack_names
        self.attack_name_id = {}
        self.attack_name_id['legitimate'] = 0
        for i,attack_name in enumerate(attack_names):
            self.attack_name_id[attack_name] = i+1

        X_adv_all = np.concatenate(X_adv_list)
        X_leg_all = X[:len(X_adv_all)]

        self.X_detect = X_detect = np.concatenate([X_leg_all, X_adv_all])
        Y_label_adv = Y_label[:len(X_adv_list[0])]

        detection_db_path = os.path.join(self.task_dir, "detection_db_%s_clip_%s.json" % (attack_string_hash, clip))

        if os.path.isfile(detection_db_path):
            self.db = TinyDB(detection_db_path)
            self.query = Query()
            print ("Loaded an existing detection dataset.")
            return
        else:
            print ("Preparing the detection dataset...")

        # 1. Split Train and Test
        random.seed(1234)
        length = len(X_detect)
        train_ratio = 0.5
        train_idx = random.sample(range(length), int(train_ratio*length))
        train_test_seq = [1 if idx in train_idx else 0 for idx in range(length) ]

        # 2. Tag the misclassified examples, both legitimate and adversarial.
        # TODO: Differentiate the successful examples between targeted and non-targeted.
        misclassified_seq = list(np.argmax(Y_label[:len(X_leg_all)], axis=1) != np.argmax(Y_pred[:len(X_leg_all)], axis=1))
        for Y_adv_pred in Y_adv_pred_list:
            misclassified_seq_adv = list(np.argmax(Y_adv_pred, axis=1) != np.argmax(Y_label_adv, axis=1))
            misclassified_seq += misclassified_seq_adv

        success_adv_seq = [False] * len(X_leg_all)
        for i, Y_adv_pred in enumerate(Y_adv_pred_list):
            attack_name = attack_names[i]
            if 'targeted=ll' in attack_name:
                success_adv_seq_attack = list(np.argmax(Y_adv_pred, axis=1) == np.argmax(Y_test_target_ll, axis=1))
            elif 'targeted=next' in attack_name:
                success_adv_seq_attack = list(np.argmax(Y_adv_pred, axis=1) == np.argmax(Y_test_target_next, axis=1))
            else:
                # The same as misclassified.
                success_adv_seq_attack = list(np.argmax(Y_adv_pred, axis=1) != np.argmax(Y_label_adv, axis=1))
            success_adv_seq += success_adv_seq_attack

        # 3. Tag the attack ID, 0 as legitimate.
        attack_id_seq = [0]*len(X_leg_all)
        for i,attack_name in enumerate(attack_names):
            attack_id_seq += [i+1]*len(X_adv_list[0])

        assert len(X_detect) == len(train_test_seq) == len(misclassified_seq) == len(attack_id_seq)

        self.db = TinyDB(detection_db_path)
        self.query = Query()

        for i in range(len(X_detect)):
            attack_id = attack_id_seq[i]
            misclassified = 1 if misclassified_seq[i] == True else 0
            success = 1 if success_adv_seq[i] == True else 0
            train = train_test_seq[i]
            rec = {'index': i, 'attack_id': attack_id, 'misclassified': misclassified, 'success': success, 'train': train}
            self.db.insert(rec)

    def get_data_from_db_records(self, recs):
        if len(recs) == 0:
            return None, None
        X_idx = [rec['index'] for rec in recs]
        X = self.X_detect[np.array(X_idx)]
        Y = np.array([1 if rec['attack_id']>0 else 0 for rec in recs])
        return X, Y

    def get_training_testing_data(self, train = True):
        db = self.db
        query = self.query

        recs = db.search(query.train == 1)
        X_train, Y_train = self.get_data_from_db_records(recs)

        recs = db.search(query.train == 0)
        X_test, Y_test = self.get_data_from_db_records(recs)

        return X_train, Y_train, X_test, Y_test

    def get_adversarial_data(self, only_testing, success, attack_name=None, include_legitimate=False):
        db = self.db
        query = self.query

        conditions_and = []
        if only_testing:
            conditions_and.append(query.train == 0)

        if attack_name is None:
            conditions_and.append(query.attack_id > 0)
        else:
            attack_id = self.get_attack_id(attack_name)
            conditions_and.append(query.attack_id == attack_id)

        if success:
            conditions_and.append(query.success == 1)
        else:
            conditions_and.append(query.success == 0)

        conditions = reduce(lambda a,b:a&b, conditions_and)
        # print ("conditions: %s " % conditions)

        recs = db.search(conditions)

        if include_legitimate:
            if only_testing:
                conditions = (query.attack_id == 0) & (query.train == 0)
            else:
                conditions = query.attack_id == 0
            # print ("additional conditions: %s " % conditions)
            recs += db.search(conditions)

        return self.get_data_from_db_records(recs)

    def get_sae_testing_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=True, success=True, attack_name=attack_name)

    def get_sae_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=False, success=True, attack_name=attack_name)

    def get_fae_testing_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=True, success=False, attack_name=attack_name)

    def get_fae_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=False, success=False, attack_name=attack_name)

    def get_all_non_fae_testing_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=True, success=True, attack_name=attack_name, include_legitimate=True)

    def get_all_non_fae_data(self, attack_name=None):
        return self.get_adversarial_data(only_testing=False, success=True, attack_name=attack_name, include_legitimate=True)

    def get_detector_by_name(self, detector_name):
        model = self.model
        detector = None

        if detector_name.startswith('FeatureSqueezing'):
            detector = FeatureSqueezingDetector(model, detector_name)
        elif detector_name.startswith('MagNet'):
            if self.dataset_name == 'MNIST':
                detector = MagNetDetectorMNIST(model, detector_name)
            elif self.dataset_name == "CIFAR-10":
                detector = MagNetDetectorCIFAR(model, detector_name)

        return detector

    def evaluate_detections(self, params_str):
        X_train, Y_train, X_test, Y_test = self.get_training_testing_data()

        # Example: --detection "FeatureSqueezing?distance_measure=l1&squeezers=median_smoothing_2,bit_depth_4;"
        detector_names = [ele.strip() for ele in params_str.split(';') if ele.strip()!= '']

        dataset_name = self.dataset_name
        csv_fpath = "./detection_%s_saes.csv" % dataset_name
        fieldnames = ['detector', 'threshold', 'fpr'] + self.attack_names + ['overall']
        to_csv = []

        for detector_name in detector_names:
            detector = self.get_detector_by_name(detector_name)
            if detector is None:
                print ("Skipped an unknown detector [%s]" % detector_name.split('?')[0])
                continue
            detector.train(X_train, Y_train)
            Y_test_pred, Y_test_pred_score = detector.test(X_test)

            accuracy, tpr, fpr, tp, ap = evalulate_detection_test(Y_test, Y_test_pred)
            fprs, tprs, thresholds = roc_curve(Y_test, Y_test_pred_score)
            roc_auc = auc(fprs, tprs)

            print ("Detector: %s" % detector_name)
            print ("Accuracy: %f\tTPR: %f\tFPR: %f\tROC-AUC: %f" % (accuracy, tpr, fpr, roc_auc))

            rec = {}
            rec['detector'] = detector_name
            if hasattr(detector, 'threshold'):
                rec['threshold'] = detector.threshold
            else:
                rec['threshold'] = None
            rec['fpr'] = fpr
            overall_detection_rate_saes = 0
            nb_saes = 0
            for attack_name in self.attack_names:
                # No adversarial examples for training for the current detection methods.
                # X_sae, Y_sae = self.get_sae_testing_data(attack_name)
                if FLAGS.detection_train_test_mode:
                    X_sae, Y_sae = self.get_sae_testing_data(attack_name)
                else:
                    X_sae, Y_sae = self.get_sae_data(attack_name)
                Y_test_pred, Y_test_pred_score = detector.test(X_sae)
                _, tpr, _, tp, ap = evalulate_detection_test(Y_sae, Y_test_pred)
                print ("Detection rate on SAEs: %.4f \t %3d/%3d \t %s" % (tpr, tp, ap, attack_name))
                overall_detection_rate_saes += tpr * len(Y_sae)
                nb_saes += len(Y_sae)
                rec[attack_name] = tpr
                # print ("overall_detection_rate_saes/nb_saes: %d/%d" % (overall_detection_rate_saes, nb_saes))

            print ("Overall detection rate on SAEs: %f (%d/%d)" % (overall_detection_rate_saes/nb_saes, overall_detection_rate_saes, nb_saes))
            rec['overall'] = float(overall_detection_rate_saes/nb_saes)
            to_csv.append(rec)

            # No adversarial examples for training for the current detection methods.
            # X_sae_all, Y_sae_all = self.get_sae_testing_data()
            print ("### Excluding FAEs:")
            if FLAGS.detection_train_test_mode:
                X_nfae_all, Y_nfae_all = self.get_all_non_fae_testing_data()
            else:
                X_nfae_all, Y_nfae_all = self.get_all_non_fae_data()
            Y_pred, Y_pred_score = detector.test(X_nfae_all)
            _, tpr, _, tp, ap = evalulate_detection_test(Y_nfae_all, Y_pred)
            fprs, tprs, thresholds = roc_curve(Y_nfae_all, Y_pred_score)

            # print ("threshold\tfpr\ttpr")
            # for i, threshold  in enumerate(thresholds):
            #     print ("%.4f\t%.4f\t%.4f" % (threshold, fprs[i], tprs[i]))

            roc_auc = auc(fprs, tprs)
            print ("Overall TPR: %f\tROC-AUC: %f" % (tpr, roc_auc))

            # FAEs
            if FLAGS.detection_train_test_mode:
                X_fae, Y_fae = self.get_fae_testing_data()
            else:
                X_fae, Y_fae = self.get_fae_data()
            Y_test_pred, Y_test_pred_score = detector.test(X_fae)
            _, tpr, _, tp, ap = evalulate_detection_test(Y_fae, Y_test_pred)
            print ("Overall detection rate on FAEs: %.4f \t %3d/%3d" % (tpr, tp, ap))

        write_to_csv(to_csv, csv_fpath, fieldnames)<|endoftext|>


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

"""
    tickets.trains
    ~~~~~~~~~~~~~~

    Train tickets model.
"""

import os
import re
import sys
import json
from datetime import datetime
from collections import OrderedDict
from .utils import colored, exit_after_echo

import requests
from prettytable import PrettyTable
from requests.exceptions import ConnectionError

__all__ = ['TrainTicketsQuery']

QUERY_URL = 'https://kyfw.12306.cn/otn/lcxxcx/query'

# ERROR MSG
FROM_STATION_NOT_FOUND = 'From station not found.'
TO_STATION_NOT_FOUND = 'To station not found.'
INVALID_DATE = 'Invalid query date.'
NETWORK_CONNECTION_FAIL = 'Network connection failed.'

class TrainsCollection(object):

    """A set of raw datas from a query."""

    headers = '车次 车站 时间 历时 商务 一等 二等 软卧 硬卧 软座 硬座 无座'.split()

    def __init__(self, rows, opts):
        self._rows = rows
        self._opts = opts

    def __repr__(self):
        return '<TrainsCollection size={}>'.format(len(self))

    def __iter__(self):
        i = 0
        while True:
            if i < len(self):
                yield self[i]
            else:
                yield next(self)
            i += 1

    def __len__(self):
        return len(self._rows)

    def _get_duration(self, row):
        duration = row.get('lishi').replace(':', '小时') + '分钟'
        # take 0 hour , only show minites
        if duration.startswith('00'):
            return duration[4:]
        # take <10 hours, show 1 bit
        if duration.startswith('0'):
            return duration[1:]
        return duration

    @property
    def trains(self):
        """Filter rows according to `headers`"""
        for row in self._rows:
            train_no = row.get('station_train_code')
            initial = train_no[0].lower()
            if not self._opts or initial in self._opts:
                train = [
                    # Column: '车次'
                    train_no,
                    # Column: '车站'
                    ''.join([
                        colored.green(row.get('from_station_name')),
                        '\n',
                        colored.red(row.get('to_station_name')),
                    ]),
                    # Column: '时间'
                    ''.join([
                        colored.green(row.get('start_time')),
                        '\n',
                        colored.red(row.get('arrive_time')),
                    ]),
                    # Column: '历时'
                    self._get_duration(row),
                    # Column: '商务'
                    row.get('swz_num'),
                    # Column: '一等'
                    row.get('zy_num'),
                    # Column: '二等'
                    row.get('ze_num'),
                    # Column: '软卧'
                    row.get('rw_num'),
                    # Column: '硬卧'
                    row.get('yw_num'),
                    # Column: '软座'
                    row.get('rz_num'),
                    # Column: '硬座'
                    row.get('yz_num'),
                    # Column: '无座'
                    row.get('wz_num')
                ]
                yield train

    def pretty_print(self):
        """Use `PrettyTable` to perform formatted outprint."""
        pt = PrettyTable()
        if len(self) == 0:
            pt._set_field_names(['Sorry,'])
            pt.add_row([TRAIN_NOT_FOUND])
        else:
            pt._set_field_names(self.headers)
            for train in self.trains:
                pt.add_row(train)
        print(pt)

class TrainTicketsQuery(object):

    """Docstring for TrainTicketsCollection. """

    def __init__(self, from_station, to_station, date, opts=None):

        self.from_station = from_station
        self.to_station = to_station
        self.date = date
        self.opts = opts

    def __repr__(self):
        return 'TrainTicketsQuery from={} to={} date={}'.format(
            self._from_station, self._to_station, self._date
        )

    @property
    def stations(self):
        filepath = os.path.join(
            os.path.dirname(__file__),
            'datas', 'stations.dat'
        )
        d = {}
        with open(filepath, 'r') as f:
            for line in f.readlines():
                name, telecode = line.split()
                d.setdefault(name, telecode)
        return d

    @property
    def _from_station_telecode(self):
        code = self.stations.get(self.from_station)
        if not code:
            exit_after_echo(FROM_STATION_NOT_FOUND)
        return code

    @property
    def _to_station_telecode(self):
        code = self.stations.get(self.to_station)
        if not code:
            exit_after_echo(TO_STATION_NOT_FOUND)
        return code

    @property
    def _valid_date(self):
        date = re.sub(r'[-/\:,]+', '', self.date)
        try:
            date = datetime.strptime(self.date, '%Y%m%d')
        except ValueError:
            exit_after_echo(INVALID_DATE)
        diff = date - datetime.today()
        if diff.days not in range(-1, 50):
            exit_after_echo(INVALID_DATE)
        return datetime.strftime(date, '%Y-%m-%d')

    def _build_params(self):
        """Have no idea why wrong params order can't get data.
        So, use `OrderedDict` here.
        """
        d = OrderedDict()
        d['purpose_codes'] = 'ADULT'
        d['queryDate'] = self._valid_date
        d['from_station'] = self._from_station_telecode
        d['to_station'] = self._to_station_telecode
        return d

    def query(self):

        params = self._build_params()

        try:
            r = requests.get(QUERY_URL, params=params, verify=False)
        except ConnectionError:
            exit_after_echo(NETWORK_CONNECTION_FAIL)

        try:
            rows = r.json()['data']['datas']
        except KeyError:
            rows = []

        return TrainsCollection(rows, self.opts)
<commit_msg>small fixes<commit_after># -*- coding: utf-8 -*-

"""
    tickets.trains
    ~~~~~~~~~~~~~~

    Train tickets model.
"""

import os
import re
import sys
import json
from datetime import datetime
from collections import OrderedDict
from .utils import colored, exit_after_echo

import requests
from prettytable import PrettyTable
from requests.exceptions import ConnectionError

__all__ = ['TrainTicketsQuery']

QUERY_URL = 'https://kyfw.12306.cn/otn/lcxxcx/query'

# ERROR MSG
FROM_STATION_NOT_FOUND = 'From station not found.'
TO_STATION_NOT_FOUND = 'To station not found.'
INVALID_DATE = 'Invalid query date.'
NETWORK_CONNECTION_FAIL = 'Network connection failed.'
TRAIN_NOT_FOUND = 'No result.'

class TrainsCollection(object):

    """A set of raw datas from a query."""

    headers = '车次 车站 时间 历时 商务 一等 二等 软卧 硬卧 软座 硬座 无座'.split()

    def __init__(self, rows, opts):
        self._rows = rows
        self._opts = opts

    def __repr__(self):
        return '<TrainsCollection size={}>'.format(len(self))

    def __iter__(self):
        i = 0
        while True:
            if i < len(self):
                yield self[i]
            else:
                yield next(self)
            i += 1

    def __len__(self):
        return len(self._rows)

    def _get_duration(self, row):
        duration = row.get('lishi').replace(':', '小时') + '分钟'
        # take 0 hour , only show minites
        if duration.startswith('00'):
            return duration[4:]
        # take <10 hours, show 1 bit
        if duration.startswith('0'):
            return duration[1:]
        return duration

    @property
    def trains(self):
        """Filter rows according to `headers`"""
        for row in self._rows:
            train_no = row.get('station_train_code')
            initial = train_no[0].lower()
            if not self._opts or initial in self._opts:
                train = [
                    # Column: '车次'
                    train_no,
                    # Column: '车站'
                    ''.join([
                        colored.green(row.get('from_station_name')),
                        '\n',
                        colored.red(row.get('to_station_name')),
                    ]),
                    # Column: '时间'
                    ''.join([
                        colored.green(row.get('start_time')),
                        '\n',
                        colored.red(row.get('arrive_time')),
                    ]),
                    # Column: '历时'
                    self._get_duration(row),
                    # Column: '商务'
                    row.get('swz_num'),
                    # Column: '一等'
                    row.get('zy_num'),
                    # Column: '二等'
                    row.get('ze_num'),
                    # Column: '软卧'
                    row.get('rw_num'),
                    # Column: '硬卧'
                    row.get('yw_num'),
                    # Column: '软座'
                    row.get('rz_num'),
                    # Column: '硬座'
                    row.get('yz_num'),
                    # Column: '无座'
                    row.get('wz_num')
                ]
                yield train

    def pretty_print(self):
        """Use `PrettyTable` to perform formatted outprint."""
        pt = PrettyTable()
        if len(self) == 0:
            pt._set_field_names(['Sorry,'])
            pt.add_row([TRAIN_NOT_FOUND])
        else:
            pt._set_field_names(self.headers)
            for train in self.trains:
                pt.add_row(train)
        print(pt)

class TrainTicketsQuery(object):

    """Docstring for TrainTicketsCollection. """

    def __init__(self, from_station, to_station, date, opts=None):

        self.from_station = from_station
        self.to_station = to_station
        self.date = date
        self.opts = opts

    def __repr__(self):
        return 'TrainTicketsQuery from={} to={} date={}'.format(
            self._from_station, self._to_station, self._date
        )

    @property
    def stations(self):
        filepath = os.path.join(
            os.path.dirname(__file__),
            'datas', 'stations.dat'
        )
        d = {}
        with open(filepath, 'r') as f:
            for line in f.readlines():
                name, telecode = line.split()
                d.setdefault(name, telecode)
        return d

    @property
    def _from_station_telecode(self):
        code = self.stations.get(self.from_station)
        if not code:
            exit_after_echo(FROM_STATION_NOT_FOUND)
        return code

    @property
    def _to_station_telecode(self):
        code = self.stations.get(self.to_station)
        if not code:
            exit_after_echo(TO_STATION_NOT_FOUND)
        return code

    @property
    def _valid_date(self):
        date = re.sub(r'[-/\:,]+', '', self.date)
        try:
            date = datetime.strptime(self.date, '%Y%m%d')
        except ValueError:
            exit_after_echo(INVALID_DATE)
        diff = date - datetime.today()
        if diff.days not in range(-1, 50):
            exit_after_echo(INVALID_DATE)
        return datetime.strftime(date, '%Y-%m-%d')

    def _build_params(self):
        """Have no idea why wrong params order can't get data.
        So, use `OrderedDict` here.
        """
        d = OrderedDict()
        d['purpose_codes'] = 'ADULT'
        d['queryDate'] = self._valid_date
        d['from_station'] = self._from_station_telecode
        d['to_station'] = self._to_station_telecode
        return d

    def query(self):

        params = self._build_params()

        try:
            r = requests.get(QUERY_URL, params=params, verify=False)
        except ConnectionError:
            exit_after_echo(NETWORK_CONNECTION_FAIL)

        try:
            rows = r.json()['data']['datas']
        except KeyError:
            rows = []

        return TrainsCollection(rows, self.opts)
<|endoftext|>


<commit_before>import datetime
from django.db import models
from django.core.exceptions import ValidationError
from django.contrib.auth.models import User
from taggit.managers import TaggableManager
from preferences.models import Preferences

from epl.custommodels import IntegerRangeField, FloatRangeField
from util.file_validator import FileValidator
from timemap.constants import BRANCH_NAME_LEN, BRANCH_DESCRIPTION_LEN, STORY_TITLE_LEN, \
                              STORY_DESCRIPTION_LEN, STORY_TEXT_LEN, MAP_BASE_FOLDER_LEN, \
                              MAP_TITLE_LEN, MAP_AUTHOR_LEN, UPLOAD_EXTENSIONS, \
                              UPLOAD_MIME_TYPES

from util.email import emailer, email_template

class Branch(models.Model):

    class Meta:
        verbose_name_plural = "Branches"

    name = models.CharField(max_length=BRANCH_NAME_LEN)
    description = models.TextField(max_length=BRANCH_DESCRIPTION_LEN)
    start_year = IntegerRangeField(min_value=1900, max_value=3000)
    end_year = IntegerRangeField(min_value=1900, max_value=3000, blank=True, null=True)
    floor_plan = models.FileField(upload_to="floor_plans", blank=True, null=True)
    latitude_help = "Latitude range : -90:90"
    latitude = FloatRangeField(min_value=-90, max_value=90, help_text=latitude_help)
    longitude_help = "Longitude range : -180:180"
    longitude = FloatRangeField(min_value=-180, max_value=180, help_text=longitude_help)

    def clean(self):
        if self.end_year and self.start_year > self.end_year:
            raise ValidationError("End year must occur after start year")

    def __unicode__(self):
        return self.name

def media_upload_to(instance, filename):
    filename = filename.split('/')[-1]
    return instance.CONTENT_TYPE_DICT[instance.content_type]+ "/" + filename

class Story(models.Model):
    TEXT = "T"
    LINK = "L"
    IMAGE = "I"
    PDF = "P"
    AUDIO = "A"
    VIDEO = "V"
    CONTENT_TYPE_CHOICES = (
        (TEXT, 'text'),
        (LINK, 'link'),
        (IMAGE, 'image'),
        (PDF, 'pdf'),
        (AUDIO, 'audio'),
        (VIDEO, 'video'),
    )
    CONTENT_TYPE_DICT = dict(CONTENT_TYPE_CHOICES)

    class Meta:
        verbose_name_plural = "Stories"

    title = models.CharField(max_length=STORY_TITLE_LEN)
    description = models.TextField(max_length=STORY_DESCRIPTION_LEN)
    story_text = models.TextField(max_length=STORY_TEXT_LEN, blank=True)
    link_url = models.URLField(blank=True, error_messages={'invalid': "Please input a valid URL (for example: http://www.example.com)."})
    media_file = models.FileField(upload_to=media_upload_to,
                                  blank=True,
                                  validators=[FileValidator(allowed_extensions=UPLOAD_EXTENSIONS,
                                                           allowed_mimetypes=UPLOAD_MIME_TYPES)])
    year = IntegerRangeField(min_value=1900, max_value=3000)
    month = IntegerRangeField(min_value=1, max_value=12, blank=True, null=True)
    day = IntegerRangeField(min_value=1, max_value=31, blank=True, null=True)
    branch = models.ForeignKey('Branch')
    keywords = TaggableManager(verbose_name="keywords",
                               help_text=("A comma-separated list of keywords"),
                               blank=True)
    user = models.ForeignKey(User)
    public_approved = models.BooleanField(default=False)
    content_type = models.CharField(max_length=1,
                                    choices=CONTENT_TYPE_CHOICES,
                                    default=TEXT)

    def clean(self):
        try:
            day = self.day if self.day else 1
            month = self.month if self.month else 1
            date = "%s/%s/%s" % (day, month, self.year)
            datetime.datetime.strptime(date, "%d/%m/%Y")
        except ValueError:
            #TODO: Should make the resulting error clearer
            raise ValidationError("Please enter a valid date.")

    def __unicode__(self):
        return self.title

class Map(models.Model):

    class Meta:
        verbose_name_plural = "Maps"

    base_folder = models.CharField(max_length=MAP_BASE_FOLDER_LEN)
    title = models.CharField(max_length=MAP_TITLE_LEN)
    author = models.CharField(max_length=MAP_AUTHOR_LEN)
    published = IntegerRangeField(min_value=1900, max_value=3000)
    start_year = IntegerRangeField(min_value=1900, max_value=3000)
    end_year = IntegerRangeField(min_value=1900, max_value=3000)

    def clean(self):
        if self.start_year > self.end_year:
            raise ValidationError("End year must occur after start year.")

    def __unicode__(self):
        return self.title

class TimemapPreferences(Preferences):

    class Meta:
        verbose_name_plural = "Timemap Preferences"

    __module__ = 'preferences.models'
    initial_timemap_date = models.DateField(default=datetime.date(2013, 1, 1))

# Signal setup
from django.dispatch.dispatcher import receiver
from django.db.models.signals import pre_save, pre_delete

@receiver(pre_save)
def validate_model(sender, **kwargs):
    """
    Force a clean call when certain models are saved in order to do
    keep model constrains
    """
    if sender in [Branch, Story] and 'raw' in kwargs and not kwargs['raw']:
        kwargs['instance'].full_clean()

@receiver(pre_delete)
def story_delete(sender, instance, **kwargs):
    """
    Delete media files when stories are deleted
    """
    if sender in [Story] and instance.media_file:
        instance.media_file.delete(False)
<commit_msg>Create unique file names for uploads<commit_after>import uuid
import datetime
from django.db import models
from django.core.exceptions import ValidationError
from django.contrib.auth.models import User
from taggit.managers import TaggableManager
from preferences.models import Preferences

from epl.custommodels import IntegerRangeField, FloatRangeField
from util.file_validator import FileValidator
from timemap.constants import BRANCH_NAME_LEN, BRANCH_DESCRIPTION_LEN, STORY_TITLE_LEN, \
                              STORY_DESCRIPTION_LEN, STORY_TEXT_LEN, MAP_BASE_FOLDER_LEN, \
                              MAP_TITLE_LEN, MAP_AUTHOR_LEN, UPLOAD_EXTENSIONS, \
                              UPLOAD_MIME_TYPES

from util.email import emailer, email_template

class Branch(models.Model):

    class Meta:
        verbose_name_plural = "Branches"

    name = models.CharField(max_length=BRANCH_NAME_LEN)
    description = models.TextField(max_length=BRANCH_DESCRIPTION_LEN)
    start_year = IntegerRangeField(min_value=1900, max_value=3000)
    end_year = IntegerRangeField(min_value=1900, max_value=3000, blank=True, null=True)
    floor_plan = models.FileField(upload_to="floor_plans", blank=True, null=True)
    latitude_help = "Latitude range : -90:90"
    latitude = FloatRangeField(min_value=-90, max_value=90, help_text=latitude_help)
    longitude_help = "Longitude range : -180:180"
    longitude = FloatRangeField(min_value=-180, max_value=180, help_text=longitude_help)

    def clean(self):
        if self.end_year and self.start_year > self.end_year:
            raise ValidationError("End year must occur after start year")

    def __unicode__(self):
        return self.name

def media_upload_to(instance, filename):
    ext = filename.split('.')[-1]
    filename = "%s.%s" % (uuid.uuid4(), ext)
    return instance.CONTENT_TYPE_DICT[instance.content_type]+ "/" + filename

class Story(models.Model):
    TEXT = "T"
    LINK = "L"
    IMAGE = "I"
    PDF = "P"
    AUDIO = "A"
    VIDEO = "V"
    CONTENT_TYPE_CHOICES = (
        (TEXT, 'text'),
        (LINK, 'link'),
        (IMAGE, 'image'),
        (PDF, 'pdf'),
        (AUDIO, 'audio'),
        (VIDEO, 'video'),
    )
    CONTENT_TYPE_DICT = dict(CONTENT_TYPE_CHOICES)

    class Meta:
        verbose_name_plural = "Stories"

    title = models.CharField(max_length=STORY_TITLE_LEN)
    description = models.TextField(max_length=STORY_DESCRIPTION_LEN)
    story_text = models.TextField(max_length=STORY_TEXT_LEN, blank=True)
    link_url = models.URLField(blank=True, error_messages={'invalid': "Please input a valid URL (for example: http://www.example.com)."})
    media_file = models.FileField(upload_to=media_upload_to,
                                  blank=True,
                                  validators=[FileValidator(allowed_extensions=UPLOAD_EXTENSIONS,
                                                           allowed_mimetypes=UPLOAD_MIME_TYPES)])
    year = IntegerRangeField(min_value=1900, max_value=3000)
    month = IntegerRangeField(min_value=1, max_value=12, blank=True, null=True)
    day = IntegerRangeField(min_value=1, max_value=31, blank=True, null=True)
    branch = models.ForeignKey('Branch')
    keywords = TaggableManager(verbose_name="keywords",
                               help_text=("A comma-separated list of keywords"),
                               blank=True)
    user = models.ForeignKey(User)
    public_approved = models.BooleanField(default=False)
    content_type = models.CharField(max_length=1,
                                    choices=CONTENT_TYPE_CHOICES,
                                    default=TEXT)

    def clean(self):
        try:
            day = self.day if self.day else 1
            month = self.month if self.month else 1
            date = "%s/%s/%s" % (day, month, self.year)
            datetime.datetime.strptime(date, "%d/%m/%Y")
        except ValueError:
            #TODO: Should make the resulting error clearer
            raise ValidationError("Please enter a valid date.")

    def __unicode__(self):
        return self.title

class Map(models.Model):

    class Meta:
        verbose_name_plural = "Maps"

    base_folder = models.CharField(max_length=MAP_BASE_FOLDER_LEN)
    title = models.CharField(max_length=MAP_TITLE_LEN)
    author = models.CharField(max_length=MAP_AUTHOR_LEN)
    published = IntegerRangeField(min_value=1900, max_value=3000)
    start_year = IntegerRangeField(min_value=1900, max_value=3000)
    end_year = IntegerRangeField(min_value=1900, max_value=3000)

    def clean(self):
        if self.start_year > self.end_year:
            raise ValidationError("End year must occur after start year.")

    def __unicode__(self):
        return self.title

class TimemapPreferences(Preferences):

    class Meta:
        verbose_name_plural = "Timemap Preferences"

    __module__ = 'preferences.models'
    initial_timemap_date = models.DateField(default=datetime.date(2013, 1, 1))

# Signal setup
from django.dispatch.dispatcher import receiver
from django.db.models.signals import pre_save, pre_delete

@receiver(pre_save)
def validate_model(sender, **kwargs):
    """
    Force a clean call when certain models are saved in order to do
    keep model constrains
    """
    if sender in [Branch, Story] and 'raw' in kwargs and not kwargs['raw']:
        kwargs['instance'].full_clean()

@receiver(pre_delete)
def story_delete(sender, instance, **kwargs):
    """
    Delete media files when stories are deleted
    """
    if sender in [Story] and instance.media_file:
        instance.media_file.delete(False)
<|endoftext|>