zhaolei
16f134e328
核心功能: - 6 层检测管线: L0 结构分析 → L1 ClamAV → L2 文档深度检测 → L3 关键词匹配 → L4 YARA → L5 PE/二进制 - FastAPI + PostgreSQL + Redis + Docker Compose 架构 - 客户端 HMAC-SHA256 签名鉴权 + 限流 + 幂等性 - 控制台 UI (登录/Dashboard/诊断中心/策略管理/任务列表/测试中心/产品介绍页) - 三层诊断定时器 (error_scan 30s / db_aggregate 30s / container_check 5min) 并发与稳定性修复: - 诊断定时器 Redis 时间槽位锁 (多 worker 同一槽位只有一个执行) - JSONL 存储 fcntl 跨进程排他锁 + 容错解码 + 原子重写 - 移除 SQLite fallback, 统一 PostgreSQL - API_CLIENTS 脱敏为环境变量, .env 已 .gitignore
108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
"""数据库初始化脚本。
|
||
|
||
创建所有表结构 → 生成默认管理员账号 → 载入初始策略版本。
|
||
部署时 ``docker compose run --rm api python -m scripts.init_db`` 执行一次。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import sys
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
|
||
# 确保项目根目录在 sys.path 中
|
||
ROOT = Path(__file__).resolve().parent.parent
|
||
sys.path.insert(0, str(ROOT))
|
||
|
||
import os
|
||
|
||
os.environ.setdefault("MESHSAFE_MASTER_KEY",
|
||
os.environ.get("MESHSAFE_MASTER_KEY", "Dt0ktRs8SjEH8VxWZJDX125NV1qiII88HoiXZPhWers="))
|
||
os.environ.setdefault("MESHSAFE_ADMIN_PASSWORD",
|
||
os.environ.get("MESHSAFE_ADMIN_PASSWORD", "admin123"))
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.db import Base, engine
|
||
from app.models import ConsoleUser, PolicyVersion
|
||
from app.security.crypto import hash_password
|
||
from app.services.policy import PolicyStore
|
||
|
||
logger = logging.getLogger("init_db")
|
||
|
||
|
||
def main() -> None:
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||
|
||
logger.info("=== MeshSafe 数据库初始化 ===")
|
||
|
||
# 1) 创建所有表
|
||
logger.info("步骤 1: 创建表结构 ...")
|
||
Base.metadata.create_all(bind=engine)
|
||
logger.info(" ✓ 所有表已创建 / 校验存在")
|
||
|
||
with Session(engine) as db:
|
||
# 2) 默认管理员（首次部署创建，已存在则跳过）
|
||
logger.info("步骤 2: 管理员账号 ...")
|
||
existing = db.execute(select(ConsoleUser).where(ConsoleUser.username == "admin")).scalar_one_or_none()
|
||
if existing:
|
||
logger.info(" ✓ 默认管理员已存在（username=admin），跳过")
|
||
else:
|
||
admin_pw = os.environ["MESHSAFE_ADMIN_PASSWORD"]
|
||
pw_hash = hash_password(admin_pw)
|
||
user = ConsoleUser(
|
||
username="admin",
|
||
password_hash=pw_hash,
|
||
role="admin",
|
||
is_active=True,
|
||
must_change_password=True,
|
||
)
|
||
db.add(user)
|
||
db.commit()
|
||
logger.info(" ✓ 默认管理员已创建（username=admin, 密码见环境变量 MESHSAFE_ADMIN_PASSWORD）")
|
||
|
||
# 3) 载入初始策略版本（从 policy.yaml）
|
||
logger.info("步骤 3: 载入初始策略版本 ...")
|
||
existing_pv = db.execute(
|
||
select(PolicyVersion).order_by(PolicyVersion.created_at.desc())
|
||
).scalar_one_or_none()
|
||
|
||
policy_path = os.environ.get("MESHSAFE_POLICY_FILE", "config/policy.yaml")
|
||
policy_file = ROOT / policy_path
|
||
if not policy_file.exists():
|
||
logger.warning(" 策略文件不存在：%s，跳过", policy_file)
|
||
return
|
||
|
||
store = PolicyStore(policy_file=str(policy_file))
|
||
policy = store.current()
|
||
|
||
# 生成 version 和 sha256
|
||
content_json = json.dumps(policy.raw, sort_keys=True, ensure_ascii=False)
|
||
content_hash = hashlib.sha256(content_json.encode("utf-8")).hexdigest()
|
||
|
||
version_name = f"init-{policy.version}-{content_hash[:8]}"
|
||
|
||
if existing_pv and existing_pv.sha256 == content_hash:
|
||
logger.info(" ✓ 当前策略已是最新版本，跳过（sha256=%s）", content_hash[:12])
|
||
else:
|
||
pv = PolicyVersion(
|
||
version=version_name,
|
||
content=policy.raw,
|
||
sha256=content_hash,
|
||
)
|
||
db.add(pv)
|
||
db.commit()
|
||
logger.info(" ✓ 初始策略版本已载入（version=%s, rules=%d, sha256=%s）",
|
||
version_name, len(policy.zero_tolerance_rules), content_hash[:12])
|
||
|
||
logger.info("=== 初始化完成 ===")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|