db.py

import datetime
import psycopg2
from config import DB_PARAMS


def get_connection():
    return psycopg2.connect(**DB_PARAMS)


def get_user(conn, name: str):
    try:
        conn.rollback()
        cur = conn.cursor()
        cur.execute("SELECT user_id, password FROM users WHERE name = %s", (name,))
        return cur.fetchone()
    except Exception:
        conn.rollback()
        raise


def create_user(conn, name: str, password_hash: str) -> int:
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO users (name, password) VALUES (%s, %s) RETURNING user_id",
        (name, password_hash)
    )
    uid = cur.fetchone()[0]
    conn.commit()
    return uid


def get_admin(conn, name: str):
    try:
        conn.rollback()
        cur = conn.cursor()
        cur.execute(
            "SELECT admin_id, password FROM administrators WHERE name = %s", (name,)
        )
        return cur.fetchone()
    except Exception:
        conn.rollback()
        raise


def save_scan_results(conn, user_id: int, networks: list):
    cur = conn.cursor()
    scan_time = datetime.datetime.now()
    for net in networks:
        cur.execute(
            """
            INSERT INTO security_system (
                user_id, ssid, scan_time,
                signal_strength, encryption_type, auth_type, security_status
            ) VALUES (%s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT DO NOTHING
            """,
            (
                user_id,
                net.get("SSID", ""),
                scan_time,
                net.get("signal", 0),
                net.get("encryption", ""),
                net.get("auth", ""),
                net.get("security", "Неизвестно"),
            ),
        )
    conn.commit()

def write_log(conn, user_id, admin_id, ssid: str, description: str):
    cur = conn.cursor()
    cur.execute(
        """
        SELECT scan_time FROM security_system
        WHERE user_id = %s AND ssid = %s
        ORDER BY scan_time DESC LIMIT 1
        """,
        (user_id, ssid),
    )
    row = cur.fetchone()
    log_time = row[0] if row else datetime.datetime.now()

    cur.execute(
        """
        INSERT INTO logs (user_id, admin_id, ssid, scan_time, description)
        VALUES (%s, %s, %s, %s, %s)
        """,
        (user_id, admin_id, ssid, log_time, description),
    )
    conn.commit()


def fetch_all_logs(conn) -> list:
    cur = conn.cursor()
    cur.execute(
        """
        SELECT log_id, user_id, admin_id, ssid, scan_time, description
        FROM logs
        ORDER BY log_id DESC
        """
    )
    return cur.fetchall()


def delete_log(conn, log_id: int, admin_id: int, ssid: str, scan_time):
    cur = conn.cursor()
    cur.execute("DELETE FROM logs WHERE log_id = %s", (log_id,))
    cur.execute(
        """
        INSERT INTO logs (user_id, admin_id, ssid, scan_time, description)
        VALUES (%s, %s, %s, %s, %s)
        """,
        (None, admin_id, ssid, scan_time, f"Удалён администратором (ID {admin_id})"),
    )
    conn.commit()


def update_log_description(conn, log_id: int, admin_id: int, new_desc: str):
    cur = conn.cursor()
    cur.execute(
        "UPDATE logs SET description = %s, admin_id = %s WHERE log_id = %s",
        (new_desc, admin_id, log_id),
    )
    conn.commit()


_SECURITY_PRIORITY = {
    "Небезопасно":        1,
    "Слабо":              2,
    "Удовлетворительно":  3,
    "Неизвестно":         4,
    "Хорошо":             5,
    "Отлично":            6,
}

_PRIORITY_MAP = [
    ("Небезопасно",       1),
    ("Слабо",             2),
    ("Удовлетворительно", 3),
    ("Неизвестно",        4),
    ("Хорошо",            5),
    ("Отлично",           6),
]


def fetch_security_stats(conn) -> dict:
    try:
        conn.rollback()
        cur = conn.cursor()
        cur.execute("SELECT ssid, security_status FROM security_system")
        rows = cur.fetchall()
    except Exception:
        conn.rollback()
        return {}

    priority_index = {name: pri for name, pri in _PRIORITY_MAP}
    worst: dict = {}
    for ssid, status in rows:
        pri = priority_index.get(status, 4)
        if ssid not in worst or pri < worst[ssid][0]:
            worst[ssid] = (pri, status)

    result: dict = {}
    for _, (_, status) in worst.items():
        result[status] = result.get(status, 0) + 1
    return result


def fetch_scan_history(conn, limit: int = 14) -> list:
    try:
        conn.rollback()
        cur = conn.cursor()
        cur.execute(
            """
            SELECT day, COUNT(*) AS unique_networks
            FROM (
                SELECT DISTINCT scan_time::date AS day, ssid
                FROM security_system
            ) AS daily_unique
            GROUP BY day
            ORDER BY day DESC
            LIMIT %s
            """,
            (limit,),
        )
        return cur.fetchall()
    except Exception:
        conn.rollback()
        return []


#Bluetooth

def save_bt_scan_results(conn, user_id: int, devices: list):
    cur = conn.cursor()
    scan_time = datetime.datetime.now()
    for dev in devices:
        cur.execute(
            """
            INSERT INTO bt_scan_results (
                user_id, device_name, mac_address, device_type,
                rssi, signal_strength, has_name, services_count,
                security_status, scan_time
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            (
                user_id,
                dev.get("name", ""),
                dev.get("address", ""),
                dev.get("device_type", ""),
                dev.get("rssi", 0),
                dev.get("signal", 0),
                dev.get("has_name", False),
                dev.get("services", 0),
                dev.get("security", "Неизвестно"),
                scan_time,
            ),
        )
    conn.commit()


def write_bt_log(conn, user_id, admin_id, device_name: str, mac: str, description: str):
    cur = conn.cursor()
    cur.execute(
        """
        INSERT INTO bt_logs (user_id, admin_id, device_name, mac_address, scan_time, description)
        VALUES (%s, %s, %s, %s, NOW(), %s)
        """,
        (user_id, admin_id, device_name, mac, description),
    )
    conn.commit()


def fetch_all_bt_logs(conn) -> list:
    cur = conn.cursor()
    cur.execute(
        """
        SELECT log_id, user_id, admin_id, device_name, mac_address, scan_time, description
        FROM bt_logs
        ORDER BY log_id DESC
        """
    )
    return cur.fetchall()


def delete_bt_log(conn, log_id: int, admin_id: int, device_name: str, mac: str, scan_time):
    cur = conn.cursor()
    cur.execute("DELETE FROM bt_logs WHERE log_id = %s", (log_id,))
    cur.execute(
        """
        INSERT INTO bt_logs (user_id, admin_id, device_name, mac_address, scan_time, description)
        VALUES (%s, %s, %s, %s, %s, %s)
        """,
        (None, admin_id, device_name, mac, scan_time,
         f"Удалён администратором (ID {admin_id})"),
    )
    conn.commit()


def update_bt_log_description(conn, log_id: int, admin_id: int, new_desc: str):
    cur = conn.cursor()
    cur.execute(
        "UPDATE bt_logs SET description = %s, admin_id = %s WHERE log_id = %s",
        (new_desc, admin_id, log_id),
    )
    conn.commit()


def fetch_bt_security_stats(conn) -> dict:
    try:
        conn.rollback()
        cur = conn.cursor()
        cur.execute("SELECT mac_address, security_status FROM bt_scan_results")
        rows = cur.fetchall()
    except Exception:
        conn.rollback()
        return {}

    priority_index = {
        "Небезопасно": 1, "Удовлетворительно": 2,
        "Неизвестно": 3, "Безопасно": 4,
    }
    worst: dict = {}
    for mac, status in rows:
        pri = priority_index.get(status, 3)
        if mac not in worst or pri < worst[mac][0]:
            worst[mac] = (pri, status)

    result: dict = {}
    for _, (_, status) in worst.items():
        result[status] = result.get(status, 0) + 1
    return result