125 lines
4.1 KiB
Python
125 lines
4.1 KiB
Python
import sqlite3
|
|
import random
|
|
from datetime import datetime
|
|
|
|
DB_NAME = "bot.db"
|
|
|
|
|
|
def init_db():
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
# Единая таблица для всех записей
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS records (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
type TEXT NOT NULL,
|
|
phone TEXT,
|
|
qr_code TEXT,
|
|
qr_data TEXT,
|
|
created_at TEXT NOT NULL,
|
|
used_at TEXT,
|
|
usage_count INTEGER DEFAULT 0
|
|
)
|
|
""")
|
|
|
|
# Таблица белого списка
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS allowed_users (
|
|
user_id INTEGER PRIMARY KEY,
|
|
username TEXT,
|
|
added_at TEXT NOT NULL,
|
|
added_by TEXT
|
|
)
|
|
""")
|
|
|
|
# Индексы
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_records_used ON records(used_at)")
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_records_type ON records(type)")
|
|
|
|
|
|
def add_record(record_type: str, phone: str = None, qr_code: str = None, qr_data: str = None) -> bool:
|
|
"""Добавляет запись"""
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
try:
|
|
conn.execute("""
|
|
INSERT INTO records (type, phone, qr_code, qr_data, created_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""", (record_type, phone, qr_code, qr_data, datetime.now().isoformat()))
|
|
return True
|
|
except sqlite3.IntegrityError:
|
|
return False
|
|
|
|
|
|
def get_random_record():
|
|
"""Возвращает случайную запись"""
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
|
|
total_count = conn.execute("SELECT COUNT(*) as count FROM records").fetchone()["count"]
|
|
|
|
if total_count <= 50:
|
|
limit = 10
|
|
else:
|
|
limit = max(10, min(50, int(total_count * 0.2)))
|
|
|
|
oldest = conn.execute("""
|
|
SELECT *, COALESCE(used_at, created_at) as sort_date
|
|
FROM records
|
|
ORDER BY sort_date ASC
|
|
LIMIT ?
|
|
""", (limit,)).fetchall()
|
|
|
|
if not oldest:
|
|
return None
|
|
|
|
selected = random.choice(oldest)
|
|
|
|
conn.execute("""
|
|
UPDATE records SET usage_count = usage_count + 1, used_at = ?
|
|
WHERE id = ?
|
|
""", (datetime.now().isoformat(), selected["id"]))
|
|
|
|
return dict(selected)
|
|
|
|
|
|
def get_all_records():
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
rows = conn.execute("SELECT * FROM records ORDER BY created_at DESC").fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
|
|
def delete_record(record_id: int):
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
conn.execute("DELETE FROM records WHERE id = ?", (record_id,))
|
|
|
|
|
|
# ========== Функции для белого списка ==========
|
|
def add_allowed_user(user_id: int, username: str = None, added_by: str = "admin") -> bool:
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO allowed_users (user_id, username, added_at, added_by) VALUES (?, ?, ?, ?)",
|
|
(user_id, username, datetime.now().isoformat(), added_by)
|
|
)
|
|
return True
|
|
except sqlite3.IntegrityError:
|
|
return False
|
|
|
|
|
|
def remove_allowed_user(user_id: int) -> bool:
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
cursor = conn.execute("DELETE FROM allowed_users WHERE user_id = ?", (user_id,))
|
|
return cursor.rowcount > 0
|
|
|
|
|
|
def is_user_allowed(user_id: int) -> bool:
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
cursor = conn.execute("SELECT 1 FROM allowed_users WHERE user_id = ?", (user_id,))
|
|
return cursor.fetchone() is not None
|
|
|
|
|
|
def get_allowed_users():
|
|
with sqlite3.connect(DB_NAME) as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
rows = conn.execute("SELECT * FROM allowed_users ORDER BY added_at DESC").fetchall()
|
|
return [dict(row) for row in rows] |