nixfiles/modules/nixos/encryptedDrive.nix

{
  self,
  config,
  pkgs,
  lib,
  ...
}: let
  inherit (builtins) listToAttrs getAttr;
  inherit (lib.modules) mkIf mkMerge mkBefore;
  inherit (lib.options) mkOption mkEnableOption;
  inherit (lib.strings) optionalString;
  inherit (lib.attrsets) attrValues;
  inherit (lib) types;

  encryptedUSB = import "${self}/data/drives/encryptedUSB.nix";
  driveData = import "${self}/data/drives/encryptedDrive.nix";

  cfg = config.boot.encryptedDrive;
in {
  options.boot.encryptedDrive = {
    enable = mkEnableOption "encrypted drive support for servers and other devices";

    mode = mkOption {
      type = types.enum [
        "encrypted-usb"
        "password"
      ];
      default = "encrypted-usb";
    };

    allowPasswordDecrypt = mkOption {
      description = "allow fallback to decrypting with a password when using USB based auth, pass cryptsetup_password to kernel cmdline to switch to password input mode";
      type = types.bool;
      default = true;
    };
  };

  config = mkIf cfg.enable {
    boot = {
      initrd.availableKernelModules = [
        # For USB w/ Encryption Key
        "usb_storage"
        "usbcore"
        "uas"
        "rtsx_pci_sdmmc"
        # For USB Keyboards
        "usbhid"
        "hid_generic"
        # For Cryptography
        "aesni_intel"
        "cryptd"
        "crypto_simd"
      ];
      initrd.postDeviceCommands = mkIf (cfg.mode == "encrypted-usb") (mkBefore ''
        mkdir -m 0755 -p /keys
        mkdir -m 0755 -p ${encryptedUSB.mountpoint}

        ${
          if cfg.allowPasswordDecrypt
          then ''
            if grep "cryptsetup_password" /proc/cmdline; then
              USE_PASSWORD_FALLBACK=true
            else
              USE_PASSWORD_FALLBACK=false
            fi
          ''
          else ''
            USE_PASSWORD_FALLBACK=false
          ''
        }

        while !(test -b ${encryptedUSB.encryptedPath}) && [ "$USE_PASSWORD_FALLBACK" == "false" ]
        do
          ${
          if config.boot.plymouth.enable
          then ''
            ${pkgs.plymouth}/bin/plymouth display-message --text="Please Plug In USB"
          ''
          else ''
            echo "Please Plug In USB"
          ''
        }
          sleep 1
        done

        ${
          if config.boot.plymouth.enable
          then ''
            ${pkgs.plymouth}/bin/plymouth hide-message --text="Please Plug In USB"

            if [ "$USE_PASSWORD_FALLBACK" == "true" ]; then
              ${pkgs.plymouth}/bin/plymouth ask-for-password \
                --prompt="Please Enter Password" \
                --command="cryptsetup -T1 open ${driveData.encryptedPath} ${driveData.mapperName}" \
                --number-of-tries=3
            else
              ${pkgs.plymouth}/bin/plymouth ask-for-password \
                --prompt="Please Decrypt USB" \
                --command="cryptsetup -T1 open ${encryptedUSB.encryptedPath} ${encryptedUSB.preBootMapperName}" \
                --number-of-tries=3
            fi
          ''
          else ''
            if [ "$USE_PASSWORD_FALLBACK" == "true" ]; then
              echo "Please Decrypt Drive"
              cryptsetup open ${driveData.encryptedPath} ${driveData.mapperName}
            else
              echo "Please Decrypt USB"
              cryptsetup open ${encryptedUSB.encryptedPath} ${encryptedUSB.preBootMapperName}
            fi
          ''
        }

        if [ "$USE_PASSWORD_FALLBACK" == "false" ]; then
          mount -n -t ${encryptedUSB.unencryptedFSType} -o ro ${encryptedUSB.preBootMapperPath} ${encryptedUSB.mountpoint}

          cp ${encryptedUSB.encryptionKeysPath}/${config.networking.hostName}.key /keys
          chmod 0755 /keys/${config.networking.hostName}.key
          umount -f ${encryptedUSB.mountpoint}

          cryptsetup close ${encryptedUSB.preBootMapperName}
        else
          touch /keys/${config.networking.hostName}.key
        fi
      '');
      initrd.luks.devices = {
        "${driveData.mapperName}" = {
          device = "${driveData.encryptedPath}";
          keyFile =
            if cfg.mode == "encrypted-usb"
            then "/keys/${config.networking.hostName}.key"
            else null;
          preLVM = false;
          allowDiscards = true;
          # Allows decrypting with a password when key is existant on USB but invalid
          fallbackToPassword = cfg.allowPasswordDecrypt;
        };
      };
    };

    fileSystems = {
      "/" = {
        device = "${driveData.decryptedPath}";
        fsType = "${driveData.unencryptedFSType}";
      };
      "/boot" = {
        device = "${driveData.bootPath}";
        fsType = "${driveData.bootFSType}";
      };
    };
  };
}