def update_state(self, state_str=u"初始化"):
self.state["Data"]["Name"] = self.wvs_name
self.state["Data"]["Address"] = self.address
self.state["Data"]["State"] = state_str
self.state["Data"]["Timestamp"] = time.time()
common_msg.msg_wvs_state.data = self.state
msg_bus.send_msg(common_msg.msg_wvs_state)
def __agents_state_monitor(self):
while self.__running:
if len(self.agent_state_dict) > 0:
for agent_state in list(self.agent_state_dict.values()):
new_state = self.__check_state(agent_state)
if new_state == "Dead":
logger.info("Agent {0} is dead.\nAgent {1} is removed.".format(
agent_state.agent_identifier,
agent_state.agent_identifier))
agent_state.state = new_state
common_msg.msg_agent_state_update.data = agent_state.gen_json_object()
msg_bus.send_msg(common_msg.msg_agent_state_update)
self.agent_state_dict.pop(agent_state.agent_identifier)
else:
agent_state.state = new_state
self.agent_state_dict[agent_state.agent_identifier] = agent_state
common_msg.msg_agent_state_update.data = agent_state.gen_json_object()
msg_bus.send_msg(common_msg.msg_agent_state_update)
time.sleep(1)
time.sleep(STATE_UPDATE_INTERVAL)
def wvs_scan_control(data):
command = data.get("Command")
logger.info("Start scan with config: {}".format(command))
if GlobalVar.ServerState is "Stopped":
return
if command == "start":
setting = ScanSetting()
start_scan_cmd = {
"Type": "WVSCommand",
"Data": {
"Action": "StartNewScan",
"Config": { # 可选，当命令为StartNewScan时需提供该字段作为扫描参数
"StartURL": setting.start_url,
"ScanPolicy": setting.scan_policy
}
}
}
common_msg.msg_server_command.data = start_scan_cmd
msg_bus.send_msg(common_msg.msg_server_command)
GlobalVar.ScanState = "Started"
logger.info("Start scan with config: {}".format(start_scan_cmd))
elif command is "stop":
stop_scan_cmd = {
"Type": "WVSCommand",
"Data": {
"Action": "StopScan",
}
}
common_msg.msg_server_command.data = stop_scan_cmd
msg_bus.send_msg(common_msg.msg_server_command)
GlobalVar.ScanState = "Stopped"
logger.info("Stopped scan")
def __process_scan_result(self):
res_list = self.__gen_scan_result()
if res_list:
wvsstate.update_state(u"结果上报")
for res in res_list:
result = {"Type": "ScanResult", "Data": res}
common_msg.msg_scan_result_send.data = result
msg_bus.send_msg(common_msg.msg_scan_result_send)
self.__print_scan_result(res)
time.sleep(1)
else:
wvsstate.update_state(u"结果为空")
time.sleep(5)
wvsstate.update_state(u"任务结束")
def start_task(task_name):
session["TaskName"] = task_name
if task_name:
task_manager.get_task_info(task_name)
ScanResult().wvs_result_list.clear()
task_manager.save_task(task_name)
start_scan_cmd = {
"Type": "WVSCommand",
"Data": {
"Action": "StartNewScan",
"Config": { # 可选，当命令为StartNewScan时需提供该字段作为扫描参数
"StartURL": ScanSetting().start_url,
"ScanPolicy": ScanSetting().scan_policy
}
}
}
common_msg.msg_server_command.data = start_scan_cmd
msg_bus.send_msg(common_msg.msg_server_command)
logger.info("Start scan with config: {}".format(start_scan_cmd))
return redirect(url_for("monitor"))
def receive_data_handler(self, data, address):
if address not in self.agent_list:
self.agent_list.append(address)
pkg_obj = dict(json.loads(str(data, encoding='utf-8')))
try:
if pkg_obj["Type"] == "Heartbeat": #心跳包
common_msg.msg_heartbeat.data = pkg_obj["Data"]
msg_bus.send_msg(common_msg.msg_heartbeat)
# self.heartbeat_handle(data)
elif pkg_obj["Type"] == "WVSState":
common_msg.msg_wvs_state.data = pkg_obj["Data"]
msg_bus.send_msg(common_msg.msg_wvs_state)
logger.info(common_msg.msg_wvs_state.data)
# print("收到代理{}的漏扫状态数据".format(address))
elif pkg_obj["Type"] == "ScanResult":
common_msg.msg_scan_result_receive.data = pkg_obj["Data"]
msg_bus.send_msg(common_msg.msg_scan_result_receive)
self.__print_scan_result(pkg_obj["Data"])
else:
print("收到来自{}的未知类型数据".format(address))
except KeyError as e:
print("收到来自{}的未知类型数据——{}".format(address, data))
def recv_data_handler(self, data, address):
if self.cc_server is None:
self.cc_server = address
pkg_obj = dict(json.loads(str(data, encoding='utf-8')))
#针对Data中字段进行分析处理
try:
if pkg_obj["Type"] == "WVSCommand": #命令数据包
data = pkg_obj["Data"]
# logger.info(pkg_obj["Data"])
common_msg.msg_wvs_command.data = data
msg_bus.send_msg(common_msg.msg_wvs_command)
elif pkg_obj["Type"] == "AgentControl":
data = pkg_obj["Data"]
if data["Control"] == "Exit":
common_msg.msg_agent_exit.data = data
msg_bus.send_msg(common_msg.msg_agent_exit)
logger.info("Agent exit...")
else:
logger.info("收到来自{}的未知类型数据".format(address))
except KeyError as e:
logger.error("收到来自{}的未知类型数据——{}".format(address, data))
def __check_state(self, agent_state):
"""
根据前后两次时标对比判断代理状态，时标间隔大于2倍更新时间小于3倍更新时间时为离线，更长时间为Dead
:param agent_state: 代理状态对象
:return: 代理状态：Offline、Dead、Online
"""
last_time = time.time() - agent_state.timestamp
if STATE_UPDATE_INTERVAL * 2.0 < last_time <= STATE_UPDATE_INTERVAL * 5.0:
return "Offline"
elif last_time > STATE_UPDATE_INTERVAL * 5.0:
return "Dead"
else:
return "Online"
if __name__ == '__main__':
monitor = AgentStateMonitor()
print(id(monitor))
monitor.start_monitor()
common_msg.msg_heartbeat.data = {
"Name":"tste",
"Address": ("127.0.0.1", 5555),
"Timestamp": time.time(),
"State": "Online"
}
msg_bus.send_msg(common_msg.msg_heartbeat)
time.sleep(20)
