enhanced system action handling (#59)

ansibleguy · Mar 11, 2024 · a613029 · a613029
1 parent 69401e3
commit a613029
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 33 deletions.
diff --git a/docs/source/modules/system.rst b/docs/source/modules/system.rst
@@ -58,10 +58,6 @@ Examples
           ansibleguy.opnsense.system:
             action: 'update'
 
-        - name: Wait for information to be refreshed
-          ansible.builtin.pause:
-            seconds: 10
-
         - name: Start upgrade - will wait until finished
           ansibleguy.opnsense.system:
             action: 'upgrade'

diff --git a/plugins/module_utils/base/api.py b/plugins/module_utils/base/api.py
@@ -32,7 +32,7 @@ def _start(self, timeout: float) -> httpx.Client:
         return httpx.Client(
             base_url=f"https://{self.m.params['firewall']}:{self.m.params['api_port']}/api",
             auth=(self.m.params['api_key'], self.m.params['api_secret']),
-            timeout=httpx.Timeout(timeout=timeout),
+            timeout=httpx.Timeout(timeout=timeout, connect=2.0),
             transport=httpx.HTTPTransport(
                 verify=ssl_verification(module=self.m),
                 retries=self.m.params['api_retries'],

diff --git a/plugins/module_utils/helper/system.py b/plugins/module_utils/helper/system.py
@@ -4,6 +4,7 @@
 
 from ansible.module_utils.basic import AnsibleModule
 
+from ansible_collections.ansibleguy.opnsense.plugins.module_utils.base.api import Session, HTTPX_EXCEPTIONS
 from ansible_collections.ansibleguy.opnsense.plugins.module_utils.defaults.main import CONNECTION_TEST_TIMEOUT
 
 
@@ -31,20 +32,62 @@ def _wait_msg(module: AnsibleModule, msg: str):
 def wait_for_response(module: AnsibleModule) -> bool:
     timeout = time() + module.params['wait_timeout']
 
-    if module.params['action'] == 'upgrade':
-        _wait_msg(module, 'Waiting download & upgrade to finish..')
-        sleep(int(module.params['wait_timeout'] / 2))
-
-    else:
-        _wait_msg(module, 'Waiting for service to stop..')
-        sleep(10)
+    _wait_msg(module, 'Waiting for service to stop..')
+    sleep(10)
 
     while time() < timeout:
+        poll_interval_start = time()
+
         if _opn_reachable(module=module):
             _wait_msg(module, 'Got response!')
             return True
 
         _wait_msg(module, 'Waiting for response..')
-        sleep(module.params['poll_interval'])
+        poll_interval_elapsed = time() - poll_interval_start
+        if poll_interval_elapsed < module.params['poll_interval']:
+            sleep(module.params['poll_interval'] - poll_interval_elapsed)
+
+    raise TimeoutError
+
+
+def wait_for_update(module: AnsibleModule, s: Session) -> bool:
+    timeout = time() + module.params['wait_timeout']
+
+    if module.params['action'] == 'upgrade':
+        _wait_msg(module, 'Waiting for download & upgrade to finish..')
+
+    else:
+        _wait_msg(module, 'Waiting for update to finish..')
+
+    sleep(2)
+
+    while time() < timeout:
+        poll_interval_start = time()
+
+        try:
+            result = s.get({
+                'command': 'upgradestatus',
+                'module': 'core',
+                'controller': 'firmware',
+            })
+            status = result['status']
+
+            _wait_msg(module, f"Got response: {status}")
+
+            if status == 'error' and 'log' in result:
+                _wait_msg(module, f"Got error: {result['log']}")
+                return False
+
+            if status == 'done':
+                _wait_msg(module, f"Got result: {result['log']}")
+                return True
+
+        except HTTPX_EXCEPTIONS:
+            # not reachable while rebooting
+            _wait_msg(module, 'Waiting for response..')
+
+        poll_interval_elapsed = time() - poll_interval_start
+        if poll_interval_elapsed < module.params['poll_interval']:
+            sleep(module.params['poll_interval'] - poll_interval_elapsed)
 
-    return False
+    raise TimeoutError
diff --git a/plugins/modules/system.py b/plugins/modules/system.py
@@ -11,9 +11,10 @@
     module_dependency_error, MODULE_EXCEPTIONS
 
 try:
-    from ansible_collections.ansibleguy.opnsense.plugins.module_utils.base.api import single_post
+    from ansible_collections.ansibleguy.opnsense.plugins.module_utils.base.api import Session
     from ansible_collections.ansibleguy.opnsense.plugins.module_utils.defaults.main import OPN_MOD_ARGS
-    from ansible_collections.ansibleguy.opnsense.plugins.module_utils.helper.system import wait_for_response
+    from ansible_collections.ansibleguy.opnsense.plugins.module_utils.helper.system import wait_for_response, \
+        wait_for_update
 
 except MODULE_EXCEPTIONS:
     module_dependency_error()
@@ -46,24 +47,26 @@ def run_module():
     }
 
     if not module.check_mode:
-        single_post(
-            module=module,
-            cnf={
+        with Session(module=module) as s:
+            s.post({
                 'command': module.params['action'],
                 'module': 'core',
                 'controller': 'firmware',
-            }
-        )
+            })
 
-        if module.params['action'] in ['reboot', 'upgrade'] and module.params['wait']:
-            if module.params['debug']:
-                module.warn(f"Waiting for firewall to complete '{module.params['action']}'!")
+            if module.params['wait']:
+                if module.params['debug']:
+                    module.warn(f"Waiting for firewall to complete '{module.params['action']}'!")
 
-            # todo: cleaner way of handling if no upgrade is needed
-            result['failed'] = not wait_for_response(module=module)
+                try:
+                    if module.params['action'] in ['upgrade', 'update']:
+                        result['failed'] = not wait_for_update(module=module, s=s)
 
-            if result['failed']:
-                result['timeout_exceeded'] = True
+                    elif module.params['action'] == 'reboot':
+                        result['failed'] = not wait_for_response(module=module)
+
+                except TimeoutError:
+                    result['timeout_exceeded'] = True
 
     module.exit_json(**result)
 

diff --git a/tests/system.yml b/tests/system.yml
@@ -14,11 +14,6 @@
       ansibleguy.opnsense.system:
         action: 'update'
 
-    - name: Wait for information to be refreshed
-      ansible.builtin.pause:
-        seconds: 10
-      when: not ansible_check_mode
-
     - name: Upgrade with wait
       ansibleguy.opnsense.system:
         action: 'upgrade'