/* SPDX-License-Identifier: LGPL-2.1-or-later */

#include <sys/stat.h>
#include <sysexits.h>

#include "sd-messages.h"

#include "alloc-util.h"
#include "build.h"
#include "conf-files.h"
#include "constants.h"
#include "crypto-util.h"
#include "errno-util.h"
#include "fd-util.h"
#include "fileio.h"
#include "format-table.h"
#include "fs-util.h"
#include "hexdecoct.h"
#include "log.h"
#include "main-func.h"
#include "mkdir.h"
#include "options.h"
#include "parse-util.h"
#include "pretty-print.h"
#include "set.h"
#include "sort-util.h"
#include "string-util.h"
#include "strv.h"
#include "tmpfile-util.h"
#include "tpm2-util.h"

static char *arg_tpm2_device = NULL;
static bool arg_early = false;
static bool arg_graceful = false;

STATIC_DESTRUCTOR_REGISTER(arg_tpm2_device, freep);

#define TPM2_SRK_PEM_PERSISTENT_PATH "/var/lib/systemd/tpm2-srk-public-key.pem"
#define TPM2_SRK_PEM_RUNTIME_PATH "/run/systemd/tpm2-srk-public-key.pem"

#define TPM2_SRK_TPM2B_PUBLIC_PERSISTENT_PATH "/var/lib/systemd/tpm2-srk-public-key.tpm2b_public"
#define TPM2_SRK_TPM2B_PUBLIC_RUNTIME_PATH "/run/systemd/tpm2-srk-public-key.tpm2b_public"

static int help(void) {
        _cleanup_free_ char *link = NULL;
        _cleanup_(table_unrefp) Table *options = NULL;
        int r;

        r = terminal_urlify_man("systemd-tpm2-setup", "8", &link);
        if (r < 0)
                return log_oom();

        r = option_parser_get_help_table(&options);
        if (r < 0)
                return r;

        printf("%s [OPTIONS...]\n"
               "\n%sSet up the TPM2 Storage Root Key (SRK), and initialize NvPCRs.%s\n"
               "\n%sOptions:%s\n",
               program_invocation_short_name,
               ansi_highlight(),
               ansi_normal(),
               ansi_underline(),
               ansi_normal());

        r = table_print_or_warn(options);
        if (r < 0)
                return r;

        printf("\nSee the %s for details.\n", link);
        return 0;
}

static int parse_argv(int argc, char *argv[]) {
        int r;

        assert(argc >= 0);
        assert(argv);

        OptionParser opts = { argc, argv };

        FOREACH_OPTION_OR_RETURN(c, &opts)
                switch (c) {

                OPTION_COMMON_HELP:
                        return help();

                OPTION_COMMON_VERSION:
                        return version();

                OPTION_LONG("tpm2-device", "PATH", "Pick TPM2 device"):
                        if (streq(opts.arg, "list"))
                                return tpm2_list_devices(/* legend= */ true, /* quiet= */ false);

                        if (free_and_strdup(&arg_tpm2_device, streq(opts.arg, "auto") ? NULL : opts.arg) < 0)
                                return log_oom();
                        break;

                OPTION_LONG("early", "BOOL", "Store SRK public key in /run/ rather than /var/lib/"):
                        r = parse_boolean(opts.arg);
                        if (r < 0)
                                return log_error_errno(r, "Failed to parse --early= argument: %s", opts.arg);

                        arg_early = r;
                        break;

                OPTION_LONG("graceful", NULL, "Exit gracefully if no TPM2 device is found"):
                        arg_graceful = true;
                        break;
                }

        if (option_parser_get_n_args(&opts) != 0)
                return log_error_errno(SYNTHETIC_ERRNO(EINVAL), "This program expects no argument.");

        return 1;
}

struct public_key_data {
        EVP_PKEY *pkey;         /* as OpenSSL object */
        TPM2B_PUBLIC *public;   /* in TPM2 format */
        void *fingerprint;
        size_t fingerprint_size;
        char *fingerprint_hex;
        char *path;
};

static void public_key_data_done(struct public_key_data *d) {
        assert(d);

        if (d->pkey) {
                sym_EVP_PKEY_free(d->pkey);
                d->pkey = NULL;
        }
        if (d->public) {
                Esys_Freep(&d->public);
                d->public = NULL;
        }
        d->fingerprint = mfree(d->fingerprint);
        d->fingerprint_size = 0;
        d->fingerprint_hex = mfree(d->fingerprint_hex);
        d->path = mfree(d->path);
}

static int public_key_make_fingerprint(struct public_key_data *d) {
        int r;

        assert(d);
        assert(d->pkey);
        assert(!d->fingerprint);
        assert(!d->fingerprint_hex);

        r = pubkey_fingerprint(d->pkey, sym_EVP_sha256(), &d->fingerprint, &d->fingerprint_size);
        if (r < 0)
                return log_error_errno(r, "Failed to calculate fingerprint of public key: %m");

        d->fingerprint_hex = hexmem(d->fingerprint, d->fingerprint_size);
        if (!d->fingerprint_hex)
                return log_oom();

        return 0;
}

static int load_public_key_disk(const char *path, struct public_key_data *ret) {
        _cleanup_(public_key_data_done) struct public_key_data data = {};
        _cleanup_free_ char *blob = NULL;
        size_t blob_size;
        int r;

        assert(path);
        assert(ret);

        r = read_full_file(path, &blob, &blob_size);
        if (r < 0) {
                if (r != -ENOENT)
                        return log_error_errno(r, "Failed to read '%s': %m", path);

                log_debug("SRK public key file '%s' does not exist.", path);
        } else {
                log_debug("Loaded SRK public key from '%s'.", path);

                r = openssl_pubkey_from_pem(blob, blob_size, &data.pkey);
                if (r < 0)
                        return log_error_errno(r, "Failed to parse SRK public key file '%s': %m", path);

                r = public_key_make_fingerprint(&data);
                if (r < 0)
                        return r;

                log_debug("Loaded SRK public key fingerprint: %s", data.fingerprint_hex);
        }

        data.path = strdup(path);
        if (!data.path)
                return log_oom();

        *ret = data;
        data = (struct public_key_data) {};

        return 0;
}

static int load_public_key_tpm2(struct public_key_data *ret) {
        _cleanup_(public_key_data_done) struct public_key_data data = {};
        _cleanup_(tpm2_context_unrefp) Tpm2Context *c = NULL;
        int r;

        assert(ret);

        r = tpm2_context_new_or_warn(arg_tpm2_device, &c);
        if (r < 0)
                return r;

        r = tpm2_get_or_create_srk(
                        c,
                        /* session= */ NULL,
                        &data.public,
                        /* ret_name= */ NULL,
                        /* ret_qname= */ NULL,
                        NULL);
        if (r == -EDEADLK)
                return r;
        if (r < 0)
                return log_error_errno(r, "Failed to get or create SRK: %m");
        if (r > 0)
                log_info("New SRK generated and stored in the TPM.");
        else
                log_info("SRK already stored in the TPM.");

        r = tpm2_tpm2b_public_to_openssl_pkey(data.public, &data.pkey);
        if (r < 0)
                return log_error_errno(r, "Failed to convert TPM2 SRK public key to OpenSSL public key: %m");

        r = public_key_make_fingerprint(&data);
        if (r < 0)
                return r;

        log_info("SRK fingerprint is %s.", data.fingerprint_hex);

        *ret = data;
        data = (struct public_key_data) {};

        return 0;
}

static int setup_srk(void) {
        int r;

        _cleanup_(public_key_data_done) struct public_key_data runtime_key = {}, persistent_key = {}, tpm2_key = {};

        r = load_public_key_disk(TPM2_SRK_PEM_RUNTIME_PATH, &runtime_key);
        if (r < 0)
                return r;

        if (!arg_early) {
                r = load_public_key_disk(TPM2_SRK_PEM_PERSISTENT_PATH, &persistent_key);
                if (r < 0)
                        return r;

                if (runtime_key.pkey && persistent_key.pkey &&
                    memcmp_nn(runtime_key.fingerprint, runtime_key.fingerprint_size,
                              persistent_key.fingerprint, persistent_key.fingerprint_size) != 0) {

                        /* One of those days we might want to add a stricter policy option here, that refuses
                         * to boot when the SRK changes. For now, let's just warn and proceed, in order not
                         * to break OS images that are moved around PCs. */

                        log_notice("Saved persistent SRK (%s) and runtime SRK differ (fingerprint %s vs. %s), updating persistent SRK.",
                                   persistent_key.path, persistent_key.fingerprint_hex, runtime_key.fingerprint_hex);

                        public_key_data_done(&persistent_key);
                }
        }

        r = load_public_key_tpm2(&tpm2_key);
        if (r == -EDEADLK) {
                log_struct_errno(LOG_INFO, r,
                                 LOG_MESSAGE("Insufficient permissions to access TPM, not generating SRK."),
                                 LOG_MESSAGE_ID(SD_MESSAGE_SRK_ENROLLMENT_NEEDS_AUTHORIZATION_STR));
                return EX_PROTOCOL; /* Special return value which means "Insufficient permissions to access TPM,
                                     * cannot generate SRK". This isn't really an error when called at boot. */;
        }
        if (r < 0)
                return r;

        assert(tpm2_key.pkey);

        if (runtime_key.pkey) {
                if (memcmp_nn(tpm2_key.fingerprint, tpm2_key.fingerprint_size,
                              runtime_key.fingerprint, runtime_key.fingerprint_size) != 0)
                        return log_error_errno(SYNTHETIC_ERRNO(ENOTRECOVERABLE),
                                               "Saved runtime SRK differs from TPM SRK, refusing.");

                if (arg_early) {
                        log_info("SRK saved in '%s' matches SRK in TPM2.", runtime_key.path);
                        return 0;
                }
        }

        if (persistent_key.pkey) {
                if (memcmp_nn(tpm2_key.fingerprint, tpm2_key.fingerprint_size,
                              persistent_key.fingerprint, persistent_key.fingerprint_size) == 0) {
                        log_info("SRK saved in '%s' matches SRK in TPM2.", persistent_key.path);
                        return 0;
                }

                /* As above, we probably want a stricter policy option here, one day. */

                log_notice("Saved persistent SRK (%s) and TPM SRK differ (fingerprint %s vs. %s), updating persistent SRK.",
                           persistent_key.path, persistent_key.fingerprint_hex, tpm2_key.fingerprint_hex);

                public_key_data_done(&persistent_key);
        }

        const char *pem_path = arg_early ? TPM2_SRK_PEM_RUNTIME_PATH : TPM2_SRK_PEM_PERSISTENT_PATH;
        (void) mkdir_parents(pem_path, 0755);

        /* Write out public key (note that we only do that as a help to the user, we don't make use of this ever */
        _cleanup_(unlink_and_freep) char *t = NULL;
        _cleanup_fclose_ FILE *f = NULL;
        r = fopen_tmpfile_linkable(pem_path, O_WRONLY|O_CLOEXEC, &t, &f);
        if (r < 0)
                return log_error_errno(r, "Failed to open SRK public key file '%s' for writing: %m", pem_path);

        if (sym_PEM_write_PUBKEY(f, tpm2_key.pkey) <= 0)
                return log_error_errno(SYNTHETIC_ERRNO(EIO), "Failed to write SRK public key file '%s'.", pem_path);

        if (fchmod(fileno(f), 0444) < 0)
                return log_error_errno(errno, "Failed to adjust access mode of SRK public key file '%s' to 0444: %m", pem_path);

        r = flink_tmpfile(f, t, pem_path, LINK_TMPFILE_SYNC|LINK_TMPFILE_REPLACE);
        if (r < 0)
                return log_error_errno(r, "Failed to move SRK public key file to '%s': %m", pem_path);

        f = safe_fclose(f);
        t = mfree(t);

        log_info("SRK public key saved to '%s' in PEM format.", pem_path);

        const char *tpm2b_public_path = arg_early ? TPM2_SRK_TPM2B_PUBLIC_RUNTIME_PATH : TPM2_SRK_TPM2B_PUBLIC_PERSISTENT_PATH;
        (void) mkdir_parents(tpm2b_public_path, 0755);

        /* Now also write this out in TPM2B_PUBLIC format */
        r = fopen_tmpfile_linkable(tpm2b_public_path, O_WRONLY|O_CLOEXEC, &t, &f);
        if (r < 0)
                return log_error_errno(r, "Failed to open SRK public key file '%s' for writing: %m", tpm2b_public_path);

        _cleanup_free_ void *marshalled = NULL;
        size_t marshalled_size = 0;
        r = tpm2_marshal_public(tpm2_key.public, &marshalled, &marshalled_size);
        if (r < 0)
                return log_error_errno(r, "Failed to marshal TPM2_PUBLIC key.");

        if (fwrite(marshalled, 1, marshalled_size, f) != marshalled_size)
                return log_error_errno(SYNTHETIC_ERRNO(EIO),
                                       "Failed to write SRK public key file '%s'.", tpm2b_public_path);

        if (fchmod(fileno(f), 0444) < 0)
                return log_error_errno(errno, "Failed to adjust access mode of SRK public key file '%s' to 0444: %m", tpm2b_public_path);

        r = flink_tmpfile(f, t, tpm2b_public_path, LINK_TMPFILE_SYNC|LINK_TMPFILE_REPLACE);
        if (r < 0)
                return log_error_errno(r, "Failed to move SRK public key file to '%s': %m", tpm2b_public_path);

        t = mfree(t);

        log_info("SRK public key saved to '%s' in TPM2B_PUBLIC format.", tpm2b_public_path);
        return 0;
}

typedef struct SetupNvPCRContext {
        Tpm2Context *tpm2_context;
        struct iovec anchor_secret;
        size_t n_already, n_anchored, n_failed, n_skipped;
        bool nv_space_exhausted; /* Set once the TPM ran out of NV index space, so we skip the rest. */
        Set *done;
} SetupNvPCRContext;

static void setup_nvpcr_context_done(SetupNvPCRContext *c) {
        assert(c);

        iovec_done_erase(&c->anchor_secret);
        c->tpm2_context = tpm2_context_unref(c->tpm2_context);
        c->done = set_free(c->done);
}

static int setup_nvpcr_one(
                SetupNvPCRContext *c,
                const char *name) {
        int r;

        assert(c);
        assert(name);
        assert(c->tpm2_context);

        if (set_contains(c->done, name))
                return 0;

        /* If the TPM already ran out of NV index space, don't even try to allocate further NvPCRs. As
         * NvPCRs are processed in order of priority (most important first), everything still pending is
         * less important than what already didn't fit. */
        if (c->nv_space_exhausted) {
                log_debug("TPM NV index space already exhausted, skipping allocation of NvPCR '%s'.", name);
                c->n_skipped++;
                return 0;
        }

        r = tpm2_nvpcr_initialize(c->tpm2_context, /* session= */ NULL, name, &c->anchor_secret);
        if (r == -EUNATCH) {
                assert(!iovec_is_set(&c->anchor_secret));

                /* If we get EUNATCH this means we actually need to initialize this NvPCR
                 * now, and haven't provided the anchor secret yet. Hence acquire it now. */

                r = tpm2_nvpcr_acquire_anchor_secret(&c->anchor_secret, /* sync_secondary= */ !arg_early);
                if (r < 0)
                        return log_error_errno(r, "Failed to acquire anchor secret: %m");

                r = tpm2_nvpcr_initialize(c->tpm2_context, /* session= */ NULL, name, &c->anchor_secret);
        }
        if (r == -EOPNOTSUPP) {
                c->n_failed++;
                return log_struct_errno(LOG_ERR, r,
                                        LOG_MESSAGE("The TPM does not correctly support NV indexes in NT_EXTEND mode, unable to allocate NvPCR '%s': %m", name),
                                        LOG_MESSAGE_ID(SD_MESSAGE_TPM_NVPCR_UNSUPPORTED_STR));
        }
        if (r == -ENOSPC) {
                /* The TPM's NV index space is exhausted. Remember this so we skip the remaining (less
                 * important) NvPCRs, and report it gracefully at the end rather than failing the boot.
                 * Logged at notice level, not error. */
                c->nv_space_exhausted = true;
                c->n_skipped++;
                return log_struct_errno(LOG_NOTICE, r,
                                        LOG_MESSAGE("The TPM's NV index space is exhausted, skipping allocation of NvPCR '%s' and any less important ones: %m", name),
                                        LOG_MESSAGE_ID(SD_MESSAGE_TPM_NVINDEX_EXHAUSTED_STR));
        }
        if (r < 0) {
                c->n_failed++;
                return log_error_errno(r, "Failed to extend NvPCR index with anchor secret: %m");
        }

        if (r > 0)
                c->n_anchored++;
        else
                c->n_already++;

        if (set_put_strdup(&c->done, name) < 0)
                return log_oom();

        return 0;
}

typedef struct NvPCREntry {
        const char *name;       /* points into the strv 'l' in setup_nvpcr() */
        uint64_t priority;
} NvPCREntry;

static int nvpcr_entry_compare(const NvPCREntry *a, const NvPCREntry *b) {
        int r;

        assert(a);
        assert(b);

        /* Lower priority value means more important, hence allocate it first. Break ties by name
         * (ascending) so the resulting order is fully deterministic. */
        r = CMP(a->priority, b->priority);
        if (r != 0)
                return r;

        return strcmp(a->name, b->name);
}

static int setup_nvpcr(void) {
        _cleanup_(setup_nvpcr_context_done) SetupNvPCRContext c = {};
        int r;

        _cleanup_strv_free_ char **l = NULL;
        r = conf_files_list_nulstr(
                        &l,
                        ".nvpcr",
                        /* root= */ NULL,
                        CONF_FILES_REGULAR|CONF_FILES_BASENAME|CONF_FILES_FILTER_MASKED|CONF_FILES_TRUNCATE_SUFFIX|CONF_FILES_WARN,
                        CONF_PATHS_NULSTR("nvpcr"));
        if (r < 0)
                return log_error_errno(r, "Failed to find .nvpcr files: %m");

        /* Pair each NvPCR name with its allocation priority, then sort, so that we allocate the most
         * important NvPCRs first. This matters when the TPM's NV index space is too small to fit all of
         * them: we then degrade gracefully, skipping the least important NvPCRs. */
        size_t n = strv_length(l);
        _cleanup_free_ NvPCREntry *entries = new(NvPCREntry, n);
        if (!entries)
                return log_oom();

        for (size_t i = 0; i < n; i++) {
                uint64_t priority = TPM2_NVPCR_PRIORITY_DEFAULT;

                r = tpm2_nvpcr_get_index(l[i], /* ret_nv_index= */ NULL, &priority);
                if (r < 0)
                        /* If we can't read the definition, assume the default priority and let the
                         * per-NvPCR error path in setup_nvpcr_one() report the actual problem. */
                        log_warning_errno(r, "Failed to read priority of NvPCR '%s', assuming default priority %" PRIu64 ": %m", l[i], priority);

                entries[i] = (NvPCREntry) {
                        .name = l[i],
                        .priority = priority,
                };
        }

        typesafe_qsort(entries, n, nvpcr_entry_compare);

        int ret = 0;
        FOREACH_ARRAY(e, entries, n) {
                if (!c.tpm2_context) {
                        /* Inability to contact the TPM shall be fatal for us */
                        r = tpm2_context_new_or_warn(arg_tpm2_device, &c.tpm2_context);
                        if (r < 0)
                                return r;
                }

                log_debug("Setting up NvPCR '%s' (priority %" PRIu64 ").", e->name, e->priority);

                /* But if we fail to initialize some NvPCR, we go on */
                RET_GATHER(ret, setup_nvpcr_one(&c, e->name));
        }

        if (c.n_already > 0 && c.n_anchored == 0 && !arg_early)
                /* If we didn't anchor anything right now, but we anchored something earlier, then it might
                 * have happened in the initrd, and thus the anchor ID was not committed to /var/ or the ESP
                 * yet. Hence, let's explicitly do so now, to catch up. */
                RET_GATHER(ret, tpm2_nvpcr_acquire_anchor_secret(/* ret= */ NULL, /* sync_secondary= */ true));

        if (c.nv_space_exhausted)
                log_notice("Skipped %zu lowest-priority NvPCR(s) because the TPM's NV index space is exhausted, proceeding anyway.", c.n_skipped);

        if (c.n_failed > 0)
                log_warning("%zu NvPCRs failed to initialize, proceeding anyway.", c.n_failed);

        if (c.n_anchored > 0) {
                if (c.n_already == 0)
                        log_info("%zu NvPCRs initialized.", c.n_anchored);
                else
                        log_info("%zu NvPCRs initialized. (%zu NvPCRs were already initialized.)", c.n_anchored, c.n_already);
        } else if (c.n_already > 0)
                log_info("%zu NvPCRs already initialized.", c.n_already);
        else if (c.n_failed == 0 && c.n_skipped == 0)
                log_debug("No NvPCRs defined, nothing initialized.");

        /* Turn some errors into recognizable ones, which we can catch with
         * SuccessExitStatus= in the service unit file. */
        if (ret == -EOPNOTSUPP)
                return EX_UNAVAILABLE;   /* e.g. no NvPCR support in TPM */
        if (ret == -ENOSPC)
                return EX_CANTCREAT;     /* NV index space on TPM exhausted */

        return ret;
}

static int run(int argc, char *argv[]) {
        int r;

        log_setup();

        r = parse_argv(argc, argv);
        if (r <= 0)
                return r;

        if (arg_graceful && !tpm2_is_mostly_supported()) {
                log_notice("No complete TPM2 support detected, exiting gracefully.");
                return EXIT_SUCCESS;
        }

        r = dlopen_libcrypto(LOG_ERR);
        if (r < 0)
                return r;

        umask(0022);

        /* Execute both jobs, and then return unlisted errors preferably, and listed errors
         * (i.e. EX_UNAVAILABLE, EX_CANTCREAT, EX_PROTOCOL) otherwise. */
        r = setup_srk();
        int k = setup_nvpcr();
        if (r < 0)
                return r;
        if (k < 0)
                return k;
        return r != EXIT_SUCCESS ? r : k;
}

DEFINE_MAIN_FUNCTION_WITH_POSITIVE_FAILURE(run);
