StellarFrontier/stellarwebapp/routes/auth.py

from flask import Blueprint, redirect, url_for, session, request, flash
from flask_login import login_user, logout_user, login_required, current_user
from models import db, User
import requests
import os

auth_bp = Blueprint('auth', __name__)

DISCORD_API_BASE_URL = 'https://discord.com/api'
AUTHORIZATION_BASE_URL = DISCORD_API_BASE_URL + '/oauth2/authorize'
TOKEN_URL = DISCORD_API_BASE_URL + '/oauth2/token'

@auth_bp.route('/login')
def login():
    client_id = os.getenv('DISCORD_CLIENT_ID')
    redirect_uri = os.getenv('DISCORD_REDIRECT_URI')
    scope = 'identify'
    discord_login_url = f"{AUTHORIZATION_BASE_URL}?response_type=code&client_id={client_id}&scope={scope}&redirect_uri={redirect_uri}&prompt=consent"
    return redirect(discord_login_url)

@auth_bp.route('/callback')
def callback():
    code = request.args.get('code')
    if not code:
        flash("Error: No code provided.", "danger")
        return redirect(url_for('main.index'))

    data = {
        'client_id': os.getenv('DISCORD_CLIENT_ID'),
        'client_secret': os.getenv('DISCORD_CLIENT_SECRET'),
        'grant_type': 'authorization_code',
        'code': code,
        'redirect_uri': os.getenv('DISCORD_REDIRECT_URI'),
        'scope': 'identify'
    }
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded'
    }

    response = requests.post(TOKEN_URL, data=data, headers=headers)
    token_json = response.json()
    
    if 'access_token' not in token_json:
        flash("Error: Failed to retrieve access token.", "danger")
        return redirect(url_for('main.index'))

    access_token = token_json['access_token']
    user_headers = {
        'Authorization': f"Bearer {access_token}"
    }
    user_response = requests.get(f"{DISCORD_API_BASE_URL}/users/@me", headers=user_headers)
    user_data = user_response.json()

    discord_id = user_data['id']
    username = user_data['username']

    user = User.query.filter_by(discord_id=discord_id).first()
    if not user:
        user = User(discord_id=discord_id, username=username)
        db.session.add(user)
        db.session.commit()
    else:
        user.username = username
        db.session.commit()

    login_user(user)
    flash(f"Logged in as {username}!", "success")
    return redirect(url_for('main.index'))

@auth_bp.route('/logout')
@login_required
def logout():
    logout_user()
    flash("You have been logged out.", "info")
    return redirect(url_for('main.index'))