# PyQuarantine-Milter is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# PyQuarantine-Milter is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with PyQuarantineMilter.  If not, see <http://www.gnu.org/licenses/>.
#

import json
import logging
import os

from calendar import timegm
from datetime import datetime
from glob import glob
from shutil import copyfileobj
from time import gmtime

from pyquarantine import mailer


class BaseQuarantine(object):
    "Quarantine base class"
    def __init__(self, global_config, config, configtest=False):
        self.name = config["name"]
        self.global_config = global_config
        self.config = config
        self.logger = logging.getLogger(__name__)

    def add(self, queueid, mailfrom, recipients, fp):
        "Add e-mail to quarantine."
        fp.seek(0)
        return ""

    def find(self, mailfrom=None, recipients=None, older_than=None):
        "Find e-mails in quarantine."
        return
    
    def get_metadata(self, quarantine_id):
        "Return metadata of quarantined e-mail."
        return

    def delete(self, quarantine_id, recipient=None):
        "Delete e-mail from quarantine."
        return

    def release(self, quarantine_id, recipient=None):
        "Release e-mail from quarantine."
        return


class FileQuarantine(BaseQuarantine):
    "Quarantine class to store mails on filesystem."
    def __init__(self, global_config, config, configtest=False):
        super(FileQuarantine, self).__init__(global_config, config, configtest)

        # check if mandatory options are present in config
        for option in ["directory"]:
            if option not in self.config.keys() and option in self.global_config.keys():
                self.config[option] = self.global_config[option]
            if option not in self.config.keys():
                raise RuntimeError("mandatory option '{}' not present in config section '{}' or 'global'".format(option, self.name))
        self.directory = self.config["directory"]

        # check if quarantine directory exists and is writable
        if not os.path.isdir(self.directory) or not os.access(self.directory, os.W_OK):
            raise RuntimeError("file quarantine directory '{}' does not exist or is not writable".format(self.directory))
        self._metadata_suffix = ".metadata"

    def _save_datafile(self, quarantine_id, fp):
        datafile = os.path.join(self.directory, quarantine_id)
        try:
            with open(datafile, "wb") as f:
                copyfileobj(fp, f)
        except IOError as e:
            raise RuntimeError("unable save data file: {}".format(e))

    def _save_metafile(self, quarantine_id, metadata):
        metafile = os.path.join(self.directory, "{}{}".format(quarantine_id, self._metadata_suffix))
        try:
            with open(metafile, "w") as f:
                json.dump(metadata, f, indent=2)
        except IOError as e:
            raise RuntimeError("unable to save metadata file: {}".format(e))

    def _remove(self, quarantine_id):
        datafile = os.path.join(self.directory, quarantine_id)
        metafile = "{}{}".format(datafile, self._metadata_suffix)

        try:
            os.remove(metafile)
        except IOError as e:
            raise RuntimeError("unable to remove metadata file: {}".format(e))

        try:
            os.remove(datafile)
        except IOError as e:
            raise RuntimeError("unable to remove data file: {}".format(e))

    def add(self, queueid, mailfrom, recipients, fp):
        "Add e-mail to file quarantine and return quarantine-id."
        super(FileQuarantine, self).add(queueid, mailfrom, recipients, fp)
        quarantine_id = "{}_{}".format(datetime.now().strftime("%Y%m%d%H%M%S"), queueid)

        # save mail
        self._save_datafile(quarantine_id, fp)

        # save metadata
        metadata = {
            "from": mailfrom,
            "recipients": recipients,
            "date": timegm(gmtime()),
            "queue_id": queueid
        }
        try:
            self._save_metafile(quarantine_id, metadata)
        except RuntimeError as e:
            datafile = os.path.join(self.directory, quarantine_id)
            os.remove(datafile)
            raise e

        return quarantine_id

    def get_metadata(self, quarantine_id):
        "Return metadata of quarantined e-mail."
        super(FileQuarantine, self).get_metadata(quarantine_id)

        metafile = os.path.join(self.directory, "{}{}".format(quarantine_id, self._metadata_suffix))
        if not os.path.isfile(metafile):
            raise RuntimeError("invalid quarantine id '{}'".format(quarantine_id))

        try:
            with open(metafile, "r") as f:
                metadata = json.load(f)
        except IOError as e:
            raise RuntimeError("unable to read metadata file: {}".format(e))
        except json.JSONDecodeError as e:
            raise RuntimeError("invalid meta file '{}': {}".format(metafile, e))

        return metadata

    def find(self, mailfrom=None, recipients=None, older_than=None):
        "Find e-mails in quarantine."
        super(FileQuarantine, self).find(mailfrom, recipients, older_than)
        if type(mailfrom) == str: mailfrom = [mailfrom]
        if type(recipients) == str: recipients = [recipients]

        emails = {}
        metafiles = glob(os.path.join(self.directory, "*{}".format(self._metadata_suffix)))
        for metafile in metafiles:
            if not os.path.isfile(metafile): continue

            quarantine_id = os.path.basename(metafile[:-len(self._metadata_suffix)])
            metadata = self.get_metadata(quarantine_id)
            if older_than != None:
                if timegm(gmtime()) - metadata["date"] < (older_than * 24 * 3600):
                    continue

            if mailfrom != None:
                if metadata["from"] not in mailfrom:
                    continue

            if recipients != None:
                if len(recipients) == 1 and recipients[0] not in metadata["recipients"]:
                    continue
                elif len(set(recipients + metadata["recipients"])) == len(recipients + metadata["recipients"]):
                    continue

            emails[quarantine_id] = metadata

        return emails

    def delete(self, quarantine_id, recipient=None):
        "Delete e-mail in quarantine."
        super(FileQuarantine, self).delete(quarantine_id, recipient)

        try:
            metadata = self.get_metadata(quarantine_id)
        except RuntimeError as e:
            raise RuntimeError("unable to delete e-mail: {}".format(e))

        if recipient == None:
            self._remove(quarantine_id)
        else:
            if recipient not in metadata["recipients"]:
                raise RuntimeError("invalid recipient '{}'".format(recipient))

            metadata["recipients"].remove(recipient)
            if not metadata["recipients"]:
                self._remove(quarantine_id)
            else:
                self._save_metafile(quarantine_id, metadata)

    def release(self, quarantine_id, recipient=None):
        "Release e-mail from quarantine."
        super(FileQuarantine, self).release(quarantine_id, recipient)

        try:
            metadata = self.get_metadata(quarantine_id)
        except RuntimeError as e:
            raise RuntimeError("unable to release e-mail: {}".format(e))

        datafile = os.path.join(self.directory, quarantine_id)
        if recipient != None:
            if recipient not in metadata["recipients"]:
                raise RuntimeError("invalid recipient '{}'".format(recipient))
            recipients = [recipient]
        else:
            recipients = metadata["recipients"]

        try:
            with open(datafile, "rb") as f:
                mail = f.read()
        except IOError as e:
            raise RuntimeError("unable to read data file: {}".format(e))

        for recipient in recipients:
            try:
                mailer.smtp_send(self.config["smtp_host"], self.config["smtp_port"], metadata["from"], recipient, mail)
            except Exception as e:
                raise RuntimeError("error while sending e-mail to '{}': {}".format(recipient, e))

            self.delete(quarantine_id, recipient)


# list of quarantine types and their related quarantine classes
TYPES = {"file": FileQuarantine}