1 import io
2 import textwrap
3 import unittest
4 from email import message_from_string, message_from_bytes
5 from email.message import EmailMessage
6 from email.generator import Generator, BytesGenerator
7 from email.headerregistry import Address
8 from email import policy
9 from test.test_email import TestEmailBase, parameterize
10
11
12 @parameterize
13 class ESC[4;38;5;81mTestGeneratorBase:
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
220 class ESC[4;38;5;81mTestGenerator(ESC[4;38;5;149mTestGeneratorBase, ESC[4;38;5;149mTestEmailBase):
221
222 msgfunc = staticmethod(message_from_string)
223 genclass = Generator
224 ioclass = io.StringIO
225 typ = str
226
227
228 class ESC[4;38;5;81mTestBytesGenerator(ESC[4;38;5;149mTestGeneratorBase, ESC[4;38;5;149mTestEmailBase):
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
317 if __name__ == '__main__':
318 unittest.main()