Travis Vasceannie
3acb779569
- Deleted .env.example file as it is no longer needed. - Added .gitignore to manage ignored files and directories. - Introduced CLAUDE.md for AI provider integration documentation. - Created dev.sh for development setup and scripts. - Updated Dockerfile and Dockerfile.production for improved build processes. - Added multiple test files and directories for comprehensive testing. - Introduced new utility and service files for enhanced functionality. - Organized codebase with new directories and files for better maintainability.
314 lines
10 KiB
Python
314 lines
10 KiB
Python
"""
|
|
Security Manager for Discord Voice Chat Quote Bot
|
|
Essential security features: rate limiting, permissions, authentication
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import secrets
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
from typing import Any, Dict, Optional, Set, Tuple
|
|
|
|
import discord
|
|
import jwt
|
|
import redis.asyncio as redis
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SecurityLevel(Enum):
|
|
PUBLIC = "public"
|
|
USER = "user"
|
|
MODERATOR = "moderator"
|
|
ADMIN = "admin"
|
|
OWNER = "owner"
|
|
|
|
|
|
class RateLimitType(Enum):
|
|
COMMAND = "command"
|
|
API = "api"
|
|
UPLOAD = "upload"
|
|
|
|
|
|
@dataclass
|
|
class RateLimitConfig:
|
|
requests: int
|
|
window: int # seconds
|
|
burst: int = 0
|
|
|
|
|
|
class SecurityManager:
|
|
"""Core security management with rate limiting and permissions"""
|
|
|
|
def __init__(self, redis_client: redis.Redis, config: Dict[str, Any]):
|
|
self.redis = redis_client
|
|
self.config = config
|
|
|
|
# Rate limiting
|
|
self.rate_limits = {
|
|
"command": RateLimitConfig(requests=30, window=60, burst=5),
|
|
"api": RateLimitConfig(requests=100, window=60, burst=10),
|
|
"upload": RateLimitConfig(requests=5, window=300, burst=2),
|
|
}
|
|
|
|
# Authentication
|
|
self.jwt_secret = config.get("jwt_secret", secrets.token_urlsafe(32))
|
|
self.session_timeout = 3600
|
|
|
|
# Permissions
|
|
self.role_permissions = {
|
|
"owner": {"*"},
|
|
"admin": {"bot.configure", "users.manage", "quotes.manage"},
|
|
"moderator": {"quotes.moderate", "users.timeout"},
|
|
"user": {"quotes.create", "quotes.view"},
|
|
}
|
|
|
|
self._initialized = False
|
|
|
|
async def initialize(self):
|
|
"""Initialize security manager"""
|
|
if self._initialized:
|
|
return
|
|
|
|
try:
|
|
logger.info("Initializing security manager...")
|
|
self._initialized = True
|
|
logger.info("Security manager initialized")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize security: {e}")
|
|
raise
|
|
|
|
async def check_rate_limit(
|
|
self, limit_type: RateLimitType, user_id: int, guild_id: Optional[int] = None
|
|
) -> Tuple[bool, Dict]:
|
|
"""Check if request is within rate limits"""
|
|
try:
|
|
config = self.rate_limits.get(limit_type.value)
|
|
if not config:
|
|
return True, {}
|
|
|
|
rate_key = f"rate:{limit_type.value}:user:{user_id}"
|
|
current_time = int(time.time())
|
|
|
|
# Get usage from Redis
|
|
usage_data = await self.redis.get(rate_key)
|
|
if usage_data:
|
|
usage = json.loads(usage_data)
|
|
# Clean old entries
|
|
window_start = current_time - config.window
|
|
usage["requests"] = [r for r in usage["requests"] if r >= window_start]
|
|
else:
|
|
usage = {"requests": [], "burst_used": 0}
|
|
|
|
# Check limits
|
|
request_count = len(usage["requests"])
|
|
if request_count >= config.requests:
|
|
if config.burst > 0 and usage["burst_used"] < config.burst:
|
|
usage["burst_used"] += 1
|
|
else:
|
|
return False, {"rate_limited": True, "retry_after": config.window}
|
|
|
|
# Record request
|
|
usage["requests"].append(current_time)
|
|
|
|
# Store updated usage
|
|
await self.redis.setex(rate_key, config.window + 60, json.dumps(usage))
|
|
|
|
return True, {"remaining": max(0, config.requests - request_count)}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Rate limit error: {e}")
|
|
return True, {} # Fail open
|
|
|
|
async def validate_permissions(
|
|
self, user_id: int, guild_id: int, permission: str
|
|
) -> bool:
|
|
"""Validate user permissions"""
|
|
try:
|
|
user_permissions = await self._get_user_permissions(user_id, guild_id)
|
|
return permission in user_permissions or "*" in user_permissions
|
|
except Exception as e:
|
|
logger.error(f"Permission validation error: {e}")
|
|
return False
|
|
|
|
async def create_session(self, user_id: int, guild_id: int) -> str:
|
|
"""Create JWT session token"""
|
|
try:
|
|
session_id = secrets.token_urlsafe(32)
|
|
expires = int(time.time()) + self.session_timeout
|
|
|
|
payload = {
|
|
"user_id": user_id,
|
|
"guild_id": guild_id,
|
|
"session_id": session_id,
|
|
"exp": expires,
|
|
}
|
|
|
|
token = jwt.encode(payload, self.jwt_secret, algorithm="HS256")
|
|
|
|
# Store session
|
|
session_key = f"session:{user_id}:{session_id}"
|
|
await self.redis.setex(
|
|
session_key,
|
|
self.session_timeout,
|
|
json.dumps({"user_id": user_id, "guild_id": guild_id}),
|
|
)
|
|
|
|
return token
|
|
except Exception as e:
|
|
logger.error(f"Session creation error: {e}")
|
|
raise
|
|
|
|
async def authenticate_request(self, token: str) -> Optional[Dict]:
|
|
"""Authenticate JWT token"""
|
|
try:
|
|
payload = jwt.decode(token, self.jwt_secret, algorithms=["HS256"])
|
|
|
|
# Validate session exists
|
|
session_key = f"session:{payload['user_id']}:{payload['session_id']}"
|
|
session_data = await self.redis.get(session_key)
|
|
return payload if session_data else None
|
|
|
|
except jwt.InvalidTokenError:
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"Authentication error: {e}")
|
|
return None
|
|
|
|
async def log_security_event(
|
|
self, event_type: str, user_id: int, severity: str, message: str
|
|
):
|
|
"""Log security event"""
|
|
try:
|
|
event_data = {
|
|
"type": event_type,
|
|
"user_id": user_id,
|
|
"severity": severity,
|
|
"message": message,
|
|
"timestamp": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
event_key = f"security_event:{int(time.time())}:{secrets.token_hex(4)}"
|
|
await self.redis.setex(event_key, 86400 * 7, json.dumps(event_data))
|
|
|
|
if severity in ["high", "critical"]:
|
|
logger.critical(f"SECURITY: {event_type} - {message}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Security event logging error: {e}")
|
|
|
|
async def _get_user_permissions(self, user_id: int, guild_id: int) -> Set[str]:
|
|
"""Get user permissions based on roles"""
|
|
try:
|
|
# Default user permissions
|
|
permissions = set(self.role_permissions["user"])
|
|
|
|
# Get cached role or determine from Discord
|
|
role_key = f"user_role:{user_id}:{guild_id}"
|
|
cached_role = await self.redis.get(role_key)
|
|
|
|
if cached_role:
|
|
user_role = cached_role.decode()
|
|
else:
|
|
user_role = await self._determine_user_role(user_id, guild_id)
|
|
await self.redis.setex(role_key, 300, user_role) # 5 min cache
|
|
|
|
# Add role permissions
|
|
role_perms = self.role_permissions.get(user_role, set())
|
|
permissions.update(role_perms)
|
|
|
|
return permissions
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error getting permissions: {e}")
|
|
return set(self.role_permissions["user"])
|
|
|
|
async def _determine_user_role(self, user_id: int, guild_id: int) -> str:
|
|
"""Determine user role (simplified implementation)"""
|
|
# This would integrate with Discord API to check actual roles
|
|
# For now, return basic role determination
|
|
|
|
owner_ids = self.config.get("owner_ids", [])
|
|
if user_id in owner_ids:
|
|
return "owner"
|
|
|
|
admin_ids = self.config.get("admin_ids", [])
|
|
if user_id in admin_ids:
|
|
return "admin"
|
|
|
|
return "user"
|
|
|
|
async def check_health(self) -> Dict[str, Any]:
|
|
"""Check security system health"""
|
|
try:
|
|
active_sessions = len(await self.redis.keys("session:*"))
|
|
recent_events = len(await self.redis.keys("security_event:*"))
|
|
|
|
return {
|
|
"initialized": self._initialized,
|
|
"active_sessions": active_sessions,
|
|
"recent_security_events": recent_events,
|
|
"rate_limits_configured": len(self.rate_limits),
|
|
}
|
|
except Exception as e:
|
|
return {"error": str(e), "healthy": False}
|
|
|
|
|
|
# Decorators for Discord commands
|
|
def require_permissions(*permissions):
|
|
"""Require specific permissions for command"""
|
|
|
|
def decorator(func):
|
|
async def wrapper(self, interaction: discord.Interaction, *args, **kwargs):
|
|
security = getattr(self.bot, "security_manager", None)
|
|
if not security:
|
|
await interaction.response.send_message(
|
|
"Security unavailable", ephemeral=True
|
|
)
|
|
return
|
|
|
|
for permission in permissions:
|
|
if not await security.validate_permissions(
|
|
interaction.user.id, interaction.guild_id, permission
|
|
):
|
|
await interaction.response.send_message(
|
|
f"Missing permission: {permission}", ephemeral=True
|
|
)
|
|
return
|
|
|
|
return await func(self, interaction, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
def rate_limit(limit_type: RateLimitType):
|
|
"""Rate limit decorator for commands"""
|
|
|
|
def decorator(func):
|
|
async def wrapper(self, interaction: discord.Interaction, *args, **kwargs):
|
|
security = getattr(self.bot, "security_manager", None)
|
|
if not security:
|
|
return await func(self, interaction, *args, **kwargs)
|
|
|
|
allowed, info = await security.check_rate_limit(
|
|
limit_type, interaction.user.id, interaction.guild_id
|
|
)
|
|
|
|
if not allowed:
|
|
await interaction.response.send_message(
|
|
f"Rate limited. Try again in {info.get('retry_after', 60)}s",
|
|
ephemeral=True,
|
|
)
|
|
return
|
|
|
|
return await func(self, interaction, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|