90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
import os
|
|
from flask import flash, redirect, url_for
|
|
from flask_dance.consumer import oauth_authorized
|
|
from flask_dance.consumer.storage.sqla import SQLAlchemyStorage
|
|
from flask_dance.contrib.google import make_google_blueprint
|
|
from flask_login import current_user, login_user
|
|
from sqlalchemy.orm.exc import NoResultFound
|
|
from models import db, User, OAuth
|
|
|
|
# Get Google OAuth credentials from environment variables
|
|
GOOGLE_CLIENT_ID = os.environ.get("GOOGLE_CLIENT_ID")
|
|
GOOGLE_CLIENT_SECRET = os.environ.get("GOOGLE_CLIENT_SECRET")
|
|
|
|
# Create the Google OAuth blueprint
|
|
google_bp = make_google_blueprint(
|
|
client_id=GOOGLE_CLIENT_ID,
|
|
client_secret=GOOGLE_CLIENT_SECRET,
|
|
scope=["profile", "email"],
|
|
storage=SQLAlchemyStorage(OAuth, db.session, user=current_user)
|
|
)
|
|
|
|
# Register a callback for when the OAuth login is successful
|
|
@oauth_authorized.connect_via(google_bp)
|
|
def google_logged_in(blueprint, token):
|
|
if not token:
|
|
flash("Failed to log in with Google.", "danger")
|
|
return False
|
|
|
|
# Get the user's Google account info
|
|
resp = blueprint.session.get("/oauth2/v1/userinfo")
|
|
if not resp.ok:
|
|
flash("Failed to fetch user info from Google.", "danger")
|
|
return False
|
|
|
|
google_info = resp.json()
|
|
google_user_id = str(google_info["id"])
|
|
|
|
# Find this OAuth token in the database, or create it
|
|
query = OAuth.query.filter_by(
|
|
provider=blueprint.name,
|
|
provider_user_id=google_user_id
|
|
)
|
|
try:
|
|
oauth = query.one()
|
|
except NoResultFound:
|
|
oauth = OAuth(
|
|
provider=blueprint.name,
|
|
provider_user_id=google_user_id,
|
|
token=token
|
|
)
|
|
|
|
# Now, check to see if we already have a user with this email
|
|
if oauth.user:
|
|
# If this OAuth token already has an associated user, log them in
|
|
login_user(oauth.user)
|
|
flash("Successfully signed in with Google.", "success")
|
|
else:
|
|
# Check if we have a user with this email already
|
|
user = User.query.filter_by(email=google_info["email"]).first()
|
|
if user:
|
|
# If we do, link the OAuth token to the existing user
|
|
oauth.user = user
|
|
db.session.add(oauth)
|
|
db.session.commit()
|
|
login_user(user)
|
|
flash("Successfully signed in with Google.", "success")
|
|
else:
|
|
# If we don't, create a new user
|
|
username = google_info["email"].split("@")[0]
|
|
# Make sure username is unique
|
|
base_username = username
|
|
count = 1
|
|
while User.query.filter_by(username=username).first():
|
|
username = f"{base_username}{count}"
|
|
count += 1
|
|
|
|
user = User(
|
|
username=username,
|
|
email=google_info["email"],
|
|
is_oauth_user=True
|
|
)
|
|
# No password is set for OAuth users
|
|
oauth.user = user
|
|
db.session.add_all([user, oauth])
|
|
db.session.commit()
|
|
login_user(user)
|
|
flash("Successfully signed in with Google.", "success")
|
|
|
|
# Prevent Flask-Dance from creating another token
|
|
return False |