business-assistant/scripts/tests/test_qcc_proxy.py

import contextlib
import hashlib
import io
import json
import os
from pathlib import Path
import tempfile
import threading
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import patch
from urllib.parse import parse_qs, urlsplit

from scripts import test_qcc_proxy as probe


class QccProbeTests(unittest.TestCase):
    def test_signature_refreshes_for_each_request(self):
        for timestamp in (1700000000, 1700000861):
            with patch.object(probe.time, 'time', return_value=timestamp):
                result = probe.probe('search', self.values, 'synthetic', 2)
            self.assertTrue(result['ok'])
            headers = self.calls[-1][2]
            self.assertEqual(headers['Timespan'], str(timestamp))
            self.assertEqual(headers['Token'], hashlib.md5(
                ('fake-key' + str(timestamp) + 'fake-secret').encode()).hexdigest().upper())
        self.assertNotEqual(self.calls[0][2]['Token'], self.calls[1][2]['Token'])

    def test_clock_diagnostics_and_invalid_dates(self):
        date = 'Tue, 14 Nov 2023 22:27:41 GMT'
        result = probe.clock_diagnostics('1700000000', [date], .35)
        self.assertEqual(result['response_date_minus_request_seconds'], 861)
        self.assertTrue(result['clock_skew_suspected'])
        self.assertFalse(probe.clock_diagnostics('1700000861', [date], .35)['clock_skew_suspected'])
        self.assertTrue(probe.clock_diagnostics('1700001722', [date], .35)['clock_skew_suspected'])
        for dates in ([], ['bad'], [date, date]):
            self.assertNotIn('clock_skew_suspected', probe.clock_diagnostics('1700000000', dates, .35))

    def test_env_file_only_ignores_stale_credentials(self):
        with tempfile.TemporaryDirectory() as folder:
            path = Path(folder) / '.env'
            path.write_text('QCC_KEY=file-key\nQCC_SECRET_KEY=file-secret\n', encoding='utf-8')
            with patch.dict(os.environ, {'QCC_KEY': 'old-key', 'QCC_SECRET_KEY': 'old-secret'}):
                self.assertEqual(probe.load_settings(path)['QCC_KEY'], 'old-key')
                values = probe.load_settings(path, env_file_only=True)
                self.assertEqual(values['QCC_KEY'], 'file-key')
                self.assertEqual(values['QCC_SECRET_KEY'], 'file-secret')

    def setUp(self):
        self.calls = []
        self.status = 200
        self.broken = False
        self.raw_body = None
        self.payloads = {
            'fuzzySearch': {'Status': '200', 'Result': [{'Name': '合成企业', 'KeyNo': 'synthetic'}]},
            'basicDetails': {'Status': '200', 'Result': {'Name': '合成企业', 'KeyNo': 'synthetic'}},
            'honor': {'Status': '200', 'Result': {'VerifyResult': 0, 'Data': None}},
        }
        test = self
        class Handler(BaseHTTPRequestHandler):
            protocol_version = 'HTTP/1.1'
            def log_message(self, *_):
                pass
            def do_GET(self):
                query = parse_qs(urlsplit(self.path).query)
                test.calls.append((self.path, query, dict(self.headers)))
                self.send_response(test.status)
                if test.status == 302:
                    self.send_header('Location', '/other')
                route = query.get('servertype', ['fuzzySearch'])[0]
                body = test.raw_body if test.raw_body is not None else json.dumps(test.payloads[route]).encode()
                self.send_header('Content-Type', 'text/html' if test.raw_body else 'application/json')
                self.send_header('X-Request-ID', 'synthetic-trace')
                if test.broken:
                    self.send_header('Transfer-Encoding', 'chunked')
                    self.end_headers()
                    self.wfile.write(b'broken\r\n')
                    self.close_connection = True
                else:
                    self.send_header('Content-Length', str(len(body)))
                    self.end_headers()
                    self.wfile.write(body)
        self.server = ThreadingHTTPServer(('127.0.0.1', 0), Handler)
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self.values = {'QCC_KEY': 'fake-key', 'QCC_SECRET_KEY': 'fake-secret'}
        for name, route in zip(probe.URL_KEYS.values(), self.payloads):
            self.values[name] = f'http://127.0.0.1:{self.server.server_port}/proxy/?servertype={route}&proxyToken=fake%2B%26token'

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join()

    def test_three_get_requests_signatures_query_and_no_sensitive_output(self):
        with patch.dict(os.environ, {'HTTP_PROXY': 'http://invalid.example:1'}):
            for api in probe.URL_KEYS:
                result = probe.probe(api, self.values, '合成企业', 2)
                self.assertTrue(result['ok'], result)
                self.assertNotIn('fake', json.dumps(result))
                self.assertNotIn('合成企业', json.dumps(result, ensure_ascii=False))
        self.assertEqual(len(self.calls), 3)
        for path, query, headers in self.calls:
            self.assertIn('proxyToken=fake%2B%26token&key=', path)
            self.assertEqual(headers['Connection'], 'keep-alive')
            self.assertEqual(headers['Token'], hashlib.md5(('fake-key' + headers['Timespan'] + 'fake-secret').encode()).hexdigest().upper())
            self.assertEqual(query['key'], ['fake-key'])
        self.assertEqual(self.calls[0][1]['searchKey'], ['合成企业'])
        self.assertEqual(self.calls[1][1]['keyword'], ['合成企业'])
        self.assertEqual(self.calls[2][1]['pageSize'], ['20'])

    def test_business_failure_and_empty_distinction(self):
        self.payloads['fuzzySearch'] = {'Status': '201'}
        self.assertEqual(probe.probe('search', self.values, 'synthetic', 2)['code'], 'no_matches')
        self.payloads['basicDetails'] = {'Status': '201'}
        result = probe.probe('basic', self.values, 'synthetic', 2)
        self.assertFalse(result['ok'])
        self.assertEqual(result['code'], 'provider_failure')
        self.assertFalse(probe.validate_payload('honor', {'Status': '200', 'Result': {'VerifyResult': 0, 'Data': [{}]}})['ok'])

    def test_invalid_url_and_case_are_not_sent(self):
        original = self.values['QCC_SEARCH_URL']
        for value in (original.replace('servertype', 'Servertype'), original.replace('proxyToken=', 'proxyToken-'), original.replace('127.0', '127. 0')):
            self.values['QCC_SEARCH_URL'] = value
            self.assertFalse(probe.probe('search', self.values, 'synthetic', 2)['ok'])
        self.assertEqual(self.calls, [])

    def test_redirect_http_error_and_broken_response(self):
        for status in (302, 500):
            self.status = status
            self.assertEqual(probe.probe('search', self.values, 'synthetic', 2)['code'], 'http_error')
        self.assertEqual(len(self.calls), 2)
        self.status, self.broken = 200, True
        self.assertEqual(probe.probe('search', self.values, 'synthetic', 2)['code'], 'response_incomplete')

    def test_env_and_cli(self):
        with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ, {}, clear=True):
            path = Path(directory) / '.env'
            path.write_text('export QCC_KEY="fake-key" # note\nQCC_SECRET_KEY=fake-secret\n', encoding='utf-8-sig')
            self.assertEqual(probe.load_settings(path)['QCC_KEY'], 'fake-key')
        with patch.object(probe, 'load_settings', return_value=self.values), patch('sys.argv', ['test_qcc_proxy.py', '--company', '合成企业']), contextlib.redirect_stdout(io.StringIO()) as output:
            self.assertEqual(probe.main(), 0)
        self.assertEqual(set(json.loads(output.getvalue())['checks']), {'search', 'basic', 'honor'})

    def test_timeout_closes_and_does_not_leak(self):
        with patch.object(probe, 'HTTPConnection') as conn:
            conn.return_value.request.side_effect = TimeoutError('fake-secret')
            result = probe.probe('search', self.values, 'synthetic', 2)
            self.assertEqual(result['code'], 'timeout')
            self.assertNotIn('fake', json.dumps(result))
            conn.return_value.close.assert_called_once()

    def test_400_html_body_is_visible_and_sensitive_echoes_are_redacted(self):
        self.status = 400
        timestamp = '1700000000'
        signature = hashlib.md5(('fake-key' + timestamp + 'fake-secret').encode()).hexdigest().upper()
        self.raw_body = ('<html>Bad Request: invalid upstream request fake-key fake-secret '
                         'fake+&token fake%2b%26token 合成企业 ' + signature + '</html>').encode()
        with patch.object(probe.time, 'time', return_value=int(timestamp)):
            result = probe.probe('search', self.values, '合成企业', 2)
        self.assertIn('invalid upstream request', result['error_body_preview'])
        self.assertEqual(result['response_headers']['X-Request-ID'], ['synthetic-trace'])
        self.assertIn('elapsed_seconds', result)
        for value in ('fake-key', 'fake-secret', 'fake+&token', 'fake%2b%26token', '合成企业', signature):
            self.assertNotIn(value, json.dumps(result, ensure_ascii=False))

    def test_provider_message_visible_result_omitted_and_preview_bounded(self):
        self.payloads['fuzzySearch'] = {'Status': '115', 'Message': '身份验证错误 fake-key',
                                      'Result': {'Name': 'private-company'}, 'OrderNumber': 'test-order'}
        result = probe.probe('search', self.values, '合成企业', 2)
        self.assertIn('身份验证错误', result['error_body_preview'])
        self.assertNotIn('private-company', result['error_body_preview'])
        self.assertNotIn('fake-key', result['error_body_preview'])
        self.raw_body = b'x' * 3000
        self.status = 400
        result = probe.probe('search', self.values, 'synthetic', 2)
        self.assertEqual(len(result['error_body_preview']), 2000)
        self.assertTrue(result['error_body_truncated'])

    def test_success_never_outputs_response_body(self):
        result = probe.probe('search', self.values, '合成企业', 2)
        self.assertTrue(result['ok'])
        self.assertNotIn('error_body_preview', result)
        self.assertNotIn('合成企业', json.dumps(result, ensure_ascii=False))

    def test_direct_mode_matches_proxy_business_parameters_and_headers(self):
        with patch.object(probe.time, 'time', return_value=1700000000):
            proxy = probe.probe('search', self.values, '合成企业', 2)
            self.values['QCC_SEARCH_URL'] = f'http://127.0.0.1:{self.server.server_port}/FuzzySearch/GetList'
            direct = probe.probe('search', self.values, '合成企业', 2, mode='direct')
        self.assertTrue(proxy['ok'])
        self.assertTrue(direct['ok'])
        self.assertTrue(self.calls[1][0].startswith('/FuzzySearch/GetList?key='))
        for key in ('key', 'searchKey', 'pageIndex'):
            self.assertEqual(self.calls[0][1][key], self.calls[1][1][key])
        for key in ('Token', 'Timespan', 'Connection', 'Accept'):
            self.assertEqual(self.calls[0][2][key], self.calls[1][2][key])
        self.assertFalse(probe.probe('search', self.values, '合成企业', 2)['ok'])

    def test_direct_mode_rejects_proxy_url_before_request(self):
        self.assertEqual(probe.probe('search', self.values, 'synthetic', 2, mode='direct')['code'],
                         'direct_url_must_not_have_query')
        self.assertEqual(self.calls, [])