124 lines
4.9 KiB
Python
124 lines
4.9 KiB
Python
import sqlite3
|
|
from datetime import datetime
|
|
import json
|
|
|
|
class Database:
|
|
def __init__(self, db_path="bot.db"):
|
|
self.db_path = db_path
|
|
self.init_db()
|
|
|
|
def get_connection(self):
|
|
return sqlite3.connect(self.db_path)
|
|
|
|
def init_db(self):
|
|
with self.get_connection() as conn:
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
user_id INTEGER PRIMARY KEY,
|
|
username TEXT,
|
|
full_name TEXT,
|
|
language_code TEXT,
|
|
first_seen TIMESTAMP,
|
|
is_banned INTEGER DEFAULT 0
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER,
|
|
content_type TEXT,
|
|
text TEXT,
|
|
json_data TEXT,
|
|
timestamp TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users(user_id)
|
|
)
|
|
""")
|
|
conn.commit()
|
|
|
|
def log_user(self, user):
|
|
with self.get_connection() as conn:
|
|
# Check if user already exists to preserve is_banned status if we were using INSERT OR REPLACE
|
|
# But INSERT OR REPLACE would overwrite is_banned if we don't include it.
|
|
# Better to use INSERT OR IGNORE and then UPDATE or a more complex query.
|
|
conn.execute("""
|
|
INSERT INTO users (user_id, username, full_name, language_code, first_seen)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
ON CONFLICT(user_id) DO UPDATE SET
|
|
username=excluded.username,
|
|
full_name=excluded.full_name,
|
|
language_code=excluded.language_code
|
|
""", (user.id, user.username, user.full_name, user.language_code, datetime.now()))
|
|
conn.commit()
|
|
|
|
def ban_user(self, user_id):
|
|
with self.get_connection() as conn:
|
|
conn.execute("UPDATE users SET is_banned = 1 WHERE user_id = ?", (user_id,))
|
|
conn.commit()
|
|
|
|
def unban_user(self, user_id):
|
|
with self.get_connection() as conn:
|
|
conn.execute("UPDATE users SET is_banned = 0 WHERE user_id = ?", (user_id,))
|
|
conn.commit()
|
|
|
|
def is_user_banned(self, user_id):
|
|
with self.get_connection() as conn:
|
|
cursor = conn.execute("SELECT is_banned FROM users WHERE user_id = ?", (user_id,))
|
|
row = cursor.fetchone()
|
|
return bool(row[0]) if row else False
|
|
|
|
def log_message(self, message):
|
|
self.log_user(message.from_user)
|
|
with self.get_connection() as conn:
|
|
content_type = message.content_type
|
|
text = message.text or message.caption or ""
|
|
try:
|
|
json_data = message.model_dump_json()
|
|
except Exception:
|
|
# Fallback if model_dump_json fails due to non-serializable types
|
|
try:
|
|
import json
|
|
# Use model_dump to get a dict, then serialize manually
|
|
# exclude_none=True can help, and we can handle unknown types
|
|
data_dict = message.model_dump(exclude_none=True)
|
|
def custom_serializer(obj):
|
|
return str(obj)
|
|
json_data = json.dumps(data_dict, default=custom_serializer)
|
|
except Exception as e:
|
|
json_data = json.dumps({"error": f"Failed to serialize message: {str(e)}"})
|
|
|
|
cursor = conn.execute("""
|
|
INSERT INTO messages (user_id, content_type, text, json_data, timestamp)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""", (message.from_user.id, content_type, text, json_data, datetime.now()))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
|
|
def get_stats(self):
|
|
with self.get_connection() as conn:
|
|
cursor = conn.execute("SELECT COUNT(*) FROM users")
|
|
user_count = cursor.fetchone()[0]
|
|
return {"user_count": user_count}
|
|
|
|
def get_recent_messages(self, limit=10):
|
|
with self.get_connection() as conn:
|
|
cursor = conn.execute("""
|
|
SELECT m.id, u.full_name, m.text, m.timestamp
|
|
FROM messages m
|
|
JOIN users u ON m.user_id = u.user_id
|
|
ORDER BY m.timestamp DESC
|
|
LIMIT ?
|
|
""", (limit,))
|
|
return cursor.fetchall()
|
|
|
|
def get_message_details(self, msg_id):
|
|
with self.get_connection() as conn:
|
|
cursor = conn.execute("""
|
|
SELECT m.id, u.full_name, u.username, u.user_id, m.content_type, m.timestamp, m.text
|
|
FROM messages m
|
|
JOIN users u ON m.user_id = u.user_id
|
|
WHERE m.id = ?
|
|
""", (msg_id,))
|
|
return cursor.fetchone()
|
|
|
|
db = Database()
|