use pydantic settings

Author: thorinaboenke

src/attackmate_api_server/auth_utils.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import secrets
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, Optional
@@ -8,11 +7,11 @@
 from fastapi import Depends, HTTPException, status
 from fastapi.security import APIKeyHeader
 from passlib.context import CryptContext
+from attackmate_api_server.config import settings
 
 load_dotenv()
 
 
-TOKEN_EXPIRE_MINUTES = int(os.getenv('TOKEN_EXPIRE_MINUTES', 30))
 API_KEY_HEADER_NAME = 'X-Auth-Token'
 api_key_header_scheme = APIKeyHeader(name=API_KEY_HEADER_NAME, auto_error=True)
 pwd_context = CryptContext(schemes=['argon2'], deprecated='auto')
@@ -31,15 +30,15 @@ def verify_password(plain_password: str, hashed_password: str) -> bool:
 
 
 def get_user_hash(username: str) -> Optional[str]:
-    """Fetches the hashed password from environment variables."""
-    env_var_name = f'USER_{username.upper()}_HASH'
-    return os.getenv(env_var_name)
+    """Fetches the hashed password from the centralized settings."""
+    logger.info(f"user hashses from settings: {settings}")
+    return settings.users.get(username.lower())
 
 
 def create_access_token(username: str) -> str:
     """Creates a new token, stores it, and returns the token string."""
     token = secrets.token_urlsafe(32)
-    expires = datetime.now(timezone.utc) + timedelta(minutes=TOKEN_EXPIRE_MINUTES)
+    expires = datetime.now(timezone.utc) + timedelta(minutes=settings.token_expire_minutes)
     # TODO locking needed for multi-threaded access, smth like with token_lock ?
     ACTIVE_TOKENS[token] = {'username': username, 'expires': expires}
     logger.info(f"Created new token for user '{username}' expiring at {expires}")
@@ -53,7 +52,7 @@ def renew_token_expiry(token: str) -> bool:
     """
     token_data = ACTIVE_TOKENS.get(token)
     if token_data:
-        token_data['expires'] = datetime.now(timezone.utc) + timedelta(minutes=TOKEN_EXPIRE_MINUTES)
+        token_data['expires'] = datetime.now(timezone.utc) + timedelta(minutes=settings.token_expire_minutes)
         logger.debug(f"Renewed token expiry for user '{token_data['username']}'")
         return True
     return False

src/attackmate_api_server/config.py

@@ -0,0 +1,29 @@
+
+from typing import Dict, Optional
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic import Field
+
+
+class Settings(BaseSettings):
+    # App Settings
+    # Environment Var: ATTACKMATE_CONFIG_PATH
+    attackmate_config_path: Optional[str] = Field(default=None, alias="ATTACKMATE_CONFIG_PATH")
+    token_expire_minutes: int = 30
+
+    # SSL Settings
+    ssl_key_file: str = "key.pem"
+    ssl_cert_file: str = "cert.pem"
+
+    # User Hashes: Key is username (lowercase), Value is argon2 hash
+    # Environment Var: users='{"user1": "hash1"}'
+    users: Dict[str, str] = {}
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        # Allows extra env vars without crashing
+        extra="ignore"
+    )
+
+
+settings = Settings()

src/attackmate_api_server/create_hashes.py

@@ -1,17 +1,18 @@
-
 from passlib.context import CryptContext
+import json
 
 pwd_context = CryptContext(schemes=['argon2'], deprecated='auto')
 
-
+# Define your users here: { 'username': 'plain_text_password' }
 users = {
     'testuser': 'testuser',
+    'admin': 'change_me_immediately',
+    'analyst': 'secure_password_123'
 }
 
-env_content = ''
-print('\nCopy the following lines into your .env file:\n')
-for username, plain_password in users.items():
-    hashed_password = pwd_context.hash(plain_password)
-    env_line = f"USER_{username.upper()}_HASH=\"{hashed_password}\""
-    print(env_line)
-    env_content += env_line + '\n'
+hashed_users = {username: pwd_context.hash(password) for username, password in users.items()}
+
+print("\n--- Copy the line below into your .env file ---")
+print("# AttackMate API User Credentials")
+print(f"USERS='{json.dumps(hashed_users)}'")
+print("\n--- End of Environment Variables ---")

src/attackmate_api_server/main.py

@@ -13,13 +13,13 @@
 
 import attackmate_api_server.state as state
 from attackmate_api_server.routers import commands, instances, playbooks
-
+from attackmate_api_server.config import settings
 from attackmate_api_server.auth_utils import create_access_token, get_user_hash, verify_password
 from attackmate_api_server.schemas import TokenResponse
 
 CERT_DIR = os.path.dirname(os.path.abspath(__file__))
-KEY_FILE = os.path.join(CERT_DIR, 'key.pem')
-CERT_FILE = os.path.join(CERT_DIR, 'cert.pem')
+KEY_FILE = os.path.join(CERT_DIR, settings.ssl_key_file)
+CERT_FILE = os.path.join(CERT_DIR, settings.ssl_cert_file)
 
 # Initialize loggers
 logger = initialize_api_logger(debug=True, append_logs=False)
@@ -31,7 +31,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     logger.info('AttackMate API starting up (lifespan)...')
     try:
         # Load global config on startup and assign to the variable in state.py
-        config_path = os.getenv('ATTACKMATE_CONFIG_PATH')
+        config_path = settings.attackmate_config_path
         if config_path:
             loaded_config = parse_config(config_file=config_path, logger=logger)
             logger.info(f'Loading AttackMate configuration from specified path: {config_path}')