<commit_before>import slimplectic_GGL as ggl, numpy as np

###########################################
# Interface class for the nonconservative #
#         variational integrator          #
###########################################

class GalerkinGaussLobatto(object):

  def __init__(self, t, q_list, v_list, mod_list = False):
    # Validate inputs
    assert type(t) is str, "String input required."
    assert len(q_list) == len(v_list), "Unequal number of coordinates and velocities."
    self._num_dof = len(q_list)
    assert type(q_list) is list, "List input required."
    assert type(v_list) is list, "List input required."

    for ii in range(len(q_list)):
      assert type(q_list[ii]) is str, "String input required."
      assert type(v_list[ii]) is str, "String input required."

    # Make sympy variables
    self.t = ggl.Symbol(t)
    self.q = [ggl.Symbol(qq) for qq in q_list]
    self.v = [ggl.Symbol(vv) for vv in v_list]

    # Double the sympy variables
    self.qp, self.qm = ggl.q_Generate_pm(self.q)
    self.vp, self.vm = ggl.q_Generate_pm(self.v)

    #keep track of which variables are periodic and need to be modded
    if mod_list:
        self.modlist = mod_list
    if not mod_list:
        self.modlist = []
        for q in q_list:
            self.modlist.append(False)

  def keys(self):
    return self.__dict__.keys()

  def discretize(self, L, K, order, method='explicit', verbose=False):
    """Generate the nonconservative variational integrator maps"""
    self._qi_soln_map, self._q_np1_map, self._pi_np1_map, self._qdot_n_map = ggl.Gen_GGL_NC_VI_Map(self.t, \
                              self.q, self.qp, self.qm, \
                              self.v, self.vp, self.vm, \
                              L, K, order, method=method, verbose=verbose)

  def integrate(self, q0_list, pi0_list, t):
    """Numerical integration from given initial data"""

    # Check if total Lagrangian is discretized already
    if not hasattr(self, '_qi_soln_map'):
        raise AttributeError("Run `discretize` to discretize the total Lagrangian.")

    # Validate input
    assert type(q0_list) in [list, np.ndarray], "List or numpy array input required."
    assert type(pi0_list) in [list, np.ndarray], "List or numpy array input required."

    # Allocate memory for solutions
    t_len = t.size
    q_list_soln = np.zeros((t_len, self._num_dof))
    pi_list_soln = np.zeros((t_len, self._num_dof))

    # Set initial data
    q_list_soln[0,:] = q0_list
    # mod the value of any periodic variables that have mod value specified
    # This prevents NC evolution error due to increasing roundoff error as
    # any cyclic periodic variables become large.
    for jj,mod in enumerate(self.modlist):
        if mod:
            q_list_soln[0,jj] = q_list_soln[0,jj]%mod
    pi_list_soln[0,:] = pi0_list

    # Perform the integration at fixed time steps
    dt = t[1]-t[0]
    for ii in range(1, t_len):
        args = [q_list_soln[ii-1], pi_list_soln[ii-1], t[ii-1], dt]
        qi_sol = self._qi_soln_map(*args)
        q_list_soln[ii] = self._q_np1_map(qi_sol, *args)
        # mod the value of any periodic variables that have mod value specified
        # This prevents NC evolution error due to increasing roundoff error as
        # any cyclic periodic variables become large.
        for jj,mod in enumerate(self.modlist):
            if mod:
                q_list_soln[ii,jj] = q_list_soln[ii,jj]%mod
        pi_list_soln[ii] = self._pi_np1_map(qi_sol, *args)

    # Return the numerical solutions
    return q_list_soln.T, pi_list_soln.T

  def __call__(self, qi_list, vi_list, t):
    return self.integrate(qi_list, vi_list, t)

##########################################
# Class for 4th order Runge-Kutta method #
##########################################

class RungeKutta4(object):

  def __init__(self):
    self._b = np.array([1./6., 1./3., 1./3., 1./6.])
    self._k = [[],[],[],[]]

  def _iter(self, tn, yn_list, f, h):
    self._k[0] = f(tn, yn_list)
    self._k[1] = f(tn+h/2., yn_list+h/2.*self._k[0])
    self._k[2] = f(tn+h/2., yn_list+h/2.*self._k[1])
    self._k[3] = f(tn+h, yn_list+h*self._k[2])
    return yn_list + h*np.sum(self._b[ii]*self._k[ii] for ii in range(4))

  def integrate(self, q0_list, v0_list, t, f):
    y0_list = np.hstack([q0_list, v0_list])
    h = t[1]-t[0]
    ans = np.zeros((t.size, len(y0_list)))
    ans[0,:] = y0_list
    for ii, tt in enumerate(t[:-1]):
      ans[ii+1,:] = self._iter(tt, ans[ii], f, h)
    out = ans.T
    return out[:len(q0_list)], out[len(q0_list):]

  def __call__(self, q0, v0, t, f):
    return self.integrate(q0, v0, t, f)

##########################################
# Class for 2nd order Runge-Kutta method #
##########################################

class RungeKutta2(object):

  def __init__(self):
    self._b = np.array([0., 1.])
    self._k = [[],[]]

  def _iter(self, tn, yn_list, f, h):
    self._k[0] = f(tn, yn_list)
    self._k[1] = f(tn+h/2., yn_list+h/2.*self._k[0])
    return yn_list + h*np.sum(self._b[ii]*self._k[ii] for ii in range(2))

  def integrate(self, q0_list, v0_list, t, f):
    y0_list = np.hstack([q0_list, v0_list])
    h = t[1]-t[0]
    ans = np.zeros((t.size, len(y0_list)))
    ans[0,:] = y0_list
    for ii, tt in enumerate(t[:-1]):
      ans[ii+1,:] = self._iter(tt, ans[ii], f, h)
    #return ans.T
    out = ans.T
    return out[:len(q0_list)], out[len(q0_list):]

  def __call__(self, q0, v0, t, f):
    return self.integrate(q0, v0, t, f)
<commit_msg>added docstrings for slimplectic.py<commit_after>import slimplectic_GGL as ggl, numpy as np

###########################################
# Interface class for the nonconservative #
#         variational integrator          #
###########################################

class GalerkinGaussLobatto(object):

  def __init__(self, t, q_list, v_list, mod_list = False):
    """GalerkinGaussLobatto class:
    args: t : string for generating the sympy symbol the independent "time" variable
          q_list : list of strings for generating symbols for dof q
          v_list : list of strings for generating symbols for \dot{q}
          mod_list : list of values to mod the q values by for periodic variables
                     if not periodic, the value is set to False (default).

    methods: discretize : creates the variational integrator maps
             integrate: applies the maps created by discretize
    """
    # Validate inputs
    assert type(t) is str, "String input required."
    assert len(q_list) == len(v_list), "Unequal number of coordinates and velocities."
    self._num_dof = len(q_list)
    assert type(q_list) is list, "List input required."
    assert type(v_list) is list, "List input required."

    for ii in range(len(q_list)):
      assert type(q_list[ii]) is str, "String input required."
      assert type(v_list[ii]) is str, "String input required."

    # Make sympy variables
    self.t = ggl.Symbol(t)
    self.q = [ggl.Symbol(qq) for qq in q_list]
    self.v = [ggl.Symbol(vv) for vv in v_list]

    # Double the sympy variables
    self.qp, self.qm = ggl.q_Generate_pm(self.q)
    self.vp, self.vm = ggl.q_Generate_pm(self.v)

    #keep track of which variables are periodic and need to be modded
    if mod_list:
        self.modlist = mod_list
    if not mod_list:
        self.modlist = []
        for q in q_list:
            self.modlist.append(False)

  def keys(self):
    return self.__dict__.keys()

  def discretize(self, L, K, order, method='explicit', verbose=False):
    """Generate the nonconservative variational integrator maps
    by setting the methods _qi_soln_map, _q_np1_map, _pi_np1_map, _qdot_n_map
    args: L : sympy expression for the conservative Lagrangian. Should be in
              terms of the self.q and self.v variables
          K : sympy expression for the nonconservative potential. Should be in
              terms of the self.qp/m and self.vp/m variables
          order: integer order (r) of the GGL method (r+2 is the total order of the method)
          method: string 'implicit' or 'explicit' evaluation (default explicit)
          verbose: Boolean. True to output mapping expressions (default False)
    output: none
    """
    self._qi_soln_map, self._q_np1_map, self._pi_np1_map, self._qdot_n_map = ggl.Gen_GGL_NC_VI_Map(self.t, \
                              self.q, self.qp, self.qm, \
                              self.v, self.vp, self.vm, \
                              L, K, order, method=method, verbose=verbose)

  def integrate(self, q0_list, pi0_list, t):
    """Numerical integration from given initial data
    args: q0_list: list of initial q values (floats)
          pi0_list: list of initial pi (nonconservative momentum) values (floats)
          t: list of t values (floats) over which to integrate the system.
    output: q_list_soln.T, pi_list_soln.T the integrated q and pi arrays at each time.
    """

    # Check if total Lagrangian is discretized already
    if not hasattr(self, '_qi_soln_map'):
        raise AttributeError("Run `discretize` to discretize the total Lagrangian.")

    # Validate input
    assert type(q0_list) in [list, np.ndarray], "List or numpy array input required."
    assert type(pi0_list) in [list, np.ndarray], "List or numpy array input required."

    # Allocate memory for solutions
    t_len = t.size
    q_list_soln = np.zeros((t_len, self._num_dof))
    pi_list_soln = np.zeros((t_len, self._num_dof))

    # Set initial data
    q_list_soln[0,:] = q0_list
    # mod the value of any periodic variables that have mod value specified
    # This prevents NC evolution error due to increasing roundoff error as
    # any cyclic periodic variables become large.
    for jj,mod in enumerate(self.modlist):
        if mod:
            q_list_soln[0,jj] = q_list_soln[0,jj]%mod
    pi_list_soln[0,:] = pi0_list

    # Perform the integration at fixed time steps
    dt = t[1]-t[0]
    for ii in range(1, t_len):
        args = [q_list_soln[ii-1], pi_list_soln[ii-1], t[ii-1], dt]
        qi_sol = self._qi_soln_map(*args)
        q_list_soln[ii] = self._q_np1_map(qi_sol, *args)
        # mod the value of any periodic variables that have mod value specified
        # This prevents NC evolution error due to increasing roundoff error as
        # any cyclic periodic variables become large.
        for jj,mod in enumerate(self.modlist):
            if mod:
                q_list_soln[ii,jj] = q_list_soln[ii,jj]%mod
        pi_list_soln[ii] = self._pi_np1_map(qi_sol, *args)

    # Return the numerical solutions
    return q_list_soln.T, pi_list_soln.T

  def __call__(self, qi_list, vi_list, t):
    return self.integrate(qi_list, vi_list, t)

##########################################
# Class for 4th order Runge-Kutta method #
##########################################

class RungeKutta4(object):

  def __init__(self):
    self._b = np.array([1./6., 1./3., 1./3., 1./6.])
    self._k = [[],[],[],[]]

  def _iter(self, tn, yn_list, f, h):
    self._k[0] = f(tn, yn_list)
    self._k[1] = f(tn+h/2., yn_list+h/2.*self._k[0])
    self._k[2] = f(tn+h/2., yn_list+h/2.*self._k[1])
    self._k[3] = f(tn+h, yn_list+h*self._k[2])
    return yn_list + h*np.sum(self._b[ii]*self._k[ii] for ii in range(4))

  def integrate(self, q0_list, v0_list, t, f):
    y0_list = np.hstack([q0_list, v0_list])
    h = t[1]-t[0]
    ans = np.zeros((t.size, len(y0_list)))
    ans[0,:] = y0_list
    for ii, tt in enumerate(t[:-1]):
      ans[ii+1,:] = self._iter(tt, ans[ii], f, h)
    out = ans.T
    return out[:len(q0_list)], out[len(q0_list):]

  def __call__(self, q0, v0, t, f):
    return self.integrate(q0, v0, t, f)

##########################################
# Class for 2nd order Runge-Kutta method #
##########################################

class RungeKutta2(object):

  def __init__(self):
    self._b = np.array([0., 1.])
    self._k = [[],[]]

  def _iter(self, tn, yn_list, f, h):
    self._k[0] = f(tn, yn_list)
    self._k[1] = f(tn+h/2., yn_list+h/2.*self._k[0])
    return yn_list + h*np.sum(self._b[ii]*self._k[ii] for ii in range(2))

  def integrate(self, q0_list, v0_list, t, f):
    y0_list = np.hstack([q0_list, v0_list])
    h = t[1]-t[0]
    ans = np.zeros((t.size, len(y0_list)))
    ans[0,:] = y0_list
    for ii, tt in enumerate(t[:-1]):
      ans[ii+1,:] = self._iter(tt, ans[ii], f, h)
    #return ans.T
    out = ans.T
    return out[:len(q0_list)], out[len(q0_list):]

  def __call__(self, q0, v0, t, f):
    return self.integrate(q0, v0, t, f)
<|endoftext|>


<commit_before>from .base import Module, extract_value

class Repository(Module):
    _name = 'repository'

    @extract_value('repositories')
    def init(self):
        return self._request('init')

    def add(self):
        #TODO repository::add
        raise NotImplementedError

    def edit(self):
        #TODO repository::edit
        raise NotImplementedError

    def delete_simulate(self, *ids):
        return self._request('deleteSimulate', {
            'repositories': [{'id': id} for id in ids]
        })['effects']

    def delete(self, *ids):
        """Deletes a specific repository or group of repositories.

        :param *ids: the id of specified repository

        :return: dict containing the id of the deleted repository
        """

        return self._request('delete', {
            'repositories': [{'id': id} for id in ids]
        })['repositories']

    def import_(self, id, file):
        """Import a repository file.

        :param id: the id of the uploaded repository
        :param file: the file is expected to be an absolute path to a tar.gz

        :return: returns the remote repository details just imported
        """

        return self._request('import', {
            'id': id,
            'file': file
        })

    def export(self, id):
        return self._request('export', {
            'id': id
        }, parse=False).content

    def download(self):
        #TODO repository::download
        raise NotImplementedError

    def get_remote_repositories(self, host):
        """Retrieve repositories available within the database.

        :param host: attribute host is expected to be a valid host IP address

        :return: return repositories available in database
        """

        # return self._request('getRemoteRepositories', {
        #     'repositories': [{'host': host} for host in hosts]
        # })

        # decide which one

        return self._request('getRemoteRepositories', {
            'host': host
        })

    def download_remote_repositories(self):
        #TODO repository::downloadRemoteRepositories
        raise NotImplementedError

    def validate_add(self, type, data_format, correlation, remote_type, nessus_type,
                     download_format='v2', org=None, description=None, enable_trend=True,
                     remote_def=None, nessus_def=None):
        """Validate adding a repository.

        :param type: the repository type
        :param data_format: repository data format should be IPv4
        :param correlation: need to be a valid repository correlation
        :param remote_type: the schedule type for the repository
        :param nessus_type: the Nessus schedule type for the repository
        :param download_format: need to be a valid download format for Nessus v2
        :param org: dict of organization id
        :param description: repository description
        :param enable_trend: default is set to true
        :param remote_def: the schedule definition for the repository
        :param nessus_def: the Nessus schedule definition for the repository

        :return: return valid added repository information
        """

        #include correlation LCELib section?

        if org is not None:
            org = [{'id': id} for id in org]

        return self._request('validateAdd', {
            'type': type,
            'organizations': org,
            'dataFormat': data_format,
            'correlation': correlation,
            'enableTrending': enable_trend,
            'remoteScheduleType': remote_type,
            'remoteScheduleDefinition': remote_def,
            'description': description,
            'downloadFormat': download_format,
            'nessusScheduleType': nessus_type,
            'nessusScheduleDefinition': nessus_def
        })

    def validate_edit(self, repos_id, type, name, description, ip, correlation,
                      remote_id, remote_ip, remote_type, nessus_type, download_format='v2',
                      org=None, remote_def=None, nessus_def=None):
        """Validate editing a repository.

        :param id: the repository id
        :param type: the repository type
        :param name: the repository name
        :param description: the repository description
        :param ip: valid IP address in IPv4 or IPv6 depending on data_format
        :param correlation: need to be a valid repository correlation
        :param remote_id: the repository remote id
        :param remote_ip: the repository remote ip
        :param remote_type: the schedule type for the repository (default set to never)
        :param nessus_type: Nessus schedule type (default set to never)
        :param download_format: need to be a valid download format for Nessus v2
        :param org: dict of organization id
        :param remote_def: schedule definition for the repository
        :param nessus_def: Nessus schedule definition for the repository

        :return: return valid edited repository information
        """

        #include correlation LCELib section?

        if org is not None:
            org = [{'id': id} for id in org]

        return self._request('validateEdit', {
            'repID': repos_id,
            'type': type,
            'name': name,
            'organizations': org,
            'description': description,
            'ipRange': ip,
            'correlation': correlation,
            'remoteID': remote_id,
            'remoteIP': remote_ip,
            'downloadFormat': download_format,
            'remoteScheduleType': remote_type,
            'remoteScheduleDefinition': remote_def,
            'nessusScheduleType': nessus_type,
            'nessusScheduleDefinition': nessus_def
        })

    def generate_nessus_file(self):
        #TODO repository::generateNessusFile
        raise NotImplementedError
<commit_msg>comment changes<commit_after>from .base import Module, extract_value

class Repository(Module):
    _name = 'repository'

    @extract_value('repositories')
    def init(self):
        return self._request('init')

    def add(self):
        #TODO repository::add
        raise NotImplementedError

    def edit(self):
        #TODO repository::edit
        raise NotImplementedError

    def delete_simulate(self, *ids):
        return self._request('deleteSimulate', {
            'repositories': [{'id': id} for id in ids]
        })['effects']

    def delete(self, *ids):
        """Deletes a specific repository or group of repositories.

        :param *ids: the id of specified repository

        :return: dict containing the id of the deleted repository
        """

        return self._request('delete', {
            'repositories': [{'id': id} for id in ids]
        })['repositories']

    def import_(self, id, file):
        """Import a repository file.

        :param id: the id of the uploaded repository
        :param file: the file is expected to be an absolute path to a tar.gz

        :return: returns the remote repository details just imported
        """

        return self._request('import', {
            'id': id,
            'file': file
        })

    def export(self, id):
        return self._request('export', {
            'id': id
        }, parse=False).content

    def download(self):
        #TODO repository::download
        raise NotImplementedError

    def get_remote_repositories(self, host):
        """Retrieve repositories available within the database.

        :param host: attribute host is expected to be a valid host IP address

        :return: return repositories available in database
        """

        # return self._request('getRemoteRepositories', {
        #     'repositories': [{'host': host} for host in hosts]
        # })

        # note: decide which one

        return self._request('getRemoteRepositories', {
            'host': host
        })

    def download_remote_repositories(self):
        #TODO repository::downloadRemoteRepositories
        raise NotImplementedError

    def validate_add(self, type, data_format, correlation, remote_type, nessus_type,
                     download_format='v2', org=None, description=None, enable_trend=True,
                     remote_def=None, nessus_def=None):
        """Validate adding a repository.

        :param type: the repository type
        :param data_format: repository data format should be IPv4
        :param correlation: need to be a valid repository correlation
        :param remote_type: the schedule type for the repository
        :param nessus_type: the Nessus schedule type for the repository
        :param download_format: need to be a valid download format for Nessus v2
        :param org: dict of organization id
        :param description: repository description
        :param enable_trend: default is set to true
        :param remote_def: the schedule definition for the repository
        :param nessus_def: the Nessus schedule definition for the repository

        :return: return valid added repository information
        """

        #include correlation LCELib section?

        if org is not None:
            org = [{'id': id} for id in org]

        return self._request('validateAdd', {
            'type': type,
            'organizations': org,
            'dataFormat': data_format,
            'correlation': correlation,
            'enableTrending': enable_trend,
            'remoteScheduleType': remote_type,
            'remoteScheduleDefinition': remote_def,
            'description': description,
            'downloadFormat': download_format,
            'nessusScheduleType': nessus_type,
            'nessusScheduleDefinition': nessus_def
        })

    def validate_edit(self, repos_id, type, name, description, ip, correlation,
                      remote_id, remote_ip, remote_type, nessus_type, download_format='v2',
                      org=None, remote_def=None, nessus_def=None):
        """Validate editing a repository.

        :param repos_id: the repository id
        :param type: the repository type
        :param name: the repository name
        :param description: the repository description
        :param ip: valid IP address in IPv4 or IPv6 depending on data_format
        :param correlation: need to be a valid repository correlation
        :param remote_id: the repository remote id
        :param remote_ip: the repository remote ip
        :param remote_type: the schedule type for the repository (default set to never)
        :param nessus_type: Nessus schedule type (default set to never)
        :param download_format: need to be a valid download format for Nessus v2
        :param org: dict of organization id
        :param remote_def: schedule definition for the repository
        :param nessus_def: Nessus schedule definition for the repository

        :return: return valid edited repository information
        """

        #include correlation LCELib section?

        if org is not None:
            org = [{'id': id} for id in org]

        return self._request('validateEdit', {
            'repID': repos_id,
            'type': type,
            'name': name,
            'organizations': org,
            'description': description,
            'ipRange': ip,
            'correlation': correlation,
            'remoteID': remote_id,
            'remoteIP': remote_ip,
            'downloadFormat': download_format,
            'remoteScheduleType': remote_type,
            'remoteScheduleDefinition': remote_def,
            'nessusScheduleType': nessus_type,
            'nessusScheduleDefinition': nessus_def
        })

    def generate_nessus_file(self):
        #TODO repository::generateNessusFile
        raise NotImplementedError
<|endoftext|>


<commit_before>'''
slots

A Python library to perform simple multi-armed bandit analyses.

Scenarios:
    - Run MAB test on simulated data (N bandits), default epsilon-greedy test.
        mab = slots.MAB(probs = [0.1,0.15,0.05])
        mab.run(trials = 10000)
        mab.best  # Bandit with highest probability after T trials

    - Run MAB test on "real" payout data (probabilites unknown).
        mab = slots.MAB(payouts = [0,0,0,1,0,0,0,0,0,....])
        mab.run(trials = 10000) # Max is length of payouts
'''

import numpy as np

class MAB(object):
    '''
    Multi-armed bandit test class.
    '''

    def __init__(self, num_bandits=3, probs=None, payouts=None, live=False,
                 stop_criterion={'criterion': 'regret', 'value': 1.0}):
        '''
        Parameters
        ----------
        num_bandits : int
            default is 3
        probs : np.array of floats
            payout probabilities
        payouts : np.array of floats
            If `live` is True, `payouts` should be an N*T array of payout
            amount per pull (floats) for N bandits and T trials
        live : bool
        stop_criterion : dict
            Stopping criterion (str) and threshold value (float).
        '''

        self.choices = []

        if not probs:
            if payouts is None:
                self.bandits = Bandits(probs=[np.random.rand() for x in
                                       range(num_bandits)],
                                       payouts=np.ones(num_bandits))
            else:
                if live:
                    self.bandits = Bandits(live=True, payouts=payouts,
                                           probs=None)
                else:
                    # Not sure why anyone would do this
                    self.bandits = Bandits(probs=[np.random.rand() for x in
                                           range(len(payouts))],
                                           payouts=payouts)
                num_bandits = len(payouts)
        else:
            if payouts:
                self.bandits = Bandits(probs=probs, payouts=payouts)
                num_bandits = len(payouts)
            else:
                self.bandits = Bandits(probs=probs,
                                       payouts=np.ones(len(probs)))
                num_bandits = len(probs)

        self.wins = np.zeros(num_bandits)
        self.pulls = np.zeros(num_bandits)

        # Set the stopping criteria
        self.criteria = {'regret': self.regret_met}
        self.criterion = stop_criterion.get('criterion', 'regret')
        self.stop_value = stop_criterion.get('value', 0.1)

        # Bandit selection strategies
        self.strategies = ['eps_greedy', 'softmax', 'ucb', 'bayesian']

    def run(self, trials=100, strategy=None, parameters=None):
        '''
        Run MAB test with T trials.

        Parameters
        ----------
        trials : int
            Number of trials to run.
        strategy : str
            Name of selected strategy.
        parameters : dict
            Parameters for selected strategy.

        Available strategies:
            - Epsilon-greedy ("eps_greedy")
            - Softmax ("softmax")
            - Upper confidence bound ("ucb")

        Returns
        -------
        None
        '''

        if trials < 1:
            raise Exception('MAB.run: Number of trials cannot be less than 1!')
        if not strategy:
            strategy = 'eps_greedy'
        else:
            if strategy not in self.strategies:
                raise Exception('MAB,run: Strategy name invalid. Choose from:'
                                ' {}'.format(', '.join(self.strategies)))

        # Run strategy
        for n in range(trials):
            self._run(strategy, parameters)

    def _run(self, strategy, parameters=None):
        '''
        Run single trial of MAB strategy.

        Parameters
        ----------
        strategy : function
        parameters : dict

        Returns
        -------
        None
        '''

        choice = self.run_strategy(strategy, parameters)
        self.choices.append(choice)
        payout = self.bandits.pull(choice)
        if payout is None:
            print('Trials exhausted. No more values for bandit', choice)
            return None
        else:
            self.wins[choice] += payout
        self.pulls[choice] += 1

    def run_strategy(self, strategy, parameters):
        '''
        Run the selected strategy and retrun bandit choice.

        Parameters
        ----------
        strategy : str
            Name of MAB strategy.
        parameters : dict
            Strategy function parameters

        Returns
        -------
        int
            Bandit arm choice index
        '''

        return self.__getattribute__(strategy)(params=parameters)

# ###### ----------- MAB strategies ---------------------------------------####
    def max_mean(self):
        """
        Pick the bandit with the current best observed proportion of winning.

        Returns
        -------
        int
            Index of chosen bandit
        """

        return np.argmax(self.wins / (self.pulls + 0.1))

    def bayesian(self, params=None):
        '''
        Run the Bayesian Bandit algorithm which utilizes a beta distribution
        for exploration and exploitation.

        Parameters
        ----------
        params : None
            For API consistency, this function can take a parameters argument,
            but it is ignored.

        Returns
        -------
        int
            Index of chosen bandit
        '''
        p_success_arms = [
            np.random.beta(self.wins[i] + 1, self.pulls[i] - self.wins[i] + 1)
            for i in range(len(self.wins))
            ]

        return np.array(p_success_arms).argmax()

    def eps_greedy(self, params):
        '''
        Run the epsilon-greedy strategy and update self.max_mean()

        Parameters
        ----------
        Params : dict
            Epsilon

        Returns
        -------
        int
            Index of chosen bandit
        '''

        if params and type(params) == dict:
            eps = params.get('epsilon')
        else:
            eps = 0.1

        r = np.random.rand()

        if r < eps:
            return np.random.choice(list(set(range(len(self.wins))) -
                                    {self.max_mean()}))
        else:
            return self.max_mean()

    def softmax(self, params):
        '''
        Run the softmax selection strategy.

        Parameters
        ----------
        Params : dict
            Tau

        Returns
        -------
        int
            Index of chosen bandit
        '''

        default_tau = 0.1

        if params and type(params) == dict:
            tau = params.get('tau')
            try:
                float(tau)
            except ValueError:
                'slots: softmax: Setting tau to default'
                tau = default_tau
        else:
            tau = default_tau

        # Handle cold start. Not all bandits tested yet.
        if True in (self.pulls < 3):
            return np.random.choice(range(len(self.pulls)))
        else:
            payouts = self.wins / (self.pulls + 0.1)
            norm = sum(np.exp(payouts/tau))

        ps = np.exp(payouts/tau)/norm

        # Randomly choose index based on CMF
        cmf = [sum(ps[:i+1]) for i in range(len(ps))]

        rand = np.random.rand()

        found = False
        found_i = None
        i = 0
        while not found:
            if rand < cmf[i]:
                found_i = i
                found = True
            else:
                i += 1

        return found_i

    def ucb(self, params=None):
        '''
        Run the upper confidence bound MAB selection strategy.

        This is the UCB1 algorithm described in
        https://homes.di.unimi.it/~cesabian/Pubblicazioni/ml-02.pdf

        Parameters
        ----------
        params : None
            For API consistency, this function can take a parameters argument,
            but it is ignored.

        Returns
        -------
        int
            Index of chosen bandit
        '''

        # UCB = j_max(payout_j + sqrt(2ln(n_tot)/n_j))

        # Handle cold start. Not all bandits tested yet.
        if True in (self.pulls < 3):
            return np.random.choice(range(len(self.pulls)))
        else:
            n_tot = sum(self.pulls)
            payouts = self.wins / (self.pulls + 0.1)
            ubcs = payouts + np.sqrt(2*np.log(n_tot)/self.pulls)

            return np.argmax(ubcs)

    # ###------------------------------------------------------------------####

    def best(self):
        '''
        Return current 'best' choice of bandit.

        Returns
        -------
        int
            Index of bandit
        '''

        if len(self.choices) < 1:
            print('slots: No trials run so far.')
            return None
        else:
            return np.argmax(self.wins/(self.pulls+0.1))

    def est_payouts(self):
        '''
        Calculate current estimate of average payout for each bandit.

        Returns
        -------
        array of floats or None
        '''

        if len(self.choices) < 1:
            print('slots: No trials run so far.')
            return None
        else:
            return self.wins/(self.pulls+0.1)

    def regret(self):
        '''
        Calculate expected regret, where expected regret is
        maximum optimal reward - sum of collected rewards, i.e.

        expected regret = T*max_k(mean_k) - sum_(t=1-->T) (reward_t)

        Returns
        -------
        float
        '''

        return (sum(self.pulls)*np.max(np.nan_to_num(self.wins/self.pulls)) -
                sum(self.wins)) / sum(self.pulls)

    def crit_met(self):
        '''
        Determine if stopping criterion has been met.

        Returns
        -------
        bool
        '''

        return self.criteria[self.criterion](self.stop_value)

    def regret_met(self, threshold=None):
        '''
        Determine if regret criterion has been met.

        Parameters
        ----------
        threshold : float

        Returns
        -------
        bool
        '''

        if not threshold:
            return self.regret() <= self.stop_value
        elif self.regret() <= threshold:
            return True
        else:
            return False

    # ## ------------ Online bandit testing ------------------------------ ####
    def online_trial(self, bandit=None, payout=None, strategy='eps_greedy',
                     parameters=None):
        '''
        Update the bandits with the results of the previous live, online trial.
            Next run a the selection algorithm. If the stopping criteria is
            met, return the best arm estimate. Otherwise return the next arm to
            try.

        Parameters
        ----------
        bandit : int
            Bandit index
        payout : float
            Payout value
        strategy : string
            Name of update strategy
        parameters : dict
            Parameters for update strategy function

        Returns
        -------
        dict
            Format: {'new_trial': boolean, 'choice': int, 'best': int}
        '''

        if bandit and payout:
            self.update(bandit=bandit, payout=payout)
        else:
            raise Exception('slots.online_trial: bandit and/or payout value'
                            ' missing.')

        if self.crit_met():
            return {'new_trial': False, 'choice': self.best(),
                    'best': self.best()}
        else:
            return {'new_trial': True,
                    'choice': self.run_strategy(strategy, parameters),
                    'best': self.best()}

    def update(self, bandit=None, payout=None):
        '''
        Update bandit trials and payouts for given bandit.

        Parameters
        ----------
        bandit : int
            Bandit index
        payout : float

        Returns
        -------
        None
        '''

        self.pulls[bandit] += 1
        self.wins[bandit] += payout
        self.bandits.payouts[bandit] += payout

class Bandits():
    '''
    Bandit class.
    '''

    def __init__(self, probs, payouts, live=False):
        '''
        Instantiate Bandit class, determining
            - Probabilities of bandit payouts
            - Bandit payouts

        Parameters
        ----------
        probs: array of floats
            Probabilities of bandit payouts
        payouts : array of floats
            Amount of bandit payouts. If `live` is True, `payouts` should be an
            N*T array of floats giving the payout amount per pull for N bandits
            and T trials.
        live : bool
        '''

        if not live:
            # Only use arrays of equal length
            if len(probs) != len(payouts):
                raise Exception('Bandits.__init__: Probability and payouts '
                                'arrays of different lengths!')
            self.probs = probs
            self.payouts = payouts
            self.live = False
        else:
            self.live = True
            self.probs = None
            self.payouts = payouts

    def pull(self, i):
        '''
        Return the payout from a single pull of the bandit i's arm.

        Parameters
        ----------
        i : int
            Index of bandit.

        Returns
        -------
        float or None
        '''

        if self.live:
            if len(self.payouts[i]) > 0:
                return self.payouts[i].pop()
            else:
                return None
        else:
            if np.random.rand() < self.probs[i]:
                return self.payouts[i]
            else:
                return 0.0

    def info(self):
        pass
<commit_msg>Changed default settings to "live first" (online is default use case) Fixed instatiation of n online bandits<commit_after>'''
slots

A Python library to perform simple multi-armed bandit analyses.

Scenarios:
    - Run MAB test on simulated data (N bandits), default epsilon-greedy test.
        mab = slots.MAB(probs = [0.1,0.15,0.05])
        mab.run(trials = 10000)
        mab.best  # Bandit with highest probability after T trials

    - Run MAB test on "real" payout data (probabilites unknown).
        mab = slots.MAB(payouts = [0,0,0,1,0,0,0,0,0,....])
        mab.run(trials = 10000) # Max is length of payouts
'''

import numpy as np

class MAB(object):
    '''
    Multi-armed bandit test class.
    '''

    def __init__(self, num_bandits=3, probs=None, payouts=None, live=True,
                 stop_criterion={'criterion': 'regret', 'value': 0.1}):
        '''
        Parameters
        ----------
        num_bandits : int
            default is 3
        probs : np.array of floats
            payout probabilities
        payouts : np.array of floats
            If `live` is True, `payouts` should be None.
        live : bool
            Whether the use is for a live, online trial.
        stop_criterion : dict
            Stopping criterion (str) and threshold value (float).
        '''

        self.choices = []

        if not probs:
            if not payouts:
                if live:
                    self.bandits = Bandits(live=True,
                                           payouts=np.zeros(num_bandits),
                                           probs=None)
                else:
                    self.bandits = Bandits(probs=[np.random.rand() for x in
                                           range(num_bandits)],
                                           payouts=np.ones(num_bandits),
                                           live=False)
            else:

                self.bandits = Bandits(probs=[np.random.rand() for x in
                                       range(len(payouts))],
                                       payouts=payouts,
                                       live=False)
                num_bandits = len(payouts)
        else:
            if payouts:
                self.bandits = Bandits(probs=probs, payouts=payouts,
                                       live=False)
                num_bandits = len(payouts)
            else:
                self.bandits = Bandits(probs=probs,
                                       payouts=np.ones(len(probs)),
                                       live=False)
                num_bandits = len(probs)

        self.wins = np.zeros(num_bandits)
        self.pulls = np.zeros(num_bandits)

        # Set the stopping criteria
        self.criteria = {'regret': self.regret_met}
        self.criterion = stop_criterion.get('criterion', 'regret')
        self.stop_value = stop_criterion.get('value', 0.1)

        # Bandit selection strategies
        self.strategies = ['eps_greedy', 'softmax', 'ucb', 'bayesian']

    def run(self, trials=100, strategy=None, parameters=None):
        '''
        Run MAB test with T trials.

        Parameters
        ----------
        trials : int
            Number of trials to run.
        strategy : str
            Name of selected strategy.
        parameters : dict
            Parameters for selected strategy.

        Available strategies:
            - Epsilon-greedy ("eps_greedy")
            - Softmax ("softmax")
            - Upper confidence bound ("ucb")

        Returns
        -------
        None
        '''

        if trials < 1:
            raise Exception('MAB.run: Number of trials cannot be less than 1!')
        if not strategy:
            strategy = 'eps_greedy'
        else:
            if strategy not in self.strategies:
                raise Exception('MAB,run: Strategy name invalid. Choose from:'
                                ' {}'.format(', '.join(self.strategies)))

        # Run strategy
        for n in range(trials):
            self._run(strategy, parameters)

    def _run(self, strategy, parameters=None):
        '''
        Run single trial of MAB strategy.

        Parameters
        ----------
        strategy : function
        parameters : dict

        Returns
        -------
        None
        '''

        choice = self.run_strategy(strategy, parameters)
        self.choices.append(choice)
        payout = self.bandits.pull(choice)
        if payout is None:
            print('Trials exhausted. No more values for bandit', choice)
            return None
        else:
            self.wins[choice] += payout
        self.pulls[choice] += 1

    def run_strategy(self, strategy, parameters):
        '''
        Run the selected strategy and retrun bandit choice.

        Parameters
        ----------
        strategy : str
            Name of MAB strategy.
        parameters : dict
            Strategy function parameters

        Returns
        -------
        int
            Bandit arm choice index
        '''

        return self.__getattribute__(strategy)(params=parameters)

# ###### ----------- MAB strategies ---------------------------------------####
    def max_mean(self):
        """
        Pick the bandit with the current best observed proportion of winning.

        Returns
        -------
        int
            Index of chosen bandit
        """

        return np.argmax(self.wins / (self.pulls + 0.1))

    def bayesian(self, params=None):
        '''
        Run the Bayesian Bandit algorithm which utilizes a beta distribution
        for exploration and exploitation.

        Parameters
        ----------
        params : None
            For API consistency, this function can take a parameters argument,
            but it is ignored.

        Returns
        -------
        int
            Index of chosen bandit
        '''
        p_success_arms = [
            np.random.beta(self.wins[i] + 1, self.pulls[i] - self.wins[i] + 1)
            for i in range(len(self.wins))
            ]

        return np.array(p_success_arms).argmax()

    def eps_greedy(self, params):
        '''
        Run the epsilon-greedy strategy and update self.max_mean()

        Parameters
        ----------
        Params : dict
            Epsilon

        Returns
        -------
        int
            Index of chosen bandit
        '''

        if params and type(params) == dict:
            eps = params.get('epsilon')
        else:
            eps = 0.1

        r = np.random.rand()

        if r < eps:
            return np.random.choice(list(set(range(len(self.wins))) -
                                    {self.max_mean()}))
        else:
            return self.max_mean()

    def softmax(self, params):
        '''
        Run the softmax selection strategy.

        Parameters
        ----------
        Params : dict
            Tau

        Returns
        -------
        int
            Index of chosen bandit
        '''

        default_tau = 0.1

        if params and type(params) == dict:
            tau = params.get('tau')
            try:
                float(tau)
            except ValueError:
                'slots: softmax: Setting tau to default'
                tau = default_tau
        else:
            tau = default_tau

        # Handle cold start. Not all bandits tested yet.
        if True in (self.pulls < 3):
            return np.random.choice(range(len(self.pulls)))
        else:
            payouts = self.wins / (self.pulls + 0.1)
            norm = sum(np.exp(payouts/tau))

        ps = np.exp(payouts/tau)/norm

        # Randomly choose index based on CMF
        cmf = [sum(ps[:i+1]) for i in range(len(ps))]

        rand = np.random.rand()

        found = False
        found_i = None
        i = 0
        while not found:
            if rand < cmf[i]:
                found_i = i
                found = True
            else:
                i += 1

        return found_i

    def ucb(self, params=None):
        '''
        Run the upper confidence bound MAB selection strategy.

        This is the UCB1 algorithm described in
        https://homes.di.unimi.it/~cesabian/Pubblicazioni/ml-02.pdf

        Parameters
        ----------
        params : None
            For API consistency, this function can take a parameters argument,
            but it is ignored.

        Returns
        -------
        int
            Index of chosen bandit
        '''

        # UCB = j_max(payout_j + sqrt(2ln(n_tot)/n_j))

        # Handle cold start. Not all bandits tested yet.
        if True in (self.pulls < 3):
            return np.random.choice(range(len(self.pulls)))
        else:
            n_tot = sum(self.pulls)
            payouts = self.wins / (self.pulls + 0.1)
            ubcs = payouts + np.sqrt(2*np.log(n_tot)/self.pulls)

            return np.argmax(ubcs)

    # ###------------------------------------------------------------------####

    def best(self):
        '''
        Return current 'best' choice of bandit.

        Returns
        -------
        int
            Index of bandit
        '''

        if len(self.choices) < 1:
            print('slots: No trials run so far.')
            return None
        else:
            return np.argmax(self.wins/(self.pulls+0.1))

    def est_payouts(self):
        '''
        Calculate current estimate of average payout for each bandit.

        Returns
        -------
        array of floats or None
        '''

        if len(self.choices) < 1:
            print('slots: No trials run so far.')
            return None
        else:
            return self.wins/(self.pulls+0.1)

    def regret(self):
        '''
        Calculate expected regret, where expected regret is
        maximum optimal reward - sum of collected rewards, i.e.

        expected regret = T*max_k(mean_k) - sum_(t=1-->T) (reward_t)

        Returns
        -------
        float
        '''

        return (sum(self.pulls)*np.max(np.nan_to_num(self.wins/self.pulls)) -
                sum(self.wins)) / sum(self.pulls)

    def crit_met(self):
        '''
        Determine if stopping criterion has been met.

        Returns
        -------
        bool
        '''

        if True in (self.pulls < 3):
            return False
        else:
            return self.criteria[self.criterion](self.stop_value)

    def regret_met(self, threshold=None):
        '''
        Determine if regret criterion has been met.

        Parameters
        ----------
        threshold : float

        Returns
        -------
        bool
        '''

        if not threshold:
            return self.regret() <= self.stop_value
        elif self.regret() <= threshold:
            return True
        else:
            return False

    # ## ------------ Online bandit testing ------------------------------ ####
    def online_trial(self, bandit=None, payout=None, strategy='eps_greedy',
                     parameters=None):
        '''
        Update the bandits with the results of the previous live, online trial.
            Next run a the selection algorithm. If the stopping criteria is
            met, return the best arm estimate. Otherwise return the next arm to
            try.

        Parameters
        ----------
        bandit : int
            Bandit index
        payout : float
            Payout value
        strategy : string
            Name of update strategy
        parameters : dict
            Parameters for update strategy function

        Returns
        -------
        dict
            Format: {'new_trial': boolean, 'choice': int, 'best': int}
        '''

        if bandit is not None and payout is not None:
            self.update(bandit=bandit, payout=payout)
        else:
            raise Exception('slots.online_trial: bandit and/or payout value'
                            ' missing.')

        if self.crit_met():
            return {'new_trial': False, 'choice': self.best(),
                    'best': self.best()}
        else:
            return {'new_trial': True,
                    'choice': self.run_strategy(strategy, parameters),
                    'best': self.best()}

    def update(self, bandit, payout):
        '''
        Update bandit trials and payouts for given bandit.

        Parameters
        ----------
        bandit : int
            Bandit index
        payout : float

        Returns
        -------
        None
        '''

        self.choices.append(bandit)
        self.pulls[bandit] += 1
        self.wins[bandit] += payout
        self.bandits.payouts[bandit] += payout

class Bandits():
    '''
    Bandit class.
    '''

    def __init__(self, probs, payouts, live=True):
        '''
        Instantiate Bandit class, determining
            - Probabilities of bandit payouts
            - Bandit payouts

        Parameters
        ----------
        probs: array of floats
            Probabilities of bandit payouts
        payouts : array of floats
            Amount of bandit payouts. If `live` is True, `payouts` should be an
            N length array of zeros.
        live : bool
        '''

        if not live:
            # Only use arrays of equal length
            if len(probs) != len(payouts):
                raise Exception('Bandits.__init__: Probability and payouts '
                                'arrays of different lengths!')
            self.probs = probs
            self.payouts = payouts
            self.live = False
        else:
            self.live = True
            self.probs = None
            self.payouts = payouts

    def pull(self, i):
        '''
        Return the payout from a single pull of the bandit i's arm.

        Parameters
        ----------
        i : int
            Index of bandit.

        Returns
        -------
        float or None
        '''

        if self.live:
            if len(self.payouts[i]) > 0:
                return self.payouts[i].pop()
            else:
                return None
        else:
            if np.random.rand() < self.probs[i]:
                return self.payouts[i]
            else:
                return 0.0

    def info(self):
        pass
<|endoftext|>


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

from datetime import datetime, timedelta
from dateutil.parser import parse
from dateutil.tz import tzutc

from c7n.actions import BaseAction
from c7n.filters import ValueFilter
from c7n.manager import resources
from c7n.query import QueryResourceManager
from c7n.utils import local_session, type_schema

@resources.register('iam-group')
class Group(QueryResourceManager):

    resource_type = 'aws.iam.group'

@resources.register('iam-role')
class Role(QueryResourceManager):

    resource_type = 'aws.iam.role'

@resources.register('iam-user')
class User(QueryResourceManager):

    resource_type = 'aws.iam.user'

@User.filter_registry.register('policy')
class UserAttachedPolicy(ValueFilter):

    schema = type_schema('policy', rinherit=ValueFilter.schema)

    def process(self, resources, event=None):

        def _user_policies(resource):
            client = local_session(self.manager.session_factory).client('iam')
            resource['AttachedPolicies'] = client.list_attached_user_policies(
                UserName=resource['UserName'])['AttachedPolicies']

        with self.executor_factory(max_workers=2) as w:
            query_resources = [
                r for r in resources if 'AttachedPolicies' not in r]
            self.log.debug("Querying %d users policies" % len(query_resources))
            list(w.map(_user_policies, query_resources))

        matched = []
        for r in resources:
            for p in r['AttachedPolicies']:
                if self.match(p):
                    matched.append(r)
                    break
        return matched

@User.filter_registry.register('access-key')
class UserAccessKey(ValueFilter):

    schema = type_schema('access-key', rinherit=ValueFilter.schema)

    def process(self, resources, event=None):

        def _user_keys(resource):
            client = local_session(self.manager.session_factory).client('iam')
            resource['AccessKeys'] = client.list_access_keys(
                UserName=resource['UserName'])['AccessKeyMetadata']

        with self.executor_factory(max_workers=2) as w:
            query_resources = [
                r for r in resources if 'AccessKeys' not in r]
            self.log.debug("Querying %d users' api keys" % len(query_resources))
            list(w.map(_user_keys, query_resources))

        matched = []
        for r in resources:
            for p in r['AccessKeys']:
                if self.match(p):
                    matched.append(r)
                    break
        return matched

@User.action_registry.register('remove-keys')
class UserRemoveAccessKey(BaseAction):

    schema = type_schema(
        'remove-keys', age={'type': 'number'}, disable={'type': 'boolean'})

    def process(self, resources):
        client = local_session(self.manager.session_factory).client('iam')

        age = self.data.get('age')
        disable = self.data.get('disable')

        if age:
            threshold_date = datetime(tz=tzutc()) - timedelta(age)

        for r in resources:
            if 'AccessKeys' not in r:
                r['AccessKeys'] = client.list_access_keys(
                    UserName=r['UserName'])['AccessKeyMetadata']
            keys = r['AccessKeys']
            for k in keys:
                if age:
                    if not parse(k['CreateDate']) < threshold_date:
                        continue
                if disable:
                    client.update_access_key(
                        UserName=r['UserName'],
                        AccessKeyId=k['AccessKeyId'],
                        Status='Inactive')
                else:
                    client.delete_access_key(
                        UserName=r['UserName'],
                        AccessKeyId=k['AccessKeyId'])

@resources.register('iam-policy')
class Policy(QueryResourceManager):

    resource_type = 'aws.iam.policy'

@resources.register('iam-profile')
class InstanceProfile(QueryResourceManager):

    resource_type = 'aws.iam.instance-profile'

@resources.register('iam-certificate')
class ServerCerficate(QueryResourceManager):

    resource_type = 'aws.iam.server-certificate'

<commit_msg>Added 'mfa-devices' filter for iam-users. This allows filtering over mfa-devices, however at the moment it seems that it does not allow the filtering of iam-users who do NOT have mfa-devices (e.g. mfa-devices==null).<commit_after># Copyright 2016 Capital One Services, LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from datetime import datetime, timedelta
from dateutil.parser import parse
from dateutil.tz import tzutc

from c7n.actions import BaseAction
from c7n.filters import ValueFilter
from c7n.manager import resources
from c7n.query import QueryResourceManager
from c7n.utils import local_session, type_schema

@resources.register('iam-group')
class Group(QueryResourceManager):

    resource_type = 'aws.iam.group'

@resources.register('iam-role')
class Role(QueryResourceManager):

    resource_type = 'aws.iam.role'

@resources.register('iam-user')
class User(QueryResourceManager):

    resource_type = 'aws.iam.user'

@User.filter_registry.register('policy')
class UserAttachedPolicy(ValueFilter):

    schema = type_schema('policy', rinherit=ValueFilter.schema)

    def process(self, resources, event=None):

        def _user_policies(resource):
            client = local_session(self.manager.session_factory).client('iam')
            resource['AttachedPolicies'] = client.list_attached_user_policies(
                UserName=resource['UserName'])['AttachedPolicies']

        with self.executor_factory(max_workers=2) as w:
            query_resources = [
                r for r in resources if 'AttachedPolicies' not in r]
            self.log.debug("Querying %d users policies" % len(query_resources))
            list(w.map(_user_policies, query_resources))

        matched = []
        for r in resources:
            for p in r['AttachedPolicies']:
                if self.match(p):
                    matched.append(r)
                    break
        return matched

@User.filter_registry.register('access-key')
class UserAccessKey(ValueFilter):

    schema = type_schema('access-key', rinherit=ValueFilter.schema)

    def process(self, resources, event=None):

        def _user_keys(resource):
            client = local_session(self.manager.session_factory).client('iam')
            resource['AccessKeys'] = client.list_access_keys(
                UserName=resource['UserName'])['AccessKeyMetadata']

        with self.executor_factory(max_workers=2) as w:
            query_resources = [
                r for r in resources if 'AccessKeys' not in r]
            self.log.debug("Querying %d users' api keys" % len(query_resources))
            list(w.map(_user_keys, query_resources))

        matched = []
        for r in resources:
            for p in r['AccessKeys']:
                if self.match(p):
                    matched.append(r)
                    break
        return matched

# Mfa-device filter for iam-users
@User.filter_registry.register('mfa-device')
class UserMfaDevice(ValueFilter):

    schema = type_schema('mfa-device', rinherit=ValueFilter.schema)

    def process(self, resources, event=None):

        def _user_mfa_devices(resource):
            client = local_session(self.manager.session_factory).client('iam')
            resource['MfaDevices'] = client.list_mfa_devices(
                UserName=resource['UserName'])['MFADevices']

        with self.executor_factory(max_workers=2) as w:
            query_resources = [
                r for r in resources if 'MfaDevices' not in r]
            self.log.debug("Querying %d users' mfa devices" % len(query_resources))
            list(w.map(_user_mfa_devices, query_resources))

        matched = []
        for r in resources:
            for p in r['MfaDevices']:
                if self.match(p):
                    matched.append(r)
                    break
        return matched

@User.action_registry.register('remove-keys')
class UserRemoveAccessKey(BaseAction):

    schema = type_schema(
        'remove-keys', age={'type': 'number'}, disable={'type': 'boolean'})

    def process(self, resources):
        client = local_session(self.manager.session_factory).client('iam')

        age = self.data.get('age')
        disable = self.data.get('disable')

        if age:
            threshold_date = datetime(tz=tzutc()) - timedelta(age)

        for r in resources:
            if 'AccessKeys' not in r:
                r['AccessKeys'] = client.list_access_keys(
                    UserName=r['UserName'])['AccessKeyMetadata']
            keys = r['AccessKeys']
            for k in keys:
                if age:
                    if not parse(k['CreateDate']) < threshold_date:
                        continue
                if disable:
                    client.update_access_key(
                        UserName=r['UserName'],
                        AccessKeyId=k['AccessKeyId'],
                        Status='Inactive')
                else:
                    client.delete_access_key(
                        UserName=r['UserName'],
                        AccessKeyId=k['AccessKeyId'])

@resources.register('iam-policy')
class Policy(QueryResourceManager):

    resource_type = 'aws.iam.policy'

@resources.register('iam-profile')
class InstanceProfile(QueryResourceManager):

    resource_type = 'aws.iam.instance-profile'

@resources.register('iam-certificate')
class ServerCerficate(QueryResourceManager):

    resource_type = 'aws.iam.server-certificate'
<|endoftext|>


<commit_before>import functools
import sys
import contextlib
from collections import defaultdict
from inspect import Signature
from typing import TYPE_CHECKING, TypeVar, Type, Generic, Callable, Optional

import progressbar
import logging
import time

from ..misc.plugins import PluginVendor, VendorPreset
from ..misc.ux import deprecated

if TYPE_CHECKING:
    from ..knowledge_base import KnowledgeBase
    import angr
    from ..project import Project
    from typing_extensions import ParamSpec
    AnalysisParams = ParamSpec("AnalysisParams")

l = logging.getLogger(name=__name__)

class AnalysisLogEntry:
    def __init__(self, message, exc_info=False):
        if exc_info:
            (e_type, value, traceback) = sys.exc_info()
            self.exc_type = e_type
            self.exc_value = value
            self.exc_traceback = traceback
        else:
            self.exc_type = None
            self.exc_value = None
            self.exc_traceback = None

        self.message = message

    def __getstate__(self):
        return str(self.__dict__.get("exc_type")), \
               str(self.__dict__.get("exc_value")), \
               str(self.__dict__.get("exc_traceback")), \
               self.message

    def __setstate__(self, s):
        self.exc_type, self.exc_value, self.exc_traceback, self.message = s

    def __repr__(self):
        if self.exc_type is None:
            msg_str = repr(self.message)
            if len(msg_str) > 70:
                msg_str = msg_str[:66] + '...'
                if msg_str[0] in ('"', "'"):
                    msg_str += msg_str[0]
            return '<AnalysisLogEntry %s>' % msg_str
        else:
            msg_str = repr(self.message)
            if len(msg_str) > 40:
                msg_str = msg_str[:36] + '...'
                if msg_str[0] in ('"', "'"):
                    msg_str += msg_str[0]
            return '<AnalysisLogEntry %s with %s: %s>' % (msg_str, self.exc_type.__name__, self.exc_value)

A = TypeVar("A", bound="Analysis")
class AnalysesHub(PluginVendor):
    """
    This class contains functions for all the registered and runnable analyses,
    """
    def __init__(self, project):
        super(AnalysesHub, self).__init__()
        self.project = project

    @deprecated()
    def reload_analyses(self): # pylint: disable=no-self-use
        return

    def _init_plugin(self, plugin_cls: Type[A]) -> "AnalysisFactory[A]":
        return AnalysisFactory(self.project, plugin_cls)

    def __getstate__(self):
        s = super(AnalysesHub, self).__getstate__()
        return (s, self.project)

    def __setstate__(self, sd):
        s, self.project = sd
        super(AnalysesHub, self).__setstate__(s)

    def __getitem__(self, plugin_cls: "Type[A]") -> "AnalysisFactory[A]":
        return AnalysisFactory(self.project, plugin_cls)

class AnalysisFactory(Generic[A]):
    def __init__(self, project: "Project", analysis_cls: Type[A]):
        self._project = project
        self._analysis_cls = analysis_cls
        self.__doc__ = ''
        self.__doc__ += analysis_cls.__doc__ or ''
        self.__doc__ += analysis_cls.__init__.__doc__ or ''
        self.__call__.__func__.__signature__ = Signature.from_callable(analysis_cls.__init__)

    def prep(self,
              fail_fast=False,
              kb: Optional["KnowledgeBase"] = None,
              progress_callback: Optional[Callable] = None,
              show_progressbar: bool = False) -> Type[A]:

        @functools.wraps(self._analysis_cls.__init__)
        def wrapper(*args, **kwargs):
            oself = object.__new__(self._analysis_cls)
            oself.named_errors = {}
            oself.errors = []
            oself.log = []

            oself._fail_fast = fail_fast
            oself._name = self._analysis_cls.__name__
            oself.project = self._project
            oself.kb = kb or self._project.kb
            oself._progress_callback = progress_callback

            oself._show_progressbar = show_progressbar
            oself.__init__(*args, **kwargs)
            return oself

        return wrapper # type: ignore

    def __call__(self, *args, **kwargs):
        fail_fast = kwargs.pop('fail_fast', False)
        kb = kwargs.pop('kb', self._project.kb)
        progress_callback = kwargs.pop('progress_callback', None)
        show_progressbar = kwargs.pop('show_progressbar', False)

        w = self.prep(fail_fast=fail_fast,
                  kb=kb,
                  progress_callback=progress_callback,
                  show_progressbar=show_progressbar)

        r = w(*args, **kwargs)
        # clean up so that it's always pickleable
        r._progressbar = None
        return r

class StatusBar(progressbar.widgets.WidgetBase):
    """
    Implements a progressbar component for displaying raw text.
    """
    def __init__(self, width: Optional[int]=40):
        super().__init__()
        self.status: str = ""
        self.width = width

    def __call__(self, progress, data, **kwargs):  # pylint:disable=unused-argument
        if self.width is None:
            return self.status
        if len(self.status) < self.width:
            return self.status.ljust(self.width, " ")
        else:
            return self.status[:self.width]

class Analysis:
    """
    This class represents an analysis on the program.

    :ivar project:  The project for this analysis.
    :type project:  angr.Project
    :ivar KnowledgeBase kb: The knowledgebase object.
    :ivar _progress_callback: A callback function for receiving the progress of this analysis. It only takes
                                        one argument, which is a float number from 0.0 to 100.0 indicating the current
                                        progress.
    :ivar bool _show_progressbar: If a progressbar should be shown during the analysis. It's independent from
                                    _progress_callback.
    :ivar progressbar.ProgressBar _progressbar: The progress bar object.
    """

    project: "Project" = None
    kb: 'KnowledgeBase' = None
    _fail_fast = None
    _name = None
    errors = []
    named_errors = defaultdict(list)
    _progress_callback = None
    _show_progressbar = False
    _progressbar = None
    _statusbar: Optional[StatusBar] = None

    _PROGRESS_WIDGETS = [
        progressbar.Percentage(),
        ' ',
        progressbar.Bar(),
        ' ',
        progressbar.Timer(),
        ' ',
        progressbar.ETA(),
        ' ',
        StatusBar(),
    ]

    @contextlib.contextmanager
    def _resilience(self, name=None, exception=Exception):
        try:
            yield
        except exception:  # pylint:disable=broad-except
            if self._fail_fast:
                raise
            else:
                error = AnalysisLogEntry("exception occurred", exc_info=True)
                l.error("Caught and logged %s with resilience: %s", error.exc_type.__name__, error.exc_value)
                if name is None:
                    self.errors.append(error)
                else:
                    self.named_errors[name].append(error)

    def _initialize_progressbar(self):
        """
        Initialize the progressbar.
        :return: None
        """

        self._progressbar = progressbar.ProgressBar(widgets=Analysis._PROGRESS_WIDGETS, max_value=10000 * 100).start()
        self._statusbar = self._progressbar.widgets[-1]

    def _update_progress(self, percentage, text=None, **kwargs):
        """
        Update the progress with a percentage, including updating the progressbar as well as calling the progress
        callback.

        :param float percentage:    Percentage of the progressbar. from 0.0 to 100.0.
        :param kwargs:              Other parameters that will be passed to the progress_callback handler.
        :return: None
        """

        if self._show_progressbar:
            if self._progressbar is None:
                self._initialize_progressbar()

            self._progressbar.update(percentage * 10000)

        if text is not None and self._statusbar is not None:
            self._statusbar.status = text

        if self._progress_callback is not None:
            self._progress_callback(percentage, text=text, **kwargs)  # pylint:disable=not-callable

    def _finish_progress(self):
        """
        Mark the progressbar as finished.
        :return: None
        """

        if self._show_progressbar:
            if self._progressbar is None:
                self._initialize_progressbar()
            if self._progressbar is not None:
                self._progressbar.finish()
                # Remove the progressbar object so it will not be pickled
                self._progressbar = None

        if self._progress_callback is not None:
            self._progress_callback(100.0)  # pylint:disable=not-callable

    @staticmethod
    def _release_gil(ctr, freq, sleep_time=0.001):
        """
        Periodically calls time.sleep() and releases the GIL so other threads (like, GUI threads) have a much better
        chance to be scheduled, and other critical components (like the GUI) can be kept responsiveness.

        This is, of course, a hack before we move all computational intensive tasks to pure C++ implementations.

        :param int ctr:     A number provided by the caller.
        :param int freq:    How frequently time.sleep() should be called. time.sleep() is called when ctr % freq == 0.
        :param sleep_time:  Number (or fraction) of seconds to sleep.
        :return:            None
        """

        if ctr != 0 and ctr % freq == 0:
            time.sleep(sleep_time)

    def __getstate__(self):
        d = dict(self.__dict__)
        if "_progressbar" in d:
            del d['_progressbar']
        if '_progress_callback' in d:
            del d['_progress_callback']
        if '_statusbar' in d:
            del d['_statusbar']
        return d

    def __setstate__(self, state):
        self.__dict__.update(state)

    def __repr__(self):
        return '<%s Analysis Result at %#x>' % (self._name, id(self))

default_analyses = VendorPreset()
AnalysesHub.register_preset('default', default_analyses)
<commit_msg>un-break resilience with named errors; closes #3145<commit_after>import functools
import sys
import contextlib
from collections import defaultdict
from inspect import Signature
from typing import TYPE_CHECKING, TypeVar, Type, Generic, Callable, Optional

import progressbar
import logging
import time

from ..misc.plugins import PluginVendor, VendorPreset
from ..misc.ux import deprecated

if TYPE_CHECKING:
    from ..knowledge_base import KnowledgeBase
    import angr
    from ..project import Project
    from typing_extensions import ParamSpec
    AnalysisParams = ParamSpec("AnalysisParams")

l = logging.getLogger(name=__name__)

class AnalysisLogEntry:
    def __init__(self, message, exc_info=False):
        if exc_info:
            (e_type, value, traceback) = sys.exc_info()
            self.exc_type = e_type
            self.exc_value = value
            self.exc_traceback = traceback
        else:
            self.exc_type = None
            self.exc_value = None
            self.exc_traceback = None

        self.message = message

    def __getstate__(self):
        return str(self.__dict__.get("exc_type")), \
               str(self.__dict__.get("exc_value")), \
               str(self.__dict__.get("exc_traceback")), \
               self.message

    def __setstate__(self, s):
        self.exc_type, self.exc_value, self.exc_traceback, self.message = s

    def __repr__(self):
        if self.exc_type is None:
            msg_str = repr(self.message)
            if len(msg_str) > 70:
                msg_str = msg_str[:66] + '...'
                if msg_str[0] in ('"', "'"):
                    msg_str += msg_str[0]
            return '<AnalysisLogEntry %s>' % msg_str
        else:
            msg_str = repr(self.message)
            if len(msg_str) > 40:
                msg_str = msg_str[:36] + '...'
                if msg_str[0] in ('"', "'"):
                    msg_str += msg_str[0]
            return '<AnalysisLogEntry %s with %s: %s>' % (msg_str, self.exc_type.__name__, self.exc_value)

A = TypeVar("A", bound="Analysis")
class AnalysesHub(PluginVendor):
    """
    This class contains functions for all the registered and runnable analyses,
    """
    def __init__(self, project):
        super(AnalysesHub, self).__init__()
        self.project = project

    @deprecated()
    def reload_analyses(self): # pylint: disable=no-self-use
        return

    def _init_plugin(self, plugin_cls: Type[A]) -> "AnalysisFactory[A]":
        return AnalysisFactory(self.project, plugin_cls)

    def __getstate__(self):
        s = super(AnalysesHub, self).__getstate__()
        return (s, self.project)

    def __setstate__(self, sd):
        s, self.project = sd
        super(AnalysesHub, self).__setstate__(s)

    def __getitem__(self, plugin_cls: "Type[A]") -> "AnalysisFactory[A]":
        return AnalysisFactory(self.project, plugin_cls)

class AnalysisFactory(Generic[A]):
    def __init__(self, project: "Project", analysis_cls: Type[A]):
        self._project = project
        self._analysis_cls = analysis_cls
        self.__doc__ = ''
        self.__doc__ += analysis_cls.__doc__ or ''
        self.__doc__ += analysis_cls.__init__.__doc__ or ''
        self.__call__.__func__.__signature__ = Signature.from_callable(analysis_cls.__init__)

    def prep(self,
              fail_fast=False,
              kb: Optional["KnowledgeBase"] = None,
              progress_callback: Optional[Callable] = None,
              show_progressbar: bool = False) -> Type[A]:

        @functools.wraps(self._analysis_cls.__init__)
        def wrapper(*args, **kwargs):
            oself = object.__new__(self._analysis_cls)
            oself.named_errors = defaultdict(list)
            oself.errors = []
            oself.log = []

            oself._fail_fast = fail_fast
            oself._name = self._analysis_cls.__name__
            oself.project = self._project
            oself.kb = kb or self._project.kb
            oself._progress_callback = progress_callback

            oself._show_progressbar = show_progressbar
            oself.__init__(*args, **kwargs)
            return oself

        return wrapper # type: ignore

    def __call__(self, *args, **kwargs):
        fail_fast = kwargs.pop('fail_fast', False)
        kb = kwargs.pop('kb', self._project.kb)
        progress_callback = kwargs.pop('progress_callback', None)
        show_progressbar = kwargs.pop('show_progressbar', False)

        w = self.prep(fail_fast=fail_fast,
                  kb=kb,
                  progress_callback=progress_callback,
                  show_progressbar=show_progressbar)

        r = w(*args, **kwargs)
        # clean up so that it's always pickleable
        r._progressbar = None
        return r

class StatusBar(progressbar.widgets.WidgetBase):
    """
    Implements a progressbar component for displaying raw text.
    """
    def __init__(self, width: Optional[int]=40):
        super().__init__()
        self.status: str = ""
        self.width = width

    def __call__(self, progress, data, **kwargs):  # pylint:disable=unused-argument
        if self.width is None:
            return self.status
        if len(self.status) < self.width:
            return self.status.ljust(self.width, " ")
        else:
            return self.status[:self.width]

class Analysis:
    """
    This class represents an analysis on the program.

    :ivar project:  The project for this analysis.
    :type project:  angr.Project
    :ivar KnowledgeBase kb: The knowledgebase object.
    :ivar _progress_callback: A callback function for receiving the progress of this analysis. It only takes
                                        one argument, which is a float number from 0.0 to 100.0 indicating the current
                                        progress.
    :ivar bool _show_progressbar: If a progressbar should be shown during the analysis. It's independent from
                                    _progress_callback.
    :ivar progressbar.ProgressBar _progressbar: The progress bar object.
    """

    project: "Project" = None
    kb: 'KnowledgeBase' = None
    _fail_fast = None
    _name = None
    errors = []
    named_errors = defaultdict(list)
    _progress_callback = None
    _show_progressbar = False
    _progressbar = None
    _statusbar: Optional[StatusBar] = None

    _PROGRESS_WIDGETS = [
        progressbar.Percentage(),
        ' ',
        progressbar.Bar(),
        ' ',
        progressbar.Timer(),
        ' ',
        progressbar.ETA(),
        ' ',
        StatusBar(),
    ]

    @contextlib.contextmanager
    def _resilience(self, name=None, exception=Exception):
        try:
            yield
        except exception:  # pylint:disable=broad-except
            if self._fail_fast:
                raise
            else:
                error = AnalysisLogEntry("exception occurred", exc_info=True)
                l.error("Caught and logged %s with resilience: %s", error.exc_type.__name__, error.exc_value)
                if name is None:
                    self.errors.append(error)
                else:
                    self.named_errors[name].append(error)

    def _initialize_progressbar(self):
        """
        Initialize the progressbar.
        :return: None
        """

        self._progressbar = progressbar.ProgressBar(widgets=Analysis._PROGRESS_WIDGETS, max_value=10000 * 100).start()
        self._statusbar = self._progressbar.widgets[-1]

    def _update_progress(self, percentage, text=None, **kwargs):
        """
        Update the progress with a percentage, including updating the progressbar as well as calling the progress
        callback.

        :param float percentage:    Percentage of the progressbar. from 0.0 to 100.0.
        :param kwargs:              Other parameters that will be passed to the progress_callback handler.
        :return: None
        """

        if self._show_progressbar:
            if self._progressbar is None:
                self._initialize_progressbar()

            self._progressbar.update(percentage * 10000)

        if text is not None and self._statusbar is not None:
            self._statusbar.status = text

        if self._progress_callback is not None:
            self._progress_callback(percentage, text=text, **kwargs)  # pylint:disable=not-callable

    def _finish_progress(self):
        """
        Mark the progressbar as finished.
        :return: None
        """

        if self._show_progressbar:
            if self._progressbar is None:
                self._initialize_progressbar()
            if self._progressbar is not None:
                self._progressbar.finish()
                # Remove the progressbar object so it will not be pickled
                self._progressbar = None

        if self._progress_callback is not None:
            self._progress_callback(100.0)  # pylint:disable=not-callable

    @staticmethod
    def _release_gil(ctr, freq, sleep_time=0.001):
        """
        Periodically calls time.sleep() and releases the GIL so other threads (like, GUI threads) have a much better
        chance to be scheduled, and other critical components (like the GUI) can be kept responsiveness.

        This is, of course, a hack before we move all computational intensive tasks to pure C++ implementations.

        :param int ctr:     A number provided by the caller.
        :param int freq:    How frequently time.sleep() should be called. time.sleep() is called when ctr % freq == 0.
        :param sleep_time:  Number (or fraction) of seconds to sleep.
        :return:            None
        """

        if ctr != 0 and ctr % freq == 0:
            time.sleep(sleep_time)

    def __getstate__(self):
        d = dict(self.__dict__)
        if "_progressbar" in d:
            del d['_progressbar']
        if '_progress_callback' in d:
            del d['_progress_callback']
        if '_statusbar' in d:
            del d['_statusbar']
        return d

    def __setstate__(self, state):
        self.__dict__.update(state)

    def __repr__(self):
        return '<%s Analysis Result at %#x>' % (self._name, id(self))

default_analyses = VendorPreset()
AnalysesHub.register_preset('default', default_analyses)
<|endoftext|>


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

"""Django signal handlers for modoboa_postfix_autoreply."""

from __future__ import unicode_literals

from django.conf.urls import url
from django.db.models import signals
from django.dispatch import receiver
from django.utils.translation import ugettext as _

from modoboa.admin import models as admin_models, signals as admin_signals
from modoboa.core import signals as core_signals
from modoboa.transport import models as tr_models

from . import forms, models

@receiver(signals.post_save, sender=admin_models.Domain)
def manage_transport_entry(sender, instance, **kwargs):
    """Create or update a transport entry for this domain."""
    if kwargs.get("created"):
        tr_models.Transport.objects.get_or_create(
            pattern="autoreply.{}".format(instance), service="autoreply"
        )
        return
    oldname = getattr(instance, "oldname", "None")
    if oldname is None or oldname == instance.name:
        return
    tr_models.Transport.objects.filter(
        pattern="autoreply.{}".format(oldname)).update(
            pattern="autoreply.{}".format(instance.name))
    qset = (
        admin_models.AliasRecipient.objects
        .select_related("alias", "r_mailbox")
        .filter(
            alias__domain=instance, alias__internal=True,
            address__contains="@autoreply")
    )
    for alr in qset:
        alr.address = alr.address.replace(oldname, instance.name)
        alr.save()

@receiver(signals.post_delete, sender=admin_models.Domain)
def delete_transport_entry(sender, instance, **kwargs):
    """Delete a transport entry."""
    tr_models.Transport.objects.filter(
        pattern="autoreply.{}".format(instance)).delete()

@receiver(signals.post_save, sender=admin_models.Mailbox)
def rename_autoreply_alias(sender, instance, **kwargs):
    """Rename AR alias if needed."""
    old_address = getattr(instance, "old_full_address", None)
    if old_address is None or old_address == instance.full_address:
        return
    admin_models.AliasRecipient.objects.filter(
        address__contains="{}@autoreply".format(old_address)).update(
            address="{}@autoreply.{}".format(
                instance.full_address, instance.domain))

@receiver(signals.post_delete, sender=admin_models.Mailbox)
def delete_autoreply_alias(sender, instance, **kwargs):
    """Delete alias."""
    admin_models.AliasRecipient.objects.filter(
        address="{}@autoreply.{}".format(
            instance.full_address, instance.domain)).delete()

@receiver(signals.post_save, sender=models.ARmessage)
def manage_autoreply_alias(sender, instance, **kwargs):
    """Create or delete the alias."""
    ar_alias_address = "{}@autoreply.{}".format(
        instance.mbox.full_address, instance.mbox.domain)
    admin_models.Alias.objects.get(
        address=instance.mbox.full_address, domain=instance.mbox.domain,
        internal=True)
    alias, created = admin_models.Alias.objects.get_or_create(
        address=instance.mbox.full_address, domain=instance.mbox.domain,
        internal=True)
    if instance.enabled:
        admin_models.AliasRecipient.objects.get_or_create(
            alias=alias, address=ar_alias_address)
    else:
        admin_models.AliasRecipient.objects.filter(
            address=ar_alias_address).delete()

@receiver(core_signals.extra_uprefs_routes)
def extra_routes(sender, **kwargs):
    """Add extra routes."""
    from . import views
    return [
        url(r'^user/autoreply/$', views.autoreply, name="autoreply")
    ]

@receiver(core_signals.extra_static_content)
def extra_js(sender, caller, st_type, user, **kwargs):
    """Add static content."""
    if caller != "user_index" or st_type != "js":
        return ""
    return """function autoreply_cb() {
    $('.datefield').datetimepicker({
        format: 'YYYY-MM-DD HH:mm:ss',
        language: '%(lang)s'
    });
}
""" % {"lang": user.language}

@receiver(core_signals.extra_user_menu_entries)
def menu(sender, location, user, **kwargs):
    """Inject new menu entries."""
    if location != "uprefs_menu" or not hasattr(user, "mailbox"):
        return []
    return [
        {"name": "autoreply",
         "class": "ajaxnav",
         "url": "autoreply/",
         "label": _("Auto-reply message")}
    ]

@receiver(admin_signals.extra_account_forms)
def extra_account_form(sender, user, account, **kwargs):
    """Add autoreply form to the account edition form."""
    result = []
    if user.role in ("SuperAdmins", "DomainAdmins"):
        if hasattr(account, "mailbox"):
            extraform = {
                "id": "auto_reply_message",
                "title": _("Auto reply"),
                "cls": forms.ARmessageForm,
                "new_args": [account.mailbox]
            }
            result.append(extraform)
    return result

@receiver(admin_signals.get_account_form_instances)
def fill_account_tab(sender, user, account, **kwargs):
    """Return form instance."""
    condition = (
        user.role not in ("SuperAdmins", "DomainAdmins") or
        not hasattr(account, "mailbox"))
    if condition:
        return {}
    return {"auto_reply_message": account.mailbox.armessage_set.first()}
<commit_msg>Fixed datetimepicker init.<commit_after># -*- coding: utf-8 -*-

"""Django signal handlers for modoboa_postfix_autoreply."""

from __future__ import unicode_literals

from django.conf.urls import url
from django.db.models import signals
from django.dispatch import receiver
from django.utils.translation import ugettext as _

from modoboa.admin import models as admin_models, signals as admin_signals
from modoboa.core import signals as core_signals
from modoboa.transport import models as tr_models

from . import forms, models

@receiver(signals.post_save, sender=admin_models.Domain)
def manage_transport_entry(sender, instance, **kwargs):
    """Create or update a transport entry for this domain."""
    if kwargs.get("created"):
        tr_models.Transport.objects.get_or_create(
            pattern="autoreply.{}".format(instance), service="autoreply"
        )
        return
    oldname = getattr(instance, "oldname", "None")
    if oldname is None or oldname == instance.name:
        return
    tr_models.Transport.objects.filter(
        pattern="autoreply.{}".format(oldname)).update(
            pattern="autoreply.{}".format(instance.name))
    qset = (
        admin_models.AliasRecipient.objects
        .select_related("alias", "r_mailbox")
        .filter(
            alias__domain=instance, alias__internal=True,
            address__contains="@autoreply")
    )
    for alr in qset:
        alr.address = alr.address.replace(oldname, instance.name)
        alr.save()

@receiver(signals.post_delete, sender=admin_models.Domain)
def delete_transport_entry(sender, instance, **kwargs):
    """Delete a transport entry."""
    tr_models.Transport.objects.filter(
        pattern="autoreply.{}".format(instance)).delete()

@receiver(signals.post_save, sender=admin_models.Mailbox)
def rename_autoreply_alias(sender, instance, **kwargs):
    """Rename AR alias if needed."""
    old_address = getattr(instance, "old_full_address", None)
    if old_address is None or old_address == instance.full_address:
        return
    admin_models.AliasRecipient.objects.filter(
        address__contains="{}@autoreply".format(old_address)).update(
            address="{}@autoreply.{}".format(
                instance.full_address, instance.domain))

@receiver(signals.post_delete, sender=admin_models.Mailbox)
def delete_autoreply_alias(sender, instance, **kwargs):
    """Delete alias."""
    admin_models.AliasRecipient.objects.filter(
        address="{}@autoreply.{}".format(
            instance.full_address, instance.domain)).delete()

@receiver(signals.post_save, sender=models.ARmessage)
def manage_autoreply_alias(sender, instance, **kwargs):
    """Create or delete the alias."""
    ar_alias_address = "{}@autoreply.{}".format(
        instance.mbox.full_address, instance.mbox.domain)
    admin_models.Alias.objects.get(
        address=instance.mbox.full_address, domain=instance.mbox.domain,
        internal=True)
    alias, created = admin_models.Alias.objects.get_or_create(
        address=instance.mbox.full_address, domain=instance.mbox.domain,
        internal=True)
    if instance.enabled:
        admin_models.AliasRecipient.objects.get_or_create(
            alias=alias, address=ar_alias_address)
    else:
        admin_models.AliasRecipient.objects.filter(
            address=ar_alias_address).delete()

@receiver(core_signals.extra_uprefs_routes)
def extra_routes(sender, **kwargs):
    """Add extra routes."""
    from . import views
    return [
        url(r'^user/autoreply/$', views.autoreply, name="autoreply")
    ]

@receiver(core_signals.extra_static_content)
def extra_js(sender, caller, st_type, user, **kwargs):
    """Add static content."""
    if caller != "user_index" or st_type != "js":
        return ""
    return """function autoreply_cb() {
    $('.datefield').datetimepicker({
        format: 'YYYY-MM-DD HH:mm:ss',
        locale: '%(lang)s'
    });
}
""" % {"lang": user.language}

@receiver(core_signals.extra_user_menu_entries)
def menu(sender, location, user, **kwargs):
    """Inject new menu entries."""
    if location != "uprefs_menu" or not hasattr(user, "mailbox"):
        return []
    return [
        {"name": "autoreply",
         "class": "ajaxnav",
         "url": "autoreply/",
         "label": _("Auto-reply message")}
    ]

@receiver(admin_signals.extra_account_forms)
def extra_account_form(sender, user, account, **kwargs):
    """Add autoreply form to the account edition form."""
    result = []
    if user.role in ("SuperAdmins", "DomainAdmins"):
        if hasattr(account, "mailbox"):
            extraform = {
                "id": "auto_reply_message",
                "title": _("Auto reply"),
                "cls": forms.ARmessageForm,
                "new_args": [account.mailbox]
            }
            result.append(extraform)
    return result

@receiver(admin_signals.get_account_form_instances)
def fill_account_tab(sender, user, account, **kwargs):
    """Return form instance."""
    condition = (
        user.role not in ("SuperAdmins", "DomainAdmins") or
        not hasattr(account, "mailbox"))
    if condition:
        return {}
    return {"auto_reply_message": account.mailbox.armessage_set.first()}
<|endoftext|>