text: """
  Utilities to manage salt bootstrap and other stuff
  """
  import copy
  import os
  import logging
  import itertools
  import threading
  
  import dask_ec2
  from dask_ec2.utils import retry
  from dask_ec2.exceptions import DaskEc2Exception, RetriesExceededException
  
  logger = logging.getLogger(__name__)
  
  
  class Response(dict):
      """Response from a Salt Command
      Format is based on salt output but is flattened by compute_id, example:
      ```
      {'compute1': [state_1_output, state_2_output], 'compute2': [state_1_output]}
      ```
      Where `state_x_output` is: `{'name': 'my_state', 'result': True}`
      """
  
      def __init__(self, *args, **kwargs):
          super(Response, self).__init__(*args, **kwargs)
  
      @classmethod
      def from_dict(cls, data):
          self = cls()
          data = data["return"]
          for item in data:
              for minion_id, states in item.items():
                  self[minion_id] = states
          return self
  
      def aggregate_by(self, field='result', validation=True):
          """
          Useful when the Command module returns a dictionary, for example state.sls
          Default values are for salt module `state.sls`
          """
          try:
              ret = Response()
              for minion_id, values in self.items():
                  inner_values = []
                  if type(values) == dict:
                      # Assumes: depth=1 going to flat
                      for key, value in values.items():
                          value['name'] = key
                          inner_values.append(value)
                  elif type(self[minion_id]) == list:
                      inner_values = self[minion_id]
  
                  successful = [action for action in inner_values if action[field] == validation]
                  failed = [action for action in inner_values if action[field] != validation]
                  summary = {'successful': successful, 'failed': failed}
                  ret[minion_id] = summary
                  return ret
          except TypeError:
              logger.debug("Error with salt state.  Printing full returned output")
              logger.debug(self)
  
      def aggregated_to_table(self, agg=None):
          """
          From an aggregated response it returns a list of 3 columns:
          1. node id
          2. successful states
          3. failed states
          """
          ret = []
          for node_id, data in self.items():
              if agg:
                  ret.append([node_id, agg(data["successful"]), agg(data["failed"])])
              else:
                  ret.append([node_id, data["successful"], data["failed"]])
          return ret
  
      def aggregated_success(self):
          """
          From an aggregated return True if all the states where successful
          """
          ret = True
          for node_id, data in self.items():
              failed = data["failed"]
              ret = ret & (len(failed) == 0)
          return ret
  
      def group_by_id(self, ignore_fields=None, sort=True):
          if ignore_fields:
              clone = copy.deepcopy(self)
              for id_ in clone:
                  for field in ignore_fields:
                      del clone[id_][field]
          else:
              clone = self
  
          items = sorted(clone.items(), key=lambda x: x[1])
          groups = []
          for key, group in itertools.groupby(items, key=lambda x: x[1]):
              groups.append((key, [item[0] for item in group]))
  
          if sort:
              groups = sorted(groups, key=lambda x: len(x[1]), reverse=True)
          return groups
  
  
  def install_salt_master(cluster):
      master = cluster.instances[0].ssh_client
      dask_ec2_src = os.path.realpath(os.path.dirname(dask_ec2.__file__))
      templates_src = os.path.join(dask_ec2_src, "templates")
  
      @retry(retries=3, wait=0)
      def __install_salt_master():
          cmd = "curl -sS -L https://bootstrap.saltstack.com | sh -s -- -d -X -M -N stable"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
          return True
  
      try:
          __install_salt_master()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't bootstrap salt-master. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __install_salt_api():
          cmd = "curl -L https://bootstrap.saltstack.com | sh -s -- -d -X -M -N -P -L -p salt-api stable"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
  
      try:
          __install_salt_api()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't bootstrap salt-api. Error is above (maybe try again)" % e.last_exception)
  
      @retry(retries=3, wait=0)
      def __setup_salt_master():
          local = os.path.join(templates_src, "auto_accept.conf")
          remote = "/etc/salt/master.d/auto_accept.conf"
          master.put(local, remote, sudo=True)
  
      try:
          __setup_salt_master()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't setup salt-master settings. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __apt_installs():
          cmd = "apt-get install -y python-pip libssl-dev libffi-dev python-dev"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
  
      try:
          __apt_installs()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't install ubuntu dependencies. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __upgrade_pip():
          cmd = "pip install --upgrade pip packaging appdirs six"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"])
  
      try:
          __upgrade_pip()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't upgrade pip. Error is above (maybe try again)" % e.last_exception)
  
      @retry(retries=3, wait=0)
      def __install_salt_rest_api():
          cmd = "pip install cherrypy"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
  
      try:
          __install_salt_rest_api()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't install CherryPy. Error is above (maybe try again)" % e.last_exception)
  
      @retry(retries=3, wait=0)
      def __install_pyopensll():
          cmd = "pip install PyOpenSSL==16.2.0"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"])
  
      try:
          __install_pyopensll()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't install PyOpenSSL. Error is above (maybe try again)" % e.last_exception)
  
      @retry(retries=3, wait=0)
      def __create_ssl_cert():
          cmd = "salt-call --local tls.create_self_signed_cert"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
  
      try:
          __create_ssl_cert()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't generate SSL certificate. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __setup_rest_cherrypy():
          local = os.path.join(templates_src, "rest_cherrypy.conf")
          remote = "/etc/salt/master.d/rest_cherrypy.conf"
          master.put(local, remote, sudo=True)
  
      try:
          __setup_rest_cherrypy()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't setup salt-rest server. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __setup_salt_external_auth():
          local = os.path.join(templates_src, "external_auth.conf")
          remote = "/etc/salt/master.d/external_auth.conf"
          master.put(local, remote, sudo=True)
  
      try:
          __setup_salt_external_auth()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't setup salt external auth system. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __create_saltdev_user():
          cmd = "id -u saltdev &>/dev/null || useradd -p $(openssl passwd -1 saltdev) saltdev"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
  
      try:
          __create_saltdev_user()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't create 'saltdev' user. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __restart_salt_master():
          cmd = "service salt-master restart"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
  
      try:
          __restart_salt_master()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't restart salt-master service. Error is above (maybe try again)" %
                                 e.last_exception)
  
      @retry(retries=3, wait=0)
      def __restart_salt_api():
          cmd = "service salt-api restart"
          ret = master.exec_command(cmd, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
  
      try:
          __restart_salt_api()
      except RetriesExceededException as e:
          raise DaskEc2Exception("%s\nCouldn't restart salt-api service. Error is above (maybe try again)" %
                                 e.last_exception)
  
  
  def async_cmd(results, instance, command):
      client = instance.ssh_client
  
      @retry(retries=3, wait=0)
      def __remote_cmd():
          ret = client.exec_command(command, sudo=True)
          if ret["exit_code"] != 0:
              raise Exception(ret["stderr"].decode('utf-8'))
          return ret
  
      try:
          results[instance.ip] = __remote_cmd()
      except RetriesExceededException:
          results[instance.ip] = False
  
  
  def async_upload(results, instance, local, remote):
      client = instance.ssh_client
  
      @retry(retries=3, wait=0)
      def __remote_upload():
          client.put(local, remote, sudo=True)
          return True
  
      try:
          results[instance.ip] = __remote_upload()
      except RetriesExceededException:
          results[instance.ip] = False
  
  
  def install_salt_minion(cluster):
      dask_ec2_src = os.path.realpath(os.path.dirname(dask_ec2.__file__))
      templates_src = os.path.join(dask_ec2_src, "templates")
  
      logger.debug("Installing salt-minion on all the nodes")
      results, threads = {}, []
      master_ip = cluster.instances[0].ip
      for i, instance in enumerate(cluster.instances):
          minion_id = "node-{}".format(i)
          cmd = "curl -L https://bootstrap.saltstack.com | sh -s -- "
          cmd += "-d -X -P -L -A {master_ip} -i {minion_id} stable".format(master_ip=master_ip, minion_id=minion_id)
          t = threading.Thread(target=async_cmd, args=(results, instance, cmd))
          t.start()
          threads.append(t)
      for t in threads:
          t.join()
  
      all_ok = all([r is False for r in results])
      if not all_ok:
          failed_nodes = []
          for minion_ip, minion_data in results.items():
              if minion_data is False:
                  failed_nodes.append(minion_ip)
          if failed_nodes:
              raise DaskEc2Exception("Error bootstrapping salt-minion at nodes: %s (maybe try again)" % failed_nodes)
  
      logger.debug("Configuring salt-mine on the salt minions")
      results, threads = {}, []
      for i, instance in enumerate(cluster.instances):
          local = os.path.join(templates_src, "mine_functions.conf")
          remote = "/etc/salt/minion.d/mine.conf"
          t = threading.Thread(target=async_upload, args=(results, instance, local, remote))
          t.start()
          threads.append(t)
      for t in threads:
          t.join()
  
      all_ok = all([r is False for r in results])
      if not all_ok:
          failed_nodes = []
          for minion_ip, minion_data in results.items():
              if minion_data is False:
                  failed_nodes.append(minion_ip)
          if failed_nodes:
              raise DaskEc2Exception("Error configuring the salt-mine in the salt-minion at nodes: %s (maybe try again)" %
                                     failed_nodes)
  
      logger.debug("Restarting the salt-minion service")
      results, threads = {}, []
      for i, instance in enumerate(cluster.instances):
          cmd = "service salt-minion restart"
          t = threading.Thread(target=async_cmd, args=(results, instance, cmd))
          t.start()
          threads.append(t)
      for t in threads:
          t.join()
  
      all_ok = all([r is False for r in results])
      if not all_ok:
          failed_nodes = []
          for minion_ip, minion_data in results.items():
              if minion_data is False:
                  failed_nodes.append(minion_ip)
          if failed_nodes:
              raise DaskEc2Exception("Error restarting the salt-minion at nodes: %s (maybe try again)" % failed_nodes)
  
  
  def upload_formulas(cluster):
      dask_ec2_src = os.path.realpath(os.path.dirname(dask_ec2.__file__))
      src_salt_root = os.path.join(dask_ec2_src, "formulas", "salt")
      src_pillar_root = os.path.join(dask_ec2_src, "formulas", "pillar")
      dst_salt_root = "/srv/salt"
      dst_pillar_root = "/srv/pillar"
  
      client = cluster.instances[0].ssh_client
      client.put(src_salt_root, dst_salt_root, sudo=True)
      client.put(src_pillar_root, dst_pillar_root, sudo=True)
  
  
  def upload_pillar(cluster, name, data):
      import os
      import yaml
      import tempfile
  
      master = cluster.instances[0].ssh_client
      f = tempfile.NamedTemporaryFile("w", delete=False)
      try:
          yaml.safe_dump(data, f, default_flow_style=False)
          f.close()
          local = f.name
          remote = "/srv/pillar/{}".format(name)
          master.put(local, remote, sudo=True)
      finally:
          os.remove(f.name)
  