updates

Author: ogenstad

infrahub_sdk/ctl/cli_commands.py

@@ -29,7 +29,7 @@
 from ..ctl.generator import run as run_generator
 from ..ctl.menu import app as menu_app
 from ..ctl.object import app as object_app
-from ..ctl.render import list_jinja2_transforms
+from ..ctl.render import list_jinja2_transforms, print_template_errors
 from ..ctl.repository import app as repository_app
 from ..ctl.repository import get_repository_config
 from ..ctl.schema import app as schema_app
@@ -44,12 +44,7 @@
 from ..exceptions import GraphQLError, ModuleImportError
 from ..schema import MainSchemaTypesAll, SchemaRoot
 from ..template import Jinja2Template
-from ..template.exceptions import (
-    JinjaTemplateError,
-    JinjaTemplateNotFoundError,
-    JinjaTemplateSyntaxError,
-    JinjaTemplateUndefinedError,
-)
+from ..template.exceptions import JinjaTemplateError
 from ..utils import get_branch, write_to_file
 from ..yaml import SchemaFile
 from .exporter import dump
@@ -172,50 +167,22 @@ async def run(
         raise typer.Abort(f"Unable to Load the method {method} in the Python script at {script}")
 
     client = initialize_client(
-        branch=branch, timeout=timeout, max_concurrent_execution=concurrent, identifier=module_name
+        branch=branch,
+        timeout=timeout,
+        max_concurrent_execution=concurrent,
+        identifier=module_name,
     )
     func = getattr(module, method)
     await func(client=client, log=log, branch=branch, **variables_dict)
 
 
 async def render_jinja2_template(template_path: Path, variables: dict[str, Any], data: dict[str, Any]) -> str:
-    if not template_path.is_file():
-        console.print(f"[red]Unable to locate the template at {template_path}")
-        raise typer.Exit(1)
-
     variables["data"] = data
     jinja_template = Jinja2Template(template=Path(template_path), template_directory=Path())
     try:
         rendered_tpl = await jinja_template.render(variables=variables)
-    except JinjaTemplateNotFoundError as exc:
-        console.print("[red]An error occurred while rendering the jinja template")
-        console.print("")
-        if exc.base_template:
-            console.print(f"Base template: [yellow]{exc.base_template}")
-        console.print(f"Missing template: [yellow]{exc.filename}")
-        raise typer.Exit(1) from exc
-
-    except JinjaTemplateUndefinedError as exc:
-        console.print("[red]An error occurred while rendering the jinja template")
-        for error in exc.errors:
-            console.print(f"[yellow]{error.frame.filename} on line {error.frame.lineno}\n")
-            console.print(error.syntax)
-        console.print("")
-        console.print(exc.message)
-        raise typer.Exit(1) from exc
-    except JinjaTemplateSyntaxError as exc:
-        console.print("[red]A syntax error was encountered within the template")
-        console.print("")
-        if exc.filename:
-            console.print(f"Filename: [yellow]{exc.filename}")
-        console.print(f"Line number: [yellow]{exc.lineno}")
-        console.print()
-        console.print(exc.message)
-        raise typer.Exit(1) from exc
     except JinjaTemplateError as exc:
-        console.print("[red]An error occurred while rendering the jinja template")
-        console.print("")
-        console.print(f"[yellow]{exc.message}")
+        print_template_errors(error=exc, console=console)
         raise typer.Exit(1) from exc
 
     return rendered_tpl
@@ -244,7 +211,11 @@ async def _run_transform(
 
     try:
         response = execute_graphql_query(
-            query=query_name, variables_dict=variables, branch=branch, debug=debug, repository_config=repository_config
+            query=query_name,
+            variables_dict=variables,
+            branch=branch,
+            debug=debug,
+            repository_config=repository_config,
         )
 
         # TODO: response is a dict and can't be printed to the console in this way.
@@ -427,7 +398,10 @@ def version() -> None:
 
 @app.command(name="info")
 @catch_exception(console=console)
-def info(detail: bool = typer.Option(False, help="Display detailed information."), _: str = CONFIG_PARAM) -> None:  # noqa: PLR0915
+def info(  # noqa: PLR0915
+    detail: bool = typer.Option(False, help="Display detailed information."),
+    _: str = CONFIG_PARAM,
+) -> None:
     """Display the status of the Python SDK."""
 
     info: dict[str, Any] = {
@@ -493,10 +467,14 @@ def info(detail: bool = typer.Option(False, help="Display detailed information."
         infrahub_info = Table(show_header=False, box=None)
         if info["user_info"]:
             infrahub_info.add_row("User:", info["user_info"]["AccountProfile"]["display_label"])
-            infrahub_info.add_row("Description:", info["user_info"]["AccountProfile"]["description"]["value"])
+            infrahub_info.add_row(
+                "Description:",
+                info["user_info"]["AccountProfile"]["description"]["value"],
+            )
             infrahub_info.add_row("Status:", info["user_info"]["AccountProfile"]["status"]["label"])
             infrahub_info.add_row(
-                "Number of Groups:", str(info["user_info"]["AccountProfile"]["member_of_groups"]["count"])
+                "Number of Groups:",
+                str(info["user_info"]["AccountProfile"]["member_of_groups"]["count"]),
             )
 
             if groups := info["groups"]:

infrahub_sdk/ctl/render.py

@@ -1,6 +1,12 @@
 from rich.console import Console
 
 from ..schema.repository import InfrahubRepositoryConfig
+from ..template.exceptions import (
+    JinjaTemplateError,
+    JinjaTemplateNotFoundError,
+    JinjaTemplateSyntaxError,
+    JinjaTemplateUndefinedError,
+)
 
 
 def list_jinja2_transforms(config: InfrahubRepositoryConfig) -> None:
@@ -9,3 +15,36 @@ def list_jinja2_transforms(config: InfrahubRepositoryConfig) -> None:
 
     for transform in config.jinja2_transforms:
         console.print(f"{transform.name} ({transform.template_path})")
+
+
+def print_template_errors(error: JinjaTemplateError, console: Console) -> None:
+    if isinstance(error, JinjaTemplateNotFoundError):
+        console.print("[red]An error occurred while rendering the jinja template")
+        console.print("")
+        if error.base_template:
+            console.print(f"Base template: [yellow]{error.base_template}")
+        console.print(f"Missing template: [yellow]{error.filename}")
+        return
+
+    if isinstance(error, JinjaTemplateUndefinedError):
+        console.print("[red]An error occurred while rendering the jinja template")
+        for current_error in error.errors:
+            console.print(f"[yellow]{current_error.frame.filename} on line {current_error.frame.lineno}\n")
+            console.print(current_error.syntax)
+        console.print("")
+        console.print(error.message)
+        return
+
+    if isinstance(error, JinjaTemplateSyntaxError):
+        console.print("[red]A syntax error was encountered within the template")
+        console.print("")
+        if error.filename:
+            console.print(f"Filename: [yellow]{error.filename}")
+        console.print(f"Line number: [yellow]{error.lineno}")
+        console.print()
+        console.print(error.message)
+        return
+
+    console.print("[red]An error occurred while rendering the jinja template")
+    console.print("")
+    console.print(f"[yellow]{error.message}")

infrahub_sdk/template/__init__.py

@@ -5,37 +5,48 @@
 from typing import Any, Callable, NoReturn
 
 import jinja2
-from jinja2 import meta
+from jinja2 import meta, nodes
 from jinja2.sandbox import SandboxedEnvironment
 from netutils.utils import jinja2_convenience_function
 from rich.syntax import Syntax
 from rich.traceback import Traceback
 
-from .constants import CURRATED_NETUTILS_FILTERS
 from .exceptions import (
     JinjaTemplateError,
     JinjaTemplateNotFoundError,
     JinjaTemplateOperationViolationError,
     JinjaTemplateSyntaxError,
     JinjaTemplateUndefinedError,
 )
+from .filters import AVAILABLE_FILTERS
 from .models import UndefinedJinja2Error
 
 netutils_filters = jinja2_convenience_function()
 
 
 class Jinja2Template:
-    def __init__(self, template: str | Path, template_directory: Path | None = None) -> None:
+    def __init__(
+        self,
+        template: str | Path,
+        template_directory: Path | None = None,
+        filters: dict[str, Callable] | None = None,
+    ) -> None:
         self.is_string_based = isinstance(template, str)
         self.is_file_based = isinstance(template, Path)
         self._template = str(template)
         self._template_directory = template_directory
         self._environment: jinja2.Environment | None = None
-        self._filters: dict[str, Callable] = {}
 
-        self._filters.update(
-            {name: jinja_filter for name, jinja_filter in netutils_filters.items() if name in CURRATED_NETUTILS_FILTERS}
-        )
+        self._available_filters = [filter_definition.name for filter_definition in AVAILABLE_FILTERS]
+        self._trusted_filters = [
+            filter_definition.name for filter_definition in AVAILABLE_FILTERS if filter_definition.trusted
+        ]
+
+        self._filters = filters or {}
+        for user_filter in self._filters:
+            self._available_filters.append(user_filter)
+            self._trusted_filters.append(user_filter)
+
         self._template_definition: jinja2.Template | None = None
 
     def get_environment(self) -> jinja2.Environment:
@@ -78,6 +89,21 @@ def get_variables(self) -> list[str]:
 
         return sorted(meta.find_undeclared_variables(template))
 
+    def validate_filters(self, restricted: bool = True) -> None:
+        allowed_list = self._available_filters
+        if restricted:
+            allowed_list = self._trusted_filters
+
+        env = self.get_environment()
+        template_source = self._template
+        if self.is_file_based and env.loader:
+            template_source = env.loader.get_source(env, self._template)[0]
+
+        template = env.parse(template_source)
+        for node in template.find_all(nodes.Filter):
+            if node.name not in allowed_list:
+                raise JinjaTemplateOperationViolationError(f"The '{node.name}' filter isn't allowed to be used")
+
     async def render(self, variables: dict[str, Any]) -> str:
         template = self.get_template()
         try:
@@ -101,17 +127,34 @@ async def render(self, variables: dict[str, Any]) -> str:
 
     def _get_string_based_environment(self) -> jinja2.Environment:
         env = SandboxedEnvironment(enable_async=True, undefined=jinja2.StrictUndefined)
-        env.filters.update(self._filters)
+        self._set_filters(env=env)
         self._environment = env
         return self._environment
 
     def _get_file_based_environment(self) -> jinja2.Environment:
         template_loader = jinja2.FileSystemLoader(searchpath=str(self._template_directory))
-        env = jinja2.Environment(loader=template_loader, trim_blocks=True, lstrip_blocks=True, enable_async=True)
-        env.filters.update(self._filters)
+        env = jinja2.Environment(
+            loader=template_loader,
+            trim_blocks=True,
+            lstrip_blocks=True,
+            enable_async=True,
+        )
+        self._set_filters(env=env)
         self._environment = env
         return self._environment
 
+    def _set_filters(self, env: jinja2.Environment) -> None:
+        for default_filter in list(env.filters.keys()):
+            if default_filter not in self._available_filters:
+                del env.filters[default_filter]
+
+        # Add filters from netutils
+        env.filters.update(
+            {name: jinja_filter for name, jinja_filter in netutils_filters.items() if name in self._available_filters}
+        )
+        # Add user supplied filters
+        env.filters.update(self._filters)
+
     def _get_string_based_template(self) -> jinja2.Template:
         env = self.get_environment()
         self._template_definition = env.from_string(self._template)
@@ -150,7 +193,10 @@ def _identify_faulty_jinja_code(traceback: Traceback, nbr_context_lines: int = 3
             code,
             lexer_name,
             line_numbers=True,
-            line_range=(frame.lineno - nbr_context_lines, frame.lineno + nbr_context_lines),
+            line_range=(
+                frame.lineno - nbr_context_lines,
+                frame.lineno + nbr_context_lines,
+            ),
             highlight_lines={frame.lineno},
             code_width=88,
             theme=traceback.theme,