Catch bare exceptions, shorten id, add md5 fallback

Author: hssyoo

awscli/telemetry.py

@@ -10,7 +10,6 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-import hashlib
 import io
 import os
 import sqlite3
@@ -22,6 +21,8 @@
 from functools import cached_property
 from pathlib import Path
 
+from botocore.compat import get_md5
+from botocore.exceptions import MD5UnavailableError
 from botocore.useragent import UserAgentComponent
 
 from awscli.compat import is_windows
@@ -30,10 +31,22 @@
 _CACHE_DIR = Path.home() / '.aws' / 'cli' / 'cache'
 _DATABASE_FILENAME = 'session.db'
 _SESSION_LENGTH_SECONDS = 60 * 30
+_SESSION_ID_LENGTH = 12
 
 _CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
 
+def _get_checksum():
+    hashlib_params = {"usedforsecurity": False}
+    try:
+        checksum = get_md5(**hashlib_params)
+    except MD5UnavailableError:
+        import hashlib
+
+        checksum = hashlib.sha256(**hashlib_params)
+    return checksum
+
+
 @dataclass
 class CLISessionData:
     key: str
@@ -179,17 +192,19 @@ def sweep(self, timestamp):
 
 class CLISessionGenerator:
     def generate_session_id(self, host_id, tty, timestamp):
-        return self._generate_md5_hash(host_id, tty, timestamp)
+        return self._generate_checksum(host_id, tty, timestamp)
 
     def generate_cache_key(self, host_id, tty):
-        return self._generate_md5_hash(host_id, tty)
+        return self._generate_checksum(host_id, tty)
 
-    def _generate_md5_hash(self, *args):
+    def _generate_checksum(self, *args):
+        checksum = _get_checksum()
         str_to_hash = ""
         for arg in args:
             if arg is not None:
                 str_to_hash += str(arg)
-        return hashlib.md5(str_to_hash.encode('utf-8')).hexdigest()
+        checksum.update(str_to_hash.encode('utf-8'))
+        return checksum.hexdigest()[:_SESSION_ID_LENGTH]
 
 
 class CLISessionOrchestrator:
@@ -273,7 +288,16 @@ def _get_cli_session_orchestrator():
 
 
 def add_session_id_component_to_user_agent_extra(session, orchestrator=None):
-    cli_session_orchestrator = orchestrator or _get_cli_session_orchestrator()
-    add_component_to_user_agent_extra(
-        session, UserAgentComponent("sid", cli_session_orchestrator.session_id)
-    )
+    try:
+        cli_session_orchestrator = (
+            orchestrator or _get_cli_session_orchestrator()
+        )
+        add_component_to_user_agent_extra(
+            session,
+            UserAgentComponent("sid", cli_session_orchestrator.session_id),
+        )
+    except Exception:
+        # Ideally, the AWS CLI should never throw if the session id
+        # can't be generated since it's not critical for users. Issues
+        # with session dada should instead be caught server-side.
+        pass

tests/functional/test_telemetry.py

@@ -14,6 +14,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from botocore.exceptions import MD5UnavailableError
 from botocore.session import Session
 
 from awscli.telemetry import (
@@ -182,14 +183,24 @@ def test_generate_session_id(self, session_generator):
             'my-tty',
             1000000000,
         )
-        assert session_id == 'd949713b13ee3fb52983b04316e8e6b5'
+        assert session_id == 'd949713b13ee'
 
     def test_generate_cache_key(self, session_generator):
         cache_key = session_generator.generate_cache_key(
             'my-hostname',
             'my-tty',
         )
-        assert cache_key == 'b1ca2be0ffac12f172933b6777e06f2c'
+        assert cache_key == 'b1ca2be0ffac'
+
+    @patch('awscli.telemetry.get_md5')
+    def test_checksum_fips_fallback(self, patched_get_md5, session_generator):
+        patched_get_md5.side_effect = MD5UnavailableError()
+        session_id = session_generator.generate_session_id(
+            'my-hostname',
+            'my-tty',
+            1000000000,
+        )
+        assert session_id == '183b154db015'
 
 
 @skip_if_windows
@@ -211,7 +222,7 @@ def test_session_id_gets_cached(
         orchestrator = CLISessionOrchestrator(
             session_generator, session_writer, session_reader, session_sweeper
         )
-        assert orchestrator.session_id == '881cea8546fa4888970cce8d133c3bf9'
+        assert orchestrator.session_id == '881cea8546fa'
 
         session_data = session_reader.read(orchestrator.cache_key)
         assert session_data.key == orchestrator.cache_key
@@ -299,3 +310,15 @@ def test_add_session_id_component_to_user_agent_extra():
     orchestrator.session_id = 'my-session-id'
     add_session_id_component_to_user_agent_extra(session, orchestrator)
     assert session.user_agent_extra == 'sid/my-session-id'
+
+
+def test_entrypoint_catches_bare_exceptions():
+    class FakeOrchestrator(CLISessionOrchestrator):
+        def __init__(self):
+            pass
+
+        def session_id(self):
+            raise Exception()
+
+    session = MagicMock(Session)
+    add_session_id_component_to_user_agent_extra(session, FakeOrchestrator())