Coverage for kwai/core/mail/smtp_mailer.py: 78%

1"""Module that defines a mailer service with SMTP.""" 

2import smtplib 

3import ssl 

4from email.headerregistry import Address 

5from email.message import EmailMessage 

6 

7from .mailer import Mailer, MailerException 

8from .message import Message 

9 

10 

11class SmtpMailer(Mailer): 

12 """A mailer service that uses an SMTP server.""" 

13 

14 def __init__( 

15 self, host: str, port: int | None = None, ssl_: bool = False, tls: bool = True 

16 ): 

17 self._host = host 

18 if port is None: 

19 self._port = 456 if ssl_ else 587 

20 else: 

21 self._port = port 

22 self._ssl = ssl_ 

23 self._server = None 

24 self._tls = tls 

25 

26 def __del__(self): 

27 """Disconnect from the mailer.""" 

28 self.disconnect() 

29 

30 def connect(self): 

31 """Connect to the SMTP server.""" 

32 context = ssl.create_default_context() 

33 if self._ssl: 

34 self._server = smtplib.SMTP_SSL(self._host, self._port, context=context) 

35 else: 

36 self._server = smtplib.SMTP(self._host, self._port) 

37 

38 if self._tls: 

39 try: 

40 self._server.starttls(context=context) 

41 except smtplib.SMTPException as exc: 

42 raise MailerException("STARTTLS failed.") from exc 

43 

44 def login(self, user: str, password: str): 

45 """Login to the SMTP server.""" 

46 assert self._server is not None, "There is no connection to an SMTP server!" 

47 

48 try: 

49 self._server.login(user, password) 

50 except smtplib.SMTPException as exc: 

51 raise MailerException("Login failed.") from exc 

52 

53 def disconnect(self): 

54 """Disconnect from the SMTP server.""" 

55 if self._server: 

56 try: 

57 self._server.quit() 

58 except smtplib.SMTPException: 

59 pass # Ignore exceptions here 

60 self._server = None 

61 

62 def send(self, message: Message): 

63 """Send a message to the SMTP server.""" 

64 assert self._server is not None, "There is no connection to an SMTP server!" 

65 

66 email_message = EmailMessage() 

67 

68 email_message["Subject"] = message.subject 

69 email_message.set_content(message.text) 

70 for header_name, header_value in message.headers.items(): 

71 email_message.add_header(header_name, header_value) 

72 

73 email_message["From"] = Address( 

74 display_name=message.recipients.from_.name, 

75 addr_spec=str(message.recipients.from_.email), 

76 ) 

77 

78 email_message["To"] = [ 

79 Address(display_name=to.name, addr_spec=str(to.email)) 

80 for to in message.recipients.to 

81 ] 

82 if len(message.recipients.cc) > 0: 

83 email_message["Cc"] = [ 

84 Address(display_name=cc.name, addr_spec=str(cc.email)) 

85 for cc in message.recipients.cc 

86 ] 

87 if len(message.recipients.bcc) > 0: 

88 email_message["Bcc"] = [ 

89 Address(display_name=bcc.name, addr_spec=str(bcc.email)) 

90 for bcc in message.recipients.bcc 

91 ] 

92 

93 if len(message.html) > 0: 

94 email_message.add_alternative(message.html, subtype="html") 

95 

96 try: 

97 self._server.send_message(email_message) 

98 except smtplib.SMTPException as exc: 

99 raise MailerException("Sendmail failed") from exc