From 833c8517065f2e8f0ad5228760d2d37793503bc1 Mon Sep 17 00:00:00 2001
From: Roberto Dip <me@roperzh.com>
Date: Wed, 26 Jul 2023 12:04:11 -0300
Subject: [PATCH] fetch the correct hiera information based on the node env
 (#12944)

For https://github.com/fleetdm/fleet/issues/12897
---
 changes/12897-puppet-envs                     |  1 +
 .../fleetdm/lib/puppet/util/fleet_client.rb   | 20 ++++++++++---------
 ee/tools/puppet/fleetdm/metadata.json         |  2 +-
 3 files changed, 13 insertions(+), 10 deletions(-)
 create mode 100644 changes/12897-puppet-envs

diff --git a/changes/12897-puppet-envs b/changes/12897-puppet-envs
new file mode 100644
index 0000000000..8d3727d73c
--- /dev/null
+++ b/changes/12897-puppet-envs
@@ -0,0 +1 @@
+* Allow the puppet module to read different Fleet URL/token combinations for different environments
diff --git a/ee/tools/puppet/fleetdm/lib/puppet/util/fleet_client.rb b/ee/tools/puppet/fleetdm/lib/puppet/util/fleet_client.rb
index 603b0ae19c..38f1d76973 100644
--- a/ee/tools/puppet/fleetdm/lib/puppet/util/fleet_client.rb
+++ b/ee/tools/puppet/fleetdm/lib/puppet/util/fleet_client.rb
@@ -25,13 +25,6 @@ module Puppet::Util
     end
 
     def initialize
-      node_name = Puppet[:node_name_value]
-      node = Puppet::Node.new(node_name)
-      compiler = Puppet::Parser::Compiler.new(node)
-      scope = Puppet::Parser::Scope.new(compiler)
-      lookup_invocation = Puppet::Pops::Lookup::Invocation.new(scope, {}, {}, nil)
-      @host = Puppet::Pops::Lookup.lookup('fleetdm::host', nil, '', false, nil, lookup_invocation)
-      @token = Puppet::Pops::Lookup.lookup('fleetdm::token', nil, '', false, nil, lookup_invocation)
       @cache = {}
       @cache_mutex = Mutex.new
     end
@@ -114,6 +107,15 @@ module Puppet::Util
     private
 
     def req(method: :get, path: '', body: nil, headers: {}, cached: false)
+      node_name = Puppet[:node_name_value]
+      node = Puppet::Node.new(node_name)
+      node.environment = Puppet.lookup(:current_environment).name.to_s
+      compiler = Puppet::Parser::Compiler.new(node)
+      scope = Puppet::Parser::Scope.new(compiler)
+      lookup_invocation = Puppet::Pops::Lookup::Invocation.new(scope, {}, {}, nil)
+      host = Puppet::Pops::Lookup.lookup('fleetdm::host', nil, '', false, nil, lookup_invocation)
+      token = Puppet::Pops::Lookup.lookup('fleetdm::token', nil, '', false, nil, lookup_invocation)
+
       if cached
         @cache_mutex.synchronize do
           unless @cache[path].nil?
@@ -123,7 +125,7 @@ module Puppet::Util
       end
 
       out = { 'error' => '' }
-      uri = URI.parse("#{@host}#{path}")
+      uri = URI.parse("#{host}#{path}")
       uri.path.squeeze! '/'
       uri.path.chomp! '/'
 
@@ -139,7 +141,7 @@ module Puppet::Util
         throw "HTTP method #{method} not implemented"
       end
 
-      headers['Authorization'] = "Bearer #{@token}"
+      headers['Authorization'] = "Bearer #{token}"
       headers.each { |key, value| request[key] = value }
       request.body = body.to_json if body
 
diff --git a/ee/tools/puppet/fleetdm/metadata.json b/ee/tools/puppet/fleetdm/metadata.json
index 17e0d1f44e..f648792730 100644
--- a/ee/tools/puppet/fleetdm/metadata.json
+++ b/ee/tools/puppet/fleetdm/metadata.json
@@ -1,6 +1,6 @@
 {
   "name": "fleetdm-fleetdm",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "author": "Fleet Device Management Inc",
   "summary": "MDM management and profile assignment using FleetDM",
   "license": "proprietary",