@@ -17,40 +17,104 @@ class SecurityError(Exception):
1717 pass
1818
1919
20+ class AuthWrapper :
21+ """
22+ Wrapper to provide consistent interface for both FileAccessToken and API Key authentication
23+ """
24+ def __init__ (self , domain , wp_path , auth_type , source_obj = None ):
25+ self .domain = domain
26+ self .wp_path = wp_path
27+ self .auth_type = auth_type # 'file_token' or 'api_key'
28+ self .source_obj = source_obj # Original FileAccessToken or AIScannerSettings object
29+
30+
2031def validate_access_token (token , scan_id ):
2132 """
22- Implement proper token validation
23- - Check token format
24- - Verify token hasn't expired
25- - Confirm token is for the correct scan
26- - Log access attempts
33+ Validate authentication token - accepts BOTH file access tokens and API keys
34+
35+ Authentication Flow:
36+ 1. Try FileAccessToken (temporary token for active scans)
37+ 2. If not found, try API Key (for post-scan file operations)
38+
39+ Returns: (AuthWrapper object or None, error_message or None)
2740 """
2841 try :
2942 if not token or not token .startswith ('cp_' ):
3043 logging .writeToFile (f'[API] Invalid token format: { token [:20 ] if token else "None" } ...' )
3144 return None , "Invalid token format"
3245
33- # Find the token in database
46+ # OPTION 1: Try FileAccessToken first (for active scans)
3447 try :
3548 file_token = FileAccessToken .objects .get (
3649 token = token ,
3750 scan_history__scan_id = scan_id ,
3851 is_active = True
3952 )
40-
53+
4154 if file_token .is_expired ():
42- logging .writeToFile (f'[API] Token expired for scan { scan_id } ' )
43- return None , "Token expired"
44-
45- logging .writeToFile (f'[API] Token validated successfully for scan { scan_id } ' )
46- return file_token , None
47-
55+ logging .writeToFile (f'[API] File token expired for scan { scan_id } , trying API key fallback...' )
56+ # Don't return here - fall through to try API key
57+ else :
58+ logging .writeToFile (f'[API] File token validated successfully for scan { scan_id } ' )
59+ return AuthWrapper (
60+ domain = file_token .domain ,
61+ wp_path = file_token .wp_path ,
62+ auth_type = 'file_token' ,
63+ source_obj = file_token
64+ ), None
65+
4866 except FileAccessToken .DoesNotExist :
49- logging .writeToFile (f'[API] Token not found for scan { scan_id } ' )
67+ logging .writeToFile (f'[API] File token not found for scan { scan_id } , trying API key fallback...' )
68+ # Fall through to try API key
69+
70+ # OPTION 2: Try API Key (for post-scan file operations)
71+ try :
72+ from .models import AIScannerSettings , ScanHistory
73+
74+ # Find API key in settings
75+ scanner_settings = AIScannerSettings .objects .get (
76+ api_key = token
77+ )
78+
79+ logging .writeToFile (f'[API] Found API key for admin: { scanner_settings .admin .userName } ' )
80+
81+ # Get the scan to verify it belongs to this admin and get domain/path
82+ try :
83+ scan = ScanHistory .objects .get (
84+ scan_id = scan_id ,
85+ admin = scanner_settings .admin
86+ )
87+
88+ # Get wp_path from website
89+ try :
90+ website = Websites .objects .get (domain = scan .domain )
91+ wp_path = website .path
92+
93+ logging .writeToFile (f'[API] API key validated successfully for scan { scan_id } , domain { scan .domain } ' )
94+
95+ return AuthWrapper (
96+ domain = scan .domain ,
97+ wp_path = wp_path ,
98+ auth_type = 'api_key' ,
99+ source_obj = scanner_settings
100+ ), None
101+
102+ except Websites .DoesNotExist :
103+ logging .writeToFile (f'[API] Website not found for domain: { scan .domain } ' )
104+ return None , "Website not found"
105+
106+ except ScanHistory .DoesNotExist :
107+ logging .writeToFile (f'[API] Scan { scan_id } not found or does not belong to API key owner' )
108+ return None , "Scan not found or access denied"
109+
110+ except AIScannerSettings .DoesNotExist :
111+ logging .writeToFile (f'[API] API key not found in settings' )
50112 return None , "Invalid token"
51-
113+
52114 except Exception as e :
53115 logging .writeToFile (f'[API] Token validation error: { str (e )} ' )
116+ import traceback
117+ logging .writeToFile (f'[API] Traceback: { traceback .format_exc ()} ' )
54118 return None , "Token validation failed"
55119
56120
0 commit comments