<commit_before># This should cover all the syntactical constructs that we hope to support.
# Headings refer to https://docs.python.org/3/reference/expressions.html,
# and are selected whenever they incur dataflow.
# Intended sources should be the variable `SOURCE` and intended sinks should be
# arguments to the function `SINK` (see python/ql/test/experimental/dataflow/testConfig.qll).
#
# Functions whose name ends with "_with_local_flow" will also be tested for local flow.
#
# All functions starting with "test_" should run and print a source (sources are defined in testConfig.qll).

# These are defined so that we can evaluate the test code.
SOURCE = "source"
def SINK(x):
    print(x)

def test_tuple_with_local_flow():
    x = (3, SOURCE)
    y = x[1]
    SINK(y) # Flow missing

# 6.2.1. Identifiers (Names)
def test_names():
    x = SOURCE
    SINK(x)

# 6.2.2. Literals
def test_string_literal():
    x = "source"
    SINK(x)

def test_bytes_literal():
    x = b"source"
    SINK(x)

def test_integer_literal():
    x = 42
    SINK(x)

def test_floatnumber_literal():
    x = 42.0
    SINK(x)

def test_imagnumber_literal():
    x = 42j
    SINK(x) # Flow missing

# 6.2.3. Parenthesized forms
def test_parenthesized_form():
    x = (SOURCE)
    SINK(x) # Flow missing

# 6.2.5. List displays
def test_list_display():
    x = [SOURCE]
    SINK(x[0]) # Flow missing

def test_list_comprehension():
    x = [SOURCE for y in [3]]
    SINK(x[0]) # Flow missing

def test_nested_list_display():
    x = [* [SOURCE]]
    SINK(x[0]) # Flow missing

# 6.2.6. Set displays
def test_set_display():
    x = {SOURCE}
    SINK(x.pop()) # Flow missing

def test_set_comprehension():
    x = {SOURCE for y in [3]}
    SINK(x.pop()) # Flow missing

def test_nested_set_display():
    x = {* {SOURCE}}
    SINK(x.pop()) # Flow missing

# 6.2.7. Dictionary displays
def test_dict_display():
    x = {"s": SOURCE}
    SINK(x["s"]) # Flow missing

def test_dict_comprehension():
    x = {y: SOURCE for y in ["s"]}
    SINK(x["s"]) # Flow missing

def test_nested_dict_display():
    x = {** {"s": SOURCE}}
    SINK(x["s"]) # Flow missing

# 6.2.8. Generator expressions
def test_generator():
    x = (SOURCE for y in [3])
    SINK([*x][0]) # Flow missing

# 6.2.9. Yield expressions
def gen(x):
    yield x

def test_yield():
    g = gen(SOURCE)
    SINK(next(g)) # Flow missing

def gen_from(x):
    yield from gen(x)

def test_yield_from():
    g = gen_from(SOURCE)
    SINK(next(g)) # Flow missing

# a statement rather than an expression, but related to generators
def test_for():
    for x in gen(SOURCE):
        SINK(x) # Flow missing

# 6.2.9.1. Generator-iterator methods
def test___next__():
    g = gen(SOURCE)
    SINK(g.__next__()) # Flow missing

def gen2(x):
    m = yield x # argument of `send` has to flow to value of `yield x` (and so to `m`)
    yield m

def test_send():
    g = gen2(3)
    n = next(g)
    SINK(g.send(SOURCE)) # Flow missing

def gen_ex(x):
    try:
        yield 3
    except:
        yield x # `x` has to flow to call to `throw`

def test_throw():
    g = gen_ex(SOURCE)
    n = next(g)
    SINK(g.throw(TypeError)) # Flow missing

# no `test_close` as `close` involves no data flow

# 6.2.9.3. Asynchronous generator functions
async def agen(x):
    yield x

# 6.2.9.4. Asynchronous generator-iterator methods

# helper to run async test functions
def runa(a):
    import asyncio
    asyncio.run(a)

async def atest___anext__():
    g = agen(SOURCE)
    SINK(await g.__anext__()) # Flow missing

def test___anext__():
    runa(atest___anext__())

async def agen2(x):
    m = yield x # argument of `send` has to flow to value of `yield x` (and so to `m`)
    yield m

async def atest_asend():
    g = agen2(3)
    n = await g.__anext__()
    SINK(await g.asend(SOURCE)) # Flow missing

def test_asend():
    runa(atest_asend())

async def agen_ex(x):
    try:
        yield 3
    except:
        yield x # `x` has to flow to call to `athrow`

async def atest_athrow():
    g = agen_ex(SOURCE)
    n = await g.__anext__()
    SINK(await g.athrow(TypeError)) # Flow missing

def test_athrow():
    runa(atest_athrow())

# 6.3.1. Attribute references
class C:
    a = SOURCE

def test_attribute_reference():
    SINK(C.a) # Flow missing

# overriding __getattr__ should be tested by the class coverage tests

# 6.3.2. Subscriptions
# This does not constitute dataflow (but could be taint flow)
def example_subscription_string():
    SINK("source"[0]) # Flow not expected

def test_subscription_tuple():
    SINK((SOURCE,)[0]) # Flow missing

def test_subscription_list():
    SINK([SOURCE][0]) # Flow missing

def test_subscription_mapping():
    SINK({"s":SOURCE}["s"]) # Flow missing

# overriding __getitem__ should be tested by the class coverage tests

# 6.3.3. Slicings
l = [SOURCE]

def test_slicing():
    s = l[0:1:1]
    SINK(s[0]) # Flow missing

# The grammar seems to allow `l[0:1:1, 0:1]`, but the interpreter does not like it

# 6.3.4. Calls
def f(a, b):
    return b

def test_call_positional():
    SINK(f(3, SOURCE))

def test_call_keyword():
    SINK(f(3, b=SOURCE)) # Flow missing

def test_call_unpack_iterable():
    SINK(f(3, *[SOURCE])) # Flow missing

def test_call_unpack_mapping():
    SINK(f(3, **{"b": SOURCE})) # Flow missing

def f_extra_pos(a, *b):
    return b[0]

def test_call_extra_pos():
    SINK(f_extra_pos(3, SOURCE)) # Flow missing

def f_extra_keyword(a, **b):
    return b["b"]

def test_call_extra_keyword():
    SINK(f_extra_keyword(3, b=SOURCE)) # Flow missing

# return the name of the first extra keyword argument
def f_extra_keyword_flow(**a):
    return [*a][0]

# call the function with our source as the name of the keyword arguemnt
def test_call_extra_keyword_flow():
    SINK(f_extra_keyword_flow(**{SOURCE: None})) # Flow missing

# 6.12. Assignment expressions
def test_assignment_expression():
    x = 3
    SINK(x := SOURCE) # Flow missing

# 6.13. Conditional expressions
def test_conditional_true():
    SINK(SOURCE if True else 3) # Flow missing

def test_conditional_false():
    SINK(3 if False else SOURCE) # Flow missing

# Condition is evaluated first, so x is SOURCE once chosen
def test_conditional_evaluation_true():
    x = 3
    SINK(x if (SOURCE == (x := SOURCE)) else 3) # Flow missing

# Condition is evaluated first, so x is SOURCE once chosen
def test_conditional_evaluation_false():
    x = 3
    SINK(3 if (3 == (x := SOURCE)) else x) # Flow missing

# 6.14. Lambdas
def test_lambda():
    f = lambda x : x
    SINK(f(SOURCE))
<commit_msg>Python: Fix missing flow annotation<commit_after># This should cover all the syntactical constructs that we hope to support.
# Headings refer to https://docs.python.org/3/reference/expressions.html,
# and are selected whenever they incur dataflow.
# Intended sources should be the variable `SOURCE` and intended sinks should be
# arguments to the function `SINK` (see python/ql/test/experimental/dataflow/testConfig.qll).
#
# Functions whose name ends with "_with_local_flow" will also be tested for local flow.
#
# All functions starting with "test_" should run and print a source (sources are defined in testConfig.qll).

# These are defined so that we can evaluate the test code.
SOURCE = "source"
def SINK(x):
    print(x)

def test_tuple_with_local_flow():
    x = (3, SOURCE)
    y = x[1]
    SINK(y) # Flow missing

# 6.2.1. Identifiers (Names)
def test_names():
    x = SOURCE
    SINK(x)

# 6.2.2. Literals
def test_string_literal():
    x = "source"
    SINK(x)

def test_bytes_literal():
    x = b"source"
    SINK(x)

def test_integer_literal():
    x = 42
    SINK(x)

def test_floatnumber_literal():
    x = 42.0
    SINK(x)

def test_imagnumber_literal():
    x = 42j
    SINK(x) # Flow missing

# 6.2.3. Parenthesized forms
def test_parenthesized_form():
    x = (SOURCE)
    SINK(x)

# 6.2.5. List displays
def test_list_display():
    x = [SOURCE]
    SINK(x[0]) # Flow missing

def test_list_comprehension():
    x = [SOURCE for y in [3]]
    SINK(x[0]) # Flow missing

def test_nested_list_display():
    x = [* [SOURCE]]
    SINK(x[0]) # Flow missing

# 6.2.6. Set displays
def test_set_display():
    x = {SOURCE}
    SINK(x.pop()) # Flow missing

def test_set_comprehension():
    x = {SOURCE for y in [3]}
    SINK(x.pop()) # Flow missing

def test_nested_set_display():
    x = {* {SOURCE}}
    SINK(x.pop()) # Flow missing

# 6.2.7. Dictionary displays
def test_dict_display():
    x = {"s": SOURCE}
    SINK(x["s"]) # Flow missing

def test_dict_comprehension():
    x = {y: SOURCE for y in ["s"]}
    SINK(x["s"]) # Flow missing

def test_nested_dict_display():
    x = {** {"s": SOURCE}}
    SINK(x["s"]) # Flow missing

# 6.2.8. Generator expressions
def test_generator():
    x = (SOURCE for y in [3])
    SINK([*x][0]) # Flow missing

# 6.2.9. Yield expressions
def gen(x):
    yield x

def test_yield():
    g = gen(SOURCE)
    SINK(next(g)) # Flow missing

def gen_from(x):
    yield from gen(x)

def test_yield_from():
    g = gen_from(SOURCE)
    SINK(next(g)) # Flow missing

# a statement rather than an expression, but related to generators
def test_for():
    for x in gen(SOURCE):
        SINK(x) # Flow missing

# 6.2.9.1. Generator-iterator methods
def test___next__():
    g = gen(SOURCE)
    SINK(g.__next__()) # Flow missing

def gen2(x):
    m = yield x # argument of `send` has to flow to value of `yield x` (and so to `m`)
    yield m

def test_send():
    g = gen2(3)
    n = next(g)
    SINK(g.send(SOURCE)) # Flow missing

def gen_ex(x):
    try:
        yield 3
    except:
        yield x # `x` has to flow to call to `throw`

def test_throw():
    g = gen_ex(SOURCE)
    n = next(g)
    SINK(g.throw(TypeError)) # Flow missing

# no `test_close` as `close` involves no data flow

# 6.2.9.3. Asynchronous generator functions
async def agen(x):
    yield x

# 6.2.9.4. Asynchronous generator-iterator methods

# helper to run async test functions
def runa(a):
    import asyncio
    asyncio.run(a)

async def atest___anext__():
    g = agen(SOURCE)
    SINK(await g.__anext__()) # Flow missing

def test___anext__():
    runa(atest___anext__())

async def agen2(x):
    m = yield x # argument of `send` has to flow to value of `yield x` (and so to `m`)
    yield m

async def atest_asend():
    g = agen2(3)
    n = await g.__anext__()
    SINK(await g.asend(SOURCE)) # Flow missing

def test_asend():
    runa(atest_asend())

async def agen_ex(x):
    try:
        yield 3
    except:
        yield x # `x` has to flow to call to `athrow`

async def atest_athrow():
    g = agen_ex(SOURCE)
    n = await g.__anext__()
    SINK(await g.athrow(TypeError)) # Flow missing

def test_athrow():
    runa(atest_athrow())

# 6.3.1. Attribute references
class C:
    a = SOURCE

def test_attribute_reference():
    SINK(C.a) # Flow missing

# overriding __getattr__ should be tested by the class coverage tests

# 6.3.2. Subscriptions
# This does not constitute dataflow (but could be taint flow)
def example_subscription_string():
    SINK("source"[0]) # Flow not expected

def test_subscription_tuple():
    SINK((SOURCE,)[0]) # Flow missing

def test_subscription_list():
    SINK([SOURCE][0]) # Flow missing

def test_subscription_mapping():
    SINK({"s":SOURCE}["s"]) # Flow missing

# overriding __getitem__ should be tested by the class coverage tests

# 6.3.3. Slicings
l = [SOURCE]

def test_slicing():
    s = l[0:1:1]
    SINK(s[0]) # Flow missing

# The grammar seems to allow `l[0:1:1, 0:1]`, but the interpreter does not like it

# 6.3.4. Calls
def f(a, b):
    return b

def test_call_positional():
    SINK(f(3, SOURCE))

def test_call_keyword():
    SINK(f(3, b=SOURCE)) # Flow missing

def test_call_unpack_iterable():
    SINK(f(3, *[SOURCE])) # Flow missing

def test_call_unpack_mapping():
    SINK(f(3, **{"b": SOURCE})) # Flow missing

def f_extra_pos(a, *b):
    return b[0]

def test_call_extra_pos():
    SINK(f_extra_pos(3, SOURCE)) # Flow missing

def f_extra_keyword(a, **b):
    return b["b"]

def test_call_extra_keyword():
    SINK(f_extra_keyword(3, b=SOURCE)) # Flow missing

# return the name of the first extra keyword argument
def f_extra_keyword_flow(**a):
    return [*a][0]

# call the function with our source as the name of the keyword arguemnt
def test_call_extra_keyword_flow():
    SINK(f_extra_keyword_flow(**{SOURCE: None})) # Flow missing

# 6.12. Assignment expressions
def test_assignment_expression():
    x = 3
    SINK(x := SOURCE) # Flow missing

# 6.13. Conditional expressions
def test_conditional_true():
    SINK(SOURCE if True else 3) # Flow missing

def test_conditional_false():
    SINK(3 if False else SOURCE) # Flow missing

# Condition is evaluated first, so x is SOURCE once chosen
def test_conditional_evaluation_true():
    x = 3
    SINK(x if (SOURCE == (x := SOURCE)) else 3) # Flow missing

# Condition is evaluated first, so x is SOURCE once chosen
def test_conditional_evaluation_false():
    x = 3
    SINK(3 if (3 == (x := SOURCE)) else x) # Flow missing

# 6.14. Lambdas
def test_lambda():
    f = lambda x : x
    SINK(f(SOURCE))
<|endoftext|>


<commit_before>""" soil_moisture_maps.py
Draw the maps of soil moisture from the outputs of the TOPKAPI simulations
"""
import numpy as np
import pylab as pl
import matplotlib.numerix.ma as M

import pytopkapi.pretreatment as pm
import pytopkapi.utils as ut

from ConfigParser import SafeConfigParser
config = SafeConfigParser()

def run(ini_file='plot_soil_moisture_maps.ini'):

    config.read(ini_file)
    print 'Read the file ',ini_file

    file_global_param=config.get('files','file_global_param')
    file_cell_param=config.get('files','file_cell_param')
    file_sim=config.get('files','file_sim')

    path_out=config.get('paths','path_out')

    fac_L=config.getfloat('calib_params','fac_L')
    fac_Ks=config.getfloat('calib_params','fac_Ks')
    fac_n_o=config.getfloat('calib_params','fac_n_o')
    fac_n_c=config.getfloat('calib_params','fac_n_c')

    t1=config.getfloat('flags','t1')
    t2=config.getfloat('flags','t2')
    variable=config.getfloat('flags','variable')

    ##~~~~~~PROGRAM~~~~~##
    # Create the folder if it doesn't exist
    ut.check_folder_exist(path_out)

    #~~~~Read Global parameters file
    X,Dt,alpha_s,alpha_o,alpha_c,A_thres,W_min,W_max\
      =pm.read_global_parameters(file_global_param)
    #~~~~Read Cell parameters file
    ar_cell_label,ar_coorx,ar_coory,ar_lambda,ar_Xc,ar_dam,ar_tan_beta,ar_tan_beta_channel,ar_L0,ar_Ks0,\
    ar_theta_r,ar_theta_s,ar_n_o0,ar_n_c0,\
    ar_cell_down,ar_pVs_t0,ar_Vo_t0,ar_Qc_t0,ar_kc\
        =pm.read_cell_parameters(file_cell_param)
    #~~~~Number of cell in the catchment
    nb_cell=len(ar_cell_label)
    #~~~~Computation of cell order
    ar_label_sort=pm.sort_cell(ar_cell_label,ar_cell_down)
    #~~~~Computation of upcells
    li_cell_up=pm.direct_up_cell(ar_cell_label,ar_cell_down,ar_label_sort)
    #~~~~Computation of drained area
    ar_A_drained=pm.drained_area(ar_label_sort,li_cell_up,X)
    #~~~~Modifies the values of the parameters
    ar_L=ar_L0*fac_L
    ar_Ks=ar_Ks0*fac_Ks
    ar_n_o=ar_n_o0*fac_n_o
    ar_n_c=ar_n_c0*fac_n_c
    #~~~~Computation of model parameters from physical parameters
    ar_Vsm, ar_b_s, ar_b_o, ar_W, ar_b_c\
      =pm.compute_cell_param(X,ar_Xc,Dt,alpha_s,alpha_o,alpha_c,nb_cell,\
                              A_thres,W_max,W_min,\
                              ar_lambda,ar_tan_beta,ar_tan_beta_channel,ar_L,\
                              ar_Ks,ar_theta_r,ar_theta_s,ar_n_o,ar_n_c,\
                              ar_A_drained)

    #Read of data from the outputs of TOPKAPI in hdf5 format
    ndar_Vs=np.array(ut.read_one_array_hdf(file_sim,'/Soil/','V_s'))

    #Assign the variables
    if variable==1:
        im_out=path_out+'field_Vs_'
        tab=ndar_Vs
    elif variable==2:
        im_out=path_out+'field_Vo_'
        tab=ndar_Vo
    elif variable==3:
        im_out=path_out+'field_Vo_bin_'
        tab=ndar_Vo
        tab[tab>0]=1.
    elif variable==4:
        im_out=path_out+'field_SWI_'
        tab=ndar_Vs/ar_Vsm*100.

    # Plot the maps
    for t in np.arange(int(t1),int(t2+1)):
        print 'Map time-step ', t
        image_out=im_out+ut.string(t,len(str(t2)))+'.png'
        field_map_ndar(tab,t,ar_coorx,ar_coory,X,image_out,variable)

def field_map_ndar(ndar_field,t,ar_coorx,ar_coory,X,image_out,variable):

    ar_field=ndar_field[t,:]
    max_val=int(np.max(ndar_field))
    if variable==4:
        max_val=100.
    xmin=min(ar_coorx);xmax=max(ar_coorx)
    ymin=min(ar_coory);ymax=max(ar_coory)
    step=X
    nx=(xmax-xmin)/step+1
    ny=(ymax-ymin)/step+1

    ar_indx=np.array((ar_coorx-xmin)/step,int)
    ar_indy=np.array((ar_coory-ymin)/step,int)

    ar_map=np.ones((ny,nx))*-99.9
    ar_map[ar_indy,ar_indx]=ar_field

    ar_map2 = M.masked_where(ar_map <0, ar_map)
    ut.check_file_exist(image_out)

    pl.clf()
    pl.axes(axisbg='gray')
    pl.imshow(ar_map2, cmap=pl.cm.RdBu,
              interpolation='Nearest', origin='lower', vmax=max_val, vmin=0)
    pl.title('time step= '+ut.string(t,len(str(t))))
    pl.colorbar()
    pl.savefig(image_out)

<commit_msg>BUG: Don't use deprecated matplotlib.numerix<commit_after>""" soil_moisture_maps.py
Draw the maps of soil moisture from the outputs of the TOPKAPI simulations
"""
import numpy as np
import pylab as pl
import numpy.ma as M

import pytopkapi.pretreatment as pm
import pytopkapi.utils as ut

from ConfigParser import SafeConfigParser
config = SafeConfigParser()

def run(ini_file='plot_soil_moisture_maps.ini'):

    config.read(ini_file)
    print 'Read the file ',ini_file

    file_global_param=config.get('files','file_global_param')
    file_cell_param=config.get('files','file_cell_param')
    file_sim=config.get('files','file_sim')

    path_out=config.get('paths','path_out')

    fac_L=config.getfloat('calib_params','fac_L')
    fac_Ks=config.getfloat('calib_params','fac_Ks')
    fac_n_o=config.getfloat('calib_params','fac_n_o')
    fac_n_c=config.getfloat('calib_params','fac_n_c')

    t1=config.getfloat('flags','t1')
    t2=config.getfloat('flags','t2')
    variable=config.getfloat('flags','variable')

    ##~~~~~~PROGRAM~~~~~##
    # Create the folder if it doesn't exist
    ut.check_folder_exist(path_out)

    #~~~~Read Global parameters file
    X,Dt,alpha_s,alpha_o,alpha_c,A_thres,W_min,W_max\
      =pm.read_global_parameters(file_global_param)
    #~~~~Read Cell parameters file
    ar_cell_label,ar_coorx,ar_coory,ar_lambda,ar_Xc,ar_dam,ar_tan_beta,ar_tan_beta_channel,ar_L0,ar_Ks0,\
    ar_theta_r,ar_theta_s,ar_n_o0,ar_n_c0,\
    ar_cell_down,ar_pVs_t0,ar_Vo_t0,ar_Qc_t0,ar_kc\
        =pm.read_cell_parameters(file_cell_param)
    #~~~~Number of cell in the catchment
    nb_cell=len(ar_cell_label)
    #~~~~Computation of cell order
    ar_label_sort=pm.sort_cell(ar_cell_label,ar_cell_down)
    #~~~~Computation of upcells
    li_cell_up=pm.direct_up_cell(ar_cell_label,ar_cell_down,ar_label_sort)
    #~~~~Computation of drained area
    ar_A_drained=pm.drained_area(ar_label_sort,li_cell_up,X)
    #~~~~Modifies the values of the parameters
    ar_L=ar_L0*fac_L
    ar_Ks=ar_Ks0*fac_Ks
    ar_n_o=ar_n_o0*fac_n_o
    ar_n_c=ar_n_c0*fac_n_c
    #~~~~Computation of model parameters from physical parameters
    ar_Vsm, ar_b_s, ar_b_o, ar_W, ar_b_c\
      =pm.compute_cell_param(X,ar_Xc,Dt,alpha_s,alpha_o,alpha_c,nb_cell,\
                              A_thres,W_max,W_min,\
                              ar_lambda,ar_tan_beta,ar_tan_beta_channel,ar_L,\
                              ar_Ks,ar_theta_r,ar_theta_s,ar_n_o,ar_n_c,\
                              ar_A_drained)

    #Read of data from the outputs of TOPKAPI in hdf5 format
    ndar_Vs=np.array(ut.read_one_array_hdf(file_sim,'/Soil/','V_s'))

    #Assign the variables
    if variable==1:
        im_out=path_out+'field_Vs_'
        tab=ndar_Vs
    elif variable==2:
        im_out=path_out+'field_Vo_'
        tab=ndar_Vo
    elif variable==3:
        im_out=path_out+'field_Vo_bin_'
        tab=ndar_Vo
        tab[tab>0]=1.
    elif variable==4:
        im_out=path_out+'field_SWI_'
        tab=ndar_Vs/ar_Vsm*100.

    # Plot the maps
    for t in np.arange(int(t1),int(t2+1)):
        print 'Map time-step ', t
        image_out=im_out+ut.string(t,len(str(t2)))+'.png'
        field_map_ndar(tab,t,ar_coorx,ar_coory,X,image_out,variable)

def field_map_ndar(ndar_field,t,ar_coorx,ar_coory,X,image_out,variable):

    ar_field=ndar_field[t,:]
    max_val=int(np.max(ndar_field))
    if variable==4:
        max_val=100.
    xmin=min(ar_coorx);xmax=max(ar_coorx)
    ymin=min(ar_coory);ymax=max(ar_coory)
    step=X
    nx=(xmax-xmin)/step+1
    ny=(ymax-ymin)/step+1

    ar_indx=np.array((ar_coorx-xmin)/step,int)
    ar_indy=np.array((ar_coory-ymin)/step,int)

    ar_map=np.ones((ny,nx))*-99.9
    ar_map[ar_indy,ar_indx]=ar_field

    ar_map2 = M.masked_where(ar_map <0, ar_map)
    ut.check_file_exist(image_out)

    pl.clf()
    pl.axes(axisbg='gray')
    pl.imshow(ar_map2, cmap=pl.cm.RdBu,
              interpolation='Nearest', origin='lower', vmax=max_val, vmin=0)
    pl.title('time step= '+ut.string(t,len(str(t))))
    pl.colorbar()
    pl.savefig(image_out)

<|endoftext|>


<commit_before>"""
Overview
========
Retrieve public feed data from Hybrid Analysis
"""
import argparse

from stoq.args import StoqArgs
from stoq.plugins import StoqWorkerPlugin

class HybridAnalysis(StoqWorkerPlugin):

    def __init__(self):
        super().__init__()

    def activate(self, stoq):
        self.stoq = stoq

        parser = argparse.ArgumentParser()
        parser = StoqArgs(parser)
        options = parser.parse_args(self.stoq.argv[2:])

        super().activate(options=options)

    def scan(self, payload=None, **kwargs):
        super().scan()
        url = "https://www.hybrid-analysis.com/feed?json"
        raw_results = self.stoq.get_file(url)
        scan_results = self.stoq.loads(raw_results)

        if scan_results['data']:
            return scan_results['data']

        return None
<commit_msg>Update hybridanalysis.py<commit_after>#   Copyright 2014-2015 PUNCH Cyber Analytics Group
#
#   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.

"""
Overview
========
Retrieve public feed data from Hybrid Analysis
"""
import argparse

from stoq.args import StoqArgs
from stoq.plugins import StoqWorkerPlugin

class HybridAnalysis(StoqWorkerPlugin):

    def __init__(self):
        super().__init__()

    def activate(self, stoq):
        self.stoq = stoq

        parser = argparse.ArgumentParser()
        parser = StoqArgs(parser)
        options = parser.parse_args(self.stoq.argv[2:])

        super().activate(options=options)

    def scan(self, payload=None, **kwargs):
        super().scan()
        url = "https://www.hybrid-analysis.com/feed?json"
        raw_results = self.stoq.get_file(url)
        scan_results = self.stoq.loads(raw_results)

        if scan_results['data']:
            return scan_results['data']

        return None
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright (c) 2014-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
#

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import os
import sys
import tempfile
import subprocess
import xml.etree.ElementTree as ET
import getopt
import shutil
from . import metadata
from .simple_puller import SimplePuller
import glob

from sets import Set
from os.path import join
from os.path import abspath

ROOT_SCREENSHOT_DIR = '/sdcard/screenshots'
OLD_ROOT_SCREENSHOT_DIR = '/data/data/'

def usage():
    print >>sys.stderr, "usage: ./scripts/screenshot_tests/pull_screenshots com.facebook.apk.name.tests [--generate-png]"
    return

def generate_html(dir):
    root = ET.parse(join(dir, 'metadata.xml')).getroot()
    alternate = False
    index_html = abspath(join(dir, "index.html"))
    with open(index_html, "w") as html:
        html.write('<!DOCTYPE html>')
        html.write('')
        html.write('')
        html.write('')
        html.write('')
        html.write('')
        html.write('<link rel="stylesheet" href="https://ajax.googleapis.com/ajax/libs/jqueryui/1.11.3/themes/smoothness/jquery-ui.css" />')
        html.write('<link rel="stylesheet" href="default.css">')
        html.write('')

        html.write('')
        for screenshot in root.iter('screenshot'):
            alternate = not alternate
            html.write('
' % ('alternate' if alternate else ''))
            html.write('
%s
' % (screenshot.find('name').text))
            html.write('<button class="view_dump" data-name="%s">Dump view hierarchy</button>' % (screenshot.find('name').text))

            extras = screenshot.find('extras')
            if extras is not None:
                str = ""
                for node in extras:
                    if node.text is not None:
                        str = str + "*****" + node.tag + "*****\n\n" + node.text + "\n\n\n"
                if str != "":
                    extra_html = '<button class="extra" data="%s">Extra info</button>' % str
                    html.write(extra_html.encode('utf-8').strip())

            description = screenshot.find('description')
            if description is not None:
                html.write('
%s
' % description.text)

            error = screenshot.find('error')
            if error is not None:
                html.write('
%s
' % error.text)
            else:
                write_image(dir, html, screenshot)

            html.write('
')

        html.write('')
        return index_html

def write_image(dir, html, screenshot):
    html.write('
')
    for y in xrange(int(screenshot.find('tile_height').text)):
        html.write('
')
        for x in xrange(int(screenshot.find('tile_width').text)):
            html.write('')
            image_file = "./" + screenshot.find('name').text

            if x != 0 or y != 0:
                image_file += "_%d_%d" % (x, y)

            image_file += ".png"

            if os.path.exists(join(dir, image_file)):
                html.write('' % image_file)

            html.write('	')
        html.write('
')
    html.write('
')

def test_for_wkhtmltoimage():
    if subprocess.call(['which', 'wkhtmltoimage']) != 0:
        raise RuntimeError("""Could not find wkhtmltoimage in your path, we need this for generating pngs
Download an appropriate version from:
    http://wkhtmltopdf.org/downloads.html""")

def generate_png(path_to_html, path_to_png):
    test_for_wkhtmltoimage()
    subprocess.check_call(['wkhtmltoimage', path_to_html, path_to_png], stdout=sys.stdout)

def copy_assets(destination):
    """Copy static assets required for rendering the HTML"""
    thisdir = os.path.dirname(__file__)
    os.symlink(abspath(join(thisdir, "default.css")), join(destination, "default.css"))
    os.symlink(abspath(join(thisdir, "default.js")), join(destination, "default.js"))
    os.symlink(abspath(join(thisdir, "background.png")), join(destination, "background.png"))

def pull_metadata(package, dir, adb_puller):
    metadata_file = '%s/%s/screenshots-default/metadata.xml' % (ROOT_SCREENSHOT_DIR, package)
    old_metadata_file = '%s/%s/app_screenshots-default/metadata.xml' % (OLD_ROOT_SCREENSHOT_DIR, package)

    if adb_puller.remote_file_exists(metadata_file):
        adb_puller.pull(metadata_file, join(dir, 'metadata.xml'))
    elif adb_puller.remote_file_exists(old_metadata_file):
        adb_puller.pull(old_metadata_file, join(dir, 'metadata.xml'))
    else:
        create_empty_metadata_file(dir)

def create_empty_metadata_file(dir):
    with open(join(dir, 'metadata.xml'), 'w') as out:
        out.write("""<?xml version="1.0" encoding="UTF-8"?>
<screenshots>
</screenshots>""")

def pull_images(dir, adb_puller):
    root = ET.parse(join(dir, 'metadata.xml')).getroot()
    for s in root.iter('screenshot'):
        filename_nodes = s.findall('absolute_file_name')
        for filename_node in filename_nodes:
            adb_puller.pull(filename_node.text, join(dir, os.path.basename(filename_node.text)))
        dump_node = s.find('view_hierarchy')
        if dump_node is not None:
            adb_puller.pull(dump_node.text, join(dir, os.path.basename(dump_node.text)))

def pull_all(package, dir, adb_puller):
    pull_metadata(package, dir, adb_puller=adb_puller)
    pull_images(dir, adb_puller=adb_puller)

def pull_filtered(package, dir, adb_puller, filter_name_regex=None):
    pull_metadata(package, dir, adb_puller=adb_puller)
    metadata.filter_screenshots(join(dir, 'metadata.xml'), name_regex=filter_name_regex)
    pull_images(dir, adb_puller=adb_puller)

def pull_screenshots(process,
                     adb_puller,
                     temp_dir=None,
                     filter_name_regex=None,
                     opt_generate_png=None):
    temp_dir = temp_dir or tempfile.mkdtemp(prefix='screenshots')
    copy_assets(temp_dir)

    pull_filtered(process, adb_puller=adb_puller, dir=temp_dir, filter_name_regex=filter_name_regex)
    path_to_html = generate_html(temp_dir)

    if opt_generate_png:
        generate_png(path_to_html, opt_generate_png)
        shutil.rmtree(temp_dir)
    else:
        print('Open the following url in a browser to view the results: ')
        print('  file://%s' % path_to_html)

def _check_output(args, **kwargs):
    with open(os.devnull) as f:
        kwargs['stderr'] = f
        return subprocess.check_output(args, **kwargs)

def parse_package_line(line):
    """The line looks like this:
    package: name='com.facebook.testing.tests' versionCode='1' versionName=''"""

    for word in line.split():
        if word.startswith("name='"):
            return word[len("name='"):-1]

def get_aapt_bin():
    """Find the binary for aapt from $ANDROID_SDK"""
    android_sdk = os.environ.get('ANDROID_SDK') or os.environ.get('ANDROID_HOME')
    build_tools = os.path.join(android_sdk, 'build-tools')

    all = list(glob.glob(os.path.join(build_tools, '*/aapt')))
    bad = list(glob.glob(os.path.join(build_tools, 'android-*/aapt')))
    good = list(Set(all) - Set(bad))

    good.sort()
    bad.sort()

    if len(good) == 0:
        return bad[-1]

    return good[-1]

def get_package(apk):
    output = _check_output([get_aapt_bin(), 'dump', 'badging', apk], stderr=os.devnull)
    for line in output.split('\n'):
        if line.startswith('package:'):
            return parse_package_line(line)

def main(argv):
    try:
        opt_list, rest_args = getopt.gnu_getopt(
            argv[1:],
            "eds:",
            ["generate-png=", "filter-name-regex=", "apk"])
    except getopt.GetoptError, err:
        usage()
        return 2

    if len(rest_args) != 1:
        usage()
        return 2

    process = rest_args[0]  # something like com.facebook.places.tests

    opts = dict(opt_list)

    if "--apk" in opts:
        # treat process as an apk instead
        process = get_package(process)

    puller_args = []
    if "-e" in opts:
        puller_args.append("-e")

    if "-d" in opts:
        puller_args.append("-d")

    if "-s" in opts:
        puller_args += ["-s", opts["-s"]]

    return pull_screenshots(process,
                            filter_name_regex=opts.get('--filter-name-regex'),
                            opt_generate_png=opts.get('--generate-png'),
                            adb_puller=SimplePuller(puller_args))

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<commit_msg>extract copyasset<commit_after>#!/usr/bin/env python
#
# Copyright (c) 2014-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
#

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import os
import sys
import tempfile
import subprocess
import xml.etree.ElementTree as ET
import getopt
import shutil
from . import metadata
from .simple_puller import SimplePuller
import glob

from sets import Set
from os.path import join
from os.path import abspath

ROOT_SCREENSHOT_DIR = '/sdcard/screenshots'
OLD_ROOT_SCREENSHOT_DIR = '/data/data/'

def usage():
    print >>sys.stderr, "usage: ./scripts/screenshot_tests/pull_screenshots com.facebook.apk.name.tests [--generate-png]"
    return

def generate_html(dir):
    root = ET.parse(join(dir, 'metadata.xml')).getroot()
    alternate = False
    index_html = abspath(join(dir, "index.html"))
    with open(index_html, "w") as html:
        html.write('<!DOCTYPE html>')
        html.write('')
        html.write('')
        html.write('')
        html.write('')
        html.write('')
        html.write('<link rel="stylesheet" href="https://ajax.googleapis.com/ajax/libs/jqueryui/1.11.3/themes/smoothness/jquery-ui.css" />')
        html.write('<link rel="stylesheet" href="default.css">')
        html.write('')

        html.write('')
        for screenshot in root.iter('screenshot'):
            alternate = not alternate
            html.write('
' % ('alternate' if alternate else ''))
            html.write('
%s
' % (screenshot.find('name').text))
            html.write('<button class="view_dump" data-name="%s">Dump view hierarchy</button>' % (screenshot.find('name').text))

            extras = screenshot.find('extras')
            if extras is not None:
                str = ""
                for node in extras:
                    if node.text is not None:
                        str = str + "*****" + node.tag + "*****\n\n" + node.text + "\n\n\n"
                if str != "":
                    extra_html = '<button class="extra" data="%s">Extra info</button>' % str
                    html.write(extra_html.encode('utf-8').strip())

            description = screenshot.find('description')
            if description is not None:
                html.write('
%s
' % description.text)

            error = screenshot.find('error')
            if error is not None:
                html.write('
%s
' % error.text)
            else:
                write_image(dir, html, screenshot)

            html.write('
')

        html.write('')
        return index_html

def write_image(dir, html, screenshot):
    html.write('
')
    for y in xrange(int(screenshot.find('tile_height').text)):
        html.write('
')
        for x in xrange(int(screenshot.find('tile_width').text)):
            html.write('')
            image_file = "./" + screenshot.find('name').text

            if x != 0 or y != 0:
                image_file += "_%d_%d" % (x, y)

            image_file += ".png"

            if os.path.exists(join(dir, image_file)):
                html.write('' % image_file)

            html.write('	')
        html.write('
')
    html.write('
')

def test_for_wkhtmltoimage():
    if subprocess.call(['which', 'wkhtmltoimage']) != 0:
        raise RuntimeError("""Could not find wkhtmltoimage in your path, we need this for generating pngs
Download an appropriate version from:
    http://wkhtmltopdf.org/downloads.html""")

def generate_png(path_to_html, path_to_png):
    test_for_wkhtmltoimage()
    subprocess.check_call(['wkhtmltoimage', path_to_html, path_to_png], stdout=sys.stdout)

def copy_assets(destination):
    """Copy static assets required for rendering the HTML"""
    _copy_asset("default.css", destination)
    _copy_asset("default.js", destination)
    _copy_asset("background.png", destination)

def _copy_asset(filename, destination):
    thisdir = os.path.dirname(__file__)
    os.symlink(abspath(join(thisdir, filename)), join(destination, filename))

def pull_metadata(package, dir, adb_puller):
    metadata_file = '%s/%s/screenshots-default/metadata.xml' % (ROOT_SCREENSHOT_DIR, package)
    old_metadata_file = '%s/%s/app_screenshots-default/metadata.xml' % (OLD_ROOT_SCREENSHOT_DIR, package)

    if adb_puller.remote_file_exists(metadata_file):
        adb_puller.pull(metadata_file, join(dir, 'metadata.xml'))
    elif adb_puller.remote_file_exists(old_metadata_file):
        adb_puller.pull(old_metadata_file, join(dir, 'metadata.xml'))
    else:
        create_empty_metadata_file(dir)

def create_empty_metadata_file(dir):
    with open(join(dir, 'metadata.xml'), 'w') as out:
        out.write("""<?xml version="1.0" encoding="UTF-8"?>
<screenshots>
</screenshots>""")

def pull_images(dir, adb_puller):
    root = ET.parse(join(dir, 'metadata.xml')).getroot()
    for s in root.iter('screenshot'):
        filename_nodes = s.findall('absolute_file_name')
        for filename_node in filename_nodes:
            adb_puller.pull(filename_node.text, join(dir, os.path.basename(filename_node.text)))
        dump_node = s.find('view_hierarchy')
        if dump_node is not None:
            adb_puller.pull(dump_node.text, join(dir, os.path.basename(dump_node.text)))

def pull_all(package, dir, adb_puller):
    pull_metadata(package, dir, adb_puller=adb_puller)
    pull_images(dir, adb_puller=adb_puller)

def pull_filtered(package, dir, adb_puller, filter_name_regex=None):
    pull_metadata(package, dir, adb_puller=adb_puller)
    metadata.filter_screenshots(join(dir, 'metadata.xml'), name_regex=filter_name_regex)
    pull_images(dir, adb_puller=adb_puller)

def pull_screenshots(process,
                     adb_puller,
                     temp_dir=None,
                     filter_name_regex=None,
                     opt_generate_png=None):
    temp_dir = temp_dir or tempfile.mkdtemp(prefix='screenshots')
    copy_assets(temp_dir)

    pull_filtered(process, adb_puller=adb_puller, dir=temp_dir, filter_name_regex=filter_name_regex)
    path_to_html = generate_html(temp_dir)

    if opt_generate_png:
        generate_png(path_to_html, opt_generate_png)
        shutil.rmtree(temp_dir)
    else:
        print('Open the following url in a browser to view the results: ')
        print('  file://%s' % path_to_html)

def _check_output(args, **kwargs):
    with open(os.devnull) as f:
        kwargs['stderr'] = f
        return subprocess.check_output(args, **kwargs)

def parse_package_line(line):
    """The line looks like this:
    package: name='com.facebook.testing.tests' versionCode='1' versionName=''"""

    for word in line.split():
        if word.startswith("name='"):
            return word[len("name='"):-1]

def get_aapt_bin():
    """Find the binary for aapt from $ANDROID_SDK"""
    android_sdk = os.environ.get('ANDROID_SDK') or os.environ.get('ANDROID_HOME')
    build_tools = os.path.join(android_sdk, 'build-tools')

    all = list(glob.glob(os.path.join(build_tools, '*/aapt')))
    bad = list(glob.glob(os.path.join(build_tools, 'android-*/aapt')))
    good = list(Set(all) - Set(bad))

    good.sort()
    bad.sort()

    if len(good) == 0:
        return bad[-1]

    return good[-1]

def get_package(apk):
    output = _check_output([get_aapt_bin(), 'dump', 'badging', apk], stderr=os.devnull)
    for line in output.split('\n'):
        if line.startswith('package:'):
            return parse_package_line(line)

def main(argv):
    try:
        opt_list, rest_args = getopt.gnu_getopt(
            argv[1:],
            "eds:",
            ["generate-png=", "filter-name-regex=", "apk"])
    except getopt.GetoptError, err:
        usage()
        return 2

    if len(rest_args) != 1:
        usage()
        return 2

    process = rest_args[0]  # something like com.facebook.places.tests

    opts = dict(opt_list)

    if "--apk" in opts:
        # treat process as an apk instead
        process = get_package(process)

    puller_args = []
    if "-e" in opts:
        puller_args.append("-e")

    if "-d" in opts:
        puller_args.append("-d")

    if "-s" in opts:
        puller_args += ["-s", opts["-s"]]

    return pull_screenshots(process,
                            filter_name_regex=opts.get('--filter-name-regex'),
                            opt_generate_png=opts.get('--generate-png'),
                            adb_puller=SimplePuller(puller_args))

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<|endoftext|>


<commit_before>from django.apps import apps
from django.conf import settings
from django.db import models
from django.db.models import Q
from django.utils.translation import get_language
from haystack import indexes, signals
from munigeo.models import AdministrativeDivisionType, Street

from .models.unit import CONTRACT_TYPES_TRANSLATED

ADMIN_DIV_TYPES = ("sub_district", "neighborhood", "postcode_area")

class DeleteOnlySignalProcessor(signals.BaseSignalProcessor):
    """
    Delete models from index automatically.

    Use the settings key DISABLE_HAYSTACK_SIGNAL_PROCESSOR to
    disable.
    """

    settings_key = "DISABLE_HAYSTACK_SIGNAL_PROCESSOR"

    def handle_delete(self, sender, instance, **kwargs):
        if not getattr(settings, self.settings_key, False):
            super().handle_delete(sender, instance, **kwargs)

    def setup(self):
        models.signals.post_delete.connect(self.handle_delete)
        # TODO: ?
        # Efficient would be going through all backends & collecting all models
        # being used, then hooking up signals only for those.

    def teardown(self):
        # Naive (listen to all model saves).
        models.signals.post_delete.disconnect(self.handle_delete)

class ServiceMapBaseIndex(indexes.SearchIndex, indexes.Indexable):
    text = indexes.CharField(document=True, use_template=True)
    name = indexes.CharField(model_attr="name", boost=1.125)
    name_sort = indexes.CharField(model_attr="name")
    autosuggest = indexes.EdgeNgramField(model_attr="name")
    autosuggest_exact = indexes.CharField(model_attr="name", boost=1.125)
    extra_searchwords = indexes.CharField()
    autosuggest_extra_searchwords = indexes.CharField()
    public = indexes.BooleanField()
    suggest = indexes.CharField()

    def __init__(self, *args, **kwargs):
        super(*args, **kwargs)
        self.model = None

    def get_model(self):
        return self.model

    def get_updated_field(self):
        return "last_modified_time"

    def prepare_public(self, obj):
        return True

    def get_extra_searchwords(self, obj):
        values = [
            category.name for category in obj.keywords.filter(language=get_language())
        ]
        try:
            if obj.address_zip:
                values.append(obj.address_zip)
        except AttributeError:
            pass
        try:
            if obj.contract_type is not None:
                self.get_contract_type(obj, values)
        except AttributeError:
            pass
        return values

    def prepare_extra_searchwords(self, obj):
        return " ".join(self.get_extra_searchwords(obj))

    def prepare_autosuggest_extra_searchwords(self, obj):
        return self.prepare_extra_searchwords(obj)

    def prepare_suggest(self, obj):
        return dict(name=None, service=[], location=[])

    def get_contract_type(self, obj, values):
        contract_type = CONTRACT_TYPES_TRANSLATED[obj.contract_type]
        if contract_type:
            values.append(contract_type[1])

class UnitIndex(ServiceMapBaseIndex):
    text = indexes.CharField(document=True, use_template=False)
    name = indexes.CharField(boost=1.125)
    autosuggest = indexes.EdgeNgramField()
    municipality = indexes.CharField(model_attr="municipality_id", null=True)
    services = indexes.MultiValueField()
    root_department = indexes.CharField(null=True)
    suggest = indexes.CharField()
    address = indexes.CharField(use_template=False)

    def read_queryset(self, using=None):
        return self.get_model().search_objects

    def __init__(self, *args, **kwargs):
        super(*args, **kwargs)
        self._street_name_cache = None
        self.model = apps.get_model(app_label="services", model_name="Unit")

    def prepare_public(self, obj):
        return obj.public

    def _get_street_name_cache(self):
        if self._street_name_cache is None:
            self._street_name_cache = set(Street.objects.values_list("name", flat=True))
        return self._street_name_cache

    def _word_matches_street_address(self, word):
        return word in self._get_street_name_cache()

    def prepare_name(self, obj):
        return obj.name

    def prepare_autosuggest(self, obj):
        return obj.name

    def prepare_text(self, obj):
        values = [
            obj.name,
            obj.service_names(),
            obj.service_keywords(),
            obj.highlight_names(),
        ]
        if obj.municipality:
            values.append(obj.municipality.name)
        if obj.department:
            values.append(obj.department.top_departments())
        if obj.address_zip:
            values.append(obj.address_zip)
        if obj.contract_type is not None:
            self.get_contract_type(obj, values)
        return " ".join(values)

    def prepare_services(self, obj):
        return [ow.id for ow in obj.services.all()]

    def prepare_address(self, obj):
        return obj.street_address

    def prepare_root_department(self, obj):
        if obj.root_department is not None:
            return str(obj.root_department.uuid)

    def prepare_suggest(self, obj):
        terms_by_field = {
            "service": list(set((s.name for s in obj.services.all()))),
            "keyword": self.get_extra_searchwords(obj),
            "location": [],
        }
        if obj.municipality:
            terms_by_field["location"].append(obj.municipality.name)

        combined = [term for terms in terms_by_field.values() for term in terms]
        name = obj.name
        combined.append(name)
        terms_by_field["name"] = name
        terms_by_field["combined"] = " ".join(combined)
        return terms_by_field

class ServiceIndex(ServiceMapBaseIndex):
    def __init__(self, *args, **kwargs):
        super(*args, **kwargs)
        self.model = apps.get_model(app_label="services", model_name="Service")

class ServiceNodeIndex(ServiceMapBaseIndex):
    def __init__(self, *args, **kwargs):
        super(*args, **kwargs)
        self.model = apps.get_model(app_label="services", model_name="ServiceNode")

    def index_queryset(self, using=None):
        manager = self.get_model().objects
        # Decision: exclude top level tree nodes (where the level is 0);
        # they are too broad and are not good results for full
        # text queries which are usually trying to be somewhat
        # specific.
        unique_ids = (
            # The query below ensures that duplicate servicenodes
            # are only indexed once. They are servicenodes which
            # have the exact same service reference.
            #
            # Note the empty order_by clause which prevents
            # default ordering from interfering with the grouping.
            manager.exclude(service_reference__isnull=True)
            .values("service_reference")
            .annotate(id=models.Min("id"))
            .values_list("id", flat=True)
            .order_by()
        )

        return manager.filter(
            Q(id__in=unique_ids)
            | Q(
                # We have to separately add all nodes with null references
                # and level > 0.
                Q(level__gt=0)
                & Q(service_reference__isnull=True)
            )
        )

class AddressIndex(indexes.SearchIndex, indexes.Indexable):
    text = indexes.CharField(use_template=False, document=True)
    address = indexes.CharField(use_template=False)
    number = indexes.CharField(use_template=False)
    autosuggest = indexes.EdgeNgramField(use_template=False)
    autosuggest_exact = indexes.CharField(use_template=False)
    public = indexes.BooleanField()

    def prepare_public(self, obj):
        return True

    def get_model(self):
        return apps.get_model("munigeo", "Address")

    def prepare_text(self, obj):
        return ""

    def prepare_number(self, obj):
        return obj.number

    def prepare_address(self, obj):
        number_end = ""
        letter = ""
        if obj.number_end:
            number_end = "-" + obj.number_end
        if obj.letter:
            letter = obj.letter
        return "{street} {number}{number_end}{letter} {municipality}".format(
            street=obj.street,
            number=obj.number,
            number_end=number_end,
            letter=letter,
            municipality=obj.street.municipality,
        )

    def prepare_autosuggest(self, obj):
        return self.prepare_address(obj)

    def prepare_autosuggest_exact(self, obj):
        return None

    def get_updated_field(self):
        return "modified_at"

class AdministrativeDivisionIndex(indexes.SearchIndex, indexes.Indexable):
    text = indexes.CharField(use_template=False, document=True)
    name = indexes.CharField(use_template=False)
    autosuggest = indexes.EdgeNgramField(use_template=False)
    autosuggest_exact = indexes.CharField(use_template=False)
    public = indexes.BooleanField()

    def prepare_public(self, obj):
        return True

    def get_model(self):
        return apps.get_model("munigeo", "AdministrativeDivision")

    def prepare_text(self, obj):
        return obj.name

    def prepare_name(self, obj):
        return obj.name

    def prepare_autosuggest(self, obj):
        return self.prepare_name(obj)

    def prepare_autosuggest_exact(self, obj):
        return None

    def get_updated_field(self):
        return "modified_at"

    @staticmethod
    def indexed_types():
        return AdministrativeDivisionType.objects.filter(type__in=ADMIN_DIV_TYPES)

    def index_queryset(self, using=None):
        return self.get_model().objects.filter(type__in=self.indexed_types())
<commit_msg>Explicitly call super init-method<commit_after>from django.apps import apps
from django.conf import settings
from django.db import models
from django.db.models import Q
from django.utils.translation import get_language
from haystack import indexes, signals
from munigeo.models import AdministrativeDivisionType, Street

from .models.unit import CONTRACT_TYPES_TRANSLATED

ADMIN_DIV_TYPES = ("sub_district", "neighborhood", "postcode_area")

class DeleteOnlySignalProcessor(signals.BaseSignalProcessor):
    """
    Delete models from index automatically.

    Use the settings key DISABLE_HAYSTACK_SIGNAL_PROCESSOR to
    disable.
    """

    settings_key = "DISABLE_HAYSTACK_SIGNAL_PROCESSOR"

    def handle_delete(self, sender, instance, **kwargs):
        if not getattr(settings, self.settings_key, False):
            super().handle_delete(sender, instance, **kwargs)

    def setup(self):
        models.signals.post_delete.connect(self.handle_delete)
        # TODO: ?
        # Efficient would be going through all backends & collecting all models
        # being used, then hooking up signals only for those.

    def teardown(self):
        # Naive (listen to all model saves).
        models.signals.post_delete.disconnect(self.handle_delete)

class ServiceMapBaseIndex(indexes.SearchIndex, indexes.Indexable):
    text = indexes.CharField(document=True, use_template=True)
    name = indexes.CharField(model_attr="name", boost=1.125)
    name_sort = indexes.CharField(model_attr="name")
    autosuggest = indexes.EdgeNgramField(model_attr="name")
    autosuggest_exact = indexes.CharField(model_attr="name", boost=1.125)
    extra_searchwords = indexes.CharField()
    autosuggest_extra_searchwords = indexes.CharField()
    public = indexes.BooleanField()
    suggest = indexes.CharField()

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = None

    def get_model(self):
        return self.model

    def get_updated_field(self):
        return "last_modified_time"

    def prepare_public(self, obj):
        return True

    def get_extra_searchwords(self, obj):
        values = [
            category.name for category in obj.keywords.filter(language=get_language())
        ]
        try:
            if obj.address_zip:
                values.append(obj.address_zip)
        except AttributeError:
            pass
        try:
            if obj.contract_type is not None:
                self.get_contract_type(obj, values)
        except AttributeError:
            pass
        return values

    def prepare_extra_searchwords(self, obj):
        return " ".join(self.get_extra_searchwords(obj))

    def prepare_autosuggest_extra_searchwords(self, obj):
        return self.prepare_extra_searchwords(obj)

    def prepare_suggest(self, obj):
        return dict(name=None, service=[], location=[])

    def get_contract_type(self, obj, values):
        contract_type = CONTRACT_TYPES_TRANSLATED[obj.contract_type]
        if contract_type:
            values.append(contract_type[1])

class UnitIndex(ServiceMapBaseIndex):
    text = indexes.CharField(document=True, use_template=False)
    name = indexes.CharField(boost=1.125)
    autosuggest = indexes.EdgeNgramField()
    municipality = indexes.CharField(model_attr="municipality_id", null=True)
    services = indexes.MultiValueField()
    root_department = indexes.CharField(null=True)
    suggest = indexes.CharField()
    address = indexes.CharField(use_template=False)

    def read_queryset(self, using=None):
        return self.get_model().search_objects

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._street_name_cache = None
        self.model = apps.get_model(app_label="services", model_name="Unit")

    def prepare_public(self, obj):
        return obj.public

    def _get_street_name_cache(self):
        if self._street_name_cache is None:
            self._street_name_cache = set(Street.objects.values_list("name", flat=True))
        return self._street_name_cache

    def _word_matches_street_address(self, word):
        return word in self._get_street_name_cache()

    def prepare_name(self, obj):
        return obj.name

    def prepare_autosuggest(self, obj):
        return obj.name

    def prepare_text(self, obj):
        values = [
            obj.name,
            obj.service_names(),
            obj.service_keywords(),
            obj.highlight_names(),
        ]
        if obj.municipality:
            values.append(obj.municipality.name)
        if obj.department:
            values.append(obj.department.top_departments())
        if obj.address_zip:
            values.append(obj.address_zip)
        if obj.contract_type is not None:
            self.get_contract_type(obj, values)
        return " ".join(values)

    def prepare_services(self, obj):
        return [ow.id for ow in obj.services.all()]

    def prepare_address(self, obj):
        return obj.street_address

    def prepare_root_department(self, obj):
        if obj.root_department is not None:
            return str(obj.root_department.uuid)

    def prepare_suggest(self, obj):
        terms_by_field = {
            "service": list(set((s.name for s in obj.services.all()))),
            "keyword": self.get_extra_searchwords(obj),
            "location": [],
        }
        if obj.municipality:
            terms_by_field["location"].append(obj.municipality.name)

        combined = [term for terms in terms_by_field.values() for term in terms]
        name = obj.name
        combined.append(name)
        terms_by_field["name"] = name
        terms_by_field["combined"] = " ".join(combined)
        return terms_by_field

class ServiceIndex(ServiceMapBaseIndex):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = apps.get_model(app_label="services", model_name="Service")

class ServiceNodeIndex(ServiceMapBaseIndex):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = apps.get_model(app_label="services", model_name="ServiceNode")

    def index_queryset(self, using=None):
        manager = self.get_model().objects
        # Decision: exclude top level tree nodes (where the level is 0);
        # they are too broad and are not good results for full
        # text queries which are usually trying to be somewhat
        # specific.
        unique_ids = (
            # The query below ensures that duplicate servicenodes
            # are only indexed once. They are servicenodes which
            # have the exact same service reference.
            #
            # Note the empty order_by clause which prevents
            # default ordering from interfering with the grouping.
            manager.exclude(service_reference__isnull=True)
            .values("service_reference")
            .annotate(id=models.Min("id"))
            .values_list("id", flat=True)
            .order_by()
        )

        return manager.filter(
            Q(id__in=unique_ids)
            | Q(
                # We have to separately add all nodes with null references
                # and level > 0.
                Q(level__gt=0)
                & Q(service_reference__isnull=True)
            )
        )

class AddressIndex(indexes.SearchIndex, indexes.Indexable):
    text = indexes.CharField(use_template=False, document=True)
    address = indexes.CharField(use_template=False)
    number = indexes.CharField(use_template=False)
    autosuggest = indexes.EdgeNgramField(use_template=False)
    autosuggest_exact = indexes.CharField(use_template=False)
    public = indexes.BooleanField()

    def prepare_public(self, obj):
        return True

    def get_model(self):
        return apps.get_model("munigeo", "Address")

    def prepare_text(self, obj):
        return ""

    def prepare_number(self, obj):
        return obj.number

    def prepare_address(self, obj):
        number_end = ""
        letter = ""
        if obj.number_end:
            number_end = "-" + obj.number_end
        if obj.letter:
            letter = obj.letter
        return "{street} {number}{number_end}{letter} {municipality}".format(
            street=obj.street,
            number=obj.number,
            number_end=number_end,
            letter=letter,
            municipality=obj.street.municipality,
        )

    def prepare_autosuggest(self, obj):
        return self.prepare_address(obj)

    def prepare_autosuggest_exact(self, obj):
        return None

    def get_updated_field(self):
        return "modified_at"

class AdministrativeDivisionIndex(indexes.SearchIndex, indexes.Indexable):
    text = indexes.CharField(use_template=False, document=True)
    name = indexes.CharField(use_template=False)
    autosuggest = indexes.EdgeNgramField(use_template=False)
    autosuggest_exact = indexes.CharField(use_template=False)
    public = indexes.BooleanField()

    def prepare_public(self, obj):
        return True

    def get_model(self):
        return apps.get_model("munigeo", "AdministrativeDivision")

    def prepare_text(self, obj):
        return obj.name

    def prepare_name(self, obj):
        return obj.name

    def prepare_autosuggest(self, obj):
        return self.prepare_name(obj)

    def prepare_autosuggest_exact(self, obj):
        return None

    def get_updated_field(self):
        return "modified_at"

    @staticmethod
    def indexed_types():
        return AdministrativeDivisionType.objects.filter(type__in=ADMIN_DIV_TYPES)

    def index_queryset(self, using=None):
        return self.get_model().objects.filter(type__in=self.indexed_types())
<|endoftext|>


<commit_before># This file is part of the GOsa project.
#
#  http://gosa-project.org
#
# Copyright:
#  (C) 2016 GONICUS GmbH, Germany, http://www.gonicus.de
#
# See the LICENSE file in the project's top-level directory for details.
import pytest
from tornado.concurrent import Future
from tornado.testing import AsyncTestCase, gen_test
from unittest import mock
from gosa.plugins.goto.client_service import *

class MQTTHandlerMock(mock.MagicMock):
    callback = None

    def set_subscription_callback(self, cb):
        self.callback = cb

    def simulate_message(self, topic, message):
        self.callback(topic, message)

class GotoClientServiceTestCase(AsyncTestCase):

    def setUp(self):
        super(GotoClientServiceTestCase, self).setUp()
        self.service = ClientService()
        # use the mocked handler instead of the real one
        self.service.mqtt = MQTTHandlerMock()
        self.service.serve()

    def __announce_client(self):
        e = EventMaker()
        netinfo = []
        more = []
        netinfo.append(
            e.NetworkDevice(
                e.Name("eth0"),
                e.IPAddress("127.0.0.1"),
                e.IPv6Address("::1/128"),
                e.MAC("00:00:00:00:01"),
                e.Netmask("255.255.255.0"),
                e.Broadcast("127.0.0.255")))
        more.append(e.NetworkInformation(*netinfo))

        info = e.Event(
            e.ClientAnnounce(
                e.Id("fake_client_uuid"),
                e.Name("Testclient"),
                *more))

        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))

    def __announce_client_caps(self):
        e = EventMaker()
        # Assemble capabilities
        more = []
        caps = []
        caps.append(
            e.ClientMethod(
                e.Name("fakeCommand"),
                e.Path("path"),
                e.Signature("signature"),
                e.Documentation("This is a fake command for testing purpose")))
        more.append(e.ClientCapabilities(*caps))

        info = e.Event(
            e.ClientSignature(
                e.Id('fake_client_uuid'),
                e.Name('Testclient'),
                *more))
        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))

    def __set_client_offline(self, uuid):
        utcnow = datetime.datetime.utcnow() - datetime.timedelta(seconds=1500)
        e = EventMaker()
        # set client offline
        with mock.patch("gosa.plugins.goto.client_service.datetime.datetime") as m:
            m.utcnow.return_value = utcnow
            # send a client ping to store the faked current time
            ping = e.Event(e.ClientPing(e.Id(uuid)))
            self.service.mqtt.simulate_message("net.example/client/%s" % uuid, etree.tostring(ping))

        # run garbage collection
        self.service._ClientService__gc()

    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemSetStatus")
    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemGetStatus", return_value="")
    def test_clients(self, mocked_get_status, mocked_set_status):
        assert self.service.getClients() == {}

        # announce fake client via MQTT
        self.__announce_client()

        clients = self.service.getClients()
        assert len(clients) == 1
        assert "fake_client_uuid" in clients
        assert clients["fake_client_uuid"]["name"] == "Testclient"

        net = self.service.getClientNetInfo('fake_client_uuid')
        assert "eth0" in net
        assert net["eth0"]['IPAddress'] == "127.0.0.1"

        assert self.service.getClientNetInfo('unknown_client_uuid') == []

        assert self.service.getClientMethods('unknown_client_uuid') == []
        assert self.service.getClientMethods('fake_client_uuid') == {}

        self.__announce_client_caps()
        methods = self.service.getClientMethods('fake_client_uuid')
        assert 'fakeCommand' in methods

        self.__set_client_offline('fake_client_uuid')
        assert self.service.getClientMethods('fake_client_uuid') == []

        # disconnect client
        e = EventMaker()
        offline = e.Event(e.ClientLeave(
            e.Id('fake_client_uuid')
        ))
        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(offline))
        assert self.service.getClientMethods('fake_client_uuid') == []

    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemSetStatus")
    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemGetStatus", return_value="")
    @gen_test
    def test_clientDispatch(self, mocked_get_status, mocked_set_status):
        # announce fake client via MQTT
        self.__announce_client()
        self.__announce_client_caps()

        with pytest.raises(JSONRPCException):
            yield self.service.clientDispatch('unknown_client', 'fakeCommand')

        with pytest.raises(JSONRPCException):
            yield self.service.clientDispatch('fake_client_uuid', 'unknown_command')

        with mock.patch("gosa.plugins.goto.client_service.MQTTServiceProxy") as m:
            future = Future()
            future.set_result("result")
            m.return_value.fakeCommand.return_value = future
            res = yield self.service.clientDispatch('fake_client_uuid', 'fakeCommand')
            assert res == "result"

        self.__set_client_offline('fake_client_uuid')
        with pytest.raises(JSONRPCException):
            yield self.service.clientDispatch('fake_client_uuid', 'fakeCommand')

    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemSetStatus")
    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemGetStatus", return_value="")
    def test_users(self, mocked_get_status, mocked_set_status):
        # setup users
        self.__announce_client()
        self.__announce_client_caps()

        # setup session
        e = EventMaker()
        info = e.Event(
            e.UserSession(
                e.Id('fake_client_uuid'),
                e.User(e.Name('tester'))))

        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))

        assert self.service.getUserSessions()[0] == 'fake_client_uuid'

        assert self.service.getUserSessions("unknown_client_uuid") == []
        assert self.service.getUserSessions("fake_client_uuid")[0] == 'tester'

        assert self.service.getUserClients("tester") == ['fake_client_uuid']
        assert self.service.getUserClients("unknown_user") == []

        # notify user
        with mock.patch.object(self.service, "clientDispatch") as m:
            self.service.notifyUser('unknown_user', 'Title', 'Message')
            assert not m.called

            self.service.notifyUser('tester', 'Title', 'Message')
            m.assert_called_with('fake_client_uuid', 'notify', 'tester', 'Title', 'Message', 10, 'dialog-information')
            m.reset_mock()

            self.service.notifyUser('tester', 'Title', 'Message', icon=None)
            m.assert_called_with('fake_client_uuid', 'notify', 'tester', 'Title', 'Message', 10, '_no_icon_')

            m.reset_mock()
            m.side_effect = Exception("test")
            self.service.notifyUser('tester', 'Title', 'Message')
            m.side_effect = None

            with mock.patch("gosa.plugins.goto.client_service.JsonRpcHandler.user_sessions_available", return_value=True):
                self.service.notifyUser('tester', 'Title', 'Message')
                assert self.service.mqtt.send_event.called

            m.reset_mock()
            self.service.mqtt.send_event.reset_mock()

            # all users
            self.service.notifyUser(None, 'Title', 'Message')
            m.assert_called_with('fake_client_uuid', 'notify_all', 'Title', 'Message', 10, 'dialog-information')

            with mock.patch("gosa.plugins.goto.client_service.JsonRpcHandler.user_sessions_available", return_value=True):
                self.service.notifyUser(None, 'Title', 'Message')
                assert self.service.mqtt.send_event.called

            m.reset_mock()
            m.side_effect = Exception("test")
            self.service.notifyUser(None, 'Title', 'Message')

        # remove the user
        info = e.Event(
            e.UserSession(
                e.Id('fake_client_uuid'), e.User()))
        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))
        assert len(self.service.getUserSessions("fake_client_uuid")) == 0

    @mock.patch("gosa.plugins.goto.client_service.ObjectProxy")
    def test_systemGetStatus(self, mocked_proxy):
        mocked_proxy.return_value.deviceStatus = "O"
        with mock.patch("gosa.plugins.goto.client_service.PluginRegistry.getInstance") as mocked_index:
            mocked_index.return_value.search.return_value = []

            with pytest.raises(ValueError):
                self.service.systemGetStatus('some_uuid')

            mocked_index.return_value.search.return_value = [{'_uuid': ["some_uuid"]}]

            assert self.service.systemGetStatus('some_uuid') == "O"

    @mock.patch("gosa.plugins.goto.client_service.ObjectProxy")
    def test_systemSetStatus(self, mocked_proxy):
        with mock.patch("gosa.plugins.goto.client_service.PluginRegistry.getInstance") as mocked_index:
            mocked_index.return_value.search.return_value = []

            with pytest.raises(ValueError):
                self.service.systemSetStatus('some_uuid', "+O")

            mocked_index.return_value.search.return_value = [{'_uuid': 'some_uuid'}]
            self.service.systemSetStatus('some_uuid', "-O")
            mocked_proxy.return_value.system_Online is False

            with pytest.raises(ValueError):
                self.service.systemSetStatus('some_uuid', "+X")

            self.service.systemSetStatus('some_uuid', "+U")
            mocked_proxy.return_value.status_UpdateInProgress is True

            self.service.systemSetStatus('some_uuid', "-U")
            mocked_proxy.return_value.status_UpdateInProgress is False

    @mock.patch("gosa.plugins.goto.client_service.ObjectProxy")
    def test_joinClient(self, mocked_proxy):

        with mock.patch("gosa.plugins.goto.client_service.PluginRegistry.getInstance") as mocked_reg:
            mocked_index = mocked_reg.return_value
            with pytest.raises(ValueError):
                self.service.joinClient('tester', 'wrong_uuid', '00:00:00:00:00:01')

            # device exists
            # mocked_index.search.return_value = [1]
            # with pytest.raises(GOtoException):
            #     self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01')

            # user not unique
            mocked_index.search.side_effect = [[], [1, 2]]
            with pytest.raises(GOtoException):
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01')

            mocked_index.search.side_effect = [[], [{'dn': 'fake-manager-dn'}]]
            self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01')
            assert mocked_proxy.return_value.commit.called

            mocked_proxy.reset_mock()
            mocked_index.search.side_effect = None

            # with info
            with pytest.raises(ValueError):
                # wrong value type
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': '&/(&'})

            with pytest.raises(ValueError):
                # wrong deviceType
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': 'devices',
                                                                                'deviceType': 'unknown'})

            mocked_index.search.side_effect = [[], [], [{'dn': 'fake-manager-dn'}]]
            with pytest.raises(ValueError):
                # wrong owner
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': 'devices',
                                                                                'deviceType': 'terminal',
                                                                                'owner': 'tester'})
            mocked_index.search.side_effect = [['tester-dn'], [], [{'dn': 'fake-manager-dn'}]]
            self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': 'devices',
                                                                                'deviceType': 'terminal',
                                                                                'owner': 'tester'})
            assert mocked_proxy.return_value.commit.called
            assert mocked_proxy.return_value.owner == "tester"
            assert mocked_proxy.return_value.deviceType == "terminal"
            assert mocked_proxy.return_value.ou == "devices"

    def test_listeners(self):
        callback1 = mock.MagicMock()
        callback2 = mock.MagicMock()
        self.service.register_listener('testMethod1', callback1)
        self.service.register_listener('testMethod2', callback2)

        self.service.unregister_listener('testMethod3', callback1)
        self.service.unregister_listener('testMethod2', callback1)
        self.service.unregister_listener('testMethod2', callback2)

        self.service.notify_listeners("id1", "testMethod1", "state1")
        self.service.notify_listeners("id2", "testMethod2", "state2")

        callback1.assert_called_with("id1", "testMethod1", "state1")
        assert not callback2.called

    def test_eventProcessor(self):
        with mock.patch.object(self.service.log, "debug") as m:
            self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", "{}")
            assert m.called

        e = EventMaker()
        event = e.Event(e.UnknownEvent())
        with mock.patch.object(self.service.log, "error") as m:
            self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(event))
            assert m.called
<commit_msg>mock configureUsers<commit_after># This file is part of the GOsa project.
#
#  http://gosa-project.org
#
# Copyright:
#  (C) 2016 GONICUS GmbH, Germany, http://www.gonicus.de
#
# See the LICENSE file in the project's top-level directory for details.
import pytest
from tornado.concurrent import Future
from tornado.testing import AsyncTestCase, gen_test
from unittest import mock
from gosa.plugins.goto.client_service import *

class MQTTHandlerMock(mock.MagicMock):
    callback = None

    def set_subscription_callback(self, cb):
        self.callback = cb

    def simulate_message(self, topic, message):
        self.callback(topic, message)

class GotoClientServiceTestCase(AsyncTestCase):

    def setUp(self):
        super(GotoClientServiceTestCase, self).setUp()
        self.service = ClientService()
        # use the mocked handler instead of the real one
        self.service.mqtt = MQTTHandlerMock()
        self.service.serve()

    def __announce_client(self):
        e = EventMaker()
        netinfo = []
        more = []
        netinfo.append(
            e.NetworkDevice(
                e.Name("eth0"),
                e.IPAddress("127.0.0.1"),
                e.IPv6Address("::1/128"),
                e.MAC("00:00:00:00:01"),
                e.Netmask("255.255.255.0"),
                e.Broadcast("127.0.0.255")))
        more.append(e.NetworkInformation(*netinfo))

        info = e.Event(
            e.ClientAnnounce(
                e.Id("fake_client_uuid"),
                e.Name("Testclient"),
                *more))

        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))

    def __announce_client_caps(self):
        e = EventMaker()
        # Assemble capabilities
        more = []
        caps = []
        caps.append(
            e.ClientMethod(
                e.Name("fakeCommand"),
                e.Path("path"),
                e.Signature("signature"),
                e.Documentation("This is a fake command for testing purpose")))
        more.append(e.ClientCapabilities(*caps))

        info = e.Event(
            e.ClientSignature(
                e.Id('fake_client_uuid'),
                e.Name('Testclient'),
                *more))
        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))

    def __set_client_offline(self, uuid):
        utcnow = datetime.datetime.utcnow() - datetime.timedelta(seconds=1500)
        e = EventMaker()
        # set client offline
        with mock.patch("gosa.plugins.goto.client_service.datetime.datetime") as m:
            m.utcnow.return_value = utcnow
            # send a client ping to store the faked current time
            ping = e.Event(e.ClientPing(e.Id(uuid)))
            self.service.mqtt.simulate_message("net.example/client/%s" % uuid, etree.tostring(ping))

        # run garbage collection
        self.service._ClientService__gc()

    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemSetStatus")
    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemGetStatus", return_value="")
    def test_clients(self, mocked_get_status, mocked_set_status):
        assert self.service.getClients() == {}

        # announce fake client via MQTT
        self.__announce_client()

        clients = self.service.getClients()
        assert len(clients) == 1
        assert "fake_client_uuid" in clients
        assert clients["fake_client_uuid"]["name"] == "Testclient"

        net = self.service.getClientNetInfo('fake_client_uuid')
        assert "eth0" in net
        assert net["eth0"]['IPAddress'] == "127.0.0.1"

        assert self.service.getClientNetInfo('unknown_client_uuid') == []

        assert self.service.getClientMethods('unknown_client_uuid') == []
        assert self.service.getClientMethods('fake_client_uuid') == {}

        self.__announce_client_caps()
        methods = self.service.getClientMethods('fake_client_uuid')
        assert 'fakeCommand' in methods

        self.__set_client_offline('fake_client_uuid')
        assert self.service.getClientMethods('fake_client_uuid') == []

        # disconnect client
        e = EventMaker()
        offline = e.Event(e.ClientLeave(
            e.Id('fake_client_uuid')
        ))
        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(offline))
        assert self.service.getClientMethods('fake_client_uuid') == []

    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemSetStatus")
    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemGetStatus", return_value="")
    @gen_test
    def test_clientDispatch(self, mocked_get_status, mocked_set_status):
        # announce fake client via MQTT
        self.__announce_client()
        self.__announce_client_caps()

        with pytest.raises(JSONRPCException):
            yield self.service.clientDispatch('unknown_client', 'fakeCommand')

        with pytest.raises(JSONRPCException):
            yield self.service.clientDispatch('fake_client_uuid', 'unknown_command')

        with mock.patch("gosa.plugins.goto.client_service.MQTTServiceProxy") as m:
            future = Future()
            future.set_result("result")
            m.return_value.fakeCommand.return_value = future
            res = yield self.service.clientDispatch('fake_client_uuid', 'fakeCommand')
            assert res == "result"

        self.__set_client_offline('fake_client_uuid')
        with pytest.raises(JSONRPCException):
            yield self.service.clientDispatch('fake_client_uuid', 'fakeCommand')

    @mock.patch("gosa.plugins.goto.client_service.ClientService.configureUsers")
    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemSetStatus")
    @mock.patch("gosa.plugins.goto.client_service.ClientService.systemGetStatus", return_value="")
    def test_users(self, mocked_get_status, mocked_set_status, mocked_configure):
        # setup users
        self.__announce_client()
        self.__announce_client_caps()

        # setup session
        e = EventMaker()
        info = e.Event(
            e.UserSession(
                e.Id('fake_client_uuid'),
                e.User(e.Name('tester'))))

        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))

        assert self.service.getUserSessions()[0] == 'fake_client_uuid'

        assert self.service.getUserSessions("unknown_client_uuid") == []
        assert self.service.getUserSessions("fake_client_uuid")[0] == 'tester'

        assert self.service.getUserClients("tester") == ['fake_client_uuid']
        assert self.service.getUserClients("unknown_user") == []
        mocked_configure.assert_called_with('fake_client_uuid', ["tester"])

        # notify user
        with mock.patch.object(self.service, "clientDispatch") as m:
            self.service.notifyUser('unknown_user', 'Title', 'Message')
            assert not m.called

            self.service.notifyUser('tester', 'Title', 'Message')
            m.assert_called_with('fake_client_uuid', 'notify', 'tester', 'Title', 'Message', 10, 'dialog-information')
            m.reset_mock()

            self.service.notifyUser('tester', 'Title', 'Message', icon=None)
            m.assert_called_with('fake_client_uuid', 'notify', 'tester', 'Title', 'Message', 10, '_no_icon_')

            m.reset_mock()
            m.side_effect = Exception("test")
            self.service.notifyUser('tester', 'Title', 'Message')
            m.side_effect = None

            with mock.patch("gosa.plugins.goto.client_service.JsonRpcHandler.user_sessions_available", return_value=True):
                self.service.notifyUser('tester', 'Title', 'Message')
                assert self.service.mqtt.send_event.called

            m.reset_mock()
            self.service.mqtt.send_event.reset_mock()

            # all users
            self.service.notifyUser(None, 'Title', 'Message')
            m.assert_called_with('fake_client_uuid', 'notify_all', 'Title', 'Message', 10, 'dialog-information')

            with mock.patch("gosa.plugins.goto.client_service.JsonRpcHandler.user_sessions_available", return_value=True):
                self.service.notifyUser(None, 'Title', 'Message')
                assert self.service.mqtt.send_event.called

            m.reset_mock()
            m.side_effect = Exception("test")
            self.service.notifyUser(None, 'Title', 'Message')

        # remove the user
        info = e.Event(
            e.UserSession(
                e.Id('fake_client_uuid'), e.User()))
        self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(info))
        assert len(self.service.getUserSessions("fake_client_uuid")) == 0

    @mock.patch("gosa.plugins.goto.client_service.ObjectProxy")
    def test_systemGetStatus(self, mocked_proxy):
        mocked_proxy.return_value.deviceStatus = "O"
        with mock.patch("gosa.plugins.goto.client_service.PluginRegistry.getInstance") as mocked_index:
            mocked_index.return_value.search.return_value = []

            with pytest.raises(ValueError):
                self.service.systemGetStatus('some_uuid')

            mocked_index.return_value.search.return_value = [{'_uuid': ["some_uuid"]}]

            assert self.service.systemGetStatus('some_uuid') == "O"

    @mock.patch("gosa.plugins.goto.client_service.ObjectProxy")
    def test_systemSetStatus(self, mocked_proxy):
        with mock.patch("gosa.plugins.goto.client_service.PluginRegistry.getInstance") as mocked_index:
            mocked_index.return_value.search.return_value = []

            with pytest.raises(ValueError):
                self.service.systemSetStatus('some_uuid', "+O")

            mocked_index.return_value.search.return_value = [{'_uuid': 'some_uuid'}]
            self.service.systemSetStatus('some_uuid', "-O")
            mocked_proxy.return_value.system_Online is False

            with pytest.raises(ValueError):
                self.service.systemSetStatus('some_uuid', "+X")

            self.service.systemSetStatus('some_uuid', "+U")
            mocked_proxy.return_value.status_UpdateInProgress is True

            self.service.systemSetStatus('some_uuid', "-U")
            mocked_proxy.return_value.status_UpdateInProgress is False

    @mock.patch("gosa.plugins.goto.client_service.ObjectProxy")
    def test_joinClient(self, mocked_proxy):

        with mock.patch("gosa.plugins.goto.client_service.PluginRegistry.getInstance") as mocked_reg:
            mocked_index = mocked_reg.return_value
            with pytest.raises(ValueError):
                self.service.joinClient('tester', 'wrong_uuid', '00:00:00:00:00:01')

            # device exists
            # mocked_index.search.return_value = [1]
            # with pytest.raises(GOtoException):
            #     self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01')

            # user not unique
            mocked_index.search.side_effect = [[], [1, 2]]
            with pytest.raises(GOtoException):
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01')

            mocked_index.search.side_effect = [[], [{'dn': 'fake-manager-dn'}]]
            self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01')
            assert mocked_proxy.return_value.commit.called

            mocked_proxy.reset_mock()
            mocked_index.search.side_effect = None

            # with info
            with pytest.raises(ValueError):
                # wrong value type
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': '&/(&'})

            with pytest.raises(ValueError):
                # wrong deviceType
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': 'devices',
                                                                                'deviceType': 'unknown'})

            mocked_index.search.side_effect = [[], [], [{'dn': 'fake-manager-dn'}]]
            with pytest.raises(ValueError):
                # wrong owner
                self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': 'devices',
                                                                                'deviceType': 'terminal',
                                                                                'owner': 'tester'})
            mocked_index.search.side_effect = [['tester-dn'], [], [{'dn': 'fake-manager-dn'}]]
            self.service.joinClient('tester', 'fff0c8ad-d26b-4b6d-8e8e-75e054614dd9', '00:00:00:00:00:01', info={'ou': 'devices',
                                                                                'deviceType': 'terminal',
                                                                                'owner': 'tester'})
            assert mocked_proxy.return_value.commit.called
            assert mocked_proxy.return_value.owner == "tester"
            assert mocked_proxy.return_value.deviceType == "terminal"
            assert mocked_proxy.return_value.ou == "devices"

    def test_listeners(self):
        callback1 = mock.MagicMock()
        callback2 = mock.MagicMock()
        self.service.register_listener('testMethod1', callback1)
        self.service.register_listener('testMethod2', callback2)

        self.service.unregister_listener('testMethod3', callback1)
        self.service.unregister_listener('testMethod2', callback1)
        self.service.unregister_listener('testMethod2', callback2)

        self.service.notify_listeners("id1", "testMethod1", "state1")
        self.service.notify_listeners("id2", "testMethod2", "state2")

        callback1.assert_called_with("id1", "testMethod1", "state1")
        assert not callback2.called

    def test_eventProcessor(self):
        with mock.patch.object(self.service.log, "debug") as m:
            self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", "{}")
            assert m.called

        e = EventMaker()
        event = e.Event(e.UnknownEvent())
        with mock.patch.object(self.service.log, "error") as m:
            self.service.mqtt.simulate_message("net.example/client/fake_client_uuid", etree.tostring(event))
            assert m.called
<|endoftext|>


<commit_before>from django import template
from django.contrib.auth.models import User
from django.utils.safestring import mark_safe
import re
import urllib

register = template.Library()

username_re = re.compile('@[0-9a-zA-Z]+')
hashtag_re = re.compile('#[^\s]+')

@register.filter
def buglise(s):
    s = unicode(s)

    usernames = set(User.objects.values_list('username', flat=True))
    def replace_username(match):
        username = match.group(0)[1:]
        if username.lower() == 'all':
            return '@all'
        if username in usernames:
            return '@%s' % (username, username)
        else:
            return '@' + username

    s = username_re.sub(replace_username, s)

    s = hashtag_re.sub(
        lambda m: '%s' % (
            urllib.quote(m.group(0)),
            m.group(0),
        ),
        s
    )
    print s
    return mark_safe(s)

<commit_msg>Removed debugging<commit_after>from django import template
from django.contrib.auth.models import User
from django.utils.safestring import mark_safe
import re
import urllib

register = template.Library()

username_re = re.compile('@[0-9a-zA-Z]+')
hashtag_re = re.compile('#[^\s]+')

@register.filter
def buglise(s):
    s = unicode(s)

    usernames = set(User.objects.values_list('username', flat=True))
    def replace_username(match):
        username = match.group(0)[1:]
        if username.lower() == 'all':
            return '@all'
        if username in usernames:
            return '@%s' % (username, username)
        else:
            return '@' + username

    s = username_re.sub(replace_username, s)

    s = hashtag_re.sub(
        lambda m: '%s' % (
            urllib.quote(m.group(0)),
            m.group(0),
        ),
        s
    )
    return mark_safe(s)

<|endoftext|>