We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Allow other OIDC providers for authentication
**Is your feature request related to a problem? Please describe.**
Currently authentication using OIDC is only available to Google identity federation

**Describe the solution you'd like**
To be able to use other OIDC providers for authentication 



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `timesketch/lib/google_auth.py`
Content:
```
1 # Copyright 2018 Google Inc. All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """
15 Based on the example code from the public Google IAP documentation:
16 https://cloud.google.com/iap/docs/signed-headers-howto
17 """
18 
19 from __future__ import unicode_literals
20 
21 import time
22 import json
23 import hashlib
24 import os
25 import six
26 
27 # six.moves is a dynamically-created namespace that doesn't actually
28 # exist and therefore pylint can't statically analyze it.
29 # pylint: disable-msg=import-error
30 from six.moves.urllib import parse as urlparse
31 
32 import jwt
33 import requests
34 
35 from flask import url_for
36 from flask import current_app
37 from flask import session
38 
39 from timesketch.lib.definitions import HTTP_STATUS_CODE_OK
40 
41 
42 CSRF_KEY = 'google_oauth2_csrf_token'
43 AUTH_URI = 'https://accounts.google.com/o/oauth2/v2/auth'
44 DISCOVERY_URL = 'https://accounts.google.com/.well-known/openid-configuration'
45 
46 
47 class JwtValidationError(Exception):
48     """Raised when a JSON Web Token cannot be validated."""
49 
50 
51 class JwtKeyError(Exception):
52     """Raised when there is a problem with the public key used for signing."""
53 
54 
55 class JwtFetchError(Exception):
56     """Raised when there is a problem with the public key used for signing."""
57 
58 
59 class DiscoveryDocumentError(Exception):
60     """Raised when there is a problem with the discovery document."""
61 
62 
63 def _fetch_public_keys(url):
64     """Fetch public keys used for verifying signatures.
65 
66     Args:
67         url: URL where keys can be fetched.
68 
69     Raises:
70         JwTKeyError if keys cannot be fetched.
71 
72     Returns:
73         HTTP response.
74     """
75     try:
76         resp = requests.get(url)
77     except requests.exceptions.RequestException as e:
78         raise JwtKeyError('Cannot fetch public keys: {}'.format(e)) from e
79     if resp.status_code != HTTP_STATUS_CODE_OK:
80         raise JwtKeyError(
81             'Cannot fetch public keys: {}'.format(resp.status_code))
82     return resp.json()
83 
84 
85 def _fetch_oauth2_discovery_document():
86     """Fetch Google OAuth2 discovery document.
87 
88     Raises:
89         DiscoveryDocumentError if document cannot be fetched.
90 
91     Returns:
92         HTTP response.
93     """
94     try:
95         resp = requests.get(DISCOVERY_URL)
96     except requests.exceptions.RequestException as e:
97         raise DiscoveryDocumentError(
98             'Cannot fetch discovery document: {}'.format(e)) from e
99     if resp.status_code != HTTP_STATUS_CODE_OK:
100         raise DiscoveryDocumentError(
101             'Cannot fetch discovery_document: {}'.format(resp.status_code))
102     return resp.json()
103 
104 
105 def _generate_random_token():
106     """Generate random string to use as CSRF and nonce tokens.
107 
108     Returns:
109         Random string.
110     """
111     return hashlib.sha256(os.urandom(1024)).hexdigest()
112 
113 
114 def get_oauth2_authorize_url(hosted_domain=None):
115     """Generate an authorization URL for Google's OAuth2 service.
116 
117     Args:
118         hosted_domain: Optional GSuite domain to limit access to.
119 
120     Returns:
121         Authorization URL.
122     """
123     csrf_token = _generate_random_token()
124     nonce = _generate_random_token()
125     redirect_uri = url_for(
126         'user_views.google_openid_connect',
127         _scheme='https',
128         _external=True
129     )
130     scopes = ('openid', 'email', 'profile')
131 
132     # Add the generated CSRF token to the client session for later validation.
133     session[CSRF_KEY] = csrf_token
134 
135     # Generate authorization URL
136     params = dict(
137         client_id=current_app.config.get('GOOGLE_OIDC_CLIENT_ID'),
138         scope=' '.join(scopes),
139         response_type='code',
140         access_type='online',  # Because we don't need a refresh token.
141         state=csrf_token,
142         nonce=nonce,  # Enable replay attack protection attack.
143         redirect_uri=redirect_uri
144     )
145     if hosted_domain:
146         params['hd'] = hosted_domain
147 
148     urlencoded_params = urlparse.urlencode(params)
149     google_authorization_url = '{}?{}'.format(AUTH_URI, urlencoded_params)
150     return google_authorization_url
151 
152 
153 def get_encoded_jwt_over_https(code):
154     """Fetch a JSON Web Token (JWT) using a authentication code.
155 
156     Args:
157         code: Authentication code obtained from an OAuth2 flow.
158 
159     Raises:
160         JwtFetchError if JWT cannot be fetched.
161 
162     Returns:
163         Encoded JWT.
164     """
165 
166     discovery_document = get_oauth2_discovery_document()
167     redirect_uri = url_for(
168         'user_views.google_openid_connect',
169         _scheme='https',
170         _external=True
171     )
172     post_data = {
173         'code': code,
174         'client_id': current_app.config.get('GOOGLE_OIDC_CLIENT_ID'),
175         'client_secret': current_app.config.get('GOOGLE_OIDC_CLIENT_SECRET'),
176         'redirect_uri': redirect_uri,
177         'grant_type': 'authorization_code'
178     }
179     token_url = discovery_document.get('token_endpoint')
180     try:
181         response = requests.post(token_url, data=post_data)
182         encoded_jwt = response.json().get('id_token')
183     except requests.exceptions.RequestException as e:
184         raise JwtFetchError(
185             'Cannot fetch JWT: {}'.format(e)) from e
186     if response.status_code != HTTP_STATUS_CODE_OK:
187         raise JwtFetchError(
188             'Cannot fetch JWT: {}'.format(response.status_code))
189 
190     if not encoded_jwt:
191         raise JwtFetchError('Cannot fetch JWT: Missing id_token in response')
192 
193     return encoded_jwt
194 
195 
196 def decode_jwt(encoded_jwt, public_key, algorithm, expected_audience):
197     """Decode a JSON Web Token (JWT).
198 
199     Args:
200         encoded_jwt: The contents of the X-Goog-IAP-JWT-Assertion header.
201         public_key: Key to verify signature of the JWT.
202         algorithm: Algorithm used for the key. E.g. ES256, RS256
203         expected_audience: Expected audience in the JWT.
204 
205     Returns:
206         Decoded JWT as a dict object.
207 
208     Raises:
209         JwtValidationError: if the JWT token cannot be decoded.
210     """
211     try:
212         decoded_jwt = jwt.decode(
213             jwt=encoded_jwt, key=public_key, algorithms=[algorithm],
214             audience=expected_audience)
215         return decoded_jwt
216     except (jwt.exceptions.InvalidTokenError,
217             jwt.exceptions.InvalidKeyError) as e:
218         raise JwtValidationError('JWT validation error: {}'.format(e)) from e
219 
220     return None
221 
222 
223 def validate_jwt(decoded_jwt, expected_issuer, expected_domain=None):
224     """Decode and validate a JSON Web token (JWT).
225 
226     Cloud IAP:
227     https://cloud.google.com/iap/docs/signed-headers-howto
228 
229     Google OpenID Connect:
230     https://developers.google.com/identity/protocols/OpenIDConnect
231 
232     Args:
233         decoded_jwt: A dict object containing the decoded JWT token.
234         expected_issuer: Expected issuer of the JWT.
235         expected_domain: Expected GSuite domain in the JWT (optional).
236 
237     Raises:
238         JwtValidationError: If unable to validate the JWT.
239     """
240     # Make sure the token is not created in the future or has expired.
241     try:
242         now = int(time.time())
243         issued_at = decoded_jwt['iat']
244         if isinstance(issued_at, six.string_types):
245             issued_at = int(issued_at, 10)
246 
247         expires_at = decoded_jwt['exp']
248         if isinstance(expires_at, six.string_types):
249             expires_at = int(expires_at, 10)
250 
251         if issued_at > now:
252             raise JwtValidationError('Token was issued in the future')
253         if expires_at < now:
254             raise JwtValidationError('Token has expired')
255     except KeyError as e:
256         raise JwtValidationError('Missing timestamp: {}'.format(e)) from e
257 
258     # Check that the issuer of the token is correct.
259     try:
260         issuer = decoded_jwt['iss']
261         if issuer != expected_issuer:
262             raise JwtValidationError('Wrong issuer: {}'.format(issuer))
263     except KeyError as e:
264         raise JwtValidationError('Missing issuer') from e
265 
266     if 'email' not in decoded_jwt:
267         raise JwtValidationError('Missing email field in token')
268 
269     if expected_domain:
270         try:
271             domain = decoded_jwt['hd']
272             if not domain == expected_domain:
273                 raise JwtValidationError('Wrong domain: {}'.format(domain))
274         except KeyError as e:
275             raise JwtValidationError('Missing domain: {}'.format(e)) from e
276 
277 
278 def get_public_key_for_jwt(encoded_jwt, url):
279     """Get public key for JWT in order to verify the signature.
280 
281     The keys get cached in order to limit the amount of network round trips.
282 
283     Args:
284         encoded_jwt: Base64 encoded JWT.
285         url: URL where keys can be fetched.
286 
287     Raises:
288         JwTKeyError if keys cannot be fetched.
289 
290     Returns:
291         Key as string.
292     """
293     # Get the Key ID from the JWT header.
294     key_id = jwt.get_unverified_header(encoded_jwt).get('kid')
295     if not key_id:
296         raise JwtKeyError('Missing key ID field in token header')
297     key_cache = get_public_key_for_jwt.key_cache
298     key = key_cache.get(key_id)
299     if key:
300         return key
301 
302     # Re-fetch the key file.
303     keys_json = _fetch_public_keys(url)
304     if 'keys' in keys_json:
305         _new_keys_dict = {}
306         for key_dict in keys_json['keys']:
307             public_key = jwt.algorithms.RSAAlgorithm.from_jwk(
308                 json.dumps(key_dict))
309             _new_keys_dict[key_dict['kid']] = public_key
310         key_cache = _new_keys_dict
311     else:
312         key_cache = keys_json
313     get_public_key_for_jwt.key_cache = key_cache
314     key = key_cache.get(key_id)
315     if not key:
316         raise JwtKeyError('IAP public key {!r} not found'.format(key_id))
317 
318     return key
319 
320 
321 def get_oauth2_discovery_document():
322     """Get Google OAuth2 discovery document.
323 
324     The document is cached in order to limit the amount of network round trips
325     and is set to expire in 12 hours from when it was fetched.
326 
327     Returns:
328         Discovery document as a dictionary.
329     """
330     now = int(time.time())
331     cache = get_oauth2_discovery_document.cache
332     discovery_document = cache.get('current')
333     if discovery_document:
334         # Check if the document has expired.
335         created_at = discovery_document['created_at']
336         expires_at = created_at + 12*60*60  # 12 hours in seconds
337         if now < expires_at:
338             return discovery_document['document']
339 
340     # Re-fetch the discovery document.
341     new_discovery_document = _fetch_oauth2_discovery_document()
342     cache = {
343         'current': {
344             'created_at': now,
345             'document': new_discovery_document
346         }
347     }
348     get_oauth2_discovery_document.cache = cache
349     return new_discovery_document
350 
351 
352 # Used to cache public keys.
353 get_public_key_for_jwt.key_cache = {}
354 
355 # Used to cache discovery documents.
356 get_oauth2_discovery_document.cache = {}
357
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/timesketch/lib/google_auth.py b/timesketch/lib/google_auth.py
--- a/timesketch/lib/google_auth.py
+++ b/timesketch/lib/google_auth.py
@@ -38,12 +38,10 @@
 
 from timesketch.lib.definitions import HTTP_STATUS_CODE_OK
 
-
 CSRF_KEY = 'google_oauth2_csrf_token'
 AUTH_URI = 'https://accounts.google.com/o/oauth2/v2/auth'
 DISCOVERY_URL = 'https://accounts.google.com/.well-known/openid-configuration'
 
-
 class JwtValidationError(Exception):
     """Raised when a JSON Web Token cannot be validated."""
 
@@ -91,8 +89,10 @@
     Returns:
         HTTP response.
     """
+    discovery_url = current_app.config.get(
+        'GOOGLE_OIDC_DISCOVERY_URL', DISCOVERY_URL)
     try:
-        resp = requests.get(DISCOVERY_URL)
+        resp = requests.get(discovery_url)
     except requests.exceptions.RequestException as e:
         raise DiscoveryDocumentError(
             'Cannot fetch discovery document: {}'.format(e)) from e
@@ -120,6 +120,7 @@
     Returns:
         Authorization URL.
     """
+    auth_uri = current_app.config.get('GOOGLE_OIDC_AUTH_URL', AUTH_URI)
     csrf_token = _generate_random_token()
     nonce = _generate_random_token()
     redirect_uri = url_for(
@@ -146,7 +147,7 @@
         params['hd'] = hosted_domain
 
     urlencoded_params = urlparse.urlencode(params)
-    google_authorization_url = '{}?{}'.format(AUTH_URI, urlencoded_params)
+    google_authorization_url = '{}?{}'.format(auth_uri, urlencoded_params)
     return google_authorization_url
 
 
@@ -199,7 +200,9 @@
     Args:
         encoded_jwt: The contents of the X-Goog-IAP-JWT-Assertion header.
         public_key: Key to verify signature of the JWT.
-        algorithm: Algorithm used for the key. E.g. ES256, RS256
+        algorithm: Algorithm used for the key. E.g. ES256, RS256. If the
+            GOOGLE_OIDC_ALGORITHM is set in the config, it will overwrite
+            the algorithm used here.
         expected_audience: Expected audience in the JWT.
 
     Returns:
@@ -208,9 +211,12 @@
     Raises:
         JwtValidationError: if the JWT token cannot be decoded.
     """
+    chosen_algorithm = current_app.config.get(
+        'GOOGLE_OIDC_ALGORITHM', algorithm)
     try:
         decoded_jwt = jwt.decode(
-            jwt=encoded_jwt, key=public_key, algorithms=[algorithm],
+            jwt=encoded_jwt, key=public_key,
+            algorithms=[chosen_algorithm],
             audience=expected_audience)
         return decoded_jwt
     except (jwt.exceptions.InvalidTokenError,

{"golden_diff": "diff --git a/timesketch/lib/google_auth.py b/timesketch/lib/google_auth.py\n--- a/timesketch/lib/google_auth.py\n+++ b/timesketch/lib/google_auth.py\n@@ -38,12 +38,10 @@\n \n from timesketch.lib.definitions import HTTP_STATUS_CODE_OK\n \n-\n CSRF_KEY = 'google_oauth2_csrf_token'\n AUTH_URI = 'https://accounts.google.com/o/oauth2/v2/auth'\n DISCOVERY_URL = 'https://accounts.google.com/.well-known/openid-configuration'\n \n-\n class JwtValidationError(Exception):\n     \"\"\"Raised when a JSON Web Token cannot be validated.\"\"\"\n \n@@ -91,8 +89,10 @@\n     Returns:\n         HTTP response.\n     \"\"\"\n+    discovery_url = current_app.config.get(\n+        'GOOGLE_OIDC_DISCOVERY_URL', DISCOVERY_URL)\n     try:\n-        resp = requests.get(DISCOVERY_URL)\n+        resp = requests.get(discovery_url)\n     except requests.exceptions.RequestException as e:\n         raise DiscoveryDocumentError(\n             'Cannot fetch discovery document: {}'.format(e)) from e\n@@ -120,6 +120,7 @@\n     Returns:\n         Authorization URL.\n     \"\"\"\n+    auth_uri = current_app.config.get('GOOGLE_OIDC_AUTH_URL', AUTH_URI)\n     csrf_token = _generate_random_token()\n     nonce = _generate_random_token()\n     redirect_uri = url_for(\n@@ -146,7 +147,7 @@\n         params['hd'] = hosted_domain\n \n     urlencoded_params = urlparse.urlencode(params)\n-    google_authorization_url = '{}?{}'.format(AUTH_URI, urlencoded_params)\n+    google_authorization_url = '{}?{}'.format(auth_uri, urlencoded_params)\n     return google_authorization_url\n \n \n@@ -199,7 +200,9 @@\n     Args:\n         encoded_jwt: The contents of the X-Goog-IAP-JWT-Assertion header.\n         public_key: Key to verify signature of the JWT.\n-        algorithm: Algorithm used for the key. E.g. ES256, RS256\n+        algorithm: Algorithm used for the key. E.g. ES256, RS256. If the\n+            GOOGLE_OIDC_ALGORITHM is set in the config, it will overwrite\n+            the algorithm used here.\n         expected_audience: Expected audience in the JWT.\n \n     Returns:\n@@ -208,9 +211,12 @@\n     Raises:\n         JwtValidationError: if the JWT token cannot be decoded.\n     \"\"\"\n+    chosen_algorithm = current_app.config.get(\n+        'GOOGLE_OIDC_ALGORITHM', algorithm)\n     try:\n         decoded_jwt = jwt.decode(\n-            jwt=encoded_jwt, key=public_key, algorithms=[algorithm],\n+            jwt=encoded_jwt, key=public_key,\n+            algorithms=[chosen_algorithm],\n             audience=expected_audience)\n         return decoded_jwt\n     except (jwt.exceptions.InvalidTokenError,\n", "issue": "Allow other OIDC providers for authentication\n**Is your feature request related to a problem? Please describe.**\r\nCurrently authentication using OIDC is only available to Google identity federation\r\n\r\n**Describe the solution you'd like**\r\nTo be able to use other OIDC providers for authentication \r\n\r\n\n", "before_files": [{"content": "# Copyright 2018 Google Inc. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nBased on the example code from the public Google IAP documentation:\nhttps://cloud.google.com/iap/docs/signed-headers-howto\n\"\"\"\n\nfrom __future__ import unicode_literals\n\nimport time\nimport json\nimport hashlib\nimport os\nimport six\n\n# six.moves is a dynamically-created namespace that doesn't actually\n# exist and therefore pylint can't statically analyze it.\n# pylint: disable-msg=import-error\nfrom six.moves.urllib import parse as urlparse\n\nimport jwt\nimport requests\n\nfrom flask import url_for\nfrom flask import current_app\nfrom flask import session\n\nfrom timesketch.lib.definitions import HTTP_STATUS_CODE_OK\n\n\nCSRF_KEY = 'google_oauth2_csrf_token'\nAUTH_URI = 'https://accounts.google.com/o/oauth2/v2/auth'\nDISCOVERY_URL = 'https://accounts.google.com/.well-known/openid-configuration'\n\n\nclass JwtValidationError(Exception):\n    \"\"\"Raised when a JSON Web Token cannot be validated.\"\"\"\n\n\nclass JwtKeyError(Exception):\n    \"\"\"Raised when there is a problem with the public key used for signing.\"\"\"\n\n\nclass JwtFetchError(Exception):\n    \"\"\"Raised when there is a problem with the public key used for signing.\"\"\"\n\n\nclass DiscoveryDocumentError(Exception):\n    \"\"\"Raised when there is a problem with the discovery document.\"\"\"\n\n\ndef _fetch_public_keys(url):\n    \"\"\"Fetch public keys used for verifying signatures.\n\n    Args:\n        url: URL where keys can be fetched.\n\n    Raises:\n        JwTKeyError if keys cannot be fetched.\n\n    Returns:\n        HTTP response.\n    \"\"\"\n    try:\n        resp = requests.get(url)\n    except requests.exceptions.RequestException as e:\n        raise JwtKeyError('Cannot fetch public keys: {}'.format(e)) from e\n    if resp.status_code != HTTP_STATUS_CODE_OK:\n        raise JwtKeyError(\n            'Cannot fetch public keys: {}'.format(resp.status_code))\n    return resp.json()\n\n\ndef _fetch_oauth2_discovery_document():\n    \"\"\"Fetch Google OAuth2 discovery document.\n\n    Raises:\n        DiscoveryDocumentError if document cannot be fetched.\n\n    Returns:\n        HTTP response.\n    \"\"\"\n    try:\n        resp = requests.get(DISCOVERY_URL)\n    except requests.exceptions.RequestException as e:\n        raise DiscoveryDocumentError(\n            'Cannot fetch discovery document: {}'.format(e)) from e\n    if resp.status_code != HTTP_STATUS_CODE_OK:\n        raise DiscoveryDocumentError(\n            'Cannot fetch discovery_document: {}'.format(resp.status_code))\n    return resp.json()\n\n\ndef _generate_random_token():\n    \"\"\"Generate random string to use as CSRF and nonce tokens.\n\n    Returns:\n        Random string.\n    \"\"\"\n    return hashlib.sha256(os.urandom(1024)).hexdigest()\n\n\ndef get_oauth2_authorize_url(hosted_domain=None):\n    \"\"\"Generate an authorization URL for Google's OAuth2 service.\n\n    Args:\n        hosted_domain: Optional GSuite domain to limit access to.\n\n    Returns:\n        Authorization URL.\n    \"\"\"\n    csrf_token = _generate_random_token()\n    nonce = _generate_random_token()\n    redirect_uri = url_for(\n        'user_views.google_openid_connect',\n        _scheme='https',\n        _external=True\n    )\n    scopes = ('openid', 'email', 'profile')\n\n    # Add the generated CSRF token to the client session for later validation.\n    session[CSRF_KEY] = csrf_token\n\n    # Generate authorization URL\n    params = dict(\n        client_id=current_app.config.get('GOOGLE_OIDC_CLIENT_ID'),\n        scope=' '.join(scopes),\n        response_type='code',\n        access_type='online',  # Because we don't need a refresh token.\n        state=csrf_token,\n        nonce=nonce,  # Enable replay attack protection attack.\n        redirect_uri=redirect_uri\n    )\n    if hosted_domain:\n        params['hd'] = hosted_domain\n\n    urlencoded_params = urlparse.urlencode(params)\n    google_authorization_url = '{}?{}'.format(AUTH_URI, urlencoded_params)\n    return google_authorization_url\n\n\ndef get_encoded_jwt_over_https(code):\n    \"\"\"Fetch a JSON Web Token (JWT) using a authentication code.\n\n    Args:\n        code: Authentication code obtained from an OAuth2 flow.\n\n    Raises:\n        JwtFetchError if JWT cannot be fetched.\n\n    Returns:\n        Encoded JWT.\n    \"\"\"\n\n    discovery_document = get_oauth2_discovery_document()\n    redirect_uri = url_for(\n        'user_views.google_openid_connect',\n        _scheme='https',\n        _external=True\n    )\n    post_data = {\n        'code': code,\n        'client_id': current_app.config.get('GOOGLE_OIDC_CLIENT_ID'),\n        'client_secret': current_app.config.get('GOOGLE_OIDC_CLIENT_SECRET'),\n        'redirect_uri': redirect_uri,\n        'grant_type': 'authorization_code'\n    }\n    token_url = discovery_document.get('token_endpoint')\n    try:\n        response = requests.post(token_url, data=post_data)\n        encoded_jwt = response.json().get('id_token')\n    except requests.exceptions.RequestException as e:\n        raise JwtFetchError(\n            'Cannot fetch JWT: {}'.format(e)) from e\n    if response.status_code != HTTP_STATUS_CODE_OK:\n        raise JwtFetchError(\n            'Cannot fetch JWT: {}'.format(response.status_code))\n\n    if not encoded_jwt:\n        raise JwtFetchError('Cannot fetch JWT: Missing id_token in response')\n\n    return encoded_jwt\n\n\ndef decode_jwt(encoded_jwt, public_key, algorithm, expected_audience):\n    \"\"\"Decode a JSON Web Token (JWT).\n\n    Args:\n        encoded_jwt: The contents of the X-Goog-IAP-JWT-Assertion header.\n        public_key: Key to verify signature of the JWT.\n        algorithm: Algorithm used for the key. E.g. ES256, RS256\n        expected_audience: Expected audience in the JWT.\n\n    Returns:\n        Decoded JWT as a dict object.\n\n    Raises:\n        JwtValidationError: if the JWT token cannot be decoded.\n    \"\"\"\n    try:\n        decoded_jwt = jwt.decode(\n            jwt=encoded_jwt, key=public_key, algorithms=[algorithm],\n            audience=expected_audience)\n        return decoded_jwt\n    except (jwt.exceptions.InvalidTokenError,\n            jwt.exceptions.InvalidKeyError) as e:\n        raise JwtValidationError('JWT validation error: {}'.format(e)) from e\n\n    return None\n\n\ndef validate_jwt(decoded_jwt, expected_issuer, expected_domain=None):\n    \"\"\"Decode and validate a JSON Web token (JWT).\n\n    Cloud IAP:\n    https://cloud.google.com/iap/docs/signed-headers-howto\n\n    Google OpenID Connect:\n    https://developers.google.com/identity/protocols/OpenIDConnect\n\n    Args:\n        decoded_jwt: A dict object containing the decoded JWT token.\n        expected_issuer: Expected issuer of the JWT.\n        expected_domain: Expected GSuite domain in the JWT (optional).\n\n    Raises:\n        JwtValidationError: If unable to validate the JWT.\n    \"\"\"\n    # Make sure the token is not created in the future or has expired.\n    try:\n        now = int(time.time())\n        issued_at = decoded_jwt['iat']\n        if isinstance(issued_at, six.string_types):\n            issued_at = int(issued_at, 10)\n\n        expires_at = decoded_jwt['exp']\n        if isinstance(expires_at, six.string_types):\n            expires_at = int(expires_at, 10)\n\n        if issued_at > now:\n            raise JwtValidationError('Token was issued in the future')\n        if expires_at < now:\n            raise JwtValidationError('Token has expired')\n    except KeyError as e:\n        raise JwtValidationError('Missing timestamp: {}'.format(e)) from e\n\n    # Check that the issuer of the token is correct.\n    try:\n        issuer = decoded_jwt['iss']\n        if issuer != expected_issuer:\n            raise JwtValidationError('Wrong issuer: {}'.format(issuer))\n    except KeyError as e:\n        raise JwtValidationError('Missing issuer') from e\n\n    if 'email' not in decoded_jwt:\n        raise JwtValidationError('Missing email field in token')\n\n    if expected_domain:\n        try:\n            domain = decoded_jwt['hd']\n            if not domain == expected_domain:\n                raise JwtValidationError('Wrong domain: {}'.format(domain))\n        except KeyError as e:\n            raise JwtValidationError('Missing domain: {}'.format(e)) from e\n\n\ndef get_public_key_for_jwt(encoded_jwt, url):\n    \"\"\"Get public key for JWT in order to verify the signature.\n\n    The keys get cached in order to limit the amount of network round trips.\n\n    Args:\n        encoded_jwt: Base64 encoded JWT.\n        url: URL where keys can be fetched.\n\n    Raises:\n        JwTKeyError if keys cannot be fetched.\n\n    Returns:\n        Key as string.\n    \"\"\"\n    # Get the Key ID from the JWT header.\n    key_id = jwt.get_unverified_header(encoded_jwt).get('kid')\n    if not key_id:\n        raise JwtKeyError('Missing key ID field in token header')\n    key_cache = get_public_key_for_jwt.key_cache\n    key = key_cache.get(key_id)\n    if key:\n        return key\n\n    # Re-fetch the key file.\n    keys_json = _fetch_public_keys(url)\n    if 'keys' in keys_json:\n        _new_keys_dict = {}\n        for key_dict in keys_json['keys']:\n            public_key = jwt.algorithms.RSAAlgorithm.from_jwk(\n                json.dumps(key_dict))\n            _new_keys_dict[key_dict['kid']] = public_key\n        key_cache = _new_keys_dict\n    else:\n        key_cache = keys_json\n    get_public_key_for_jwt.key_cache = key_cache\n    key = key_cache.get(key_id)\n    if not key:\n        raise JwtKeyError('IAP public key {!r} not found'.format(key_id))\n\n    return key\n\n\ndef get_oauth2_discovery_document():\n    \"\"\"Get Google OAuth2 discovery document.\n\n    The document is cached in order to limit the amount of network round trips\n    and is set to expire in 12 hours from when it was fetched.\n\n    Returns:\n        Discovery document as a dictionary.\n    \"\"\"\n    now = int(time.time())\n    cache = get_oauth2_discovery_document.cache\n    discovery_document = cache.get('current')\n    if discovery_document:\n        # Check if the document has expired.\n        created_at = discovery_document['created_at']\n        expires_at = created_at + 12*60*60  # 12 hours in seconds\n        if now < expires_at:\n            return discovery_document['document']\n\n    # Re-fetch the discovery document.\n    new_discovery_document = _fetch_oauth2_discovery_document()\n    cache = {\n        'current': {\n            'created_at': now,\n            'document': new_discovery_document\n        }\n    }\n    get_oauth2_discovery_document.cache = cache\n    return new_discovery_document\n\n\n# Used to cache public keys.\nget_public_key_for_jwt.key_cache = {}\n\n# Used to cache discovery documents.\nget_oauth2_discovery_document.cache = {}\n", "path": "timesketch/lib/google_auth.py"}], "after_files": [{"content": "# Copyright 2018 Google Inc. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nBased on the example code from the public Google IAP documentation:\nhttps://cloud.google.com/iap/docs/signed-headers-howto\n\"\"\"\n\nfrom __future__ import unicode_literals\n\nimport time\nimport json\nimport hashlib\nimport os\nimport six\n\n# six.moves is a dynamically-created namespace that doesn't actually\n# exist and therefore pylint can't statically analyze it.\n# pylint: disable-msg=import-error\nfrom six.moves.urllib import parse as urlparse\n\nimport jwt\nimport requests\n\nfrom flask import url_for\nfrom flask import current_app\nfrom flask import session\n\nfrom timesketch.lib.definitions import HTTP_STATUS_CODE_OK\n\nCSRF_KEY = 'google_oauth2_csrf_token'\nAUTH_URI = 'https://accounts.google.com/o/oauth2/v2/auth'\nDISCOVERY_URL = 'https://accounts.google.com/.well-known/openid-configuration'\n\nclass JwtValidationError(Exception):\n    \"\"\"Raised when a JSON Web Token cannot be validated.\"\"\"\n\n\nclass JwtKeyError(Exception):\n    \"\"\"Raised when there is a problem with the public key used for signing.\"\"\"\n\n\nclass JwtFetchError(Exception):\n    \"\"\"Raised when there is a problem with the public key used for signing.\"\"\"\n\n\nclass DiscoveryDocumentError(Exception):\n    \"\"\"Raised when there is a problem with the discovery document.\"\"\"\n\n\ndef _fetch_public_keys(url):\n    \"\"\"Fetch public keys used for verifying signatures.\n\n    Args:\n        url: URL where keys can be fetched.\n\n    Raises:\n        JwTKeyError if keys cannot be fetched.\n\n    Returns:\n        HTTP response.\n    \"\"\"\n    try:\n        resp = requests.get(url)\n    except requests.exceptions.RequestException as e:\n        raise JwtKeyError('Cannot fetch public keys: {}'.format(e)) from e\n    if resp.status_code != HTTP_STATUS_CODE_OK:\n        raise JwtKeyError(\n            'Cannot fetch public keys: {}'.format(resp.status_code))\n    return resp.json()\n\n\ndef _fetch_oauth2_discovery_document():\n    \"\"\"Fetch Google OAuth2 discovery document.\n\n    Raises:\n        DiscoveryDocumentError if document cannot be fetched.\n\n    Returns:\n        HTTP response.\n    \"\"\"\n    discovery_url = current_app.config.get(\n        'GOOGLE_OIDC_DISCOVERY_URL', DISCOVERY_URL)\n    try:\n        resp = requests.get(discovery_url)\n    except requests.exceptions.RequestException as e:\n        raise DiscoveryDocumentError(\n            'Cannot fetch discovery document: {}'.format(e)) from e\n    if resp.status_code != HTTP_STATUS_CODE_OK:\n        raise DiscoveryDocumentError(\n            'Cannot fetch discovery_document: {}'.format(resp.status_code))\n    return resp.json()\n\n\ndef _generate_random_token():\n    \"\"\"Generate random string to use as CSRF and nonce tokens.\n\n    Returns:\n        Random string.\n    \"\"\"\n    return hashlib.sha256(os.urandom(1024)).hexdigest()\n\n\ndef get_oauth2_authorize_url(hosted_domain=None):\n    \"\"\"Generate an authorization URL for Google's OAuth2 service.\n\n    Args:\n        hosted_domain: Optional GSuite domain to limit access to.\n\n    Returns:\n        Authorization URL.\n    \"\"\"\n    auth_uri = current_app.config.get('GOOGLE_OIDC_AUTH_URL', AUTH_URI)\n    csrf_token = _generate_random_token()\n    nonce = _generate_random_token()\n    redirect_uri = url_for(\n        'user_views.google_openid_connect',\n        _scheme='https',\n        _external=True\n    )\n    scopes = ('openid', 'email', 'profile')\n\n    # Add the generated CSRF token to the client session for later validation.\n    session[CSRF_KEY] = csrf_token\n\n    # Generate authorization URL\n    params = dict(\n        client_id=current_app.config.get('GOOGLE_OIDC_CLIENT_ID'),\n        scope=' '.join(scopes),\n        response_type='code',\n        access_type='online',  # Because we don't need a refresh token.\n        state=csrf_token,\n        nonce=nonce,  # Enable replay attack protection attack.\n        redirect_uri=redirect_uri\n    )\n    if hosted_domain:\n        params['hd'] = hosted_domain\n\n    urlencoded_params = urlparse.urlencode(params)\n    google_authorization_url = '{}?{}'.format(auth_uri, urlencoded_params)\n    return google_authorization_url\n\n\ndef get_encoded_jwt_over_https(code):\n    \"\"\"Fetch a JSON Web Token (JWT) using a authentication code.\n\n    Args:\n        code: Authentication code obtained from an OAuth2 flow.\n\n    Raises:\n        JwtFetchError if JWT cannot be fetched.\n\n    Returns:\n        Encoded JWT.\n    \"\"\"\n\n    discovery_document = get_oauth2_discovery_document()\n    redirect_uri = url_for(\n        'user_views.google_openid_connect',\n        _scheme='https',\n        _external=True\n    )\n    post_data = {\n        'code': code,\n        'client_id': current_app.config.get('GOOGLE_OIDC_CLIENT_ID'),\n        'client_secret': current_app.config.get('GOOGLE_OIDC_CLIENT_SECRET'),\n        'redirect_uri': redirect_uri,\n        'grant_type': 'authorization_code'\n    }\n    token_url = discovery_document.get('token_endpoint')\n    try:\n        response = requests.post(token_url, data=post_data)\n        encoded_jwt = response.json().get('id_token')\n    except requests.exceptions.RequestException as e:\n        raise JwtFetchError(\n            'Cannot fetch JWT: {}'.format(e)) from e\n    if response.status_code != HTTP_STATUS_CODE_OK:\n        raise JwtFetchError(\n            'Cannot fetch JWT: {}'.format(response.status_code))\n\n    if not encoded_jwt:\n        raise JwtFetchError('Cannot fetch JWT: Missing id_token in response')\n\n    return encoded_jwt\n\n\ndef decode_jwt(encoded_jwt, public_key, algorithm, expected_audience):\n    \"\"\"Decode a JSON Web Token (JWT).\n\n    Args:\n        encoded_jwt: The contents of the X-Goog-IAP-JWT-Assertion header.\n        public_key: Key to verify signature of the JWT.\n        algorithm: Algorithm used for the key. E.g. ES256, RS256. If the\n            GOOGLE_OIDC_ALGORITHM is set in the config, it will overwrite\n            the algorithm used here.\n        expected_audience: Expected audience in the JWT.\n\n    Returns:\n        Decoded JWT as a dict object.\n\n    Raises:\n        JwtValidationError: if the JWT token cannot be decoded.\n    \"\"\"\n    chosen_algorithm = current_app.config.get(\n        'GOOGLE_OIDC_ALGORITHM', algorithm)\n    try:\n        decoded_jwt = jwt.decode(\n            jwt=encoded_jwt, key=public_key,\n            algorithms=[chosen_algorithm],\n            audience=expected_audience)\n        return decoded_jwt\n    except (jwt.exceptions.InvalidTokenError,\n            jwt.exceptions.InvalidKeyError) as e:\n        raise JwtValidationError('JWT validation error: {}'.format(e)) from e\n\n    return None\n\n\ndef validate_jwt(decoded_jwt, expected_issuer, expected_domain=None):\n    \"\"\"Decode and validate a JSON Web token (JWT).\n\n    Cloud IAP:\n    https://cloud.google.com/iap/docs/signed-headers-howto\n\n    Google OpenID Connect:\n    https://developers.google.com/identity/protocols/OpenIDConnect\n\n    Args:\n        decoded_jwt: A dict object containing the decoded JWT token.\n        expected_issuer: Expected issuer of the JWT.\n        expected_domain: Expected GSuite domain in the JWT (optional).\n\n    Raises:\n        JwtValidationError: If unable to validate the JWT.\n    \"\"\"\n    # Make sure the token is not created in the future or has expired.\n    try:\n        now = int(time.time())\n        issued_at = decoded_jwt['iat']\n        if isinstance(issued_at, six.string_types):\n            issued_at = int(issued_at, 10)\n\n        expires_at = decoded_jwt['exp']\n        if isinstance(expires_at, six.string_types):\n            expires_at = int(expires_at, 10)\n\n        if issued_at > now:\n            raise JwtValidationError('Token was issued in the future')\n        if expires_at < now:\n            raise JwtValidationError('Token has expired')\n    except KeyError as e:\n        raise JwtValidationError('Missing timestamp: {}'.format(e)) from e\n\n    # Check that the issuer of the token is correct.\n    try:\n        issuer = decoded_jwt['iss']\n        if issuer != expected_issuer:\n            raise JwtValidationError('Wrong issuer: {}'.format(issuer))\n    except KeyError as e:\n        raise JwtValidationError('Missing issuer') from e\n\n    if 'email' not in decoded_jwt:\n        raise JwtValidationError('Missing email field in token')\n\n    if expected_domain:\n        try:\n            domain = decoded_jwt['hd']\n            if not domain == expected_domain:\n                raise JwtValidationError('Wrong domain: {}'.format(domain))\n        except KeyError as e:\n            raise JwtValidationError('Missing domain: {}'.format(e)) from e\n\n\ndef get_public_key_for_jwt(encoded_jwt, url):\n    \"\"\"Get public key for JWT in order to verify the signature.\n\n    The keys get cached in order to limit the amount of network round trips.\n\n    Args:\n        encoded_jwt: Base64 encoded JWT.\n        url: URL where keys can be fetched.\n\n    Raises:\n        JwTKeyError if keys cannot be fetched.\n\n    Returns:\n        Key as string.\n    \"\"\"\n    # Get the Key ID from the JWT header.\n    key_id = jwt.get_unverified_header(encoded_jwt).get('kid')\n    if not key_id:\n        raise JwtKeyError('Missing key ID field in token header')\n    key_cache = get_public_key_for_jwt.key_cache\n    key = key_cache.get(key_id)\n    if key:\n        return key\n\n    # Re-fetch the key file.\n    keys_json = _fetch_public_keys(url)\n    if 'keys' in keys_json:\n        _new_keys_dict = {}\n        for key_dict in keys_json['keys']:\n            public_key = jwt.algorithms.RSAAlgorithm.from_jwk(\n                json.dumps(key_dict))\n            _new_keys_dict[key_dict['kid']] = public_key\n        key_cache = _new_keys_dict\n    else:\n        key_cache = keys_json\n    get_public_key_for_jwt.key_cache = key_cache\n    key = key_cache.get(key_id)\n    if not key:\n        raise JwtKeyError('IAP public key {!r} not found'.format(key_id))\n\n    return key\n\n\ndef get_oauth2_discovery_document():\n    \"\"\"Get Google OAuth2 discovery document.\n\n    The document is cached in order to limit the amount of network round trips\n    and is set to expire in 12 hours from when it was fetched.\n\n    Returns:\n        Discovery document as a dictionary.\n    \"\"\"\n    now = int(time.time())\n    cache = get_oauth2_discovery_document.cache\n    discovery_document = cache.get('current')\n    if discovery_document:\n        # Check if the document has expired.\n        created_at = discovery_document['created_at']\n        expires_at = created_at + 12*60*60  # 12 hours in seconds\n        if now < expires_at:\n            return discovery_document['document']\n\n    # Re-fetch the discovery document.\n    new_discovery_document = _fetch_oauth2_discovery_document()\n    cache = {\n        'current': {\n            'created_at': now,\n            'document': new_discovery_document\n        }\n    }\n    get_oauth2_discovery_document.cache = cache\n    return new_discovery_document\n\n\n# Used to cache public keys.\nget_public_key_for_jwt.key_cache = {}\n\n# Used to cache discovery documents.\nget_oauth2_discovery_document.cache = {}\n", "path": "timesketch/lib/google_auth.py"}]}