[v2] Implement update command

.changes/next-release/feature-update-59722.json

@@ -0,0 +1,5 @@
+{
+  "type": "feature",
+  "category": "``update``",
+  "description": "Adds the ``update`` command, which downloads and installs the latest AWS CLI version."
+}

awscli/autocomplete/local/indexer.py

@@ -32,6 +32,7 @@ class ModelIndexer:
         'login',
         'logout',
         'agent-toolkit',
+        'update',
     ]
 
     _CREATE_CMD_TABLE = """\

awscli/clidriver.py

@@ -93,6 +93,7 @@
 )
 HISTORY_RECORDER = get_global_history_recorder()
 METADATA_FILENAME = 'metadata.json'
+INSTALL_FILENAME = 'install.json'
 _NO_AUTO_PROMPT_ARGS = ['help', '--version']
 _CLI_AUTO_PROMPT_OPTION = '--cli-auto-prompt'
 _NO_CLI_AUTO_PROMPT_OPTION = '--no-cli-auto-prompt'
@@ -166,16 +167,17 @@ def resolve_auto_prompt_mode(args, session):
         return 'off'
 
 
-def _get_distribution_source():
-    metadata_file = os.path.join(
-        os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data'),
-        METADATA_FILENAME,
-    )
-    metadata = {}
-    if os.path.isfile(metadata_file):
-        with open(metadata_file) as f:
-            metadata = json.load(f)
-    return metadata.get('distribution_source', 'other')
+def get_distribution_source():
+    data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
+    for name in (INSTALL_FILENAME, METADATA_FILENAME):
+        path = os.path.join(data_dir, name)
+        if not os.path.isfile(path):
+            continue
+        with open(path) as f:
+            data = json.load(f)
+        if 'distribution_source' in data:
+            return data['distribution_source']
+    return 'other'
 
 
 def _get_distribution():
@@ -199,7 +201,7 @@ def _get_linux_distribution():
 
 def _add_distribution_source_to_user_agent(session):
     add_metadata_component_to_user_agent_extra(
-        session, 'installer', _get_distribution_source()
+        session, 'installer', get_distribution_source()
     )
 
 
@@ -554,7 +556,7 @@ def _cli_version(self):
                 f' exec-env/{os.environ.get("AWS_EXECUTION_ENV")}'
             )
 
-        version_string += f' {_get_distribution_source()}/{platform.machine()}'
+        version_string += f' {get_distribution_source()}/{platform.machine()}'
 
         if linux_distribution := _get_distribution():
             version_string += f'.{linux_distribution}'

awscli/customizations/update.py

@@ -0,0 +1,279 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import ctypes
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+
+from awscli.botocore.awsrequest import AWSRequest
+from awscli.botocore.httpsession import URLLib3Session
+from awscli.clidriver import (
+    INSTALL_FILENAME,
+    get_distribution_source,
+)
+from awscli.compat import is_windows
+from awscli.customizations.commands import BasicCommand
+from awscli.customizations.utils import uni_print
+
+_DOWNLOAD_BASE_URL = 'https://awscli.amazonaws.com'
+
+_SUPPORTED_SOURCES = ('exe', 'script-exe', 'update-exe')
+
+
+def download_with_retry(url, dest, retries=1, session=None):
+    session = session or URLLib3Session()
+    uni_print(f"Downloading {url}\n")
+    request = AWSRequest(method='GET', url=url).prepare()
+    for attempt in range(retries + 1):
+        try:
+            response = session.send(request)
+            if response.status_code != 200:
+                raise UpdateError(
+                    f"unexpected HTTP status {response.status_code}"
+                )
+            with open(dest, 'wb') as out:
+                out.write(response.content)
+            return
+        except Exception as exc:
+            if attempt == retries:
+                raise UpdateError(f"failed to download {url}: {exc}")
+            uni_print(f"download failed ({exc}); retrying...\n", sys.stderr)
+
+
+class BaseUpdateCommand(BasicCommand):
+    NAME = 'update'
+    DESCRIPTION = (
+        'Update the AWS CLI to the latest version.\n\n'
+        'Note that ``update`` is only supported if the '
+        'current CLI instance was installed using an official '
+        'installer, install script, or the ``update`` command. '
+        'Other distribution mechanisms such as source or container '
+        'images are not supported.'
+    )
+    SYNOPSIS = 'aws update'
+    ARG_TABLE = []
+
+    _no_color = False
+
+    def __init__(
+        self, session, source=None, install_metadata=None, downloader=None
+    ):
+        super().__init__(session)
+        self._source = get_distribution_source() if source is None else source
+        self._install_metadata = (
+            self._read_install_json()
+            if install_metadata is None
+            else install_metadata
+        )
+        self._download = downloader or download_with_retry
+
+    def _read_install_json(self):
+        import awscli
+
+        path = os.path.join(
+            os.path.dirname(os.path.abspath(awscli.__file__)),
+            'data',
+            INSTALL_FILENAME,
+        )
+        if not os.path.isfile(path):
+            return {}
+        with open(path) as f:
+            return json.load(f)
+
+    def _run_main(self, parsed_args, parsed_globals):
+        source = self._source
+        if source not in _SUPPORTED_SOURCES:
+            raise UpdateError(
+                f"Detected distribution source: {source}. "
+                f"`aws update` is only supported on AWS CLI instances "
+                f"installed using an official installer, install script, "
+                f"or the `aws update` command."
+            )
+        uni_print(f"Updating AWS CLI (source: {source})\n")
+        self._no_color = parsed_globals.color == 'off'
+        self._do_update()
+        return 0
+
+    def _do_update(self):
+        raise NotImplementedError
+
+
+class UnixUpdateCommand(BaseUpdateCommand):
+    SCRIPT_URL = f'{_DOWNLOAD_BASE_URL}/v2/install.sh'
+    SYSTEM_INSTALL_DIR = '/usr/local/aws-cli'
+
+    def __init__(self, session, is_elevated=None, runner=None, **kwargs):
+        super().__init__(session, **kwargs)
+        self._is_elevated = is_elevated
+        self._run_update = runner or self._run_install
+
+    def _do_update(self):
+        install_dir = self._install_metadata.get('install_dir')
+        if not install_dir:
+            raise UpdateError(
+                'Install-time metadata could not be found. Reinstall '
+                'using the install script or an official installer '
+                'and try again.'
+            )
+        bin_dir = self._install_metadata.get('bin_dir')
+        is_system = self._is_system_install(self._install_metadata)
+        if is_system:
+            self._assert_elevated()
+        with tempfile.TemporaryDirectory() as tmp:
+            script_path = os.path.join(tmp, 'install.sh')
+            self._download(self.SCRIPT_URL, script_path)
+            env = os.environ.copy()
+            env['AWS_CLI_DISTRIBUTION_SOURCE_OVERRIDE'] = 'update-exe'
+            if self._no_color:
+                env['NO_COLOR'] = '1'
+            cmd = ['bash', script_path]
+            if is_system:
+                cmd.append('--system')
+            else:
+                env['XDG_DATA_HOME'] = os.path.dirname(install_dir)
+                if bin_dir:
+                    env['XDG_BIN_HOME'] = bin_dir
+                else:
+                    env['XDG_BIN_HOME'] = os.path.join(tmp, 'bin')
+                    env['AWS_CLI_NO_BIN_DIR'] = '1'
+            uni_print('Running install script...\n')
+            try:
+                self._run_update(cmd, env)
+            except subprocess.CalledProcessError as exc:
+                raise UpdateError(
+                    f"install script failed with exit code {exc.returncode}"
+                )
+
+    def _run_install(self, cmd, env):
+        subprocess.run(cmd, env=env, check=True)
+
+    def _is_system_install(self, install_metadata):
+        if 'script_install' in install_metadata:
+            return install_metadata['script_install'].get('system', False)
+        if not sys.executable:
+            return False
+        aws_bin = os.path.realpath(sys.executable)
+        return aws_bin.startswith(self.SYSTEM_INSTALL_DIR + os.sep)
+
+    def _assert_elevated(self):
+        elevated = (
+            os.geteuid() == 0
+            if self._is_elevated is None
+            else self._is_elevated
+        )
+        if not elevated:
+            raise UpdateError(
+                'Updating a system-wide AWS CLI installation requires root.'
+            )
+
+
+class WindowsUpdateCommand(BaseUpdateCommand):
+    SCRIPT_URL = f'{_DOWNLOAD_BASE_URL}/v2/install.ps1'
+
+    def __init__(
+        self,
+        session,
+        is_elevated=None,
+        runner=None,
+        powershell_path=None,
+        **kwargs,
+    ):
+        super().__init__(session, **kwargs)
+        self._is_elevated = is_elevated
+        self._run_update = runner or self._run_install
+        self._powershell_path = powershell_path
+
+    def _do_update(self):
+        is_system = self._is_system_install(self._install_metadata)
+        if is_system:
+            self._assert_elevated()
+
+        tmp = tempfile.mkdtemp()
+        script_path = os.path.join(tmp, 'install.ps1')
+        self._download(self.SCRIPT_URL, script_path)
+
+        wrapper_path = os.path.join(tmp, 'aws-update.cmd')
+        ps_exe = self._powershell_path or self._find_powershell()
+        ps_args = f'-NoProfile -File "{script_path}"'
+        if is_system:
+            ps_args += ' -System'
+
+        with open(wrapper_path, 'w') as f:
+            # Windows acquires a lock when running an exe process, preventing
+            # an update in-place. The workaround is to launch a detached CMD
+            # subprocess and exit the parent early. Use ping to wait before
+            # running the installation to ensure the parent isn't holding onto
+            # the lock.
+            f.write('@echo off\n')
+            f.write('set AWS_CLI_DISTRIBUTION_SOURCE_OVERRIDE=update-exe\n')
+            if self._no_color:
+                f.write('set NO_COLOR=1\n')
+            f.write('ping -n 3 127.0.0.1 >nul 2>&1\n')
+            f.write(f'"{ps_exe}" {ps_args}\n')
+
+        self._run_update(['cmd', '/c', wrapper_path])
+        uni_print(
+            'Update started. This process will exit and the '
+            'update will complete shortly.\n'
+        )
+
+    def _run_install(self, cmd):
+        subprocess.Popen(
+            cmd,
+            creationflags=subprocess.DETACHED_PROCESS
+            | subprocess.CREATE_NEW_PROCESS_GROUP,
+        )
+
+    def _find_powershell(self):
+        path = shutil.which('powershell')
+        if not path:
+            raise UpdateError('powershell.exe not found on PATH.')
+        return path
+
+    def _is_system_install(self, install_metadata):
+        if 'script_install' in install_metadata:
+            return install_metadata['script_install'].get('system', False)
+        # AWS CLI is always installed to 'C:\Program Files\Amazon\AWSCLIV2'
+        # for all users. 'ProgramW6432' always points to 'C:\Program Files'
+        # while 'ProgramFiles' dynamically points to 'C:\Program Files' or
+        # 'C:\Program Files (x86)', depending on the architecture.
+        # Prefer 'ProgramW6432' and only read 'ProgramFiles' as a fallback.
+        program_files = os.environ.get('ProgramW6432') or os.environ.get(
+            'ProgramFiles'
+        )
+        if not program_files:
+            return False
+        canonical = os.path.join(
+            program_files, 'Amazon', 'AWSCLIV2', 'aws.exe'
+        )
+        buf = ctypes.create_unicode_buffer(32768)
+        ctypes.windll.kernel32.GetModuleFileNameW(None, buf, len(buf))
+        return os.path.normcase(buf.value) == os.path.normcase(canonical)
+
+    def _assert_elevated(self):
+        elevated = (
+            bool(ctypes.windll.shell32.IsUserAnAdmin())
+            if self._is_elevated is None
+            else self._is_elevated
+        )
+        if not elevated:
+            raise UpdateError(
+                'Updating a system-wide AWS CLI install requires an '
+                'elevated shell. Re-run from an Administrator prompt.'
+            )
+
+
+UpdateCommand = WindowsUpdateCommand if is_windows else UnixUpdateCommand
+
+
+def register_update_command(event_handlers):
+    event_handlers.register(
+        'building-command-table.main', UpdateCommand.add_command
+    )
+
+
+class UpdateError(Exception):
+    pass

awscli/handlers_registry.py

@@ -677,6 +677,7 @@ class CommandTableOp(enum.Enum):
         ('awscli.customizations.history', 'register_history_commands'),
         ('awscli.customizations.devcommands', 'register_dev_commands'),
         ('awscli.customizations.login', 'register_login_cmds'),
+        ('awscli.customizations.update', 'register_update_command'),
     ],
     'building-command-table.polly': [
         ('awscli.customizations.removals', 'register_removals')
@@ -909,4 +910,10 @@ class CommandTableOp(enum.Enum):
         'agenttoolkit',
         'agent-toolkit',
     ),
+    (
+        CommandTableOp.ADD,
+        'update',
+        'awscli.customizations.update',
+        'UpdateCommand',
+    ),
 ]

exe/assets/install

@@ -136,6 +136,19 @@ create_bin_symlinks() {
   ln -sf "$CURRENT_AWS_COMPLETER_EXE" "$BIN_AWS_COMPLETER_EXE"
 }
 
+write_install_json() {
+  if [ "${AWS_CLI_SKIP_INSTALL_JSON:-}" = "1" ]; then
+    return 0
+  fi
+  install_json="$INSTALL_DIST_DIR/awscli/data/install.json"
+  cat > "$install_json" <<EOF
+{
+  "install_dir": "$ROOT_INSTALL_DIR",
+  "bin_dir": "$BIN_DIR"
+}
+EOF
+}
+
 die() {
 	err_msg="$1"
 	echo "$err_msg" >&2
@@ -148,6 +161,7 @@ main() {
   check_preexisting_install
   create_install_dir
   create_bin_symlinks
+  write_install_json
   echo "You can now run: $BIN_AWS_EXE --version"
   exit 0
 }