feat: only allow whitelisted users to add bot to groups

When the bot is added to a group, check if the user who added it is
in the whitelist. If not, send a message explaining and leave the group
immediately. This prevents unauthorized users from bypassing DM access
controls by adding the bot to a group.

1 files changed, 61 insertions, 7 deletions

diff --git a/Omni/Agent/Telegram.hs b/Omni/Agent/Telegram.hs
index 5dcf914..418e589 100644
--- a/Omni/Agent/Telegram.hs
+++ b/Omni/Agent/Telegram.hs
@@ -33,6 +33,7 @@ module Omni.Agent.Telegram
     sendMessageReturningId,
     editMessage,
     sendTypingAction,
+    leaveChat,
 
     -- * Media (re-exported from Media)
     getFile,
@@ -173,6 +174,11 @@ telegramSystemPrompt =
 
 getUpdates :: Types.TelegramConfig -> Int -> IO [Types.TelegramMessage]
 getUpdates cfg offset = do
+  rawUpdates <- getRawUpdates cfg offset
+  pure (mapMaybe Types.parseUpdate rawUpdates)
+
+getRawUpdates :: Types.TelegramConfig -> Int -> IO [Aeson.Value]
+getRawUpdates cfg offset = do
   let url =
         Text.unpack (Types.tgApiBaseUrl cfg)
           <> "/bot"
@@ -194,8 +200,7 @@ getUpdates cfg offset = do
       let body = HTTP.getResponseBody response
       case Aeson.decode body of
         Just (Aeson.Object obj) -> case KeyMap.lookup "result" obj of
-          Just (Aeson.Array updates) ->
-            pure (mapMaybe Types.parseUpdate (toList updates))
+          Just (Aeson.Array updates) -> pure (toList updates)
           _ -> pure []
         _ -> pure []
 
@@ -303,6 +308,26 @@ sendTypingAction cfg chatId = do
   _ <- try @SomeException (HTTP.httpLBS req)
   pure ()
 
+leaveChat :: Types.TelegramConfig -> Int -> IO ()
+leaveChat cfg chatId = do
+  let url =
+        Text.unpack (Types.tgApiBaseUrl cfg)
+          <> "/bot"
+          <> Text.unpack (Types.tgBotToken cfg)
+          <> "/leaveChat"
+      body =
+        Aeson.object
+          [ "chat_id" .= chatId
+          ]
+  req0 <- HTTP.parseRequest url
+  let req =
+        HTTP.setRequestMethod "POST"
+          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
+          <| HTTP.setRequestBodyLBS (Aeson.encode body)
+          <| req0
+  _ <- try @SomeException (HTTP.httpLBS req)
+  pure ()
+
 runTelegramBot :: Types.TelegramConfig -> Provider.Provider -> IO ()
 runTelegramBot tgConfig provider = do
   putText "Starting Telegram bot..."
@@ -329,11 +354,40 @@ runTelegramBot tgConfig provider = do
 
   forever <| do
     offset <- readTVarIO offsetVar
-    messages <- getUpdates tgConfig offset
-    forM_ messages <| \msg -> do
-      atomically (writeTVar offsetVar (Types.tmUpdateId msg + 1))
-      handleMessage tgConfig provider engineCfg botName msg
-    when (null messages) <| threadDelay 1000000
+    rawUpdates <- getRawUpdates tgConfig offset
+    forM_ rawUpdates <| \rawUpdate -> do
+      case Types.parseBotAddedToGroup botName rawUpdate of
+        Just addedEvent -> do
+          atomically (writeTVar offsetVar (Types.bagUpdateId addedEvent + 1))
+          handleBotAddedToGroup tgConfig addedEvent
+        Nothing -> case Types.parseUpdate rawUpdate of
+          Just msg -> do
+            atomically (writeTVar offsetVar (Types.tmUpdateId msg + 1))
+            handleMessage tgConfig provider engineCfg botName msg
+          Nothing -> do
+            let updateId = getUpdateId rawUpdate
+            forM_ updateId <| \uid -> atomically (writeTVar offsetVar (uid + 1))
+    when (null rawUpdates) <| threadDelay 1000000
+
+getUpdateId :: Aeson.Value -> Maybe Int
+getUpdateId (Aeson.Object obj) = case KeyMap.lookup "update_id" obj of
+  Just (Aeson.Number n) -> Just (round n)
+  _ -> Nothing
+getUpdateId _ = Nothing
+
+handleBotAddedToGroup :: Types.TelegramConfig -> Types.BotAddedToGroup -> IO ()
+handleBotAddedToGroup tgConfig addedEvent = do
+  let addedBy = Types.bagAddedByUserId addedEvent
+      chatId = Types.bagChatId addedEvent
+      firstName = Types.bagAddedByFirstName addedEvent
+  if Types.isUserAllowed tgConfig addedBy
+    then do
+      putText <| "Bot added to group " <> tshow chatId <> " by authorized user " <> firstName <> " (" <> tshow addedBy <> ")"
+      sendMessage tgConfig chatId "hello! i'm ready to help."
+    else do
+      putText <| "Bot added to group " <> tshow chatId <> " by UNAUTHORIZED user " <> firstName <> " (" <> tshow addedBy <> ") - leaving"
+      sendMessage tgConfig chatId "sorry, you're not authorized to add me to groups."
+      leaveChat tgConfig chatId
 
 handleMessage ::
   Types.TelegramConfig ->