1import io
2import textwrap
3import unittest
4from email import message_from_string, message_from_bytes
5from email.message import EmailMessage
6from email.generator import Generator, BytesGenerator
7from email.headerregistry import Address
8from email import policy
9from test.test_email import TestEmailBase, parameterize
10
11
12@parameterize
13class TestGeneratorBase:
14
15    policy = policy.default
16
17    def msgmaker(self, msg, policy=None):
18        policy = self.policy if policy is None else policy
19        return self.msgfunc(msg, policy=policy)
20
21    refold_long_expected = {
22        0: textwrap.dedent("""\
23            To: whom_it_may_concern@example.com
24            From: nobody_you_want_to_know@example.com
25            Subject: We the willing led by the unknowing are doing the
26             impossible for the ungrateful. We have done so much for so long with so little
27             we are now qualified to do anything with nothing.
28
29            None
30            """),
31        40: textwrap.dedent("""\
32            To: whom_it_may_concern@example.com
33            From:
34             nobody_you_want_to_know@example.com
35            Subject: We the willing led by the
36             unknowing are doing the impossible for
37             the ungrateful. We have done so much
38             for so long with so little we are now
39             qualified to do anything with nothing.
40
41            None
42            """),
43        20: textwrap.dedent("""\
44            To:
45             whom_it_may_concern@example.com
46            From:
47             nobody_you_want_to_know@example.com
48            Subject: We the
49             willing led by the
50             unknowing are doing
51             the impossible for
52             the ungrateful. We
53             have done so much
54             for so long with so
55             little we are now
56             qualified to do
57             anything with
58             nothing.
59
60            None
61            """),
62        }
63    refold_long_expected[100] = refold_long_expected[0]
64
65    refold_all_expected = refold_long_expected.copy()
66    refold_all_expected[0] = (
67            "To: whom_it_may_concern@example.com\n"
68            "From: nobody_you_want_to_know@example.com\n"
69            "Subject: We the willing led by the unknowing are doing the "
70              "impossible for the ungrateful. We have done so much for "
71              "so long with so little we are now qualified to do anything "
72              "with nothing.\n"
73              "\n"
74              "None\n")
75    refold_all_expected[100] = (
76            "To: whom_it_may_concern@example.com\n"
77            "From: nobody_you_want_to_know@example.com\n"
78            "Subject: We the willing led by the unknowing are doing the "
79                "impossible for the ungrateful. We have\n"
80              " done so much for so long with so little we are now qualified "
81                "to do anything with nothing.\n"
82              "\n"
83              "None\n")
84
85    length_params = [n for n in refold_long_expected]
86
87    def length_as_maxheaderlen_parameter(self, n):
88        msg = self.msgmaker(self.typ(self.refold_long_expected[0]))
89        s = self.ioclass()
90        g = self.genclass(s, maxheaderlen=n, policy=self.policy)
91        g.flatten(msg)
92        self.assertEqual(s.getvalue(), self.typ(self.refold_long_expected[n]))
93
94    def length_as_max_line_length_policy(self, n):
95        msg = self.msgmaker(self.typ(self.refold_long_expected[0]))
96        s = self.ioclass()
97        g = self.genclass(s, policy=self.policy.clone(max_line_length=n))
98        g.flatten(msg)
99        self.assertEqual(s.getvalue(), self.typ(self.refold_long_expected[n]))
100
101    def length_as_maxheaderlen_parm_overrides_policy(self, n):
102        msg = self.msgmaker(self.typ(self.refold_long_expected[0]))
103        s = self.ioclass()
104        g = self.genclass(s, maxheaderlen=n,
105                          policy=self.policy.clone(max_line_length=10))
106        g.flatten(msg)
107        self.assertEqual(s.getvalue(), self.typ(self.refold_long_expected[n]))
108
109    def length_as_max_line_length_with_refold_none_does_not_fold(self, n):
110        msg = self.msgmaker(self.typ(self.refold_long_expected[0]))
111        s = self.ioclass()
112        g = self.genclass(s, policy=self.policy.clone(refold_source='none',
113                                                      max_line_length=n))
114        g.flatten(msg)
115        self.assertEqual(s.getvalue(), self.typ(self.refold_long_expected[0]))
116
117    def length_as_max_line_length_with_refold_all_folds(self, n):
118        msg = self.msgmaker(self.typ(self.refold_long_expected[0]))
119        s = self.ioclass()
120        g = self.genclass(s, policy=self.policy.clone(refold_source='all',
121                                                      max_line_length=n))
122        g.flatten(msg)
123        self.assertEqual(s.getvalue(), self.typ(self.refold_all_expected[n]))
124
125    def test_crlf_control_via_policy(self):
126        source = "Subject: test\r\n\r\ntest body\r\n"
127        expected = source
128        msg = self.msgmaker(self.typ(source))
129        s = self.ioclass()
130        g = self.genclass(s, policy=policy.SMTP)
131        g.flatten(msg)
132        self.assertEqual(s.getvalue(), self.typ(expected))
133
134    def test_flatten_linesep_overrides_policy(self):
135        source = "Subject: test\n\ntest body\n"
136        expected = source
137        msg = self.msgmaker(self.typ(source))
138        s = self.ioclass()
139        g = self.genclass(s, policy=policy.SMTP)
140        g.flatten(msg, linesep='\n')
141        self.assertEqual(s.getvalue(), self.typ(expected))
142
143    def test_set_mangle_from_via_policy(self):
144        source = textwrap.dedent("""\
145            Subject: test that
146             from is mangled in the body!
147
148            From time to time I write a rhyme.
149            """)
150        variants = (
151            (None, True),
152            (policy.compat32, True),
153            (policy.default, False),
154            (policy.default.clone(mangle_from_=True), True),
155            )
156        for p, mangle in variants:
157            expected = source.replace('From ', '>From ') if mangle else source
158            with self.subTest(policy=p, mangle_from_=mangle):
159                msg = self.msgmaker(self.typ(source))
160                s = self.ioclass()
161                g = self.genclass(s, policy=p)
162                g.flatten(msg)
163                self.assertEqual(s.getvalue(), self.typ(expected))
164
165    def test_compat32_max_line_length_does_not_fold_when_none(self):
166        msg = self.msgmaker(self.typ(self.refold_long_expected[0]))
167        s = self.ioclass()
168        g = self.genclass(s, policy=policy.compat32.clone(max_line_length=None))
169        g.flatten(msg)
170        self.assertEqual(s.getvalue(), self.typ(self.refold_long_expected[0]))
171
172    def test_rfc2231_wrapping(self):
173        # This is pretty much just to make sure we don't have an infinite
174        # loop; I don't expect anyone to hit this in the field.
175        msg = self.msgmaker(self.typ(textwrap.dedent("""\
176            To: nobody
177            Content-Disposition: attachment;
178             filename="afilenamelongenoghtowraphere"
179
180            None
181            """)))
182        expected = textwrap.dedent("""\
183            To: nobody
184            Content-Disposition: attachment;
185             filename*0*=us-ascii''afilename;
186             filename*1*=longenoghtowraphere
187
188            None
189            """)
190        s = self.ioclass()
191        g = self.genclass(s, policy=self.policy.clone(max_line_length=33))
192        g.flatten(msg)
193        self.assertEqual(s.getvalue(), self.typ(expected))
194
195    def test_rfc2231_wrapping_switches_to_default_len_if_too_narrow(self):
196        # This is just to make sure we don't have an infinite loop; I don't
197        # expect anyone to hit this in the field, so I'm not bothering to make
198        # the result optimal (the encoding isn't needed).
199        msg = self.msgmaker(self.typ(textwrap.dedent("""\
200            To: nobody
201            Content-Disposition: attachment;
202             filename="afilenamelongenoghtowraphere"
203
204            None
205            """)))
206        expected = textwrap.dedent("""\
207            To: nobody
208            Content-Disposition:
209             attachment;
210             filename*0*=us-ascii''afilenamelongenoghtowraphere
211
212            None
213            """)
214        s = self.ioclass()
215        g = self.genclass(s, policy=self.policy.clone(max_line_length=20))
216        g.flatten(msg)
217        self.assertEqual(s.getvalue(), self.typ(expected))
218
219
220class TestGenerator(TestGeneratorBase, TestEmailBase):
221
222    msgfunc = staticmethod(message_from_string)
223    genclass = Generator
224    ioclass = io.StringIO
225    typ = str
226
227
228class TestBytesGenerator(TestGeneratorBase, TestEmailBase):
229
230    msgfunc = staticmethod(message_from_bytes)
231    genclass = BytesGenerator
232    ioclass = io.BytesIO
233    typ = lambda self, x: x.encode('ascii')
234
235    def test_cte_type_7bit_handles_unknown_8bit(self):
236        source = ("Subject: Maintenant je vous présente mon "
237                 "collègue\n\n").encode('utf-8')
238        expected = ('Subject: Maintenant je vous =?unknown-8bit?q?'
239                    'pr=C3=A9sente_mon_coll=C3=A8gue?=\n\n').encode('ascii')
240        msg = message_from_bytes(source)
241        s = io.BytesIO()
242        g = BytesGenerator(s, policy=self.policy.clone(cte_type='7bit'))
243        g.flatten(msg)
244        self.assertEqual(s.getvalue(), expected)
245
246    def test_cte_type_7bit_transforms_8bit_cte(self):
247        source = textwrap.dedent("""\
248            From: foo@bar.com
249            To: Dinsdale
250            Subject: Nudge nudge, wink, wink
251            Mime-Version: 1.0
252            Content-Type: text/plain; charset="latin-1"
253            Content-Transfer-Encoding: 8bit
254
255            oh là là, know what I mean, know what I mean?
256            """).encode('latin1')
257        msg = message_from_bytes(source)
258        expected =  textwrap.dedent("""\
259            From: foo@bar.com
260            To: Dinsdale
261            Subject: Nudge nudge, wink, wink
262            Mime-Version: 1.0
263            Content-Type: text/plain; charset="iso-8859-1"
264            Content-Transfer-Encoding: quoted-printable
265
266            oh l=E0 l=E0, know what I mean, know what I mean?
267            """).encode('ascii')
268        s = io.BytesIO()
269        g = BytesGenerator(s, policy=self.policy.clone(cte_type='7bit',
270                                                       linesep='\n'))
271        g.flatten(msg)
272        self.assertEqual(s.getvalue(), expected)
273
274    def test_smtputf8_policy(self):
275        msg = EmailMessage()
276        msg['From'] = "Páolo <főo@bar.com>"
277        msg['To'] = 'Dinsdale'
278        msg['Subject'] = 'Nudge nudge, wink, wink \u1F609'
279        msg.set_content("oh là là, know what I mean, know what I mean?")
280        expected = textwrap.dedent("""\
281            From: Páolo <főo@bar.com>
282            To: Dinsdale
283            Subject: Nudge nudge, wink, wink \u1F609
284            Content-Type: text/plain; charset="utf-8"
285            Content-Transfer-Encoding: 8bit
286            MIME-Version: 1.0
287
288            oh là là, know what I mean, know what I mean?
289            """).encode('utf-8').replace(b'\n', b'\r\n')
290        s = io.BytesIO()
291        g = BytesGenerator(s, policy=policy.SMTPUTF8)
292        g.flatten(msg)
293        self.assertEqual(s.getvalue(), expected)
294
295    def test_smtp_policy(self):
296        msg = EmailMessage()
297        msg["From"] = Address(addr_spec="foo@bar.com", display_name="Páolo")
298        msg["To"] = Address(addr_spec="bar@foo.com", display_name="Dinsdale")
299        msg["Subject"] = "Nudge nudge, wink, wink"
300        msg.set_content("oh boy, know what I mean, know what I mean?")
301        expected = textwrap.dedent("""\
302            From: =?utf-8?q?P=C3=A1olo?= <foo@bar.com>
303            To: Dinsdale <bar@foo.com>
304            Subject: Nudge nudge, wink, wink
305            Content-Type: text/plain; charset="utf-8"
306            Content-Transfer-Encoding: 7bit
307            MIME-Version: 1.0
308
309            oh boy, know what I mean, know what I mean?
310            """).encode().replace(b"\n", b"\r\n")
311        s = io.BytesIO()
312        g = BytesGenerator(s, policy=policy.SMTP)
313        g.flatten(msg)
314        self.assertEqual(s.getvalue(), expected)
315
316
317if __name__ == '__main__':
318    unittest.main()
319