-
Notifications
You must be signed in to change notification settings - Fork 4.3k
[v2] Add session id to user agent string #9498
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
10 commits
Select commit
Hold shift + click to select a range
904421e
Add session id to user agent string
hssyoo c9c6253
Use uuid4 instead of hostname
hssyoo f8109bd
Catch bare exceptions, shorten id, add md5 fallback
hssyoo 6aab28d
Move mkdir for cache dir
hssyoo f1f74e4
Fix typo
hssyoo f7d1d1c
Mock orchestrator
hssyoo a5cf8d3
Check table before writing host id
hssyoo b7f3215
Check result before grabbing index val
hssyoo f9efcd3
Actually just fix the test
hssyoo 5086487
Fix indent
hssyoo File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,5 @@ | ||
{ | ||
"type": "enhancement", | ||
"category": "User Agent", | ||
"description": "Append session id to user agent string" | ||
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,311 @@ | ||
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. | ||
# | ||
# Licensed under the Apache License, Version 2.0 (the "License"). You | ||
# may not use this file except in compliance with the License. A copy of | ||
# the License is located at | ||
# | ||
# http://aws.amazon.com/apache2.0/ | ||
# | ||
# or in the "license" file accompanying this file. This file is | ||
# 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 io | ||
import os | ||
import sqlite3 | ||
import sys | ||
import threading | ||
import time | ||
import uuid | ||
from dataclasses import dataclass | ||
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 | ||
from awscli.utils import add_component_to_user_agent_extra | ||
|
||
_CACHE_DIR = Path.home() / '.aws' / 'cli' / 'cache' | ||
_DATABASE_FILENAME = 'session.db' | ||
_SESSION_LENGTH_SECONDS = 60 * 30 | ||
_SESSION_ID_LENGTH = 12 | ||
|
||
|
||
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 | ||
session_id: str | ||
timestamp: int | ||
|
||
|
||
class CLISessionDatabaseConnection: | ||
_CREATE_TABLE = """ | ||
CREATE TABLE IF NOT EXISTS session ( | ||
key TEXT PRIMARY KEY, | ||
session_id TEXT NOT NULL, | ||
timestamp INTEGER NOT NULL | ||
) | ||
""" | ||
_CREATE_HOST_ID_TABLE = """ | ||
CREATE TABLE IF NOT EXISTS host_id ( | ||
key INTEGER PRIMARY KEY, | ||
id TEXT UNIQUE NOT NULL | ||
) | ||
""" | ||
_CHECK_HOST_ID = """ | ||
SELECT COUNT(*) FROM host_id | ||
""" | ||
_INSERT_HOST_ID = """ | ||
INSERT OR IGNORE INTO host_id ( | ||
key, id | ||
) VALUES (?, ?) | ||
""" | ||
_ENABLE_WAL = 'PRAGMA journal_mode=WAL' | ||
|
||
def __init__(self, connection=None): | ||
self._connection = connection or sqlite3.connect( | ||
_CACHE_DIR / _DATABASE_FILENAME, | ||
check_same_thread=False, | ||
isolation_level=None, | ||
) | ||
self._ensure_cache_dir() | ||
self._ensure_database_setup() | ||
|
||
def execute(self, query, *parameters): | ||
try: | ||
return self._connection.execute(query, *parameters) | ||
except sqlite3.OperationalError: | ||
# Process timed out waiting for database lock. | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
# Return any empty `Cursor` object instead of | ||
# raising an exception. | ||
return sqlite3.Cursor(self._connection) | ||
|
||
def _ensure_cache_dir(self): | ||
_CACHE_DIR.mkdir(parents=True, exist_ok=True) | ||
|
||
def _ensure_database_setup(self): | ||
self._create_session_table() | ||
self._create_host_id_table() | ||
self._ensure_host_id() | ||
self._try_to_enable_wal() | ||
|
||
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): | ||
cur = self.execute(self._CHECK_HOST_ID) | ||
host_id_ct = cur.fetchone()[0] | ||
if host_id_ct == 0: | ||
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) | ||
except sqlite3.Error: | ||
# This is just a performance enhancement so it is optional. Not all | ||
# systems will have a sqlite compiled with the WAL enabled. | ||
pass | ||
|
||
|
||
class CLISessionDatabaseWriter: | ||
_WRITE_RECORD = """ | ||
INSERT OR REPLACE INTO session ( | ||
key, session_id, timestamp | ||
) VALUES (?, ?, ?) | ||
""" | ||
|
||
def __init__(self, connection): | ||
self._connection = connection | ||
|
||
def write(self, data): | ||
self._connection.execute( | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
self._WRITE_RECORD, | ||
( | ||
data.key, | ||
data.session_id, | ||
data.timestamp, | ||
), | ||
) | ||
|
||
|
||
class CLISessionDatabaseReader: | ||
_READ_RECORD = """ | ||
SELECT * | ||
FROM session | ||
WHERE key = ? | ||
""" | ||
_READ_HOST_ID = """ | ||
SELECT id | ||
FROM host_id | ||
WHERE key = 0 | ||
""" | ||
|
||
def __init__(self, connection): | ||
self._connection = connection | ||
|
||
def read(self, key): | ||
cursor = self._connection.execute(self._READ_RECORD, (key,)) | ||
result = cursor.fetchone() | ||
if result is None: | ||
return | ||
return CLISessionData(*result) | ||
|
||
def read_host_id(self): | ||
cursor = self._connection.execute(self._READ_HOST_ID) | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
return cursor.fetchone()[0] | ||
|
||
|
||
class CLISessionDatabaseSweeper: | ||
_DELETE_RECORDS = """ | ||
DELETE FROM session | ||
WHERE timestamp < ? | ||
""" | ||
|
||
def __init__(self, connection): | ||
self._connection = connection | ||
|
||
def sweep(self, timestamp): | ||
try: | ||
self._connection.execute(self._DELETE_RECORDS, (timestamp,)) | ||
except Exception: | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
# This is just a background cleanup task. No need to | ||
# handle it or direct to stderr. | ||
return | ||
|
||
|
||
class CLISessionGenerator: | ||
def generate_session_id(self, host_id, tty, timestamp): | ||
return self._generate_checksum(host_id, tty, timestamp) | ||
|
||
def generate_cache_key(self, host_id, tty): | ||
return self._generate_checksum(host_id, tty) | ||
|
||
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) | ||
checksum.update(str_to_hash.encode('utf-8')) | ||
return checksum.hexdigest()[:_SESSION_ID_LENGTH] | ||
|
||
|
||
class CLISessionOrchestrator: | ||
def __init__(self, generator, writer, reader, sweeper): | ||
self._generator = generator | ||
self._writer = writer | ||
self._reader = reader | ||
self._sweeper = sweeper | ||
|
||
self._sweep_cache() | ||
|
||
@cached_property | ||
def cache_key(self): | ||
return self._generator.generate_cache_key(self._host_id, self._tty) | ||
|
||
@cached_property | ||
def _session_id(self): | ||
return self._generator.generate_session_id( | ||
self._host_id, self._tty, self._timestamp | ||
) | ||
|
||
@cached_property | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
def session_id(self): | ||
if (cached_data := self._reader.read(self.cache_key)) is not None: | ||
# Cache hit, but session id is expired. Generate new id and update. | ||
if ( | ||
cached_data.timestamp + _SESSION_LENGTH_SECONDS | ||
< self._timestamp | ||
): | ||
cached_data.session_id = self._session_id | ||
# Always update the timestamp to last used. | ||
cached_data.timestamp = self._timestamp | ||
self._writer.write(cached_data) | ||
return cached_data.session_id | ||
# Cache miss, generate and write new record. | ||
session_id = self._session_id | ||
session_data = CLISessionData( | ||
self.cache_key, session_id, self._timestamp | ||
) | ||
self._writer.write(session_data) | ||
return session_id | ||
|
||
@cached_property | ||
def _tty(self): | ||
# os.ttyname is only available on Unix platforms. | ||
if is_windows: | ||
return | ||
try: | ||
return os.ttyname(sys.stdin.fileno()) | ||
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, etc. | ||
return | ||
|
||
@cached_property | ||
def _host_id(self): | ||
return self._reader.read_host_id() | ||
|
||
@cached_property | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
def _timestamp(self): | ||
return int(time.time()) | ||
|
||
def _sweep_cache(self): | ||
t = threading.Thread( | ||
target=self._sweeper.sweep, | ||
args=(self._timestamp - _SESSION_LENGTH_SECONDS,), | ||
daemon=True, | ||
) | ||
t.start() | ||
|
||
|
||
def _get_cli_session_orchestrator(): | ||
conn = CLISessionDatabaseConnection() | ||
return CLISessionOrchestrator( | ||
CLISessionGenerator(), | ||
CLISessionDatabaseWriter(conn), | ||
CLISessionDatabaseReader(conn), | ||
CLISessionDatabaseSweeper(conn), | ||
) | ||
|
||
|
||
def add_session_id_component_to_user_agent_extra(session, orchestrator=None): | ||
try: | ||
cli_session_orchestrator = ( | ||
orchestrator or _get_cli_session_orchestrator() | ||
) | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
add_component_to_user_agent_extra( | ||
session, | ||
UserAgentComponent("sid", cli_session_orchestrator.session_id), | ||
) | ||
except Exception: | ||
aemous marked this conversation as resolved.
Show resolved
Hide resolved
|
||
# 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 data should instead be caught server-side. | ||
pass |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.