diff --git a/config.py b/config.py index a7cc381..7139443 100644 --- a/config.py +++ b/config.py @@ -4,5 +4,13 @@ TELEGRAM_TOKEN = os.environ.get("SECRETARX_TG_TOKEN", None) DB_PATH = "data.db" CHROMEDRIVER_PATH = "chromedriver" +INTERVALS = { + 3600 * 0.5: 30, + 3600 * 1.0: 60, + 3600 * 1.5: 150, + 3600 * 2.0: 300, + 3600 * 3.0: 1800, +} + assert TELEGRAM_TOKEN diff --git a/tgbot.py b/tgbot.py index a05ce7c..9acd630 100644 --- a/tgbot.py +++ b/tgbot.py @@ -5,6 +5,7 @@ import time from threading import Thread import sqlite3 from xclient import XClient +import config class TelegramBot: @@ -52,6 +53,18 @@ class TelegramBot: def callback_remove_slot(call): self.callback_remove_slot(call) + def update_user_access_token_callback(self, chat_id): + def callback(access_token): + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "UPDATE users SET access_token = ? WHERE chat_id = ?", + (access_token, chat_id), + ) + conn.commit() + + return callback + def init_db(self): with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() @@ -61,7 +74,8 @@ class TelegramBot: id INTEGER PRIMARY KEY AUTOINCREMENT, chat_id INTEGER UNIQUE, username TEXT, - password TEXT + password TEXT, + access_token TEXT ) """) conn.commit() @@ -71,12 +85,13 @@ class TelegramBot: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute( - "SELECT username, password FROM users WHERE chat_id = ?", (chat_id,) + "SELECT username, password, access_token FROM users WHERE chat_id = ?", (chat_id,) ) result = cursor.fetchone() if result: - username, password = result - self.xclients[chat_id] = XClient(username, password) + username, password, access_token = result + self.xclients[chat_id] = XClient(username, password, on_access_token_change=self.update_user_access_token_callback(chat_id)) + self.xclients[chat_id].access_token = access_token else: self.xclients[chat_id] = None @@ -410,14 +425,17 @@ class TelegramBot: now = datetime.now().replace(tzinfo=timezone.utc) time_until_slot = (slot_time - now).total_seconds() - if time_until_slot < 3600: # If less than 1 hour - return 300 # 5 minutes - if time_until_slot < 3600 * 2: # If less than 2 hour - return 600 # 10 minutes - if time_until_slot < 3600 * 3: # If less than 3 hour - return 900 # 15 minutes - else: - return 1800 # 30 minutes + intervals = sorted(lambda x: x[0], [(k,v) for k,v in config.INTERVALS.items()]) + + last_wait = 1800 + + for (threshold, wait) in intervals: + if time_until_slot < threshold: + return wait + + last_wait = wait + + return last_wait def poll_periodically(self): # Continuously poll the watchlist at dynamically adjusted intervals diff --git a/xclient.py b/xclient.py index aff7212..eee83af 100644 --- a/xclient.py +++ b/xclient.py @@ -11,11 +11,12 @@ from selenium.webdriver.chrome.service import Service class XClient: - def __init__(self, username, password): + def __init__(self, username, password, on_access_token_change=lambda x: x): self.username = username self.password = password self.access_token = None self.id = -1 + self.on_access_token_change = on_access_token_change self.session = requests.Session() # Default headers for session that mimicks browser @@ -280,6 +281,7 @@ class XClient: def login(self): self.access_token = self.fetch_access_token() + self.on_access_token_change(self.access_token) def is_logged_in(self): # Send request to 'https://backbone-web-api.production.delft.delcom.nl/auth?cf=0 with access token as bearer, check status code