forked from dpkp/kafka-python
-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #74 from Yelp/u/dpopes/msk-auth-support
Support MSK IAM Authentication
- Loading branch information
Showing
5 changed files
with
391 additions
and
3 deletions.
There are no files selected for viewing
This file contains 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 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 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,213 @@ | ||
| import datetime | ||
| import hashlib | ||
| import hmac | ||
| import json | ||
| import string | ||
|
|
||
| from kafka.errors import IllegalArgumentError | ||
| from kafka.vendor.six.moves import urllib | ||
|
|
||
|
|
||
| class AwsMskIamClient: | ||
| UNRESERVED_CHARS = string.ascii_letters + string.digits + '-._~' | ||
|
|
||
| def __init__(self, host, boto_session): | ||
| """ | ||
| Arguments: | ||
| host (str): The hostname of the broker. | ||
| boto_session (botocore.BotoSession) the boto session | ||
| """ | ||
| self.algorithm = 'AWS4-HMAC-SHA256' | ||
| self.expires = '900' | ||
| self.hashfunc = hashlib.sha256 | ||
| self.headers = [ | ||
| ('host', host) | ||
| ] | ||
| self.version = '2020_10_22' | ||
|
|
||
| self.service = 'kafka-cluster' | ||
| self.action = '{}:Connect'.format(self.service) | ||
|
|
||
| now = datetime.datetime.utcnow() | ||
| self.datestamp = now.strftime('%Y%m%d') | ||
| self.timestamp = now.strftime('%Y%m%dT%H%M%SZ') | ||
|
|
||
| self.host = host | ||
| self.boto_session = boto_session | ||
|
|
||
| # This will raise if the region can't be determined | ||
| # Do this during init instead of waiting for failures downstream | ||
| if self.region: | ||
| pass | ||
|
|
||
| @property | ||
| def access_key(self): | ||
| return self.boto_session.get_credentials().access_key | ||
|
|
||
| @property | ||
| def secret_key(self): | ||
| return self.boto_session.get_credentials().secret_key | ||
|
|
||
| @property | ||
| def token(self): | ||
| return self.boto_session.get_credentials().token | ||
|
|
||
| @property | ||
| def region(self): | ||
| # Try to get the region information from the broker hostname | ||
| for host in self.host.split(','): | ||
| if 'amazonaws.com' in host: | ||
| return host.split('.')[-3] | ||
|
|
||
| # If the region can't be determined from hostname, try the boto session | ||
| # This will only have a value if: | ||
| # - `AWS_DEFAULT_REGION` environment variable is set | ||
| # - `~/.aws/config` region variable is set | ||
| region = self.boto_session.get_config_variable('region') | ||
| if region: | ||
| return region | ||
|
|
||
| # Otherwise give up | ||
| raise IllegalArgumentError('Could not determine region from broker host(s) or aws configuration') | ||
|
|
||
| @property | ||
| def _credential(self): | ||
| return '{0.access_key}/{0._scope}'.format(self) | ||
|
|
||
| @property | ||
| def _scope(self): | ||
| return '{0.datestamp}/{0.region}/{0.service}/aws4_request'.format(self) | ||
|
|
||
| @property | ||
| def _signed_headers(self): | ||
| """ | ||
| Returns (str): | ||
| An alphabetically sorted, semicolon-delimited list of lowercase | ||
| request header names. | ||
| """ | ||
| return ';'.join(sorted(k.lower() for k, _ in self.headers)) | ||
|
|
||
| @property | ||
| def _canonical_headers(self): | ||
| """ | ||
| Returns (str): | ||
| A newline-delited list of header names and values. | ||
| Header names are lowercased. | ||
| """ | ||
| return '\n'.join(map(':'.join, self.headers)) + '\n' | ||
|
|
||
| @property | ||
| def _canonical_request(self): | ||
| """ | ||
| Returns (str): | ||
| An AWS Signature Version 4 canonical request in the format: | ||
| <Method>\n | ||
| <Path>\n | ||
| <CanonicalQueryString>\n | ||
| <CanonicalHeaders>\n | ||
| <SignedHeaders>\n | ||
| <HashedPayload> | ||
| """ | ||
| # The hashed_payload is always an empty string for MSK. | ||
| hashed_payload = self.hashfunc(b'').hexdigest() | ||
| return '\n'.join(( | ||
| 'GET', | ||
| '/', | ||
| self._canonical_querystring, | ||
| self._canonical_headers, | ||
| self._signed_headers, | ||
| hashed_payload, | ||
| )) | ||
|
|
||
| @property | ||
| def _canonical_querystring(self): | ||
| """ | ||
| Returns (str): | ||
| A '&'-separated list of URI-encoded key/value pairs. | ||
| """ | ||
| params = [] | ||
| params.append(('Action', self.action)) | ||
| params.append(('X-Amz-Algorithm', self.algorithm)) | ||
| params.append(('X-Amz-Credential', self._credential)) | ||
| params.append(('X-Amz-Date', self.timestamp)) | ||
| params.append(('X-Amz-Expires', self.expires)) | ||
| if self.token: | ||
| params.append(('X-Amz-Security-Token', self.token)) | ||
| params.append(('X-Amz-SignedHeaders', self._signed_headers)) | ||
|
|
||
| return '&'.join(self._uriencode(k) + '=' + self._uriencode(v) for k, v in params) | ||
|
|
||
| @property | ||
| def _signing_key(self): | ||
| """ | ||
| Returns (bytes): | ||
| An AWS Signature V4 signing key generated from the secret_key, date, | ||
| region, service, and request type. | ||
| """ | ||
| key = self._hmac(('AWS4' + self.secret_key).encode('utf-8'), self.datestamp) | ||
| key = self._hmac(key, self.region) | ||
| key = self._hmac(key, self.service) | ||
| key = self._hmac(key, 'aws4_request') | ||
| return key | ||
|
|
||
| @property | ||
| def _signing_str(self): | ||
| """ | ||
| Returns (str): | ||
| A string used to sign the AWS Signature V4 payload in the format: | ||
| <Algorithm>\n | ||
| <Timestamp>\n | ||
| <Scope>\n | ||
| <CanonicalRequestHash> | ||
| """ | ||
| canonical_request_hash = self.hashfunc(self._canonical_request.encode('utf-8')).hexdigest() | ||
| return '\n'.join((self.algorithm, self.timestamp, self._scope, canonical_request_hash)) | ||
|
|
||
| def _uriencode(self, msg): | ||
| """ | ||
| Arguments: | ||
| msg (str): A string to URI-encode. | ||
| Returns (str): | ||
| The URI-encoded version of the provided msg, following the encoding | ||
| rules specified: https://github.com/aws/aws-msk-iam-auth#uriencode | ||
| """ | ||
| return urllib.parse.quote(msg, safe=self.UNRESERVED_CHARS) | ||
|
|
||
| def _hmac(self, key, msg): | ||
| """ | ||
| Arguments: | ||
| key (bytes): A key to use for the HMAC digest. | ||
| msg (str): A value to include in the HMAC digest. | ||
| Returns (bytes): | ||
| An HMAC digest of the given key and msg. | ||
| """ | ||
| return hmac.new(key, msg.encode('utf-8'), digestmod=self.hashfunc).digest() | ||
|
|
||
| def first_message(self): | ||
| """ | ||
| Returns (bytes): | ||
| An encoded JSON authentication payload that can be sent to the | ||
| broker. | ||
| """ | ||
| signature = hmac.new( | ||
| self._signing_key, | ||
| self._signing_str.encode('utf-8'), | ||
| digestmod=self.hashfunc, | ||
| ).hexdigest() | ||
| msg = { | ||
| 'version': self.version, | ||
| 'host': self.host, | ||
| 'user-agent': 'kafka-python', | ||
| 'action': self.action, | ||
| 'x-amz-algorithm': self.algorithm, | ||
| 'x-amz-credential': self._credential, | ||
| 'x-amz-date': self.timestamp, | ||
| 'x-amz-signedheaders': self._signed_headers, | ||
| 'x-amz-expires': self.expires, | ||
| 'x-amz-signature': signature, | ||
| } | ||
| if self.token: | ||
| msg['x-amz-security-token'] = self.token | ||
|
|
||
| return json.dumps(msg, separators=(',', ':')).encode('utf-8') |
This file contains 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 |
|---|---|---|
| @@ -1 +1 @@ | ||
| __version__ = '1.4.7.post4' | ||
| __version__ = '1.4.7.post5' |
Oops, something went wrong.