def sync_by_category(self, category):
dist_collection = 'asset_{}'.format(category)
for data in conn(category).find({"task_id": self.task_id}):
query = {"scope_id": self.scope_id, category: data[category]}
del data["_id"]
data["scope_id"] = self.scope_id
old = conn(dist_collection).find_one(query)
if old is None:
data["save_date"] = utils.curr_date_obj()
data["update_date"] = data["save_date"]
logger.info("sync {}, insert {} {} -> {}".format(
category, data[category], self.task_id, self.scope_id))
conn(dist_collection).insert_one(data)
if old and self.update_flag:
curr_date = utils.curr_date_obj()
data["save_date"] = old.get("save_date", curr_date)
data["update_date"] = curr_date
if category == 'ip':
if data.get("domain") and old.get("domain"):
old["domain"].extend(data["domain"])
data["domain"] = list(set(old["domain"]))
logger.info("sync {}, replace {} {} -> {}".format(
category, data[category], self.task_id, self.scope_id))
conn(dist_collection).find_one_and_replace(query, data)
def async_site_info(self, site_info_list):
"""
用来同步发现的 site 中的信息，仅仅在监控阶段使用
"""
new_site_info_list = []
for site_info in site_info_list:
curr_date_obj = utils.curr_date_obj()
query = {"site": site_info["site"], "scope_id": self.scope_id}
data = utils.conn_db('asset_site').find_one(query)
if data:
continue
new_site_info_list.append(site_info)
site_info["save_date"] = curr_date_obj
site_info["update_date"] = curr_date_obj
site_info["scope_id"] = self.scope_id
utils.conn_db('asset_site').insert_one(site_info)
new_asset_map = {
"site": new_site_info_list[:10],
"ip": self.ip_info_list[:10],
"task_name": self.task_name
}
new_asset_counter = {
"site": len(new_site_info_list),
"ip": len(self.ip_info_list)
}
if len(self.ip_info_list) > 0:
utils.message_push(asset_map=new_asset_map,
asset_counter=new_asset_counter)
def post(self):
"""
添加指纹信息
"""
args = self.parse_args(add_fingerprint_fields)
human_rule = args.pop('human_rule')
name = args.pop('name')
rule_map = parse_human_rule(human_rule)
if rule_map is None:
return utils.build_ret(ErrorMsg.RuleInvalid, {"rule": human_rule})
data = {
"name": name,
"rule": rule_map,
"human_rule": transform_rule_map(rule_map),
"update_date": utils.curr_date_obj()
}
utils.conn_db('fingerprint').insert_one(data)
finger_id = str(data.pop('_id'))
data.pop('update_date')
return utils.build_ret(ErrorMsg.Success, {
"_id": finger_id,
"data": data
})
def async_ip_info(self):
new_ip_info_list = []
for ip_info in self.ip_info_list:
curr_ip = ip_info["ip"]
curr_date_obj = utils.curr_date_obj()
# 新发现的IP ，直接入资产集合
if curr_ip not in self.asset_ip_info_map:
asset_ip_info = ip_info.copy()
asset_ip_info["scope_id"] = self.scope_id
asset_ip_info["domain"] = []
asset_ip_info["save_date"] = curr_date_obj
asset_ip_info["update_date"] = curr_date_obj
utils.conn_db('asset_ip').insert_one(asset_ip_info)
utils.conn_db('ip').insert_one(ip_info)
new_ip_info_list.append(ip_info)
continue
# 保存新发现的端口
new_port_info_list = []
for port_info in ip_info["port_info"]:
ip_port = "{}:{}".format(curr_ip, port_info["port_id"])
if ip_port in self.asset_ip_port_set:
continue
new_port_info_list.append(port_info)
if new_port_info_list:
asset_ip_info = self.asset_ip_info_map[curr_ip]
asset_ip_info["port_info"].extend(new_port_info_list)
update_info = dict()
update_info["update_date"] = utils.curr_date_obj()
update_info["port_info"] = asset_ip_info["port_info"]
query = {"_id": asset_ip_info["_id"]}
utils.conn_db('asset_ip').update_one(query,
{"$set": update_info})
# 只是保存新发现的端口
ip_info["port_info"] = new_port_info_list
utils.conn_db('ip').insert_one(ip_info)
new_ip_info_list.append(ip_info)
continue
self.ip_info_list = new_ip_info_list
logger.info("found new ip_info {}".format(len(self.ip_info_list)))
def add_site_to_scope(site, scope_id):
fetch_site_data = services.fetch_site([site])
web_analyze_data = services.web_analyze([site])
finger = web_analyze_data.get(site, [])
curr_date = utils.curr_date_obj()
if fetch_site_data:
item = fetch_site_data[0]
item["finger"] = finger
item["screenshot"] = ""
item["scope_id"] = scope_id
item["save_date"] = curr_date
item["update_date"] = curr_date
utils.conn_db('asset_site').insert_one(item)