Merge pull request #28 from trendmicro/fix/mypy_error

Fixed mypy generic type error

src/pytmv1/__init__.py

@@ -64,6 +64,7 @@
     ScriptType,
     Severity,
     Status,
+    TaskAction,
 )
 from .model.request import (
     AccountRequest,
@@ -114,7 +115,6 @@
     SandboxSubmissionStatusResp,
     SandboxSubmitUrlTaskResp,
     SubmitFileToSandboxResp,
-    TaskAction,
     TerminateProcessTaskResp,
     TextResp,
 )

src/pytmv1/core.py

@@ -10,6 +10,7 @@
 from pydantic import AnyHttpUrl, TypeAdapter
 from requests import PreparedRequest, Request, Response
 
+from . import utils
 from .__about__ import __version__
 from .adapter import HTTPAdapter
 from .exception import (
@@ -20,7 +21,7 @@
     ServerTextError,
 )
 from .model.common import Error, MsError, MsStatus
-from .model.enum import Api, HttpMethod, Status
+from .model.enum import Api, HttpMethod, Status, TaskAction
 from .model.request import EndpointRequest
 from .model.response import (
     MR,
@@ -40,7 +41,6 @@
     S,
     SandboxSubmissionStatusResp,
     T,
-    TaskAction,
     TextResp,
 )
 from .result import multi_result, result
@@ -304,8 +304,10 @@ def _parse_data(raw_response: Response, class_: Type[R]) -> R:
                 etag=raw_response.headers.get("ETag", ""),
             )
         if class_ == BaseTaskResp:
-            resp_class = task_action(raw_response.json()["action"]).class_
-            class_ = resp_class if resp_class else class_
+            resp_class: Type[BaseTaskResp] = utils.task_action_resp_class(
+                TaskAction(raw_response.json()["action"])
+            )
+            class_ = resp_class if issubclass(resp_class, class_) else class_
         return class_(**raw_response.json())
     if "application" in content_type and class_ == BytesResp:
         log.debug("Parsing binary response")
@@ -319,10 +321,6 @@ def _parse_data(raw_response: Response, class_: Type[R]) -> R:
     raise ParseModelError(class_.__name__, raw_response)
 
 
-def task_action(action_name: str) -> TaskAction:
-    return next(filter(lambda ta: action_name == ta.action, TaskAction))
-
-
 def _parse_html(html: str) -> str:
     log.info("Parsing html response [Html=%s]", html)
     soup = BeautifulSoup(html, "html.parser")

src/pytmv1/model/enum.py

@@ -299,3 +299,31 @@ class Status(str, Enum):
     RUNNING = "running"
     SUCCEEDED = "succeeded"
     WAIT_FOR_APPROVAL = "waitForApproval"
+
+
+class TaskAction(str, Enum):
+    COLLECT_FILE = "collectFile"
+    COLLECT_EVIDENCE = "collectEvidence"
+    COLLECT_NETWORK_ANALYSIS_PACKAGE = "collectNetworkAnalysisPackage"
+    ISOLATE_ENDPOINT = "isolate"
+    ISOLATE_ENDPOINT_MULTIPLE = "isolateForMultiple"
+    RESTORE_ENDPOINT = "restoreIsolate"
+    RESTORE_ENDPOINT_MULTIPLE = "restoreIsolateForMultiple"
+    TERMINATE_PROCESS = "terminateProcess"
+    DUMP_PROCESS_MEMORY = "dumpProcessMemory"
+    QUARANTINE_MESSAGE = "quarantineMessage"
+    DELETE_MESSAGE = "deleteMessage"
+    RESTORE_MESSAGE = "restoreMessage"
+    BLOCK_SUSPICIOUS = "block"
+    REMOVE_SUSPICIOUS = "restoreBlock"
+    RESET_PASSWORD = "resetPassword"
+    SUBMIT_SANDBOX = "submitSandbox"
+    ENABLE_ACCOUNT = "enableAccount"
+    DISABLE_ACCOUNT = "disableAccount"
+    FORCE_SIGN_OUT = "forceSignOut"
+    REMOTE_SHELL = "remoteShell"
+    RUN_INVESTIGATION_KIT = "runInvestigationKit"
+    RUN_CUSTOM_SCRIPT = "runCustomScript"
+    RUN_CUSTOM_SCRIPT_MULTIPLE = "runCustomScriptForMultiple"
+    RUN_OS_QUERY = "runOsquery"
+    RUN_YARA_RULES = "runYaraRules"

src/pytmv1/model/response.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-from enum import Enum
-from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
+from typing import Any, Dict, Generic, List, Optional, TypeVar, Union
 
 from pydantic import Field, model_validator
 
@@ -282,35 +281,3 @@ class TerminateProcessTaskResp(BaseTaskResp):
 
 class TextResp(BaseResponse):
     text: str
-
-
-class TaskAction(Enum):
-    COLLECT_FILE = ("collectFile", CollectFileTaskResp)
-    COLLECT_EVIDENCE = ("collectEvidence", None)
-    COLLECT_NETWORK_ANALYSIS_PACKAGE = ("collectNetworkAnalysisPackage", None)
-    ISOLATE_ENDPOINT = ("isolate", EndpointTaskResp)
-    ISOLATE_ENDPOINT_MULTIPLE = ("isolateForMultiple", None)
-    RESTORE_ENDPOINT = ("restoreIsolate", EndpointTaskResp)
-    RESTORE_ENDPOINT_MULTIPLE = ("restoreIsolateForMultiple", None)
-    TERMINATE_PROCESS = ("terminateProcess", TerminateProcessTaskResp)
-    DUMP_PROCESS_MEMORY = ("dumpProcessMemory", None)
-    QUARANTINE_MESSAGE = ("quarantineMessage", EmailMessageTaskResp)
-    DELETE_MESSAGE = ("deleteMessage", EmailMessageTaskResp)
-    RESTORE_MESSAGE = ("restoreMessage", EmailMessageTaskResp)
-    BLOCK_SUSPICIOUS = ("block", BlockListTaskResp)
-    REMOVE_SUSPICIOUS = ("restoreBlock", BlockListTaskResp)
-    RESET_PASSWORD = ("resetPassword", AccountTaskResp)
-    SUBMIT_SANDBOX = ("submitSandbox", SandboxSubmitUrlTaskResp)
-    ENABLE_ACCOUNT = ("enableAccount", AccountTaskResp)
-    DISABLE_ACCOUNT = ("disableAccount", AccountTaskResp)
-    FORCE_SIGN_OUT = ("forceSignOut", AccountTaskResp)
-    REMOTE_SHELL = ("remoteShell", None)
-    RUN_INVESTIGATION_KIT = ("runInvestigationKit", None)
-    RUN_CUSTOM_SCRIPT = ("runCustomScript", CustomScriptTaskResp)
-    RUN_CUSTOM_SCRIPT_MULTIPLE = ("runCustomScriptForMultiple", None)
-    RUN_OS_QUERY = ("runOsquery", None)
-    RUN_YARA_RULES = ("runYaraRules", None)
-
-    def __init__(self, action: str, class_: Optional[Type[T]]):
-        self.action = action
-        self.class_ = class_

src/pytmv1/utils.py

@@ -1,15 +1,44 @@
 import base64
 import re
-from typing import Any, Dict, List, Optional, Pattern
+from typing import Any, Dict, List, Optional, Pattern, Type
 
-from .model.enum import QueryOp, SearchMode
+from .model.enum import QueryOp, SearchMode, TaskAction
 from .model.request import ObjectRequest, SuspiciousObjectRequest
+from .model.response import (
+    AccountTaskResp,
+    BaseTaskResp,
+    BlockListTaskResp,
+    CollectFileTaskResp,
+    CustomScriptTaskResp,
+    EmailMessageTaskResp,
+    EndpointTaskResp,
+    SandboxSubmitUrlTaskResp,
+    TerminateProcessTaskResp,
+)
 
 MAC_ADDRESS_PATTERN: Pattern[str] = re.compile(
     "^([0-9A-Fa-f]{2}[:-]){5}([0-9A-Fa-f]{2})$"
 )
 GUID_PATTERN: Pattern[str] = re.compile("^(\\w+-+){1,5}\\w+$")
 
+TASK_ACTION_MAP: Dict[TaskAction, Type[BaseTaskResp]] = {
+    TaskAction.COLLECT_FILE: CollectFileTaskResp,
+    TaskAction.ISOLATE_ENDPOINT: EndpointTaskResp,
+    TaskAction.RESTORE_ENDPOINT: EndpointTaskResp,
+    TaskAction.TERMINATE_PROCESS: TerminateProcessTaskResp,
+    TaskAction.QUARANTINE_MESSAGE: EmailMessageTaskResp,
+    TaskAction.DELETE_MESSAGE: EmailMessageTaskResp,
+    TaskAction.RESTORE_MESSAGE: EmailMessageTaskResp,
+    TaskAction.BLOCK_SUSPICIOUS: BlockListTaskResp,
+    TaskAction.REMOVE_SUSPICIOUS: BlockListTaskResp,
+    TaskAction.RESET_PASSWORD: AccountTaskResp,
+    TaskAction.SUBMIT_SANDBOX: SandboxSubmitUrlTaskResp,
+    TaskAction.ENABLE_ACCOUNT: AccountTaskResp,
+    TaskAction.DISABLE_ACCOUNT: AccountTaskResp,
+    TaskAction.FORCE_SIGN_OUT: AccountTaskResp,
+    TaskAction.RUN_CUSTOM_SCRIPT: CustomScriptTaskResp,
+}
+
 
 def _build_query(
     op: QueryOp, header: str, fields: Dict[str, str]
@@ -122,3 +151,9 @@ def tmv1_activity_query(op: QueryOp, fields: Dict[str, str]) -> Dict[str, str]:
 
 def filter_query(op: QueryOp, fields: Dict[str, str]) -> Dict[str, str]:
     return _build_query(op, "filter", fields)
+
+
+def task_action_resp_class(
+    task_action: TaskAction,
+) -> Type[BaseTaskResp]:
+    return TASK_ACTION_MAP.get(task_action, BaseTaskResp)