def _init_db(self):
"""
初始化数据库
"""
model_prefix = self.target_host + "_" + str(self.target_port)
self.new_scan_model = NewRequestModel(model_prefix)
self.new_scan_model.reset_unscanned_item()
report_model = ReportModel(model_prefix)
Communicator().set_internal_shared("report_model", report_model)
def _init_db(self):
"""
初始化数据库
"""
model_prefix = self.target_host + "_" + str(self.target_port)
self.new_scan_model = NewRequestModel(
model_prefix, multiplexing_conn=True)
self.new_scan_model.reset_unscanned_item()
report_model = ReportModel(model_prefix, multiplexing_conn=True)
Communicator().set_internal_shared("report_model", report_model)
Communicator().set_internal_shared("failed_task_set", set())
def get_urls(self, host_port, page=0, status=0):
"""
获取指定状态的的url列表
Parameters:
page - int, 获取的页数，每页10条
status - int, url的状态 未扫描：0, 已扫描：1, 正在扫描：2, 扫描中出现错误: 3
Returns:
total, urls - total为数据总数, int类型，urls为已扫描的url, list类型, item形式为tuple (url对应id, url字符串)
"""
try:
model = NewRequestModel(host_port, create_table=False, multiplexing_conn=True)
except exceptions.TableNotExist as e:
raise e
return model.get_urls(page, status)
def clean_target(self, host, port, url_only=False):
"""
清空目标对应的数据库，同时重置预处理lru
Parameters:
host - str, 目标host
port - int, 目标port
url_only - bool, 是否仅清空url
Raises:
exceptions.DatabaseError - 数据库出错时引发此异常
"""
table_prefix = host + "_" + str(port)
if url_only:
NewRequestModel(table_prefix,
multiplexing_conn=True).truncate_table()
else:
NewRequestModel(table_prefix, multiplexing_conn=True).drop_table()
ReportModel(table_prefix, multiplexing_conn=True).drop_table()
self.config_model.delete(table_prefix)
Communicator().set_clean_lru([table_prefix])
def clean_target(self, host, port, url_only=False):
"""
清空目标对应的数据库，同时重置预处理lru
Parameters:
host - str, 目标host
port - int, 目标port
url_only - bool, 是否仅清空url
Raises:
exceptions.DatabaseError - 数据库出错时引发此异常
"""
host_port = common.concat_host(host, port)
if self._scanner_info.get_scanner_id(host_port) is not None:
raise exceptions.TargetIsScanning
if url_only:
NewRequestModel(host_port, multiplexing_conn=True).truncate_table()
else:
NewRequestModel(host_port, multiplexing_conn=True).drop_table()
ReportModel(host_port, multiplexing_conn=True).drop_table()
self._config.del_config(host_port)
Communicator().set_clean_lru([host_port])
class Scanner(base.BaseModule):
def __init__(self, **kwargs):
"""
初始化
"""
# kwargs 参数初始化
self.target_host = kwargs["host"]
self.target_port = kwargs["port"]
self._init_scan_config()
# 用于记录失败请求并标记
self.failed_task_set = set()
Communicator().set_internal_shared("failed_task_set",
self.failed_task_set)
self.module_id = Communicator().get_module_name().split("_")[-1]
Communicator().set_value("max_concurrent_request", 1)
Communicator().set_value(
"request_interval",
Config().get_config("scanner.min_request_interval"))
self._init_db()
self._init_plugin()
# 更新运行时配置
self._update_scan_config()
Logger().info("Start scanning target host:{}:{}".format(
self.target_host, self.target_port))
def _init_scan_config(self):
"""
获取缓存的扫描配置
"""
config_model = ConfigModel(table_prefix="",
use_async=True,
create_table=False,
multiplexing_conn=True)
host_port = self.target_host + "_" + str(self.target_port)
config = config_model.get(host_port)
if config is None:
raise exceptions.GetRuntimeConfigFail
self.scan_config = json.loads(config)
def _save_scan_config(self):
"""
存储当前扫描目标配置
"""
config_model = ConfigModel(table_prefix="",
use_async=True,
create_table=False,
multiplexing_conn=True)
host_port = self.target_host + "_" + str(self.target_port)
config_model.update(host_port, json.dumps(self.scan_config))
def _update_scan_config(self):
"""
更新当前运行的扫描配置
"""
config_model = ConfigModel(table_prefix="",
use_async=True,
create_table=False,
multiplexing_conn=True)
host_port = self.target_host + "_" + str(self.target_port)
self.scan_config = json.loads(config_model.get(host_port))
for plugin_name in self.scan_config["scan_plugin_status"]:
self.plugin_loaded[plugin_name].set_enable(
self.scan_config["scan_plugin_status"][plugin_name]["enable"])
self.plugin_loaded[plugin_name].set_white_url_reg(
self.scan_config["white_url_reg"])
self.plugin_loaded[plugin_name].set_scan_proxy(
self.scan_config["scan_proxy"])
Logger().debug(
"Update scanner config to version {}, new config json is {}".
format(self.scan_config["version"], json.dumps(self.scan_config)))
def _init_plugin(self):
"""
初始化扫描插件
"""
self.plugin_loaded = {}
plugin_import_path = "plugin.scanner"
for plugin_name in self.scan_config["scan_plugin_status"].keys():
try:
plugin_module = __import__(plugin_import_path,
fromlist=[plugin_name])
except Exception as e:
Logger().error("Error in load plugin: {}".format(plugin_name),
exc_info=e)
else:
plugin_instance = getattr(plugin_module,
plugin_name).ScanPlugin()
if isinstance(plugin_instance,
scan_plugin_base.ScanPluginBase):
self.plugin_loaded[plugin_name] = plugin_instance
Logger().debug(
"scanner plugin: {} load success!".format(plugin_name))
else:
Logger().warning(
"scanner plugin {} not inherit class ScanPluginBase!".
format(plugin_name))
if len(self.plugin_loaded) == 0:
Logger().error("No scanner plugin detected, scanner exit!")
raise exceptions.NoPluginError
def _init_db(self):
"""
初始化数据库
"""
model_prefix = self.target_host + "_" + str(self.target_port)
self.new_scan_model = NewRequestModel(model_prefix,
multiplexing_conn=True)
self.new_scan_model.reset_unscanned_item()
report_model = ReportModel(model_prefix, multiplexing_conn=True)
Communicator().set_internal_shared("report_model", report_model)
def _exit(self, signame, loop):
loop.stop()
def run(self):
"""
模块主函数，启动协程
"""
try:
loop = asyncio.get_event_loop()
loop.run_until_complete(self.async_run())
except RuntimeError:
Logger().info("Scanner process has been killed!")
except Exception as e:
Logger().error("Scanner exit with unknow error!", exc_info=e)
async def async_run(self):
"""
协程主函数
"""
# 注册信号处理
loop = asyncio.get_event_loop()
for signame in {'SIGINT', 'SIGTERM'}:
loop.add_signal_handler(
getattr(signal, signame),
functools.partial(self._exit, signame, loop))
# 初始化context
await audit_tools.context.Context().async_init()
# 启动插件
plugin_tasks = []
for plugin_name in self.plugin_loaded:
plugin_tasks.append(
loop.create_task(self.plugin_loaded[plugin_name].async_run()))
# 启动获取扫描结果队列的协程
task_fetch_rasp_result = loop.create_task(self._fetch_from_queue())
# 执行获取新扫描任务
await self._fetch_new_scan()
# 结束所有协程任务，reset共享内存
task_fetch_rasp_result.cancel()
await asyncio.wait({task_fetch_rasp_result})
for task in plugin_tasks:
task.cancel()
await asyncio.wait(set(plugin_tasks),
return_when=asyncio.ALL_COMPLETED)
Communicator().reset_all_value()
async def _fetch_from_queue(self):
"""
获取扫描请求的RaspResult, 并分发给扫描插件
"""
queue_name = "rasp_result_queue_" + self.module_id
sleep_interval = 0.1
continuously_sleep = 0
Logger().debug("Fetch task is running, use queue: " + queue_name)
while True:
if Communicator().get_value(
"config_version") > self.scan_config["version"]:
self._update_scan_config()
try:
data = Communicator().get_data_nowait(queue_name)
Logger().debug("From rasp_result_queue got data: " + str(data))
result_receiver.RaspResultReceiver().add_result(data)
Logger().debug("Send data to rasp_result receiver: {}".format(
data.get_request_id()))
continuously_sleep = 0
except exceptions.QueueEmpty:
if continuously_sleep < 10:
continuously_sleep += 1
await asyncio.sleep(sleep_interval * continuously_sleep)
async def _fetch_new_scan(self):
"""
获取非扫描请求（新扫描任务），并分发给插件
"""
# 扫描插件任务队列最大值
scan_queue_max = 300
# 已扫描的任务数量
self.scan_num = 0
# 扫描队列数量
self.scan_queue_remaining = 0
# 下次获取任务数量
self.fetch_count = 20
# 待标记的已扫描的最大id
self.mark_id = 0
while True:
try:
await self._fetch_task_from_db()
except exceptions.DatabaseError as e:
Logger().error("Database error occured when fetch scan task.",
exc_info=e)
except asyncio.CancelledError as e:
raise e
except Exception as e:
Logger().error(
"Unexpected error occured when fetch scan task.",
exc_info=e)
if self.scan_queue_remaining == 0:
continue
await self._check_scan_progress()
# 调整每次获取的扫描任务数
if self.scan_queue_remaining + self.fetch_count > scan_queue_max:
self.fetch_count = scan_queue_max - self.scan_queue_remaining
elif self.fetch_count < 5:
self.fetch_count = 5
async def _fetch_task_from_db(self):
"""
从数据库中获取当前扫描目标的非扫描请求（新扫描任务）
"""
await self.new_scan_model.mark_result(self.mark_id,
list(self.failed_task_set))
self.failed_task_set.clear()
sleep_interval = 1
continuously_sleep = 0
while True:
data_list = await self.new_scan_model.get_new_scan(self.fetch_count
)
data_count = len(data_list)
Logger().debug("Fetch {} task from db.".format(data_count))
if data_count > 0 or self.scan_queue_remaining > 0:
for item in data_list:
for plugin_name in self.plugin_loaded:
# item 格式: {"id": id, "data":rasp_result_json}
self.plugin_loaded[plugin_name].add_task(item)
Logger().debug("Send task with id: {} to plugins.".format(
item["id"]))
self.scan_queue_remaining += data_count
return
else:
Logger().debug("No url need scan, fetch task sleep {}s".format(
sleep_interval * continuously_sleep))
if continuously_sleep < 10:
continuously_sleep += 1
await asyncio.sleep(sleep_interval * continuously_sleep)
async def _check_scan_progress(self):
"""
监测扫描进度，给出下次获取的任务量
"""
sleep_interval = 1
sleep_count = 0
while True:
await asyncio.sleep(sleep_interval)
sleep_count += 1
scan_num_list = []
scan_id_list = []
for plugin_name in self.plugin_loaded:
plugin_ins = self.plugin_loaded[plugin_name]
plugin_scan_num, plugin_last_id = plugin_ins.get_scan_progress(
)
scan_num_list.append(plugin_scan_num)
scan_id_list.append(plugin_last_id)
plugin_scan_min_num = min(scan_num_list)
plugin_scan_min_id = min(scan_id_list)
finish_count = plugin_scan_min_num - self.scan_num
if sleep_count > 20:
# 20个sleep内未扫描完成，每次最大获取任务量减半
self.scan_queue_remaining -= finish_count
self.scan_num = plugin_scan_min_num
sleep_count = 0
Logger().debug(
"Finish scan num: {}, remain task: {}, max scanned id: {}, decrease task fetch_count."
.format(finish_count, self.scan_queue_remaining,
plugin_scan_min_id))
elif sleep_count > 10:
# 10-20个sleep内完成一半以上，每次最大获取任务量不变
if self.scan_queue_remaining < finish_count * 2:
self.fetch_count = finish_count
break
elif self.scan_queue_remaining == finish_count:
# 10个sleep内完成，每次最大获取任务量加倍
self.fetch_count = self.scan_queue_remaining * 2
break
self.scan_queue_remaining -= finish_count
self.scan_num = plugin_scan_min_num
self.mark_id = plugin_scan_min_id
Logger().debug(
"Finish scan num: {}, remain task: {}, max scanned id: {}".format(
finish_count, self.scan_queue_remaining, plugin_scan_min_id))
async def get_all_target(self):
"""
获取数据库中存在的所有目标主机的列表
Returns:
list, item为dict,格式为：
正在扫描的item:
{
"id": 1, // 扫描任务id
"pid": 64067, // 扫描进程pid
"host": "127.0.0.1", // 扫描的目标主机
"port": 8005, // 扫描的目标端口
"cancel": 0, // 是否正在取消
"pause": 0, // 是否被暂停
"cpu": "0.0%", // cpu占用
"mem": "10.51 M", // 内存占用
"total": 5, // 当前url总数
"scanned": 2, // 扫描的url数量
"concurrent_request": 10, // 当前并发数
"request_interval": 0, // 当前请求间隔
"config": {...}, // 配置信息
"last_time": 1563182956 // 最近一次获取到新url的时间
}
未在扫描的item:
{
"host": "127.0.0.1", // 扫描的目标主机
"port": 8005, // 扫描的目标端口
"total": 5, // 当前url总数
"scanned": 2, // 扫描的url数量
"config": {...}, // 配置信息
"last_time": 1563182956 // 最近一次获取到新url的时间
}
Raises:
exceptions.DatabaseError - 数据库错误引发此异常
"""
tables = BaseModel().get_tables()
result = {}
for table_name in tables:
if table_name.endswith("_ResultList"):
host_port = table_name[:-11]
else:
continue
host_port_split = host_port.split("_")
host = "".join(host_port_split[:-1])
port = host_port_split[-1]
result[host_port] = {"host": host, "port": port}
running_info = await self.get_running_info()
for scanner_id in running_info:
host_port = running_info[scanner_id]["host"] + "_" + str(
running_info[scanner_id]["port"])
result[host_port] = running_info[scanner_id]
result[host_port]["id"] = scanner_id
result_list = []
for host_port in result:
new_request_model = NewRequestModel(host_port,
multiplexing_conn=True)
result[host_port][
"last_time"] = await new_request_model.get_last_time()
if result[host_port].get("id", None) is None:
total, scanned = await new_request_model.get_scan_count()
result[host_port]["total"] = total
result[host_port]["scanned"] = scanned
target_config = self.config_model.get(host_port)
if target_config is None:
target_config = self.config_model.get("default")
result[host_port]["config"] = json.loads(target_config)
result_list.append(result[host_port])
result_list.sort(key=(lambda k: k["last_time"]), reverse=True)
return result_list
async def get_all_target(self, page=1):
"""
获取数据库中存在的所有目标主机的列表
Parameters:
page - int, 获取的页码，每页10个主机
Returns:
list, item为dict,格式为：
正在扫描的item:
{
"id": 1, // 扫描任务id
"pid": 64067, // 扫描进程pid
"host": "127.0.0.1", // 扫描的目标主机
"port": 8005, // 扫描的目标端口
"cancel": 0, // 是否正在取消
"pause": 0, // 是否被暂停
"cpu": "0.0%", // cpu占用
"mem": "10.51 M", // 内存占用
"rasp_result_timeout": 0, // 获取rasp-agent结果超时数量
"waiting_rasp_request": 0, // 等待中的rasp-agent结果数量
"dropped_rasp_result": 0, // 收到的无效rasp-agent结果数量
"send_request": 0, // 已发送测试请求
"failed_request": 0, // 发生错误的测试请求
"total": 5, // 当前url总数
"scanned": 2, // 扫描的url数量
"concurrent_request": 10, // 当前并发数
"request_interval": 0, // 当前请求间隔
"config": {...}, // 配置信息
"last_time": 1563182956 // 最近一次获取到新url的时间
}
未在扫描的item:
{
"host": "127.0.0.1", // 扫描的目标主机
"port": 8005, // 扫描的目标端口
"total": 5, // 当前url总数
"scanned": 2, // 扫描的url数量
"config": {...}, // 配置信息
"last_time": 1563182956 // 最近一次获取到新url的时间
}
Raises:
exceptions.DatabaseError - 数据库错误引发此异常
"""
tables = BaseModel(multiplexing_conn=True).get_tables()
Logger().debug("Got current tables: {}".format(", ".join(tables)))
result_tables = []
result = {}
for table_name in tables:
if table_name.lower().endswith("_resultlist"):
host_port = table_name[:-11]
result_tables.append(host_port)
total_target = len(result_tables)
if page <= 0 or (page - 1) * 10 > total_target:
page = 1
result_tables = result_tables[(page - 1) * 10:page * 10]
for host_port in result_tables:
host_port_split = host_port.split("_")
host = "_".join(host_port_split[:-1])
port = host_port_split[-1]
result[host_port] = {
"host": host,
"port": port
}
running_info = await self.get_running_info()
for scanner_id in running_info:
host_port = running_info[scanner_id]["host"] + \
"_" + str(running_info[scanner_id]["port"])
if host_port in result_tables:
result[host_port] = running_info[scanner_id]
result[host_port]["id"] = scanner_id
config_model = ConfigModel(
table_prefix="", use_async=True, create_table=True, multiplexing_conn=True)
config_result = config_model.get_list(list(result.keys()))
result_list = []
for host_port in result:
new_request_model = NewRequestModel(
host_port, multiplexing_conn=True)
result[host_port]["last_time"] = await new_request_model.get_last_time()
if result[host_port].get("id", None) is None:
total, scanned, failed = await new_request_model.get_scan_count()
result[host_port]["total"] = total
result[host_port]["scanned"] = scanned
result[host_port]["failed"] = failed
if config_result[host_port] is None:
result[host_port]["config"] = self.default_config
else:
result[host_port]["config"] = json.loads(
config_result[host_port])
result_list.append(result[host_port])
# result_list.sort(key=(lambda k:k["last_time"]), reverse=True)
return result_list, total_target
