ZK-AI
/
ollama_proxy_server
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update main.py

This commit is contained in:
Saifeddine ALOUI 2024-01-22 09:29:56 +01:00 committed by GitHub
parent 416b6d4a65
commit 34274ee0a7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
ollama_proxy_server/main.py
View File

@ -30,8 +30,11 @@ def get_authorized_users(filename):
lines = f.readlines() lines = f.readlines()
authorized_users = {} authorized_users = {}
for line in lines: for line in lines:
user, key = line.strip().split(':') try:
authorized_users[user] = key user, key = line.strip().split(':')
authorized_users[user] = key
except:
ASCIIColors.red(f"User entry broken:{line.strip()}")
return authorized_users return authorized_users
@ -98,7 +101,13 @@ def main():
if not self._validate_user_and_key(): if not self._validate_user_and_key():
ASCIIColors.red(f'User is not authorized') ASCIIColors.red(f'User is not authorized')
client_ip, client_port = self.client_address client_ip, client_port = self.client_address
self.add_access_log_entry(event='rejected', user="unknown", ip_address=client_ip, access="Denied", server="None", nb_queued_requests_on_server=-1, error="Authentication failed") # Extract the bearer token from the headers
auth_header = self.headers.get('Authorization')
if not auth_header or not auth_header.startswith('Bearer '):
self.add_access_log_entry(event='rejected', user="unknown", ip_address=client_ip, access="Denied", server="None", nb_queued_requests_on_server=-1, error="Authentication failed")
else:
token = auth_header.split(' ')[1]
self.add_access_log_entry(event='rejected', user=token, ip_address=client_ip, access="Denied", server="None", nb_queued_requests_on_server=-1, error="Authentication failed")
self.send_response(403) self.send_response(403)
self.end_headers() self.end_headers()
return return