Skip to content

Commit 7782def

Browse files
committed
Add --apn-id global arg for PRM user-agent attribution
Add a new --apn-id global CLI argument that appends an AWS Partner Network identifier to the User-Agent header in the format APN_1.1/<value>$ as required by the PRM onboarding guide. The user provides the full identifier including the type prefix (e.g. pc_PRODUCTCODE for product codes, ra_PARTNERID for referral agents) to support current and future APN ID types. The value is resolved from (highest precedence first): 1. --apn-id CLI flag 2. AWS_APN_ID environment variable 3. apn_id setting in the shared config file profile The handler is registered on the session-initialized event rather than top-level-args-parsed to ensure the profile has been applied to the session before reading config file values. Validation uses a simple character set check (no regex) to avoid any ReDoS concerns: 1-255 ASCII alphanumeric chars plus _ and -. Fixes: #10476
1 parent b0a12eb commit 7782def

5 files changed

Lines changed: 224 additions & 0 deletions

File tree

awscli/customizations/globalargs.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ def register_parse_global_args(cli):
4242
resolve_cli_connect_timeout,
4343
unique_id='resolve-cli-connect-timeout',
4444
)
45+
cli.register(
46+
'session-initialized',
47+
resolve_apn_id,
48+
unique_id='resolve-apn-id',
49+
)
4550

4651

4752
def resolve_types(parsed_args, **kwargs):
@@ -133,3 +138,53 @@ def _update_default_client_config(session, arg_name, arg_value):
133138
if current_default_config is not None:
134139
new_default_config = current_default_config.merge(new_default_config)
135140
session.set_default_client_config(new_default_config)
141+
142+
143+
# Allowed characters for APN ID values: ASCII letters, digits,
144+
# underscore, and hyphen. Simple set check avoids regex entirely.
145+
_APN_ID_ALLOWED_CHARS = frozenset(
146+
'abcdefghijklmnopqrstuvwxyz' 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' '0123456789' '_-'
147+
)
148+
149+
150+
def _is_valid_apn_id(value):
151+
"""Validate an APN ID value: 1-255 chars from the allowed set."""
152+
if not value or len(value) > 255:
153+
return False
154+
return all(c in _APN_ID_ALLOWED_CHARS for c in value)
155+
156+
157+
def resolve_apn_id(parsed_args, session, **kwargs):
158+
"""Resolve the APN ID from CLI flag, env var, or config file.
159+
160+
Registered on 'session-initialized' so the profile is already set
161+
and config file values are accessible.
162+
"""
163+
apn_id = getattr(parsed_args, 'apn_id', None)
164+
if apn_id is None:
165+
apn_id = os.environ.get('AWS_APN_ID')
166+
if apn_id is None:
167+
try:
168+
scoped_config = session.get_scoped_config()
169+
apn_id = scoped_config.get('apn_id')
170+
except Exception:
171+
pass
172+
if not apn_id:
173+
return
174+
if not _is_valid_apn_id(apn_id):
175+
raise ParamValidationError(
176+
"Bad value for --apn-id %r: must be 1-255 characters and may "
177+
"only contain letters, digits, '_', and '-'." % apn_id
178+
)
179+
# Format defined by the AWS PRM onboarding guide:
180+
# https://docs.aws.amazon.com/PRM/latest/aws-prm-onboarding-guide/user-agent-string.html
181+
# The trailing '$' is a required end delimiter, not a regex anchor.
182+
# The user provides the full identifier including the type prefix
183+
# (e.g. "pc_PRODUCTCODE" or "ra_ATTRIBUTIONID").
184+
marker = 'APN_1.1/%s$' % apn_id
185+
existing = session.user_agent_extra or ''
186+
if marker in existing.split():
187+
return
188+
session.user_agent_extra = (
189+
'%s %s' % (existing, marker) if existing else marker
190+
)

awscli/data/cli.json

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,10 @@
9797
"enhanced"
9898
],
9999
"help": "<p>The formatting style for error output. By default, errors are displayed in enhanced format.</p>"
100+
},
101+
"apn-id": {
102+
"dest": "apn_id",
103+
"help": "<p>An AWS Partner Network identifier used for Partner Revenue Measurement (PRM) attribution. The value should include the type prefix (e.g. <code>pc_PRODUCTCODE</code> for product codes or <code>ra_PARTNERID</code> for referral agents). The value is appended to the User-Agent header of all requests as <code>APN_1.1/&lt;value&gt;$</code>. Overrides the <code>AWS_APN_ID</code> environment variable and the <code>apn_id</code> setting in the shared config file. See <a href=\"https://docs.aws.amazon.com/PRM/latest/aws-prm-onboarding-guide/user-agent-string.html\">the PRM onboarding guide</a> for the canonical format.</p>"
100104
}
101105
}
102106
}

awscli/examples/global_options.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,3 +116,7 @@
116116
* enhanced
117117

118118

119+
``--apn-id`` (string)
120+
121+
An AWS Partner Network identifier used for Partner Revenue Measurement (PRM) attribution. The value should include the type prefix (e.g. ``pc_PRODUCTCODE`` for product codes or ``ra_PARTNERID`` for referral agents). The value is appended to the User-Agent header of all requests as ``APN_1.1/<value>$``. Overrides the ``AWS_APN_ID`` environment variable and the ``apn_id`` setting in the shared config file. See `the PRM onboarding guide <https://docs.aws.amazon.com/PRM/latest/aws-prm-onboarding-guide/user-agent-string.html>`__ for the canonical format.
122+

awscli/examples/global_synopsis.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,4 @@
1717
[--cli-auto-prompt]
1818
[--no-cli-auto-prompt]
1919
[--cli-error-format <value>]
20+
[--apn-id <value>]

tests/unit/customizations/test_globalargs.py

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,3 +194,163 @@ def test_cli_connect_timeout_for_blocking(self):
194194
self.assertEqual(
195195
session.get_default_client_config().connect_timeout, None
196196
)
197+
198+
199+
class TestResolveApnId(unittest.TestCase):
200+
def _make_session(self, config_value=None, user_agent_extra=''):
201+
session = mock.Mock()
202+
session.user_agent_extra = user_agent_extra
203+
if config_value is not None:
204+
session.get_scoped_config.return_value = {'apn_id': config_value}
205+
else:
206+
session.get_scoped_config.return_value = {}
207+
return session
208+
209+
def test_no_apn_id_leaves_user_agent_unchanged(self):
210+
session = self._make_session(user_agent_extra='botocore/1.0')
211+
parsed_args = FakeParsedArgs()
212+
with mock.patch.dict(os.environ, {}, clear=True):
213+
globalargs.resolve_apn_id(parsed_args, session)
214+
self.assertEqual(session.user_agent_extra, 'botocore/1.0')
215+
216+
def test_cli_flag_appends_apn_marker(self):
217+
session = self._make_session(user_agent_extra='botocore/1.0')
218+
parsed_args = FakeParsedArgs(apn_id='pc_abc123')
219+
globalargs.resolve_apn_id(parsed_args, session)
220+
self.assertEqual(
221+
session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_abc123$'
222+
)
223+
224+
def test_env_var_used_when_flag_missing(self):
225+
session = self._make_session(user_agent_extra='botocore/1.0')
226+
parsed_args = FakeParsedArgs()
227+
with mock.patch.dict(os.environ, {'AWS_APN_ID': 'pc_env123'}):
228+
globalargs.resolve_apn_id(parsed_args, session)
229+
self.assertEqual(
230+
session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_env123$'
231+
)
232+
233+
def test_config_file_used_when_flag_and_env_missing(self):
234+
session = self._make_session(
235+
config_value='pc_cfg123', user_agent_extra='botocore/1.0'
236+
)
237+
parsed_args = FakeParsedArgs()
238+
with mock.patch.dict(os.environ, {}, clear=True):
239+
globalargs.resolve_apn_id(parsed_args, session)
240+
self.assertEqual(
241+
session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_cfg123$'
242+
)
243+
244+
def test_cli_flag_overrides_env_and_config(self):
245+
session = self._make_session(
246+
config_value='pc_from-config', user_agent_extra='botocore/1.0'
247+
)
248+
parsed_args = FakeParsedArgs(apn_id='pc_from-flag')
249+
with mock.patch.dict(os.environ, {'AWS_APN_ID': 'pc_from-env'}):
250+
globalargs.resolve_apn_id(parsed_args, session)
251+
self.assertEqual(
252+
session.user_agent_extra,
253+
'botocore/1.0 APN_1.1/pc_from-flag$',
254+
)
255+
256+
def test_env_overrides_config(self):
257+
session = self._make_session(
258+
config_value='pc_from-config', user_agent_extra='botocore/1.0'
259+
)
260+
parsed_args = FakeParsedArgs()
261+
with mock.patch.dict(os.environ, {'AWS_APN_ID': 'ra_from-env'}):
262+
globalargs.resolve_apn_id(parsed_args, session)
263+
self.assertEqual(
264+
session.user_agent_extra,
265+
'botocore/1.0 APN_1.1/ra_from-env$',
266+
)
267+
268+
def test_appends_when_user_agent_extra_is_empty(self):
269+
session = self._make_session(user_agent_extra='')
270+
parsed_args = FakeParsedArgs(apn_id='pc_abc123')
271+
globalargs.resolve_apn_id(parsed_args, session)
272+
self.assertEqual(session.user_agent_extra, 'APN_1.1/pc_abc123$')
273+
274+
def test_appends_when_user_agent_extra_is_none(self):
275+
session = self._make_session(user_agent_extra=None)
276+
session.user_agent_extra = None
277+
parsed_args = FakeParsedArgs(apn_id='pc_abc123')
278+
globalargs.resolve_apn_id(parsed_args, session)
279+
self.assertEqual(session.user_agent_extra, 'APN_1.1/pc_abc123$')
280+
281+
def test_idempotent_when_marker_already_present(self):
282+
session = self._make_session(
283+
user_agent_extra='botocore/1.0 APN_1.1/pc_abc123$'
284+
)
285+
parsed_args = FakeParsedArgs(apn_id='pc_abc123')
286+
globalargs.resolve_apn_id(parsed_args, session)
287+
self.assertEqual(
288+
session.user_agent_extra, 'botocore/1.0 APN_1.1/pc_abc123$'
289+
)
290+
291+
def test_empty_string_treated_as_unset(self):
292+
session = self._make_session(user_agent_extra='botocore/1.0')
293+
parsed_args = FakeParsedArgs(apn_id='')
294+
with mock.patch.dict(os.environ, {}, clear=True):
295+
globalargs.resolve_apn_id(parsed_args, session)
296+
self.assertEqual(session.user_agent_extra, 'botocore/1.0')
297+
298+
def test_rejects_value_with_whitespace(self):
299+
session = self._make_session(user_agent_extra='botocore/1.0')
300+
parsed_args = FakeParsedArgs(apn_id='bad value')
301+
with self.assertRaises(ParamValidationError):
302+
globalargs.resolve_apn_id(parsed_args, session)
303+
304+
def test_rejects_value_with_special_chars(self):
305+
session = self._make_session(user_agent_extra='botocore/1.0')
306+
parsed_args = FakeParsedArgs(apn_id='bad/value')
307+
with self.assertRaises(ParamValidationError):
308+
globalargs.resolve_apn_id(parsed_args, session)
309+
310+
def test_rejects_value_with_dot(self):
311+
session = self._make_session(user_agent_extra='botocore/1.0')
312+
parsed_args = FakeParsedArgs(apn_id='1.0')
313+
with self.assertRaises(ParamValidationError):
314+
globalargs.resolve_apn_id(parsed_args, session)
315+
316+
def test_rejects_value_over_255_chars(self):
317+
session = self._make_session(user_agent_extra='botocore/1.0')
318+
parsed_args = FakeParsedArgs(apn_id='a' * 256)
319+
with self.assertRaises(ParamValidationError):
320+
globalargs.resolve_apn_id(parsed_args, session)
321+
322+
def test_accepts_max_length_value(self):
323+
session = self._make_session(user_agent_extra='botocore/1.0')
324+
long_id = 'a' * 255
325+
parsed_args = FakeParsedArgs(apn_id=long_id)
326+
globalargs.resolve_apn_id(parsed_args, session)
327+
self.assertEqual(
328+
session.user_agent_extra,
329+
'botocore/1.0 APN_1.1/%s$' % long_id,
330+
)
331+
332+
def test_accepts_dash_and_underscore(self):
333+
session = self._make_session(user_agent_extra='botocore/1.0')
334+
parsed_args = FakeParsedArgs(apn_id='pc_Pc-1_test')
335+
globalargs.resolve_apn_id(parsed_args, session)
336+
self.assertEqual(
337+
session.user_agent_extra,
338+
'botocore/1.0 APN_1.1/pc_Pc-1_test$',
339+
)
340+
341+
def test_accepts_ra_prefix(self):
342+
session = self._make_session(user_agent_extra='botocore/1.0')
343+
parsed_args = FakeParsedArgs(apn_id='ra_MyPartner123')
344+
globalargs.resolve_apn_id(parsed_args, session)
345+
self.assertEqual(
346+
session.user_agent_extra,
347+
'botocore/1.0 APN_1.1/ra_MyPartner123$',
348+
)
349+
350+
def test_handler_is_registered(self):
351+
cli = mock.Mock()
352+
globalargs.register_parse_global_args(cli)
353+
registered = [
354+
c.kwargs.get('unique_id') for c in cli.register.call_args_list
355+
]
356+
self.assertIn('resolve-apn-id', registered)

0 commit comments

Comments
 (0)