From 5b531e0cc3ba2be5264a2dc60c3ec31e462557ef Mon Sep 17 00:00:00 2001 From: Aaron Ang <67321817+aaron-ang@users.noreply.github.com> Date: Thu, 9 Jan 2025 00:07:10 -0800 Subject: [PATCH] chore: fix types, add dep, expose admin scripts --- .gitignore | 1 - admin/adhoc.py | 26 ++++++++++++++ admin/announcement.py | 80 +++++++++++++++++++++++++++++++++++++++++++ admin/commands.py | 36 +++++++++++++++++++ requirements.txt | 1 + src/bot.py | 1 - src/finder.py | 6 ++-- utils/db.py | 6 ++-- 8 files changed, 149 insertions(+), 8 deletions(-) create mode 100644 admin/adhoc.py create mode 100644 admin/announcement.py create mode 100644 admin/commands.py diff --git a/.gitignore b/.gitignore index b5939ab..87eff85 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,3 @@ .env .vscode/ __pycache__/ -admin/ diff --git a/admin/adhoc.py b/admin/adhoc.py new file mode 100644 index 0000000..65a3188 --- /dev/null +++ b/admin/adhoc.py @@ -0,0 +1,26 @@ +import sys + +sys.path.append("./") +from utils.db import Database +from src.bot import Environment + + +def update(env: Environment): + database = Database(env) + + for course in database.get_all_courses(): + for user in course["users"]: + database.update_subscription_status(user, course["name"], True) + + for user in database.get_all_users(): + if "last_subscription" not in user: + database.update_subscription_status(user["user"], "", False) + + +if __name__ == "__main__": + try: + update(Environment.DEV) + except Exception as e: + print(e) + else: + print("Data updated successfully!") diff --git a/admin/announcement.py b/admin/announcement.py new file mode 100644 index 0000000..ec73ebf --- /dev/null +++ b/admin/announcement.py @@ -0,0 +1,80 @@ +"""One-time script to send out announcements to users""" + +import os +import sys +import asyncio +from datetime import datetime +import telegram +from telegram import Message +from dotenv import load_dotenv + +sys.path.append("./") +from utils.db import Database +from utils.constants import Environment, TimeConstants + + +async def send_maintenance_announcement(): + """Send maintenance message to all users""" + announcement = "Terrier Alert has resumed service. Thank you for your patience!" + await broadcast_message(announcement) + + +async def send_live_announcement(): + """Send announcement to all users when service goes back live""" + announcement = ( + "Terrier Alert will be unavailable until further notice as we are upgrading our systems to " + "integrate with the new course search. If you are interested in contributing or maintaining the project, " + "please use the /feedback command to get in touch with us (and specify your Telegram username). " + "Thank you for your patience!" + ) + # announcement = ( + # "Thank you for using Terrier Alert!\n" + # f"*Release Notes ({datetime.now().strftime('%B %-d, %Y')})*\n" + # # "*What's 🆕*\n" + # # "• 🚀 /resubscribe: Received a notification but failed to secure your spot? " + # # "Use this command to quickly subscribe to the same class!\n" + # # "• 🔍 Scraping logic: In light of classes that may reopen, " + # # "Terrier Alert will ignore Closed/Restricted classes (instead of sending a notification) " + # # "but will continue to monitor them for openings\n" + # # "• 🏫 Added *CGS, SPH, SED* to the list of schools\n" + # "*Bug Fixes*\n" + # "• 🔧 Fixed: Some users were not able to subscribe to classes. " + # "We have since resolved this issue and added additional error handling for more visibility in the future. " + # "Apologies for any inconvenience caused!\n" + # ) + await broadcast_message(announcement) + + +async def broadcast_message(message): + """Send message to all users""" + for user in DB.get_all_users(): + try: + msg: Message = await BOT.send_message( + user["user"], + message, + parse_mode="Markdown", + write_timeout=TimeConstants.TIMEOUT_SECONDS, + ) + await msg.pin() + except Exception as e: + print(f"Error sending message to {user['user']}: {e}") + + +async def main(env: Environment): + global BOT, DB + load_dotenv() + bot_token = os.getenv( + "TELEGRAM_TOKEN" if env == Environment.PROD else "TEST_TELEGRAM_TOKEN" + ) + BOT = telegram.Bot(bot_token) + DB = Database(env) + await send_live_announcement() + + +if __name__ == "__main__": + try: + asyncio.run(main(Environment.PROD)) + except Exception as e: + print(e) + else: + print("Announcement sent successfully.") diff --git a/admin/commands.py b/admin/commands.py new file mode 100644 index 0000000..2812767 --- /dev/null +++ b/admin/commands.py @@ -0,0 +1,36 @@ +"""Set bot commands and descriptions""" + +import os +import asyncio +from telegram import Bot, BotCommand +from dotenv import load_dotenv + +COMMANDS = [ + BotCommand("start", "Start the bot"), + BotCommand("subscribe", "Subscribe to a course"), + # BotCommand("register", "Register for a subscribed course"), + BotCommand("resubscribe", "Resubscribe to last subscribed course"), + BotCommand("unsubscribe", "Unsubscribe from a course"), + BotCommand("help", "Important information"), + BotCommand("feedback", "Report bugs and submit feedback"), + BotCommand("about", "Tech stack and source code"), +] + + +async def main(): + load_dotenv() + BOT_TOKENS = [os.getenv("TELEGRAM_TOKEN"), os.getenv("TEST_TELEGRAM_TOKEN")] + for BOT_TOKEN in BOT_TOKENS: + bot = Bot(token=BOT_TOKEN) + successs = await bot.set_my_commands(COMMANDS) + if not successs: + raise Exception(f"Failed to set commands in Bot: {BOT_TOKEN}") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except Exception as e: + print(e) + else: + print("Succesfully set bot commands.") diff --git a/requirements.txt b/requirements.txt index 8ff78f9..e970f35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ certifi +curl_cffi pendulum pymongo python-dotenv diff --git a/src/bot.py b/src/bot.py index 507f760..3adecc6 100644 --- a/src/bot.py +++ b/src/bot.py @@ -15,7 +15,6 @@ Message, Update, constants, - error, ) from telegram.ext import ( ApplicationBuilder, diff --git a/src/finder.py b/src/finder.py index 8ff41f7..97901cb 100644 --- a/src/finder.py +++ b/src/finder.py @@ -42,7 +42,7 @@ timeout: Optional[pendulum.DateTime] = None -def setup_chrome_options(env: str) -> webdriver.ChromeOptions: +def setup_chrome_options(env: Environment) -> webdriver.ChromeOptions: """Configure Chrome options based on environment.""" options = webdriver.ChromeOptions() options.add_argument("--no-sandbox") @@ -59,7 +59,7 @@ def setup_chrome_options(env: str) -> webdriver.ChromeOptions: def init_driver( - env: str, wait_timeout: int = 30 + env: Environment, wait_timeout: int = 30 ) -> Tuple[webdriver.Chrome, WebDriverWait]: """Initialize Chrome driver with appropriate configuration.""" options = setup_chrome_options(env) @@ -71,7 +71,7 @@ def init_driver( return driver, wait -async def register_course(env: str, user_cache: dict, query: CallbackQuery): +async def register_course(env: Environment, user_cache: dict, query: CallbackQuery): """Regiser a course for the user, to be called by `bot.py`""" # Create one driver per task driver, wait = init_driver(env, wait_timeout=10) diff --git a/utils/db.py b/utils/db.py index db9bebe..3ff832b 100644 --- a/utils/db.py +++ b/utils/db.py @@ -15,7 +15,7 @@ class Database: - def __init__(self, env: str) -> None: + def __init__(self, env: Environment): mongo_client = MongoClient(os.getenv("MONGO_URL"), tlsCAFile=certifi.where()) mongo_db = mongo_client.get_database(f"{env}_db") self.env = env @@ -67,7 +67,7 @@ def update_subscription_time(self, uid: str, time: pendulum.DateTime) -> None: ) def update_subscription_status( - self, uid: str, course_name: str, is_subscribed: bool + self, uid: str, last_subscribed: str, is_subscribed: bool ) -> None: """Update user's subscription status.""" self.user_collection.update_one( @@ -75,7 +75,7 @@ def update_subscription_status( { "$set": { IS_SUBSCRIBED: is_subscribed, - LAST_SUBSCRIPTION: course_name, + LAST_SUBSCRIPTION: last_subscribed, } }, upsert=True,