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

#!/usr/bin/env python3
"""In-pod key transport. Only invoke through key-custody.py (captured output)."""
import base64
import fcntl
import json
import os
from pathlib import Path
import secrets
import stat
import sys

FILES = ('enckey', 'private.pem', 'public.pem')
LIMIT = 65536


class Refused(Exception):
    pass


def validate(values):
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import rsa
    if set(values) != set(FILES) or len(values['enckey']) != 96:
        raise Refused()
    if any(not value or len(value) > LIMIT for value in values.values()):
        raise Refused()
    private = serialization.load_pem_private_key(values['private.pem'], password=None)
    public = serialization.load_pem_public_key(values['public.pem'])
    if (not isinstance(private, rsa.RSAPrivateKey) or private.key_size < 2048
            or not isinstance(public, rsa.RSAPublicKey)
            or private.public_key().public_numbers() != public.public_numbers()):
        raise Refused()


def read(directory, name):
    fd = os.open(name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory)
    try:
        info = os.fstat(fd)
        forbidden = 0o022 if name == 'public.pem' else 0o027
        if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or info.st_mode & forbidden:
            raise Refused()
        with os.fdopen(os.dup(fd), 'rb') as stream:
            data = stream.read(LIMIT + 1)
        after = os.stat(name, dir_fd=directory, follow_symlinks=False)
        if len(data) > LIMIT or any(getattr(info, k) != getattr(after, k)
                for k in ('st_dev', 'st_ino', 'st_size', 'st_mtime_ns', 'st_ctime_ns')):
            raise Refused()
        return data
    finally:
        os.close(fd)


def install(directory, values, *, uid=0, gid=999):
    """Verify every existing file before creating anything; never replace."""
    missing = []
    for name, value in values.items():
        try:
            if read(directory, name) != value:
                raise Refused()
        except FileNotFoundError:
            missing.append(name)
    for name in missing:
        temp = '.key-custody-' + secrets.token_hex(12)
        fd = os.open(temp, 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.fchown(stream.fileno(), uid, gid)
                os.fchmod(stream.fileno(), 0o644 if name == 'public.pem' else 0o640)
                os.fsync(stream.fileno())
            os.link(temp, name, src_dir_fd=directory, dst_dir_fd=directory, follow_symlinks=False)
            os.unlink(temp, dir_fd=directory)
            os.fsync(directory)
        finally:
            try:
                os.unlink(temp, dir_fd=directory)
            except FileNotFoundError:
                pass
    if any(read(directory, name) != value for name, value in values.items()):
        raise Refused()


def operate(request, root=Path('/etc/privacyidea'), *, uid=0, gid=999):
    if any(p.is_symlink() for p in (root, *root.parents)):
        raise Refused()
    directory = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
    try:
        # The admitted service uses fsGroup 999, which may make its PVC group-writable.
        if os.fstat(directory).st_mode & 0o002:
            raise Refused()
        fcntl.flock(directory, fcntl.LOCK_EX | fcntl.LOCK_NB)
        if request == {'action': 'initialize-audit'}:
            # This explicit first-bootstrap action never creates/replaces enckey.
            enckey = read(directory, 'enckey')
            if len(enckey) != 96:
                raise Refused()
            for name in ('private.pem', 'public.pem'):
                try:
                    os.stat(name, dir_fd=directory, follow_symlinks=False)
                except FileNotFoundError:
                    continue
                raise Refused()
            from cryptography.hazmat.primitives import serialization
            from cryptography.hazmat.primitives.asymmetric import rsa
            private = rsa.generate_private_key(public_exponent=65537, key_size=2048)
            values = {'enckey': enckey,
                      'private.pem': private.private_bytes(serialization.Encoding.PEM,
                          serialization.PrivateFormat.TraditionalOpenSSL, serialization.NoEncryption()),
                      'public.pem': private.public_key().public_bytes(serialization.Encoding.PEM,
                          serialization.PublicFormat.SubjectPublicKeyInfo)}
            validate(values)
            install(directory, values, uid=uid, gid=gid)
            return {'result': 'initialized'}
        if request == {'action': 'export'}:
            values = {name: read(directory, name) for name in FILES}
            validate(values)
            return {name: base64.b64encode(value).decode() for name, value in values.items()}
        if set(request) == {'action', 'keys'} and request['action'] == 'restore':
            values = {name: base64.b64decode(value, validate=True) for name, value in request['keys'].items()}
            validate(values)
            install(directory, values, uid=uid, gid=gid)
            return {'result': 'restored'}
        raise Refused()
    finally:
        os.close(directory)


if __name__ == '__main__':
    try:
        raw = sys.stdin.buffer.read(4 * LIMIT + 1)
        if len(raw) > 4 * LIMIT:
            raise Refused()
        print(json.dumps(operate(json.loads(raw))))
    except Exception:
        # Transport may carry key bytes; diagnostics must never serialize them.
        print('Key custody operation refused or failed.', file=sys.stderr)
        sys.exit(1)