github.com/merlinepedra/gopphish-attack@v0.9.0/mailer/mockmailer.go (about)

     1  package mailer
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"io"
     7  	"time"
     8  
     9  	"github.com/gophish/gomail"
    10  )
    11  
    12  // errHostUnreachable is a mock error to represent a host
    13  // being unreachable
    14  var errHostUnreachable = errors.New("host unreachable")
    15  
    16  // errDialerUnavailable is a mock error to represent a dialer
    17  // being unavailable (perhaps an error getting the dialer config
    18  // or a database error)
    19  var errDialerUnavailable = errors.New("dialer unavailable")
    20  
    21  // mockDialer keeps track of calls to Dial
    22  type mockDialer struct {
    23  	dialCount int
    24  	dial      func() (Sender, error)
    25  }
    26  
    27  // newMockDialer returns a new instance of the mockDialer with the default
    28  // dialer set.
    29  func newMockDialer() *mockDialer {
    30  	md := &mockDialer{}
    31  	md.dial = md.defaultDial
    32  	return md
    33  }
    34  
    35  // defaultDial simply returns a mockSender
    36  func (md *mockDialer) defaultDial() (Sender, error) {
    37  	return newMockSender(), nil
    38  }
    39  
    40  // unreachableDial is to simulate network error conditions in which
    41  // a host is unavailable.
    42  func (md *mockDialer) unreachableDial() (Sender, error) {
    43  	return nil, errHostUnreachable
    44  }
    45  
    46  // Dial increments the internal dial count. Otherwise, it's a no-op for the mock client.
    47  func (md *mockDialer) Dial() (Sender, error) {
    48  	md.dialCount++
    49  	return md.dial()
    50  }
    51  
    52  // setDial sets the Dial function for the mockDialer
    53  func (md *mockDialer) setDial(dial func() (Sender, error)) {
    54  	md.dial = dial
    55  }
    56  
    57  // mockSender is a mock gomail.Sender used for testing.
    58  type mockSender struct {
    59  	messages    []*mockMessage
    60  	status      string
    61  	send        func(*mockMessage) error
    62  	messageChan chan *mockMessage
    63  	resetCount  int
    64  }
    65  
    66  func newMockSender() *mockSender {
    67  	ms := &mockSender{
    68  		status:      "ehlo",
    69  		messageChan: make(chan *mockMessage),
    70  	}
    71  	ms.send = ms.defaultSend
    72  	return ms
    73  }
    74  
    75  func (ms *mockSender) setSend(send func(*mockMessage) error) {
    76  	ms.send = send
    77  }
    78  
    79  func (ms *mockSender) defaultSend(mm *mockMessage) error {
    80  	ms.messageChan <- mm
    81  	return nil
    82  }
    83  
    84  // Send just appends the provided message record to the internal slice
    85  func (ms *mockSender) Send(from string, to []string, msg io.WriterTo) error {
    86  	mm := newMockMessage(from, to, msg)
    87  	ms.messages = append(ms.messages, mm)
    88  	ms.status = "sent"
    89  	return ms.send(mm)
    90  }
    91  
    92  // Close is a noop for the mock client
    93  func (ms *mockSender) Close() error {
    94  	ms.status = "closed"
    95  	close(ms.messageChan)
    96  	return nil
    97  }
    98  
    99  // Reset sets the status to "Reset". In practice, this would reset the connection
   100  // to the same state as if the client had just sent an EHLO command.
   101  func (ms *mockSender) Reset() error {
   102  	ms.status = "reset"
   103  	ms.resetCount++
   104  	return nil
   105  }
   106  
   107  // mockMessage holds the information sent via a call to MockClient.Send()
   108  type mockMessage struct {
   109  	from         string
   110  	to           []string
   111  	message      []byte
   112  	sendAt       time.Time
   113  	backoffCount int
   114  	getdialer    func() (Dialer, error)
   115  	err          error
   116  	finished     bool
   117  }
   118  
   119  func newMockMessage(from string, to []string, msg io.WriterTo) *mockMessage {
   120  	buff := &bytes.Buffer{}
   121  	msg.WriteTo(buff)
   122  	mm := &mockMessage{
   123  		from:    from,
   124  		to:      to,
   125  		message: buff.Bytes(),
   126  		sendAt:  time.Now(),
   127  	}
   128  	mm.getdialer = mm.defaultDialer
   129  	return mm
   130  }
   131  
   132  func (mm *mockMessage) setDialer(dialer func() (Dialer, error)) {
   133  	mm.getdialer = dialer
   134  }
   135  
   136  func (mm *mockMessage) defaultDialer() (Dialer, error) {
   137  	return newMockDialer(), nil
   138  }
   139  
   140  func (mm *mockMessage) errorDialer() (Dialer, error) {
   141  	return nil, errDialerUnavailable
   142  }
   143  
   144  func (mm *mockMessage) GetDialer() (Dialer, error) {
   145  	return mm.getdialer()
   146  }
   147  
   148  func (mm *mockMessage) Backoff(reason error) error {
   149  	mm.backoffCount++
   150  	mm.err = reason
   151  	return nil
   152  }
   153  
   154  func (mm *mockMessage) Error(err error) error {
   155  	mm.err = err
   156  	mm.finished = true
   157  	return nil
   158  }
   159  
   160  func (mm *mockMessage) Finish() error {
   161  	mm.finished = true
   162  	return nil
   163  }
   164  
   165  func (mm *mockMessage) Generate(message *gomail.Message) error {
   166  	message.SetHeaders(map[string][]string{
   167  		"From": {mm.from},
   168  		"To":   mm.to,
   169  	})
   170  	message.SetBody("text/html", string(mm.message))
   171  	return nil
   172  }
   173  
   174  func (mm *mockMessage) Success() error {
   175  	mm.finished = true
   176  	return nil
   177  }