feat: convert TagColor class to Enum

and improve plugin type hinting

Author: jstucke

src/analysis/plugin/plugin.py

@@ -2,13 +2,15 @@
 
 import abc
 import time
-import typing
+from typing import TYPE_CHECKING, Type, TypeVar, final
 
 import pydantic
 import semver
 from pydantic import BaseModel, ConfigDict, field_validator
 
-if typing.TYPE_CHECKING:
+from helperFunctions.tag import TagColor  # noqa: TCH001
+
+if TYPE_CHECKING:
     import io
 
 
@@ -31,8 +33,7 @@ class Tag(BaseModel):
     #: In FACT_core this is shown as tooltip
     value: str
     #: The color of the tag
-    #: See :py:class:`helperFunctions.tag.TagColor`.
-    color: str
+    color: TagColor
     #: Whether or not the tag should be shown in parent files.
     propagate: bool = False
 
@@ -54,7 +55,7 @@ class MetaData(BaseModel):
         description: str
         #: Pydantic model of the object returned by :py:func:`analyse`.
         # Note that we cannot allow pydantic dataclasses because they lack the `schema` method
-        Schema: typing.Type
+        Schema: Type[BaseModel]
         #: The version of the plugin.
         #: It MUST be a `semver <https://semver.org/>`_ version.
         #: Here is a quick summary how semver relates to plugins.
@@ -67,13 +68,13 @@ class MetaData(BaseModel):
         version: semver.Version
         #: The version of the backing analysis system.
         #: E.g. for yara plugins this would be the yara version.
-        system_version: typing.Optional[str] = None
+        system_version: str | None = None
         #: A list of all plugins that this plugin depends on
-        dependencies: typing.List = pydantic.Field(default_factory=list)
+        dependencies: list[str] = pydantic.Field(default_factory=list)
         #: List of mimetypes that should not be processed
-        mime_blacklist: list = pydantic.Field(default_factory=list)
+        mime_blacklist: list[str] = pydantic.Field(default_factory=list)
         #: List of mimetypes that should be processed
-        mime_whitelist: list = pydantic.Field(default_factory=list)
+        mime_whitelist: list[str] = pydantic.Field(default_factory=list)
         #: The analysis in not expected to take longer than timeout seconds on any given file
         #: and will be aborted if the timeout is reached.
         timeout: int = 300
@@ -90,7 +91,7 @@ def __init__(self, metadata: MetaData):
         self.metadata: AnalysisPluginV0.MetaData = metadata
 
     # The type MetaData.Schema
-    Schema = typing.TypeVar('Schema')
+    Schema = TypeVar('Schema')
 
     def summarize(self, result: Schema) -> list[str]:  # noqa: ARG002
         """
@@ -115,7 +116,7 @@ def analyze(
         self,
         file_handle: io.FileIO,
         virtual_file_path: dict,
-        analyses: dict[str, pydantic.BaseModel],
+        analyses: dict[str, BaseModel],
     ) -> Schema:
         """Analyze a file.
         May return None if nothing was found.
@@ -127,8 +128,8 @@ def analyze(
         :return: The analysis if anything was found.
         """
 
-    @typing.final
-    def get_analysis(self, file_handle: io.FileIO, virtual_file_path: dict, analyses: dict[str, dict]) -> dict:
+    @final
+    def get_analysis(self, file_handle: io.FileIO, virtual_file_path: dict, analyses: dict[str, BaseModel]) -> dict:
         start_time = time.time()
         result = self.analyze(file_handle, virtual_file_path, analyses)
 

src/helperFunctions/tag.py

@@ -1,6 +1,9 @@
-class TagColor:
+from enum import Enum
+
+
+class TagColor(str, Enum):
     """
-    A class containing the different colors the tags may have. `TagColor.ALL` contains a list of all colors.
+    A class containing the different colors the tags may have.
     """
 
     GRAY = 'secondary'
@@ -11,4 +14,6 @@ class TagColor:
     RED = 'danger'
     LIGHT = 'light'
     DARK = 'dark'
-    ALL = [GRAY, BLUE, GREEN, LIGHT_BLUE, ORANGE, RED, LIGHT, DARK]  # noqa: RUF012
+
+    def __str__(self) -> str:
+        return self.value

src/plugins/analysis/device_tree/code/device_tree.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING
 
 from semver import Version
 
@@ -13,6 +13,8 @@
 if TYPE_CHECKING:
     import io
 
+    from pydantic import BaseModel
+
 
 class AnalysisPlugin(AnalysisPluginV0):
     def __init__(self):
@@ -39,7 +41,7 @@ def analyze(
         self,
         file_handle: io.FileIO,
         virtual_file_path: dict,
-        analyses: Dict[str, dict],
+        analyses: dict[str, BaseModel],
     ) -> Schema:
         del virtual_file_path, analyses
 

src/plugins/analysis/example_plugin/code/example_plugin.py

@@ -1,38 +1,40 @@
-import io
+from __future__ import annotations
+
 from pathlib import Path
+from typing import TYPE_CHECKING
 
-import pydantic
-from pydantic import Field
+from pydantic import BaseModel, Field
 from semver import Version
 
 from analysis.plugin.plugin import AnalysisFailedError, AnalysisPluginV0
 
+if TYPE_CHECKING:
+    import io
+
 
 class AnalysisPlugin(AnalysisPluginV0):
-    class Schema(pydantic.BaseModel):
+    class Schema(BaseModel):
         """Here goes the toplevel description of the plugin result"""
 
-        # fmt: off
         number: int = Field(
             description=(
-                'This is a description of the field number.\n'
-                'In an actual plugin all fields must have a description.'
+                'This is a description of the field "number".\n'
+                'In an actual plugin all fields should have a description.'
             ),
         )
-        # fmt: on
         name: str
         first_byte: str
         virtual_file_path: dict
         dependant_analysis: dict
 
     def __init__(self):
         metadata = self.MetaData(
+            # mandatory fields:
             name='ExamplePlugin',
             description='An example description',
             version=Version(0, 0, 0),
-            Schema=AnalysisPlugin.Schema,
-            # Note that you don't have to set these fields,
-            # they are just here to show that you can.
+            Schema=self.Schema,
+            # optional fields:
             system_version=None,
             mime_blacklist=[],
             mime_whitelist=[],
@@ -41,11 +43,11 @@ def __init__(self):
         )
         super().__init__(metadata=metadata)
 
-    def summarize(self, result):
+    def summarize(self, result: Schema) -> list[str]:
         del result
         return ['big-file', 'binary']
 
-    def analyze(self, file_handle: io.FileIO, virtual_file_path: dict, analyses: dict) -> Schema:
+    def analyze(self, file_handle: io.FileIO, virtual_file_path: dict, analyses: dict[str, BaseModel]) -> Schema:
         first_byte = file_handle.read(1)
         if first_byte == b'\xff':
             raise AnalysisFailedError('reason for fail')

src/web_interface/filter.py

@@ -302,7 +302,7 @@ def render_analysis_tags(tags, uid=None, root_uid=None, size=14):
             for key, tag in sorted(tags[plugin_name].items(), key=_sort_tags_key):
                 if key == 'root_uid':
                     continue
-                color = tag['color'] if tag['color'] in TagColor.ALL else TagColor.BLUE
+                color = tag['color'] if tag['color'] in set(TagColor) else TagColor.BLUE
                 output += render_template(
                     'generic_view/tags.html',
                     color=color,
@@ -465,9 +465,7 @@ def render_changed_text_files(changed_text_files: dict) -> str:
                 f'  {inner}\n'
                 '</div>\n'
             )
-        elements.append(
-            f'<div class="list-group-item border-top" style="padding: 0 0 0 25px">\n' f'{element}\n' '</div>'
-        )
+        elements.append(f'<div class="list-group-item border-top" style="padding: 0 0 0 25px">\n{element}\n</div>')
     return '\n'.join(elements)