business-assistant/scripts/tests/test_proxy_auth.py

"""Proxy recovery against loopback HTTP only; all credentials are synthetic."""

from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
import hashlib
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
import os
from pathlib import Path
import tempfile
import threading
from types import SimpleNamespace
import unittest
from unittest.mock import patch
from urllib.parse import parse_qs, urlsplit

from src.proxy_auth import ProxyAuthError, ProxyTokens, proxy_identity, replace_token, token_rejected
from src.step1_knowledge_prepare import summarize
from src.step3_question_answer import qcc
from src.step3_question_answer.streaming import stream_transport
from scripts import deploy_check


class TokenStateTests(unittest.TestCase):
    url = 'http://example.invalid:10011/proxy/?servertype=deepseek&proxyToken=synthetic-old'

    def client(self, login=None):
        client = SimpleNamespace(config=SimpleNamespace(username='synthetic-user', password='synthetic-password'),
                                 login_url='http://example.invalid/login', token='synthetic-short')
        client._login = login or (lambda: None)
        return client

    def test_only_definite_outer_rejections(self):
        for value in ({'code': 212}, {'code': '212'}):
            self.assertTrue(token_rejected(value))
        for value in ({'code': 201}, {'Status': '115'}, {'code': 212, 'choices': []},
                      {'code': 212, 'Status': '200'}, {'message': 'token invalid'}, [], None):
            self.assertFalse(token_rejected(value))

    def test_diagnostics_expose_only_fixed_auth_codes(self):
        from src.step4_api.logging import safe_code
        for code in ('proxy_token_invalid', 'proxy_token_refresh_failed'):
            self.assertEqual(deploy_check.safe_code(RuntimeError(code)), code)
            self.assertEqual(safe_code(code), code)
            self.assertEqual(safe_code('model_' + code), 'model_' + code)
        self.assertEqual(deploy_check.safe_code(RuntimeError('synthetic-token-in-url')), 'request_failed')

    def test_replace_preserves_signed_query_and_encodes_token(self):
        url = self.url + '&raw=%2f+%20&x=1&x=2&blank=&flag'
        updated = replace_token(url, 'synthetic+/&=?')
        self.assertEqual(updated, url.replace('synthetic-old', 'synthetic%2B%2F%26%3D%3F'))
        self.assertEqual(proxy_identity(updated)[1], 'synthetic+/&=?')
        self.assertIsNone(proxy_identity(url + '&proxyToken=duplicate'))

    def test_singleflight_and_origin_isolation(self):
        calls = []
        auth = ProxyTokens(lambda: self.client(lambda: calls.append('login')))
        barrier = threading.Barrier(8)
        def refresh(_):
            barrier.wait()
            auth.refresh(self.url)
            return auth.url(self.url)
        with ThreadPoolExecutor(max_workers=8) as pool:
            urls = list(pool.map(refresh, range(8)))
        self.assertEqual(calls, ['login'])
        self.assertEqual(set(urls), {self.url.replace('synthetic-old', 'synthetic-short')})
        for other in (self.url.replace('http:', 'https:'), self.url.replace(':10011', ':10012'),
                      self.url.replace('example.invalid', 'other.invalid')):
            self.assertEqual(auth.url(other), other)
        self.assertEqual(repr(auth), '<ProxyTokens>')

    def test_failure_cooldown_and_later_recovery(self):
        now, calls = [10], []
        def factory():
            calls.append(1)
            if len(calls) == 1:
                raise RuntimeError('synthetic-password')
            return self.client()
        auth = ProxyTokens(factory, lambda: now[0])
        for _ in range(2):
            with self.assertRaisesRegex(ProxyAuthError, '^proxy_token_refresh_failed$'):
                auth.refresh(self.url)
        self.assertEqual(calls, [1])
        now[0] += 31
        auth.refresh(self.url)
        self.assertIn('synthetic-short', auth.url(self.url))

    def test_missing_login_same_or_malformed_token_fails_closed(self):
        for changes in ({'login_url': None}, {'token': 'synthetic-old'}, {'token': 'bad\nvalue'},
                        {'token': ''}, {'token': None}, {'config': SimpleNamespace(username='', password='')}):
            with self.subTest(changes=changes):
                client = self.client()
                for name, value in changes.items():
                    setattr(client, name, value)
                auth = ProxyTokens(lambda: client)
                with self.assertRaisesRegex(ProxyAuthError, '^proxy_token_refresh_failed$'):
                    auth.refresh(self.url)
                self.assertEqual(auth.url(self.url), self.url)

    def test_fresh_rejection_cools_down_without_relogin(self):
        calls, now = [], [10]
        auth = ProxyTokens(lambda: self.client(lambda: calls.append(1)), lambda: now[0])
        auth.refresh(self.url)
        new = auth.url(self.url)
        auth.reject(new)
        with self.assertRaisesRegex(ProxyAuthError, '^proxy_token_refresh_failed$'):
            auth.refresh(new)
        self.assertEqual(calls, [1])


class RecoveryHttpTests(unittest.TestCase):
    def setUp(self):
        self.calls, self.logins = [], []
        self.valid = 'synthetic-short'
        self.mode, self.login_mode, self.rejection_status = 'normal', 'normal', 200
        self.first_received = threading.Event()
        test = self

        class Handler(BaseHTTPRequestHandler):
            protocol_version = 'HTTP/1.1'

            def log_message(self, *_):
                pass

            def handle(self):
                try:
                    super().handle()
                except (ConnectionAbortedError, ConnectionResetError, BrokenPipeError):
                    pass

            def envelope(self, value, status=200):
                body = json.dumps(value).encode()
                self.send_response(status)
                self.send_header('Content-Length', str(len(body)))
                self.end_headers()
                self.wfile.write(body)

            def do_POST(self):
                body = self.rfile.read(int(self.headers['Content-Length']))
                if self.path == '/oauth/api/user/login':
                    test.logins.append((self.path, parse_qs(body.decode()), dict(self.headers)))
                    if test.login_mode == 'redirect':
                        self.send_response(302)
                        self.send_header('Location', '/do-not-follow')
                        self.send_header('Content-Length', '0')
                        self.end_headers()
                    else:
                        self.envelope({'code': 201 if test.login_mode == 'denied' else 200,
                                       'message': test.valid, 'content': None})
                    return
                self.proxy(json.loads(body))

            def do_GET(self):
                self.proxy(None)

            def proxy(self, payload):
                query = parse_qs(urlsplit(self.path).query)
                test.calls.append((self.path, dict(self.headers), payload))
                if test.mode == 'permission':
                    self.envelope({'code': 201})
                elif test.mode == 'upstream_auth':
                    self.envelope({'Status': '115'})
                elif test.mode == 'http401':
                    self.envelope({'message': 'unauthorized'}, 401)
                elif test.mode == 'http500':
                    self.envelope({'code': 212}, 500)
                elif test.mode == 'truncated_rejection':
                    body = b'{"code":212}'
                    self.send_response(401)
                    self.send_header('Content-Length', str(len(body) + 10))
                    self.end_headers()
                    self.wfile.write(body)
                    self.close_connection = True
                elif test.mode == 'invalid_json':
                    self.send_response(200)
                    self.send_header('Content-Length', '1')
                    self.end_headers()
                    self.wfile.write(b'{')
                elif test.mode == 'always_invalid' or query['proxyToken'][0] != test.valid:
                    if test.mode == 'whitespace' and payload is not None:
                        self.send_response(200)
                        self.send_header('Transfer-Encoding', 'chunked')
                        self.end_headers()
                        self.chunk(b' \n ')
                        self.chunk(b'{"code":212}')
                        self.wfile.write(b'0\r\n\r\n')
                    else:
                        self.envelope({'code': '212', 'message': 'synthetic rejection'}, test.rejection_status)
                elif payload is None:
                    service = query['servertype'][0]
                    company = {'KeyNo': 'synthetic-id', 'Name': '合成企业'}
                    result = ([company] if service == 'fuzzySearch' else company
                              if service == 'basicDetails' else {'VerifyResult': 0, 'Data': []})
                    self.envelope({'Status': '200', 'Result': result})
                elif payload.get('stream'):
                    self.send_response(200)
                    self.send_header('Content-Type', 'text/event-stream')
                    self.send_header('Transfer-Encoding', 'chunked')
                    self.end_headers()
                    self.chunk(b'data: {"choices":[{"delta":{"content":"first"}}]}\n\n')
                    if test.mode == 'partial':
                        self.close_connection = True
                        return
                    if test.mode == 'incremental' and not test.first_received.wait(2):
                        self.close_connection = True
                        return
                    self.chunk(b'data: {"choices":[{"delta":{"content":"OK"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n')
                    self.wfile.write(b'0\r\n\r\n')
                else:
                    self.envelope({'choices': [{'message': {'content': 'OK'}}]})

            def chunk(self, body):
                self.wfile.write(f'{len(body):x}\r\n'.encode() + body + b'\r\n')
                self.wfile.flush()

        self.server = ThreadingHTTPServer(('127.0.0.1', 0), Handler)
        self.thread = threading.Thread(target=lambda: self.server.serve_forever(poll_interval=.02), daemon=True)
        self.thread.start()
        self.temp = tempfile.TemporaryDirectory()
        self.path = Path(self.temp.name) / '.env'
        base = f'http://127.0.0.1:{self.server.server_port}'
        query = '?servertype={}&proxyToken=synthetic-old&raw=%2f+%20&x=1&x=2'
        lines = ['QCC_KEY=synthetic-key', 'QCC_SECRET_KEY=synthetic-secret', 'QCC_TIMEOUT_SECONDS=2',
                 'DEEPSEEK_API_KEY=synthetic-api-key', 'DEEPSEEK_MODEL=synthetic-model', 'SUMMARY_TIMEOUT_SECONDS=2',
                 'DMS_BASE_URL=http://127.0.0.1', f'AUTH_PORT={self.server.server_port}',
                 f'DMS_PORT={self.server.server_port}', 'DMS_USERNAME=synthetic-user', 'DMS_PASSWORD=synthetic-password',
                 'DEEPSEEK_BASE_URL=' + base + '/proxy/chat/completions' + query.format('deepseek')]
        for name, service in (('SEARCH', 'fuzzySearch'), ('BASIC', 'basicDetails'), ('HONOR', 'honor')):
            lines.append('QCC_' + name + '_URL=' + base + '/proxy/' + query.format(service))
        self.path.write_text('\n'.join(lines), encoding='utf-8')
        self.original = self.path.read_bytes()
        self.env = patch.dict(os.environ, {'NO_PROXY': '*'}, clear=True)
        self.env.start()
        self.qconfig, self.mconfig = qcc.load_config(self.path), summarize.load_config(self.path)

    def tearDown(self):
        self.assertEqual(self.path.read_bytes(), self.original, '.env must stay byte-identical')
        self.env.stop()
        self.server.shutdown()
        self.server.server_close()
        self.thread.join()
        self.temp.cleanup()

    def model(self):
        return summarize.transport({'model': self.mconfig.model, 'messages': [{'role': 'user', 'content': 'synthetic'}]}, self.mconfig)

    def test_qcc_login_shared_with_other_interfaces_and_model(self):
        self.assertIs(self.qconfig.proxy_auth, self.mconfig.proxy_auth)
        client = qcc.QccClient(self.qconfig)
        company = client.search('合成企业')['candidates'][0]
        client.basic(company)
        self.assertTrue(client.honors(company)['complete'])
        self.assertEqual(self.model()['choices'][0]['message']['content'], 'OK')
        self.assertEqual(len(self.logins), 1)
        self.assertEqual(len(self.calls), 5)
        first, replay = self.calls[:2]
        self.assertEqual(first[0].replace('synthetic-old', 'synthetic-short'), replay[0])
        for name in ('Token', 'Timespan'):
            self.assertEqual(first[1][name], replay[1][name])
        expected = hashlib.md5(('synthetic-key' + replay[1]['Timespan'] + 'synthetic-secret').encode()).hexdigest().upper()
        self.assertEqual(replay[1]['Token'], expected)
        self.assertEqual(self.logins[0][1], {'userName': ['synthetic-user'], 'password': ['synthetic-password'],
                                         'clientId': ['0'], 'serviceId': ['0']})
        self.assertNotIn('Token', self.logins[0][2])
        self.assertEqual(self.calls[-1][1]['Authorization'], 'Bearer synthetic-api-key')
        for path, headers, body in self.calls:
            self.assertIn('raw=%2f+%20&x=1&x=2', path)
            self.assertNotIn('synthetic-password', str(headers) + str(body) + path)

    def test_valid_long_token_needs_no_dms(self):
        self.valid = 'synthetic-old'
        with patch('src.step5_session_analysis.dms.load_dms_config', side_effect=AssertionError('no login expected')):
            self.model()
            qcc.QccClient(self.qconfig).search('合成企业')
        self.assertEqual(self.logins, [])

    def test_healthy_json_is_strictly_decoded_only_once(self):
        self.valid = 'synthetic-old'
        with patch.object(summarize, 'strict_json', wraps=summarize.strict_json) as decode:
            self.model()
        self.assertEqual(decode.call_count, 1)

    def test_model_json_replays_exact_payload_once(self):
        self.model()
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(self.calls[0][2], self.calls[1][2])
        self.assertEqual(len(self.logins), 1)
        self.valid = 'synthetic-next'
        self.model()
        self.assertEqual(len(self.calls), 4)
        self.assertEqual(len(self.logins), 2)
        self.assertIn('proxyToken=synthetic-next', self.calls[-1][0])

    def test_stream_recovery_keeps_first_delta_incremental(self):
        self.mode = 'incremental'
        parts = []
        def consume(text):
            parts.append(text)
            self.first_received.set()
        result = stream_transport({'model': 'synthetic-model'}, self.mconfig, consume)
        self.assertEqual(parts, ['first', 'OK'])
        self.assertEqual(result['choices'][0]['message']['content'], 'firstOK')
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(len(self.logins), 1)
        self.assertEqual(self.calls[0][2], self.calls[1][2])

    def test_whitespace_fragmented_rejection(self):
        self.mode = 'whitespace'
        self.assertEqual(self.model()['choices'][0]['message']['content'], 'OK')
        self.assertEqual(len(self.logins), 1)

    def test_http_auth_requires_explicit_212(self):
        for status in (401, 403):
            self.rejection_status = status
            self.valid = 'synthetic-' + str(status)
            self.model()
            self.assertIn('proxyToken=' + self.valid, self.calls[-1][0])
        self.assertEqual(len(self.logins), 2)
        # QCC uses the same outer rejection contract for HTTP auth failures.
        self.valid = 'synthetic-qcc'
        qcc.QccClient(self.qconfig).search('合成企业')
        self.assertEqual(len(self.logins), 3)

    def test_unrelated_failures_do_not_log_in(self):
        for mode, model_code, qcc_code in (
            ('permission', None, 'provider_failure'), ('upstream_auth', None, 'qcc_auth_failed'),
            ('http401', 'http_401', 'http_401'), ('http500', 'http_500', 'http_500'),
            ('invalid_json', 'invalid_response_json', 'invalid_json')):
            self.mode = mode
            if model_code is None:
                self.model()  # Original provider schema validation belongs to the caller.
            else:
                with self.assertRaisesRegex(summarize.RequestFailure, '^' + model_code + '$'):
                    self.model()
            with self.assertRaisesRegex(qcc.QccError, '^' + qcc_code + '$'):
                qcc.QccClient(self.qconfig).search('合成企业')
        self.assertEqual(self.logins, [])
        self.assertEqual(len(self.calls), 10)

    def test_second_rejection_stops_and_cools_down(self):
        self.mode = 'always_invalid'
        with self.assertRaisesRegex(summarize.RequestFailure, '^proxy_token_invalid$') as failure:
            self.model()
        self.assertFalse(failure.exception.retryable)
        self.assertEqual(len(self.calls), 2)
        with self.assertRaisesRegex(qcc.QccError, '^proxy_token_refresh_failed$'):
            qcc.QccClient(self.qconfig).search('合成企业')
        self.assertEqual(len(self.logins), 1)
        self.assertEqual(len(self.calls), 3)

    def test_truncated_auth_envelope_is_not_a_definite_rejection(self):
        self.mode = 'truncated_rejection'
        with self.assertRaisesRegex(summarize.RequestFailure, '^response_incomplete$'):
            self.model()
        with self.assertRaisesRegex(qcc.QccError, '^response_incomplete$'):
            qcc.QccClient(self.qconfig).search('合成企业')
        self.assertEqual(self.logins, [])
        self.assertEqual(len(self.calls), 2)

    def test_qcc_second_rejection_is_bounded(self):
        self.mode = 'always_invalid'
        with self.assertRaisesRegex(qcc.QccError, '^proxy_token_invalid$'):
            qcc.QccClient(self.qconfig).search('合成企业')
        self.assertEqual(len(self.calls), 2)
        self.assertEqual(len(self.logins), 1)

    def test_login_failure_and_redirect_are_safe_and_not_replayed(self):
        for mode in ('denied', 'redirect'):
            self.login_mode = mode
            # Isolate the two scenarios' shared caches using an explicit test env.
            config = summarize.load_config(self.path, {})
            with self.assertRaisesRegex(summarize.RequestFailure, '^proxy_token_refresh_failed$') as failure:
                summarize.transport({}, config)
            self.assertFalse(failure.exception.retryable)
        self.assertEqual(len(self.logins), 2)
        self.assertEqual(len(self.calls), 2)

    def test_partial_stream_is_never_replayed(self):
        self.valid, self.mode = 'synthetic-old', 'partial'
        parts = []
        with self.assertRaisesRegex(summarize.RequestFailure, '^response_incomplete$'):
            stream_transport({}, self.mconfig, parts.append)
        self.assertEqual(parts, ['first'])
        self.assertEqual(len(self.calls), 1)
        self.assertEqual(self.logins, [])

    def test_cancellation_after_login_prevents_replay(self):
        def check():
            if self.logins:
                raise summarize.RequestFailure('stream_cancelled', False)
        with self.assertRaisesRegex(summarize.RequestFailure, '^stream_cancelled$'):
            stream_transport({}, self.mconfig, lambda _: None, check)
        self.assertEqual(len(self.calls), 1)
        self.assertEqual(len(self.logins), 1)

    def test_direct_qcc_envelope_does_not_use_proxy_auth(self):
        config = replace(self.qconfig, search_url=qcc.DEFAULT_SEARCH_URL)
        with self.assertRaisesRegex(qcc.QccError, '^provider_failure$'):
            qcc.QccClient(config, send=lambda *_: {'code': 212}).search('合成企业')
        self.assertEqual(self.logins, [])


if __name__ == '__main__':
    unittest.main()