<commit_before>import logging
log = logging.getLogger(__name__)
from collections import OrderedDict
from copy import copy, deepcopy
import numpy as np
from numpy import array  # Needed for eval. Do not remove.

class SweepPoints(list):
    """
    This class is used to create sweep points for any measurement.
    The SweepPoints object is a list of dictionaries of the form:
        [
            # 1st sweep dimension
            {param_name0: (values, unit, label),
             param_name1: (values, unit, label),
            ...
             param_nameN: (values, unit, label)},

            # 2nd sweep dimension
            {param_name0: (values, unit, label),
             param_name1: (values, unit, label),
            ...
             param_nameN: (values, unit, label)},

             .
             .
             .

            # D-th sweep dimension
            {param_name0: (values, unit, label),
             param_name1: (values, unit, label),
            ...
             param_nameN: (values, unit, label)},
        ]

    Example how to use this class to create a 2D sweep for 3 qubits, where
    the first (hard) sweep is over amplitudes and the :

    sp = SweepPoints()
    for qb in ['qb1', 'qb2', 'qb3']:
        sp.add_sweep_parameter(f'lengths_{qb}', np.linspace(10e-9, 1e-6, 80),
        's', 'Pulse length, $L$')
    sp.add_sweep_dimension()
    for qb in ['qb1', 'qb2', 'qb3']:
        sp.add_sweep_parameter(f'amps_{qb}', np.linspace(0, 1, 20),
        'V', 'Pulse amplitude, $A$')
    """
    def __init__(self, param=None, values=None, unit='', label=None,
                 dimension=-1, min_length=0):
        """
        Create a SweepPoints instance.
        :param param: list of dicts, repr of SweepPoints instance, or name of
            a sweep parameter
        :param values: list or array of numeric values of the sweep parameter
            specified by param
        :param unit: string specifying the unit of the sweep parameter
            specified by param
        :param label: str specifying the (latex style) label/name of the sweep
            parameter specified by param
        :param dimension: sweep dimension where sweep parameter param should be
            added
        :param min_length: minimum number of sweep dimensions to create
        """
        super().__init__()
        if isinstance(param, list):
            self.add_dict_list(param)
        elif isinstance(param, str):
            if values is not None:
                self.add_sweep_parameter(param, values, unit, label,
                                         dimension)
            else:
                self.add_dict_list(eval(param))

        while len(self) < min_length:
            self.add_sweep_dimension()

    def __getitem__(self, i):
        """
        Overloading of List.__getitem__ to ensure type SweepPoints is preserved.
        :param i: element or slice
        :return: element or new SweepPoints instance
        """
        if isinstance(i, str):
            new_data = self.get_sweep_params_property('values', 'all', i)
        else:
            new_data = super().__getitem__(i)
            if isinstance(i, slice):
                new_data = self.__class__(new_data)
        return new_data

    def add_dict_list(self, dict_list):
        """
        Append the dicts in dict_list to self.
        :param dict_list: list of dictionaries in the format of this class, or
            in the legacy format {param_name: {'values': ...,
                                               'unit': ...,
                                               'label': ...}}
        """
        for d in deepcopy(dict_list):
            if len(d) == 0 or isinstance(list(d.values())[0], tuple):
                # assume that dicts have the same format as this class
                self.append(d)
            else:
                # import from a list of sweep dicts in the old format
                self.append({k: (v['values'],
                                 v.get('unit', ''),
                                 v.get('label', k))
                             for k, v in d.items()})

    def add_sweep_parameter(self, param_name, values, unit='', label=None,
                            dimension=-1):
        """
        Adds sweep points to a given dimension.
        :param param_name: (str) parameter name
        :param values: (list or numpy array) sweep values
        :param unit: (optional str) unit of the values (default: '')
        :param label: (optional str) label e.g. for plots (default: param_name)
        :dim: the dimension to which the point should be added (default:
            last dimension)
        """
        if label is None:
            label = param_name
        assert self.find_parameter(param_name) is None, \
            f'A sweep parameter with name "{param_name}" already exists.'
        while len(self) == 0 or (dimension >= 0 and dimension >= len(self)):
            self.add_sweep_dimension()
        assert self.length(dimension) in [0, len(values)], \
            'Number of values has to match the length of existing sweep ' \
            'points.'
        self[dimension].update({param_name: (values, unit, label)})

    def add_sweep_dimension(self):
        self.append(dict())

    def get_sweep_dimension(self, dimension='all', pop=False):
        """
        Returns the sweep dict of the sweep dimension specified by dimension.
        :param dimension: int specifying a sweep dimension or
            the string 'all'
        :param pop: bool specifying whether to pop (True) or get(False) the
            sweep dimension.
        :return: self if dimension == 'all', else self[dimension]
        """
        if dimension == 'all':
            if pop:
                to_return = deepcopy(self)
                self.clear()
                return to_return
            else:
                return self
        else:
            if len(self) < dimension:
                raise ValueError(f'Dimension {dimension} not found.')
            to_return = self[dimension]
            if pop:
                self[dimension] = {}
            return to_return

    def get_sweep_params_description(self, param_names, dimension='all',
                                     pop=False):
        """
        Get the sweep tuples for the sweep parameters param_names if they are
        found in the sweep dimension dict specified by dimension.
        :param param_names: string or list of strings corresponding to keys in
            the dictionaries in self. Can also be 'all'
        :param dimension: 'all' or int specifying a sweep dimension
        :param pop: bool specifying whether to pop (True) or get(False) the
            sweep parameters.
        :return:
            If the param_names are found in self or self[dimension]:
            if param_names == 'all': list with all the sweep tuples
                in the sweep dimension dict specified by dimension.
            if param_names is string: string with the sweep tuples of each
                param_names in the sweep dimension dict specified by dimension.
            if param_names is list: list with the sweep tuples of each
                param_names in the sweep dimension dict specified by dimension.
            if param_names is None: string corresponding to the
                first sweep parameter in the sweep dimension dict
            If none of param_names are found, raises KeyError.
        """
        sweep_points_dim = self.get_sweep_dimension(
            dimension, pop=pop and param_names == 'all')
        is_list = True
        if param_names != 'all' and not isinstance(param_names, list):
            param_names = [param_names]
            is_list = False

        sweep_param_values = []
        if isinstance(sweep_points_dim, list):
            for sweep_dim_dict in sweep_points_dim:
                if param_names == 'all':
                    sweep_param_values += list(sweep_dim_dict.values())
                else:
                    for pn in param_names:
                        if pn in sweep_dim_dict:
                            sweep_param_values += [sweep_dim_dict.pop(pn) if pop
                                                   else sweep_dim_dict[pn]]
        else:  # it is a dict
            if param_names == 'all':
                sweep_param_values += list(sweep_points_dim.values())
            else:
                for pn in param_names:
                    if pn in sweep_points_dim:
                        sweep_param_values += [sweep_points_dim.pop(pn) if pop
                                               else sweep_points_dim[pn]]
        if len(sweep_param_values) == 0:
            s = "sweep points" if dimension == "all" else f'sweep dimension ' \
                                                          f'{dimension}'
            raise KeyError(f'{param_names} not found in {s}.')

        if is_list:
            return sweep_param_values
        else:
            return sweep_param_values[0]

    def get_sweep_params_property(self, property, dimension='all',
                                  param_names=None):
        """
        Get a property of the sweep parameters param_names in self.
        :param property: str with the name of a sweep param property. Can be
            "values", "unit", "label."
        :param dimension: 'all' or int specifying a sweep dimension
            (default 'all')
        :param param_names: None, or string or list of strings corresponding to
            keys in the sweep dimension specified by dimension.
            Can also be 'all'
        :return:
            if param_names == 'all': list with the property of all
                param_names in the sweep dimension dict specified by dimension.
            if param_names is string: the property of the sweep parameter
                specified in param_names in the sweep dimension dict specified
                by dimension.
            if param_names is list: list with the property of each
                param_names in the sweep dimension dict specified by dimension.
            if param_names is None: property corresponding to the
                first sweep parameter in the sweep dimension dict
        """
        properties_dict = {'values': 0, 'unit': 1, 'label': 2}
        sweep_points_dim = self.get_sweep_dimension(dimension)

        if param_names is None:
            if isinstance(sweep_points_dim, list):
                for sweep_dim_dict in sweep_points_dim:
                    if len(sweep_dim_dict) == 0:
                        return [] if property == 'values' else ''
                    else:
                        return next(iter(sweep_dim_dict.values()))[
                            properties_dict[property]]
            else:
                if len(sweep_points_dim) == 0:
                    return [] if property == 'values' else ''
                else:
                    return next(iter(sweep_points_dim.values()))[
                        properties_dict[property]]
        elif param_names == 'all' or isinstance(param_names, list):
            return [pnd[properties_dict[property]] for pnd in
                    self.get_sweep_params_description(param_names,
                                                      dimension)]
        else:
            return self.get_sweep_params_description(
                param_names, dimension)[properties_dict[property]]

    def get_meas_obj_sweep_points_map(self, measurement_objects):
        """
        Constructs the measurement-objects-sweep-points map as the dict
        {mobj_name: [sweep_param_name_0, ..., sweep_param_name_n]}

        If a sweep dimension has only one sweep parameter name (dict with only
        one key), then it assumes all mobjs use that sweep parameter name.

        If the sweep dimension has more than one sweep parameter name (dict with
        several keys), then:
            - first adds to the list for each mobj only those sweep
            param names that contain the mobj_name.
            - if some parameters remain that do not contain any of the
            measurement_objects names, it is assumed that all
            measurement_objects used them so they will be added for each mobj

        :param measurement_objects: list of strings to be used as keys in the
            returned dictionary. These are the measured object names.
            Can also be list of measurement object instances with a name
            attribute, in which case this function gets the list of names.
        :return: dict of the form
         {mobj_name: [sweep_param_name_0, ..., sweep_param_name_n]}
        """

        # Ensure measurement_objects is a list
        if isinstance(measurement_objects, list):
            measurement_objects = copy(measurement_objects)
        else:
            measurement_objects = [measurement_objects]

        for i, mobj in enumerate(measurement_objects):
            if hasattr(mobj, 'name'):
                # A list of measurement object instances with a
                # name attribute was provided
                measurement_objects[i] = mobj.name

        sweep_points_map = {mobjn: [] for mobjn in measurement_objects}

        for dim, d in enumerate(self):
            # d is a dictionary with keys == sweep parameter names
            if len(d) == 1:
                # Only one sweep parameter.
                # Assume all mobjs use the same param_name given by the key
                # of d.
                for mobjn in measurement_objects:
                    sweep_points_map[mobjn] += [next(iter(d))]
            else:
                all_pars = []
                for mobjn in measurement_objects:
                    # Find all sweep param names that contain the mobj name
                    pars = [k for k in list(d) if k.startswith(mobjn)]
                    all_pars += pars
                    if len(pars):
                        # Append found sweep param names to the sweep_points_map
                        # for this mobj
                        sweep_points_map[mobjn] += pars
                # Find the remaining sweep paramenter names in this dimention
                # that do not contain any of the mobj names. These are assumed
                # to be used by all measurement_objects and will be appended
                # in sweep_points_map for each mobj.
                remaining_pars = [k for k in list(d) if k not in all_pars]
                if len(remaining_pars):
                    for mobjn in measurement_objects:
                        sweep_points_map[mobjn] += remaining_pars
        return sweep_points_map

    def length(self, dimension='all'):
        """
        Returns the number of sweep points in a given sweep dimension (after a
        sanity checking).

        :param dimension: ('all' or int) sweep dimension (default: 'all').

        :return: (int) number of sweep points in the given dimension
        """

        if dimension == 'all':
            return [self.length(d) for d in range(len(self))]

        if len(self) == 0 or (dimension >= 0 and dimension >= len(self)):
            return 0
        n = 0
        for p in self[dimension].values():
            if n == 0:
                n = len(p[0])
            elif n != len(p[0]):
                raise ValueError('The lengths of the sweep points are not '
                                 'consistent.')
        return n

    def update(self, sweep_points):
        """
        Updates the sweep dictionaries of all dimensions with the sweep
        dictionaries passed as sweep_points. Non-exisiting
        parameters and required additional dimensions are added if needed.

        :param sweep_points: (SweepPoints) a SweepPoints object containing
            the sweep points to be updated.

        :return:
        """
        while len(self) < len(sweep_points):
            self.add_sweep_dimension()
        for d, u in zip(self, sweep_points):
            d.update(u)

    def find_parameter(self, param_name):
        """
        Returns the index of the first dimension in which a given sweep
        parameter is found.

        :param param_name: (str) name of the sweep parameter

        :return: (int or None) the index of the first dimension in which the
            parameter if found or None if no parameter with the given name
            exists.
        """
        for dim in range(len(self)):
            if param_name in self[dim]:
                return dim
        return None

    def subset(self, i, dimension=0):
        """
        Returns a new SweepPoints object with one of the dimensions reduced
        to a subset of the sweep values. The other dimensions are unchanged.
        :param i: (list) indices of the sweep values that should be
            contained in the subset.
        :param dimension: (int, default 0) index of the dimension that
            should be reduced
        """
        sp = SweepPoints(self)
        for k, v in sp[dimension].items():
            sp[dimension][k] = (np.array(v[0])[i], v[1], v[2])
        return sp

    def remove_sweep_parameter(self, param_name):
        """
        Removes a the sweep parameter with a given name from the SweepPoints
        object. If the parameter is not found, a warning is issued.
        :param param_name: (str) name of the sweep parameter to remove
        """
        dim = self.find_parameter(param_name)
        if dim is None:
            log.warning(f"remove_sweep_parameter: Sweep parameter "
                        f"{param_name} not found.")
        else:
            del self[dim][param_name]
<commit_msg>SweepPoints.get_meas_obj_sweep_points_map: proper fix to distinguish between, for ex 'qb1' and 'qb12' in the sweep parameter names<commit_after>import logging
log = logging.getLogger(__name__)
from collections import OrderedDict
from copy import copy, deepcopy
import numpy as np
from numpy import array  # Needed for eval. Do not remove.

class SweepPoints(list):
    """
    This class is used to create sweep points for any measurement.
    The SweepPoints object is a list of dictionaries of the form:
        [
            # 1st sweep dimension
            {param_name0: (values, unit, label),
             param_name1: (values, unit, label),
            ...
             param_nameN: (values, unit, label)},

            # 2nd sweep dimension
            {param_name0: (values, unit, label),
             param_name1: (values, unit, label),
            ...
             param_nameN: (values, unit, label)},

             .
             .
             .

            # D-th sweep dimension
            {param_name0: (values, unit, label),
             param_name1: (values, unit, label),
            ...
             param_nameN: (values, unit, label)},
        ]

    Example how to use this class to create a 2D sweep for 3 qubits, where
    the first (hard) sweep is over amplitudes and the :

    sp = SweepPoints()
    for qb in ['qb1', 'qb2', 'qb3']:
        sp.add_sweep_parameter(f'lengths_{qb}', np.linspace(10e-9, 1e-6, 80),
        's', 'Pulse length, $L$')
    sp.add_sweep_dimension()
    for qb in ['qb1', 'qb2', 'qb3']:
        sp.add_sweep_parameter(f'amps_{qb}', np.linspace(0, 1, 20),
        'V', 'Pulse amplitude, $A$')
    """
    def __init__(self, param=None, values=None, unit='', label=None,
                 dimension=-1, min_length=0):
        """
        Create a SweepPoints instance.
        :param param: list of dicts, repr of SweepPoints instance, or name of
            a sweep parameter
        :param values: list or array of numeric values of the sweep parameter
            specified by param
        :param unit: string specifying the unit of the sweep parameter
            specified by param
        :param label: str specifying the (latex style) label/name of the sweep
            parameter specified by param
        :param dimension: sweep dimension where sweep parameter param should be
            added
        :param min_length: minimum number of sweep dimensions to create
        """
        super().__init__()
        if isinstance(param, list):
            self.add_dict_list(param)
        elif isinstance(param, str):
            if values is not None:
                self.add_sweep_parameter(param, values, unit, label,
                                         dimension)
            else:
                self.add_dict_list(eval(param))

        while len(self) < min_length:
            self.add_sweep_dimension()

    def __getitem__(self, i):
        """
        Overloading of List.__getitem__ to ensure type SweepPoints is preserved.
        :param i: element or slice
        :return: element or new SweepPoints instance
        """
        if isinstance(i, str):
            new_data = self.get_sweep_params_property('values', 'all', i)
        else:
            new_data = super().__getitem__(i)
            if isinstance(i, slice):
                new_data = self.__class__(new_data)
        return new_data

    def add_dict_list(self, dict_list):
        """
        Append the dicts in dict_list to self.
        :param dict_list: list of dictionaries in the format of this class, or
            in the legacy format {param_name: {'values': ...,
                                               'unit': ...,
                                               'label': ...}}
        """
        for d in deepcopy(dict_list):
            if len(d) == 0 or isinstance(list(d.values())[0], tuple):
                # assume that dicts have the same format as this class
                self.append(d)
            else:
                # import from a list of sweep dicts in the old format
                self.append({k: (v['values'],
                                 v.get('unit', ''),
                                 v.get('label', k))
                             for k, v in d.items()})

    def add_sweep_parameter(self, param_name, values, unit='', label=None,
                            dimension=-1):
        """
        Adds sweep points to a given dimension.
        :param param_name: (str) parameter name
        :param values: (list or numpy array) sweep values
        :param unit: (optional str) unit of the values (default: '')
        :param label: (optional str) label e.g. for plots (default: param_name)
        :dim: the dimension to which the point should be added (default:
            last dimension)
        """
        if label is None:
            label = param_name
        assert self.find_parameter(param_name) is None, \
            f'A sweep parameter with name "{param_name}" already exists.'
        while len(self) == 0 or (dimension >= 0 and dimension >= len(self)):
            self.add_sweep_dimension()
        assert self.length(dimension) in [0, len(values)], \
            'Number of values has to match the length of existing sweep ' \
            'points.'
        self[dimension].update({param_name: (values, unit, label)})

    def add_sweep_dimension(self):
        self.append(dict())

    def get_sweep_dimension(self, dimension='all', pop=False):
        """
        Returns the sweep dict of the sweep dimension specified by dimension.
        :param dimension: int specifying a sweep dimension or
            the string 'all'
        :param pop: bool specifying whether to pop (True) or get(False) the
            sweep dimension.
        :return: self if dimension == 'all', else self[dimension]
        """
        if dimension == 'all':
            if pop:
                to_return = deepcopy(self)
                self.clear()
                return to_return
            else:
                return self
        else:
            if len(self) < dimension:
                raise ValueError(f'Dimension {dimension} not found.')
            to_return = self[dimension]
            if pop:
                self[dimension] = {}
            return to_return

    def get_sweep_params_description(self, param_names, dimension='all',
                                     pop=False):
        """
        Get the sweep tuples for the sweep parameters param_names if they are
        found in the sweep dimension dict specified by dimension.
        :param param_names: string or list of strings corresponding to keys in
            the dictionaries in self. Can also be 'all'
        :param dimension: 'all' or int specifying a sweep dimension
        :param pop: bool specifying whether to pop (True) or get(False) the
            sweep parameters.
        :return:
            If the param_names are found in self or self[dimension]:
            if param_names == 'all': list with all the sweep tuples
                in the sweep dimension dict specified by dimension.
            if param_names is string: string with the sweep tuples of each
                param_names in the sweep dimension dict specified by dimension.
            if param_names is list: list with the sweep tuples of each
                param_names in the sweep dimension dict specified by dimension.
            if param_names is None: string corresponding to the
                first sweep parameter in the sweep dimension dict
            If none of param_names are found, raises KeyError.
        """
        sweep_points_dim = self.get_sweep_dimension(
            dimension, pop=pop and param_names == 'all')
        is_list = True
        if param_names != 'all' and not isinstance(param_names, list):
            param_names = [param_names]
            is_list = False

        sweep_param_values = []
        if isinstance(sweep_points_dim, list):
            for sweep_dim_dict in sweep_points_dim:
                if param_names == 'all':
                    sweep_param_values += list(sweep_dim_dict.values())
                else:
                    for pn in param_names:
                        if pn in sweep_dim_dict:
                            sweep_param_values += [sweep_dim_dict.pop(pn) if pop
                                                   else sweep_dim_dict[pn]]
        else:  # it is a dict
            if param_names == 'all':
                sweep_param_values += list(sweep_points_dim.values())
            else:
                for pn in param_names:
                    if pn in sweep_points_dim:
                        sweep_param_values += [sweep_points_dim.pop(pn) if pop
                                               else sweep_points_dim[pn]]
        if len(sweep_param_values) == 0:
            s = "sweep points" if dimension == "all" else f'sweep dimension ' \
                                                          f'{dimension}'
            raise KeyError(f'{param_names} not found in {s}.')

        if is_list:
            return sweep_param_values
        else:
            return sweep_param_values[0]

    def get_sweep_params_property(self, property, dimension='all',
                                  param_names=None):
        """
        Get a property of the sweep parameters param_names in self.
        :param property: str with the name of a sweep param property. Can be
            "values", "unit", "label."
        :param dimension: 'all' or int specifying a sweep dimension
            (default 'all')
        :param param_names: None, or string or list of strings corresponding to
            keys in the sweep dimension specified by dimension.
            Can also be 'all'
        :return:
            if param_names == 'all': list with the property of all
                param_names in the sweep dimension dict specified by dimension.
            if param_names is string: the property of the sweep parameter
                specified in param_names in the sweep dimension dict specified
                by dimension.
            if param_names is list: list with the property of each
                param_names in the sweep dimension dict specified by dimension.
            if param_names is None: property corresponding to the
                first sweep parameter in the sweep dimension dict
        """
        properties_dict = {'values': 0, 'unit': 1, 'label': 2}
        sweep_points_dim = self.get_sweep_dimension(dimension)

        if param_names is None:
            if isinstance(sweep_points_dim, list):
                for sweep_dim_dict in sweep_points_dim:
                    if len(sweep_dim_dict) == 0:
                        return [] if property == 'values' else ''
                    else:
                        return next(iter(sweep_dim_dict.values()))[
                            properties_dict[property]]
            else:
                if len(sweep_points_dim) == 0:
                    return [] if property == 'values' else ''
                else:
                    return next(iter(sweep_points_dim.values()))[
                        properties_dict[property]]
        elif param_names == 'all' or isinstance(param_names, list):
            return [pnd[properties_dict[property]] for pnd in
                    self.get_sweep_params_description(param_names,
                                                      dimension)]
        else:
            return self.get_sweep_params_description(
                param_names, dimension)[properties_dict[property]]

    def get_meas_obj_sweep_points_map(self, measurement_objects):
        """
        Constructs the measurement-objects-sweep-points map as the dict
        {mobj_name: [sweep_param_name_0, ..., sweep_param_name_n]}

        If a sweep dimension has only one sweep parameter name (dict with only
        one key), then it assumes all mobjs use that sweep parameter name.

        If the sweep dimension has more than one sweep parameter name (dict with
        several keys), then:
            - first adds to the list for each mobj only those sweep
            param names that contain the mobj_name.

            ! Currently this only does the correct thing if each sweep
            parameter name contains only one measurement object name. So for
            example, qb1qb9_amplitude will not be found for both qb1 and qb9 !

            - if some parameters remain that do not contain any of the
            measurement_objects names, it is assumed that all
            measurement_objects used them so they will be added for each mobj

        :param measurement_objects: list of strings to be used as keys in the
            returned dictionary. These are the measured object names.
            Can also be list of measurement object instances with a name
            attribute, in which case this function gets the list of names.
        :return: dict of the form
         {mobj_name: [sweep_param_name_0, ..., sweep_param_name_n]}
        """

        # Ensure measurement_objects is a list
        if isinstance(measurement_objects, list):
            measurement_objects = copy(measurement_objects)
        else:
            measurement_objects = [measurement_objects]

        for i, mobj in enumerate(measurement_objects):
            if hasattr(mobj, 'name'):
                # A list of measurement object instances with a
                # name attribute was provided
                measurement_objects[i] = mobj.name

        # sort measurement_objects (which are not strings) by length
        # This is needed in the else statement below in order to distinguish
        # between, for example, 'qb1' and 'qb13' (or 'TWPA' and 'TWPA_A')
        measurement_objects.sort(key=len)

        sweep_points_map = {mobjn: [] for mobjn in measurement_objects}

        for dim, d in enumerate(self):
            # d is a dictionary with keys == sweep parameter names
            if len(d) == 1:
                # Only one sweep parameter.
                # Assume all mobjs use the same param_name given by the key
                # of d.
                for mobjn in measurement_objects:
                    sweep_points_map[mobjn] += [next(iter(d))]
            else:
                all_pars = []
                for i, mobjn in enumerate(measurement_objects):
                    # Find all sweep param names that contain the mobj name
                    pars = []
                    for k in list(d):
                        if mobjn in k:
                            # k really contains mobjn if none of the other,
                            # longer mobjn in measurement_objects are in k
                            longer_mobjns_not_in_k = all([
                                omon not in k for omon in
                                measurement_objects[i+1:]])
                            if longer_mobjns_not_in_k:
                                pars += [k]
                    if len(pars):
                        # Append found sweep param names to the sweep_points_map
                        # for this mobj
                        sweep_points_map[mobjn] += pars

                    # Collect all found pars, to be used below
                    all_pars += pars

                # Find the remaining sweep parameter names in this dimension
                # that do not contain any of the mobj names. These are assumed
                # to be used by all measurement_objects and will be appended
                # in sweep_points_map for each mobj.
                remaining_pars = [k for k in list(d) if k not in all_pars]
                if len(remaining_pars):
                    for mobjn in measurement_objects:
                        sweep_points_map[mobjn] += remaining_pars

        return sweep_points_map

    def length(self, dimension='all'):
        """
        Returns the number of sweep points in a given sweep dimension (after a
        sanity checking).

        :param dimension: ('all' or int) sweep dimension (default: 'all').

        :return: (int) number of sweep points in the given dimension
        """

        if dimension == 'all':
            return [self.length(d) for d in range(len(self))]

        if len(self) == 0 or (dimension >= 0 and dimension >= len(self)):
            return 0
        n = 0
        for p in self[dimension].values():
            if n == 0:
                n = len(p[0])
            elif n != len(p[0]):
                raise ValueError('The lengths of the sweep points are not '
                                 'consistent.')
        return n

    def update(self, sweep_points):
        """
        Updates the sweep dictionaries of all dimensions with the sweep
        dictionaries passed as sweep_points. Non-exisiting
        parameters and required additional dimensions are added if needed.

        :param sweep_points: (SweepPoints) a SweepPoints object containing
            the sweep points to be updated.

        :return:
        """
        while len(self) < len(sweep_points):
            self.add_sweep_dimension()
        for d, u in zip(self, sweep_points):
            d.update(u)

    def find_parameter(self, param_name):
        """
        Returns the index of the first dimension in which a given sweep
        parameter is found.

        :param param_name: (str) name of the sweep parameter

        :return: (int or None) the index of the first dimension in which the
            parameter if found or None if no parameter with the given name
            exists.
        """
        for dim in range(len(self)):
            if param_name in self[dim]:
                return dim
        return None

    def subset(self, i, dimension=0):
        """
        Returns a new SweepPoints object with one of the dimensions reduced
        to a subset of the sweep values. The other dimensions are unchanged.
        :param i: (list) indices of the sweep values that should be
            contained in the subset.
        :param dimension: (int, default 0) index of the dimension that
            should be reduced
        """
        sp = SweepPoints(self)
        for k, v in sp[dimension].items():
            sp[dimension][k] = (np.array(v[0])[i], v[1], v[2])
        return sp

    def remove_sweep_parameter(self, param_name):
        """
        Removes a the sweep parameter with a given name from the SweepPoints
        object. If the parameter is not found, a warning is issued.
        :param param_name: (str) name of the sweep parameter to remove
        """
        dim = self.find_parameter(param_name)
        if dim is None:
            log.warning(f"remove_sweep_parameter: Sweep parameter "
                        f"{param_name} not found.")
        else:
            del self[dim][param_name]
<|endoftext|>


<commit_before>from django.shortcuts import render
from django.shortcuts import redirect
from django.contrib import auth
from django.utils.translation import ugettext
from django.contrib import messages
from django.contrib.auth.decorators import login_required

from accounts.helpers import allauth as allauth_helpers
from . import forms as accounts_forms
from account.common_helpers import formrenderer

def login(request):
    openid_providers, form = allauth_helpers.get_openid_providers(request)
    if form:
        return allauth_helpers.openid_login(request, form)
    else:
        return render(request, 'accounts/login.html', {
            'auth_providers': allauth_helpers.get_auth_providers(request),
            'openid_providers': openid_providers,
        })

def logout(request):
    auth.logout(request)
    return redirect('index')

@login_required
def settings(request):
    if request.method == 'POST':
        form = accounts_forms.SettingsForm(request.POST, instance=request.user)
        if form.is_valid():
            form.save()
            messages.success(request, ugettext("Nustatymai buvo sėkimingai išsaugoti."))
            return redirect('accounts_settings')
    else:
        form = accounts_forms.SettingsForm(instance=request.user)
    return render(request, 'accounts/settings.html', {
        'form': formrenderer.render(request, form, title=ugettext("Profilio nustatymai"), submit=ugettext("Saugoti")),
    })
<commit_msg>Correct typo in views.py<commit_after>from django.shortcuts import render
from django.shortcuts import redirect
from django.contrib import auth
from django.utils.translation import ugettext
from django.contrib import messages
from django.contrib.auth.decorators import login_required

from account.helpers import allauth as allauth_helpers
from . import forms as accounts_forms
from account.common_helpers import formrenderer

def login(request):
    openid_providers, form = allauth_helpers.get_openid_providers(request)
    if form:
        return allauth_helpers.openid_login(request, form)
    else:
        return render(request, 'accounts/login.html', {
            'auth_providers': allauth_helpers.get_auth_providers(request),
            'openid_providers': openid_providers,
        })

def logout(request):
    auth.logout(request)
    return redirect('index')

@login_required
def settings(request):
    if request.method == 'POST':
        form = accounts_forms.SettingsForm(request.POST, instance=request.user)
        if form.is_valid():
            form.save()
            messages.success(request, ugettext("Nustatymai buvo sėkimingai išsaugoti."))
            return redirect('accounts_settings')
    else:
        form = accounts_forms.SettingsForm(instance=request.user)
    return render(request, 'accounts/settings.html', {
        'form': formrenderer.render(request, form, title=ugettext("Profilio nustatymai"), submit=ugettext("Saugoti")),
    })
<|endoftext|>


<commit_before><commit_msg>Added function to assay findings admin which will check if the description is a link (aka starts with http://) and make it a clickable URL if that is the case<commit_after><|endoftext|>


<commit_before>#
# extract voter_id and passwords for a particular email address
# may return many rows, if they all have the same email address
#
# python extract-passwords-for-email.py <election_uuid> <email_address>
#

from django.core.management import setup_environ
import settings, sys, csv

setup_environ(settings)

from helios.models import *

election_uuid = sys.argv[1]
email = sys.argv[2]

csv_output = csv.writer(sys.stdout)

voters = Election.objects.get(uuid=election_uuid).voter_set.filter(voter_email=email)

for voter in voters:
    csv_output.writerow([voter.voter_id, voter.voter_password])

<commit_msg>more on password extraction<commit_after>#
# extract voter_id and passwords for a particular email address
# may return many rows, if they all have the same email address
#
# python extract-passwords-for-email.py <election_uuid> <email_address>
#

from django.core.management import setup_environ
import settings, sys, csv

setup_environ(settings)

from helios.models import *

election_uuid = sys.argv[1]
email = sys.argv[2]

csv_output = csv.writer(sys.stdout)

voters = Election.objects.get(uuid=election_uuid).voter_set.filter(voter_email=email)

for voter in voters:
    csv_output.writerow([voter.voter_login_id, voter.voter_password])

<|endoftext|>


<commit_before>import io
import unittest
import xml.etree.ElementTree as ET

from cumulusci.core.exceptions import TaskOptionsError
from cumulusci.tasks.salesforce.tests.util import create_task
from cumulusci.tasks.salesforce.metadata_etl import AddValueSetEntries

VALUESET_XML = """<?xml version="1.0" encoding="UTF-8"?>
<StandardValueSet xmlns="http://soap.sforce.com/2006/04/metadata">
    <sorted>false</sorted>
    <standardValue>
        <fullName>Value</fullName>
        <default>true</default>
        Value
    </standardValue>
    <standardValue>
        <fullName>Other</fullName>
        <default>false</default>
        Other
    </standardValue>
</StandardValueSet>
"""

class test_AddValueSetEntries(unittest.TestCase):
    def test_adds_entry(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [
                    {"fullName": "Test", "label": "Label"},
                    {"fullName": "Test_2", "label": "Label 2"},
                ],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test']", namespaces))
            == 0
        )
        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces))
            == 0
        )

        result = task._transform_entity(root, "ValueSet")

        entry = result.findall(".//sf:standardValue[sf:fullName='Test']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"

        entry = result.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label 2"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"

    def test_adds_entry__opportunitystage(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "OpportunityStage",
                "entries": [
                    {
                        "fullName": "Test",
                        "label": "Label",
                        "closed": True,
                        "won": True,
                        "forecastCategory": "Omitted",
                        "probability": 100,
                    }
                ],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test']", namespaces))
            == 0
        )
        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces))
            == 0
        )

        result = task._transform_entity(root, "OpportunityStage")

        entry = result.findall(".//sf:standardValue[sf:fullName='Test']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"
        closed = entry[0].findall(f".//sf:closed", namespaces)
        assert len(closed) == 1
        assert closed[0].text == "true"
        won = entry[0].findall(f".//sf:won", namespaces)
        assert len(won) == 1
        assert won[0].text == "true"
        forecastCategory = entry[0].findall(f".//sf:forecastCategory", namespaces)
        assert len(forecastCategory) == 1
        assert forecastCategory[0].text == "Omitted"
        probability = entry[0].findall(f".//sf:probability", namespaces)
        assert len(probability) == 1
        assert probability[0].text == "100"

    def test_adds_entry__casestatus(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "CaseStatus",
                "entries": [{"fullName": "Test", "label": "Label", "closed": True}],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test']", namespaces))
            == 0
        )
        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces))
            == 0
        )

        result = task._transform_entity(root, "CaseStatus")

        entry = result.findall(".//sf:standardValue[sf:fullName='Test']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"
        closed = entry[0].findall(f".//sf:closed", namespaces)
        assert len(closed) == 1
        assert closed[0].text == "true"

    def test_does_not_add_existing_entry(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [{"fullName": "Value", "label": "Label"}],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Value']", namespaces))
            == 1
        )

        result = task._transform_entity(root, "ValueSet")

        assert (
            len(result.findall(".//sf:standardValue[sf:fullName='Value']", namespaces))
            == 1
        )

    def test_raises_exception_missing_values(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [{"fullName": "Value"}],
            },
        )
        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))

        with self.assertRaises(TaskOptionsError):
            task._transform_entity(root, "ValueSet")

        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [{"label": "Value"}],
            },
        )

        with self.assertRaises(TaskOptionsError):
            task._transform_entity(root, "ValueSet")

    def test_raises_exception_missing_values__opportunitystage(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "OpportunityStage",
                "entries": [{"fullName": "Value"}],
            },
        )
        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))

        with self.assertRaises(TaskOptionsError):
            task._transform_entity(root, "OpportunityStage")

    def test_raises_exception_missing_values__casestatus(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "CaseStatus",
                "entries": [{"fullName": "Value"}],
            },
        )
        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))

        with self.assertRaises(TaskOptionsError):
            task._transform_entity(root, "CaseStatus")
<commit_msg>Correct error-handling tests<commit_after>import io
import unittest
import xml.etree.ElementTree as ET

from cumulusci.core.exceptions import TaskOptionsError
from cumulusci.tasks.salesforce.tests.util import create_task
from cumulusci.tasks.salesforce.metadata_etl import AddValueSetEntries

VALUESET_XML = """<?xml version="1.0" encoding="UTF-8"?>
<StandardValueSet xmlns="http://soap.sforce.com/2006/04/metadata">
    <sorted>false</sorted>
    <standardValue>
        <fullName>Value</fullName>
        <default>true</default>
        Value
    </standardValue>
    <standardValue>
        <fullName>Other</fullName>
        <default>false</default>
        Other
    </standardValue>
</StandardValueSet>
"""

class test_AddValueSetEntries(unittest.TestCase):
    def test_adds_entry(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [
                    {"fullName": "Test", "label": "Label"},
                    {"fullName": "Test_2", "label": "Label 2"},
                ],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test']", namespaces))
            == 0
        )
        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces))
            == 0
        )

        result = task._transform_entity(root, "ValueSet")

        entry = result.findall(".//sf:standardValue[sf:fullName='Test']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"

        entry = result.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label 2"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"

    def test_adds_entry__opportunitystage(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "OpportunityStage",
                "entries": [
                    {
                        "fullName": "Test",
                        "label": "Label",
                        "closed": True,
                        "won": True,
                        "forecastCategory": "Omitted",
                        "probability": 100,
                    }
                ],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test']", namespaces))
            == 0
        )
        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces))
            == 0
        )

        result = task._transform_entity(root, "OpportunityStage")

        entry = result.findall(".//sf:standardValue[sf:fullName='Test']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"
        closed = entry[0].findall(f".//sf:closed", namespaces)
        assert len(closed) == 1
        assert closed[0].text == "true"
        won = entry[0].findall(f".//sf:won", namespaces)
        assert len(won) == 1
        assert won[0].text == "true"
        forecastCategory = entry[0].findall(f".//sf:forecastCategory", namespaces)
        assert len(forecastCategory) == 1
        assert forecastCategory[0].text == "Omitted"
        probability = entry[0].findall(f".//sf:probability", namespaces)
        assert len(probability) == 1
        assert probability[0].text == "100"

    def test_adds_entry__casestatus(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "CaseStatus",
                "entries": [{"fullName": "Test", "label": "Label", "closed": True}],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test']", namespaces))
            == 0
        )
        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Test_2']", namespaces))
            == 0
        )

        result = task._transform_entity(root, "CaseStatus")

        entry = result.findall(".//sf:standardValue[sf:fullName='Test']", namespaces)
        assert len(entry) == 1
        label = entry[0].findall(f".//sf:label", namespaces)
        assert len(label) == 1
        assert label[0].text == "Label"
        default = entry[0].findall(f".//sf:default", namespaces)
        assert len(default) == 1
        assert default[0].text == "false"
        closed = entry[0].findall(f".//sf:closed", namespaces)
        assert len(closed) == 1
        assert closed[0].text == "true"

    def test_does_not_add_existing_entry(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [{"fullName": "Value", "label": "Label"}],
            },
        )

        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))
        namespaces = {"sf": "http://soap.sforce.com/2006/04/metadata"}

        assert (
            len(root.findall(".//sf:standardValue[sf:fullName='Value']", namespaces))
            == 1
        )

        result = task._transform_entity(root, "ValueSet")

        assert (
            len(result.findall(".//sf:standardValue[sf:fullName='Value']", namespaces))
            == 1
        )

    def test_raises_exception_missing_values(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [{"fullName": "Value"}],
            },
        )
        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))

        with self.assertRaises(TaskOptionsError):
            task._transform_entity(root, "ValueSet")

        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "bar,foo",
                "entries": [{"label": "Value"}],
            },
        )

        with self.assertRaises(TaskOptionsError):
            task._transform_entity(root, "ValueSet")

    def test_raises_exception_missing_values__opportunitystage(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "OpportunityStage",
                "entries": [{"fullName": "Value", "label": "Value"}],
            },
        )
        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))

        with self.assertRaises(TaskOptionsError) as err:
            task._transform_entity(root, "OpportunityStage")
            assert "OpportunityStage" in err

    def test_raises_exception_missing_values__casestatus(self):
        task = create_task(
            AddValueSetEntries,
            {
                "unmanaged": False,
                "api_version": "47.0",
                "api_names": "CaseStatus",
                "entries": [{"fullName": "Value", "label": "Value"}],
            },
        )
        root = ET.ElementTree(file=io.StringIO(VALUESET_XML))

        with self.assertRaises(TaskOptionsError) as err:
            task._transform_entity(root, "CaseStatus")
            assert "CaseStatus" in err
<|endoftext|>


<commit_before>
import inspect
from functools import wraps
from stoplight.exceptions import *
from stoplight.rule import *

def validate(**rules):
    """Validation endpoint decorator

    This decorator allows validation of input from user
    API endpoints. This allows separation of business logic
    and avoids convoluted validation logic.

    Typical use is as follows:

    @validate({'bird_id': val_bird_id(allow_none=True), 404)
    def get_one(self, bird_id):
        lookup_bird_data(bird_id)

    In this example, bird_id will be passed to val_bird_id
    and validated. If the validation function throws the
    ValidationFailed exception, the specified code will be returned
    in the response and the resultant function will never actually
    be called.
    """
    def _validate(f):
        @wraps(f)
        def wrapper(*args, **kwargs):

            funcparams = inspect.getargspec(f)

            # Holds the list of validated values. Only
            # these values are passed to the endpoint
            outargs = dict()

            # Create dictionary that maps parameters passed
            # to their values passed
            param_values = dict(zip(funcparams.args, args))

            # Bring in kwargs so that we can validate as well
            param_values.update(kwargs)

            for param, rule in rules.items():

                # Where can we get the value? It's either
                # the getter on the rule or we default
                # to verifying parameters.
                getval = rule.getter or param_values.get

                # Call the validation function, passing
                # the value was retrieved from the getter
                try:
                    value = getval(param)

                    # Ensure that this validation function
                    # did not return a funciton. This
                    # checks that the user did not forget to
                    # execute the outer function of a closure
                    # in the rule declaration

                    resp = rule.vfunc(value)

                    if inspect.isfunction(resp):
                        msg = 'Validation function returned a function.'
                        raise ValidationProgrammingError(msg)

                    # If this is a param rule, add the
                    # param to the list of out args
                    if rule.getter is None:
                        outargs[param] = value

                except ValidationFailed as ex:
                    rule.errfunc()
                    return

            assert funcparams.args[0] == 'self'

            # Validation was successful, call the wrapped function
            return f(args[0], **outargs)
        return wrapper
    return _validate

def validation_function(func):
    """Decorator for creating a validation function"""
    @wraps(func)
    def inner(none_ok=False, empty_ok=False):
        def wrapper(value, **kwargs):
            if none_ok and value is None:
                return

            if not none_ok and value is None:
                msg = 'None value not permitted'
                raise ValidationFailed(msg)

            if empty_ok and value == '':
                return

            if not empty_ok and value == '':
                msg = 'Empty value not permitted'
                raise ValidationFailed(msg)
            func(value)
        return wrapper
    return inner
<commit_msg>Adds whitespace<commit_after>
import inspect
from functools import wraps
from stoplight.exceptions import *
from stoplight.rule import *

def validate(**rules):
    """Validation endpoint decorator

    This decorator allows validation of input from user
    API endpoints. This allows separation of business logic
    and avoids convoluted validation logic.

    Typical use is as follows:

    @validate({'bird_id': val_bird_id(allow_none=True), 404)
    def get_one(self, bird_id):
        lookup_bird_data(bird_id)

    In this example, bird_id will be passed to val_bird_id
    and validated. If the validation function throws the
    ValidationFailed exception, the specified code will be returned
    in the response and the resultant function will never actually
    be called.
    """
    def _validate(f):
        @wraps(f)
        def wrapper(*args, **kwargs):

            funcparams = inspect.getargspec(f)

            # Holds the list of validated values. Only
            # these values are passed to the endpoint
            outargs = dict()

            # Create dictionary that maps parameters passed
            # to their values passed
            param_values = dict(zip(funcparams.args, args))

            # Bring in kwargs so that we can validate as well
            param_values.update(kwargs)

            for param, rule in rules.items():

                # Where can we get the value? It's either
                # the getter on the rule or we default
                # to verifying parameters.
                getval = rule.getter or param_values.get

                # Call the validation function, passing
                # the value was retrieved from the getter
                try:
                    value = getval(param)

                    # Ensure that this validation function
                    # did not return a funciton. This
                    # checks that the user did not forget to
                    # execute the outer function of a closure
                    # in the rule declaration

                    resp = rule.vfunc(value)

                    if inspect.isfunction(resp):
                        msg = 'Validation function returned a function.'
                        raise ValidationProgrammingError(msg)

                    # If this is a param rule, add the
                    # param to the list of out args
                    if rule.getter is None:
                        outargs[param] = value

                except ValidationFailed as ex:
                    rule.errfunc()
                    return

            assert funcparams.args[0] == 'self'

            # Validation was successful, call the wrapped function
            return f(args[0], **outargs)
        return wrapper
    return _validate

def validation_function(func):
    """Decorator for creating a validation function"""
    @wraps(func)
    def inner(none_ok=False, empty_ok=False):
        def wrapper(value, **kwargs):
            if none_ok and value is None:
                return

            if not none_ok and value is None:
                msg = 'None value not permitted'
                raise ValidationFailed(msg)

            if empty_ok and value == '':
                return

            if not empty_ok and value == '':
                msg = 'Empty value not permitted'
                raise ValidationFailed(msg)

            func(value)
        return wrapper
    return inner
<|endoftext|>


<commit_before>from tensorflow.contrib.learn.python.learn.estimators._sklearn import accuracy_score
from tensorflow.contrib import learn
from tensorflow.contrib import layers
#from tensorflow.contrib import metrics
from tensorflow.contrib import framework
from tensorflow.contrib.learn import monitors
from tensorflow.python.platform import tf_logging as logging
import tensorflow as tf
import datetime
import numpy as np

tf.set_random_seed(0)

def learning_rate(lr, i):
    max_learning_rate = tf.constant(0.003, dtype=tf.float32)
    min_learning_rate = tf.constant(0.0001, dtype=tf.float32)
    decay_speed = tf.constant(2000.0, dtype=tf.float32)
    lr = tf.add(min_learning_rate, tf.mul(tf.sub(max_learning_rate, min_learning_rate), tf.exp(tf.div(tf.to_float(-i),decay_speed))))
    return lr

### Download and load MNIST data.

mnist = learn.datasets.load_dataset('mnist')

#hello
### Convolutional network
def conv_model(X, Y_, mode):
    XX = tf.reshape(X, [-1, 28, 28, 1])
    Y1 = layers.conv2d(XX,  num_outputs=6,  kernel_size=[6, 6], normalizer_fn=layers.batch_norm)
    Y2 = layers.conv2d(Y1, num_outputs=12, kernel_size=[5, 5], stride=2, normalizer_fn=layers.batch_norm)
    Y3 = layers.conv2d(Y2, num_outputs=24, kernel_size=[4, 4], stride=2, normalizer_fn=layers.batch_norm)
    Y4 = layers.flatten(Y3)
    Y5 = layers.relu(Y4, 200, normalizer_fn=layers.batch_norm)
    #Y6 = layers.dropout(Y5, keep_prob=0.75, is_training=(mode ==learn.ModeKeys.TRAIN))
    Ylogits = layers.linear(Y5, 10)
    #prediction = tf.arg_max(tf.nn.softmax(Ylogits), 1)
    ## dense to one-hot
    Y__ = tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, Y__))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam", learning_rate_decay_fn=learning_rate)
    #train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam")

    # 10000 iterations with batch normalization on all layers, no dropout : final accuracy 99.42%
    # 10000 iterations with batch normalization on all layers, dropout on fully connected layer : final accuracy 99.28%
    return Ylogits, loss, train_op

def conv_model2(X, Y_, mode):
    XX = tf.reshape(X, [-1, 28, 28, 1])
    Y1 = layers.conv2d(XX,  num_outputs=6,  kernel_size=[6, 6], normalizer_fn=layers.batch_norm)
    Y2 = layers.conv2d(Y1, num_outputs=12, kernel_size=[5, 5], stride=2, normalizer_fn=layers.batch_norm)
    Y3 = layers.conv2d(Y2, num_outputs=24, kernel_size=[4, 4], stride=2, normalizer_fn=layers.batch_norm)
    Y4 = layers.flatten(Y3)
    Y5 = layers.relu(Y4, 200, normalizer_fn=layers.batch_norm)
    Ylogits = layers.linear(Y5, 10)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam", learning_rate_decay_fn=learning_rate)
    # 10000 iterations with batch normalization on all layers, 4 conv layers, 2 normal layers, no dropout : final accuracy 99.38%
    # (layer parameters 6x6x8s1, 5x5x16s2, 4x4x24s2, 3x3x32s1, 400, 200)
    # 10000 iterations with batch normalization on all layers, 3 conv layers, 2 normal layers, no dropout : final accuracy 99.27%
    # (layer parameters 6x6x8s1, 5x5x16s2, 4x4x32s2, 400, 200)
    # 10000 iterations with batch normalization on all layers, 3 conv layers, 2 normal layers, no dropout : final accuracy 99.33%
    # (layer parameters 6x6x6s1, 5x5x12s2, 4x4x24s2, 200, 50)
    return Ylogits, loss, train_op

def softmax_model(X, Y_, mode):
    Ylogits = layers.linear(X, 10)
    #prediction = tf.nn.softmax(Ylogits)
    Y__ = tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, Y__))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam")
    return Ylogits, loss, train_op

def dense_model(X, Y_, mode):
    Y1 = layers.fully_connected(X, 200, normalizer_fn=layers.batch_norm)
    Ylogits = layers.linear(Y1, 10)
    #prediction = tf.nn.softmax(Ylogits)
    Y__ = tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, Y__))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam")
    return Ylogits, loss, train_op

logging.set_verbosity(logging.INFO)

# Training and predicting
dense_train_labels = mnist.train.labels.astype(np.int64)
dense_test_labels = mnist.test.labels.astype(np.int64)
#mon = monitors.ValidationMonitor(mnist.test.images, dense_test_labels, every_n_steps=100)
classifier = learn.Classifier(model_fn=conv_model, n_classes=10, model_dir="logs/run"+str(datetime.datetime.now().timestamp()))
classifier.fit(mnist.train.images, dense_train_labels, batch_size=100, steps=10000)
score = accuracy_score(mnist.test.labels, classifier.predict(mnist.test.images))
print('Accuracy: {0:f}'.format(score))

## code graveyard

#classifier = learn.TensorFlowEstimator(model_fn=conv_model, n_classes=10, batch_size=100, optimizer="Adam", learning_rate=learning_rate, verbose=2)
#one_hot_train_labels = learn.datasets.mnist.dense_to_one_hot(mnist.train.labels, 10).astype(np.float32)
#one_hot_test_labels = learn.datasets.mnist.dense_to_one_hot(mnist.test.labels, 10).astype(np.float32)<commit_msg>removed dummy commit<commit_after>from tensorflow.contrib.learn.python.learn.estimators._sklearn import accuracy_score
from tensorflow.contrib import learn
from tensorflow.contrib import layers
#from tensorflow.contrib import metrics
from tensorflow.contrib import framework
from tensorflow.contrib.learn import monitors
from tensorflow.python.platform import tf_logging as logging
import tensorflow as tf
import datetime
import numpy as np

tf.set_random_seed(0)

def learning_rate(lr, i):
    max_learning_rate = tf.constant(0.003, dtype=tf.float32)
    min_learning_rate = tf.constant(0.0001, dtype=tf.float32)
    decay_speed = tf.constant(2000.0, dtype=tf.float32)
    lr = tf.add(min_learning_rate, tf.mul(tf.sub(max_learning_rate, min_learning_rate), tf.exp(tf.div(tf.to_float(-i),decay_speed))))
    return lr

### Download and load MNIST data.

mnist = learn.datasets.load_dataset('mnist')

### Convolutional network
def conv_model(X, Y_, mode):
    XX = tf.reshape(X, [-1, 28, 28, 1])
    Y1 = layers.conv2d(XX,  num_outputs=6,  kernel_size=[6, 6], normalizer_fn=layers.batch_norm)
    Y2 = layers.conv2d(Y1, num_outputs=12, kernel_size=[5, 5], stride=2, normalizer_fn=layers.batch_norm)
    Y3 = layers.conv2d(Y2, num_outputs=24, kernel_size=[4, 4], stride=2, normalizer_fn=layers.batch_norm)
    Y4 = layers.flatten(Y3)
    Y5 = layers.relu(Y4, 200, normalizer_fn=layers.batch_norm)
    #Y6 = layers.dropout(Y5, keep_prob=0.75, is_training=(mode ==learn.ModeKeys.TRAIN))
    Ylogits = layers.linear(Y5, 10)
    #prediction = tf.arg_max(tf.nn.softmax(Ylogits), 1)
    ## dense to one-hot
    Y__ = tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, Y__))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam", learning_rate_decay_fn=learning_rate)
    #train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam")

    # 10000 iterations with batch normalization on all layers, no dropout : final accuracy 99.42%
    # 10000 iterations with batch normalization on all layers, dropout on fully connected layer : final accuracy 99.28%
    return Ylogits, loss, train_op

def conv_model2(X, Y_, mode):
    XX = tf.reshape(X, [-1, 28, 28, 1])
    Y1 = layers.conv2d(XX,  num_outputs=6,  kernel_size=[6, 6], normalizer_fn=layers.batch_norm)
    Y2 = layers.conv2d(Y1, num_outputs=12, kernel_size=[5, 5], stride=2, normalizer_fn=layers.batch_norm)
    Y3 = layers.conv2d(Y2, num_outputs=24, kernel_size=[4, 4], stride=2, normalizer_fn=layers.batch_norm)
    Y4 = layers.flatten(Y3)
    Y5 = layers.relu(Y4, 200, normalizer_fn=layers.batch_norm)
    Ylogits = layers.linear(Y5, 10)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam", learning_rate_decay_fn=learning_rate)
    # 10000 iterations with batch normalization on all layers, 4 conv layers, 2 normal layers, no dropout : final accuracy 99.38%
    # (layer parameters 6x6x8s1, 5x5x16s2, 4x4x24s2, 3x3x32s1, 400, 200)
    # 10000 iterations with batch normalization on all layers, 3 conv layers, 2 normal layers, no dropout : final accuracy 99.27%
    # (layer parameters 6x6x8s1, 5x5x16s2, 4x4x32s2, 400, 200)
    # 10000 iterations with batch normalization on all layers, 3 conv layers, 2 normal layers, no dropout : final accuracy 99.33%
    # (layer parameters 6x6x6s1, 5x5x12s2, 4x4x24s2, 200, 50)
    return Ylogits, loss, train_op

def softmax_model(X, Y_, mode):
    Ylogits = layers.linear(X, 10)
    #prediction = tf.nn.softmax(Ylogits)
    Y__ = tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, Y__))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam")
    return Ylogits, loss, train_op

def dense_model(X, Y_, mode):
    Y1 = layers.fully_connected(X, 200, normalizer_fn=layers.batch_norm)
    Ylogits = layers.linear(Y1, 10)
    #prediction = tf.nn.softmax(Ylogits)
    Y__ = tf.one_hot(tf.cast(Y_, tf.int32), 10, dtype=tf.float32)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(Ylogits, Y__))
    train_op = layers.optimize_loss(loss, framework.get_global_step(), 0.003, "Adam")
    return Ylogits, loss, train_op

logging.set_verbosity(logging.INFO)

# Training and predicting
dense_train_labels = mnist.train.labels.astype(np.int64)
dense_test_labels = mnist.test.labels.astype(np.int64)
#mon = monitors.ValidationMonitor(mnist.test.images, dense_test_labels, every_n_steps=100)
classifier = learn.Classifier(model_fn=conv_model, n_classes=10, model_dir="logs/run"+str(datetime.datetime.now().timestamp()))
classifier.fit(mnist.train.images, dense_train_labels, batch_size=100, steps=10000)
score = accuracy_score(mnist.test.labels, classifier.predict(mnist.test.images))
print('Accuracy: {0:f}'.format(score))

## code graveyard

#classifier = learn.TensorFlowEstimator(model_fn=conv_model, n_classes=10, batch_size=100, optimizer="Adam", learning_rate=learning_rate, verbose=2)
#one_hot_train_labels = learn.datasets.mnist.dense_to_one_hot(mnist.train.labels, 10).astype(np.float32)
#one_hot_test_labels = learn.datasets.mnist.dense_to_one_hot(mnist.test.labels, 10).astype(np.float32)<|endoftext|>


<commit_before>import json
import logging
try:
    from gevent import socket
    from gevent import Greenlet as Concurrent
except ImportError:
    import socket
    from threading import Thread as Concurrent
    print('Cannot find gevent, using threading')
import sys
import re

class ProxyType:
    NONE = 0
    SOCKS = 1
    CONNECT = 2

class Pipe(Concurrent):

    def __init__(self):
        Concurrent.__init__(self)

    def setSockPair(self, sockIn, sockOut):
        self.sockIn = sockIn
        self.sockOut = sockOut

    def pipeData(self):
        try:
            while True:
                data = self.sockIn.recv(65536)
                if data == b'' or data == '' or data == None:
                    break
                self.sockOut.send(data)
        except Exception:
            logging.info('Pipe end')
        finally:
            self.sockIn.close()
            self.sockOut.close()

    def run(self):
        self.pipeData()

    def _run(self):
        self.run()

class Tunnel(Pipe):

    def __init__(self, config, client):
        Pipe.__init__(self)
        self.config = config
        self.client = client

    def run(self):

        def recvFully(sock, byteslen):
            buf = b''
            while byteslen != 0:
                t += sock.recv(byteslen)
                if t == b'':
                    raise Exception('End connection in socksHandshake')
                buf += t
                byteslen -= len(t)
            return buf

        def socksHandshake():
            raise Exception('Not implemented yet')  #   TODO:

        def connectHandshake():

            def recvHeaderFully(sock):
                endBytes = b'\r\n\r\n'
                idx = 0
                buf = b''
                while True:
                    t = self.client.recv(65536)
                    if t == b'':
                        raise Exception('End connection in connectHandshake')
                    try:
                        idx = buf.index(endBytes, idx)
                        idx += 4
                        return buf[ : idx], buf[idx : ]
                    except ValueError:
                        idx = len(buf) - len(endBytes) + 1

            request, clientData = recvHeaderFully(self.client)
            method, host, protocol = re.split(b'\\s+', request[ : request.index(b'\r\n')])    #   Although not meets RFC, it's no matter.  Because it's a local proxy.
            if method.upper() != b'CONNECT':
                raise Exception('Not a CONNECT(HTTPS) proxy.')
            self.hostname, self.port = host.split(b':')
            self.port = int(self.port)
            self.parent.send(request)
            response, parentData = recvHeaderFully(self.parent)
            self.client.send(response)
            return clientData, parentData

        def startPipe():
            if clientData != b'':
                self.parent.send(clientData)
            if parentData != b'':
                self.client.send(parentData)
            pipe = Pipe()
            pipe.setSockPair(client, parent)
            pipe.start()
            self.setSockPair(parent, client)
            self.pipeData()

        def sslHandshake(clientData, parentData):

            def pipeHello(paramList):
                sockIn, inData, sockOut = paramList
                missDataLen = 5 - len(inData)
                if missDataLen > 0:
                    inData += recvFully(sockIn, missDataLen)
                helloLen = ord(inData[3]) << 8 | ord(inData[4])
                missDataLen = 5 + helloLen - len(inData)
                if missDataLen > 0:
                    inData += recvFully(sockIn, missDataLen)
                packetLen = 5 + helloLen
                sockOut.send(inData[ : packetLen])
                return inData[packetLen : ]

            clientData = pipeHello(self.client, clientData, self.parent)
            parentData = pipeHello(self.parent, parentData, self,client)
            raise Exception('Not implemented yet')  #   TODO:
            return clientData, parentData

        self.parent = socket.socket()
        try:
            if self.config['parentProxyHost'] != None:
                self.parent.connect((self.config['parentProxyHost'], self.config['parentProxyPort']))
            if self.config['parentProxyType'] == ProxyType.NONE:
                raise Exception('Not implemented yet')  #   TODO: Implement as local socks5 proxy
                clientData = b''
                parentData = b''
            elif self.config['parentProxyType'] == ProxyType.SOCKS:
                socksHandshake()
                clientData = b''
                parentData = b''
            elif self.config['parentProxyType'] == ProxyType.CONNECT:
                clientData, parentData = connectHandshake()
            else:
                assert(False)
            if clientData == b'':
                clientData = self.client.recv(65536)
            if clientData[0] != '\x22':    #   Not SSL
                startPipe()
            clientData, parentData = sslHandshake(clientData, parentData)
            startPipe()
        except Exception:
            logging.exception('Exception in Tunnel.run:')
        finally:
            self.client.close()
            self.parent.close()

class Main(object):

    def __init__(self, config):
        logging.basicConfig(filename = config['logFilename'], level = getattr(logging, config['logLevel']))
        config['parentProxyType'] = config['parentProxyType'].lower()
        if config['parentProxyType'] == 'none':
            config['parentProxyType'] = ProxyType.NONE
        elif config['parentProxyType'] == 'socks':
            config['parentProxyType'] = ProxyType.SOCKS
        elif config['parentProxyType'] == 'socks':
            config['parentProxyType'] = ProxyType.CONNECT
        else:
            logging.error('Unknown parentProxyType %s', config['parentProxyType'])
            raise Exception('Unknown parentProxyType %s' % (config['parentProxyType'], ))
        self.config = config

    def start(self):
        server = socket.socket()
        server.bind((config['proxyHost'], config['proxyPort']))
        server.listen(50)
        while True:
            try:
                client, addr = server.accept()
                Tunnel(self.config, client).start()
            except Exception:
                logging.exception('Exception in Main.start:')

if __name__ == '__main__':
    Main(json.loads(open(sys.argv[1], 'r').read())).start()
<commit_msg>start check ssl certification process<commit_after>import json
import logging
try:
    from gevent import socket
    from gevent import Greenlet as Concurrent
except ImportError:
    import socket
    from threading import Thread as Concurrent
    print('Cannot find gevent, using threading')
import sys
import re

class ProxyType:
    NONE = 0
    SOCKS = 1
    CONNECT = 2

class Pipe(Concurrent):

    def __init__(self):
        Concurrent.__init__(self)

    def setSockPair(self, sockIn, sockOut):
        self.sockIn = sockIn
        self.sockOut = sockOut

    def pipeData(self):
        try:
            while True:
                data = self.sockIn.recv(65536)
                if data == b'' or data == '' or data == None:
                    break
                self.sockOut.send(data)
        except Exception:
            logging.info('Pipe end')
        finally:
            self.sockIn.close()
            self.sockOut.close()

    def run(self):
        self.pipeData()

    def _run(self):
        self.run()

class Tunnel(Pipe):

    def __init__(self, config, client):
        Pipe.__init__(self)
        self.config = config
        self.client = client

    def run(self):

        def recvFully(sock, byteslen):
            buf = b''
            while byteslen != 0:
                t += sock.recv(byteslen)
                if t == b'':
                    raise Exception('End connection in socksHandshake')
                buf += t
                byteslen -= len(t)
            return buf

        def socksHandshake():
            raise Exception('Not implemented yet')  #   TODO:

        def connectHandshake():

            def recvHeaderFully(sock):
                endBytes = b'\r\n\r\n'
                idx = 0
                buf = b''
                while True:
                    t = self.client.recv(65536)
                    if t == b'':
                        raise Exception('End connection in connectHandshake')
                    try:
                        idx = buf.index(endBytes, idx)
                        idx += 4
                        return buf[ : idx], buf[idx : ]
                    except ValueError:
                        idx = len(buf) - len(endBytes) + 1

            request, clientData = recvHeaderFully(self.client)
            method, host, protocol = re.split(b'\\s+', request[ : request.index(b'\r\n')])    #   Although not meets RFC, it's no matter.  Because it's a local proxy.
            if method.upper() != b'CONNECT':
                raise Exception('Not a CONNECT(HTTPS) proxy.')
            self.hostname, self.port = host.split(b':')
            self.port = int(self.port)
            self.parent.send(request)
            response, parentData = recvHeaderFully(self.parent)
            self.client.send(response)
            return clientData, parentData

        def startClientParentPipe():
            if clientData != b'':
                self.parent.send(clientData)
            pipe = Pipe()
            pipe.setSockPair(self.client, self.parent)
            pipe.start()

        def startParentClientPipe():
            if parentData != b'':
                self.client.send(parentData)
            self.setSockPair(self.parent, self.client)
            self.pipeData()

        def startPipe():
            startClientParentPipe()
            startParentClientPipe()

        def sslCheckCertification(packet):
            raise Exception('Not implemented yet')  #   TODO:

        def sslGetPacket(sock, data):
            missDataLen = 5 - len(data)
            if missDataLen > 0:
                data += recvFully(sock, missDataLen)
            helloLen = ord(data[3]) << 8 | ord(data[4])
            missDataLen = 5 + helloLen - len(data)
            if missDataLen > 0:
                data += recvFully(sock, missDataLen)
            packetLen = 5 + helloLen
            return data[ : packetLen], data[packetLen : ]

        self.parent = socket.socket()
        try:
            if self.config['parentProxyHost'] != None:
                self.parent.connect((self.config['parentProxyHost'], self.config['parentProxyPort']))
            if self.config['parentProxyType'] == ProxyType.NONE:
                raise Exception('Not implemented yet')  #   TODO: Implement as local socks5 proxy
                clientData = b''
                parentData = b''
            elif self.config['parentProxyType'] == ProxyType.SOCKS:
                socksHandshake()
                clientData = b''
                parentData = b''
            elif self.config['parentProxyType'] == ProxyType.CONNECT:
                clientData, parentData = connectHandshake()
            else:
                assert(False)
            if clientData == b'':
                clientData = self.client.recv(65536)
            if clientData[0] != '\x16':    #   Not SSL Handshake
                startPipe()
                return
            startClientParentPipe()
            while True:
                packet, parentData = sslGetPacket(self.parent, parentData)
                if parentData[0] == '\x17':    #   Start SSL Application Data
                    self.client.send(packet)
                    break
                parentData = sslCheckCertification(packet)
                self.client.send(packet)
            startParentClientPipe()
        except Exception:
            logging.exception('Exception in Tunnel.run:')
        finally:
            self.client.close()
            self.parent.close()

class Main(object):

    def __init__(self, config):
        logging.basicConfig(filename = config['logFilename'], level = getattr(logging, config['logLevel']))
        config['parentProxyType'] = config['parentProxyType'].lower()
        if config['parentProxyType'] == 'none':
            config['parentProxyType'] = ProxyType.NONE
        elif config['parentProxyType'] == 'socks':
            config['parentProxyType'] = ProxyType.SOCKS
        elif config['parentProxyType'] == 'socks':
            config['parentProxyType'] = ProxyType.CONNECT
        else:
            logging.error('Unknown parentProxyType %s', config['parentProxyType'])
            raise Exception('Unknown parentProxyType %s' % (config['parentProxyType'], ))
        self.config = config

    def start(self):
        server = socket.socket()
        server.bind((config['proxyHost'], config['proxyPort']))
        server.listen(50)
        while True:
            try:
                client, addr = server.accept()
                Tunnel(self.config, client).start()
            except Exception:
                logging.exception('Exception in Main.start:')

if __name__ == '__main__':
    Main(json.loads(open(sys.argv[1], 'r').read())).start()
<|endoftext|>