#!/usr/bin/ruby
#
# Setup local filesystems for Puavo hosts

require 'fileutils'
require 'getoptlong'
require 'highline/import'
require 'json'
require 'open3'
require 'pathname'
require 'securerandom'
require 'tempfile'
require 'tmpdir'

GPT_PARTTYPE_EFI   = 'c12a7328-f81f-11d2-ba4b-00a0c93ec93b'
GPT_PARTTYPE_LINUX = '4f68bce3-e8cd-4db1-96e7-fbcaf984b709'
GPT_PARTTYPE_RAID  = 'a19d880f-05fc-4d3b-a006-743f0f84911e'
GPT_PARTTYPE_SWAP  = '0657fd6d-a4ab-43c4-84e5-0933c84b4f4f'

GPT_PARTLABEL_PUAVO_INSTALLED_WINDOWS_C = 'Puavo_Installed_Windows_C'

# Why 20G: https://learn.microsoft.com/en-us/windows-hardware/manufacture/desktop/configure-uefigpt-based-hard-drive-partitions?view=windows-11#windows-partition
DUALBOOT_LAPTOP_WINDOWS_MIN_SIZE_G = 20
DUALBOOT_LAPTOP_PUAVOOS_MIN_SIZE_G = 70 # ~60G + 10G home
DUALBOOT_LAPTOP_DISK_MIN_SIZE_G = DUALBOOT_LAPTOP_WINDOWS_MIN_SIZE_G + DUALBOOT_LAPTOP_PUAVOOS_MIN_SIZE_G

def run(*cmd)
  system(*cmd) or raise "Error running command: '#{ cmd.inspect }'"
end

def run_and_get_lines(*cmd)
  stdout, stderr, status = Open3.capture3(*cmd)
  raise "Command failed: #{ cmd.join(' ') }: #{ stderr.strip }" \
    unless status.success?
  stdout.lines(chomp: true)
end

def tpm2_available?
  system(
    'grep', '-q', '2',
    '/sys/class/tpm/tpm0/tpm_version_major'
  )
end

# XXX this should be elsewhere, other scripts might find this useful
module PuavoFacts
  Puavodir = '/etc/puavo'

  FactPaths = {
                'kerberos_master'   => 'kerberos/master',
                'kerberos_realm'    => 'kerberos/realm',
                'kerberos_toprealm' => 'kerberos/toprealm',
                'ldap_base'         => 'ldap/base',
                'ldap_binddn'       => 'ldap/dn',
                'ldap_bindpw'       => 'ldap/password',
                'ldap_master'       => 'ldap/master',
                'puavo_domain'      => 'domain',
                'puavo_hostname'    => 'hostname',
                'puavo_hosttype'    => 'hosttype',
                'puavo_topdomain'   => 'topdomain',
              }

  def self.all
    Hash[ FactPaths.map { |name, path| [ name, get(name) ] } ]
  end

  def self.get(name)
    File.read("#{ Puavodir }/#{ FactPaths[ name ] }").chomp \
      or raise "Could not read a fact '#{ name }'"
  end
end

class LuksEnrollmentPolicy
  # The format of these arguments follow the systemd-cryptenroll specifications.
  # See "man systemd-cryptenroll" or "https://www.freedesktop.org/software/systemd/man/latest/systemd-cryptenroll.html".
  attr_accessor :specific_pcrs, :public_key_pcrs, :public_key_path, :specific_arguments

  def initialize(specific_pcrs, public_key_pcrs, public_key_path, specific_arguments=[])
    @specific_pcrs = specific_pcrs
    @public_key_pcrs = public_key_pcrs
    @public_key_path = public_key_path
    @specific_arguments = specific_arguments
  end

  def arguments
    specific_pcrs_option_value = @specific_pcrs.join("+")
    public_key_pcrs_option_value = @public_key_pcrs.join("+")

    base_arguments = [
      '--tpm2-device=auto',
      "--tpm2-pcrs=#{ specific_pcrs_option_value }",
      "--tpm2-public-key-pcrs=#{ public_key_pcrs_option_value }",
      "--tpm2-public-key=#{ public_key_path }"
    ]

    base_arguments + @specific_arguments
  end
end

class LuksPartitionHandler
  attr_accessor :partition, :encryption_device_name, :policy

  LUKS_FORMAT_OPTIONS = [
    '--type=luks2',
    '--cipher=aes-xts-plain64',
    '--key-size=256',
    '--pbkdf=argon2id',
    '--iter-time=1000',
    '--batch-mode'
  ].freeze

  def initialize(partition, encryption_device_name, policy)
    @partition = partition
    @encryption_device_name = encryption_device_name
    @policy = policy
  end

  def encryption_device()
    "/dev/mapper/#{ @encryption_device_name }"
  end

  def generate_control_key()
    # Generate a control key as a hex string for UTF-8 safety.
    SecureRandom.hex(32) # 256 bits of entropy, 64 hex characters
  end

  def enroll(control_key)
    require 'tempfile'

    Tempfile.create('puavo-luks-enroll-key', '/tmp') do |control_key_file|
      control_key_file.chmod(0600)
      control_key_file.write(control_key.to_s)
      control_key_file.flush

      command = [ 'systemd-cryptenroll' ]
      command += @policy.arguments
      command += [ "--unlock-key-file=#{ control_key_file.path }" ]
      command += [ "/dev/#{ @partition }" ]

      ok = system(*command)
      ok or raise 'Failed to enroll LUKS encryption device'
    end
  end

  def setup(predefined_control_key=nil)
    control_key = predefined_control_key || self.generate_control_key()

    puts 'Creating LUKS partition...'
    IO.popen([ 'cryptsetup', 'luksFormat',
               *LUKS_FORMAT_OPTIONS,
               "--key-file=-",
               "/dev/#{ @partition }" ], 'w') do |f|
      f.print(control_key)
    end

    exit_status = $?
    unless exit_status.success?
      raise 'Failed to format LUKS encryption device'
    end

    puts 'Opening the LUKS partition for installing the system...'
    IO.popen([ 'cryptsetup', 'luksOpen',
               "/dev/#{ @partition }",
               @encryption_device_name ], 'w') do |f|
      f.print(control_key)
    end

    exit_status = $?
    unless exit_status.success?
      raise 'Failed to open the created LUKS encryption device'
    end

    if @policy
      puts 'Enrolling the LUKS partition for TPM based unlocking...'
      self.enroll(control_key)
    end

    return control_key
  end

  def add_key(control_key, new_key)
    return if new_key.nil? || new_key.empty?

    Tempfile.create('puavo-luks-control-key-') do |control_key_file|
      control_key_file.write(control_key)
      control_key_file.flush
      File.chmod(0600, control_key_file.path)

      Tempfile.create('puavo-luks-new-key-') do |new_key_file|
        new_key_file.write(new_key)
        new_key_file.flush
        File.chmod(0600, new_key_file.path)

        run('cryptsetup', 'luksAddKey',
            "--key-file=#{ control_key_file.path }",
            "--new-keyfile=#{ new_key_file.path }",
            "/dev/#{ @partition }")
      end
    end
  end
end

class BootVaultBuilder
  DEFAULT_SIZE        = '32M'
  DEFAULT_IMAGE_NAME  = 'vault.img'
  DEFAULT_LUKS_NAME   = 'puavo-boot-vault'

  CONTROL_KEY_PROPERTY = 'recovery.key'
  VERSION_PROPERTY      = 'version'

  VERSION = 1

  def initialize(control_key:, recovery_key:, encryption_policy:, device:, size: DEFAULT_SIZE, image_name: DEFAULT_IMAGE_NAME, luks_name: DEFAULT_LUKS_NAME)
    @control_key = control_key
    @recovery_key = recovery_key
    @encryption_policy = encryption_policy
    @size = size
    @image_name = image_name
    @luks_name = luks_name
    @device = device
  end

  def build(device_keys_temporary_directory)
    Dir.mktmpdir('puavo-boot-vault-build-') do |build_directory|
      image_path = create_image(build_directory)
      setup_luks_container(image_path, device_keys_temporary_directory)
      move_to_efi(image_path)
    end
  end

  private

  def create_image(build_directory)
    image_path = File.join(build_directory, @image_name)
    run('truncate', '--size', @size, image_path)
    File.chmod(0600, image_path)
    image_path
  end

  def setup_luks_container(image_path, device_keys_temporary_directory)
    loop_device = IO.popen([ 'losetup', '--find', '--show', image_path ]).read.strip
    raise "Failed to allocate loop device" if loop_device.empty?
    begin
      handler = LuksPartitionHandler.new(
        File.basename(loop_device),
        @luks_name,
        @encryption_policy
      )
      handler.setup(@control_key)

      # If the recovery key is set, add it
      if @recovery_key && !@recovery_key.empty?
        handler.add_key(@control_key, @recovery_key)
      end

      run('mkfs.ext4', '-q', handler.encryption_device)

      Dir.mktmpdir('puavo-boot-vault-') do |mountpoint|
        begin
          run('mount', handler.encryption_device, mountpoint)

          control_key_property = File.join(mountpoint, CONTROL_KEY_PROPERTY)
          File.write(control_key_property, @control_key)
          File.chmod(0600, control_key_property)

          version_property = File.join(mountpoint, VERSION_PROPERTY)
          File.write(version_property, VERSION.to_s)

          # Copy all device keys (Secure Boot keys, TPM keys, etc.) into the vault
          copy_device_keys(mountpoint, device_keys_temporary_directory)
        ensure
          run('umount', mountpoint)
        end
      end
    ensure
      run('cryptsetup', 'close', @luks_name)
      run('losetup', '--detach', loop_device)
    end
  end

  def move_to_efi(image_path)
    efi_partition = DiskHandler.wait_one_gpt_partition_by_parttype(
      @device,
      GPT_PARTTYPE_EFI,
      try_count: 30
    )
    Dir.mktmpdir('puavo-efi-mount-') do |mountpoint|
      begin
        run('mount', "/dev/#{ efi_partition }", mountpoint)
        FileUtils.mkdir_p File.join(mountpoint, "EFI", "puavo")
        FileUtils.mv(
          image_path,
          File.join(mountpoint, "EFI", "puavo", @image_name))
        run('sync')
      ensure
        run('umount', mountpoint)
      end
    end
  end

  def copy_device_keys(mountpoint, device_keys_temporary_directory)
    # Copy all device key files from the temporary directory to the vault
    FileUtils.cp_r(Pathname(device_keys_temporary_directory).children, mountpoint)
  end
end

module DiskCleaner
  def self.active_swap_paths
    run_and_get_lines('swapon', '--show=NAME', '--noheadings', '--raw')
  end

  def self.block_device_number(device_path)
    stat = File.stat(device_path)
    raise "'#{ device_path }' is not a block device" unless stat.blockdev?
    "#{ stat.rdev_major }:#{ stat.rdev_minor }"
  end

  def self.device_dependency_numbers(device_path)
    run_and_get_lines('lsblk', '--inverse', '--list', '--noheadings',
                      '--output', 'MAJ:MIN', device_path) \
      .map(&:strip)
  end

  def self.swap_backing_device(swap_path)
    return swap_path if File.blockdev?(swap_path)

    # The swap is a file. Find the block device containing its filesystem.
    sources = run_and_get_lines('findmnt', '--first-only', '--noheadings',
                '--evaluate', '--nofsroot', '--output', 'SOURCE', '--target',
                swap_path)
    source = sources.first

    raise "Could not determine block device backing swap '#{ swap_path }'" \
      unless source && File.blockdev?(source)

    source
  end

  def self.swap_uses_device?(swap_path, target_number)
    backing_device = swap_backing_device(swap_path)
    device_dependency_numbers(backing_device).include?(target_number)
  end

  def self.active_swaps_on_device(device_path)
    target_number = block_device_number(device_path)
    active_swap_paths.select do |swap_path|
      swap_uses_device?(swap_path, target_number)
    end
  end

  def self.deactivate_swaps_on(device_path)
    active_swaps_on_device(device_path).each do |swap_path|
      run('swapoff', swap_path)
    end

    remaining = active_swaps_on_device(device_path)
    return if remaining.empty?

    raise "Swap remains active on '#{ device_path }': " \
          "#{ remaining.join(', ') }"
  end

  def self.deactivate_holders(device_paths)
    blkdeactivate_opts = [ '--dmoptions',  'retry', '--errors',
                           '--lvmoptions', 'retry', '--umount' ]

    # check three times
    3.times do
      holders = device_holders(device_paths)
      return if holders.empty?

      holders.each do |holder|
        # blkdeactivate may have made another holder disappear, so check
        next unless File.blockdev?(holder)
        run('blkdeactivate', *blkdeactivate_opts, holder)
      end

      run('udevadm', 'settle')
    end

    remaining_holders = device_holders(device_paths)
    return if remaining_holders.empty?

    raise "Could not deactivate holders of #{ device_paths.join(', ') }: " \
          "#{ remaining_holders.join(', ') }"
  end

  def self.device_holders_for_dev(device_path)
    begin
      stat = File.stat(device_path)
      raise "'#{ device_path }' is not a block device" unless stat.blockdev?
      holders_dir =
        "/sys/dev/block/#{ stat.rdev_major }:#{ stat.rdev_minor }/holders"
      Dir.children(holders_dir).map { |kernel_name| "/dev/#{ kernel_name }" }
    rescue Errno::ENOENT
      # we do not need to consider devices that disappear due to deactivations
      []
    end
  end

  def self.device_holders(device_paths)
    device_paths.map { |dev| device_holders_for_dev(dev) }.flatten.uniq
  end

  def self.reset_disk(device)
    device_path = "/dev/#{ device }"

    puts "Resetting disk device #{ device_path }..."

    run('udevadm', 'settle')

    old_partition_paths \
      = DiskHandler::get_partitions(device).map do |partition|
          "/dev/#{ partition }"
        end

    deactivate_swaps_on(device_path)

    old_device_paths = [ device_path, *old_partition_paths ]
    unmount_filesystems_on(old_device_paths)

    # free the devices (and subdevices)
    deactivate_holders(old_device_paths)
    check_no_mountpoints(old_device_paths)

    still_existing_partition_paths \
      = old_partition_paths.select { |path| File.exist?(path) }
    run('wipefs', '--all', '--force', '--lock=yes',
        *still_existing_partition_paths) \
      unless still_existing_partition_paths.empty?

    # Erase all signatures and both GPT headers from the device.
    run('sgdisk', '--zap-all', device_path) rescue true # may fail but redundant
    run('wipefs', '--all', '--lock=yes', device_path)

    DiskHandler::reread_partition_table(device)

    remaining_partitions = DiskHandler::get_partitions(device)
    raise "Partitions remain on '#{ device_path }' after reset: " \
          "#{ remaining_partitions.join(', ') }" \
      unless remaining_partitions.empty?

    puts 'Reset done.'
  end

  def self.check_no_mountpoints(device_paths)
    mountpoints = mountpoints_for_devices(device_paths)
    return if mountpoints.empty?
    raise "Devices are still mounted at: #{ mountpoints.join(', ') }"
  end

  def self.mountpoints_for_device(device_path)
    stdout, stderr, status = Open3.capture3(
      'findmnt', '--json', '--list', '--output', 'TARGET',
      '--source', device_path
    )

    # findmnt uses exit status 1 when nothing matched
    return [] if status.exitstatus == 1 && stdout.empty? && stderr.empty?

    unless status.success?
      raise "Could not look up mounts for '#{ device_path }': #{ stderr.strip }"
    end

    JSON.parse(stdout)
      .fetch('filesystems', [])
      .map { |filesystem| filesystem.fetch('target') }
  end

  def self.mountpoints_for_devices(device_paths)
    device_paths.map do |device_path|
      mountpoints_for_device(device_path)
    end.flatten.uniq
  end

  def self.unmount_filesystems_on(device_paths)
    mountpoints = mountpoints_for_devices(device_paths)
    # Unmount nested mountpoints first.
    mountpoints.sort_by { |mountpoint| -mountpoint.length }.each do |mountpoint|
      run('umount', mountpoint)
    end
  end
end

module DiskHandler
  Filesystems = {
    'bootserver' => {
      '.EFI'                  => { 'size' => '3072M', 'type' => 'fat'     },
      '.SWAP'                 => { 'size' => '16G',   'type' => 'swap'    },
      'home'                  => {                    'type' => 'btrfssv' },
      'imageoverlays'         => {                    'type' => 'btrfssv' },
      'images'                => {                    'type' => 'btrfssv' },
      'images/puavo-pkg'      => {                    'type' => 'btrfssv' },
      'state'                 => {                    'type' => 'btrfssv' },
      'state/var/lib/docker'  => {                    'type' => 'btrfssv' },
      'state/var/lib/flatpak' => {                    'type' => 'btrfssv' },
      'state/var/log'         => {                    'type' => 'btrfssv' },
    },
    'diskinstaller' => {
      '.EFI'                  => { 'size' => '768M',  'type' => 'fat'     },
      'installimages'         => {                    'type' => 'btrfssv' },
    },
    'exam' => {
      '.EFI'                  => { 'size' => '256M',  'type' => 'fat'     },
      'examimages'            => {                    'type' => 'btrfssv' },
    },
    'laptop' => {
      '.EFI'                  => { 'size' => '3072M', 'type' => 'fat'     },
      '.SWAP'                 => { 'size' => '4G',    'type' => 'swap'    },
      'home'                  => {                    'type' => 'btrfssv' },
      'imageoverlays'         => {                    'type' => 'btrfssv' },
      'images'                => {                    'type' => 'btrfssv' },
      'images/puavo-pkg'      => {                    'type' => 'btrfssv' },
      'state'                 => {                    'type' => 'btrfssv' },
      'state/var/lib/docker'  => {                    'type' => 'btrfssv' },
      'state/var/lib/flatpak' => {                    'type' => 'btrfssv' },
      'state/var/log'         => {                    'type' => 'btrfssv' },
    },
  }

  EncryptionPolicyPublicKeyPath = "/etc/puavo-conf/tpm2-pcr-public-key.pem"

  EncryptionPresets = {
    'none' => {
      'label' => 'No Encryption',
      'policy' => nil,
    },
    'tpm-with-pin' => {
      'label' => 'PIN Based TPM Unlock',
      'request_pin' => true,
      'policy' => LuksEnrollmentPolicy.new(
        [],
        [ '11:sha256' ],
        EncryptionPolicyPublicKeyPath,
        []
      )
    },
    'tpm-with-recovery-key' => {
      'label' => 'Automatic TPM Unlock',
      'policy' => LuksEnrollmentPolicy.new(
        [],
        [ '11:sha256' ],
        EncryptionPolicyPublicKeyPath,
        []
      )
    }
  }

  EncryptionDeviceName = "puavo-cryptroot"

  def self.setup_raid_if_possible
    disks_for_mirroring = nil

    loop do
      disks_for_mirroring = QueryDiskInfo::list_disks_for_mirroring()
      break if disks_for_mirroring
      puts "\nYou are doing a bootserver installation but there are no\n" \
              + "two or more identically sized disks available for\n"    \
              + "setting up mirroring raid configuration.\n"
      msg = 'Install without raid mirroring? (type "yes" if this is okay)'
      answer = UI::ask_with_default(msg, 'no', UI::InstallWithoutRaidPrompt)
      return false if answer == 'yes'
    end

    loop do
      puts "\nThe following disks can be used for mirroring RAID:\n"
      QueryDiskInfo::show_disk_devices(disks_for_mirroring)
      msg = "\nUse #{ disks_for_mirroring.join('+') }" \
              + " in mirroring RAID setup?\n" \
              + '(type "yes" to destroy data and setup RAID, "no" otherwise) '

      answer = HighLine.ask(msg) { |q| q.whitespace = nil }
      case answer.strip
        when 'yes'
          break
        when 'no'
          puts 'NOT setting up RAID now.'
          return false
        else
          puts 'Please answer "yes" or "no".'
      end
    end

    disks_for_mirroring.each { |disk| DiskCleaner::reset_disk(disk)     }
    disks_for_mirroring.each { |disk| raid_partition(disk) }

    spare_args = (disks_for_mirroring.count > 2) \
                   ? [ "--spare-devices=#{ disks_for_mirroring.count - 2 }" ] \
                   : []

    raid_members = disks_for_mirroring.map do |device|
                     partition = wait_one_gpt_partition_by_parttype(
                       device, GPT_PARTTYPE_RAID, try_count: 30)
                     "/dev/#{ partition }"
                   end

    run('mdadm', '--create', '--run', '--verbose', '/dev/md0',
                 '--level=mirror', '--raid-devices=2', *spare_args,
                 *raid_members)

    puts ''

    return true
  end

  def self.check_tpm_enrollment_feasibility
    unless File.exist?(EncryptionPolicyPublicKeyPath)
      warn 'TPM based unlocking is not possible, because the public key ' \
            + "does not exist at #{ EncryptionPolicyPublicKeyPath }"
      exit(1)
    end

    unless tpm2_available?
      warn 'TPM based unlocking is not possible, because TPM 2.x is not available'
      exit(1)
    end
  end

  def self.reread_partition_table(device)
    device_path = "/dev/#{ device }"

    5.times do
      if system('partprobe', device_path, { :err => '/dev/null' }) then
        run('udevadm', 'settle')
        return
      end
      sleep(1)
    end

    raise "Could not re-read the partition table on '#{ device_path }'"
  end

  def self.reset_device(fs_conf)
    DiskCleaner::reset_disk(fs_conf['device'])
  end

  def self.setup_device(device, gpt_spec, create_efi:, create_swap:)
    device_path = "/dev/#{ device }"
    stdout, stderr, status = Open3.capture3('gdisk', device_path,
                                            stdin_data: gpt_spec)
    unless status.success? then
      raise "gdisk failed on #{ device_path }:\n#{ stderr }"
    end

    reread_partition_table(device)

    create_efi_filesystem(device) if create_efi
    clear_swap_partition(device)  if create_swap
  end

  def self.gpt_format(partitioning_type, filesystems, windows_size=nil)
    begin
      efi_size = filesystems['.EFI']['size']
      raise "EFI size not found" \
        unless efi_size.kind_of?(String) && !efi_size.empty?
    rescue StandardError => e
      raise "Could not determine EFI size: #{ e.message }"
    end

    swap_size = nil
    begin
      if filesystems.has_key?('.SWAP') then
        swap_size = filesystems['.SWAP']['size']
        raise "Unexpected swap size" \
          unless swap_size.kind_of?(String) && !swap_size.empty?
      end
    rescue StandardError => e
      raise "Could not determine swap size: #{ e.message }"
    end

    p_index = 0
    slices = [ "2\no\ny\n" ]    # init

    if partitioning_type != :on_raid then
      # EFI boot partition
      slices << "n\n#{ p_index+=1 }\n\n+#{ efi_size }\nef00\n"
      if swap_size then
        slices << "n\n#{ p_index+=1 }\n\n+#{ swap_size }\n8200\n"
      end
    end

    case partitioning_type
      when :basic_boot
        slices << "n\n#{ p_index+=1 }\n\n\n8304\n"  # puavo partition
      when :on_raid
        slices << "n\n#{ p_index+=1 }\n\n\n8304\n"  # puavo partition
      when :raid_boot
        slices << "n\n#{ p_index+=1 }\n\n\nfd00\n"  # RAID partition
      when :with_windows
        # Microsoft reserved partition:
        slices << "n\n#{ p_index+=1 }\n\n+16M\n0c01\n"
        # Microsoft basic data partition:
        slices << "n\n#{ p_index+=1 }\n\n+#{ windows_size }\n0700\nc\n#{ p_index }\n#{ GPT_PARTLABEL_PUAVO_INSTALLED_WINDOWS_C }\n"
        slices << "n\n#{ p_index+=1 }\n\n+768M\n2700\n"     # Windows RE
        slices << "n\n#{ p_index+=1 }\n\n\n8304\nw\ny\n"    # puavo partition
    end

    slices << "w\ny\n"
    return slices.join('')
  end

  def self.raid_partition(device)
    # we only support bootserver hosttype with RAID
    fs_conf = DiskHandler::Filesystems['bootserver']
    setup_device(device, gpt_format(:raid_boot, fs_conf),
                 create_efi: true, create_swap: true)
  end

  def self.get_partitions(device)
    device_path = "/dev/#{ device }"

    stdout, stderr, status = Open3.capture3(
      'lsblk', '--list', '--noheadings', '--output', 'TYPE,NAME', device_path,
    )
    raise "Could not list partitions on '#{ device_path }': #{ stderr.strip }" \
      unless status.success?

    stdout.lines.filter_map do |line|
      type, name = line.strip.split(/\s+/, 2)
      name if type&.casecmp?('part')
    end
  end

  def self.get_partition_by_partnum(device, partnum)
    partitions = get_partitions(device)

    raise "Device '#{ device }' does not have any partitions" \
      if partitions.empty?
    raise "Partition number '#{ partnum }' not found on '#{ device }'" \
      if partitions.length < partnum

    partitions[partnum - 1]
  end

  def self.wait_partition_by_partnum(device, partnum, try_count:,
                                     try_interval: 1)
    try_count.times do
      warn "Looking for partition '#{ partnum }' on device '#{ device }'..."
      return get_partition_by_partnum(device, partnum) \
        rescue sleep(try_interval)
    end
    raise "Partition '#{ partnum }' not found on device '#{ device }'"
  end

  def self.get_gpt_partitions_by_parttype(device, parttype)
    cmd = [ 'lsblk', '-l', '-n', '-oPARTTYPE,NAME', "/dev/#{ device }" ]
    IO.popen(cmd).readlines.map do |line|
      line.strip.split(' ')
    end.select do |lineparts|
      lineparts.length > 1 && lineparts[0].downcase == parttype.downcase
    end.map do |lineparts|
      lineparts[1..].join(' ')
    end
  end

  def self.get_one_gpt_partition_by_parttype(device, parttype)
    partitions = get_gpt_partitions_by_parttype(device, parttype)

    raise "GPT partition of type '#{ parttype }' not found on '#{ device }'" \
      if partitions.empty?
    raise "Multiple GPT partitions of type '#{ parttype }' found on '#{ device }'" \
      if partitions.length > 1

    partitions[0]
  end

  def self.wait_one_gpt_partition_by_parttype(device, parttype, try_count:, try_interval: 1)
    try_count.times do
      warn "Looking for GPT partition of type '#{ parttype }' on device '#{ device }'..."
      return get_one_gpt_partition_by_parttype(device, parttype) \
        rescue sleep(try_interval)
    end
    raise "GPT partition of type '#{ parttype }' not found on device '#{ device }'"
  end

  def self.get_puavo_partition(device)
    # It should not be possible to have multiple LVM partitions,
    # because at this stage, we have already created a new partition
    # table. If partition is not found in 30secs, then it's most
    # probably a progamming error; check gdisk/fdisk inputs.
    return wait_one_gpt_partition_by_parttype(device, GPT_PARTTYPE_LINUX,
                                              try_count: 30)
  end

  def self.setup_luks_partition(fs_conf, device_keys_temporary_directory)
    encryption = fs_conf['encryption']
    # The primary partition does not get a TPM enrollment policy.
    # It is unlocked exclusively via the control key stored in the
    # boot vault, which has its own TPM policy. Enrolling TPM on the
    # primary partition would weaken security by allowing it to be
    # unlocked without the Boot Trust Manager.
    handler = LuksPartitionHandler.new(
      fs_conf['puavo_partition'], EncryptionDeviceName, nil)
    control_key = handler.setup()

    # If the recovery key is set, add it
    if encryption && \
       encryption['recovery_key'] && \
       !encryption['recovery_key'].empty?
      handler.add_key(control_key, encryption['recovery_key'])
    end

    # Create the encrypted boot vault that will store the LUKS key
    # for controlling the LUKS partition during the boot process.
    vault = BootVaultBuilder.new(
      control_key: control_key,
      recovery_key: encryption['recovery_key'],
      encryption_policy: encryption['policy'],
      device: fs_conf['device']
    )

    vault.build(device_keys_temporary_directory)

    fs_conf['install_target'] = handler.encryption_device
  end

  def self.create_device_keys(device_keys_temporary_directory)
    script_path = '/usr/lib/puavo-ltsp-install/puavo-create-device-keys'
    run(script_path, device_keys_temporary_directory)
  end

  def self.finalize_secure_boot_keys(device_keys_temporary_directory)
    # Copy only Secure Boot key files from temporary directory into
    # the installation. TPM keys remain only in the boot vault.
    final_addons_directory = '/images/boot/addons'
    FileUtils.mkdir_p(final_addons_directory)

    secure_boot_files = [
      'secure-boot.priv', 'secure-boot.pem', 'db.auth', # DB private key, certificate and auth file
      'pk.pem', 'pk.auth',                              # PK certificate and auth file
      'kek.pem', 'kek.auth',                            # KEK certificate and auth file
    ]
    secure_boot_files.each do |filename|
      src = File.join(device_keys_temporary_directory, filename)
      FileUtils.cp(src, final_addons_directory)
    end

    # Lastly, install the device Secure Boot keys for the
    # installer so the UKI command-line signer can use them.
    run('/etc/puavo-conf/scripts/setup_device_secure_boot_keys',
        device_keys_temporary_directory)
  end

  def self.request_pin_enrollment(fs_conf)
    puts 'Requesting PIN enrollment on next boot...'
    run('/usr/lib/puavo-ltsp-install/puavo-change-unlock-pin')
  end

  def self.do_fs_setup(conf, vgname, raid_setup_done)
    fs_conf = conf.clone

    Dir.mktmpdir('puavo-device-keys-') do |device_keys_temporary_directory|
      fs_setup_phases = [ :reset_device ]
      fs_setup_phases += [ :wipe_device ] if conf['wipe']
      fs_setup_phases += [ :create_partitions ]

      # Assign the install target, which might be overwritten by LUKS methods
      fs_setup_phases += [ :select_default_install_target ]

      # Create device keys (Secure Boot keys, TPM keys) for permanent installations.
      # Boot vault also requires these keys to be present.
      hosttype = fs_conf['hosttype']
      if hosttype != 'diskinstaller' then
        fs_setup_phases += [ :create_device_keys ]
      end

      encryption = fs_conf['encryption']

      if encryption && encryption['policy'] then
        check_tpm_enrollment_feasibility
        fs_setup_phases += [ :setup_luks_partition ]
      end

      fs_setup_phases += [ :puavo_filesystems, :mount_filesystems ]

      # Finalize secure boot keys after filesystems are mounted
      if hosttype != 'diskinstaller' then
        fs_setup_phases += [ :finalize_secure_boot_keys ]
      end

      # If PIN-based encryption was selected, request PIN enrollment
      # on after installation instead of enrolling it now. This avoids
      # requiring the PIN later during the installation process.
      if encryption && encryption['request_pin'] then
        fs_setup_phases += [ :request_pin_enrollment ]
      end

      fs_setup_phases.each do |fn_sym|
        args = case fn_sym
                 when :mount_filesystems, :puavo_filesystems
                   [ fs_conf, vgname ]
                 when :create_device_keys,
                      :finalize_secure_boot_keys
                   [ device_keys_temporary_directory ]
                 when :create_partitions
                   [ fs_conf, raid_setup_done ]
                 when :setup_luks_partition
                   [ fs_conf, device_keys_temporary_directory ]
                 else
                   [ fs_conf ]
               end
        method(fn_sym).call(*args)
      end
    end
  end

  def self.create_partitions(fs_conf, raid_setup_done)
    if fs_conf['windows_size'] then
      # raid_setup_done == false because we must be on laptop
      gpt_spec = gpt_format(:with_windows, fs_conf['filesystems'],
                            fs_conf['windows_size'])
    else
      partitioning_scheme = raid_setup_done ? :on_raid : :basic_boot
      gpt_spec = gpt_format(partitioning_scheme, fs_conf['filesystems'])
    end

    create_efi  = !raid_setup_done && fs_conf['filesystems'].has_key?('.EFI')
    create_swap = !raid_setup_done && fs_conf['filesystems'].has_key?('.SWAP')
    setup_device(fs_conf['device'], gpt_spec,
                 create_efi:  create_efi,
                 create_swap: create_swap)
    fs_conf['puavo_partition'] = get_puavo_partition(fs_conf['device'])
  end

  def self.btrfs_subvolumes(hosttype)
    Filesystems[ hosttype ] \
      .select { |name, attrs| attrs['type'] == 'btrfssv' }
  end

  def self.select_default_install_target(fs_conf)
    fs_conf['install_target'] = "/dev/#{ fs_conf['puavo_partition'] }"
  end

  def self.create_efi_filesystem(device)
    efi_partition = wait_one_gpt_partition_by_parttype(device,
                      GPT_PARTTYPE_EFI, try_count: 30)
    run('mkfs.fat', '-F32', "/dev/#{ efi_partition }")
  end

  def self.clear_swap_partition(device)
    swap_partition = wait_one_gpt_partition_by_parttype(device,
                       GPT_PARTTYPE_SWAP, try_count: 30)
    run('wipefs', '--all', "/dev/#{ swap_partition }")
  end

  def self.puavo_filesystems(fs_conf, vgname)
    install_target_dev = fs_conf['install_target']

    # clean up possible confusing mess from the partition
    run('dd', 'if=/dev/zero', "of=#{ install_target_dev }", 'count=2K',
        { :err => '/dev/null' })

    run('mkfs.btrfs', '-L', vgname, install_target_dev)
    btrfs_mount_dir = "/.#{ vgname }"
    begin
      FileUtils.mkdir_p(btrfs_mount_dir)
      begin
        run('mount', install_target_dev, btrfs_mount_dir)
        btrfs_subvolumes( fs_conf['hosttype'] ).keys.each do |name|
          subvolume_path = File.join(btrfs_mount_dir, name)
          FileUtils.mkdir_p( File.dirname(subvolume_path) )
          run('btrfs', 'subvolume', 'create', subvolume_path)
        end
      ensure
        run('umount', btrfs_mount_dir)
      end
    ensure
      FileUtils.remove_dir(btrfs_mount_dir)
    end
  end

  def self.mount_filesystems(fs_conf, vgname)
    install_target_dev = fs_conf['install_target']

    btrfs_subvolumes( fs_conf['hosttype'] ).each do |name, attrs|
      mnt_path = "/#{ name }"
      FileUtils.mkdir_p(mnt_path)
      run('mount', '-o', "subvol=#{ name }", install_target_dev, mnt_path)
    end
  end

  def self.wipe(device)
    run('nwipe', '--autonuke',
                 '--method=quick',
                 '--nowait',
                 '--verify=off',
                 "/dev/#{ device }")
  end

  def self.wipe_device(fs_conf); wipe(fs_conf['device']); end
end

module QueryDiskInfo
  def self.disk_device_regexp()
    if UI::loopback_only then
      return /\Aloop[0-9]+\z/
    else
      # possible choices here should sort so that "md0" is the first
      # if it exists
      return /\A((md|mmcblk)[0-9]+|nvme[0-9]+n[0-9]+|[sv]d[a-z]|xvd[a-z])\z/
    end
  end

  def self.list_disks_for_mirroring()
    # return a list of devices if there are multiple disks with the same size,
    # otherwise return nil
    current_system_devices = get_current_system_devices()

    devices_by_size = {}
    IO.readlines('/proc/partitions').each do |line|
      disk_size, disk_name = * line.split(' ')[2,3]
      next unless disk_size && disk_name
      next unless disk_name.match(disk_device_regexp)
      next if disk_name.match(/\Amd[0-9]+\z/)       # no raid devices
      next if current_system_devices.include?(disk_name)
      (devices_by_size[disk_size.to_i] ||= []) << "#{ disk_name }"
    end

    duplicate_sized_disks = devices_by_size.select do |size, disks|
                              disks.count > 1
                            end
    return nil if duplicate_sized_disks.empty?
    biggest_common_disksize = duplicate_sized_disks.keys.sort.last
    return devices_by_size[biggest_common_disksize].sort
  end

  def self.get_blkdev_size_g(device)
    cmd = [ 'lsblk', '-d', '-b', '-l', '-n', '-oSIZE', "/dev/#{ device }" ]
    IO.popen(cmd) do |io|
      output = io.read
      io.close
      raise "lsblk failed" unless $?.success?
      output.to_i / 1024 ** 3
    end
  end

  def self.ask_device()
    choosable_disk_devices = []

    begin
      current_system_devices = get_current_system_devices()
    rescue StandardError => e
      raise "Could not find the current system device: #{ e.message }"
    end

    chosen_disk_device = nil
    until choosable_disk_devices.include?(chosen_disk_device) do
      # these should sort here so that "md0" is the first if it exists
      all_disk_devices = IO.readlines('/proc/partitions') \
                           .map { |s| s.split(' ')[3] }   \
                           .select { |s| s && s.match(disk_device_regexp) } \
                           .sort

      # do not allow installation to the devices where the currently running
      # system resides
      choosable_disk_devices = all_disk_devices - current_system_devices

      if choosable_disk_devices.empty? then
        raise 'Could not find any disk device where we could install to'
      end

      puts '-----'
      puts "We can install to the following disk devices:\n"
      show_disk_devices(choosable_disk_devices)

      if chosen_disk_device.nil? then
        chosen_disk_device = choosable_disk_devices[0]
      end

      puts "\nPossible disk devices are: #{ choosable_disk_devices.join(' ') }"
      chosen_disk_device = UI::ask_with_default('Choose a disk device:',
                                                chosen_disk_device,
                                                UI::DiskDevicePrompt)
    end

    chosen_disk_device
  end

  def self.ask_confirmation(device)
    if device.nil? || device.empty? then
      return false
    end

    puts %Q{\nYou are going to install to device "/dev/#{ device }":}
    UI::colormsg("IF YOU PROCEED, ALL DATA ON THAT DEVICE IS DESTROYED!\n",
                 HighLine::RED)

    prompt = %q{Are you sure you want to proceed (write "yes" if this is okay?)}

    UI::ask_with_default(prompt, 'no', UI::WritePartitionsPrompt) == 'yes'
  end

  def self.ask_windows_setup(hosttype, device)
    return nil unless hosttype == 'laptop'

    30.times do # Invalid preseeds shall not cause infinite loops.
      prompt = "\nSetup Windows partitions? (yes/no)"
      setup_windows_answer = UI::ask_with_default(prompt, 'no',
                                                  UI::SetupWindowsPrompt)
      case setup_windows_answer
      when 'yes'
        disk_size_g = get_blkdev_size_g(device)
        unless disk_size_g >= DUALBOOT_LAPTOP_DISK_MIN_SIZE_G then
          UI::colormsg("Windows partitions can be setup only on #{ DUALBOOT_LAPTOP_DISK_MIN_SIZE_G }G or larger disks.", HighLine::RED)
          next
        end

        windows_max_size_g = disk_size_g - DUALBOOT_LAPTOP_PUAVOOS_MIN_SIZE_G
        windows_size_tuple = UI::ask_size_with_default(
                               "Size of Windows Local Disk (C:) partition:",
                               "#{ [40, windows_max_size_g].min }G",
                               UI::WindowsSizePrompt)

        windows_size_g = Size::to_bytes(windows_size_tuple) / 1024 ** 3
        unless windows_size_g >= DUALBOOT_LAPTOP_WINDOWS_MIN_SIZE_G then
          UI::colormsg("Minimum Windows Local Disk (C:) partition size is #{ DUALBOOT_LAPTOP_WINDOWS_MIN_SIZE_G }G.", HighLine::RED)
          next
        end

        unless windows_size_g <= windows_max_size_g then
          UI::colormsg("Maximum Windows Local Disk (C:) partition size on this disk (#{ disk_size_g }G) is #{ windows_max_size_g }G.", HighLine::RED)
          next
        end

        return windows_size_tuple.join()
      when 'no'
        return nil
      end
    end
    raise 'Windows setup prompts exceeded.'
  end

  # Layout portable alphabet (modhex) used for recovery key generation.
  RecoveryKeyAlphabet = 'cbdefghijklnrtuv'

  # Generated recovery key export location.
  RecoveryKeyExportPath = '/dev/shm/puavo-recovery-key'

  def self.generate_recovery_key()
    # Generate a 64 character key (256 bits) from 32 bytes
    SecureRandom.random_bytes(32).each_byte.map { |byte|
      first_index = byte >> 4
      second_index = byte & 0x0f
      RecoveryKeyAlphabet[first_index] + RecoveryKeyAlphabet[second_index]
    }.join
  end

  def self.generate_and_export_recovery_key()
    recovery_key = generate_recovery_key()

    File.open(RecoveryKeyExportPath, 'w', 0600) do |file|
      file.chmod(0600)
      file.write(recovery_key)
    end

    UI::colormsg('Recovery key generated. Record it now.', HighLine::YELLOW)
    UI::colormsg(recovery_key, HighLine::GREEN)
    puts "Also written to #{ RecoveryKeyExportPath } for capture."

    recovery_key
  end

  def self.ask_encryption_setup()
    unless tpm2_available?
      warn 'TPM 2.x is not available, disk encryption is disabled'
      return DiskHandler::EncryptionPresets['none'].clone
    end

    prompt = %q{Select the encryption preset:}
    choices = DiskHandler::EncryptionPresets.map {
      |key, preset| [key, preset['label']] }
    preset_key = UI::ask_multiple_choices(
      prompt, choices, UI::EncryptionPresetPrompt)

    unless DiskHandler::EncryptionPresets.keys.include?(preset_key) then
      raise "Unrecognized encryption preset '#{ preset_key }'"
    end

    # Use a clone so we can safely mutate
    encryption = DiskHandler::EncryptionPresets[preset_key].clone

    # If encryption is enabled, offer an optional generated recovery
    # key, in order to ensure portability between keyboard layouts
    if encryption['policy']
      add_key_answer = UI::ask_with_default(
        "Generate a recovery key? (y/N)", 'no', UI::RecoveryKeyPrompt)
      if add_key_answer.downcase.start_with?('y')
        encryption['recovery_key'] = generate_and_export_recovery_key()
      end
    end

    return encryption
  end

  def self.ask_device_with_confirmation(hosttype, raid_setup_done)
    device       = nil
    do_it        = false
    encryption   = nil
    windows_size = nil
    wipe         = false

    until do_it do
      UI::reset_questions()

      if raid_setup_done then
        device = 'md0'
      else
        device = QueryDiskInfo::ask_device()
      end

      windows_size = QueryDiskInfo::ask_windows_setup(hosttype, device)

      msg = "\nIt is possible to wipe disk/partition before installing,\n" \
              + "but this can take a rather long time.\n"                  \
              + "Should we wipe the disk/partition before installing? (yes/no)"
      while true do
        wipe_answer = UI::ask_with_default(msg, 'no', UI::WipePartitionPrompt)
        case wipe_answer
          when 'yes'
            wipe = true
            break
          when 'no'
            wipe = false
            break
        end
      end

      encryption = QueryDiskInfo::ask_encryption_setup()

      do_it = QueryDiskInfo::ask_confirmation(device)
    end

    {
      'device'          => device,
      'encryption'      => encryption,
      'filesystems'     => DiskHandler::Filesystems[hosttype],
      'hosttype'        => hosttype,
      'windows_size'    => windows_size,
      'wipe'            => wipe,
    }
  end

  def self.get_current_system_devices
    cmd = '/usr/lib/puavo-core/puavo-get-boot-disks'
    stdout_str, stderr_str, status = Open3.capture3(cmd)
    raise stderr_str unless status.success?
    return stdout_str.split("\n").map { |s| s.delete_prefix('/dev/') }
  end

  def self.show_disk_devices(disk_devices)
    sleep(1.5)

    disk_devices.each do |device|
      print "\n[#{ device }]\n"
      system("gdisk -l /dev/#{ device } | sed 's/^/  /' | more -e") \
        or raise "Error running fdisk for device '#{ device }'"
    end
  end
end

module Size
  UNIT_BYTES = {
    'K' => 1024 ** 1,
    'M' => 1024 ** 2,
    'G' => 1024 ** 3,
  }

  def self.parse(str)
    unit = str[-1]
    if not UNIT_BYTES.include?(unit)
      raise ArgumentError, "invalid unit, expected one from #{UNIT_BYTES.keys}"
    end

    number = str[0..-2]
    value = Integer(number)
    if value < 0
      raise ArgumentError, "invalid size, expected a non-negative value"
    end

    [value, unit]
  end

  def self.to_bytes(size)
    value, unit = size
    value * UNIT_BYTES[unit]
  end

end

module UI
  DiskDevicePrompt,
  EncryptionPresetPrompt,
  InstallWithoutRaidPrompt,
  PartitionPrompt,
  SetupWindowsPrompt,
  UnpartitionedSpacePrompt,
  WindowsSizePrompt,
  WipePartitionPrompt,
  WritePartitionsPrompt,
  RecoveryKeyPrompt,
    = *(1..11)

  @ask_forced      = {}
  @defaults        = {}
  @force_defaults  = {}
  @hosttype        = nil
  @loopback_only   = false

  def self.hosttype         ; @hosttype         end
  def self.loopback_only    ; @loopback_only    end

  def self.reset_questions
    @ask_forced = @force_defaults.clone
  end

  def self.help_and_exit
    puts <<-EOF
puavo-setup-filesystems [OPTIONS]

-h, --help                        show help

    --ask-disk-device             ask with default for disk device
    --ask-encryption              ask with default for encryption preset
    --ask-imageoverlay-size       ask with default for imageoverlay size
    --ask-install-without-raid    ask with default for install without raid
    --ask-partition               ask with default for partition
    --ask-setup-windows           ask with default for setup windows
    --ask-unpartitioned-space     ask with default for unpartitioned space
    --ask-windows-size            ask with default for windows size
    --ask-wipe-partition          ask with default for wipe partition
    --ask-write-partitions        ask with default for write partitions
    --ask-recovery-key            ask with default for recovery key
    --force-disk-device           force value for disk device
    --force-encryption            force value for encryption preset
    --force-imageoverlay-size     force value for imageoverlay size
    --force-install-without-raid  force value for install without raid
    --force-partition             force value for partition
    --force-setup-windows         force value for setup windows
    --force-unpartitioned-space   force value for unpartitioned space
    --force-windows-size          force value for windows size
    --force-wipe-partition        force value for wipe partition
    --force-write-partitions      force value for write partitions
    --force-recovery-key          force value for recovery key
    --hosttype                    force value for hosttype
    --loopback-only               use only loopback devices

EOF
  end

  def self.parse_args
    begin
      opts = GetoptLong.new(
        [ '--help', '-h'                , GetoptLong::NO_ARGUMENT       ],
        [ '--ask-disk-device'           , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-encryption'            , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-imageoverlay-size'     , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-install-without-raid'  , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-partition'             , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-setup-windows'         , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-unpartitioned-space'   , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-windows-size'          , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-wipe-partition'        , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-write-partitions'      , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--ask-recovery-key'          , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-disk-device'         , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-encryption'          , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-imageoverlay-size'   , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-install-without-raid', GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-partition'           , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-setup-windows'       , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-unpartitioned-space' , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-windows-size'        , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-wipe-partition'      , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-write-partitions'    , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--force-recovery-key'        , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--hosttype'                  , GetoptLong::REQUIRED_ARGUMENT ],
        [ '--loopback-only'             , GetoptLong::NO_ARGUMENT       ],
      )

      opts.each do |opt, arg|
        case opt
        when '--help'
          help_and_exit
          exit(0)
        when '--ask-disk-device'
          @defaults[DiskDevicePrompt] = arg
        when '--ask-encryption'
          @defaults[EncryptionPresetPrompt] = arg
        when '--ask-install-without-raid'
          @defaults[InstallWithoutRaidPrompt] = arg
        when '--ask-partition'
          @defaults[PartitionPrompt] = arg
        when '--ask-setup-windows'
          @defaults[SetupWindowsPrompt] = arg
        when '--ask-unpartitioned-space'
          @defaults[UnpartitionedSpacePrompt] = arg
        when '--ask-windows-size'
          @defaults[WindowsSizePrompt] = arg
        when '--ask-wipe-partition'
          @defaults[WipePartitionPrompt] = arg
        when '--ask-write-partitions'
          @defaults[WritePartitionsPrompt] = arg
        when '--ask-recovery-key'
          @defaults[RecoveryKeyPrompt] = arg
        when '--force-disk-device'
          @defaults[DiskDevicePrompt] = arg
          @force_defaults[DiskDevicePrompt] = true
        when '--force-encryption'
          @defaults[EncryptionPresetPrompt] = arg
          @force_defaults[EncryptionPresetPrompt] = true
        when '--force-install-without-raid'
          @defaults[InstallWithoutRaidPrompt] = arg
          @force_defaults[InstallWithoutRaidPrompt] = true
        when '--force-partition'
          @defaults[PartitionPrompt] = arg
          @force_defaults[PartitionPrompt] = true
        when '--force-setup-windows'
          @defaults[SetupWindowsPrompt] = arg
          @force_defaults[SetupWindowsPrompt] = true
        when '--force-unpartitioned-space'
          @defaults[UnpartitionedSpacePrompt] = arg
          @force_defaults[UnpartitionedSpacePrompt] = true
        when '--force-windows-size'
          @defaults[WindowsSizePrompt] = arg
          @force_defaults[WindowsSizePrompt] = true
        when '--force-wipe-partition'
          @defaults[WipePartitionPrompt] = arg
          @force_defaults[WipePartitionPrompt] = true
        when '--force-write-partitions'
          @defaults[WritePartitionsPrompt] = arg
          @force_defaults[WritePartitionsPrompt] = true
        when '--force-recovery-key'
          @defaults[RecoveryKeyPrompt] = arg
          @force_defaults[RecoveryKeyPrompt] = true
        when '--hosttype'
          @hosttype = arg
        when '--loopback-only'
          @loopback_only = true
        end
      end
    rescue GetoptLong::InvalidOption => e
      help_and_exit
      exit(1)
    end
  end

  def self.ask_with_default(question, default, prompt_key=nil)
    if @defaults.has_key?(prompt_key) && @defaults[prompt_key] != 'default' then
      default = @defaults[prompt_key]
    end
    if @ask_forced.has_key?(prompt_key) then
      # Do not return forced twice... in case it did not work,
      # we need user input.
      @ask_forced.delete(prompt_key)
      return default
    end
    prompt = "#{ question } [#{ default }] "
    answer = HighLine.ask(prompt) { |q| q.whitespace = nil }

    return answer.match(/\A\n\z/)  ? default \
         : answer.match(/\A\s+\z/) ? ''      \
         : answer.strip
  end

  def self.ask_multiple_choices(question, choices, prompt_key=nil)
    default = nil
    if @defaults.has_key?(prompt_key) && @defaults[prompt_key] != 'default'
      default = @defaults[prompt_key]
    end
    if @ask_forced.has_key?(prompt_key) then
      # Do not return forced twice... in case it did not work,
      # we need user input.
      @ask_forced.delete(prompt_key)
      return default
    end
    cli = HighLine.new
    answer = cli.choose do |menu|
      menu.prompt = question
      choices.each do |key, label|
        menu.choice(label) { key }
      end
    end
    answer
  end

  def self.ask_size_with_default(question, default, prompt_key=nil)
    while true do
      str = ask_with_default(question, default, prompt_key)
      begin
        size = Size::parse(str)
      rescue ArgumentError => e
        colormsg(e.message, HighLine::RED)
        next
      end
      break
    end
    size
  end

  def self.colormsg(msg, color)
    HighLine.say(HighLine.new.color(msg, color))
  end
end

UI::parse_args

hosttype = UI::hosttype || PuavoFacts::get('puavo_hosttype')
unless DiskHandler::Filesystems[hosttype] then
  warn "No filesystem scheme defined for hosttype '#{ hosttype }'," \
         + ' try --hosttype'
  exit(1)
end

vgname = (hosttype == 'diskinstaller') ? 'puavoinstaller' : \
         (hosttype == 'exam')          ? 'puavoexam'      : 'puavo'

raid_setup_done = false
if hosttype == 'bootserver' then
  raid_setup_done = DiskHandler::setup_raid_if_possible()
end

case hosttype
  when 'bootserver', 'diskinstaller', 'exam', 'laptop'
    conf = QueryDiskInfo::ask_device_with_confirmation(hosttype, raid_setup_done)
    DiskHandler::do_fs_setup(conf, vgname, raid_setup_done)
  else
    puts "Hosttype '#{ hosttype }' does not need to setup local filesystems."
end
