def main():
parser = parse_option()
if len(sys.argv) < 2:
logger.error("WatchAD must run with an action.")
parser.print_help()
sys.exit(1)
options, args = parser.parse_args()
if options.install:
if not options.domain or not options.server or not options.username or not options.password:
logger.error(
"WatchAD install action must provide domain, server, user and password params."
)
sys.exit(1)
install(domain=options.domain,
server=options.server,
user=options.username,
password=options.password)
elif options.check:
check()
elif options.start:
start()
elif options.restart:
restart()
elif options.stop:
stop()
elif options.status:
status()
def check_mq_connection() -> bool:
c = Consumer()
if not c.check_connection():
logger.error("Can't connect to the MQ, please reconfirm the settings.")
return False
logger.info("Connect to the MQ successfully, OK.")
return True
def stop():
logger.info("Stopping the WatchAD detect engine ...")
stop_rsp = subprocess.call(
"supervisorctl -c {root_dir}/supervisor.conf stop all".format(
root_dir=project_dir),
shell=True,
env={
"WATCHAD_ENGINE_DIR": project_dir,
"WATCHAD_ENGINE_NUM": "5"
})
if stop_rsp == 0:
logger.info("Stopped detection processes.")
else:
logger.error("Stop failed.")
shutdown_rsp = subprocess.call(
"supervisorctl -c {root_dir}/supervisor.conf shutdown".format(
root_dir=project_dir),
shell=True,
env={
"WATCHAD_ENGINE_DIR": project_dir,
"WATCHAD_ENGINE_NUM": "5"
})
if shutdown_rsp == 0:
logger.info("Shutdown WatchAD.")
else:
logger.error("Shutdown WatchAD failed.")
def wait_log_in_database(self, computer_name, record_number):
"""
因为消息队列和入库ES是分开进行的，所以可能会出现当消费到某条日志时，ES还没入库，所以需要检查同步
"""
count = 0
query = {
"query":
get_must_statement(
get_term_statement("computer_name", computer_name),
get_term_statement("record_number", record_number)),
"_source":
False,
"size":
1
}
while True:
try:
rsp = self.es.search(body=query,
index=ElasticConfig.event_log_index,
doc_type=ElasticConfig.event_log_doc_type,
request_timeout=100)
if rsp.get("error"):
logger.error(rsp.get("error").get("reason"))
break
if len(rsp["hits"]["hits"]) > 0:
return rsp["hits"]["hits"][0]["_id"]
time.sleep(2)
# 最多等5次，即 2 * 5 = 10秒
if count == 10:
break
count += 1
except Exception as e:
logger.error("es wait_log_in_database search error: " + str(e))
break
def check_mongo_connection() -> bool:
mongo = MongoHelper(MongoConfig.uri)
if not mongo.check_connection():
logger.error(
"Can't connect to the MongoDB, please reconfirm the settings.")
return False
logger.info("Connect to the MongoDB successfully, OK.")
return True
def start():
if not check():
sys.exit(-1)
logger.info("Starting the WatchAD detect engine ...")
rsp = subprocess.call("supervisord -c {root_dir}/supervisor.conf".format(root_dir=project_dir),
shell=True,
env={"WATCHAD_ENGINE_DIR": project_dir, "WATCHAD_ENGINE_NUM": "5"})
if rsp == 0:
logger.info("Started!")
else:
logger.error("Start failed.")
def multi_search(self, body, index, doc_type):
try:
rsp = self.es.msearch(body=body,
index=index,
doc_type=doc_type,
request_timeout=100)
if rsp.get("error"):
logger.error(rsp.get("error").get("reason"))
return
return rsp
except Exception as e:
logger.error("es msearch error: " + str(e))
def stop():
logger.info("Stopping the WatchAD detect engine ...")
rsp = subprocess.call(
"supervisorctl -c {root_dir}/supervisor.conf shutdown".format(
root_dir=project_dir),
shell=True,
env={
"WATCHAD_ENGINE_DIR": project_dir,
"ENV_WATCHAD_ENGINE_NUM": 5
})
if rsp == 0:
logger.info("Stopped!")
else:
logger.error("Stop failed.")
def check_es_template() -> bool:
"""
检查ES模板安装状态
"""
logger.info("Check the elasticsearch index template.")
es = ElasticHelper()
for name, temp in template_map.items():
if es.exists_template(name=name):
logger.info("template \"{name}\" ---> exist.".format(name=name))
else:
logger.info(
"template \"{name}\" ---> not exist.".format(name=name))
logger.error("Check the elasticsearch template fail.")
return False
logger.info("Check the elasticsearch template successfully, OK.")
return True
