Spaces:
Sleeping
Sleeping
Update auth.py
Browse files
auth.py
CHANGED
@@ -1,655 +1,656 @@
|
|
1 |
-
import sqlite3
|
2 |
-
|
3 |
-
import
|
4 |
-
import
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
from
|
9 |
-
from
|
10 |
-
from fastapi
|
11 |
-
from
|
12 |
-
from
|
13 |
-
from
|
14 |
-
import
|
15 |
-
|
16 |
-
import
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
logging
|
23 |
-
|
24 |
-
|
25 |
-
)
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
#
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
#
|
61 |
-
SUBSCRIPTION_TIERS
|
62 |
-
|
63 |
-
|
64 |
-
"
|
65 |
-
"
|
66 |
-
"
|
67 |
-
|
68 |
-
"
|
69 |
-
"
|
70 |
-
"
|
71 |
-
"
|
72 |
-
"
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
"
|
78 |
-
"
|
79 |
-
"
|
80 |
-
|
81 |
-
"
|
82 |
-
"
|
83 |
-
"
|
84 |
-
"
|
85 |
-
"
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
"
|
91 |
-
"
|
92 |
-
"
|
93 |
-
|
94 |
-
"
|
95 |
-
"
|
96 |
-
"
|
97 |
-
"
|
98 |
-
"
|
99 |
-
|
100 |
-
|
101 |
-
}
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
os.
|
110 |
-
|
111 |
-
|
112 |
-
conn
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
#
|
120 |
-
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
188 |
-
|
189 |
-
|
190 |
-
|
191 |
-
#
|
192 |
-
|
193 |
-
|
194 |
-
#
|
195 |
-
#
|
196 |
-
#
|
197 |
-
#
|
198 |
-
#
|
199 |
-
#
|
200 |
-
#
|
201 |
-
#
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
salt
|
213 |
-
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
if
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
228 |
-
|
229 |
-
|
230 |
-
|
231 |
-
|
232 |
-
|
233 |
-
|
234 |
-
|
235 |
-
|
236 |
-
|
237 |
-
|
238 |
-
|
239 |
-
|
240 |
-
|
241 |
-
|
242 |
-
|
243 |
-
|
244 |
-
|
245 |
-
|
246 |
-
|
247 |
-
|
248 |
-
|
249 |
-
|
250 |
-
|
251 |
-
|
252 |
-
|
253 |
-
|
254 |
-
|
255 |
-
|
256 |
-
|
257 |
-
|
258 |
-
|
259 |
-
|
260 |
-
|
261 |
-
|
262 |
-
|
263 |
-
|
264 |
-
|
265 |
-
|
266 |
-
|
267 |
-
|
268 |
-
|
269 |
-
"
|
270 |
-
"
|
271 |
-
|
272 |
-
|
273 |
-
|
274 |
-
|
275 |
-
|
276 |
-
|
277 |
-
|
278 |
-
|
279 |
-
|
280 |
-
#
|
281 |
-
#
|
282 |
-
|
283 |
-
|
284 |
-
|
285 |
-
|
286 |
-
|
287 |
-
|
288 |
-
|
289 |
-
|
290 |
-
|
291 |
-
|
292 |
-
|
293 |
-
|
294 |
-
|
295 |
-
|
296 |
-
|
297 |
-
|
298 |
-
logger.info(f"
|
299 |
-
|
300 |
-
|
301 |
-
|
302 |
-
|
303 |
-
|
304 |
-
|
305 |
-
|
306 |
-
|
307 |
-
|
308 |
-
|
309 |
-
|
310 |
-
|
311 |
-
|
312 |
-
|
313 |
-
|
314 |
-
|
315 |
-
|
316 |
-
|
317 |
-
|
318 |
-
|
319 |
-
|
320 |
-
|
321 |
-
|
322 |
-
|
323 |
-
|
324 |
-
|
325 |
-
return User
|
326 |
-
|
327 |
-
|
328 |
-
|
329 |
-
|
330 |
-
|
331 |
-
|
332 |
-
|
333 |
-
|
334 |
-
|
335 |
-
|
336 |
-
|
337 |
-
|
338 |
-
|
339 |
-
|
340 |
-
|
341 |
-
|
342 |
-
|
343 |
-
|
344 |
-
|
345 |
-
|
346 |
-
|
347 |
-
|
348 |
-
|
349 |
-
|
350 |
-
"
|
351 |
-
|
352 |
-
|
353 |
-
|
354 |
-
|
355 |
-
|
356 |
-
|
357 |
-
|
358 |
-
|
359 |
-
|
360 |
-
|
361 |
-
|
362 |
-
|
363 |
-
|
364 |
-
|
365 |
-
|
366 |
-
|
367 |
-
|
368 |
-
|
369 |
-
|
370 |
-
|
371 |
-
|
372 |
-
|
373 |
-
|
374 |
-
|
375 |
-
|
376 |
-
|
377 |
-
|
378 |
-
|
379 |
-
|
380 |
-
|
381 |
-
|
382 |
-
|
383 |
-
|
384 |
-
|
385 |
-
|
386 |
-
|
387 |
-
|
388 |
-
|
389 |
-
|
390 |
-
|
391 |
-
|
392 |
-
|
393 |
-
|
394 |
-
|
395 |
-
|
396 |
-
|
397 |
-
|
398 |
-
|
399 |
-
|
400 |
-
|
401 |
-
|
402 |
-
|
403 |
-
|
404 |
-
|
405 |
-
|
406 |
-
|
407 |
-
|
408 |
-
|
409 |
-
cursor
|
410 |
-
|
411 |
-
|
412 |
-
|
413 |
-
|
414 |
-
|
415 |
-
|
416 |
-
|
417 |
-
|
418 |
-
|
419 |
-
|
420 |
-
|
421 |
-
|
422 |
-
|
423 |
-
|
424 |
-
|
425 |
-
|
426 |
-
|
427 |
-
|
428 |
-
|
429 |
-
|
430 |
-
|
431 |
-
|
432 |
-
|
433 |
-
|
434 |
-
|
435 |
-
|
436 |
-
|
437 |
-
|
438 |
-
|
439 |
-
|
440 |
-
|
441 |
-
|
442 |
-
|
443 |
-
|
444 |
-
|
445 |
-
|
446 |
-
|
447 |
-
|
448 |
-
|
449 |
-
|
450 |
-
|
451 |
-
|
452 |
-
|
453 |
-
|
454 |
-
|
455 |
-
|
456 |
-
|
457 |
-
|
458 |
-
|
459 |
-
|
460 |
-
|
461 |
-
|
462 |
-
|
463 |
-
|
464 |
-
|
465 |
-
|
466 |
-
|
467 |
-
|
468 |
-
|
469 |
-
|
470 |
-
|
471 |
-
|
472 |
-
|
473 |
-
|
474 |
-
"
|
475 |
-
"
|
476 |
-
"
|
477 |
-
"
|
478 |
-
"
|
479 |
-
|
480 |
-
|
481 |
-
|
482 |
-
|
483 |
-
|
484 |
-
|
485 |
-
|
486 |
-
|
487 |
-
|
488 |
-
|
489 |
-
|
490 |
-
|
491 |
-
|
492 |
-
|
493 |
-
|
494 |
-
|
495 |
-
|
496 |
-
|
497 |
-
|
498 |
-
|
499 |
-
|
500 |
-
|
501 |
-
|
502 |
-
|
503 |
-
|
504 |
-
|
505 |
-
|
506 |
-
|
507 |
-
|
508 |
-
|
509 |
-
|
510 |
-
|
511 |
-
|
512 |
-
|
513 |
-
|
514 |
-
|
515 |
-
|
516 |
-
|
517 |
-
|
518 |
-
|
519 |
-
|
520 |
-
|
521 |
-
|
522 |
-
|
523 |
-
|
524 |
-
|
525 |
-
|
526 |
-
|
527 |
-
|
528 |
-
|
529 |
-
|
530 |
-
user.
|
531 |
-
|
532 |
-
|
533 |
-
c.
|
534 |
-
|
535 |
-
|
536 |
-
|
537 |
-
|
538 |
-
|
539 |
-
|
540 |
-
|
541 |
-
|
542 |
-
|
543 |
-
|
544 |
-
if user
|
545 |
-
|
546 |
-
|
547 |
-
|
548 |
-
|
549 |
-
|
550 |
-
|
551 |
-
|
552 |
-
|
553 |
-
|
554 |
-
|
555 |
-
|
556 |
-
|
557 |
-
|
558 |
-
|
559 |
-
if
|
560 |
-
|
561 |
-
|
562 |
-
|
563 |
-
|
564 |
-
|
565 |
-
|
566 |
-
|
567 |
-
|
568 |
-
|
569 |
-
|
570 |
-
|
571 |
-
c.
|
572 |
-
|
573 |
-
|
574 |
-
|
575 |
-
|
576 |
-
|
577 |
-
|
578 |
-
|
579 |
-
|
580 |
-
|
581 |
-
|
582 |
-
|
583 |
-
|
584 |
-
|
585 |
-
|
586 |
-
user.
|
587 |
-
|
588 |
-
|
589 |
-
|
590 |
-
c.
|
591 |
-
|
592 |
-
|
593 |
-
|
594 |
-
|
595 |
-
|
596 |
-
|
597 |
-
|
598 |
-
|
599 |
-
|
600 |
-
|
601 |
-
|
602 |
-
|
603 |
-
|
604 |
-
|
605 |
-
|
606 |
-
|
607 |
-
|
608 |
-
|
609 |
-
|
610 |
-
|
611 |
-
|
612 |
-
|
613 |
-
|
614 |
-
|
615 |
-
|
616 |
-
|
617 |
-
c
|
618 |
-
|
619 |
-
|
620 |
-
|
621 |
-
|
622 |
-
|
623 |
-
|
624 |
-
|
625 |
-
|
626 |
-
|
627 |
-
|
628 |
-
|
629 |
-
|
630 |
-
|
631 |
-
|
632 |
-
"
|
633 |
-
"
|
634 |
-
"
|
635 |
-
"
|
636 |
-
|
637 |
-
|
638 |
-
|
639 |
-
|
640 |
-
|
641 |
-
|
642 |
-
|
643 |
-
|
644 |
-
|
645 |
-
|
646 |
-
|
647 |
-
|
648 |
-
|
649 |
-
|
650 |
-
"
|
651 |
-
"
|
652 |
-
"
|
653 |
-
"
|
654 |
-
|
|
|
655 |
return plans
|
|
|
1 |
+
import sqlite3
|
2 |
+
|
3 |
+
import uuid
|
4 |
+
import os
|
5 |
+
import logging
|
6 |
+
from datetime import datetime, timedelta
|
7 |
+
import hashlib # Use hashlib instead of jwt
|
8 |
+
from passlib.hash import bcrypt
|
9 |
+
from dotenv import load_dotenv
|
10 |
+
from fastapi import Depends, HTTPException
|
11 |
+
from fastapi.security import OAuth2PasswordBearer
|
12 |
+
from pydantic import BaseModel
|
13 |
+
from typing import Optional
|
14 |
+
from fastapi import HTTPException, status
|
15 |
+
import jwt
|
16 |
+
from jose import JWTError
|
17 |
+
import sqlite3
|
18 |
+
|
19 |
+
# Load environment variables
|
20 |
+
load_dotenv()
|
21 |
+
|
22 |
+
# Configure logging
|
23 |
+
logging.basicConfig(
|
24 |
+
level=logging.INFO,
|
25 |
+
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
26 |
+
)
|
27 |
+
logger = logging.getLogger('auth')
|
28 |
+
|
29 |
+
# Security configuration
|
30 |
+
SECRET_KEY = os.getenv("JWT_SECRET", "your-secret-key-for-development-only")
|
31 |
+
ALGORITHM = "HS256"
|
32 |
+
JWT_EXPIRATION_DELTA = timedelta(days=1) # Token valid for 1 day
|
33 |
+
# Database path from environment variable or default
|
34 |
+
# Fix the incorrect DB_PATH
|
35 |
+
DB_PATH = os.getenv("DB_PATH", os.path.join(os.path.dirname(__file__), "data/user_data.db"))
|
36 |
+
|
37 |
+
# FastAPI OAuth2 scheme
|
38 |
+
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
39 |
+
|
40 |
+
# Pydantic models for FastAPI
|
41 |
+
class User(BaseModel):
|
42 |
+
id: str
|
43 |
+
email: str
|
44 |
+
subscription_tier: str = "free_tier"
|
45 |
+
subscription_expiry: Optional[datetime] = None
|
46 |
+
api_calls_remaining: int = 5
|
47 |
+
last_reset_date: Optional[datetime] = None
|
48 |
+
|
49 |
+
class UserCreate(BaseModel):
|
50 |
+
email: str
|
51 |
+
password: str
|
52 |
+
|
53 |
+
class Token(BaseModel):
|
54 |
+
access_token: str
|
55 |
+
token_type: str
|
56 |
+
|
57 |
+
class TokenData(BaseModel):
|
58 |
+
user_id: Optional[str] = None
|
59 |
+
|
60 |
+
# Subscription tiers and limits
|
61 |
+
# Update the SUBSCRIPTION_TIERS dictionary
|
62 |
+
SUBSCRIPTION_TIERS = {
|
63 |
+
"free_tier": {
|
64 |
+
"price": 0,
|
65 |
+
"currency": "INR",
|
66 |
+
"features": ["basic_document_analysis", "basic_risk_assessment"],
|
67 |
+
"limits": {
|
68 |
+
"document_size_mb": 5,
|
69 |
+
"documents_per_month": 3,
|
70 |
+
"video_size_mb": 0,
|
71 |
+
"audio_size_mb": 0,
|
72 |
+
"daily_api_calls": 10, # <-- Add this
|
73 |
+
"max_document_size_mb": 5 # <-- Add this
|
74 |
+
}
|
75 |
+
},
|
76 |
+
"standard_tier": {
|
77 |
+
"price": 799,
|
78 |
+
"currency": "INR",
|
79 |
+
"features": ["basic_document_analysis", "basic_risk_assessment", "video_analysis", "audio_analysis", "chatbot"],
|
80 |
+
"limits": {
|
81 |
+
"document_size_mb": 20,
|
82 |
+
"documents_per_month": 20,
|
83 |
+
"video_size_mb": 100,
|
84 |
+
"audio_size_mb": 50,
|
85 |
+
"daily_api_calls": 100, # <-- Add this
|
86 |
+
"max_document_size_mb": 20 # <-- Add this
|
87 |
+
}
|
88 |
+
},
|
89 |
+
"premium_tier": {
|
90 |
+
"price": 1499,
|
91 |
+
"currency": "INR",
|
92 |
+
"features": ["basic_document_analysis", "basic_risk_assessment", "video_analysis", "audio_analysis", "chatbot", "detailed_risk_assessment", "contract_clause_analysis"],
|
93 |
+
"limits": {
|
94 |
+
"document_size_mb": 50,
|
95 |
+
"documents_per_month": 999999,
|
96 |
+
"video_size_mb": 500,
|
97 |
+
"audio_size_mb": 200,
|
98 |
+
"daily_api_calls": 1000, # <-- Add this
|
99 |
+
"max_document_size_mb": 50 # <-- Add this
|
100 |
+
}
|
101 |
+
}
|
102 |
+
}
|
103 |
+
|
104 |
+
# Database connection management
|
105 |
+
def get_db_connection():
|
106 |
+
"""Create and return a database connection with proper error handling"""
|
107 |
+
try:
|
108 |
+
# Ensure the directory exists
|
109 |
+
db_dir = os.path.dirname(DB_PATH)
|
110 |
+
os.makedirs(db_dir, exist_ok=True)
|
111 |
+
|
112 |
+
conn = sqlite3.connect(DB_PATH)
|
113 |
+
conn.row_factory = sqlite3.Row # Return rows as dictionaries
|
114 |
+
return conn
|
115 |
+
except sqlite3.Error as e:
|
116 |
+
logger.error(f"Database connection error: {e}")
|
117 |
+
raise Exception(f"Database connection failed: {e}")
|
118 |
+
|
119 |
+
# Database setup
|
120 |
+
# In the init_auth_db function, update the CREATE TABLE statement to match our schema
|
121 |
+
def init_auth_db():
|
122 |
+
"""Initialize the authentication database with required tables"""
|
123 |
+
try:
|
124 |
+
conn = get_db_connection()
|
125 |
+
c = conn.cursor()
|
126 |
+
|
127 |
+
# Create users table with the correct schema
|
128 |
+
c.execute('''
|
129 |
+
CREATE TABLE IF NOT EXISTS users (
|
130 |
+
id TEXT PRIMARY KEY,
|
131 |
+
email TEXT UNIQUE NOT NULL,
|
132 |
+
hashed_password TEXT NOT NULL,
|
133 |
+
password TEXT,
|
134 |
+
subscription_tier TEXT DEFAULT 'free_tier',
|
135 |
+
is_active BOOLEAN DEFAULT 1,
|
136 |
+
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
137 |
+
api_calls_remaining INTEGER DEFAULT 10,
|
138 |
+
last_reset_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
139 |
+
)
|
140 |
+
''')
|
141 |
+
|
142 |
+
# Create subscriptions table
|
143 |
+
c.execute('''
|
144 |
+
CREATE TABLE IF NOT EXISTS subscriptions (
|
145 |
+
id TEXT PRIMARY KEY,
|
146 |
+
user_id TEXT,
|
147 |
+
tier TEXT,
|
148 |
+
plan_id TEXT,
|
149 |
+
status TEXT,
|
150 |
+
created_at TIMESTAMP,
|
151 |
+
expires_at TIMESTAMP,
|
152 |
+
paypal_subscription_id TEXT,
|
153 |
+
FOREIGN KEY (user_id) REFERENCES users (id)
|
154 |
+
)
|
155 |
+
''')
|
156 |
+
|
157 |
+
# Create usage stats table
|
158 |
+
c.execute('''
|
159 |
+
CREATE TABLE IF NOT EXISTS usage_stats (
|
160 |
+
id TEXT PRIMARY KEY,
|
161 |
+
user_id TEXT,
|
162 |
+
month INTEGER,
|
163 |
+
year INTEGER,
|
164 |
+
analyses_used INTEGER,
|
165 |
+
FOREIGN KEY (user_id) REFERENCES users (id)
|
166 |
+
)
|
167 |
+
''')
|
168 |
+
|
169 |
+
# Create tokens table for refresh tokens
|
170 |
+
c.execute('''
|
171 |
+
CREATE TABLE IF NOT EXISTS refresh_tokens (
|
172 |
+
user_id TEXT,
|
173 |
+
token TEXT,
|
174 |
+
expires_at TIMESTAMP,
|
175 |
+
FOREIGN KEY (user_id) REFERENCES users (id)
|
176 |
+
)
|
177 |
+
''')
|
178 |
+
|
179 |
+
conn.commit()
|
180 |
+
logger.info("Database initialized successfully")
|
181 |
+
except Exception as e:
|
182 |
+
logger.error(f"Database initialization error: {e}")
|
183 |
+
raise
|
184 |
+
finally:
|
185 |
+
if conn:
|
186 |
+
conn.close()
|
187 |
+
|
188 |
+
# Initialize the database
|
189 |
+
init_auth_db()
|
190 |
+
|
191 |
+
# Password hashing with bcrypt
|
192 |
+
# Update the password hashing and verification functions to use a more reliable method
|
193 |
+
|
194 |
+
# Replace these functions
|
195 |
+
# Remove these conflicting functions
|
196 |
+
# def hash_password(password):
|
197 |
+
# """Hash a password using bcrypt"""
|
198 |
+
# return bcrypt.hash(password)
|
199 |
+
#
|
200 |
+
# def verify_password(plain_password, hashed_password):
|
201 |
+
# """Verify a password against its hash"""
|
202 |
+
# return bcrypt.verify(plain_password, hashed_password)
|
203 |
+
|
204 |
+
# Keep only these improved functions
|
205 |
+
def hash_password(password):
|
206 |
+
"""Hash a password using bcrypt"""
|
207 |
+
# Use a more direct approach to avoid bcrypt version issues
|
208 |
+
import bcrypt
|
209 |
+
# Convert password to bytes if it's not already
|
210 |
+
if isinstance(password, str):
|
211 |
+
password = password.encode('utf-8')
|
212 |
+
# Generate salt and hash
|
213 |
+
salt = bcrypt.gensalt()
|
214 |
+
hashed = bcrypt.hashpw(password, salt)
|
215 |
+
# Return as string for storage
|
216 |
+
return hashed.decode('utf-8')
|
217 |
+
|
218 |
+
def verify_password(plain_password, hashed_password):
|
219 |
+
"""Verify a password against its hash"""
|
220 |
+
import bcrypt
|
221 |
+
# Convert inputs to bytes if they're not already
|
222 |
+
if isinstance(plain_password, str):
|
223 |
+
plain_password = plain_password.encode('utf-8')
|
224 |
+
if isinstance(hashed_password, str):
|
225 |
+
hashed_password = hashed_password.encode('utf-8')
|
226 |
+
|
227 |
+
try:
|
228 |
+
# Use direct bcrypt verification
|
229 |
+
return bcrypt.checkpw(plain_password, hashed_password)
|
230 |
+
except Exception as e:
|
231 |
+
logger.error(f"Password verification error: {e}")
|
232 |
+
return False
|
233 |
+
|
234 |
+
# User registration
|
235 |
+
def register_user(email, password):
|
236 |
+
try:
|
237 |
+
conn = get_db_connection()
|
238 |
+
c = conn.cursor()
|
239 |
+
|
240 |
+
# Check if user already exists
|
241 |
+
c.execute("SELECT * FROM users WHERE email = ?", (email,))
|
242 |
+
if c.fetchone():
|
243 |
+
return False, "Email already registered"
|
244 |
+
|
245 |
+
# Create new user
|
246 |
+
user_id = str(uuid.uuid4())
|
247 |
+
|
248 |
+
# Add more detailed logging
|
249 |
+
logger.info(f"Registering new user with email: {email}")
|
250 |
+
hashed_pw = hash_password(password)
|
251 |
+
logger.info(f"Password hashed successfully: {bool(hashed_pw)}")
|
252 |
+
|
253 |
+
c.execute("""
|
254 |
+
INSERT INTO users
|
255 |
+
(id, email, hashed_password, subscription_tier, api_calls_remaining, last_reset_date)
|
256 |
+
VALUES (?, ?, ?, ?, ?, ?)
|
257 |
+
""", (user_id, email, hashed_pw, "free_tier", 5, datetime.now()))
|
258 |
+
|
259 |
+
conn.commit()
|
260 |
+
logger.info(f"User registered successfully: {email}")
|
261 |
+
|
262 |
+
# Verify the user was actually stored
|
263 |
+
c.execute("SELECT * FROM users WHERE email = ?", (email,))
|
264 |
+
stored_user = c.fetchone()
|
265 |
+
logger.info(f"User verification after registration: {bool(stored_user)}")
|
266 |
+
|
267 |
+
access_token = create_access_token(user_id)
|
268 |
+
return True, {
|
269 |
+
"user_id": user_id,
|
270 |
+
"access_token": access_token,
|
271 |
+
"token_type": "bearer"
|
272 |
+
}
|
273 |
+
except Exception as e:
|
274 |
+
logger.error(f"User registration error: {e}")
|
275 |
+
return False, f"Registration failed: {str(e)}"
|
276 |
+
finally:
|
277 |
+
if conn:
|
278 |
+
conn.close()
|
279 |
+
|
280 |
+
# User login
|
281 |
+
# Fix the authenticate_user function
|
282 |
+
# In the authenticate_user function, update the password verification to use hashed_password
|
283 |
+
def authenticate_user(email, password):
|
284 |
+
"""Authenticate a user and return user data with tokens"""
|
285 |
+
try:
|
286 |
+
conn = get_db_connection()
|
287 |
+
c = conn.cursor()
|
288 |
+
|
289 |
+
# Get user by email
|
290 |
+
c.execute("SELECT * FROM users WHERE email = ? AND is_active = 1", (email,))
|
291 |
+
user = c.fetchone()
|
292 |
+
|
293 |
+
if not user:
|
294 |
+
logger.warning(f"User not found: {email}")
|
295 |
+
return None
|
296 |
+
|
297 |
+
# Add debug logging for password verification
|
298 |
+
logger.info(f"Verifying password for user: {email}")
|
299 |
+
logger.info(f"Stored hashed password: {user['hashed_password'][:20]}...")
|
300 |
+
|
301 |
+
try:
|
302 |
+
# Check if password verification works
|
303 |
+
is_valid = verify_password(password, user['hashed_password'])
|
304 |
+
logger.info(f"Password verification result: {is_valid}")
|
305 |
+
|
306 |
+
if not is_valid:
|
307 |
+
logger.warning(f"Password verification failed for user: {email}")
|
308 |
+
return None
|
309 |
+
except Exception as e:
|
310 |
+
logger.error(f"Password verification error: {e}")
|
311 |
+
return None
|
312 |
+
|
313 |
+
# Update last login time if column exists
|
314 |
+
try:
|
315 |
+
c.execute("UPDATE users SET last_login = ? WHERE id = ?",
|
316 |
+
(datetime.now(), user['id']))
|
317 |
+
conn.commit()
|
318 |
+
except sqlite3.OperationalError:
|
319 |
+
# last_login column might not exist
|
320 |
+
pass
|
321 |
+
|
322 |
+
# Convert sqlite3.Row to dict to use get() method
|
323 |
+
user_dict = dict(user)
|
324 |
+
|
325 |
+
# Create and return a User object
|
326 |
+
return User(
|
327 |
+
id=user_dict['id'],
|
328 |
+
email=user_dict['email'],
|
329 |
+
subscription_tier=user_dict.get('subscription_tier', 'free_tier'),
|
330 |
+
subscription_expiry=None, # Handle this properly if needed
|
331 |
+
api_calls_remaining=user_dict.get('api_calls_remaining', 5),
|
332 |
+
last_reset_date=user_dict.get('last_reset_date')
|
333 |
+
)
|
334 |
+
except Exception as e:
|
335 |
+
logger.error(f"Login error: {e}")
|
336 |
+
return None
|
337 |
+
finally:
|
338 |
+
if conn:
|
339 |
+
conn.close()
|
340 |
+
|
341 |
+
# Token generation and validation - completely replaced
|
342 |
+
def create_access_token(user_id):
|
343 |
+
"""Create a new access token for a user"""
|
344 |
+
try:
|
345 |
+
# Create a JWT token with user_id and expiration
|
346 |
+
expiration = datetime.now() + JWT_EXPIRATION_DELTA
|
347 |
+
|
348 |
+
# Create a token payload
|
349 |
+
payload = {
|
350 |
+
"sub": user_id,
|
351 |
+
"exp": expiration.timestamp()
|
352 |
+
}
|
353 |
+
|
354 |
+
# Generate the JWT token
|
355 |
+
token = jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)
|
356 |
+
|
357 |
+
logger.info(f"Created access token for user: {user_id}")
|
358 |
+
return token
|
359 |
+
except Exception as e:
|
360 |
+
logger.error(f"Token creation error: {e}")
|
361 |
+
return None
|
362 |
+
|
363 |
+
|
364 |
+
def update_auth_db_schema():
|
365 |
+
"""Update the authentication database schema with any missing columns"""
|
366 |
+
try:
|
367 |
+
conn = get_db_connection()
|
368 |
+
c = conn.cursor()
|
369 |
+
|
370 |
+
# Check if tier column exists in subscriptions table
|
371 |
+
c.execute("PRAGMA table_info(subscriptions)")
|
372 |
+
columns = [column[1] for column in c.fetchall()]
|
373 |
+
|
374 |
+
# Add tier column if it doesn't exist
|
375 |
+
if "tier" not in columns:
|
376 |
+
logger.info("Adding 'tier' column to subscriptions table")
|
377 |
+
c.execute("ALTER TABLE subscriptions ADD COLUMN tier TEXT")
|
378 |
+
conn.commit()
|
379 |
+
logger.info("Database schema updated successfully")
|
380 |
+
|
381 |
+
conn.close()
|
382 |
+
except Exception as e:
|
383 |
+
logger.error(f"Database schema update error: {e}")
|
384 |
+
raise HTTPException(
|
385 |
+
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
386 |
+
detail=f"Database schema update error: {str(e)}"
|
387 |
+
)
|
388 |
+
|
389 |
+
# Add this to your get_current_user function
|
390 |
+
async def get_current_user(token: str = Depends(oauth2_scheme)):
|
391 |
+
credentials_exception = HTTPException(
|
392 |
+
status_code=status.HTTP_401_UNAUTHORIZED,
|
393 |
+
detail="Could not validate credentials",
|
394 |
+
headers={"WWW-Authenticate": "Bearer"},
|
395 |
+
)
|
396 |
+
try:
|
397 |
+
# Decode the JWT token
|
398 |
+
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
|
399 |
+
user_id: str = payload.get("sub")
|
400 |
+
if user_id is None:
|
401 |
+
logger.error("Token missing 'sub' field")
|
402 |
+
raise credentials_exception
|
403 |
+
except Exception as e:
|
404 |
+
logger.error(f"Token validation error: {str(e)}")
|
405 |
+
raise credentials_exception
|
406 |
+
|
407 |
+
# Get user from database
|
408 |
+
conn = get_db_connection()
|
409 |
+
cursor = conn.cursor()
|
410 |
+
cursor.execute("SELECT id, email, subscription_tier, is_active FROM users WHERE id = ?", (user_id,))
|
411 |
+
user_data = cursor.fetchone()
|
412 |
+
conn.close()
|
413 |
+
|
414 |
+
if user_data is None:
|
415 |
+
logger.error(f"User not found: {user_id}")
|
416 |
+
raise credentials_exception
|
417 |
+
|
418 |
+
user = User(
|
419 |
+
id=user_data[0],
|
420 |
+
email=user_data[1],
|
421 |
+
subscription_tier=user_data[2],
|
422 |
+
is_active=bool(user_data[3])
|
423 |
+
)
|
424 |
+
|
425 |
+
return user
|
426 |
+
|
427 |
+
async def get_current_active_user(current_user: User = Depends(get_current_user)):
|
428 |
+
"""Get the current active user"""
|
429 |
+
return current_user
|
430 |
+
|
431 |
+
def create_user_subscription(email, tier):
|
432 |
+
"""Create a subscription for a user"""
|
433 |
+
try:
|
434 |
+
# Get user by email
|
435 |
+
conn = get_db_connection()
|
436 |
+
c = conn.cursor()
|
437 |
+
|
438 |
+
# Get user ID
|
439 |
+
c.execute("SELECT id FROM users WHERE email = ?", (email,))
|
440 |
+
user_data = c.fetchone()
|
441 |
+
|
442 |
+
if not user_data:
|
443 |
+
return False, "User not found"
|
444 |
+
|
445 |
+
user_id = user_data['id']
|
446 |
+
|
447 |
+
# Check if tier is valid
|
448 |
+
valid_tiers = ["standard_tier", "premium_tier"]
|
449 |
+
if tier not in valid_tiers:
|
450 |
+
return False, f"Invalid tier: {tier}. Must be one of {valid_tiers}"
|
451 |
+
|
452 |
+
# Create subscription
|
453 |
+
subscription_id = str(uuid.uuid4())
|
454 |
+
created_at = datetime.now()
|
455 |
+
expires_at = created_at + timedelta(days=30) # 30-day subscription
|
456 |
+
|
457 |
+
# Insert subscription
|
458 |
+
c.execute("""
|
459 |
+
INSERT INTO subscriptions
|
460 |
+
(id, user_id, tier, status, created_at, expires_at)
|
461 |
+
VALUES (?, ?, ?, ?, ?, ?)
|
462 |
+
""", (subscription_id, user_id, tier, "active", created_at, expires_at))
|
463 |
+
|
464 |
+
# Update user's subscription tier
|
465 |
+
c.execute("""
|
466 |
+
UPDATE users
|
467 |
+
SET subscription_tier = ?
|
468 |
+
WHERE id = ?
|
469 |
+
""", (tier, user_id))
|
470 |
+
|
471 |
+
conn.commit()
|
472 |
+
|
473 |
+
return True, {
|
474 |
+
"id": subscription_id,
|
475 |
+
"user_id": user_id,
|
476 |
+
"tier": tier,
|
477 |
+
"status": "active",
|
478 |
+
"created_at": created_at.isoformat(),
|
479 |
+
"expires_at": expires_at.isoformat()
|
480 |
+
}
|
481 |
+
except Exception as e:
|
482 |
+
logger.error(f"Subscription creation error: {e}")
|
483 |
+
return False, f"Failed to create subscription: {str(e)}"
|
484 |
+
finally:
|
485 |
+
if conn:
|
486 |
+
conn.close()
|
487 |
+
|
488 |
+
def get_user(user_id: str):
|
489 |
+
"""Get user by ID"""
|
490 |
+
try:
|
491 |
+
conn = get_db_connection()
|
492 |
+
c = conn.cursor()
|
493 |
+
|
494 |
+
# Get user
|
495 |
+
c.execute("SELECT * FROM users WHERE id = ? AND is_active = 1", (user_id,))
|
496 |
+
user_data = c.fetchone()
|
497 |
+
|
498 |
+
if not user_data:
|
499 |
+
return None
|
500 |
+
|
501 |
+
# Convert to User model
|
502 |
+
user_dict = dict(user_data)
|
503 |
+
|
504 |
+
# Handle datetime conversions if needed
|
505 |
+
if user_dict.get("subscription_expiry") and isinstance(user_dict["subscription_expiry"], str):
|
506 |
+
user_dict["subscription_expiry"] = datetime.fromisoformat(user_dict["subscription_expiry"])
|
507 |
+
if user_dict.get("last_reset_date") and isinstance(user_dict["last_reset_date"], str):
|
508 |
+
user_dict["last_reset_date"] = datetime.fromisoformat(user_dict["last_reset_date"])
|
509 |
+
|
510 |
+
return User(
|
511 |
+
id=user_dict['id'],
|
512 |
+
email=user_dict['email'],
|
513 |
+
subscription_tier=user_dict['subscription_tier'],
|
514 |
+
subscription_expiry=user_dict.get('subscription_expiry'),
|
515 |
+
api_calls_remaining=user_dict.get('api_calls_remaining', 5),
|
516 |
+
last_reset_date=user_dict.get('last_reset_date')
|
517 |
+
)
|
518 |
+
except Exception as e:
|
519 |
+
logger.error(f"Get user error: {e}")
|
520 |
+
return None
|
521 |
+
finally:
|
522 |
+
if conn:
|
523 |
+
conn.close()
|
524 |
+
|
525 |
+
def check_subscription_access(user: User, feature: str, file_size_mb: Optional[float] = None):
|
526 |
+
"""Check if the user has access to the requested feature and file size"""
|
527 |
+
# Check if subscription is expired
|
528 |
+
if user.subscription_tier != "free_tier" and user.subscription_expiry and user.subscription_expiry < datetime.now():
|
529 |
+
# Downgrade to free tier if subscription expired
|
530 |
+
user.subscription_tier = "free_tier"
|
531 |
+
user.api_calls_remaining = SUBSCRIPTION_TIERS["free_tier"]["daily_api_calls"]
|
532 |
+
with get_db_connection() as conn:
|
533 |
+
c = conn.cursor()
|
534 |
+
c.execute("""
|
535 |
+
UPDATE users
|
536 |
+
SET subscription_tier = ?, api_calls_remaining = ?
|
537 |
+
WHERE id = ?
|
538 |
+
""", (user.subscription_tier, user.api_calls_remaining, user.id))
|
539 |
+
conn.commit()
|
540 |
+
|
541 |
+
# Reset API calls if needed
|
542 |
+
user = reset_api_calls_if_needed(user)
|
543 |
+
|
544 |
+
# Check if user has API calls remaining
|
545 |
+
if user.api_calls_remaining <= 0:
|
546 |
+
raise HTTPException(
|
547 |
+
status_code=429,
|
548 |
+
detail="API call limit reached for today. Please upgrade your subscription or try again tomorrow."
|
549 |
+
)
|
550 |
+
|
551 |
+
# Check if feature is available in user's subscription tier
|
552 |
+
tier_features = SUBSCRIPTION_TIERS[user.subscription_tier]["features"]
|
553 |
+
if feature not in tier_features:
|
554 |
+
raise HTTPException(
|
555 |
+
status_code=403,
|
556 |
+
detail=f"The {feature} feature is not available in your {user.subscription_tier} subscription. Please upgrade to access this feature."
|
557 |
+
)
|
558 |
+
|
559 |
+
# Check file size limit if applicable
|
560 |
+
if file_size_mb:
|
561 |
+
max_size = SUBSCRIPTION_TIERS[user.subscription_tier]["max_document_size_mb"]
|
562 |
+
if file_size_mb > max_size:
|
563 |
+
raise HTTPException(
|
564 |
+
status_code=413,
|
565 |
+
detail=f"File size exceeds the {max_size}MB limit for your {user.subscription_tier} subscription. Please upgrade or use a smaller file."
|
566 |
+
)
|
567 |
+
|
568 |
+
# Decrement API calls remaining
|
569 |
+
user.api_calls_remaining -= 1
|
570 |
+
with get_db_connection() as conn:
|
571 |
+
c = conn.cursor()
|
572 |
+
c.execute("""
|
573 |
+
UPDATE users
|
574 |
+
SET api_calls_remaining = ?
|
575 |
+
WHERE id = ?
|
576 |
+
""", (user.api_calls_remaining, user.id))
|
577 |
+
conn.commit()
|
578 |
+
|
579 |
+
return True
|
580 |
+
|
581 |
+
def reset_api_calls_if_needed(user: User):
|
582 |
+
"""Reset API call counter if it's a new day"""
|
583 |
+
today = datetime.now().date()
|
584 |
+
if user.last_reset_date is None or user.last_reset_date.date() < today:
|
585 |
+
tier_limits = SUBSCRIPTION_TIERS[user.subscription_tier]
|
586 |
+
user.api_calls_remaining = tier_limits["daily_api_calls"]
|
587 |
+
user.last_reset_date = datetime.now()
|
588 |
+
# Update the user in the database
|
589 |
+
with get_db_connection() as conn:
|
590 |
+
c = conn.cursor()
|
591 |
+
c.execute("""
|
592 |
+
UPDATE users
|
593 |
+
SET api_calls_remaining = ?, last_reset_date = ?
|
594 |
+
WHERE id = ?
|
595 |
+
""", (user.api_calls_remaining, user.last_reset_date, user.id))
|
596 |
+
conn.commit()
|
597 |
+
|
598 |
+
return user
|
599 |
+
|
600 |
+
def login_user(email, password):
|
601 |
+
"""Login a user with email and password"""
|
602 |
+
try:
|
603 |
+
# Authenticate user
|
604 |
+
user = authenticate_user(email, password)
|
605 |
+
if not user:
|
606 |
+
return False, "Incorrect username or password"
|
607 |
+
|
608 |
+
# Create access token
|
609 |
+
access_token = create_access_token(user.id)
|
610 |
+
|
611 |
+
# Create refresh token
|
612 |
+
refresh_token = str(uuid.uuid4())
|
613 |
+
expires_at = datetime.now() + timedelta(days=30)
|
614 |
+
|
615 |
+
# Store refresh token
|
616 |
+
conn = get_db_connection()
|
617 |
+
c = conn.cursor()
|
618 |
+
c.execute("INSERT INTO refresh_tokens VALUES (?, ?, ?)",
|
619 |
+
(user.id, refresh_token, expires_at))
|
620 |
+
conn.commit()
|
621 |
+
|
622 |
+
# Get subscription info
|
623 |
+
c.execute("SELECT * FROM subscriptions WHERE user_id = ? AND status = 'active'", (user.id,))
|
624 |
+
subscription = c.fetchone()
|
625 |
+
|
626 |
+
# Convert subscription to dict if it exists, otherwise set to None
|
627 |
+
subscription_dict = dict(subscription) if subscription else None
|
628 |
+
|
629 |
+
conn.close()
|
630 |
+
|
631 |
+
return True, {
|
632 |
+
"user_id": user.id,
|
633 |
+
"email": user.email,
|
634 |
+
"access_token": access_token,
|
635 |
+
"refresh_token": refresh_token,
|
636 |
+
"subscription": subscription_dict
|
637 |
+
}
|
638 |
+
except Exception as e:
|
639 |
+
logger.error(f"Login error: {e}")
|
640 |
+
return False, f"Login failed: {str(e)}"
|
641 |
+
|
642 |
+
|
643 |
+
def get_subscription_plans():
|
644 |
+
"""
|
645 |
+
Returns a list of available subscription plans based on SUBSCRIPTION_TIERS.
|
646 |
+
"""
|
647 |
+
plans = []
|
648 |
+
for tier, details in SUBSCRIPTION_TIERS.items():
|
649 |
+
plans.append({
|
650 |
+
"tier": tier,
|
651 |
+
"price": details["price"],
|
652 |
+
"currency": details["currency"],
|
653 |
+
"features": details["features"],
|
654 |
+
"limits": details["limits"]
|
655 |
+
})
|
656 |
return plans
|