module Beaker::HostPrebuiltSteps

def additive_hash_merge h1, h2
  merged_hash = {}
  normalized_h2 = h2.each_with_object({}) { |(k, v), h| h[k.to_s.upcase] = v; }
  h1.each_pair do |key, _val|
    normalized_key = key.to_s.upcase
    if normalized_h2.has_key?(normalized_key)
      merged_hash[key] = [h1[key], normalized_h2[normalized_key]]
      merged_hash[key] = merged_hash[key].uniq # remove dupes
    end
  end
  merged_hash
end

def apt_get_update hosts
  block_on hosts do |host|
    host.exec(Command.new("apt-get update")) if /ubuntu|debian|cumulus/.match?(host[:platform])
  end
end

def check_and_install_packages_if_needed host, package_list
  package_list.each do |string|
    alternatives = string.split('|')
    next if alternatives.any? { |pkg| host.check_for_package pkg }
    install_one_of_packages host, alternatives
  end
end

def construct_env host, opts
  env = additive_hash_merge(host[:host_env], opts[:host_env])
  env.each_key do |key|
    separator = host['pathseparator']
    separator = ':' if key == 'PATH' && (not host.is_powershell?)
    env[key] = env[key].join(separator)
  end
  env
end

def copy_file_to_remote(host, file_path, file_content)
  block_on host do |host|
    Tempfile.open 'beaker' do |tempfile|
      File.open(tempfile.path, 'w') { |file| file.puts file_content }
      host.do_scp_to(tempfile.path, file_path, @options)
    end
  end
end

def copy_ssh_to_root host, opts
  logger = opts[:logger]
  block_on host do |host|
    logger.debug "Give root a copy of current user's keys, on #{host.name}"
    if host['platform'].include?('windows') and host.is_cygwin?
      host.exec(Command.new('cp -r .ssh /cygdrive/c/Users/Administrator/.'))
      host.exec(Command.new('chown -R Administrator /cygdrive/c/Users/Administrator/.ssh'))
    elsif host['platform'].include?('windows') and not host.is_cygwin?
      # from https://www.microsoft.com/resources/documentation/windows/xp/all/proddocs/en-us/xcopy.mspx?mfr=true:
      #    /i : If Source is a directory or contains wildcards and Destination
      #      does not exist, xcopy assumes destination specifies a directory
      #      name and creates a new directory. Then, xcopy copies all specified
      #      files into the new directory. By default, xcopy prompts you to
      #      specify whether Destination is a file or a directory.
      #
      #    /y : Suppresses prompting to confirm that you want to overwrite an
      #      existing destination file.
      host.exec(Command.new("if exist .ssh (xcopy .ssh C:\\Users\\Administrator\\.ssh /s /e /y /i)"))
    elsif host['platform'].include?('osx')
      host.exec(Command.new('sudo cp -r .ssh /var/root/.'), { :pty => true })
    elsif /(free|open)bsd/.match?(host['platform']) || host['platform'].include?('solaris-11')
      host.exec(Command.new('sudo cp -r .ssh /root/.'), { :pty => true })
    elsif host['platform'].include?('solaris-10')
      host.exec(Command.new('sudo cp -r .ssh /.'), { :pty => true })
    else
      host.exec(Command.new('sudo su -c "cp -r .ssh /root/."'), { :pty => true })
    end
    host.exec(Command.new('sudo fixfiles restore /root')) if host.selinux_enabled?
  end
end

def disable_se_linux host, opts
  logger = opts[:logger]
  block_on host do |host|
    if /centos|el-|redhat|fedora|eos/.match?(host['platform'])
      logger.debug("Disabling se_linux on #{host.name}")
      host.exec(Command.new("sudo su -c \"setenforce 0\""), { :pty => true })
    else
      logger.warn("Attempting to disable SELinux on non-supported platform: #{host.name}: #{host['platform']}")
    end
  end
end

def disable_updates hosts, opts
  logger = opts[:logger]
  hosts.each do |host|
    next if host['platform'].include?('netscaler')
    logger.notify "Disabling updates.puppetlabs.com by modifying hosts file to resolve updates to 127.0.0.1 on #{host}"
    set_etc_hosts(host, "127.0.0.1\tupdates.puppetlabs.com\n")
  end
end

def enable_root_login host, opts
  logger = opts[:logger]
  block_on host do |host|
    logger.debug "Update sshd_config to allow root login"
    if host['platform'].include?('osx')
      # If osx > 10.10 use '/private/etc/ssh/sshd_config', else use '/etc/sshd_config'
      ssh_config_file = '/private/etc/ssh/sshd_config'
      ssh_config_file = '/etc/sshd_config' if /^osx-10\.(9|10)/i.match?(host['platform'])
      host.exec(Command.new("sudo sed -i '' 's/#PermitRootLogin no/PermitRootLogin Yes/g' #{ssh_config_file}"))
      host.exec(Command.new("sudo sed -i '' 's/#PermitRootLogin yes/PermitRootLogin Yes/g' #{ssh_config_file}"))
    elsif host['platform'].include?('freebsd')
      host.exec(Command.new("sudo sed -i -e 's/#PermitRootLogin no/PermitRootLogin yes/g' /etc/ssh/sshd_config"), { :pty => true })
    elsif host['platform'].include?('openbsd')
      host.exec(Command.new("sudo perl -pi -e 's/^PermitRootLogin no/PermitRootLogin yes/' /etc/ssh/sshd_config"), { :pty => true })
    elsif host['platform'].include?('solaris-10')
      host.exec(Command.new("sudo gsed -i -e 's/#PermitRootLogin no/PermitRootLogin yes/g' /etc/ssh/sshd_config"), { :pty => true })
    elsif host['platform'].include?('solaris-11')
      host.exec(Command.new("if grep \"root::::type=role\" /etc/user_attr; then sudo rolemod -K type=normal root; else echo \"root user already type=normal\"; fi"), { :pty => true })
      host.exec(Command.new("sudo gsed -i -e 's/PermitRootLogin no/PermitRootLogin yes/g' /etc/ssh/sshd_config"), { :pty => true })
    elsif host['platform'].include?('f5') || host.is_powershell?
      # interacting with f5 should using tmsh
      logger.warn("Attempting to enable root login non-supported platform: #{host.name}: #{host['platform']}")
    elsif host.is_cygwin?
      host.exec(Command.new("sed -ri 's/^#?PermitRootLogin /PermitRootLogin yes/' /etc/sshd_config"), { :pty => true })
    else
      host.exec(Command.new("sudo su -c \"sed -ri 's/^#?PermitRootLogin no|^#?PermitRootLogin yes/PermitRootLogin yes/' /etc/ssh/sshd_config\""), { :pty => true })
    end
    # restart sshd
    if /debian|ubuntu|cumulus/.match?(host['platform'])
      host.exec(Command.new("sudo su -c \"service ssh restart\""), { :pty => true })
    elsif /amazon|arch|(centos|el|redhat)-[789]|fedora-(1[4-9]|2[0-9]|3[0-9])/.match?(host['platform'])
      host.exec(Command.new("sudo -E systemctl restart sshd.service"), { :pty => true })
    elsif /centos|el-|redhat|fedora|eos/.match?(host['platform'])
      host.exec(Command.new("sudo -E /sbin/service sshd reload"), { :pty => true })
    elsif /(free|open)bsd/.match?(host['platform'])
      host.exec(Command.new("sudo /etc/rc.d/sshd restart"))
    elsif host['platform'].include?('solaris')
      host.exec(Command.new("sudo -E svcadm restart network/ssh"), { :pty => true })
    else
      logger.warn("Attempting to update ssh on non-supported platform: #{host.name}: #{host['platform']}")
    end
  end
end

def get_domain_name(host)
  domain = nil
  search = nil
  resolv_conf = if host['platform'].include?('windows')
                  if host.is_cygwin?
                    host.exec(Command.new("cat /cygdrive/c/Windows/System32/drivers/etc/hosts")).stdout
                  else
                    host.exec(Command.new('type C:\Windows\System32\drivers\etc\hosts')).stdout
                  end
                else
                  host.exec(Command.new("cat /etc/resolv.conf")).stdout
                end
  resolv_conf.each_line do |line|
    if (match = /^\s*domain\s+(\S+)/.match(line))
      domain = match[1]
    elsif (match = /^\s*search\s+(\S+)/.match(line))
      search = match[1]
    end
  end
  return_value ||= domain
  return_value ||= search
  return unless return_value
  return_value.gsub(/\.$/, '')
end

def hack_etc_hosts hosts, _opts
  etc_hosts = "127.0.0.1\tlocalhost localhost.localdomain\n"
  hosts.each do |host|
    ip = host['vm_ip'] || host['ip'].to_s
    hostname = host[:vmhostname] || host.name
    domain = get_domain_name(host)
    etc_hosts += "#{ip}\t#{hostname}.#{domain} #{hostname}\n"
  end
  hosts.each do |host|
    set_etc_hosts(host, etc_hosts)
  end
end

def host_packages(host)
  case host['platform']
  when /amazon/
    AMAZON2023_PACKAGES
  when /el-8/
    RHEL8_PACKAGES
  when /el-9/
    RHEL9_PACKAGES
  when /sles-10/
    SLES10_PACKAGES
  when /opensuse|sles-/
    SLES_PACKAGES
  when /debian/
    DEBIAN_PACKAGES
  when /cumulus/
    CUMULUS_PACKAGES
  when /windows/
    if host.is_cygwin?
      raise RuntimeError, "cygwin is not installed on #{host}" if !host.cygwin_installed?
      WINDOWS_PACKAGES
    else
      PSWINDOWS_PACKAGES
    end
  when /freebsd/
    FREEBSD_PACKAGES
  when /openbsd/
    OPENBSD_PACKAGES
  when /solaris-10/
    SOLARIS10_PACKAGES
  when /solaris-1[1-9]/
    SOLARIS11_PACKAGES
  when /archlinux/
    ARCHLINUX_PACKAGES
  when /fedora/
    FEDORA_PACKAGES
  else
    if !/aix|solaris|osx-|f5-|netscaler|cisco_/.match?(host['platform'])
      UNIX_PACKAGES
    else
      []
    end
  end
end

def install_one_of_packages host, packages
  error = nil
  packages.each do |pkg|
    begin
      return host.install_package pkg
    rescue Beaker::Host::CommandFailure => e
      error = e
    end
  end
  raise error
end

def package_proxy host, opts
  logger = opts[:logger]
  block_on host do |host|
    logger.debug("enabling proxy support on #{host.name}")
    case host['platform']
    when /ubuntu/, /debian/, /cumulus/
      host.exec(Command.new("echo 'Acquire::http::Proxy \"#{opts[:package_proxy]}/\";' >> /etc/apt/apt.conf.d/10proxy"))
    when /amazon/, /^el-/, /centos/, /fedora/, /redhat/, /eos/
      host.exec(Command.new("echo 'proxy=#{opts[:package_proxy]}/' >> /etc/yum.conf"))
    when /solaris-11/
      host.exec(Command.new("/usr/bin/pkg unset-publisher solaris || :"))
      host.exec(Command.new("/usr/bin/pkg set-publisher -g %s solaris" % opts[:package_proxy]))
    else
      logger.debug("Attempting to enable package manager proxy support on non-supported platform: #{host.name}: #{host['platform']}")
    end
  end
end

def set_env host, opts
  logger = opts[:logger]
  block_on host do |host|
    skip_msg = host.skip_set_env?
    unless skip_msg.nil?
      logger.debug(skip_msg)
      next
    end
    env = construct_env(host, opts)
    logger.debug("setting local environment on #{host.name}")
    env['CYGWIN'] = 'nodosfilewarning' if host['platform'].include?('windows') && host.is_cygwin?
    host.ssh_permit_user_environment
    host.ssh_set_user_environment(env)
    # close the host to re-establish the connection with the new sshd settings
    host.close
    # print out the working env
    if host.is_powershell?
      host.exec(Command.new("SET"))
    else
      host.exec(Command.new("cat #{host[:ssh_env_file]}"))
    end
  end
end

def set_etc_hosts(host, etc_hosts)
  if host['platform'].include?('freebsd')
    host.echo_to_file(etc_hosts, '/etc/hosts')
  elsif ((host['platform'].include?('windows')) and not host.is_cygwin?)
    host.exec(Command.new("echo '#{etc_hosts}' >> C:\\Windows\\System32\\drivers\\etc\\hosts"))
  else
    host.exec(Command.new("echo '#{etc_hosts}' >> /etc/hosts"))
  end
  # AIX must be configured to prefer local DNS over external
  return unless host['platform'].include?('aix')
  aix_netsvc = '/etc/netsvc.conf'
  aix_local_resolv = 'hosts = local, bind'
  return if host.exec(Command.new("grep '#{aix_local_resolv}' #{aix_netsvc}"), :accept_all_exit_codes => true).exit_code == 0
  host.exec(Command.new("echo '#{aix_local_resolv}' >> #{aix_netsvc}"))
end

def sync_root_keys host, opts
  # JJM This step runs on every system under test right now.  We're anticipating
  # issues on Windows and maybe Solaris.  We will likely need to filter this step
  # but we're deliberately taking the approach of "assume it will work, fix it
  # when reality dictates otherwise"
  logger = opts[:logger]
  block_on host do |host|
    logger.notify "Sync root authorized_keys from github on #{host.name}"
    # Allow all exit code, as this operation is unlikely to cause problems if it fails.
    if /solaris|eos/.match?(host['platform'])
      host.exec(Command.new(ROOT_KEYS_SYNC_CMD % "bash"), :accept_all_exit_codes => true)
    elsif host['platform'].include?('aix')
      host.exec(Command.new(ROOT_KEYS_SYNC_CMD_AIX % "env PATH=/usr/gnu/bin:$PATH bash"), :accept_all_exit_codes => true)
    else
      host.exec(Command.new(ROOT_KEYS_SYNC_CMD % "env PATH=\"/usr/gnu/bin:$PATH\" bash"), :accept_all_exit_codes => true)
    end
  end
rescue => e
  report_and_raise(logger, e, "sync_root_keys")
end

def timesync host, opts
  logger = opts[:logger]
  ntp_server = opts[:ntp_server] ? opts[:ntp_server] : NTPSERVER
  block_on host do |host|
    logger.notify "Update system time sync for '#{host.name}'"
    if host['platform'].include? 'windows'
      # The exit code of 5 is for Windows 2008 systems where the w32tm /register command
      # is not actually necessary.
      host.exec(Command.new("w32tm /register"), :acceptable_exit_codes => [0, 5])
      host.exec(Command.new("net start w32time"), :acceptable_exit_codes => [0, 2])
      host.exec(Command.new("w32tm /config /manualpeerlist:#{ntp_server} /syncfromflags:manual /update"))
      host.exec(Command.new("w32tm /resync"))
      logger.notify "NTP date succeeded on #{host}"
    else
      if /amazon|el-[89]|fedora/.match?(host['platform'])
        ntp_command = "chronyc add server #{ntp_server} prefer trust;chronyc makestep;chronyc burst 1/2"
      elsif /opensuse-|sles-/.match?(host['platform'])
        ntp_command = "sntp #{ntp_server}"
      elsif host['platform'].include?('cisco_nexus')
        ntp_server = host.exec(Command.new("getent hosts #{NTPSERVER} | head -n1 |cut -d \" \" -f1"), :acceptable_exit_codes => [0]).stdout
        ntp_command = "sudo -E sh -c 'export DCOS_CONTEXT=2;/isan/bin/ntpdate -u -t 20 #{ntp_server}'"
      else
        ntp_command = "ntpdate -u -t 20 #{ntp_server}"
      end
      success = false
      try = 0
      until try >= TRIES
        try += 1
        if host.exec(Command.new(ntp_command), :accept_all_exit_codes => true).exit_code == 0
          success = true
          break
        end
        sleep SLEEPWAIT
      end
      raise "NTP date was not successful after #{try} tries" unless success
      logger.notify "NTP date succeeded on #{host} after #{try} tries"
    end
  end
  nil
rescue => e
  report_and_raise(logger, e, "timesync (--ntp)")
end

def validate_host host, opts
  logger = opts[:logger]
  block_on host do |host|
    check_and_install_packages_if_needed(host, host_packages(host))
  end
rescue => e
  report_and_raise(logger, e, "validate")
end