key-cape/scripts/test_openbao_negative_probe.py

import base64
import json
import os
import sys
import unittest

sys.path.insert(0, os.path.dirname(__file__))
import openbao_negative_probe as probe  # noqa: E402

TOKEN_URL = 'https://kc.example/token'
BAO = 'https://bao.example'
SECRET = 'not-a-real-secret'
SECRETS = {cid: SECRET for _, cid, _ in probe.CLAIM_PROBES + [probe.EXPIRY_PROBE, probe.CODING_AGENT]}


def fake_jwt(client_id, exp=1000):
    claims = dict({'aud': 'secrets-engine-openbao', 'sub': 'service:secrets-engine', 'tenant': 'tenant:platform', 'roles': ['secrets-engine'], 'scope': 'openbao:login', 'principal_type': 'service'}, exp=exp, iat=exp-60, cid=client_id, iss='https://kc.coulomb.social')
    field = client_id.removeprefix('openbao-probe-')
    if field in {'aud', 'sub', 'tenant', 'roles', 'scope'}:
        claims[field] = ['probe'] if field == 'roles' else 'probe'
    if client_id == 'human':
        claims['principal_type'] = 'human'
    if client_id == 'codex-railiance-platform':
        claims['roles'] = ['coding-agent']
    payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip('=')
    return f'h.{payload}.s'


class FakeBao:
    """KeyCape and OpenBao in one: refuses every token unless told to accept it."""

    def __init__(self, accept=(), login_status=None):
        self.accept = set(accept)
        self.login_status = login_status
        self.revoked = 0
        self.expiry_presentations = 0

    def __call__(self, url, form=None, body=None, headers=None, timeout=15):
        if url == TOKEN_URL:
            cid = base64.b64decode(headers['Authorization'].split()[1]).decode().split(':')[0]
            return 200, json.dumps({'access_token': fake_jwt(cid)}).encode()
        if url.endswith('/revoke-self'):
            self.revoked += 1
            return 204, b''
        cid = json.loads(base64.urlsafe_b64decode(body['jwt'].split('.')[1] + '==')).get('cid', 'human')
        if self.login_status:
            return self.login_status, b''
        if cid == 'openbao-probe-expiry':
            self.expiry_presentations += 1
        if cid in self.accept or (cid == 'openbao-probe-expiry' and self.expiry_presentations == 1):
            return 200, json.dumps({'auth': {'client_token': 'bao-token'}}).encode()
        reason = {'openbao-probe-expiry': 'token is expired',
                  'human': 'claim "principal_type" does not match'}.get(cid, 'claim "sub" does not match any bound value')
        return 400, json.dumps({'errors': [reason]}).encode()


def run(fake, **kw):
    kw.setdefault('human_token', fake_jwt('human'))
    return probe.run(fake, SECRETS, TOKEN_URL, BAO, 'keycape-services', 'secrets-engine-openbao',
                     sleep=lambda s: None, now=lambda: 0, **kw)


class ProbeTest(unittest.TestCase):
    def test_all_refused_with_classes(self):
        receipt = run(FakeBao())
        self.assertTrue(receipt['all_refused'])
        classes = {c['case']: c['class'] for c in receipt['cases']}
        self.assertEqual(classes['expired'], 'expired')
        self.assertEqual(classes['principal_type'], 'not_service_principal')
        self.assertEqual(classes['sub'], 'claim_mismatch')
        self.assertEqual(len(receipt['cases']), 8)

    def test_accepted_token_fails_and_is_revoked(self):
        fake = FakeBao(accept={'openbao-probe-tenant'})
        receipt = run(fake)
        self.assertFalse(receipt['all_refused'])
        tenant = next(c for c in receipt['cases'] if c['case'] == 'tenant')
        self.assertEqual((tenant['refused'], tenant['class'], tenant['revoked']), (False, 'accepted', True))
        self.assertEqual(fake.revoked, 2)

    def test_transport_failure_is_never_a_refusal(self):
        receipt = run(FakeBao(login_status=503))
        self.assertFalse(receipt['all_refused'])
        self.assertTrue(all(c['class'] == 'inconclusive' for c in receipt['cases']))

    def test_missing_inputs_are_inconclusive(self):
        receipt = run(FakeBao(), human_token=None)
        pt = next(c for c in receipt['cases'] if c['case'] == 'principal_type')
        self.assertEqual(pt['class'], 'inconclusive')
        self.assertFalse(receipt['all_refused'])

    def test_receipt_carries_no_values(self):
        text = json.dumps(run(FakeBao(accept={'openbao-probe-aud'})))
        for value in (SECRET, 'bao-token', 'does not match', 'h.'):
            self.assertNotIn(value, text)

    def test_expiry_waits_past_leeway(self):
        waits = []
        probe.run(FakeBao(), SECRETS, TOKEN_URL, BAO, 'm', 'r', human_token=fake_jwt('human'),
                  leeway=150, sleep=waits.append, now=lambda: 900)
        self.assertEqual(waits, [1000 + 150 + 10 - 900])

    def test_wrong_expiry_reason_is_not_evidence(self):
        fake = FakeBao()
        def wrong(url, **kw):
            result = fake(url, **kw)
            if url.endswith('/login') and fake.expiry_presentations > 1:
                return 400, json.dumps({'errors': ['claim sub does not match']}).encode()
            return result
        receipt = run(wrong)
        self.assertFalse(receipt['all_refused'])
        expired = receipt['cases'][-1]
        self.assertTrue(expired['refused'])
        self.assertFalse(expired['verified'])

    def test_generic_denial_is_not_evidence(self):
        fake = FakeBao()
        def deny(url, **kw):
            result = fake(url, **kw)
            if result[0] == 400:
                return 403, b'{"errors":["permission denied"]}'
            return result
        self.assertFalse(run(deny)['all_refused'])

    def test_expiry_requires_valid_control(self):
        fake = FakeBao()
        def reject_control(url, **kw):
            result = fake(url, **kw)
            if url.endswith('/login') and fake.expiry_presentations == 1:
                return 400, b'{"errors":["claim sub mismatch"]}'
            return result
        receipt = run(reject_control)
        self.assertFalse(receipt['all_refused'])
        self.assertFalse(receipt['expiry_positive_control_verified'])

    def test_unconfirmed_control_revoke_fails(self):
        fake = FakeBao()
        def bad_revoke(url, **kw):
            if url.endswith('/revoke-self'):
                return 500, b''
            return fake(url, **kw)
        self.assertFalse(run(bad_revoke)['all_refused'])

    def test_each_probe_has_exactly_one_claim_defect(self):
        for case, cid, _ in probe.CLAIM_PROBES:
            self.assertEqual(probe.check_claims(fake_jwt(cid), case, lambda: 0), 1000)
            with self.assertRaises(probe.Inconclusive):
                probe.check_claims(fake_jwt(cid), 'expired', lambda: 0)
        self.assertEqual(probe.check_claims(fake_jwt('openbao-probe-expiry'), 'expired', lambda: 0), 1000)

    def test_audience_alias_breaks_probe_isolation(self):
        token = fake_jwt('openbao-probe-sub')
        claims = json.loads(base64.urlsafe_b64decode(token.split('.')[1] + '=='))
        claims['aud'] = 'openbao-probe-sub'
        payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip('=')
        with self.assertRaises(probe.Inconclusive):
            probe.check_claims(f'h.{payload}.s', 'sub', lambda: 0)

    def test_redirects_are_not_followed(self):
        self.assertIsNone(probe.NoRedirect().redirect_request(None, None, 302, '', {}, 'https://other.example'))

    def test_not_yet_valid_does_not_count_as_expired(self):
        self.assertNotEqual(probe.classify('token not valid yet: nbf'), 'expired')

    def test_malformed_refusal_is_inconclusive(self):
        fake = FakeBao()
        def malformed(url, **kw):
            result = fake(url, **kw)
            if result[0] == 400:
                return 400, b'{"errors":null}'
            return result
        receipt = run(malformed)
        self.assertFalse(receipt['all_refused'])
        self.assertTrue(all(c['class'] == 'inconclusive' for c in receipt['cases'] if c['case'] != 'expired'))

    def test_expiry_client_ttl_is_bounded(self):
        token = fake_jwt('openbao-probe-expiry')
        claims = json.loads(base64.urlsafe_b64decode(token.split('.')[1] + '=='))
        claims['iat'] = claims['exp'] - 900
        payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip('=')
        with self.assertRaises(probe.Inconclusive):
            probe.check_claims(f'h.{payload}.s', 'expired', lambda: 0)

    def test_classify(self):
        self.assertEqual(probe.classify('error validating token: token is expired'), 'expired')
        self.assertEqual(probe.classify('aud claim does not match any bound audience'), 'claim_mismatch')
        self.assertEqual(probe.classify('permission denied'), 'refused_other')


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