From 7964429e04b6cf053ff4dee1f345c3e284f73f80 Mon Sep 17 00:00:00 2001
From: Shane Lee <shane.d.lee@gmail.com>
Date: Fri, 1 Mar 2024 14:09:37 -0700
Subject: [PATCH] Add the ability to bootstrap a specific version of chocolatey

Added a state to bootstrap and unbootstrap chocolatey
Add changelog
Add and update tests
---
 changelog/64722.added.md                      |   3 +
 salt/modules/chocolatey.py                    |  44 ++++-
 salt/states/chocolatey.py                     | 181 +++++++++++++++++-
 .../functional/modules/test_chocolatey.py     |  45 +++++
 .../states/chocolatey/test_bootstrap.py       | 100 ++++++++++
 .../test_post_2.0.py}                         |   2 +-
 .../test_pre_2.0.py}                          |   2 +-
 tests/pytests/unit/modules/test_chocolatey.py |  32 +++-
 8 files changed, 396 insertions(+), 13 deletions(-)
 create mode 100644 changelog/64722.added.md
 create mode 100644 tests/pytests/functional/modules/test_chocolatey.py
 create mode 100644 tests/pytests/functional/states/chocolatey/test_bootstrap.py
 rename tests/pytests/functional/states/{test_chocolatey_latest.py => chocolatey/test_post_2.0.py} (98%)
 rename tests/pytests/functional/states/{test_chocolatey_1_2_1.py => chocolatey/test_pre_2.0.py} (98%)

diff --git a/changelog/64722.added.md b/changelog/64722.added.md
new file mode 100644
index 00000000000..6a05d2f59d6
--- /dev/null
+++ b/changelog/64722.added.md
@@ -0,0 +1,3 @@
+Added the ability to pass a version of chocolatey to install to the
+chocolatey.bootstrap function. Also added states to bootstrap and
+unbootstrap chocolatey.
diff --git a/salt/modules/chocolatey.py b/salt/modules/chocolatey.py
index ade057da95c..d618a175572 100644
--- a/salt/modules/chocolatey.py
+++ b/salt/modules/chocolatey.py
@@ -14,6 +14,7 @@ from requests.structures import CaseInsensitiveDict
 
 import salt.utils.data
 import salt.utils.platform
+import salt.utils.win_dotnet
 from salt.exceptions import (
     CommandExecutionError,
     CommandNotFoundError,
@@ -126,16 +127,26 @@ def _find_chocolatey():
     raise CommandExecutionError(err)
 
 
-def chocolatey_version():
+def chocolatey_version(refresh=False):
     """
     Returns the version of Chocolatey installed on the minion.
 
+    Args:
+
+        refresh (bool):
+            Refresh the cached version of chocolatey
+
+            .. versionadded:: 3008.0
+
     CLI Example:
 
     .. code-block:: bash
 
         salt '*' chocolatey.chocolatey_version
     """
+    if refresh:
+        __context__.pop("chocolatey._version", False)
+
     if "chocolatey._version" in __context__:
         return __context__["chocolatey._version"]
 
@@ -147,7 +158,7 @@ def chocolatey_version():
     return __context__["chocolatey._version"]
 
 
-def bootstrap(force=False, source=None):
+def bootstrap(force=False, source=None, version=None):
     """
     Download and install the latest version of the Chocolatey package manager
     via the official bootstrap.
@@ -166,6 +177,11 @@ def bootstrap(force=False, source=None):
         and .NET requirements must already be met on the target. This shouldn't
         be a problem on Windows versions 2012/8 and later
 
+    .. note::
+        If you're installing chocolatey version 2.0+ the system requires .NET
+        4.8. Installing this requires a reboot, therefore this module will not
+        automatically install .NET 4.8.
+
     Args:
 
         force (bool):
@@ -182,6 +198,12 @@ def bootstrap(force=False, source=None):
 
             .. versionadded:: 3001
 
+        version (str):
+            The version of chocolatey to install. The latest version is
+            installed if this value is ``None``. Default is ``None``
+
+            .. versionadded:: 3008.0
+
     Returns:
         str: The stdout of the Chocolatey installation script
 
@@ -198,6 +220,9 @@ def bootstrap(force=False, source=None):
 
         # To bootstrap Chocolatey from a file on C:\\Temp
         salt '*' chocolatey.bootstrap source=C:\\Temp\\chocolatey.nupkg
+
+        # To bootstrap Chocolatey version 1.4.0
+        salt '*' chocolatey.bootstrap version=1.4.0
     """
     # Check if Chocolatey is already present in the path
     try:
@@ -272,7 +297,7 @@ def bootstrap(force=False, source=None):
     # Check that .NET v4.0+ is installed
     # Windows 7 / Windows Server 2008 R2 and below do not come with at least
     # .NET v4.0 installed
-    if not __utils__["dotnet.version_at_least"](version="4"):
+    if not salt.utils.win_dotnet.version_at_least(version="4"):
         # It took until .NET v4.0 for Microsoft got the hang of making
         # installers, this should work under any version of Windows
         url = "http://download.microsoft.com/download/1/B/E/1BE39E79-7E39-46A3-96FF-047F95396215/dotNetFx40_Full_setup.exe"
@@ -336,10 +361,21 @@ def bootstrap(force=False, source=None):
             f"Failed to find Chocolatey installation script: {script}"
         )
 
+    # You tell the chocolatey install script which version to install by setting
+    # an environment variable
+    if version:
+        env = {"chocolateyVersion": version}
+    else:
+        env = None
+
     # Run the Chocolatey bootstrap
     log.debug("Installing Chocolatey: %s", script)
     result = __salt__["cmd.script"](
-        script, cwd=os.path.dirname(script), shell="powershell", python_shell=True
+        source=script,
+        cwd=os.path.dirname(script),
+        shell="powershell",
+        python_shell=True,
+        env=env,
     )
     if result["retcode"] != 0:
         err = "Bootstrapping Chocolatey failed: {}".format(result["stderr"])
diff --git a/salt/states/chocolatey.py b/salt/states/chocolatey.py
index 5e49113607e..9c8f1951cad 100644
--- a/salt/states/chocolatey.py
+++ b/salt/states/chocolatey.py
@@ -11,7 +11,7 @@ Manage Windows Packages using Chocolatey
 
 import salt.utils.data
 import salt.utils.versions
-from salt.exceptions import SaltInvocationError
+from salt.exceptions import CommandExecutionError, SaltInvocationError
 
 
 def __virtual__():
@@ -513,3 +513,182 @@ def source_present(
     ret["changes"] = salt.utils.data.compare_dicts(pre_install, post_install)
 
     return ret
+
+
+def bootstrapped(name, force=False, source=None, version=None):
+    """
+    .. versionadded:: 3008.0
+
+    Ensure chocolatey is installed on the system.
+
+    You can't upgrade an existing installation with this state. You must use
+    chocolatey to upgrade chocolatey.
+
+    For example:
+
+    .. code-block:: bash
+
+        choco upgrade chocolatey --version 2.2.0
+
+    Args:
+
+        name (str):
+            The name of the state that installs chocolatey. Required for all
+            states. This is ignored.
+
+        force (bool):
+            Run the bootstrap process even if Chocolatey is found in the path.
+
+            .. note::
+                If chocolatey is already installed this will just re-run the
+                install script over the existing version. The ``version``
+                parameter is ignored.
+
+        source (str):
+            The location of the ``.nupkg`` file or ``.ps1`` file to run from an
+            alternate location. This can be one of the following types of URLs:
+
+            - salt://
+            - http(s)://
+            - ftp://
+            - file:// - A local file on the system
+
+        version (str):
+            The version of chocolatey to install. The latest version is
+            installed if this value is ``None``. Default is ``None``
+
+    Example:
+
+    .. code-block:: yaml
+
+        # Bootstrap the latest version of chocolatey
+        bootstrap_chocolatey:
+          chocolatey.bootstrapped
+
+        # Bootstrap the latest version of chocolatey
+        # If chocolatey is already present, re-run the install script
+        bootstrap_chocolatey:
+          chocolatey.bootstrapped:
+            - force: True
+
+        # Bootstrap Chocolatey version 1.4.0
+        bootstrap_chocolatey:
+          chocolatey.bootstrapped:
+            - version: 1.4.0
+
+        # Bootstrap Chocolatey from a local file
+        bootstrap_chocolatey:
+          chocolatey.bootstrapped:
+            - source: C:\\Temp\\chocolatey.nupkg
+
+        # Bootstrap Chocolatey from a file on the salt master
+        bootstrap_chocolatey:
+          chocolatey.bootstrapped:
+            - source: salt://Temp/chocolatey.nupkg
+    """
+    ret = {"name": name, "result": True, "changes": {}, "comment": ""}
+
+    try:
+        old = __salt__["chocolatey.chocolatey_version"]()
+    except CommandExecutionError:
+        old = None
+
+    # Try to predict what will happen
+    if old:
+        if force:
+            ret["comment"] = (
+                f"Chocolatey {old} will be reinstalled\n"
+                'Use "choco upgrade chocolatey --version 2.1.0" to change the version'
+            )
+        else:
+            # You can't upgrade chocolatey using the install script, you have to use
+            # chocolatey itself
+            ret["comment"] = (
+                f"Chocolatey {old} is already installed.\n"
+                'Use "choco upgrade chocolatey --version 2.1.0" to change the version'
+            )
+            return ret
+
+    else:
+        if version is None:
+            ret["comment"] = "The latest version of Chocolatey will be installed"
+        else:
+            ret["comment"] = f"Chocolatey {version} will be installed"
+
+    if __opts__["test"]:
+        ret["result"] = None
+        return ret
+
+    __salt__["chocolatey.bootstrap"](force=force, source=source, version=version)
+
+    try:
+        new = __salt__["chocolatey.chocolatey_version"](refresh=True)
+    except CommandExecutionError:
+        new = None
+
+    if new is None:
+        ret["comment"] = f"Failed to install chocolatey {new}"
+        ret["result"] = False
+    else:
+        if salt.utils.versions.version_cmp(old, new) == 0:
+            ret["comment"] = f"Re-installed chocolatey {new}"
+        else:
+            ret["comment"] = f"Installed chocolatey {new}"
+            ret["changes"] = {"old": old, "new": new}
+
+    return ret
+
+
+def unbootstrapped(name):
+    """
+    .. versionadded:: 3008.0
+
+    Ensure chocolatey is removed from the system.
+
+    Args:
+
+        name (str):
+            The name of the state that uninstalls chocolatey. Required for all
+            states. This is ignored.
+
+    Example:
+
+    .. code-block:: yaml
+
+        # Uninstall chocolatey
+        uninstall_chocolatey:
+          chocolatey.unbootstrapped
+
+    """
+    ret = {"name": name, "result": True, "changes": {}, "comment": ""}
+
+    try:
+        old = __salt__["chocolatey.chocolatey_version"]()
+    except CommandExecutionError:
+        old = None
+
+    if old is None:
+        ret["comment"] = "Chocolatey not found on this system"
+        return ret
+
+    ret["comment"] = f"Chocolatey {old} will be removed"
+
+    if __opts__["test"]:
+        ret["result"] = None
+        return ret
+
+    __salt__["chocolatey.unbootstrap"]()
+
+    try:
+        new = __salt__["chocolatey.chocolatey_version"](refresh=True)
+    except CommandExecutionError:
+        new = None
+
+    if new is None:
+        ret["comment"] = f"Uninstalled chocolatey {old}"
+        ret["changes"] = {"new": new, "old": old}
+    else:
+        ret["comment"] = f"Failed to uninstall chocolatey {old}\nFound version {new}"
+        ret["result"] = False
+
+    return ret
diff --git a/tests/pytests/functional/modules/test_chocolatey.py b/tests/pytests/functional/modules/test_chocolatey.py
new file mode 100644
index 00000000000..d11cdf6cbed
--- /dev/null
+++ b/tests/pytests/functional/modules/test_chocolatey.py
@@ -0,0 +1,45 @@
+import pytest
+
+from salt.exceptions import CommandExecutionError
+
+pytestmark = [
+    pytest.mark.destructive,
+    pytest.mark.skip_unless_on_windows,
+    pytest.mark.slow_test,
+    pytest.mark.windows_whitelisted,
+]
+
+
+@pytest.fixture(scope="module")
+def chocolatey(modules):
+    return modules.chocolatey
+
+
+@pytest.fixture()
+def clean(chocolatey):
+    chocolatey.unbootstrap()
+    try:
+        chocolatey_version = chocolatey.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version is None
+    yield
+    chocolatey.unbootstrap()
+
+
+def test_bootstrap(chocolatey, clean):
+    chocolatey.bootstrap()
+    try:
+        chocolatey_version = chocolatey.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version is not None
+
+
+def test_bootstrap_version(chocolatey, clean):
+    chocolatey.bootstrap(version="1.4.0")
+    try:
+        chocolatey_version = chocolatey.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version == "1.4.0"
diff --git a/tests/pytests/functional/states/chocolatey/test_bootstrap.py b/tests/pytests/functional/states/chocolatey/test_bootstrap.py
new file mode 100644
index 00000000000..0d37d67c665
--- /dev/null
+++ b/tests/pytests/functional/states/chocolatey/test_bootstrap.py
@@ -0,0 +1,100 @@
+import pytest
+
+from salt.exceptions import CommandExecutionError
+
+pytestmark = [
+    pytest.mark.destructive_test,
+    pytest.mark.skip_unless_on_windows,
+    pytest.mark.slow_test,
+    pytest.mark.windows_whitelisted,
+]
+
+
+@pytest.fixture(scope="module")
+def chocolatey(states):
+    yield states.chocolatey
+
+
+@pytest.fixture(scope="module")
+def chocolatey_mod(modules):
+    yield modules.chocolatey
+
+
+@pytest.fixture()
+def clean(chocolatey_mod):
+    chocolatey_mod.unbootstrap()
+    try:
+        chocolatey_version = chocolatey_mod.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version is None
+    yield
+    chocolatey_mod.unbootstrap()
+
+
+@pytest.fixture()
+def installed(chocolatey_mod):
+    chocolatey_mod.bootstrap(force=True)
+    try:
+        chocolatey_version = chocolatey_mod.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version is not None
+    yield
+    chocolatey_mod.unbootstrap()
+
+
+def test_bootstrapped(chocolatey, chocolatey_mod, clean):
+    ret = chocolatey.bootstrapped(name="junk name")
+    assert "Installed chocolatey" in ret.comment
+    assert ret.result is True
+    try:
+        chocolatey_version = chocolatey_mod.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version is not None
+
+
+def test_bootstrapped_test_true(chocolatey, clean):
+    ret = chocolatey.bootstrapped(name="junk name", test=True)
+    assert ret.result is None
+    assert ret.comment == "The latest version of Chocolatey will be installed"
+
+
+def test_bootstrapped_version(chocolatey, chocolatey_mod, clean):
+    ret = chocolatey.bootstrapped(name="junk_name", version="1.4.0")
+    assert ret.comment == "Installed chocolatey 1.4.0"
+    assert ret.result is True
+    try:
+        chocolatey_version = chocolatey_mod.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version == "1.4.0"
+
+
+def test_bootstrapped_version_test_true(chocolatey, chocolatey_mod, clean):
+    ret = chocolatey.bootstrapped(name="junk_name", version="1.4.0", test=True)
+    assert ret.comment == "Chocolatey 1.4.0 will be installed"
+
+
+def test_unbootstrapped_installed(chocolatey, chocolatey_mod, installed):
+    ret = chocolatey.unbootstrapped(name="junk_name")
+    assert "Uninstalled chocolatey" in ret.comment
+    assert ret.result is True
+    try:
+        chocolatey_version = chocolatey_mod.chocolatey_version(refresh=True)
+    except CommandExecutionError:
+        chocolatey_version = None
+    assert chocolatey_version is None
+
+
+def test_unbootstrapped_installed_test_true(chocolatey, chocolatey_mod, installed):
+    ret = chocolatey.unbootstrapped(name="junk_name", test=True)
+    assert "will be removed" in ret.comment
+    assert ret.result is None
+
+
+def test_unbootstrapped_clean(chocolatey, chocolatey_mod, clean):
+    ret = chocolatey.unbootstrapped(name="junk_name")
+    assert ret.comment == "Chocolatey not found on this system"
+    assert ret.result is True
diff --git a/tests/pytests/functional/states/test_chocolatey_latest.py b/tests/pytests/functional/states/chocolatey/test_post_2.0.py
similarity index 98%
rename from tests/pytests/functional/states/test_chocolatey_latest.py
rename to tests/pytests/functional/states/chocolatey/test_post_2.0.py
index 41ba0df5b38..5e62c86d1db 100644
--- a/tests/pytests/functional/states/test_chocolatey_latest.py
+++ b/tests/pytests/functional/states/chocolatey/test_post_2.0.py
@@ -1,5 +1,5 @@
 """
-Functional tests for chocolatey state
+Functional tests for chocolatey state with Chocolatey 2.0+
 """
 
 import os
diff --git a/tests/pytests/functional/states/test_chocolatey_1_2_1.py b/tests/pytests/functional/states/chocolatey/test_pre_2.0.py
similarity index 98%
rename from tests/pytests/functional/states/test_chocolatey_1_2_1.py
rename to tests/pytests/functional/states/chocolatey/test_pre_2.0.py
index 0e9972df17e..0d00d313854 100644
--- a/tests/pytests/functional/states/test_chocolatey_1_2_1.py
+++ b/tests/pytests/functional/states/chocolatey/test_pre_2.0.py
@@ -1,5 +1,5 @@
 """
-Functional tests for chocolatey state
+Functional tests for chocolatey state with Chocolatey < 2.0
 """
 
 import os
diff --git a/tests/pytests/unit/modules/test_chocolatey.py b/tests/pytests/unit/modules/test_chocolatey.py
index 8dd630793f1..0f6207e2382 100644
--- a/tests/pytests/unit/modules/test_chocolatey.py
+++ b/tests/pytests/unit/modules/test_chocolatey.py
@@ -7,14 +7,10 @@ import os
 import pytest
 
 import salt.modules.chocolatey as chocolatey
-import salt.utils
-import salt.utils.platform
 from tests.support.mock import MagicMock, patch
 
 pytestmark = [
-    pytest.mark.skipif(
-        not salt.utils.platform.is_windows(), reason="Not a Windows system"
-    )
+    pytest.mark.skip_unless_on_windows,
 ]
 
 
@@ -68,7 +64,7 @@ def test__clear_context(choco_path):
     }
     with patch.dict(chocolatey.__context__, context):
         chocolatey._clear_context()
-        # Did it clear all chocolatey items from __context__P?
+        # Did it clear all chocolatey items from __context__?
         assert chocolatey.__context__ == {}
 
 
@@ -333,3 +329,27 @@ def test_list_windowsfeatures_post_2_0_0():
         chocolatey.list_windowsfeatures()
         expected_call = [choco_path, "search", "--source", "windowsfeatures"]
         mock_run.assert_called_with(expected_call, python_shell=False)
+
+
+def test_chocolatey_version():
+    context = {
+        "chocolatey._version": "0.9.9",
+    }
+    with patch.dict(chocolatey.__context__, context):
+        result = chocolatey.chocolatey_version()
+        expected = "0.9.9"
+        assert result == expected
+
+
+def test_chocolatey_version_refresh():
+    context = {"chocolatey._version": "0.9.9"}
+    mock_find = MagicMock(return_value="some_path")
+    mock_run = MagicMock(return_value="2.2.0")
+    with (
+        patch.dict(chocolatey.__context__, context),
+        patch.object(chocolatey, "_find_chocolatey", mock_find),
+        patch.dict(chocolatey.__salt__, {"cmd.run": mock_run}),
+    ):
+        result = chocolatey.chocolatey_version(refresh=True)
+        expected = "2.2.0"
+        assert result == expected