# -*- coding: utf-8 -*-

# Copyright 2008-2016 Mir Calculate. http://www.calculate-linux.org
#
#  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 sys
import ldap
from utils.common import _error
from collections import defaultdict
from ldif import LDIFParser, LDIFWriter
import cStringIO

from calculate.lib.cl_lang import setLocalTranslate

setLocalTranslate('cl_lib3', sys.modules[__name__])


class ldapFun(_error):
    """
    Объект для работы с LDAP сервером

    подключение к серверу и поиск данных
    """

    def __init__(self, dnUser, password, host="localhost"):
        self.conLdap = False
        # Получаем соединение с LDAP
        try:
            self.conLdap = self.__ldapConnect(dnUser, password, host)
        except ldap.LDAPError, e:
            self.setError(e[0]['desc'])

    def __ldapConnect(self, dnUser, password, host):
        """Соединение с LDAP сервером"""
        con_ldap = ldap.initialize('ldap://%s' % host)
        con_ldap.simple_bind_s(dnUser, password)
        return con_ldap

    def ldapSearch(self, base_dn, search_scope, search_filter, retrieve_attrs):
        try:
            ldap_result_id = self.conLdap.search(base_dn, search_scope,
                                                 search_filter,
                                                 retrieve_attrs)
            result_set = []
            while 1:
                result_type, result_data = self.conLdap.result(
                    ldap_result_id, 0)
                if isinstance(result_data, list) and not result_data:
                    break
                else:
                    if result_type == ldap.RES_SEARCH_ENTRY:
                        result_set.append(result_data)
        except ldap.NO_SUCH_OBJECT:
            return []
        except Exception:
            return False
        return result_set


class LDAPConnectError(Exception):
    pass

class LDAPBadSearchFilter(LDAPConnectError):
    pass

class LDAPConnect(ldapFun):
    """
    Объект работающий с исключениями
    """

    def setError(self, message):
        raise LDAPConnectError(message)

    def ldap_search(self, base_dn, search_scope=ldap.SCOPE_BASE,
                    search_filter='(objectClass=*)',
                    retrieve_attrs=None):
        try:
            ldap_result_id = self.conLdap.search(
                base_dn, search_scope, search_filter, retrieve_attrs)
            while 1:
                result_type, result_data = self.conLdap.result(
                    ldap_result_id, 0)
                if isinstance(result_data, list) and not result_data:
                    break
                else:
                    if result_type == ldap.RES_SEARCH_ENTRY:
                        yield result_data
        except ldap.NO_SUCH_OBJECT:
            pass
        except ldap.LDAPError, e:
            error = e[0]['desc']
            if "Bad search filter" in error:
                raise LDAPBadSearchFilter(error)
            raise LDAPConnectError(error)
        except Exception as e:
            raise LDAPConnectError(str(e))

    def ldap_dump(self, fobj, base_dn, search_scope=ldap.SCOPE_SUBTREE,
                  search_filter='(objectClass=*)', retrieve_attrs=None):
        """
        вернуть dump ldif
        :param dn:
        :param search_scope:
        :param search_filter:
        :param retrieve_attrs:
        :return:
        """
        try:
            dn_list = self.conLdap.search_s(
                base_dn, search_scope, search_filter, retrieve_attrs)
            writer = LDIFWriter(fobj)
            for dn, f in dn_list:
                writer.unparse(dn, f)
        except ldap.NO_SUCH_OBJECT:
            pass
        except ldap.LDAPError, e:
            error = e[0]['desc']
            if "Bad search filter" in error:
                raise LDAPBadSearchFilter(error)
            raise LDAPConnectError(error)
        except Exception as e:
            raise LDAPConnectError(str(e))

    def ldap_simple_search(self, base_dn, search_filter, attr, first=True):
        """
        Простой поиск аттрибутов
        :param base_dn:
        :param search_filter:
        :param attr:
        :param first: вернуть первый из списка
        :return:
        """
        for entry in self.ldap_search(base_dn, ldap.SCOPE_ONELEVEL,
                                      search_filter, [attr]):
            attrs = entry[0][1]
            yield attrs[attr][0]

    def ldap_modify_attrs(self, base_dn, attrs):
        """Модифицирует аттрибуты DN"""
        attrs = list(attrs)
        if attrs:
            try:
                self.conLdap.modify_s(base_dn, attrs)
            except ldap.LDAPError, e:
                raise LDAPConnectError(e[0]['desc'])

    def ldap_modify_dn(self, base_dn, new_dn):
        """
        Изменить ветку
        :param base_dn: предыдущее название
        :param new_dn: новое название
        :return:
        """
        try:
            self.conLdap.modrdn_s(base_dn, new_dn)
        except ldap.LDAPError, e:
            raise LDAPConnectError(e[0]['desc'])

    def ldap_remove_dn(self, base_dn):
        """
        Удалить указанный dn
        :param base_dn: удаляемый dn
        :return:
        """
        try:
            self.conLdap.delete_s(base_dn)
        except ldap.LDAPError, e:
            raise LDAPConnectError(e[0]['desc'])



            # @adminConnectLdap
            # def modifyElemDN(self, relDN, newFirstDn):
            #    """Изменяет основной элемент DN (uid, cn и др.)"""
            #    DN = self.addDN(relDN,self.baseDN)
            #    try:
            #        self.conLdap.modrdn_s(DN, newFirstDn)
            #    except ldap.LDAPError, e:
            #        self.printERROR(e[0]['desc'])
            #        return False
            #    return True

            # @adminConnectLdap
            # def delDN(self, relDN):
            #    """Удаляет одиночный DN"""
            #    DN = self.addDN(relDN,self.baseDN)
            #    try:
            #        self.conLdap.delete_s(DN)
            #    except ldap.LDAPError, e:
            #        self.printERROR(e[0]['desc'])
            #        return False
            #    return True


class ldapUser(_error):
    """Получение данных для пользователя из LDAP"""
    # Данные из /etc/ldap.conf
    _dictData = {}
    # Объект LDAP
    ldapObj = False
    # Подключение к LDAP
    conLdap = False

    def addDN(self, *arg):
        """
        Append text DN elements
        """
        return ",".join(x for x in arg if x)

    def getDataInLdapConf(self, bindData=True, cache=True):
        """Получение данных из /etc/ldap.conf"""
        data = [("host", 'host'),
                ("usersDN", 'nss_base_passwd'),
                ("groupsDN", 'nss_base_group')]
        if bindData:
            data += [("bindDn", 'binddn'), ("bindPw", 'bindpw')]
        names_data = [x[0] for x in data]
        # Данные из кеша, если он есть
        if (cache and self._dictData and
                    set(names_data) <= set(self._dictData.keys())):
            return self._dictData
        file_name = "/etc/ldap.conf"
        get_str_list = lambda x: reduce(lambda x, y: [x, y.upper()], ([x] * 2))
        workdata = map(lambda x: (x[0], get_str_list(x[1]), len(x[1])), data)
        dict_data = defaultdict(list)
        delimeter = (" ", "\t")
        try:
            for line in open(file_name):
                for name, keys, lenKey in workdata:
                    if (name not in dict_data.keys() and
                            any(line.startswith(x) for x in keys) and
                                len(line) > lenKey):
                        spl = line[lenKey]
                        if spl in delimeter:
                            param_value = line.rpartition(spl)[2]
                            if name in ("usersDN", "groupsDN"):
                                dict_data[name].append(
                                    param_value.partition('?')[0].strip())
                            else:
                                dict_data[name].append(param_value.strip())
        except Exception:
            # self.setError(_("Can not open %s")%fileName)
            return False
        if set(dict_data.keys()) == set(names_data):
            # Кеширование данных
            if cache:
                self._dictData.clear()
                self._dictData.update(dict_data)
            return dict_data
        else:
            return {}

    def getBindConnectData(self):
        """Получение данных для соединения с LDAP bind пользователем"""
        configdata = self.getDataInLdapConf()
        if configdata:
            bind_dn = configdata["bindDn"][0]
            bind_pw = configdata["bindPw"][0]
            host = configdata["host"][0]
            return bind_dn, bind_pw, host
        return False

    def getUsersDN(self):
        """Получение DN пользователей"""
        configdata = self.getDataInLdapConf(bindData=False)
        if configdata:
            return self._dictData["usersDN"][0]
        return False

    def getHost(self):
        """Получение LDAP хоста"""
        configdata = self.getDataInLdapConf(bindData=False)
        if configdata:
            return configdata["host"][0]
        return False

    def getGroupsDN(self):
        """Получение списка DN групп"""
        configdata = self.getDataInLdapConf(bindData=False)
        if configdata:
            return self._dictData["groupsDN"]
        return False

    def connectLdap(self):
        """
        Connect to LDAP
        """
        connectData = self.getBindConnectData()
        if not connectData:
            return {}
        bindDn, bindPw, host = connectData
        self.getUsersDN()
        # Соединяемся с LDAP
        return self.ldapConnect(bindDn, bindPw, host)

    def getUserLdapInfo(self, user_name, shadowAttr=False):
        """Выдаем информацию о пользователе из LDAP"""
        if not self.connectLdap():
            return False
        users_dn = self.getUsersDN()
        groups_dn = self.getGroupsDN()
        search_user = self.ldapObj.ldapSearch(users_dn, ldap.SCOPE_ONELEVEL,
                                              "uid=%s" % user_name, None)
        if not search_user:
            return False
        convert_dict = {'uid': ('user', 'uidNumber'),
                        'gid': ('user', 'gidNumber'),
                        'fullName': ('user', 'cn'),
                        'mail': ('user', 'mail'),
                        'jid': ('user', 'registeredAddress'),
                        'home': ('user', 'homeDirectory'),
                        'group': ('group', 'cn')}
        if shadowAttr:
            convert_dict.update({'loginShell': ('user', 'loginShell'),
                                 'shadowLastChange': (
                                     'user', 'shadowLastChange'),
                                 'shadowMin': ('user', 'shadowMin'),
                                 'shadowMax': ('user', 'shadowMax'),
                                 'shadowWarning': ('user', 'shadowWarning'),
                                 'shadowExpire': ('user', 'shadowExpire'),
                                 'shadowFlag': ('user', 'shadowFlag'),
                                 'groups': ('group', 'memberUid')})
        list_user_attr = [k for k, v in convert_dict.items() if v[0] == "user"]
        list_group_attr = [k for k, v in convert_dict.items()
                           if v[0] == "group"]
        uid = ""
        gid = ""
        dict_out = {}
        for dict_attr in list_user_attr:
            ldap_attr = convert_dict[dict_attr][1]
            if ldap_attr in search_user[0][0][1]:
                dict_out[dict_attr] = search_user[0][0][1][ldap_attr][0]
            else:
                dict_out[dict_attr] = ""
            if dict_attr == 'uid':
                uid = dict_out[dict_attr]
            if dict_attr == 'gid':
                gid = dict_out[dict_attr]
        if gid:
            for dict_attr in list_group_attr:
                search_group = []
                ldap_attr = convert_dict[dict_attr][1]
                if dict_attr == "group":
                    for groupDN in groups_dn:
                        search_group = self.ldapObj.ldapSearch(
                            groupDN, ldap.SCOPE_ONELEVEL, "gidNumber=%s" % gid,
                            None)
                        if search_group:
                            break
                    if search_group:
                        data = search_group[0][0][1]
                        if ldap_attr in data:
                            dict_out[dict_attr] = data[ldap_attr][0]
                        else:
                            dict_out[dict_attr] = ""
                    else:
                        dict_out[dict_attr] = ""
                elif dict_attr == "groups":
                    user_groups_data = []
                    for groupDN in groups_dn:
                        search_group = self.ldapObj.ldapSearch(
                            groupDN, ldap.SCOPE_ONELEVEL,
                            "%s=%s" % (ldap_attr, user_name),
                            ["cn", "gidNumber"])
                        if search_group:
                            user_groups_data.extend(
                                [(x[0][1]["cn"][0], x[0][1]["gidNumber"][0])
                                 for x in search_group])
                    dict_out[dict_attr] = user_groups_data
        if uid and gid:
            return dict_out
        else:
            return {}

    def ldapConnect(self, bind_dn, bind_pw, host):
        """Подключение к LDAP"""
        if not self.ldapObj:
            ldap_obj = ldapFun(bind_dn, bind_pw, host)
            if ldap_obj.getError():
                ldap_obj.clearErrors()
                return False
            # Устанавливаем у объекта соединение и объект LDAP функций
            self.ldapObj = ldap_obj
            self.conLdap = ldap_obj.conLdap
        return True


class LDIFError(Exception):
    pass


class LDIFAdd(LDIFParser):
    """
    Добавление LDIF в базу
    """

    def __init__(self, ldif_data, ldap_connect):
        self.ldap_connect = ldap_connect
        LDIFParser.__init__(self, cStringIO.StringIO(ldif_data))

    def handle(self, dn, entry, controls=None):
        # (self, dn, entry, *args):
        try:
            self.ldap_connect.add_s(dn, entry.items())
        except ldap.LDAPError as e:
            raise LDIFError(e[0]['desc'])
        except Exception as e:
            print str(e)
            raise LDIFError(_("An error in the LDIF file"))
