# Copyright (C) 2001-2007 Python Software Foundation

# Contact: email-sig@python.org

# email package unit tests

import os

import sys

import time

import base64

import difflib

import unittest

import warnings

from cStringIO import StringIO

import email

from email.charset import Charset

from email.header import Header, decode_header, make_header

from email.parser import Parser, HeaderParser

from email.generator import Generator, DecodedGenerator

from email.message import Message

from email.mime.application import MIMEApplication

from email.mime.audio import MIMEAudio

from email.mime.text import MIMEText

from email.mime.image import MIMEImage

from email.mime.base import MIMEBase

from email.mime.message import MIMEMessage

from email.mime.multipart import MIMEMultipart

from email import utils

from email import errors

from email import encoders

from email import iterators

from email import base64mime

from email import quoprimime

from test.test_support import findfile, run_unittest

from email.test import __file__ as landmark

NL = '\n'

EMPTYSTRING = ''

SPACE = ' '

def openfile(filename, mode='r'):

    path = os.path.join(os.path.dirname(landmark), 'data', filename)

    return open(path, mode)

# Base test class

class TestEmailBase(unittest.TestCase):

    def ndiffAssertEqual(self, first, second):

        """Like failUnlessEqual except use ndiff for readable output."""

        if first <> second:

            sfirst = str(first)

            ssecond = str(second)

            diff = difflib.ndiff(sfirst.splitlines(), ssecond.splitlines())

            fp = StringIO()

            print >> fp, NL, NL.join(diff)

            raise self.failureException, fp.getvalue()

    def _msgobj(self, filename):

        fp = openfile(findfile(filename))

        try:

            msg = email.message_from_file(fp)

        finally:

            fp.close()

        return msg

# Test various aspects of the Message class's API

class TestMessageAPI(TestEmailBase):

    def test_get_all(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_20.txt')

        eq(msg.get_all('cc'), ['ccc@zzz.org', 'ddd@zzz.org', 'eee@zzz.org'])

        eq(msg.get_all('xx', 'n/a'), 'n/a')

    def test_getset_charset(self):

        eq = self.assertEqual

        msg = Message()

        eq(msg.get_charset(), None)

        charset = Charset('iso-8859-1')

        msg.set_charset(charset)

        eq(msg['mime-version'], '1.0')

        eq(msg.get_content_type(), 'text/plain')

        eq(msg['content-type'], 'text/plain; charset="iso-8859-1"')

        eq(msg.get_param('charset'), 'iso-8859-1')

        eq(msg['content-transfer-encoding'], 'quoted-printable')

        eq(msg.get_charset().input_charset, 'iso-8859-1')

        # Remove the charset

        msg.set_charset(None)

        eq(msg.get_charset(), None)

        eq(msg['content-type'], 'text/plain')

        # Try adding a charset when there's already MIME headers present

        msg = Message()

        msg['MIME-Version'] = '2.0'

        msg['Content-Type'] = 'text/x-weird'

        msg['Content-Transfer-Encoding'] = 'quinted-puntable'

        msg.set_charset(charset)

        eq(msg['mime-version'], '2.0')

        eq(msg['content-type'], 'text/x-weird; charset="iso-8859-1"')

        eq(msg['content-transfer-encoding'], 'quinted-puntable')

    def test_set_charset_from_string(self):

        eq = self.assertEqual

        msg = Message()

        msg.set_charset('us-ascii')

        eq(msg.get_charset().input_charset, 'us-ascii')

        eq(msg['content-type'], 'text/plain; charset="us-ascii"')

    def test_set_payload_with_charset(self):

        msg = Message()

        charset = Charset('iso-8859-1')

        msg.set_payload('This is a string payload', charset)

        self.assertEqual(msg.get_charset().input_charset, 'iso-8859-1')

    def test_get_charsets(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_08.txt')

        charsets = msg.get_charsets()

        eq(charsets, [None, 'us-ascii', 'iso-8859-1', 'iso-8859-2', 'koi8-r'])

        msg = self._msgobj('msg_09.txt')

        charsets = msg.get_charsets('dingbat')

        eq(charsets, ['dingbat', 'us-ascii', 'iso-8859-1', 'dingbat',

                      'koi8-r'])

        msg = self._msgobj('msg_12.txt')

        charsets = msg.get_charsets()

        eq(charsets, [None, 'us-ascii', 'iso-8859-1', None, 'iso-8859-2',

                      'iso-8859-3', 'us-ascii', 'koi8-r'])

    def test_get_filename(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_04.txt')

        filenames = [p.get_filename() for p in msg.get_payload()]

        eq(filenames, ['msg.txt', 'msg.txt'])

        msg = self._msgobj('msg_07.txt')

        subpart = msg.get_payload(1)

        eq(subpart.get_filename(), 'dingusfish.gif')

    def test_get_filename_with_name_parameter(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_44.txt')

        filenames = [p.get_filename() for p in msg.get_payload()]

        eq(filenames, ['msg.txt', 'msg.txt'])

    def test_get_boundary(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_07.txt')

        # No quotes!

        eq(msg.get_boundary(), 'BOUNDARY')

    def test_set_boundary(self):

        eq = self.assertEqual

        # This one has no existing boundary parameter, but the Content-Type:

        # header appears fifth.

        msg = self._msgobj('msg_01.txt')

        msg.set_boundary('BOUNDARY')

        header, value = msg.items()[4]

        eq(header.lower(), 'content-type')

        eq(value, 'text/plain; charset="us-ascii"; boundary="BOUNDARY"')

        # This one has a Content-Type: header, with a boundary, stuck in the

        # middle of its headers.  Make sure the order is preserved; it should

        # be fifth.

        msg = self._msgobj('msg_04.txt')

        msg.set_boundary('BOUNDARY')

        header, value = msg.items()[4]

        eq(header.lower(), 'content-type')

        eq(value, 'multipart/mixed; boundary="BOUNDARY"')

        # And this one has no Content-Type: header at all.

        msg = self._msgobj('msg_03.txt')

        self.assertRaises(errors.HeaderParseError,

                          msg.set_boundary, 'BOUNDARY')

    def test_get_decoded_payload(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_10.txt')

        # The outer message is a multipart

        eq(msg.get_payload(decode=True), None)

        # Subpart 1 is 7bit encoded

        eq(msg.get_payload(0).get_payload(decode=True),

           'This is a 7bit encoded message.\n')

        # Subpart 2 is quopri

        eq(msg.get_payload(1).get_payload(decode=True),

           '\xa1This is a Quoted Printable encoded message!\n')

        # Subpart 3 is base64

        eq(msg.get_payload(2).get_payload(decode=True),

           'This is a Base64 encoded message.')

        # Subpart 4 has no Content-Transfer-Encoding: header.

        eq(msg.get_payload(3).get_payload(decode=True),

           'This has no Content-Transfer-Encoding: header.\n')

    def test_get_decoded_uu_payload(self):

        eq = self.assertEqual

        msg = Message()

        msg.set_payload('begin 666 -\n+:&5L;&\\@=V]R;&0 \n \nend\n')

        for cte in ('x-uuencode', 'uuencode', 'uue', 'x-uue'):

            msg['content-transfer-encoding'] = cte

            eq(msg.get_payload(decode=True), 'hello world')

        # Now try some bogus data

        msg.set_payload('foo')

        eq(msg.get_payload(decode=True), 'foo')

    def test_decoded_generator(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_07.txt')

        fp = openfile('msg_17.txt')

        try:

            text = fp.read()

        finally:

            fp.close()

        s = StringIO()

        g = DecodedGenerator(s)

        g.flatten(msg)

        eq(s.getvalue(), text)

    def test__contains__(self):

        msg = Message()

        msg['From'] = 'Me'

        msg['to'] = 'You'

        # Check for case insensitivity

        self.failUnless('from' in msg)

        self.failUnless('From' in msg)

        self.failUnless('FROM' in msg)

        self.failUnless('to' in msg)

        self.failUnless('To' in msg)

        self.failUnless('TO' in msg)

    def test_as_string(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_01.txt')

        fp = openfile('msg_01.txt')

        try:

            text = fp.read()

        finally:

            fp.close()

        eq(text, msg.as_string())

        fullrepr = str(msg)

        lines = fullrepr.split('\n')

        self.failUnless(lines[0].startswith('From '))

        eq(text, NL.join(lines[1:]))

    def test_bad_param(self):

        msg = email.message_from_string("Content-Type: blarg; baz; boo\n")

        self.assertEqual(msg.get_param('baz'), '')

    def test_missing_filename(self):

        msg = email.message_from_string("From: foo\n")

        self.assertEqual(msg.get_filename(), None)

    def test_bogus_filename(self):

        msg = email.message_from_string(

        "Content-Disposition: blarg; filename\n")

        self.assertEqual(msg.get_filename(), '')

    def test_missing_boundary(self):

        msg = email.message_from_string("From: foo\n")

        self.assertEqual(msg.get_boundary(), None)

    def test_get_params(self):

        eq = self.assertEqual

        msg = email.message_from_string(

            'X-Header: foo=one; bar=two; baz=three\n')

        eq(msg.get_params(header='x-header'),

           [('foo', 'one'), ('bar', 'two'), ('baz', 'three')])

        msg = email.message_from_string(

            'X-Header: foo; bar=one; baz=two\n')

        eq(msg.get_params(header='x-header'),

           [('foo', ''), ('bar', 'one'), ('baz', 'two')])

        eq(msg.get_params(), None)

        msg = email.message_from_string(

            'X-Header: foo; bar="one"; baz=two\n')

        eq(msg.get_params(header='x-header'),

           [('foo', ''), ('bar', 'one'), ('baz', 'two')])

    def test_get_param_liberal(self):

        msg = Message()

        msg['Content-Type'] = 'Content-Type: Multipart/mixed; boundary = "CPIMSSMTPC06p5f3tG"'

        self.assertEqual(msg.get_param('boundary'), 'CPIMSSMTPC06p5f3tG')

    def test_get_param(self):

        eq = self.assertEqual

        msg = email.message_from_string(

            "X-Header: foo=one; bar=two; baz=three\n")

        eq(msg.get_param('bar', header='x-header'), 'two')

        eq(msg.get_param('quuz', header='x-header'), None)

        eq(msg.get_param('quuz'), None)

        msg = email.message_from_string(

            'X-Header: foo; bar="one"; baz=two\n')

        eq(msg.get_param('foo', header='x-header'), '')

        eq(msg.get_param('bar', header='x-header'), 'one')

        eq(msg.get_param('baz', header='x-header'), 'two')

        # XXX: We are not RFC-2045 compliant!  We cannot parse:

        # msg["Content-Type"] = 'text/plain; weird="hey; dolly? [you] @ <\\"home\\">?"'

        # msg.get_param("weird")

        # yet.

    def test_get_param_funky_continuation_lines(self):

        msg = self._msgobj('msg_22.txt')

        self.assertEqual(msg.get_payload(1).get_param('name'), 'wibble.JPG')

    def test_get_param_with_semis_in_quotes(self):

        msg = email.message_from_string(

            'Content-Type: image/pjpeg; name="Jim&&Jill"\n')

        self.assertEqual(msg.get_param('name'), 'Jim&&Jill')

        self.assertEqual(msg.get_param('name', unquote=False),

                         '"Jim&&Jill"')

    def test_has_key(self):

        msg = email.message_from_string('Header: exists')

        self.failUnless(msg.has_key('header'))

        self.failUnless(msg.has_key('Header'))

        self.failUnless(msg.has_key('HEADER'))

        self.failIf(msg.has_key('headeri'))

    def test_set_param(self):

        eq = self.assertEqual

        msg = Message()

        msg.set_param('charset', 'iso-2022-jp')

        eq(msg.get_param('charset'), 'iso-2022-jp')

        msg.set_param('importance', 'high value')

        eq(msg.get_param('importance'), 'high value')

        eq(msg.get_param('importance', unquote=False), '"high value"')

        eq(msg.get_params(), [('text/plain', ''),

                              ('charset', 'iso-2022-jp'),

                              ('importance', 'high value')])

        eq(msg.get_params(unquote=False), [('text/plain', ''),

                                       ('charset', '"iso-2022-jp"'),

                                       ('importance', '"high value"')])

        msg.set_param('charset', 'iso-9999-xx', header='X-Jimmy')

        eq(msg.get_param('charset', header='X-Jimmy'), 'iso-9999-xx')

    def test_del_param(self):

        eq = self.assertEqual

        msg = self._msgobj('msg_05.txt')

        eq(msg.get_params(),

           [('multipart/report', ''), ('report-type', 'delivery-status'),

            ('boundary', 'D1690A7AC1.996856090/mail.example.com')])

        old_val = msg.get_param("report-type")

        msg.del_param("report-type")

        eq(msg.get_params(),

           [('multipart/report', ''),

            ('boundary', 'D1690A7AC1.996856090/mail.example.com')])

        msg.set_param("report-type", old_val)

        eq(msg.get_params(),

           [('multipart/report', ''),

            ('boundary', 'D1690A7AC1.996856090/mail.example.com'),

            ('report-type', old_val)])

    def test_del_param_on_other_header(self):

        msg = Message()

        msg.add_header('Content-Disposition', 'attachment', filename='bud.gif')

        msg.del_param('filename', 'content-disposition')

        self.assertEqual(msg['content-disposition'], 'attachment')

    def test_set_type(self):

        eq = self.assertEqual

        msg = Message()

        self.assertRaises(ValueError, msg.set_type, 'text')

        msg.set_type('text/plain')

        eq(msg['content-type'], 'text/plain')

        msg.set_param('charset', 'us-ascii')

        eq(msg['content-type'], 'text/plain; charset="us-ascii"')

        msg.set_type('text/html')

        eq(msg['content-type'], 'text/html; charset="us-ascii"')

    def test_set_type_on_other_header(self):

        msg = Message()

        msg['X-Content-Type'] = 'text/plain'

        msg.set_type('application/octet-stream', 'X-Content-Type')

        self.assertEqual(msg['x-content-type'], 'application/octet-stream')

    def test_get_content_type_missing(self):

        msg = Message()

        self.assertEqual(msg.get_content_type(), 'text/plain')

    def test_get_content_type_missing_with_default_type(self):

        msg = Message()

        msg.set_default_type('message/rfc822')

        self.assertEqual(msg.get_content_type(), 'message/rfc822')

    def test_get_content_type_from_message_implicit(self):

        msg = self._msgobj('msg_30.txt')

        self.assertEqual(msg.get_payload(0).get_content_type(),

                         'message/rfc822')

    def test_get_content_type_from_message_explicit(self):

        msg = self._msgobj('msg_28.txt')

        self.assertEqual(msg.get_payload(0).get_content_type(),

                         'message/rfc822')

    def test_get_content_type_from_message_text_plain_implicit(self):

        msg = self._msgobj('msg_03.txt')

        self.assertEqual(msg.get_content_type(), 'text/plain')

    def test_get_content_type_from_message_text_plain_explicit(self):

        msg = self._msgobj('msg_01.txt')

        self.assertEqual(msg.get_content_type(), 'text/plain')

    def test_get_content_maintype_missing(self):

        msg = Message()

        self.assertEqual(msg.get_content_maintype(), 'text')

    def test_get_content_maintype_missing_with_default_type(self):

        msg = Message()

        msg.set_default_type('message/rfc822')

        self.assertEqual(msg.get_content_maintype(), 'message')

    def test_get_content_maintype_from_message_implicit(self):

        msg = self._msgobj('msg_30.txt')

        self.assertEqual(msg.get_payload(0).get_content_maintype(), 'message')

    def test_get_content_maintype_from_message_explicit(self):

        msg = self._msgobj('msg_28.txt')

        self.assertEqual(msg.get_payload(0).get_content_maintype(), 'message')

    def test_get_content_maintype_from_message_text_plain_implicit(self):

        msg = self._msgobj('msg_03.txt')

        self.assertEqual(msg.get_content_maintype(), 'text')

    def test_get_content_maintype_from_message_text_plain_explicit(self):

        msg = self._msgobj('msg_01.txt')

        self.assertEqual(msg.get_content_maintype(), 'text')

    def test_get_content_subtype_missing(self):

        msg = Message()

        self.assertEqual(msg.get_content_subtype(), 'plain')

    def test_get_content_subtype_missing_with_default_type(self):

        msg = Message()

        msg.set_default_type('message/rfc822')

        self.assertEqual(msg.get_content_subtype(), 'rfc822')

    def test_get_content_subtype_from_message_implicit(self):

        msg = self._msgobj('msg_30.txt')

        self.assertEqual(msg.get_payload(0).get_content_subtype(), 'rfc822')

    def test_get_content_subtype_from_message_explicit(self):

        msg = self._msgobj('msg_28.txt')

        self.assertEqual(msg.get_payload(0).get_content_subtype(), 'rfc822')

    def test_get_content_subtype_from_message_text_plain_implicit(self):

        msg = self._msgobj('msg_03.txt')

        self.assertEqual(msg.get_content_subtype(), 'plain')

    def test_get_content_subtype_from_message_text_plain_explicit(self):

        msg = self._msgobj('msg_01.txt')

        self.assertEqual(msg.get_content_subtype(), 'plain')

    def test_get_content_maintype_error(self):

        msg = Message()

        msg['Content-Type'] = 'no-slash-in-this-string'

        self.assertEqual(msg.get_content_maintype(), 'text')

    def test_get_content_subtype_error(self):

        msg = Message()

        msg['Content-Type'] = 'no-slash-in-this-string'

        self.assertEqual(msg.get_content_subtype(), 'plain')

    def test_replace_header(self):

        eq = self.assertEqual

        msg = Message()

        msg.add_header('First', 'One')

        msg.add_header('Second', 'Two')

        msg.add_header('Third', 'Three')

        eq(msg.keys(), ['First', 'Second', 'Third'])

        eq(msg.values(), ['One', 'Two', 'Three'])

        msg.replace_header('Second', 'Twenty')

        eq(msg.keys(), ['First', 'Second', 'Third'])

        eq(msg.values(), ['One', 'Twenty', 'Three'])

        msg.add_header('First', 'Eleven')

        msg.replace_header('First', 'One Hundred')

        eq(msg.keys(), ['First', 'Second', 'Third', 'First'])

        eq(msg.values(), ['One Hundred', 'Twenty', 'Three', 'Eleven'])

        self.assertRaises(KeyError, msg.replace_header, 'Fourth', 'Missing')

    def test_broken_base64_payload(self):

        x = 'AwDp0P7//y6LwKEAcPa/6Q=9'

        msg = Message()

        msg['content-type'] = 'audio/x-midi'

        msg['content-transfer-encoding'] = 'base64'

        msg.set_payload(x)

        self.assertEqual(msg.get_payload(decode=True), x)

# Test the email.encoders module

class TestEncoders(unittest.TestCase):

    def test_encode_empty_payload(self):

        eq = self.assertEqual

        msg = Message()

        msg.set_charset('us-ascii')

        eq(msg['content-transfer-encoding'], '7bit')

    def test_default_cte(self):

        eq = self.assertEqual

        msg = MIMEText('hello world')

        eq(msg['content-transfer-encoding'], '7bit')

    def test_default_cte(self):

        eq = self.assertEqual

        # With no explicit _charset its us-ascii, and all are 7-bit

        msg = MIMEText('hello world')

        eq(msg['content-transfer-encoding'], '7bit')

        # Similar, but with 8-bit data

        msg = MIMEText('hello \xf8 world')

        eq(msg['content-transfer-encoding'], '8bit')

        # And now with a different charset

        msg = MIMEText('hello \xf8 world', _charset='iso-8859-1')

        eq(msg['content-transfer-encoding'], 'quoted-printable')

# Test long header wrapping

class TestLongHeaders(TestEmailBase):

    def test_split_long_continuation(self):

        eq = self.ndiffAssertEqual

        msg = email.message_from_string("""\

Subject: bug demonstration

\t12345678911234567892123456789312345678941234567895123456789612345678971234567898112345678911234567892123456789112345678911234567892123456789

\tmore text