157 lines
9.5 KiB
Python
157 lines
9.5 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import secrets
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
import urllib.request
|
|
import urllib.error
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from auth import CredentialStore, initial_record, verify_password
|
|
from server import Server
|
|
|
|
class StoreTests(unittest.TestCase):
|
|
def test_persistence_and_no_plaintext(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path=Path(directory)/'auth.json'
|
|
store=CredentialStore(path)
|
|
password,token=secrets.token_urlsafe(32),secrets.token_urlsafe(32)
|
|
store.setup(password,token)
|
|
raw=path.read_text()
|
|
self.assertNotIn(password,raw);self.assertNotIn(token,raw)
|
|
self.assertEqual(path.stat().st_mode&0o777,0o600)
|
|
restored=CredentialStore(path)
|
|
self.assertTrue(verify_password(password,restored.read()['password']))
|
|
with self.assertRaises(ValueError):restored.setup(password,token)
|
|
new_password=secrets.token_urlsafe(32)
|
|
restored.change('password',password,new_password)
|
|
record=CredentialStore(path).read()
|
|
self.assertTrue(verify_password(new_password,record['password']))
|
|
self.assertFalse(verify_password(password,record['password']))
|
|
self.assertEqual(record['api_token_hash'],hashlib.sha256(token.encode()).hexdigest())
|
|
def test_write_failure_keeps_previous_credentials(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store=CredentialStore(Path(directory)/'auth.json')
|
|
password=secrets.token_urlsafe(32)
|
|
before=store.setup(password,secrets.token_urlsafe(32))
|
|
with patch('auth.os.replace',side_effect=OSError('disk unavailable')):
|
|
with self.assertRaises(OSError):store.change('token',password,secrets.token_urlsafe(32))
|
|
self.assertEqual(store.read(),before)
|
|
def test_legacy_password_preserved_api_disabled(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path=Path(directory)/'auth.json'
|
|
password=secrets.token_urlsafe(32)
|
|
old=initial_record(password,secrets.token_urlsafe(32))['password']
|
|
path.write_text(json.dumps(old))
|
|
store=CredentialStore(path)
|
|
self.assertTrue(verify_password(password,store.read()['password']))
|
|
self.assertIsNone(store.read()['api_token_hash'])
|
|
store.change('token',password,secrets.token_urlsafe(32))
|
|
self.assertTrue(verify_password(password,store.read()['password']))
|
|
def test_bad_credentials_do_not_mutate(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store=CredentialStore(Path(directory)/'auth.json')
|
|
p,t=secrets.token_urlsafe(32),secrets.token_urlsafe(32)
|
|
before=store.setup(p,t)
|
|
for kind,current,new in [('token','wrong',secrets.token_urlsafe(32)),('password',p,'short'),('token',p,p),('password',p,t),('token',p,t)]:
|
|
with self.assertRaises(ValueError):store.change(kind,current,new)
|
|
self.assertEqual(store.read(),before)
|
|
def test_simultaneous_setup_has_one_winner(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
store=CredentialStore(Path(directory)/'auth.json')
|
|
def setup(_):
|
|
try:store.setup(secrets.token_urlsafe(32),secrets.token_urlsafe(32));return True
|
|
except ValueError:return False
|
|
with ThreadPoolExecutor(max_workers=2) as pool:self.assertEqual(sum(pool.map(setup,range(2))),1)
|
|
|
|
class AccessAPITests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.directory=tempfile.TemporaryDirectory()
|
|
self.server=Server(0,state_dir=self.directory.name)
|
|
self.thread=threading.Thread(target=self.server.serve_forever);self.thread.start()
|
|
self.base=f'http://127.0.0.1:{self.server.server_port}'
|
|
self.password=secrets.token_urlsafe(32);self.token=secrets.token_urlsafe(32)
|
|
def tearDown(self):
|
|
self.server.shutdown();self.server.endpoint.close();self.server.server_close();self.thread.join();self.directory.cleanup()
|
|
def request(self,path,body=None,cookie=None,token=None,headers=None):
|
|
h={**(headers or {})}
|
|
if body is not None:h.update({'Content-Type':'application/json','X-Athena-Deck':'1'})
|
|
if cookie:h['Cookie']=cookie
|
|
if token:h['Authorization']='Bearer '+token
|
|
req=urllib.request.Request(self.base+path,data=json.dumps(body).encode() if body is not None else None,headers=h)
|
|
try:r=urllib.request.urlopen(req)
|
|
except urllib.error.HTTPError as e:r=e
|
|
with r:
|
|
raw=r.read()
|
|
return r.status,json.loads(raw),r.headers
|
|
def setup_login(self):
|
|
code,_,_=self.request('/api/v1/auth/setup',dict(password=self.password,api_token=self.token))
|
|
self.assertEqual(code,200)
|
|
return self.login(self.password)
|
|
def login(self,password):
|
|
code,_,headers=self.request('/api/v1/login',{'password':password})
|
|
self.assertEqual(code,200)
|
|
return headers['Set-Cookie'].split(';')[0]
|
|
def test_prompt_enhancer_management_requires_login_and_validates_actions(self):
|
|
self.assertEqual(self.request('/api/v1/prompt-enhancers')[0],401)
|
|
cookie=self.setup_login()
|
|
code,state,_=self.request('/api/v1/prompt-enhancers',cookie=cookie)
|
|
self.assertEqual(code,200)
|
|
self.assertEqual(set(state['models']),{'t2i','i2i'})
|
|
self.assertEqual(self.request('/api/v1/prompt-enhancers/install',{'task':'unknown'},cookie=cookie)[0],400)
|
|
self.assertEqual(self.request('/api/v1/profiles/prompt-enhancer',{'id':'missing','revision':0,'prompt_enhancer':{'t2i':None,'i2i':None,'device':'auto'}},cookie=cookie)[0],400)
|
|
def test_hub_credentials_are_session_only_and_never_returned(self):
|
|
cookie=self.setup_login();path='/api/v1/huggingface';secret='hf_'+'a'*30
|
|
for value in ({}, {'token':self.token}):
|
|
self.assertEqual(self.request(path,**value)[0],401)
|
|
self.assertEqual(self.request(path,{'token':secret},**value)[0],401)
|
|
code,body,_=self.request(path,{'token':secret},cookie=cookie)
|
|
self.assertEqual((code,body),(200,{'configured':True}))
|
|
self.assertEqual(self.request(path,cookie=cookie)[1],{'configured':True})
|
|
self.assertEqual(self.request(path,{'token':None},cookie=cookie)[1],{'configured':False})
|
|
def test_first_run_and_setup_guard(self):
|
|
self.assertEqual(self.request('/api/v1/status')[0],401)
|
|
self.assertFalse(self.request('/api/v1/auth/status')[1]['initialized'])
|
|
self.assertEqual(self.request('/api/v1/auth/setup',{'password':self.password,'api_token':self.token},headers={'Origin':'http://evil.test'})[0],403)
|
|
cookie=self.setup_login()
|
|
self.assertEqual(self.request('/api/v1/status',cookie=cookie)[0],200)
|
|
self.assertEqual(self.request('/api/v1/auth/setup',{'password':self.password,'api_token':self.token})[0],400)
|
|
def test_password_rotation_revokes_all_sessions_not_api(self):
|
|
first=self.setup_login();second=self.login(self.password)
|
|
new=secrets.token_urlsafe(32)
|
|
self.assertEqual(self.request('/api/v1/auth/password',dict(current_password=self.password,new_password=new),cookie=first)[0],200)
|
|
for cookie in (first,second):self.assertEqual(self.request('/api/v1/status',cookie=cookie)[0],401)
|
|
self.assertEqual(self.request('/api/v1/login',{'password':self.password})[0],401)
|
|
self.login(new)
|
|
self.assertEqual(self.request('/api/v1/status',token=self.token)[0],200)
|
|
def test_token_rotation_scopes_and_session_survival(self):
|
|
cookie=self.setup_login()
|
|
self.assertEqual(self.request('/api/v1/status',token=self.token)[0],200)
|
|
for path,body in [('/api/v1/network',None),('/api/v1/auth/password',dict(current_password=self.password,new_password=secrets.token_urlsafe(32))),('/api/v1/auth/token',dict(current_password=self.password,new_token=secrets.token_urlsafe(32))),('/api/v1/network/disconnect',{})]:
|
|
self.assertEqual(self.request(path,body,token=self.token)[0],401)
|
|
new=secrets.token_urlsafe(32)
|
|
self.assertEqual(self.request('/api/v1/auth/token',dict(current_password=self.password,new_token=new),cookie=cookie)[0],200)
|
|
self.assertEqual(self.request('/api/v1/status',token=self.token)[0],401)
|
|
self.assertEqual(self.request('/api/v1/status',token=new)[0],200)
|
|
self.assertEqual(self.request('/api/v1/status',cookie=cookie)[0],200)
|
|
status=self.request('/api/v1/auth/status',cookie=cookie)[1]
|
|
raw=json.dumps(status)
|
|
for secret in (new,self.token,self.password,'api_token_hash','salt'):self.assertNotIn(secret,raw)
|
|
def test_api_client_without_browser_headers(self):
|
|
self.setup_login()
|
|
req=urllib.request.Request(self.base+'/api/v1/status',method='GET',headers={'Authorization':'Bearer '+self.token})
|
|
with urllib.request.urlopen(req) as r:self.assertEqual(json.load(r)['name'],'Athena Deck')
|
|
def test_wrong_password_and_logout(self):
|
|
cookie=self.setup_login()
|
|
self.assertEqual(self.request('/api/v1/auth/token',dict(current_password='wrong',new_token=secrets.token_urlsafe(32)),cookie=cookie)[0],400)
|
|
self.assertEqual(self.request('/api/v1/status',token=self.token)[0],200)
|
|
self.assertEqual(self.request('/api/v1/logout',{},cookie=cookie)[0],200)
|
|
self.assertEqual(self.request('/api/v1/status',cookie=cookie)[0],401)
|
|
|
|
if __name__=='__main__':unittest.main()
|