Spaces:
Sleeping
Sleeping
Update auth.py
Browse files
auth.py
CHANGED
@@ -1,10 +1,6 @@
|
|
1 |
import os
|
2 |
import uuid
|
3 |
import logging
|
4 |
-
import secrets
|
5 |
-
import smtplib
|
6 |
-
from email.mime.multipart import MIMEMultipart
|
7 |
-
from email.mime.text import MIMEText
|
8 |
from datetime import datetime, timedelta
|
9 |
from urllib.parse import quote_plus
|
10 |
from typing import List, Optional, Any
|
@@ -19,11 +15,16 @@ from pydantic import BaseModel, EmailStr, Field, validator
|
|
19 |
from pymongo import MongoClient
|
20 |
import os.path
|
21 |
|
22 |
-
#
|
|
|
|
|
|
|
23 |
try:
|
24 |
import markdown
|
25 |
except ImportError:
|
26 |
markdown = None
|
|
|
|
|
27 |
|
28 |
load_dotenv()
|
29 |
|
@@ -37,6 +38,8 @@ MONGO_URL = os.getenv("CONNECTION_STRING").replace("${PASSWORD}", password)
|
|
37 |
client = MongoClient(MONGO_URL)
|
38 |
db = client.users_database
|
39 |
users_collection = db.users
|
|
|
|
|
40 |
|
41 |
# OAuth2 setup
|
42 |
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
@@ -44,7 +47,7 @@ oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
|
44 |
# Create an APIRouter instance
|
45 |
router = APIRouter()
|
46 |
|
47 |
-
#
|
48 |
class User(BaseModel):
|
49 |
name: str = Field(..., min_length=3, max_length=50)
|
50 |
email: EmailStr
|
@@ -96,12 +99,7 @@ class LoginResponse(Token):
|
|
96 |
class TokenData(BaseModel):
|
97 |
email: Optional[str] = None
|
98 |
|
99 |
-
#
|
100 |
-
class OTPLogin(BaseModel):
|
101 |
-
email: EmailStr
|
102 |
-
otp: str
|
103 |
-
|
104 |
-
# ------------------ Password Hashing ------------------
|
105 |
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
106 |
|
107 |
def verify_password(plain_password: str, hashed_password: str) -> bool:
|
@@ -110,7 +108,6 @@ def verify_password(plain_password: str, hashed_password: str) -> bool:
|
|
110 |
def get_password_hash(password: str) -> str:
|
111 |
return pwd_context.hash(password)
|
112 |
|
113 |
-
# ------------------ User & Auth Helpers ------------------
|
114 |
def get_user(email: str) -> Optional[dict]:
|
115 |
return users_collection.find_one({"email": email})
|
116 |
|
@@ -148,7 +145,7 @@ def get_current_user(token: str = Depends(oauth2_scheme)) -> dict:
|
|
148 |
except JWTError:
|
149 |
raise HTTPException(status_code=401, detail="Invalid token")
|
150 |
|
151 |
-
#
|
152 |
AVATAR_DIR = "avatars"
|
153 |
if not os.path.exists(AVATAR_DIR):
|
154 |
os.makedirs(AVATAR_DIR)
|
@@ -176,7 +173,7 @@ def save_avatar_file(file: UploadFile) -> str:
|
|
176 |
file.file.close()
|
177 |
return file_path
|
178 |
|
179 |
-
#
|
180 |
def generate_otp(length=6):
|
181 |
"""Generate a numeric OTP of specified length."""
|
182 |
digits = "0123456789"
|
@@ -234,9 +231,44 @@ def send_email_func(sender_email, sender_password, receiver_email, subject, body
|
|
234 |
except Exception as e:
|
235 |
logger.error(f"Error sending email: {e}")
|
236 |
raise e
|
237 |
-
#
|
238 |
|
239 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
240 |
|
241 |
@router.post("/signup", response_model=Token)
|
242 |
async def signup(
|
@@ -244,8 +276,20 @@ async def signup(
|
|
244 |
name: str = Form(...),
|
245 |
email: EmailStr = Form(...),
|
246 |
password: str = Form(...),
|
|
|
247 |
avatar: Optional[UploadFile] = File(None)
|
248 |
):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
249 |
try:
|
250 |
_ = User(name=name, email=email, password=password)
|
251 |
except Exception as e:
|
@@ -272,87 +316,21 @@ async def signup(
|
|
272 |
"token_type": "bearer"
|
273 |
}
|
274 |
|
275 |
-
@router.post("/
|
276 |
-
async def
|
277 |
-
|
278 |
-
OTP-based login endpoint.
|
279 |
-
The user must supply their email and the OTP received via email.
|
280 |
-
"""
|
281 |
-
user = get_user(otp_data.email)
|
282 |
if not user:
|
283 |
-
logger.warning(f"
|
284 |
-
raise HTTPException(status_code=401, detail="
|
285 |
-
|
286 |
-
# Check if OTP exists and is valid
|
287 |
-
stored_otp = user.get("otp")
|
288 |
-
otp_expires = user.get("otp_expires")
|
289 |
-
if not stored_otp or not otp_expires:
|
290 |
-
raise HTTPException(status_code=400, detail="OTP not generated. Please request a new OTP.")
|
291 |
-
|
292 |
-
# Check if OTP is expired (using ISO format stored in DB)
|
293 |
-
if datetime.utcnow() > datetime.fromisoformat(otp_expires):
|
294 |
-
raise HTTPException(status_code=400, detail="OTP has expired. Please request a new OTP.")
|
295 |
-
|
296 |
-
if otp_data.otp != stored_otp:
|
297 |
-
raise HTTPException(status_code=401, detail="Invalid OTP.")
|
298 |
-
|
299 |
-
# OTP is valid. Remove OTP fields.
|
300 |
-
users_collection.update_one(
|
301 |
-
{"email": otp_data.email},
|
302 |
-
{"$unset": {"otp": "", "otp_expires": ""}}
|
303 |
-
)
|
304 |
-
logger.info(f"User logged in with OTP: {otp_data.email}")
|
305 |
return {
|
306 |
-
"access_token": create_access_token(
|
307 |
-
"refresh_token": create_refresh_token(
|
308 |
"token_type": "bearer",
|
309 |
"name": user["name"],
|
310 |
"avatar": user.get("avatar")
|
311 |
}
|
312 |
|
313 |
-
@router.post("/send_otp")
|
314 |
-
async def send_otp(receiver_email: EmailStr = Body(..., embed=True)):
|
315 |
-
"""
|
316 |
-
Generates an OTP, stores it in the user's document with an expiration,
|
317 |
-
and sends it via email to the provided receiver_email.
|
318 |
-
(For demonstration purposes, the OTP is returned. Remove in production.)
|
319 |
-
"""
|
320 |
-
user = get_user(receiver_email)
|
321 |
-
if not user:
|
322 |
-
raise HTTPException(status_code=404, detail="User not found. Please sign up first.")
|
323 |
-
|
324 |
-
sender_email = os.getenv("SENDER_EMAIL")
|
325 |
-
sender_password = os.getenv("SENDER_PASSWORD")
|
326 |
-
if not sender_email or not sender_password:
|
327 |
-
raise HTTPException(status_code=500, detail="Email sender credentials are not configured.")
|
328 |
-
|
329 |
-
otp = generate_otp(6)
|
330 |
-
# Set OTP to expire in 5 minutes
|
331 |
-
otp_expires = (datetime.utcnow() + timedelta(minutes=1)).isoformat()
|
332 |
-
subject = "Your One-Time Password (OTP)"
|
333 |
-
body = f"""\
|
334 |
-
# Your OTP
|
335 |
-
|
336 |
-
Here is your one-time password (OTP):
|
337 |
-
|
338 |
-
`{otp}`
|
339 |
-
|
340 |
-
Please use the code above to proceed with your login.
|
341 |
-
"""
|
342 |
-
try:
|
343 |
-
send_email_func(sender_email, sender_password, receiver_email, subject, body)
|
344 |
-
except Exception as e:
|
345 |
-
raise HTTPException(status_code=500, detail=f"Error sending OTP email: {str(e)}")
|
346 |
-
|
347 |
-
# Update user document with OTP and expiration
|
348 |
-
users_collection.update_one(
|
349 |
-
{"email": receiver_email},
|
350 |
-
{"$set": {"otp": otp, "otp_expires": otp_expires}}
|
351 |
-
)
|
352 |
-
|
353 |
-
# For demonstration, return the OTP (remove in production)
|
354 |
-
return {"message": "OTP sent successfully", "otp": otp}
|
355 |
-
|
356 |
@router.get("/user/data")
|
357 |
async def get_user_data(request: Request, current_user: dict = Depends(get_current_user)):
|
358 |
return {
|
|
|
1 |
import os
|
2 |
import uuid
|
3 |
import logging
|
|
|
|
|
|
|
|
|
4 |
from datetime import datetime, timedelta
|
5 |
from urllib.parse import quote_plus
|
6 |
from typing import List, Optional, Any
|
|
|
15 |
from pymongo import MongoClient
|
16 |
import os.path
|
17 |
|
18 |
+
# ----- OTP and Email Imports -----
|
19 |
+
import smtplib
|
20 |
+
from email.mime.multipart import MIMEMultipart
|
21 |
+
from email.mime.text import MIMEText
|
22 |
try:
|
23 |
import markdown
|
24 |
except ImportError:
|
25 |
markdown = None
|
26 |
+
import secrets
|
27 |
+
# ----------------------------------
|
28 |
|
29 |
load_dotenv()
|
30 |
|
|
|
38 |
client = MongoClient(MONGO_URL)
|
39 |
db = client.users_database
|
40 |
users_collection = db.users
|
41 |
+
# New collection to store OTP records
|
42 |
+
otp_collection = db.otp_verifications
|
43 |
|
44 |
# OAuth2 setup
|
45 |
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
|
|
47 |
# Create an APIRouter instance
|
48 |
router = APIRouter()
|
49 |
|
50 |
+
# Pydantic models
|
51 |
class User(BaseModel):
|
52 |
name: str = Field(..., min_length=3, max_length=50)
|
53 |
email: EmailStr
|
|
|
99 |
class TokenData(BaseModel):
|
100 |
email: Optional[str] = None
|
101 |
|
102 |
+
# Password hashing
|
|
|
|
|
|
|
|
|
|
|
103 |
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
104 |
|
105 |
def verify_password(plain_password: str, hashed_password: str) -> bool:
|
|
|
108 |
def get_password_hash(password: str) -> str:
|
109 |
return pwd_context.hash(password)
|
110 |
|
|
|
111 |
def get_user(email: str) -> Optional[dict]:
|
112 |
return users_collection.find_one({"email": email})
|
113 |
|
|
|
145 |
except JWTError:
|
146 |
raise HTTPException(status_code=401, detail="Invalid token")
|
147 |
|
148 |
+
# Setup for avatar file saving
|
149 |
AVATAR_DIR = "avatars"
|
150 |
if not os.path.exists(AVATAR_DIR):
|
151 |
os.makedirs(AVATAR_DIR)
|
|
|
173 |
file.file.close()
|
174 |
return file_path
|
175 |
|
176 |
+
# ----- OTP and Email Functions -----
|
177 |
def generate_otp(length=6):
|
178 |
"""Generate a numeric OTP of specified length."""
|
179 |
digits = "0123456789"
|
|
|
231 |
except Exception as e:
|
232 |
logger.error(f"Error sending email: {e}")
|
233 |
raise e
|
234 |
+
# ------------------------------------
|
235 |
|
236 |
+
# ----- Auth Endpoints -----
|
237 |
+
|
238 |
+
@router.post("/send_otp")
|
239 |
+
async def send_otp(receiver_email: EmailStr = Body(..., embed=True)):
|
240 |
+
"""
|
241 |
+
Generates an OTP and sends it via email to the provided receiver_email.
|
242 |
+
(For demonstration purposes, the OTP is returned. Remove it in production.)
|
243 |
+
"""
|
244 |
+
sender_email = os.getenv("SENDER_EMAIL")
|
245 |
+
sender_password = os.getenv("SENDER_PASSWORD")
|
246 |
+
if not sender_email or not sender_password:
|
247 |
+
raise HTTPException(status_code=500, detail="Email sender credentials are not configured.")
|
248 |
+
|
249 |
+
otp = generate_otp(6)
|
250 |
+
subject = "Your One-Time Password (OTP)"
|
251 |
+
body = f"""\
|
252 |
+
# Your OTP
|
253 |
+
|
254 |
+
Here is your one-time password (OTP):
|
255 |
+
|
256 |
+
`{otp}`
|
257 |
+
|
258 |
+
Please use the code above to verify your email and proceed with account creation.
|
259 |
+
"""
|
260 |
+
try:
|
261 |
+
send_email_func(sender_email, sender_password, receiver_email, subject, body)
|
262 |
+
# Store OTP record with a 10-minute expiration
|
263 |
+
otp_record = {
|
264 |
+
"email": receiver_email,
|
265 |
+
"otp": otp,
|
266 |
+
"expires_at": datetime.utcnow() + timedelta(minutes=1)
|
267 |
+
}
|
268 |
+
otp_collection.update_one({"email": receiver_email}, {"$set": otp_record}, upsert=True)
|
269 |
+
except Exception as e:
|
270 |
+
raise HTTPException(status_code=500, detail=f"Error sending OTP email: {str(e)}")
|
271 |
+
return {"message": "OTP sent successfully", "otp": otp} # Remove OTP from response in production
|
272 |
|
273 |
@router.post("/signup", response_model=Token)
|
274 |
async def signup(
|
|
|
276 |
name: str = Form(...),
|
277 |
email: EmailStr = Form(...),
|
278 |
password: str = Form(...),
|
279 |
+
otp: str = Form(...),
|
280 |
avatar: Optional[UploadFile] = File(None)
|
281 |
):
|
282 |
+
# Verify OTP for this email
|
283 |
+
otp_record = otp_collection.find_one({"email": email})
|
284 |
+
if not otp_record:
|
285 |
+
raise HTTPException(status_code=400, detail="No OTP sent to this email. Please request an OTP first.")
|
286 |
+
if otp_record["otp"] != otp:
|
287 |
+
raise HTTPException(status_code=400, detail="Invalid OTP provided.")
|
288 |
+
if datetime.utcnow() > otp_record["expires_at"]:
|
289 |
+
raise HTTPException(status_code=400, detail="OTP has expired. Please request a new one.")
|
290 |
+
# Remove the OTP record after successful verification
|
291 |
+
otp_collection.delete_one({"email": email})
|
292 |
+
|
293 |
try:
|
294 |
_ = User(name=name, email=email, password=password)
|
295 |
except Exception as e:
|
|
|
316 |
"token_type": "bearer"
|
317 |
}
|
318 |
|
319 |
+
@router.post("/login", response_model=LoginResponse)
|
320 |
+
async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends()):
|
321 |
+
user = authenticate_user(form_data.username, form_data.password)
|
|
|
|
|
|
|
|
|
322 |
if not user:
|
323 |
+
logger.warning(f"Failed login attempt for: {form_data.username}")
|
324 |
+
raise HTTPException(status_code=401, detail="Incorrect username or password")
|
325 |
+
logger.info(f"User logged in: {user['email']}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
326 |
return {
|
327 |
+
"access_token": create_access_token(user["email"]),
|
328 |
+
"refresh_token": create_refresh_token(user["email"]),
|
329 |
"token_type": "bearer",
|
330 |
"name": user["name"],
|
331 |
"avatar": user.get("avatar")
|
332 |
}
|
333 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
334 |
@router.get("/user/data")
|
335 |
async def get_user_data(request: Request, current_user: dict = Depends(get_current_user)):
|
336 |
return {
|