Fixed tmv1-query not working for activity

trendmicro · Feb 15, 2024 · 408459c · 408459c
1 parent b7eee2a
commit 408459c
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 6 deletions.
diff --git a/src/pytmv1/api/email.py b/src/pytmv1/api/email.py
@@ -116,7 +116,7 @@ def get_activity_count(
                 top,
                 SearchMode.COUNT_ONLY,
             ),
-            headers=utils.tmv1_query(op, fields),
+            headers=utils.tmv1_activity_query(op, fields),
         )
 
     def list_activity(
@@ -161,7 +161,7 @@ def list_activity(
                 top,
                 SearchMode.DEFAULT,
             ),
-            headers=utils.tmv1_query(op, fields),
+            headers=utils.tmv1_activity_query(op, fields),
         )
 
     def consume_activity(
@@ -210,5 +210,5 @@ def consume_activity(
                 top,
                 SearchMode.DEFAULT,
             ),
-            headers=utils.tmv1_query(op, fields),
+            headers=utils.tmv1_activity_query(op, fields),
         )
diff --git a/src/pytmv1/api/endpoint.py b/src/pytmv1/api/endpoint.py
@@ -113,7 +113,7 @@ def get_activity_count(
                 top,
                 SearchMode.COUNT_ONLY,
             ),
-            headers=utils.tmv1_query(op, fields),
+            headers=utils.tmv1_activity_query(op, fields),
         )
 
     def list_data(
@@ -176,7 +176,7 @@ def list_activity(
                 top,
                 SearchMode.DEFAULT,
             ),
-            headers=utils.tmv1_query(op, fields),
+            headers=utils.tmv1_activity_query(op, fields),
         )
 
     def consume_data(
@@ -249,5 +249,5 @@ def consume_activity(
                 top,
                 SearchMode.DEFAULT,
             ),
-            headers=utils.tmv1_query(op, fields),
+            headers=utils.tmv1_activity_query(op, fields),
         )
diff --git a/src/pytmv1/utils.py b/src/pytmv1/utils.py
@@ -23,6 +23,18 @@ def _build_query(
     )
 
 
+def _build_activity_query(
+    op: QueryOp, fields: Dict[str, str]
+) -> Dict[str, str]:
+    return filter_none(
+        {
+            "TMV1-Query": (" " + op + " ").join(
+                [f'{k}:"{v}"' for k, v in fields.items()]
+            )
+        }
+    )
+
+
 def _b64_encode(value: Optional[str]) -> Optional[str]:
     return base64.b64encode(value.encode()).decode() if value else None
 
@@ -104,5 +116,9 @@ def tmv1_query(op: QueryOp, fields: Dict[str, str]) -> Dict[str, str]:
     return _build_query(op, "TMV1-Query", fields)
 
 
+def tmv1_activity_query(op: QueryOp, fields: Dict[str, str]) -> Dict[str, str]:
+    return _build_activity_query(op, fields)
+
+
 def filter_query(op: QueryOp, fields: Dict[str, str]) -> Dict[str, str]:
     return _build_query(op, "filter", fields)
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
@@ -72,6 +72,13 @@ def test_build_query():
     assert result == {"TMV1-Query": "dpt eq '443' and src eq '1.1.1.1'"}
 
 
+def test__build_activity_query():
+    result = utils._build_activity_query(
+        QueryOp.AND, {"dpt": "443", "src": "1.1.1.1"}
+    )
+    assert result == {"TMV1-Query": 'dpt:"443" and src:"1.1.1.1"'}
+
+
 def test_filter_query():
     assert utils.filter_query(
         QueryOp.AND, {"fileName": "test.sh", "fileType": "bash"}