| #!/usr/bin/env python3 |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| """ |
| Doris Security Management Module |
| Implements enterprise-level authentication, authorization, SQL security validation and data masking functionality |
| """ |
| |
| from __future__ import annotations |
| |
| import asyncio |
| import hashlib |
| import re |
| import time |
| from collections.abc import Callable, Mapping |
| from contextvars import ContextVar |
| from contextvars import Token as ContextToken |
| from dataclasses import dataclass, field |
| from datetime import datetime |
| from enum import Enum |
| from typing import TYPE_CHECKING, Any |
| |
| import sqlparse |
| from sqlparse.sql import Statement |
| from sqlparse.tokens import Keyword, Name |
| |
| from .auth_credentials import ( |
| EMPTY_CREDENTIAL, |
| BearerCredentials, |
| normalize_bearer_credentials, |
| ) |
| from .config import ( |
| DorisConfig, |
| EffectiveAuthConfig, |
| get_effective_auth_config, |
| ) |
| from .datetime_utils import utc_now |
| from .logger import get_logger |
| |
| if TYPE_CHECKING: |
| from ..auth.doris_oauth_provider import DorisOAuthProvider |
| from ..auth.jwt_manager import JWTManager |
| from ..auth.oauth_provider import OAuthAuthenticationProvider |
| from ..auth.token_manager import ( |
| DatabaseConfig as TokenDatabaseConfig, |
| ) |
| from ..auth.token_manager import ( |
| TokenInfo, |
| TokenManager, |
| ) |
| from .db import DorisConnectionManager |
| |
| # Global ContextVar for auth_context - must be a single instance shared across all modules |
| # This allows token-bound database configuration to work correctly in concurrent requests |
| mcp_auth_context_var: ContextVar[AuthContext | None] = ContextVar( |
| "mcp_auth_context", |
| default=None, |
| ) |
| |
| RESERVED_DORIS_OAUTH_PREFIX = "doa_" |
| # Backward-compatible name; the value is a public token namespace, not a secret. |
| RESERVED_DORIS_OAUTH_TOKEN_PREFIX = RESERVED_DORIS_OAUTH_PREFIX |
| ANONYMOUS_PRINCIPAL = "anonymous" |
| ANONYMOUS_SESSION_ID = "anonymous_session" |
| |
| |
| class SecurityLevel(Enum): |
| """Security level enumeration""" |
| |
| PUBLIC = "public" |
| INTERNAL = "internal" |
| CONFIDENTIAL = "confidential" |
| # Bandit audit: classification wire value, never credential material. |
| SECRET = "secret" # nosec B105 |
| |
| |
| @dataclass |
| class AuthContext: |
| """Authentication context for audit and session tracking""" |
| |
| token_id: str = "" # Token identifier for audit logging |
| user_id: str = "" # User identifier |
| roles: list[str] = field(default_factory=list) # User roles |
| permissions: list[str] = field(default_factory=list) # User permissions |
| security_level: SecurityLevel = field( |
| default_factory=lambda: SecurityLevel.INTERNAL |
| ) # Security level |
| client_ip: str = "unknown" # Client IP address |
| session_id: str = "" # Session identifier |
| login_time: datetime = field(default_factory=utc_now) |
| last_activity: datetime | None = None |
| token: str = "" # Raw token for token-bound database configuration |
| auth_method: str = "" # anonymous, token, jwt, external_oauth, doris_oauth |
| doris_user: str = "" |
| oauth_client_id: str = "" |
| oauth_scopes: list[str] = field(default_factory=list) |
| oauth_token_id: str = "" |
| oauth_issuer: str = "" |
| oauth_resource: str = "" |
| oauth_audiences: list[str] = field(default_factory=list) |
| pool_key: str = "" |
| semantic_tools_enabled: bool = False |
| semantic_resources_enabled: bool = False |
| doris_oauth_child_tools_enabled: bool = False |
| doris_oauth_child_tool_allowlist: tuple[str, ...] = field( |
| default_factory=tuple |
| ) |
| doris_oauth_db_tools_enabled: bool = False |
| doris_oauth_db_tool_allowlist: tuple[str, ...] = field(default_factory=tuple) |
| doris_oauth_query_tools_enabled: bool = False |
| doris_oauth_query_tool_allowlist: tuple[str, ...] = field(default_factory=tuple) |
| doris_oauth_explain_tools_enabled: bool = False |
| doris_oauth_explain_tool_allowlist: tuple[str, ...] = field(default_factory=tuple) |
| |
| |
| def get_current_auth_context() -> AuthContext | None: |
| """Return the current request auth context.""" |
| return mcp_auth_context_var.get() |
| |
| |
| def set_current_auth_context( |
| auth_context: AuthContext, |
| ) -> ContextToken[AuthContext | None]: |
| """Set auth context and return the ContextVar token for reset.""" |
| return mcp_auth_context_var.set(auth_context) |
| |
| |
| def clear_current_auth_context() -> ContextToken[AuthContext | None]: |
| """Force-clear auth context and return the ContextVar token.""" |
| return mcp_auth_context_var.set(None) |
| |
| |
| def reset_auth_context(token: ContextToken[AuthContext | None]) -> None: |
| """Reset auth context using the token returned by set_current_auth_context.""" |
| mcp_auth_context_var.reset(token) |
| |
| |
| @dataclass |
| class ValidationResult: |
| """Validation result""" |
| |
| is_valid: bool |
| error_message: str | None = None |
| risk_level: str = "low" |
| blocked_operations: list[str] = field(default_factory=list) |
| |
| |
| @dataclass |
| class MaskingRule: |
| """Data masking rule""" |
| |
| column_pattern: str |
| algorithm: str |
| parameters: dict[str, Any] |
| security_level: SecurityLevel |
| |
| |
| MaskingAlgorithm = Callable[[str, dict[str, Any]], str] |
| |
| |
| class DorisSecurityManager: |
| """Doris security manager |
| |
| Provides complete security control functionality, including authentication, authorization, SQL security validation and data masking |
| """ |
| |
| def __init__( |
| self, |
| config: DorisConfig, |
| connection_manager: DorisConnectionManager | None = None, |
| ) -> None: |
| self.config = config |
| self.logger = get_logger(__name__) |
| self.connection_manager = connection_manager |
| |
| # Initialize security components |
| self.auth_provider = AuthenticationProvider(config, self) |
| self.authz_provider = AuthorizationProvider(config) |
| self.sql_validator = SQLSecurityValidator(config) |
| self.masking_processor = DataMaskingProcessor(config) |
| |
| # Security rule configuration |
| self.blocked_keywords = self._load_blocked_keywords() |
| self.sensitive_tables = self._load_sensitive_tables() |
| self.masking_rules = self._load_masking_rules() |
| |
| # Track initialization state |
| self._initialized = False |
| self._token_database_validation_cache: dict[str, float] = {} |
| self._token_database_validation_locks: dict[str, asyncio.Lock] = {} |
| |
| async def initialize(self) -> None: |
| """Initialize security manager components""" |
| if self._initialized: |
| return |
| |
| try: |
| # Initialize authentication provider (for JWT setup) |
| await self.auth_provider.initialize() |
| |
| self._initialized = True |
| self.logger.info("DorisSecurityManager initialized successfully") |
| |
| except Exception as e: |
| self.logger.error(f"Failed to initialize DorisSecurityManager: {e}") |
| raise |
| |
| def _get_effective_auth_config(self) -> EffectiveAuthConfig: |
| return get_effective_auth_config(self.config) |
| |
| async def shutdown(self) -> None: |
| """Shutdown security manager components""" |
| try: |
| await self.auth_provider.shutdown() |
| self._token_database_validation_cache.clear() |
| self._token_database_validation_locks.clear() |
| self._initialized = False |
| self.logger.info("DorisSecurityManager shutdown completed") |
| |
| except Exception as e: |
| self.logger.error(f"Error during DorisSecurityManager shutdown: {e}") |
| raise |
| |
| def _load_blocked_keywords(self) -> set[str]: |
| """Load blocked SQL keywords from configuration""" |
| return set(self.config.security.blocked_keywords) |
| |
| def _load_sensitive_tables(self) -> dict[str, SecurityLevel]: |
| """Load sensitive table configuration""" |
| default_tables = { |
| "user_info": SecurityLevel.CONFIDENTIAL, |
| "payment_records": SecurityLevel.SECRET, |
| "employee_data": SecurityLevel.CONFIDENTIAL, |
| "public_reports": SecurityLevel.PUBLIC, |
| } |
| |
| for table_name, level in self.config.security.sensitive_tables.items(): |
| try: |
| default_tables[table_name] = SecurityLevel(level.lower()) |
| except ValueError: |
| default_tables[table_name] = SecurityLevel.INTERNAL |
| return default_tables |
| |
| def _load_masking_rules(self) -> list[MaskingRule]: |
| """Load data masking rules""" |
| default_rules = [ |
| MaskingRule( |
| column_pattern=r".*phone.*|.*mobile.*", |
| algorithm="phone_mask", |
| parameters={"mask_char": "*", "keep_prefix": 3, "keep_suffix": 4}, |
| security_level=SecurityLevel.INTERNAL, |
| ), |
| MaskingRule( |
| column_pattern=r".*email.*", |
| algorithm="email_mask", |
| parameters={"mask_char": "*"}, |
| security_level=SecurityLevel.INTERNAL, |
| ), |
| MaskingRule( |
| column_pattern=r".*id_card.*|.*identity.*", |
| algorithm="id_mask", |
| parameters={"mask_char": "*", "keep_prefix": 6, "keep_suffix": 4}, |
| security_level=SecurityLevel.CONFIDENTIAL, |
| ), |
| ] |
| |
| # Load custom rules from configuration |
| for rule_config in self.config.security.masking_rules: |
| default_rules.append(MaskingRule(**rule_config)) |
| |
| return default_rules |
| |
| async def authenticate_request( |
| self, |
| auth_input: BearerCredentials | Mapping[str, Any], |
| ) -> AuthContext: |
| """Validate request authentication information |
| |
| Tries authentication methods in normalized effective config order. |
| Any one method succeeding allows access |
| If all methods are disabled, returns anonymous context |
| """ |
| effective_auth = self._get_effective_auth_config() |
| credentials = normalize_bearer_credentials(auth_input) |
| legacy_auth_info: dict[str, Any] | None = None |
| if isinstance(auth_input, Mapping): |
| legacy_auth_info = dict(auth_input) |
| legacy_auth_type = ( |
| str(legacy_auth_info.get("type") or "") if legacy_auth_info else "" |
| ) |
| |
| if not effective_auth.auth_methods: |
| if legacy_auth_type and legacy_auth_info is not None: |
| return await self.auth_provider.authenticate(legacy_auth_info) |
| self.logger.debug("All authentication methods are disabled") |
| return AuthContext( |
| token_id=ANONYMOUS_PRINCIPAL, |
| user_id=ANONYMOUS_PRINCIPAL, |
| roles=[ANONYMOUS_PRINCIPAL], |
| permissions=["read"], |
| security_level=SecurityLevel.PUBLIC, |
| client_ip=credentials.client_ip, |
| session_id=ANONYMOUS_SESSION_ID, |
| auth_method=ANONYMOUS_PRINCIPAL, |
| pool_key="global", |
| ) |
| |
| last_error = None |
| |
| for auth_method in effective_auth.auth_methods: |
| try: |
| if auth_method == "doris_oauth": |
| return await self.auth_provider.authenticate_doris_oauth( |
| credentials |
| ) |
| if auth_method == "token": |
| return await self.auth_provider.authenticate_token(credentials) |
| if auth_method == "jwt": |
| return await self.auth_provider.authenticate_jwt(credentials) |
| if auth_method == "external_oauth": |
| return await self.auth_provider.authenticate_oauth(credentials) |
| except Exception as e: |
| self.logger.debug(f"{auth_method} authentication failed: {e}") |
| last_error = e |
| if auth_method == "doris_oauth" and credentials.token.startswith( |
| RESERVED_DORIS_OAUTH_TOKEN_PREFIX |
| ): |
| raise |
| |
| # All enabled authentication methods failed |
| from ..auth.oauth_token_validation import ( |
| OAuthAccessTokenValidationError, |
| ) |
| |
| if isinstance(last_error, OAuthAccessTokenValidationError): |
| raise last_error |
| error_message = ( |
| f"Authentication failed: {str(last_error)}" |
| if last_error |
| else "No authentication method succeeded" |
| ) |
| self.logger.warning( |
| f"Authentication failed for client {credentials.client_ip}: {error_message}" |
| ) |
| raise ValueError(error_message) |
| |
| async def authorize_resource_access( |
| self, auth_context: AuthContext, resource_uri: str |
| ) -> bool: |
| """Validate resource access permissions""" |
| return await self.authz_provider.check_permission( |
| auth_context, resource_uri, "read" |
| ) |
| |
| async def validate_sql_security( |
| self, sql: str, auth_context: AuthContext |
| ) -> ValidationResult: |
| """Validate SQL query security""" |
| return await self.sql_validator.validate(sql, auth_context) |
| |
| async def apply_data_masking( |
| self, data: list[dict[str, Any]], auth_context: AuthContext |
| ) -> list[dict[str, Any]]: |
| """Apply data masking processing""" |
| return await self.masking_processor.process(data, auth_context) |
| |
| # OAuth-specific methods |
| def get_oauth_authorization_url(self) -> tuple[str, str]: |
| """Get OAuth authorization URL |
| |
| Returns: |
| Tuple of (authorization_url, state) |
| """ |
| if not self.auth_provider.oauth_provider: |
| raise ValueError("OAuth is not enabled") |
| return self.auth_provider.oauth_provider.get_authorization_url() |
| |
| async def handle_oauth_callback(self, code: str, state: str) -> AuthContext: |
| """Handle OAuth callback |
| |
| Args: |
| code: Authorization code from OAuth provider |
| state: State parameter for CSRF protection |
| |
| Returns: |
| AuthContext for authenticated user |
| """ |
| if not self.auth_provider.oauth_provider: |
| raise ValueError("OAuth is not enabled") |
| return await self.auth_provider.oauth_provider.handle_callback(code, state) |
| |
| def get_oauth_provider_info(self) -> dict[str, Any]: |
| """Get OAuth provider information |
| |
| Returns: |
| OAuth provider information |
| """ |
| if not self.auth_provider.oauth_provider: |
| return {"enabled": False} |
| return self.auth_provider.oauth_provider.get_provider_info() |
| |
| # Token management methods |
| async def create_token( |
| self, |
| token_id: str, |
| expires_hours: int | None = None, |
| description: str = "", |
| custom_token: str | None = None, |
| database_config: TokenDatabaseConfig | None = None, |
| ) -> str: |
| """Create a new API access token |
| |
| Args: |
| token_id: Unique token identifier for audit and management |
| expires_hours: Token expiration in hours (None for no expiration) |
| description: Token description for management purposes |
| custom_token: Custom token string (if None, generates random token) |
| database_config: Optional database configuration for this token |
| |
| Returns: |
| Generated token string |
| """ |
| if not self.auth_provider.token_manager: |
| raise ValueError("Token manager not initialized") |
| |
| return await self.auth_provider.token_manager.create_token( |
| token_id=token_id, |
| expires_hours=expires_hours, |
| description=description, |
| custom_token=custom_token, |
| database_config=database_config, |
| ) |
| |
| async def revoke_token(self, token_id: str) -> bool: |
| """Revoke a token by token ID |
| |
| Args: |
| token_id: Token ID to revoke |
| |
| Returns: |
| True if token was revoked successfully |
| """ |
| if not self.auth_provider.token_manager: |
| raise ValueError("Token manager not initialized") |
| |
| return await self.auth_provider.token_manager.revoke_token(token_id) |
| |
| async def list_tokens(self) -> list[dict[str, Any]]: |
| """List all tokens (without sensitive data) |
| |
| Returns: |
| List of token information |
| """ |
| if not self.auth_provider.token_manager: |
| raise ValueError("Token manager not initialized") |
| |
| return await self.auth_provider.token_manager.list_tokens() |
| |
| async def cleanup_expired_tokens(self) -> int: |
| """Remove expired tokens and return count |
| |
| Returns: |
| Number of expired tokens removed |
| """ |
| if not self.auth_provider.token_manager: |
| return 0 |
| |
| return await self.auth_provider.token_manager.cleanup_expired_tokens() |
| |
| def get_token_stats(self) -> dict[str, Any]: |
| """Get token statistics |
| |
| Returns: |
| Token statistics dictionary |
| """ |
| if not self.auth_provider.token_manager: |
| return {"error": "Token manager not initialized"} |
| |
| return self.auth_provider.token_manager.get_token_stats() |
| |
| async def _validate_token_database_config( |
| self, |
| token: str, |
| token_info: TokenInfo, |
| ) -> None: |
| """Validate a token database route with a short successful-result cache. |
| |
| Authentication runs for every MCP request, including pings and capability |
| discovery. Reconnecting to Doris on each request creates avoidable pool |
| pressure, so successful validations are cached by token and database |
| configuration fingerprint. Failures are never cached. |
| |
| Args: |
| token: Raw authentication token |
| token_info: TokenInfo object from token validation |
| |
| Raises: |
| ValueError: If database configuration is invalid or connection fails |
| """ |
| if not self.connection_manager: |
| self.logger.warning( |
| "Connection manager not available for immediate database validation" |
| ) |
| return |
| |
| cache_key = self._token_database_validation_cache_key(token, token_info) |
| ttl_seconds = max( |
| 0, |
| int(self.config.security.token_db_validation_ttl_seconds), |
| ) |
| now = time.monotonic() |
| if self._token_database_validation_cache.get(cache_key, 0.0) > now: |
| self.logger.debug( |
| "Using cached database validation for token %s", |
| token_info.token_id, |
| ) |
| return |
| |
| lock = self._token_database_validation_locks.setdefault( |
| cache_key, |
| asyncio.Lock(), |
| ) |
| async with lock: |
| now = time.monotonic() |
| if self._token_database_validation_cache.get(cache_key, 0.0) > now: |
| return |
| try: |
| success, config_source = ( |
| await self.connection_manager.configure_for_token(token) |
| ) |
| if not success: |
| raise ValueError("Database configuration validation failed") |
| if ttl_seconds: |
| self._token_database_validation_cache[cache_key] = ( |
| time.monotonic() + ttl_seconds |
| ) |
| self.logger.info( |
| "Database configuration validated successfully for token %s " |
| "(source: %s)", |
| token_info.token_id, |
| config_source, |
| ) |
| except Exception as e: |
| self._token_database_validation_cache.pop(cache_key, None) |
| error_msg = ( |
| "Database configuration validation failed for token " |
| f"{token_info.token_id}: {str(e) or type(e).__name__}" |
| ) |
| self.logger.error(error_msg) |
| raise ValueError(error_msg) from e |
| |
| @staticmethod |
| def _token_database_validation_cache_key( |
| token: str, |
| token_info: TokenInfo, |
| ) -> str: |
| """Return a non-secret cache key that changes with bound DB credentials.""" |
| database_config = token_info.database_config |
| route = ( |
| getattr(database_config, "host", ""), |
| getattr(database_config, "port", ""), |
| getattr(database_config, "user", ""), |
| getattr(database_config, "password", ""), |
| getattr(database_config, "database", ""), |
| getattr(database_config, "charset", ""), |
| ) |
| digest = hashlib.sha256() |
| digest.update(token.encode("utf-8")) |
| for value in route: |
| digest.update(b"\0") |
| digest.update(str(value).encode("utf-8")) |
| return digest.hexdigest() |
| |
| |
| class AuthenticationProvider: |
| """Authentication provider""" |
| |
| def __init__( |
| self, |
| config: DorisConfig, |
| security_manager: DorisSecurityManager | None = None, |
| ) -> None: |
| self.config = config |
| self.logger = get_logger(__name__) |
| self.session_cache: dict[str, AuthContext] = {} |
| self.jwt_manager: JWTManager | None = None |
| self.oauth_provider: OAuthAuthenticationProvider | None = None |
| self.doris_oauth_provider: DorisOAuthProvider | None = None |
| self.token_manager: TokenManager | None = None |
| self.security_manager = security_manager |
| self.effective_auth = get_effective_auth_config(config) |
| |
| # Initialize authentication providers based on individual switches |
| auth_methods_enabled = [] |
| |
| # Initialize Token manager if enabled |
| if self.effective_auth.enable_token_auth: |
| self._initialize_token_manager() |
| auth_methods_enabled.append("Token") |
| |
| # Initialize JWT manager if enabled |
| if self.effective_auth.enable_jwt_auth: |
| self._initialize_jwt_manager() |
| auth_methods_enabled.append("JWT") |
| |
| # Initialize OAuth provider if enabled |
| if self.effective_auth.enable_external_oauth_auth: |
| self._initialize_oauth_provider() |
| auth_methods_enabled.append("OAuth") |
| |
| if self.effective_auth.enable_doris_oauth_auth: |
| self._initialize_doris_oauth_provider() |
| auth_methods_enabled.append("Doris OAuth") |
| |
| if auth_methods_enabled: |
| self.logger.info( |
| f"Authentication enabled with methods: {', '.join(auth_methods_enabled)}" |
| ) |
| else: |
| self.logger.info( |
| "All authentication methods are disabled - anonymous access allowed" |
| ) |
| |
| def _initialize_jwt_manager(self) -> None: |
| """Initialize JWT manager""" |
| try: |
| from ..auth.jwt_manager import JWTManager |
| |
| self.jwt_manager = JWTManager(self.config) |
| self.logger.info("JWT manager initialized") |
| except ImportError as e: |
| self.logger.error(f"Failed to import JWT manager: {e}") |
| raise |
| except Exception as e: |
| self.logger.error(f"Failed to initialize JWT manager: {e}") |
| raise |
| |
| def _initialize_token_manager(self) -> None: |
| """Initialize Token manager""" |
| try: |
| from ..auth.token_manager import TokenManager |
| |
| self.token_manager = TokenManager(self.config) |
| self.logger.info("Token manager initialized") |
| except ImportError as e: |
| self.logger.error(f"Failed to import Token manager: {e}") |
| raise |
| except Exception as e: |
| self.logger.error(f"Failed to initialize Token manager: {e}") |
| raise |
| |
| def _initialize_oauth_provider(self) -> None: |
| """Initialize OAuth provider""" |
| try: |
| from ..auth.oauth_provider import OAuthAuthenticationProvider |
| |
| self.oauth_provider = OAuthAuthenticationProvider(self.config) |
| self.logger.info("OAuth provider initialized") |
| except ImportError as e: |
| self.logger.error(f"Failed to import OAuth provider: {e}") |
| raise |
| except Exception as e: |
| self.logger.error(f"Failed to initialize OAuth provider: {e}") |
| raise |
| |
| def _initialize_doris_oauth_provider(self) -> None: |
| """Initialize Doris-backed OAuth provider shell.""" |
| try: |
| from ..auth.doris_oauth_provider import DorisOAuthProvider |
| |
| self.doris_oauth_provider = DorisOAuthProvider(self.config) |
| self.logger.info("Doris OAuth provider initialized") |
| except ImportError as e: |
| self.logger.error(f"Failed to import Doris OAuth provider: {e}") |
| raise |
| except Exception as e: |
| self.logger.error(f"Failed to initialize Doris OAuth provider: {e}") |
| raise |
| |
| def configure_doris_oauth( |
| self, |
| connection_manager: DorisConnectionManager, |
| ) -> None: |
| """Inject the Phase 2 connection manager after it is created.""" |
| if self.doris_oauth_provider: |
| self.doris_oauth_provider.configure_connection_manager(connection_manager) |
| |
| async def initialize(self) -> None: |
| """Initialize authentication provider asynchronously""" |
| if self.jwt_manager: |
| success = await self.jwt_manager.initialize() |
| if not success: |
| raise RuntimeError("Failed to initialize JWT manager") |
| self.logger.info("JWT authentication provider initialized successfully") |
| |
| if self.token_manager: |
| # Token manager doesn't need async initialization, just log success |
| self.logger.info("Token authentication provider initialized successfully") |
| |
| if self.oauth_provider: |
| success = await self.oauth_provider.initialize() |
| if not success: |
| raise RuntimeError("Failed to initialize OAuth provider") |
| self.logger.info("OAuth authentication provider initialized successfully") |
| |
| async def shutdown(self) -> None: |
| """Shutdown authentication provider""" |
| if self.jwt_manager: |
| await self.jwt_manager.shutdown() |
| self.logger.info("JWT authentication provider shutdown completed") |
| |
| if self.token_manager: |
| # Token manager doesn't need async shutdown, just log |
| self.logger.info("Token authentication provider shutdown completed") |
| |
| if self.oauth_provider: |
| await self.oauth_provider.shutdown() |
| self.logger.info("OAuth authentication provider shutdown completed") |
| |
| if self.doris_oauth_provider: |
| await self.doris_oauth_provider.shutdown() |
| self.logger.info("Doris OAuth authentication provider shutdown completed") |
| |
| async def authenticate(self, auth_info: dict[str, Any]) -> AuthContext: |
| """Legacy direct authentication entrypoint. |
| |
| Runtime HTTP auth goes through DorisSecurityManager.authenticate_request() |
| and EffectiveAuthConfig. This method is kept for existing direct callers |
| that pass an explicit auth_info["type"]. |
| """ |
| auth_type = str(auth_info.get("type") or "").strip().lower() |
| credentials = normalize_bearer_credentials(auth_info) |
| if auth_type == "token": |
| if not self.effective_auth.enable_token_auth: |
| raise ValueError("Token authentication is not enabled") |
| if not self.token_manager: |
| raise ValueError("Token manager is not initialized") |
| return await self.authenticate_token(credentials) |
| if auth_type == "basic": |
| raise ValueError( |
| "Basic authentication is not supported; use a configured " |
| "Bearer or OAuth authentication provider" |
| ) |
| if auth_type == "jwt": |
| return await self.authenticate_jwt(credentials) |
| if auth_type == "oauth": |
| if "code" in auth_info and "state" in auth_info: |
| if not self.effective_auth.enable_external_oauth_auth: |
| raise ValueError("OAuth authentication is not enabled") |
| if not self.oauth_provider: |
| raise ValueError("OAuth provider not initialized") |
| auth_context = await self.oauth_provider.handle_callback( |
| auth_info["code"], |
| auth_info["state"], |
| ) |
| auth_context.auth_method = "external_oauth" |
| auth_context.token = EMPTY_CREDENTIAL |
| auth_context.pool_key = "global" |
| return auth_context |
| return await self.authenticate_oauth(credentials) |
| if auth_type == "doris_oauth": |
| return await self.authenticate_doris_oauth(credentials) |
| raise ValueError(f"Unsupported authentication type: {auth_type or '<missing>'}") |
| |
| async def authenticate_token( |
| self, |
| credentials: BearerCredentials, |
| ) -> AuthContext: |
| """Perform token authentication""" |
| if not self.effective_auth.enable_token_auth: |
| raise ValueError("Token authentication is not enabled") |
| return await self._authenticate_token(credentials) |
| |
| async def authenticate_jwt( |
| self, |
| credentials: BearerCredentials, |
| ) -> AuthContext: |
| """Perform JWT authentication""" |
| if not self.effective_auth.enable_jwt_auth: |
| raise ValueError("JWT authentication is not enabled") |
| return await self._authenticate_jwt(credentials) |
| |
| async def authenticate_oauth( |
| self, |
| credentials: BearerCredentials, |
| ) -> AuthContext: |
| """Perform OAuth authentication""" |
| if not self.effective_auth.enable_external_oauth_auth: |
| raise ValueError("OAuth authentication is not enabled") |
| return await self._authenticate_oauth(credentials) |
| |
| async def authenticate_doris_oauth( |
| self, |
| credentials: BearerCredentials, |
| ) -> AuthContext: |
| """Authenticate a Doris OAuth doa_ access token.""" |
| if not self.effective_auth.enable_doris_oauth_auth: |
| raise ValueError("Doris OAuth authentication is not enabled") |
| if not credentials.is_bearer or not credentials.token.startswith( |
| RESERVED_DORIS_OAUTH_TOKEN_PREFIX |
| ): |
| raise ValueError("Missing Doris OAuth access token") |
| if not self.doris_oauth_provider: |
| raise ValueError("Doris OAuth provider is not initialized") |
| return await self.doris_oauth_provider.authenticate_access_token(credentials) |
| |
| async def _authenticate_jwt( |
| self, |
| credentials: BearerCredentials, |
| ) -> AuthContext: |
| """JWT authentication""" |
| if not self.jwt_manager: |
| raise ValueError("JWT manager not initialized") |
| if not credentials.is_bearer: |
| raise ValueError("Missing JWT token") |
| |
| try: |
| # Use JWT middleware for authentication |
| from ..auth.auth_middleware import AuthMiddleware |
| |
| middleware = AuthMiddleware(self.jwt_manager) |
| return await middleware.authenticate_request(credentials) |
| |
| except Exception as e: |
| self.logger.error(f"JWT authentication failed: {e}") |
| raise ValueError(f"JWT authentication failed: {str(e)}") |
| |
| async def _authenticate_oauth( |
| self, |
| credentials: BearerCredentials, |
| ) -> AuthContext: |
| """OAuth authentication""" |
| if not self.oauth_provider: |
| raise ValueError("OAuth provider not initialized") |
| if not credentials.is_bearer: |
| raise ValueError("Missing external OAuth access token") |
| |
| auth_context = await self.oauth_provider.authenticate_with_token( |
| credentials.token |
| ) |
| auth_context.auth_method = "external_oauth" |
| auth_context.token = EMPTY_CREDENTIAL |
| auth_context.pool_key = "global" |
| return auth_context |
| |
| async def _authenticate_token( |
| self, |
| credentials: BearerCredentials, |
| ) -> AuthContext: |
| """Token authentication""" |
| if not self.token_manager: |
| raise ValueError("Token manager not initialized") |
| if not credentials.is_static_token: |
| raise ValueError("Missing authentication token") |
| token = credentials.token |
| |
| try: |
| # Validate token using TokenManager |
| validation_result = await self.token_manager.validate_token(token) |
| |
| if not validation_result.is_valid: |
| raise ValueError( |
| f"Token validation failed: {validation_result.error_message}" |
| ) |
| |
| token_info = validation_result.token_info |
| if token_info is None: |
| raise ValueError("Token validation did not return token information") |
| |
| # Immediately validate database configuration for this token |
| if self.security_manager: |
| await self.security_manager._validate_token_database_config( |
| token, token_info |
| ) |
| |
| return AuthContext( |
| token_id=token_info.token_id, |
| user_id=token_info.token_id, # Use token_id as user_id for token auth |
| roles=["token_user"], # Default role for token users |
| permissions=["read", "write"], # Default permissions for token users |
| security_level=SecurityLevel.INTERNAL, |
| client_ip=credentials.client_ip, |
| session_id=credentials.session_id or f"session_{token_info.token_id}", |
| login_time=utc_now(), |
| last_activity=token_info.last_used, |
| token=token, # Store raw token for token-bound database configuration |
| auth_method="token", |
| pool_key=f"static_token:{token_info.token_id}", |
| ) |
| |
| except Exception as e: |
| self.logger.error(f"Token authentication failed: {e}") |
| raise ValueError(f"Token authentication failed: {str(e)}") |
| |
| |
| class AuthorizationProvider: |
| """Authorization provider""" |
| |
| def __init__(self, config: DorisConfig | Mapping[str, Any]) -> None: |
| self.config = config |
| self.logger = get_logger(__name__) |
| self.permission_cache: dict[str, bool] = {} |
| |
| # Load sensitive tables configuration |
| self.sensitive_tables = self._load_sensitive_tables() |
| |
| def _load_sensitive_tables(self) -> dict[str, SecurityLevel]: |
| """Load sensitive table configuration""" |
| default_tables = { |
| "user_info": SecurityLevel.CONFIDENTIAL, |
| "payment_records": SecurityLevel.SECRET, |
| "employee_data": SecurityLevel.CONFIDENTIAL, |
| "public_reports": SecurityLevel.PUBLIC, |
| } |
| |
| if isinstance(self.config, Mapping): |
| config_tables = self.config.get("sensitive_tables", {}) |
| # Convert string values to SecurityLevel enum |
| for table_name, level in config_tables.items(): |
| if isinstance(level, str): |
| try: |
| default_tables[table_name] = SecurityLevel(level.lower()) |
| except ValueError: |
| default_tables[table_name] = SecurityLevel.INTERNAL |
| else: |
| default_tables[table_name] = level |
| return default_tables |
| else: |
| return default_tables |
| |
| async def check_permission( |
| self, auth_context: AuthContext, resource_uri: str, action: str |
| ) -> bool: |
| """Check permissions""" |
| # Parse resource information |
| resource_info = self._parse_resource_uri(resource_uri) |
| |
| # First check security level - this is mandatory |
| if not await self._check_security_level_permission(auth_context, resource_info): |
| return False |
| |
| # Then check role-based permissions |
| if await self._check_role_permission(auth_context, resource_info, action): |
| return True |
| |
| # Finally check user-based permissions |
| if await self._check_user_permission(auth_context, resource_info, action): |
| return True |
| |
| return False |
| |
| def _parse_resource_uri(self, uri: str) -> dict[str, str]: |
| """Parse resource URI""" |
| parts = uri.split("/") |
| if len(parts) >= 3: |
| return { |
| "type": parts[2], # table, view, etc. |
| "name": parts[3] if len(parts) > 3 else "", |
| "schema": parts[4] if len(parts) > 4 else "default", |
| } |
| return {"type": "unknown", "name": "", "schema": "default"} |
| |
| async def _check_role_permission( |
| self, auth_context: AuthContext, resource_info: dict[str, str], action: str |
| ) -> bool: |
| """Check role-based permissions""" |
| # Role permission mapping |
| role_permissions = { |
| "data_analyst": {"table": ["read"], "view": ["read"]}, |
| "data_admin": { |
| "table": ["read", "write", "admin"], |
| "view": ["read", "write", "admin"], |
| }, |
| } |
| |
| for role in auth_context.roles: |
| role_perms = role_permissions.get(role, {}) |
| resource_perms = role_perms.get(resource_info["type"], []) |
| if action in resource_perms: |
| return True |
| |
| return False |
| |
| async def _check_user_permission( |
| self, auth_context: AuthContext, resource_info: dict[str, str], action: str |
| ) -> bool: |
| """Check user-based permissions""" |
| # User-specific permission check |
| if "admin" in auth_context.permissions: |
| return True |
| |
| if action == "read" and "read_data" in auth_context.permissions: |
| return True |
| |
| return False |
| |
| async def _check_security_level_permission( |
| self, auth_context: AuthContext, resource_info: dict[str, str] |
| ) -> bool: |
| """Check security level permissions""" |
| # Get resource security level |
| resource_security_level = self._get_resource_security_level(resource_info) |
| |
| # Check if user security level is sufficient |
| security_hierarchy = { |
| SecurityLevel.PUBLIC: 0, |
| SecurityLevel.INTERNAL: 1, |
| SecurityLevel.CONFIDENTIAL: 2, |
| SecurityLevel.SECRET: 3, |
| } |
| |
| user_level = security_hierarchy.get(auth_context.security_level, 0) |
| resource_level = security_hierarchy.get(resource_security_level, 0) |
| |
| # User must have higher or equal security level to access resource |
| return user_level >= resource_level |
| |
| def _get_resource_security_level( |
| self, resource_info: dict[str, str] |
| ) -> SecurityLevel: |
| """Get resource security level""" |
| # Get table security level from configuration |
| table_name = resource_info.get("name", "") |
| |
| # Use the loaded sensitive tables |
| sensitive_tables = self.sensitive_tables |
| |
| # Convert string values to SecurityLevel enum if needed |
| security_level = sensitive_tables.get(table_name, SecurityLevel.INTERNAL) |
| if isinstance(security_level, str): |
| try: |
| security_level = SecurityLevel(security_level.lower()) |
| except ValueError: |
| security_level = SecurityLevel.INTERNAL |
| |
| return security_level |
| |
| |
| class SQLSecurityValidator: |
| """SQL security validator""" |
| |
| def __init__(self, config: DorisConfig | Mapping[str, Any]) -> None: |
| self.config = config |
| self.logger = get_logger(__name__) |
| self.blocked_keywords: set[str] |
| self.max_query_complexity: int |
| self.enable_security_check: bool |
| |
| # Handle DorisConfig object or dictionary configuration |
| if isinstance(config, Mapping): |
| # Dictionary configuration |
| self.blocked_keywords = { |
| str(keyword) for keyword in config.get("blocked_keywords", []) |
| } |
| self.max_query_complexity = int(config.get("max_query_complexity", 100)) |
| self.enable_security_check = bool(config.get("enable_security_check", True)) |
| else: |
| # DorisConfig object with security attribute - unified source from config |
| self.blocked_keywords = set(config.security.blocked_keywords) |
| self.max_query_complexity = config.security.max_query_complexity |
| self.enable_security_check = config.security.enable_security_check |
| |
| async def validate(self, sql: str, auth_context: AuthContext) -> ValidationResult: |
| """Validate SQL query security""" |
| # If security check is disabled, always return valid |
| if not self.enable_security_check: |
| self.logger.debug("SQL security check is disabled, allowing all queries") |
| return ValidationResult(is_valid=True) |
| |
| try: |
| # SECURITY FIX: Parse ALL SQL statements, not just the first one |
| # This prevents bypassing security checks by injecting additional statements |
| all_statements = sqlparse.parse(sql) |
| |
| if not all_statements: |
| return ValidationResult( |
| is_valid=False, |
| error_message="Empty or invalid SQL statement", |
| risk_level="medium", |
| ) |
| |
| # SECURITY FIX: Validate each statement individually |
| for idx, parsed in enumerate(all_statements): |
| # Skip empty statements (e.g., from trailing semicolons) |
| if not parsed.tokens or str(parsed).strip() == "": |
| continue |
| |
| self.logger.debug( |
| "Validating SQL statement %s/%s", |
| idx + 1, |
| len(all_statements), |
| ) |
| |
| # Check blocked operations first (more specific) |
| keyword_result = await self._check_blocked_keywords(parsed) |
| if not keyword_result.is_valid: |
| keyword_result.error_message = ( |
| f"Statement {idx + 1}: {keyword_result.error_message}" |
| ) |
| return keyword_result |
| |
| # Check SQL injection risks |
| injection_result = await self._check_sql_injection(sql, parsed) |
| if not injection_result.is_valid: |
| injection_result.error_message = ( |
| f"Statement {idx + 1}: {injection_result.error_message}" |
| ) |
| return injection_result |
| |
| # Check query complexity |
| complexity_result = await self._check_query_complexity(parsed) |
| if not complexity_result.is_valid: |
| complexity_result.error_message = ( |
| f"Statement {idx + 1}: {complexity_result.error_message}" |
| ) |
| return complexity_result |
| |
| # Check table access permissions |
| table_result = await self._check_table_access(parsed, auth_context) |
| if not table_result.is_valid: |
| table_result.error_message = ( |
| f"Statement {idx + 1}: {table_result.error_message}" |
| ) |
| return table_result |
| |
| # The legacy validator remains in use below the formal Query |
| # runtime. Reuse the canonical fail-closed allowlist so unknown |
| # Doris statements cannot pass merely because they are absent |
| # from blocked_keywords. |
| from .query_runtime import QueryRuntimeFailure, ReadOnlySQLGuard |
| |
| try: |
| ReadOnlySQLGuard.validate(sql) |
| except QueryRuntimeFailure as exc: |
| return ValidationResult( |
| is_valid=False, |
| error_message=f"Read-only SQL policy violation: {exc}", |
| risk_level="high", |
| ) |
| |
| return ValidationResult(is_valid=True) |
| |
| except Exception as e: |
| self.logger.error(f"SQL security validation failed: {e}") |
| return ValidationResult( |
| is_valid=False, |
| error_message=f"SQL parsing error: {str(e)}", |
| risk_level="high", |
| ) |
| |
| async def _check_sql_injection( |
| self, sql: str, parsed: Statement |
| ) -> ValidationResult: |
| """Check SQL injection risks with improved pattern detection |
| |
| FIX for Issue #62 Bug 2: Improved patterns to reduce false positives |
| Now better distinguishes between legitimate SQL (like BETWEEN...AND) and injection attempts |
| """ |
| # Improved injection patterns that are more specific and less prone to false positives |
| injection_patterns = [ |
| # Stacked queries with dangerous operations (true injection risk) |
| r";\s*(DROP|DELETE|TRUNCATE|ALTER|CREATE|INSERT|UPDATE)\s+", |
| # UNION-based injection (but allow legitimate UNION queries) |
| # Only flag if UNION is followed by suspicious patterns like SELECT with WHERE 1=1 |
| r"UNION\s+(ALL\s+)?SELECT\s+.*\s+(WHERE|AND|OR)\s+\d+\s*=\s*\d+", |
| r"UNION\s+(ALL\s+)?SELECT\s+.*\b(PASSWORD|SECRET|TOKEN|CREDENTIAL|ADMIN)\b", |
| # Boolean tautology injection. BETWEEN clauses are cleaned below before this is applied. |
| r"\b(OR|AND)\s+\d+\s*=\s*\d+\b", |
| # Boolean-based blind injection with comments (true injection pattern) |
| r"(WHERE|AND|OR)\s+\d+\s*=\s*\d+\s*(--|#|/\*)", |
| # Quote-based injection attempts (but not in legitimate strings) |
| r"(WHERE|AND|OR)\s+(['\"])[^\2]*\2\s*=\s*\2[^\2]*\2", |
| # Time-based blind injection |
| r"(SLEEP|WAITFOR|BENCHMARK)\s*\(", |
| # System stored procedure injection |
| r"(EXEC|EXECUTE|SP_|XP_)\s*\(", |
| # Script injection attempts |
| r"<\s*(SCRIPT|JAVASCRIPT|VBSCRIPT)", |
| ] |
| |
| # FIX: Don't flag legitimate SQL functions and keywords |
| # These patterns are too broad and cause false positives: |
| # - REMOVED: r"(char|ascii|substring|concat)\s*\(" - These are legitimate SQL functions |
| # - REMOVED: r"(\s|^)(or|and)\s+\d+\s*=\s*\d+" - This flags BETWEEN...AND constructs |
| # - REMOVED: r"(\s|^)(or|and)\s+['\"].*['\"]" - This is too broad |
| |
| sql_upper = sql.upper() |
| |
| # Special case: Allow BETWEEN...AND which is legitimate SQL |
| # This prevents false positives like "WHERE dt BETWEEN '2025-01-01' AND '2025-01-31'" |
| if "BETWEEN" in sql_upper and "AND" in sql_upper: |
| # This is likely a BETWEEN clause, not injection |
| # Check if AND appears in a BETWEEN context |
| between_pattern = r"BETWEEN\s+[^\s]+\s+AND\s+[^\s]+" |
| if re.search(between_pattern, sql_upper, re.IGNORECASE): |
| # Remove BETWEEN clauses before checking other patterns |
| sql_cleaned = re.sub( |
| between_pattern, "BETWEEN_CLAUSE", sql_upper, flags=re.IGNORECASE |
| ) |
| sql_to_check = sql_cleaned |
| else: |
| sql_to_check = sql_upper |
| else: |
| sql_to_check = sql_upper |
| |
| for pattern in injection_patterns: |
| if re.search(pattern, sql_to_check, re.IGNORECASE): |
| self.logger.warning( |
| f"Potential SQL injection pattern detected: {pattern}" |
| ) |
| return ValidationResult( |
| is_valid=False, |
| error_message="Potential SQL injection risk detected", |
| risk_level="high", |
| ) |
| |
| # Check suspicious quotes and comments (with improved detection) |
| if self._has_suspicious_quotes_or_comments(sql): |
| return ValidationResult( |
| is_valid=False, |
| error_message="Suspicious quote or comment pattern detected", |
| risk_level="medium", |
| ) |
| |
| return ValidationResult(is_valid=True) |
| |
| def _has_suspicious_quotes_or_comments(self, sql: str) -> bool: |
| """Check suspicious quote and comment patterns with improved detection |
| |
| FIX for Issue #62 Bug 2: Improved detection to reduce false positives |
| Now distinguishes between legitimate comments/strings and injection attempts |
| """ |
| try: |
| # Use sqlparse to parse the SQL and distinguish between code and comments/strings |
| import sqlparse |
| from sqlparse.tokens import Comment, String |
| |
| # Parse the SQL |
| parsed = sqlparse.parse(sql) |
| if not parsed: |
| # If parsing fails, be conservative |
| return True |
| |
| statement = parsed[0] |
| |
| # Check for unmatched quotes ONLY in non-string tokens |
| # This prevents false positives from legitimate string content |
| non_string_content = [] |
| |
| for token in statement.flatten(): |
| if token.ttype in (String.Single, String.Double): |
| # Skip string content - quotes inside strings are legitimate |
| continue |
| elif token.ttype in (Comment.Single, Comment.Multi): |
| # Comments are generally OK, but check for suspicious injection patterns |
| comment_value = str(token).lower() |
| comment_body = re.sub(r"^(--|#|/\*)\s*", "", comment_value).strip() |
| truncation_pattern = r"^(and|or|union|select|drop|delete|insert|update|exec|execute)\b" |
| sensitive_pattern = r"\b(admin|credential|password|secret|token)\b" |
| if re.search(truncation_pattern, comment_body) or re.search( |
| sensitive_pattern, comment_body |
| ): |
| self.logger.warning( |
| f"Suspicious SQL keyword in comment: {token}" |
| ) |
| return True |
| # Normal comments are OK |
| continue |
| else: |
| # Accumulate non-string, non-comment content |
| non_string_content.append(str(token)) |
| |
| # Check for unmatched quotes in non-string content |
| non_string_text = "".join(non_string_content) |
| single_quotes = non_string_text.count("'") |
| double_quotes = non_string_text.count('"') |
| |
| # Only flag if there are unmatched quotes in actual SQL code (not in strings) |
| if single_quotes % 2 != 0 or double_quotes % 2 != 0: |
| return True |
| |
| # FIX: Don't flag legitimate SQL comments |
| # Comments are OK as long as they don't contain dangerous patterns (already checked above) |
| |
| return False |
| |
| except Exception as e: |
| self.logger.debug(f"SQL parsing error in quote/comment check: {e}") |
| # On parsing error, fall back to conservative check |
| # But be more lenient than before |
| return False # Don't flag on parse errors to reduce false positives |
| |
| async def _check_blocked_keywords(self, parsed: Statement) -> ValidationResult: |
| """Check blocked keywords""" |
| blocked_operations = [] |
| |
| # Check all tokens in the parsed statement |
| for token in parsed.flatten(): |
| # Check if token is a keyword (including DML/DDL) or name that matches blocked operations |
| if ( |
| token.ttype is Keyword |
| or token.ttype is Name |
| or (token.ttype and str(token.ttype).startswith("Token.Keyword")) |
| ): |
| token_value = token.value.upper().strip() |
| if token_value in self.blocked_keywords: |
| blocked_operations.append(token_value) |
| # Also check for DDL/DML keywords in token values |
| elif hasattr(token, "value") and token.value: |
| token_value = token.value.upper().strip() |
| for blocked_keyword in self.blocked_keywords: |
| if blocked_keyword in token_value: |
| blocked_operations.append(blocked_keyword) |
| |
| if blocked_operations: |
| return ValidationResult( |
| is_valid=False, |
| error_message=f"Contains blocked operations: {', '.join(set(blocked_operations))}", |
| risk_level="high", |
| blocked_operations=list(set(blocked_operations)), |
| ) |
| |
| return ValidationResult(is_valid=True) |
| |
| async def _check_query_complexity(self, parsed: Statement) -> ValidationResult: |
| """Check query complexity""" |
| complexity_score = 0 |
| |
| # Calculate complexity score |
| for token in parsed.flatten(): |
| if token.ttype is Keyword: |
| keyword = token.value.upper() |
| if keyword in ["JOIN", "INNER", "LEFT", "RIGHT", "FULL"]: |
| complexity_score += 10 |
| elif keyword in ["UNION", "INTERSECT", "EXCEPT"]: |
| complexity_score += 15 |
| elif keyword in ["GROUP BY", "ORDER BY", "HAVING"]: |
| complexity_score += 5 |
| elif keyword in ["SUBQUERY", "EXISTS", "IN"]: |
| complexity_score += 8 |
| |
| if complexity_score > self.max_query_complexity: |
| return ValidationResult( |
| is_valid=False, |
| error_message=f"Query complexity too high (score: {complexity_score}, limit: {self.max_query_complexity})", |
| risk_level="medium", |
| ) |
| |
| return ValidationResult(is_valid=True) |
| |
| async def _check_table_access( |
| self, parsed: Statement, auth_context: AuthContext |
| ) -> ValidationResult: |
| """Check table access permissions""" |
| # Extract table names from query |
| tables = self._extract_table_names(parsed) |
| |
| # Check access permissions for each table |
| unauthorized_tables = [] |
| for table in tables: |
| # Should call authorization provider to check permissions |
| # Simplified implementation, assume some tables require special permissions |
| if ( |
| table.lower() in ["sensitive_data", "admin_logs"] |
| and "admin" not in auth_context.roles |
| ): |
| unauthorized_tables.append(table) |
| |
| if unauthorized_tables: |
| return ValidationResult( |
| is_valid=False, |
| error_message=f"No access to tables: {', '.join(unauthorized_tables)}", |
| risk_level="high", |
| ) |
| |
| return ValidationResult(is_valid=True) |
| |
| def _extract_table_names(self, parsed: Statement) -> list[str]: |
| """Extract table names from SQL statement""" |
| tables = [] |
| |
| # Simplified table name extraction logic |
| tokens = list(parsed.flatten()) |
| for i, token in enumerate(tokens): |
| if token.ttype is Keyword and token.value.upper() == "FROM": |
| # Find table name after FROM |
| for j in range(i + 1, len(tokens)): |
| next_token = tokens[j] |
| if next_token.ttype is Name: |
| tables.append(next_token.value) |
| break |
| elif next_token.ttype is Keyword: |
| break |
| |
| return tables |
| |
| |
| class DataMaskingProcessor: |
| """Data masking processor""" |
| |
| def __init__(self, config: DorisConfig | Mapping[str, Any]) -> None: |
| self.config = config |
| self.logger = get_logger(__name__) |
| self.masking_algorithms: dict[str, MaskingAlgorithm] = ( |
| self._init_masking_algorithms() |
| ) |
| self.masking_rules = self._load_masking_rules() |
| |
| def _load_masking_rules(self) -> list[MaskingRule]: |
| """Load data masking rules""" |
| default_rules = [ |
| MaskingRule( |
| column_pattern=r".*phone.*|.*mobile.*", |
| algorithm="phone_mask", |
| parameters={"mask_char": "*", "keep_prefix": 3, "keep_suffix": 4}, |
| security_level=SecurityLevel.INTERNAL, |
| ), |
| MaskingRule( |
| column_pattern=r".*email.*", |
| algorithm="email_mask", |
| parameters={"mask_char": "*"}, |
| security_level=SecurityLevel.INTERNAL, |
| ), |
| MaskingRule( |
| column_pattern=r".*id_card.*|.*identity.*", |
| algorithm="id_mask", |
| parameters={"mask_char": "*", "keep_prefix": 6, "keep_suffix": 4}, |
| security_level=SecurityLevel.CONFIDENTIAL, |
| ), |
| ] |
| |
| # Load custom rules from configuration |
| if isinstance(self.config, Mapping): |
| custom_rules = self.config.get("masking_rules", []) |
| for rule_config in custom_rules: |
| if isinstance(rule_config, dict): |
| # Convert string security level to enum |
| if "security_level" in rule_config and isinstance( |
| rule_config["security_level"], str |
| ): |
| try: |
| rule_config["security_level"] = SecurityLevel( |
| rule_config["security_level"].lower() |
| ) |
| except ValueError: |
| rule_config["security_level"] = SecurityLevel.INTERNAL |
| default_rules.append(MaskingRule(**rule_config)) |
| elif isinstance(rule_config, MaskingRule): |
| default_rules.append(rule_config) |
| |
| return default_rules |
| |
| def _init_masking_algorithms(self) -> dict[str, MaskingAlgorithm]: |
| """Initialize masking algorithms""" |
| return { |
| "phone_mask": self._mask_phone, |
| "email_mask": self._mask_email, |
| "id_mask": self._mask_id_card, |
| "name_mask": self._mask_name, |
| "partial_mask": self._mask_partial, |
| } |
| |
| async def process( |
| self, data: list[dict[str, Any]], auth_context: AuthContext |
| ) -> list[dict[str, Any]]: |
| """Process data masking""" |
| if not data: |
| return data |
| |
| # Get applicable masking rules |
| applicable_rules = self._get_applicable_rules(auth_context) |
| |
| masked_data: list[dict[str, Any]] = [] |
| for row in data: |
| masked_row: dict[str, Any] = {} |
| for column, value in row.items(): |
| masked_value = await self._apply_masking_rules( |
| column, value, applicable_rules |
| ) |
| masked_row[column] = masked_value |
| masked_data.append(masked_row) |
| |
| return masked_data |
| |
| def _get_applicable_rules(self, auth_context: AuthContext) -> list[MaskingRule]: |
| """Get applicable masking rules""" |
| applicable_rules: list[MaskingRule] = [] |
| |
| for rule in self.masking_rules: |
| # Decide whether to apply masking rules based on user security level |
| if self._should_apply_rule(rule, auth_context): |
| applicable_rules.append(rule) |
| |
| return applicable_rules |
| |
| def _should_apply_rule(self, rule: MaskingRule, auth_context: AuthContext) -> bool: |
| """Determine whether masking rule should be applied""" |
| # Admin users can see original data |
| if "admin" in auth_context.roles: |
| return False |
| |
| # Decide based on security level |
| security_hierarchy = { |
| SecurityLevel.PUBLIC: 0, |
| SecurityLevel.INTERNAL: 1, |
| SecurityLevel.CONFIDENTIAL: 2, |
| SecurityLevel.SECRET: 3, |
| } |
| |
| user_level = security_hierarchy.get(auth_context.security_level, 0) |
| rule_level = security_hierarchy.get(rule.security_level, 0) |
| |
| # Apply masking if user level is less than or equal to rule level |
| return user_level <= rule_level |
| |
| async def _apply_masking_rules( |
| self, column: str, value: Any, rules: list[MaskingRule] |
| ) -> Any: |
| """Apply masking rules""" |
| if value is None: |
| return value |
| |
| for rule in rules: |
| if re.match(rule.column_pattern, column, re.IGNORECASE): |
| algorithm = self.masking_algorithms.get(rule.algorithm) |
| if algorithm: |
| return algorithm(str(value), rule.parameters) |
| |
| return value |
| |
| def _mask_phone(self, value: str, params: dict[str, Any]) -> str: |
| """Phone number masking""" |
| if len(value) < 7: |
| return value |
| |
| mask_char = str(params.get("mask_char", "*")) |
| keep_prefix = int(params.get("keep_prefix", 3)) |
| keep_suffix = int(params.get("keep_suffix", 4)) |
| |
| if len(value) <= keep_prefix + keep_suffix: |
| return mask_char * len(value) |
| |
| prefix = value[:keep_prefix] |
| suffix = value[-keep_suffix:] |
| middle_length = len(value) - keep_prefix - keep_suffix |
| |
| return prefix + mask_char * middle_length + suffix |
| |
| def _mask_email(self, value: str, params: dict[str, Any]) -> str: |
| """Email masking""" |
| if "@" not in value: |
| return value |
| |
| mask_char = str(params.get("mask_char", "*")) |
| local, domain = value.split("@", 1) |
| |
| if len(local) <= 2: |
| masked_local = mask_char * len(local) |
| else: |
| masked_local = local[0] + mask_char * (len(local) - 2) + local[-1] |
| |
| return f"{masked_local}@{domain}" |
| |
| def _mask_id_card(self, value: str, params: dict[str, Any]) -> str: |
| """ID card number masking""" |
| if len(value) < 10: |
| return value |
| |
| mask_char = str(params.get("mask_char", "*")) |
| keep_prefix = int(params.get("keep_prefix", 6)) |
| keep_suffix = int(params.get("keep_suffix", 4)) |
| |
| if len(value) <= keep_prefix + keep_suffix: |
| return mask_char * len(value) |
| |
| prefix = value[:keep_prefix] |
| suffix = value[-keep_suffix:] |
| middle_length = len(value) - keep_prefix - keep_suffix |
| |
| return prefix + mask_char * middle_length + suffix |
| |
| def _mask_name(self, value: str, params: dict[str, Any]) -> str: |
| """Name masking""" |
| if len(value) <= 1: |
| return value |
| |
| mask_char = str(params.get("mask_char", "*")) |
| |
| if len(value) == 2: |
| return value[0] + mask_char |
| else: |
| return value[0] + mask_char * (len(value) - 2) + value[-1] |
| |
| def _mask_partial(self, value: str, params: dict[str, Any]) -> str: |
| """Partial masking""" |
| mask_char = str(params.get("mask_char", "*")) |
| mask_ratio = float(params.get("mask_ratio", 0.5)) |
| |
| mask_length = int(len(value) * mask_ratio) |
| start_pos = (len(value) - mask_length) // 2 |
| |
| result = list(value) |
| for i in range(start_pos, start_pos + mask_length): |
| if i < len(result): |
| result[i] = mask_char |
| |
| return "".join(result) |
