Metadata-Version: 2.4
Name: tasklite
Version: 0.0.7
Summary: A tool for integrating with other systems
Author-email: Frank <frank@example.com>
License: MIT
Project-URL: Homepage, https://github.com/yourname/your-sdk
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: pydantic>=2.7.4
Requires-Dist: SQLAlchemy>=2.0.30
Requires-Dist: PyMySQL>=1.1.1

# tasklite

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

## 工具集成说明

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

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

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

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

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

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

## 安装

```shell
pip install tasklite
```

依赖要求：Python >= 3.8

## 快速开始

```python
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 执行日志（调试用） |

```python
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`。

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

---

#### `update_task_progress(progress: float) -> bool`

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

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

```python
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`，传空字符串则仅更新进度 |

```python
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` |

```python
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"` |

```python
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]` | 流程节点列表 |

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

---

#### `save_task_node_edge(edge: dto.TaskFlowEdge) -> bool`

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

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

```python
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` | 日志对象 |

```python
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` | 漏洞信息对象 |

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

---

## 完整示例

```python
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()
```
