Refactored OAT package API

README.md

@@ -129,7 +129,7 @@ Supported APIs
 | `oat.get_pipeline`                             | [Get pipeline settings](https://automation.trendmicro.com/xdr/api-v3#tag/Observed-Attack-Techniques-Pipeline/paths/~1v3.0~1oat~1dataPipelines~1%7Bid%7D/get)                                              |
 | `oat.delete_pipelines`                         | [Unregister from data pipeline](https://automation.trendmicro.com/xdr/api-v3#tag/Observed-Attack-Techniques-Pipeline/paths/~1v3.0~1oat~1dataPipelines~1delete/post)                                       |
 | `oat.list_packages/consume_packages`           | [Get Observed Attack Techniques event packages](https://automation.trendmicro.com/xdr/api-v3#tag/Observed-Attack-Techniques-Pipeline/paths/~1v3.0~1oat~1dataPipelines~1%7Bid%7D~1packages/get)            |
-| `oat.download_package`                         | [Get Observed Attack Techniques package](https://automation.trendmicro.com/xdr/api-v3#tag/Observed-Attack-Techniques-Pipeline/paths/~1v3.0~1oat~1dataPipelines~1%7Bid%7D~1packages~1%7BpackageId%7D/get)  |
+| `oat.get_package`                              | [Get Observed Attack Techniques package](https://automation.trendmicro.com/xdr/api-v3#tag/Observed-Attack-Techniques-Pipeline/paths/~1v3.0~1oat~1dataPipelines~1%7Bid%7D~1packages~1%7BpackageId%7D/get)  |
 | **Sandbox Analysis**                           |                                                                                                                                                                                                           |
 | `sandbox.submit_file`                          | [Submit file to sandbox](https://automation.trendmicro.com/xdr/api-v3#tag/Sandbox-Analysis/paths/~1v3.0~1sandbox~1files~1analyze/post)                                                                    |
 | `sandbox.submit_url`                           | [Submit URLs to sandbox](https://automation.trendmicro.com/xdr/api-v3#tag/Sandbox-Analysis/paths/~1v3.0~1sandbox~1urls~1analyze/post)                                                                     |

src/pytmv1/__init__.py

@@ -47,8 +47,6 @@
     ApiStatus,
     DetectionType,
     EntityType,
-    EventID,
-    EventSubID,
     Iam,
     IntegrityLevel,
     InvestigationResult,
@@ -101,6 +99,7 @@
     GetApiKeyResp,
     GetEmailActivitiesCountResp,
     GetEndpointActivitiesCountResp,
+    GetOatPackageResp,
     GetPipelineResp,
     ListAlertNoteResp,
     ListAlertsResp,
@@ -168,15 +167,14 @@
     "Entity",
     "EntityType",
     "Error",
-    "EventID",
-    "EventSubID",
     "ExceptionObject",
     "ScriptType",
     "GetAlertResp",
     "GetAlertNoteResp",
     "GetApiKeyResp",
     "GetEmailActivitiesCountResp",
     "GetEndpointActivitiesCountResp",
+    "GetOatPackageResp",
     "GetPipelineResp",
     "HostInfo",
     "Iam",

src/pytmv1/api/oat.py

@@ -5,8 +5,8 @@
 from ..model.common import OatEvent, OatPackage
 from ..model.enum import Api, HttpMethod, OatRiskLevel, QueryOp
 from ..model.response import (
-    BytesResp,
     ConsumeLinkableResp,
+    GetOatPackageResp,
     GetPipelineResp,
     ListOatPackagesResp,
     ListOatPipelinesResp,
@@ -230,19 +230,19 @@ def list_pipelines(self) -> Result[ListOatPipelinesResp]:
         """
         return self._core.send(ListOatPipelinesResp, Api.LIST_OAT_PIPELINE)
 
-    def download_package(
+    def get_package(
         self, pipeline_id: str, package_id: str
-    ) -> Result[BytesResp]:
+    ) -> Result[GetOatPackageResp]:
         """Retrieves the specified Observed Attack Techniques package.
 
         :param pipeline_id: Pipeline ID.
         :type pipeline_id: str
         :param package_id: Package ID.
         :type package_id: str
-        :return: Result[BytesResp]
+        :return: Result[GetOatPackageResp]
         """
         return self._core.send(
-            BytesResp,
+            GetOatPackageResp,
             Api.DOWNLOAD_OAT_PACKAGE.value.format(pipeline_id, package_id),
         )
 

src/pytmv1/core.py

@@ -33,6 +33,7 @@
     GetAlertNoteResp,
     GetAlertResp,
     GetApiKeyResp,
+    GetOatPackageResp,
     GetPipelineResp,
     MultiApiKeyResp,
     MultiResp,
@@ -299,10 +300,14 @@ def _parse_data(raw_response: Response, class_: Type[R]) -> R:
         return class_(**raw_response.headers)
     if raw_response.status_code == 204 and class_ == NoContentResp:
         return class_()
-    if "json" in content_type:
+    if "text" in content_type and class_ == TextResp:
+        return class_.model_construct(text=raw_response.text)
+    if "json" in content_type or "text" in content_type:
         log.debug("Parsing json response [Class=%s]", class_.__name__)
         if class_ in [MultiResp, MultiUrlResp, MultiApiKeyResp]:
             return class_(items=raw_response.json())
+        if class_ == GetOatPackageResp:
+            return class_(package=raw_response.json())
         if class_ in [
             GetAlertResp,
             GetApiKeyResp,
@@ -325,8 +330,6 @@ def _parse_data(raw_response: Response, class_: Type[R]) -> R:
     ) and class_ == BytesResp:
         log.debug("Parsing binary response")
         return class_.model_construct(content=raw_response.content)
-    if "text" in content_type and class_ == TextResp:
-        return class_.model_construct(text=raw_response.text)
     raise ParseModelError(class_.__name__, raw_response)
 
 

src/pytmv1/model/common.py

@@ -13,8 +13,6 @@
     ApiStatus,
     DetectionType,
     EntityType,
-    EventID,
-    EventSubID,
     Iam,
     IntegrityLevel,
     InvestigationResult,
@@ -33,9 +31,8 @@
     Status,
 )
 
-CFG = ConfigDict(
-    alias_generator=to_camel, populate_by_name=True, protected_namespaces=()
-)
+CFG = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+CFG_NAMESPACE = ConfigDict(**CFG, protected_namespaces=())
 
 
 class BaseModel(PydanticBaseModel):
@@ -57,6 +54,8 @@ class Account(BaseModel):
 
 
 class Alert(BaseConsumable):
+    model_config = CFG_NAMESPACE
+
     id: str
     schema_version: str
     status: AlertStatus
@@ -151,9 +150,10 @@ def allow_empty_enum_string(cls, value: Any) -> Union[Any, None]:
 
 
 class EmailActivity(BaseConsumable):
+    event_source_type: Optional[str] = None
     mail_msg_subject: Optional[str] = None
     mail_msg_id: Optional[str] = None
-    msg_uuid: Optional[str] = None
+    msg_uuid: str
     mailbox: Optional[str] = None
     mail_sender_ip: Optional[str] = None
     mail_from_addresses: List[str] = Field(default=[])
@@ -171,11 +171,12 @@ class EmailActivity(BaseConsumable):
 class EndpointActivity(BaseConsumable):
     dpt: Optional[int] = None
     dst: Optional[str] = None
-    endpoint_guid: Optional[str] = None
+    endpoint_guid: str
     endpoint_host_name: Optional[str] = None
     endpoint_ip: List[str] = Field(default=[])
-    event_id: Optional[EventID] = None
-    event_sub_id: Optional[EventSubID] = None
+    event_id: Optional[str] = None
+    event_source_type: Optional[str] = None
+    event_sub_id: Optional[str] = None
     object_integrity_level: Optional[IntegrityLevel] = None
     object_true_type: Optional[int] = None
     object_sub_true_type: Optional[int] = None
@@ -213,6 +214,15 @@ class EndpointActivity(BaseConsumable):
     tags: List[str] = Field(default=[])
     uuid: Optional[str] = None
 
+    @field_validator("event_sub_id", mode="before")
+    @classmethod
+    def map_event_sub_id(
+        cls, value: Optional[Union[int, str]]
+    ) -> Optional[str]:
+        if value is not None and isinstance(value, int):
+            return str(value)
+        return value
+
     @field_validator("object_integrity_level", mode="before")
     @classmethod
     def map_object_integrity_level(
@@ -375,36 +385,45 @@ def values(self) -> List[int]:
 
 
 class OatEndpoint(BaseModel):
-    endpoint_name: str
-    agent_guid: str
+    endpoint_name: Optional[str] = None
+    name: Optional[str] = None
+    agent_guid: Optional[str] = None
+    guid: Optional[str] = None
     ips: List[str]
 
 
 class OatObject(BaseModel):
-    type: str
     field: str
+    type: str
     value: Union[int, str, List[str]]
+    master: Optional[bool] = None
+    risk_level: Optional[OatRiskLevel] = None
 
 
 class OatFilter(BaseModel):
     id: str
+    unique_id: Optional[str] = None
     name: str
     description: Optional[str] = None
-    mitre_tactic_ids: List[str]
-    mitre_technique_ids: List[str]
+    mitre_tactic_ids: Optional[List[str]] = None
+    tactics: Optional[List[str]] = None
+    mitre_technique_ids: Optional[List[str]] = None
+    techniques: Optional[List[str]] = None
     highlighted_objects: List[OatObject]
-    risk_level: OatRiskLevel
-    type: str
+    risk_level: Optional[OatRiskLevel] = None
+    level: Optional[OatRiskLevel] = None
+    type: DetectionType
 
 
 class OatEvent(BaseConsumable):
-    source: OatDataSource
-    uuid: str
+    source: Optional[OatDataSource] = None
+    uuid: Optional[str] = None
     filters: List[OatFilter]
     endpoint: Optional[OatEndpoint] = None
     entity_type: OatEntityType
     entity_name: str
-    detected_date_time: str
+    detected_date_time: Optional[str] = None
+    detection_time: Optional[str] = None
     ingested_date_time: Optional[str] = None
     detail: Union[EndpointActivity, EmailActivity]
 

src/pytmv1/model/enum.py

@@ -128,71 +128,6 @@ class EntityType(str, Enum):
     AWS_LAMBDA = "awsLambda"
 
 
-class EventID(str, Enum):
-    EVENT_PROCESS = "1"
-    EVENT_FILE = "2"
-    EVENT_CONNECTIO = "3"
-    EVENT_DNS = "4"
-    EVENT_REGISTRY = "5"
-    EVENT_ACCOUNT = "6"
-    EVENT_INTERNET = "7"
-    XDR_EVENT_MODIFIED_PROCESS = "8"
-    EVENT_WINDOWS_HOOK = "9"
-    EVENT_WINDOWS_EVENT = "10"
-    EVENT_AMSI = "11"
-    EVENT_WMI = "12"
-    TELEMETRY_MEMORY = "13"
-    TELEMETRY_BM = "14"
-
-
-class EventSubID(int, Enum):
-    TELEMETRY_NONE = 0
-    XDR_PROCESS_OPEN = 1
-    XDR_PROCESS_CREATE = 2
-    XDR_PROCESS_TERMINATE = 3
-    XDR_PROCESS_LOAD_IMAGE = 4
-    TELEMETRY_PROCESS_EXECUTE = 5
-    TELEMETRY_PROCESS_CONNECT = 6
-    TELEMETRY_PROCESS_TRACME = 7
-    XDR_FILE_CREATE = 101
-    XDR_FILE_OPEN = 102
-    XDR_FILE_DELETE = 103
-    XDR_FILE_SET_SECURITY = 104
-    XDR_FILE_COPY = 105
-    XDR_FILE_MOVE = 106
-    XDR_FILE_CLOSE = 107
-    TELEMETRY_FILE_MODIFY_TIMESTAMP = 108
-    TELEMETRY_FILE_MODIFY = 109
-    XDR_CONNECTION_CONNECT = 201
-    XDR_CONNECTION_LISTEN = 202
-    XDR_CONNECTION_CONNECT_INBOUND = 203
-    XDR_CONNECTION_CONNECT_OUTBOUND = 204
-    XDR_DNS_QUERY = 301
-    XDR_REGISTRY_CREATE = 401
-    XDR_REGISTRY_SET = 402
-    XDR_REGISTRY_DELETE = 403
-    XDR_REGISTRY_RENAME = 404
-    XDR_ACCOUNT_ADD = 501
-    XDR_ACCOUNT_DELETE = 502
-    XDR_ACCOUNT_IMPERSONATE = 503
-    XDR_ACCOUNT_MODIFY = 504
-    XDR_INTERNET_OPEN = 601
-    XDR_INTERNET_CONNECT = 602
-    XDR_INTERNET_DOWNLOAD = 603
-    XDR_MODIFIED_PROCESS_CREATE_REMOTETHREAD = 701
-    XDR_MODIFIED_PROCESS_WRITE_MEMORY = 702
-    TELEMETRY_MODIFIED_PROCESS_WRITE_PROCESS = 703
-    TELEMETRY_MODIFIED_PROCESS_READ_PROCESS = 704
-    TELEMETRY_MODIFIED_PROCESS_WRITE_PROCESS_NAME = 705
-    XDR_WINDOWS_HOOK_SET = 801
-    XDR_AMSI_EXECUTE = 901
-    TELEMETRY_MEMORY_MODIFY = 1001
-    TELEMETRY_MEMORY_MODIFY_PERMISSION = 1002
-    TELEMETRY_MEMORY_READ = 1003
-    TELEMETRY_BM_INVOKE = 1101
-    TELEMETRY_BM_INVOKE_API = 1102
-
-
 class ScriptType(str, Enum):
     POWERSHELL = "powershell"
     BASH = "bash"

src/pytmv1/model/response.py

@@ -169,6 +169,10 @@ class GetApiKeyResp(BaseResponse):
     etag: str
 
 
+class GetOatPackageResp(BaseResponse):
+    package: OatEvent
+
+
 class GetPipelineResp(BaseResponse):
     data: OatPipeline
     etag: str

tests/integration/test_oat.py

@@ -1,6 +1,6 @@
 from pytmv1 import (
-    BytesResp,
     EndpointActivity,
+    GetOatPackageResp,
     NoContentResp,
     OatRiskLevel,
     ResultCode,
@@ -95,14 +95,28 @@ def test_list_pipelines(client):
     assert result.response.items[0].description == "siemhost1"
 
 
-def test_download_package(client):
-    result = client.oat.download_package(
+def test_get_package(client):
+    result = client.oat.get_package(
         "83df1ed3-84e7-4e6d-98b5-d79468cccba1",
         "2024073012-774c3fb6-f777-4ce1-8564-39885e7d41a4",
     )
-    assert isinstance(result.response, BytesResp)
+    assert isinstance(result.response, GetOatPackageResp)
     assert result.result_code == ResultCode.SUCCESS
-    assert result.response.content
+    assert result.response.package.detection_time == "2024-07-30T12:01:01Z"
+    assert (
+        result.response.package.endpoint.guid
+        == "ab673395-9bf9-49fc-b8ac-7fa15467d20a"
+    )
+    assert result.response.package.filters[0].tactics == ["TA0002"]
+    assert result.response.package.filters[0].techniques == ["T1059.004"]
+    assert (
+        result.response.package.detail.event_source_type
+        == "EVENT_SOURCE_TELEMETRY"
+    )
+    assert (
+        result.response.package.detail.uuid
+        == "eb5b2977-3bde-45ea-b493-05984bab5d0f"
+    )
 
 
 def test_list_packages(client):

tests/unit/test_core.py

@@ -164,9 +164,9 @@ def test_parse_data_with_bytes():
 
 def test_parse_data_with_html_is_failed():
     raw_response = Response()
-    raw_response.headers = {"Content-Type": "text/html"}
+    raw_response.status_code = 204
     with pytest.raises(ParseModelError):
-        core_m._parse_data(raw_response, NoContentResp)
+        core_m._parse_data(raw_response, AddAlertNoteResp)
 
 
 def test_parse_data_with_json():