diff --git a/cpapi/api_exceptions.py b/cpapi/api_exceptions.py index 38f3746..6b6145a 100644 --- a/cpapi/api_exceptions.py +++ b/cpapi/api_exceptions.py @@ -12,3 +12,7 @@ class APIClientException(APIException): def __init__(self, value): APIException.__init__(self, value, None) + +class TimeoutException(APIException): + def __init__(self, value): + APIException.__init__(self, value, None) \ No newline at end of file diff --git a/cpapi/mgmt_api.py b/cpapi/mgmt_api.py index 3c1bb00..62972a4 100644 --- a/cpapi/mgmt_api.py +++ b/cpapi/mgmt_api.py @@ -13,7 +13,7 @@ import sys # compatible import for python 2 and 3 -from .api_exceptions import APIException, APIClientException +from .api_exceptions import APIException, APIClientException, TimeoutException from .api_response import APIResponse from cpapi.utils import get_massage_from_io_error, compatible_loads @@ -248,7 +248,7 @@ def login_as_root(self, domain=None, payload=None): except (WindowsError) as err: raise APIClientException("Could not login as root:\n" + str(type(err)) + " - " + str(err)) - def api_call(self, command, payload=None, sid=None, wait_for_task=True): + def api_call(self, command, payload=None, sid=None, wait_for_task=True, timeout=-1): """ performs a web-service API request to the management server @@ -260,9 +260,12 @@ def api_call(self, command, payload=None, sid=None, wait_for_task=True): and will not return until the task is completed. when wait_for_task=False, it is up to the user to call the "show-task" API and check the status of the command. + :param timeout: Optional positive timeout (in seconds) before stop waiting for the task even if not completed. :return: APIResponse object :side-effects: updates the class's uid and server variables """ + timeout_start = time.time() + self.check_fingerprint() if payload is None: payload = {} @@ -351,9 +354,9 @@ def api_call(self, command, payload=None, sid=None, wait_for_task=True): # If we want to wait for the task to end, wait for it if wait_for_task is True and res.success and command != "show-task": if "task-id" in res.data: - res = self.__wait_for_task(res.data["task-id"]) + res = self.__wait_for_task(res.data["task-id"], timeout=(timeout - time.time() + timeout_start)) elif "tasks" in res.data: - res = self.__wait_for_tasks(res.data["tasks"]) + res = self.__wait_for_tasks(res.data["tasks"], timeout=(timeout - time.time() + timeout_start)) return res @@ -478,7 +481,7 @@ def get_server_fingerprint(self): conn.close() return fingerprint_hash - def __wait_for_task(self, task_id): + def __wait_for_task(self, task_id, timeout=-1): """ When the server needs to perform an API call that may take a long time (e.g. run-script, install-policy, publish), the server responds with a 'task-id'. @@ -487,15 +490,22 @@ def __wait_for_task(self, task_id): The function will return when the task (and its sub-tasks) are no longer in-progress. :param task_id: The task identifier. + :param timeout: Optional positive timeout (in seconds) that will end the task even if not completed. :return: APIResponse object (response of show-task command). :raises APIException """ task_complete = False task_result = None + task_start = time.time() in_progress = "in progress" - # As long as there is a task in progress + # As long as there is a task in progress or the timeout isn't expired (and is positive) while not task_complete: + + # If timeout parameter was set and valid and timeout did expire, raise exception + if timeout >= 0 and time.time() - task_start > timeout: + raise TimeoutException("Timeout reached when waiting for task to complete") + # Check the status of the task task_result = self.api_call("show-task", {"task-id": task_id, "details-level": "full"}, self.sid, False) @@ -526,13 +536,14 @@ def __wait_for_task(self, task_id): self.check_tasks_status(task_result) return task_result - def __wait_for_tasks(self, task_objects): + def __wait_for_tasks(self, task_objects, timeout=-1): """ The version of __wait_for_task function for the collection of tasks :param task_objects: A list of task objects :return: APIResponse object (response of show-task command). """ + timeout_start = time.time() # A list of task ids to be retrieved tasks = [] @@ -540,7 +551,7 @@ def __wait_for_tasks(self, task_objects): # Retrieve the taskId and wait for the task to be completed task_id = task_obj["task-id"] tasks.append(task_id) - self.__wait_for_task(task_id) + self.__wait_for_task(task_id, timeout=(timeout - time.time() + timeout_start)) task_result = self.api_call("show-task", {"task-id": tasks, "details-level": "full"}, self.sid, False) @@ -557,7 +568,7 @@ def check_tasks_status(task_result): :return: """ for task in task_result.data["tasks"]: - if task["status"] == "failed" or task["status"] == "partially succeeded": + if task["status"] == "failed" or task["status"] == "partially succeeded" or task["status"] == "in progress": task_result.set_success_status(False) break diff --git a/setup.py b/setup.py index 2abdaed..9f6ca91 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ setup( name="cpapi", - version="1.1.1", + version="1.1.2", author="API team", author_email="api_team@checkpoint.com", description="Check Point Management API SDK",