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

import sys
import socket
import web
from libs import iredutils
from libs.amavisd import core, log as amavisdlog

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

# Import backend related modules.
if cfg.general.get('backend') == 'ldap':
    from libs.ldaplib import admin as adminlib
elif cfg.general.get('backend') == 'mysql':
    from libs.mysql import admin as adminlib

# If msgs.quar_type != 'Q' (SQL), we can't get mail body.
class Quarantine(core.AmavisdWrap):
    def getRawMessage(self, mail_id=None,):
        """Query SQL server and return raw mail message."""

        # TODO Check permission.
        if mail_id is None:
            return (False, 'INVALID_MAILID')

        mail_id = str(mail_id)

        try:
            records = self.db.select(
                'quarantine',
                what='mail_text',
                where='mail_id = %s' % web.sqlquote(mail_id),
                order='chunk_ind ASC',
            )

            if len(records) == 0:
                return (False, 'INVALID_MAILID')

            # Combine mail_text as RAW mail message.
            message = ''
            for i in list(records):
                for j in i.mail_text:
                    message += j

            return (True, message)
        except Exception, e:
            return (False, str(e))

    def getRecordsOfQuarantinedMails(self, cur_page=1, sizelimit=session.pageSizeLimit,):
        """Return (total, records)"""

        self.cur_page = int(cur_page)
        self.sizelimit = sizelimit

        # Pre-defined values.
        self.count = 0
        self.records = []
        # Reversed domain name if need to sql query server with domain names.
        allReversedDomainNames = []

        self.sql_append_selection = ''

        # Query SQL.
        # If not domainGlobalAdmin, we should list all controlled domains in query.
        if session.get('domainGlobalAdmin') is not True:
            if cfg.general.get('backend') == 'ldap':
                adminLib = adminlib.Admin()
                result_all_domains = adminLib.getManagedDomains(mail=session.get('username'), attrs=['domainName'],)
                if result_all_domains[0] is True:
                    for i in result_all_domains[1]:
                        domain = i[1].get('domainName', [''])[0]

                        reversedDomain = domain.split('.')
                        reversedDomain.reverse()
                        allReversedDomainNames += ['.'.join(reversedDomain)]

            elif cfg.general.get('backend') == 'mysql':
                adminLib = adminlib.Admin()
                result_all_domains = adminLib.getManagedDomains(session.get('username'), domainNameOnly=True)

                allDomains = []
                if result_all_domains[0] is True:
                    allDomains += result_all_domains[1]

                amavisdLogLib = amavisdlog.Log()
                allReversedDomainNames = amavisdLogLib.reverseDomainNames(allDomains)

            if len(allReversedDomainNames) > 0:
                self.sql_append_selection = ' AND recip.domain IN ' + web.sqlquote(allReversedDomainNames)
            else:
                return (0, {})

        # Get number of total records. SQL table: amavisd.msgs
        try:
            # Refer to templates/default/macros/amavisd.html for more detail
            # about msgs.content (content type, spam status), msgs.quar_type
            # (quarantine type).
            result = self.db.query("""
                                   SELECT COUNT(msgs.mail_id) AS total
                                   FROM msgs
                                   LEFT JOIN msgrcpt ON msgs.mail_id = msgrcpt.mail_id
                                   LEFT JOIN maddr AS sender ON msgs.sid = sender.id
                                   LEFT JOIN maddr AS recip ON msgrcpt.rid = recip.id
                                   WHERE msgs.content IN ('S', 's', 'Y')
                                   AND msgs.quar_type IN ('Q', 'F')
                                   %s
                                   """ % (self.sql_append_selection))
            self.count = result[0].total or 0
        except Exception, e:
            pass

        # Get records of quarantined mails.
        try:
            # msgs.content:
            #   - S: spam(kill)
            #   - s: prior to 2.7.0 the CC_SPAMMY was logged as 's', now 'Y' is used.
            # msgs.quar_type:
            #   - Q: sql
            #   - F: file
            result = self.db.query(
                '''
                SELECT
                    msgs.mail_id, msgs.secret_id, msgs.subject, msgs.time_iso, msgs.size,
                    sender.email as sender_email,
                    recip.email as recipient
                FROM msgs
                LEFT JOIN msgrcpt ON msgs.mail_id = msgrcpt.mail_id
                LEFT JOIN maddr AS sender ON msgs.sid = sender.id
                LEFT JOIN maddr AS recip ON msgrcpt.rid = recip.id
                WHERE
                    msgs.content IN ('S', 's', 'Y')
                    AND msgs.quar_type in ('Q', 'F') %s
                ORDER BY msgs.time_num DESC
                LIMIT %d
                OFFSET %d
                ''' % (self.sql_append_selection, self.sizelimit, (self.cur_page-1) * self.sizelimit,)
            )
            self.records = iredutils.convertAmavisdRecords(result)
        except Exception, e:
            pass

        # Delete old quarantined mails from table 'msgs'. It will also
        # delete records in table 'quarantine'.
        try:
            if iredutils.AMAVISD_REMOVE_QUARANTINED_IN_DAYS > 0:
                self.db.delete(
                    'msgs',
                    where="""content IN ('S', 's', 'V', 'Y') \
                            AND quar_type IN ('Q', 'F') \
                            AND FROM_UNIXTIME(time_num) < DATE_SUB(NOW(), INTERVAL %d DAY)
                    """ % (iredutils.AMAVISD_REMOVE_QUARANTINED_IN_DAYS,),
                )
        except:
            pass

        return (self.count, self.records)

    def releaseQuarantinedMails(self, records=[]):
        '''Release quarantined mails.

        >>> records = [
            {'mail_id': 'xxx',
             'secret_id': 'yyy',
             'requested_by': session.get('username'),
            }
            ...
        ]
        '''

        # Refer to amavisd doc 'README.protocol' for more detail:
        #   - Releasing a message from a quarantine

        # Don't waste time if no records is specified.
        if len(records) == 0:
            return (True,)

        # Pre-defined variables.
        self.releasedMailIDs = []

        # Get amavisd server address and quarantine port.
        self.amavisdServer = cfg.amavisd.get('server', '127.0.0.1')
        self.quarantinePort = int(cfg.amavisd.get('quarantine_port', 9998))

        # Create socket.
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.connect((self.amavisdServer, self.quarantinePort))
        except Exception, e:
            return (False, str(e))

        # Generate commands from dict, used for socket communication.
        # Note: We need to update Amavisd SQL database after mail was released
        #       with success, so do NOT send all release requests in ONE socket
        #       command although it will get better performance (a little).
        for record in records:
            cmd = ''

            # Skip record without 'mail_id'.
            if not 'mail_id' in record:
                continue

            for k in record.keys():
                if record[k] is not None and record[k] != '':
                    cmd += '%s=%s\r\n' % (k, record[k])

            try:
                s.send('request=release\r\n' + cmd + '\r\n')
                response = s.recv(1024)

                self.releasedMailIDs += [record.get('mail_id', 'NOT-EXIST')]
            except Exception, e:
                pass

        # Close socket.
        try:
            s.close()
        except Exception, e:
            return (False, str(e))

        # Update Amavisd SQL database.
        if len(self.releasedMailIDs) == 0:
            return (True, )

        try:
            # Update Amavisd SQL database.
            #
            #   - Update msgs.content to 'C' (Clean)
            #       UPDATE msgs \
            #       SET msgs.content = 'C' \
            #       WHERE msgs.mail_id IN ('xxx', 'yyy', ..)
            #
            #   - Delete records in 'quarantine':
            #       DELETE FROM quarantine \
            #       WHERE quarantine.partition_tag = msgs.partition_tag \
            #       AND quarantine.mail_id = msgs.mail_id
            #
            self.db.update(
                'msgs',
                where='mail_id IN ' + web.sqlquote(self.releasedMailIDs),
                content='C',
            )

            self.db.delete(
                'quarantine',
                where='mail_id IN ' + web.sqlquote(self.releasedMailIDs),
            )
            return (True, )
        except Exception, e:
            return (False, str(e))