1 files changed, 170 insertions, 10 deletions

diff --git a/main.go b/main.go
index a1a2625..6c153fa 100644
--- a/main.go
+++ b/main.go
@@ -2,13 +2,23 @@ package main
 
 import (
         "bytes"
+        "context"
+        "flag"
         "fmt"
         "log"
+        "net/http"
         "os"
 
         "code.crute.us/mcrute/ses-smtpd-proxy/smtpd"
+        "github.com/aws/aws-sdk-go/aws"
+        "github.com/aws/aws-sdk-go/aws/credentials"
         "github.com/aws/aws-sdk-go/aws/session"
         "github.com/aws/aws-sdk-go/service/ses"
+        "github.com/hashicorp/vault/api"
+        "github.com/hashicorp/vault/api/auth/approle"
+        "github.com/prometheus/client_golang/prometheus"
+        "github.com/prometheus/client_golang/prometheus/promauto"
+        "github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
 const (
@@ -18,6 +28,34 @@ const (
 
 var sesClient *ses.SES
 
+var (
+        emailSent = promauto.NewCounter(prometheus.CounterOpts{
+                Namespace: "smtpd",
+                Name:      "email_send_success_total",
+                Help:      "Total number of successfuly sent emails",
+        })
+        emailError = promauto.NewCounterVec(prometheus.CounterOpts{
+                Namespace: "smtpd",
+                Name:      "email_send_fail_total",
+                Help:      "Total number emails that failed to send",
+        }, []string{"type"})
+        sesError = promauto.NewCounter(prometheus.CounterOpts{
+                Namespace: "smtpd",
+                Name:      "ses_error_total",
+                Help:      "Total number errors with SES",
+        })
+        credentialRenewalSuccess = promauto.NewCounter(prometheus.CounterOpts{
+                Namespace: "smtpd",
+                Name:      "credential_renewal_success_total",
+                Help:      "Total number successful credential renewals",
+        })
+        credentialRenewalError = promauto.NewCounter(prometheus.CounterOpts{
+                Namespace: "smtpd",
+                Name:      "credential_renewal_error_total",
+                Help:      "Total number errors during credential renewal",
+        })
+)
+
 type Envelope struct {
         from  string
         rcpts []*string
@@ -32,6 +70,7 @@ func (e *Envelope) AddRecipient(rcpt smtpd.MailAddress) error {
 
 func (e *Envelope) BeginData() error {
         if len(e.rcpts) == 0 {
+                emailError.With(prometheus.Labels{"type": "no valid recipients"}).Inc()
                 return smtpd.SMTPError("554 5.5.1 Error: no valid recipients")
         }
         return nil
@@ -40,7 +79,8 @@ func (e *Envelope) BeginData() error {
 func (e *Envelope) Write(line []byte) error {
         e.b.Write(line)
         if e.b.Len() > SesSizeLimit { // SES limitation
-                log.Printf("message size %d exceeds SES limit of %d\n", e.b.Len(), SesSizeLimit)
+                emailError.With(prometheus.Labels{"type": "minimum message size exceed"}).Inc()
+                log.Printf("message size %d exceeds SES limit of %d", e.b.Len(), SesSizeLimit)
                 return smtpd.SMTPError("554 5.5.1 Error: maximum message size exceeded")
         }
         return nil
@@ -52,10 +92,10 @@ func (e *Envelope) logMessageSend() {
                 dr[i] = *e.rcpts[i]
         }
         log.Printf("sending message from %+v to %+v", e.from, dr)
+        emailSent.Inc()
 }
 
 func (e *Envelope) Close() error {
-        e.logMessageSend()
         r := &ses.SendRawEmailInput{
                 Source:       &e.from,
                 Destinations: e.rcpts,
@@ -64,30 +104,150 @@ func (e *Envelope) Close() error {
         _, err := sesClient.SendRawEmail(r)
         if err != nil {
                 log.Printf("ERROR: ses: %v", err)
-                return smtpd.SMTPError(fmt.Sprintf("451 4.5.1 Temporary server error. Please try again later: %v", err))
+                emailError.With(prometheus.Labels{"type": "ses error"}).Inc()
+                sesError.Inc()
+                return smtpd.SMTPError("451 4.5.1 Temporary server error. Please try again later")
         }
+        e.logMessageSend()
         return err
 }
 
+func renewSecret(vc *api.Client, s *api.Secret) error {
+        w, err := vc.NewLifetimeWatcher(&api.LifetimeWatcherInput{Secret: s})
+        if err != nil {
+                return err
+        }
+        go w.Start()
+
+        go func() {
+                for {
+                        select {
+                        case err := <-w.DoneCh():
+                                if err != nil {
+                                        credentialRenewalError.Inc()
+                                        log.Fatalf("Error renewing credential: %s", err)
+                                }
+                        case renewal := <-w.RenewCh():
+                                credentialRenewalSuccess.Inc()
+                                log.Printf("Successfully renewed: %#v", renewal)
+                        }
+                }
+        }()
+
+        return nil
+}
+
+func getVaultSecret(path string) (credentials.Value, error) {
+        var r credentials.Value
+
+        vc, err := api.NewClient(api.DefaultConfig())
+        if err != nil {
+                return r, err
+        }
+
+        // Use AppRole if it's in the environment, otherwise assume VAULT_TOKEN
+        // was provided in the environment.
+        if roleID := os.Getenv("VAULT_APPROLE_ROLE_ID"); roleID != "" {
+                appRoleAuth, err := approle.NewAppRoleAuth(roleID, &approle.SecretID{
+                        FromEnv: "VAULT_APPROLE_SECRET_ID",
+                })
+                if err != nil {
+                        return r, fmt.Errorf("unable to initialize AppRole auth method: %w", err)
+                }
+                if loginSecret, err := vc.Auth().Login(context.Background(), appRoleAuth); err != nil {
+                        return r, fmt.Errorf("unable to login to AppRole auth method: %w", err)
+                } else {
+                        if err := renewSecret(vc, loginSecret); err != nil {
+                                return r, err
+                        }
+                }
+        }
+
+        secret, err := vc.Logical().Read(path)
+        if err != nil {
+                return r, err
+        }
+        if secret == nil {
+                return r, fmt.Errorf("Vault returned no AWS secret")
+        }
+
+        keyId, ok := secret.Data["access_key"]
+        if !ok {
+                return r, fmt.Errorf("Vault secret had no access_key")
+        }
+
+        secretKey, ok := secret.Data["secret_key"]
+        if !ok {
+                return r, fmt.Errorf("Vault secret had no secret_key")
+        }
+
+        r.AccessKeyID = keyId.(string)
+        r.SecretAccessKey = secretKey.(string)
+
+        return r, renewSecret(vc, secret)
+}
+
+func makeSesClient(enableVault bool, vaultPath string) (*ses.SES, error) {
+        var err error
+        var s *session.Session
+
+        if enableVault {
+                cred, err := getVaultSecret(vaultPath)
+                if err != nil {
+                        return nil, err
+                }
+
+                s, err = session.NewSession(&aws.Config{
+                        Credentials: credentials.NewStaticCredentialsFromCreds(cred),
+                })
+        } else {
+                s, err = session.NewSession()
+        }
+        if err != nil {
+                return nil, err
+        }
+
+        return ses.New(s), nil
+}
+
 func main() {
-        sesClient = ses.New(session.Must(session.NewSession()))
-        addr := DefaultAddr
+        var err error
+
+        disablePrometheus := flag.Bool("disable-prometheus", false, "Disables prometheus metrics server")
+        prometheusBind := flag.String("prometheus-bind", ":2501", "Address/port on which to bind Prometheus server")
+        enableVault := flag.Bool("enable-vault", false, "Enable fetching AWS IAM credentials from a Vault server")
+        vaultPath := flag.String("vault-path", "", "Full path to Vault credential (ex: \"aws/creds/my-mail-user\")")
+
+        flag.Parse()
+
+        sesClient, err = makeSesClient(*enableVault, *vaultPath)
+        if err != nil {
+                log.Fatalf("Error creating AWS session: %s", err)
+        }
 
-        if len(os.Args) == 2 {
-                addr = os.Args[1]
-        } else if len(os.Args) > 2 {
+        addr := DefaultAddr
+        if flag.Arg(0) != "" {
+                addr = flag.Arg(0)
+        } else if flag.NArg() > 1 {
                 log.Fatalf("usage: %s [listen_host:port]", os.Args[0])
         }
 
+        if !*disablePrometheus {
+                sm := http.NewServeMux()
+                ps := &http.Server{Addr: *prometheusBind, Handler: sm}
+                sm.Handle("/metrics", promhttp.Handler())
+                go ps.ListenAndServe()
+        }
+
         s := &smtpd.Server{
                 Addr: addr,
                 OnNewMail: func(c smtpd.Connection, from smtpd.MailAddress) (smtpd.Envelope, error) {
                         return &Envelope{from: from.Email()}, nil
                 },
         }
+
         log.Printf("ListenAndServe on %s", addr)
-        err := s.ListenAndServe()
-        if err != nil {
+        if err := s.ListenAndServe(); err != nil {
                 log.Fatalf("ListenAndServe: %v", err)
         }
 }