nova-web/web/account.py

import os
import flask
import requests

from dotenv import load_dotenv
from requests_oauthlib import OAuth2Session

load_dotenv()

os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '0' # change this to '1' if you're running this on localhost!

API_URL = 'https://discordapp.com/api'
CLIENT_ID = str(os.environ['DISCORD_CLIENT_ID'])
CLIENT_SECRET = os.environ['DISCORD_CLIENT_SECRET']
REDIRECT_URI = 'https://nova-oss.com/callback/discord'
# REDIRECT_URI = 'http://localhost:2211/callback/discord'
SCOPES = ['identify']

def get_user(discord_id: int) -> dict:
    res = requests.get(
        url=f'http://localhost:2333/users?discord_id={discord_id}',
        timeout=3,
        headers={
            'Content-Type': 'application/json',
            'Authorization': os.environ['CORE_API_KEY']
        }
    ).json()

    return res

def register(app):
    @app.route('/logout')
    def logout():
        flask.session = {}
        return flask.redirect('/')

    @app.route('/account')
    def account_view():
        try:
            discord_id = flask.session['discord_id']
        except KeyError:
            return flask.redirect('/login')

        discord_account = OAuth2Session(CLIENT_ID, token=flask.session['discord_token'])
        discord_user = discord_account.get(f'{API_URL}/users/@me').json()

        db_user = get_user(discord_id)

        user = {
            'username': discord_user.get('username'),
            'display_name': discord_user.get('global_name'),
            'avatar': f'https://cdn.discordapp.com/avatars/{discord_user.get("id")}/{discord_user.get("avatar")}.png',
            'role': db_user.get('role'),
            'credits': db_user.get('credits'),
            'api_key': db_user.get('api_key')
        }

        return flask.render_template('account.html', user=user)

    @app.route('/login')
    def login_view():
        if flask.session.get('discord_token'):
            return flask.redirect('/account')

        return flask.render_template('login.html')

    @app.route('/auth/discord', methods=['POST', 'GET'])
    def auth_view():
        """Redirect to Discord's OAuth2 authorization page."""

        oauth = OAuth2Session(CLIENT_ID, redirect_uri=REDIRECT_URI, scope=SCOPES)
        login_url, state = oauth.authorization_url(f'{API_URL}/oauth2/authorize')
        flask.session['state'] = state

        return flask.redirect(login_url)

    @app.route('/callback/discord')
    def callback():
        """Handle the OAuth2 callback and log the user in."""
        if flask.request.args.get('code'): # OAUTH CALLBACK/REDIRECT
            discord_account = OAuth2Session(CLIENT_ID, redirect_uri=REDIRECT_URI, state=flask.session.get('state'), scope=SCOPES)
            token = discord_account.fetch_token(
                f'{API_URL}/oauth2/token',
                client_secret=CLIENT_SECRET,
                authorization_response=flask.request.url,
            )

            flask.session['discord_token'] = token
            flask.session['discord_id'] = discord_account.get(f'{API_URL}/users/@me').json()['id']

            return flask.redirect('/account')

        if flask.session.get('discord_token'): # ALREADY LOGGED IN
            print('Already logged in.')
            discord_account = OAuth2Session(CLIENT_ID, token=flask.session['discord_token'])

            try:
                profile = discord_account.get(f'{API_URL}/users/@me').json()
            except Exception as e:
                return f'Your token has expired. Please try logging in again.'

            flask.session['discord_id'] = profile['id']

            return flask.redirect('/') # DONE