mominah commited on
Commit
3c28e31
·
verified ·
1 Parent(s): c760cc8

Update auth.py

Browse files
Files changed (1) hide show
  1. auth.py +69 -91
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
- # Optionally import markdown to convert markdown text to HTML
 
 
 
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
- # ------------------ Pydantic Models ------------------
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
- # Model for OTP-based login
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
- # ------------------ Avatar File Saving ------------------
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
- # ------------------ OTP and Email Functions ------------------
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
- # ------------------ Auth Endpoints ------------------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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("/login_otp", response_model=LoginResponse)
276
- async def login_otp(otp_data: OTPLogin):
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"Login OTP attempt for non-existent email: {otp_data.email}")
284
- raise HTTPException(status_code=401, detail="User not found")
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(otp_data.email),
307
- "refresh_token": create_refresh_token(otp_data.email),
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 {