trestle.core.jinja.tags

class MDCleanInclude(TrestleJinjaExtension):
    """Inject the parameter of the tag as the resulting content."""

    tags = {'md_clean_include'}

    def __init__(self, environment: Environment) -> None:
        """Ensure enviroment is set and carried into class vars."""
        super().__init__(environment)

    def parse(self, parser):
        """Execute parsing of md token and return nodes.

        Security Note: This method re-parses included markdown content using a
        SandboxedEnvironment to prevent Server-Side Template Injection (SSTI) attacks.
        The sandbox restricts access to unsafe Python attributes, preventing arbitrary
        code execution even if untrusted OSCAL data contains Jinja syntax.
        """
        kwargs = None
        expected_heading_level = None
        count = 0
        while parser.stream.current.type != lexer.TOKEN_BLOCK_END:
            count = count + 1
            if count > self.max_tag_parse:
                raise err.TrestleError('Unexpected Jinja tag structure provided, please review docs.')
            token = parser.stream.current
            if token.test('name:md_clean_include'):
                parser.stream.expect(lexer.TOKEN_NAME)
                markdown_source = parser.stream.expect(lexer.TOKEN_STRING)
            elif kwargs is not None:
                arg = token.value
                next(parser.stream)
                parser.stream.expect(lexer.TOKEN_ASSIGN)
                token = parser.stream.current
                exp = self.parse_expression(parser)
                kwargs[arg] = exp.value
            else:
                if parser.stream.look().type == lexer.TOKEN_ASSIGN:
                    kwargs = {}
                continue
        md_content, _, _ = self.environment.loader.get_source(self.environment, markdown_source.value)
        fm = frontmatter.loads(md_content)
        content = fm.content
        content += '\n\n'
        if kwargs is not None:
            expected_heading_level = kwargs.get('heading_level')
        if expected_heading_level is not None:
            content = adjust_heading_level(content, expected_heading_level)

        # SECURITY FIX: Use sandboxed environment for re-parsing to prevent SSTI
        # Create a sandboxed environment if the current one isn't already sandboxed
        parse_env = self.environment
        if not isinstance(self.environment, SandboxedEnvironment):
            parse_env = SandboxedEnvironment(
                loader=self.environment.loader,
                extensions=self.environment.extensions,
                trim_blocks=self.environment.trim_blocks,
                autoescape=self.environment.autoescape,
            )

        local_parser = Parser(parse_env, content)
        top_level_output = local_parser.parse()

        return top_level_output.body

class MDDatestamp(TrestleJinjaExtension):
    """Inject the parameter of the tag as the resulting content."""

    tags = {'md_datestamp'}

    def __init__(self, environment: Environment) -> None:
        """Ensure enviroment is set and carried into class vars."""
        super().__init__(environment)

    def parse(self, parser):
        """Execute parsing of md token and return nodes.

        Security Note: This method re-parses the datestamp using a SandboxedEnvironment
        to prevent potential SSTI attacks, though the risk is minimal for date strings.
        """
        kwargs = None
        count = 0
        while parser.stream.current.type != lexer.TOKEN_BLOCK_END:
            count = count + 1
            token = parser.stream.current
            if count > self.max_tag_parse:
                raise err.TrestleError(f'Unexpected Jinja tag structure provided at token {token.value}')
            if token.test('name:md_datestamp'):
                parser.stream.expect(lexer.TOKEN_NAME)
            elif kwargs is not None:
                arg = token.value
                next(parser.stream)
                parser.stream.expect(lexer.TOKEN_ASSIGN)
                token = parser.stream.current
                exp = self.parse_expression(parser)
                kwargs[arg] = exp.value
            else:
                if parser.stream.look().type == lexer.TOKEN_ASSIGN or parser.stream.look().type == lexer.TOKEN_STRING:
                    kwargs = {}
                continue

        if kwargs is not None:
            if 'format' in kwargs and type(kwargs['format'] is str):
                date_string = date.today().strftime(kwargs['format'])
            else:
                date_string = date.today().strftime(markdown_const.JINJA_DATESTAMP_FORMAT)
            if 'newline' in kwargs and kwargs['newline'] is False:
                pass
            else:
                date_string += '\n\n'
        else:
            date_string = date.today().strftime(markdown_const.JINJA_DATESTAMP_FORMAT) + '\n\n'

        # SECURITY FIX: Use sandboxed environment for re-parsing to prevent SSTI
        # Create a sandboxed environment if the current one isn't already sandboxed
        parse_env = self.environment
        if not isinstance(self.environment, SandboxedEnvironment):
            parse_env = SandboxedEnvironment(
                loader=self.environment.loader,
                extensions=self.environment.extensions,
                trim_blocks=self.environment.trim_blocks,
                autoescape=self.environment.autoescape,
            )

        local_parser = Parser(parse_env, date_string)
        datestamp_output = local_parser.parse()

        return datestamp_output.body

class MDSectionInclude(TrestleJinjaExtension):
    """Inject the parameter of the tag as the resulting content."""

    tags = {'mdsection_include'}

    def __init__(self, environment: Environment) -> None:
        """Ensure enviroment is set and carried into class vars."""
        super().__init__(environment)

    def parse(self, parser):
        """Execute parsing of md token and return nodes.

        Security Note: This method re-parses included markdown content using a
        SandboxedEnvironment to prevent Server-Side Template Injection (SSTI) attacks.
        The sandbox restricts access to unsafe Python attributes, preventing arbitrary
        code execution even if untrusted OSCAL data contains Jinja syntax.
        """
        kwargs = None
        expected_heading_level = None
        count = 0
        while parser.stream.current.type != lexer.TOKEN_BLOCK_END:
            count = count + 1
            if count > self.max_tag_parse:
                raise err.TrestleError('Unexpected Jinja tag structure provided, please review docs.')
            token = parser.stream.current
            if token.test('name:mdsection_include'):
                parser.stream.expect(lexer.TOKEN_NAME)
                markdown_source = parser.stream.expect(lexer.TOKEN_STRING)
                section_title = parser.stream.expect(lexer.TOKEN_STRING)
            elif kwargs is not None:
                arg = token.value
                next(parser.stream)
                parser.stream.expect(lexer.TOKEN_ASSIGN)
                token = parser.stream.current
                exp = self.parse_expression(parser)
                kwargs[arg] = exp.value
            else:
                if parser.stream.look().type == lexer.TOKEN_ASSIGN:
                    kwargs = {}
                continue
        # Use the established environment to source the file
        md_content, _, _ = self.environment.loader.get_source(self.environment, markdown_source.value)
        fm = frontmatter.loads(md_content)
        if not fm.metadata == {}:
            logger.warning('Non zero metadata on MD section include - ignoring')
        full_md = docs_markdown_node.DocsMarkdownNode.build_tree_from_markdown(fm.content.split('\n'))
        md_section = full_md.get_node_for_key(section_title.value, strict_matching=True)
        # adjust
        if kwargs is not None:
            expected_heading_level = kwargs.get('heading_level')
        if expected_heading_level is not None:
            level = md_section.get_node_header_lvl()
            delta = int(expected_heading_level) - level
            if not delta == 0:
                md_section.change_header_level_by(delta)
        if not md_section:
            raise err.TrestleError(
                f'Unable to retrieve section "{section_title.value}"" from {markdown_source.value} jinja template.'
            )

        # SECURITY FIX: Use sandboxed environment for re-parsing to prevent SSTI
        # Create a sandboxed environment if the current one isn't already sandboxed
        parse_env = self.environment
        if not isinstance(self.environment, SandboxedEnvironment):
            parse_env = SandboxedEnvironment(
                loader=self.environment.loader,
                extensions=self.environment.extensions,
                trim_blocks=self.environment.trim_blocks,
                autoescape=self.environment.autoescape,
            )

        local_parser = Parser(parse_env, md_section.content.raw_text)
        top_level_output = local_parser.parse()

        return top_level_output.body