146 lines
5.1 KiB
Python
146 lines
5.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 numbers (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
phone TEXT UNIQUE NOT NULL,
|
||
bonus_card 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_created ON numbers(created_at)")
|
||
conn.execute("CREATE INDEX IF NOT EXISTS idx_used ON numbers(used_at)")
|
||
|
||
|
||
# ========== Функции для работы с номерами и бонусными картами ==========
|
||
def add_number(phone: str, bonus_card: str = None) -> bool:
|
||
"""Добавляет номер и опционально бонусную карту"""
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO numbers (phone, bonus_card, created_at) VALUES (?, ?, ?)",
|
||
(phone, bonus_card, datetime.now().isoformat())
|
||
)
|
||
return True
|
||
except sqlite3.IntegrityError:
|
||
return False
|
||
|
||
|
||
def update_bonus_card(phone: str, bonus_card: str) -> bool:
|
||
"""Обновляет бонусную карту для существующего номера"""
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
cursor = conn.execute(
|
||
"UPDATE numbers SET bonus_card = ? WHERE phone = ?",
|
||
(bonus_card, phone)
|
||
)
|
||
return cursor.rowcount > 0
|
||
|
||
|
||
def get_random_from_oldest():
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.row_factory = sqlite3.Row
|
||
|
||
total_count = conn.execute("SELECT COUNT(*) as count FROM numbers").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 numbers ORDER BY sort_date ASC LIMIT ?",
|
||
(limit,)
|
||
).fetchall()
|
||
|
||
if not oldest:
|
||
return None
|
||
|
||
selected = random.choice(oldest)
|
||
|
||
conn.execute(
|
||
"UPDATE numbers SET usage_count = usage_count + 1, used_at = ? WHERE id = ?",
|
||
(datetime.now().isoformat(), selected["id"])
|
||
)
|
||
|
||
return dict(selected)
|
||
|
||
|
||
def get_all_numbers():
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.row_factory = sqlite3.Row
|
||
rows = conn.execute("SELECT * FROM numbers ORDER BY created_at DESC").fetchall()
|
||
return [dict(row) for row in rows]
|
||
|
||
|
||
def delete_number(record_id: int):
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.execute("DELETE FROM numbers WHERE id = ?", (record_id,))
|
||
|
||
|
||
def get_bonus_card(record_id: int) -> str | None:
|
||
"""Получить бонусную карту по ID записи"""
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.row_factory = sqlite3.Row
|
||
result = conn.execute("SELECT bonus_card FROM numbers WHERE id = ?", (record_id,)).fetchone()
|
||
return result["bonus_card"] if result else None
|
||
|
||
|
||
def get_bonus_card_by_phone(phone: str) -> str | None:
|
||
"""Получить бонусную карту по номеру телефона"""
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.row_factory = sqlite3.Row
|
||
result = conn.execute("SELECT bonus_card FROM numbers WHERE phone = ?", (phone,)).fetchone()
|
||
return result["bonus_card"] if result else None
|
||
|
||
|
||
# ========== Функции для белого списка ==========
|
||
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] |