WuYanBo
/
CodeReview
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
CodeReview/backend/alembic/versions/006_add_agent_tables.py

256 lines
11 KiB
Python
Raw Blame History

"""Add agent audit tables
Revision ID: 006_add_agent_tables
Revises: 5fc1cc05d5d0
Create Date: 2024-01-15 10:00:00.000000
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision = '006_add_agent_tables'
down_revision = '5fc1cc05d5d0'
branch_labels = None
depends_on = None
def upgrade() -> None:
# 创建 agent_tasks 表
op.create_table(
'agent_tasks',
sa.Column('id', sa.String(36), primary_key=True),
sa.Column('project_id', sa.String(36), sa.ForeignKey('projects.id', ondelete='CASCADE'), nullable=False),
# 任务基本信息
sa.Column('name', sa.String(255), nullable=True),
sa.Column('description', sa.Text(), nullable=True),
sa.Column('task_type', sa.String(50), default='agent_audit'),
# 任务配置
sa.Column('audit_scope', sa.JSON(), nullable=True),
sa.Column('target_vulnerabilities', sa.JSON(), nullable=True),
sa.Column('verification_level', sa.String(50), default='sandbox'),
# 分支信息
sa.Column('branch_name', sa.String(255), nullable=True),
# 排除模式
sa.Column('exclude_patterns', sa.JSON(), nullable=True),
# 文件范围
sa.Column('target_files', sa.JSON(), nullable=True),
# LLM 配置
sa.Column('llm_config', sa.JSON(), nullable=True),
# Agent 配置
sa.Column('agent_config', sa.JSON(), nullable=True),
sa.Column('max_iterations', sa.Integer(), default=50),
sa.Column('token_budget', sa.Integer(), default=100000),
sa.Column('timeout_seconds', sa.Integer(), default=1800),
# 状态
sa.Column('status', sa.String(20), default='pending'),
sa.Column('current_phase', sa.String(50), nullable=True),
sa.Column('current_step', sa.String(255), nullable=True),
sa.Column('error_message', sa.Text(), nullable=True),
# 进度统计
sa.Column('total_files', sa.Integer(), default=0),
sa.Column('indexed_files', sa.Integer(), default=0),
sa.Column('analyzed_files', sa.Integer(), default=0),
sa.Column('total_chunks', sa.Integer(), default=0),
# Agent 统计
sa.Column('total_iterations', sa.Integer(), default=0),
sa.Column('tool_calls_count', sa.Integer(), default=0),
sa.Column('tokens_used', sa.Integer(), default=0),
# 发现统计
sa.Column('findings_count', sa.Integer(), default=0),
sa.Column('verified_count', sa.Integer(), default=0),
sa.Column('false_positive_count', sa.Integer(), default=0),
# 严重程度统计
sa.Column('critical_count', sa.Integer(), default=0),
sa.Column('high_count', sa.Integer(), default=0),
sa.Column('medium_count', sa.Integer(), default=0),
sa.Column('low_count', sa.Integer(), default=0),
# 质量评分
sa.Column('quality_score', sa.Float(), default=0.0),
sa.Column('security_score', sa.Float(), default=0.0),
# 审计计划
sa.Column('audit_plan', sa.JSON(), nullable=True),
# 时间戳
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
sa.Column('updated_at', sa.DateTime(timezone=True), onupdate=sa.func.now()),
sa.Column('started_at', sa.DateTime(timezone=True), nullable=True),
sa.Column('completed_at', sa.DateTime(timezone=True), nullable=True),
# 创建者
sa.Column('created_by', sa.String(36), sa.ForeignKey('users.id'), nullable=False),
)
# 创建 agent_tasks 索引
op.create_index('ix_agent_tasks_project_id', 'agent_tasks', ['project_id'])
op.create_index('ix_agent_tasks_status', 'agent_tasks', ['status'])
op.create_index('ix_agent_tasks_created_by', 'agent_tasks', ['created_by'])
op.create_index('ix_agent_tasks_created_at', 'agent_tasks', ['created_at'])
# 创建 agent_events 表
op.create_table(
'agent_events',
sa.Column('id', sa.String(36), primary_key=True),
sa.Column('task_id', sa.String(36), sa.ForeignKey('agent_tasks.id', ondelete='CASCADE'), nullable=False),
# 事件信息
sa.Column('event_type', sa.String(50), nullable=False),
sa.Column('phase', sa.String(50), nullable=True),
# 事件内容
sa.Column('message', sa.Text(), nullable=True),
# 工具调用相关
sa.Column('tool_name', sa.String(100), nullable=True),
sa.Column('tool_input', sa.JSON(), nullable=True),
sa.Column('tool_output', sa.JSON(), nullable=True),
sa.Column('tool_duration_ms', sa.Integer(), nullable=True),
# 关联的发现
sa.Column('finding_id', sa.String(36), nullable=True),
# Token 消耗
sa.Column('tokens_used', sa.Integer(), default=0),
# 元数据
sa.Column('metadata', sa.JSON(), nullable=True),
# 序号
sa.Column('sequence', sa.Integer(), default=0),
# 时间戳
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
)
# 创建 agent_events 索引
op.create_index('ix_agent_events_task_id', 'agent_events', ['task_id'])
op.create_index('ix_agent_events_event_type', 'agent_events', ['event_type'])
op.create_index('ix_agent_events_sequence', 'agent_events', ['sequence'])
op.create_index('ix_agent_events_created_at', 'agent_events', ['created_at'])
# 创建 agent_findings 表
op.create_table(
'agent_findings',
sa.Column('id', sa.String(36), primary_key=True),
sa.Column('task_id', sa.String(36), sa.ForeignKey('agent_tasks.id', ondelete='CASCADE'), nullable=False),
# 漏洞基本信息
sa.Column('vulnerability_type', sa.String(100), nullable=False),
sa.Column('severity', sa.String(20), nullable=False),
sa.Column('title', sa.String(500), nullable=False),
sa.Column('description', sa.Text(), nullable=True),
# 位置信息
sa.Column('file_path', sa.String(500), nullable=True),
sa.Column('line_start', sa.Integer(), nullable=True),
sa.Column('line_end', sa.Integer(), nullable=True),
sa.Column('column_start', sa.Integer(), nullable=True),
sa.Column('column_end', sa.Integer(), nullable=True),
sa.Column('function_name', sa.String(255), nullable=True),
sa.Column('class_name', sa.String(255), nullable=True),
# 代码片段
sa.Column('code_snippet', sa.Text(), nullable=True),
sa.Column('code_context', sa.Text(), nullable=True),
# 数据流信息
sa.Column('source', sa.Text(), nullable=True),
sa.Column('sink', sa.Text(), nullable=True),
sa.Column('dataflow_path', sa.JSON(), nullable=True),
# 验证信息
sa.Column('status', sa.String(30), default='new'),
sa.Column('is_verified', sa.Boolean(), default=False),
sa.Column('verification_method', sa.Text(), nullable=True),
sa.Column('verification_result', sa.JSON(), nullable=True),
sa.Column('verified_at', sa.DateTime(timezone=True), nullable=True),
# PoC
sa.Column('has_poc', sa.Boolean(), default=False),
sa.Column('poc_code', sa.Text(), nullable=True),
sa.Column('poc_description', sa.Text(), nullable=True),
sa.Column('poc_steps', sa.JSON(), nullable=True),
# 修复建议
sa.Column('suggestion', sa.Text(), nullable=True),
sa.Column('fix_code', sa.Text(), nullable=True),
sa.Column('fix_description', sa.Text(), nullable=True),
sa.Column('references', sa.JSON(), nullable=True),
# AI 解释
sa.Column('ai_explanation', sa.Text(), nullable=True),
sa.Column('ai_confidence', sa.Float(), nullable=True),
# XAI
sa.Column('xai_what', sa.Text(), nullable=True),
sa.Column('xai_why', sa.Text(), nullable=True),
sa.Column('xai_how', sa.Text(), nullable=True),
sa.Column('xai_impact', sa.Text(), nullable=True),
# 关联规则
sa.Column('matched_rule_code', sa.String(100), nullable=True),
sa.Column('matched_pattern', sa.Text(), nullable=True),
# CVSS 评分
sa.Column('cvss_score', sa.Float(), nullable=True),
sa.Column('cvss_vector', sa.String(100), nullable=True),
# 元数据
sa.Column('metadata', sa.JSON(), nullable=True),
sa.Column('tags', sa.JSON(), nullable=True),
# 去重标识
sa.Column('fingerprint', sa.String(64), nullable=True),
# 时间戳
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
sa.Column('updated_at', sa.DateTime(timezone=True), onupdate=sa.func.now()),
)
# 创建 agent_findings 索引
op.create_index('ix_agent_findings_task_id', 'agent_findings', ['task_id'])
op.create_index('ix_agent_findings_vulnerability_type', 'agent_findings', ['vulnerability_type'])
op.create_index('ix_agent_findings_severity', 'agent_findings', ['severity'])
op.create_index('ix_agent_findings_file_path', 'agent_findings', ['file_path'])
op.create_index('ix_agent_findings_status', 'agent_findings', ['status'])
op.create_index('ix_agent_findings_fingerprint', 'agent_findings', ['fingerprint'])
def downgrade() -> None:
# 删除索引和表
op.drop_index('ix_agent_findings_fingerprint', 'agent_findings')
op.drop_index('ix_agent_findings_status', 'agent_findings')
op.drop_index('ix_agent_findings_file_path', 'agent_findings')
op.drop_index('ix_agent_findings_severity', 'agent_findings')
op.drop_index('ix_agent_findings_vulnerability_type', 'agent_findings')
op.drop_index('ix_agent_findings_task_id', 'agent_findings')
op.drop_table('agent_findings')
op.drop_index('ix_agent_events_created_at', 'agent_events')
op.drop_index('ix_agent_events_sequence', 'agent_events')
op.drop_index('ix_agent_events_event_type', 'agent_events')
op.drop_index('ix_agent_events_task_id', 'agent_events')
op.drop_table('agent_events')
op.drop_index('ix_agent_tasks_created_at', 'agent_tasks')
op.drop_index('ix_agent_tasks_created_by', 'agent_tasks')
op.drop_index('ix_agent_tasks_status', 'agent_tasks')
op.drop_index('ix_agent_tasks_project_id', 'agent_tasks')
op.drop_table('agent_tasks')
Reference in New Issue View Git Blame Copy Permalink