#!/usr/bin/env ruby

##################################################################
# This part of the code might be running on Ruby versions other
# than 2.0. Testing on multiple Ruby versions is required for
# changes to this part of the code.
##################################################################
require 'json'

class Proxy
  instance_methods.each do |m|
    undef_method m unless m =~ /(^__|^send$|^object_id$)/
  end

  def initialize(*targets)
    @targets = targets
  end

  protected

  def method_missing(name, *args, &block)
    @targets.map do |target|
      target.__send__(name, *args, &block)
    end
  end
end

require 'tmpdir'
require 'logger'

log_file_path = "#{Dir.tmpdir()}/codedeploy-agent.update.log"

if($stdout.isatty)
  # if we are being run in a terminal, log to stdout and the log file.
  @log = Logger.new(Proxy.new(File.open(log_file_path, 'a+'), $stdout))
else
  # keep at most 2MB of old logs rotating out 1MB at a time
  @log = Logger.new(log_file_path, 2, 1048576)
  # make sure anything coming out of ruby ends up in the log file
  $stdout.reopen(log_file_path, 'a+')
  $stderr.reopen(log_file_path, 'a+')
end

@log.level = Logger::INFO

require 'net/http'
require 'json'

TOKEN_PATH = '/latest/api/token'
DOCUMENT_PATH = '/latest/dynamic/instance-identity/document'

class IMDSV2
  def self.region
    doc['region'].strip
  end

  private
  def self.http_request(request)
    Net::HTTP.start('169.254.169.254', 80, :read_timeout => 120, :open_timeout => 120) do |http|
      response = http.request(request)
      if response.code.to_i != 200
        raise "HTTP error from metadata service: #{response.message}, code #{response.code}"
      end
      return response.body
    end
  end

  def self.put_request(path)
    request = Net::HTTP::Put.new(path)
    request['X-aws-ec2-metadata-token-ttl-seconds'] = '21600'
    http_request(request)
  end

  def self.get_request(path, token = nil)
    request = Net::HTTP::Get.new(path)
    unless token.nil?
      request['X-aws-ec2-metadata-token'] = token
    end
    http_request(request)
  end

  def self.doc
    begin
      token = put_request(TOKEN_PATH)
      JSON.parse(get_request(DOCUMENT_PATH, token).strip)
    rescue
      JSON.parse(get_request(DOCUMENT_PATH).strip)
    end
  end
end

require 'set'
VALID_TYPES = Set.new ['rpm','zypper','deb','msi']

begin
  require 'fileutils'
  require 'openssl'
  require 'open-uri'
  require 'uri'
  require 'getoptlong'
  require 'tempfile'

  def usage
    print <<EOF

update [--sanity-check] [--proxy http://hostname:port] [--upgrade|--downgrade] <package-type>
   --sanity-check [optional]
   --proxy [optional]
   --upgrade | --downgrade [optional]
   package-type: #{VALID_TYPES.to_a.join(', ')}, or auto

Installs fetches the latest package version of the specified type and
installs it. rpms are installed with yum; debs are installed using gdebi.

This program is invoked automatically to update the agent once per day using
the same package manager the codedeploy-agent is initially installed with.

To use this script for a hands free install on any system specify a package
type of 'auto'. This will detect if yum or gdebi is present on the system
and select the one present if possible. If both rpm and deb package managers
are detected the automatic detection will abort
When using the automatic setup, if the system has apt-get but not gdebi,
the gdebi will be installed using apt-get first.

If --sanity-check is specified, the install script will wait for 3 minutes post installation
to check for a running agent.

To use a HTTP proxy, specify --proxy followed by the proxy server
defined by http://hostname:port

If --upgrade is specified, the script will only update the agent if a newer version is available.
Downgrades will not be respected.

If --downgrade is specified, the script will only update the agent if an older version of the
agent is marked as current. Upgrades will be ignored.

This install script needs Ruby version 2.x installed as a prerequisite.
Currently recommended Ruby versions are 2.0.0, 2.1.8, 2.2.4, 2.3, 2.4, 2.5, 2.6 and 2.7.
If multiple Ruby versions are installed, the default ruby version will be used.
If the default ruby version does not satisfy requirement, the newest version will be used.
If you do not have a supported Ruby version installed, please install one of them first.

EOF
  end

  def supported_ruby_versions
    ['2.7', '2.6', '2.5', '2.4', '2.3', '2.2', '2.1', '2.0']
  end

  # check ruby version, only version 2.x works
  def check_ruby_version_and_symlink
    @log.info("Starting Ruby version check.")
    actual_ruby_version = RUBY_VERSION.split('.').map{|s|s.to_i}[0,2]
    
    supported_ruby_versions.each do |version|
      if ((actual_ruby_version <=> version.split('.').map{|s|s.to_i}) == 0)
        return File.join(RbConfig::CONFIG["bindir"], RbConfig::CONFIG["RUBY_INSTALL_NAME"] + RbConfig::CONFIG["EXEEXT"])
      end
    end

    supported_ruby_versions.each do |version|
      if(File.exist?("/usr/bin/ruby#{version}"))
        return "/usr/bin/ruby#{version}"
      elsif (File.symlink?("/usr/bin/ruby#{version}"))
          @log.error("The symlink /usr/bin/ruby#{version} exists, but it's linked to a non-existent directory or non-executable file.")
          exit(1)
      end
    end

    unsupported_ruby_version_error    
    exit(1)
  end

  def unsupported_ruby_version_error
    @log.error("Current running Ruby version for "+ENV['USER']+" is "+RUBY_VERSION+", but Ruby version 2.x needs to be installed.")
    @log.error('If you already have the proper Ruby version installed, please either create a symlink to /usr/bin/ruby2.x,') 
    @log.error( "or run this install script with right interpreter. Otherwise please install Ruby 2.x for "+ENV['USER']+" user.")
    @log.error('You can get more information by running the script with --help option.')
  end

  def parse_args()
    if (ARGV.length > 4)
      usage
      @log.error('Too many arguments.')
      exit(1)
    elsif (ARGV.length < 1)
      usage
      @log.error('Expected package type as argument.')
      exit(1)
    end

    @sanity_check = false
    @reexeced = false
    @http_proxy = nil
    @downgrade = false
    @upgrade = false
    @target_version_arg = nil

    @args = Array.new(ARGV)
    opts = GetoptLong.new(['--sanity-check', GetoptLong::NO_ARGUMENT], ['--help', GetoptLong::NO_ARGUMENT],
                          ['--re-execed', GetoptLong::NO_ARGUMENT], ['--proxy', GetoptLong::OPTIONAL_ARGUMENT],
                          ['--downgrade', GetoptLong::NO_ARGUMENT], ['--upgrade', GetoptLong::NO_ARGUMENT],
                          ['-v', '--version', GetoptLong::OPTIONAL_ARGUMENT])
    opts.each do |opt, args|
      case opt
      when '--sanity-check'
        @sanity_check = true
      when '--help'
        usage
      when '--re-execed'
        @reexeced = true
      when '--downgrade'
        @downgrade = true
      when '--upgrade'
        @upgrade = true
      when '--proxy'
        if (args != '')
          @http_proxy = args
        end
      when '-v' || '--version'
        @target_version_arg = args
      end
    end

    if (@upgrade and @downgrade)
      usage
      @log.error('Cannot provide both --upgrade and --downgrade')
      exit(1)
    elsif (!@upgrade and !@downgrade)
      #Default to allowing both if one if neither is specified
      @upgrade = true
      @downgrade = true
    end


    if (ARGV.length < 1)
      usage
      @log.error('Expected package type as argument.')
      exit(1)
    end
    @type = ARGV.shift.downcase;
  end

  def force_ruby2x(ruby_interpreter_path)
    # change interpreter when symlink /usr/bin/ruby2.x exists, but running with non-supported ruby version
    actual_ruby_version = RUBY_VERSION.split('.').map{|s|s.to_i}
    left_bound = '2.0.0'.split('.').map{|s|s.to_i}
    right_bound = '2.7.0'.split('.').map{|s|s.to_i}
    if (actual_ruby_version <=> left_bound) < 0
      if(!@reexeced)
        @log.info("The current Ruby version is not 2.x! Restarting the installer with #{ruby_interpreter_path}")
        exec("#{ruby_interpreter_path}", __FILE__, '--re-execed' , *@args)
      else
        unsupported_ruby_version_error
        exit(1)
      end
    elsif ((actual_ruby_version <=> right_bound) > 0)
      @log.warn("The Ruby version in #{ruby_interpreter_path} is "+RUBY_VERSION+", . Attempting to install anyway.")
    end
  end

  def is_windows?
    is_windows = false

    begin
      require 'rbconfig'
      is_windows = (RbConfig::CONFIG['host_os'] =~ /mswin|mingw|cygwin/)
    rescue
    end

    is_windows
  end

  LOCAL_SERVICE_REGISTRY_KEY = 'S-1-5-19'
  def is_current_user_local_admin_windows?
    is_admin = false

    begin
      require 'win32/registry'

      # Best way to determine if admin which works on windows going all the way back to XP is
      # to check the LOCAL SERVICE account reg key
      Win32::Registry::HKEY_USERS.open(LOCAL_SERVICE_REGISTRY_KEY) {|reg| }
      is_admin = true
    rescue
    end

    is_admin
  end

  if (is_windows?)
    if (!is_current_user_local_admin_windows?)
      @log.error('Must run as user with Administrator privileges to update agent')
      exit(1)
    end
  else
    if (Process.uid != 0)
      @log.error('Must run as root to install packages')
      exit(1)
    end
  end

  parse_args()

  if @type == 'help'
    usage
    exit(0)
  end

  ########## Force running as Ruby 2.x or fail here       ##########
  ruby_interpreter_path = check_ruby_version_and_symlink
  force_ruby2x(ruby_interpreter_path)

  def run_command(*args)
    exit_ok = system(*args)
    $stdout.flush
    $stderr.flush
    @log.debug("Exit code: #{$?.exitstatus}")
    return exit_ok
  end

  def get_ec2_metadata_region
    begin
      return IMDSV2.region
    rescue => error
      @log.warn("Could not get region from EC2 metadata service at '#{error.message}'")
      return nil
    end
  end

  def get_region
    @log.info('Checking AWS_REGION environment variable for region information...')
    region = ENV['AWS_REGION']
    return region if region

    @log.info('Checking EC2 metadata service for region information...')
    region = get_ec2_metadata_region
    return region if region

    @log.info('Using fail-safe default region: us-east-1')
    return 'us-east-1'
  end

  def get_s3_uri(key)
    if (REGION == 'us-east-1')
      URI.parse("https://#{BUCKET}.s3.amazonaws.com/#{key}")
    elsif (REGION.split("-")[0] == 'cn')
      URI.parse("https://#{BUCKET}.s3.#{REGION}.amazonaws.com.cn/#{key}")
    else
      URI.parse("https://#{BUCKET}.s3.#{REGION}.amazonaws.com/#{key}")
    end
  end

  def get_package_from_s3(key, package_file)
    @log.info("Downloading package from BUCKET #{BUCKET} and key #{key}...")

    uri = get_s3_uri(key)

    # stream package file to disk
    begin
      uri.open(:ssl_verify_mode => OpenSSL::SSL::VERIFY_PEER, :redirect => true, :read_timeout => 120, :proxy => @http_proxy) do |s3|
        package_file.write(s3.read)
      end
    rescue OpenURI::HTTPError
      @log.error("Could not find package to download at '#{uri.to_s}'")
      exit(1)
    end
  end

  def setup_windows_certificates
    app_root_folder = File.join(ENV['PROGRAMDATA'], "Amazon/CodeDeploy")
    cert_dir = File.expand_path(File.join(app_root_folder, 'certs'))
    @log.info("Setting up windows certificates from cert directory #{cert_dir}")
    ENV['AWS_SSL_CA_DIRECTORY'] = File.join(cert_dir, 'ca-bundle.crt')
    ENV['SSL_CERT_FILE'] = File.join(cert_dir, 'ca-bundle.crt')
  end

  def get_version_file_from_s3
    @log.info("Downloading version file from BUCKET #{BUCKET} and key #{VERSION_FILE_KEY}...")

    uri = get_s3_uri(VERSION_FILE_KEY)

    begin
      require 'json'

      if (is_windows?)
        setup_windows_certificates
      end

      version_string = uri.read(:ssl_verify_mode => OpenSSL::SSL::VERIFY_PEER, :redirect => true, :read_timeout => 120, :proxy => @http_proxy)
      JSON.parse(version_string)
    rescue OpenURI::HTTPError
      @log.error("Could not find version file to download at '#{uri.to_s}'")
      exit(1)
    end
  end

  def install_from_s3(package_key, install_cmd, post_install_arguments=[])
    package_base_name = File.basename(package_key)
    package_extension = File.extname(package_base_name)
    package_name = File.basename(package_base_name, package_extension)
    package_file = File.new(File.join("#{Dir.tmpdir}","#{package_name}#{package_extension}"), "wb")

    get_package_from_s3(package_key, package_file)
    package_file.close

    install_cmd << package_file_path(package_file)
    install_cmd.concat(post_install_arguments)
    @log.info("Executing `#{install_cmd.join(" ")}`...")

    if (!run_command(*install_cmd))
      @log.error("Error installing #{package_file_path(package_file)}.")
      exit(1)
    end
  end

  def package_file_path(package_file)
    package_file_path = File.expand_path(package_file.path)

    if (is_windows?)
      #Flip slashes because in the command line shell it can only handle backwards slashes in windows
      package_file_path.gsub('/','\\')
    else
      package_file_path
    end
  end

  def do_sanity_check(cmd)
    if @sanity_check
      @log.info("Waiting for 3 minutes before I check for a running agent")
      sleep(3 * 60)
      res = run_command(cmd, 'codedeploy-agent', 'status')
      if (res.nil? || res == false)
        @log.info("No codedeploy agent seems to be running. Starting the agent.")
        run_command(cmd, 'codedeploy-agent', 'start-no-update')
      end
    end
  end

  @log.info("Starting update check.")

  if (@type == 'auto')
    @log.info('Attempting to automatically detect supported package manager type for system...')

    has_yum = run_command('which yum >/dev/null 2>/dev/null')
    has_apt_get = run_command('which apt-get >/dev/null 2>/dev/null')
    has_gdebi = run_command('which gdebi >/dev/null 2>/dev/null')
    has_zypper = run_command('which zypper >/dev/null 2>/dev/null')

    if (has_yum && (has_apt_get || has_gdebi))
      @log.error('Detected both supported rpm and deb package managers. Please specify which package type to use manually.')
      exit(1)
    end

    if(has_yum)
      @type = 'rpm'
    elsif(is_windows?)
      @type = 'msi'
    elsif(has_zypper)
      @type = 'zypper'
    elsif(has_gdebi)
      @type = 'deb'
    elsif(has_apt_get)
      @type = 'deb'

      @log.warn('apt-get found but no gdebi. Installing gdebi with `apt-get install gdebi-core -y`...')
      #use -y to answer yes to confirmation prompts
      if(!run_command('/usr/bin/apt-get', 'install', 'gdebi-core', '-y'))
        @log.error('Could not install gdebi.')
        exit(1)
      end
    else
      @log.error('Could not detect any supported package managers.')
      exit(1)
    end
  end

  unless VALID_TYPES.include? @type
    @log.error("Unsupported package type '#{@type}'")
    exit(1)
  end

  REGION = get_region
  BUCKET = "aws-codedeploy-#{REGION}"
  VERSION_FILE_KEY = 'latest/LATEST_VERSION'

  NO_AGENT_INSTALLED_REPORTED_WINDOWS_VERSION = 'No Agent Installed'
  def running_agent_version_windows
    installed_agent_versions_cmd_output = `wmic product where "name like 'CodeDeploy Host Agent'" get version`
    installed_agent_versions = installed_agent_versions_cmd_output.lines
      .collect{|line| line.strip}
      .reject{|line| line == 'Version'}
      .reject{|line| line.empty?}

    agent_version = installed_agent_versions.first
    #Example Agent Version Outputted from the above command: 1.0.1.1231
    if (/[0-9].[0-9].[0-9].[0-9]+/ =~ agent_version)
      return agent_version
    end

    NO_AGENT_INSTALLED_REPORTED_WINDOWS_VERSION
  end

  def upgrade_or_install_required?(target_version, running_version)
    running_version_numbers = version_numbers(running_version)
    @log.info("running_version_numbers: #{running_version_numbers}")

    if running_version_numbers == 'No running version' then return true end

    # detect returns the first number for which block is true, otherwise return nil
    version_numbers(target_version).zip(running_version_numbers).detect do |target_version_number, running_version_number|
      target_version_number.to_i > running_version_number.to_i
    end
  end

  def version_numbers(version)
    if match = version.match(/^.*(\d+)\.(\d+)[.-](\d+)\.(\d+).*$/i)
      match.captures
    else
      'No running version'
    end
  end

  def running_version(type)
    case type
    when 'rpm','zypper'
      `rpm --query codedeploy-agent`.strip
    when 'deb'
      running_agent = `dpkg --status codedeploy-agent`
      running_agent_info = running_agent.split
      version_index = running_agent_info.index('Version:')
      if !version_index.nil?
        running_agent_info[version_index + 1]
      else
        'No running version'
      end
    when 'msi'
      running_agent_version_windows
    else
      @log.error("Unsupported package type '#{@type}'")
      exit(1)
    end
  end

  def target_version(type)
    file_type = type == 'zypper' ? 'rpm' : type
    get_version_file_from_s3[file_type]
  end

  def install_command(type, upgrade_or_install_required)
    case @type
    when 'rpm'
      if upgrade_or_install_required
        ['/usr/bin/yum', '--assumeyes', 'localinstall']
      else
        ['/usr/bin/yum', '--assumeyes', 'downgrade']
      end
    when 'deb'
      if upgrade_or_install_required
        #use --option to not overwrite config files unless they have not been changed
        ['/usr/bin/gdebi', '--non-interactive', '--option=Dpkg::Options::=--force-confdef', '--option=Dkpg::Options::=--force-confold']
      else
        ['/usr/bin/dpkg', '--install']
      end
    when 'zypper'
      if upgrade_or_install_required
        ['/usr/bin/zypper', '--non-interactive', 'install']
      else
        ['/usr/bin/zypper', '--non-interactive', 'install', '--oldpackage']
      end
    when 'msi'
      ['msiexec','/quiet','/i']
    else
      @log.error("Unsupported package type '#{@type}'")
      exit(1)
    end
  end

  def pre_installation_steps(type, running_version)
    if type == 'msi'
      unless running_version == NO_AGENT_INSTALLED_REPORTED_WINDOWS_VERSION
        @log.info('Uninstalling old versions of the agent')
        uninstall_command_succeeded = system('wmic product where "name like \'CodeDeploy Host Agent\'" call uninstall /nointeractive')
        unless uninstall_command_succeeded
          @log.warn('Uninstalling existing agent failed')
        end
      end
    end
  end

  def post_install_arguments(type)
    if type == 'msi'
      ['/L*V',"#{Dir.tmpdir()}/codedeploy-agent.msi_installer.log"]
    else
      []
    end
  end

  running_version = running_version(@type)
  target_version = @target_version_arg
  if target_version.nil?
    target_version = target_version(@type)
  end
  if target_version.include? running_version
    @log.info("Running version, #{running_version}, matches target version, #{target_version}, skipping install")
  else
    if upgrade_or_install_required?(target_version, running_version)
      if @upgrade
        @log.info("Running version, #{running_version}, less than target version, #{target_version}, updating agent")
      else
        @log.info("New version available but only checking for downgrades. Skipping install.")
        exit 0;
      end
    else
      if @downgrade
        @log.info("Running version, #{running_version}, greater than target version, #{target_version}, rolling back agent")
      else
        @log.info("Older version available but only checking for upgrades. Skipping install.")
        exit 0;
      end
    end

    pre_installation_steps(@type, running_version)

    install_cmd = install_command(@type, upgrade_or_install_required?(target_version, running_version))
    post_install_args = post_install_arguments(@type)
    install_from_s3(target_version, install_cmd, post_install_args)

    unless @type == 'msi'
      do_sanity_check('/sbin/service')
    end
  end

  @log.info("Update check complete.")
  @log.info("Stopping updater.")

rescue SystemExit => e
  # don't log exit() as an error
  raise e
rescue Exception => e
  # make sure all unhandled exceptions are logged to the log
  @log.error("Unhandled exception: #{e.inspect}")
  e.backtrace.each do |line|
    @log.error("  at " + line)
  end
  exit(1)
end