Merge pull request #170 from Syslifters/library-improvements

Library improvements

CHANGELOG.md

@@ -1,3 +1,6 @@
+# 0.18
+* Several adaptions to allow reptor be used as library
+
 # 0.17
 * Respect finding order in Project (and exportfindings)
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "reptor"
-version = "0.17"
+version = "0.18"
 authors = [
   { name="Richard Schwabe" },
   { name="Aron Molnar", email="[email protected]" }

reptor/lib/plugins/ToolBase.py

@@ -18,6 +18,7 @@
 from reptor.models.FindingTemplate import FindingTemplate
 from reptor.models.Note import NoteTemplate
 from reptor.models.ProjectDesign import ProjectDesign
+from reptor.utils.django_tags import custom_django_tags
 
 from .Base import Base
 
@@ -356,9 +357,10 @@ def format_note_template(self, note_templates: typing.List[NoteTemplate], level=
         formatted_input = ""
         for note_template in note_templates:
             if note_template.template:
-                note_template.text = render_to_string(
-                    f"{note_template.template}.md", note_template.template_data
-                )
+                with custom_django_tags():
+                    note_template.text = render_to_string(
+                        f"{note_template.template}.md", note_template.template_data
+                    )
             if note_template.title:
                 formatted_input += f"{'#' * level} {note_template.title}\n\n"
             if note_template.text:
@@ -567,9 +569,10 @@ def generate_findings(self) -> typing.List[Finding]:
                     if finding_data.value:
                         if finding_data.type in ["markdown", "string", "cvss"]:
                             # Render template
-                            finding_data.value = Template(finding_data.value).render(
-                                django_context
-                            )
+                            with custom_django_tags():
+                                finding_data.value = Template(finding_data.value).render(
+                                    django_context
+                                )
                     elif finding_data.type in [
                         "list",
                         "enum",

reptor/lib/reptor.py

@@ -17,7 +17,6 @@
 from reptor.lib.logger import ReptorAdapter, reptor_logger
 from reptor.lib.pluginmanager import PluginManager
 from reptor.lib.plugins.PluginMeta import PluginMeta
-from reptor.utils.django_tags import setup_django_tags
 from reptor.utils.markdown import convert_markdown_to_console
 
 root_logger = logging.getLogger("root")
@@ -63,7 +62,11 @@ def __init__(
         token: typing.Optional[str] = None,
         project_id: typing.Optional[str] = None,
     ) -> None:
-        signal.signal(signal.SIGINT, signal_handler)
+        try:
+            signal.signal(signal.SIGINT, signal_handler)
+        except ValueError:
+            # not possible in thread. ignore
+            pass
 
         # Set encoding for stdin utf-8
         try:
@@ -75,9 +78,6 @@ def __init__(
         self._config = Config()
         self._config.load_config(server=server, token=token, project_id=project_id)
 
-        # Setup Django tags
-        setup_django_tags()
-
         self.plugin_manager = PluginManager(self)
 
         if self.get_config().get_log_file():

reptor/utils/django_tags.py

@@ -1,6 +1,8 @@
 import re
+from contextlib import contextmanager
+from unittest import mock
 
-from django.template import base as django_base
+import django
 from django.template.base import Token, TokenType
 
 # HMTL comment tags
@@ -12,8 +14,8 @@
 COMMENT_TAG_START = "<!--{#"
 COMMENT_TAG_END = "#}-->"
 
-
-def setup_django_tags():
+@contextmanager
+def custom_django_tags():
     def create_token(self, token_string, position, lineno, in_tag):
         """
         Convert the given token string into a new Token object and return it.
@@ -26,10 +28,10 @@ def create_token(self, token_string, position, lineno, in_tag):
             # len(BLOCK_TAG_START) would permit BLOCK_TAG_START to be
             # different, but it's not likely that the TAG_START values will
             # change anytime soon.
-            token_start = token_string[0 : len(django_base.BLOCK_TAG_START)]
-            if token_start == django_base.BLOCK_TAG_START:
+            token_start = token_string[0 : len(django.template.base.BLOCK_TAG_START)]
+            if token_start == django.template.base.BLOCK_TAG_START:
                 content = token_string[
-                    len(django_base.BLOCK_TAG_START) : -len(django_base.BLOCK_TAG_END)
+                    len(django.template.base.BLOCK_TAG_START) : -len(django.template.base.BLOCK_TAG_END)
                 ].strip()
                 if self.verbatim:
                     # Then a verbatim block is being processed.
@@ -43,21 +45,23 @@ def create_token(self, token_string, position, lineno, in_tag):
                 return Token(TokenType.BLOCK, content, position, lineno)
             if not self.verbatim:
                 content = token_string[
-                    len(django_base.BLOCK_TAG_START) : -len(django_base.BLOCK_TAG_END)
+                    len(django.template.base.BLOCK_TAG_START) : -len(django.template.base.BLOCK_TAG_END)
                 ].strip()
-                if token_start == django_base.VARIABLE_TAG_START:
+                if token_start == django.template.base.VARIABLE_TAG_START:
                     return Token(TokenType.VAR, content, position, lineno)
                 # BLOCK_TAG_START was handled above.
-                assert token_start == django_base.COMMENT_TAG_START
+                assert token_start == django.template.base.COMMENT_TAG_START
                 return Token(TokenType.COMMENT, content, position, lineno)
         return Token(TokenType.TEXT, token_string, position, lineno)
 
+    with mock.patch('django.template.base.tag_re', HTML_REGEX), \
+        mock.patch('django.template.base.Lexer.create_token', create_token), \
+        mock.patch('django.template.base.BLOCK_TAG_START', BLOCK_TAG_START), \
+        mock.patch('django.template.base.BLOCK_TAG_END', BLOCK_TAG_END), \
+        mock.patch('django.template.base.VARIABLE_TAG_START', VARIABLE_TAG_START), \
+        mock.patch('django.template.base.VARIABLE_TAG_END', VARIABLE_TAG_END), \
+        mock.patch('django.template.base.COMMENT_TAG_START', COMMENT_TAG_START), \
+        mock.patch('django.template.base.COMMENT_TAG_END', COMMENT_TAG_END):
+        yield
     # Monkey patch
-    django_base.tag_re = HTML_REGEX
-    django_base.Lexer.create_token = create_token
-    django_base.BLOCK_TAG_START = BLOCK_TAG_START
-    django_base.BLOCK_TAG_END = BLOCK_TAG_END
-    django_base.VARIABLE_TAG_START = VARIABLE_TAG_START
-    django_base.VARIABLE_TAG_END = VARIABLE_TAG_END
-    django_base.COMMENT_TAG_START = COMMENT_TAG_START
-    django_base.COMMENT_TAG_END = COMMENT_TAG_END
+