diff --git a/main.py b/main.py index b1f651d..65ce865 100755 --- a/main.py +++ b/main.py @@ -2,6 +2,7 @@ """ParkerBot""" import argparse +import asyncio import datetime import html import os @@ -14,7 +15,7 @@ from google.auth.transport.requests import Request from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build from googleapiclient import errors -from nio import HttpClient, RoomMessageText, SyncResponse, UploadResponse +from nio import AsyncClient, RoomMessageText, SyncResponse, UploadResponse DATA_DIR = os.getenv("DATA_DIR", "./") DB_PATH = os.path.join(DATA_DIR, "parkerbot.sqlite3") @@ -182,18 +183,18 @@ def get_video_info(youtube, video_id): """Check whether a YouTube video is music and return its title and channel.""" try: video_details = youtube.videos().list(id=video_id, part="snippet").execute() - + # Check if the video actually exists/is accessible if not video_details.get("items"): return False, "[Video unavailable or private]", "" snippet = video_details["items"][0]["snippet"] - + # Check if the video category is Music (10) or Entertainment (24) is_music = snippet.get("categoryId") in ("10", "24") title = snippet.get("title", "[Unknown Title]") channel = snippet.get("channelTitle", "[Unknown Channel]") - + return is_music, title, channel except errors.HttpError as error: @@ -204,14 +205,14 @@ def get_video_info(youtube, video_id): return False, "[Error fetching title]", "" -def send_intro_message(client, sender, room_id): +async def send_intro_message(client, sender, room_id): """Sends introduction message in reply to sender, in room with room_id.""" intro_message = ( f"Hi {sender}, I'm ParkerBot! I generate YouTube playlists from links " "sent to this channel. You can find my source code here: " "https://git.abdulocra.cy/abdulocracy/parkerbot" ) - client.room_send( + await client.room_send( room_id=room_id, message_type="m.room.message", content={"msgtype": "m.text", "body": intro_message}, @@ -219,11 +220,11 @@ def send_intro_message(client, sender, room_id): # TODO: Figure out how to properly send GIF, this is broken as shit. with open("./parker.gif", "rb") as gif_file: - response = client.upload(gif_file, content_type="image/gif") + response = await client.upload(gif_file, content_type="image/gif") if isinstance(response, UploadResponse): print("Image was uploaded successfully to server. ") gif_uri = response.content_uri - client.room_send( + await client.room_send( room_id=room_id, message_type="m.room.message", content={ @@ -237,29 +238,29 @@ def send_intro_message(client, sender, room_id): print(f"Failed to upload image. Failure response: {response}") -def send_playlist_of_week(client, sender, room_id, playlist_id): +async def send_playlist_of_week(client, sender, room_id, playlist_id): """Sends playlist of the week in reply to sender, in room with room_id.""" playlist_link = f"https://www.youtube.com/playlist?list={playlist_id}" reply_msg = f"{sender}, here's the playlist of the week: {playlist_link}" - client.room_send( + await client.room_send( room_id=room_id, message_type="m.room.message", content={"msgtype": "m.text", "body": reply_msg}, ) -def send_playlist_of_all(client, sender, room_id, playlist_id): +async def send_playlist_of_all(client, sender, room_id, playlist_id): """Sends playlist of all time in reply to sender, in room with room_id.""" playlist_link = f"https://www.youtube.com/playlist?list={playlist_id}" reply_msg = f"{sender}, here's the playlist of all time: {playlist_link}" - client.room_send( + await client.room_send( room_id=room_id, message_type="m.room.message", content={"msgtype": "m.text", "body": reply_msg}, ) -def message_callback(conn, cursor, youtube, client, room, event): +async def message_callback(conn, cursor, youtube, client, room, event): """Event handler for received messages.""" sender = event.sender if sender != MATRIX_USER: @@ -273,24 +274,23 @@ def message_callback(conn, cursor, youtube, client, room, event): ) timestamp_sec = datetime.datetime.fromtimestamp( - event.server_timestamp / 1000, - datetime.UTC, # millisec to sec + event.server_timestamp / 1000, datetime.UTC # millisec to sec ) current_time = datetime.datetime.now(datetime.UTC) - + # Account for up to 5 minutes of clock drift recent = abs(current_time - timestamp_sec) < datetime.timedelta(minutes=5) if body == "!parkerbot" and recent: - send_intro_message(client, sender, room.room_id) + await send_intro_message(client, sender, room.room_id) return if body == "!week" and recent: - send_playlist_of_week(client, sender, room.room_id, playlist_id) + await send_playlist_of_week(client, sender, room.room_id, playlist_id) return if body == "!all" and recent: - send_playlist_of_all(client, sender, room.room_id, all_playlist_id) + await send_playlist_of_all(client, sender, room.room_id, all_playlist_id) return youtube_link_pattern = ( @@ -301,31 +301,29 @@ def message_callback(conn, cursor, youtube, client, room, event): for link in youtube_links: video_id = link.split("v=")[-1].split("&")[0].split("/")[-1] - + # Safely fetch the category check, title, and channel is_music_vid, title, channel = get_video_info(youtube, video_id) - + # Send the title to the channel so people know what the link is # Only do this for recent messages to prevent spam during backwards-sync if recent: plain_text = f"{title}" if channel: plain_text += f" - {channel}" - + # Escape HTML characters to prevent broken rendering in Matrix escaped_text = html.escape(plain_text) - html_text = ( - f"{escaped_text}" - ) - - client.room_send( + html_text = f"{escaped_text}" + + await client.room_send( room_id=room.room_id, message_type="m.room.message", content={ - "msgtype": "m.text", + "msgtype": "m.text", "body": plain_text, "format": "org.matrix.custom.html", - "formatted_body": html_text, + "formatted_body": html_text }, ) @@ -382,7 +380,7 @@ def record_message(conn, cursor, sender, link, timestamp): return cursor.fetchone()[0] -def sync_callback(response): +async def sync_callback(response): """Saves Matrix sync token.""" with open(TOKEN_PATH, "w", encoding="utf-8") as f: f.write(response.next_batch) @@ -397,9 +395,9 @@ def load_sync_token(): return None -def get_client(conn, cursor, youtube): +async def get_client(conn, cursor, youtube): """Returns configured and logged in Matrix client.""" - client = HttpClient(MATRIX_SERVER, MATRIX_USER) + client = AsyncClient(MATRIX_SERVER, MATRIX_USER) client.add_event_callback( lambda room, event: message_callback( conn, cursor, youtube, client, room, event @@ -407,23 +405,23 @@ def get_client(conn, cursor, youtube): RoomMessageText, ) client.add_response_callback(sync_callback, SyncResponse) - print(client.login(MATRIX_PASSWORD)) + print(await client.login(MATRIX_PASSWORD)) return client -def backwards_sync(conn, cursor, youtube, client, room, start_token): +async def backwards_sync(conn, cursor, youtube, client, room, start_token): """Fetch and process historical messages from a given room.""" print("Starting to process channel log...") from_token = start_token room_id = room.room_id while True: # Fetch room messages - response = client.room_messages(room_id, from_token, direction="b") + response = await client.room_messages(room_id, from_token, direction="b") # Process each message for event in response.chunk: if isinstance(event, RoomMessageText): - message_callback(conn, cursor, youtube, client, room, event) + await message_callback(conn, cursor, youtube, client, room, event) # Break if there are no more messages to fetch if not response.end or response.end == from_token: @@ -433,30 +431,23 @@ def backwards_sync(conn, cursor, youtube, client, room, start_token): from_token = response.end -def main(): +async def main(): """Get DB and Matrix client ready, and start syncing.""" args = parse_arguments() conn, cursor = connect_db() define_tables(conn, cursor) youtube = get_authenticated_service() - client = get_client(conn, cursor, youtube) + client = await get_client(conn, cursor, youtube) sync_token = load_sync_token() # This is incredibly dumb and most probably will exceed your YouTube API quota. if args.backwards_sync: - init_sync = client.sync(30000) - room = client.room_resolve_alias(MATRIX_ROOM) - backwards_sync(conn, cursor, youtube, client, room, init_sync.next_batch) + init_sync = await client.sync(30000) + room = await client.room_resolve_alias(MATRIX_ROOM) + await backwards_sync(conn, cursor, youtube, client, room, init_sync.next_batch) - print("Started syncing...") - while True: - try: - client.sync(timeout=30000, full_state=True, since=sync_token) - sync_token = load_sync_token() - except Exception as e: - print(f"Sync error: {e}") - time.sleep(5) + await client.sync_forever(30000, full_state=True, since=sync_token) if __name__ == "__main__": - main() + asyncio.run(main())