Merge pull request #19 from BukeLy/feat/user-whitelist-security

feat: add user whitelist security module

Author: BukeLy

agent-sdk-client/config.py

@@ -28,20 +28,27 @@ def extract_command(text: Optional[str]) -> Optional[str]:
     return command
 
 
-def _load_config(config_path: Path = DEFAULT_CONFIG_PATH) -> tuple[list[str], dict[str, str]]:
-    """Load agent/local commands from TOML config file."""
+def _load_config(config_path: Path = DEFAULT_CONFIG_PATH) -> tuple[list[str], dict[str, str], list[int | str]]:
+    """Load commands and security config from TOML config file.
+
+    Returns:
+        Tuple of (agent_commands, local_commands, user_whitelist).
+    """
     if not config_path.exists():
-        return [], {}
+        return [], {}, ['all']
 
     try:
         with config_path.open('rb') as f:
             data = tomllib.load(f)
+
+        # Load agent commands
         agent_commands = data.get('agent_commands', {}).get('commands', [])
         if not isinstance(agent_commands, list):
             logger.warning("Agent commands config is not a list; ignoring configuration")
             agent_commands = []
         agent_commands = [cmd for cmd in agent_commands if isinstance(cmd, str)]
 
+        # Load local commands
         local_commands_raw = data.get('local_commands', {})
         if not isinstance(local_commands_raw, dict):
             logger.warning("Local commands config is not a table; ignoring configuration")
@@ -52,10 +59,28 @@ def _load_config(config_path: Path = DEFAULT_CONFIG_PATH) -> tuple[list[str], di
             if isinstance(name, str) and isinstance(value, str)
         }
 
-        return agent_commands, local_commands
+        # Load security whitelist
+        security = data.get('security', {})
+        whitelist = security.get('user_whitelist', ['all'])
+        if not isinstance(whitelist, list):
+            logger.warning("user_whitelist is not a list; using default ['all']")
+            whitelist = ['all']
+        else:
+            validated = []
+            for item in whitelist:
+                if item == 'all':
+                    validated.append('all')
+                elif isinstance(item, int):
+                    validated.append(item)
+                else:
+                    logger.warning(f"Invalid whitelist entry: {item}; skipping")
+            whitelist = validated if validated else ['all']
+
+        return agent_commands, local_commands, whitelist
+
     except (OSError, tomllib.TOMLDecodeError) as exc:  # pragma: no cover - defensive logging
-        logger.warning("Failed to load command configuration: %s", exc)
-    return [], {}
+        logger.warning("Failed to load configuration: %s", exc)
+    return [], {}, ['all']
 
 
 @dataclass
@@ -68,18 +93,20 @@ class Config:
     queue_url: str
     agent_commands: list[str]
     local_commands: dict[str, str]
+    user_whitelist: list[int | str]
 
     @classmethod
     def from_env(cls, config_path: Optional[Path] = None) -> 'Config':
         """Load configuration from environment variables."""
-        agent_cmds, local_cmds = _load_config(config_path or DEFAULT_CONFIG_PATH)
+        agent_cmds, local_cmds, whitelist = _load_config(config_path or DEFAULT_CONFIG_PATH)
         return cls(
             telegram_token=os.getenv('TELEGRAM_BOT_TOKEN', ''),
             agent_server_url=os.getenv('AGENT_SERVER_URL', ''),
             auth_token=os.getenv('SDK_CLIENT_AUTH_TOKEN', 'default-token'),
             queue_url=os.getenv('QUEUE_URL', ''),
             agent_commands=agent_cmds,
             local_commands=local_cmds,
+            user_whitelist=whitelist,
         )
 
     def get_command(self, text: Optional[str]) -> Optional[str]:

agent-sdk-client/config.toml

@@ -8,3 +8,9 @@ commands = [
 [local_commands]
 # Local-only commands handled by the client
 help = "Hello World"
+
+[security]
+# User IDs allowed to add bot to groups and send private messages.
+# Use ["all"] to allow everyone (default behavior).
+# Example with specific users: user_whitelist = [123456789, 987654321]
+user_whitelist = ["all"]

agent-sdk-client/handler.py

@@ -12,6 +12,7 @@
 from telegram import Bot, Update
 
 from config import Config
+from security import is_user_allowed, should_leave_group
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
@@ -172,11 +173,35 @@ def lambda_handler(event: dict, context: Any) -> dict:
         logger.debug('Ignoring non-update webhook')
         return {'statusCode': 200}
 
+    # Handle my_chat_member event (bot added to group)
+    if update.my_chat_member:
+        if should_leave_group(update, config.user_whitelist):
+            chat_id = update.my_chat_member.chat.id
+            inviter_id = update.my_chat_member.from_user.id
+            asyncio.run(bot.leave_chat(chat_id))
+            logger.info(
+                f"Left unauthorized group",
+                extra={'chat_id': chat_id, 'inviter_id': inviter_id},
+            )
+            _send_metric('SecurityBlock.UnauthorizedGroup')
+        return {'statusCode': 200}
+
     message = update.message or update.edited_message
     if not message or not message.text:
         logger.debug('Ignoring webhook without text message')
         return {'statusCode': 200}
 
+    # Check private message whitelist
+    if message.chat.type == 'private':
+        user_id = message.from_user.id if message.from_user else None
+        if user_id and not is_user_allowed(user_id, config.user_whitelist):
+            logger.info(
+                f"Blocked private message from unauthorized user",
+                extra={'user_id': user_id},
+            )
+            _send_metric('SecurityBlock.UnauthorizedPrivate')
+            return {'statusCode': 200}
+
     cmd = config.get_command(message.text)
     if cmd and config.is_local_command(cmd):
         _handle_local_command(bot, message, config, cmd)

agent-sdk-client/security.py

@@ -0,0 +1,42 @@
+"""Security module for Telegram Bot access control."""
+from telegram import Update
+
+
+def is_user_allowed(user_id: int, whitelist: list[int | str]) -> bool:
+    """Check if user is in whitelist.
+
+    Args:
+        user_id: Telegram user ID to check.
+        whitelist: List of allowed user IDs, or ['all'] to allow everyone.
+
+    Returns:
+        True if user is allowed, False otherwise.
+    """
+    if 'all' in whitelist:
+        return True
+    return user_id in whitelist
+
+
+def should_leave_group(update: Update, whitelist: list[int | str]) -> bool:
+    """Check if bot should leave a group based on who added it.
+
+    Args:
+        update: Telegram Update object with my_chat_member event.
+        whitelist: List of allowed user IDs who can add bot to groups.
+
+    Returns:
+        True if bot should leave (added by unauthorized user), False otherwise.
+    """
+    if not update.my_chat_member:
+        return False
+
+    member_update = update.my_chat_member
+    old_status = member_update.old_chat_member.status
+    new_status = member_update.new_chat_member.status
+
+    # Bot being added to group (status changed from left/kicked to member/administrator)
+    if old_status in ('left', 'kicked') and new_status in ('member', 'administrator'):
+        inviter_id = member_update.from_user.id
+        return not is_user_allowed(inviter_id, whitelist)
+
+    return False