[CHEF-18664] Updated the shell-init command to work with habitat and omnibus installations

Gemfile

@@ -11,7 +11,8 @@ group :test do
   gem "rspec-mocks", "~> 3.8"
   gem "cookstyle"
   gem "chefstyle"
-  gem "test-kitchen"
+  gem "faraday_middleware"
+  gem "chef-test-kitchen-enterprise", git: "https://github.com/chef/chef-test-kitchen-enterprise", branch: "main"
   gem "simplecov", require: false
 end
 

chef-cli.gemspec

@@ -50,5 +50,5 @@ Gem::Specification.new do |gem|
   gem.add_dependency "diff-lcs", ">= 1.0", "< 1.4" # 1.4 changes the output
   gem.add_dependency "pastel", "~> 0.7" # used for policyfile differ
   gem.add_dependency "license-acceptance", ">= 1.0.11", "< 3"
-  gem.add_dependency "chef-licensing", ">= 1.0.2"
+  gem.add_dependency "chef-licensing", "~> 1.0"
 end

habitat/plan.ps1

@@ -47,9 +47,9 @@ function Invoke-Build {
         bundle install
 
         gem build chef-cli.gemspec
-	    Write-BuildLine " ** Using gem to  install"
-	    gem install chef-cli-*.gem --no-document
-
+        Write-BuildLine " ** Using gem to  install"
+        gem install chef-cli-*.gem --no-document
+        ruby ./post-bundle-install.rb
 
         If ($lastexitcode -ne 0) { Exit $lastexitcode }
     } finally {

habitat/plan.sh

@@ -9,6 +9,7 @@ pkg_build_deps=(
     core/sed
     core/gcc
     core/libarchive
+    core/git
     )
 pkg_bin_dirs=(bin)
 
@@ -42,6 +43,7 @@ do_build() {
     bundle config --local silence_root_warning 1
     bundle install
     gem build chef-cli.gemspec
+    ruby ./post-bundle-install.rb
 }
 do_install() {
    export GEM_HOME="$pkg_prefix/vendor"

lib/chef-cli/command/env.rb

@@ -85,7 +85,8 @@ def gem_environment
       end
 
       def paths
-        omnibus_env["PATH"].split(File::PATH_SEPARATOR)
+        env = habitat_install? ? habitat_env : omnibus_env
+        env["PATH"].split(File::PATH_SEPARATOR)
       rescue OmnibusInstallNotFound
         ENV["PATH"].split(File::PATH_SEPARATOR)
       end

lib/chef-cli/command/exec.rb

@@ -30,7 +30,8 @@ class Exec < ChefCLI::Command::Base
       def run(params)
         # Set ENV directly on the "parent" process (us) before running #exec to
         # ensure the custom PATH is honored when finding the command to exec
-        omnibus_env.each { |var, value| ENV[var] = value }
+        env = habitat_install? ? habitat_env : omnibus_env
+        env.each { |var, value| ENV[var] = value }
         exec(*params)
         raise "Exec failed without an exception, your ruby is buggy" # should never get here
       end

lib/chef-cli/command/license.rb

@@ -75,6 +75,8 @@ def run(params)
         else
           ChefCLI::Licensing::Base.send(remaining_args[0])
         end
+      rescue ChefLicensing::LicenseKeyFetcher::LicenseKeyNotFetchedError
+        ui.msg("License key not fetched. Please try again.")
       end
 
       def debug?

lib/chef-cli/command/shell_init.rb

@@ -29,6 +29,7 @@ module Mixlib
 module ChefCLI
 
   class ShellCompletionTemplateContext
+    include ChefCLI::Helpers
 
     def commands
       ChefCLI.commands_map.command_specs.inject({}) do |cmd_info, (_key, cmd_spec)|
@@ -40,6 +41,10 @@ def commands
     def get_binding
       binding
     end
+
+    def habitat?
+      habitat_install?
+    end
   end
 
   module Command
@@ -103,7 +108,7 @@ def run(argv)
           return 1
         end
 
-        env = omnibus_env.dup
+        env = (habitat_install? ? habitat_env : omnibus_env).dup
         path = env.delete("PATH")
         export(shell_name, "PATH", path)
         env.each do |var_name, value|

lib/chef-cli/completions/bash.sh.erb

@@ -2,4 +2,4 @@ _chef_comp() {
     local COMMANDS="<%= commands.keys.join(' ')-%>"
     COMPREPLY=($(compgen -W "$COMMANDS" -- ${COMP_WORDS[COMP_CWORD]} ))
 }
-complete -F _chef_comp chef
+complete -F _chef_comp <%= habitat? ? "chef-cli" : "chef" %>

lib/chef-cli/completions/chef.fish.erb

@@ -5,5 +5,5 @@
 set -l chef_commands <%= commands.keys.join(' ') %>;
 
 <% commands.each do |command, desc| -%>
-complete -c chef -f -n "not __fish_seen_subcommand_from $chef_commands" -a <%= command %> -d "<%= desc %>";
+complete -c <%= habitat? ? "chef-cli" : "chef" %> -f -n "not __fish_seen_subcommand_from $chef_commands" -a <%= command %> -d "<%= desc %>";
 <% end -%>

lib/chef-cli/completions/zsh.zsh.erb

@@ -17,5 +17,5 @@ function _chef() {
   fi
 }
 
-compdef _chef chef
+compdef _chef <%= habitat? ? "chef-cli" : "chef" %>
 

lib/chef-cli/dist.rb

@@ -10,6 +10,9 @@ class Dist
     CLI_PRODUCT = "Chef CLI".freeze
     CLI_GEM = "chef-cli".freeze
 
+    CHEF_DKE_PKG_NAME = "chef/chef-development-kit-enterprise".freeze
+    HAB_PKG_NAME = "chef/chef-cli".freeze
+
     # the name of the overall infra product
     INFRA_PRODUCT = "Chef Infra".freeze
 

lib/chef-cli/helpers.rb

@@ -60,6 +60,21 @@ def omnibus_install?
       File.exist?(expected_omnibus_root) && File.exist?(File.join(expected_omnibus_root, "version-manifest.json"))
     end
 
+    # The habitat version of the chef-cli can be installed with standalone or chef-development-kit-enterprise
+    # This method checks if the habitat version of chef-cli is installed as standalone
+    def habitat_standalone?
+      @hab_standalone ||= (hab_pkg_installed?(ChefCLI::Dist::HAB_PKG_NAME) && !habitat_chef_dke?)
+    end
+
+    # This method checks if the habitat version of chef-cli is installed with chef-development-kit-enterprise
+    def habitat_chef_dke?
+      @hab_dke ||= hab_pkg_installed?(ChefCLI::Dist::CHEF_DKE_PKG_NAME)
+    end
+
+    def habitat_install?
+      habitat_chef_dke? || habitat_standalone?
+    end
+
     def omnibus_root
       @omnibus_root ||= omnibus_expand_path(expected_omnibus_root)
     end
@@ -110,29 +125,30 @@ def git_windows_bin_dir
       @git_windows_bin_dir ||= File.expand_path(File.join(omnibus_root, "embedded", "git", "usr", "bin"))
     end
 
-    # #
-    # # environment vars for habitat
-    # #
-    # def habitat_env
-    #   @habitat_env ||=
-    #   begin
-    #     # Define the necessary paths for the Habitat environment
-    #     # Custom GEM_HOME within Habitat
-    #     pkg_prefix = get_pkg_prefix
-    #     vendor_dir = File.join(pkg_prefix, "vendor")
-    #     path = [
-    #       File.join(pkg_prefix, "bin"),
-    #       ENV["PATH"].split(File::PATH_SEPARATOR), # Preserve existing PATH
-    #     ].flatten.uniq
-
-    #     {
-    #     "PATH" => path.join(File::PATH_SEPARATOR),
-    #     "GEM_ROOT" => Gem.default_dir, # Default directory for gems
-    #     "GEM_HOME" => vendor_dir,  # GEM_HOME pointing to the vendor directory
-    #     "GEM_PATH" => vendor_dir,  # GEM_PATH also pointing to the vendor directory
-    #     }
-    #   end
-    # end
+    #
+    # environment vars for habitat
+    #
+    def habitat_env
+      @habitat_env ||=
+      begin
+        # Define the necessary paths for the Habitat environment
+        # If it is a chef-dke installation, we will use the chef-dke bin path.
+        # Otherwise, we will use the chef-cli bin path.
+        bin_pkg_prefix = get_pkg_prefix(habitat_chef_dke? ? ChefCLI::Dist::CHEF_DKE_PKG_NAME : ChefCLI::Dist::HAB_PKG_NAME)
+        vendor_dir = File.join(get_pkg_prefix(ChefCLI::Dist::HAB_PKG_NAME), "vendor")
+        path = [
+          File.join(bin_pkg_prefix, "bin"),
+          ENV["PATH"].split(File::PATH_SEPARATOR), # Preserve existing PATH
+        ].flatten.uniq
+
+        {
+        "PATH" => path.join(File::PATH_SEPARATOR),
+        "GEM_ROOT" => Gem.default_dir, # Default directory for gems
+        "GEM_HOME" => vendor_dir,  # GEM_HOME pointing to the vendor directory
+        "GEM_PATH" => vendor_dir,  # GEM_PATH also pointing to the vendor directory
+        }
+      end
+    end
 
     #
     # environment vars for omnibus
@@ -153,17 +169,15 @@ def omnibus_env
         end
     end
 
-    # def get_pkg_prefix
-    #   pkg_origin = "chef"
-    #   pkg_name = "#{pkg_origin}/chef-cli" # Your origin and package name
-    #   path = `hab pkg path #{pkg_name}`.strip
+    def get_pkg_prefix(pkg_name)
+      path = `hab pkg path #{pkg_name}`.strip
 
-    #   if $?.success? && !path.empty?
-    #     path
-    #   else
-    #     raise "Failed to get pkg_prefix for #{pkg_name}: #{path}"
-    #   end
-    # end
+      if $?.success? && !path.empty?
+        path
+      else
+        raise "Failed to get pkg_prefix for #{pkg_name}: #{path}"
+      end
+    end
 
     def omnibus_expand_path(*paths)
       dir = File.expand_path(File.join(paths))
@@ -209,5 +223,14 @@ def reset!
     def macos?
       !!(RUBY_PLATFORM =~ /darwin/)
     end
+
+    # @return [Boolean] Checks if a habitat package is installed.
+    # If habitat itself is not installed, this method will return false.
+    #
+    # @api private
+    #
+    def hab_pkg_installed?(pkg_name)
+      `hab pkg list #{pkg_name} 2>/dev/null`.include?(pkg_name) rescue false
+    end
   end
 end

lib/chef-cli/licensing/config.rb

@@ -22,6 +22,5 @@
   config.chef_product_name = "workstation"
   config.chef_entitlement_id = "x6f3bc76-a94f-4b6c-bc97-4b7ed2b045c0"
   config.chef_executable_name = "chef"
-  # config.license_server_url = "https://services.chef.io/licensing"
-  config.license_server_url   = "https://licensing-acceptance.chef.co/License"
+  config.license_server_url = "https://services.chef.io/licensing"
 end

post-bundle-install.rb

@@ -0,0 +1,30 @@
+#!/usr/bin/env ruby
+
+gem_home = Gem.paths.home
+
+puts "fixing bundle installed gems in #{gem_home}"
+
+# Install gems from git repos.  This makes the assumption that there is a <gem_name>.gemspec and
+# you can simply gem build + gem install the resulting gem, so nothing fancy.  This does not use
+# rake install since we need --conservative --minimal-deps in order to not install duplicate gems.
+#
+#
+puts "gem path #{gem_home}"
+
+Dir["#{gem_home}/bundler/gems/*"].each do |gempath|
+  puts "#{gempath}"
+  matches = File.basename(gempath).match(/.*-[A-Fa-f0-9]{12}/)
+  next unless matches
+
+  gem_name = File.basename(Dir["#{gempath}/*.gemspec"].first, ".gemspec")
+  # FIXME: should strip any valid ruby platform off of the gem_name if it matches
+
+  next unless gem_name
+
+  puts "re-installing #{gem_name}..."
+
+  Dir.chdir(gempath) do
+    system("gem build #{gem_name}.gemspec") or raise "gem build failed"
+    system("gem install #{gem_name}*.gem --conservative --minimal-deps --no-document") or raise "gem install failed"
+  end
+end

spec/unit/command/env_spec.rb

@@ -35,6 +35,7 @@
 
   describe "when running from within an omnibus install" do
     before do
+      allow(command_instance).to receive(:habitat_install?).and_return false
       allow(command_instance).to receive(:omnibus_install?).and_return true
       allow(command_instance).to receive(:omnibus_embedded_bin_dir).and_return(omnibus_embedded_bin_dir)
       allow(command_instance).to receive(:omnibus_bin_dir).and_return(omnibus_bin_dir)
@@ -57,6 +58,7 @@
   end
   describe "when running locally" do
     before do
+      allow(command_instance).to receive(:habitat_install?).and_return false
       allow(command_instance).to receive(:omnibus_install?).and_return false
       command_instance.ui = ui
     end

spec/unit/command/exec_spec.rb

@@ -57,6 +57,7 @@ def run_command
     let(:ruby_path) { File.join(fixtures_path, "eg_omnibus_dir/valid/embedded/bin/ruby") }
 
     before do
+      allow(command_instance).to receive(:habitat_install?).and_return(false)
       allow(Gem).to receive(:ruby).and_return(ruby_path)
 
       # Using a fake path separator to keep to prevent people from accidentally