toru-vault 0.1.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

toru_vault/vault.py

@@ -1,28 +1,18 @@
 #!/usr/bin/env python3
 import os
 import logging
-import time
-import json
-import tempfile
 import stat
-import atexit
-import secrets as pysecrets
-from typing import Dict, Optional, Tuple
 from bitwarden_sdk import BitwardenClient, DeviceType, client_settings_from_dict
-from cryptography.fernet import Fernet
-from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
-import base64
-from .lazy_dict import LazySecretsDict
 
-# Try importing keyring - it might not be available in container environments
-try:
-    import keyring
-    _KEYRING_AVAILABLE = True
-except ImportError:
-    _KEYRING_AVAILABLE = False
+from .in_memory import (
+    _KEYRING_AVAILABLE,
+    create_secrets_dict
+)
+from .in_env import (
+    set_env_vars,
+    load_secrets_env
+)
 
-# Setup minimal logging
 logger = logging.getLogger(__name__)
 
 # Constants for keyring storage
@@ -30,164 +20,7 @@ _KEYRING_SERVICE_NAME = "bitwarden_vault"
 _KEYRING_BWS_TOKEN_KEY = "bws_token"
 _KEYRING_ORG_ID_KEY = "organization_id"
 _KEYRING_STATE_FILE_KEY = "state_file"
-
-# Secure cache configuration
-_SECRET_CACHE_TIMEOUT = 300  # 5 minutes
-_secrets_cache: Dict[str, Tuple[float, Dict[str, str]]] = {}
-
-def _generate_encryption_key(salt: bytes = None) -> Tuple[bytes, bytes]:
-    """
-    Generate an encryption key for securing the cache
-    
-    Args:
-        salt (bytes, optional): Salt for key derivation
-        
-    Returns:
-        Tuple[bytes, bytes]: Key and salt
-    """
-    if salt is None:
-        salt = os.urandom(16)
-    
-    # Generate a key from the machine-specific information and random salt
-    machine_id = _get_machine_id()
-    password = machine_id.encode()
-    
-    kdf = PBKDF2HMAC(
-        algorithm=hashes.SHA256(),
-        length=32,
-        salt=salt,
-        iterations=100000,
-    )
-    key = base64.urlsafe_b64encode(kdf.derive(password))
-    return key, salt
-
-def _get_machine_id() -> str:
-    """Get a unique identifier for the current machine"""
-    # Try platform-specific methods to get a machine ID
-    machine_id = ""
-    
-    if os.path.exists('/etc/machine-id'):
-        with open('/etc/machine-id', 'r') as f:
-            machine_id = f.read().strip()
-    elif os.path.exists('/var/lib/dbus/machine-id'):
-        with open('/var/lib/dbus/machine-id', 'r') as f:
-            machine_id = f.read().strip()
-    elif os.name == 'nt':  # Windows
-        import subprocess
-        try:
-            result = subprocess.run(['wmic', 'csproduct', 'get', 'UUID'], capture_output=True, text=True)
-            if result.returncode == 0:
-                machine_id = result.stdout.strip().split('\n')[-1].strip()
-        except (FileNotFoundError, subprocess.SubprocessError):
-            pass
-    
-    # Fallback if we couldn't get a machine ID
-    if not machine_id:
-        # Use a combination of hostname and a persisted random value
-        import socket
-        hostname = socket.gethostname()
-        
-        # Create a persistent random ID
-        id_file = os.path.join(tempfile.gettempdir(), '.vault_machine_id')
-        if os.path.exists(id_file):
-            try:
-                with open(id_file, 'r') as f:
-                    random_id = f.read().strip()
-            except Exception:
-                random_id = pysecrets.token_hex(16)
-        else:
-            random_id = pysecrets.token_hex(16)
-            try:
-                # Try to save it with restricted permissions
-                with open(id_file, 'w') as f:
-                    f.write(random_id)
-                os.chmod(id_file, stat.S_IRUSR | stat.S_IWUSR)  # 0600 permissions
-            except Exception:
-                pass
-                
-        machine_id = f"{hostname}-{random_id}"
-    
-    return machine_id
-
-def _encrypt_secrets(secrets_dict: Dict[str, str]) -> Optional[str]:
-    """
-    Encrypt secrets dictionary
-    
-    Args:
-        secrets_dict (Dict[str, str]): Dictionary of secrets
-        
-    Returns:
-        Optional[str]: Encrypted data or None if encryption fails
-    """
-    try:
-        key, salt = _generate_encryption_key()
-        if not key:
-            return None
-            
-        # Encrypt the serialized secrets
-        f = Fernet(key)
-        encrypted_data = f.encrypt(json.dumps(secrets_dict).encode())
-        
-        # Store along with the salt
-        return base64.urlsafe_b64encode(salt).decode() + ":" + encrypted_data.decode()
-    except Exception as e:
-        logger.warning(f"Failed to encrypt secrets: {e}")
-        return None
-
-def _decrypt_secrets(encrypted_data: str) -> Optional[Dict[str, str]]:
-    """
-    Decrypt secrets
-    
-    Args:
-        encrypted_data (str): Encrypted data
-        
-    Returns:
-        Optional[Dict[str, str]]: Decrypted secrets dictionary or None if decryption fails
-    """
-    try:
-        # Split salt and encrypted data
-        salt_b64, encrypted = encrypted_data.split(":", 1)
-        salt = base64.urlsafe_b64decode(salt_b64)
-        
-        # Regenerate the key with the same salt
-        key, _ = _generate_encryption_key(salt)
-        if not key:
-            return None
-            
-        # Decrypt the data
-        f = Fernet(key)
-        decrypted_data = f.decrypt(encrypted.encode())
-        
-        return json.loads(decrypted_data.decode())
-    except Exception as e:
-        logger.warning(f"Failed to decrypt secrets: {e}")
-        return None
-
-def _secure_state_file(state_path: str) -> None:
-    """
-    Ensure the state file has secure permissions
-    
-    Args:
-        state_path (str): Path to the state file
-    """
-    try:
-        if os.path.exists(state_path):
-            if os.name == 'posix':  # Linux/Mac
-                os.chmod(state_path, stat.S_IRUSR | stat.S_IWUSR)  # 0600 permissions
-            elif os.name == 'nt':  # Windows
-                import subprocess
-                subprocess.run(['icacls', state_path, '/inheritance:r', '/grant:r', f'{os.getlogin()}:(F)'], 
-                               capture_output=True)
-    except Exception as e:
-        logger.warning(f"Could not set secure permissions on state file: {e}")
-
-def _clear_cache() -> None:
-    """Clear the secrets cache on exit"""
-    global _secrets_cache
-    _secrets_cache = {}
-    
-# Register the cache clearing function to run on exit
-atexit.register(_clear_cache)
+_KEYRING_PROJECT_ID_KEY = "project_id"
 
 def _get_from_keyring_or_env(key, env_var):
     """
@@ -205,56 +38,59 @@ def _get_from_keyring_or_env(key, env_var):
     # Try keyring first if available
     if _KEYRING_AVAILABLE:
         try:
+            import keyring
             value = keyring.get_password(_KEYRING_SERVICE_NAME, key)
         except Exception as e:
             logger.warning(f"Failed to get {key} from keyring: {e}")
     
-    # Fall back to environment variable
     if not value:
         value = os.getenv(env_var)
     
     return value
 
-def set_to_keyring(key, value):
+
+
+def _secure_state_file(state_path: str) -> None:
     """
-    Set a value to keyring
+    Ensure the state file has secure permissions
     
     Args:
-        key (str): Key in keyring
-        value (str): Value to store
-    
-    Returns:
-        bool: True if successful, False otherwise
+        state_path (str): Path to the state file
     """
-    if not _KEYRING_AVAILABLE:
-        return False
-    
     try:
-        keyring.set_password(_KEYRING_SERVICE_NAME, key, value)
-        return True
+        if os.path.exists(state_path):
+            if os.name == 'posix':  # Linux/Mac
+                os.chmod(state_path, stat.S_IRUSR | stat.S_IWUSR)  # 0600 permissions
+            elif os.name == 'nt':  # Windows
+                import subprocess
+                # /inheritance:r - Removes all inherited ACEs (Access Control Entries).
+                # /grant:r - Grants specified user rights, replacing any previous explicit ACEs for that user.
+                # <os.getlogin()>:(F) - Grants the current user (F)ull control.
+                result = subprocess.run(['icacls', state_path, '/inheritance:r', '/grant:r', f'{os.getlogin()}:(F)'], 
+                               capture_output=True)
+                if result.returncode != 0:
+                    raise Exception(f"Could not set secure permissions on state file: {result.stderr.decode()}")
+        
     except Exception as e:
-        logger.warning(f"Failed to set {key} to keyring: {e}")
-        return False
+        logger.warning(f"Could not set secure permissions on state file: {e}")
 
 def _initialize_client():
     """
     Initialize the Bitwarden client
     """
-    # Get environment variables with defaults
     api_url = os.getenv("API_URL", "https://api.bitwarden.com")
     identity_url = os.getenv("IDENTITY_URL", "https://identity.bitwarden.com")
     
-    # Get BWS_TOKEN from keyring or environment variable
     bws_token = _get_from_keyring_or_env(_KEYRING_BWS_TOKEN_KEY, "BWS_TOKEN")
-    
-    # Get STATE_FILE from keyring or environment variable
     state_path = _get_from_keyring_or_env(_KEYRING_STATE_FILE_KEY, "STATE_FILE")
+    org_id = _get_from_keyring_or_env(_KEYRING_ORG_ID_KEY, "ORGANIZATION_ID")
     
-    # Validate required environment variables
     if not bws_token:
         raise ValueError("BWS_TOKEN not found in keyring or environment variable")
     if not state_path:
         raise ValueError("STATE_FILE not found in keyring or environment variable")
+    if not org_id:
+        raise ValueError("ORGANIZATION_ID not found in keyring or environment variable")
         
     # Ensure state file directory exists
     state_dir = os.path.dirname(state_path)
@@ -267,10 +103,8 @@ def _initialize_client():
         except Exception as e:
             logger.warning(f"Could not create state directory with secure permissions: {e}")
     
-    # Secure the state file
     _secure_state_file(state_path)
     
-    # Create and initialize the client
     client = BitwardenClient(
         client_settings_from_dict({
             "apiUrl": api_url,
@@ -280,102 +114,40 @@ def _initialize_client():
         })
     )
     
-    # Authenticate with the Secrets Manager Access Token
     client.auth().login_access_token(bws_token, state_path)
     
+    del bws_token
+    del org_id
+    
     return client
 
-def _load_secrets(project_id=None):
+def env_load(project_id=None, override=False):
     """
-    Load secrets from Bitwarden
+    Load all secrets related to the project into environmental variables.
     
     Args:
-        project_id (str): Project ID to filter secrets
-    
-    Returns:
-        dict: Dictionary of secrets with their names as keys
+        project_id (str, optional): Project ID to filter secrets. If None, will try to get from keyring or PROJECT_ID environment variable
+        override (bool, optional): Whether to override existing environment variables
     """
-    # Initialize client with credentials from environment or keyring
     try:
         client = _initialize_client()
     except Exception as e:
         logger.error(f"Failed to initialize Bitwarden client: {e}")
-        return {}
-    
-    # Get ORGANIZATION_ID from keyring or environment variable
+        return
     organization_id = _get_from_keyring_or_env(_KEYRING_ORG_ID_KEY, "ORGANIZATION_ID")
     if not organization_id:
         logger.error("ORGANIZATION_ID not found in keyring or environment variable")
-        return {}
-    
-    # Get secrets from BitWarden
-    try:
-        # Sync secrets to ensure we have the latest
-        client.secrets().sync(organization_id, None)
-        
-        # Initialize empty secrets dictionary
-        secrets = {}
-        
-        # Retrieve all secrets
-        all_secrets = client.secrets().list(organization_id)
-        
-        # Validate response format
-        if not hasattr(all_secrets, 'data') or not hasattr(all_secrets.data, 'data'):
-            return {}
-        
-        # We need to collect all secret IDs first
-        secret_ids = []
-        for secret in all_secrets.data.data:
-            secret_ids.append(secret.id)
-        
-        # If we have secret IDs, fetch their values
-        if secret_ids:
-            # Get detailed information for all secrets by their IDs
-            secrets_detailed = client.secrets().get_by_ids(secret_ids)
-            
-            # Validate response format
-            if not hasattr(secrets_detailed, 'data') or not hasattr(secrets_detailed.data, 'data'):
-                return {}
-            
-            # Process each secret
-            for secret in secrets_detailed.data.data:
-                # Extract the project ID
-                secret_project_id = getattr(secret, 'project_id', None)
-                
-                # Check if this secret belongs to the specified project
-                if project_id and secret_project_id is not None and project_id != str(secret_project_id):
-                    continue
-                
-                # Add the secret to our dictionary
-                secrets[secret.key] = secret.value
-        
-        # Update the cache with encryption
-        encrypted_data = _encrypt_secrets(secrets)
-        if encrypted_data:
-            _secrets_cache[f"{organization_id}:{project_id or ''}"] = (time.time(), encrypted_data)
-        else:
-            _secrets_cache[f"{organization_id}:{project_id or ''}"] = (time.time(), secrets.copy())
-        
-        return secrets
-    except Exception as e:
-        logger.error(f"Error loading secrets: {e}")
-        raise
-
-def env_load(project_id=None, override=False):
-    """
-    Load all secrets related to the project into environmental variables.
+        return
     
-    Args:
-        project_id (str, optional): Project ID to filter secrets
-        override (bool, optional): Whether to override existing environment variables
-    """
-    # Get all secrets from BWS
-    secrets = _load_secrets(project_id)
+    # If project_id is not provided, try to get it from keyring or environment variable
+    if project_id is None:
+        project_id = _get_from_keyring_or_env(_KEYRING_PROJECT_ID_KEY, "PROJECT_ID")
     
-    # Set environment variables
-    for key, value in secrets.items():
-        if override or key not in os.environ:
-            os.environ[key] = value
+    secrets = load_secrets_env(client, organization_id, project_id)
+
+    set_env_vars(secrets, override)
+
+    del secrets
 
 def env_load_all(override=False):
     """
@@ -384,192 +156,106 @@ def env_load_all(override=False):
     Args:
         override (bool, optional): Whether to override existing environment variables
     """
-    # Get ORGANIZATION_ID from keyring or environment variable
     organization_id = _get_from_keyring_or_env(_KEYRING_ORG_ID_KEY, "ORGANIZATION_ID")
+    if not organization_id:
+        logger.error("ORGANIZATION_ID not found in keyring or environment variable")
+        return
     
-    # Initialize Bitwarden client
     try:
         client = _initialize_client()
+        from .in_env import load_secrets_env_all
+        secrets = load_secrets_env_all(client, organization_id)
+        set_env_vars(secrets, override)
+        del secrets
     except Exception as e:
-        logger.error(f"Failed to initialize Bitwarden client: {e}")
+        logger.error(f"Failed to load all secrets: {e}")
         return
-    
-    try:
-        # Sync to ensure we have the latest data
-        client.secrets().sync(organization_id, None)
-        
-        # Get all projects
-        projects_response = client.projects().list(organization_id)
-        
-        # Validate response format
-        if not hasattr(projects_response, 'data') or not hasattr(projects_response.data, 'data'):
-            logger.warning(f"No projects found in organization {organization_id}")
-            return
-        
-        # Process each project
-        for project in projects_response.data.data:
-            if hasattr(project, 'id'):
-                project_id = str(project.id)
-                
-                # Load environment variables for this project
-                try:
-                    # Get the secrets for this project and set them as environment variables
-                    env_load(project_id=project_id, override=override)
-                    logger.info(f"Loaded secrets from project: {getattr(project, 'name', project_id)}")
-                except Exception as e:
-                    logger.warning(f"Failed to load secrets from project {project_id}: {e}")
-                
-    except Exception as e:
-        logger.error(f"Failed to load all secrets into environment variables: {e}")
 
-def get(project_id=None, refresh=False, use_keyring=True):
+def get(project_id=None, use_keyring=True):
     """
-    Return a dictionary of all project secrets
+    Return a dictionary of all project secrets with JIT decryption
     
     Args:
-        project_id (str, optional): Project ID to filter secrets
-        refresh (bool, optional): Force refresh the secrets cache
+        project_id (str, optional): Project ID to filter secrets. If None, will try to get from keyring or PROJECT_ID environment variable
         use_keyring (bool, optional): Whether to use system keyring (True) or in-memory encryption (False)
         
     Returns:
-        dict: Dictionary of secrets with their names as keys, using lazy loading
+        dict: Dictionary of secrets with their names as keys, using lazy loading with JIT decryption
     """
-    # Get ORGANIZATION_ID from keyring or environment variable
+    try:
+        client = _initialize_client()
+    except Exception as e:
+        logger.error(f"Failed to initialize Bitwarden client: {e}")
+        return {}
+        
     organization_id = _get_from_keyring_or_env(_KEYRING_ORG_ID_KEY, "ORGANIZATION_ID")
+    if not organization_id:
+        logger.error("ORGANIZATION_ID not found in keyring or environment variable")
+        return {}
     
-    # Build the service name for keyring storage
-    service_name = f"vault_{organization_id or 'default'}"
+    # If project_id is not provided, try to get it from keyring or environment variable
+    if project_id is None:
+        project_id = _get_from_keyring_or_env(_KEYRING_PROJECT_ID_KEY, "PROJECT_ID")
     
-    # Function to either fetch from keyring or decrypt from cache based on container flag
-    def _load_decrypted_secrets():
-        # Check if we need to force a refresh
-        if refresh:
-            return _load_secrets(project_id)
-        
-        # Otherwise try to use cached values first
-        cache_key = f"{organization_id}:{project_id or ''}"
-        current_time = time.time()
-        
-        if cache_key in _secrets_cache:
-            timestamp, encrypted_secrets = _secrets_cache[cache_key]
-            
-            # If cache hasn't expired
-            if current_time - timestamp < _SECRET_CACHE_TIMEOUT:
-                # If we have encryption, try to decrypt
-                if encrypted_secrets:
-                    decrypted_secrets = _decrypt_secrets(encrypted_secrets)
-                    if decrypted_secrets:
-                        return decrypted_secrets
-                # Otherwise return the unencrypted data (backward compatibility)
-                elif isinstance(encrypted_secrets, dict):
-                    return encrypted_secrets.copy()
-        
-        # If we couldn't get from cache, load fresh
-        return _load_secrets(project_id)
+    from .in_memory import load_secrets_memory
+    all_secrets = load_secrets_memory(client, organization_id, project_id)
     
-    # Get all secrets and their keys
-    all_secrets = _load_decrypted_secrets()
+    # Get all secret keys - values will be decrypted JIT when accessed
     secret_keys = set(all_secrets.keys())
     
-    # Store secrets in keyring if available and requested
-    keyring_usable = _KEYRING_AVAILABLE and use_keyring
-    if keyring_usable:
-        for key, value in all_secrets.items():
-            keyring.set_password(service_name, key, value)
-    
-    # When keyring is unavailable or not requested (likely in container)
-    if not keyring_usable:
-        # Create a dictionary of cached secrets for container mode
-        container_secrets = {}
-        encrypted_data = None
-        cache_key = f"{organization_id}:{project_id or ''}"
-        
-        # If we have a cached encrypted version, use that
-        if cache_key in _secrets_cache:
-            _, encrypted_data = _secrets_cache[cache_key]
-            
-        # Create getter function for container mode
-        def _container_getter(key):
-            if key in container_secrets:
-                return container_secrets[key]
-            
-            # If not in memory cache, check if we have pre-loaded decrypted secrets
-            if all_secrets and key in all_secrets:
-                container_secrets[key] = all_secrets[key]
-                return container_secrets[key]
-                
-            # Otherwise, try to decrypt from cache
-            if encrypted_data and not isinstance(encrypted_data, dict):
-                decrypted = _decrypt_secrets(encrypted_data)
-                if decrypted and key in decrypted:
-                    container_secrets[key] = decrypted[key]
-                    return container_secrets[key]
-                
-            # If all else fails, load from API
-            fresh_secrets = _load_secrets(project_id)
-            if key in fresh_secrets:
-                container_secrets[key] = fresh_secrets[key]
-                return container_secrets[key]
-                
-            return None
-        
-        # Create the lazy dictionary with container getter
-        return LazySecretsDict(secret_keys, _container_getter)
-    else:
-        # Create getter function for keyring mode
-        def _keyring_getter(key):
-            return keyring.get_password(service_name, key)
-            
-        # Create setter function for keyring mode
-        def _keyring_setter(key, value):
-            keyring.set_password(service_name, key, value)
-            
-        # Create deleter function for keyring mode
-        def _keyring_deleter(key):
-            keyring.delete_password(service_name, key)
-        
-        # Create the lazy dictionary with keyring getter/setter/deleter
-        return LazySecretsDict(secret_keys, _keyring_getter, _keyring_setter, _keyring_deleter)
-
+    return create_secrets_dict(
+        secret_keys, 
+        organization_id, 
+        project_id or "", 
+        all_secrets,
+        use_keyring
+    )
 
-def get_all(refresh=False, use_keyring=True):
+def get_all(use_keyring=True):
     """
-    Return a combined dictionary of secrets from all projects that user has access to
+    Return a combined dictionary of secrets from all projects that user has access to with JIT decryption
     
     Args:
-        refresh (bool, optional): Force refresh the secrets cache
         use_keyring (bool, optional): Whether to use system keyring (True) or in-memory encryption (False)
         
     Returns:
-        dict: Dictionary of secrets with their names as keys, using lazy loading
+        dict: Dictionary of secrets with their names as keys, using lazy loading with JIT decryption
     """
-    # Initialize the client
-    client = _initialize_client()
-    if not client:
-        logging.error("Failed to initialize Bitwarden client")
+    try:
+        client = _initialize_client()
+    except Exception as e:
+        logger.error(f"Failed to initialize Bitwarden client: {e}")
         return {}
-
-    # Get organization ID
-    organization_id = _get_from_keyring_or_env("ORGANIZATION_ID", "ORGANIZATION_ID")
+    organization_id = _get_from_keyring_or_env(_KEYRING_ORG_ID_KEY, "ORGANIZATION_ID")
     if not organization_id:
-        logging.error("Organization ID not found in keyring or environment variables")
+        logger.error("Organization ID not found in keyring or environment variables")
         return {}
 
     try:
-        # Get all projects
-        projects = client.projects_api.get_projects(organization_id)
-        if not projects or not projects.data:
-            logging.warning("No projects found in the organization")
+        projects_response = client.projects().list(organization_id)
+        
+        # Validate response format
+        if not hasattr(projects_response, 'data') or not hasattr(projects_response.data, 'data'):
+            logger.warning(f"No projects found in organization {organization_id}")
             return {}
 
         # Create a combined dictionary with all secrets
         all_secrets = {}
-        for project in projects.data:
-            project_secrets = get(project.id, refresh=refresh, use_keyring=use_keyring)
-            # Update the combined dictionary (note: this will overwrite duplicate keys)
+        project_ids = []
+        
+        # First collect all project IDs
+        for project in projects_response.data.data:
+            if hasattr(project, 'id'):
+                project_ids.append(str(project.id))
+        
+        # Create merged dictionary of all secrets with JIT decryption
+        for project_id in project_ids:
+            # Get secrets for this project
+            project_secrets = get(project_id, use_keyring=use_keyring)
+            # Update the combined dictionary (this will overwrite duplicate keys)
             all_secrets.update(project_secrets)
 
         return all_secrets
     except Exception as e:
-        logging.error(f"Error retrieving projects: {str(e)}")
+        logger.error(f"Error retrieving projects: {str(e)}")
         return {}