#!/usr/bin/env python
# Copyright 2019 VMware, Inc.
# All rights reserved. -- VMware Confidential

# Provides methods which are needed on vsphere-ui minor upgrade

import sys
import os
from os.path import join, abspath, dirname, isfile

import logging
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.StreamHandler())

sys.path.extend(os.environ['VMWARE_PYTHON_PATH'].split(':'))

import datetime
import warnings
import cis.defaults
import xml.etree.ElementTree as ET

from cis.defaults import *
from cis.cisreglib import AuthzClient, CisregOptionParser, \
   LookupServiceClient, SsoClient, VmafdClient, _get_syscfg_info
from cis.vecs import vmafd_machine_id, VecsEntry
from cis.utils import FileBuffer, setupLogging, is_linux
from cis.vpxdevent_lib import get_default_connect_urls
from identity.vmkeystore import VmKeyStore

from pyVmomi import vim, VmomiSupport, SoapStubAdapter, dataservice
from pyVim import sso
from tempfile import NamedTemporaryFile
import identity.vmafdpy3

VMAFD_PATH = os.path.join(get_component_home_dir('vmafd'), 'lib64')
sys.path.append(VMAFD_PATH)
# Suppress warnings coming from vmafd and identity imports.
with warnings.catch_warnings():
   warnings.simplefilter('ignore', RuntimeWarning)
H5_CLIENT_SOLUTION_USER = "vsphere-webclient"
AUTHZ_TYPE_ID = 'cs.authorization'
EP_AUTHZ_PROTOCOL = 'vmomi'
EP_AUTHZ_TYPE_ID = 'com.vmware.cis.authorization.server'
VSPHERE_CLIENT_AUTHZ_ROLE = 'vSphere Client Solution User'
MACHINE_NAME = 'machine'
VSPHERE_CLIENT_ROLE_ID = 1003
COMPONENT_NAME = "vsphere-ui"

# Creates a solution user role with a set of privileges and assigns it to the vsphere-webclient user
def create_vsphere_client_solution_user_role_and_assign():
   try:
      logging.info('Creating vSphere Client Solution User role and assigning to vsphere-webclient solution user')
      vmafdClient = identity.vmafdpy3.client('localhost')
      username = vmafdClient.GetMachineName()
      password = vmafdClient.GetMachinePassword()
      machine_id = vmafd_machine_id()
      local_nodeid = vmafdClient.GetLDU()
      logging.info('create_vsphere_client_solution_user_role_and_assign - The LDU is: %s' % local_nodeid)
      solution_user_id = '%s-%s' % (H5_CLIENT_SOLUTION_USER, machine_id)
      logging.info('create_vsphere_client_solution_user_role_and_assign - Connecting to Lookup Service')
      ls_url, domain_name = _get_syscfg_info()
      ls_obj = LookupServiceClient(ls_url, retry_count=1)
      logging.info('create_vsphere_client_solution_user_role_and_assign - Getting STS endpoint')
      sts_url, sts_cert_data = ls_obj.get_sts_endpoint_data()
      cert = None
      key = None
      with VmKeyStore('VKS') as ks:
         ks.load(MACHINE_NAME)
         cert = ks.get_certificate(MACHINE_NAME)
         key = ks.get_key(MACHINE_NAME)
      sso_client = SsoClient(sts_url, sts_cert_data, None, None, cert=cert, key=key)
      try:
         ls_obj.set_sso_client(sso_client)
         authz_endpoints = ls_obj.get_service_endpoints(
            AUTHZ_TYPE_ID, ep_protocol=EP_AUTHZ_PROTOCOL,
            ep_type=EP_AUTHZ_TYPE_ID, local_nodeid=local_nodeid)

         role_data_list = [{'priv_ids': ['System.View', 'System.Read', 'System.Anonymous', 'Task.Create', 'Task.Update', 'Extension.Register', 'Extension.Unregister', 'Extension.Update'], 'name': 'vSphere Client Solution User', 'description': 'This role entitles minimum required privileges for vSphere Client Solution User', 'id': '1003'}]
         logging.info('create_vsphere_client_solution_user_role_and_assign - Logging into Authz')
         authz_client = AuthzClientWrapper(authz_endpoints[0].url, sso_client)
         role_obj = authz_client.get_role(role_id=VSPHERE_CLIENT_ROLE_ID)
         if role_obj is None:
            logging.info('create_vsphere_client_solution_user_role_and_assign - Creating vSphere Client Solution User role in Authz')
            authz_client.load_roles(role_data_list)
         else:
            logging.info('create_vsphere_client_solution_user_role_and_assign - Updating vSphere Client Solution User role in Authz')
            for role_data in role_data_list:
               if int(role_data['id']) == VSPHERE_CLIENT_ROLE_ID:
                  role_privilege_ids = role_obj.privilegeId
                  for privilegeId in role_data['priv_ids']:
                     if privilegeId not in role_privilege_ids:
                        role_privilege_ids.append(privilegeId)
                  authz_client.update_role(role_id=VSPHERE_CLIENT_ROLE_ID, priv_ids=role_privilege_ids)
                  break
         logging.info('create_vsphere_client_solution_user_role_and_assign - Assigning vSphere Client Solution User role to vsphere-webclient solution user')
         authz_client.set_permission(domain_name, VSPHERE_CLIENT_AUTHZ_ROLE, solution_user_id, is_group=False)
      finally:
         sso_client.cleanup()
   except Exception as ex:
      logging.info("Failed to create a vSphere Client Solution User role and assign to vsphere-webclient solution user. Exception is: %s" % str(ex))

def get_config_dir():
   return os.path.join(dirname(__file__), "config")

# ---
# Wrapper used to add additional functionality to the AuthzClient
# This functionality is copied from the main branch where this logic is present in the AuthzClient itself
# This logic should be removed once the AuthzClient is updated in vsphere67 branch.
# -
class AuthzClientWrapper(AuthzClient):
   def __init__(self, authz_url, sso_client):
      super().__init__(authz_url, sso_client)

   def get_role(self, role_id):
      '''
      Get role by id.
      '''
      for role_obj in self._authz_service.GetRoles():
         if role_obj.id == role_id:
            return role_obj
      return None

   def update_role(self, role_id, priv_ids):
      '''
      Update role.
      '''
      role_obj = self.get_role(role_id)
      if role_obj is not None:
         new_role_obj = dataservice.accesscontrol.Role(id=role_id, name=role_obj.name, \
                                                       description=role_obj.description, privilegeId=priv_ids)
         self._authz_service.UpdateRole(new_role_obj)


# ----
# Creates a new vim.Extension object for the vSphere Client extension.
# [out]: A new vim.Extension spec with populated data.
# [throws]: IOError, ParseError, Exception
# ----
def create_vsphere_client_extension():
    extension_xml_file = os.path.join(get_config_dir(), 'extension.xml')
    extensionXml = ET.parse(extension_xml_file)
    extRoot = extensionXml.getroot().find('extension')

    name = extRoot.find('name').text
    desc = vim.Description(label=(name + '.label'), summary=(name + '.summary'))
    company = extRoot.find('company').text
    ver = extRoot.find('version').text

    # Fill tasks type information.
    tasks = [vim.Extension.TaskTypeInfo(taskID=task.get('id')) for task in extRoot.findall('./tasks/task')]

    # Fill faults type information.
    faults = [vim.Extension.FaultTypeInfo(faultID=fault.get('id')) for fault in extRoot.findall('./faults/fault')]

    # Fill resources information.
    resources = []
    l10n_dir = os.path.join(get_config_dir(), 'l10n')
    subdirs = next(os.walk(l10n_dir))[1]
    for locale in subdirs:
       locale_dir = os.path.join(l10n_dir, locale)
       vmsg_file_names = next(os.walk(locale_dir))[2]
       for vmsg_file_name in vmsg_file_names:
          entries=[]
          vmsg_file = os.path.join(locale_dir, vmsg_file_name)
          resource_file_buffer = FileBuffer()
          resource_file_buffer.readFile(vmsg_file)
          for line in resource_file_buffer.getBufferContents():
             if '=' in line:
                prop = line.split('=')
                entries.append(vim.KeyValue(key=prop[0].strip(),value=prop[1].strip()))

          resources.append(
             vim.Extension.ResourceInfo(
                data=entries,
                locale=locale,
                module=vmsg_file_name[:-5]
             ))

    return vim.Extension(description=desc, key=name, company=company, version=ver,
                       taskList=tasks, faultList=faults, resourceList = resources,
                       lastHeartbeatTime=datetime.datetime.now())

# Registers the 'com.vmware.vsphere.client' extension in the vCenter Server
def register_vsphere_client_extension():
   try:
      cert_tmp_file = NamedTemporaryFile(dir=get_cis_tmp_dir(), delete=True)
      cert_tmp_file.close()
      crt_key_tmp_file = NamedTemporaryFile(dir=get_cis_tmp_dir(), delete=True)
      crt_key_tmp_file.close()

      vecs_entry = VecsEntry('vpxd-extension')
      cert = cert_tmp_file.name
      vecs_entry.get_cert('vpxd-extension', cert)
      cert_key = crt_key_tmp_file.name
      vecs_entry.get_key('vpxd-extension', cert_key)

      sts_url, vc_endpoints = get_default_connect_urls()
      url = vc_endpoints[0]

      authenticator = sso.SsoAuthenticator(sts_url=sts_url)
      token = authenticator.get_hok_saml_assertion(cert, cert_key, token_duration=86400)
      vmodlVersion = VmomiSupport.newestVersions.Get('vim')

      conn = SoapStubAdapter(
         url=url,
         version=vmodlVersion,
         samlToken=token
      )

      serviceInstance = vim.ServiceInstance('ServiceInstance', conn)
      content = serviceInstance.content
      em = content.extensionManager
      sm = content.sessionManager

      if not sm.currentSession:
         with conn.requestModifier(
                 lambda request: sso.add_saml_context(request, token, cert_key)):
            try:
               conn.samlToken = token
               sm.LoginByToken()
            except Exception as err:
               logging.info("VC login failure. Exception is %s" % str(err))
            finally:
               conn.samlToken = None

      ext_spec = create_vsphere_client_extension()
      try:
         ext = em.FindExtension(ext_spec.key)
         if ext:
            em.UpdateExtension(extension=ext_spec)
            logging.info("Updated extension %s" % str(ext_spec.key))
         else:
            em.RegisterExtension(extension=ext_spec)
            logging.info("Registered extension %s" % str(ext_spec.key))
      except Exception as ex:
         logging.info("Failed to Update/Register extension %s: %s" % (ext_spec.key, str(ex)))
      finally:
         sm.Logout()
   except Exception as ex:
      logging.info("Failed to register extension 'com.vmware.vsphere.client'. Exception is: %s" % str(ex))