63 lines
3.2 KiB
Python
63 lines
3.2 KiB
Python
import json
|
|
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
ROOT=Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0,str(ROOT/'scripts'))
|
|
from control import Client,APIError
|
|
|
|
|
|
class ClientTests(unittest.TestCase):
|
|
def test_redirect_never_forwards_bearer_to_another_origin(self):
|
|
received=[]
|
|
class Target(BaseHTTPRequestHandler):
|
|
def log_message(self,*args):pass
|
|
def do_GET(self):
|
|
received.append(self.headers.get('Authorization'))
|
|
self.send_response(200);self.end_headers();self.wfile.write(b'{}')
|
|
target=ThreadingHTTPServer(('127.0.0.1',0),Target)
|
|
class Redirect(BaseHTTPRequestHandler):
|
|
def log_message(self,*args):pass
|
|
def do_GET(self):
|
|
self.send_response(302)
|
|
self.send_header('Location','http://127.0.0.1:'+str(target.server_port)+'/capture')
|
|
self.end_headers()
|
|
redirect=ThreadingHTTPServer(('127.0.0.1',0),Redirect)
|
|
threads=[threading.Thread(target=s.serve_forever,daemon=True) for s in (target,redirect)]
|
|
for t in threads:t.start()
|
|
try:
|
|
with tempfile.TemporaryDirectory() as root:
|
|
token=Path(root)/'token';token.write_text('test-only-token');token.chmod(0o600)
|
|
client=Client({'url':'http://127.0.0.1:'+str(redirect.server_port),'token_file':str(token)})
|
|
with self.assertRaises(APIError) as result:client.call('/api/state')
|
|
self.assertEqual(result.exception.status,302)
|
|
self.assertEqual(received,[])
|
|
finally:
|
|
for s in (redirect,target):s.shutdown();s.server_close()
|
|
for t in threads:t.join()
|
|
|
|
def test_ssh_json_is_explicit_utf8_and_preserves_chinese(self):
|
|
from types import SimpleNamespace
|
|
client=Client({'transport':'ssh','identity_file':'example-key','ssh_host':'observer@entry.example.test',
|
|
'remote_client':'/home/observer/agent-workbench/scripts/control.py','remote_config':'/home/observer/private/config.json'})
|
|
with patch('control.subprocess.run',return_value=SimpleNamespace(returncode=0,stdout=json.dumps({'text':'中文回复'}))) as run:
|
|
self.assertEqual(client.call('/api/state')['text'],'中文回复')
|
|
self.assertEqual(run.call_args.kwargs['encoding'],'utf-8')
|
|
self.assertIn('python3 -X utf8',run.call_args.args[0][-1])
|
|
|
|
def test_ssh_preserves_conflict_status_for_reconciliation(self):
|
|
from types import SimpleNamespace
|
|
client=Client({'transport':'ssh','identity_file':'example-key','ssh_host':'observer@entry.example.test',
|
|
'remote_client':'/home/observer/agent-workbench/scripts/control.py','remote_config':'/home/observer/private/config.json'})
|
|
with patch('control.subprocess.run',return_value=SimpleNamespace(returncode=2,stdout='{"api_error":409}')):
|
|
with self.assertRaises(APIError) as result:client.call('/api/state')
|
|
self.assertEqual(result.exception.status,409)
|
|
|
|
|
|
if __name__=='__main__':unittest.main()
|