From ae835a195f5ed06004604503ecaf11d20cce6128 Mon Sep 17 00:00:00 2001 From: Masa Date: Wed, 3 May 2017 13:15:26 -0700 Subject: [PATCH] Enable JWT authentication in addition to old certificate authentication (#40) --- apns2/client.py | 27 +++++++---- apns2/credentials.py | 96 ++++++++++++++++++++++++++++++++++++++++ setup.py | 6 ++- test/eckey.pem | 8 ++++ test/test_client.py | 4 +- test/test_credentials.py | 56 +++++++++++++++++++++++ 6 files changed, 185 insertions(+), 12 deletions(-) create mode 100644 apns2/credentials.py create mode 100644 test/eckey.pem create mode 100644 test/test_credentials.py diff --git a/apns2/client.py b/apns2/client.py index b49f867..be47fed 100644 --- a/apns2/client.py +++ b/apns2/client.py @@ -3,11 +3,12 @@ import json import logging from enum import Enum -from hyper import HTTP20Connection -from hyper.tls import init_context - from .errors import ConnectionFailed, exception_class_for_reason +# We don't generally need to know about the Credentials subclasses except to +# keep the old API, where APNsClient took a cert_file +from .credentials import CertificateCredentials + class NotificationPriority(Enum): Immediate = '10' @@ -30,18 +31,22 @@ class APNsClient(object): DEFAULT_PORT = 443 ALTERNATIVE_PORT = 2197 - def __init__(self, cert_file, use_sandbox=False, use_alternative_port=False, proto=None, json_encoder=None, password=None): - ssl_context = init_context() - ssl_context.load_cert_chain(cert_file, password=password) - self._init_connection(use_sandbox, use_alternative_port, ssl_context, proto) + def __init__(self, credentials, use_sandbox=False, use_alternative_port=False, proto=None, json_encoder=None, password=None): + if credentials is None or isinstance(credentials, str): + self.__credentials = CertificateCredentials(credentials) + else: + self.__credentials = credentials + self._init_connection(use_sandbox, use_alternative_port, proto) + self.__json_encoder = json_encoder self.__max_concurrent_streams = None self.__previous_server_max_concurrent_streams = None - def _init_connection(self, use_sandbox, use_alternative_port, ssl_context, proto): + def _init_connection(self, use_sandbox, use_alternative_port, proto): server = self.SANDBOX_SERVER if use_sandbox else self.LIVE_SERVER port = self.ALTERNATIVE_PORT if use_alternative_port else self.DEFAULT_PORT - self._connection = HTTP20Connection(server, port, ssl_context=ssl_context, force_proto=proto or 'h2') + self._connection = self.__credentials.create_connection(server, port, + proto) def send_notification(self, token_hex, notification, topic, priority=NotificationPriority.Immediate, expiration=None, collapse_id=None): @@ -62,6 +67,10 @@ class APNsClient(object): if expiration is not None: headers['apns-expiration'] = '%d' % expiration + auth_header = self.__credentials.get_authorization_header(topic) + if auth_header is not None: + headers['authorization'] = auth_header + if collapse_id is not None: headers['apns-collapse-id'] = collapse_id diff --git a/apns2/credentials.py b/apns2/credentials.py new file mode 100644 index 0000000..a5bbdbd --- /dev/null +++ b/apns2/credentials.py @@ -0,0 +1,96 @@ + +from hyper import HTTP20Connection +from hyper.tls import init_context + +import jwt + +# For creating and comparing the time for the JWT token +import time + + +DEFAULT_TOKEN_LIFETIME = 3600 +DEFAULT_TOKEN_ENCRYPTION_ALGORITHM = 'ES256' + + +# Abstract Base class. This should not be instantiated directly. +class Credentials(object): + + def __init__(self, ssl_context=None): + self.__ssl_context = ssl_context + + # Creates a connection with the credentials, if available or necessary. + def create_connection(self, server, port, proto): + # self.__ssl_context may be none, and that's fine. + return HTTP20Connection(server, port, + ssl_context=self.__ssl_context, + force_proto=proto or 'h2') + + def get_authorization_header(self, topic): + return None + + +# Credentials subclass for certificate authentication +class CertificateCredentials(Credentials): + def __init__(self, cert_file): + ssl_context = init_context() + ssl_context.load_cert_chain(cert_file) + super(CertificateCredentials, self).__init__(ssl_context) + + +# Credentials subclass for JWT token based authentication +class TokenCredentials(Credentials): + def __init__(self, auth_key_path, auth_key_id, team_id, + encryption_algorithm=None, token_lifetime=None): + self.__auth_key = self._get_signing_key(auth_key_path) + self.__auth_key_id = auth_key_id + self.__team_id = team_id + self.__encryption_algorithm = DEFAULT_TOKEN_ENCRYPTION_ALGORITHM if \ + encryption_algorithm is None else \ + encryption_algorithm + self.__token_lifetime = DEFAULT_TOKEN_LIFETIME if \ + token_lifetime is None else token_lifetime + + # Dictionary of {topic: (issue time, ascii decoded token)} + self.__topicTokens = {} + + # Use the default constructor because we don't have an SSL context + super(TokenCredentials, self).__init__() + + def get_tokens(self): + return [val[1] for val in self.__topicTokens] + + def get_authorization_header(self, topic): + token = self._get_or_create_topic_token(topic) + return 'bearer %s' % token + + def _isExpiredToken(self, issueDate): + now = time.time() + return now < issueDate + DEFAULT_TOKEN_LIFETIME + + def _get_or_create_topic_token(self, topic): + # dict of topic to issue date and JWT token + tokenPair = self.__topicTokens.get(topic) + if tokenPair is None or self._isExpiredToken(tokenPair[0]): + # Create a new token + issuedAt = time.time() + tokenDict = {'iss': self.__team_id, + 'iat': issuedAt + } + headers = {'alg': self.__encryption_algorithm, + 'kid': self.__auth_key_id, + } + jwtToken = jwt.encode(tokenDict, self.__auth_key, + algorithm=self.__encryption_algorithm, + headers=headers).decode('ascii') + + self.__topicTokens[topic] = (issuedAt, jwtToken) + return jwtToken + else: + return tokenPair[1] + + def _get_signing_key(self, key_path): + secret = '' + if key_path: + with open(key_path) as f: + secret = f.read() + return secret diff --git a/setup.py b/setup.py index 1ef879b..b4ae0f9 100755 --- a/setup.py +++ b/setup.py @@ -2,7 +2,11 @@ from setuptools import setup -dependencies = ['hyper>=0.7'] +dependencies = [ + 'hyper>=0.7', + 'PyJWT>=1.4.0', + 'cryptography>=1.7.2', +] try: # noinspection PyUnresolvedReferences diff --git a/test/eckey.pem b/test/eckey.pem new file mode 100644 index 0000000..76226d2 --- /dev/null +++ b/test/eckey.pem @@ -0,0 +1,8 @@ +-----BEGIN EC PARAMETERS----- +BggqhkjOPQMBBw== +-----END EC PARAMETERS----- +-----BEGIN EC PRIVATE KEY----- +MHcCAQEEIM0scFXkVBBc7d8DSL9rfFB1PvET/dWQa9eWfxpgaqaBoAoGCCqGSM49 +AwEHoUQDQgAEnVxQ41VKMN6uSsSYCCdOUhCms+HT2VpUhFGML5SzYGoodKtRD/6J +YI9Rxq1lPGHMwECSaPtPf9kVDCUM6UHvhA== +-----END EC PRIVATE KEY----- diff --git a/test/test_client.py b/test/test_client.py index 2c0fb7f..7681d5a 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -32,14 +32,14 @@ class ClientTestCase(TestCase): self.mock_results = None self.next_stream_id = 0 - with patch('apns2.client.HTTP20Connection') as mock_connection_constructor, patch('apns2.client.init_context'): + with patch('apns2.credentials.HTTP20Connection') as mock_connection_constructor, patch('apns2.credentials.init_context'): self.mock_connection = MagicMock() self.mock_connection.get_response.side_effect = self.mock_get_response self.mock_connection.request.side_effect = self.mock_request self.mock_connection._conn.__enter__.return_value = self.mock_connection._conn self.mock_connection._conn.remote_settings.max_concurrent_streams = 500 mock_connection_constructor.return_value = self.mock_connection - self.client = APNsClient(cert_file=None) + self.client = APNsClient(credentials=None) @contextlib.contextmanager def mock_get_response(self, stream_id): diff --git a/test/test_credentials.py b/test/test_credentials.py new file mode 100644 index 0000000..180b9fb --- /dev/null +++ b/test/test_credentials.py @@ -0,0 +1,56 @@ +# This only tests the TokenCredentials test case, since the +# CertificateCredentials would be mocked out anyway. +# Namely: +# - timing out of the token +# - creating multiple tokens for different topics + +from unittest import TestCase, main + +import time + +from apns2.credentials import TokenCredentials + + +class TokenCredentialsTestCase(TestCase): + @classmethod + def setUpClass(cls): + cls.key_path = 'test/eckey.pem' + cls.team_id = '3Z24IP123A' + cls.key_id = '1QBCDJ9RST' + cls.topics = ('com.example.first_app', 'com.example.second_app',) + cls.token_lifetime = 0.5 + + def setUp(self): + # Create an 'ephemeral' token so we can test token timeouts. We + # want a timeout long enough to last the test, but we don't want to + # slow down the tests too much either. + self.normal_creds = TokenCredentials(self.key_path, self.key_id, + self.team_id) + self.lasting_header = self.normal_creds.get_authorization_header( + self.topics[0]) + self.expiring_creds = \ + TokenCredentials(self.key_path, self.key_id, + self.team_id, + token_lifetime=self.token_lifetime) + self.expiring_header = self.expiring_creds.get_authorization_header( + self.topics[0]) + + def test_create_multiple_topics(self): + h1 = self.normal_creds.get_authorization_header(self.topics[0]) + self.assertEqual(len(self.normal_creds.get_tokens()), 1) + h2 = self.normal_creds.get_authorization_header(self.topics[1]) + self.assertNotEqual(h1, h2) + self.assertEqual(len(self.normal_creds.get_tokens()), 2) + + def test_token_expiration(self): + # As long as the token lifetime hasn't elapsed, this should work. To + # be really careful, we should check how much time has elapsed to + # know if it fail. But, either way, we'd have to come up with a good + # lifetime for future tests... + time.sleep(self.token_lifetime) + h3 = self.expiring_creds.get_authorization_header(self.topics[0]) + self.assertNotEqual(self.expiring_header, h3) + + +if __name__ == '__main__': + main()