Telezab/telezab.py
2024-07-10 22:48:51 +05:00

560 lines
24 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from flask import Flask, request, jsonify
import telebot
from dotenv import load_dotenv
import os
import hashlib
import logging
from logging.config import dictConfig
from threading import Thread, Lock
import sqlite3
import time
# Load environment variables
load_dotenv()
# Configure logging
DEBUG_LOGGING = os.getenv('DEBUG_LOGGING', 'false').lower() == 'true'
if DEBUG_LOGGING:
log_level = 'DEBUG'
else:
log_level = 'INFO'
dictConfig({
'version': 1,
'formatters': {'default': {
'format': '[%(asctime)s] %(levelname)s in %(module)s: %(message)s',
}},
'handlers': {'wsgi': {
'class': 'logging.StreamHandler',
'stream': 'ext://flask.logging.wsgi_errors_stream',
'formatter': 'default'
}},
'root': {
'level': log_level,
'handlers': ['wsgi']
}
})
app = Flask(__name__)
# Get the token from environment variables
TOKEN = os.getenv('TELEGRAM_TOKEN')
if not TOKEN:
raise ValueError("No TELEGRAM_TOKEN found in environment variables")
ADMIN_CHAT_IDS = os.getenv('ADMIN_CHAT_IDS', '').split(',')
bot = telebot.TeleBot(TOKEN)
# Lock for database operations
db_lock = Lock()
# Initialize SQLite database
def init_db():
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
# Create events table
cursor.execute('''CREATE TABLE IF NOT EXISTS events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
hash TEXT UNIQUE,
data TEXT,
delivered BOOLEAN)''')
# Create subscriptions table with username and active flag
cursor.execute('''CREATE TABLE IF NOT EXISTS subscriptions (
chat_id INTEGER,
region_id TEXT,
username TEXT,
active BOOLEAN DEFAULT TRUE,
skip BOOLEAN DEFAULT FALSE,
UNIQUE(chat_id, region_id))''')
# Create whitelist table
cursor.execute('''CREATE TABLE IF NOT EXISTS whitelist (
chat_id INTEGER PRIMARY KEY)''')
# Create regions table with active flag
cursor.execute('''CREATE TABLE IF NOT EXISTS regions (
region_id TEXT PRIMARY KEY,
region_name TEXT,
active BOOLEAN DEFAULT TRUE)''')
# Create user events table for logging
cursor.execute('''CREATE TABLE IF NOT EXISTS user_events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id INTEGER,
username TEXT,
action TEXT,
timestamp TEXT)''')
# Insert sample regions
cursor.execute('''INSERT OR IGNORE INTO regions (region_id, region_name) VALUES
('01', 'Адыгея'),
('02', 'Башкортостан (Уфа)'),
('04', 'Алтай'),
('19', 'Республика Хакасия')''')
conn.commit()
conn.close()
# Hash the incoming data
def hash_data(data):
return hashlib.sha256(str(data).encode('utf-8')).hexdigest()
# Check if user is in whitelist
def is_whitelisted(chat_id):
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
query = 'SELECT COUNT(*) FROM whitelist WHERE chat_id = ?'
app.logger.debug(f"Executing query: {query} with chat_id={chat_id}")
cursor.execute(query, (chat_id,))
count = cursor.fetchone()[0]
conn.close()
return count > 0
# Add user to whitelist
def add_to_whitelist(chat_id):
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
query = 'INSERT OR IGNORE INTO whitelist (chat_id) VALUES (?)'
app.logger.debug(f"Executing query: {query} with chat_id={chat_id}")
cursor.execute(query, (chat_id,))
conn.commit()
conn.close()
# Remove user from whitelist
def remove_from_whitelist(chat_id):
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
query = 'DELETE FROM whitelist WHERE chat_id = ?'
app.logger.debug(f"Executing query: {query} with chat_id={chat_id}")
cursor.execute(query, (chat_id,))
conn.commit()
conn.close()
# Get list of regions
def get_regions():
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
cursor.execute('SELECT region_id, region_name FROM regions WHERE active = TRUE ORDER BY region_id')
regions = cursor.fetchall()
conn.close()
return regions
# Get list of regions a user is subscribed to
def get_user_subscribed_regions(chat_id):
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
cursor.execute('''
SELECT regions.region_id, regions.region_name
FROM subscriptions
JOIN regions ON subscriptions.region_id = regions.region_id
WHERE subscriptions.chat_id = ? AND subscriptions.active = TRUE AND subscriptions.skip = FALSE
ORDER BY regions.region_id
''', (chat_id,))
regions = cursor.fetchall()
conn.close()
return regions
# Format regions list
def format_regions_list(regions):
return '\n'.join([f"{region_id} - {region_name}" for region_id, region_name in regions])
# Log user events
def log_user_event(chat_id, username, action):
timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
query = 'INSERT INTO user_events (chat_id, username, action, timestamp) VALUES (?, ?, ?, ?)'
app.logger.debug(f"Executing query: {query} with chat_id={chat_id}, username={username}, action={action}, timestamp={timestamp}")
cursor.execute(query, (chat_id, username, action, timestamp))
conn.commit()
conn.close()
# Main menu for users
def show_main_menu(chat_id, is_whitelisted_user):
markup = telebot.types.ReplyKeyboardMarkup(one_time_keyboard=True, resize_keyboard=True)
if is_whitelisted_user:
markup.add('Подписаться', 'Отписаться', 'Мои подписки', 'Активные регионы', 'Помощь', 'Добавить регион', 'Удалить регион', 'Добавить в белый список', 'Удалить из белого списка')
else:
markup.add('Регистрация')
bot.send_message(chat_id, "Выберите действие:", reply_markup=markup)
# Handle /start command
@bot.message_handler(commands=['start'])
def handle_start(message):
chat_id = message.chat.id
username = message.from_user.username
if username:
username = f"@{username}"
else:
username = "N/A"
if is_whitelisted(chat_id):
show_main_menu(chat_id, is_whitelisted_user=True)
else:
show_main_menu(chat_id, is_whitelisted_user=False)
app.logger.info(f"User {chat_id} (@{username}) started with command /start.")
# Handle menu button presses
@bot.message_handler(func=lambda message: True)
def handle_menu_selection(message):
chat_id = message.chat.id
text = message.text.strip().lower()
if text == 'регистрация':
handle_register(message)
elif text == 'подписаться':
handle_subscribe(message)
elif text == 'отписаться':
handle_unsubscribe(message)
elif text == 'мои подписки':
handle_my_subscriptions(message)
elif text == 'активные регионы':
handle_active_regions(message)
elif text == 'помощь':
handle_help(message)
elif text == 'добавить регион' and str(chat_id) in ADMIN_CHAT_IDS:
prompt_admin_for_region(chat_id, 'add')
elif text == 'удалить регион' and str(chat_id) in ADMIN_CHAT_IDS:
prompt_admin_for_region(chat_id, 'remove')
elif text == 'добавить в белый список' and str(chat_id) in ADMIN_CHAT_IDS:
prompt_admin_for_whitelist(chat_id, 'add')
elif text == 'удалить из белого списка' and str(chat_id) in ADMIN_CHAT_IDS:
prompt_admin_for_whitelist(chat_id, 'remove')
else:
bot.send_message(chat_id, "Команда не распознана или у вас нет прав для выполнения этой команды.")
show_main_menu(chat_id, is_whitelisted(chat_id))
# Handle /subscribe command to subscribe to a region
@bot.message_handler(commands=['subscribe', 'sub'])
def handle_subscribe(message):
chat_id = message.chat.id
if not is_whitelisted(chat_id):
bot.send_message(chat_id, "Вы не авторизованы для использования этого бота.")
app.logger.info(f"Unauthorized access attempt by {chat_id}")
return
username = message.from_user.username
if username:
username = f"@{username}"
else:
username = "N/A"
regions_list = format_regions_list(get_regions())
bot.send_message(chat_id, f"Отправьте номер или номера регионов, на которые хотите подписаться (через запятую):\n{regions_list}\n\nНапишите 'отмена' для отмены.")
bot.register_next_step_handler(message, process_subscription, chat_id, username)
def process_subscription(message, chat_id, username):
if message.text.lower() == 'отмена':
bot.send_message(chat_id, "Действие отменено.")
show_main_menu(chat_id, is_whitelisted(chat_id))
return
region_ids = message.text.split(',')
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
for region_id in region_ids:
region_id = region_id.strip()
query = 'INSERT OR IGNORE INTO subscriptions (chat_id, region_id, username, active) VALUES (?, ?, ?, TRUE)'
app.logger.debug(f"Executing query: {query} with chat_id={chat_id}, region_id={region_id}, username={username}")
cursor.execute(query, (chat_id, region_id, username))
if cursor.rowcount == 0:
query = 'UPDATE subscriptions SET active = TRUE WHERE chat_id = ? AND region_id = ?'
app.logger.debug(f"Executing query: {query} with chat_id={chat_id}, region_id={region_id}")
cursor.execute(query, (chat_id, region_id))
conn.commit()
conn.close()
bot.send_message(chat_id, f"Подписка на регионы: {', '.join(region_ids)} оформлена.")
app.logger.info(f"User {chat_id} (@{username}) subscribed to regions: {', '.join(region_ids)}.")
log_user_event(chat_id, username, f"Subscribed to regions: {', '.join(region_ids)}")
show_main_menu(chat_id, is_whitelisted(chat_id))
# Handle /unsubscribe command to unsubscribe from a region
@bot.message_handler(commands=['unsubscribe'])
def handle_unsubscribe(message):
chat_id = message.chat.id
if not is_whitelisted(chat_id):
bot.send_message(chat_id, "Вы не авторизованы для использования этого бота.")
app.logger.info(f"Unauthorized access attempt by {chat_id}")
return
user_regions = get_user_subscribed_regions(chat_id)
if not user_regions:
bot.send_message(chat_id, "Вы не подписаны ни на один регион.")
else:
regions_list = format_regions_list(user_regions)
bot.send_message(chat_id, f"Отправьте номер или номера регионов, от которых хотите отписаться (через запятую):\n{regions_list}\n\nНапишите 'отмена' для отмены.")
bot.register_next_step_handler(message, process_unsubscription, chat_id)
def process_unsubscription(message, chat_id):
if message.text.lower() == 'отмена':
bot.send_message(chat_id, "Действие отменено.")
show_main_menu(chat_id, is_whitelisted(chat_id))
return
region_ids = message.text.split(',')
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
for region_id in region_ids:
region_id = region_id.strip()
query = 'UPDATE subscriptions SET active = FALSE WHERE chat_id = ? AND region_id = ?'
app.logger.debug(f"Executing query: {query} with chat_id={chat_id}, region_id={region_id}")
cursor.execute(query, (chat_id, region_id))
conn.commit()
conn.close()
bot.send_message(chat_id, f"Отписка от регионов: {', '.join(region_ids)} оформлена.")
username = message.from_user.username
if username:
username = f"@{username}"
else:
username = "N/A"
app.logger.info(f"User {chat_id} (@{username}) unsubscribed from regions: {', '.join(region_ids)}.")
log_user_event(chat_id, username, f"Unsubscribed from regions: {', '.join(region_ids)}")
show_main_menu(chat_id, is_whitelisted(chat_id))
# Handle /help command to provide instructions
@bot.message_handler(commands=['help'])
def handle_help(message):
help_text = (
"/start - Начать работу с ботом\n"
"/subscribe (/sub) - Подписаться на рассылку событий по региону. Укажите номер региона, например, /subscribe 01 - Адыгея\n"
"/unsubscribe - Отписаться от рассылки событий по региону. Укажите номер региона, например, /unsubscribe 01 - Адыгея\n"
"/register - Запросить регистрацию в боте"
)
bot.send_message(message.chat.id, help_text)
show_main_menu(message.chat.id, is_whitelisted(message.chat.id))
# Handle /register command for new user registration
@bot.message_handler(commands=['register'])
def handle_register(message):
chat_id = message.chat.id
username = message.from_user.username
if username:
username = f"@{username}"
else:
username = "N/A"
markup = telebot.types.ReplyKeyboardMarkup(one_time_keyboard=True, resize_keyboard=True)
markup.add('Подтвердить регистрацию', 'Отмена')
bot.send_message(chat_id, f"Ваш chat ID: {chat_id}, ваше имя пользователя: @{username}. Запрос на одобрение отправлен администратору.", reply_markup=markup)
log_user_event(chat_id, username, "Requested registration")
bot.register_next_step_handler(message, process_register, chat_id, username)
def process_register(message, chat_id, username):
if message.text.lower() == 'отмена':
bot.send_message(chat_id, "Регистрация отменена.")
show_main_menu(chat_id, is_whitelisted(chat_id))
return
if message.text.lower() == 'подтвердить регистрацию':
for admin_chat_id in ADMIN_CHAT_IDS:
markup = telebot.types.ReplyKeyboardMarkup(one_time_keyboard=True, resize_keyboard=True)
markup.add(f'/add_whitelist {chat_id}', 'Отмена')
bot.send_message(
admin_chat_id,
f"Пользователь {username} ({chat_id}) запрашивает регистрацию.\n"
f"Вы подтверждаете это действие?",
reply_markup=markup
)
bot.send_message(chat_id, "Запрос отправлен администратору для одобрения.")
app.logger.info(f"User {chat_id} (@{username}) requested registration.")
else:
bot.send_message(chat_id, "Некорректный выбор. Регистрация отменена.")
show_main_menu(chat_id, is_whitelisted(chat_id))
# Handle admin region management commands
def prompt_admin_for_region(chat_id, action):
if action == 'add':
bot.send_message(chat_id, "Введите ID и название региона в формате: <region_id> <region_name>")
bot.register_next_step_handler_by_chat_id(chat_id, process_add_region)
elif action == 'remove':
bot.send_message(chat_id, "Введите ID региона, который хотите сделать неактивным")
bot.register_next_step_handler_by_chat_id(chat_id, process_remove_region)
def process_add_region(message):
chat_id = message.chat.id
try:
region_id, region_name = message.text.split()[0], ' '.join(message.text.split()[1:])
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
query = 'INSERT OR IGNORE INTO regions (region_id, region_name) VALUES (?, ?)'
app.logger.debug(f"Executing query: {query} with region_id={region_id}, region_name={region_name}")
cursor.execute(query, (region_id, region_name))
conn.commit()
conn.close()
bot.send_message(chat_id, f"Регион {region_id} - {region_name} добавлен.")
app.logger.info(f"Admin {chat_id} added region {region_id} - {region_name}.")
except (IndexError, ValueError):
bot.send_message(chat_id, "Неверный формат. Используйте: <region_id> <region_name>")
show_main_menu(chat_id, is_whitelisted(chat_id))
def process_remove_region(message):
chat_id = message.chat.id
try:
region_id = message.text.split()[0]
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
query = 'UPDATE regions SET active = FALSE WHERE region_id = ?'
app.logger.debug(f"Executing query: {query} with region_id={region_id}")
cursor.execute(query, (region_id,))
query = 'UPDATE subscriptions SET active = FALSE WHERE region_id = ? AND active = TRUE'
app.logger.debug(f"Executing query: {query} with region_id={region_id}")
cursor.execute(query, (region_id,))
conn.commit()
conn.close()
bot.send_message(chat_id, f"Регион {region_id} теперь неактивен и все активные подписки обновлены.")
app.logger.info(f"Admin {chat_id} set region {region_id} to inactive and updated subscriptions.")
except IndexError:
bot.send_message(chat_id, "Неверный формат. Используйте: <region_id>")
show_main_menu(chat_id, is_whitelisted(chat_id))
# Handle admin whitelist management commands
def prompt_admin_for_whitelist(chat_id, action):
if action == 'add':
bot.send_message(chat_id, "Введите ID пользователя, которого хотите добавить в белый список")
bot.register_next_step_handler_by_chat_id(chat_id, process_add_whitelist)
elif action == 'remove':
bot.send_message(chat_id, "Введите ID пользователя, которого хотите удалить из белого списка")
bot.register_next_step_handler_by_chat_id(chat_id, process_remove_whitelist)
def process_add_whitelist(message):
chat_id = message.chat.id
try:
new_chat_id = int(message.text.split()[0])
add_to_whitelist(new_chat_id)
bot.send_message(chat_id, f"Chat ID {new_chat_id} добавлен в белый список.")
app.logger.info(f"Admin {chat_id} added {new_chat_id} to the whitelist.")
log_user_event(new_chat_id, "N/A", f"Added to whitelist by {chat_id} (@{message.from_user.username})")
except (IndexError, ValueError):
bot.send_message(chat_id, "Неверный формат. Используйте: <chat_id>")
show_main_menu(chat_id, is_whitelisted(chat_id))
def process_remove_whitelist(message):
chat_id = message.chat.id
try:
remove_chat_id = int(message.text.split()[0])
remove_from_whitelist(remove_chat_id)
bot.send_message(chat_id, f"Chat ID {remove_chat_id} удален из белого списка.")
app.logger.info(f"Admin {chat_id} removed {remove_chat_id} from the whitelist.")
log_user_event(remove_chat_id, "N/A", f"Removed from whitelist by {chat_id} (@{message.from_user.username})")
except (IndexError, ValueError):
bot.send_message(chat_id, "Неверный формат. Используйте: <chat_id>")
show_main_menu(chat_id, is_whitelisted(chat_id))
# Handle displaying active subscriptions for a user
def handle_my_subscriptions(message):
chat_id = message.chat.id
if not is_whitelisted(chat_id):
bot.send_message(chat_id, "Вы не авторизованы для использования этого бота.")
app.logger.info(f"Unauthorized access attempt by {chat_id}")
return
user_regions = get_user_subscribed_regions(chat_id)
if not user_regions:
bot.send_message(chat_id, "Вы не подписаны ни на один регион.")
else:
regions_list = format_regions_list(user_regions)
bot.send_message(chat_id, f"Ваши активные подписки:\n{regions_list}")
show_main_menu(chat_id, is_whitelisted(chat_id))
# Handle displaying all active regions
def handle_active_regions(message):
chat_id = message.chat.id
if not is_whitelisted(chat_id):
bot.send_message(chat_id, "Вы не авторизованы для использования этого бота.")
app.logger.info(f"Unauthorized access attempt by {chat_id}")
return
regions = get_regions()
if not regions:
bot.send_message(chat_id, "Нет активных регионов.")
else:
regions_list = format_regions_list(regions)
bot.send_message(chat_id, f"Активные регионы:\n{regions_list}")
show_main_menu(chat_id, is_whitelisted(chat_id))
@app.route('/webhook', methods=['POST'])
def webhook():
data = request.get_json()
app.logger.info(f"Received data: {data}")
event_hash = hash_data(data)
with db_lock:
conn = sqlite3.connect('telezab.db')
cursor = conn.cursor()
cursor.execute('SELECT COUNT(*) FROM events')
count = cursor.fetchone()[0]
if count >= 200:
query = 'DELETE FROM events WHERE id = (SELECT MIN(id) FROM events)'
app.logger.debug(f"Executing query: {query}")
cursor.execute(query)
query = 'INSERT OR IGNORE INTO events (hash, data, delivered) VALUES (?, ?, ?)'
app.logger.debug(f"Executing query: {query} with hash={event_hash}, data={data}, delivered={False}")
cursor.execute(query, (event_hash, str(data), False))
# Fetch chat_ids to send the alert
region_id = data.get("region")
query = 'SELECT chat_id, username FROM subscriptions WHERE region_id = ? AND active = TRUE AND skip = FALSE'
app.logger.debug(f"Executing query: {query} with region_id={region_id}")
cursor.execute(query, (region_id,))
results = cursor.fetchall()
# Check if the region is active
query = 'SELECT active FROM regions WHERE region_id = ?'
cursor.execute(query, (region_id,))
region_active = cursor.fetchone()[0]
if region_active:
message = format_message(data)
for chat_id, username in results:
try:
app.logger.debug(f"Sending message: {message} to chat_id={chat_id}, username={username}")
bot.send_message(chat_id, message)
app.logger.info(f"Sent alert to {chat_id} ({username}) for region {region_id}")
except telebot.apihelper.ApiTelegramException as e:
app.logger.error(f"Failed to send message to {chat_id} ({username}): {e}")
except Exception as e:
app.logger.error(f"Error sending message to {chat_id} ({username}): {e}")
conn.commit()
conn.close()
return jsonify({"status": "success"}), 200
def format_message(data):
return (f"Zabbix Alert\n"
f"Host: {data['host']}\n"
f"Item: {data['item']}\n"
f"Trigger: {data['trigger']}\n"
f"Value: {data['value']}")
def run_polling():
bot.polling(none_stop=True, interval=0)
if __name__ == '__main__':
init_db()
# Start Flask app in a separate thread
Thread(target=app.run, kwargs={'port': 5000, 'debug': False}, daemon=True).start()
# Start bot polling
run_polling()