def execute(vm, protocol, args):
""" server side """
from AVMaster import vm_manager
#logging.debug(" CS Execute")
assert vm, "null vm"
tick = 15
if args:
assert isinstance(args, int), "you must specify an int for timeout."
timeout = args
off = False
logging.debug("%s, shutting down with timeout %s." % (vm,timeout))
vm_manager.execute(vm, "executeCmd","C:/Windows/System32/shutdown.exe",["/s", "/t", "30"], timeout, False, True)
for i in range(0, timeout, tick):
sleep(tick)
if vm_manager.execute(vm, "is_powered_off"):
return True, "Stopped VM"
logging.info("Forcing shutdown")
ret = vm_manager.execute(vm, "shutdown")
logging.debug("%s, shutdown returns: %s" % (vm, ret))
for i in range(10):
if vm_manager.execute(vm, "is_powered_off"):
return True, "Stopped VM"
sleep(tick)
return False, "Cannot stop VM"
def test_formatter(selfs):
from AVCommon import logger
logger.init()
from AVCommon.logger import logging
globals()['logging']=logging
formatter = logging.Formatter('%(asctime)s -%(levelname)s- %(filename)s:%(lineno)s %(message)s')
stdout_handler = logging.StreamHandler(sys.stdout)
stdout_handler.setFormatter(formatter)
logger = logging.getLogger('')
logger.addHandler(stdout_handler)
logger.setLevel(logging.DEBUG)
logging.debug('A debug message')
logging.info('Some information')
logging.warning('A shot across the bows')
logging.debug('A debug message')
logging.info('Some information')
logging.warning('A shot across the bows')
logging_child.ClassName("calling a child")
def test_formatter(selfs):
from AVCommon import logger
logger.init()
from AVCommon.logger import logging
globals()['logging'] = logging
formatter = logging.Formatter(
'%(asctime)s -%(levelname)s- %(filename)s:%(lineno)s %(message)s'
)
stdout_handler = logging.StreamHandler(sys.stdout)
stdout_handler.setFormatter(formatter)
logger = logging.getLogger('')
logger.addHandler(stdout_handler)
logger.setLevel(logging.DEBUG)
logging.debug('A debug message')
logging.info('Some information')
logging.warning('A shot across the bows')
logging.debug('A debug message')
logging.info('Some information')
logging.warning('A shot across the bows')
logging_child.ClassName("calling a child")
def test_execute():
# vmm = VMManager()
logging.info( "TEST VMManager")
#vms=["zenovm", "noav"]
vms = ["noav"]
vm = "noav"
vm_manager.execute(vm, "executeCmd", "c:/python27/python.exe", [], 40, True, False)
def mail_summary(mail_recipients, mail_server="mail.hackingteam.com"):
logging.info("sending mail to %s" % mail_recipients)
# Import smtplib for the actual sending function
import smtplib
# Import the email modules we'll need
from email.mime.text import MIMEText
report = Report()
# Open a plain text file for reading. For this example, assume that
# the text file contains only ASCII characters.
sum = summary()
# Create a text/plain message
for recipient in mail_recipients:
msg = MIMEText(sum)
# me == the sender's email address
# you == the recipient's email address
hostname = helper.get_hostname()
msg['Subject'] = '%s@%s: %s' % (report.name, hostname,
report.timestamp)
msg['From'] = "*****@*****.**"
msg['To'] = recipient
logging.debug(" msg to: %s" % msg['To'])
# Send the message via our own SMTP server, but don't include the
# envelope header.
s = smtplib.SMTP(mail_server)
s.sendmail(msg['From'], [msg['To']], msg.as_string())
s.quit()
def execute(vm, protocol, args):
from AVMaster import vm_manager
""" server side """
clean = True # VM IS NOT INFECTED!! TEST CAN CONTINUE!!!
#logging.debug(" CS Execute")
assert vm, "null vm"
invert = "STOP_IF_CLEAN" in args if args else False
#blacklist = ['BTHSAmpPalService','CyCpIo','CyHidWin','iSCTsysTray','quickset']
dirs = ['C:Users/avtest/AppData/Roaming/Microsoft/Windows/Start Menu/Programs/Startup',
'C:/Documents and Settings/avtest/Start Menu/Programs/Startup']
names = build.names[:]
names.remove("agent")
for d in dirs:
out = vm_manager.execute(vm, "listDirectoryInGuest", d)
#logging.debug("listDirectoryInGuest: %s" % out)
for b in names:
if b in out:
logging.info("%s, found %s in %s" % (vm, b, d))
clean = False
break
ret = xor(clean is True, invert)
if clean is True:
return ret, "VM is not infected"
else:
return ret, "VM is INFECTED"
def mail_summary(mail_recipients, mail_server = "mail.hackingteam.com"):
logging.info("sending mail to %s" % mail_recipients)
# Import smtplib for the actual sending function
import smtplib
# Import the email modules we'll need
from email.mime.text import MIMEText
report = Report()
# Open a plain text file for reading. For this example, assume that
# the text file contains only ASCII characters.
sum = summary()
# Create a text/plain message
for recipient in mail_recipients:
msg = MIMEText(sum)
# me == the sender's email address
# you == the recipient's email address
hostname = helper.get_hostname()
msg['Subject'] = '%s@%s: %s' % (report.name, hostname, report.timestamp)
msg['From'] = "*****@*****.**"
msg['To'] = recipient
logging.debug(" msg to: %s" % msg['To'])
# Send the message via our own SMTP server, but don't include the
# envelope header.
s = smtplib.SMTP(mail_server)
s.sendmail(msg['From'], [msg['To']], msg.as_string())
s.quit()
def execute(vm, protocol, args):
from AVMaster import vm_manager
""" server side """
clean = True # VM IS NOT INFECTED!! TEST CAN CONTINUE!!!
#logging.debug(" CS Execute")
assert vm, "null vm"
invert = "STOP_IF_CLEAN" in args if args else False
#blacklist = ['BTHSAmpPalService','CyCpIo','CyHidWin','iSCTsysTray','quickset']
dirs = [
'C:Users/avtest/AppData/Roaming/Microsoft/Windows/Start Menu/Programs/Startup',
'C:/Documents and Settings/avtest/Start Menu/Programs/Startup'
]
names = build.names[:]
names.remove("agent")
for d in dirs:
out = vm_manager.execute(vm, "listDirectoryInGuest", d)
#logging.debug("listDirectoryInGuest: %s" % out)
for b in names:
if b in out:
logging.info("%s, found %s in %s" % (vm, b, d))
clean = False
break
ret = xor(clean is True, invert)
if clean is True:
return ret, "VM is not infected"
else:
return ret, "VM is INFECTED"
def test_execute():
# vmm = VMManager()
logging.info("TEST VMManager")
#vms=["zenovm", "noav"]
vms = ["noav"]
vm = "noav"
vm_manager.execute(vm, "executeCmd", "c:/python27/python.exe", [], 40,
True, False)
def check_blacklist(blacklist=None):
with connection() as client:
logging.debug("connected")
blacklist_server = client.blacklist()
logging.info("blacklist from server: %s" % blacklist_server)
if blacklist:
logging.info("blacklist from conf: %s" % blacklist)
report_send("+ BLACKLIST: %s" % blacklist_server)
def start(self, p):
logging.debug("- START: %s" % p.vm)
self.mq.clean(p)
r = p.send_next_command()
c = p.last_command
report.sent(p.vm, c)
logging.info("- SENT: %s" % c)
def check_blacklist(blacklist=None):
with connection() as client:
logging.debug("connected")
blacklist_server = client.blacklist()
logging.info("blacklist from server: %s" % blacklist_server)
if blacklist:
logging.info("blacklist from conf: %s" % blacklist)
report_send("+ BLACKLIST: %s" % blacklist_server)
def start(self, p):
logging.debug("- START: %s" % p.vm)
self.mq.clean(p)
r = p.send_next_command()
c = p.last_command
report.sent(p.vm, c)
logging.info("- SENT: %s" % c)
def load_procedures(self):
if os.path.exists("AVMaster/conf/procedures.yaml"):
Procedure.load_from_file("AVMaster/conf/procedures.yaml")
confs = glob.glob("AVMaster/conf/procedures/*.yaml")
for conf in confs:
logging.info("Loading conf: %s" % conf)
Procedure.load_from_file(conf)
if not Procedure.check():
raise SyntaxError("Errors in procedures")
def check():
ret = True
called = set()
system = []
try:
for name, p in Procedure.procedures.items():
for c in p.command_list:
if c.name == "CALL":
call = c.args
called.add(call)
if call not in Procedure.procedures.keys():
logging.error(
"Error in procedure: %s, call to non existant proc: %s"
% (name, call))
ret = False
if c.name == "REPORT":
calls = c.args
for c in calls:
call = None
if isinstance(c, basestring):
call = c
elif isinstance(c, dict):
call = c.keys()[0]
else:
logging.error(
"Error in procedure: %s, call to non compliant proc: %s"
% (name, call))
if call:
called.add(call)
if call not in Procedure.procedures.keys():
logging.error(
"Error in procedure: %s, call to non existant proc: %s"
% (name, call))
ret = False
procs = set(Procedure.procedures.keys())
for p in called:
if p.startswith("SYSTEM_"):
logging.warn("system proc called: %s" % p)
for p in procs.difference(called):
good_start = ["TEST", "SYSTEM", "UPDATE"]
if not any([p.startswith(g) for g in good_start]):
logging.warn("probably unused PROC: %s" % p)
else:
system.append(p)
system.sort()
logging.info("Callable Procedures: %s" % system)
return ret
except:
logging.exception("Check")
def end(self, c):
logging.debug("- END: %s" % c)
self.ended.add(c)
if self.pool:
m = self.pool.pop(0)
logging.debug("pool popped: %s, remains: %s" % (m.vm, len(self.pool)))
self.start(m)
else:
logging.info("pool is empty")
report.Report().pool = [ p.vm for p in self.pool]
logging.debug("- pool: %s" % report.Report().pool)
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 end(self, c):
logging.debug("- END: %s" % c)
self.ended.add(c)
if self.pool:
m = self.pool.pop(0)
logging.debug("pool popped: %s, remains: %s" %
(m.vm, len(self.pool)))
self.start(m)
else:
logging.info("pool is empty")
report.Report().pool = [p.vm for p in self.pool]
logging.debug("- pool: %s" % report.Report().pool)
def check():
ret = True
called = set()
system = []
try:
for name,p in Procedure.procedures.items():
for c in p.command_list:
if c.name == "CALL":
call = c.args
called.add(call)
if call not in Procedure.procedures.keys():
logging.error("Error in procedure: %s, call to non existant proc: %s" % (name, call))
ret = False
if c.name == "REPORT":
calls = c.args
for c in calls:
call = None
if isinstance(c, basestring):
call = c
elif isinstance(c, dict):
call = c.keys()[0]
else:
logging.error("Error in procedure: %s, call to non compliant proc: %s" % (name, call))
if call:
called.add(call)
if call not in Procedure.procedures.keys():
logging.error("Error in procedure: %s, call to non existant proc: %s" % (name, call))
ret = False
procs = set(Procedure.procedures.keys())
for p in called:
if p.startswith("SYSTEM_"):
logging.warn("system proc called: %s" % p)
for p in procs.difference(called):
good_start=["TEST","SYSTEM","UPDATE"]
if not any([p.startswith(g) for g in good_start]) :
logging.warn("probably unused PROC: %s" % p)
else:
system.append(p)
system.sort()
logging.info("Callable Procedures: %s" % system)
return ret
except:
logging.exception("Check")
def main():
parser = argparse.ArgumentParser(description='AVMonitor master.')
parser.add_argument('-m', '--vm', required=False, default="",
help="Virtual Machines comma separated on which executing the operation")
parser.add_argument('-v', '--verbose', action='store_true', default=False,
help="Verbose")
parser.add_argument('-r', '--procedure', type=str, default=False, required=True,
help="Procedure to execute")
parser.add_argument('-p', '--pool', type=int, required=False, default=6,
help="This is the number of parallel process (default 6)")
parser.add_argument('-d', '--redis', default="localhost",
help="redis host")
parser.add_argument('-c', '--clean', default=False, action='store_true',
help="clean redis mq")
parser.add_argument('-s', '--session', default="dsession",
help="session redis mq ")
parser.add_argument('-e', '--report', type=str, default="")
parser.add_argument('-q', '--quiet', action='store_true', default=False,
help="Quiet")
args = parser.parse_args()
if args.report:
report = args.report
else:
report = time.strftime("%y%m%d", time.localtime(time.time()))
logger.init(report, quiet=args.quiet)
from AVCommon.logger import logging
globals()['logging']=logging
logging.debug(args)
from av_master import AVMaster
from AVCommon import package
logging.info("chdir %s" % package.basedir)
os.chdir(package.basedir)
master = AVMaster(args)
master.start()
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 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_up_and_down():
# vmm = VMManager("../AVMaster/conf/vms.cfg")
vmm = VMManager()
logging.info("TEST VMManager")
#vms=["zenovm", "noav"]
vms = ["noav"]
logging.info("Testing existent methods")
for vm in vms:
if not vmm.execute(vm, "is_powered_on"):
logging.debug("powering on %s" % vm)
vmm.execute(vm, "startup")
for vm in vms:
while not vmm.execute(vm, "is_powered_on"):
logging.debug("sleeping 5 secs waiting for avg")
sleep(5)
for vm in vms:
assert vmm.execute(vm, "is_powered_on")
for vm in vms:
vmm.execute(vm, "shutdown")
for vm in vms:
while vmm.execute(vm, "is_powered_on"):
logging.debug("sleeping 5 secs waiting for avg")
sleep(5)
for vm in vms:
assert vmm.execute(vm, "is_powered_off")
logging.info("Testing non existent methods")
exp = False
try:
for vm in vms:
vmm.execute(vm, "this_method_doesnt_exists")
except:
exp = True
finally:
assert exp is True
def test_up_and_down():
# vmm = VMManager("../AVMaster/conf/vms.cfg")
vmm = VMManager()
logging.info( "TEST VMManager")
#vms=["zenovm", "noav"]
vms = ["noav"]
logging.info( "Testing existent methods")
for vm in vms:
if not vmm.execute(vm, "is_powered_on"):
logging.debug("powering on %s" % vm)
vmm.execute(vm, "startup")
for vm in vms:
while not vmm.execute(vm, "is_powered_on") :
logging.debug( "sleeping 5 secs waiting for avg")
sleep(5)
for vm in vms:
assert vmm.execute(vm, "is_powered_on")
for vm in vms:
vmm.execute(vm, "shutdown")
for vm in vms:
while vmm.execute(vm, "is_powered_on"):
logging.debug( "sleeping 5 secs waiting for avg")
sleep(5)
for vm in vms:
assert vmm.execute(vm, "is_powered_off")
logging.info( "Testing non existent methods")
exp = False
try:
for vm in vms:
vmm.execute(vm, "this_method_doesnt_exists")
except:
exp = True
finally:
assert exp is True
def execute(vm, args):
from AVAgent import build
logging.debug(" CS CHECK_EMPTY_DIR: %s,%s" % (vm, str(args)))
assert isinstance(args, list)
dirs, whitelist = args
assert isinstance(dirs, list)
assert isinstance(whitelist, list)
wl = set(whitelist)
success = True
res = []
for dir in dirs:
if not os.path.exists(dir):
success |= True
#res.append("Not existent dir: %s" % dir)
logging.info("Not existent dir: %s" % dir)
else:
l = set(os.listdir(dir))
files_remained=l.difference(wl)
if l and wl and not files_remained:
logging.debug("all the files are whitelist: %s" % l)
if not files_remained:
success |= True
#res.append("Empty dir: %s" % dir)
logging.info("Empty dir: %s" % dir)
else:
success |= False
res.append("Non empty dir %s: %s" % (dir,files_remained))
logging.info("Non empty dir %s: %s" % (dir,files_remained))
logging.debug("CHECK_EMPTY: %s, %s" % (success, res))
return success, res
def execute(vm, args):
from AVAgent import build
logging.debug(" CS CHECK_EMPTY_DIR: %s,%s" % (vm, str(args)))
assert isinstance(args, list)
dirs, whitelist = args
assert isinstance(dirs, list)
assert isinstance(whitelist, list)
wl = set(whitelist)
success = True
res = []
for dir in dirs:
if not os.path.exists(dir):
success |= True
#res.append("Not existent dir: %s" % dir)
logging.info("Not existent dir: %s" % dir)
else:
l = set(os.listdir(dir))
files_remained = l.difference(wl)
if l and wl and not files_remained:
logging.debug("all the files are whitelist: %s" % l)
if not files_remained:
success |= True
#res.append("Empty dir: %s" % dir)
logging.info("Empty dir: %s" % dir)
else:
success |= False
res.append("Non empty dir %s: %s" % (dir, files_remained))
logging.info("Non empty dir %s: %s" % (dir, files_remained))
logging.debug("CHECK_EMPTY: %s, %s" % (success, res))
return success, res
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"])
parser.add_argument('-f', '--force', action='store_true', default=False,
help="force ")
args = parser.parse_args()
report = time.strftime("%y%m%d", time.localtime(time.time()))
logger.init(".", "avagent.log")
from AVCommon.logger import logging
globals()['logging']=logging
logging.debug(args)
mq = None
procedure = None
if args.procedure and args.procedure_file:
logging.info("Procedure %s" % args.procedure)
logging.debug("pwd: %s" % os.getcwd())
path = os.getcwd()
procs = Procedure.load_from_file(args.procedure_file)
logging.debug("%s" % procs)
procedure = procs[args.procedure]
mq = MQFeedProcedure()
Procedure.check()
try:
avagent = AVAgent(args.vm, args.redis, args.session)
avagent.start_agent(mq, procedure, force=args.force)
except:
logging.exception("FATAL")
def on_answer(vm, success, answer):
""" server side """
if isinstance(answer, list) and len(answer) > 0:
logging.info("BUILD ANSWER LAST: %s" % answer[-1])
else:
logging.info("BUILD ANSWER: %s" % str(answer))
f = getattr(vmrun, cmd)
if not args:
return f(vm)
else:
return f(vm, *args)
elif cmd in vmachine_cmds:
f = getattr(vm, cmd)
if not args:
return f()
else:
return f(args)
else:
logging.error("command not found: %s" % cmd)
raise Exception("Command not found")
except AssertionError as ae:
logging.error("Assertion found: %s" % ae)
raise
except Exception as e:
logging.error("Exception found. %s" % e)
raise
if __name__ == '__main__':
logging.debug("args: %s" % str(sys.argv[1:]))
t = tuple(sys.argv[1:])
ret = execute(*t)
logging.info(ret)
def main():
parser = argparse.ArgumentParser(description='AVMonitor master.')
parser.add_argument(
'-m',
'--vm',
required=False,
default="",
help="Virtual Machines comma separated on which executing the operation"
)
parser.add_argument('-v',
'--verbose',
action='store_true',
default=False,
help="Verbose")
parser.add_argument('-r',
'--procedure',
type=str,
default=False,
required=True,
help="Procedure to execute")
parser.add_argument(
'-p',
'--pool',
type=int,
required=False,
default=6,
help="This is the number of parallel process (default 6)")
parser.add_argument('-d',
'--redis',
default="localhost",
help="redis host")
parser.add_argument('-c',
'--clean',
default=False,
action='store_true',
help="clean redis mq")
parser.add_argument('-s',
'--session',
default="dsession",
help="session redis mq ")
parser.add_argument('-e', '--report', type=str, default="")
parser.add_argument('-q',
'--quiet',
action='store_true',
default=False,
help="Quiet")
args = parser.parse_args()
if args.report:
report = args.report
else:
report = time.strftime("%y%m%d", time.localtime(time.time()))
logger.init(report, quiet=args.quiet)
from AVCommon.logger import logging
globals()['logging'] = logging
logging.debug(args)
from av_master import AVMaster
from AVCommon import package
logging.info("chdir %s" % package.basedir)
os.chdir(package.basedir)
master = AVMaster(args)
master.start()
def dispatch(self, procedure, pool=0 ):
global received
exit = False
command.context = {}
procedure.add_begin_end()
#logging.debug("- SERVER len(procedure): %s" % len(procedure))
self.num_commands = len(procedure)
report.init(procedure.name)
assert self.vms
assert self.vms[0], "please specify at least one VM"
logging.debug("self.vms: %s" % self.vms)
av_machines = OrderedDict()
p_id = 0
for vm in self.vms:
av_machines[vm] = Protocol(self, vm, procedure, id = p_id)
p_id += 1
if pool == 0:
pool = len(self.vms)
Protocol.pool = pool
self.pool_start(av_machines.values(), pool)
self.ended = set()
answered = 0
while not exit and len(self.ended) < len(self.vms):
rec = self.mq.receive_server(blocking=True, timeout=self.timeout)
if rec is not None:
c, msg = rec
try:
command_unserialize = command.unserialize(msg)
except:
logging.exception("cannot unserialize: %s" % msg)
#exit = True
continue
#command_unserialize =
logging.info("- RECEIVED %s, %s" % (c, red(command_unserialize)))
if c not in av_machines.keys():
logging.warn("A message for %s probably belongs to another test!" % c)
continue
p = av_machines[c]
try:
answer = p.receive_answer(c, command_unserialize)
except:
logging.exception("cannot receive: %s" % command_unserialize)
continue
report.received(c, command_unserialize)
if answer.success == None:
#logging.debug("- SERVER IGNORING")
continue
answered += 1
#logging.debug("- SERVER RECEIVED ANSWER: %s" % answer.success)
if answer.name == "END":
self.end(c)
logging.info("- RECEIVE END: %s, %s" % (c, self.ended))
logging.debug("self.ended: (%s/%s) %s" % (len(self.ended), len(self.vms), self.ended))
elif p.on_error != "DISABLED" and (answer.success or p.on_error == "CONTINUE"):
r = p.send_next_command()
cmd = p.last_command
report.sent(p.vm, cmd)
logging.info("- SENT: %s, %s" % (c, red(cmd)))
if not r:
logging.info("- SENDING ERROR, ENDING: %s" %c)
self.end(c)
logging.debug("self.ended: (%s/%s) %s" % (len(self.ended), len(self.vms), self.ended))
else:
# answer.success == False
# deve skippare fino al command: END_PROC
if p.on_error == "SKIP":
logging.debug("on_error == %s" % p.on_error)
r = p.send_next_call()
cmd = p.last_command
if cmd:
report.sent(p.vm, cmd)
else:
logging.info("- RECEIVE ERROR, ENDING: %s" %c)
self.end(c)
logging.debug("self.ended: (%s/%s) %s" % (len(self.ended), len(self.vms), self.ended))
elif p.on_error == "DISABLED":
logging.debug("on_error == DISABLED")
r = p.send_next_proc()
cmd = p.last_command
if cmd:
report.sent(p.vm, cmd)
else:
logging.info("- RECEIVE ERROR, ENDING: %s" %c)
self.end(c)
logging.debug("self.ended: (%s/%s) %s" % (len(self.ended), len(self.vms), self.ended))
else:
assert p.on_error == "STOP"
logging.info("- RECEIVE ERROR, STOP: %s" %c)
self.end(c)
logging.debug("self.ended: (%s/%s) %s" % (len(self.ended), len(self.vms), self.ended))
else:
logging.info("- SERVER RECEIVED empty")
exit = True
report.finish()
logging.debug("answered: %s, ended: %s, num_commands: %s" % ( answered, len(self.ended), self.num_commands))
assert len(self.ended) == len(self.vms), "answered: %s, ended: %s, num_commands: %s" % ( answered, len(self.ended), len(self.vms))
#assert answered >= (len(self.vms) * (self.num_commands)), "answered: %s, len(vms): %s, num_commands: %s" % (answered , len(self.vms), self.num_commands)
return answered
def on_answer(vm, success, answer):
""" server side """
if isinstance(answer, list) and len(answer) > 0:
logging.info("BUILD ANSWER LAST: %s" % answer[-1])
else:
logging.info("BUILD ANSWER: %s" % str(answer))
def dispatch(self, procedure, pool=0):
global received
exit = False
command.context = {}
procedure.add_begin_end()
#logging.debug("- SERVER len(procedure): %s" % len(procedure))
self.num_commands = len(procedure)
report.init(procedure.name)
assert self.vms
assert self.vms[0], "please specify at least one VM"
logging.debug("self.vms: %s" % self.vms)
av_machines = OrderedDict()
p_id = 0
for vm in self.vms:
av_machines[vm] = Protocol(self, vm, procedure, id=p_id)
p_id += 1
if pool == 0:
pool = len(self.vms)
Protocol.pool = pool
self.pool_start(av_machines.values(), pool)
self.ended = set()
answered = 0
while not exit and len(self.ended) < len(self.vms):
rec = self.mq.receive_server(blocking=True, timeout=self.timeout)
if rec is not None:
c, msg = rec
try:
command_unserialize = command.unserialize(msg)
except:
logging.exception("cannot unserialize: %s" % msg)
#exit = True
continue
#command_unserialize =
logging.info("- RECEIVED %s, %s" %
(c, red(command_unserialize)))
if c not in av_machines.keys():
logging.warn(
"A message for %s probably belongs to another test!" %
c)
continue
p = av_machines[c]
try:
answer = p.receive_answer(c, command_unserialize)
except:
logging.exception("cannot receive: %s" %
command_unserialize)
continue
report.received(c, command_unserialize)
if answer.success == None:
#logging.debug("- SERVER IGNORING")
continue
answered += 1
#logging.debug("- SERVER RECEIVED ANSWER: %s" % answer.success)
if answer.name == "END":
self.end(c)
logging.info("- RECEIVE END: %s, %s" % (c, self.ended))
logging.debug("self.ended: (%s/%s) %s" %
(len(self.ended), len(self.vms), self.ended))
elif p.on_error != "DISABLED" and (answer.success or p.on_error
== "CONTINUE"):
r = p.send_next_command()
cmd = p.last_command
report.sent(p.vm, cmd)
logging.info("- SENT: %s, %s" % (c, red(cmd)))
if not r:
logging.info("- SENDING ERROR, ENDING: %s" % c)
self.end(c)
logging.debug(
"self.ended: (%s/%s) %s" %
(len(self.ended), len(self.vms), self.ended))
else:
# answer.success == False
# deve skippare fino al command: END_PROC
if p.on_error == "SKIP":
logging.debug("on_error == %s" % p.on_error)
r = p.send_next_call()
cmd = p.last_command
if cmd:
report.sent(p.vm, cmd)
else:
logging.info("- RECEIVE ERROR, ENDING: %s" % c)
self.end(c)
logging.debug(
"self.ended: (%s/%s) %s" %
(len(self.ended), len(self.vms), self.ended))
elif p.on_error == "DISABLED":
logging.debug("on_error == DISABLED")
r = p.send_next_proc()
cmd = p.last_command
if cmd:
report.sent(p.vm, cmd)
else:
logging.info("- RECEIVE ERROR, ENDING: %s" % c)
self.end(c)
logging.debug(
"self.ended: (%s/%s) %s" %
(len(self.ended), len(self.vms), self.ended))
else:
assert p.on_error == "STOP"
logging.info("- RECEIVE ERROR, STOP: %s" % c)
self.end(c)
logging.debug(
"self.ended: (%s/%s) %s" %
(len(self.ended), len(self.vms), self.ended))
else:
logging.info("- SERVER RECEIVED empty")
exit = True
report.finish()
logging.debug("answered: %s, ended: %s, num_commands: %s" %
(answered, len(self.ended), self.num_commands))
assert len(self.ended) == len(
self.vms), "answered: %s, ended: %s, num_commands: %s" % (
answered, len(self.ended), len(self.vms))
#assert answered >= (len(self.vms) * (self.num_commands)), "answered: %s, len(vms): %s, num_commands: %s" % (answered , len(self.vms), self.num_commands)
return answered
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"])