keyless-entry

#!/usr/bin/env python3

# Copyright 2022 Jonathan Kamens <jik@kamens.us>.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

# This script temporarily or persistently configures your
# LUKS-encrypted Linux system to reboot without a key needing to be
# typed. It's based on the mechanism documented by Tobias in
# https://askubuntu.com/a/997668. The idea is to configure a
# decryption key that is stored in a static key file on disk, and then
# to configure the system to use that key file when you want to reboot
# without typing a key, and not otherwise.
#
# Run the script with one of these arguments:
#
# configure     - Set things up without enabling keyless entry
# unconfigure   - Disable if enabled, remove configuration completely
# enable-once   - Enable keyless entry just for the next reboot
# enable-always - Enable keyless entry until it's disabled
# disable       - Disable keyless entry
# recover       - Same as disable, but clean up even if not fully
#                 enabled and rebuild all initial RAM disks even if
#                 enable was only told to enable a subset of them.
#
# enable-once works by doing configuring /etc/rc.local to call the
# script with the disable argument at the end of the reboot sequence,
# so wherever you invoke the script from with enable-once needs to be
# accessible on boot and executable or it won't work.
#
# You can specify --default-kernel with enable-once or enable-always
# to only enable for the default kernel that grub will use the next
# time it boots, or you can specify --kernel-version one or more times
# to specify one or more kernel versions to enable for. Otherwise
# enable and disable impact all installed kernels. The reason why you
# would want to only enable for a subset of installed kernels is
# because it makes enable and disable faster since fewer initial RAM
# disks need to be rebuilt, and in the case of enable-once, this also
# makes the next reboot finish faster.
#
# You will be prompted to enter an existing LUKS key (passphrase)
# once for each encrypted device during configure. This will add a
# master key that should be stored on an encrypted partition to allow
# us to add/remove one-time keys every time we enable/disable keyless
# without requiring user input.  The only key that is readable when
# the source is encrypted is temporary and can't be used to unlock
# the source partition when keyless is disabled.
#
# When configured, the following files are created:
#
# /etc/keyless-entry.conf - stores persistent state
# /boot/keyless-entry - stores the decryption key
# /etc/crypttab.keyful - copy of original /etc/crypttab
# /etc/crypttab.keyless - crypttab configured for keyless entry
#
# When enable-once is called, /etc/rc.local is also modified, creating
# if it it doesn't exist. The changes to /etc/rc.local are removed
# after reboot.
#
# When enable-once, enable-always, or disable is called, the initial
# ramdisk (initrd) for each installed kernel is modified. The initrds
# will also be modified when disable is called if keyless entry was
# previously enabled.

import argparse
import configparser
import filecmp
import json
import re
import os
import random
import shlex
import shutil
import string
import subprocess
import sys

config_file = '/etc/keyless-entry.conf'
key_file_tail = '/keyless-entry'
key_file = f'/boot{key_file_tail}'
crypttab_file = '/etc/crypttab'
keyful_file = f'{crypttab_file}.keyful'
keyless_file = f'{crypttab_file}.keyless'
rootkey_file = '/etc/keyless-master'
rc_local = '/etc/rc.local'

config = None


#
# Begin GRUB parsing code
#

# The code in this section is used to determine the default kernel when
# --default-kernel is specified to enable-once or enable-always.

def dequote(string):
    return re.sub(r'^([\'\"]?)(.*)\1$', r'\2', string)


def grub_get_default(path=None):
    if path is None:
        path = '/etc/default/grub'
    with open(path) as f:
        for line in f:
            match = re.match(r'\s*GRUB_DEFAULT\s*=\s*(.*?)\s*$', line)
            if match:
                return dequote(match[1])
    return None


def grub_get_saved_entry(path=None):
    if path is None:
        path = '/boot/grub/grubenv'
    if not os.path.exists(path):
        return None
    with open(path) as f:
        for line in f:
            match = re.match(r'\s*saved_entry\s*=\s*(.*?)\s*$', line)
            if match:
                return dequote(match[1])
    return None


def grub_cfg_iter(path=None):
    if path is None:
        path = '/boot/grub/grub.cfg'
    parents = []
    menu_index = 0

    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument('title', action='store')
    arg_parser.add_argument('--class', action='store')
    arg_parser.add_argument('--users', action='store')
    arg_parser.add_argument('--unrestricted', action='store_true')
    arg_parser.add_argument('--hotkey', action='store')
    arg_parser.add_argument('--id', action='store')
    arg_parser.add_argument('old_id', nargs='?')

    with open(path) as f:
        for line in f:
            if parents:
                if re.match(r'\s*\}\s*$', line):
                    parents.pop()
                    menu_index += 1
                    continue
                match = re.match(r'\s*initrd\s+/initrd\.img-(.*?)\s*$', line)
                if match:
                    yield ([(m['menu_index'], m['title'], m['id'])
                            for m in parents], match[1])
            match = re.match(r'\s*(menuentry|submenu)\s+(.*?)\s*\{\s*$', line)
            if not match:
                continue

            entry_type = match[1]
            args = arg_parser.parse_args(shlex.split(match[2].replace(
                '$menuentry_id_option', '--id')))
            parents.append({
                'title': args.title,
                'id': args.id or args.old_id,
                'entry_type': entry_type,
                'menu_index': str(menu_index)
            })
            if entry_type == 'submenu':
                menu_index = 0


def grub_get_cfg_kernel_version(selector, path=None):
    for entries, kernel_version in grub_cfg_iter(path):
        selectors = re.split(r'\s*>\s*', selector)
        while selectors and entries:
            if selectors.pop(0) not in entries.pop(0):
                break
        else:
            if not (selectors or entries):
                return kernel_version
    return None


def grub_get_default_kernel_version(
        cfg_path=None, default_path=None, grubenv_path=None):
    default = grub_get_default(path=default_path)
    if default == 'saved':
        default = grub_get_saved_entry(path=grubenv_path)
    if not default:
        default = '0'
    return grub_get_cfg_kernel_version(default, path=cfg_path)

#
# End GRUB parsing code
#


def load_config():
    global config
    if config is None:
        config = configparser.ConfigParser()
        config.read(config_file)
    if 'settings' not in config:
        config['settings'] = {}
    if 'version' not in config['settings']:
        if config.getboolean('settings', 'configured', fallback=False):
            # No root key file, no targets stored in config
            config['settings']['version'] = '1'
            upgrade_config()
        else:
            # Root key file, targets stored in config
            config['settings']['version'] = '2'


def save_config():
    with open(f'{config_file}.new', 'w') as f:
        config.write(f)
    os.rename(f'{config_file}.new', config_file)


def upgrade_config():
    print('Upgrading configuration for new script version')
    if config['settings']['version'] == '2':
        return
    if config.getboolean('settings', 'enabled', fallback=False):
        shutil.copy(key_file, rootkey_file)
        config['settings']['v1key'] = 'true'
    else:
        shutil.move(key_file, rootkey_file)
    filesystems, _ = make_keyless_content(keyful_file)
    for i in range(len(filesystems)):
        config['settings'][f'target{i}'] = filesystems[i][0]
        config['settings'][f'source{i}'] = filesystems[i][1]
    config['settings']['version'] = '2'
    save_config()


def get_filesystems():
    return ((config['settings'][f'target{k[6:]}'],
             config['settings'][f'source{k[6:]}'])
            for k in config['settings'].keys()
            if k.startswith('source'))


def get_boot_mountpoint():
    ud = '/dev/disk/by-uuid'
    findmnt = json.loads(subprocess.check_output(('findmnt', '-J', '/boot')))
    real_device = os.path.basename(findmnt['filesystems'][0]['source'])
    links = os.listdir(ud)
    for link in links:
        try:
            target = os.readlink(f'{ud}/{link}')
        except Exception:
            continue
        if target.endswith(real_device):
            return f'{ud}/{link}'
    raise Exception(f'Could not find by-uuid device for {real_device}')


def make_keyless_content(keyful_file):
    content = ''
    filesystems = []
    boot_mountpoint = get_boot_mountpoint()
    for line in open(keyful_file):
        line = re.sub(r'#.*', '', line).strip()
        if not line:
            continue
        (target, source, keyfile, options) = line.split()
        if keyfile != 'none':
            sys.exit(
                f"Can't work when key file already specified in {keyful_file}")
        if not options:
            sys.exit(f"Can't work when options not specified in {keyful_file}")
        if 'keyscript' in options:
            sys.exit(
                f"Can't work when keyscript already specified in "
                f"{keyful_file}")
        keyfile = f'{boot_mountpoint}:{key_file_tail}'
        options += ',keyscript=/lib/cryptsetup/scripts/passdev'
        content += f'{target} {source} {keyfile} {options}\n'
        filesystems.append((target, source))
    return filesystems, content


def in_rc_local():
    try:
        next(line for line in open(rc_local)
             if f'{os.path.basename(sys.argv[0])} disable' in line)
        return True
    except Exception:
        pass
    return False


def remove_from_rc_local():
    lines = [line for line in open(rc_local)
             if f'{os.path.basename(sys.argv[0])} disable' not in line]
    if len(lines) > 1:  # Something besides the #! line
        with open(f'{rc_local}.new', 'w') as f:
            f.write(''.join(lines))
        os.chmod(f'{rc_local}.new', 0o744)
        os.rename(f'{rc_local}.new', rc_local)
    else:
        os.remove(rc_local)


def add_to_rc_local():
    with open(f'{rc_local}.new', 'w') as out:
        try:
            with open(rc_local) as _in:
                old = _in.read()
                if old and not re.match(
                        r'^#!\s*(?:/usr)?/bin/(?:env\s+)?(sh|bash)\b', old):
                    sys.exit(f"Don't know how to safely add to {rc_local} "
                             f"as it is currently formatted")
        except FileNotFoundError:
            old = '#!/bin/sh\n'
        out.write(old)
        print(f'{os.path.realpath(sys.argv[0])} disable', file=out)
    os.chmod(f'{rc_local}.new', 0o744)
    os.rename(f'{rc_local}.new', rc_local)


def configure(args):
    # Load config.
    # Abort if already configured.
    # Abort if any of the files we use unexpectedly already exist.
    # Abort if crypttab does not exist.
    # Abort if crypttab has unexpected format.
    # Create master key file.
    # Add root keys to LUKS, saving sources in config.
    # Copy crypttab to crypttab.keyful.
    # Save crypttab.keyless.
    # Set configured to true in config.
    # Save config.
    load_config()
    if 'configured' in config['settings']:
        sys.exit('Already configured')
    for file in (rootkey_file, keyful_file, keyless_file):
        if os.path.exists(file):
            sys.exit(
                f"Can't configure when {file} already unexpectedly exists")
    if not os.path.exists(crypttab_file):
        sys.exit(f'{crypttab_file} does not exist')
    # This will abort if /etc/crypttab is missing or has unexpected format.
    filesystems, keyless_content = make_keyless_content(crypttab_file)
    with open(f'{rootkey_file}.new', 'w') as f:
        os.chmod(f'{rootkey_file}.new', 0o400)
        f.write(''.join(random.choices(string.ascii_letters, k=64)))
    os.rename(f'{rootkey_file}.new', rootkey_file)
    for fs_num in range(len(filesystems)):
        (target, source) = filesystems[fs_num]
        if len(filesystems) > 1:
            print(f'Adding key to {target}')
        subprocess.check_call(
                ('cryptsetup', 'luksAddKey', source, rootkey_file))
        config['settings'][f'source{fs_num}'] = source
        config['settings'][f'target{fs_num}'] = target
    shutil.copyfile(crypttab_file, keyful_file)
    with open(f'{keyless_file}.new', 'w') as f:
        f.write(keyless_content)
    os.rename(f'{keyless_file}.new', keyless_file)
    config['settings']['configured'] = 'true'
    save_config()


def unconfigure(args):
    # Load config.
    # Abort if not configured.
    # Abort if enabled.
    # Delete crypttab.keyless.
    # Delete crypttab.keyful.
    # Remove keys from LUKS.
    # Remove key file.
    # Delete config.
    load_config()
    if not config.getboolean('settings', 'configured', fallback=False):
        sys.exit("Can't unconfigure when not configured")
    if config.getboolean('settings', 'enabled', fallback=False):
        sys.exit("Please disable before unconfiguring")
    os.remove(keyless_file)
    os.remove(keyful_file)
    if not config.getboolean('settings', 'v1key', fallback=False):
        for target, source in get_filesystems():
            print(f'Removing key from {target} ({source})')
            subprocess.check_call(
                ('cryptsetup', 'luksRemoveKey', source, rootkey_file))
    os.remove(rootkey_file)
    os.remove(config_file)


def enable_once(args):
    # Load config.
    # Abort if already enabled and in rc.local
    # enable-always if not already enabled.
    # Add disable command to rc.local.
    load_config()
    if config.getboolean('settings', 'enabled', fallback=False):
        if in_rc_local():
            sys.exit('Already enabled once')
    else:
        enable_always(args, once=True)
    add_to_rc_local()


def enable_always(args, once=False):
    # Load config.
    # Abort if not configured.
    # Abort if already enabled, unless not doing once and in rc.local, in
    # which case remove from rc.local.
    # Abort if crypttab is different from crypttab.keyful.
    # Create key_file and add to source filesystems using rootkey_file
    # Replace crypttab with crypttab.keyless.
    # Regenerate initrds.
    # Set enabled to true in config.
    # Set kernel versions in config.
    # Save config.
    load_config()
    if not config.getboolean('settings', 'configured', fallback=False):
        sys.exit('You need to configure before enabling')
    if config.getboolean('settings', 'enabled', fallback=False):
        if not once and in_rc_local():
            remove_from_rc_local()
            print('Switched from once to always')
            return
        sys.exit('Already enabled')
    if not filecmp.cmp(crypttab_file, keyful_file):
        sys.exit(f'{crypttab_file} is different from {keyful_file}, aborting')
    if not os.path.exists(crypttab_file):
        sys.exit(f'{crypttab_file} does not exist')
    with open(f'{key_file}.new', 'w') as f:
        os.chmod(f'{key_file}.new', 0o400)
        f.write(''.join(random.choices(string.ascii_letters, k=64)))
    os.rename(f'{key_file}.new', key_file)
    for target, source in get_filesystems():
        print(f'Adding key to {target} ({source})')
        subprocess.check_call(
                ('cryptsetup', 'luksAddKey', '-d', rootkey_file,
                 source, key_file))
    if 'v1key' in config['settings']:
        del config['settings']['v1key']
    shutil.copyfile(keyless_file, crypttab_file)
    if args.kernel_version[0] == 'all':
        subprocess.check_call(('update-initramfs', '-c', '-k', 'all'))
    else:
        for version in args.kernel_version:
            subprocess.check_call(('update-initramfs', '-c', '-k', version))
    config['settings']['kernel_versions'] = ','.join(args.kernel_version)
    config['settings']['enabled'] = 'true'
    save_config()


def disable(args, recover=False):
    # Load config.
    # Abort if not configured.
    # Abort if not enabled unless recover is True.
    # Abort if crypttab is different from crypttab.keyless, unless
    #   recover is True and it's the same as crypttab.keyful.
    # Remove key_file (ok to not exist if recover is True).
    # Replace crypttab with crypttab.keyful.
    # Regenerate initrds.
    # Remove enabled setting from config.
    # Remove kernel_versions setting from config.
    # Save config.
    # Remove disable command from rc.local if it's there.
    load_config()
    if not config.getboolean('settings', 'configured', fallback=False):
        sys.exit("Can't disable if not yet configured")
    if not (recover or config.getboolean('settings', 'enabled', fallback=False)):
        sys.exit("Can't disable if not enabled")
    if not (filecmp.cmp(crypttab_file, keyless_file) or
            recover and filecmp.cmp(crypttab_file, keyful_file)):
        sys.exit(f'{crypttab_file} is different from {keyless_file}, aborting')
    for target, source in get_filesystems():
        print(f'Removing key from {target} ({source})')
        try:
            subprocess.check_call(
                ('cryptsetup', 'luksRemoveKey', source, key_file))
        except subprocess.CalledProcessError:
            if not recover:
                raise
    shutil.copyfile(keyful_file, crypttab_file)
    # At some point far into the future perhaps we can remove the backward
    # compatibility here and below and assume that kernel_versions is set.
    # *shrug*
    if recover:
        kernel_versions = ['all']
    else:
        kernel_versions = re.split(r'\s*,\s*', config['settings'].get(
            'kernel_versions', 'all'))
    for version in kernel_versions:
        subprocess.check_call(('update-initramfs', '-c', '-k', version))
    try:
        os.remove(key_file)
    except FileNotFoundError:
        if not recover:
            raise
    config['settings'].pop('enabled', None)
    config['settings'].pop('kernel_versions', None)
    save_config()
    if in_rc_local():
        remove_from_rc_local()


def recover(args):
    return disable(args, recover=True)


def parse_args():
    parser = argparse.ArgumentParser(description='Manage password-free '
                                     'reboots of encrypted Linux hosts')
    subparsers = parser.add_subparsers(required=True)
    configure_parser = subparsers.add_parser(
        'configure', help='Do initial configuration required before use')
    configure_parser.set_defaults(func=configure)
    unconfigure_parser = subparsers.add_parser(
        'unconfigure', help='Undo configuration, restoring system to pristine '
        'state')
    unconfigure_parser.set_defaults(func=unconfigure)

    kernel_version_parser = argparse.ArgumentParser(add_help=False)
    group = kernel_version_parser.add_mutually_exclusive_group()
    group.add_argument('--default-kernel', action='store_true',
                       help='Only enable for the default kernel')
    group.add_argument('--kernel-version', action='append', default=[],
                       help='Only enable for the specified kernel version(s)')

    enable_once_parser = subparsers.add_parser(
        'enable-once', parents=[kernel_version_parser],
        help='Enable no passphrase for just the next boot')
    enable_once_parser.set_defaults(func=enable_once)
    enable_always_parser = subparsers.add_parser(
        'enable-always', parents=[kernel_version_parser],
        help='Enable no passphrase for all boots')
    enable_always_parser.set_defaults(func=enable_always)

    disable_parser = subparsers.add_parser(
        'disable', help='Disable previously enabled no passphrase')
    disable_parser.set_defaults(func=disable)
    recover_parser = subparsers.add_parser(
        'recover', help='Recover from partial enable')
    recover_parser.set_defaults(func=recover)

    args = parser.parse_args()
    if 'default_kernel' in args and args.default_kernel:
        args.kernel_version = [grub_get_default_kernel_version()]
        if args.kernel_version[0] is None:
            parser.error("Couldn't determine default kernel version")
    elif 'kernel_version' in args and not args.kernel_version:
        args.kernel_version = ['all']
    elif 'kernel_version' in args:
        for version in args.kernel_version:
            if ',' in version:
                parser.error("Kernel version numbers can't have commas in them")
        
    return args


def add_to_search_path(_dir):
    dirs = os.environ['PATH'].split(':')
    if _dir not in dirs:
        dirs.append(_dir)
        os.environ['PATH'] = ':'.join(dirs)


def main():
    args = parse_args()
    # Need to make sure we can find cryptsetup and update-initramfs
    add_to_search_path('/sbin')
    add_to_search_path('/usr/sbin')
    path_good = True
    for cmd in ('cryptsetup', 'update-initramfs'):
        if shutil.which(cmd) is None:
            print(f"Can't find {cmd} in search path", file=sys.stderr)
            path_good = False
    if not path_good:
        sys.exit(1)
    args.func(args)


if __name__ == '__main__':
    main()