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

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

2 

3import smtplib 

4import ssl 

5 

6from email.headerregistry import Address 

7from email.message import EmailMessage 

8 

9from .mailer import Mailer, MailerException 

10from .message import Message 

11 

12 

13class SmtpMailer(Mailer): 

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

15 

16 def __init__( 

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

18 ): 

19 self._host = host 

20 if port is None: 

21 self._port = 456 if ssl_ else 587 

22 else: 

23 self._port = port 

24 self._ssl = ssl_ 

25 self._server = None 

26 self._tls = tls 

27 

28 def __del__(self): 

29 """Disconnect from the mailer.""" 

30 self.disconnect() 

31 

32 def connect(self): 

33 """Connect to the SMTP server.""" 

34 context = ssl.create_default_context() 

35 if self._ssl: 

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

37 else: 

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

39 

40 if self._tls: 

41 try: 

42 self._server.starttls(context=context) 

43 except smtplib.SMTPException as exc: 

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

45 

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

47 """Login to the SMTP server.""" 

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

49 

50 try: 

51 self._server.login(user, password) 

52 except smtplib.SMTPException as exc: 

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

54 

55 def disconnect(self): 

56 """Disconnect from the SMTP server.""" 

57 if self._server: 

58 try: 

59 self._server.quit() 

60 except smtplib.SMTPException: 

61 pass # Ignore exceptions here 

62 self._server = None 

63 

64 def send(self, message: Message): 

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

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

67 

68 email_message = EmailMessage() 

69 

70 email_message["Subject"] = message.subject 

71 email_message.set_content(message.text) 

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

73 email_message.add_header(header_name, header_value) 

74 

75 email_message["From"] = Address( 

76 display_name=message.recipients.from_.name, 

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

78 ) 

79 

80 email_message["To"] = [ 

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

82 for to in message.recipients.to 

83 ] 

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

85 email_message["Cc"] = [ 

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

87 for cc in message.recipients.cc 

88 ] 

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

90 email_message["Bcc"] = [ 

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

92 for bcc in message.recipients.bcc 

93 ] 

94 

95 if len(message.html) > 0: 

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

97 

98 try: 

99 self._server.send_message(email_message) 

100 except smtplib.SMTPException as exc: 

101 raise MailerException("Sendmail failed") from exc