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