Improve code quality

This commit is contained in:
Sergey Petrov
2017-05-23 00:45:08 +03:00
parent e10319b78d
commit 653b37054d
3 changed files with 44 additions and 49 deletions
+8 -7
View File
@@ -14,6 +14,7 @@ class NotificationPriority(Enum):
Immediate = '10'
Delayed = '5'
RequestStream = collections.namedtuple('RequestStream', ['stream_id', 'token'])
Notification = collections.namedtuple('Notification', ['token', 'payload'])
@@ -31,7 +32,8 @@ class APNsClient(object):
DEFAULT_PORT = 443
ALTERNATIVE_PORT = 2197
def __init__(self, credentials, use_sandbox=False, use_alternative_port=False, proto=None, json_encoder=None, password=None):
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, password)
else:
@@ -45,8 +47,7 @@ class APNsClient(object):
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 = self.__credentials.create_connection(server, port,
proto)
self._connection = self.__credentials.create_connection(server, port, proto)
def send_notification(self, token_hex, notification, topic=None, priority=NotificationPriority.Immediate,
expiration=None, collapse_id=None):
@@ -92,7 +93,7 @@ class APNsClient(object):
def send_notification_batch(self, notifications, topic=None, priority=NotificationPriority.Immediate,
expiration=None, collapse_id=None):
'''
"""
Send a notification to a list of tokens in batch. Instead of sending a synchronous request
for each token, send multiple requests concurrently. This is done on the same connection,
using HTTP/2 streams (one request per stream).
@@ -104,7 +105,7 @@ class APNsClient(object):
The function returns a dictionary mapping each token to its result. The result is "Success"
if the token was sent successfully, or the string returned by APNs in the 'reason' field of
the response, if the token generated an error.
'''
"""
notification_iterator = iter(notifications)
next_notification = next(notification_iterator, None)
# Make sure we're connected to APNs, so that we receive and process the server's SETTINGS
@@ -171,10 +172,10 @@ class APNsClient(object):
self.__max_concurrent_streams = max_concurrent_streams
def connect(self):
'''
"""
Establish a connection to APNs. If already connected, the function does nothing. If the
connection fails, the function retries up to MAX_CONNECTION_RETRIES times.
'''
"""
retries = 0
while retries < MAX_CONNECTION_RETRIES:
try:
+35 -41
View File
@@ -1,29 +1,22 @@
import time
import jwt
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')
return HTTP20Connection(server, port, ssl_context=self.__ssl_context, force_proto=proto or 'h2')
def get_authorization_header(self, topic):
return None
@@ -41,15 +34,13 @@ class CertificateCredentials(Credentials):
# 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):
encryption_algorithm=DEFAULT_TOKEN_ENCRYPTION_ALGORITHM,
token_lifetime=DEFAULT_TOKEN_LIFETIME):
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
self.__encryption_algorithm = encryption_algorithm
self.__token_lifetime = token_lifetime
# Dictionary of {topic: (issue time, ascii decoded token)}
self.__topicTokens = {}
@@ -64,34 +55,37 @@ class TokenCredentials(Credentials):
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
@staticmethod
def _is_expired_token(issue_date):
return time.time() < issue_date + 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):
@staticmethod
def _get_signing_key(key_path):
secret = ''
if key_path:
with open(key_path) as f:
secret = f.read()
return secret
def _get_or_create_topic_token(self, topic):
# dict of topic to issue date and JWT token
token_pair = self.__topicTokens.get(topic)
if token_pair is None or self._is_expired_token(token_pair[0]):
# Create a new token
issued_at = time.time()
token_dict = {
'iss': self.__team_id,
'iat': issued_at,
}
headers = {
'alg': self.__encryption_algorithm,
'kid': self.__auth_key_id,
}
jwt_token = jwt.encode(token_dict, self.__auth_key,
algorithm=self.__encryption_algorithm,
headers=headers).decode('ascii')
self.__topicTokens[topic] = (issued_at, jwt_token)
return jwt_token
else:
return token_pair[1]
+1 -1
View File
@@ -110,7 +110,7 @@ class InvalidProviderToken(APNsException):
class MissingProviderToken(APNsException):
"""No provider certificate was used to connect to APNs and Authorization header was missing or no provider token
"""No provider certificate was used to connect to APNs and Authorization header was missing or no provider token
was specified. """
pass