scripts/generate_inventory.py - master - cyberknowledge/CVE_lab

scripts/generate_inventory.py
165 строк5 KB
Zeros312
Initial CVE_lab documentation: schema CVE_lab, maturity rules, LLM scenarios
26 июн 2026, 16:36
26 июн 2026, 16:360de3524
#!/usr/bin/env python3
"""Regenerate README inventory block and data-dictionary stats from PostgreSQL CVE_lab."""
from __future__ import annotations

import os
import re
from datetime import datetime, timezone
from pathlib import Path

import psycopg2
from psycopg2 import sql

ROOT = Path(__file__).resolve().parents[1]
SCHEMA = "CVE_lab"

CONN = dict(
    host=os.environ.get("PGHOST", "172.32.2.200"),
    port=int(os.environ.get("PGPORT", "5432")),
    dbname=os.environ.get("PGDATABASE", "cybz"),
    user=os.environ.get("PGUSER", "postgres"),
    password=os.environ.get("PGPASSWORD", "postgres"),
)

TABLES = [
    ("NVD", "Справочник CVE из NVD (NIST)"),
    ("BDU_FSTEK", "Уязвимости из банка данных ФСТЭК (БДУ)"),
    ("Exploits", "Агрегированные публикации с потенциальными эксплойтами/PoC"),
    ("Exploit_AI_Analysis", "LLM-анализ публикаций из Exploits"),
    ("PoCs_table", "Результаты ручного тестирования PoC/эксплойтов"),
]


def fmt_int(n: int) -> str:
    return f"{n:,}".replace(",", "\u00a0")


def fmt_gb(n_bytes: int) -> str:
    return f"{n_bytes / (1024 ** 3):.3f}"


def table_size(cur, table: str) -> int:
    cur.execute(
        """
        SELECT pg_total_relation_size(c.oid)
        FROM pg_class c
        JOIN pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname = %s AND c.relname = %s
        """,
        (SCHEMA, table),
    )
    row = cur.fetchone()
    return int(row[0]) if row else 0


def row_count(cur, table: str) -> int:
    q = sql.SQL("SELECT COUNT(*) FROM {}.{}").format(
        sql.Identifier(SCHEMA), sql.Identifier(table)
    )
    cur.execute(q)
    return int(cur.fetchone()[0])


def fill_rate(cur, table: str, column: str) -> float | None:
    q = sql.SQL(
        "SELECT COUNT(*) FILTER (WHERE {col} IS NOT NULL AND {col}::text <> '')::float "
        "/ NULLIF(COUNT(*), 0) FROM {schema}.{table}"
    ).format(
        col=sql.Identifier(column),
        schema=sql.Identifier(SCHEMA),
        table=sql.Identifier(table),
    )
    cur.execute(q)
    val = cur.fetchone()[0]
    return float(val) * 100 if val is not None else None


def inventory_markdown(cur) -> str:
    lines = [
        "## Инвентаризация таблиц (схема `CVE_lab`)",
        "",
        f"Сгенерировано: **{datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M UTC')}** "
        f"(`python scripts/generate_inventory.py`).",
        "",
        "Колонки: таблица, назначение, строк в PostgreSQL, размер в БД (ГБ).",
        "",
        "| Таблица | Назначение | Строк в БД | Размер в БД (ГБ) |",
        "|---|---|---:|---:|",
    ]
    total_bytes = 0
    for table, purpose in TABLES:
        cnt = row_count(cur, table)
        size = table_size(cur, table)
        total_bytes += size
        lines.append(
            f"| `{SCHEMA}.{table}` | {purpose} | {fmt_int(cnt)} | {fmt_gb(size)} |"
        )
    lines.extend(
        [
            "",
            f"Итого схема `{SCHEMA}`: ≈ **{fmt_gb(total_bytes)} GiB**.",
            "",
        ]
    )
    return "\n".join(lines)


def maturity_block(cur) -> str:
    cur.execute(
        f'SELECT maturity, COUNT(*) FROM "{SCHEMA}"."Exploit_AI_Analysis" '
        "GROUP BY 1 ORDER BY 2 DESC"
    )
    ai_rows = cur.fetchall()
    cur.execute(
        f'SELECT workability, COUNT(*) FROM "{SCHEMA}"."PoCs_table" '
        "GROUP BY 1 ORDER BY 2 DESC"
    )
    poc_w = cur.fetchall()
    lines = [
        "## Срезы качества данных",
        "",
        "### Зрелость (Exploit_AI_Analysis.maturity)",
        "",
        "| Значение | Записей |",
        "|---|---:|",
    ]
    for val, cnt in ai_rows:
        lines.append(f"| {val or '(пусто)'} | {fmt_int(cnt)} |")
    lines.extend(["", "### Работоспособность (PoCs_table.workability)", "", "| Значение | Записей |", "|---|---:|"])
    for val, cnt in poc_w:
        lines.append(f"| {val or '(пусто)'} | {fmt_int(cnt)} |")
    lines.append("")
    return "\n".join(lines)


def patch_readme(content: str, inventory: str, maturity: str) -> str:
    start = "<!-- inventory:start -->"
    end = "<!-- inventory:end -->"
    block = f"{start}\n{inventory}\n{maturity}{end}"
    if start in content and end in content:
        return re.sub(
            re.escape(start) + r".*?" + re.escape(end),
            block,
            content,
            flags=re.DOTALL,
        )
    marker = "## Инвентаризация таблиц"
    if marker in content:
        head, _ = content.split(marker, 1)
        return head.rstrip() + "\n\n" + block + "\n"
    return content.rstrip() + "\n\n" + block + "\n"


def main() -> None:
    conn = psycopg2.connect(**CONN)
    cur = conn.cursor()
    inv = inventory_markdown(cur)
    mat = maturity_block(cur)
    readme = (ROOT / "README.md").read_text(encoding="utf-8")
    (ROOT / "README.md").write_text(patch_readme(readme, inv, mat), encoding="utf-8")
    print("Updated README.md inventory block")
    conn.close()


if __name__ == "__main__":
    main()