diff --git a/src/analysis/plugin/plugin.py b/src/analysis/plugin/plugin.py index 798e1c6ff..5841385cd 100644 --- a/src/analysis/plugin/plugin.py +++ b/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 `_ 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) diff --git a/src/helperFunctions/tag.py b/src/helperFunctions/tag.py index 5913a39ee..1f6ccabfa 100644 --- a/src/helperFunctions/tag.py +++ b/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 diff --git a/src/plugins/analysis/device_tree/code/device_tree.py b/src/plugins/analysis/device_tree/code/device_tree.py index 0d997901c..c458426b2 100644 --- a/src/plugins/analysis/device_tree/code/device_tree.py +++ b/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 diff --git a/src/plugins/analysis/example_plugin/code/example_plugin.py b/src/plugins/analysis/example_plugin/code/example_plugin.py index 705458177..65a27d863 100644 --- a/src/plugins/analysis/example_plugin/code/example_plugin.py +++ b/src/plugins/analysis/example_plugin/code/example_plugin.py @@ -1,25 +1,27 @@ -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 @@ -27,12 +29,12 @@ class Schema(pydantic.BaseModel): 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') diff --git a/src/web_interface/filter.py b/src/web_interface/filter.py index 3ae1ad39f..2658feb26 100644 --- a/src/web_interface/filter.py +++ b/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' '\n' ) - elements.append( - f'
\n' f'{element}\n' '
' - ) + elements.append(f'
\n{element}\n
') return '\n'.join(elements)