text: # Copyright 2015-2016 Yelp Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  import abc
  import collections
  import csv
  import socket
  from typing import cast
  from typing import Collection
  from typing import Container
  from typing import DefaultDict
  from typing import Dict
  from typing import Iterable
  from typing import List
  from typing import NamedTuple
  from typing import Optional
  from typing import Sequence
  from typing import Tuple
  from typing import TypeVar
  
  import requests
  from kubernetes.client import V1Node
  from mypy_extensions import TypedDict
  
  from paasta_tools import kubernetes_tools
  from paasta_tools import marathon_tools
  from paasta_tools import mesos_tools
  from paasta_tools.mesos.exceptions import NoSlavesAvailableError
  from paasta_tools.utils import compose_job_id
  from paasta_tools.utils import DEFAULT_SOA_DIR
  from paasta_tools.utils import DeployBlacklist
  from paasta_tools.utils import get_user_agent
  from paasta_tools.utils import InstanceConfig
  from paasta_tools.utils import SystemPaastaConfig
  
  
  HaproxyBackend = TypedDict(
      'HaproxyBackend',
      {
          'pxname': str,
          'svname': str,
          'status': str,
      },
      total=False,
  )
  
  
  def retrieve_haproxy_csv(
      synapse_host: str,
      synapse_port: int,
      synapse_haproxy_url_format: str,
  ) -> Iterable[Dict[str, str]]:
      """Retrieves the haproxy csv from the haproxy web interface
  
      :param synapse_host_port: A string in host:port format that this check
                                should contact for replication information.
      :returns reader: a csv.DictReader object
      """
      synapse_uri = synapse_haproxy_url_format.format(host=synapse_host, port=synapse_port)
  
      # timeout after 1 second and retry 3 times
      haproxy_request = requests.Session()
      haproxy_request.headers.update({'User-Agent': get_user_agent()})
      haproxy_request.mount(
          'http://',
          requests.adapters.HTTPAdapter(max_retries=3),
      )
      haproxy_request.mount(
          'https://',
          requests.adapters.HTTPAdapter(max_retries=3),
      )
      haproxy_response = haproxy_request.get(synapse_uri, timeout=1)
      haproxy_data = haproxy_response.text
      reader = csv.DictReader(haproxy_data.splitlines())
      return reader
  
  
  def get_backends(
      service: str,
      synapse_host: str,
      synapse_port: int,
      synapse_haproxy_url_format: str,
  ) -> List[HaproxyBackend]:
      """Fetches the CSV from haproxy and returns a list of backends,
      regardless of their state.
  
      :param service: If None, return backends for all services, otherwise only return backends for this particular
                      service.
      :param synapse_host_port: A string in host:port format that this check
                                should contact for replication information.
      :returns backends: A list of dicts representing the backends of all
                         services or the requested service
      """
      if service:
          services = [service]
      else:
          services = None
      return get_multiple_backends(
          services, synapse_host=synapse_host, synapse_port=synapse_port,
          synapse_haproxy_url_format=synapse_haproxy_url_format,
      )
  
  
  def get_multiple_backends(
      services: Optional[Container[str]],
      synapse_host: str,
      synapse_port: int,
      synapse_haproxy_url_format: str,
  ) -> List[HaproxyBackend]:
      """Fetches the CSV from haproxy and returns a list of backends,
      regardless of their state.
  
      :param services: If None, return backends for all services, otherwise only return backends for these particular
                       services.
      :param synapse_host_port: A string in host:port format that this check
                                should contact for replication information.
      :returns backends: A list of dicts representing the backends of all
                         services or the requested service
      """
  
      reader = retrieve_haproxy_csv(synapse_host, synapse_port, synapse_haproxy_url_format=synapse_haproxy_url_format)
      backends = []
  
      for line in reader:
          # clean up two irregularities of the CSV output, relative to
          # DictReader's behavior there's a leading "# " for no good reason:
          line['pxname'] = line.pop('# pxname')
          # and there's a trailing comma on every line:
          line.pop('')
  
          # Look for the service in question and ignore the fictional
          # FRONTEND/BACKEND hosts, use starts_with so that hosts that are UP
          # with 1/X healthchecks to go before going down get counted as UP:
          ha_slave, ha_service = line['svname'], line['pxname']
          if (services is None or ha_service in services) and ha_slave not in ('FRONTEND', 'BACKEND'):
              backends.append(cast(HaproxyBackend, line))
  
      return backends
  
  
  def load_smartstack_info_for_service(
      service: str,
      namespace: str,
      blacklist: DeployBlacklist,
      system_paasta_config: SystemPaastaConfig,
      soa_dir: str=DEFAULT_SOA_DIR,
  ) -> Dict[str, Dict[str, int]]:
      """Retrieves number of available backends for given services
  
      :param service_instances: A list of tuples of (service, instance)
      :param namespaces: list of Smartstack namespaces
      :param blacklist: A list of blacklisted location tuples in the form (location, value)
      :param system_paasta_config: A SystemPaastaConfig object representing the system configuration.
      :returns: a dictionary of the form
  
      ::
  
          {
            'location_type': {
                'unique_location_name': {
                    'service.instance': <# ofavailable backends>
                },
                'other_unique_location_name': ...
            }
          }
  
      """
      service_namespace_config = marathon_tools.load_service_namespace_config(
          service=service, namespace=namespace, soa_dir=soa_dir,
      )
      discover_location_type = service_namespace_config.get_discover()
      return get_smartstack_replication_for_attribute(
          attribute=discover_location_type,
          service=service,
          namespace=namespace,
          blacklist=blacklist,
          system_paasta_config=system_paasta_config,
      )
  
  
  def get_smartstack_replication_for_attribute(
      attribute: str,
      service: str,
      namespace: str,
      blacklist: DeployBlacklist,
      system_paasta_config: SystemPaastaConfig,
  ) -> Dict[str, Dict[str, int]]:
      """Loads smartstack replication from a host with the specified attribute
  
      :param attribute: a Mesos attribute
      :param service: A service name, like 'example_service'
      :param namespace: A particular smartstack namespace to inspect, like 'main'
      :param constraints: A list of Marathon constraints to restrict which synapse hosts to query
      :param blacklist: A list of blacklisted location tuples in the form of (location, value)
      :param system_paasta_config: A SystemPaastaConfig object representing the system configuration.
      :returns: a dictionary of the form {'<unique_attribute_value>': <smartstack replication hash>}
                (the dictionary will contain keys for unique all attribute values)
      """
      replication_info = {}
      filtered_slaves = mesos_tools.get_all_slaves_for_blacklist_whitelist(
          blacklist=blacklist,
          whitelist=None,
      )
      if not filtered_slaves:
          raise NoSlavesAvailableError
  
      attribute_slave_dict = mesos_tools.get_mesos_slaves_grouped_by_attribute(
          slaves=filtered_slaves,
          attribute=attribute,
      )
  
      full_name = compose_job_id(service, namespace)
  
      for value, hosts in attribute_slave_dict.items():
          # arbitrarily choose the first host with a given attribute to query for replication stats
          synapse_host = hosts[0]['hostname']
          repl_info = get_replication_for_services(
              synapse_host=synapse_host,
              synapse_port=system_paasta_config.get_synapse_port(),
              synapse_haproxy_url_format=system_paasta_config.get_synapse_haproxy_url_format(),
              services=[full_name],
          )
          replication_info[value] = repl_info
  
      return replication_info
  
  
  def get_replication_for_all_services(
          synapse_host: str,
          synapse_port: int,
          synapse_haproxy_url_format: str,
  ) -> Dict[str, int]:
      """Returns the replication level for all services known to this synapse haproxy
  
      :param synapse_host: The host that this check should contact for replication information.
      :param synapse_port: The port number that this check should contact for replication information.
      :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
      :returns available_instance_counts: A dictionary mapping the service names
                                          to an integer number of available replicas.
      """
      backends = get_multiple_backends(
          services=None,
          synapse_host=synapse_host,
          synapse_port=synapse_port,
          synapse_haproxy_url_format=synapse_haproxy_url_format,
      )
      return collections.Counter([b['pxname'] for b in backends if backend_is_up(b)])
  
  
  def get_replication_for_services(
      synapse_host: str,
      synapse_port: int,
      synapse_haproxy_url_format: str,
      services: Collection[str],
  ) -> Dict[str, int]:
      """Returns the replication level for the provided services
  
      This check is intended to be used with an haproxy load balancer, and
      relies on the implementation details of that choice.
  
      :param synapse_host: The host that this check should contact for replication information.
      :param synapse_port: The port number that this check should contact for replication information.
      :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
      :param services: A list of strings that are the service names
                            that should be checked for replication.
  
      :returns available_instance_counts: A dictionary mapping the service names
                                    to an integer number of available
                                    replicas
      :returns None: If it cannot connect to the specified synapse host and port
      """
      backends = get_multiple_backends(
          services=services,
          synapse_host=synapse_host,
          synapse_port=synapse_port,
          synapse_haproxy_url_format=synapse_haproxy_url_format,
      )
  
      counter = collections.Counter([b['pxname'] for b in backends if backend_is_up(b)])
      return {sn: counter[sn] for sn in services}
  
  
  def backend_is_up(backend: HaproxyBackend) -> bool:
      """Returns whether a server is receiving traffic in HAProxy.
  
      :param backend: backend dict, like one of those returned by smartstack_tools.get_multiple_backends.
  
      :returns is_up: Whether the backend is in a state that receives traffic.
      """
      return str(backend['status']).startswith('UP')
  
  
  def ip_port_hostname_from_svname(svname: str) -> Tuple[str, int, str]:
      """This parses the haproxy svname that smartstack creates.
      In old versions of synapse, this is in the format ip:port_hostname.
      In versions newer than dd5843c987740a5d5ce1c83b12b258b7253784a8 it is
      hostname_ip:port
  
      :param svname: A svname, in either of the formats described above
      :returns ip_port_hostname: A tuple of ip, port, hostname.
      """
      # split into parts
      parts = set(svname.split("_"))
  
      # find those that can be split by : - this is the ip:port
      # there will only be 1 of these
      ip_ports = {part for part in parts if len(part.split(":")) == 2}
  
      # the one *not* in the list is the hostname
      hostname = parts.difference(ip_ports).pop()
  
      ip, port = ip_ports.pop().split(":")
      return ip, int(port), hostname
  
  
  def get_registered_marathon_tasks(
      synapse_host: str,
      synapse_port: int,
      synapse_haproxy_url_format: str,
      service: str,
      marathon_tasks: Iterable[marathon_tools.MarathonTask],
  ) -> List[marathon_tools.MarathonTask]:
      """Returns the marathon tasks that are registered in haproxy under a given service (nerve_ns).
  
      :param synapse_host: The host that this check should contact for replication information.
      :param synapse_port: The port that this check should contact for replication information.
      :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
      :param service: A list of strings that are the service names that should be checked for replication.
      :param marathon_tasks: A list of MarathonTask objects, whose tasks we will check for in the HAProxy status.
      """
      backends = get_multiple_backends(
          [service], synapse_host=synapse_host, synapse_port=synapse_port,
          synapse_haproxy_url_format=synapse_haproxy_url_format,
      )
      healthy_tasks = []
      for backend, task in match_backends_and_tasks(backends, marathon_tasks):
          if backend is not None and task is not None and backend['status'].startswith('UP'):
              healthy_tasks.append(task)
      return healthy_tasks
  
  
  def are_services_up_on_ip_port(
      synapse_host: str,
      synapse_port: int,
      synapse_haproxy_url_format: str,
      services: Collection[str],
      host_ip: str,
      host_port: int,
  ) -> bool:
      backends = get_multiple_backends(
          services, synapse_host=synapse_host, synapse_port=synapse_port,
          synapse_haproxy_url_format=synapse_haproxy_url_format,
      )
      backends_by_ip_port: DefaultDict[Tuple[str, int], List[HaproxyBackend]] = collections.defaultdict(list)
  
      for backend in backends:
          ip, port, _ = ip_port_hostname_from_svname(backend['svname'])
          backends_by_ip_port[ip, port].append(backend)
  
      backends_on_ip = backends_by_ip_port[host_ip, host_port]
      # any backend being up is okay because a previous backend
      # may have had the same IP and synapse only removes them completely
      # after some time
      services_with_atleast_one_backend_up = {service: False for service in services}
      for service in services:
          for be in backends_on_ip:
              if be['pxname'] == service and backend_is_up(be):
                  services_with_atleast_one_backend_up[service] = True
      return all(services_with_atleast_one_backend_up.values())
  
  
  def match_backends_and_tasks(
      backends: Iterable[HaproxyBackend],
      tasks: Iterable[marathon_tools.MarathonTask],
  ) -> List[Tuple[Optional[HaproxyBackend], Optional[marathon_tools.MarathonTask]]]:
      """Returns tuples of matching (backend, task) pairs, as matched by IP and port. Each backend will be listed exactly
      once, and each task will be listed once per port. If a backend does not match with a task, (backend, None) will
      be included. If a task's port does not match with any backends, (None, task) will be included.
  
      :param backends: An iterable of haproxy backend dictionaries, e.g. the list returned by
                       smartstack_tools.get_multiple_backends.
      :param tasks: An iterable of MarathonTask objects.
      """
  
      # { (ip, port) : [backend1, backend2], ... }
      backends_by_ip_port: DefaultDict[Tuple[str, int], List[HaproxyBackend]] = collections.defaultdict(list)
      backend_task_pairs = []
  
      for backend in backends:
          ip, port, _ = ip_port_hostname_from_svname(backend['svname'])
          backends_by_ip_port[ip, port].append(backend)
  
      for task in tasks:
          ip = socket.gethostbyname(task.host)
          for port in task.ports:
              for backend in backends_by_ip_port.pop((ip, port), [None]):
                  backend_task_pairs.append((backend, task))
  
      # we've been popping in the above loop, so anything left didn't match a marathon task.
      for backends in backends_by_ip_port.values():
          for backend in backends:
              backend_task_pairs.append((backend, None))
  
      return backend_task_pairs
  
  
  _MesosSlaveDict = TypeVar('_MesosSlaveDict', bound=Dict)  # no type has been defined in mesos_tools for these yet.
  
  
  class SmartstackHost(NamedTuple):
      hostname: str
      pool: str
  
  
  class SmartstackReplicationChecker(abc.ABC):
      """Base class for checking smartstack replication. Extendable for different frameworks.
  
      Optimized for multiple queries. Gets the list of backends from synapse-`roxy
      only once per location and reuse it in all subsequent calls of
      SmartstackReplicationChecker.get_replication_for_instance().
  
      _get_allowed_locations_and_hosts must be implemented in sub class
      """
  
      def __init__(
          self,
          system_paasta_config: SystemPaastaConfig,
      ) -> None:
          self._synapse_port = system_paasta_config.get_synapse_port()
          self._synapse_haproxy_url_format = system_paasta_config.get_synapse_haproxy_url_format()
          self._system_paasta_config = system_paasta_config
          self._cache: Dict[str, Dict[str, int]] = {}
  
      @abc.abstractmethod
      def _get_allowed_locations_and_hosts(self, instance_config: InstanceConfig) -> Dict[str, Sequence[SmartstackHost]]:
          pass
  
      def get_replication_for_instance(
          self,
          instance_config: InstanceConfig,
      ) -> Dict[str, Dict[str, int]]:
          """Returns the number of registered instances in each discoverable location.
  
          :param instance_config: An instance of MarathonServiceConfig.
          :returns: a dict {'location_type': {'service.instance': int}}
          """
          replication_info = {}
          attribute_host_dict = self._get_allowed_locations_and_hosts(instance_config)
          instance_pool = instance_config.get_pool()
          for location, hosts in attribute_host_dict.items():
              hostname = self._get_first_host_in_pool(hosts, instance_pool)
              replication_info[location] = self._get_replication_info(location, hostname, instance_config)
          return replication_info
  
      def _get_first_host_in_pool(
          self,
          hosts: Sequence[SmartstackHost],
          pool: str,
      ) -> str:
          for host in hosts:
              if host.pool == pool:
                  return host.hostname
          return hosts[0].hostname
  
      def _get_replication_info(
          self,
          location: str,
          hostname: str,
          instance_config: InstanceConfig,
      ) -> Dict[str, int]:
          """Returns service.instance and the number of instances registered in smartstack
          at the location as a dict.
  
          :param location: A string that identifies a habitat, a region and etc.
          :param hostname: A mesos slave hostname to read replication information from.
          :param instance_config: An instance of MarathonServiceConfig.
          :returns: A dict {"service.instance": number_of_instances}.
          """
          full_name = compose_job_id(instance_config.service, instance_config.instance)
          if location not in self._cache:
              self._cache[location] = get_replication_for_all_services(
                  synapse_host=hostname,
                  synapse_port=self._synapse_port,
                  synapse_haproxy_url_format=self._synapse_haproxy_url_format,
              )
          return {full_name: self._cache[location][full_name]}
  
  
  class MesosSmartstackReplicationChecker(SmartstackReplicationChecker):
      """Retrieves the number of registered instances in each discoverable location.
  
      Based on SmartstackReplicationChecker takes mesos slaves as an argument to filter
      which services are allowed to run where.
      :Example:
  
      >>> from paasta_tools.mesos_tools import get_slaves
      >>> from paasta_tools.utils import load_system_paasta_config
      >>> from paasta_tools.marathon_tools import load_marathon_service_config
      >>> from paasta_tools.smartstack_tools import SmartstackReplicationChecker
      >>>
      >>> mesos_slaves = get_slaves()
      >>> system_paasta_config = load_system_paasta_config()
      >>> instance_config = load_marathon_service_config(service='fake_service',
      ...                       instance='fake_instance', cluster='norcal-stagef')
      >>>
      >>> c = SmartstackReplicationChecker(mesos_slaves, system_paasta_config)
      >>> c.get_replication_for_instance(instance_config)
      {'uswest1-stagef': {'fake_service.fake_instance': 2}}
      >>>
      """
  
      def __init__(
          self,
          mesos_slaves: List[_MesosSlaveDict],
          system_paasta_config: SystemPaastaConfig,
      ) -> None:
          self._mesos_slaves = mesos_slaves
          super().__init__(system_paasta_config=system_paasta_config)
  
      def _get_allowed_locations_and_hosts(self, instance_config: InstanceConfig) -> Dict[str, Sequence[SmartstackHost]]:
          """Returns a dict of locations and lists of corresponding mesos slaves
          where deployment of the instance is allowed.
  
          :param instance_config: An instance of MarathonServiceConfig
          :returns: A dict {"uswest1-prod": [SmartstackHost(), SmartstackHost(), ...]}
          """
          monitoring_blacklist = instance_config.get_monitoring_blacklist(
              system_deploy_blacklist=self._system_paasta_config.get_deploy_blacklist(),
          )
          filtered_slaves = mesos_tools.filter_mesos_slaves_by_blacklist(
              slaves=self._mesos_slaves,
              blacklist=monitoring_blacklist,
              whitelist=None,
          )
          discover_location_type = marathon_tools.load_service_namespace_config(
              service=instance_config.service,
              namespace=instance_config.instance,
              soa_dir=instance_config.soa_dir,
          ).get_discover()
          attribute_to_slaves = mesos_tools.get_mesos_slaves_grouped_by_attribute(
              slaves=filtered_slaves,
              attribute=discover_location_type,
          )
          ret: Dict[str, Sequence[SmartstackHost]] = {}
          for attr, slaves in attribute_to_slaves.items():
              ret[attr] = [
                  SmartstackHost(hostname=slave['hostname'], pool=slave['attributes']['pool']) for slave in slaves
              ]
          return ret
  
  
  class KubeSmartstackReplicationChecker(SmartstackReplicationChecker):
  
      def __init__(
          self,
          nodes: Sequence[V1Node],
          system_paasta_config: SystemPaastaConfig,
      ) -> None:
          self.nodes = nodes
          super().__init__(system_paasta_config=system_paasta_config)
  
      def _get_allowed_locations_and_hosts(self, instance_config: InstanceConfig) -> Dict[str, Sequence[SmartstackHost]]:
  
          monitoring_blacklist = instance_config.get_monitoring_blacklist(
              system_deploy_blacklist=self._system_paasta_config.get_deploy_blacklist(),
          )
          filtered_nodes = kubernetes_tools.filter_nodes_by_blacklist(
              nodes=self.nodes,
              blacklist=monitoring_blacklist,
              whitelist=None,
          )
          discover_location_type = kubernetes_tools.load_service_namespace_config(
              service=instance_config.service,
              namespace=instance_config.instance,
              soa_dir=instance_config.soa_dir,
          ).get_discover()
          attribute_to_nodes = kubernetes_tools.get_nodes_grouped_by_attribute(
              nodes=filtered_nodes,
              attribute=discover_location_type,
          )
          ret: Dict[str, Sequence[SmartstackHost]] = {}
          for attr, nodes in attribute_to_nodes.items():
              ret[attr] = [
                  SmartstackHost(
                      hostname=node.metadata.labels['yelp.com/hostname'],
                      pool=node.metadata.labels['yelp.com/pool'],
                  ) for node in nodes
              ]
          return ret
  