Merge branch 'releases/0.8.1'

src/pytmv1/__about__.py

@@ -1 +1 @@
-__version__ = "0.8.0"
+__version__ = "0.8.1"

src/pytmv1/__init__.py

@@ -40,7 +40,6 @@
     EntityType,
     EventID,
     EventSubID,
-    FileType,
     Iam,
     IntegrityLevel,
     InvestigationStatus,
@@ -54,6 +53,7 @@
     SandboxAction,
     SandboxObjectType,
     ScanAction,
+    ScriptType,
     Severity,
     Status,
 )
@@ -152,7 +152,7 @@
     "EventID",
     "EventSubID",
     "ExceptionObject",
-    "FileType",
+    "ScriptType",
     "GetAlertResp",
     "GetAlertNoteResp",
     "GetApiKeyResp",

src/pytmv1/api/__init__.py

@@ -1,12 +1,12 @@
 from .account import Account
 from .alert import Alert
 from .api_key import ApiKey
-from .custom_script import CustomScript
 from .email import Email
 from .endpoint import Endpoint
 from .note import Note
 from .object import Object
 from .sandbox import Sandbox
+from .script import CustomScript
 from .system import System
 from .task import Task
 

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),
         )

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),
         )

src/pytmv1/api/note.py

@@ -20,20 +20,22 @@ class Note:
     def __init__(self, core: Core):
         self._core = core
 
-    def add(self, alert_id: str, note: str) -> Result[AddAlertNoteResp]:
+    def create(
+        self, alert_id: str, note_content: str
+    ) -> Result[AddAlertNoteResp]:
         """Adds a note to the specified Workbench alert.
 
         :param alert_id: Workbench alert id.
         :type alert_id: str
-        :param note: Value of the note.
-        :type note: str
+        :param note_content: Value of the note.
+        :type note_content: str
         :rtype: Result[AddAlertNoteResp]:
         """
         return self._core.send(
             AddAlertNoteResp,
             Api.ADD_ALERT_NOTE.value.format(alert_id),
             HttpMethod.POST,
-            json={"content": note},
+            json={"content": note_content},
         )
 
     def update(

src/pytmv1/api/custom_script.py → src/pytmv1/api/script.py

@@ -3,7 +3,7 @@
 from .. import utils
 from ..core import Core
 from ..model.common import Script
-from ..model.enum import Api, FileType, HttpMethod, QueryOp
+from ..model.enum import Api, HttpMethod, QueryOp, ScriptType
 from ..model.request import CustomScriptRequest
 from ..model.response import (
     AddCustomScriptResp,
@@ -22,22 +22,22 @@ class CustomScript:
     def __init__(self, core: Core):
         self._core = core
 
-    def add(
+    def create(
         self,
-        file_type: FileType,
-        file_name: str,
-        file_content: str,
+        script_type: ScriptType,
+        script_name: str,
+        script_content: str,
         description: Optional[str] = None,
     ) -> Result[AddCustomScriptResp]:
         """
         Uploads a custom script. Supported file extensions: .ps1, .sh.
         Note: Custom scripts must use UTF-8 encoding.
-        :param file_type: File type.
-        :type file_type: FileType
-        :param file_name: File name.
-        :type file_name: str
-        :param file_content: Plain text content of the file.
-        :type file_content: str
+        :param script_type: File type.
+        :type script_type: ScriptType
+        :param script_name: File name.
+        :type script_name: str
+        :param script_content: Plain text content of the script.
+        :type script_content: str
         :param description: Description.
         :type description: Optional[str]
         :return: Result[AddACustomScriptResp]
@@ -47,32 +47,36 @@ def add(
             Api.ADD_CUSTOM_SCRIPT,
             HttpMethod.POST,
             data=utils.filter_none(
-                {"fileType": file_type.value, "description": description}
+                {"fileType": script_type.value, "description": description}
             ),
             files={
-                "file": (file_name, bytes(file_content, "utf-8"), "text/plain")
+                "file": (
+                    script_name,
+                    bytes(script_content, "utf-8"),
+                    "text/plain",
+                )
             },
         )
 
     def update(
         self,
         script_id: str,
-        file_type: FileType,
-        file_name: str,
-        file_content: str,
+        script_type: ScriptType,
+        script_name: str,
+        script_content: str,
         description: Optional[str] = None,
     ) -> Result[NoContentResp]:
         """
         Updates a custom script. Supported file extensions: .ps1, .sh.
         Note: Custom scripts must use UTF-8 encoding.
         :param script_id: Unique string that identifies a script file.
         :type script_id: str
-        :param file_type: File type.
-        :type file_type: FileType
-        :param file_name: File name.
-        :type file_name: str
-        :param file_content: Plain text content of the file.
-        :type file_content: str
+        :param script_type: File type.
+        :type script_type: ScriptType
+        :param script_name: File name.
+        :type script_name: str
+        :param script_content: Plain text content of the file.
+        :type script_content: str
         :param description: Description.
         :type description: Optional[str]
         :return: Result[NoContentResp]
@@ -82,10 +86,14 @@ def update(
             Api.UPDATE_CUSTOM_SCRIPT.value.format(script_id),
             HttpMethod.POST,
             data=utils.filter_none(
-                {"fileType": file_type.value, "description": description}
+                {"fileType": script_type.value, "description": description}
             ),
             files={
-                "file": (file_name, bytes(file_content, "utf-8"), "text/plain")
+                "file": (
+                    script_name,
+                    bytes(script_content, "utf-8"),
+                    "text/plain",
+                )
             },
         )
 

src/pytmv1/model/common.py

@@ -13,7 +13,6 @@
     EntityType,
     EventID,
     EventSubID,
-    FileType,
     Iam,
     IntegrityLevel,
     InvestigationStatus,
@@ -23,6 +22,7 @@
     Provider,
     RiskLevel,
     ScanAction,
+    ScriptType,
     Severity,
     Status,
 )
@@ -356,7 +356,7 @@ def map_data(cls, data: Dict[str, str]) -> Dict[str, str]:
 class Script(BaseConsumable):
     id: str
     file_name: str
-    file_type: FileType
+    file_type: ScriptType
     description: Optional[str] = None
 
 

src/pytmv1/model/enum.py

@@ -171,7 +171,7 @@ class EventSubID(int, Enum):
     TELEMETRY_BM_INVOKE_API = 1102
 
 
-class FileType(str, Enum):
+class ScriptType(str, Enum):
     POWERSHELL = "powershell"
     BASH = "bash"
 

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)

tests/integration/test_network.py

@@ -21,7 +21,7 @@ def test_conn_opened_with_multi_call_single_client_is_one(
 
 
 def test_conn_opened_with_multi_processing_single_client_is_one(client):
-    threads = thread_list(lambda: client.note.add("1", "dummy note"))
+    threads = thread_list(lambda: client.note.create("1", "dummy note"))
     for t in threads:
         t.start()
     for t in threads:

tests/integration/test_script.py

@@ -2,20 +2,20 @@
     AddCustomScriptResp,
     ConsumeLinkableResp,
     CustomScriptRequest,
-    FileType,
     ListCustomScriptsResp,
     MultiResp,
     NoContentResp,
     ResultCode,
+    ScriptType,
     TextResp,
 )
 
 
 def test_add_custom_script(client):
-    result = client.script.add(
-        file_type=FileType.BASH,
-        file_name="add_script.sh",
-        file_content="#!/bin/sh\necho 'Add script'",
+    result = client.script.create(
+        script_type=ScriptType.BASH,
+        script_name="add_script.sh",
+        script_content="#!/bin/sh\necho 'Add script'",
     )
     assert isinstance(result.response, AddCustomScriptResp)
     assert result.result_code == ResultCode.SUCCESS
@@ -48,9 +48,9 @@ def test_run_custom_scripts(client):
 def test_update_custom_script(client):
     result = client.script.update(
         script_id="123",
-        file_type=FileType.BASH,
-        file_name="update_script.sh",
-        file_content="#!/bin/sh Update script",
+        script_type=ScriptType.BASH,
+        script_name="update_script.sh",
+        script_content="#!/bin/sh Update script",
     )
     assert isinstance(result.response, NoContentResp)
     assert result.result_code == ResultCode.SUCCESS

tests/integration/test_workbench.py

@@ -11,8 +11,8 @@
 )
 
 
-def test_add_note(client):
-    result = client.note.add("1", "dummy note")
+def test_create_note(client):
+    result = client.note.create("1", "dummy note")
     assert isinstance(result.response, AddAlertNoteResp)
     assert result.result_code == ResultCode.SUCCESS
     assert result.response.note_id.isdigit()

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"}