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_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_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_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_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)