From c851b9d7e20662e3b5642cf170301d618d37e846 Mon Sep 17 00:00:00 2001 From: Yeet Date: Tue, 24 Sep 2024 00:48:26 +0200 Subject: [PATCH] Google calendar support --- .gitignore | 3 + Dockerfile | 9 +- config.py | 6 ++ google_calendar.py | 215 +++++++++++++++++++++++++++++++++++++++++++++ models.py | 15 ++-- requirements.txt | 6 ++ tgbot.py | 161 +++++++++++++++++++++++++++++---- 7 files changed, 389 insertions(+), 26 deletions(-) create mode 100644 google_calendar.py diff --git a/.gitignore b/.gitignore index f7c4fbf..290e06c 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,6 @@ chromedriver.exe chromedriver data.db .envrc +google_client.json +credentials.json +pyrightconfig.json diff --git a/Dockerfile b/Dockerfile index 0510189..fd9c9a5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -10,9 +10,6 @@ WORKDIR /app # Copy the local requirements file to the container COPY requirements.txt /app/requirements.txt -# Install the required Python packages -RUN pip install --no-cache-dir -r requirements.txt - # Install Selenium and its required drivers RUN apt-get update && \ apt-get install -y --no-install-recommends \ @@ -20,9 +17,15 @@ RUN apt-get update && \ curl \ unzip \ chromium-driver \ + build-essential \ + libsqlite3-dev \ + sqlite3 \ && apt-get clean \ && rm -rf /var/lib/apt/lists/* +# Install the required Python packages +RUN pip install --no-cache-dir -r requirements.txt + # Set up Chrome WebDriver for Selenium RUN wget -q -O /app/chromedriver.zip https://chromedriver.storage.googleapis.com/114.0.5735.90/chromedriver_linux64.zip && \ unzip /app/chromedriver.zip && \ diff --git a/config.py b/config.py index 9ffc728..ee4677e 100644 --- a/config.py +++ b/config.py @@ -4,6 +4,12 @@ TELEGRAM_TOKEN = os.environ.get("SECRETARX_TG_TOKEN", None) DB_PATH = "data.db" CHROMEDRIVER_PATH = "chromedriver" + +LISTEN_PORT = 9090 +GOOGLE_CLIENT_SECRET_FILE = "credentials.json" +REDIRECT_URL = "http://localhost:9090/" +CALENDAR_NAME = "secretarx" + INTERVALS = { 1800 + 3600 * 0.0: 1800, 1800 + 3600 * 0.5: 30, diff --git a/google_calendar.py b/google_calendar.py new file mode 100644 index 0000000..10bae85 --- /dev/null +++ b/google_calendar.py @@ -0,0 +1,215 @@ +from collections.abc import Mapping +import datetime +from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler +import os.path +import socketserver +from threading import Thread +from time import sleep +from urllib.parse import urlparse, parse_qs +import multiprocessing + +from google.auth.transport.requests import Request +from google.oauth2.credentials import Credentials +from google_auth_oauthlib.flow import InstalledAppFlow +from googleapiclient.discovery import build +from googleapiclient.errors import HttpError +from config import GOOGLE_CLIENT_SECRET_FILE + +# If modifying these scopes, delete the file token.json. +SCOPES = ["https://www.googleapis.com/auth/calendar"] + + +class GoogleCalendar: + class TCPServer(socketserver.ForkingTCPServer): + flow_callback = None + + class HTTPHandler(BaseHTTPRequestHandler): + def do_GET(self): + url_parts = urlparse(self.path) + query_params = parse_qs(url_parts.query) + + if "code" not in query_params: + self.send_response(500) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write("

Missing OAuth Code

".encode("utf-8")) + return + + if "state" not in query_params: + self.send_response(500) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write("

Missing OAuth State

".encode("utf-8")) + return + + code = query_params["code"][0] + state = query_params["state"][0] + + result = self.server.flow_callback(code=code, state=state) + + if result is None: + self.send_response(200) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write("

OAuth Success

".encode("utf-8")) + return + + self.send_response(500) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write(f"

{result}

".encode("utf-8")) + return + + def __init__( + self, + http_port: int, + client_secret_file: str, + redirect_uri: str, + flow_callback: callable, + calendar_summary: str, + token_refresh_callback: callable = None, + ) -> None: + self.flow = InstalledAppFlow.from_client_secrets_file( + client_secret_file, SCOPES, redirect_uri=redirect_uri + ) + self.httpd = ThreadingHTTPServer(("", http_port), self.HTTPHandler) + + self.httpd.flow_callback = self.finish_flow + + manager = multiprocessing.Manager() + self.state_callback_kwargs = manager.dict() + self.flow_callback = flow_callback + self.calendar_summary = calendar_summary + self.token_refresh_callback = token_refresh_callback + + def finish_flow(self, state: str, code: str): + if (kwargs := self.state_callback_kwargs.get(state, None)) is None: + return "Invalid OAuth State" + + token = self.flow.fetch_token(code=code) + + self.flow_callback(token, **kwargs) + + def flow_url(self, callback_kwargs={}): + url, state = self.flow.authorization_url(prompt="consent") + self.state_callback_kwargs[state] = callback_kwargs + return url + + def service(self, token: Mapping[str, str]): + # Convert the token dict into a Credentials object + creds = Credentials( + token=token.get("access_token"), + refresh_token=token.get("refresh_token"), + token_uri=self.flow.client_config["token_uri"], + client_id=self.flow.client_config["client_id"], + client_secret=self.flow.client_config["client_secret"], + scopes=SCOPES, + ) + + if not creds or not creds.valid: + return None + + if creds and creds.expired and creds.refresh_token: + creds.refresh(Request()) + self.token_refresh_callback( + token, + { + "access_token": creds.token, + "expires_in": creds.expiry, + "refresh_token": creds.refresh_token, + "scope": creds.scopes, + "token_type": creds.token, + "expires_at": creds.expiry, + }, + ) + + service = build("calendar", "v3", credentials=creds) + + # Return the credentials object + return service + + def get_calendar_id(self, token: Mapping[str, str]): + # Create calendar with id + service = self.service(token) + + if not service: + return None + + calendars = service.calendarList().list().execute().get("items") + + for calendar in calendars: + if calendar["summary"] == self.calendar_summary: + return calendar["id"] + + calendar = ( + service.calendars() + .insert( + body={ + "summary": self.calendar_summary, + "timeZone": "Europe/Amsterdam", + } + ) + .execute() + ) + + return calendar["id"] + + def get_first_event(self, token: Mapping[str, str], start: datetime.datetime): + service = self.service(token) + + if not service: + return None + + calendar_id = self.get_calendar_id(token) + print(start.isoformat() + "Z") + + try: + events_result = ( + service.events() + .list( + calendarId=calendar_id, + timeMin=start.isoformat(), + maxResults=1, + singleEvents=True, + orderBy="startTime", + ) + .execute() + ) + except HttpError as e: + if e.resp.status == 404: + return [] + print(e) + + items = events_result.get("items", []) + + if items: + return items[0] + + return [] + + def create_event(self, token: Mapping[str, str], event: dict): + service = self.service(token) + if not service: + return None + + calendar_id = self.get_calendar_id(token) + + event = service.events().insert(calendarId=calendar_id, body=event).execute() + + return event + + def delete_event(self, token: Mapping[str, str], event_id: str): + service = self.service(token) + + if not service: + return None + + calendar_id = self.get_calendar_id(token) + + service.events().delete(calendarId=calendar_id, eventId=event_id).execute() + + def httpd_listen(self): + # Start httpd as a daemon thread + thread = Thread(target=self.httpd.serve_forever, args=()) + thread.daemon = True + thread.start() diff --git a/models.py b/models.py index 6e46cec..f421418 100644 --- a/models.py +++ b/models.py @@ -1,5 +1,6 @@ from datetime import datetime, timedelta, timezone + class Booking: def __init__( self, @@ -12,21 +13,24 @@ class Booking: **kwargs, ): self.booking_id = bookingId - self.start = datetime.strptime(startDate, "%Y-%m-%dT%H:%M:%S.%fZ").replace(tzinfo=timezone.utc) - self.end = datetime.strptime(endDate, "%Y-%m-%dT%H:%M:%S.%fZ").replace(tzinfo=timezone.utc) + self.start = datetime.strptime(startDate, "%Y-%m-%dT%H:%M:%S.%fZ").replace( + tzinfo=timezone.utc + ) + self.end = datetime.strptime(endDate, "%Y-%m-%dT%H:%M:%S.%fZ").replace( + tzinfo=timezone.utc + ) self.bookable_product_id = bookableProductId self.linked_product_id = linkedProductId self.available = isAvailable self.booking_dict = kwargs - @property def start_stamp(self): ## Get string representation of the timestamp in GMT +2 with only Day of week and month, hour and minute gmt_plus_2 = timezone(timedelta(hours=2)) return self.start.astimezone(gmt_plus_2).strftime("%A %d %B %H:%M") - + @property def end_stamp(self): gmt_plus_2 = timezone(timedelta(hours=2)) @@ -58,4 +62,5 @@ class Booking: if time_left.total_seconds() > 0: return time_left else: - return timedelta(0) # No time left if the booking has already started \ No newline at end of file + return timedelta(0) # No time left if the booking has already started + diff --git a/requirements.txt b/requirements.txt index 6744be6..f2f0115 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,9 @@ pysqlite3 selenium==4.24.0 requests==2.32.3 telebot==0.0.5 +google-api-core==2.19.2 +google-api-python-client==2.145.0 +google-auth==2.34.0 +google-auth-httplib2==0.2.0 +google-auth-oauthlib==1.2.1 +googleapis-common-protos==1.65.0 diff --git a/tgbot.py b/tgbot.py index 3372cab..c12a822 100644 --- a/tgbot.py +++ b/tgbot.py @@ -7,10 +7,13 @@ from threading import Thread import sqlite3 from xclient import XClient import config +from google_calendar import GoogleCalendar +import json logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) + class TelegramBot: def __init__(self, bot_token, db_path="bot_database.db"): self.bot = telebot.TeleBot(bot_token) @@ -19,6 +22,14 @@ class TelegramBot: self.user_bookings = {} # Store user bookings by chat ID self.watchlist = {} # Watchlist to store full slots by chat ID self.db_path = db_path + self.calendar = GoogleCalendar( + config.LISTEN_PORT, + config.GOOGLE_CLIENT_SECRET_FILE, + config.REDIRECT_URL, + self.calendar_connect_callback, + config.CALENDAR_NAME, + self.calendar_token_update_callback, + ) self.init_db() # Set up the command handlers @@ -56,17 +67,9 @@ 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 + @self.bot.message_handler(commands=["calendar"]) + def calendar_auth(message): + self.calendar_auth(message) def init_db(self): with sqlite3.connect(self.db_path) as conn: @@ -78,22 +81,81 @@ class TelegramBot: chat_id INTEGER UNIQUE, username TEXT, password TEXT, - access_token TEXT + access_token TEXT, + calendar_token TEXT ) """) conn.commit() + def calendar_connect_callback(self, calendar_token, chat_id): + calendar_token_str = json.dumps(calendar_token, sort_keys=True) + + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "UPDATE users SET calendar_token = ? WHERE chat_id = ?", + (calendar_token_str, chat_id), + ) + conn.commit() + + # Send message to user + self.bot.send_message(chat_id, "Calendar connected successfully!") + + def calendar_token_update_callback(self, old_token, new_token): + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "UPDATE users SET calendar_token = ? WHERE calendar_token = ?", + ( + json.dumps(new_token, sort_keys=True), + json.dumps(old_token, sort_keys=True), + ), + ) + conn.commit() + + 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 get_calendar_token(self, chat_id): + with sqlite3.connect(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT calendar_token FROM users WHERE chat_id = ?", (chat_id,) + ) + result = cursor.fetchone() + + if result: + return json.loads(result[0]) + + return None + def get_xclient(self, chat_id): if chat_id not in self.xclients: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute( - "SELECT username, password, access_token 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, 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] = 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 @@ -150,6 +212,12 @@ class TelegramBot: message, "Invalid format. Please use `/login username password`." ) + def calendar_auth(self, message): + url = self.calendar.flow_url({"chat_id": message.chat.id}) + + # Send url to user + self.bot.reply_to(message, f"Please visit this URL to authenticate: {url}") + def send_welcome(self, message): self.bot.reply_to( message, @@ -217,6 +285,29 @@ class TelegramBot: # Attempt to book the available slot try: xclient.make_booking(selected_slot) + + if ( + token := self.get_calendar_token(call.message.chat.id) + ) is not None: + self.calendar.create_event( + token, + { + "summary": "Gym Booking", + "start": { + "dateTime": selected_slot.start.strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ), + "timeZone": "Europe/Amsterdam", + }, + "end": { + "dateTime": selected_slot.end.strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ), + "timeZone": "Europe/Amsterdam", + }, + }, + ) + self.bot.answer_callback_query( call.id, "Slot booked successfully!" ) @@ -313,6 +404,19 @@ class TelegramBot: call.message.chat.id, f"Booking for {selected_booking.start_stamp} has been canceled.", ) + + if ( + token := self.get_calendar_token(call.message.chat.id) + ) is not None: + gmt_plus_2 = timezone(timedelta(hours=2)) + + event = self.calendar.get_first_event( + token, selected_booking.start.astimezone(gmt_plus_2) + ) + + if event: + self.calendar.delete_event(token, event["id"]) + else: self.bot.answer_callback_query( call.id, "Failed to cancel booking." @@ -407,6 +511,25 @@ class TelegramBot: # If the slot becomes available try: xclient.make_booking(slot) + if (token := self.get_calendar_token(chat_id)) is not None: + self.calendar.create_event( + token, + { + "summary": "Gym Booking", + "start": { + "dateTime": slot.start.strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ), + "timeZone": "Europe/Amsterdam", + }, + "end": { + "dateTime": slot.end.strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ), + "timeZone": "Europe/Amsterdam", + }, + }, + ) self.bot.send_message( chat_id, f"Slot {slot.start_stamp} is now available and has been booked for you.", @@ -424,16 +547,17 @@ class TelegramBot: if not self.watchlist[chat_id]: del self.watchlist[chat_id] # Remove chat_id from watchlist if empty - def calculate_polling_interval(self, slot_time): now = datetime.now().replace(tzinfo=timezone.utc) time_until_slot = (slot_time - now).total_seconds() - intervals = sorted([(k,v) for k,v in config.INTERVALS.items()], key=lambda x: x[0]) + intervals = sorted( + [(k, v) for k, v in config.INTERVALS.items()], key=lambda x: x[0] + ) last_wait = 1800 - for (threshold, wait) in intervals: + for threshold, wait in intervals: if time_until_slot < threshold: return wait @@ -462,7 +586,6 @@ class TelegramBot: logger.info(f"POLLING {self.watchlist}. Sleeping {polling_interval}") - time.sleep(polling_interval) # Wait before checking the watchlist again def run(self): @@ -471,4 +594,6 @@ class TelegramBot: thread.daemon = True thread.start() + self.calendar.httpd_listen() + self.bot.polling()