tasklite

一个用于安全测试工具集成的 Python SDK，提供任务管理、流程编排、日志记录和漏洞上报等能力。

工具集成说明

1、已有工具改造非常简单，把之前的工具转换为 cli 启动，接受一个文件路径参数。

2、需要把原来的任务输入改装一下，通过调用函数方法获取任务。

3、（可选）如果有任务调度图，可以把节点、边（执行流向）分别保存一下，如果没有可以不保存。

4、把之前的任务日志，调用函数方法保存一下。

5、把之前的漏洞信息，调用函数方法保存一下。

6、把工具编译成一个二进制文件，目标平台：Linux x86_64 架构，或者提供代码和 dockerfile 文件。

安装

pip install tasklite

依赖要求：Python >= 3.8

快速开始

from datetime import datetime
from tasklite import model, task, dto

# 1. 初始化数据库（仅需执行一次）
model.init_db()

# 2. 更新任务进度
task.update_task_progress(50.5)

# 3. 保存任务日志
log = dto.TaskLog()
log.status = "success"
log.message = "SQL 注入扫描完成"
log.level = "info"
log.execute_time = datetime.now()
task.save_task_log(log)

# 4. 上报漏洞信息
vlun = dto.TaskVlun()
vlun.name = "SQL注入漏洞"
vlun.level = "高危"
vlun.path = "/api/login"
vlun.description = "登录接口存在 SQL 注入漏洞"
task.save_task_vlun(vlun)

API 参考

model 模块

提供数据库初始化与会话管理。

init_db(echo: bool = False)

初始化数据库连接并创建所有表。必须在调用其他任何 API 之前执行。

参数	类型	默认值	说明
db_path	str	"/home/db.sqlite"	SQLite 数据库文件路径，也可以从环境变量 DB_PATH 读取
echo	bool	False	是否打印 SQL 执行日志（调试用）

model.init_db(echo=True)  # 开启 SQL 调试日志

---

dto 数据模型

所有数据模型均基于 Pydantic，支持从 ORM 对象直接转换（from_attributes=True）。

dto.Task — 任务信息

字段	类型	默认值	说明
test_type	str	""	测试任务类型：binary / web / app / code / agent / miniprogram
name	str	""	测试任务名称
target	str	""	测试目标
download_dir	Optional[str]	None	工具文件自动下载到本地的目录
upload_dir	Optional[str]	None	工具运行完成后需要上传到 COS 的文件目录
status	Optional[str]	None	任务状态：pending / running / finish / error
progress	Optional[str]	None	测试进度
ai_summary	Optional[str]	None	AI 总结
error_reason	Optional[str]	None	错误原因

dto.TaskFlowNode — 任务流程节点

字段	类型	默认值	说明
code	str	""	流程代码（唯一标识）
name	str	""	流程名称
status	Optional[str]	None	节点状态：static / pending / running / finish / error
result	Optional[str]	None	流程结果
description	Optional[str]	None	流程描述
error_reason	Optional[str]	None	错误原因

dto.TaskFlowEdge — 任务流程边（执行流向）

字段	类型	默认值	说明
from_node	str	""	起始节点代码
to_node	str	""	目标节点代码
condition	Optional[str]	None	执行条件（如 "所有用例通过"）
description	Optional[str]	None	描述信息
error_reason	Optional[str]	None	错误原因
order_no	Optional[int]	None	多条相同路径时的执行顺序
execute_time	datetime	datetime.now()	执行时间

dto.TaskLog — 任务日志

字段	类型	默认值	说明
execute_time	datetime	datetime.now()	执行时间
level	str	""	日志级别（如 info / warn / error）
status	str	""	状态：success / failed
message	str	""	日志消息内容

dto.TaskVlun — 漏洞信息

字段	类型	默认值	说明
discovery_time	datetime	datetime.now()	发现时间
level	str	""	漏洞等级
name	str	""	漏洞名称
path	str	""	漏洞路径
description	Optional[str]	None	漏洞描述
related	Optional[str]	None	相关信息（如 CVE 编号）
payload	Optional[str]	None	漏洞 payload
result	Optional[str]	None	漏洞结果
reproduce	Optional[str]	None	漏洞复现完整过程
ai_summary	Optional[str]	None	AI 总结
fix_suggestion	Optional[str]	None	修复建议

---

task 模块

提供任务管理、进度上报、日志记录和漏洞保存等核心功能。所有方法返回 bool 表示操作是否成功。

---

get_task() -> Optional[dto.Task]

获取最新的任务信息。

返回值：最新的一条 dto.Task 对象，无任务时返回 None。

t = task.get_task()
if t:
    print(f"当前任务: {t.name}")

---

update_task_progress(progress: float) -> bool

更新最新任务的执行进度。

参数	类型	说明
progress	float	任务进度，范围建议 0~100

task.update_task_progress(75.0)

---

update_flow_node_progress(progress: float, code: str, status: str) -> bool

更新指定流程节点的进度和状态。

参数	类型	说明
progress	float	节点进度
code	str	节点代码（唯一标识）
status	str	节点状态：static / pending / running / finish / error，传空字符串则仅更新进度

task.update_flow_node_progress(100.0, "scan_sql_injection", "finish")

---

save_static_task_nodes(nodes: list[dto.TaskFlowNode], is_auto_add_edge: bool) -> bool

保存静态任务流程节点（在 SaaS 平台展示执行流程）。节点的 status 会自动设置为 "static"。

参数	类型	说明
nodes	list[dto.TaskFlowNode]	按执行顺序排列的节点列表，不能为空
is_auto_add_edge	bool	是否自动按顺序生成边（from_node → to_node）。如边的信息有描述需求（如 condition），建议设为 False 并手动调用 save_task_node_edge

nodes = [
    dto.TaskFlowNode(code="step1", name="代码扫描"),
    dto.TaskFlowNode(code="step2", name="漏洞验证"),
    dto.TaskFlowNode(code="step3", name="报告生成"),
]
task.save_static_task_nodes(nodes, is_auto_add_edge=True)

---

save_task_node(node: dto.TaskFlowNode) -> bool

保存单个任务流程节点。status 不能为 "static" 或空，需调用 save_task_node_edge 手动添加边。

参数	类型	说明
node	dto.TaskFlowNode	流程节点对象，status 必填且不能为 "static"

node = dto.TaskFlowNode(
    code="scan_sql",
    name="SQL注入扫描",
    status="running",
    description="正在执行SQL注入扫描"
)
task.save_task_node(node)

---

save_task_nodes(nodes: list[dto.TaskFlowNode]) -> bool

批量保存任务流程节点。每个节点的 status 不能为 "static" 或空。

参数	类型	说明
nodes	list[dto.TaskFlowNode]	流程节点列表

task.save_task_nodes([node1, node2, node3])

---

save_task_node_edge(edge: dto.TaskFlowEdge) -> bool

保存单条任务节点间的执行流向（边）。

参数	类型	说明
edge	dto.TaskFlowEdge	边的流向信息

edge = dto.TaskFlowEdge(
    from_node="step1",
    to_node="step2",
    condition="所有用例通过",
    description="测试全部通过则部署"
)
task.save_task_node_edge(edge)

---

save_task_node_edges(edges: list[dto.TaskFlowEdge]) -> bool

批量保存任务节点间的执行流向（边）。

参数	类型	说明
edges	list[dto.TaskFlowEdge]	边的流向信息列表

---

save_task_log(log: dto.TaskLog) -> bool

保存一条任务执行日志。

参数	类型	说明
log	dto.TaskLog	日志对象

log = dto.TaskLog()
log.level = "error"
log.status = "failed"
log.message = "SQL注入扫描失败：目标无响应"
task.save_task_log(log)

---

save_task_vlun(vlun: dto.TaskVlun) -> bool

保存一条漏洞信息。

参数	类型	说明
vlun	dto.TaskVlun	漏洞信息对象

vlun = dto.TaskVlun(
    name="XSS 跨站脚本漏洞",
    level="中危",
    path="/search",
    description="搜索参数未做 HTML 实体编码",
    payload='<script>alert(1)</script>',
    fix_suggestion="对用户输入进行 HTML 实体编码"
)
task.save_task_vlun(vlun)

---

完整示例

from datetime import datetime
from tasklite import model, task, dto

def main():
    # 初始化数据库
    model.init_db(echo=False)

    # 上报进度
    task.update_task_progress(10.0)

    # 记录执行日志
    log = dto.TaskLog()
    log.level = "info"
    log.status = "success"
    log.message = "工具初始化完成"
    log.execute_time = datetime.now()
    task.save_task_log(log)

    # ... 执行核心扫描逻辑 ...

    # 上报漏洞
    vlun = dto.TaskVlun()
    vlun.name = "SQL注入漏洞"
    vlun.level = "高危"
    vlun.path = "/api/login"
    vlun.description = "登录接口 username 参数存在 SQL 注入"
    vlun.payload = "' OR 1=1 --"
    vlun.fix_suggestion = "使用参数化查询或 ORM 框架"
    task.save_task_vlun(vlun)

    # 更新进度至完成
    task.update_task_progress(100.0)

if __name__ == "__main__":
    main()