prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  Support authentication with Cloud IAP
  
  
  </issue>
  <code>
  [start of src/appengine/libs/auth.py]
  1 # Copyright 2019 Google LLC
  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 """Authentication helpers."""
  15 
  16 import collections
  17 
  18 from firebase_admin import auth
  19 from google.cloud import ndb
  20 import webapp2
  21 
  22 from base import utils
  23 from config import local_config
  24 from datastore import data_types
  25 from metrics import logs
  26 from system import environment
  27 
  28 User = collections.namedtuple('User', ['email'])
  29 
  30 
  31 class AuthError(Exception):
  32   """Auth error."""
  33 
  34 
  35 def auth_domain():
  36   """Get the auth domain."""
  37   domain = local_config.ProjectConfig().get('firebase.auth_domain')
  38   if domain:
  39     return domain
  40 
  41   return utils.get_application_id() + '.firebaseapp.com'
  42 
  43 
  44 def is_current_user_admin():
  45   """Returns whether or not the current logged in user is an admin."""
  46   if environment.is_local_development():
  47     return True
  48 
  49   user = get_current_user()
  50   if not user:
  51     return False
  52 
  53   key = ndb.Key(data_types.Admin, user.email)
  54   return bool(key.get())
  55 
  56 
  57 def get_current_user():
  58   """Get the current logged in user, or None."""
  59   if environment.is_local_development():
  60     return User('user@localhost')
  61 
  62   loas_user = environment.get_value('LOAS_PEER_USERNAME')
  63   if loas_user:
  64     return User(loas_user + '@google.com')
  65 
  66   current_request = get_current_request()
  67   oauth_email = getattr(current_request, '_oauth_email', None)
  68   if oauth_email:
  69     return User(oauth_email)
  70 
  71   cached_email = getattr(current_request, '_cached_email', None)
  72   if cached_email:
  73     return User(cached_email)
  74 
  75   session_cookie = get_session_cookie()
  76   if not session_cookie:
  77     return None
  78 
  79   try:
  80     decoded_claims = decode_claims(get_session_cookie())
  81   except AuthError:
  82     logs.log_warn('Invalid session cookie.')
  83     return None
  84 
  85   if not decoded_claims.get('email_verified'):
  86     return None
  87 
  88   email = decoded_claims.get('email')
  89   if not email:
  90     return None
  91 
  92   # We cache the email for this request if we've validated the user to make
  93   # subsequent get_current_user() calls fast.
  94   setattr(current_request, '_cached_email', email)
  95   return User(email)
  96 
  97 
  98 def create_session_cookie(id_token, expires_in):
  99   """Create a new session cookie."""
  100   try:
  101     return auth.create_session_cookie(id_token, expires_in=expires_in)
  102   except auth.AuthError:
  103     raise AuthError('Failed to create session cookie.')
  104 
  105 
  106 def get_current_request():
  107   """Get the current request."""
  108   return webapp2.get_request()
  109 
  110 
  111 def get_session_cookie():
  112   """Get the current session cookie."""
  113   return get_current_request().cookies.get('session')
  114 
  115 
  116 def revoke_session_cookie(session_cookie):
  117   """Revoke a session cookie."""
  118   decoded_claims = decode_claims(session_cookie)
  119   auth.revoke_refresh_tokens(decoded_claims['sub'])
  120 
  121 
  122 def decode_claims(session_cookie):
  123   """Decode the claims for the current session cookie."""
  124   try:
  125     return auth.verify_session_cookie(session_cookie, check_revoked=True)
  126   except (ValueError, auth.AuthError):
  127     raise AuthError('Invalid session cookie.')
  128
  [end of src/appengine/libs/auth.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/src/appengine/libs/auth.py b/src/appengine/libs/auth.py
  --- a/src/appengine/libs/auth.py
  +++ b/src/appengine/libs/auth.py
  @@ -13,12 +13,17 @@
   # limitations under the License.
   """Authentication helpers."""
   
  +from builtins import str
   import collections
  +import jwt
   
   from firebase_admin import auth
   from google.cloud import ndb
  +from googleapiclient.discovery import build
  +import requests
   import webapp2
   
  +from base import memoize
   from base import utils
   from config import local_config
   from datastore import data_types
  @@ -54,6 +59,68 @@
     return bool(key.get())
   
   
  +@memoize.wrap(memoize.FifoInMemory(1))
  +def _project_number_from_id(project_id):
  +  """Get the project number from project ID."""
  +  resource_manager = build('cloudresourcemanager', 'v1')
  +  result = resource_manager.projects().get(projectId=project_id).execute()
  +  if 'projectNumber' not in result:
  +    raise AuthError('Failed to get project number.')
  +
  +  return result['projectNumber']
  +
  +
  +@memoize.wrap(memoize.FifoInMemory(1))
  +def _get_iap_key(key_id):
  +  """Retrieves a public key from the list published by Identity-Aware Proxy,
  +  re-fetching the key file if necessary.
  +  """
  +  resp = requests.get('https://www.gstatic.com/iap/verify/public_key')
  +  if resp.status_code != 200:
  +    raise AuthError('Unable to fetch IAP keys: {} / {} / {}'.format(
  +        resp.status_code, resp.headers, resp.text))
  +
  +  result = resp.json()
  +  key = result.get(key_id)
  +  if not key:
  +    raise AuthError('Key {!r} not found'.format(key_id))
  +
  +  return key
  +
  +
  +def _validate_iap_jwt(iap_jwt):
  +  """Validate JWT assertion."""
  +  project_id = utils.get_application_id()
  +  expected_audience = '/projects/{}/apps/{}'.format(
  +      _project_number_from_id(project_id), project_id)
  +
  +  try:
  +    key_id = jwt.get_unverified_header(iap_jwt).get('kid')
  +    if not key_id:
  +      raise AuthError('No key ID.')
  +
  +    key = _get_iap_key(key_id)
  +    decoded_jwt = jwt.decode(
  +        iap_jwt,
  +        key,
  +        algorithms=['ES256'],
  +        issuer='https://cloud.google.com/iap',
  +        audience=expected_audience)
  +    return decoded_jwt['email']
  +  except (jwt.exceptions.InvalidTokenError,
  +          requests.exceptions.RequestException) as e:
  +    raise AuthError('JWT assertion decode error: ' + str(e))
  +
  +
  +def get_iap_email(current_request):
  +  """Get Cloud IAP email."""
  +  jwt_assertion = current_request.headers.get('X-Goog-IAP-JWT-Assertion')
  +  if not jwt_assertion:
  +    return None
  +
  +  return _validate_iap_jwt(jwt_assertion)
  +
  +
   def get_current_user():
     """Get the current logged in user, or None."""
     if environment.is_local_development():
  @@ -64,6 +131,10 @@
       return User(loas_user + '@google.com')
   
     current_request = get_current_request()
  +  iap_email = get_iap_email(current_request)
  +  if iap_email:
  +    return User(iap_email)
  +
     oauth_email = getattr(current_request, '_oauth_email', None)
     if oauth_email:
       return User(oauth_email)
  
verification_info: {"golden_diff": "diff --git a/src/appengine/libs/auth.py b/src/appengine/libs/auth.py\n--- a/src/appengine/libs/auth.py\n+++ b/src/appengine/libs/auth.py\n@@ -13,12 +13,17 @@\n # limitations under the License.\n \"\"\"Authentication helpers.\"\"\"\n \n+from builtins import str\n import collections\n+import jwt\n \n from firebase_admin import auth\n from google.cloud import ndb\n+from googleapiclient.discovery import build\n+import requests\n import webapp2\n \n+from base import memoize\n from base import utils\n from config import local_config\n from datastore import data_types\n@@ -54,6 +59,68 @@\n   return bool(key.get())\n \n \n+@memoize.wrap(memoize.FifoInMemory(1))\n+def _project_number_from_id(project_id):\n+  \"\"\"Get the project number from project ID.\"\"\"\n+  resource_manager = build('cloudresourcemanager', 'v1')\n+  result = resource_manager.projects().get(projectId=project_id).execute()\n+  if 'projectNumber' not in result:\n+    raise AuthError('Failed to get project number.')\n+\n+  return result['projectNumber']\n+\n+\n+@memoize.wrap(memoize.FifoInMemory(1))\n+def _get_iap_key(key_id):\n+  \"\"\"Retrieves a public key from the list published by Identity-Aware Proxy,\n+  re-fetching the key file if necessary.\n+  \"\"\"\n+  resp = requests.get('https://www.gstatic.com/iap/verify/public_key')\n+  if resp.status_code != 200:\n+    raise AuthError('Unable to fetch IAP keys: {} / {} / {}'.format(\n+        resp.status_code, resp.headers, resp.text))\n+\n+  result = resp.json()\n+  key = result.get(key_id)\n+  if not key:\n+    raise AuthError('Key {!r} not found'.format(key_id))\n+\n+  return key\n+\n+\n+def _validate_iap_jwt(iap_jwt):\n+  \"\"\"Validate JWT assertion.\"\"\"\n+  project_id = utils.get_application_id()\n+  expected_audience = '/projects/{}/apps/{}'.format(\n+      _project_number_from_id(project_id), project_id)\n+\n+  try:\n+    key_id = jwt.get_unverified_header(iap_jwt).get('kid')\n+    if not key_id:\n+      raise AuthError('No key ID.')\n+\n+    key = _get_iap_key(key_id)\n+    decoded_jwt = jwt.decode(\n+        iap_jwt,\n+        key,\n+        algorithms=['ES256'],\n+        issuer='https://cloud.google.com/iap',\n+        audience=expected_audience)\n+    return decoded_jwt['email']\n+  except (jwt.exceptions.InvalidTokenError,\n+          requests.exceptions.RequestException) as e:\n+    raise AuthError('JWT assertion decode error: ' + str(e))\n+\n+\n+def get_iap_email(current_request):\n+  \"\"\"Get Cloud IAP email.\"\"\"\n+  jwt_assertion = current_request.headers.get('X-Goog-IAP-JWT-Assertion')\n+  if not jwt_assertion:\n+    return None\n+\n+  return _validate_iap_jwt(jwt_assertion)\n+\n+\n def get_current_user():\n   \"\"\"Get the current logged in user, or None.\"\"\"\n   if environment.is_local_development():\n@@ -64,6 +131,10 @@\n     return User(loas_user + '@google.com')\n \n   current_request = get_current_request()\n+  iap_email = get_iap_email(current_request)\n+  if iap_email:\n+    return User(iap_email)\n+\n   oauth_email = getattr(current_request, '_oauth_email', None)\n   if oauth_email:\n     return User(oauth_email)\n", "issue": "Support authentication with Cloud IAP\n\n", "before_files": [{"content": "# Copyright 2019 Google LLC\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\"\"\"Authentication helpers.\"\"\"\n\nimport collections\n\nfrom firebase_admin import auth\nfrom google.cloud import ndb\nimport webapp2\n\nfrom base import utils\nfrom config import local_config\nfrom datastore import data_types\nfrom metrics import logs\nfrom system import environment\n\nUser = collections.namedtuple('User', ['email'])\n\n\nclass AuthError(Exception):\n  \"\"\"Auth error.\"\"\"\n\n\ndef auth_domain():\n  \"\"\"Get the auth domain.\"\"\"\n  domain = local_config.ProjectConfig().get('firebase.auth_domain')\n  if domain:\n    return domain\n\n  return utils.get_application_id() + '.firebaseapp.com'\n\n\ndef is_current_user_admin():\n  \"\"\"Returns whether or not the current logged in user is an admin.\"\"\"\n  if environment.is_local_development():\n    return True\n\n  user = get_current_user()\n  if not user:\n    return False\n\n  key = ndb.Key(data_types.Admin, user.email)\n  return bool(key.get())\n\n\ndef get_current_user():\n  \"\"\"Get the current logged in user, or None.\"\"\"\n  if environment.is_local_development():\n    return User('user@localhost')\n\n  loas_user = environment.get_value('LOAS_PEER_USERNAME')\n  if loas_user:\n    return User(loas_user + '@google.com')\n\n  current_request = get_current_request()\n  oauth_email = getattr(current_request, '_oauth_email', None)\n  if oauth_email:\n    return User(oauth_email)\n\n  cached_email = getattr(current_request, '_cached_email', None)\n  if cached_email:\n    return User(cached_email)\n\n  session_cookie = get_session_cookie()\n  if not session_cookie:\n    return None\n\n  try:\n    decoded_claims = decode_claims(get_session_cookie())\n  except AuthError:\n    logs.log_warn('Invalid session cookie.')\n    return None\n\n  if not decoded_claims.get('email_verified'):\n    return None\n\n  email = decoded_claims.get('email')\n  if not email:\n    return None\n\n  # We cache the email for this request if we've validated the user to make\n  # subsequent get_current_user() calls fast.\n  setattr(current_request, '_cached_email', email)\n  return User(email)\n\n\ndef create_session_cookie(id_token, expires_in):\n  \"\"\"Create a new session cookie.\"\"\"\n  try:\n    return auth.create_session_cookie(id_token, expires_in=expires_in)\n  except auth.AuthError:\n    raise AuthError('Failed to create session cookie.')\n\n\ndef get_current_request():\n  \"\"\"Get the current request.\"\"\"\n  return webapp2.get_request()\n\n\ndef get_session_cookie():\n  \"\"\"Get the current session cookie.\"\"\"\n  return get_current_request().cookies.get('session')\n\n\ndef revoke_session_cookie(session_cookie):\n  \"\"\"Revoke a session cookie.\"\"\"\n  decoded_claims = decode_claims(session_cookie)\n  auth.revoke_refresh_tokens(decoded_claims['sub'])\n\n\ndef decode_claims(session_cookie):\n  \"\"\"Decode the claims for the current session cookie.\"\"\"\n  try:\n    return auth.verify_session_cookie(session_cookie, check_revoked=True)\n  except (ValueError, auth.AuthError):\n    raise AuthError('Invalid session cookie.')\n", "path": "src/appengine/libs/auth.py"}]}