# encoding: utf-8

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

import sys
import types
import web
from libs import iredutils
from libs.policyd import core
from MySQLdb.constants import ER

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


def getSenderThrottlingSettingFromForm(account, accountType, form,):
    '''Get sender throttling setting from web from, returned as dict.

    >>> ts = getSenderThrottlingSettingFromForm(
                account='zhb@iredmail.org',
                accountType='user',
                form=form,
        )
    >>> ts = getSenderThrottlingSettingFromForm(
                account='@iredmail.org',
                accountType='domain',
                form=form,
        )

    ts result formats:
        - (True, {...})
            + True: Insert or update
            + {}: throttling columns and values
        - (False, None): Invalid data.
        - (False, 'xxx'): Delete record which account is 'xxx'.
    '''

    account = web.safestr(account)
    if not core.isThrottlingAccount(account):
        return (False, None)

    # -- Predefined values --
    # Sender                                    # -- SQL columns --
    senderEnableThrottle = False
    senderSizeOfSingleMessage = 0               # throttle._mail_size
    senderSizeOfSingleMessageCustom = 0         # - Custom @senderSizeOfSingleMessage
    senderSecondsOfTimeUnit = 0                 # throttle._time_limit
    senderSecondsOfTimeUnitCustom = 0           # - Custom @senderSecondsOfTimeUnit
    senderNumberOfMaxMessagesPerTimeUnit = 0    # throttle._count_max
    senderMaxMessageSizePerTimeUnit = 0         # throttle._quota_max
    senderPriority = session.get('priorityOfUserThrottling', '10')
    if accountType == 'domain':
        senderPriority = session.get('priorityOfDomainThrottling', '20')
    # ----

    # -- Get values from form.
    if 'senderEnableThrottle' in form.keys():
        senderEnableThrottle = True

    if senderEnableThrottle is True:
        # Enabled. Update throttling settings.
        senderSizeOfSingleMessage = web.safestr(form.get('senderSizeOfSingleMessage', '0'))
        senderSizeOfSingleMessageCustom = web.safestr(form.get('senderSizeOfSingleMessageCustom', '0'))
        if senderSizeOfSingleMessage == 'on':
            # Get custom value.
            if senderSizeOfSingleMessageCustom.isdigit():
                senderSizeOfSingleMessage = int(senderSizeOfSingleMessageCustom)
        else:
            # Get predefined value.
            if senderSizeOfSingleMessage.isdigit():
                senderSizeOfSingleMessage = int(senderSizeOfSingleMessage)

        senderSecondsOfTimeUnit = web.safestr(form.get('senderSecondsOfTimeUnit', '0'))
        senderSecondsOfTimeUnitCustom = web.safestr(form.get('senderSecondsOfTimeUnitCustom', '0'))
        if senderSecondsOfTimeUnit == 'on':
            # Get custom value.
            if senderSecondsOfTimeUnitCustom.isdigit():
                senderSecondsOfTimeUnit = int(senderSecondsOfTimeUnitCustom)
        else:
            # Get predefined value.
            if senderSecondsOfTimeUnit.isdigit():
                senderSecondsOfTimeUnit = int(senderSecondsOfTimeUnit)

        senderNumberOfMaxMessagesPerTimeUnit = web.safestr(form.get('senderNumberOfMaxMessagesPerTimeUnit', '0'))
        if senderNumberOfMaxMessagesPerTimeUnit.isdigit():
            senderNumberOfMaxMessagesPerTimeUnit = int(senderNumberOfMaxMessagesPerTimeUnit)

        senderMaxMessageSizePerTimeUnit = web.safestr(form.get('senderMaxMessageSizePerTimeUnit', '0'))
        if senderMaxMessageSizePerTimeUnit.isdigit():
            senderMaxMessageSizePerTimeUnit = int(senderMaxMessageSizePerTimeUnit)

        senderPriority = web.safestr(form.get('senderPriority', senderPriority))
        if senderPriority.isdigit():
            senderPriority = int(senderPriority)

        settings = {
            '_from': account,
            '_mail_size': senderSizeOfSingleMessage,
            '_time_limit': senderSecondsOfTimeUnit,
            '_count_max': senderNumberOfMaxMessagesPerTimeUnit,
            '_quota_max': senderMaxMessageSizePerTimeUnit,
            '_priority': senderPriority,
        }
        return (True, settings)
    else:
        # False: Delete record.
        return (False, account)

def getRecipientThrottlingSettingFromForm(account, accountType, form):
    # -- Predefined values --
    account = web.safestr(account)
    if not core.isThrottlingAccount(account):
        return (False, None)

    # Recipient
    recipientEnableThrottle = False
    recipientSecondsOfTimeUnit = 0              # throttle_rcpt._time_limit
    recipientSecondsOfTimeUnitCustom = 0        # - Custom @recipientSecondsOfTimeUnit
    recipientNumberOfMaxMessagesPerTimeUnit = 0 # throttle_rcpt._count_max
    # ----

    # -- Get values from form.
    if 'recipientEnableThrottle' in form.keys():
        recipientEnableThrottle = True

    if recipientEnableThrottle is True:
        # Enabled. Update throttling settings.
        recipientSecondsOfTimeUnit = web.safestr(form.get('recipientSecondsOfTimeUnit', '0'))
        recipientSecondsOfTimeUnitCustom = web.safestr(form.get('recipientSecondsOfTimeUnitCustom', '0'))
        if recipientSecondsOfTimeUnit == 'on':
            # Get custom value.
            if recipientSecondsOfTimeUnitCustom.isdigit():
                recipientSecondsOfTimeUnit = int(recipientSecondsOfTimeUnitCustom)
        else:
            # Get predefined value.
            if recipientSecondsOfTimeUnit.isdigit():
                recipientSecondsOfTimeUnit = int(recipientSecondsOfTimeUnit)

        recipientNumberOfMaxMessagesPerTimeUnit = web.safestr(form.get('recipientNumberOfMaxMessagesPerTimeUnit', '0'))
        if recipientNumberOfMaxMessagesPerTimeUnit.isdigit():
            recipientNumberOfMaxMessagesPerTimeUnit = int(recipientNumberOfMaxMessagesPerTimeUnit)

        settings = {
            '_rcpt': account,
            '_time_limit': recipientSecondsOfTimeUnit,
            '_count_max': recipientNumberOfMaxMessagesPerTimeUnit,
        }

        # True: Insert or update record.
        return (True, settings)
    else:
        # False: Delete record.
        return (False, account)


class Throttle(core.PolicydWrap):
    def list(self, sender=None, recipient=None):
        self.senderThrottleRecords = {}
        self.recipientThrottleRecords = {}

        if sender is not None:
            tmpResult = self.db.select(
                'throttle',
                dict(sender=str(sender)),
                where='''_from = $sender''',
                limit=1,
                #_test=True,
            )
            if len(tmpResult) == 1:
                self.senderThrottleRecords = list(tmpResult)[0]

        if recipient is not None:
            tmpResult = self.db.select(
                'throttle_rcpt',
                vars=dict(recipient=str(recipient)),
                where='''_rcpt = $recipient''',
                limit=1,
            )
            if len(tmpResult) == 1:
                self.recipientThrottleRecords = list(tmpResult)[0]

        return (True, self.senderThrottleRecords, self.recipientThrottleRecords)

    def delete(self, account, accountType='sender'):
        '''Delete throttling record.'''

        account = web.safestr(account)

        if not core.isThrottlingAccount(account):
            return (False, 'INVALID_ACCOUNT')

        if accountType == 'sender':
            dbTable = 'throttle'
            column = '_from'
        else:
            dbTable = 'throttle_rcpt'
            column = '_rcpt'

        try:
            self.db.delete(
                dbTable,
                where='''%s=%s''' % (column, web.sqlquote(account)),
            )
            return (True,)
        except Exception, e:
            return (False, str(e))

    def __addOrUpdateSenderThrottle(self, **kw):
        '''Try to INSERT new sender throttling record, if failed, try to UPDATE it.
        '''
        self.account = web.safestr(kw.get('_from', ''))
        if not core.isThrottlingAccount(self.account):
            return (False, 'INVALID_ACCOUNT')

        col_mail_size = int(kw.get('_mail_size', 0))
        col_time_limit = int(kw.get('_time_limit', 0))
        col_count_max = int(kw.get('_count_max', 0))
        col_quota_max = int(kw.get('_quota_max', 0))
        col_priority = int(kw.get('_priority', session.get('priorityOfUserThrottling', 10)))

        try:
            self.db.insert(
                'throttle',
                _from=self.account,
                _mail_size=col_mail_size,
                _time_limit=col_time_limit,
                _count_max=col_count_max,
                _quota_max=col_quota_max,
                _priority=col_priority,
                _date=iredutils.sqlUnixTimestamp,
            )
            return (True,)
        except Exception, e:
            if e.args[0] == ER.DUP_ENTRY:
                # Record already exists. Try to update it instead.
                try:
                    self.db.update(
                        'throttle',
                        vars=dict(account=self.account),
                        where='''_from = $account''',
                        _mail_size=col_mail_size,
                        _time_limit=col_time_limit,
                        _count_max=col_count_max,
                        _quota_max=col_quota_max,
                        _priority=col_priority,
                    )
                    return (True,)
                except Exception, e:
                    return (False, str(e))
            else:
                return (False, str(e))

    def __addOrUpdateRecipientThrottle(self, **kw):
        '''Try to INSERT new recipient throttling record, if failed, try to UPDATE it.
        '''
        self.account = web.safestr(kw.get('_rcpt', ''))
        if not core.isThrottlingAccount(self.account):
            return (False, 'INVALID_ACCOUNT')

        col_time_limit = int(kw.get('_time_limit', 0))
        col_count_max = int(kw.get('_count_max', 0))

        try:
            self.db.insert(
                'throttle_rcpt',
                _rcpt=self.account,
                _time_limit=col_time_limit,
                _count_max=col_count_max,
                _date=iredutils.sqlUnixTimestamp,
            )
            return (True,)
        except Exception, e:
            if e.args[0] == ER.DUP_ENTRY:
                # Record already exists. Try to update it instead.
                try:
                    self.db.update(
                        'throttle_rcpt',
                        vars=dict(account=self.account),
                        where='''_rcpt = $account''',
                        _time_limit=col_time_limit,
                        _count_max=col_count_max,
                        _date=iredutils.sqlUnixTimestamp,
                    )
                    return (True,)
                except Exception, e:
                    return (False, str(e))
            else:
                return (False, str(e))

    def updateThrottlingSetting(self, account, accountType, setting):
        # setting should be generated by
        #   - self.getSenderThrottlingSettingFromForm
        #   - self.getRecipientThrottlingSettingFromForm
        #
        # @accountType: [sender, recipient]

        self.account = web.safestr(account)
        self.accountType = web.safestr(accountType)
        if not core.isThrottlingAccount(self.account):
            return (False, 'INVALID_ACCOUNT')

        if not isinstance(setting, types.TupleType) or len(setting) != 2:
            return (False, 'INVALID_SETTING1')

        if setting[1] is None:
            return (False, 'INVALID_SETTING2')

        # Available formats here:
        #   - (True, {...})
        #   - (False, 'xxx')
        if setting[0] is True:
            # INSERT or UPDATE record.
            if self.accountType == 'sender':
                try:
                    self.__addOrUpdateSenderThrottle(**setting[1])
                except Exception, e:
                    return (False, str(e))
            else:
                try:
                    self.__addOrUpdateRecipientThrottle(**setting[1])
                except Exception, e:
                    return (False, str(e))
        else:
            # Delete record.
            if not core.isThrottlingAccount(setting[1]):
                return (False, 'INVALID_SETTING3')

            try:
                self.delete(account=self.account, accountType=self.accountType)
                return (True,)
            except Exception, e:
                return (False, str(e))