From db54dd2cdd94691a4d0055c387ceba1a69f98bf9 Mon Sep 17 00:00:00 2001
From: Luc Didry <luc@didry.org>
Date: Wed, 12 Feb 2025 16:08:16 +0100
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20=E2=80=94=20Allow=20to=20customize?=
 =?UTF-8?q?=20agent=20User-Agent=20header=20(fix=20#78)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 CHANGELOG.md      |  2 ++
 argos/agent.py    | 12 +++++++++---
 argos/commands.py |  9 +++++++--
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3ba1a1c..fbcc3b6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 
+- ✨ — Allow to customize agent User-Agent header (#78)
+
 ## 0.7.4
 
 Date: 2025-02-12
diff --git a/argos/agent.py b/argos/agent.py
index 195ffdc..70689e5 100644
--- a/argos/agent.py
+++ b/argos/agent.py
@@ -37,11 +37,17 @@ def log_failure(retry_state):
 class ArgosAgent:  # pylint: disable-msg=too-many-instance-attributes
     """The Argos agent is responsible for running the checks and reporting the results."""
 
-    def __init__(self, server: str, auth: str, max_tasks: int, wait_time: int):
+    def __init__(  # pylint: disable-msg=too-many-positional-arguments
+        self, server: str, auth: str, max_tasks: int, wait_time: int, user_agent: str
+    ):
         self.server = server
         self.max_tasks = max_tasks
         self.wait_time = wait_time
         self.auth = auth
+        if user_agent == "":
+            self.ua = user_agent
+        else:
+            self.ua = f" - {user_agent}"
         self._http_client: httpx.AsyncClient | None = None
         self._http_client_v4: httpx.AsyncClient | None = None
         self._http_client_v6: httpx.AsyncClient | None = None
@@ -53,13 +59,13 @@ class ArgosAgent:  # pylint: disable-msg=too-many-instance-attributes
     async def run(self):
         auth_header = {
             "Authorization": f"Bearer {self.auth}",
-            "User-Agent": f"Argos Panoptes agent {VERSION}",
+            "User-Agent": f"Argos Panoptes agent {VERSION}{self.ua}",
         }
         self._http_client = httpx.AsyncClient(headers=auth_header)
 
         ua_header = {
             "User-Agent": f"Argos Panoptes {VERSION} "
-            "(about: https://argos-monitoring.framasoft.org/)",
+            f"(about: https://argos-monitoring.framasoft.org/){self.ua}",
         }
         self._http_client_v4 = httpx.AsyncClient(
             headers=ua_header,
diff --git a/argos/commands.py b/argos/commands.py
index ac1309f..97c244e 100644
--- a/argos/commands.py
+++ b/argos/commands.py
@@ -92,7 +92,12 @@ def version():
     default="INFO",
     type=click.Choice(logging.LOG_LEVELS, case_sensitive=False),
 )
-def agent(server_url, auth, max_tasks, wait_time, log_level):
+@click.option(
+    "--user-agent",
+    default="",
+    help="A custom string to append to the User-Agent header",
+)
+def agent(server_url, auth, max_tasks, wait_time, log_level, user_agent):  # pylint: disable-msg=too-many-positional-arguments
     """Get and run tasks for the provided server. Will wait for new tasks.
 
     Usage: argos agent https://argos.example.org "auth-token-here"
@@ -108,7 +113,7 @@ def agent(server_url, auth, max_tasks, wait_time, log_level):
     from argos.logging import logger
 
     logger.setLevel(log_level)
-    agent_ = ArgosAgent(server_url, auth, max_tasks, wait_time)
+    agent_ = ArgosAgent(server_url, auth, max_tasks, wait_time, user_agent)
     asyncio.run(agent_.run())