1+ import httpx
2+ import logging
3+ import sys
4+
5+ from starlette .responses import JSONResponse
6+ from starlette .requests import Request
7+ from starlette .authentication import AuthCredentials , SimpleUser , AuthenticationBackend
8+ from starlette .authentication import AuthenticationError as StarletteAuthenticationError
9+
10+ from authlib .jose import jwt
11+ from authlib .common .errors import AuthlibBaseError
12+
13+ from git_issue_agent .config import settings
14+
15+ logger = logging .getLogger (__name__ )
16+ logging .basicConfig (level = settings .LOG_LEVEL , stream = sys .stdout , format = '%(levelname)s: %(message)s' )
17+
18+ def on_auth_error (request : Request , e : Exception ):
19+ status_code = 401
20+ message = "Authentication failed"
21+ headers = {"WWW-Authenticate" : "Bearer" }
22+
23+ if isinstance (e , AuthenticationError ):
24+ status_code = e .status_code
25+ message = str (e )
26+
27+ return JSONResponse (
28+ {"error" : message },
29+ status_code = status_code ,
30+ headers = headers if status_code == 401 else None
31+ )
32+
33+ class AuthenticationError (StarletteAuthenticationError ):
34+ def __init__ (self , message : str , status_code : int = 401 ):
35+ self .status_code = status_code
36+ super ().__init__ (message )
37+
38+ class BearerAuthBackend (AuthenticationBackend ):
39+ def __init__ (self ):
40+ if settings .JWKS_URI is None : # TODO implement oidc discovery in this case
41+ raise Exception ("JWKS_URI env var not set. " )
42+ self .jwks_url = settings .JWKS_URI
43+
44+ self .claims_options = {}
45+ if settings .AUDIENCE is None :
46+ logger .debug (f"AUDIENCE env var not set. No audience check will be performed. " )
47+ else :
48+ self .claims_options ["aud" ] = {"essential" : True , "value" : settings .AUDIENCE }
49+ if settings .ISSUER is None :
50+ logger .debug (f"ISSUER env var no set. No issuer check will be performed" )
51+ else :
52+ self .claims_options ["iss" ] = {"essential" : True , "value" : settings .ISSUER }
53+
54+ async def get_jwks (self ):
55+ logger .debug (f"Fetching JWKS from { self .jwks_url } )
56+ jwks = None
57+ try :
58+ async with httpx .AsyncClient () as client :
59+ response = await client .get (self .jwks_url )
60+ response .raise_for_status ()
61+ jwks = response .json () # TODO should probably use caching
62+ return jwks
63+ except httpx .HTTPStatusError as e :
64+ logger .debug (f"Could not retrieve JWKS from { self .jwks_url } { e } )
65+ return None
66+
67+ async def get_token (self , conn ):
68+ logger .debug ("Obtaining bearer token..." )
69+ headers = conn .headers
70+ logger .debug (f"Headers obtained: { headers } )
71+ auth = headers .get ("authorization" )
72+ if not auth or not auth .lower ().startswith ("bearer " ):
73+ logger .error ("Expected `Authorization: Bearer` access token; None provided. " )
74+ return None
75+ token = auth .split (" " , 1 )[1 ]
76+ return token
77+
78+ """This is run upon every received request"""
79+ async def authenticate (self , conn ):
80+ # bypass authentication for agent card
81+ if conn .scope .get ("path" ) == "/.well-known/agent.json" :
82+ logger .debug ("Bypassing authentication for public agent card path" )
83+ return None
84+
85+ # extract token
86+ token = await self .get_token (conn )
87+ if token is None :
88+ raise AuthenticationError (message = "Bearer token not found in Authorization header." )
89+
90+ # fetch jwks
91+ jwks = await self .get_jwks ()
92+
93+ try :
94+ # decode and validate claims
95+ claims = jwt .decode (s = token , key = jwks , claims_options = self .claims_options )
96+ claims .validate ()
97+ logger .debug ("Token successfully validated." )
98+
99+ # return user
100+ user = AgentUser (token = token , claims = claims )
101+ return AuthCredentials (user .scopes ()), user
102+ except AuthlibBaseError as e :
103+ logger .error (f"Token validation failed: { e } )
104+ raise AuthenticationError (f"Invalid token: { e } )
105+
106+ class AgentUser (SimpleUser ):
107+ def __init__ (self , token , claims ) -> None :
108+ super ().__init__ (username = claims .get ("sub" ))
109+ self .access_token = token
110+ self .claims = claims
111+
112+ def scopes (self ) -> list [str ]:
113+ scope = self .claims .get ("scope" , "" )
114+ return scope .split ()
115+
116+ class TokenExchanger :
117+ def __init__ (self ):
118+ if None in [settings .TOKEN_URL , settings .CLIENT_ID , settings .CLIENT_SECRET ]:
119+ raise Exception ("One of TOKEN_URL, CLIENT_ID, CLIENT_SECRET env vars not set - token exchange will not be performed" )
120+ self .token_url = settings .TOKEN_URL
121+ self .client_id = settings .CLIENT_ID
122+ self .client_secret = settings .CLIENT_SECRET
123+
124+ async def exchange (self , subject_token : str , audience : str = None , scope : str = None ) -> str :
125+ # headers
126+ headers = {'Content-Type' : 'application/x-www-form-urlencoded' }
127+ # data
128+ data = {
129+ 'grant_type' : 'urn:ietf:params:oauth:grant-type:token-exchange' ,
130+ 'subject_token_type' : 'urn:ietf:params:oauth:token-type:access_token' ,
131+ 'requested_token_type' : 'urn:ietf:params:oauth:token-type:access_token' ,
132+ 'client_id' : self .client_id ,
133+ 'client_secret' : self .client_secret ,
134+ 'subject_token' : subject_token ,
135+ }
136+ if not audience is None :
137+ data ['audience' ] = audience
138+ if not scope is None :
139+ data ['scope' ] = scope
140+ # make token endpoint call
141+ logger .debug ('Performing token exchange' )
142+ async with httpx .AsyncClient () as client :
143+ try :
144+ response = await client .post (self .token_url , data = data , headers = headers )
145+ response .raise_for_status () # raise exception if Http status error
146+ token_data = response .json ()
147+ if "access_token" in token_data :
148+ new_token = token_data ["access_token" ]
149+ logger .debug (f"Successful token exchange. Using token: { new_token } )
150+ return new_token
151+ logger .error ("Token exchange failed." )
152+ raise AuthenticationError ("Token exchange failed. Identity provider response did not include 'access_token'" )
153+ except httpx .HTTPStatusError as e :
154+ logger .error (f"Token exchange failed with status { e .response .status_code } { e } )
155+ raise AuthenticationError ("Token endpoint call failed." )
156+
157+ async def auth_headers (access_token , target_audience = None , target_scopes = None ):
158+ headers = {}
159+ if not access_token :
160+ return headers
161+ try :
162+ token_exchanger = TokenExchanger ()
163+ access_token = await token_exchanger .exchange (access_token , audience = target_audience , scope = target_scopes )
164+ except AuthenticationError as e :
165+ logging .error (f"Error performing token exchange - returning empty headers: { e } )
166+ return headers #
167+ except Exception as e :
168+ logging .debug (f"Error creating token exchanger - will passthrough token" )
169+
170+ headers ["Authorization" ] = f"Bearer { access_token }
171+ return headers
0 commit comments