feat: convert TagColor class to Enum

and improve plugin type hinting

Author: jstucke

pyproject.toml

@@ -41,7 +41,6 @@ select = [
     "A",  # flake8-builtins
     "C4",  # flake8-comprehensions
     "EXE",  # flake8-executable
-    "FA",  # flake8-future-annotations
     "ISC",  # flake8-implicit-str-concat
     "PIE",  # flake8-pie
     "T20",  # flake8-print
@@ -65,9 +64,6 @@ ignore = [
     "RUF002",
     "RUF003",
     "RUF015",
-    # pydantic only supports these from python>=3.9
-    "UP006",
-    "UP007",
     # rules may cause conflicts when used with the formatter
     "ISC001",
     "Q001",

src/analysis/plugin/plugin.py

@@ -1,15 +1,13 @@
-from __future__ import annotations
-
 import abc
+import io
 import time
-import typing
+from typing import Type, TypeVar, final
 
 import pydantic
 import semver
 from pydantic import BaseModel, ConfigDict, field_validator
 
-if typing.TYPE_CHECKING:
-    import io
+from helperFunctions.tag import TagColor
 
 
 class AnalysisFailedError(Exception):
@@ -31,8 +29,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 +51,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 +64,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 +87,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 +112,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 +124,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,7 +1,6 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Dict
+import io
 
+from pydantic import BaseModel
 from semver import Version
 
 from analysis.plugin import AnalysisPluginV0, Tag
@@ -10,9 +9,6 @@
 
 from ..internal.schema import DeviceTree, IllegalDeviceTreeError, Schema
 
-if TYPE_CHECKING:
-    import io
-
 
 class AnalysisPlugin(AnalysisPluginV0):
     def __init__(self):
@@ -39,7 +35,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,35 @@
 import io
 from pathlib import Path
 
-import pydantic
-from pydantic import Field
+from pydantic import BaseModel, Field
 from semver import Version
 
 from analysis.plugin.plugin import AnalysisFailedError, AnalysisPluginV0
 
 
 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 +38,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

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import binascii
 import ipaddress
 import json
@@ -16,7 +14,7 @@
 from re import Match
 from string import ascii_letters
 from time import localtime, strftime, struct_time, time
-from typing import TYPE_CHECKING, Any, Iterable, Union
+from typing import Any, Iterable, Union
 
 import packaging.version
 import semver
@@ -27,13 +25,11 @@
 from helperFunctions.data_conversion import make_unicode_string
 from helperFunctions.tag import TagColor
 from helperFunctions.web_interface import get_alternating_color_list
+from objects.file import FileObject
 from web_interface.file_tree.file_tree import get_icon_for_mime, get_mime_for_text_file
 from web_interface.security.authentication import user_has_privilege
 from web_interface.security.privileges import PRIVILEGES
 
-if TYPE_CHECKING:
-    from objects.file import FileObject
-
 BYTE_FORMAT_THRESHOLD = 2**10
 
 
@@ -302,7 +298,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 +461,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)