# 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 email
import logging
import re

from bs4 import BeautifulSoup
from cgi import escape
from collections import defaultdict
from email import policy
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.image import MIMEImage
from os.path import basename
from urllib.parse import quote

from pyquarantine import mailer


class BaseNotification(object):
    "Notification base class"
    notification_type = "base"

    def __init__(self, name, global_cfg, cfg, test=False):
        self.name = name
        self.logger = logging.getLogger(__name__)

    def notify(self, qid, storage_id, mailfrom, recipients, headers,
               fp, subgroups=None, named_subgroups=None, synchronous=False):
        fp.seek(0)
        pass


class EMailNotification(BaseNotification):
    "Notification class to send notifications via mail."
    notification_type = "email"
    _bad_tags = [
        "applet",
        "embed",
        "frame",
        "frameset",
        "head",
        "iframe",
        "script"
    ]
    _good_tags = [
        "a",
        "b",
        "br",
        "center",
        "div",
        "font",
        "h1",
        "h2",
        "h3",
        "h4",
        "h5",
        "h6",
        "i",
        "img",
        "li",
        "p",
        "pre",
        "span",
        "table",
        "td",
        "th",
        "tr",
        "tt",
        "u",
        "ul"
    ]
    _good_attributes = [
        "align",
        "alt",
        "bgcolor",
        "border",
        "cellpadding",
        "cellspacing",
        "class",
        "color",
        "colspan",
        "dir",
        "face",
        "headers",
        "height",
        "id",
        "name",
        "rowspan",
        "size",
        "src",
        "style",
        "title",
        "type",
        "valign",
        "value",
        "width"
    ]

    def __init__(self, name, global_cfg, cfg, test=False):
        super(EMailNotification, self).__init__(
            name, global_cfg, cfg, test)

        defaults = {
            "notification_email_replacement_img": "",
            "notification_email_strip_images": "false",
            "notification_email_parser_lib": "lxml"
        }
        # check config
        for opt in [
            "notification_email_smtp_host",
            "notification_email_smtp_port",
            "notification_email_envelope_from",
            "notification_email_from",
            "notification_email_subject",
            "notification_email_template",
                "notification_email_embedded_imgs"] + list(defaults.keys()):
            if opt in cfg:
                continue
            if opt in global_cfg:
                cfg[opt] = global_cfg[opt]
            elif opt in defaults:
                cfg[opt] = defaults[opt]
            else:
                raise RuntimeError(
                    f"mandatory option '{opt}' not present in config "
                    f"section '{self.name}' or 'global'")

        self.smtp_host = cfg["notification_email_smtp_host"]
        self.smtp_port = cfg["notification_email_smtp_port"]
        self.mailfrom = cfg["notification_email_envelope_from"]
        self.from_header = cfg["notification_email_from"]
        self.subject = cfg["notification_email_subject"]

        testvars = defaultdict(str, test="TEST")

        # test-parse from header
        try:
            self.from_header.format_map(testvars)
        except ValueError as e:
            raise RuntimeError(
                f"error parsing notification_email_from: {e}")

        # test-parse subject
        try:
            self.subject.format_map(testvars)
        except ValueError as e:
            raise RuntimeError(
                f"error parsing notification_email_subject: {e}")

        # read and parse email notification template
        try:
            self.template = open(
                cfg["notification_email_template"], "r").read()
            self.template.format_map(testvars)
        except IOError as e:
            raise RuntimeError(f"error reading template: {e}")
        except ValueError as e:
            raise RuntimeError(f"error parsing template: {e}")

        strip_images = cfg["notification_email_strip_images"].strip().upper()
        if strip_images in ["TRUE", "ON", "YES"]:
            self.strip_images = True
        elif strip_images in ["FALSE", "OFF", "NO"]:
            self.strip_images = False
        else:
            raise RuntimeError(
                "error parsing notification_email_strip_images: unknown value")

        self.parser_lib = cfg["notification_email_parser_lib"].strip()
        if self.parser_lib not in ["lxml", "html.parser"]:
            raise RuntimeError(
                "error parsing notification_email_parser_lib: unknown value")

        # read email replacement image if specified
        replacement_img = cfg["notification_email_replacement_img"].strip()
        if not self.strip_images and replacement_img:
            try:
                self.replacement_img = MIMEImage(
                    open(replacement_img, "rb").read())
            except IOError as e:
                raise RuntimeError(
                    f"error reading replacement image: {e}")
            else:
                self.replacement_img.add_header(
                    "Content-ID", "<removed_for_security_reasons>")
        else:
            self.replacement_img = None

        # read images to embed if specified
        embedded_img_paths = []
        for p in cfg["notification_email_embedded_imgs"].split(","):
            if p:
                embedded_img_paths.append(p.strip())
        self.embedded_imgs = []
        for img_path in embedded_img_paths:
            # read image
            try:
                img = MIMEImage(open(img_path, "rb").read())
            except IOError as e:
                raise RuntimeError(f"error reading image: {e}")
            else:
                filename = basename(img_path)
                img.add_header(f"Content-ID", f"<{filename}>")
                self.embedded_imgs.append(img)

    def get_email_body_soup(self, qid, msg):
        "Extract and decode email body and return it as BeautifulSoup object."
        # try to find the body part
        self.logger.debug(f"{qid}: trying to find email body")
        try:
            body = msg.get_body(preferencelist=("html", "plain"))
        except Exception as e:
            self.logger.error(
                f"{qid}: an error occured in "
                f"email.message.EmailMessage.get_body: {e}")
            body = None

        if body:
            charset = body.get_content_charset() or "utf-8"
            content = body.get_payload(decode=True)
            try:
                content = content.decode(encoding=charset, errors="replace")
            except LookupError:
                self.logger.info(
                    f"{qid}: unknown encoding '{charset}', "
                    f"falling back to UTF-8")
                content = content.decode("utf-8", errors="replace")
            content_type = body.get_content_type()
            if content_type == "text/plain":
                # convert text/plain to text/html
                self.logger.debug(
                    f"{qid}: content type is {content_type}, "
                    f"converting to text/html")
                content = re.sub(r"^(.*)$", r"\1<br/>",
                                 escape(content), flags=re.MULTILINE)
            else:
                self.logger.debug(
                    f"{qid}: content type is {content_type}")
        else:
            self.logger.error(
                f"{qid}: unable to find email body")
            content = "ERROR: unable to find email body"

        # create BeautifulSoup object
        length = len(content)
        self.logger.debug(
            f"{qid}: trying to create BeatufilSoup object with "
            f"parser lib {self.parser_lib}, "
            f"text length is {length} bytes")
        soup = BeautifulSoup(content, self.parser_lib)
        self.logger.debug(
            f"{qid}: sucessfully created BeautifulSoup object")

        return soup

    def sanitize(self, qid, soup):
        "Sanitize mail html text."
        self.logger.debug(f"{qid}: sanitizing email text")

        # completly remove bad elements
        for element in soup(EMailNotification._bad_tags):
            self.logger.debug(
                f"{qid}: removing dangerous tag '{element.name}' "
                f"and its content")
            element.extract()

        # remove not whitelisted elements, but keep their content
        for element in soup.find_all(True):
            if element.name not in EMailNotification._good_tags:
                self.logger.debug(
                    f"{qid}: removing tag '{element.name}', keep its content")
                element.replaceWithChildren()

        # remove not whitelisted attributes
        for element in soup.find_all(True):
            for attribute in list(element.attrs.keys()):
                if attribute not in EMailNotification._good_attributes:
                    if element.name == "a" and attribute == "href":
                        self.logger.debug(
                            f"{qid}: setting attribute href to '#' "
                            f"on tag '{element.name}'")
                        element["href"] = "#"
                    else:
                        self.logger.debug(
                            f"{qid}: removing attribute '{attribute}' "
                            f"from tag '{element.name}'")
                        del(element.attrs[attribute])
        return soup

    def notify(self, qid, storage_id, mailfrom, recipients, headers, fp,
               subgroups=None, named_subgroups=None, synchronous=False):
        "Notify recipients via email."
        super(
            EMailNotification,
            self).notify(
            qid,
            storage_id,
            mailfrom,
            recipients,
            headers,
            fp,
            subgroups,
            named_subgroups,
            synchronous)

        # extract body from email
        soup = self.get_email_body_soup(
            qid, email.message_from_binary_file(fp, policy=policy.default))

        # replace picture sources
        image_replaced = False
        if self.strip_images:
            self.logger.debug(
                f"{qid}: looking for images to strip")
            for element in soup("img"):
                if "src" in element.attrs.keys():
                    self.logger.debug(
                        f"{qid}: strip image: {element['src']}")
                element.extract()
        elif self.replacement_img:
            self.logger.debug(
                f"{qid}: looking for images to replace")
            for element in soup("img"):
                if "src" in element.attrs.keys():
                    self.logger.debug(
                        f"{qid}: replacing image: {element['src']}")
                element["src"] = "cid:removed_for_security_reasons"
                image_replaced = True

        # sanitizing email text of original email
        sanitized_text = self.sanitize(qid, soup)
        del soup

        # sending email notifications
        for recipient in recipients:
            self.logger.debug(
                f"{qid}: generating notification email for '{recipient}'")
            self.logger.debug(f"{qid}: parsing email template")

            # generate dict containing all template variables
            variables = defaultdict(
                    str,
                    EMAIL_HTML_TEXT=sanitized_text,
                    EMAIL_FROM=escape(headers["from"]),
                    EMAIL_ENVELOPE_FROM=escape(mailfrom),
                    EMAIL_ENVELOPE_FROM_URL=escape(quote(mailfrom)),
                    EMAIL_TO=escape(headers["to"]),
                    EMAIL_ENVELOPE_TO=escape(recipient),
                    EMAIL_ENVELOPE_TO_URL=escape(quote(recipient)),
                    EMAIL_SUBJECT=escape(headers["subject"]),
                    EMAIL_QUARANTINE_ID=storage_id)

            if subgroups:
                number = 0
                for subgroup in subgroups:
                    variables[f"SUBGROUP_{number}"] = escape(subgroup)
            if named_subgroups:
                for key, value in named_subgroups.items():
                    named_subgroups[key] = escape(value)
                variables.update(named_subgroups)

            # parse template
            htmltext = self.template.format_map(variables)

            msg = MIMEMultipart('related')
            msg["From"] = self.from_header.format_map(
                defaultdict(str, EMAIL_FROM=headers["from"]))
            msg["To"] = headers["to"]
            msg["Subject"] = self.subject.format_map(
                defaultdict(str, EMAIL_SUBJECT=headers["subject"]))
            msg["Date"] = email.utils.formatdate()
            msg.attach(MIMEText(htmltext, "html", 'UTF-8'))

            if image_replaced:
                self.logger.debug(
                    f"{qid}: attaching notification_replacement_img")
                msg.attach(self.replacement_img)

            for img in self.embedded_imgs:
                self.logger.debug(f"{qid}: attaching imgage")
                msg.attach(img)

            self.logger.debug(
                f"{qid}: sending notification email to: {recipient}")
            if synchronous:
                try:
                    mailer.smtp_send(self.smtp_host, self.smtp_port,
                                     self.mailfrom, recipient, msg.as_string())
                except Exception as e:
                    raise RuntimeError(
                        f"error while sending email to '{recipient}': {e}")
            else:
                mailer.sendmail(self.smtp_host, self.smtp_port, qid,
                                self.mailfrom, recipient, msg.as_string(),
                                "notification email")


# list of notification types and their related notification classes
TYPES = {"email": EMailNotification}