def test_MultipleMQ():
host = "localhost"
mq1 = MQStar(host)
mq2 = MQStar(host, session=mq1.session)
client, message = "c1", "HELLO"
mq1.send_server(client, message)
c, m = mq2.receive_server()
assert (c == client)
assert (eval(m) == message)
def test_avagent_pull():
host = "localhost"
vms = ["testvm_%d" % i for i in range(1)]
#command_client={ 'COMMAND_CLIENT': [{ 'SET': [ 'windows' 'whatever']}]}
procedure = """
TEST:
- START_AGENT
- SET:
- [backend, 192.168.100.201]
- [frontend, 172.20.100.204]
- [redis, 10.0.20.1]
- BUILD: [ pull, windows, silent]
- STOP_AGENT
"""
test = Procedure.load_from_yaml(procedure)
host = "localhost"
mq = MQStar(host)
mq.clean()
logging.debug("MQ session: %s" % mq.session)
#istanzia n client e manda delle procedure.
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
report = Report()
# dispatcher, inoltra e riceve i comandi della procedura test sulle vm
dispatcher = Dispatcher(mq, vms, report)
thread = threading.Thread(target=dispatcher.dispatch,
args=(test["TEST"], ))
thread.start()
#p = Process(target=dispatcher.dispatch, args=(test["TEST"],))
#p.start()
# i client vengono eseguiti asincronicamente e comunicano tramite redis al server
pool = Pool(len(vms))
r = pool.map_async(av_agent.start_agent,
((v, host, mq.session) for v in vms))
r.get(
) #notare che i results dei client non ci interessano, viaggia tutto su channel/command.
# chiusura del server
#p.join()
thread.join()
logging.debug(dispatcher.report)
logging.debug("sent: %s" % dispatcher.report.c_sent)
logging.debug("received: %s" % Report.c_received)
def test_report_meta():
yaml = """
P1:
- SLEEP: 1
P2:
- SLEEP
P3:
- SLEEP: 2
P4:
- SLEEP
TEST_REPORT:
- VM: [noav,zenoav]
- SET_SERVER:
mail_recipients: [[email protected]]
- REPORT:
- P1: ["AVtest", "MyCase"]
- P2: ["AVtest", "MyOtherCase", INVERSE ]
- P3
- P4
"""
procedures = Procedure.load_from_yaml(yaml)
vms = ["noav", "zenovm"]
# vms = ["noav"]
redis_host = "localhost"
mq = MQStar(redis_host)
mq.clean()
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
dispatcher = Dispatcher(mq, vms)
logging.info("STARTING TEST REPORT")
dispatcher.dispatch(procedures["TEST_REPORT"])
logging.info("STOPPING TEST REPORT")
r = report.Report()
# report.finish()
assert r
assert len(r.c_received) == len(vms)
assert len(r.c_sent) == len(vms)
assert len(r.current_procedure) == len(vms)
assert len(r.c_received) == len(vms), len(r.c_received)
for vm in vms:
assert vm in r.c_received
assert len(r.c_received[vm]) >= 18
def test_report_meta():
yaml = """
P1:
- SLEEP: 1
P2:
- SLEEP
P3:
- SLEEP: 2
P4:
- SLEEP
TEST_REPORT:
- VM: [noav,zenoav]
- SET_SERVER:
mail_recipients: [[email protected]]
- REPORT:
- P1: ["AVtest", "MyCase"]
- P2: ["AVtest", "MyOtherCase", INVERSE ]
- P3
- P4
"""
procedures = Procedure.load_from_yaml(yaml)
vms = ["noav", "zenovm"]
#vms = ["noav"]
redis_host = "localhost"
mq = MQStar(redis_host)
mq.clean()
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
dispatcher = Dispatcher(mq, vms)
logging.info("STARTING TEST REPORT")
dispatcher.dispatch(procedures["TEST_REPORT"])
logging.info("STOPPING TEST REPORT")
r = report.Report()
#report.finish()
assert r
assert len(r.c_received) == len(vms)
assert len(r.c_sent) == len(vms)
assert len(r.current_procedure) == len(vms)
assert len(r.c_received) == len(vms), len(r.c_received)
for vm in vms:
assert vm in r.c_received
assert len(r.c_received[vm]) >= 18
def test_avagent_pull():
host = "localhost"
vms = [ "testvm_%d" % i for i in range(1) ]
#command_client={ 'COMMAND_CLIENT': [{ 'SET': [ 'windows' 'whatever']}]}
procedure = """
TEST:
- START_AGENT
- SET:
- [backend, 192.168.100.201]
- [frontend, 172.20.100.204]
- [redis, 10.0.20.1]
- BUILD: [ pull, windows, silent]
- STOP_AGENT
"""
test = Procedure.load_from_yaml(procedure)
host = "localhost"
mq = MQStar(host)
mq.clean()
logging.debug("MQ session: %s" % mq.session)
#istanzia n client e manda delle procedure.
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
report= Report()
# dispatcher, inoltra e riceve i comandi della procedura test sulle vm
dispatcher = Dispatcher(mq, vms, report)
thread = threading.Thread(target=dispatcher.dispatch, args=(test["TEST"],))
thread.start()
#p = Process(target=dispatcher.dispatch, args=(test["TEST"],))
#p.start()
# i client vengono eseguiti asincronicamente e comunicano tramite redis al server
pool = Pool(len(vms))
r = pool.map_async(av_agent.start_agent, ( (v, host, mq.session) for v in vms) )
r.get() #notare che i results dei client non ci interessano, viaggia tutto su channel/command.
# chiusura del server
#p.join()
thread.join()
logging.debug(dispatcher.report)
logging.debug("sent: %s" % dispatcher.report.c_sent)
logging.debug("received: %s" % Report.c_received)
def test_dispatcher_server():
host = "localhost"
vms = ["noav", "zenovm"]
host = "localhost"
mq = MQStar(host)
mq.clean()
#istanzia n client e manda delle procedure.
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
dispatcher = Dispatcher(mq, vms)
test = Procedure("TEST", [("EVAL_SERVER", None, 'vm'), ("SLEEP", None, 10)])
dispatcher.dispatch(test)
def test_dispatcher_server():
host = "localhost"
vms = ["noav", "zenovm"]
host = "localhost"
mq = MQStar(host)
mq.clean()
#istanzia n client e manda delle procedure.
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
dispatcher = Dispatcher(mq, vms)
test = Procedure("TEST", [("EVAL_SERVER", None, 'vm'),
("SLEEP", None, 10)])
dispatcher.dispatch(test)
def test_avagent_create():
host = "localhost"
vms = [ "testvm_%d" % i for i in range(10) ]
test = Procedure("TEST", ["BEGIN", "START_AGENT", ("EVAL_CLIENT", None, 'self.vm'), "STOP_AGENT", "END"])
host = "localhost"
mq = MQStar(host)
mq.clean()
logging.debug("MQ session: %s" % mq.session)
agent = av_agent.AVAgent("test_1", session=mq.session)
assert agent
def notest_ProtocolEval():
host = "localhost"
mq = MQStar(host)
mq.clean()
c = "client1"
mq.add_client(c)
commands = [
"BEGIN", ("EVAL_SERVER", "dir()"), ("EVAL_SERVER", "locals()"),
("EVAL_SERVER", "__import__('os').getcwd()"), ("END", None, None)
]
procedure = Procedure("PROC", commands)
p = Protocol(mq, c, procedure)
while p.send_next_command():
logging.debug("sent command")
exit = False
while not exit:
rec = mq.receive_server(blocking=True, timeout=10)
if rec is not None:
logging.debug("- SERVER RECEIVED %s %s" % (rec, type(rec)))
c, msg = rec
answer = p.receive_answer(c, msg)
logging.debug("- SERVER RECEIVED ANSWER: ", answer.success)
if answer.name == "END" or not answer.success:
logging.debug("- SERVER RECEIVE END")
#if answer.success:
a = """('client1', ('EVAL_SERVER', True, {'self': <Command_EVAL_SERVER.Command_EVAL_SERVER object at 0x10931f810>, 'args': 'locals()'}))""" # p.send_next_command()
else:
logging.debug("- SERVER RECEIVED empty")
exit = True
def start(self):
self.load_procedures()
proc = Procedure.procedures[self.procedure]
assert proc, "cannot find the specified procedure: %s" % self.procedure
# command line vm list overrides procedures.yaml
if self.vm_names==[''] and proc.command_list and proc.command_list[0].name.startswith("VM"):
vm_command = proc.command_list.pop(0)
self.vm_names = vm_command.execute('server', None, vm_command.args)[1]
logging.info("VM override: %s" % self.vm_names)
assert self.vm_names, "No VM specified"
mq = MQStar(self.args.redis, self.args.session)
if self.args.clean:
logging.warn("cleaning mq")
mq.clean()
logging.info("mq session: %s" % mq.session)
dispatcher = Dispatcher(mq, self.vm_names)
dispatcher.dispatch(proc, pool = self.pool)
def test_MultipleMQ():
host = "localhost"
mq1 = MQStar(host)
mq2 = MQStar(host, session=mq1.session)
client, message = "c1", "HELLO"
mq1.send_server(client, message)
c, m = mq2.receive_server()
assert (c == client)
assert (eval(m) == message)
def test_blockingMQ():
global received
host = "localhost"
mq1 = MQStar(host)
mq2 = MQStar(host, session=mq1.session)
c = "client1"
mq1.add_client(c)
thread1 = threading.Thread(target=server, args=(mq1,))
thread1.start()
mq2.send_server(c, "WORKS")
mq2.send_server(c, "FINE TO THE")
time.sleep(1)
mq2.send_server(c, "STOP")
time.sleep(6)
print "RECEIVED: ", received
assert len(received) == 3, "len: %s" % len(received)
def test_ProtocolCall():
host = "localhost"
mq = MQStar(host)
mq.clean()
c = "client1"
mq.add_client(c)
yaml = """
HELLO:
- EVAL_SERVER: 3+2
BASIC:
- EVAL_SERVER: dir()
- CALL: HELLO
CALLER:
- BEGIN
- CALL: BASIC
- EVAL_SERVER: locals()
- END
"""
procedures = Procedure.load_from_yaml(yaml)
caller = Procedure.procedures["CALLER"]
basic = Procedure.procedures["BASIC"]
class D:
pass
d = D()
d.mq = mq
p = Protocol(d, c, caller)
while p.send_next_command():
logging.debug("sent command")
exit = False
answers =0
while not exit:
rec = mq.receive_server(blocking=True, timeout=10)
if rec is not None:
logging.debug("- SERVER RECEIVED %s %s" % (rec, type(rec)))
c, msg = rec
command_unserialize = command.unserialize(msg)
answer = p.receive_answer(c, command_unserialize)
logging.debug("- SERVER RECEIVED ANSWER: %s" % answer.success)
if answer.success:
answers += 1
if answer.name == "END" or not answer.success:
logging.debug("- SERVER RECEIVE END")
#if answer.success:
exit = True
else:
logging.debug("- SERVER RECEIVED empty")
exit = True
assert answers == 7, "wrong answers: %s" % answers
def test_MQClean():
host = "localhost"
mq = MQStar(host)
redis = StrictRedis(host, socket_timeout=60)
clients = ["c1", "c2", "c3"]
mq.add_clients(clients)
mq.send_client("c1", "whatever")
rkeys = redis.keys("MQ_*")
assert rkeys
mq.clean()
rkeys = redis.keys("MQ_*")
assert not rkeys
def test_MQClean():
host = "localhost"
mq = MQStar(host)
redis = StrictRedis(host, socket_timeout=60)
clients = ["c1", "c2", "c3"]
mq.add_clients(clients)
mq.send_client("c1", "whatever")
rkeys = redis.keys("MQ_*")
assert rkeys
mq.clean()
rkeys = redis.keys("MQ_*")
assert not rkeys
def test_ProtocolEval():
host = "localhost"
mq = MQStar(host)
mq.clean()
c = "client1"
mq.add_client(c)
commands = ["BEGIN", ("EVAL_SERVER", None, "dir()"),
("EVAL_SERVER", None, "locals()"),
("EVAL_SERVER", None, "__import__('os').getcwd()"),
("END", None, None)]
procedure = Procedure("PROC", commands)
class D:
pass
d = D()
d.mq = mq
p = Protocol(d, c, procedure)
while p.send_next_command():
logging.debug("sent command: %s" % p.last_command)
print("---- START RECEIVING ----")
exit = False
while not exit:
rec = mq.receive_server(blocking=True, timeout=10)
if rec is not None:
logging.debug("- SERVER RECEIVED %s %s" % (rec, type(rec)))
c, msg = rec
command_unserialize = command.unserialize(msg)
answer = p.receive_answer(c, command_unserialize)
logging.debug("- SERVER RECEIVED ANSWER: %s" % answer.success)
if answer.name == "END" or not answer.success:
logging.debug("- SERVER RECEIVE END")
#if answer.success:
a = """('client1', ('EVAL_SERVER', True, {'self': <Command_EVAL_SERVER.Command_EVAL_SERVER object at 0x10931f810>, 'args': 'locals()'}))"""# p.send_next_command()
else:
logging.debug("- SERVER RECEIVED empty")
exit = True
print("---- STOP RECEIVING ----")
def start_agent(self, mq=None, procedure=None, force=False):
if not force and check_running(self.vm):
logging.fatal("already running")
exit = True
return False
class D:
pass
d = D()
if not mq:
mq = MQStar(self.host, self.session)
d.mq = mq
self.pc = Protocol(d, self.vm)
else:
assert procedure
d.mq = mq
self.pc = Protocol(d, self.vm, procedure=procedure)
mq.protocol = self.pc
logging.debug("mq: %s pc:%s" % (mq.protocol.procedure, self.pc.procedure))
mq.add_client(self.vm)
mq.notify_connection(self.vm)
logging.info("start receiving commands")
exit = False
while not exit:
logging.debug("- CLIENT %s LISTENING" % self.vm)
received = self.pc.receive_command()
logging.debug("- CLIENT %s EXECUTED: %s" % (self.vm, received))
if received.name == 'STOP_AGENT':
exit = True
logging.info("stop receiving commands")
remove_running(self.vm)
def test_MQ():
host = "localhost"
mq = MQStar(host)
mq.clean()
clients = ["c1", "c2", "c3"]
mq.add_clients(clients)
for c in clients:
mq.send_server(c, "STARTED")
for i in range(len(clients)):
c, m = mq.receive_server()
assert c in clients
assert eval(m) == "STARTED", "Uncorrect value: %s" % m
mq.send_client(c, "END %s" % i)
for c in clients:
m = mq.receive_client(c)
print m
assert (m.startswith("END "))
def test_vm_commands():
yaml = """
TEST1:
- START_VM
TEST2:
- EXECUTE_VM: c:\\users\\avtest\\desktop\\pubsub\\started.bat
- PUSH:
- [/tmp/gggg]
- c:\\users\\avtest\\desktop
- SCREENSHOT: /tmp/maggic_path.png
TEST3:
- PUSH:
- [gggg, jojojo]
- /tmp
- c:\\users\\avtest\\desktop
- PUSH:
- [AVAgent/av_agent.py, AVAgent/build.py, AVAgent/package.py, AVAgent/rcs_client.py,
AVCommon/commands/START_AGENT.py, AVCommon/commands/STOP_AGENT.py,
AVCommon/commands/BUILD.py, AVCommon/commands/GET.py, AVCommon/commands/SET.py]
- /home/olli/AVTest
- c:\\AVTest
- PULL:
- [gggg, jojojo]
- c:\\users\\avtest\\desktop
- /tmp/cpl
TEST4:
- START_VM
- SCREENSHOT: /tmp/magic_img_path.png
- STOP_VM
TEST5:
- PUSH:
- [AVCommon/commands/client/*.py]
- /home/olli/AVTest
- C:\\AVTest
- PUSH:
- [AVAgent/*.py]
- /home/olli/AVTest
- C:\\AVTest
UPLOAD_AGENT:
- PUSH:
- [AVAgent/av_agent.py, AVAgent/build.py, AVAgent/package.py, AVAgent/rcs_client.py, AVCommon/commands/*.py]
- /home/olli/AVTest
- c:\\AVTest
UPDATE:
- REVERT
- START_VM
- SLEEP: 180
- CALL: UPLOAD_AGENT
- INTERNET: True
- SLEEP: 120
- INTERNET: False
- STOP_VM
- START_VM
- SLEEP: 180
- STOP_VM
- REFRESH_SNAPSHOT
ZLEEP:
- SLEEP: 120
T_IS:
- CHECK_INFECTION
- SLEEP:
- 10
- 20
- CHECK_SHUTDOWN
- SLEEP: 5
TEST_INTERNET:
- INTERNET: True
- SLEEP: 15
- INTERNET: False
TEST_DIR:
- PUSH:
- [gigi/gggg]
- /tmp
- C:/Users/avtest/Desktop/gigi
- SLEEP: 10
- DELETE_DIR: C:/Users/avtest/Desktop/gigi
TEST_DIR_KO:
- DELETE_DIR: C:/Users/avtest/Desktop/gigiol
TEST_STOP1:
- STOP_VM: 60
TEST_STOP:
- STOP_VM
"""
procedures = Procedure.load_from_yaml(yaml)
#vms = ["noav", "zenovm"]
vms = ["zenoav"]
redis_host = "localhost"
mq = MQStar(redis_host)
mq.clean()
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
dispatcher = Dispatcher(mq, vms)
'''
logging.info("STARTING TEST 1")
dispatcher.dispatch(procedures["TEST1"])
import time
time.sleep(200)
logging.info("STARTING TEST 2")
dispatcher.dispatch(procedures["TEST2"])
time.sleep(30)
dispatcher.dispatch(procedures["TEST3"])
time.sleep(30)
logging.info("STARTING TEST UPDATE PROCEDURE")
dispatcher.dispatch(procedures["UPDATE"])
'''
logging.info("STARTING TEST 5")
dispatcher.dispatch(procedures["TEST_STOP"])
def test_blockingMQ():
global received
host = "localhost"
mq1 = MQStar(host)
mq2 = MQStar(host, session=mq1.session)
c = "client1"
mq1.add_client(c)
thread1 = threading.Thread(target=server, args=(mq1, ))
thread1.start()
mq2.send_server(c, "WORKS")
mq2.send_server(c, "FINE TO THE")
time.sleep(1)
mq2.send_server(c, "STOP")
time.sleep(6)
print "RECEIVED: ", received
assert len(received) == 3, "len: %s" % len(received)
def test_vm_commands():
yaml = """
TEST1:
- START_VM
TEST2:
- EXECUTE_VM: c:\\users\\avtest\\desktop\\pubsub\\started.bat
- PUSH:
- [/tmp/gggg]
- c:\\users\\avtest\\desktop
- SCREENSHOT: /tmp/maggic_path.png
TEST3:
- PUSH:
- [gggg, jojojo]
- /tmp
- c:\\users\\avtest\\desktop
- PUSH:
- [AVAgent/av_agent.py, AVAgent/build.py, AVAgent/package.py, AVAgent/rcs_client.py,
AVCommon/commands/START_AGENT.py, AVCommon/commands/STOP_AGENT.py,
AVCommon/commands/BUILD.py, AVCommon/commands/GET.py, AVCommon/commands/SET.py]
- /home/olli/AVTest
- c:\\AVTest
- PULL:
- [gggg, jojojo]
- c:\\users\\avtest\\desktop
- /tmp/cpl
TEST4:
- START_VM
- SCREENSHOT: /tmp/magic_img_path.png
- STOP_VM
TEST5:
- PUSH:
- [AVCommon/commands/client/*.py]
- /home/olli/AVTest
- C:\\AVTest
- PUSH:
- [AVAgent/*.py]
- /home/olli/AVTest
- C:\\AVTest
UPLOAD_AGENT:
- PUSH:
- [AVAgent/av_agent.py, AVAgent/build.py, AVAgent/package.py, AVAgent/rcs_client.py, AVCommon/commands/*.py]
- /home/olli/AVTest
- c:\\AVTest
UPDATE:
- REVERT
- START_VM
- SLEEP: 180
- CALL: UPLOAD_AGENT
- INTERNET: True
- SLEEP: 120
- INTERNET: False
- STOP_VM
- START_VM
- SLEEP: 180
- STOP_VM
- REFRESH_SNAPSHOT
ZLEEP:
- SLEEP: 120
T_IS:
- CHECK_INFECTION
- SLEEP:
- 10
- 20
- CHECK_SHUTDOWN
- SLEEP: 5
TEST_INTERNET:
- INTERNET: True
- SLEEP: 15
- INTERNET: False
TEST_DIR:
- PUSH:
- [gigi/gggg]
- /tmp
- C:/Users/avtest/Desktop/gigi
- SLEEP: 10
- DELETE_DIR: C:/Users/avtest/Desktop/gigi
TEST_DIR_KO:
- DELETE_DIR: C:/Users/avtest/Desktop/gigiol
TEST_STOP1:
- STOP_VM: 60
TEST_STOP:
- STOP_VM
"""
procedures = Procedure.load_from_yaml(yaml)
#vms = ["noav", "zenovm"]
vms = ["zenoav"]
redis_host = "localhost"
mq = MQStar(redis_host)
mq.clean()
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
dispatcher = Dispatcher(mq, vms)
'''
logging.info("STARTING TEST 1")
dispatcher.dispatch(procedures["TEST1"])
import time
time.sleep(200)
logging.info("STARTING TEST 2")
dispatcher.dispatch(procedures["TEST2"])
time.sleep(30)
dispatcher.dispatch(procedures["TEST3"])
time.sleep(30)
logging.info("STARTING TEST UPDATE PROCEDURE")
dispatcher.dispatch(procedures["UPDATE"])
'''
logging.info("STARTING TEST 5")
dispatcher.dispatch(procedures["TEST_STOP"])
def test_avagent_get_set():
host = "localhost"
vms = [ "testvm_%d" % i for i in range(100) ]
#command_client={ 'COMMAND_CLIENT': [{ 'SET': [ 'windows' 'whatever']}]}
procedure = """
TEST:
- START_AGENT
- SET: {pippo: franco}
- SET:
backend: 192.168.100.201
frontend: 172.20.100.204
redis: 10.0.20.1
- SET:
android:
binary: {admin: false, demo: true}
melt: {}
platform: android
sign: {}
- GET: pippo
- STOP_AGENT
"""
test = Procedure.load_from_yaml(procedure)
host = "localhost"
mq = MQStar(host)
mq.clean()
logging.debug("MQ session: %s" % mq.session)
#istanzia n client e manda delle procedure.
vm_manager.vm_conf_file = "../AVMaster/conf/vms.cfg"
# dispatcher, inoltra e riceve i comandi della procedura test sulle vm
dispatcher = Dispatcher(mq, vms, timeout = 10)
thread = threading.Thread(target=dispatcher.dispatch, args=(test["TEST"],))
thread.start()
#p = Process(target=dispatcher.dispatch, args=(test["TEST"],))
#p.start()
# i client vengono eseguiti asincronicamente e comunicano tramite redis al server
#pool = Pool(len(vms))
#r = pool.map_async(av_agent.start_agent, ( (v, host, mq.session) for v in vms) )
#r.get() #notare che i results dei client non ci interessano, viaggia tutto su channel/command.
for v in vms:
t = threading.Thread(target=av_agent.start_agent_args, args=(v, host, mq.session, True) )
t.start()
#p = Process(target=av_agent.start_agent, args=( tuple([v, host, mq.session,])))
#p.start()
# chiusura del server
#p.join()
thread.join()
r = report.Report()
logging.debug("sent: %s" % r.c_sent)
logging.debug("received: %s" % r.c_received)
def test_MQ():
host = "localhost"
mq = MQStar(host)
mq.clean()
clients = ["c1", "c2", "c3"]
mq.add_clients(clients)
for c in clients:
mq.send_server(c, "STARTED")
for i in range(len(clients)):
c, m = mq.receive_server()
assert c in clients
assert eval(m) == "STARTED", "Uncorrect value: %s" % m
mq.send_client(c, "END %s" % i)
for c in clients:
m = mq.receive_client(c)
print m
assert (m.startswith("END "))
