From ee4d84fe6f0f5875e4cf1836b8300fbe2e3971f3 Mon Sep 17 00:00:00 2001
From: Artem Smotrakov <artem.smotrakov@gmail.com>
Date: Sat, 2 Jan 2021 23:43:47 +0100
Subject: [PATCH] ServiceAccount should expect errors

---
 src/google/auth.py | 40 +++++++++++++++++++++++++---------------
 1 file changed, 25 insertions(+), 15 deletions(-)

diff --git a/src/google/auth.py b/src/google/auth.py
index 9c11f8c..f291b00 100644
--- a/src/google/auth.py
+++ b/src/google/auth.py
@@ -16,30 +16,33 @@
 from rsa import pkcs1
 import ntp
 
+
 def encode_dict_to_base64(d):
     return encode_bytes_to_safe_base64(json.dumps(d).encode('utf8'))
 
+
 def encode_bytes_to_safe_base64(bytes):
     encoded = binascii.b2a_base64(bytes).replace(b'+', b'-')
     return encoded.replace(b'/', b'_').strip().decode('utf8')
 
+
 # this class builds a JWT to request an access token
 # from the Google OAuth 2.0 Authorization Server using a service account
 # see https://developers.google.com/identity/protocols/OAuth2ServiceAccount
 class JWTBuilder:
 
     def __init__(self):
-        self._header = {}
-        self._header['alg'] = 'RS256'
-        self._header['typ'] = 'JWT'
-        self._claim = {}
-        self._claim['iss'] = ''
-        self._claim['scope'] = ''
-        self._claim['aud'] = 'https://www.googleapis.com/oauth2/v4/token'
-        self._claim['exp'] = 0
-        self._claim['iat'] = 0
+        self._header = {
+            'alg': 'RS256',
+            'typ': 'JWT'}
+        self._claim = {
+            'iss': '',
+            'scope': '',
+            'aud': 'https://www.googleapis.com/oauth2/v4/token',
+            'exp': 0,
+            'iat': 0}
         self._key = None
-        self._expiration = 30 * 60 # 30 minutes, in seconds
+        self._expiration = 30 * 60  # 30 minutes, in seconds
 
     def service_account(self, email):
         self._claim['iss'] = email
@@ -65,10 +68,12 @@ def build(self):
         print('jwt: signing ...')
         signature = pkcs1.sign(to_be_signed.encode('utf8'), self._key, 'SHA-256')
         print('jwt: done')
-        print('jwt: encodeing')
+        print('jwt: encoding')
         encoded_signature = encode_bytes_to_safe_base64(signature)
         return '%s.%s' % (to_be_signed, encoded_signature)
 
+
+# the class obtains a token for accessing the Google API
 class ServiceAccount:
 
     def __init__(self):
@@ -97,13 +102,18 @@ def token(self):
         jwt = builder.build()
         print('token: jwt is done')
 
-        type = 'urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer'
-        body = 'grant_type=%s&assertion=%s' % (type, jwt)
+        grant_type = 'urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer'
+        body = 'grant_type=%s&assertion=%s' % (grant_type, jwt)
         headers = {'Content-Type': 'application/x-www-form-urlencoded'}
         response = requests.post('https://www.googleapis.com/oauth2/v4/token',
                                  data=body,
                                  headers=headers)
         if not response:
-            print('token: no response received')
+            raise Exception('token: no response received')
+
+        data = response.json()
+        if 'access_token' not in data:
+            print('response data: {}'.format(data))
+            raise Exception('token: no access token in response')
 
-        return response.json()['access_token']
+        return data['access_token']