systemd-cryptenroll-tpm.py

#!/usr/bin/env python3

"""
Automatically enroll systemd TPM unlock feature on all LUKS partitions.
"""

import sys
import json
import subprocess


TPM2_PCRS = "1+7"
WITH_PIN = True


class SystemdCryptEnrollScript:
    """
    Wrapper script to analyze all partitions, wipe old TPM keyslots and enroll new slots.
    """

    def __init__(self):
        self._crypt_fs_devs: list[str] = []

    def run(self):
        self._crypt_fs_devs = self.find_crypt_devices()
        for dev in self._crypt_fs_devs:
            self.setup(dev)

    def find_crypt_devices(self):
        info_str = subprocess.check_output(
            ["lsblk", "--fs", "--json"], encoding="utf-8"
        )
        info = json.loads(info_str)
        ret: list[str] = []
        for dev in info["blockdevices"]:
            for partition in dev.get("children", []):
                if partition["fstype"] == "crypto_LUKS":
                    ret.append(partition["name"])

        return ret

    def get_systemd_tpm2_slots(self, dev: str) -> str | None:
        info_str = subprocess.check_output(
            ["cryptsetup", "luksDump", "--dump-json-metadata", f"/dev/{dev}"],
            encoding="utf-8",
        )
        info = json.loads(info_str)
        tokens = info.get("tokens", {})
        slot_to_wipe = None
        for name, token in tokens.items():
            if token["type"] == "systemd-tpm2":
                if slot_to_wipe is not None:
                    print(
                        f"Found multiple slots to wipe for device {dev}: {name} and {slot_to_wipe}",
                        file=sys.stderr,
                    )
                    print("Exiting for safety reasons.", file=sys.stderr)
                    sys.exit(1)
                slots_to_wipe = token["keyslots"]
                if len(slots_to_wipe) != 1:
                    print(
                        f"Token {name} for device {dev} has {len(slots_to_wipe)} slots {', '.join(slots_to_wipe)}. Expected 1.",
                        file=sys.stderr,
                    )
                    print("Exiting for safety reasons.", file=sys.stderr)
                    sys.exit(1)

                slot_to_wipe = slots_to_wipe[0]

        if slot_to_wipe is not None:
            keyslots = info["keyslots"]
            assert slot_to_wipe in keyslots.keys()
            if len(keyslots) <= 1:
                print(
                    f"Found 1 slot to wipe for device {dev}: {slot_to_wipe}. However, this is the only keyslot for the device.",
                    file=sys.stderr,
                )
                print("Exiting for safety reasons.", file=sys.stderr)
                sys.exit(1)

        return slot_to_wipe

    def wipe_slot(self, dev: str, slot: str) -> None:
        print(
            f"Are you sure, that you want to wipe the systemd-tpm2 slot {slot} for device {dev}?"
        )
        print("We did some checks and it should be safe, but you never know.")
        try:
            while True:
                ret = input("[Y/n] ").strip().lower()
                if not ret or ret == "y":
                    break
                else:
                    sys.exit(0)
        except KeyboardInterrupt:
            print(file=sys.stderr)
            sys.exit(0)

        subprocess.check_call(
            ["systemd-cryptenroll", "--wipe-slot", slot, f"/dev/{dev}"]
        )

    def wipe_old(self, dev: str):
        slot_to_wipe = self.get_systemd_tpm2_slots(dev)
        if slot_to_wipe is not None:
            self.wipe_slot(dev, slot_to_wipe)

    def setup(self, dev: str) -> None:
        print(f"Setting up device {dev}")
        self.wipe_old(dev)
        subprocess.check_call(
            [
                "systemd-cryptenroll",
                f"--tpm2-pcrs={TPM2_PCRS}",
                f"--tpm2-with-pin={'true' if WITH_PIN else 'false'}",
                "--tpm2-device=auto",
                f"/dev/{dev}",
            ]
        )


def main():
    script = SystemdCryptEnrollScript()
    script.run()


if __name__ == "__main__":
    main()