Use uuid4 instead of hostname

Author: hssyoo

awscli/telemetry.py

@@ -13,11 +13,11 @@
 import hashlib
 import io
 import os
-import socket
 import sqlite3
 import sys
 import threading
 import time
+import uuid
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
@@ -49,6 +49,17 @@ class CLISessionDatabaseConnection:
           timestamp INTEGER NOT NULL
         )
     """
+    _CREATE_HOST_ID_TABLE = """
+        CREATE TABLE IF NOT EXISTS host_id (
+          key INTEGER PRIMARY KEY,
+          id TEXT UNIQUE NOT NULL
+        )
+    """
+    _INSERT_HOST_ID = """
+        INSERT OR IGNORE INTO host_id (
+            key, id
+        ) VALUES (?, ?)
+    """
     _ENABLE_WAL = 'PRAGMA journal_mode=WAL'
 
     def __init__(self, connection=None):
@@ -69,12 +80,28 @@ def execute(self, query, *parameters):
             return sqlite3.Cursor(self._connection)
 
     def _ensure_database_setup(self):
-        self._create_record_table()
+        self._create_session_table()
+        self._create_host_id_table()
+        self._ensure_host_id()
         self._try_to_enable_wal()
 
-    def _create_record_table(self):
+    def _create_session_table(self):
         self.execute(self._CREATE_TABLE)
 
+    def _create_host_id_table(self):
+        self.execute(self._CREATE_HOST_ID_TABLE)
+
+    def _ensure_host_id(self):
+        self.execute(
+            self._INSERT_HOST_ID,
+            # Hardcode `0` as primary key to ensure
+            # there's only ever 1 host id in the table.
+            (
+                0,
+                str(uuid.uuid4()),
+            ),
+        )
+
     def _try_to_enable_wal(self):
         try:
             self.execute(self._ENABLE_WAL)
@@ -111,6 +138,11 @@ class CLISessionDatabaseReader:
         FROM session
         WHERE key = ?
     """
+    _READ_HOST_ID = """
+        SELECT id
+        FROM host_id
+        WHERE key = 0
+    """
 
     def __init__(self, connection):
         self._connection = connection
@@ -122,6 +154,10 @@ def read(self, key):
             return
         return CLISessionData(*result)
 
+    def read_host_id(self):
+        cursor = self._connection.execute(self._READ_HOST_ID)
+        return cursor.fetchone()[0]
+
 
 class CLISessionDatabaseSweeper:
     _DELETE_RECORDS = """
@@ -142,11 +178,11 @@ def sweep(self, timestamp):
 
 
 class CLISessionGenerator:
-    def generate_session_id(self, hostname, tty, timestamp):
-        return self._generate_md5_hash(hostname, tty, timestamp)
+    def generate_session_id(self, host_id, tty, timestamp):
+        return self._generate_md5_hash(host_id, tty, timestamp)
 
-    def generate_cache_key(self, hostname, tty):
-        return self._generate_md5_hash(hostname, tty)
+    def generate_cache_key(self, host_id, tty):
+        return self._generate_md5_hash(host_id, tty)
 
     def _generate_md5_hash(self, *args):
         str_to_hash = ""
@@ -167,12 +203,12 @@ def __init__(self, generator, writer, reader, sweeper):
 
     @cached_property
     def cache_key(self):
-        return self._generator.generate_cache_key(self._hostname, self._tty)
+        return self._generator.generate_cache_key(self._host_id, self._tty)
 
     @cached_property
     def _session_id(self):
         return self._generator.generate_session_id(
-            self._hostname, self._tty, self._timestamp
+            self._host_id, self._tty, self._timestamp
         )
 
     @cached_property
@@ -206,12 +242,12 @@ def _tty(self):
         except (OSError, io.UnsupportedOperation):
             # Standard input was redirected to a pseudofile.
             # This can happen when running tests on IDEs or
-            # running scripts with redirected input.
+            # running scripts with redirected input, etc.
             return
 
     @cached_property
-    def _hostname(self):
-        return socket.gethostname()
+    def _host_id(self):
+        return self._reader.read_host_id()
 
     @cached_property
     def _timestamp(self):

tests/functional/test_telemetry.py

@@ -47,6 +47,14 @@ def session_conn():
             ) VALUES ('first_key', 'first_id', 5555555555)
         """
     )
+    # Overwrite host id with deterministic value for testing.
+    conn.execute(
+        """
+            INSERT OR REPLACE INTO host_id (
+                key, id
+            ) VALUES (0, 'my-hostname')
+        """
+    )
     return conn
 
 
@@ -93,11 +101,10 @@ def test_ensure_database_setup(self, session_conn):
             """
                 SELECT name
                 FROM sqlite_master
-                WHERE type='table'
-                AND name='session';
+                WHERE type='table';
             """
         )
-        assert cursor.fetchall() == [('session',)]
+        assert cursor.fetchall() == [('session',), ('host_id',)]
 
     def test_timeout_does_not_raise_exception(self, session_conn):
         class FakeConnection(sqlite3.Connection):
@@ -144,6 +151,10 @@ def test_read_nonexistent_record(self, session_reader):
         session_data = session_reader.read('bad_key')
         assert session_data is None
 
+    def test_read_host_id(self, session_reader):
+        host_id = session_reader.read_host_id()
+        assert host_id == 'my-hostname'
+
 
 class TestCLISessionDatabaseSweeper:
     def test_sweep(self, expired_data, session_reader, session_sweeper):
@@ -184,13 +195,11 @@ def test_generate_cache_key(self, session_generator):
 @skip_if_windows
 @patch('sys.stdin')
 @patch('time.time', return_value=5555555555)
-@patch('socket.gethostname', return_value='my-hostname')
 @patch('os.ttyname', return_value='my-tty')
 class TestCLISessionOrchestrator:
     def test_session_id_gets_cached(
         self,
         patched_tty_name,
-        patched_hostname,
         patched_time,
         patched_stdin,
         session_sweeper,
@@ -212,7 +221,6 @@ def test_session_id_gets_cached(
     def test_cached_session_id_updated_if_expired(
         self,
         patched_tty_name,
-        patched_hostname,
         patched_time,
         patched_stdin,
         session_sweeper,
@@ -249,7 +257,6 @@ def test_cached_session_id_updated_if_expired(
     def test_cached_session_id_not_updated_if_valid(
         self,
         patched_tty_name,
-        patched_hostname,
         patched_time,
         patched_stdin,
         session_sweeper,