From f33ebcc21ca163c719986fc05985ac92699396a0 Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <brad.rydzewski@gmail.com>
Date: Sat, 1 Mar 2014 20:25:10 -0700
Subject: [PATCH] ability to use custom badges. standard badges now using
 shields.io

---
 pkg/handler/badges.go | 34 +++++++++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/pkg/handler/badges.go b/pkg/handler/badges.go
index 20415115d..62f77ae7a 100644
--- a/pkg/handler/badges.go
+++ b/pkg/handler/badges.go
@@ -7,10 +7,18 @@ import (
 	"github.com/drone/drone/pkg/database"
 )
 
+const (
+	badgeSuccess = "https://img.shields.io/badge/build-success-brightgreen.png"
+	badgeFailure = "https://img.shields.io/badge/build-failure-red.png"
+	badgeUnknown = "https://img.shields.io/badge/build-unknown-lightgray.png"
+)
+
 // Display a static badge (png format) for a specific
 // repository and an optional branch.
 // TODO this needs to implement basic caching
 func Badge(w http.ResponseWriter, r *http.Request) error {
+	successParam := r.FormValue("success")
+	failureParam := r.FormValue("failure")
 	branchParam := r.FormValue("branch")
 	hostParam := r.FormValue(":host")
 	ownerParam := r.FormValue(":owner")
@@ -30,18 +38,30 @@ func Badge(w http.ResponseWriter, r *http.Request) error {
 		branchParam = repo.DefaultBranch()
 	}
 
-	// default badge of "unknown"
-	badge := "/img/build_unknown.png"
+	var badge string
 
 	// get the latest commit from the database
 	// for the requested branch
 	commit, err := database.GetBranch(repo.ID, branchParam)
 	if err == nil {
-		switch commit.Status {
-		case "Success":
-			badge = "/img/build_success.png"
-		case "Failing", "Failure":
-			badge = "/img/build_failing.png"
+		switch {
+		case commit.Status == "Success" && len(successParam) == 0:
+			// if no success image is provided, we serve a
+			// badge using the shields.io service
+			badge = badgeSuccess
+		case commit.Status == "Success" && len(successParam) != 0:
+			// otherwise we serve the user defined success badge
+			badge = successParam
+		case commit.Status == "Failure" && len(failureParam) == 0:
+			// if no failure image is provided, we serve a
+			// badge using the shields.io service
+			badge = badgeFailure
+		case commit.Status == "Failure" && len(failureParam) != 0:
+			// otherwise we serve the user defined failure badge
+			badge = failureParam
+		default:
+			// otherwise load unknown image
+			badge = badgeUnknown
 		}
 	}