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.

2 files changed, 111 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 ->
diff --git a/Omni/Agent/Telegram/Types.hs b/Omni/Agent/Telegram/Types.hs
index d240786..aaea65b 100644
--- a/Omni/Agent/Telegram/Types.hs
+++ b/Omni/Agent/Telegram/Types.hs
@@ -19,10 +19,12 @@ module Omni.Agent.Telegram.Types
     TelegramPhoto (..),
     TelegramVoice (..),
     TelegramReplyMessage (..),
+    BotAddedToGroup (..),
     ChatType (..),
 
     -- * Parsing
     parseUpdate,
+    parseBotAddedToGroup,
     parseDocument,
     parseLargestPhoto,
     parsePhotoSize,
@@ -323,6 +325,14 @@ instance Aeson.FromJSON TelegramReplyMessage where
         <*> (v .:? "from_last_name")
         <*> (v .:? "text" .!= "")
 
+data BotAddedToGroup = BotAddedToGroup
+  { bagUpdateId :: Int,
+    bagChatId :: Int,
+    bagAddedByUserId :: Int,
+    bagAddedByFirstName :: Text
+  }
+  deriving (Show, Eq, Generic)
+
 data ChatType = Private | Group | Supergroup | Channel
   deriving (Show, Eq, Generic)
 
@@ -461,6 +471,46 @@ parseUpdate val = do
         tmReplyTo = replyTo
       }
 
+parseBotAddedToGroup :: Text -> Aeson.Value -> Maybe BotAddedToGroup
+parseBotAddedToGroup botUsername val = do
+  Aeson.Object obj <- pure val
+  updateId <- case KeyMap.lookup "update_id" obj of
+    Just (Aeson.Number n) -> Just (round n)
+    _ -> Nothing
+  Aeson.Object msgObj <- KeyMap.lookup "message" obj
+  Aeson.Object chatObj <- KeyMap.lookup "chat" msgObj
+  chatId <- case KeyMap.lookup "id" chatObj of
+    Just (Aeson.Number n) -> Just (round n)
+    _ -> Nothing
+  let chatType = case KeyMap.lookup "type" chatObj of
+        Just (Aeson.String t) -> t
+        _ -> "private"
+  guard (chatType == "group" || chatType == "supergroup")
+  Aeson.Object fromObj <- KeyMap.lookup "from" msgObj
+  addedByUserId <- case KeyMap.lookup "id" fromObj of
+    Just (Aeson.Number n) -> Just (round n)
+    _ -> Nothing
+  addedByFirstName <- case KeyMap.lookup "first_name" fromObj of
+    Just (Aeson.String s) -> Just s
+    _ -> Nothing
+  Aeson.Array newMembers <- KeyMap.lookup "new_chat_members" msgObj
+  let botWasAdded = any (isBotUser botUsername) (toList newMembers)
+  guard botWasAdded
+  pure
+    BotAddedToGroup
+      { bagUpdateId = updateId,
+        bagChatId = chatId,
+        bagAddedByUserId = addedByUserId,
+        bagAddedByFirstName = addedByFirstName
+      }
+  where
+    isBotUser :: Text -> Aeson.Value -> Bool
+    isBotUser username (Aeson.Object userObj) =
+      case KeyMap.lookup "username" userObj of
+        Just (Aeson.String u) -> Text.toLower u == Text.toLower username
+        _ -> False
+    isBotUser _ _ = False
+
 parseDocument :: Aeson.Object -> Maybe TelegramDocument
 parseDocument docObj = do
   fileId <- case KeyMap.lookup "file_id" docObj of