1 files changed, 85 insertions, 5 deletions

diff --git a/lib/leap_cli/config/macros.rb b/lib/leap_cli/config/macros.rb
index 2c1f1bd..59453b0 100644
--- a/lib/leap_cli/config/macros.rb
+++ b/lib/leap_cli/config/macros.rb
@@ -38,6 +38,59 @@ module LeapCli; module Config
       nodes[:environment => @node.environment]
     end
 
+    #
+    # returns a list of nodes that match the location name
+    # and environment of @node.
+    #
+    def nodes_near_me
+      if @node['location'] && @node['location']['name']
+        nodes_like_me['location.name' => @node.location.name]
+      else
+        nodes_like_me['location' => nil]
+      end
+    end
+
+    #
+    #
+    # picks a node out from the node list in such a way that:
+    #
+    # (1) which nodes picked which nodes is saved in secrets.json
+    # (2) when other nodes call this macro with the same node list, they are guaranteed to get a different node
+    # (3) if all the nodes in the pick_node list have been picked, remaining nodes are distributed randomly.
+    #
+    # if the node_list is empty, an exception is raised.
+    # if node_list size is 1, then that node is returned and nothing is
+    # memorized via the secrets.json file.
+    #
+    # `label` is needed to distinguish between pools of nodes for different purposes.
+    #
+    # TODO: more evenly balance after all the nodes have been picked.
+    #
+    def pick_node(label, node_list)
+      if node_list.any?
+        if node_list.size == 1
+          return node_list.values.first
+        else
+          secrets_key = "pick_node(:#{label},#{node_list.keys.sort.join(',')})"
+          secrets_value = @manager.secrets.retrieve(secrets_key, @node.environment) || {}
+          secrets_value[@node.name] ||= begin
+            node_to_pick = nil
+            node_list.each_node do |node|
+              next if secrets_value.values.include?(node.name)
+              node_to_pick = node.name
+            end
+            node_to_pick ||= secrets_value.values.shuffle.first # all picked already, so pick a random one.
+            node_to_pick
+          end
+          picked_node_name = secrets_value[@node.name]
+          @manager.secrets.set(secrets_key, secrets_value, @node.environment)
+          return node_list[picked_node_name]
+        end
+      else
+        raise ArgumentError.new('pick_node(node_list): node_list cannot be empty')
+      end
+    end
+
     ##
     ## FILES
     ##
@@ -167,9 +220,7 @@ module LeapCli; module Config
     #
     def hostnames(nodes)
       @referenced_nodes ||= ObjectList.new
-      if nodes.is_a? Config::Object
-        nodes = ObjectList.new nodes
-      end
+      nodes = listify(nodes)
       nodes.each_node do |node|
         @referenced_nodes[node.name] ||= node
       end
@@ -248,6 +299,7 @@ module LeapCli; module Config
     #
     def stunnel_client(node_list, port, options={})
       @next_stunnel_port ||= 4000
+      node_list = listify(node_list)
       hostnames(node_list) # record the hosts
       result = Config::ObjectList.new
       node_list.each_node do |node|
@@ -420,9 +472,22 @@ module LeapCli; module Config
     end
 
     #
-    # wrap something that might fail in `try`. e.g.
+    # applies a JSON partial to this node
+    #
+    def apply_partial(partial_path)
+      manager.partials(partial_path).each do |partial_data|
+        self.deep_merge!(partial_data)
+      end
+    end
+
+    #
+    # If at first you don't succeed, then it is time to give up.
     #
-    # "= try{ nodes[:services => 'tor'].first.ip_address } "
+    # try{} returns nil if anything in the block throws an exception.
+    #
+    # You can wrap something that might fail in `try`, like so.
+    #
+    #   "= try{ nodes[:services => 'tor'].first.ip_address } "
     #
     def try(&block)
       yield
@@ -430,5 +495,20 @@ module LeapCli; module Config
       nil
     end
 
+    private
+
+    #
+    # returns a node list, if argument is not already one
+    #
+    def listify(node_list)
+      if node_list.is_a? Config::ObjectList
+        node_list
+      elsif node_list.is_a? Config::Object
+        Config::ObjectList.new(node_list)
+      else
+        raise ArgumentError, 'argument must be a node or node list, not a `%s`' % node_list.class, caller
+      end
+    end
+
   end
 end; end