Syakima/CIAC-15740/forcepoint-add-refresh-command

Packs/ForcepointSecurityManagementCenter/Integrations/ForcepointSecurityManagementCenter/ForcepointSecurityManagementCenter.py

@@ -3,12 +3,13 @@
 from CommonServerPython import *  # noqa # pylint: disable=unused-wildcard-import
 from CommonServerUserPython import *  # noqa
 from smc import session
+from smc.administration.tasks import Task
 from smc.elements.network import IPList, DomainName, Host
 from smc.core.engine import Engine
 from smc.base.model import Element
 from smc.policy.layer3 import FirewallTemplatePolicy, FirewallPolicy
 from smc.policy.rule import IPv6Rule, Rule
-from smc.api.exceptions import ElementNotFound
+from smc.api.exceptions import ElementNotFound, TaskRunFailed
 import urllib3
 from typing import Any
 
@@ -20,6 +21,9 @@
 DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ"  # ISO8601 format with UTC, default in XSOAR
 DEFAULT_LIMIT = 50
 API_VERSION = "6.10"
+
+DEFAULT_POLLING_INTERVAL_IN_SECONDS = 30
+DEFAULT_POLLING_TIMEOUT_IN_SECONDS = 600
 """ CLIENT CLASS """
 
 
@@ -819,6 +823,86 @@ def list_engine_command(args: dict[str, Any]) -> CommandResults:
     )
 
 
+@polling_function(
+    "forcepoint-smc-engine-refresh",
+    poll_message="Refreshing engines:",
+    requires_polling_arg=False,
+    interval=arg_to_number(demisto.args().get("interval_in_seconds")) or DEFAULT_POLLING_INTERVAL_IN_SECONDS,
+    timeout=arg_to_number(demisto.args().get("timeout_in_seconds")) or DEFAULT_POLLING_TIMEOUT_IN_SECONDS,
+)
+def engine_refresh_command(args: dict) -> PollResult:
+    """
+    Refresh the specified engines.
+
+    This command initiates refresh tasks and polls for status until complete.
+
+    Args:
+        args: Command arguments including engine names.
+
+    Returns:
+        PollResult: Contains the refresh status or indicates to continue polling.
+    """
+    task_ids = argToList(args.get("task_ids", []))
+    engine_names = argToList(args.get("engine_name", []))
+
+    if not task_ids:
+        # first call - initiate refresh tasks
+        failed_engines = []
+
+        demisto.debug("Creating refresh tasks:")
+        for engine_name in engine_names:
+            try:
+                engine = Engine(engine_name)
+                poller = engine.refresh()
+                task_ids.append(poller.task.href)
+                demisto.debug(f"\t- (Engine: {engine_name}, Task: {poller.task.href})")
+            except TaskRunFailed:
+                failed_engines.append(engine_name)
+
+        if failed_engines:
+            raise DemistoException(
+                f"Failed to initiate refresh tasks for engines ({', '.join(failed_engines)}), refresh policy is locked."
+            )
+
+        demisto.debug("Completed task initiation.")
+        demisto.debug(f"Initiated refresh tasks for engines {', '.join(engine_names)}.")
+        demisto.debug(f"Tasks IDs: {task_ids}")
+        args["task_ids"] = task_ids
+
+    # poll for refresh status
+    completed_tasks = []
+    pending_tasks = []
+
+    demisto.debug("Restoring refresh tasks, assignments are expected to be the same as during creation:")
+    for task_id, engine_name in zip(task_ids, engine_names):
+        task = Task({"follower": task_id})
+        demisto.debug(f"\t- (Engine: {engine_name}, Task: {task_id})")
+        task = task.update_status()
+        if not task.in_progress:
+            completed_tasks.append(engine_name)
+        else:
+            pending_tasks.append((task_id, engine_name))
+
+    demisto.debug(f"Refresh status:\n{pending_tasks=}\n{completed_tasks=}")
+
+    if not pending_tasks:
+        demisto.debug("All refresh tasks are complete.")
+        return PollResult(
+            response=CommandResults(
+                readable_output=f"Existing policy on specified devices ({', '.join(engine_names)}) were refreshed successfully."
+            ),
+            continue_to_poll=False,
+        )
+
+    demisto.info(f"Refresh tasks {pending_tasks} are still in progress, continuing to poll...")
+
+    return PollResult(
+        response=CommandResults(readable_output="Refresh tasks are still in progress, continuing to poll..."),
+        continue_to_poll=True,
+        args_for_next_run=args,
+    )
+
+
 """ MAIN FUNCTION """
 
 
@@ -881,6 +965,8 @@ def main():
             return_results(delete_rule_command(demisto.args()))
         elif command == "forcepoint-smc-engine-list":
             return_results(list_engine_command(demisto.args()))
+        elif command == "forcepoint-smc-engine-refresh":
+            return_results(engine_refresh_command(demisto.args()))
     # Log exceptions and return errors
     except Exception as e:
         return_error(f"Failed to execute {command} command.\nError:\n{str(e)}")

Packs/ForcepointSecurityManagementCenter/Integrations/ForcepointSecurityManagementCenter/ForcepointSecurityManagementCenter.yml

@@ -692,6 +692,25 @@ script:
     - contextPath: ForcepointSMC.Engine.Comment
       description: The comment for the engine.
       type: String
+  - arguments:
+    - description: List of engine names to refresh.
+      isArray: true
+      name: engine_name
+      required: true
+    - description: List of polling refresh task ids.
+      isArray: true
+      name: task_id
+      required: false
+      hidden: true
+    - name: interval_in_seconds
+      description: Interval between polling attempts in seconds. To prevent search timeouts, set this value within the 60-90 second range.
+      defaultValue: '30'
+    - name: timeout_in_seconds
+      description: Timeout for polling in seconds.
+      defaultValue: '600'
+    description: Refreshes the specified engines. Use forcepoint-smc-engine-list command to list the engines in the system.
+    name: forcepoint-smc-engine-refresh
+    polling: true
   isfetch: false
   runonce: false
   script: '-'

Packs/ForcepointSecurityManagementCenter/Integrations/ForcepointSecurityManagementCenter/ForcepointSecurityManagementCenter_test.py

@@ -20,13 +20,14 @@
     list_rule_command,
     delete_rule_command,
     create_rule_command,
+    engine_refresh_command,
     IPList,
     Host,
     DomainName,
     FirewallPolicy,
     DemistoException,
 )
-from smc.api.exceptions import ElementNotFound
+from smc.api.exceptions import ElementNotFound, TaskRunFailed
 from smc.base.collection import CollectionManager
 
 
@@ -112,15 +113,32 @@ def save(self):
         return
 
 
+class mock_Task:
+    def __init__(self, href: str, in_progress: bool = True):
+        self.href = href
+        self.in_progress = in_progress
+
+    def update_status(self):
+        return self
+
+
+class mock_Poller:
+    def __init__(self, task: mock_Task):
+        self.task = task
+
+
 class mock_Engine:
-    def __init__(self, name: str, comment: str):
+    def __init__(self, name: str, comment: str = ""):
         self.name = name
         self.comment = comment
         self.pending_changes = mock_Changes()
 
     def objects(self):
         return mock_GenericRuleEntity()
 
+    def refresh(self):
+        return mock_Poller(mock_Task(href=self.name))
+
 
 class mock_Changes:
     def __init__(self):
@@ -636,3 +654,89 @@ def test_update_rule_command(mocker):
     response = update_rule_command(args)
 
     assert "The rule name to the policy name was updated successfully." in response.readable_output
+
+
+def test_engine_refresh_command_initiation(mocker):
+    """
+    Given:
+        - engine_name argument
+    When:
+        - Calling engine_refresh_command for the first time (no task_ids)
+    Then:
+        - Ensure refresh is initiated, task_ids are returned in args_for_next_run, and continue_to_poll is True
+    """
+    args = {"engine_name": ["engine1", "engine2"]}
+    mocker.patch("ForcepointSecurityManagementCenter.Engine", side_effect=lambda name: mock_Engine(name))
+    mocker.patch("ForcepointSecurityManagementCenter.Task", side_effect=lambda d: mock_Task(d.get("follower")))
+
+    result = engine_refresh_command(args)
+
+    assert result.scheduled_command is not None
+    assert result.scheduled_command._command is not None
+    assert result.scheduled_command._command == "forcepoint-smc-engine-refresh"
+    assert result.scheduled_command._args is not None
+    assert "task_ids" in result.scheduled_command._args
+    assert result.scheduled_command._args["task_ids"] == ["engine1", "engine2"]
+
+
+def test_engine_refresh_command_polling_in_progress(mocker):
+    """
+    Given:
+        - task_ids and engine_name arguments
+    When:
+        - Calling engine_refresh_command while tasks are still in progress
+    Then:
+        - Ensure continue_to_poll is True
+    """
+    args = {"engine_name": ["engine1"], "task_ids": ["engine1"]}
+    mocker.patch("ForcepointSecurityManagementCenter.Task", side_effect=lambda d: mock_Task(d.get("follower"), in_progress=True))
+
+    result = engine_refresh_command(args)
+
+    assert result.scheduled_command is not None
+    assert result.scheduled_command._command is not None
+    assert result.scheduled_command._command == "forcepoint-smc-engine-refresh"
+    assert result.scheduled_command._args is not None
+    assert "task_ids" in result.scheduled_command._args
+    assert result.scheduled_command._args["task_ids"] == ["engine1"]
+
+
+def test_engine_refresh_command_polling_complete(mocker):
+    """
+    Given:
+        - task_ids and engine_name arguments
+    When:
+        - Calling engine_refresh_command when all tasks are finished
+    Then:
+        - Ensure continue_to_poll is False and success message is returned
+    """
+    args = {"engine_name": ["engine1", "engine2"], "task_ids": ["engine1", "engine2"]}
+    mocker.patch("ForcepointSecurityManagementCenter.Task", side_effect=lambda d: mock_Task(d.get("follower"), in_progress=False))
+
+    result = engine_refresh_command(args)
+
+    assert result.scheduled_command is None
+    assert "refreshed successfully" in result.readable_output
+
+
+def test_engine_refresh_command_failure(mocker):
+    """
+    Given:
+        - engine_name argument
+    When:
+        - engine.refresh() raises TaskRunFailed
+    Then:
+        - Ensure DemistoException is raised with the correct engine names
+    """
+    args = {"engine_name": ["engine1", "engine2"]}
+
+    def mock_engine_side_effect(name):
+        engine = mock_Engine(name)
+        if name == "engine2":
+            mocker.patch.object(engine, "refresh", side_effect=TaskRunFailed("Locked"))
+        return engine
+
+    mocker.patch("ForcepointSecurityManagementCenter.Engine", side_effect=mock_engine_side_effect)
+
+    with pytest.raises(DemistoException, match=r"Failed to initiate refresh tasks for engines \(engine2\)"):
+        engine_refresh_command(args)

Packs/ForcepointSecurityManagementCenter/Integrations/ForcepointSecurityManagementCenter/README.md

@@ -994,3 +994,24 @@ Lists the engines in the system.
 >|Name|Comment|
 >|---|---|
 >| Forcepoint Engine | Forcepoint Engine element pre-populated by installer |
+
+### forcepoint-smc-engine-refresh
+
+***
+Refreshes the specified engines. Use `forcepoint-smc-engine-list` command to list the engines in the system.
+
+#### Base Command
+
+`forcepoint-smc-engine-refresh`
+
+#### Input
+
+| **Argument Name** | **Description** | **Required** |
+| --- | --- | --- |
+| engine_name | List of engine names to refresh. | Required |
+| interval_in_seconds | Interval between polling attempts in seconds. To prevent search timeouts, set this value within the 60-90 second range. Default is 30. | Optional |
+| timeout_in_seconds | Timeout for polling in seconds. Default is 600. | Optional |
+
+#### Context Output
+
+There is no context output for this command.

Packs/ForcepointSecurityManagementCenter/ReleaseNotes/1_0_9.md

@@ -0,0 +1,6 @@
+
+#### Integrations
+
+##### Forcepoint Security Management Center
+
+- Added support for ***forcepoint-smc-engine-refresh*** command that refreshes the specified engines. use ***forcepoint-smc-engine-list*** command to list the engines in the system.

Packs/ForcepointSecurityManagementCenter/pack_metadata.json

@@ -2,7 +2,7 @@
     "name": "Forcepoint Security Management Center",
     "description": "Forcepoint SMC provides unified, centralized management of all models of Forcepoint engines whether physical, virtual or cloud—across large, geographically distributed enterprise environments.",
     "support": "xsoar",
-    "currentVersion": "1.0.8",
+    "currentVersion": "1.0.9",
     "author": "Cortex XSOAR",
     "url": "https://www.forcepoint.com",
     "email": "",