From 7af76123273f464a971a44d39dc0de0e40e75b84 Mon Sep 17 00:00:00 2001
From: Thomas Legros <Thomas_Legros@trendmicro.com>
Date: Tue, 29 Oct 2024 15:02:09 +0100
Subject: [PATCH] Added support for object integrity level boundary

---
 src/pytmv1/model/common.py       | 17 +++++++++++++++++
 tests/integration/test_search.py |  5 +++++
 2 files changed, 22 insertions(+)

diff --git a/src/pytmv1/model/common.py b/src/pytmv1/model/common.py
index 29435c9..bd809f8 100644
--- a/src/pytmv1/model/common.py
+++ b/src/pytmv1/model/common.py
@@ -213,6 +213,23 @@ class EndpointActivity(BaseConsumable):
     tags: List[str] = Field(default=[])
     uuid: Optional[str] = None
 
+    @field_validator("object_integrity_level", mode="before")
+    @classmethod
+    def map_object_integrity_level(
+        cls, value: Optional[int]
+    ) -> Optional[IntegrityLevel]:
+        if value:
+            if IntegrityLevel.UNTRUSTED <= value < IntegrityLevel.LOW:
+                return IntegrityLevel.UNTRUSTED
+            if IntegrityLevel.LOW <= value < IntegrityLevel.MEDIUM:
+                return IntegrityLevel.LOW
+            if IntegrityLevel.MEDIUM <= value < IntegrityLevel.HIGH:
+                return IntegrityLevel.MEDIUM
+            if IntegrityLevel.HIGH <= value < IntegrityLevel.SYSTEM:
+                return IntegrityLevel.HIGH
+            return IntegrityLevel.SYSTEM
+        return None
+
 
 class HostInfo(BaseModel):
     name: str
diff --git a/tests/integration/test_search.py b/tests/integration/test_search.py
index cca3d19..7157efa 100755
--- a/tests/integration/test_search.py
+++ b/tests/integration/test_search.py
@@ -1,6 +1,7 @@
 from pytmv1 import (
     GetEmailActivitiesCountResp,
     GetEndpointActivitiesCountResp,
+    IntegrityLevel,
     ListEmailActivityResp,
     ListEndpointActivityResp,
     ListEndpointDataResp,
@@ -45,6 +46,10 @@ def test_list_endpoint_activities(client):
     assert result.result_code == ResultCode.SUCCESS
     assert isinstance(result.response, ListEndpointActivityResp)
     assert len(result.response.items) > 0
+    assert (
+        result.response.items[0].object_integrity_level
+        == IntegrityLevel.MEDIUM
+    )
 
 
 def test_get_endpoint_activities_count(client):