Skip to content

Commit 41bdda0

Browse files
piochelepiotrdomalessi
piochelepiotr
and
domalessi
authored
Add support for AWS IAM login (DataDog#22660)
* Support IAM authentication for MSK clusters * update * Update kafka_consumer/assets/configuration/spec.yaml Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/assets/configuration/spec.yaml Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/assets/configuration/spec.yaml Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/assets/configuration/spec.yaml Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/assets/configuration/spec.yaml Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/assets/configuration/spec.yaml Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * run ddev validate * update changelog --------- Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com>
1 parent dba1bba commit 41bdda0

10 files changed

Lines changed: 278 additions & 32 deletions

File tree

‎LICENSE-3rdparty.csv‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ PyMySQL,PyPI,MIT,"Copyright (c) 2010, 2013 PyMySQL contributors"
44
PySocks,PyPI,BSD-3-Clause,Copyright 2006 Dan-Haim. All rights reserved.
55
PyYAML,PyPI,MIT,Copyright (c) 2017-2021 Ingy döt Net
66
aerospike,PyPI,Apache-2.0,"Copyright Aerospike, Inc."
7+
aws-msk-iam-sasl-signer-python,PyPI,Apache-2.0,Copyright 2023 Amazon Managed Streaming for Apache Kafka
78
aws-requests-auth,PyPI,BSD-3-Clause,Copyright (c) David Muller.
89
azure-identity,PyPI,MIT,Copyright (c) Microsoft Corporation.
910
beautifulsoup4,PyPI,MIT,Copyright (c) Leonard Richardson

‎agent_requirements.in‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
aerospike==7.1.1; sys_platform != 'win32' and sys_platform != 'darwin'
2+
aws-msk-iam-sasl-signer-python==1.0.2
23
aws-requests-auth==0.4.3
34
azure-identity==1.24.0
45
beautifulsoup4==4.13.5

‎kafka_consumer/assets/configuration/spec.yaml‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -220,29 +220,40 @@ files:
220220
Settings for when `sasl_mechanism` is set to `OAUTHBEARER`.
221221
fleet_configurable: true
222222
options:
223+
- name: method
224+
description: |
225+
The OAuth method to use. Either `aws_msk_iam` for AWS MSK IAM authentication
226+
or `oidc` for standard OIDC authentication. Defaults to `oidc` for backwards compatibility.
227+
value:
228+
type: string
229+
example: aws_msk_iam
230+
display_default: oidc
231+
- name: aws_region
232+
description: |
233+
AWS region for MSK IAM authentication. Required when method is `aws_msk_iam` and the region
234+
cannot be automatically detected from the environment
235+
(for example, through `AWS_REGION` or instance metadata).
236+
If not specified, attempts to detect the region automatically.
237+
value:
238+
type: string
239+
example: us-west-2
223240
- name: url
224241
fleet_configurable: true
225-
required: true
226-
enabled: false
227242
description: |
228-
The token endpoint.
243+
The token endpoint. Required when method is `oidc`.
229244
value:
230245
type: string
231246
- name: client_id
232247
fleet_configurable: true
233-
required: true
234-
enabled: false
235248
description: |
236-
The client identifier.
249+
The client identifier. Required when method is `oidc`.
237250
value:
238251
type: string
239252
- name: client_secret
240253
fleet_configurable: true
241254
secret: true
242-
required: true
243-
enabled: false
244255
description: |
245-
The client secret.
256+
The client secret. Required when method is `oidc`.
246257
value:
247258
type: string
248259
- name: scope

‎kafka_consumer/changelog.d/22660.added‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Support IAM authentication for MSK clusters

‎kafka_consumer/datadog_checks/kafka_consumer/client.py‎

Lines changed: 58 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,15 @@
66
from confluent_kafka import Consumer, ConsumerGroupTopicPartitions, KafkaException, TopicPartition
77
from confluent_kafka.admin import AdminClient
88

9+
# AWS MSK IAM authentication support
10+
try:
11+
import boto3
12+
from aws_msk_iam_sasl_signer import MSKAuthTokenProvider
13+
14+
AWS_MSK_IAM_AVAILABLE = True
15+
except ImportError:
16+
AWS_MSK_IAM_AVAILABLE = False
17+
918

1019
class KafkaClient:
1120
def __init__(self, config, log) -> None:
@@ -71,16 +80,55 @@ def __get_authentication_config(self):
7180
}
7281

7382
if self.config._sasl_mechanism == "OAUTHBEARER":
74-
extras_parameters['sasl.oauthbearer.method'] = "oidc"
75-
extras_parameters["sasl.oauthbearer.client.id"] = self.config._sasl_oauth_token_provider.get("client_id")
76-
extras_parameters["sasl.oauthbearer.token.endpoint.url"] = self.config._sasl_oauth_token_provider.get("url")
77-
extras_parameters["sasl.oauthbearer.client.secret"] = self.config._sasl_oauth_token_provider.get(
78-
"client_secret"
79-
)
80-
extras_parameters["sasl.oauthbearer.scope"] = self.config._sasl_oauth_token_provider.get("scope")
81-
extras_parameters["sasl.oauthbearer.extensions"] = self.config._sasl_oauth_token_provider.get("extensions")
82-
if self.config._sasl_oauth_tls_ca_cert:
83-
extras_parameters["https.ca.location"] = self.config._sasl_oauth_tls_ca_cert
83+
# Default to 'oidc' for backwards compatibility with existing configs
84+
method = self.config._sasl_oauth_token_provider.get("method", "oidc")
85+
86+
if method == "aws_msk_iam":
87+
if not AWS_MSK_IAM_AVAILABLE:
88+
raise Exception(
89+
"AWS MSK IAM authentication requires 'aws-msk-iam-sasl-signer-python' library. "
90+
"Install it with: pip install aws-msk-iam-sasl-signer-python"
91+
)
92+
93+
def _aws_msk_iam_oauth_cb(oauth_config):
94+
"""OAuth callback that generates AWS MSK IAM authentication tokens."""
95+
try:
96+
region = self.config._sasl_oauth_token_provider.get("aws_region")
97+
if not region:
98+
region = boto3.session.Session().region_name
99+
100+
if not region:
101+
raise Exception(
102+
"AWS region could not be determined. Please specify 'aws_region' in "
103+
"sasl_oauth_token_provider configuration."
104+
)
105+
106+
auth_token, expiry_ms = MSKAuthTokenProvider.generate_auth_token(region)
107+
self.log.debug("Generated AWS MSK IAM token for region %s, expires in %s ms", region, expiry_ms)
108+
return auth_token, expiry_ms / 1000 # Convert to seconds
109+
except Exception as e:
110+
self.log.error("Failed to generate AWS MSK IAM token: %s", e)
111+
raise
112+
113+
extras_parameters['oauth_cb'] = _aws_msk_iam_oauth_cb
114+
115+
elif method == "oidc":
116+
extras_parameters['sasl.oauthbearer.method'] = "oidc"
117+
extras_parameters["sasl.oauthbearer.client.id"] = self.config._sasl_oauth_token_provider.get(
118+
"client_id"
119+
)
120+
extras_parameters["sasl.oauthbearer.token.endpoint.url"] = self.config._sasl_oauth_token_provider.get(
121+
"url"
122+
)
123+
extras_parameters["sasl.oauthbearer.client.secret"] = self.config._sasl_oauth_token_provider.get(
124+
"client_secret"
125+
)
126+
extras_parameters["sasl.oauthbearer.scope"] = self.config._sasl_oauth_token_provider.get("scope")
127+
extras_parameters["sasl.oauthbearer.extensions"] = self.config._sasl_oauth_token_provider.get(
128+
"extensions"
129+
)
130+
if self.config._sasl_oauth_tls_ca_cert:
131+
extras_parameters["https.ca.location"] = self.config._sasl_oauth_tls_ca_cert
84132

85133
for key, value in extras_parameters.items():
86134
# Do not add the value if it's not specified

‎kafka_consumer/datadog_checks/kafka_consumer/config.py‎

Lines changed: 39 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -133,14 +133,46 @@ def validate_config(self):
133133
if self._sasl_oauth_token_provider is None:
134134
raise ConfigurationError("sasl_oauth_token_provider required for OAUTHBEARER sasl")
135135

136-
if self._sasl_oauth_token_provider.get("url") is None:
137-
raise ConfigurationError("The `url` setting of `auth_token` reader is required")
138-
139-
elif self._sasl_oauth_token_provider.get("client_id") is None:
140-
raise ConfigurationError("The `client_id` setting of `auth_token` reader is required")
136+
if not isinstance(self._sasl_oauth_token_provider, dict):
137+
raise ConfigurationError(
138+
f"sasl_oauth_token_provider must be a dictionary. Got: {type(self._sasl_oauth_token_provider)}"
139+
)
141140

142-
elif self._sasl_oauth_token_provider.get("client_secret") is None:
143-
raise ConfigurationError("The `client_secret` setting of `auth_token` reader is required")
141+
# Default to 'oidc' for backwards compatibility with existing configs
142+
method = self._sasl_oauth_token_provider.get("method", "oidc")
143+
144+
if method == "aws_msk_iam":
145+
aws_region = self._sasl_oauth_token_provider.get("aws_region")
146+
if not aws_region:
147+
try:
148+
import boto3
149+
150+
detected_region = boto3.session.Session().region_name
151+
if not detected_region:
152+
self.log.warning(
153+
"AWS region cannot be detected automatically for MSK IAM authentication. "
154+
"Consider specifying 'aws_region' in sasl_oauth_token_provider configuration. "
155+
"You can also set it via AWS_REGION environment variable or AWS config file. "
156+
"Authentication will fail at runtime if the region cannot be determined."
157+
)
158+
except ImportError:
159+
raise ConfigurationError(
160+
"AWS MSK IAM authentication requires 'boto3' and 'aws-msk-iam-sasl-signer-python' "
161+
"libraries. Install them with: pip install boto3 aws-msk-iam-sasl-signer-python"
162+
)
163+
elif method == "oidc":
164+
if self._sasl_oauth_token_provider.get("url") is None:
165+
raise ConfigurationError("The `url` setting of `auth_token` reader is required")
166+
167+
if self._sasl_oauth_token_provider.get("client_id") is None:
168+
raise ConfigurationError("The `client_id` setting of `auth_token` reader is required")
169+
170+
if self._sasl_oauth_token_provider.get("client_secret") is None:
171+
raise ConfigurationError("The `client_secret` setting of `auth_token` reader is required")
172+
else:
173+
raise ConfigurationError(
174+
f"Invalid method '{method}' for sasl_oauth_token_provider. Must be 'aws_msk_iam' or 'oidc'"
175+
)
144176

145177
# If `monitor_unlisted_consumer_groups` is set to true and
146178
# using `consumer_groups`, we prioritize `monitor_unlisted_consumer_groups`

‎kafka_consumer/datadog_checks/kafka_consumer/config_models/instance.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,11 @@ class SaslOauthTokenProvider(BaseModel):
3434
arbitrary_types_allowed=True,
3535
frozen=True,
3636
)
37+
aws_region: Optional[str] = None
3738
client_id: Optional[str] = None
3839
client_secret: Optional[str] = None
3940
extensions: Optional[str] = None
41+
method: Optional[str] = None
4042
scope: Optional[str] = None
4143
tls_ca_cert: Optional[str] = None
4244
url: Optional[str] = None

‎kafka_consumer/datadog_checks/kafka_consumer/data/conf.yaml.example‎

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -188,18 +188,32 @@ instances:
188188
#
189189
# sasl_oauth_token_provider:
190190

191-
## @param url - string - required
192-
## The token endpoint.
191+
## @param method - string - optional - default: oidc
192+
## The OAuth method to use. Either `aws_msk_iam` for AWS MSK IAM authentication
193+
## or `oidc` for standard OIDC authentication. Defaults to `oidc` for backwards compatibility.
194+
#
195+
# method: aws_msk_iam
196+
197+
## @param aws_region - string - optional - default: us-west-2
198+
## AWS region for MSK IAM authentication. Required when method is `aws_msk_iam` and the region
199+
## cannot be automatically detected from the environment
200+
## (for example, through `AWS_REGION` or instance metadata).
201+
## If not specified, attempts to detect the region automatically.
202+
#
203+
# aws_region: us-west-2
204+
205+
## @param url - string - optional
206+
## The token endpoint. Required when method is `oidc`.
193207
#
194208
# url: <URL>
195209

196-
## @param client_id - string - required
197-
## The client identifier.
210+
## @param client_id - string - optional
211+
## The client identifier. Required when method is `oidc`.
198212
#
199213
# client_id: <CLIENT_ID>
200214

201-
## @param client_secret - string - required
202-
## The client secret.
215+
## @param client_secret - string - optional
216+
## The client secret. Required when method is `oidc`.
203217
#
204218
# client_secret: <CLIENT_SECRET>
205219

‎kafka_consumer/pyproject.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ license = "BSD-3-Clause"
3737
[project.optional-dependencies]
3838
deps = [
3939
"confluent-kafka==2.13.0",
40+
"aws-msk-iam-sasl-signer-python==1.0.2",
41+
"boto3==1.40.21",
4042
"fastavro==1.12.0",
4143
"protobuf==6.33.5",
4244
]

0 commit comments

Comments
 (0)