From b83db1da98713333cb956819041f4ef66d68460c Mon Sep 17 00:00:00 2001
From: Thomas Legros <Thomas_Legros@trendmicro.com>
Date: Thu, 15 Feb 2024 22:14:38 +0100
Subject: [PATCH] Fixed tmv1-query not working for activity

---
 src/pytmv1/api/email.py    |  6 +++---
 src/pytmv1/api/endpoint.py |  6 +++---
 src/pytmv1/utils.py        | 16 ++++++++++++++++
 tests/unit/test_utils.py   |  7 +++++++
 4 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/pytmv1/api/email.py b/src/pytmv1/api/email.py
index 2ac2f09..b744e62 100644
--- 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
index b5924ba..056defd 100644
--- 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
index f7cdd00..3ef18b0 100755
--- 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
index d4e1c6a..92e1068 100755
--- 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"}