Merge pull request #20 from trendmicro/fix/tmv1_activity_query

Fixed tmv1-query not working for activity
trendmicro · Feb 15, 2024 · 80ce09b · 80ce09b
2 parents b7eee2a + b83db1d
commit 80ce09b
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"}