bot: support Mastodon OAuth2 login

test callback

test mastodon callback

test callback

clean customwebhook example

bot: test oauth login

test callback url

test callback

1 files changed, 117 insertions, 4 deletions

diff --git a/bot.py b/bot.py
index e73fe1b..6198bc4 100644
--- a/bot.py
+++ b/bot.py
@@ -1,6 +1,26 @@
 #!/usr/bin/env python
 
+import asyncio
 import logging
+from dataclasses import dataclass
+from http import HTTPStatus
+from config import BOT_TOKEN, TELEGRAM_WEBHOOK_URL, HEALTHCHECK_URL, FEDI_LOGIN_CALLBACK_URL, BOT_DOMAIN, BOT_PORT
+
+import uvicorn
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse, Response
+from starlette.routing import Route
+
+
+from telegram import Update
+from telegram.ext import (
+    Application,
+    CallbackContext,
+    ContextTypes,
+    ExtBot,
+    TypeHandler,
+)
 
 from telegram.ext import (
     Application,
@@ -12,6 +32,7 @@ from telegram.ext import (
 )
 
 from callback import (
+    callback_generate_fedi_login_url,
     callback_skip_media,
     callback_location_sharing,
     callback_manual_location,
@@ -24,10 +45,12 @@ from callback import (
 )
 from command import (
     start_command,
+    fedi_login_command,
     cancel_command,
     help_command
 )
 from config import (
+    FEDI_LOGIN,
     WAIT_LOCATION,
     LOCATION_SEARCH_KEYWORD,
     LOCATION_CONFIRMATION,
@@ -36,21 +59,57 @@ from config import (
     BOT_TOKEN
 )
 
+# Enable logging
 logging.basicConfig(
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
 )
 logger = logging.getLogger(__name__)
 
 
-def main() -> None:
-    application = Application.builder().token(BOT_TOKEN).build()
+@dataclass
+class FediLoginCallbackUpdate:
+    code: str
+    state: int
+
+
+class FediLoginCallbackContext(CallbackContext[ExtBot, dict, dict, dict]):
+    @classmethod
+    def from_update(
+            cls,
+            update: object,
+            application: "Application",
+    ) -> "FediLoginCallbackContext":
+        if isinstance(update, FediLoginCallbackUpdate):
+            return cls(application=application, user_id=update.state)
+        return super().from_update(update, application)
+
+
+async def process_oauth_login_callback(update: FediLoginCallbackUpdate, context: FediLoginCallbackContext) -> None:
+    combined_payloads = update.code
+    text = "Login success, your code is: {}".format(combined_payloads)
+    print(text)
+    print(update.state)
+    await context.bot.send_message(chat_id=update.state, text=text)
+
+
+async def main() -> None:
+    context_types = ContextTypes(context=FediLoginCallbackContext)
+    # Here we set updater to None because we want our custom webhook server to handle the updates
+    # and hence we don't need an Updater instance
+    application = (
+        Application.builder().token(BOT_TOKEN).context_types(context_types).build()
+    )
 
     checkin_handler = ConversationHandler(
         entry_points=[
             CommandHandler("start", start_command),
+            CommandHandler("login", fedi_login_command),
             MessageHandler(filters.LOCATION, callback_location_sharing),
         ],
         states={
+            FEDI_LOGIN: [
+                MessageHandler(filters.TEXT & ~filters.COMMAND, callback_generate_fedi_login_url),
+            ],
             WAIT_LOCATION: [
                 MessageHandler(filters.LOCATION, callback_location_sharing),
             ],
@@ -74,10 +133,64 @@ def main() -> None:
         allow_reentry=True,
     )
 
+    # register handlers
     application.add_handler(CommandHandler("help", help_command))
     application.add_handler(checkin_handler)
-    application.run_polling()
+    application.add_handler(TypeHandler(type=FediLoginCallbackUpdate, callback=process_oauth_login_callback))
+
+    # Pass webhook settings to telegram
+    await application.bot.set_webhook(url=f"{BOT_DOMAIN}{TELEGRAM_WEBHOOK_URL}")
+
+    # Set up webserver
+    async def telegram_webhook(request: Request) -> Response:
+        """Handle incoming Telegram updates by putting them into the `update_queue`"""
+        await application.update_queue.put(
+            Update.de_json(data=await request.json(), bot=application.bot)
+        )
+        return Response()
+
+    async def fedi_oauth_login_callback(request: Request) -> PlainTextResponse:
+        """
+        Handle incoming webhook updates by also putting them into the `update_queue` if
+        the required parameters were passed correctly.
+        """
+        try:
+            code = request.query_params["code"]
+            state = int(request.query_params.get("state"))
+        except KeyError:
+            return PlainTextResponse(
+                status_code=HTTPStatus.BAD_REQUEST,
+                content="Mastodon callback request doesn't contain a valid OAuth code",
+            )
+
+        await application.update_queue.put(FediLoginCallbackUpdate(state=state, code=code))
+        return PlainTextResponse("Thank you for login! Now you can close the browser")
+
+    async def healthcheck(_: Request) -> PlainTextResponse:
+        return PlainTextResponse(content="OK")
+
+    starlette_app = Starlette(
+        routes=[
+            Route(TELEGRAM_WEBHOOK_URL, telegram_webhook, methods=["POST"]),
+            Route(HEALTHCHECK_URL, healthcheck, methods=["GET"]),
+            Route(FEDI_LOGIN_CALLBACK_URL, fedi_oauth_login_callback, methods=["POST", "GET"]),
+        ]
+    )
+    webserver = uvicorn.Server(
+        config=uvicorn.Config(
+            app=starlette_app,
+            port=BOT_PORT,
+            use_colors=False,
+            host="127.0.0.1",
+        )
+    )
+
+    # Run application and webserver together
+    async with application:
+        await application.start()
+        await webserver.serve()
+        await application.stop()
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())