From 1a76ffc1d1b21421ad8d80160c5d3dae33ee0cca Mon Sep 17 00:00:00 2001
From: Timon Wong <timon86.wang@gmail.com>
Date: Thu, 5 Dec 2019 01:10:43 +0800
Subject: [PATCH] Add signature support, fix #41

---
 config.example.yml     |  2 ++
 config/config.go       |  3 ++-
 notifier/prometheus.go | 43 +++++++++++++++++++++++++++++++++++-------
 webrouter/dingtalk.go  |  2 +-
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/config.example.yml b/config.example.yml
index 6d07bbdf..76150037 100644
--- a/config.example.yml
+++ b/config.example.yml
@@ -3,5 +3,7 @@
 targets:
   webhook1:
     url: https://oapi.dingtalk.com/robot/send?access_token=xxxxxxxxxxxx
+    # secret for signature
+    secret: this_is_secret
   webhook2:
     url: https://oapi.dingtalk.com/robot/send?access_token=xxxxxxxxxxxx
diff --git a/config/config.go b/config/config.go
index 1f2ee085..23e4fc52 100644
--- a/config/config.go
+++ b/config/config.go
@@ -61,5 +61,6 @@ func (c *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
 }
 
 type Target struct {
-	URL string `yaml:"url"`
+	URL    string `yaml:"url"`
+	Secret string `yaml:"secret"`
 }
diff --git a/notifier/prometheus.go b/notifier/prometheus.go
index 6af02e32..07def262 100644
--- a/notifier/prometheus.go
+++ b/notifier/prometheus.go
@@ -2,12 +2,18 @@ package notifier
 
 import (
 	"bytes"
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/base64"
 	"encoding/json"
 	"net/http"
+	"net/url"
+	"strconv"
 	"time"
 
 	"github.com/pkg/errors"
 
+	"github.com/timonwong/prometheus-webhook-dingtalk/config"
 	"github.com/timonwong/prometheus-webhook-dingtalk/models"
 	"github.com/timonwong/prometheus-webhook-dingtalk/template"
 )
@@ -33,30 +39,53 @@ func BuildDingTalkNotification(promMessage *models.WebhookMessage) (*models.Ding
 	return notification, nil
 }
 
-func SendDingTalkNotification(httpClient *http.Client, webhookURL string, notification *models.DingTalkNotification) (*models.DingTalkNotificationResponse, error) {
+func SendDingTalkNotification(httpClient *http.Client, target config.Target, notification *models.DingTalkNotification) (*models.DingTalkNotificationResponse, error) {
 	body, err := json.Marshal(&notification)
 	if err != nil {
 		return nil, errors.Wrap(err, "error encoding DingTalk request")
 	}
 
-	httpReq, err := http.NewRequest("POST", webhookURL, bytes.NewReader(body))
+	targetURL := target.URL
+	// Calculate signature when secret is provided
+	if target.Secret != "" {
+		timestamp := strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10)
+		stringToSign := []byte(timestamp + "\n" + target.Secret)
+
+		mac := hmac.New(sha256.New, []byte(target.Secret))
+		mac.Write(stringToSign) // nolint: errcheck
+		signature := base64.StdEncoding.EncodeToString(mac.Sum(nil))
+
+		u, err := url.Parse(targetURL)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to parse target url")
+		}
+
+		qs := u.Query()
+		qs.Set("timestamp", timestamp)
+		qs.Set("sign", signature)
+		u.RawQuery = qs.Encode()
+
+		targetURL = u.String()
+	}
+
+	httpReq, err := http.NewRequest("POST", targetURL, bytes.NewReader(body))
 	if err != nil {
 		return nil, errors.Wrap(err, "error building DingTalk request")
 	}
 	httpReq.Header.Set("Content-Type", "application/json")
 
-	req, err := httpClient.Do(httpReq)
+	resp, err := httpClient.Do(httpReq)
 	if err != nil {
 		return nil, errors.Wrap(err, "error sending notification to DingTalk")
 	}
-	defer req.Body.Close()
+	defer resp.Body.Close()
 
-	if req.StatusCode != 200 {
-		return nil, errors.Errorf("unacceptable response code %d", req.StatusCode)
+	if resp.StatusCode != 200 {
+		return nil, errors.Errorf("unacceptable response code %d", resp.StatusCode)
 	}
 
 	var robotResp models.DingTalkNotificationResponse
-	enc := json.NewDecoder(req.Body)
+	enc := json.NewDecoder(resp.Body)
 	if err := enc.Decode(&robotResp); err != nil {
 		return nil, errors.Wrap(err, "error decoding response from DingTalk")
 	}
diff --git a/webrouter/dingtalk.go b/webrouter/dingtalk.go
index 96b69175..f924da7f 100644
--- a/webrouter/dingtalk.go
+++ b/webrouter/dingtalk.go
@@ -58,7 +58,7 @@ func (rs *DingTalkResource) SendNotification(w http.ResponseWriter, r *http.Requ
 		return
 	}
 
-	robotResp, err := notifier.SendDingTalkNotification(rs.HttpClient, target.URL, notification)
+	robotResp, err := notifier.SendDingTalkNotification(rs.HttpClient, target, notification)
 	if err != nil {
 		level.Error(logger).Log("msg", "Failed to send notification", "err", err)
 		http.Error(w, "Bad Request", http.StatusBadRequest)