# Author: Zhang Huangbin <zhb@iredmail.org>

import sys
import ldap
import web
from libs import iredutils, models
from libs.ldaplib import core, attrs, iredldif, ldaputils, deltree, connUtils, decorators
from libs.policyd import throttle

cfg = web.iredconfig
session = web.config.get('_session')


class Domain(core.LDAPWrap):
    def __del__(self):
        try:
            self.conn.unbind()
        except:
            pass

    @decorators.require_global_admin
    def add(self, data):
        # msg: {key: value}
        msg = {}
        self.domain = web.safestr(data.get('domainName', '')).strip().lower()

        # Check domain name.
        if not iredutils.isDomain(self.domain):
            return (False, 'INVALID_DOMAIN_NAME')

        # Check whether domain name already exist (domainName, domainAliasName).
        connutils = connUtils.Utils()
        if connutils.isDomainExists(self.domain):
            return (False, 'ALREADY_EXISTS')

        self.dn = ldaputils.convKeywordToDN(self.domain, accountType='domain')

        self.cn = data.get('cn', None)
        ldif = iredldif.ldif_maildomain(domain=self.domain, cn=self.cn,)

        # Add domain dn.
        try:
            result = self.conn.add_s(self.dn, ldif)
            web.logger(msg="Create domain: %s." % (self.domain), domain=self.domain, event='create',)
        except ldap.ALREADY_EXISTS:
            msg[self.domain] = 'ALREADY_EXISTS'
        except ldap.LDAPError, e:
            msg[self.domain] = str(e)

        # Add default groups under domain.
        if len(attrs.DEFAULT_GROUPS) >= 1:
            for i in attrs.DEFAULT_GROUPS:
                try:
                    group_dn = 'ou=' + str(i) + ',' + str(self.dn)
                    group_ldif = iredldif.ldif_group(str(i))

                    self.conn.add_s(group_dn, group_ldif)
                except ldap.ALREADY_EXISTS:
                    pass
                except ldap.LDAPError, e:
                    msg[i] = str(e)
        else:
            pass

        if len(msg) == 0:
            return (True,)
        else:
            return (False, ldaputils.getExceptionDesc(msg))

    # List all domain admins.
    def getDomainAdmins(self, domain):
        domain = web.safestr(domain)
        dn = ldaputils.convKeywordToDN(domain, accountType='domain')
        try:
            self.domainAdmins = self.conn.search_s(
                    dn,
                    ldap.SCOPE_BASE,
                    '(&(objectClass=mailDomain)(domainName=%s))' % domain,
                    ['domainAdmin'],
                    )
            return self.domainAdmins
        except Exception, e:
            return str(e)

    # List all domains under control.
    def listAccounts(self, attrs=attrs.DOMAIN_SEARCH_ATTRS):
        result = self.getAllDomains(attrs=attrs)
        if result[0] is True:
            allDomains = result[1]
            allDomains.sort()
            return (True, allDomains)
        else:
            return result

    # Get domain default user quota: domainDefaultUserQuota.
    # - domainAccountSetting must be a dict.
    def getDomainDefaultUserQuota(self, domain, domainAccountSetting=None,):
        # Return 0 as unlimited.
        self.domain = web.safestr(domain)
        self.dn = ldaputils.convKeywordToDN(self.domain, accountType='domain')

        if domainAccountSetting is not None:
            if 'defaultQuota' in domainAccountSetting.keys():
                return int(domainAccountSetting['defaultQuota'])
            else:
                return 0
        else:
            try:
                result = self.conn.search_s(
                        self.dn,
                        ldap.SCOPE_BASE,
                        '(domainName=%s)' % self.domain,
                        ['domainName', 'accountSetting'],
                        )

                settings = ldaputils.getAccountSettingFromLdapQueryResult(result, key='domainName',)

                if 'defaultQuota' in settings[self.domain].keys():
                    return int(settings[self.domain]['defaultQuota'])
                else:
                    return 0
            except Exception, e:
                return 0

    # Delete domain.
    @decorators.require_global_admin
    def delete(self, domains=[]):
        if domains is None or len(domains) == 0:
            return (False, 'INVALID_DOMAIN_NAME')

        msg = {}
        for domain in domains:
            if not iredutils.isDomain(domain):
                continue

            dn = ldaputils.convKeywordToDN(web.safestr(domain), accountType='domain')
        
            try:
                deltree.DelTree(self.conn, dn, ldap.SCOPE_SUBTREE)
                web.logger(msg="Delete domain: %s." % (domain), domain=domain, event='delete',)
            except ldap.LDAPError, e:
                msg[domain] = str(e)

            # Delete records from SQL database: real-time used quota.
            if session.get('enableShowUsedQuota', False) is True:
                try:
                    # SQL: DELETE FROM table WHERE username LIKE '%@domain.ltd'
                    web.admindb.delete(
                        models.UsedQuota.__table__,
                        where='%s LIKE %s' % (
                            models.UsedQuota.username,
                            web.sqlquote('%%@'+domain),
                        ),
                    )
                except Exception, e:
                    pass

        if msg == {}:
            return (True,)
        else:
            return (False, ldaputils.getExceptionDesc(msg))

    @decorators.require_global_admin
    def enableOrDisableAccount(self, domains, action, attr='accountStatus',):
        if domains is None or len(domains) == 0:
            return (False, 'NO_DOMAIN_SELECTED')

        result = {}
        connutils = connUtils.Utils()
        for domain in domains:
            self.domain = web.safestr(domain)
            self.dn = ldaputils.convKeywordToDN(self.domain, accountType='domain')

            try:
                connutils.enableOrDisableAccount(
                    domain=self.domain,
                    account=self.domain,
                    dn=self.dn,
                    action=web.safestr(action).strip().lower(),
                    accountTypeInLogger='domain',
                )
            except ldap.LDAPError, e:
                result[self.domain] = str(e)

        if result == {}:
            return (True,)
        else:
            return (False, ldaputils.getExceptionDesc(result))

    # Get domain attributes & values.
    @decorators.require_domain_access
    def profile(self, domain):
        self.domain = web.safestr(domain)
        self.dn = ldaputils.convKeywordToDN(self.domain, accountType='domain')

        try:
            self.domain_profile = self.conn.search_s(
                self.dn,
                ldap.SCOPE_BASE,
                '(&(objectClass=mailDomain)(domainName=%s))' % self.domain,
                attrs.DOMAIN_ATTRS_ALL,
            )
            if len(self.domain_profile) == 1:
                return (True, self.domain_profile)
            else:
                return (False, 'NO_SUCH_DOMAIN')
        except ldap.NO_SUCH_OBJECT:
            return (False, 'NO_SUCH_OBJECT')
        except Exception, e:
            return (False, ldaputils.getExceptionDesc(e))

    # Update domain profile.
    # data = web.input()
    def update(self, profile_type, domain, data):
        self.profile_type = web.safestr(profile_type)
        self.domain = web.safestr(domain)
        self.domaindn = ldaputils.convKeywordToDN(self.domain, accountType='domain')

        connutils = connUtils.Utils()
        self.accountSetting = []
        mod_attrs = []

        # Allow normal admin to update profiles.
        if self.profile_type == 'general':
            cn = data.get('cn', None)
            mod_attrs += ldaputils.getSingleModAttr(attr='cn', value=cn, default=self.domain)

            ##################
            # Disclaimer
            #
            self.disclaimer = data.get('disclaimer', None)

            if self.disclaimer is not None and self.disclaimer != u'' and self.disclaimer != '':
                mod_attrs += [(ldap.MOD_REPLACE, 'disclaimer', self.disclaimer.encode('utf-8'))]
            else:
                mod_attrs += [(ldap.MOD_REPLACE, 'disclaimer', None)]

        elif self.profile_type == 'advanced':
            # Storage base directory.
            defaultStorageBaseDirectory = data.get('defaultStorageBaseDirectory', None)
            if defaultStorageBaseDirectory is not None and str(defaultStorageBaseDirectory) is not '':
                self.accountSetting += [ 'defaultStorageBaseDirectory:%s' % str(defaultStorageBaseDirectory) ]

            # Value type is number.
            for k in ['defaultQuota', 'minPasswordLength', 'maxPasswordLength',]:
                number = web.safestr(data.get(k, '0'))
                if number.isdigit() and number != '0':
                    self.accountSetting += [ '%s:%s' % (k, number) ]

            for k in ['defaultList',]:
                val = [web.safestr(v) for v in data.get(k, []) if iredutils.isEmail(v)]
                if len(val) >= 1:
                    self.accountSetting += ['defaultList:%s' % (','.join(val))]
        else:
            pass

        # Allow global admin to update profiles.
        if session.get('domainGlobalAdmin') is True:
            if self.profile_type == 'general':
                # Get accountStatus.
                if 'accountStatus' in data.keys():
                    accountStatus = 'active'
                else:
                    accountStatus = 'disabled'

                mod_attrs += [ (ldap.MOD_REPLACE, 'accountStatus', accountStatus) ]

                #########################
                # Domain admins
                #
                # Get admins from form data.
                self.domainAdmin = [ web.safestr(v) for v in data.get('domainAdmin') ]

                # Delete all admins first and then add admins which get from form data.
                result = connutils.addOrDelAttrValue(
                        dn=self.domaindn,
                        attr='domainAdmin',
                        value=None,
                        action='delete',
                        )
                if result[0] is False:
                    return result

                # Add admins if available.
                if self.domainAdmin != []:
                    # Assign admins.
                    result = connutils.addOrDelAttrValue(
                            dn=self.domaindn,
                            attr='domainAdmin',
                            value=self.domainAdmin,
                            action='add',
                            )
                    if result[0] is False:
                        return result
                else:
                    # No admin assigned.
                    pass

            elif self.profile_type == 'aliases':
                # Enable/Disable domain alias.
                if 'domainalias' in data.keys():
                    mod_type = 'add'
                else:
                    mod_type = 'delete'

                try:
                    result = connutils.addOrDelAttrValue(
                        dn=self.domaindn,
                        attr='enabledService',
                        value='domainalias',
                        action=mod_type,
                    )
                    if result[0] is False:
                        return result
                except Exception, e:
                    pass

                # Get domain aliases from web form and store in LDAP.
                self.domainAliasName = [
                    web.safestr(v)
                    for v in data.get('domainAliasName')
                    if iredutils.isDomain(v)
                ]
                # Remove duplicate values.
                self.domainAliasName = list(set(self.domainAliasName))
                mod_attrs += [(ldap.MOD_REPLACE, 'domainAliasName', self.domainAliasName) ]

                #
                # Sync attributes of existing user accounts for domain alias.
                #
                #   - shadowAddress: user alias
                #   - memberOfGroup: alias of mail list
                #

                # Get domain aliases from ldap with real-time querying.
                self.oldDomainAliasName = []
                try:
                    result = self.conn.search_s(
                        self.domaindn,
                        ldap.SCOPE_BASE,
                        '(&(objectClass=mailDomain)(domainName=%s))' % self.domain,
                        ['domainAliasName'],
                    )
                    self.oldDomainAliasName.extend(result[0][1].get('domainAliasName', []))
                except Exception, e:
                    pass

                # Get removed domain aliases and sync all user accounts.
                self.removedDomainAliasName = [v
                                               for v in self.oldDomainAliasName
                                               if v not in self.domainAliasName and iredutils.isDomain(v)
                                              ]

                # Get new domain aliases and sync all user accounts.
                self.newDomainAliasName = [v
                                           for v in self.domainAliasName
                                           if v not in self.oldDomainAliasName and iredutils.isDomain(v)
                                          ]

                # Remove @shadowAddress from existing users.
                for d in self.removedDomainAliasName:
                    try:
                        # Get list all users.
                        result = self.conn.search_s(
                            attrs.DN_BETWEEN_USER_AND_DOMAIN + self.domaindn,
                            ldap.SCOPE_ONELEVEL,
                            '(&(objectClass=mailUser)(|(shadowAddress=*@%s)(memberOfGroup=*@%s)))' % (d, d),
                            ['mail', 'shadowAddress', 'memberOfGroup',],
                        )

                        # Get accounts which need to be updated.
                        # Format:
                        #   {'dn':
                        #           {'shadowAddress': ['mail_of_shadowAddress', ...],},
                        #           {'memberOfGroup': ['mail_of_memberOfGroup', ...],},
                        #   }
                        #           
                        tmpUpdatedObjs = {}

                        # Get list of shadowAddress which ends with @self.removedDomainAliasName.
                        for obj in result:
                            tmpDnOfUser = obj[0]
                            tmpUpdatedObjs[tmpDnOfUser] = {}

                            # List of existing shadowAddress, memberOfGroup.
                            tmpRemovedShadowAddresses = []
                            tmpRemovedMemberOfGroups = []

                            # Get shadowAddress.
                            for addr in obj[1].get('shadowAddress', []):
                                if addr.endswith('@' + d):
                                    tmpRemovedShadowAddresses += [addr]

                            if len(tmpRemovedShadowAddresses) > 0:
                                tmpUpdatedObjs[tmpDnOfUser]['shadowAddress'] = tmpRemovedShadowAddresses

                            # Get memberOfGroup.
                            for addr in obj[1].get('memberOfGroup', []):
                                if addr.endswith('@' + d):
                                    tmpRemovedMemberOfGroups += [addr]

                            if len(tmpRemovedMemberOfGroups) > 0:
                                tmpUpdatedObjs[tmpDnOfUser]['memberOfGroup'] = tmpRemovedMemberOfGroups

                        # Removing objects from LDAP.
                        for tmpDN in tmpUpdatedObjs:
                            for tmpAttrName in ['shadowAddress', 'memberOfGroup', ]:
                                for tmpRemovedValueOfAttr in tmpUpdatedObjs[tmpDN].get(tmpAttrName, []):
                                    try:
                                        result = connutils.addOrDelAttrValue(
                                            dn=tmpDN,
                                            attr=tmpAttrName,
                                            value=tmpRemovedValueOfAttr,
                                            action='delete',
                                        )
                                    except Exception, e:
                                        pass
                    except Exception, e:
                        pass

                ################################
                # Add new values of attributes:
                #   - shadowAddress
                #   - memberOfGroup
                #
                for d in self.newDomainAliasName:
                    try:
                        # Get list all users.
                        result = self.conn.search_s(
                            attrs.DN_BETWEEN_USER_AND_DOMAIN + self.domaindn,
                            ldap.SCOPE_ONELEVEL,
                            '(objectClass=mailUser)',
                            ['mail', 'shadowAddress', 'memberOfGroup',],
                        )

                        # Get accounts which need to be updated.
                        # Format:
                        #   {'dn':
                        #           {'shadowAddress': ['mail_of_shadowAddress', ...],},
                        #           {'memberOfGroup': ['mail_of_memberOfGroup', ...],},
                        #   }
                        tmpUpdatedObjs = {}

                        for obj in result:
                            tmpDnOfUser = obj[0]
                            tmpMail = obj[1]['mail'][0]
                            tmpUsername = tmpMail.split('@')[0]

                            # Init.
                            tmpUpdatedObjs[tmpDnOfUser] = {}
                            tmpNewShadowAddresses = []
                            tmpNewMemberOfGroups = []

                            # Add new shadowAddress.
                            if tmpUsername + '@' + d not in obj[1].get('shadowAddress', []):
                                tmpNewShadowAddresses = [tmpUsername + '@' + d]

                            if len(tmpNewShadowAddresses) > 0:
                                tmpUpdatedObjs[tmpDnOfUser]['shadowAddress'] = tmpNewShadowAddresses

                            # Add new memberOfGroup.
                            #
                            # If user is already assigned to a mail list under
                            # same domain, we should add new mail list of alias
                            # domain.
                            #
                            #   - User: user@domain.ltd
                            #   - Existing mail list: memberOfGroup=all@domain.ltd
                            #   - New mail list: memberOfGroup=all@alias_domain.com
                            #
                            for tmpExistingMemberOfGroup in obj[1].get('memberOfGroup', []):
                                if tmpExistingMemberOfGroup.endswith('@' + self.domain):
                                    tmpNewMemberOfGroups += [tmpExistingMemberOfGroup.split('@')[0] + '@' + d]

                            if len(tmpNewMemberOfGroups) > 0:
                                tmpUpdatedObjs[tmpDnOfUser]['memberOfGroup'] = tmpNewMemberOfGroups

                        # Adding in LDAP.
                        for tmpDN in tmpUpdatedObjs:
                            for tmpAttrName in ['shadowAddress', 'memberOfGroup', ]:
                                for tmpAddedValueOfAttr in tmpUpdatedObjs[tmpDN].get(tmpAttrName, []):
                                    try:
                                        result = connutils.addOrDelAttrValue(
                                            dn=tmpDN,
                                            attr=tmpAttrName,
                                            value=tmpAddedValueOfAttr,
                                            action='add',
                                        )
                                    except Exception, e:
                                        pass
                    except Exception, e:
                        pass

                #
                # End Sync.
                #####################

            elif self.profile_type == 'bcc':
                # Enable/Disable recipient bcc.
                if 'recipientbcc' in data.keys():
                    mod_type = 'add'
                else:
                    mod_type = 'delete'

                result = connutils.addOrDelAttrValue(
                        dn=self.domaindn,
                        attr='enabledService',
                        value='recipientbcc',
                        action=mod_type,
                        )
                if result[0] is False:
                    return result

                # Update recipient bcc address.
                self.recipientBccAddress = data.get('recipientBccAddress', None)
                mod_attrs += ldaputils.getSingleModAttr(attr='domainRecipientBccAddress', value=self.recipientBccAddress, default=None,)

                # Enable/Disable sender bcc.
                if 'senderbcc' in data.keys():
                    mod_type = 'add'
                else:
                    mod_type = 'delete'

                result = connutils.addOrDelAttrValue(
                        dn=self.domaindn,
                        attr='enabledService',
                        value='senderbcc',
                        action=mod_type,
                        )
                if result[0] is False:
                    return result

                # Update sender bcc address.
                self.senderBccAddress = data.get('senderBccAddress', None)
                mod_attrs += ldaputils.getSingleModAttr(attr='domainSenderBccAddress', value=self.senderBccAddress, default=None,)

            elif self.profile_type == 'relay':
                self.default_mtaTransport = str(cfg.general.get('mtaTransport', 'dovecot'))
                if 'mtaTransport' in data.keys():
                    # Store current setting.
                    self.transport = data.get('mtaTransport', self.default_mtaTransport)
                    mod_attrs += [ (ldap.MOD_REPLACE, 'mtaTransport', str(self.transport)) ]
                else:
                    # Replace current setting by default transport.
                    mod_attrs += [ (ldap.MOD_REPLACE, 'mtaTransport', self.default_mtaTransport) ]

            elif self.profile_type == 'catchall':
                self.catchallAddress = set([ web.safestr(v)
                                        for v in data.get('catchallAddress', '').split(',')
                                        if iredutils.isEmail(v)
                                       ])

                # Get dn of catchall object.
                catchallDN = ldaputils.convKeywordToDN(self.domain, accountType='catchall')

                tmp_mod_attrs = [(ldap.MOD_REPLACE, 'mailForwardingAddress', list(self.catchallAddress))]

                if 'accountStatus' in data.keys():
                    # Enable catch-all.

                    try:
                        # Try to modify account first.
                        tmp_mod_attrs += [(ldap.MOD_REPLACE, 'accountStatus', 'active')]

                        self.conn.modify_s(catchallDN, tmp_mod_attrs)
                        return (True,)
                    except ldap.NO_SUCH_OBJECT:
                        # Add catch-all account.
                        catchallLdif = iredldif.ldif_catchall(self.domain, self.catchallAddress)
                        try:
                            self.conn.add_s(catchallDN, catchallLdif)
                            web.logger(
                                msg="Create catch-all account of domain: %s." % (self.domain),
                                domain=self.domain,
                                event='create',
                            )
                        except ldap.ALREADY_EXISTS:
                            pass
                        except Exception, e:
                            return (False, ldaputils.getExceptionDesc(e))
                            
                    except Exception, e:
                        return (False, ldaputils.getExceptionDesc(e))
                else:
                    if len(self.catchallAddress) == 0:
                        # Delete account if self.catchallAddress is empty.
                        try:
                            deltree.DelTree(self.conn, catchallDN, ldap.SCOPE_SUBTREE )
                            return (True,)
                        except Exception, e:
                            return (False, ldaputils.getExceptionDesc(e))
                    else:
                        # Mark as disabled.
                        try:
                            tmp_mod_attrs += [(ldap.MOD_REPLACE, 'accountStatus', 'disabled')]
                            self.conn.modify_s(catchallDN, tmp_mod_attrs)
                            return (True,)
                        except Exception, e:
                            return (False, ldaputils.getExceptionDesc(e))

            elif self.profile_type == 'throttle':
                #
                # Policyd throttling integration.
                #

                self.senderThrottlingSetting = throttle.getSenderThrottlingSettingFromForm(
                    account='@'+self.domain,
                    accountType='domain',
                    form=data,
                )

                self.recipientThrottlingSetting = throttle.getRecipientThrottlingSettingFromForm(
                    account='@'+self.domain,
                    accountType='domain',
                    form=data,
                )

                throttleLib = throttle.Throttle()
                try:
                    throttleLib.updateThrottlingSetting(
                        account='@'+self.domain,
                        accountType='sender',
                        setting=self.senderThrottlingSetting,
                    )

                    throttleLib.updateThrottlingSetting(
                        account='@'+self.domain,
                        accountType='recipient',
                        setting=self.recipientThrottlingSetting,
                    )
                except Exception, e:
                    pass

            elif self.profile_type == 'advanced':
                # Get number of account limit.
                for k in ['domainQuota', 'numberOfUsers', 'numberOfLists', 'numberOfAliases',]:
                    number = web.safestr(data.get(k, '0'))
                    if number.isdigit() and number != '0':
                        if k == 'domainQuota':
                            # Append quota unit.
                            if int(number) > 0:
                                number += ':' + web.safestr(data.get('domainQuotaUnit', 'GB'))
                            else:
                                # Unknown value, set to unlimited.
                                number = '0:GB'
                        self.accountSetting += [ '%s:%s' % (k, number) ]

                mod_attrs += [ (ldap.MOD_REPLACE, 'accountSetting', self.accountSetting) ]

                #####################
                # enabledService.
                #
                self.enabledService = [web.safestr(v)
                        for v in data.get('enabledService', [])
                        if v in attrs.DOMAIN_ENABLED_SERVICE
                        ]
                self.removed_services = [v
                        for v in attrs.DOMAIN_SERVICE_UNDER_CONTROL
                        if v not in self.enabledService
                        ]

                if len(self.enabledService) != 0:
                    for i in self.enabledService:
                        result = connutils.addOrDelAttrValue(
                                dn=self.domaindn,
                                attr='enabledService',
                                value=i,
                                action='add',
                                )
                        if result[0] is False:
                            return result

                if len(self.removed_services) != 0:
                    for i in self.removed_services:
                        result = connutils.addOrDelAttrValue(
                                dn=self.domaindn,
                                attr='enabledService',
                                value=i,
                                action='delete',
                                )
                        if result[0] is False:
                            return result

            elif self.profile_type == 'backupmx':
                domainBackupMX = web.safestr(data.get('domainBackupMX', 'no'))
                if domainBackupMX not in attrs.VALUES_DOMAIN_BACKUPMX:
                    domainBackupMX = 'no'

                mod_attrs += [ (ldap.MOD_REPLACE, 'domainBackupMX', domainBackupMX) ]
            else:
                pass

        else:
            pass

        try:
            dn = ldaputils.convKeywordToDN(self.domain, accountType='domain')
            self.conn.modify_s(dn, mod_attrs)
            return (True,)
        except Exception, e:
            return (False, ldaputils.getExceptionDesc(e))

    @decorators.require_domain_access
    def getDomainAccountSetting(self, domain,):
        result = self.getAllDomains(
            filter='(&(objectClass=mailDomain)(domainName=%s))' % domain,
            attrs=['domainName', 'accountSetting',],
        )

        if result[0] is True:
            allDomains = result[1]
        else:
            return result

        # Get accountSetting of current domain.
        try:
            allAccountSettings = ldaputils.getAccountSettingFromLdapQueryResult(allDomains, key='domainName')
            return (True, allAccountSettings.get(domain, {}))
        except Exception, e:
            return (False, ldaputils.getExceptionDesc(e))