#!/usr/bin/python3 -I
# Row 234 (#1740, #436a row a), design r6 E3: the root half of host enrollment.
#
# Installed as /usr/local/sbin/mosaic-keys-write, root, 0755. mosaic-keys (the installer key's forced
# command) runs it through `sudo -n` with one request:
#
#   list <user>                          each line of /etc/ssh/mosaic-keys/<user> with a mosaic- tag
#   add <user> <type> <base64> <tag>     appends from="<from>",restrict <type> <base64> <tag>, unless the tag exists
#   remove <user> <tag>                  drops the line with that exact tag; a missing tag changes nothing
#
# It checks the request again, whatever mosaic-keys did. Host config is root-owned and written by ops:
# /etc/mosaic-keys/users (the users it may touch), /etc/mosaic-keys/privileged (users with sudo that
# Jason approved by name) and /etc/mosaic-keys/from (the from= value). Standard library only. Exit 0
# on success, 1 on any refusal, with the reason on stderr and in syslog.
import base64
import fcntl
import grp
import hashlib
import os
import pwd
import re
import socket
import stat
import subprocess
import sys
import syslog

KEY_DIR = '/etc/ssh/mosaic-keys'
CONF_DIR = '/etc/mosaic-keys'
USER = re.compile(r'[a-z_][a-z0-9_-]{0,31}')
TAG = re.compile(r'mosaic-[0-9a-f]{32}')
KEY_TYPE = re.compile(r'ssh-ed25519|ecdsa-sha2-nistp256|ssh-rsa')
BASE64 = re.compile(r'[A-Za-z0-9+/]+={0,2}')
FROM = re.compile(r'[0-9A-Za-z.:/*?!_-]+(,[0-9A-Za-z.:/*?!_-]+)*')
# In the code, not in a file: no request may name these, whatever users says.
RESERVED_USERS = ('root', 'localadmin', 'mosaic-installer')
MIN_UID = 1000
# Each is root by another route. disk reads and writes raw block devices, libvirt and incus-admin start
# guests with host mounts (r6 N5). privileged lifts this check and nothing else.
REFUSED_GROUPS = ('sudo', 'admin', 'wheel', 'docker', 'lxd', 'disk', 'libvirt', 'incus-admin')
LOCK = '.lock'


class Refused(Exception):
    pass


def _passwd(user):
    try:
        return pwd.getpwnam(user).pw_uid
    except KeyError:
        return None


def _groups(user):
    entry = pwd.getpwnam(user)
    names = set()
    for gid in os.getgrouplist(user, entry.pw_gid):
        try:
            names.add(grp.getgrgid(gid).gr_name)
        except KeyError:
            pass
    return names


def _sudo(user):
    try:
        out = subprocess.run(
            ['/usr/bin/sudo', '-l', '-U', user],
            env={'LC_ALL': 'C', 'PATH': '/usr/bin:/bin'},
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            timeout=10,
        )
    except subprocess.TimeoutExpired:
        # No answer is not the expected line, so the caller refuses user-sudo and logs it (duckworth-0081 N2).
        return ''
    return out.stdout.decode('utf-8', 'replace')


def _log(message):
    syslog.openlog('mosaic-keys-write', 0, syslog.LOG_AUTH)
    syslog.syslog(syslog.LOG_INFO, message)


class Host:
    # What the writer reads from the machine. Tests pass their own dirs, owner and lookups.
    def __init__(self, key_dir=KEY_DIR, conf_dir=CONF_DIR, owner=0, passwd=_passwd, groups=_groups,
                 sudo=_sudo, hostname=None, log=_log):
        self.key_dir = key_dir
        self.conf_dir = conf_dir
        self.owner = owner
        self.passwd = passwd
        self.groups = groups
        self.sudo = sudo
        self.hostname = hostname if hostname is not None else socket.gethostname().split('.')[0]
        self.log = log


def fingerprint(blob64):
    digest = hashlib.sha256(base64.b64decode(blob64, validate=True)).digest()
    return 'SHA256:' + base64.b64encode(digest).decode().rstrip('=')


def parse(argv):
    # The request's shape and every pattern, before any name reaches a path.
    if not argv:
        raise Refused('request')
    verb, args = argv[0], argv[1:]
    shapes = {'list': 1, 'add': 4, 'remove': 2}
    if verb not in shapes or len(args) != shapes[verb]:
        raise Refused('request')
    if not USER.fullmatch(args[0]):
        raise Refused('user-pattern')
    req = {'verb': verb, 'user': args[0]}
    if verb == 'add':
        if not KEY_TYPE.fullmatch(args[1]) or not BASE64.fullmatch(args[2]):
            raise Refused('key-format')
        try:
            req['fingerprint'] = fingerprint(args[2])
        except ValueError:
            raise Refused('key-format')
        req['type'], req['base64'] = args[1], args[2]
    if verb in ('add', 'remove'):
        if not TAG.fullmatch(args[-1]):
            raise Refused('tag-pattern')
        req['tag'] = args[-1]
    return req


def _owned(st, host, kind):
    if st.st_uid != host.owner or st.st_mode & 0o022 or not kind(st.st_mode):
        raise Refused('not-root-owned')


def read_conf(host, name):
    # A root-owned regular file, not group or world writable, opened without following a symlink.
    try:
        fd = os.open(os.path.join(host.conf_dir, name), os.O_RDONLY | os.O_NOFOLLOW)
    except OSError:
        raise Refused('config-' + name)
    try:
        _owned(os.fstat(fd), host, stat.S_ISREG)
        with os.fdopen(os.dup(fd), 'r', encoding='utf-8') as f:
            text = f.read()
    finally:
        os.close(fd)
    return [line.strip() for line in text.splitlines() if line.strip() and not line.strip().startswith('#')]


def check_user(req, host):
    user = req['user']
    users = read_conf(host, 'users')
    privileged = read_conf(host, 'privileged')
    if user in RESERVED_USERS or user not in users:
        raise Refused('user-not-listed')
    uid = host.passwd(user)
    if uid is None:
        raise Refused('user-unknown')
    if uid == 0 or uid < MIN_UID:
        raise Refused('user-system')
    if user in privileged:
        return
    if set(REFUSED_GROUPS) & set(host.groups(user)):
        raise Refused('user-group')
    expected = 'User %s is not allowed to run sudo on %s.' % (user, host.hostname)
    out = host.sudo(user)
    if out not in (expected, expected + '\n'):
        raise Refused('user-sudo')


def read_from(host):
    lines = read_conf(host, 'from')
    if len(lines) != 1 or not FROM.fullmatch(lines[0]):
        raise Refused('config-from')
    return lines[0]


def tag_of(line):
    parts = line.split()
    return parts[-1] if parts and TAG.fullmatch(parts[-1]) else None


def read_lines(dir_fd, user, host):
    # ENOENT means an empty list: a first add finds no file (r6 N6).
    try:
        fd = os.open(user, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=dir_fd)
    except FileNotFoundError:
        return []
    except OSError:
        raise Refused('user-file')
    try:
        _owned(os.fstat(fd), host, stat.S_ISREG)
        with os.fdopen(os.dup(fd), 'r', encoding='utf-8') as f:
            return f.read().splitlines()
    finally:
        os.close(fd)


def write_lines(dir_fd, user, lines, host):
    tmp = '.' + user + '.tmp'
    try:
        st = os.stat(tmp, dir_fd=dir_fd, follow_symlinks=False)
    except FileNotFoundError:
        st = None
    if st is not None:
        # Under the flock a leftover tmp can only come from a writer that crashed.
        if not stat.S_ISREG(st.st_mode) or st.st_uid != host.owner:
            raise Refused('tmp-file')
        os.unlink(tmp, dir_fd=dir_fd)
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o644, dir_fd=dir_fd)
    try:
        os.fchmod(fd, 0o644)
        os.write(fd, ''.join(line + '\n' for line in lines).encode('utf-8'))
        os.fsync(fd)
    finally:
        os.close(fd)
    os.rename(tmp, user, src_dir_fd=dir_fd, dst_dir_fd=dir_fd)
    os.fsync(dir_fd)


def run(req, host, out):
    check_user(req, host)
    line_from = read_from(host) if req['verb'] == 'add' else None
    try:
        dir_fd = os.open(host.key_dir, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
    except OSError:
        raise Refused('key-dir')
    try:
        _owned(os.fstat(dir_fd), host, stat.S_ISDIR)
        lock = os.open(LOCK, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600, dir_fd=dir_fd)
        try:
            fcntl.flock(lock, fcntl.LOCK_EX)
            user = req['user']
            lines = read_lines(dir_fd, user, host)
            if req['verb'] == 'list':
                for line in lines:
                    if tag_of(line):
                        out.write(line + '\n')
                return 'listed'
            if req['verb'] == 'add':
                if any(tag_of(line) == req['tag'] for line in lines):
                    return 'present'
                entry = 'from="%s",restrict %s %s %s' % (line_from, req['type'], req['base64'], req['tag'])
                write_lines(dir_fd, user, lines + [entry], host)
                return 'added'
            kept = [line for line in lines if tag_of(line) != req['tag']]
            if len(kept) == len(lines):
                return 'absent'
            write_lines(dir_fd, user, kept, host)
            return 'removed'
        finally:
            os.close(lock)
    finally:
        os.close(dir_fd)


def main(argv, host, out=sys.stdout, err=sys.stderr):
    req = None
    try:
        req = parse(argv)
        result = run(req, host, out)
        code = 0
    except Refused as e:
        result = 'refused ' + str(e)
        code = 1
    except OSError as e:
        result = 'refused os-' + (e.__class__.__name__)
        code = 1
    fields = ['result=' + result]
    if req:
        fields.append('user=' + req['user'])
        if 'tag' in req:
            fields.append('tag=' + req['tag'])
        if 'fingerprint' in req:
            fields.append('key=' + req['fingerprint'])
    host.log(' '.join(fields))
    if code:
        err.write('mosaic-keys-write: ' + result + '\n')
    return code


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:], Host()))