🎣 Open-Source Phishing Toolkit
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

291 lines
7.8 KiB

  1. package mailer
  2. import (
  3. "bytes"
  4. "context"
  5. "errors"
  6. "io"
  7. "net/textproto"
  8. "reflect"
  9. "testing"
  10. "github.com/stretchr/testify/suite"
  11. )
  12. type MailerSuite struct {
  13. suite.Suite
  14. }
  15. func generateMessages(dialer Dialer) []Mail {
  16. to := []string{"to@example.com"}
  17. messageContents := []io.WriterTo{
  18. bytes.NewBuffer([]byte("First email")),
  19. bytes.NewBuffer([]byte("Second email")),
  20. }
  21. m1 := newMockMessage("first@example.com", to, messageContents[0])
  22. m2 := newMockMessage("second@example.com", to, messageContents[1])
  23. m1.setDialer(func() (Dialer, error) { return dialer, nil })
  24. messages := []Mail{m1, m2}
  25. return messages
  26. }
  27. func newMockErrorSender(err error) *mockSender {
  28. sender := newMockSender()
  29. // The sending function will send a temporary error to emulate
  30. // a backoff.
  31. sender.setSend(func(mm *mockMessage) error {
  32. if len(sender.messages) == 1 {
  33. return err
  34. }
  35. sender.messageChan <- mm
  36. return nil
  37. })
  38. return sender
  39. }
  40. func (ms *MailerSuite) TestDialHost() {
  41. ctx, cancel := context.WithCancel(context.Background())
  42. defer cancel()
  43. md := newMockDialer()
  44. md.setDial(md.unreachableDial)
  45. _, err := dialHost(ctx, md)
  46. if _, ok := err.(*ErrMaxConnectAttempts); !ok {
  47. ms.T().Fatalf("Didn't receive expected ErrMaxConnectAttempts. Got: %s", err)
  48. }
  49. e := err.(*ErrMaxConnectAttempts)
  50. if e.underlyingError != errHostUnreachable {
  51. ms.T().Fatalf("Got invalid underlying error. Expected %s Got %s\n", e.underlyingError, errHostUnreachable)
  52. }
  53. if md.dialCount != MaxReconnectAttempts {
  54. ms.T().Fatalf("Unexpected number of reconnect attempts. Expected %d, Got %d", MaxReconnectAttempts, md.dialCount)
  55. }
  56. md.setDial(md.defaultDial)
  57. _, err = dialHost(ctx, md)
  58. if err != nil {
  59. ms.T().Fatalf("Unexpected error when dialing the mock host: %s", err)
  60. }
  61. }
  62. func (ms *MailerSuite) TestMailWorkerStart() {
  63. ctx, cancel := context.WithCancel(context.Background())
  64. defer cancel()
  65. mw := NewMailWorker()
  66. go func(ctx context.Context) {
  67. mw.Start(ctx)
  68. }(ctx)
  69. sender := newMockSender()
  70. dialer := newMockDialer()
  71. dialer.setDial(func() (Sender, error) {
  72. return sender, nil
  73. })
  74. messages := generateMessages(dialer)
  75. // Send the campaign
  76. mw.Queue(messages)
  77. got := []*mockMessage{}
  78. idx := 0
  79. for message := range sender.messageChan {
  80. got = append(got, message)
  81. original := messages[idx].(*mockMessage)
  82. if original.from != message.from {
  83. ms.T().Fatalf("Invalid message received. Expected %s, Got %s", original.from, message.from)
  84. }
  85. idx++
  86. }
  87. if len(got) != len(messages) {
  88. ms.T().Fatalf("Unexpected number of messages received. Expected %d Got %d", len(got), len(messages))
  89. }
  90. }
  91. func (ms *MailerSuite) TestBackoff() {
  92. ctx, cancel := context.WithCancel(context.Background())
  93. defer cancel()
  94. mw := NewMailWorker()
  95. go func(ctx context.Context) {
  96. mw.Start(ctx)
  97. }(ctx)
  98. expectedError := &textproto.Error{
  99. Code: 400,
  100. Msg: "Temporary error",
  101. }
  102. sender := newMockErrorSender(expectedError)
  103. dialer := newMockDialer()
  104. dialer.setDial(func() (Sender, error) {
  105. return sender, nil
  106. })
  107. messages := generateMessages(dialer)
  108. // Send the campaign
  109. mw.Queue(messages)
  110. got := []*mockMessage{}
  111. for message := range sender.messageChan {
  112. got = append(got, message)
  113. }
  114. // Check that we only sent one message
  115. expectedCount := 1
  116. if len(got) != expectedCount {
  117. ms.T().Fatalf("Unexpected number of messages received. Expected %d Got %d", len(got), expectedCount)
  118. }
  119. // Check that it's the correct message
  120. originalFrom := messages[1].(*mockMessage).from
  121. if got[0].from != originalFrom {
  122. ms.T().Fatalf("Invalid message received. Expected %s, Got %s", originalFrom, got[0].from)
  123. }
  124. // Check that the first message performed a backoff
  125. backoffCount := messages[0].(*mockMessage).backoffCount
  126. if backoffCount != expectedCount {
  127. ms.T().Fatalf("Did not receive expected backoff. Got backoffCount %d, Expected %d", backoffCount, expectedCount)
  128. }
  129. // Check that there was a reset performed on the sender
  130. if sender.resetCount != expectedCount {
  131. ms.T().Fatalf("Did not receive expected reset. Got resetCount %d, expected %d", sender.resetCount, expectedCount)
  132. }
  133. }
  134. func (ms *MailerSuite) TestPermError() {
  135. ctx, cancel := context.WithCancel(context.Background())
  136. defer cancel()
  137. mw := NewMailWorker()
  138. go func(ctx context.Context) {
  139. mw.Start(ctx)
  140. }(ctx)
  141. expectedError := &textproto.Error{
  142. Code: 500,
  143. Msg: "Permanent error",
  144. }
  145. sender := newMockErrorSender(expectedError)
  146. dialer := newMockDialer()
  147. dialer.setDial(func() (Sender, error) {
  148. return sender, nil
  149. })
  150. messages := generateMessages(dialer)
  151. // Send the campaign
  152. mw.Queue(messages)
  153. got := []*mockMessage{}
  154. for message := range sender.messageChan {
  155. got = append(got, message)
  156. }
  157. // Check that we only sent one message
  158. expectedCount := 1
  159. if len(got) != expectedCount {
  160. ms.T().Fatalf("Unexpected number of messages received. Expected %d Got %d", len(got), expectedCount)
  161. }
  162. // Check that it's the correct message
  163. originalFrom := messages[1].(*mockMessage).from
  164. if got[0].from != originalFrom {
  165. ms.T().Fatalf("Invalid message received. Expected %s, Got %s", originalFrom, got[0].from)
  166. }
  167. message := messages[0].(*mockMessage)
  168. // Check that the first message did not perform a backoff
  169. expectedBackoffCount := 0
  170. backoffCount := message.backoffCount
  171. if backoffCount != expectedBackoffCount {
  172. ms.T().Fatalf("Did not receive expected backoff. Got backoffCount %d, Expected %d", backoffCount, expectedCount)
  173. }
  174. // Check that there was a reset performed on the sender
  175. if sender.resetCount != expectedCount {
  176. ms.T().Fatalf("Did not receive expected reset. Got resetCount %d, expected %d", sender.resetCount, expectedCount)
  177. }
  178. // Check that the email errored out appropriately
  179. if !reflect.DeepEqual(message.err, expectedError) {
  180. ms.T().Fatalf("Did not received expected error. Got %#v\nExpected %#v", message.err, expectedError)
  181. }
  182. }
  183. func (ms *MailerSuite) TestUnknownError() {
  184. ctx, cancel := context.WithCancel(context.Background())
  185. defer cancel()
  186. mw := NewMailWorker()
  187. go func(ctx context.Context) {
  188. mw.Start(ctx)
  189. }(ctx)
  190. expectedError := errors.New("Unexpected error")
  191. sender := newMockErrorSender(expectedError)
  192. dialer := newMockDialer()
  193. dialer.setDial(func() (Sender, error) {
  194. return sender, nil
  195. })
  196. messages := generateMessages(dialer)
  197. // Send the campaign
  198. mw.Queue(messages)
  199. got := []*mockMessage{}
  200. for message := range sender.messageChan {
  201. got = append(got, message)
  202. }
  203. // Check that we only sent one message
  204. expectedCount := 1
  205. if len(got) != expectedCount {
  206. ms.T().Fatalf("Unexpected number of messages received. Expected %d Got %d", len(got), expectedCount)
  207. }
  208. // Check that it's the correct message
  209. originalFrom := messages[1].(*mockMessage).from
  210. if got[0].from != originalFrom {
  211. ms.T().Fatalf("Invalid message received. Expected %s, Got %s", originalFrom, got[0].from)
  212. }
  213. message := messages[0].(*mockMessage)
  214. // If we get an unexpected error, this means that it's likely the
  215. // underlying connection dropped. When this happens, we expect the
  216. // connection to be re-established (see #997).
  217. // In this case, we're successfully reestablishing the connection
  218. // so we expect the backoff to occur.
  219. expectedBackoffCount := 1
  220. backoffCount := message.backoffCount
  221. if backoffCount != expectedBackoffCount {
  222. ms.T().Fatalf("Did not receive expected backoff. Got backoffCount %d, Expected %d", backoffCount, expectedBackoffCount)
  223. }
  224. // Check that the underlying connection was reestablished
  225. expectedDialCount := 2
  226. if dialer.dialCount != expectedDialCount {
  227. ms.T().Fatalf("Did not receive expected dial count. Got %d expected %d", dialer.dialCount, expectedDialCount)
  228. }
  229. // Check that the email errored out appropriately
  230. if !reflect.DeepEqual(message.err, expectedError) {
  231. ms.T().Fatalf("Did not received expected error. Got %#v\nExpected %#v", message.err, expectedError)
  232. }
  233. }
  234. func TestMailerSuite(t *testing.T) {
  235. suite.Run(t, new(MailerSuite))
  236. }