working version of user and agent auth

Author: aarora79

.gitignore

@@ -132,6 +132,8 @@ celerybeat.pid
 
 # Environments
 .env
+.env.user
+.env.agent
 .env.docker
 .venv
 env/

agents/.env.template

@@ -11,4 +11,12 @@ COGNITO_CLIENT_SECRET=your_cognito_client_secret_here
 COGNITO_USER_POOL_ID=your_cognito_user_pool_id_here
 
 # AWS Region for Cognito
-AWS_REGION=us-east-1
+AWS_REGION=us-east-1
+
+# Cognito Domain (without https:// prefix, just the domain name)
+# Example: mcp-gateway or your-custom-domain
+# COGNITO_DOMAIN=
+
+# Secret key for session cookie signing (must match registry SECRET_KEY), string of hex characters
+# To generate: python -c 'import secrets; print(secrets.token_hex(32))'
+SECRET_KEY=your-secret-key-here

agents/agent_w_auth.py

@@ -77,9 +77,58 @@
 # Get logger
 logger = logging.getLogger(__name__)
 
-def load_env_config() -> Dict[str, Optional[str]]:
+def get_auth_mode_from_args() -> bool:
     """
-    Load configuration from .env file if available.
+    Parse command line arguments to determine authentication mode.
+    This is done before loading environment variables to choose the correct .env file.
+    
+    Returns:
+        bool: True if using session cookie authentication, False for M2M authentication
+    """
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument('--use-session-cookie', action='store_true',
+                        help='Use session cookie authentication instead of M2M')
+    args, _ = parser.parse_known_args()
+    return args.use_session_cookie
+
+def print_env_file_banner(env_file_name: str, use_session_cookie: bool, file_found: bool, file_path: str = None):
+    """
+    Print a prominent banner showing which .env file is being used and why.
+    
+    Args:
+        env_file_name: Name of the .env file being used
+        use_session_cookie: Whether session cookie authentication is being used
+        file_found: Whether the .env file was found
+        file_path: Full path to the .env file if found
+    """
+    print("\n" + "="*80)
+    print("🔧 ENVIRONMENT CONFIGURATION")
+    print("="*80)
+    
+    auth_mode = "Session Cookie Authentication" if use_session_cookie else "M2M Authentication"
+    print(f"Authentication Mode: {auth_mode}")
+    print(f"Expected .env file: {env_file_name}")
+    
+    if use_session_cookie:
+        print("Reason: --use-session-cookie flag specified, using .env.user for user credentials")
+    else:
+        print("Reason: M2M authentication (default), using .env.agent for machine credentials")
+    
+    if file_found and file_path:
+        print(f"✅ Found and loaded: {file_path}")
+    else:
+        print(f"⚠️  File not found: {env_file_name}")
+        print("   Falling back to system environment variables")
+    
+    print("="*80 + "\n")
+
+def load_env_config(use_session_cookie: bool) -> Dict[str, Optional[str]]:
+    """
+    Load configuration from .env file based on authentication mode.
+    Uses .env.user for session cookie auth, .env.agent for M2M auth.
+    
+    Args:
+        use_session_cookie: True for session cookie auth (.env.user), False for M2M auth (.env.agent)
     
     Returns:
         Dict[str, Optional[str]]: Dictionary containing environment variables
@@ -92,29 +141,53 @@ def load_env_config() -> Dict[str, Optional[str]]:
         'domain': None
     }
 
+    # Choose .env file based on authentication mode
+    env_file_name = '.env.user' if use_session_cookie else '.env.agent'
+    
     if DOTENV_AVAILABLE:
+        file_found = False
+        file_path = None
+        
         # Try to load from .env file in the current directory
-        env_file = os.path.join(os.path.dirname(__file__), '.env')
+        env_file = os.path.join(os.path.dirname(__file__), env_file_name)
         if os.path.exists(env_file):
             load_dotenv(env_file)
+            file_found = True
+            file_path = env_file
             logger.info(f"Loading environment variables from {env_file}")
         else:
             # Try to load from .env file in the parent directory
-            env_file = os.path.join(os.path.dirname(__file__), '..', '.env')
+            env_file = os.path.join(os.path.dirname(__file__), '..', env_file_name)
             if os.path.exists(env_file):
                 load_dotenv(env_file)
+                file_found = True
+                file_path = env_file
                 logger.info(f"Loading environment variables from {env_file}")
             else:
                 # Try to load from current working directory
-                load_dotenv()
-                logger.info("Loading environment variables from current directory")
+                env_file = os.path.join(os.getcwd(), env_file_name)
+                if os.path.exists(env_file):
+                    load_dotenv(env_file)
+                    file_found = True
+                    file_path = env_file
+                    logger.info(f"Loading environment variables from {env_file}")
+                else:
+                    # Fallback to default .env loading
+                    load_dotenv()
+                    logger.info("Loading environment variables from default .env file")
+        
+        # Print banner showing which file is being used
+        print_env_file_banner(env_file_name, use_session_cookie, file_found, file_path)
 
         # Get values from environment
         env_config['client_id'] = os.getenv('COGNITO_CLIENT_ID')
         env_config['client_secret'] = os.getenv('COGNITO_CLIENT_SECRET')
         env_config['region'] = os.getenv('AWS_REGION')
         env_config['user_pool_id'] = os.getenv('COGNITO_USER_POOL_ID')
         env_config['domain'] = os.getenv('COGNITO_DOMAIN')
+    else:
+        # Print banner even when dotenv is not available
+        print_env_file_banner(env_file_name, use_session_cookie, False)
 
     return env_config
 
@@ -126,8 +199,11 @@ def parse_arguments() -> argparse.Namespace:
     Returns:
         argparse.Namespace: The parsed command line arguments
     """
-    # Load environment configuration first
-    env_config = load_env_config()
+    # First, determine authentication mode to choose correct .env file
+    use_session_cookie = get_auth_mode_from_args()
+    
+    # Load environment configuration using the appropriate .env file
+    env_config = load_env_config(use_session_cookie)
 
     parser = argparse.ArgumentParser(description='LangGraph MCP Client with Cognito Authentication')
 
@@ -145,7 +221,7 @@ def parse_arguments() -> argparse.Namespace:
 
     # Authentication method arguments
     parser.add_argument('--use-session-cookie', action='store_true',
-                        help='Use session cookie authentication instead of M2M')
+                        help='Use session cookie authentication instead of M2M (loads .env.user instead of .env.agent)')
     parser.add_argument('--session-cookie-file', type=str, default='~/.mcp/session_cookie',
                         help='Path to session cookie file (default: ~/.mcp/session_cookie)')
 

auth_server/cli_auth.py renamed to agents/cli_user_auth.py

@@ -28,6 +28,7 @@
 from http.server import HTTPServer, BaseHTTPRequestHandler
 from itsdangerous import URLSafeTimedSerializer
 import requests
+from dotenv import load_dotenv
 
 # Configure logging
 logging.basicConfig(
@@ -36,7 +37,18 @@
 )
 logger = logging.getLogger(__name__)
 
+# Load environment variables from .env file
+# Look for .env file in the same directory as this script
+script_dir = Path(__file__).parent
+env_file = script_dir / '.env.user'
+if env_file.exists():
+    load_dotenv(env_file)
+    logger.info(f"Loaded environment variables from {env_file}")
+else:
+    logger.warning(f"No .env file found at {env_file}")
+
 # Configuration from environment
+COGNITO_USER_POOL_ID = os.environ.get('COGNITO_USER_POOL_ID')
 COGNITO_DOMAIN = os.environ.get('COGNITO_DOMAIN')
 COGNITO_CLIENT_ID = os.environ.get('COGNITO_CLIENT_ID')
 COGNITO_CLIENT_SECRET = os.environ.get('COGNITO_CLIENT_SECRET')
@@ -45,18 +57,19 @@
 AWS_REGION = os.environ.get('AWS_REGION', 'us-east-1')
 
 # Validate required environment variables
-if not all([COGNITO_DOMAIN, COGNITO_CLIENT_ID, SECRET_KEY]):
+if not all([COGNITO_USER_POOL_ID, COGNITO_CLIENT_ID, SECRET_KEY]):
     logger.error("Missing required environment variables")
-    logger.error("Required: COGNITO_DOMAIN, COGNITO_CLIENT_ID, SECRET_KEY")
+    logger.error("Required: COGNITO_USER_POOL_ID, COGNITO_CLIENT_ID, SECRET_KEY")
     sys.exit(1)
 
-# Build full domain URL if needed
-if not COGNITO_DOMAIN.startswith('https://'):
-    # If just the domain prefix is provided, build the full URL
+# Construct the Cognito domain
+if COGNITO_DOMAIN:
+    # Use custom domain if provided
     COGNITO_DOMAIN_URL = f"https://{COGNITO_DOMAIN}.auth.{AWS_REGION}.amazoncognito.com"
 else:
-    # If full URL is provided, use as-is
-    COGNITO_DOMAIN_URL = COGNITO_DOMAIN
+    # Otherwise use user pool ID without underscores (standard format)
+    user_pool_id_wo_underscore = COGNITO_USER_POOL_ID.replace('_', '')
+    COGNITO_DOMAIN_URL = f"https://{user_pool_id_wo_underscore}.auth.{AWS_REGION}.amazoncognito.com"
 
 # OAuth endpoints
 AUTHORIZE_URL = f"{COGNITO_DOMAIN_URL}/oauth2/authorize"

auth_server/client.py renamed to agents/client.py

@@ -13,5 +13,16 @@ COGNITO_CLIENT_SECRET=your_cognito_client_secret_here
 # Cognito User Pool ID
 COGNITO_USER_POOL_ID=your_cognito_user_pool_id_here
 
+# Cognito Domain (without https:// prefix, just the domain name)
+# Example: mcp-gateway or your-custom-domain
+COGNITO_DOMAIN=your_cognito_domain_here
+
+# Secret key for session cookie signing (must match registry SECRET_KEY)
+# Generate a secure random string for production use
+# To generate: python -c 'import secrets; print(secrets.token_hex(32))'
+SECRET_KEY=your-secret-key-here
+
+USER_POOL_ID=your_cognito_user_pool_id_here
+
 # AWS Region for Cognito
 AWS_REGION=us-east-1

auth_server/.env.template

@@ -21,7 +21,8 @@ dependencies = [
     "pyjwt>=2.6.0",
     "cryptography>=40.0.0",
     "pyyaml>=6.0.0",
-    "itsdangerous>=2.0.0"
+    "httpx>=0.25.0",
+    "itsdangerous>=2.1.0"
 ]
 
 [project.optional-dependencies]

auth_server/pyproject.toml

@@ -46,9 +46,7 @@ def load_scopes_config():
 
 def map_cognito_groups_to_scopes(groups: List[str]) -> List[str]:
     """
-    Map Cognito groups to MCP scopes using scopes.yml configuration.
-    
-    Uses the same scope format as M2M tokens for consistency.
+    Map Cognito groups to MCP scopes using the same format as M2M resource server scopes.
     
     Args:
         groups: List of Cognito group names
@@ -57,28 +55,36 @@ def map_cognito_groups_to_scopes(groups: List[str]) -> List[str]:
         List of MCP scopes in resource server format
     """
     scopes = []
-    
-    # Use group mappings from scopes.yml if available
-    group_mappings = SCOPES_CONFIG.get('group_mappings', {})
+    logger.info(f"Mapping Cognito groups to MCP scopes: {groups}")
 
     for group in groups:
-        if group in group_mappings:
-            # Use configured mapping
-            scopes.extend(group_mappings[group])
-        else:
-            # Legacy fallback for backward compatibility
-            logger.debug(f"Group '{group}' not in group_mappings, using legacy mapping")
-            
-            if group == 'mcp-admin':
-                scopes.extend(['mcp-servers-unrestricted/read', 
-                              'mcp-servers-unrestricted/execute'])
-            elif group == 'mcp-user':
-                scopes.append('mcp-servers-restricted/read')
-            elif group.startswith('mcp-server-'):
-                scopes.append('mcp-servers-restricted/execute')
+        if group == 'mcp-admin':
+            # Admin gets unrestricted read and execute access
+            scopes.append('mcp-servers-unrestricted/read')
+            scopes.append('mcp-servers-unrestricted/execute')
+        elif group == 'mcp-user':
+            # Regular users get restricted read access by default
+            scopes.append('mcp-servers-restricted/read')
+        elif group.startswith('mcp-server-'):
+            # Server-specific groups grant access based on server permissions
+            # For now, grant restricted execute access for specific servers
+            # This allows access to the servers defined in the restricted scope
+            scopes.append('mcp-servers-restricted/execute')
+            
+            # Note: The actual server access control is handled by the
+            # validate_server_tool_access function which checks the scopes.yml
+            # configuration. The group names are preserved in the 'groups' field
+            # for potential future fine-grained access control.
 
     # Remove duplicates while preserving order
-    return list(dict.fromkeys(scopes))
+    seen = set()
+    unique_scopes = []
+    for scope in scopes:
+        if scope not in seen:
+            seen.add(scope)
+            unique_scopes.append(scope)
+    
+    return unique_scopes
 
 def validate_session_cookie(cookie_value: str) -> Dict[str, any]:
     """
@@ -646,6 +652,7 @@ async def validate_request(request: Request):
             if cookie_value:
                 try:
                     validation_result = validate_session_cookie(cookie_value)
+                    logger.info(f"Session cookie validation result: {validation_result}")
                     logger.info(f"Session cookie validation successful for user: {validation_result['username']}")
                 except ValueError as e:
                     logger.warning(f"Session cookie validation failed: {e}")