from starlette import status
from starlette.testclient import TestClient
from httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallback
from httpx_oauth.oauth2 import OAuth2
CLIENT_ID = "CLIENT_ID"
CLIENT_SECRET = "CLIENT_SECRET"
AUTHORIZE_ENDPOINT = "https://www.camelot.bt/authorize"
ACCESS_TOKEN_ENDPOINT = "https://www.camelot.bt/access-token"
REDIRECT_URL = "https://www.tintagel.bt/callback"
ROUTE_NAME = "callback"
client = OAuth2(CLIENT_ID, CLIENT_SECRET, AUTHORIZE_ENDPOINT,
ACCESS_TOKEN_ENDPOINT)
oauth2_authorize_callback_route_name = OAuth2AuthorizeCallback(
client, route_name=ROUTE_NAME)
oauth2_authorize_callback_redirect_url = OAuth2AuthorizeCallback(
client, redirect_url=REDIRECT_URL)
app = FastAPI()
@app.get("/authorize-route-name")
async def authorize_route_name(
access_token_state=Depends(oauth2_authorize_callback_route_name), ):
return access_token_state
@app.get("/authorize-redirect-url")
async def authorize_redirect_url(
access_token_state=Depends(oauth2_authorize_callback_redirect_url), ):
return access_token_state
def get_oauth_router(
oauth_client: BaseOAuth2,
user_db: BaseUserDatabase[models.BaseUserDB],
user_db_model: Type[models.BaseUserDB],
authenticator: Authenticator,
state_secret: str,
redirect_url: str = None,
after_register: Optional[Callable[[models.UD, Request], None]] = None,
) -> APIRouter:
"""Generate a router with the OAuth routes."""
router = APIRouter()
callback_route_name = f"{oauth_client.name}-callback"
if redirect_url is not None:
oauth2_authorize_callback = OAuth2AuthorizeCallback(
oauth_client,
redirect_url=redirect_url,
)
else:
oauth2_authorize_callback = OAuth2AuthorizeCallback(
oauth_client,
route_name=callback_route_name,
)
@router.get("/authorize")
async def authorize(
request: Request,
authentication_backend: str,
scopes: List[str] = Query(None),
):
# Check that authentication_backend exists
backend_exists = False
for backend in authenticator.backends:
if backend.name == authentication_backend:
backend_exists = True
break
if not backend_exists:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
if redirect_url is not None:
authorize_redirect_url = redirect_url
else:
authorize_redirect_url = request.url_for(callback_route_name)
state_data = {
"authentication_backend": authentication_backend,
}
state = generate_state_token(state_data, state_secret)
authorization_url = await oauth_client.get_authorization_url(
authorize_redirect_url,
state,
scopes,
)
return {"authorization_url": authorization_url}
@router.get("/callback", name=f"{oauth_client.name}-callback")
async def callback(
request: Request,
response: Response,
access_token_state=Depends(oauth2_authorize_callback),
):
token, state = access_token_state
account_id, account_email = await oauth_client.get_id_email(
token["access_token"])
try:
state_data = decode_state_token(state, state_secret)
except jwt.DecodeError:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
user = await user_db.get_by_oauth_account(oauth_client.name,
account_id)
new_oauth_account = models.BaseOAuthAccount(
oauth_name=oauth_client.name,
access_token=token["access_token"],
expires_at=token.get("expires_at"),
refresh_token=token.get("refresh_token"),
account_id=account_id,
account_email=account_email,
)
if not user:
user = await user_db.get_by_email(account_email)
if user:
# Link account
user.oauth_accounts.append(new_oauth_account) # type: ignore
await user_db.update(user)
else:
# Create account
password = generate_password()
user = user_db_model(
email=account_email,
hashed_password=get_password_hash(password),
oauth_accounts=[new_oauth_account],
)
await user_db.create(user)
if after_register:
await run_handler(after_register, user, request)
else:
# Update oauth
updated_oauth_accounts = []
for oauth_account in user.oauth_accounts: # type: ignore
if oauth_account.account_id == account_id:
updated_oauth_accounts.append(new_oauth_account)
else:
updated_oauth_accounts.append(oauth_account)
user.oauth_accounts = updated_oauth_accounts # type: ignore
await user_db.update(user)
if not user.is_active:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ErrorCode.LOGIN_BAD_CREDENTIALS,
)
# Authenticate
for backend in authenticator.backends:
if backend.name == state_data["authentication_backend"]:
return await backend.get_login_response(
cast(models.BaseUserDB, user), response)
return router
def get_oauth_router(
oauth_client: BaseOAuth2,
backend: AuthenticationBackend,
get_user_manager: UserManagerDependency[models.UC, models.UD],
state_secret: SecretType,
redirect_url: str = None,
) -> APIRouter:
"""Generate a router with the OAuth routes."""
router = APIRouter()
callback_route_name = f"oauth:{oauth_client.name}.{backend.name}.callback"
if redirect_url is not None:
oauth2_authorize_callback = OAuth2AuthorizeCallback(
oauth_client,
redirect_url=redirect_url,
)
else:
oauth2_authorize_callback = OAuth2AuthorizeCallback(
oauth_client,
route_name=callback_route_name,
)
@router.get(
"/authorize",
name=f"oauth:{oauth_client.name}.{backend.name}.authorize",
response_model=OAuth2AuthorizeResponse,
)
async def authorize(
request: Request, scopes: List[str] = Query(None)
) -> OAuth2AuthorizeResponse:
if redirect_url is not None:
authorize_redirect_url = redirect_url
else:
authorize_redirect_url = request.url_for(callback_route_name)
state_data: Dict[str, str] = {}
state = generate_state_token(state_data, state_secret)
authorization_url = await oauth_client.get_authorization_url(
authorize_redirect_url,
state,
scopes,
)
return OAuth2AuthorizeResponse(authorization_url=authorization_url)
@router.get(
"/callback",
name=callback_route_name,
description="The response varies based on the authentication backend used.",
responses={
status.HTTP_400_BAD_REQUEST: {
"model": ErrorModel,
"content": {
"application/json": {
"examples": {
"INVALID_STATE_TOKEN": {
"summary": "Invalid state token.",
"value": None,
},
ErrorCode.LOGIN_BAD_CREDENTIALS: {
"summary": "User is inactive.",
"value": {"detail": ErrorCode.LOGIN_BAD_CREDENTIALS},
},
}
}
},
},
},
)
async def callback(
request: Request,
response: Response,
access_token_state: Tuple[OAuth2Token, str] = Depends(
oauth2_authorize_callback
),
user_manager: BaseUserManager[models.UC, models.UD] = Depends(get_user_manager),
strategy: Strategy[models.UC, models.UD] = Depends(backend.get_strategy),
):
token, state = access_token_state
account_id, account_email = await oauth_client.get_id_email(
token["access_token"]
)
try:
decode_jwt(state, state_secret, [STATE_TOKEN_AUDIENCE])
except jwt.DecodeError:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
new_oauth_account = models.BaseOAuthAccount(
oauth_name=oauth_client.name,
access_token=token["access_token"],
expires_at=token.get("expires_at"),
refresh_token=token.get("refresh_token"),
account_id=account_id,
account_email=account_email,
)
user = await user_manager.oauth_callback(new_oauth_account, request)
if not user.is_active:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ErrorCode.LOGIN_BAD_CREDENTIALS,
)
# Authenticate
return await backend.login(strategy, user, response)
return router