From fe2068697cb90e5ceafe31f8db48b6965fa5d4c1 Mon Sep 17 00:00:00 2001
From: Jussi Kukkonen <jkukkonen@google.com>
Date: Fri, 19 Apr 2024 17:53:34 +0300
Subject: [PATCH] Support app-specific user-agents

* application user-agent can be set with UpdaterConfig object
* Setting will affect the default fetcher only
* the application user-agent will be prefixed to the ngclient
  default user-agent

Signed-off-by: Jussi Kukkonen <jkukkonen@google.com>
---
 tuf/ngclient/_internal/requests_fetcher.py | 10 ++++++++--
 tuf/ngclient/config.py                     |  4 ++++
 tuf/ngclient/updater.py                    | 13 ++++++++++---
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/tuf/ngclient/_internal/requests_fetcher.py b/tuf/ngclient/_internal/requests_fetcher.py
index 1994729f..c931b85a 100644
--- a/tuf/ngclient/_internal/requests_fetcher.py
+++ b/tuf/ngclient/_internal/requests_fetcher.py
@@ -10,7 +10,7 @@
 # can be moved out of _internal once sigstore-python 1.0 is not relevant.
 
 import logging
-from typing import Dict, Iterator, Tuple
+from typing import Dict, Iterator, Optional, Tuple
 from urllib import parse
 
 # Imports
@@ -35,7 +35,10 @@ class RequestsFetcher(FetcherInterface):
     """
 
     def __init__(
-        self, socket_timeout: int = 30, chunk_size: int = 400000
+        self,
+        socket_timeout: int = 30,
+        chunk_size: int = 400000,
+        app_user_agent: Optional[str] = None,
     ) -> None:
         # http://docs.python-requests.org/en/master/user/advanced/#session-objects:
         #
@@ -56,6 +59,7 @@ def __init__(
         # Default settings
         self.socket_timeout: int = socket_timeout  # seconds
         self.chunk_size: int = chunk_size  # bytes
+        self.app_user_agent = app_user_agent
 
     def _fetch(self, url: str) -> Iterator[bytes]:
         """Fetch the contents of HTTP/HTTPS url from a remote server.
@@ -138,6 +142,8 @@ def _get_session(self, url: str) -> requests.Session:
             self._sessions[session_index] = session
 
             ua = f"tuf/{tuf.__version__} {session.headers['User-Agent']}"
+            if self.app_user_agent is not None:
+                ua = f"{self.app_user_agent} {ua}"
             session.headers["User-Agent"] = ua
 
             logger.debug("Made new session %s", session_index)
diff --git a/tuf/ngclient/config.py b/tuf/ngclient/config.py
index 943018fc..8019c4d2 100644
--- a/tuf/ngclient/config.py
+++ b/tuf/ngclient/config.py
@@ -5,6 +5,7 @@
 
 from dataclasses import dataclass
 from enum import Flag, unique
+from typing import Optional
 
 
 @unique
@@ -39,6 +40,8 @@ class UpdaterConfig:
         envelope_type: Configures deserialization and verification mode of TUF
             metadata. Per default, it is treated as traditional canonical JSON
             -based TUF Metadata.
+        app_user_agent: Application user agent, e.g. "MyApp/1.0.0". This will be
+            prefixed to ngclient user agent when the default fetcher is used.
     """
 
     max_root_rotations: int = 32
@@ -49,3 +52,4 @@ class UpdaterConfig:
     targets_max_length: int = 5000000  # bytes
     prefix_targets_with_hash: bool = True
     envelope_type: EnvelopeType = EnvelopeType.METADATA
+    app_user_agent: Optional[str] = None
diff --git a/tuf/ngclient/updater.py b/tuf/ngclient/updater.py
index 666e54d3..145074aa 100644
--- a/tuf/ngclient/updater.py
+++ b/tuf/ngclient/updater.py
@@ -93,11 +93,15 @@ def __init__(
         else:
             self._target_base_url = _ensure_trailing_slash(target_base_url)
 
-        # Read trusted local root metadata
-        data = self._load_local_metadata(Root.type)
-        self._fetcher = fetcher or requests_fetcher.RequestsFetcher()
         self.config = config or UpdaterConfig()
 
+        if fetcher is not None:
+            self._fetcher = fetcher
+        else:
+            self._fetcher = requests_fetcher.RequestsFetcher(
+                app_user_agent=self.config.app_user_agent
+            )
+
         supported_envelopes = [EnvelopeType.METADATA, EnvelopeType.SIMPLE]
         if self.config.envelope_type not in supported_envelopes:
             raise ValueError(
@@ -105,6 +109,9 @@ def __init__(
                 f"got '{self.config.envelope_type}'"
             )
 
+        # Read trusted local root metadata
+        data = self._load_local_metadata(Root.type)
+
         self._trusted_set = trusted_metadata_set.TrustedMetadataSet(
             data, self.config.envelope_type
         )