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 }