From b76916ae98f8fb2c83ca9155a025001870c2c51c Mon Sep 17 00:00:00 2001
From: erius <egor_kapralov2003@mail.ru>
Date: Thu, 31 Oct 2024 01:26:44 +0300
Subject: [PATCH] Added whitelist, Implemented env configuration

---
 .env.example     |  3 +++
 .gitignore       |  1 +
 bot.py           | 61 ++++++++++++++++++++++++++++++++++++++----------
 requirements.txt | 12 ++++++++++
 4 files changed, 65 insertions(+), 12 deletions(-)
 create mode 100644 .env.example
 create mode 100644 requirements.txt

diff --git a/.env.example b/.env.example
new file mode 100644
index 0000000..ff601c1
--- /dev/null
+++ b/.env.example
@@ -0,0 +1,3 @@
+BOT_TOKEN='TOKEN'
+YA_TOKEN='TOKEN'
+WHITELISTED=user1,user2,user3
diff --git a/.gitignore b/.gitignore
index 18bfee6..178b645 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
 .venv
+.env
 
diff --git a/bot.py b/bot.py
index 4acf154..641018a 100644
--- a/bot.py
+++ b/bot.py
@@ -1,6 +1,10 @@
-import logging, requests, json
+import logging
+import requests
+import json
+import os
 from telegram import Update, ReplyKeyboardMarkup
-from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, MessageHandler, filters
+from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, MessageHandler, ConversationHandler, filters
+from dotenv import load_dotenv
 
 logging.basicConfig(
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
@@ -8,8 +12,12 @@ logging.basicConfig(
 )
 
 logger = logging.getLogger(__name__)
-bot_token = '7636065770:AAF8UV5Vb2F7aRuO8EYTyv0xGBmd3dQeonQ'
-ya_token = 'y0_AgAAAABiFlEmAAbJ0QAAAAEWkS6NAABACUMU8JdJWoZeypr0qWs05ayGjQ'
+
+load_dotenv()
+bot_token = os.getenv('BOT_TOKEN')
+ya_token = os.getenv('YA_TOKEN')
+whitelisted = os.getenv('WHITELISTED')
+
 ya_url = 'https://partner2.yandex.ru/api/statistics2/get.json'
 msg_to_period = {
     "Сегодня": "today",
@@ -18,13 +26,20 @@ msg_to_period = {
 }
 
 
-async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    reply_keyboard = [[ "Сегодня", "Вчера", "Месяц" ]] 
+async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+    if update.effective_user.username not in whitelisted:
+        await update.message.reply_text("Запрет")
+        return ConversationHandler.END
+    reply_keyboard = [["Сегодня", "Вчера", "Месяц"]]
     await update.message.reply_text("Какую стату вывести", reply_markup=ReplyKeyboardMarkup(reply_keyboard))
+    return 0
 
-async def ya_api(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+
+async def ya_api(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
     reply = await fetch_data(update.message.text)
     await update.message.reply_text(reply)
+    return 0
+
 
 async def fetch_data(msg_period: str) -> str:
     response = requests.get(ya_url, headers={'Authorization': ya_token}, params={
@@ -33,7 +48,6 @@ async def fetch_data(msg_period: str) -> str:
         'field': ['shows', 'hits_render', 'partner_wo_nds', 'cpmv_partner_wo_nds']
     })
     json_data = json.loads(response.text)
-    logger.info(json_data)
     data = json_data['data']['totals']['2'][0]
     return f"Данные за {msg_period}\n\n\
     Подборы рекламы: {data['hits_render']}\n\
@@ -42,10 +56,33 @@ async def fetch_data(msg_period: str) -> str:
     Стоимость 1к видимых показов: {data['cpmv_partner_wo_nds']} руб."
 
 
+async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+    await update.message.reply_text("домой")
+    return ConversationHandler.END
+
+
 if __name__ == '__main__':
+    if ya_token is None:
+        logger.fatal("YA_TOKEN env variable is required")
+        exit(1)
+    if bot_token is None:
+        logger.fatal("BOT_TOKEN env variable is required")
+        exit(1)
+    if whitelisted is None:
+        logger.fatal("WHITELISTED env varialbe is required")
+        exit(1)
+
+    whitelisted = whitelisted.split(',')
+
     application = ApplicationBuilder().token(bot_token).build()
-    
-    application.add_handler(CommandHandler('start', start))
-    application.add_handler(MessageHandler(filters.Regex("^(Сегодня|Вчера|Месяц)$"), ya_api))
-    
+
+    conv_handler = ConversationHandler(
+        entry_points=[CommandHandler('start', start)],
+        states={
+            0: [MessageHandler(filters.Regex("^(Сегодня|Вчера|Месяц)$"), ya_api)]
+        },
+        fallbacks=[CommandHandler('cancel', cancel)]
+    )
+    application.add_handler(conv_handler)
+
     application.run_polling()
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..241bfd3
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,12 @@
+anyio==4.6.2.post1
+certifi==2024.8.30
+charset-normalizer==3.4.0
+h11==0.14.0
+httpcore==1.0.6
+httpx==0.27.2
+idna==3.10
+python-dotenv==1.0.1
+python-telegram-bot==21.6
+requests==2.32.3
+sniffio==1.3.1
+urllib3==2.2.3