Fix Pydantic warnings (#3557)

Co-authored-by: stephanie0x00 <[email protected]> Co-authored-by: Jan Klopper <[email protected]>
minvws · Sep 23, 2024 · 1ba9bc6 · 1ba9bc6
1 parent 4463cbd
commit 1ba9bc6
Show file tree

Hide file tree

Showing 39 changed files with 117 additions and 109 deletions.
diff --git a/boefjes/boefjes/api.py b/boefjes/boefjes/api.py
@@ -59,7 +59,7 @@ class StatusEnum(str, Enum):
 
 class File(BaseModel):
     name: str | None = None
-    content: str = Field(..., contentEncoding="base64")
+    content: str = Field(json_schema_extra={"contentEncoding": "base64"})
     tags: list[str] | None = None
 
 
@@ -171,6 +171,6 @@ def create_boefje_meta(task, plugin: PluginType) -> BoefjeMeta:
         input_ooi=input_ooi,
         arguments=arguments,
         organization=organization,
-        environment=get_environment_settings(task.data, plugin.schema),
+        environment=get_environment_settings(task.data, plugin.boefje_schema),
     )
     return boefje_meta
diff --git a/boefjes/boefjes/clients/bytes_client.py b/boefjes/boefjes/clients/bytes_client.py
@@ -74,7 +74,7 @@ def _get_token(self) -> str:
 
     @retry_with_login
     def save_boefje_meta(self, boefje_meta: BoefjeMeta) -> None:
-        response = self._session.post("/bytes/boefje_meta", content=boefje_meta.json(), headers=self.headers)
+        response = self._session.post("/bytes/boefje_meta", content=boefje_meta.model_dump_json(), headers=self.headers)
 
         self._verify_response(response)
 
@@ -87,7 +87,9 @@ def get_boefje_meta(self, boefje_meta_id: str) -> BoefjeMeta:
 
     @retry_with_login
     def save_normalizer_meta(self, normalizer_meta: NormalizerMeta) -> None:
-        response = self._session.post("/bytes/normalizer_meta", content=normalizer_meta.json(), headers=self.headers)
+        response = self._session.post(
+            "/bytes/normalizer_meta", content=normalizer_meta.model_dump_json(), headers=self.headers
+        )
 
         self._verify_response(response)
 

diff --git a/boefjes/boefjes/clients/scheduler_client.py b/boefjes/boefjes/clients/scheduler_client.py
@@ -76,7 +76,7 @@ def pop_item(self, queue: str) -> Task | None:
         return TypeAdapter(Task | None).validate_json(response.content)
 
     def push_item(self, queue_id: str, p_item: Task) -> None:
-        response = self._session.post(f"/queues/{queue_id}/push", content=p_item.json())
+        response = self._session.post(f"/queues/{queue_id}/push", content=p_item.model_dump_json())
         self._verify_response(response)
 
     def patch_task(self, task_id: uuid.UUID, status: TaskStatus) -> None:

diff --git a/boefjes/boefjes/dependencies/plugins.py b/boefjes/boefjes/dependencies/plugins.py
@@ -156,7 +156,7 @@ def schema(self, plugin_id: str) -> dict | None:
         try:
             boefje = self.plugin_storage.boefje_by_id(plugin_id)
 
-            return boefje.schema
+            return boefje.boefje_schema
         except PluginNotFound:
             return self.local_repo.schema(plugin_id)
 

diff --git a/boefjes/boefjes/job_handler.py b/boefjes/boefjes/job_handler.py
@@ -117,7 +117,7 @@ def handle(self, boefje_meta: BoefjeMeta) -> None:
             boefje_meta.arguments["input"] = ooi.serialize()
 
         boefje_meta.runnable_hash = plugin.runnable_hash
-        boefje_meta.environment = get_environment_settings(boefje_meta, plugin.schema)
+        boefje_meta.environment = get_environment_settings(boefje_meta, plugin.boefje_schema)
 
         mime_types = _default_mime_types(boefje_meta.boefje).union(plugin.produces)
 

diff --git a/boefjes/boefjes/katalogus/plugins.py b/boefjes/boefjes/katalogus/plugins.py
@@ -130,11 +130,11 @@ class BoefjeIn(BaseModel):
     scan_level: int = 1
     consumes: set[str] = Field(default_factory=set)
     produces: set[str] = Field(default_factory=set)
-    schema: dict | None = None
+    boefje_schema: dict | None = None
     oci_image: str | None = None
     oci_arguments: list[str] = Field(default_factory=list)
 
-    @field_validator("schema")
+    @field_validator("boefje_schema")
     @classmethod
     def json_schema_valid(cls, schema: dict | None) -> dict | None:
         if schema is not None:

diff --git a/boefjes/boefjes/models.py b/boefjes/boefjes/models.py
@@ -29,12 +29,12 @@ class Boefje(Plugin):
     scan_level: int = 1
     consumes: set[str] = Field(default_factory=set)
     produces: set[str] = Field(default_factory=set)
-    schema: dict | None = None
+    boefje_schema: dict | None = None
     runnable_hash: str | None = None
     oci_image: str | None = None
     oci_arguments: list[str] = Field(default_factory=list)
 
-    @field_validator("schema")
+    @field_validator("boefje_schema")
     @classmethod
     def json_schema_valid(cls, schema: dict) -> dict:
         if schema is not None:

diff --git a/boefjes/boefjes/plugins/models.py b/boefjes/boefjes/plugins/models.py
@@ -59,7 +59,7 @@ class BoefjeResource:
 
     def __init__(self, path: Path, package: str):
         self.path = path
-        self.boefje: Boefje = Boefje.parse_file(path / BOEFJE_DEFINITION_FILE)
+        self.boefje: Boefje = Boefje.model_validate_json(path.joinpath(BOEFJE_DEFINITION_FILE).read_text())
         self.boefje.runnable_hash = get_runnable_hash(self.path)
         self.boefje.produces = self.boefje.produces.union(set(_default_mime_types(self.boefje)))
         self.module: Runnable | None = None
@@ -69,7 +69,7 @@ def __init__(self, path: Path, package: str):
 
         if (path / SCHEMA_FILE).exists():
             try:
-                self.boefje.schema = json.load((path / SCHEMA_FILE).open())
+                self.boefje.boefje_schema = json.load((path / SCHEMA_FILE).open())
             except JSONDecodeError as e:
                 raise ModuleException("Invalid schema file") from e
             except SchemaError as e:
@@ -81,7 +81,7 @@ class NormalizerResource:
 
     def __init__(self, path: Path, package: str):
         self.path = path
-        self.normalizer = Normalizer.parse_file(path / NORMALIZER_DEFINITION_FILE)
+        self.normalizer = Normalizer.model_validate_json(path.joinpath(NORMALIZER_DEFINITION_FILE).read_text())
         self.normalizer.consumes.append(f"normalizer/{self.normalizer.id}")
         self.module = get_runnable_module_from_package(package, ENTRYPOINT_NORMALIZERS, parameter_count=2)
 

diff --git a/boefjes/boefjes/sql/organisation_storage.py b/boefjes/boefjes/sql/organisation_storage.py
@@ -30,7 +30,7 @@ def get_all(self) -> dict[str, Organisation]:
         return {organisation.id: self.to_organisation(organisation) for organisation in query.all()}
 
     def create(self, organisation: Organisation) -> None:
-        logger.info("Saving organisation: %s", organisation.json())
+        logger.info("Saving organisation: %s", organisation.model_dump_json())
 
         organisation_in_db = self.to_organisation_in_db(organisation)
         self.session.add(organisation_in_db)

diff --git a/boefjes/boefjes/sql/plugin_storage.py b/boefjes/boefjes/sql/plugin_storage.py
@@ -35,7 +35,7 @@ def normalizer_by_id(self, normalizer_id: str) -> Normalizer:
         return self.to_normalizer(instance)
 
     def create_boefje(self, boefje: Boefje) -> None:
-        logger.info("Saving plugin: %s", boefje.json())
+        logger.info("Saving plugin: %s", boefje.model_dump_json())
 
         boefje_in_db = self.to_boefje_in_db(boefje)
         self.session.add(boefje_in_db)
@@ -49,13 +49,14 @@ def update_boefje(self, boefje_id: str, data: dict) -> None:
         if instance.static:
             raise NotAllowed(f"Plugin with id '{boefje_id}' is static, so updating it is not allowed")
 
+        field_mapping = {"boefje_schema": "schema"}  # since Boefje.boefje_schema is the same as BoefjeInDB.schema
         for key, value in data.items():
-            setattr(instance, key, value)
+            setattr(instance, field_mapping.get(key, key), value)
 
         self.session.add(instance)
 
     def create_normalizer(self, normalizer: Normalizer) -> None:
-        logger.info("Saving plugin: %s", normalizer.json())
+        logger.info("Saving plugin: %s", normalizer.model_dump_json())
 
         normalizer_in_db = self.to_normalizer_in_db(normalizer)
         self.session.add(normalizer_in_db)
@@ -110,7 +111,7 @@ def to_boefje_in_db(boefje: Boefje) -> BoefjeInDB:
             scan_level=str(boefje.scan_level),
             consumes=boefje.consumes,
             produces=boefje.produces,
-            schema=boefje.schema,
+            schema=boefje.boefje_schema,
             oci_image=boefje.oci_image,
             oci_arguments=boefje.oci_arguments,
             version=boefje.version,
@@ -141,7 +142,7 @@ def to_boefje(boefje_in_db: BoefjeInDB) -> Boefje:
             scan_level=int(boefje_in_db.scan_level),
             consumes=boefje_in_db.consumes,
             produces=boefje_in_db.produces,
-            schema=boefje_in_db.schema,
+            boefje_schema=boefje_in_db.schema,
             oci_image=boefje_in_db.oci_image,
             oci_arguments=boefje_in_db.oci_arguments,
             version=boefje_in_db.version,

diff --git a/boefjes/tests/integration/test_api.py b/boefjes/tests/integration/test_api.py
@@ -32,7 +32,7 @@ def test_filter_plugins(test_client, organisation):
     boefje = Boefje(
         id="test_plugin", name="My test boefje", static=False, oci_image="ghcr.io/minvws/openkat/nmap:latest"
     )
-    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.json())
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())
     assert response.status_code == 201
 
     response = test_client.get(
@@ -43,19 +43,19 @@ def test_filter_plugins(test_client, organisation):
 
 def test_cannot_add_plugin_reserved_id(test_client, organisation):
     boefje = Boefje(id="dns-records", name="My test boefje", static=False)
-    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.json())
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())
     assert response.status_code == 400
     assert response.json() == {"message": "Plugin id 'dns-records' is already used"}
 
     normalizer = Normalizer(id="kat_nmap_normalize", name="My test normalizer")
-    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.json())
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.model_dump_json())
     assert response.status_code == 400
     assert response.json() == {"message": "Plugin id 'kat_nmap_normalize' is already used"}
 
 
 def test_add_boefje(test_client, organisation):
     boefje = Boefje(id="test_plugin", name="My test boefje", static=False)
-    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.json())
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())
     assert response.status_code == 201
 
     response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", json={"a": "b"})
@@ -64,7 +64,7 @@ def test_add_boefje(test_client, organisation):
     response = test_client.get(f"/v1/organisations/{organisation.id}/plugins/?plugin_type=boefje")
     assert len(response.json()) == 45
 
-    boefje_dict = boefje.dict()
+    boefje_dict = boefje.model_dump()
     boefje_dict["consumes"] = list(boefje_dict["consumes"])
     boefje_dict["produces"] = list(boefje_dict["produces"])
 
@@ -74,7 +74,7 @@ def test_add_boefje(test_client, organisation):
 
 def test_delete_boefje(test_client, organisation):
     boefje = Boefje(id="test_plugin", name="My test boefje", static=False)
-    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.json())
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())
     assert response.status_code == 201
 
     response = test_client.delete(f"/v1/organisations/{organisation.id}/boefjes/test_plugin")
@@ -85,19 +85,19 @@ def test_delete_boefje(test_client, organisation):
 
 def test_add_normalizer(test_client, organisation):
     normalizer = Normalizer(id="test_normalizer", name="My test normalizer", static=False)
-    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.json())
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.model_dump_json())
     assert response.status_code == 201
 
     response = test_client.get(f"/v1/organisations/{organisation.id}/plugins/?plugin_type=normalizer")
     assert len(response.json()) == 56
 
     response = test_client.get(f"/v1/organisations/{organisation.id}/plugins/test_normalizer")
-    assert response.json() == normalizer.dict()
+    assert response.json() == normalizer.model_dump()
 
 
 def test_delete_normalizer(test_client, organisation):
     normalizer = Normalizer(id="test_normalizer", name="My test normalizer", static=False)
-    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.json())
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.model_dump_json())
     assert response.status_code == 201
 
     response = test_client.delete(f"/v1/organisations/{organisation.id}/normalizers/test_normalizer")
@@ -110,15 +110,15 @@ def test_update_plugins(test_client, organisation):
     normalizer = Normalizer(id="norm_id", name="My test normalizer")
     boefje = Boefje(id="test_plugin", name="My test boefje", description="123")
 
-    test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.json())
+    test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())
     test_client.patch(f"/v1/organisations/{organisation.id}/boefjes/{boefje.id}", json={"description": "4"})
     test_client.patch(f"/v1/organisations/{organisation.id}/plugins/{boefje.id}", json={"enabled": True})
 
     response = test_client.get(f"/v1/organisations/{organisation.id}/plugins/{boefje.id}")
     assert response.json()["description"] == "4"
     assert response.json()["enabled"] is True
 
-    test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.json())
+    test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=normalizer.model_dump_json())
     test_client.patch(f"/v1/organisations/{organisation.id}/normalizers/{normalizer.id}", json={"version": "v1.2"})
 
     response = test_client.get(f"/v1/organisations/{organisation.id}/plugins/{normalizer.id}")
@@ -127,17 +127,19 @@ def test_update_plugins(test_client, organisation):
 
 def test_cannot_create_boefje_with_invalid_schema(test_client, organisation):
     boefje = Boefje(id="test_plugin", name="My test boefje", description="123").model_dump(mode="json")
-    boefje["schema"] = {"$schema": 3}
+    boefje["boefje_schema"] = {"$schema": 3}
 
     r = test_client.post(f"/v1/organisations/{organisation.id}/plugins", json=boefje)
     assert r.status_code == 400
 
 
 def test_update_boefje_schema(test_client, organisation):
     boefje = Boefje(id="test_plugin", name="My test boefje", description="123")
-    test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.json())
+    test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())
 
-    r = test_client.patch(f"/v1/organisations/{organisation.id}/boefjes/{boefje.id}", json={"schema": {"$schema": 3}})
+    r = test_client.patch(
+        f"/v1/organisations/{organisation.id}/boefjes/{boefje.id}", json={"boefje_schema": {"$schema": 3}}
+    )
     assert r.status_code == 400
 
     valid_schema = {
@@ -151,16 +153,20 @@ def test_update_boefje_schema(test_client, organisation):
         },
         "required": [],
     }
-    r = test_client.patch(f"/v1/organisations/{organisation.id}/boefjes/{boefje.id}", json={"schema": valid_schema})
+    r = test_client.patch(
+        f"/v1/organisations/{organisation.id}/boefjes/{boefje.id}", json={"boefje_schema": valid_schema}
+    )
     assert r.status_code == 204
 
     schema = test_client.get(f"/v1/organisations/{organisation.id}/plugins/{boefje.id}/schema.json").json()
     assert schema == valid_schema
 
     api_boefje = test_client.get(f"/v1/organisations/{organisation.id}/plugins/{boefje.id}").json()
-    assert api_boefje["schema"] == valid_schema
+    assert api_boefje["boefje_schema"] == valid_schema
 
-    r = test_client.patch(f"/v1/organisations/{organisation.id}/boefjes/dns-records", json={"schema": valid_schema})
+    r = test_client.patch(
+        f"/v1/organisations/{organisation.id}/boefjes/dns-records", json={"boefje_schema": valid_schema}
+    )
     assert r.status_code == 404
 
 
@@ -221,7 +227,7 @@ def test_clone_settings(test_client, organisation):
     # Add the second organisation
     new_org_id = "org2"
     org2 = Organisation(id=new_org_id, name="Second test Organisation")
-    test_client.post("/v1/organisations/", content=org2.json())
+    test_client.post("/v1/organisations/", content=org2.model_dump_json())
     test_client.put(f"/v1/organisations/{new_org_id}/{plug}/settings", json={"test_key": "second value"})
 
     # Show that the second organisation has no settings and dns-records is not enabled

diff --git a/boefjes/tests/plugins/test_leakix.py b/boefjes/tests/plugins/test_leakix.py
@@ -1,13 +1,12 @@
-from pydantic import parse_obj_as
+from pydantic import TypeAdapter
 
 from boefjes.plugins.kat_leakix.normalize import run
 from octopoes.models.types import OOIType
 from tests.loading import get_dummy_data
 
 
 def test_output():
-    input_ooi = parse_obj_as(
-        OOIType,
+    input_ooi = TypeAdapter(OOIType).validate_python(
         {
             "object_type": "HostnameHTTPURL",
             "network": "Network|internet",

diff --git a/bytes/bytes/api/models.py b/bytes/bytes/api/models.py
@@ -9,7 +9,7 @@ class RawResponse(BaseModel):
 
 class File(BaseModel):
     name: str
-    content: str = Field(..., contentEncoding="base64")
+    content: str = Field(json_schema_extra={"contentEncoding": "base64"})
     tags: list[str] = Field(default_factory=list)
 
 

diff --git a/bytes/bytes/api/router.py b/bytes/bytes/api/router.py
@@ -285,7 +285,7 @@ def get_raw_count_per_mime_type(
 
 def ignore_arguments_key(meta_repository: MetaDataRepository, query_filter: RawDataFilter):
     """Helper to not cache based on the stateful meta_repository, but only use the query parameters as a key."""
-    return query_filter.json()
+    return query_filter.model_dump_json()
 
 
 @cached(

diff --git a/bytes/bytes/database/sql_meta_repository.py b/bytes/bytes/database/sql_meta_repository.py
@@ -56,7 +56,7 @@ def get_boefje_meta_by_id(self, boefje_meta_id: uuid.UUID) -> BoefjeMeta:
         return to_boefje_meta(boefje_meta_in_db)
 
     def get_boefje_meta(self, query_filter: BoefjeMetaFilter) -> list[BoefjeMeta]:
-        logger.debug("Querying boefje meta: %s", query_filter.json())
+        logger.debug("Querying boefje meta: %s", query_filter.model_dump_json())
 
         query = self.session.query(BoefjeMetaInDB).filter(BoefjeMetaInDB.organization == query_filter.organization)
 
@@ -86,7 +86,7 @@ def get_normalizer_meta_by_id(self, normalizer_meta_id: uuid.UUID) -> Normalizer
         return to_normalizer_meta(normalizer_meta_in_db)
 
     def get_normalizer_meta(self, query_filter: NormalizerMetaFilter) -> list[NormalizerMeta]:
-        logger.debug("Querying normalizer meta: %s", query_filter.json())
+        logger.debug("Querying normalizer meta: %s", query_filter.model_dump_json())
 
         query = self.session.query(NormalizerMetaInDB)
 
@@ -134,7 +134,7 @@ def save_raw(self, raw: RawData) -> uuid.UUID:
         return raw_file_in_db.id
 
     def get_raw(self, query_filter: RawDataFilter) -> list[RawDataMeta]:
-        logger.debug("Querying raw data: %s", query_filter.json())
+        logger.debug("Querying raw data: %s", query_filter.model_dump_json())
         query = self.session.query(RawFileInDB)
         query = query_filter.apply(query)
 
@@ -177,7 +177,7 @@ def get_raw_file_count_per_organization(self) -> dict[str, int]:
         return {organization_id: count for organization_id, count in query}
 
     def get_raw_file_count_per_mime_type(self, query_filter: RawDataFilter) -> dict[str, int]:
-        logger.debug("Querying count raw data per mime type: %s", query_filter.json())
+        logger.debug("Querying count raw data per mime type: %s", query_filter.model_dump_json())
         query = self.session.query(func.unnest(RawFileInDB.mime_types), func.count()).group_by(
             func.unnest(RawFileInDB.mime_types)
         )