<commit_before>'''Constants and methods for astronomy & astrophysics'''
from __future__ import division
from __future__ import print_function
from past.utils import old_div

# Astronomy utilities module

import numpy as np
import scipy as sc
from scipy import integrate

# constants for astronomy
rsun_cm = 6.955e10
rsun_cm_s = 'cm'
lsun_erg_s = 3.839e33
lsun_erg_s_unit = 'erg/s'
msun_g = 1.9891e33
msun_unit = 'erg s^-1'
mass_earth_g=5.9742e27

au_cm=1.49598e13
au_cm_unit='cm'

radiation_constant = 7.566e-15
radiation_constant_unit = 'erg cm^-3 K^-4'
speed_light = 2.998e10
speed_light_unit = 'cm s^-1'
grav_const = 6.67259e-8
grav_const_unit = 'cm^3 g^-1 s^-2'

# constants for physics
mass_H_atom=1.674e-24
mass_H_atom_unit='g'
mass_electron=9.10938291E-28
mass_electron_unit='g'
planck_constant_h=6.62606957E-27
planck_constant_h_unit='erg s'
atomic_mass_unit=1.660538921e-24
atomic_mass_unit_unit='g'
boltzmann_constant=1.3806488e-16
boltzmann_constant_unit='erg K^-1'
avogadro_constant=6.02214179e23
avogadro_constant_unit='mol^-1'

def visc_mol_sol(T,rho,X):
    '''
    Molecular plasma viscosity (Spitzer 1962)

    Parameters
    ----------
    X : float
       H mass fraction
    T : float
       temperature in K
    rho : float
       density in cgs

    Returns
    -------
    nu : float
       molecular diffusivity in [cm**2/s]

    Notes
    -----
    According to Eq 22 in Schatzman (1977). Assume log Lambda = 15.
    (see Table 5.1), a H/He mix (for different mix use Eq. 5.54 in
    Spitzer text book)

    Examples
    --------
    see astronomy.visc_rad_kap_sc

    '''
    visc_mol = 1.84e-17*(1.+7.*X)*(old_div(T**2.5,rho))
    return visc_mol

def visc_rad_kap_sc(T,rho,X):
    '''
    Radiative viscosity (Thomas, 1930) for e- scattering opacity

    Parameters
    ----------
    X : float
       H mass fraction
    T : float
       temperature in K
    rho : float
       density in cgs

    Returns
    -------
    nu : float
       radiative diffusivity in [cm**2/s]

    Examples
    --------
    >>> In [1]: import astronomy as ast
    >>> In [2]: l = 100*1.e5 # 100km
    >>> In [3]: v = 1.e5     # typical velocity
    >>> In [4]: T   = 90.e6  # temperature
    >>> In [5]: X   = 0.001  # H mass fraction
    >>> In [6]: rho = 100.   # density
    >>> In [7]: nu = ast.visc_rad_kap_sc(T,rho,X)
    >>> In [8]: Re=v*l/nu
    >>> In [9]: print "Re_rad = "+str('%g'%Re)
    >>> Re_rad = 4.43512e+08

    Notes
    -----
    Eqn. 14' in Schatzman, 1977, assume electron scattering opacity
    kappa_sc = 0.2*(1+X), Kippenhahn (2nd edn, Eqn 17.2)

    '''
    kappa = 0.2*(1.+X)
    nu_rad = 6.88e-26*(old_div(T**4,(kappa*rho**2)))
    return nu_rad

def Gamma1_gasrad(beta):
    '''
    Gamma1 for a mix of ideal gas and radiation

    Hansen & Kawaler, page 177, Eqn. 3.110

    Parameters
    ----------
    beta : float
        Gas pressure fraction Pgas/(Pgas+Prad)

    '''
    Gamma3minus1 = (old_div(2.,3.))*(4.-3.*beta)/(8.-7.*beta)
    Gamma1 = beta + (4.-3.*beta) * Gamma3minus1
    return Gamma1

def Pgas(rho,T,mu):
    '''
    P = R/mu * rho * T

    Parameters
    ----------
    mu : float
        Mean molecular weight
    rho : float
        Density [cgs]
    T : float
        Temperature [K]

    '''
    R = old_div(boltzmann_constant, atomic_mass_unit)
    return (old_div(R,mu)) * rho * T

def Prad(T,mu):
    '''
    P = a/3 * T**4

    Parameters
    ----------
    a : float
        Radiation constant [erg cm^-3 K^-4].
    T : float
        Temperature [K].

    '''
    return (old_div(radiation_constant,3.)) * T**4

def mimf_ferrario(mi):
    ''' Curvature MiMf from Ferrario etal. 2005MNRAS.361.1131.'''

    mf=-0.00012336*mi**6+0.003160*mi**5-0.02960*mi**4+\
      0.12350*mi**3-0.21550*mi**2+0.19022*mi+0.46575
    return mf

def core_mass_L(MH):
    '''
    Core-mass luminosity relationship from Bloecker (1993).

    Parameters
    ----------
    MH
        Core mass in Msun.

    Returns
    -------
    L
        Luminosity in Lsun

    '''
    return 62200*(MH-0.487)

def imf(m):
    '''
    Returns
    -------
    N(M)dM
        for given mass according to Kroupa IMF, vectorization
        available via vimf()

    '''

    m1 = 0.08; m2 = 0.50
    a1 = 0.30; a2 = 1.30; a3 = 2.3
    const2 = m1**-a1 -m1**-a2
    const3 = m2**-a2 -m2**-a3

    if m < 0.08:
        alpha = 0.3
        const = -const2 -const3
    elif m < 0.50:
        alpha = 1.3
        const = -const3
    else:
        alpha = 2.3
        const = 0.0
    # print m,alpha, const, m**-alpha + const
    return m**-alpha + const

vimf = np.vectorize(imf)

def int_imf_dm(m1,m2,m,imf,bywhat='bymass',integral='normal'):
    '''
    Integrate IMF between m1 and m2.

    Parameters
    ----------
    m1 : float
        Min mass
    m2 : float
        Max mass
    m : float
        Mass array
    imf : float
        IMF array
    bywhat : string, optional
        'bymass' integrates the mass that goes into stars of
        that mass interval; or 'bynumber' which integrates the number
        of stars in that mass interval.  The default is 'bymass'.
    integrate : string, optional
        'normal' uses sc.integrate.trapz; 'cum' returns cumulative
        trapezoidal integral.  The default is 'normal'.

    '''

    ind_m = (m >= min(m1,m2)) & (m <= max(m1,m2))
    if integral is 'normal':
        int_func = sc.integrate.trapz
    elif integral is 'cum':
        int_func = sc.integrate.cumtrapz
    else:
        print("Error in int_imf_dm: don't know how to integrate")
        return 0

    if bywhat is 'bymass':
        return int_func(m[ind_m]*imf[ind_m],m[ind_m])
    elif bywhat is 'bynumber':
        return int_func(imf[ind_m],m[ind_m])
    else:
        print("Error in int_imf_dm: don't know by what to integrate")
        return 0

def am_orb(m1,m2,a,e):
    '''
    orbital angular momentum.

    e.g Ge etal2010

    Parameters
    ----------
    m1, m2 : float
        Masses of both stars in Msun.
    A : float
        Separation in Rsun.
    e : float
        Eccentricity

    '''

    a_cm  = a * rsun_cm
    m1_g = m1 * msun_g
    m2_g = m2 * msun_g

    J_orb=np.sqrt(grav_const*a_cm*(old_div((m1_g**2*m2_g**2),(m1_g+m2_g))))*(1-e**2)
    return J_orb

def mass_loss_loon05(L,Teff):
    '''
    mass loss rate van Loon etal (2005).

    Parameters
    ----------
    L : float
        L in L_sun.
    Teff : float
        Teff in K.

    Returns
    -------
    Mdot
        Mdot in Msun/yr

    Notes
    -----
    ref: van Loon etal 2005, A&A 438, 273

    '''

    Mdot = -5.65 + np.log10(old_div(L,10.**4)) -6.3*np.log10(old_div(Teff,3500.))
    return Mdot

def energ_orb(m1,m2,r):
    '''
    Parameters
    ----------
    m1, m2 : float
        M in Msun.
    r : float
        Distance in Rsun.

    Returns
    -------
    Epot
        Epot in erg.

    '''
    epo = -grav_const * m1 * m2 * msun_g**2 / (r * rsun_cm)
    return epo

def period(A,M1,M2):
    """
    calculate binary period from separation.

    Parameters
    ----------
    A : float
        separation A Rsun.
    M1, M2 : float
        M in Msun.

    Returns
    -------
    p
        period in days.

    """

    A *= rsun_cm
    print(A)
    velocity = np.sqrt(grav_const*msun_g*(M1+M2)/A)
    print(old_div(velocity,1.e5))

    p = 2.*np.pi * A / velocity

    p  /= (60*60*24.)
    return p

def escape_velocity(M,R):
    """
    escape velocity.

    Parameters
    ----------
    M : float
        Mass in solar masses.
    R : float
        Radius in solar radiu.

    Returns
    -------
    v_escape
        in km/s.

    """

    ve = np.sqrt(2.*grav_const*M*msun_g/(R*rsun_cm))
    ve = ve*1.e-5
    return ve

def Nasv(macs,T):
    '''
    Returns
    -------
    Na*<sigma v>
        for MACS [mb] at T [K].

    '''

    Na = avogadro_constant
    k  = boltzmann_constant
    vtherm=(2.*k*T/mass_H_atom)**0.5

    s  = macs*1.e-27
    Nasv = s*vtherm*Na
    return Nasv

def macs(nasv,T):
    '''
    Returns
    -------
    MACS
        [mb] at T [K] from Na*<sigma v>.

    '''

    Na = avogadro_constant
    k  = boltzmann_constant
    vtherm=(2.*k*T/mass_H_atom)**0.5

    s      = old_div(nasv,(vtherm*Na))
    macs   = s*1.e27
    return macs

def mu_e(X):
    '''
    mean molecular weight per free electron, assuming full ionisation, and
    approximating mu_i/Z_i ~ 2 for all elements heavier then Helium.

    (Kippenhahn & Weigert, Ch 13.1, Eq. 13.8)

    Parameters
    ----------
    X : float
        Mass fraction of H.

    '''

    try:
        mu_e = old_div(2.,(1.+X))
    except TypeError:
        X=np.array([X])
        mu_e = old_div(2.,(1.+X))

    return mu_e

def mu(X,Z,A):
    '''
    mean molecular weight assuming full ionisation.

    (Kippenhahn & Weigert, Ch 13.1, Eq. 13.6)

    Parameters
    ----------
    X : float
        Mass fraction vector.
    Z : float
        Charge number vector.
    A : float
        Mass number vector.

    '''

    if not isinstance(Z,np.ndarray):
        Z = np.array(Z)
    if not isinstance(A,np.ndarray):
        A = np.array(A)
    if not isinstance(X,np.ndarray):
        X = np.array(X)

    try:
        mu = old_div(1.,sum(X*(1.+Z)/A))
    except TypeError:
        X=np.array([X])
        A=np.array([A])
        Z=np.array([Z])
        mu = old_div(1.,sum(X*(1.+Z)/A))

    return mu

def Trho_idrad(rho,mu):
    '''
    T(rho) that separates P_rad from P_gas dominated regions.

    Kippenhahn & Weigert, Eq. 16.10

    Parameters
    ----------
    rho : float
        Density array [cgs].
    mu : float
        Mean molecular weight.

    '''

    T = 3.2E7 * (old_div(rho,mu))**(old_div(1.,3.))
    return T

def Trho_iddeg(rho,mu,mu_e):
    '''
    T(rho) that separates ideal gas and degenerate pressure dominated regions.

    Kippenhahn & Weigert, Eq. 16.6

    Parameters
    ----------
    rho : float
        Density array [cgs].
    mu : float
        Mean molecular weight.
    mu_e : float
        Mean molecular weight per free electron.

    '''

    T = 1.207E5 * rho**(old_div(2.,3.)) * mu / mu_e**(old_div(5.,3.))
    return T

<commit_msg>Update astronomy.py<commit_after>'''Constants and methods for astronomy & astrophysics'''
from __future__ import division
from __future__ import print_function
from past.utils import old_div

# Astronomy utilities module

import numpy as np
import scipy as sc
from scipy import integrate

# constants for astronomy
rsun_cm = 6.955e10
rsun_cm_s = 'cm'
lsun_erg_s = 3.839e33
lsun_erg_s_unit = 'erg/s'
msun_g = 1.9891e33
msun_unit = 'erg s^-1'
mass_earth_g=5.9742e27

au_cm=1.49598e13
au_cm_unit='cm'

radiation_constant = 7.566e-15
radiation_constant_unit = 'erg cm^-3 K^-4'
speed_light = 2.998e10
speed_light_unit = 'cm s^-1'
grav_const = 6.67259e-8
grav_const_unit = 'cm^3 g^-1 s^-2'

# constants for physics
mass_H_atom=1.674e-24
mass_H_atom_unit='g'
mass_electron=9.10938291E-28
mass_electron_unit='g'
planck_constant_h=6.62606957E-27
planck_constant_h_unit='erg s'
atomic_mass_unit=1.660538921e-24
atomic_mass_unit_unit='g'
boltzmann_constant=1.3806488e-16
boltzmann_constant_unit='erg K^-1'
avogadro_constant=6.02214179e23
avogadro_constant_unit='mol^-1'

def visc_mol_sol(T,rho,X):
    '''
    Molecular plasma viscosity (Spitzer 1962)

    Parameters
    ----------
    X : float
       H mass fraction
    T : float
       temperature in K
    rho : float
       density in cgs

    Returns
    -------
    nu : float
       molecular diffusivity in [cm**2/s]

    Notes
    -----
    According to Eq 22 in Schatzman (1977). Assume log Lambda = 15.
    (see Table 5.1), a H/He mix (for different mix use Eq. 5.54 in
    Spitzer text book)

    Examples
    --------
    see astronomy.visc_rad_kap_sc

    '''
    visc_mol = 1.84e-17*(1.+7.*X)*(old_div(T**2.5,rho))
    return visc_mol

def visc_rad_kap_sc(T,rho,X):
    '''
    Radiative viscosity (Thomas, 1930) for e- scattering opacity

    Parameters
    ----------
    X : float
       H mass fraction
    T : float
       temperature in K
    rho : float
       density in cgs

    Returns
    -------
    nu : float
       radiative diffusivity in [cm**2/s]

    Examples
    --------
    >>> In [1]: import astronomy as ast
    >>> In [2]: l = 100*1.e5 # 100km
    >>> In [3]: v = 1.e5     # typical velocity
    >>> In [4]: T   = 90.e6  # temperature
    >>> In [5]: X   = 0.001  # H mass fraction
    >>> In [6]: rho = 100.   # density
    >>> In [7]: nu = ast.visc_rad_kap_sc(T,rho,X)
    >>> In [8]: Re=v*l/nu
    >>> In [9]: print "Re_rad = "+str('%g'%Re)
    >>> Re_rad = 4.43512e+08

    Notes
    -----
    Eqn. 14' in Schatzman, 1977, assume electron scattering opacity
    kappa_sc = 0.2*(1+X), Kippenhahn (2nd edn, Eqn 17.2)

    '''
    kappa = 0.2*(1.+X)
    nu_rad = 6.88e-26*(old_div(T**4,(kappa*rho**2)))
    return nu_rad

def Gamma1_gasrad(beta):
    '''
    Gamma1 for a mix of ideal gas and radiation

    Hansen & Kawaler, page 177, Eqn. 3.110

    Parameters
    ----------
    beta : float
        Gas pressure fraction Pgas/(Pgas+Prad)

    '''
    Gamma3minus1 = (old_div(2.,3.))*(4.-3.*beta)/(8.-7.*beta)
    Gamma1 = beta + (4.-3.*beta) * Gamma3minus1
    return Gamma1

def Pgas(rho,T,mu):
    '''
    P = R/mu * rho * T

    Parameters
    ----------
    mu : float
        Mean molecular weight
    rho : float
        Density [cgs]
    T : float
        Temperature [K]

    '''
    R = old_div(boltzmann_constant, atomic_mass_unit)
    return (old_div(R,mu)) * rho * T

def Prad(T,mu):
    '''
    P = a/3 * T**4

    Parameters
    ----------
    a : float
        Radiation constant [erg cm^-3 K^-4].
    T : float
        Temperature [K].

    '''
    return (old_div(radiation_constant,3.)) * T**4

def mimf_ferrario(mi):
    ''' Curvature MiMf from Ferrario etal. 2005MNRAS.361.1131.'''

    mf=-0.00012336*mi**6+0.003160*mi**5-0.02960*mi**4+\
      0.12350*mi**3-0.21550*mi**2+0.19022*mi+0.46575
    return mf

def core_mass_L(MH):
    '''
    Core-mass luminosity relationship from Bloecker (1993).

    Parameters
    ----------
    MH : float
        Core mass in Msun.

    Returns
    -------
    L
        Luminosity in Lsun

    '''
    return 62200*(MH-0.487)

def imf(m):
    '''
    Returns
    -------
    N(M)dM
        for given mass according to Kroupa IMF, vectorization
        available via vimf()

    '''

    m1 = 0.08; m2 = 0.50
    a1 = 0.30; a2 = 1.30; a3 = 2.3
    const2 = m1**-a1 -m1**-a2
    const3 = m2**-a2 -m2**-a3

    if m < 0.08:
        alpha = 0.3
        const = -const2 -const3
    elif m < 0.50:
        alpha = 1.3
        const = -const3
    else:
        alpha = 2.3
        const = 0.0
    # print m,alpha, const, m**-alpha + const
    return m**-alpha + const

vimf = np.vectorize(imf)

def int_imf_dm(m1,m2,m,imf,bywhat='bymass',integral='normal'):
    '''
    Integrate IMF between m1 and m2.

    Parameters
    ----------
    m1 : float
        Min mass
    m2 : float
        Max mass
    m : float
        Mass array
    imf : float
        IMF array
    bywhat : string, optional
        'bymass' integrates the mass that goes into stars of
        that mass interval; or 'bynumber' which integrates the number
        of stars in that mass interval.  The default is 'bymass'.
    integrate : string, optional
        'normal' uses sc.integrate.trapz; 'cum' returns cumulative
        trapezoidal integral.  The default is 'normal'.

    '''

    ind_m = (m >= min(m1,m2)) & (m <= max(m1,m2))
    if integral is 'normal':
        int_func = sc.integrate.trapz
    elif integral is 'cum':
        int_func = sc.integrate.cumtrapz
    else:
        print("Error in int_imf_dm: don't know how to integrate")
        return 0

    if bywhat is 'bymass':
        return int_func(m[ind_m]*imf[ind_m],m[ind_m])
    elif bywhat is 'bynumber':
        return int_func(imf[ind_m],m[ind_m])
    else:
        print("Error in int_imf_dm: don't know by what to integrate")
        return 0

def am_orb(m1,m2,a,e):
    '''
    orbital angular momentum.

    e.g Ge etal2010

    Parameters
    ----------
    m1, m2 : float
        Masses of both stars in Msun.
    A : float
        Separation in Rsun.
    e : float
        Eccentricity

    '''

    a_cm  = a * rsun_cm
    m1_g = m1 * msun_g
    m2_g = m2 * msun_g

    J_orb=np.sqrt(grav_const*a_cm*(old_div((m1_g**2*m2_g**2),(m1_g+m2_g))))*(1-e**2)
    return J_orb

def mass_loss_loon05(L,Teff):
    '''
    mass loss rate van Loon etal (2005).

    Parameters
    ----------
    L : float
        L in L_sun.
    Teff : float
        Teff in K.

    Returns
    -------
    Mdot
        Mdot in Msun/yr

    Notes
    -----
    ref: van Loon etal 2005, A&A 438, 273

    '''

    Mdot = -5.65 + np.log10(old_div(L,10.**4)) -6.3*np.log10(old_div(Teff,3500.))
    return Mdot

def energ_orb(m1,m2,r):
    '''
    Parameters
    ----------
    m1, m2 : float
        M in Msun.
    r : float
        Distance in Rsun.

    Returns
    -------
    Epot
        Epot in erg.

    '''
    epo = -grav_const * m1 * m2 * msun_g**2 / (r * rsun_cm)
    return epo

def period(A,M1,M2):
    """
    calculate binary period from separation.

    Parameters
    ----------
    A : float
        separation A Rsun.
    M1, M2 : float
        M in Msun.

    Returns
    -------
    p
        period in days.

    """

    A *= rsun_cm
    print(A)
    velocity = np.sqrt(grav_const*msun_g*(M1+M2)/A)
    print(old_div(velocity,1.e5))

    p = 2.*np.pi * A / velocity

    p  /= (60*60*24.)
    return p

def escape_velocity(M,R):
    """
    escape velocity.

    Parameters
    ----------
    M : float
        Mass in solar masses.
    R : float
        Radius in solar radiu.

    Returns
    -------
    v_escape
        in km/s.

    """

    ve = np.sqrt(2.*grav_const*M*msun_g/(R*rsun_cm))
    ve = ve*1.e-5
    return ve

def Nasv(macs,T):
    '''
    Returns
    -------
    Na*<sigma v>
        for MACS [mb] at T [K].

    '''

    Na = avogadro_constant
    k  = boltzmann_constant
    vtherm=(2.*k*T/mass_H_atom)**0.5

    s  = macs*1.e-27
    Nasv = s*vtherm*Na
    return Nasv

def macs(nasv,T):
    '''
    Returns
    -------
    MACS
        [mb] at T [K] from Na*<sigma v>.

    '''

    Na = avogadro_constant
    k  = boltzmann_constant
    vtherm=(2.*k*T/mass_H_atom)**0.5

    s      = old_div(nasv,(vtherm*Na))
    macs   = s*1.e27
    return macs

def mu_e(X):
    '''
    mean molecular weight per free electron, assuming full ionisation, and
    approximating mu_i/Z_i ~ 2 for all elements heavier then Helium.

    (Kippenhahn & Weigert, Ch 13.1, Eq. 13.8)

    Parameters
    ----------
    X : float
        Mass fraction of H.

    '''

    try:
        mu_e = old_div(2.,(1.+X))
    except TypeError:
        X=np.array([X])
        mu_e = old_div(2.,(1.+X))

    return mu_e

def mu(X,Z,A):
    '''
    mean molecular weight assuming full ionisation.

    (Kippenhahn & Weigert, Ch 13.1, Eq. 13.6)

    Parameters
    ----------
    X : float
        Mass fraction vector.
    Z : float
        Charge number vector.
    A : float
        Mass number vector.

    '''

    if not isinstance(Z,np.ndarray):
        Z = np.array(Z)
    if not isinstance(A,np.ndarray):
        A = np.array(A)
    if not isinstance(X,np.ndarray):
        X = np.array(X)

    try:
        mu = old_div(1.,sum(X*(1.+Z)/A))
    except TypeError:
        X=np.array([X])
        A=np.array([A])
        Z=np.array([Z])
        mu = old_div(1.,sum(X*(1.+Z)/A))

    return mu

def Trho_idrad(rho,mu):
    '''
    T(rho) that separates P_rad from P_gas dominated regions.

    Kippenhahn & Weigert, Eq. 16.10

    Parameters
    ----------
    rho : float
        Density array [cgs].
    mu : float
        Mean molecular weight.

    '''

    T = 3.2E7 * (old_div(rho,mu))**(old_div(1.,3.))
    return T

def Trho_iddeg(rho,mu,mu_e):
    '''
    T(rho) that separates ideal gas and degenerate pressure dominated regions.

    Kippenhahn & Weigert, Eq. 16.6

    Parameters
    ----------
    rho : float
        Density array [cgs].
    mu : float
        Mean molecular weight.
    mu_e : float
        Mean molecular weight per free electron.

    '''

    T = 1.207E5 * rho**(old_div(2.,3.)) * mu / mu_e**(old_div(5.,3.))
    return T

<|endoftext|>


<commit_before>"""
Main set of commands known by the atxcf bot.
"""

from PriceNetwork import _get_price_network
import PriceNetwork
import settings
from settings import get_settings_option
import memcached_client

import coinmarketcap

import string
import sys
import threading
import time

_mutex = threading.Lock()
class Mutex(object):
    def __init__(self):
        pass
    def __enter__(self):
        _mutex.acquire()
    def __exit__(self, exc_type, exc_val, exc_tb):
        _mutex.release()

class CmdError(PriceNetwork.PriceNetworkError):
    pass

def get_symbols():
    """
    Returns all asset symbols known by the bot.
    """
    with Mutex():
        pn = _get_price_network()
        return sorted(pn.get_symbols())

def _do_get_price(value, trade_pair_str, get_last=False):
    asset_strs = string.split(trade_pair_str,"/",1)
    if len(asset_strs) != 2:
        raise CmdError("Invalid trade pair %s" % trade_pair_str)
    asset_strs = [cur.strip() for cur in asset_strs]

    with Mutex():
        pn = _get_price_network()
        price = pn.get_price(asset_strs[0], asset_strs[1], value, get_last)
        if not price:
            price = float('NaN')

        return price

def get_price(*args, **kwargs):
    """
    Returns price dependings on args:
    - 1 == len(args) -> from/to pair string (aka trade_pair_str)
    - 2 == len(args) -> (value, trade_pair_str)
    - 3 == len(args) -> (value, from_asset, to_asset)

    Supports a keyword arg 'get_last' that flags whether we should
    just return the last price cached or if we should try to get an
    updated value.
    """
    value = 1.0
    trade_pair_str = ""
    if len(args) == 1:
        # treat args as a from/to pair with amount == 1
        value = 1.0
        trade_pair_str = args[0]
    elif len(args) == 2:
        # treat args as a pair of (value, trade_pair_str)
        value = float(args[0])
        trade_pair_str = args[1]
    elif len(args) == 3:
        # treat args as a triple of (value, from_asset, to_asset)
        value = float(args[0])
        from_asset = args[1].strip()
        to_asset = args[2].strip()
        trade_pair_str = "%s/%s" % (from_asset, to_asset)
    else:
        raise CmdError("Invalid argument list for command get_price: %s" % str(args))
    get_last = True
    if "get_last" in kwargs:
        get_last = kwargs["get_last"]
    return _do_get_price(value, trade_pair_str, get_last)

_updater_thread = None
def keep_prices_updated():
    """
    Launches a thread to keep prices in the cache updated.
    """
    global _updater_thread
    if not _updater_thread:
        def updater():
            interval = get_settings_option("price_update_interval", 60)
            last_prices = {}
            while True:
                for mkt in get_markets():
                    last_price = None
                    if mkt in last_prices:
                        last_price = last_prices[mkt]
                    price = get_price(mkt, get_last=False)
                    if last_price != price:
                        print "updater: ", time.time(), mkt, price
                        last_prices[mkt] = price
                time.sleep(interval)
        print "Launching price updater thread"
        _updater_thread = threading.Thread(target=updater)
        _updater_thread.daemon = True
        _updater_thread.start()

def get_markets():
    """
    Returns all markets known by the bot.
    """
    with Mutex():
        pn = _get_price_network()
        return pn.get_markets()

def get_top_coins(top=10):
    """
    Returns the top market cap coinz....
    """
    key = "top_coins_" + str(top)
    if memcached_client.has_key(key):
        return memcached_client.get(key)
    top_coins = [coinmarketcap.short(name) for name in coinmarketcap.top(int(top))]
    memcached_client.set(key, top_coins, expire=3600) # expire every hour
    return top_coins

def is_cmd(cmd):
    """
    Returns whether specified command is valid.
    """
    cmd_obj = getattr(sys.modules[__name__], cmd)
    is_callable = hasattr(cmd_obj, "__call__")
    is_class = isinstance(cmd_obj, type)
    is_public = not cmd.startswith('_')
    return is_callable and is_public and not is_class

def get_commands():
    """
    Returns all bot commands.

    All the functions in this module are considered bot commands
    unless their names begin with a '_' character.
    """
    return [cmd for cmd in dir(sys.modules[__name__]) if is_cmd(cmd)]

def get_help(*args):
    """
    Returns help for the specified command, which is just it's docstring.
    """
    if len(args) > 0:
        cmd_name = args[0]
        if not cmd_name in get_commands():
            raise CmdError("Invalid command %s" % cmd_name)
        cmd_module = sys.modules[__name__]
        return getattr(cmd_module, cmd_name).__doc__
    else:
        return get_help.__doc__

def exit():
    """
    Exits the process.
    """
    sys.exit()

def _run_cmd(*args):
    """
    Helper function to run a command given command args including the command
    as the first argument.
    """
    if len(args) > 0:
        if not is_cmd(args[0]):
            raise CmdError("Invalid command %s" % args[0])
        cmd = getattr(sys.modules[__name__], args[0])
        return cmd(*args[1:])
    else:
        raise CmdError("No command specified")
<commit_msg>Making the default to get the last cached price or not specified by the settings file.<commit_after>"""
Main set of commands known by the atxcf bot.
"""

from PriceNetwork import _get_price_network
import PriceNetwork
import settings
from settings import get_settings_option
import memcached_client

import coinmarketcap

import string
import sys
import threading
import time

_mutex = threading.Lock()
class Mutex(object):
    def __init__(self):
        pass
    def __enter__(self):
        _mutex.acquire()
    def __exit__(self, exc_type, exc_val, exc_tb):
        _mutex.release()

class CmdError(PriceNetwork.PriceNetworkError):
    pass

def get_symbols():
    """
    Returns all asset symbols known by the bot.
    """
    with Mutex():
        pn = _get_price_network()
        return sorted(pn.get_symbols())

def _do_get_price(value, trade_pair_str, get_last=False):
    asset_strs = string.split(trade_pair_str,"/",1)
    if len(asset_strs) != 2:
        raise CmdError("Invalid trade pair %s" % trade_pair_str)
    asset_strs = [cur.strip() for cur in asset_strs]

    with Mutex():
        pn = _get_price_network()
        price = pn.get_price(asset_strs[0], asset_strs[1], value, get_last)
        if not price:
            price = float('NaN')

        return price

def get_price(*args, **kwargs):
    """
    Returns price dependings on args:
    - 1 == len(args) -> from/to pair string (aka trade_pair_str)
    - 2 == len(args) -> (value, trade_pair_str)
    - 3 == len(args) -> (value, from_asset, to_asset)

    Supports a keyword arg 'get_last' that flags whether we should
    just return the last price cached or if we should try to get an
    updated value.
    """
    value = 1.0
    trade_pair_str = ""
    if len(args) == 1:
        # treat args as a from/to pair with amount == 1
        value = 1.0
        trade_pair_str = args[0]
    elif len(args) == 2:
        # treat args as a pair of (value, trade_pair_str)
        value = float(args[0])
        trade_pair_str = args[1]
    elif len(args) == 3:
        # treat args as a triple of (value, from_asset, to_asset)
        value = float(args[0])
        from_asset = args[1].strip()
        to_asset = args[2].strip()
        trade_pair_str = "%s/%s" % (from_asset, to_asset)
    else:
        raise CmdError("Invalid argument list for command get_price: %s" % str(args))
    get_last = get_settings_option("get_last_price", True)
    if "get_last" in kwargs:
        get_last = kwargs["get_last"]
    return _do_get_price(value, trade_pair_str, get_last)

_updater_thread = None
def keep_prices_updated():
    """
    Launches a thread to keep prices in the cache updated.
    """
    global _updater_thread
    if not _updater_thread:
        def updater():
            interval = get_settings_option("price_update_interval", 60)
            last_prices = {}
            while True:
                for mkt in get_markets():
                    last_price = None
                    if mkt in last_prices:
                        last_price = last_prices[mkt]
                    price = get_price(mkt, get_last=False)
                    if last_price != price:
                        print "updater: ", time.time(), mkt, price
                        last_prices[mkt] = price
                time.sleep(interval)
        print "Launching price updater thread"
        _updater_thread = threading.Thread(target=updater)
        _updater_thread.daemon = True
        _updater_thread.start()

def get_markets():
    """
    Returns all markets known by the bot.
    """
    with Mutex():
        pn = _get_price_network()
        return pn.get_markets()

def get_top_coins(top=10):
    """
    Returns the top market cap coinz....
    """
    key = "top_coins_" + str(top)
    if memcached_client.has_key(key):
        return memcached_client.get(key)
    top_coins = [coinmarketcap.short(name) for name in coinmarketcap.top(int(top))]
    memcached_client.set(key, top_coins, expire=3600) # expire every hour
    return top_coins

def is_cmd(cmd):
    """
    Returns whether specified command is valid.
    """
    cmd_obj = getattr(sys.modules[__name__], cmd)
    is_callable = hasattr(cmd_obj, "__call__")
    is_class = isinstance(cmd_obj, type)
    is_public = not cmd.startswith('_')
    return is_callable and is_public and not is_class

def get_commands():
    """
    Returns all bot commands.

    All the functions in this module are considered bot commands
    unless their names begin with a '_' character.
    """
    return [cmd for cmd in dir(sys.modules[__name__]) if is_cmd(cmd)]

def get_help(*args):
    """
    Returns help for the specified command, which is just it's docstring.
    """
    if len(args) > 0:
        cmd_name = args[0]
        if not cmd_name in get_commands():
            raise CmdError("Invalid command %s" % cmd_name)
        cmd_module = sys.modules[__name__]
        return getattr(cmd_module, cmd_name).__doc__
    else:
        return get_help.__doc__

def exit():
    """
    Exits the process.
    """
    sys.exit()

def _run_cmd(*args):
    """
    Helper function to run a command given command args including the command
    as the first argument.
    """
    if len(args) > 0:
        if not is_cmd(args[0]):
            raise CmdError("Invalid command %s" % args[0])
        cmd = getattr(sys.modules[__name__], args[0])
        return cmd(*args[1:])
    else:
        raise CmdError("No command specified")
<|endoftext|>


<commit_before>from threading import Event, Lock
from collections import defaultdict
import logging
import io

import numpy as np
from scipy.io import wavfile
from pyaudio import PyAudio, paContinue, paComplete

class AudioSink:
    """ Asynchronous wav player based on portaudio

    Only works for 16 bits low-endian mono wav. The rate defaults to 16000.
    """

    player = None
    logger = logging.getLogger('AWavPlayer')

    # List of sound data to be mixed and played.
    _queue = None
    # The queue is used to communicate between two threads, so we need a lock.
    _queue_lock = None
    # Synchronisation primitive to resume the worker's activity
    _worker_wakeup = None

    def __init__(self, rate=16000):
        self.player = PyAudio()
        self._queue = defaultdict(list)
        self._queue_lock = Lock()
        self._worker_wakeup = Event()

        # This will patiently wait for sounds to be sent via #add.
        self.player.open(
            format=8,
            channels=1,
            rate=rate,
            output=True,
            stream_callback=self._worker
        )

    def add(self, sound, owner=None):
        """ Send sound data to be mixed with currently playing sounds """
        _, data = wavfile.read(io.BytesIO(sound))
        with self._queue_lock:
            self._queue[owner].append(data)
            # In case the worker was waiting for sounds.
            self._worker_wakeup.set()

    def remove(self, owner):
        """ Stop playing all sounds in a category as soon as possible """
        with self._queue_lock:
            del(self._queue[owner])

    def _worker(self, in_data, frame_count, time_info, status):
        """ Function called by PyAudio each time it can play more sound """
        # The queue is also used by #add, and since this callback runs into
        # a different thread, we need to lock it while we process it.
        self._queue_lock.acquire()

        if not self._queue:
            # Wait until #add has put a sound into the queue.
            self.logger.debug('Worker\'s queue is empty')
            self._worker_wakeup.clear()
            # Release the lock so #add can update it.
            self._queue_lock.release()
            self._worker_wakeup.wait()
            # If we were woken up and the queue is still empty it means
            # we've been woken up by #close.
            if not self._queue:
                return b'', paComplete
            self._queue_lock.acquire()

        # Chunks of each sound of the size requested by portaudio
        # in frame_count. They'll be mixed together later.
        chunks = list()
        # Since we'll iterate on it, we'll avoid mutating the queue
        # and instead create a new one which will become the new queue
        # once the iteration has finished.
        new_queue = defaultdict(list)
        for owner, sounds in self._queue.items():
            for sound in sounds:
                length = len(sound)
                if length <= frame_count:
                    # We reached the end of a sound, so pad it if needed, as
                    # portaudio would otherwise stop working if it isn't the
                    # right size.
                    diff = frame_count - length
                    sound = np.pad(sound, (0, diff), 'constant')
                else:
                    # We didn't reach the end of that sound, so put what's left
                    # into the new queue.
                    new_queue[owner].append(sound[frame_count:])

                chunks.append(sound[:frame_count])

        self._queue = new_queue

        self._queue_lock.release()

        # Put all chunks into a single array and lower their volume so their
        # sum doesn't saturate the output.
        stack = np.stack(chunks, axis=0) // len(chunks)
        # Don't forget to force the int16 type else #tobytes() won't output
        # bytes with the correct format and the audio output will be garbled.
        data = np.sum(stack, axis=0, dtype=np.int16)

        return data.tobytes(), paContinue

    def close(self):
        """ Graceful shutdown """
        self._worker_wakeup.set()
        self.player.terminate()
<commit_msg>Update AudioSink's logger name.<commit_after>from threading import Event, Lock
from collections import defaultdict
import logging
import io

import numpy as np
from scipy.io import wavfile
from pyaudio import PyAudio, paContinue, paComplete

class AudioSink:
    """ Asynchronous wav player based on portaudio

    Only works for 16 bits low-endian mono wav. The rate defaults to 16000.
    """

    player = None
    logger = logging.getLogger('AudioSink')

    # List of sound data to be mixed and played.
    _queue = None
    # The queue is used to communicate between two threads, so we need a lock.
    _queue_lock = None
    # Synchronisation primitive to resume the worker's activity
    _worker_wakeup = None

    def __init__(self, rate=16000):
        self.player = PyAudio()
        self._queue = defaultdict(list)
        self._queue_lock = Lock()
        self._worker_wakeup = Event()

        # This will patiently wait for sounds to be sent via #add.
        self.player.open(
            format=8,
            channels=1,
            rate=rate,
            output=True,
            stream_callback=self._worker
        )

    def add(self, sound, owner=None):
        """ Send sound data to be mixed with currently playing sounds """
        _, data = wavfile.read(io.BytesIO(sound))
        with self._queue_lock:
            self._queue[owner].append(data)
            # In case the worker was waiting for sounds.
            self._worker_wakeup.set()

    def remove(self, owner):
        """ Stop playing all sounds in a category as soon as possible """
        with self._queue_lock:
            del(self._queue[owner])

    def _worker(self, in_data, frame_count, time_info, status):
        """ Function called by PyAudio each time it can play more sound """
        # The queue is also used by #add, and since this callback runs into
        # a different thread, we need to lock it while we process it.
        self._queue_lock.acquire()

        if not self._queue:
            # Wait until #add has put a sound into the queue.
            self.logger.debug('Worker\'s queue is empty')
            self._worker_wakeup.clear()
            # Release the lock so #add can update it.
            self._queue_lock.release()
            self._worker_wakeup.wait()
            # If we were woken up and the queue is still empty it means
            # we've been woken up by #close.
            if not self._queue:
                return b'', paComplete
            self._queue_lock.acquire()

        # Chunks of each sound of the size requested by portaudio
        # in frame_count. They'll be mixed together later.
        chunks = list()
        # Since we'll iterate on it, we'll avoid mutating the queue
        # and instead create a new one which will become the new queue
        # once the iteration has finished.
        new_queue = defaultdict(list)
        for owner, sounds in self._queue.items():
            for sound in sounds:
                length = len(sound)
                if length <= frame_count:
                    # We reached the end of a sound, so pad it if needed, as
                    # portaudio would otherwise stop working if it isn't the
                    # right size.
                    diff = frame_count - length
                    sound = np.pad(sound, (0, diff), 'constant')
                else:
                    # We didn't reach the end of that sound, so put what's left
                    # into the new queue.
                    new_queue[owner].append(sound[frame_count:])

                chunks.append(sound[:frame_count])

        self._queue = new_queue

        self._queue_lock.release()

        # Put all chunks into a single array and lower their volume so their
        # sum doesn't saturate the output.
        stack = np.stack(chunks, axis=0) // len(chunks)
        # Don't forget to force the int16 type else #tobytes() won't output
        # bytes with the correct format and the audio output will be garbled.
        data = np.sum(stack, axis=0, dtype=np.int16)

        return data.tobytes(), paContinue

    def close(self):
        """ Graceful shutdown """
        self._worker_wakeup.set()
        self.player.terminate()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import print_function

# Standard library imports
import sys
import os

# IronPython imports
from System.Diagnostics import *
from System.IO import *

# Deadline imports
from Deadline.Events import *
from Deadline.Scripting import *
from Deadline.Slaves import *

# Setup system path
this_path = RepositoryUtils.GetEventPluginDirectory("autocpenv")
packages_path = os.path.join(this_path, 'packages')
if packages_path not in sys.path:
    sys.path.insert(1, packages_path)

import cpenv
from cpenv import mappings

def GetDeadlineEventListener():
    return AutoCpenv()

def CleanupDeadlineEventListener(eventListener):
    eventListener.Cleanup()

class AutoCpenv(DeadlineEventListener):
    '''
    Listen for OnSlaveStartingJob events then activate the configured cpenv
    environment and modules.
    '''

    def __init__(self):
        self.OnJobSubmittedCallback += self.OnJobSubmitted
        self._log_prefix = ''

    def Cleanup(self):
        del self.OnJobSubmittedCallback

    def log(self, message):
        '''Wraps LogInfo to add a prefix to all log messages.'''

        self.LogInfo('{}{}'.format(self._log_prefix, message))

    def OnJobSubmitted(self, job):
        '''
        Responsible for resolving modules and setting a Job's environment
        variables.
        '''

        job_plugin = job.JobPlugin
        plugin_mapping = self.GetConfigEntry('plugin_mapping')

        self.log('Checking job ExtraInfo for requirements...')
        requirements = job.GetJobExtraInfoKeyValue('cpenv_requirements')
        if requirements:
            requirements = requirements.split()
            self.log('Found {} requirements...'.format(len(requirements)))
            for requirement in requirements:
                self.log('  ' + requirement)
            return

        # read config and setup cpenv
        configure_autocpenv(self.log)

        self.log('Attempting to find requirements...')
        requirements = return_first_result(
            (self.resolve_from_job_extra_info, (job,)),
            (self.resolve_from_job_environment, (job,)),
            (self.resolve_from_environment, (job,)),
            (self.resolve_from_job_scenefile, (job,)),
            (self.resolve_from_job_plugin, (plugin_mapping, job_plugin)),
        )

        forced_requirements = self.resolve_from_job_plugin(self.GetConfigEntry('forced_plugin_mapping'), job_plugin)
        if forced_requirements:
            self.log("Forced requirements (these will be appended to the job requirements):")
            self.log(" ".join(forced_requirements))
            requirements.extend(forced_requirements)
        else:
            self.log("No forced requirements found in the Forced Plugin Mapping settings...")

        if not requirements:
            self.log('Error! Failed 4 attempts to find requirements...')
            return
        else:
            self.log('Found {} requirements...'.format(len(requirements)))
            for requirement in requirements:
                self.log('  ' + requirement)

        self.log('Setting JobExtranInfo cpenv_requirements')
        job.SetJobExtraInfoKeyValue(
            'cpenv_requirements',
            ' '.join(requirements)
        )

        self.log('Saving Job.')
        RepositoryUtils.SaveJob(job)

    def resolve_from_job_extra_info(self, job):
        '''Checks to see if the job was submitted with cpenv_requirements.'''

        self.log('Checking job extra info for module requirements...')

        requirements = job.GetJobExtraInfoKeyValue('cpenv_requirements')
        if not requirements:
            self.log('Job has no cpenv requirements...')
            return

        return requirements.split()

    def resolve_from_job_environment(self, job):
        '''Resolve modules from job CPENV_ACTIVE_MODULES variable.'''

        self.log('Checking job environment for module requirements...')

        requirements = job.GetJobEnvironmentKeyValue('CPENV_ACTIVE_MODULES')
        if not requirements:
            self.log('  CPENV_ACTIVE_MODULES not set...')
            return

        return split_path(requirements)

    def resolve_from_environment(self, job):
        '''Checks the environment of the local machine that's submitting the
        job for the CPENV_ACTIVE_MODULES variable.'''

        self.log('Checking local environment for module requirements...')

        requirements = os.getenv('CPENV_ACTIVE_MODULES')
        if not requirements:
            self.log('  CPENV_ACTIVE_MODULES not set...')
            return

        return split_path(requirements)

    def resolve_from_job_scenefile(self, job):
        '''Attempt to resolve modules from the jobs scene_file.

        Walks up the path until a .cpenv file is found.
        '''

        self.log('Checking job scene file for module requirements...')

        scene_file = job.GetJobPluginInfoKeyValue('SceneFile')
        if not scene_file:
            self.log('  SceneFile not set...')
            return

        requirement = os.path.dirname(scene_file)
        try:
            resolved = cpenv.resolve([requirement])
            return [m.qual_name for m in resolved]
        except cpenv.ResolveError:
            pass

    def resolve_from_job_plugin(self, plugin_mapping, job_plugin):
        '''Attempt to resolve modules using the Autocpenv event plugins
        configured plugin mapping.
        '''

        self.log('Checking job {} for module requirements...'.format(plugin_mapping))

        plugin_mapping = plugin_mapping_to_dict(plugin_mapping)
        requirements = plugin_mapping.get(job_plugin, None)
        return requirements

class EventLogReporter(cpenv.Reporter):

    def __init__(self, log):
        super(EventLogReporter, self).__init__()
        self.log = log

    def start_resolve(self, requirements):
        self.log('- Resolving requirements...')

    def resolve_requirement(self, requirement, module_spec):
        self.log('  %s - %s' % (module_spec.real_name, module_spec.path))

    def end_resolve(self, resolved, unresolved):
        if unresolved:
            self.log('  Failed to resolve %s' % ', '.join(unresolved))

    def start_localize(self, module_specs):
        self.log('- Localizing modules...')

    def start_progress(self, label, max_size, data):

        if 'download' in label.lower():
            spec = data['module_spec']
            self.log(
                '  Downloading %s from %s...' %
                (spec.qual_name, spec.repo.name)
            )
        elif 'upload' in label.lower():
            module = data['module']
            to_repo = data['to_repo']
            self.log(
                '  Uploading %s to %s...' %
                (module.qual_name, to_repo.name)
            )
        else:
            self.log('  ' + label)

def return_first_result(*funcs):
    for func, args in funcs:
        result = func(*args)
        if result:
            return result

def split_path(path, pathsep=os.pathsep):
    return [part for part in path.split(pathsep) if part]

def plugin_mapping_to_dict(plugin_mapping):
    '''Convert the plugin_mapping config entry to a dictionary'''

    d = {}

    for line in plugin_mapping.split(';'):
        plugin, env_paths = line.split('=')
        d.setdefault(plugin, env_paths.split())

    return d

def configure_autocpenv(log_method=None):
    '''Read config from EventPlugin and setup cpenv.'''

    log_method = log_method or print
    log_method('Configuring cpenv...')

    autocpenv = RepositoryUtils.GetEventPluginDirectory('autocpenv')
    if autocpenv not in sys.path:
        sys.path.insert(1, autocpenv)
        sys.path.insert(1, autocpenv + '/packages')

    config = RepositoryUtils.GetEventPluginConfig('autocpenv')
    home_path = config.GetConfigEntry('cpenv_home')
    if home_path:
        cpenv.set_home_path(home_path)

    # Initialize ShotgunRepo for requirement lookups
    if config.GetConfigEntry('ShotgunRepo_enable'):
        repo = cpenv.ShotgunRepo(
            name='autocpenv_shotgun',
            base_url=config.GetConfigEntry('ShotgunRepo_base_url'),
            script_name=config.GetConfigEntry('ShotgunRepo_script_name'),
            api_key=config.GetConfigEntry('ShotgunRepo_api_key'),
        )
        log_method('Adding ShotgunRepo: ' + repo.base_url)
        cpenv.add_repo(repo)

    # Setup reporting
    cpenv.set_reporter(EventLogReporter(log_method))

def GlobalJobPreLoad(plugin):
    '''Execute this method in your GlobalJobPreLoad.py file.

    See GlobalJobPreLoad.py for reference. Or if you are not using a
    GlobalJobPreLoad script simply copy the file to <repo>/custom/plugins.
    '''

    # Get job cpenv requirements
    job = plugin.GetJob()
    requirements = job.GetJobExtraInfoKeyValue('cpenv_requirements')
    if not requirements:
        plugin.LogInfo('Job has no cpenv requirements...')
        return

    # Read config from autocpenv EventPlugin
    configure_autocpenv(plugin.LogInfo)

    # Use cpenv to resolve requirements stored in ExtraInfo
    resolved = cpenv.resolve(requirements.split())
    localizer = cpenv.Localizer(cpenv.get_repo('home'))
    localized = localizer.localize(resolved, overwrite=False)
    activator = cpenv.Activator(localizer)
    env = activator.combine_modules(localized)

    # Get existing environment
    plugin.LogInfo('Collecting process and job environment variables...')
    proc_env = {}
    for key in env.keys():
        proc_env_var = plugin.GetProcessEnvironmentVariable(key)
        if proc_env_var:
            proc_env[key] = proc_env_var

    plugin.LogInfo('Merging environment variables...')
    new_job_env = mappings.dict_to_env(mappings.join_dicts(
        proc_env,
        env,
    ))
    new_job_env = mappings.expand_envvars(new_job_env)

    plugin.LogInfo('Setting process environment variables...')
    for k, v in new_job_env.items():
        plugin.SetProcessEnvironmentVariable(k, v)
<commit_msg>Refactoring changes for forced requirements and how it merges with requirements, fixed a bug if there are no requirements.<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function

# Standard library imports
import sys
import os

# IronPython imports
from System.Diagnostics import *
from System.IO import *

# Deadline imports
from Deadline.Events import *
from Deadline.Scripting import *
from Deadline.Slaves import *

# Setup system path
this_path = RepositoryUtils.GetEventPluginDirectory("autocpenv")
packages_path = os.path.join(this_path, 'packages')
if packages_path not in sys.path:
    sys.path.insert(1, packages_path)

import cpenv
from cpenv import mappings

def GetDeadlineEventListener():
    return AutoCpenv()

def CleanupDeadlineEventListener(eventListener):
    eventListener.Cleanup()

class AutoCpenv(DeadlineEventListener):
    '''
    Listen for OnSlaveStartingJob events then activate the configured cpenv
    environment and modules.
    '''

    def __init__(self):
        self.OnJobSubmittedCallback += self.OnJobSubmitted
        self._log_prefix = ''

    def Cleanup(self):
        del self.OnJobSubmittedCallback

    def log(self, message):
        '''Wraps LogInfo to add a prefix to all log messages.'''

        self.LogInfo('{}{}'.format(self._log_prefix, message))

    def OnJobSubmitted(self, job):
        '''
        Responsible for resolving modules and setting a Job's environment
        variables.
        '''

        job_plugin = job.JobPlugin
        plugin_mapping = self.GetConfigEntry('plugin_mapping')

        self.log('Checking job ExtraInfo for requirements...')
        requirements = job.GetJobExtraInfoKeyValue('cpenv_requirements')
        if requirements:
            requirements = requirements.split()
            self.log('Found {} requirements...'.format(len(requirements)))
            for requirement in requirements:
                self.log('  ' + requirement)
            return

        # read config and setup cpenv
        configure_autocpenv(self.log)

        self.log('Attempting to find requirements...')
        requirements = return_first_result(
            (self.resolve_from_job_extra_info, (job,)),
            (self.resolve_from_job_environment, (job,)),
            (self.resolve_from_environment, (job,)),
            (self.resolve_from_job_scenefile, (job,)),
            (self.resolve_from_job_plugin, (plugin_mapping, job_plugin)),
        )
        if not requirements:
            self.log('Did not find any job requirements...')

        self.log('Attempting to find forced requirements...')
        forced_requirements = self.resolve_from_forced_plugin_mappings(job_plugin)
        if not forced_requirements:
            self.log('Did not find any forced requirements...')

        # if we have both requirements and forced requirements we merge the lists together and remove duplicates
        if forced_requirements and requirements:
            requirements.extend(forced_requirements)
            # remove duplicates
            requirements = list(dict.fromkeys(requirements))
        else:
            # we don't have requirements, use the forced requirements instead
            requirements = forced_requirements

        if requirements:
            # print the merged requirements we found...
            self.log('Found {} requirements...'.format(len(requirements)))
            for requirement in requirements:
                self.log('  ' + requirement)
        else:
            # we dont have any job requirements or forced requirements, do nothing...
            self.log('Did not find any requirements for this job...')
            return

        # set the job cpenv requirements on the jobExtraInfo key
        self.log('Setting JobExtranInfo cpenv_requirements')
        job.SetJobExtraInfoKeyValue(
            'cpenv_requirements',
            ' '.join(requirements)
        )

        self.log('Saving Job.')
        RepositoryUtils.SaveJob(job)

    def resolve_from_job_extra_info(self, job):
        '''Checks to see if the job was submitted with cpenv_requirements.'''

        self.log('Checking job extra info for module requirements...')

        requirements = job.GetJobExtraInfoKeyValue('cpenv_requirements')
        if not requirements:
            self.log('Job has no cpenv requirements...')
            return

        return requirements.split()

    def resolve_from_job_environment(self, job):
        '''Resolve modules from job CPENV_ACTIVE_MODULES variable.'''

        self.log('Checking job environment for module requirements...')

        requirements = job.GetJobEnvironmentKeyValue('CPENV_ACTIVE_MODULES')
        if not requirements:
            self.log('  CPENV_ACTIVE_MODULES not set...')
            return

        return split_path(requirements)

    def resolve_from_environment(self, job):
        '''Checks the environment of the local machine that's submitting the
        job for the CPENV_ACTIVE_MODULES variable.'''

        self.log('Checking local environment for module requirements...')

        requirements = os.getenv('CPENV_ACTIVE_MODULES')
        if not requirements:
            self.log('  CPENV_ACTIVE_MODULES not set...')
            return

        return split_path(requirements)

    def resolve_from_job_scenefile(self, job):
        '''Attempt to resolve modules from the jobs scene_file.

        Walks up the path until a .cpenv file is found.
        '''

        self.log('Checking job scene file for module requirements...')

        scene_file = job.GetJobPluginInfoKeyValue('SceneFile')
        if not scene_file:
            self.log('  SceneFile not set...')
            return

        requirement = os.path.dirname(scene_file)
        try:
            resolved = cpenv.resolve([requirement])
            return [m.qual_name for m in resolved]
        except cpenv.ResolveError:
            pass

    def resolve_from_job_plugin(self, plugin_mapping, job_plugin):
        '''Attempt to resolve modules using the Autocpenv event plugins
        configured plugin mapping.
        '''

        self.log('Checking job {} for module requirements...'.format(plugin_mapping))
        if not plugin_mapping:
            self.log('Job plugin has no module requirements...')
            return None
        plugin_mapping = plugin_mapping_to_dict(plugin_mapping)
        requirements = plugin_mapping.get(job_plugin, None)
        return requirements

    def resolve_from_forced_plugin_mappings(self, job_plugin):
        '''Attempt to resolve modules from the autocpenv forced plugin mappings setting'''
        forced_requirements = self.resolve_from_job_plugin(self.GetConfigEntry('forced_plugin_mapping'), job_plugin)
        if forced_requirements:
            self.log("Forced requirements (these will be appended to the job requirements):")
            self.log(" ".join(forced_requirements))
            return forced_requirements
        else:
            self.log("No forced requirements found in the Forced Plugin Mapping settings...")
            return None

class EventLogReporter(cpenv.Reporter):

    def __init__(self, log):
        super(EventLogReporter, self).__init__()
        self.log = log

    def start_resolve(self, requirements):
        self.log('- Resolving requirements...')

    def resolve_requirement(self, requirement, module_spec):
        self.log('  %s - %s' % (module_spec.real_name, module_spec.path))

    def end_resolve(self, resolved, unresolved):
        if unresolved:
            self.log('  Failed to resolve %s' % ', '.join(unresolved))

    def start_localize(self, module_specs):
        self.log('- Localizing modules...')

    def start_progress(self, label, max_size, data):

        if 'download' in label.lower():
            spec = data['module_spec']
            self.log(
                '  Downloading %s from %s...' %
                (spec.qual_name, spec.repo.name)
            )
        elif 'upload' in label.lower():
            module = data['module']
            to_repo = data['to_repo']
            self.log(
                '  Uploading %s to %s...' %
                (module.qual_name, to_repo.name)
            )
        else:
            self.log('  ' + label)

def return_first_result(*funcs):
    for func, args in funcs:
        result = func(*args)
        if result:
            return result

def split_path(path, pathsep=os.pathsep):
    return [part for part in path.split(pathsep) if part]

def plugin_mapping_to_dict(plugin_mapping):
    '''Convert the plugin_mapping config entry to a dictionary'''

    d = {}

    for line in plugin_mapping.split(';'):
        plugin, env_paths = line.split('=')
        d.setdefault(plugin, env_paths.split())

    return d

def configure_autocpenv(log_method=None):
    '''Read config from EventPlugin and setup cpenv.'''

    log_method = log_method or print
    log_method('Configuring cpenv...')

    autocpenv = RepositoryUtils.GetEventPluginDirectory('autocpenv')
    if autocpenv not in sys.path:
        sys.path.insert(1, autocpenv)
        sys.path.insert(1, autocpenv + '/packages')

    config = RepositoryUtils.GetEventPluginConfig('autocpenv')
    home_path = config.GetConfigEntry('cpenv_home')
    if home_path:
        cpenv.set_home_path(home_path)

    # Initialize ShotgunRepo for requirement lookups
    if config.GetConfigEntry('ShotgunRepo_enable'):
        repo = cpenv.ShotgunRepo(
            name='autocpenv_shotgun',
            base_url=config.GetConfigEntry('ShotgunRepo_base_url'),
            script_name=config.GetConfigEntry('ShotgunRepo_script_name'),
            api_key=config.GetConfigEntry('ShotgunRepo_api_key'),
        )
        log_method('Adding ShotgunRepo: ' + repo.base_url)
        cpenv.add_repo(repo)

    # Setup reporting
    cpenv.set_reporter(EventLogReporter(log_method))

def GlobalJobPreLoad(plugin):
    '''Execute this method in your GlobalJobPreLoad.py file.

    See GlobalJobPreLoad.py for reference. Or if you are not using a
    GlobalJobPreLoad script simply copy the file to <repo>/custom/plugins.
    '''

    # Get job cpenv requirements
    job = plugin.GetJob()
    requirements = job.GetJobExtraInfoKeyValue('cpenv_requirements')
    if not requirements:
        plugin.LogInfo('Job has no cpenv requirements...')
        #return

    # Read config from autocpenv EventPlugin
    configure_autocpenv(plugin.LogInfo)

    # Use cpenv to resolve requirements stored in ExtraInfo
    resolved = cpenv.resolve(requirements.split())
    localizer = cpenv.Localizer(cpenv.get_repo('home'))
    localized = localizer.localize(resolved, overwrite=False)
    activator = cpenv.Activator(localizer)
    env = activator.combine_modules(localized)

    # Get existing environment
    plugin.LogInfo('Collecting process and job environment variables...')
    proc_env = {}
    for key in env.keys():
        proc_env_var = plugin.GetProcessEnvironmentVariable(key)
        if proc_env_var:
            proc_env[key] = proc_env_var

    plugin.LogInfo('Merging environment variables...')
    new_job_env = mappings.dict_to_env(mappings.join_dicts(
        proc_env,
        env,
    ))
    new_job_env = mappings.expand_envvars(new_job_env)

    plugin.LogInfo('Setting process environment variables...')
    for k, v in new_job_env.items():
        plugin.SetProcessEnvironmentVariable(k, v)
<|endoftext|>


<commit_before>#!/usr/bin/env python
import requests, json, sys, time
from pprint import pprint
import autofocus_config

AF_APIKEY = autofocus_config.AF_APIKEY

# Useful information:
#
# * We're not doing any input validation in the client itself. We pass
#   the data to the API, and rely on 4XX errors to communicate invalid
#   requests. At the time this was written, the API DOES NOT validate search
#   values. So you can offer invalid IPs, such as 592.99.1.1 and it will
#   not balk. The result set will be empty, naturally.

_base_url = "https://autofocus.paloaltonetworks.com/api/v0.9"
_headers = {"Content-Type" : "application/json"}

class NotLoaded(object):
    pass

class AFClientError(Exception):
    def __init__(self, message, resp = None):
        super(AFClientError, self).__init__(self, message)
        self.resp = resp
        self.message = message

class AFServerError(Exception):
    def __init__(self, message, resp):
        super(AFServerError, self).__init__(self, message)
        self.resp = resp
        self.message = message

class AutoFocusAPI(object):

    page_size = 50
    search_operator = "all"

    def __init__(self, **kwargs):
        for k,v in kwargs.items():
            setattr(self, k, v)

    def __repr__(self):
        return self.__dict__.__str__()

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

    @classmethod
    def _api_request(cls, path, post_data = {}, params = {}):

        if not AF_APIKEY:
            raise Exception("AF_APIKEY is not set. Library requires an APIKEY to be set")

        post_data["apiKey"] = AF_APIKEY

        resp = requests.post(_base_url + path, params = params, headers=_headers, data=json.dumps(post_data))

        if resp.status_code >= 400 and resp.status_code < 500:
            raise AFClientError(resp._content, resp)

        if resp.status_code >= 500 and resp.status_code < 600:
            raise AFServerError(resp._content, resp)

        return resp

    @classmethod
    def _api_search_request(cls, path, post_data):

        post_data["size"]   = cls.page_size
        post_data["from"]   = 0
        post_data["sort"]   = {
            "create_date": {
                "order": "desc"
            }
        }
        post_data['scope'] = "Global"
        total_results = 0

        while True:

            init_query_time = time.time()
            init_query_resp = cls._api_request(path, post_data = post_data)
            init_query_data = init_query_resp.json()
            post_data['from'] += post_data['size']
            af_cookie = init_query_data['af_cookie']

            resp = {}
            prev_resp = {}
            i = 0
            while True:

                i += 1

                request_url = "/" + path.split("/")[1] + "/results/" + af_cookie

                # Try our request. Check for AF Cookie going away, which is weird
                # Catch it and add more context. Then throw it up again
                try:
                    resp = cls._api_request(request_url).json()
                except AFClientError as e:
                    if "AF Cookie Not Found" in e.message:
                        raise AFClientError("Auto Focus Cookie has gone away after %d queries taking %f seconds" \
                                        % (i, time.time() - init_query_time), e.resp)
                    else:
                        raise e

                # If we've gotten our bucket size worth of data, or the query has complete
                if len(resp.get('hits', [])) == post_data['size'] \
                        or resp.get('af_complete_percentage', 100) == 100:
                    break

                prev_resp = resp

                continue

            if not resp.get('hits', None):
                raise StopIteration()

            hit_length = len(resp['hits'])

            yield resp

    @classmethod
    def _api_search(cls, *args, **kwargs):

        args = list(args)

        # Classes that inherit from AutoFocusAPI need to pass a search path as the
        # first arg to the protected _api_search method
        path = args.pop(0)

        if len(args) == 1:
            post_data = {
                "query" : json.loads(args[0])
            }
        else:
            # Just do an or here, we'll do validation below
            if "field" in kwargs or "value" in kwargs:
                args.append(kwargs)

            post_data = {
                "query": {
                    "operator": cls.search_operator,
                    "children": []
                }
            }

            for kwarg in args:

                # Check and make sure field and value are passed to search - req'd
                for arg in ('field', 'value'):
                    if arg not in kwarg:
                        raise Exception

                # Build the searching paramaters to be passed to the _api_request method
                # _api_request will add the additional data needed to form a valid request
                post_data['query']['children'].append({
                    "field": kwarg['field'],
                    "operator": kwarg.get('operator', "is"),
                    "value": kwarg['value']
                })

        for res in cls._api_search_request(path, post_data = post_data):
            for hit in res['hits']:
                yield hit

class AFTag(AutoFocusAPI):

    def __init__(self, **kwargs):

        self.comments = kwargs.get("comments", NotLoaded())
        self.refs = kwargs.get("refs", NotLoaded())
        self.review = kwargs.get("review", NotLoaded())
        self.support_id = kwargs.get("support_id", NotLoaded())

        for k,v in kwargs.items():
            setattr(self, k, v)

    def __getattribute__(self, attr):

        value = object.__getattribute__(self, attr)

        # Not offered in the list controller, have to call get to lazyload:
        #      comments, refs, review, support_id
        if attr in ('comments', 'refs', 'review', 'support_id') and type(value) is NotLoaded:

            # Reloading the data via the get method
            self = AFTag.get(self.public_tag_name)
            value = object.__getattribute__(self, attr)

        return value

    @classmethod
    def list(cls, *args, **kwargs):

        kwargs['scope'] = kwargs.get("scope", "Visible")
        kwargs['sortBy'] = kwargs.get("sortBy", "name")
        kwargs['order'] = kwargs.get("order", "asc")
        kwargs['pageSize'] = kwargs.get("pageSize", 1000)
        kwargs['pageNum'] = kwargs.get("pageNum", 0)

        resp = cls._api_request("/tags/", params = kwargs).json()
        results = []

        for tag in resp['tags']:
            results.append(AFTag(**tag))

        return results

    @classmethod
    def get(cls, tag_name):

        resp = cls._api_request("/tag/" + tag_name).json()

        return AFTag(**resp['tag'])

class AFSession(AutoFocusAPI):

    @classmethod
    def search(cls, *args, **kwargs):

        for res in cls._api_search("/sessions/search", *args, **kwargs):
            yield AFSession(**res['_source'])

class AFSample(AutoFocusAPI):

    def get_analyses(self, sections = ["file"], platforms = ["win7", "winxp"]):

        resp = self.__class__._api_request("/sample/" + self.sha256 + "/analysis", \
                    post_data = { "sections" : sections, "platforms" : platforms }).json()

        return resp

    @classmethod
    def search(cls, *args, **kwargs):

        for res in cls._api_search("/samples/search", *args, **kwargs):
            yield AFSample(**res['_source'])

if __name__ == "__main__":

    i = 0
    for sample in AFSample.search(field = "sample.malware", value = "1", operator = "is"):
        i += 1
        pprint(sample.__dict__)
        pprint(sample.get_analyses())

    print "%d results" % (i,)
<commit_msg>Preserving response object and renaming previously use "resp" to "resp_data", which more accurately represents the value in the variable.<commit_after>#!/usr/bin/env python
import requests, json, sys, time
from pprint import pprint
import autofocus_config

AF_APIKEY = autofocus_config.AF_APIKEY

# Useful information:
#
# * We're not doing any input validation in the client itself. We pass
#   the data to the API, and rely on 4XX errors to communicate invalid
#   requests. At the time this was written, the API DOES NOT validate search
#   values. So you can offer invalid IPs, such as 592.99.1.1 and it will
#   not balk. The result set will be empty, naturally.

_base_url = "https://autofocus.paloaltonetworks.com/api/v0.9"
_headers = {"Content-Type" : "application/json"}

class NotLoaded(object):
    pass

class AFClientError(Exception):
    def __init__(self, message, resp = None):
        super(AFClientError, self).__init__(self, message)
        self.resp = resp
        self.message = message

class AFServerError(Exception):
    def __init__(self, message, resp):
        super(AFServerError, self).__init__(self, message)
        self.resp = resp
        self.message = message

class AutoFocusAPI(object):

    page_size = 50
    search_operator = "all"

    def __init__(self, **kwargs):
        for k,v in kwargs.items():
            setattr(self, k, v)

    def __repr__(self):
        return self.__dict__.__str__()

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

    @classmethod
    def _api_request(cls, path, post_data = {}, params = {}):

        if not AF_APIKEY:
            raise Exception("AF_APIKEY is not set. Library requires an APIKEY to be set")

        post_data["apiKey"] = AF_APIKEY

        resp = requests.post(_base_url + path, params = params, headers=_headers, data=json.dumps(post_data))

        if resp.status_code >= 400 and resp.status_code < 500:
            raise AFClientError(resp._content, resp)

        if resp.status_code >= 500 and resp.status_code < 600:
            raise AFServerError(resp._content, resp)

        return resp

    @classmethod
    def _api_search_request(cls, path, post_data):

        post_data["size"]   = cls.page_size
        post_data["from"]   = 0
        post_data["sort"]   = {
            "create_date": {
                "order": "desc"
            }
        }
        post_data['scope'] = "Global"
        total_results = 0

        while True:

            init_query_time = time.time()
            init_query_resp = cls._api_request(path, post_data = post_data)
            init_query_data = init_query_resp.json()
            post_data['from'] += post_data['size']
            af_cookie = init_query_data['af_cookie']

            resp_data = {}
            prev_resp_data = {}
            i = 0
            while True:

                i += 1

                request_url = "/" + path.split("/")[1] + "/results/" + af_cookie

                # Try our request. Check for AF Cookie going away, which is weird
                # Catch it and add more context. Then throw it up again
                try:
                    resp = cls._api_request(request_url)
                    resp_data = resp.json()
                except AFClientError as e:
                    if "AF Cookie Not Found" in e.message:
                        raise AFClientError("Auto Focus Cookie has gone away after %d queries taking %f seconds. Server said percent complete was at %f, last query." \
                                        % (i, time.time() - init_query_time, prev_resp_data['af_complete_percentage']), e.resp)
                    else:
                        raise e

                # If we've gotten our bucket size worth of data, or the query has complete
                if len(resp_data.get('hits', [])) == post_data['size'] \
                        or resp_data.get('af_complete_percentage', 100) == 100:
                    break

                prev_resp_data = resp_data

                continue

            if not resp_data.get('hits', None):
                raise StopIteration()

            hit_length = len(resp_data['hits'])

            yield resp_data

    @classmethod
    def _api_search(cls, *args, **kwargs):

        args = list(args)

        # Classes that inherit from AutoFocusAPI need to pass a search path as the
        # first arg to the protected _api_search method
        path = args.pop(0)

        if len(args) == 1:
            post_data = {
                "query" : json.loads(args[0])
            }
        else:
            # Just do an or here, we'll do validation below
            if "field" in kwargs or "value" in kwargs:
                args.append(kwargs)

            post_data = {
                "query": {
                    "operator": cls.search_operator,
                    "children": []
                }
            }

            for kwarg in args:

                # Check and make sure field and value are passed to search - req'd
                for arg in ('field', 'value'):
                    if arg not in kwarg:
                        raise Exception

                # Build the searching paramaters to be passed to the _api_request method
                # _api_request will add the additional data needed to form a valid request
                post_data['query']['children'].append({
                    "field": kwarg['field'],
                    "operator": kwarg.get('operator', "is"),
                    "value": kwarg['value']
                })

        for res in cls._api_search_request(path, post_data = post_data):
            for hit in res['hits']:
                yield hit

class AFTag(AutoFocusAPI):

    def __init__(self, **kwargs):

        self.comments = kwargs.get("comments", NotLoaded())
        self.refs = kwargs.get("refs", NotLoaded())
        self.review = kwargs.get("review", NotLoaded())
        self.support_id = kwargs.get("support_id", NotLoaded())

        for k,v in kwargs.items():
            setattr(self, k, v)

    def __getattribute__(self, attr):

        value = object.__getattribute__(self, attr)

        # Not offered in the list controller, have to call get to lazyload:
        #      comments, refs, review, support_id
        if attr in ('comments', 'refs', 'review', 'support_id') and type(value) is NotLoaded:

            # Reloading the data via the get method
            self = AFTag.get(self.public_tag_name)
            value = object.__getattribute__(self, attr)

        return value

    @classmethod
    def list(cls, *args, **kwargs):

        kwargs['scope'] = kwargs.get("scope", "Visible")
        kwargs['sortBy'] = kwargs.get("sortBy", "name")
        kwargs['order'] = kwargs.get("order", "asc")
        kwargs['pageSize'] = kwargs.get("pageSize", 1000)
        kwargs['pageNum'] = kwargs.get("pageNum", 0)

        resp = cls._api_request("/tags/", params = kwargs).json()
        results = []

        for tag in resp['tags']:
            results.append(AFTag(**tag))

        return results

    @classmethod
    def get(cls, tag_name):

        resp = cls._api_request("/tag/" + tag_name).json()

        return AFTag(**resp['tag'])

class AFSession(AutoFocusAPI):

    @classmethod
    def search(cls, *args, **kwargs):

        for res in cls._api_search("/sessions/search", *args, **kwargs):
            yield AFSession(**res['_source'])

class AFSample(AutoFocusAPI):

    def get_analyses(self, sections = ["file"], platforms = ["win7", "winxp"]):

        resp = self.__class__._api_request("/sample/" + self.sha256 + "/analysis", \
                    post_data = { "sections" : sections, "platforms" : platforms }).json()

        return resp

    @classmethod
    def search(cls, *args, **kwargs):

        for res in cls._api_search("/samples/search", *args, **kwargs):
            yield AFSample(**res['_source'])

if __name__ == "__main__":

    i = 0
    for sample in AFSample.search(field = "sample.malware", value = "1", operator = "is"):
        i += 1
        pprint(sample.__dict__)
        pprint(sample.get_analyses())

    print "%d results" % (i,)
<|endoftext|>


<commit_before>import os
import sys
import json
import time
import random
import datetime
from selenium import webdriver
from selenium.webdriver.common.keys import Keys

times = {}
last = ''
sys.stdout.write('Started.')##just so we know it is running.
if len(sys.argv) < 2:
    sys.argv.append('config.json')
with open(sys.argv[1]) as fp:
    data = json.load(fp)
OFFSET = random.randint(int(data['randomoffset']) * -1, int(data['randomoffset']))
while True:
    with open(sys.argv[1]) as fp:
        data = json.load(fp)
    now = datetime.datetime.now() + datetime.timedelta(minutes=OFFSET)
    if now.strftime('%Y-%m-%d') not in data['vacations'] \
            and now.strftime('%A') in data['workdays'] \
            and now.strftime('%H:%M') in data['times'] \
            and now.strftime('%H:%M') != last:
        driver = webdriver.Firefox()
        driver.get("https://ezlmappdc1f.adp.com/ezLaborManagerNet/Login/Login.aspx")
        if 'Client Login' in driver.title:
            elem = driver.find_element_by_id('txtClientName')
            elem.send_keys(data['clientname'])
            elem.send_keys(Keys.ENTER)
        if 'Login' in driver.title:
            elem = driver.find_element_by_id('txtUserID')
            elem.send_keys(data['username'])
            elem = driver.find_element_by_id('txtPassword')
            elem.send_keys(data['password'])
            elem.send_keys(Keys.ENTER)
            if 'Home' in driver.title:
                if data['times'][now.strftime('%H:%M')] == 'in':
                    try:
                        elem = driver.find_element_by_class_name('btnClockIn_1')
                        elem.click()
                        sys.stdout.write('{0}\tClockIn\tSuccess\tOK'.format(now) + os.linesep)
                    except Exception as e:
                        sys.stderr.write('{0}\tClockIn\tError\t{1}'.format(now, str(e)) + os.linesep)
                elif data['times'][now.strftime('%H:%M')] == 'out':
                    try:
                        elem = driver.find_element_by_class_name('btnClockOut_1')
                        elem.click()
                        sys.stdout.write('{0}\tClockOut\tSuccess\tOK'.format(now) + os.linesep)
                    except Exception as e:
                        sys.stderr.write('{0}\tClockOut\tError\t{1}'.format(now, str(e)) + os.linesep)
                else:
                    print('No Command Sent')
                OFFSET = random.randint(int(data['randomoffset']) * -1, int(data['randomoffset']))
                last = now.strftime('%H:%M')
            else:
                sys.stderr.write('{0}\tLogin\tError\tCould not login to ezLaborManager' + os.linesep)
        else:
            sys.stderr.write('{0}\tClientLogin\tError\tCould not login to ezLaborManager'.format(now) + os.linesep)
        driver.close()
    time.sleep(30)
<commit_msg>Adding more robust logging<commit_after>import os
import sys
import json
import time
import random
import logging
import datetime
import argparse
from selenium import webdriver
from selenium.webdriver.common.keys import Keys

if __name__ == '__main__':
    times = {}
    last = ''
    parser = argparse.ArgumentParser()
    parser.add_argument('config',
                        help='Config file (default: config.json)',
                        default='config.json')
    args = parser.parse_args()
    sys.stdout.write('Started\n')  # Just so we know it is running.
    with open(args.config) as fp:
        data = json.load(fp)
    OFFSET = random.randint(int(data['randomoffset']) * -1,
                            int(data['randomoffset']))  # Calculates Offset for login time
    while True:
        with open(args.config) as fp:
            data = json.load(fp)
        now = datetime.datetime.now() + datetime.timedelta(minutes=OFFSET)  # Gets current date and applies offset
        if now.strftime('%Y-%m-%d') not in data['vacations'] \
                and now.strftime('%A') in data['workdays'] \
                and now.strftime('%H:%M') in data['times'] \
                and now.strftime('%H:%M') != last:
            driver = webdriver.Firefox()
            driver.get("https://ezlmappdc1f.adp.com/ezLaborManagerNet/Login/Login.aspx")  # Goes to Client Login page
            if 'Client Login' in driver.title:
                logging.debug('Logging into client')
                elem = driver.find_element_by_id('txtClientName')
                elem.send_keys(data['clientname'])
                elem.send_keys(Keys.ENTER)
            if 'Login' in driver.title:
                elem = driver.find_element_by_xpath('//*[@id="lblClientName"]')
                logging.debug('Client: {0}'.format(elem.text))
                logging.debug('Logging into user')
                elem = driver.find_element_by_id('txtUserID')
                elem.send_keys(data['username'])
                elem = driver.find_element_by_id('txtPassword')
                elem.send_keys(data['password'])
                elem.send_keys(Keys.ENTER)
                if 'Home' in driver.title:
                    elem = driver.find_element_by_xpath('//*[@id="lblName"]')
                    logging.debug('User: {0}'.format(elem.text))
                    if data['times'][now.strftime('%H:%M')] == 'in':
                        try:
                            elem = driver.find_element_by_class_name('btnClockIn_1')
                            elem.click()
                            sys.stdout.write('{0}\tClockIn\tSuccess\tOK'.format(now) + os.linesep)
                        except Exception as e:
                            sys.stderr.write('{0}\tClockIn\tError\t{1}'.format(now, str(e)) + os.linesep)
                    elif data['times'][now.strftime('%H:%M')] == 'out':
                        try:
                            elem = driver.find_element_by_class_name('btnClockOut_1')
                            elem.click()
                            sys.stdout.write('{0}\tClockOut\tSuccess\tOK'.format(now) + os.linesep)
                        except Exception as e:
                            sys.stderr.write('{0}\tClockOut\tError\t{1}'.format(now, str(e)) + os.linesep)
                    else:
                        print('No Command Sent')
                    OFFSET = random.randint(int(data['randomoffset']) * -1, int(data['randomoffset']))
                    last = now.strftime('%H:%M')
                else:
                    sys.stderr.write('{0}\tLogin\tError\tCould not login to ezLaborManager' + os.linesep)
            else:
                sys.stderr.write('{0}\tClientLogin\tError\tCould not login to ezLaborManager'.format(now) + os.linesep)
            driver.close()
        time.sleep(30)
<|endoftext|>