Update main.py

main.py

@@ -2,8 +2,9 @@ from huggingface_hub import InferenceClient
 
 import random
 
-from flask import Flask, request, jsonify, make_response
-from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
+from flask import Flask, request, jsonify, redirect, url_for
+from flask_jwt_extended import JWTManager, create_access_token
+from flask_dance.contrib.google import make_google_blueprint, google
 from werkzeug.security import generate_password_hash, check_password_hash
 import sqlite3
 
@@ -22,6 +23,12 @@ def home():
 app.config['JWT_SECRET_KEY'] = 123456
 jwt = JWTManager(app)
 
+# Setup Google OAuth
+app.config["GOOGLE_OAUTH_CLIENT_ID"] = "your_google_client_id"
+app.config["GOOGLE_OAUTH_CLIENT_SECRET"] = "your_google_client_secret"
+google_bp = make_google_blueprint(scope=["profile", "email"])
+app.register_blueprint(google_bp, url_prefix="/login")
+
 # Create SQLite database
 conn = sqlite3.connect('users.db', check_same_thread=False)
 c = conn.cursor()
@@ -33,30 +40,71 @@ c.execute('''CREATE TABLE IF NOT EXISTS user_details
              FOREIGN KEY(user_id) REFERENCES users(id))''')
 conn.commit()
 
-# Endpoint for user registration
+# Endpoint for user registration and Google signup
 @app.route('/register', methods=['POST'])
 def register():
     data = request.get_json()
     username = data.get('username')
     password = data.get('password')
+    google_token = data.get('google_token')
+
+    if google_token:
+        # User signing up with Google
+        resp = google.get("/oauth2/v2/userinfo")
+        if not resp.ok:
+            return jsonify({"message": "Failed to fetch Google user info"}), 400
+        google_user_info = resp.json()
+        google_id = google_user_info["id"]
+        username = google_user_info["email"]
+        password = None  # No password for Google signup
+
+        # Check if Google user already exists in the database
+        existing_user = c.execute("SELECT * FROM users WHERE google_id=?", (google_id,)).fetchone()
+        if existing_user:
+            access_token = create_access_token(identity=username, expires_delta=False)
+            return jsonify({"access_token": access_token}), 200
+
+    # User signing up with username and password
     if not username or not password:
         return jsonify({"message": "Missing username or password"}), 400
+
     hashed_password = generate_password_hash(password)
     try:
-        c.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, hashed_password))
+        c.execute("INSERT INTO users (username, password, google_id) VALUES (?, ?, ?)", (username, hashed_password, google_id))
         conn.commit()
-        return jsonify({"message": "User created successfully"}), 201
+        access_token = create_access_token(identity=username, expires_delta=False)
+        return jsonify({"access_token": access_token}), 201
     except sqlite3.IntegrityError:
         return jsonify({"message": "Username already exists"}), 400
 
-# Endpoint for user login
+# Endpoint for user login and Google login
 @app.route('/login', methods=['POST'])
 def login():
     data = request.get_json()
     username = data.get('username')
     password = data.get('password')
+    google_token = data.get('google_token')
+
+    if google_token:
+        # User logging in with Google
+        resp = google.get("/oauth2/v2/userinfo")
+        if not resp.ok:
+            return jsonify({"message": "Failed to fetch Google user info"}), 400
+        google_user_info = resp.json()
+        google_id = google_user_info["id"]
+
+        # Check if Google user exists in the database
+        user = c.execute("SELECT * FROM users WHERE google_id=?", (google_id,)).fetchone()
+        if user:
+            access_token = create_access_token(identity=user[1], expires_delta=False)
+            return jsonify({"access_token": access_token}), 200
+        else:
+            return jsonify({"message": "User not found"}), 404
+
+    # User logging in with username and password
     if not username or not password:
         return jsonify({"message": "Missing username or password"}), 400
+
     user = c.execute("SELECT * FROM users WHERE username=?", (username,)).fetchone()
     if user and check_password_hash(user[2], password):
         access_token = create_access_token(identity=username, expires_delta=False)