net-kingdom/sso-mfa/k8s/privacyidea/key-custody.py

#!/usr/bin/env python3
"""Captured key export/restore for the admitted privacyIDEA PVC profile."""
import argparse
import base64
import importlib.util
import json
import os
from pathlib import Path
import re
import secrets
import sys

HERE = Path(__file__).resolve().parent


def load(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


common = load('service_custody', HERE.parents[1] / 'bootstrap/apply-service-secrets.py')
remote = load('key_custody_remote', HERE / 'key-custody-remote.py')
Refused = common.Refused
LOCAL = {'enckey': 'pi.enc', 'private.pem': 'private.pem', 'public.pem': 'public.pem'}
TARGETS = {'privacyidea-enckey': ('enckey',), 'privacyidea-auditkeys': ('private.pem', 'public.pem')}
RECOVERY_POD = 'privacyidea-key-recovery'
SLEEP = ['python3', '-c', 'import time; time.sleep(3600)']
IMAGE = re.compile(r'ghcr\.io/gpappsoft/privacyidea-docker@sha256:[a-f0-9]{64}\Z')
UID = re.compile(r'[a-zA-Z0-9-]{1,128}\Z')


def get(kind, name=None, missing=False):
    args = ['kubectl', 'get', kind]
    if name:
        args.append(name)
    args += ['-n', 'mfa', '-o', 'json']
    if missing:
        args.append('--ignore-not-found')
    raw = common.run(args)
    return json.loads(raw) if raw.strip() else None


def meta(obj, kind, name=None):
    if not isinstance(obj, dict) or obj.get('kind') != kind:
        raise Refused()
    m = obj.get('metadata', {})
    if (m.get('namespace') != 'mfa' or (name and m.get('name') != name)
            or not UID.fullmatch(str(m.get('uid', ''))) or m.get('deletionTimestamp')):
        raise Refused()
    return m


def mounts_data(spec, container):
    volumes = {v['name']: v for v in spec.get('volumes', [])}
    return any(v.get('mountPath') == '/etc/privacyidea' and not v.get('subPath')
               and not v.get('readOnly') and volumes.get(v['name'], {}).get('persistentVolumeClaim', {}).get('claimName') == 'privacyidea-data'
               for v in container.get('volumeMounts', []))


def pod_uses_data(pod):
    return any(v.get('persistentVolumeClaim', {}).get('claimName') == 'privacyidea-data'
               for v in pod.get('spec', {}).get('volumes', []))


def context(action, expected_uid=None):
    deployment = get('deployment', 'privacyidea')
    dm = meta(deployment, 'Deployment', 'privacyidea')
    template = deployment['spec']['template']['spec']
    container = next(c for c in template['containers'] if c['name'] == 'privacyidea')
    if template.get('securityContext', {}).get('fsGroup') != 999 or not mounts_data(template, container):
        raise Refused()
    pvc = get('pvc', 'privacyidea-data')
    pm = meta(pvc, 'PersistentVolumeClaim', 'privacyidea-data')
    if pvc.get('status', {}).get('phase') != 'Bound' or (expected_uid and pm['uid'] != expected_uid):
        raise Refused()
    listing = get('pods')
    if listing.get('kind') not in ('PodList', 'List'):
        raise Refused()
    users = [pod for pod in listing['items'] if pod_uses_data(pod)]
    if len(users) != 1:
        raise Refused()
    pod = users[0]; podmeta = meta(pod, 'Pod')
    spec = pod['spec']
    if pod.get('status', {}).get('phase') != 'Running':
        raise Refused()
    if action == 'restore':
        if (deployment['spec'].get('replicas', 1) != 0
                or any(deployment.get('status', {}).get(k, 0) for k in ('replicas', 'readyReplicas', 'availableReplicas', 'updatedReplicas'))
                or deployment.get('status', {}).get('observedGeneration', 0) < dm.get('generation', 1)):
            raise Refused()
        if (podmeta['name'] != RECOVERY_POD or podmeta.get('annotations', {}).get('net-kingdom/recovery-pvc-uid') != expected_uid
                or spec.get('automountServiceAccountToken') is not False
                or any(spec.get(k) for k in ('hostNetwork', 'hostPID', 'hostIPC', 'initContainers', 'ephemeralContainers'))
                or len(spec.get('containers', [])) != 1 or len(spec.get('volumes', [])) != 1):
            raise Refused()
        c = spec['containers'][0]
        if (c.get('name') != 'recovery' or not IMAGE.fullmatch(c.get('image', ''))
                or c.get('command') != SLEEP or c.get('args') or c.get('env') or c.get('envFrom')
                or c.get('securityContext', {}).get('privileged')
                or not mounts_data(spec, c)):
            raise Refused()
    else:
        c = next(c for c in spec['containers'] if c['name'] == 'privacyidea')
        if not mounts_data(spec, c):
            raise Refused()
        owners = podmeta.get('ownerReferences', [])
        if len(owners) != 1 or owners[0].get('kind') != 'ReplicaSet' or not owners[0].get('controller'):
            raise Refused()
        rs = get('replicaset', owners[0]['name']); rm = meta(rs, 'ReplicaSet', owners[0]['name'])
        if rm['uid'] != owners[0]['uid'] or not any(o.get('uid') == dm['uid'] and o.get('kind') == 'Deployment' and o.get('controller') for o in rm.get('ownerReferences', [])):
            raise Refused()
    if not any(s.get('name') == c['name'] and s.get('ready') for s in pod['status'].get('containerStatuses', [])):
        raise Refused()
    return {'pod': podmeta['name'], 'pod_uid': podmeta['uid'], 'pvc_uid': pm['uid'],
            'deployment_uid': dm['uid'], 'container': c['name']}


def invoke(ctx, request):
    source = (HERE / 'key-custody-remote.py').read_text()
    raw = common.run(['kubectl', 'exec', '-i', '-n', 'mfa', ctx['pod'], '-c', ctx['container'],
                      '--', 'python3', '-c', source], json.dumps(request).encode())
    return json.loads(raw)


def read_backups():
    result = {}
    for name in TARGETS:
        obj = get('secret', name, missing=True)
        result[name] = common.decode_secret(json.dumps(obj), 'mfa', name) if obj is not None else None
    return result


def verify_backups(backups, values):
    for name, keys in TARGETS.items():
        old = backups[name]
        if old is not None and any(old['data'].get(k) != base64.b64encode(values[k]).decode() for k in keys):
            raise Refused()


def local_values(directory):
    result = {name: common.read_private(directory, local) for name, local in LOCAL.items()}
    remote.validate(result)
    return result


def save_local(directory, values):
    # All differences are checked before creating any file. Existing copies survive.
    missing = []
    for name, local in LOCAL.items():
        try:
            if common.read_private(directory, local) != values[name]:
                raise Refused()
        except FileNotFoundError:
            missing.append((name, local))
    for name, local in missing:
        temporary = '.key-export-' + secrets.token_hex(12)
        fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=directory)
        try:
            with os.fdopen(fd, 'wb') as stream:
                stream.write(values[name]); stream.flush(); os.fsync(stream.fileno())
            os.link(temporary, local, src_dir_fd=directory, dst_dir_fd=directory, follow_symlinks=False)
            os.unlink(temporary, dir_fd=directory); os.fsync(directory)
        finally:
            try:
                os.unlink(temporary, dir_fd=directory)
            except FileNotFoundError:
                pass
    if local_values(directory) != values:
        raise Refused()


def perform(action, root, expected_uid=None, initialize=False):
    if action == 'restore' and (not expected_uid or not UID.fullmatch(expected_uid)):
        raise Refused()
    with common.custody(root, ('privacyidea',)) as directories:
        directory = directories['privacyidea']
        values = local_values(directory) if action == 'restore' else None
        ctx = context(action, expected_uid)
        if action == 'restore':
            result = invoke(ctx, {'action': 'restore', 'keys': {k: base64.b64encode(v).decode() for k, v in values.items()}})
            if result != {'result': 'restored'} or context(action, expected_uid) != ctx:
                raise Refused()
        else:
            backups = read_backups()
            if initialize:
                if any(value is not None for value in backups.values()):
                    raise Refused()
                # Independent local recovery copies also prohibit first-bootstrap creation.
                for name in LOCAL.values():
                    try:
                        os.stat(name, dir_fd=directory, follow_symlinks=False)
                    except FileNotFoundError:
                        continue
                    raise Refused()
                if invoke(ctx, {'action': 'initialize-audit'}) != {'result': 'initialized'}:
                    raise Refused()
            encoded = invoke(ctx, {'action': 'export'})
            values = {k: base64.b64decode(v, validate=True) for k, v in encoded.items()}
            remote.validate(values)
            if context(action) != ctx:
                raise Refused()
            verify_backups(backups, values)
            save_local(directory, values)
            for name, keys in TARGETS.items():
                common.publish('mfa', name, {k: values[k] for k in keys}, backups[name], set())
    print('Key custody operation verified. Retain independent recovery copies; native service acceptance remains required.')


def render(image, pvc_uid):
    if not image or not IMAGE.fullmatch(image) or not pvc_uid or not UID.fullmatch(pvc_uid):
        raise Refused()
    return {'apiVersion': 'v1', 'kind': 'Pod',
            'metadata': {'name': RECOVERY_POD, 'namespace': 'mfa',
                         'annotations': {'net-kingdom/recovery-pvc-uid': pvc_uid}},
            'spec': {'restartPolicy': 'Never', 'automountServiceAccountToken': False,
                     'securityContext': {'runAsUser': 0, 'runAsGroup': 999, 'fsGroup': 999},
                     'containers': [{'name': 'recovery', 'image': image, 'command': SLEEP,
                         'securityContext': {'allowPrivilegeEscalation': False, 'readOnlyRootFilesystem': True,
                             'capabilities': {'drop': ['ALL'], 'add': ['CHOWN', 'DAC_OVERRIDE', 'FOWNER']}},
                         'resources': {'requests': {'cpu': '10m', 'memory': '64Mi'}, 'limits': {'cpu': '100m', 'memory': '128Mi'}},
                         'volumeMounts': [{'name': 'data', 'mountPath': '/etc/privacyidea'}]}],
                     'volumes': [{'name': 'data', 'persistentVolumeClaim': {'claimName': 'privacyidea-data'}}]}}


def main():
    class QuietParser(argparse.ArgumentParser):
        def error(self, message):
            raise Refused()
    parser = QuietParser(description=__doc__)
    parser.add_argument('action', choices=['export', 'restore', 'render-pod'])
    parser.add_argument('root', nargs='?')
    parser.add_argument('--pvc-uid')
    parser.add_argument('--image')
    parser.add_argument('--initialize-audit-keys', action='store_true')
    args = parser.parse_args()
    if args.action == 'render-pod':
        if args.root or args.initialize_audit_keys:
            raise Refused()
        print(json.dumps(render(args.image, args.pvc_uid), indent=2))
    else:
        if not args.root or args.image or (args.action == 'restore' and args.initialize_audit_keys) or (args.action == 'export' and args.pvc_uid):
            raise Refused()
        perform(args.action, args.root, args.pvc_uid, args.initialize_audit_keys)


if __name__ == '__main__':
    try:
        main()
    except Exception:
        print('ERROR: key custody refused or failed; diagnostics suppressed. A failed write may have partial effects. '
              'Retain copies and reconcile target state before retrying.', file=sys.stderr)
        sys.exit(1)