Skip to content

feat: add user whitelist security module #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BukeLy merged 2 commits into from
Jan 11, 2026
Merged

feat: add user whitelist security module #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
192ae6b
docs: 添加 Claude Agent SDK plugin 文档
BukeLy Jan 7, 2026
328de52
feat: add user whitelist security module
BukeLy Jan 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 34 additions & 7 deletions agent-sdk-client/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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':
Copy link

Copilot AI Jan 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The validation logic checks for the string 'all' (line 71) but doesn't validate that it's actually the string 'all' - it just checks equality. This means if someone provides user_whitelist = ["all", "ALL", "All"], only the lowercase "all" would be recognized. Consider normalizing the string to lowercase before comparison, or documenting that only lowercase "all" is valid. Similarly, the documentation in config.toml should clarify that "all" must be lowercase.

Suggested change
if item == 'all':
if isinstance(item, str) and item.lower() == 'all':

Copilot uses AI. Check for mistakes.
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']
Comment on lines +64 to +77
Copy link

Copilot AI Jan 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The whitelist validation logic has an edge case issue. If a user provides a whitelist that contains only invalid entries (e.g., strings other than 'all', floats, etc.), all entries will be skipped, resulting in an empty validated list. The code then falls back to ['all'], which effectively disables security even though the user explicitly configured a whitelist. This could be a security issue as it silently fails open rather than closed. Consider either rejecting the configuration entirely with an error, or keeping the empty list to deny all access when all entries are invalid.

Suggested change
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']
raw_whitelist = security.get('user_whitelist', None)
if raw_whitelist is None:
# No whitelist configured; default to allow all
whitelist = ['all']
elif not isinstance(raw_whitelist, list):
logger.warning("user_whitelist is not a list; using default ['all']")
whitelist = ['all']
else:
validated: list[int | str] = []
for item in raw_whitelist:
if item == 'all':
validated.append('all')
elif isinstance(item, int):
validated.append(item)
else:
logger.warning(f"Invalid whitelist entry: {item}; skipping")
if not validated:
# Explicit whitelist provided but no valid entries; deny all access
logger.warning(
"user_whitelist is empty after validation; denying all access"
)
whitelist = validated

Copilot uses AI. Check for mistakes.

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
Expand All @@ -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]:
Expand Down
6 changes: 6 additions & 0 deletions agent-sdk-client/config.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
26 changes: 26 additions & 0 deletions agent-sdk-client/handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

Receives Telegram webhook, writes to SQS, returns 200 immediately.
"""
import asyncio
import json
import logging
from typing import Any
Expand All @@ -11,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)
Expand Down Expand Up @@ -169,11 +171,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))
Copy link

Copilot AI Jan 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using asyncio.run() in a Lambda handler that's already in a synchronous context can create problems. The bot.leave_chat() call is async, but wrapping it in asyncio.run() creates a new event loop each time. This works but is inefficient. Additionally, there's no error handling if the leave_chat operation fails (e.g., network issues, bot lacks permissions). Consider either:

  1. Wrapping the entire operation in a try-except block to handle potential exceptions
  2. Using a synchronous approach if the telegram library supports it
  3. At minimum, catching and logging any exceptions from the leave_chat call to prevent the Lambda from failing
Suggested change
asyncio.run(bot.leave_chat(chat_id))
try:
asyncio.run(bot.leave_chat(chat_id))
except Exception:
logger.exception(
"Failed to leave unauthorized group",
extra={'chat_id': chat_id, 'inviter_id': inviter_id},
)

Copilot uses AI. Check for mistakes.
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}
Comment on lines +192 to +201
Copy link

Copilot AI Jan 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The security whitelist check only applies to private messages (chat.type == 'private'), but group messages are not filtered by user whitelist. This means that even if a bot leaves a group when added by an unauthorized user, if the bot is added to a group by an authorized user, all members of that group can interact with the bot regardless of whether they're in the whitelist. Consider whether this is the intended behavior or if group messages should also be filtered by the sender's user_id against the whitelist.

Suggested change
# 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}
# Check message whitelist (applies to both private and group chats)
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):
chat_type = getattr(message.chat, "type", None)
metric_name = (
'SecurityBlock.UnauthorizedPrivate'
if chat_type == 'private'
else 'SecurityBlock.UnauthorizedGroupMessage'
)
logger.info(
"Blocked message from unauthorized user",
extra={
'user_id': user_id,
'chat_id': getattr(message.chat, "id", None),
'chat_type': chat_type,
},
)
_send_metric(metric_name)
return {'statusCode': 200}

Copilot uses AI. Check for mistakes.

cmd = config.get_command(message.text)
if cmd and config.is_local_command(cmd):
_handle_local_command(bot, message, config, cmd)
Expand Down
42 changes: 42 additions & 0 deletions agent-sdk-client/security.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
"""
Copy link

Copilot AI Jan 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The function doesn't explicitly handle the case of an empty whitelist ([]). If the whitelist is empty (which shouldn't happen due to the fallback to ['all'] in config loading, but could happen if called programmatically), the function would return False for all users, which is safe but might be unexpected. Consider adding explicit handling or documentation for this edge case.

Suggested change
"""
"""
# Explicitly handle empty whitelist: deny access to all users.
if not whitelist:
return False

Copilot uses AI. Check for mistakes.
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)

Copy link

Copilot AI Jan 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The function only checks for status transitions from 'left' or 'kicked' to 'member' or 'administrator', but doesn't handle the case where the bot's status is upgraded from 'member' to 'administrator' by someone other than the original inviter. If an unauthorized user promotes the bot to admin in a group where it was invited by an authorized user, the bot won't leave. Consider whether permission changes should also be validated against the whitelist.

Suggested change
# Bot being promoted to administrator by a (potentially) unauthorized user
if old_status not in ('administrator', 'creator') and new_status == 'administrator':
inviter_id = member_update.from_user.id
return not is_user_allowed(inviter_id, whitelist)

Copilot uses AI. Check for mistakes.
return False
Comment on lines +1 to +42
Copy link

Copilot AI Jan 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new security functionality lacks test coverage. The repository has comprehensive tests for config loading (test_command_config.py), but no tests are included for:

  1. The is_user_allowed function with various whitelist configurations
  2. The should_leave_group function with different chat member status transitions
  3. Integration of security checks in the handler
  4. Edge cases like empty whitelists, mixed valid/invalid entries, etc.

Consider adding tests for the new security module to ensure the whitelist logic works correctly.

Copilot uses AI. Check for mistakes.
Loading
Loading