net-kingdom/sso-mfa/k8s/keycape/approval-clients-rollout.py

#!/usr/bin/env python3
"""Contained KeyCape approval-client config/image cutover and acceptance."""
from __future__ import annotations
import argparse
import base64
import copy
import importlib.util
import json
import os
from pathlib import Path
import secrets
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import jwt
import yaml

PLATFORM = Path('/home/worsch/railiance-platform')
KEYCAPE = Path('/home/worsch/key-cape')
sys.path.insert(0, str(PLATFORM / 'scripts'))
from state_hub_preflight_lane import LaneError, assert_cluster, bao, command, data
from keycape_approval_custody import require
spec = importlib.util.spec_from_file_location('pin', Path(__file__).with_name('openbao-client-config.py'))
pin = importlib.util.module_from_spec(spec); spec.loader.exec_module(pin)
ISSUER = 'https://kc.coulomb.social'
IMAGE = 'forgejo.coulomb.social/coulomb/key-cape@sha256:7ff54c54e63ee172ae9e6e7fd2da96e427352f712343d74626ee6fe0f6f82611'
PRIOR_IMAGE = 'forgejo.coulomb.social/coulomb/key-cape:main-153258b'
IDS = ('secrets-engine-approval', 'approval-engine-operator')
SECRET_NAMES = ('keycape-secrets-engine-approval-client', 'keycape-approval-engine-operator-client')


def registrations():
    source = yaml.safe_load((KEYCAPE / 'config/service-clients.example.yaml').read_text())['clients']
    result = [next(c for c in source if c['clientId'] == name) for name in IDS]
    for i, c in enumerate(result):
        require(c['audience'] == 'approval-engine' and c['tenant'] == 'tenant:platform'
                and c['grantTypes'] == ['client_credentials'] and c['clientType'] == 'confidential'
                and c['tokenLifetime'] == '15m', 'registration_contract_drift')
        scopes = ['approval:read', 'approval:consume'] if i == 0 else [
            'approval:create', 'approval:read', 'approval:approve', 'approval:revoke',
            'approval:supersede', 'approval:observe', 'approval:emit']
        require(c['allowedScopes'] == scopes and c['roles'] == (["secrets-engine"] if i == 0 else ['approval-operator'])
                and c['serviceSubject'] == ('service:secrets-engine' if i == 0 else 'service:approval-engine-operator')
                and c['secretRef'] == ('env:KEYCAPE_SECRETS_ENGINE_APPROVAL_CLIENT_SECRET' if i == 0 else 'env:KEYCAPE_APPROVAL_ENGINE_OPERATOR_CLIENT_SECRET'), 'registration_authority_drift')
    return result


def replacement(secret, clients):
    raw, config, _ = pin.issuer_document(secret)
    require(config['authelia'].get('issuer') == pin.VERIFIED_UPSTREAM_ISSUER, 'verified_issuer_pin_required')
    require(isinstance(config.get('clients'), list) and config['clients'], 'client_sequence_required')
    existing = [c['clientId'] for c in config['clients']]
    require(len(existing) == len(set(existing)) and not set(existing).intersection(IDS), 'existing_client_registration_requires_reconciliation')
    root = yaml.compose(raw)
    node = next(v for k, v in root.value if k.value == 'clients')
    require(isinstance(node, yaml.nodes.SequenceNode) and not node.flow_style, 'block_client_sequence_required')
    index = node.end_mark.index
    # Insert at the start of the next top-level line, preserving every old byte.
    line_start = raw.rfind('\n', 0, index) + 1
    if not raw[line_start:index].strip():
        index = line_start
    indent = ' ' * node.start_mark.column
    fragment = yaml.safe_dump(clients, sort_keys=False, width=120)
    fragment = ''.join(indent + line if line.strip() else line for line in fragment.splitlines(True))
    addition = ('' if index == 0 or raw[index - 1] == '\n' else '\n') + fragment
    updated = raw[:index] + addition + raw[index:]
    expected = copy.deepcopy(config); expected['clients'].extend(clients)
    require(yaml.load(updated, Loader=pin.UniqueLoader) == expected, 'unrelated_configuration_changed')
    return base64.b64encode(updated.encode()).decode()


def candidate_spec(deployment):
    result = copy.deepcopy(deployment['spec'])
    require(result.get('replicas', 1) == 1, 'single_replica_required')
    container = next(c for c in result['template']['spec']['containers'] if c['name'] == 'keycape')
    require(container['image'] == PRIOR_IMAGE, 'prior_image_drift')
    patch = yaml.safe_load((KEYCAPE / 'docs/approval-clients-deployment.patch.yaml').read_text())['spec']
    approved = patch['template']['spec']['containers'][0]
    require(approved['image'] == IMAGE and patch['strategy']['type'] == 'Recreate', 'candidate_source_drift')
    env = container.setdefault('env', [])
    require(not {x['name'] for x in env}.intersection(x['name'] for x in approved['env']), 'existing_candidate_env')
    for item, client, secret_name in zip(approved['env'], registrations(), SECRET_NAMES):
        require(item == {'name': client['secretRef'].removeprefix('env:'), 'valueFrom': {
            'secretKeyRef': {'name': secret_name, 'key': 'client-secret'}}}, 'env_delivery_contract_drift')
    env.extend(approved['env']); container['image'] = IMAGE
    container['readinessProbe']['httpGet'] = {'path': '/readyz', 'port': 8080}
    result['strategy'] = {'type': 'Recreate'}
    return result


def patch_object(kube, kind, obj, path, value, dry=False):
    patch = [{'op': 'test', 'path': '/metadata/uid', 'value': obj['metadata']['uid']},
             {'op': 'test', 'path': '/metadata/resourceVersion', 'value': obj['metadata']['resourceVersion']},
             {'op': 'replace', 'path': path, 'value': value}]
    args = kube + ['-n', 'sso', 'patch', kind, obj['metadata']['name'], '--type=json', '--patch-file=/dev/stdin', '-o', 'json']
    if dry:
        args += ['--dry-run=server']
    return data(command(args, payload=patch))


def get(kube, kind, name):
    return data(command(kube + ['-n', 'sso', 'get', kind, name, '-o', 'json']))


def ready(kube, image, timeout=150):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        dep = get(kube, 'deployment', 'keycape')
        status = dep.get('status', {})
        if (status.get('observedGeneration') == dep['metadata']['generation']
                and status.get('updatedReplicas') == status.get('readyReplicas') == status.get('availableReplicas') == 1
                and status.get('replicas') == 1):
            podlist = data(command(kube + ['-n', 'sso', 'get', 'pods', '-l', 'app.kubernetes.io/name=keycape', '-o', 'json']))['items']
            if len(podlist) == 1 and not podlist[0]['metadata'].get('deletionTimestamp'):
                containers = podlist[0].get('status', {}).get('containerStatuses', [])
                if any(c['name'] == 'keycape' and c['ready'] and c['image'] == image for c in containers):
                    return {'deployment_uid': dep['metadata']['uid'], 'generation': dep['metadata']['generation'],
                            'pod_uid': podlist[0]['metadata']['uid'], 'single_ready_replica': True}
        time.sleep(3)
    raise LaneError('keycape_readiness_timeout')


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        return None


def http(path, fields=None, credentials=None):
    headers = {'Accept': 'application/json'}
    if credentials:
        headers['Authorization'] = 'Basic ' + base64.b64encode(':'.join(credentials).encode()).decode()
    body = None
    if fields is not None:
        body = urllib.parse.urlencode(fields).encode()
        headers['Content-Type'] = 'application/x-www-form-urlencoded'
    request = urllib.request.Request(ISSUER + path, data=body, headers=headers)
    opener = urllib.request.build_opener(NoRedirect())
    try:
        with opener.open(request, timeout=20) as response:
            return response.status, json.load(response)
    except urllib.error.HTTPError as error:
        return error.code, json.load(error)


def denied(result, status, feature):
    code, body = result
    require(code == status and body.get('error') == 'invalid_profile_usage'
            and body.get('feature') == feature, 'keycape_denial_inconclusive')


def acceptance(kube, receipt):
    code, discovery = http('/.well-known/openid-configuration')
    require(code == 200 and discovery['issuer'] == ISSUER
            and discovery['token_endpoint'] == ISSUER + '/token'
            and discovery['jwks_uri'].startswith(ISSUER + '/'), 'discovery_binding_mismatch')
    code, jwks = http(discovery['jwks_uri'].removeprefix(ISSUER))
    require(code == 200, 'jwks_failed')
    receipt['clients'] = []
    for client, secret_name in zip(registrations(), SECRET_NAMES):
        secret = get(kube, 'secret', secret_name)
        credential = base64.b64decode(secret['data']['client-secret'], validate=True).decode()
        scopes = ' '.join(client['allowedScopes'])
        result, response = http('/token', {'grant_type': 'client_credentials', 'scope': scopes}, (client['clientId'], credential))
        require(result == 200 and response.get('token_type') == 'Bearer' and response.get('expires_in') == 900, 'service_issuance_failed')
        token = response['access_token']; header = jwt.get_unverified_header(token)
        keys = [key for key in jwks['keys'] if key['kid'] == header.get('kid')]
        require(header.get('alg') == 'RS256' and len(keys) == 1, 'signing_key_selection_failed')
        claims = jwt.decode(token, jwt.PyJWK.from_dict(keys[0]).key, algorithms=['RS256'],
                            issuer=ISSUER, audience='approval-engine',
                            options={'require': ['exp', 'iat', 'sub', 'iss', 'aud']})
        require(claims['sub'] == client['serviceSubject'] and claims['tenant'] == 'tenant:platform'
                and claims['roles'] == client['roles'] and claims['exp'] - claims['iat'] == 900
                and claims['principal_type'] == 'service'
                and set(claims['scope'].split()) == set(client['allowedScopes']), 'exact_claims_mismatch')
        excessive = 'approval:approve' if client['clientId'] == IDS[0] else 'approval:consume'
        denied(http('/token', {'grant_type': 'client_credentials', 'scope': excessive},
                    (client['clientId'], credential)), 400, 'scope')
        denied(http('/token', {'grant_type': 'client_credentials', 'scope': scopes},
                    (client['clientId'], secrets.token_urlsafe(48))), 401, 'Authorization')
        args = kube + ['-n', 'sso', 'exec', 'deployment/keycape', '--', '/keycape', 'verify-client',
            '-issuer', ISSUER, '-client-id', client['clientId'], '-audience', 'approval-engine', '-scope', scopes,
            '-secret-env', client['secretRef'].removeprefix('env:'), '-expect-subject', client['serviceSubject'],
            '-expect-tenant', 'tenant:platform', '-expect-roles', ','.join(client['roles']), '-deny-scope', excessive]
        command(args)
        receipt['clients'].append({'client_id': client['clientId'], 'live_jwks_signature_verified': True,
            'exact_claims_verified': True, 'lifetime_seconds': 900, 'excess_scope_denied': True,
            'wrong_secret_denied': True, 'pod_verify_client_passed': True,
            'real_predecessor_rotation_tested': False, 'observed_wall_clock_expiry': False})
    query = urllib.parse.urlencode({'client_id': 'openbao-admin', 'response_type': 'code',
        'redirect_uri': 'http://localhost:8250/oidc/callback', 'scope': 'openid approval:consume',
        'code_challenge_method': 'S256', 'code_challenge': 'A' * 43, 'state': secrets.token_urlsafe(32)})
    denied(http('/authorize?' + query), 400, 'scope')
    receipt['human_client_consume_denied'] = True


def rollout(kube, receipt, recovery_path):
    assert_cluster(kube)
    before = get(kube, 'secret', 'keycape-config')
    deployment = get(kube, 'deployment', 'keycape')
    require(before['metadata']['uid'] == '2e94519d-1550-41c7-9701-2efe47fe1fd3'
            and deployment['metadata']['uid'] == '99ddd83c-cb3f-4847-bcf8-35f1aa87627f', 'owner_object_identity_mismatch')
    config = replacement(before, registrations()); candidate = candidate_spec(deployment)
    patch_object(kube, 'secret', before, '/data/config.yaml', config, dry=True)
    patch_object(kube, 'deployment', deployment, '/spec', candidate, dry=True)
    previous = {'deployment': deployment, 'secret': before}
    recovery_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    require(recovery_path.parent.stat().st_mode & 0o077 == 0, 'private_recovery_directory_required')
    fd = os.open(recovery_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, 'w') as out:
        json.dump(previous, out)
    receipt['protected_recovery_retained'] = True
    changed = None; deployed = None
    try:
        changed = patch_object(kube, 'secret', before, '/data/config.yaml', config)
        require(changed['data'] == dict(before['data'], **{'config.yaml': config}), 'secret_readback_mismatch')
        deployed = patch_object(kube, 'deployment', deployment, '/spec', candidate)
        receipt.update(ready(kube, IMAGE))
        after = get(kube, 'secret', 'keycape-config')
        require(after['data'] == changed['data'], 'configuration_changed_during_rollout')
        acceptance(kube, receipt)
        receipt.update(status='service_acceptance_passed_pending_fresh_human_login', image=IMAGE,
            config_resource_version=after['metadata']['resourceVersion'], signing_key_unchanged=True,
            unrelated_config_bytes_preserved=True, existing_human_login_after=False)
    except Exception:
        # Read after uncertain API outcomes too; never assume a timeout means no write.
        now = get(kube, 'secret', 'keycape-config')
        require(now['metadata']['uid'] == before['metadata']['uid'], 'rollback_config_identity_drift')
        expected_data = dict(before['data'], **{'config.yaml': config})
        require(now['data'] in (before['data'], expected_data), 'rollback_config_drift_requires_owner_reconcile')
        if now['data'] == expected_data:
            patch_object(kube, 'secret', now, '/data/config.yaml', before['data']['config.yaml'])
        now = get(kube, 'deployment', 'keycape')
        require(now['metadata']['uid'] == deployment['metadata']['uid'], 'rollback_deployment_identity_drift')
        expected_spec = deployed['spec'] if deployed else candidate
        require(now['spec'] in (deployment['spec'], expected_spec), 'rollback_deployment_drift_requires_owner_reconcile')
        if now['spec'] == expected_spec:
            patch_object(kube, 'deployment', now, '/spec', deployment['spec'])
            ready(kube, PRIOR_IMAGE)
        receipt['compatible_pair_restored'] = True
        raise


def run(args, receipt):
    require(Path.home().parent.name == '.warden-attended-login' and not os.environ.get('BAO_TOKEN')
            and not os.environ.get('VAULT_TOKEN'), 'attended_warden_envelope_required')
    identity = data(bao(['token', 'lookup', '-format=json']))['data']
    require('platform-admin' in identity['policies'] and 'root' not in identity['policies'], 'attended_platform_admin_required')
    kube = ['kubectl', '--kubeconfig', args.kubeconfig]
    if args.action == 'rollout':
        receipt['existing_human_login_before'] = True
        rollout(kube, receipt, Path(args.recovery))
    else:
        assert_cluster(kube); receipt.update(ready(kube, IMAGE)); acceptance(kube, receipt)
        receipt.update(status='service_and_existing_human_login_acceptance_passed', existing_human_login_after=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['rollout', 'verify-after-login'])
    parser.add_argument('--kubeconfig', required=True)
    parser.add_argument('--recovery')
    parser.add_argument('--receipt', required=True)
    args = parser.parse_args()
    receipt = {'schema': 'netkingdom.keycape-approval-rollout.v1', 'status': 'failed', 'values_emitted': False}
    fd = os.open(args.receipt, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    try:
        run(args, receipt)
    except Exception as exc:
        receipt['error'] = str(exc) if isinstance(exc, (LaneError, pin.IssuerPinError)) else 'contained_operation_failed'
    finally:
        with os.fdopen(fd, 'w') as out:
            json.dump(receipt, out, indent=2); out.write('\n')
    return 0 if receipt['status'] != 'failed' else 1

if __name__ == '__main__':
    raise SystemExit(main())