github.com/edermi/gophish_mods@v0.7.0/models/template.go (about) 1 package models 2 3 import ( 4 "errors" 5 "time" 6 7 log "github.com/gophish/gophish/logger" 8 "github.com/jinzhu/gorm" 9 ) 10 11 // Template models hold the attributes for an email template to be sent to targets 12 type Template struct { 13 Id int64 `json:"id" gorm:"column:id; primary_key:yes"` 14 UserId int64 `json:"-" gorm:"column:user_id"` 15 Name string `json:"name"` 16 Subject string `json:"subject"` 17 Text string `json:"text"` 18 HTML string `json:"html" gorm:"column:html"` 19 ModifiedDate time.Time `json:"modified_date"` 20 Attachments []Attachment `json:"attachments"` 21 } 22 23 // ErrTemplateNameNotSpecified is thrown when a template name is not specified 24 var ErrTemplateNameNotSpecified = errors.New("Template name not specified") 25 26 // ErrTemplateMissingParameter is thrown when a needed parameter is not provided 27 var ErrTemplateMissingParameter = errors.New("Need to specify at least plaintext or HTML content") 28 29 // Validate checks the given template to make sure values are appropriate and complete 30 func (t *Template) Validate() error { 31 switch { 32 case t.Name == "": 33 return ErrTemplateNameNotSpecified 34 case t.Text == "" && t.HTML == "": 35 return ErrTemplateMissingParameter 36 } 37 // Test that the variables used in the template 38 // validate with no issues 39 td := struct { 40 Result 41 URL string 42 TrackingURL string 43 Tracker string 44 From string 45 }{ 46 Result{ 47 BaseRecipient: BaseRecipient{ 48 Email: "foo@bar.com", 49 FirstName: "Foo", 50 LastName: "Bar", 51 Position: "Test", 52 }, 53 }, 54 "http://foo.bar", 55 "http://foo.bar/track", 56 "<img src='http://foo.bar/track", 57 "John Doe <foo@bar.com>", 58 } 59 _, err := ExecuteTemplate(t.HTML, td) 60 if err != nil { 61 return err 62 } 63 _, err = ExecuteTemplate(t.Text, td) 64 if err != nil { 65 return err 66 } 67 return nil 68 } 69 70 // GetTemplates returns the templates owned by the given user. 71 func GetTemplates(uid int64) ([]Template, error) { 72 ts := []Template{} 73 err := db.Where("user_id=?", uid).Find(&ts).Error 74 if err != nil { 75 log.Error(err) 76 return ts, err 77 } 78 for i, _ := range ts { 79 // Get Attachments 80 err = db.Where("template_id=?", ts[i].Id).Find(&ts[i].Attachments).Error 81 if err == nil && len(ts[i].Attachments) == 0 { 82 ts[i].Attachments = make([]Attachment, 0) 83 } 84 if err != nil && err != gorm.ErrRecordNotFound { 85 log.Error(err) 86 return ts, err 87 } 88 } 89 return ts, err 90 } 91 92 // GetTemplate returns the template, if it exists, specified by the given id and user_id. 93 func GetTemplate(id int64, uid int64) (Template, error) { 94 t := Template{} 95 err := db.Where("user_id=? and id=?", uid, id).Find(&t).Error 96 if err != nil { 97 log.Error(err) 98 return t, err 99 } 100 101 // Get Attachments 102 err = db.Where("template_id=?", t.Id).Find(&t.Attachments).Error 103 if err != nil && err != gorm.ErrRecordNotFound { 104 log.Error(err) 105 return t, err 106 } 107 if err == nil && len(t.Attachments) == 0 { 108 t.Attachments = make([]Attachment, 0) 109 } 110 return t, err 111 } 112 113 // GetTemplateByName returns the template, if it exists, specified by the given name and user_id. 114 func GetTemplateByName(n string, uid int64) (Template, error) { 115 t := Template{} 116 err := db.Where("user_id=? and name=?", uid, n).Find(&t).Error 117 if err != nil { 118 log.Error(err) 119 return t, err 120 } 121 122 // Get Attachments 123 err = db.Where("template_id=?", t.Id).Find(&t.Attachments).Error 124 if err != nil && err != gorm.ErrRecordNotFound { 125 log.Error(err) 126 return t, err 127 } 128 if err == nil && len(t.Attachments) == 0 { 129 t.Attachments = make([]Attachment, 0) 130 } 131 return t, err 132 } 133 134 // PostTemplate creates a new template in the database. 135 func PostTemplate(t *Template) error { 136 // Insert into the DB 137 if err := t.Validate(); err != nil { 138 return err 139 } 140 err = db.Save(t).Error 141 if err != nil { 142 log.Error(err) 143 return err 144 } 145 146 // Save every attachment 147 for i := range t.Attachments { 148 t.Attachments[i].TemplateId = t.Id 149 err := db.Save(&t.Attachments[i]).Error 150 if err != nil { 151 log.Error(err) 152 return err 153 } 154 } 155 return nil 156 } 157 158 // PutTemplate edits an existing template in the database. 159 // Per the PUT Method RFC, it presumes all data for a template is provided. 160 func PutTemplate(t *Template) error { 161 if err := t.Validate(); err != nil { 162 return err 163 } 164 // Delete all attachments, and replace with new ones 165 err = db.Where("template_id=?", t.Id).Delete(&Attachment{}).Error 166 if err != nil && err != gorm.ErrRecordNotFound { 167 log.Error(err) 168 return err 169 } 170 if err == gorm.ErrRecordNotFound { 171 err = nil 172 } 173 for i, _ := range t.Attachments { 174 t.Attachments[i].TemplateId = t.Id 175 err := db.Save(&t.Attachments[i]).Error 176 if err != nil { 177 log.Error(err) 178 return err 179 } 180 } 181 182 // Save final template 183 err = db.Where("id=?", t.Id).Save(t).Error 184 if err != nil { 185 log.Error(err) 186 return err 187 } 188 return nil 189 } 190 191 // DeleteTemplate deletes an existing template in the database. 192 // An error is returned if a template with the given user id and template id is not found. 193 func DeleteTemplate(id int64, uid int64) error { 194 // Delete attachments 195 err := db.Where("template_id=?", id).Delete(&Attachment{}).Error 196 if err != nil { 197 log.Error(err) 198 return err 199 } 200 201 // Finally, delete the template itself 202 err = db.Where("user_id=?", uid).Delete(Template{Id: id}).Error 203 if err != nil { 204 log.Error(err) 205 return err 206 } 207 return nil 208 }