text: # (C) Datadog, Inc. 2019-present
  # All rights reserved
  # Licensed under Simplified BSD License (see LICENSE)
  import json
  from collections import defaultdict
  from typing import Any, Dict, Iterator, List, Set
  
  from pyVmomi import vim
  from six import iteritems
  
  from datadog_checks.base.log import CheckLoggingAdapter
  from datadog_checks.base.utils.http import RequestsWrapper
  from datadog_checks.vsphere.config import VSphereConfig
  from datadog_checks.vsphere.constants import ALL_RESOURCES_WITH_METRICS
  from datadog_checks.vsphere.types import ResourceTags, TagAssociation
  
  from .api import APIResponseError, smart_retry
  
  MOR_TYPE_MAPPING_FROM_STRING = {
      'HostSystem': vim.HostSystem,
      'VirtualMachine': vim.VirtualMachine,
      'Datacenter': vim.Datacenter,
      'Datastore': vim.Datastore,
      'ClusterComputeResource': vim.ClusterComputeResource,
  }
  
  MOR_TYPE_MAPPING_TO_STRING = {v: k for k, v in iteritems(MOR_TYPE_MAPPING_FROM_STRING)}
  
  
  class VSphereRestAPI(object):
      """
      Abstraction class over the vSphere REST api
      """
  
      def __init__(self, config, log):
          # type: (VSphereConfig, CheckLoggingAdapter) -> None
          self.config = config
          self.log = log
          self._client = VSphereRestClient(config, log)
          self.smart_connect()
  
      def smart_connect(self):
          # type: () -> None
          """
          Connect to vSphere client.
          """
          self._client.connect_session()
  
      def make_batch(self, mors):
          # type: (List[vim.ManagedEntity]) -> Iterator[List[vim.ManagedEntity]]
          batch = []  # type: List[vim.ManagedEntity]
          size = self.config.batch_tags_collector_size
          for mor in mors:
              if len(batch) == size:
                  yield batch
                  batch = []
              batch.append(mor)
          if batch:
              yield batch
  
      @smart_retry
      def get_resource_tags_for_mors(self, mors):
          # type: (List[vim.ManagedEntity]) -> ResourceTags
          """
          Get resource tags.
  
          Response structure:
  
              {
                  <RESOURCE_TYPE>: {
                      <RESOURCE_MOR_ID>: ['<CATEGORY_NAME>:<TAG_NAME>', ...]
                  },
                  ...
              }
          """
          tag_associations = []
          for mors_batch in self.make_batch(mors):
              batch_tag_associations = self._client.tagging_tag_association_list_attached_tags_on_objects(mors_batch)
              tag_associations.extend(batch_tag_associations)
  
          self.log.debug("Fetched tag associations: %s", tag_associations)
  
          # Initialise resource_tags
          resource_tags = {
              resource_type: defaultdict(list) for resource_type in ALL_RESOURCES_WITH_METRICS
          }  # type: ResourceTags
  
          all_tag_ids = set()
          for tag_asso in tag_associations:
              all_tag_ids.update(tag_asso["tag_ids"])
  
          tags = self._get_tags(all_tag_ids)
  
          for tag_asso in tag_associations:
              mor_id = tag_asso["object_id"]["id"]
              mor_type = MOR_TYPE_MAPPING_FROM_STRING[tag_asso["object_id"]["type"]]
              mor_tag_ids = tag_asso["tag_ids"]
              for mor_tag_id in mor_tag_ids:
                  if mor_tag_id not in tags:
                      self.log.debug("MOR tag id '%s' was not found in response, ignoring.", mor_tag_id)
                      continue
                  resource_tags[mor_type][mor_id].append(tags[mor_tag_id])
  
          self.log.debug("Result resource tags: %s", resource_tags)
          return resource_tags
  
      def _get_tags(self, tag_ids):
          # type: (Set[str]) -> Dict[str, str]
          """
          Create tags using vSphere tags prefix + vSphere tag category name as key and vSphere tag name as value.
  
              <TAGS_PREFIX><TAG_CATEGORY>:<TAG_NAME>
  
          Examples:
              - os_type:windows
              - application_name:my_app
  
          Taging best practices:
          https://www.vmware.com/content/dam/digitalmarketing/vmware/en/pdf/techpaper/performance/tagging-vsphere67-perf.pdf
          """
          tags = {}
          categories = {}
          for tag_id in tag_ids:
              tag = self._client.tagging_tags_get(tag_id)
              category_id = tag["category_id"]
              if category_id not in categories:
                  categories[category_id] = self._client.tagging_category_get(category_id)
              category_name = categories[category_id]["name"]
              tags[tag_id] = "{}{}:{}".format(self.config.tags_prefix, category_name, tag['name'])
          return tags
  
  
  class VSphereRestClient(object):
      """
      Custom vSphere Rest API Client
      """
  
      JSON_REQUEST_HEADERS = {'Content-Type': 'application/json'}
  
      def __init__(self, config, log):
          # type: (VSphereConfig, CheckLoggingAdapter) -> None
          self.log = log
          http_config = {
              'username': config.username,
              'password': config.password,
              'tls_ca_cert': config.ssl_capath,
              'tls_verify': config.ssl_verify,
              'tls_ignore_warning': config.tls_ignore_warning,
          }
          self._api_base_url = "https://{}/rest/com/vmware/cis/".format(config.hostname)
          self._http = RequestsWrapper(http_config, {})
  
      def connect_session(self):
          # type: () -> None
          session_token = self.session_create()
  
          if not session_token:
              raise APIResponseError("Failed to retrieve session token")
  
          self._http.options['headers']['vmware-api-session-id'] = session_token
  
      def session_create(self):
          # type: () -> str
          """
          Create session token
          Doc:
          https://vmware.github.io/vsphere-automation-sdk-rest/6.5/operations/com/vmware/cis/session.create-operation.html
          """
          session_token = self._request_json("session", method="post", extra_headers=self.JSON_REQUEST_HEADERS)
          return session_token
  
      def tagging_category_get(self, category_id):
          # type: (str) -> Dict[str, Any]
          """
          Get one category
          Doc:
          https://vmware.github.io/vsphere-automation-sdk-rest/6.5/operations/com/vmware/cis/tagging/category.get-operation.html
          """
          return self._request_json("tagging/category/id:{}".format(category_id))
  
      def tagging_tags_get(self, tag_id):
          # type: (str) -> Dict[str, Any]
          """
          Get one tag
          Doc:
          https://vmware.github.io/vsphere-automation-sdk-rest/6.5/operations/com/vmware/cis/tagging/tag.get-operation.html
          """
          return self._request_json("tagging/tag/id:{}".format(tag_id))
  
      def tagging_tag_association_list_attached_tags_on_objects(self, mors):
          # type: (List[vim.ManagedEntity]) -> List[TagAssociation]
          """
          Get all tags identifiers for a given set of objects
          Doc:
          https://vmware.github.io/vsphere-automation-sdk-rest/6.5/operations/com/vmware/cis/tagging/tag_association.list_attached_tags_on_objects-operation.html
  
          This operation was added in vSphere API 6.5
          """
          payload = {"object_ids": [{"id": mor._moId, "type": MOR_TYPE_MAPPING_TO_STRING[type(mor)]} for mor in mors]}
          tag_associations = self._request_json(
              "tagging/tag-association?~action=list-attached-tags-on-objects",
              method="post",
              data=json.dumps(payload),
              extra_headers=self.JSON_REQUEST_HEADERS,
          )
          return tag_associations
  
      def _request_json(self, endpoint, method='get', **options):
          # type: (str, str, **Any) -> Any
          url = self._api_base_url + endpoint
          resp = getattr(self._http, method)(url, **options)
          resp.raise_for_status()
  
          data = resp.json()
          if 'value' not in data:
              raise APIResponseError("Missing `value` element in response for url: {}".format(url))
  
          return data['value']
  