add lockdown state and mod (#199)

* add lockdown state and mod

* run pre

* add examples

* remove def for name

* fix docs

* update docs

* tests

* reformat

* fix conflict

* fix whitespace

* fix docs

Author: cmcmarrow

src/saltext/vmware/modules/esxi.py

@@ -2413,3 +2413,103 @@ def exit_maintenance_mode(host, timeout=0, catch_task_error=True, service_instan
     mode = in_maintenance_mode(host_ref, service_instance)
     mode["changes"] = mode["maintenanceMode"] == "normal"
     return mode
+
+
+def in_lockdown_mode(host, service_instance=None):
+    """
+    Check if host is in lockdown mode.
+
+    host
+        Host IP or HostSystem/ManagedObjectReference (required).
+
+    service_instance
+        Use this vCenter service connection instance instead of creating a new one (optional).
+
+    .. code-block:: bash
+
+        salt '*' vmware_esxi.in_lockdown_mode '10.288.6.117'
+    """
+    if isinstance(host, vim.HostSystem):
+        host_ref = host
+    else:
+        if service_instance is None:
+            service_instance = get_service_instance(opts=__opts__, pillar=__pillar__)
+        host_ref = utils_esxi.get_host(host, service_instance)
+    mode = "normal"
+    if host_ref.config.adminDisabled:
+        mode = "inLockdown"
+    return {"lockdownMode": mode}
+
+
+def lockdown_mode(host, catch_task_error=True, service_instance=None):
+    """
+    Put host into lockdown mode.
+
+    host
+        Host IP or HostSystem/ManagedObjectReference (required).
+
+    catch_task_error
+        If False and task failed then a salt exception will be thrown (optional).
+
+    service_instance
+        Use this vCenter service connection instance instead of creating a new one (optional).
+
+    .. code-block:: bash
+
+        salt '*' vmware_esxi.lockdown_mode '10.288.6.117'
+    """
+    if isinstance(host, vim.HostSystem):
+        host_ref = host
+    else:
+        if service_instance is None:
+            service_instance = get_service_instance(opts=__opts__, pillar=__pillar__)
+        host_ref = utils_esxi.get_host(host, service_instance)
+    mode = in_lockdown_mode(host_ref)
+    if mode["lockdownMode"] == "inLockdown":
+        mode["changes"] = False
+        return mode
+    try:
+        host_ref.EnterLockdownMode()
+    except salt.exceptions.SaltException as exc:
+        if not catch_task_error:
+            raise exc
+    mode = in_lockdown_mode(host_ref, service_instance)
+    mode["changes"] = mode["lockdownMode"] == "inLockdown"
+    return mode
+
+
+def exit_lockdown_mode(host, catch_task_error=True, service_instance=None):
+    """
+    Put host out of lockdown mode.
+
+    host
+        Host IP or HostSystem/ManagedObjectReference (required).
+
+    catch_task_error
+        If False and task failed then a salt exception will be thrown (optional).
+
+    service_instance
+        Use this vCenter service connection instance instead of creating a new one (optional).
+
+    .. code-block:: bash
+
+        salt '*' vmware_esxi.exit_lockdown_mode '10.288.6.117'
+    """
+    if isinstance(host, vim.HostSystem):
+        host_ref = host
+    else:
+        if service_instance is None:
+            service_instance = get_service_instance(opts=__opts__, pillar=__pillar__)
+        host_ref = utils_esxi.get_host(host, service_instance)
+    mode = in_lockdown_mode(host_ref)
+    if mode["lockdownMode"] == "normal":
+        mode["changes"] = False
+        return mode
+    try:
+        host_ref.ExitLockdownMode()
+    except salt.exceptions.SaltException as exc:
+        if not catch_task_error:
+            raise exc
+    mode = in_lockdown_mode(host_ref, service_instance)
+    mode["changes"] = mode["lockdownMode"] == "normal"
+    return mode

src/saltext/vmware/states/esxi.py

@@ -739,3 +739,111 @@ def maintenance_mode(
             "comment"
         ] = f"Failed to put host {str(name)} in {'Maintenance' if enter_maintenance_mode else 'Normal'} mode."
     return ret
+
+
+def lockdown_mode(
+    name,
+    enter_lockdown_mode,
+    datacenter_name=None,
+    cluster_name=None,
+    get_all_hosts=False,
+    service_instance=None,
+):
+    """
+    Pust a hosts into or out of lockdown.
+
+    name
+        IP of single host or list of host_names. If wanting to get a cluster just past an empty list (required).
+
+    enter_lockdown_mode
+        If True, put host into lockdown mode.
+        If False, put host out of lockdown mode (required)
+
+    datacenter_name
+        The datacenter name. Default is None (optional).
+
+    host_names
+        The host_names to be retrieved. Default is None (optional).
+
+    cluster_name
+        The cluster name - used to restrict the hosts retrieved. Only used if
+        the datacenter is set.  This argument is optional (optional).
+
+    get_all_hosts
+        Specifies whether to retrieve all hosts in the container.
+        Default value is False (optional).
+
+    service_instance
+        The Service Instance Object from which to obtain the hosts (optional).
+
+    .. code-block:: bash
+
+        salt '*' vmware_esxi.lockdown_mode '10.288.6.117'
+    .. code-block:: yaml
+
+        Lockdown Mode:
+          vmware_esxi.lockdown_mode:
+            - host: '10.288.6.117'
+            - enter_lockdown_mode: true
+    """
+    ret = {"name": name, "changes": {}, "result": True, "comment": ""}
+    if not isinstance(name, str):
+        host_refs = utils_esxi.get_hosts(
+            service_instance=service_instance,
+            datacenter_name=datacenter_name,
+            host_names=name,
+            cluster_name=cluster_name,
+            get_all_hosts=get_all_hosts,
+        )
+    else:
+        if isinstance(name, vim.HostSystem):
+            host_refs = (name,)
+        else:
+            if service_instance is None:
+                service_instance = get_service_instance(opts=__opts__, pillar=__pillar__)
+            host_refs = (utils_esxi.get_host(name, service_instance),)
+
+    for ref in host_refs:
+        # check that host is not all ready in lock state.
+        host_state = __salt__["vmware_esxi.in_lockdown_mode"](
+            host=ref, service_instance=service_instance
+        )
+        if (host_state["lockdownMode"] == "inLockdown") == enter_lockdown_mode:
+            ret[
+                "comment"
+            ] += f"{ref.name} already in {'Lockdown' if enter_lockdown_mode else 'Normal'} mode.\n"
+            continue
+
+        if __opts__["test"]:
+            ret["result"] = None
+            ret["changes"].setdefault("new", []).append(
+                f"{ref.name} will enter {'Lockdown' if enter_lockdown_mode else 'Normal'} mode."
+            )
+            continue
+
+        if enter_lockdown_mode:
+            host_state = __salt__["vmware_esxi.lockdown_mode"](
+                host=ref,
+                catch_task_error=True,
+                service_instance=service_instance,
+            )
+        else:
+            host_state = __salt__["vmware_esxi.exit_lockdown_mode"](
+                host=ref, catch_task_error=True, service_instance=service_instance
+            )
+        ref_results = (host_state["lockdownMode"] == "inLockdown") == enter_lockdown_mode
+        if ret["result"]:
+            ret["result"] = ref_results
+        if ref_results:
+            ret["changes"].setdefault("new", []).append(
+                f"{ref.name} entered {'Lockdown' if enter_lockdown_mode else 'Normal'} mode."
+            )
+        else:
+            ret[
+                "comment"
+            ] += f"Failed to put host {ref.name} in {'Lockdown' if enter_lockdown_mode else 'Normal'} mode.\n"
+    if ret["result"]:
+        ret["comment"] += f"Task was successfully!\n"
+    elif ret["result"] is None:
+        ret["comment"] += "These options are set to change."
+    return ret

tests/integration/modules/test_esxi.py

@@ -684,3 +684,29 @@ def test_maintenance_mode(service_instance):
 
     ret = esxi.in_maintenance_mode(host, service_instance)
     assert ret == dict(maintenanceMode="normal")
+
+
+def test_lockdown_mode(service_instance):
+    hosts = list(esxi.get(service_instance=service_instance))
+    assert hosts
+    host = hosts[0]
+    ret = esxi.in_lockdown_mode(host, service_instance)
+    assert ret == dict(lockdownMode="normal")
+
+    try:
+        for i in range(3):
+            ret = esxi.lockdown_mode(host, service_instance=service_instance)
+            assert ret == dict(lockdownMode="inLockdown", changes=not i)
+    except Exception as e:
+        esxi.exit_lockdown_mode(host, service_instance=service_instance)
+        raise e
+
+    ret = esxi.in_lockdown_mode(host, service_instance)
+    assert ret == dict(lockdownMode="inLockdown")
+
+    for i in range(3):
+        ret = esxi.exit_lockdown_mode(host, service_instance=service_instance)
+        assert ret == dict(lockdownMode="normal", changes=not i)
+
+    ret = esxi.in_lockdown_mode(host, service_instance)
+    assert ret == dict(lockdownMode="normal")

tests/integration/states/test_esxi.py

@@ -312,3 +312,38 @@ def test_maintenance_mode_dry_run(service_instance, dry_run):
     assert ret["result"] is None
     assert ret["changes"]
     assert ret["comment"] == "These options are set to change."
+
+
+def test_lockdown_mode(service_instance):
+    hosts = list(esxi_mod.get(service_instance=service_instance))
+    assert hosts
+    host = hosts[0]
+    try:
+        for i in range(3):
+            ret = esxi.lockdown_mode(host, True, service_instance=service_instance)
+            assert ret["result"]
+            if not i:
+                assert ret["changes"]
+            else:
+                assert not ret["changes"]
+    except Exception as e:
+        esxi.lockdown_mode(host, False, service_instance=service_instance)
+        raise e
+
+    for i in range(3):
+        ret = esxi.lockdown_mode(host, False, service_instance=service_instance)
+        assert ret["result"]
+        if not i:
+            assert ret["changes"]
+        else:
+            assert not ret["changes"]
+
+
+def test_lockdown_mode_dry_run(service_instance, dry_run):
+    hosts = list(esxi_mod.get(service_instance=service_instance))
+    assert hosts
+    host = hosts[0]
+    ret = esxi.lockdown_mode(host, True, service_instance=service_instance)
+    assert ret["result"] is None
+    assert ret["changes"]
+    assert ret["comment"] == "These options are set to change."