storage.py - master - Ged52/DomainInfo

storage.py
402 строки13 KB
Oleg Korshunov
+ загрузка/выгрузка базы данных на сервер
13 июл 2026, 13:14
13 июл 2026, 13:14de50768
import json
import sqlite3
from datetime import datetime, timezone, timedelta
from typing import Any, Dict, Optional, List, Tuple

from flask import jsonify, request

from auth_utils import hash_password
from domain_service import lookup_domain, lookup_domains, parse_domains
from pathlib import Path
from typing import Any

DB_PATH = Path(__file__).parent / "domains.db"

CUSTOM_FIELDS = ("responsible", "organization", "hoster", "login", "notes", "esia_verified", "esia_date", "group")


def _now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _connect() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    with _connect() as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS domains (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                domain TEXT NOT NULL UNIQUE,
                responsible TEXT NOT NULL DEFAULT '',
                organization TEXT NOT NULL DEFAULT '',
                hoster TEXT NOT NULL DEFAULT '',
                login TEXT NOT NULL DEFAULT '',
                notes TEXT NOT NULL DEFAULT '',
                `group` TEXT,
                esia_verified INTEGER NOT NULL DEFAULT 0,
                esia_date TEXT,
                expiry TEXT,
                days_left INTEGER,
                registrar TEXT,
                dns_provider TEXT,
                ns_records TEXT NOT NULL DEFAULT '[]',
                whois_status TEXT NOT NULL DEFAULT 'pending',
                whois_error TEXT,
                checked_at TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT NOT NULL UNIQUE,
                password_hash TEXT NOT NULL,
                role TEXT NOT NULL DEFAULT 'user',
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS user_groups (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                `group` TEXT NOT NULL,
                can_view INTEGER NOT NULL DEFAULT 1,
                can_edit INTEGER NOT NULL DEFAULT 0,
                created_at TEXT NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
                UNIQUE(user_id, `group`)
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS user_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                token TEXT NOT NULL UNIQUE,
                expires_at TEXT NOT NULL,
                created_at TEXT NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
            )
            """
        )
        admin_hash = hash_password("password")
        conn.execute(
            """
            INSERT OR IGNORE INTO users (username, password_hash, role, created_at, updated_at)
            VALUES ('admin', ?, 'admin', ?, ?)
            """,
            (admin_hash, _now(), _now())
        )


def _row_to_dict(row: sqlite3.Row) -> Dict:
    data = dict(row)
    try:
        data["ns_records"] = json.loads(data.get("ns_records") or "[]")
    except json.JSONDecodeError:
        data["ns_records"] = []
    return data


def list_domains() -> List[Dict]:
    with _connect() as conn:
        rows = conn.execute("SELECT * FROM domains ORDER BY COALESCE(`group`, ''), domain COLLATE NOCASE").fetchall()
    return [_row_to_dict(row) for row in rows]


def get_domain(domain_id: int) -> Optional[Dict]:
    with _connect() as conn:
        row = conn.execute("SELECT * FROM domains WHERE id = ?", (domain_id,)).fetchone()
    return _row_to_dict(row) if row else None


def add_domains(domain_names: List[str]) -> Tuple[List[Dict], int]:
    now = _now()
    added: List[Dict] = []
    skipped = 0

    with _connect() as conn:
        for name in domain_names:
            existing = conn.execute(
                "SELECT id FROM domains WHERE domain = ?", (name,)
            ).fetchone()
            if existing:
                skipped += 1
                continue

            cursor = conn.execute(
                """
                INSERT INTO domains (domain, created_at, updated_at)
                VALUES (?, ?, ?)
                """,
                (name, now, now),
            )
            row = conn.execute(
                "SELECT * FROM domains WHERE id = ?", (cursor.lastrowid,)
            ).fetchone()
            added.append(_row_to_dict(row))

    return added, skipped


def update_domain(domain_id: int, fields: Dict) -> Optional[Dict]:
    allowed = set(CUSTOM_FIELDS)
    updates = {}
    for k, v in fields.items():
        if k in allowed:
            if k == "esia_verified":
                updates[k] = int(v) if v else 0
            elif k == "esia_date":
                updates[k] = v.strip() if v else ""
            elif k == "group":
                updates[k] = v.strip() if v else ""
            else:
                updates[k] = (v or "").strip()
    if not updates:
        return get_domain(domain_id)

    updates["updated_at"] = _now()
    set_clause = ", ".join(f"`{key}` = ?" for key in updates)
    values = list(updates.values()) + [domain_id]

    with _connect() as conn:
        conn.execute(f"UPDATE domains SET {set_clause} WHERE id = ?", values)

    return get_domain(domain_id)


def update_whois(domain_id: int, whois_data: Dict) -> Optional[Dict]:
    now = _now()
    ns_records = json.dumps(whois_data.get("ns_records") or [], ensure_ascii=False)

    with _connect() as conn:
        conn.execute(
            """
            UPDATE domains SET
                expiry = ?,
                days_left = ?,
                registrar = ?,
                dns_provider = ?,
                ns_records = ?,
                whois_status = ?,
                whois_error = ?,
                checked_at = ?,
                updated_at = ?
            WHERE id = ?
            """,
            (
                whois_data.get("expiry"),
                whois_data.get("days_left"),
                whois_data.get("registrar"),
                whois_data.get("dns_provider"),
                ns_records,
                whois_data.get("status", "ok"),
                whois_data.get("error"),
                now,
                now,
                domain_id,
            ),
        )

    return get_domain(domain_id)


def delete_domain(domain_id: int) -> bool:
    with _connect() as conn:
        cursor = conn.execute("DELETE FROM domains WHERE id = ?", (domain_id,))
    return cursor.rowcount > 0


def get_domains_by_ids(domain_ids: Optional[List[int]] = None) -> List[Dict]:
    if domain_ids:
        placeholders = ",".join("?" * len(domain_ids))
        query = f"SELECT * FROM domains WHERE id IN ({placeholders}) ORDER BY COALESCE(`group`, ''), domain COLLATE NOCASE"
        params: tuple[Any, ...] = tuple(domain_ids)
    else:
        query = "SELECT * FROM domains ORDER BY COALESCE(`group`, ''), domain COLLATE NOCASE"
        params = ()

    with _connect() as conn:
        rows = conn.execute(query, params).fetchall()
    return [_row_to_dict(row) for row in rows]


def get_user_by_username(username: str) -> Optional[Dict]:
    with _connect() as conn:
        row = conn.execute("SELECT * FROM users WHERE username = ?", (username,)).fetchone()
    return _row_to_dict(row) if row else None


def get_user_by_id(user_id: int) -> Optional[Dict]:
    with _connect() as conn:
        row = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
    return _row_to_dict(row) if row else None


def create_user(username: str, password_hash: str) -> Dict:
    now = _now()
    with _connect() as conn:
        cursor = conn.execute(
            "INSERT INTO users (username, password_hash, created_at, updated_at) VALUES (?, ?, ?, ?)",
            (username, password_hash, now, now)
        )
        row = conn.execute("SELECT * FROM users WHERE id = ?", (cursor.lastrowid,)).fetchone()
    return _row_to_dict(row)


def get_user_groups(user_id: int) -> List[Dict]:
    with _connect() as conn:
        rows = conn.execute(
            "SELECT * FROM user_groups WHERE user_id = ?",
            (user_id,)
        ).fetchall()
    return [dict(row) for row in rows]


def get_user_permissions(user_id: int, group: str) -> Optional[Dict]:
    with _connect() as conn:
        row = conn.execute(
            "SELECT * FROM user_groups WHERE user_id = ? AND `group` = ?",
            (user_id, group)
        ).fetchone()
    if row:
        return {"can_view": bool(row["can_view"]), "can_edit": bool(row["can_edit"])}
    return None


def set_user_permissions(user_id: int, group: str, can_view: bool = True, can_edit: bool = False) -> Dict:
    now = _now()
    with _connect() as conn:
        cursor = conn.execute(
            """
            INSERT INTO user_groups (user_id, `group`, can_view, can_edit, created_at)
            VALUES (?, ?, ?, ?, ?)
            ON CONFLICT(user_id, `group`) DO UPDATE SET
                can_view = excluded.can_view,
                can_edit = excluded.can_edit,
                updated_at = ?
            """,
            (user_id, group, 1 if can_view else 0, 1 if can_edit else 0, now, now)
        )
        row = conn.execute(
            "SELECT * FROM user_groups WHERE user_id = ? AND `group` = ?",
            (user_id, group)
        ).fetchone()
    return dict(row) if row else {}


def delete_user_permissions(user_id: int, group: str) -> bool:
    with _connect() as conn:
        cursor = conn.execute(
            "DELETE FROM user_groups WHERE user_id = ? AND `group` = ?",
            (user_id, group)
        )
    return cursor.rowcount > 0


def get_all_users() -> List[Dict]:
    with _connect() as conn:
        rows = conn.execute("SELECT id, username, role, created_at FROM users ORDER BY username").fetchall()
    return [dict(row) for row in rows]


def update_user(user_id: int, username: Optional[str] = None, password_hash: Optional[str] = None, role: Optional[str] = None) -> Optional[Dict]:
    updates = {}
    if username is not None:
        updates["username"] = username.strip()
    if password_hash is not None:
        updates["password_hash"] = password_hash
    if role is not None:
        updates["role"] = role.strip()
    if not updates:
        return get_user_by_id(user_id)
    
    updates["updated_at"] = _now()
    set_clause = ", ".join(f"{key} = ?" for key in updates)
    values = list(updates.values()) + [user_id]
    
    with _connect() as conn:
        conn.execute(f"UPDATE users SET {set_clause} WHERE id = ?", values)
    
    return get_user_by_id(user_id)


def delete_user(user_id: int) -> bool:
    with _connect() as conn:
        cursor = conn.execute("DELETE FROM users WHERE id = ?", (user_id,))
    return cursor.rowcount > 0


def create_session(user_id: int, token: str, expires_at: str) -> Dict:
    now = _now()
    with _connect() as conn:
        cursor = conn.execute(
            "INSERT INTO user_sessions (user_id, token, expires_at, created_at) VALUES (?, ?, ?, ?)",
            (user_id, token, expires_at, now)
        )
        row = conn.execute("SELECT * FROM user_sessions WHERE id = ?", (cursor.lastrowid,)).fetchone()
    return _row_to_dict(row)


def get_session(token: str) -> Optional[Dict]:
    with _connect() as conn:
        row = conn.execute(
            "SELECT * FROM user_sessions WHERE token = ? AND expires_at > ?",
            (token, _now())
        ).fetchone()
    return _row_to_dict(row) if row else None


def delete_session(token: str) -> bool:
    with _connect() as conn:
        cursor = conn.execute("DELETE FROM user_sessions WHERE token = ?", (token,))
    return cursor.rowcount > 0


def delete_user_sessions(user_id: int) -> None:
    with _connect() as conn:
        conn.execute("DELETE FROM user_sessions WHERE user_id = ?", (user_id,))


def reset_admin_password(password_hash: str) -> None:
    with _connect() as conn:
        conn.execute(
            "UPDATE users SET password_hash = ? WHERE username = 'admin'",
            (password_hash,)
        )


def import_db(data: bytes) -> None:
    import shutil
    db_path = DB_PATH
    backups_dir = db_path.parent / "backups"
    backups_dir.mkdir(exist_ok=True)
    backup_path = backups_dir / f"db_backup_{_now().replace(':', '')}.db"
    
    if db_path.exists():
        shutil.copy2(str(db_path), str(backup_path))
    
    with open(db_path, "wb") as f:
        f.write(data)
    
    init_db()


def export_db() -> bytes:
    db_path = DB_PATH
    if not db_path.exists():
        return b""
    with open(db_path, "rb") as f:
        return f.read()