roster/src/variable_substitution.py

# variable_substitution.py
#
# Copyright 2025 Pavel Baksy
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: GPL-3.0-or-later

import re
from typing import Dict, List, Tuple, Set
from .models import HttpRequest, Environment


class VariableSubstitution:
    """Handles variable substitution with {{variable_name}} syntax."""

    # Pattern to match {{variable_name}} where variable_name is alphanumeric + underscore
    VARIABLE_PATTERN = re.compile(r'\{\{(\w+)\}\}')

    @staticmethod
    def find_variables(text: str) -> List[str]:
        """
        Extract all {{variable_name}} references from text.

        Args:
            text: The text to search for variable references

        Returns:
            List of variable names found (without the {{ }} delimiters)
        """
        if not text:
            return []
        return VariableSubstitution.VARIABLE_PATTERN.findall(text)

    @staticmethod
    def substitute(text: str, variables: Dict[str, str]) -> Tuple[str, List[str]]:
        """
        Replace {{variable_name}} with values from variables dict.

        Args:
            text: The text containing variable placeholders
            variables: Dictionary mapping variable names to their values

        Returns:
            Tuple of (substituted_text, list_of_undefined_variables)
        """
        if not text:
            return text, []

        undefined_vars = []

        def replace_var(match):
            var_name = match.group(1)
            if var_name in variables:
                value = variables[var_name]
                # Check if value is empty/None - treat as undefined for warning purposes
                if not value:
                    undefined_vars.append(var_name)
                    return ""
                return value
            else:
                # Variable not defined - track it and replace with empty string
                undefined_vars.append(var_name)
                return ""

        substituted = VariableSubstitution.VARIABLE_PATTERN.sub(replace_var, text)
        return substituted, undefined_vars

    @staticmethod
    def substitute_request(request: HttpRequest, environment: Environment) -> Tuple[HttpRequest, Set[str]]:
        """
        Substitute variables in all request fields (URL, headers, body).

        Args:
            request: The HTTP request with variable placeholders
            environment: The environment containing variable values

        Returns:
            Tuple of (new_request_with_substitutions, set_of_undefined_variables)
        """
        all_undefined = set()

        # Substitute URL
        new_url, url_undefined = VariableSubstitution.substitute(
            request.url, environment.variables
        )
        all_undefined.update(url_undefined)

        # Substitute headers (both keys and values)
        new_headers = {}
        for key, value in request.headers.items():
            new_key, key_undefined = VariableSubstitution.substitute(
                key, environment.variables
            )
            new_value, value_undefined = VariableSubstitution.substitute(
                value, environment.variables
            )
            all_undefined.update(key_undefined)
            all_undefined.update(value_undefined)
            new_headers[new_key] = new_value

        # Substitute body
        new_body, body_undefined = VariableSubstitution.substitute(
            request.body, environment.variables
        )
        all_undefined.update(body_undefined)

        # Create new HttpRequest with substituted values
        new_request = HttpRequest(
            method=request.method,
            url=new_url,
            headers=new_headers,
            body=new_body,
            syntax=request.syntax
        )

        return new_request, all_undefined

    @staticmethod
    def substitute_excluding_variables(text: str, variables: Dict[str, str], excluded_vars: List[str]) -> Tuple[str, List[str]]:
        """
        Replace {{variable_name}} with values, but skip excluded variables (leave as {{var}}).

        Args:
            text: The text containing variable placeholders
            variables: Dictionary mapping variable names to their values
            excluded_vars: List of variable names to skip (leave as placeholders)

        Returns:
            Tuple of (substituted_text, list_of_undefined_variables)
        """
        if not text:
            return text, []

        undefined_vars = []

        def replace_var(match):
            var_name = match.group(1)

            # Skip excluded variables - leave as {{var_name}}
            if var_name in excluded_vars:
                return match.group(0)  # Return original {{var_name}}

            # Substitute non-excluded variables
            if var_name in variables:
                value = variables[var_name]
                # Check if value is empty/None - treat as undefined for warning purposes
                if not value:
                    undefined_vars.append(var_name)
                    return ""
                return value
            else:
                # Variable not defined - track it and replace with empty string
                undefined_vars.append(var_name)
                return ""

        substituted = VariableSubstitution.VARIABLE_PATTERN.sub(replace_var, text)
        return substituted, undefined_vars

    @staticmethod
    def substitute_request_excluding_sensitive(
        request: HttpRequest,
        environment: Environment,
        sensitive_variables: List[str]
    ) -> Tuple[HttpRequest, Set[str], bool]:
        """
        Substitute variables in request, but exclude sensitive variables (leave as {{var}}).

        This is useful for saving requests to history without exposing sensitive values.
        Sensitive variables remain in their {{variable_name}} placeholder form.

        Args:
            request: The HTTP request with variable placeholders
            environment: The environment containing variable values
            sensitive_variables: List of variable names to exclude from substitution

        Returns:
            Tuple of (redacted_request, set_of_undefined_variables, has_sensitive_vars)
            - redacted_request: Request with only non-sensitive variables substituted
            - set_of_undefined_variables: Variables that were referenced but not defined
            - has_sensitive_vars: True if any sensitive variables were found and redacted
        """
        all_undefined = set()

        # Substitute URL (excluding sensitive variables)
        new_url, url_undefined = VariableSubstitution.substitute_excluding_variables(
            request.url, environment.variables, sensitive_variables
        )
        all_undefined.update(url_undefined)

        # Substitute headers (both keys and values, excluding sensitive variables)
        new_headers = {}
        for key, value in request.headers.items():
            new_key, key_undefined = VariableSubstitution.substitute_excluding_variables(
                key, environment.variables, sensitive_variables
            )
            new_value, value_undefined = VariableSubstitution.substitute_excluding_variables(
                value, environment.variables, sensitive_variables
            )
            all_undefined.update(key_undefined)
            all_undefined.update(value_undefined)
            new_headers[new_key] = new_value

        # Substitute body (excluding sensitive variables)
        new_body, body_undefined = VariableSubstitution.substitute_excluding_variables(
            request.body, environment.variables, sensitive_variables
        )
        all_undefined.update(body_undefined)

        # Create new HttpRequest with redacted values
        redacted_request = HttpRequest(
            method=request.method,
            url=new_url,
            headers=new_headers,
            body=new_body,
            syntax=request.syntax
        )

        # Check if any sensitive variables were actually used in the request
        all_vars_in_request = set()
        all_vars_in_request.update(VariableSubstitution.find_variables(request.url))
        all_vars_in_request.update(VariableSubstitution.find_variables(request.body))
        for key, value in request.headers.items():
            all_vars_in_request.update(VariableSubstitution.find_variables(key))
            all_vars_in_request.update(VariableSubstitution.find_variables(value))

        # Determine if any sensitive variables were actually present
        has_sensitive_vars = bool(
            set(sensitive_variables) & all_vars_in_request
        )

        return redacted_request, all_undefined, has_sensitive_vars