Skip to content

Commit 9f101d4

Browse files
committed
add aiscanner file patcher
1 parent 0100d3e commit 9f101d4

File tree

1 file changed

+79
-15
lines changed

1 file changed

+79
-15
lines changed

aiScanner/api.py

Lines changed: 79 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -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+
2031
def 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

Comments
 (0)