Docker Metrics API

main.go

package main

import (
	"encoding/json"
	"flag"
	"github.com/PuerkitoBio/goquery"
	"github.com/pmylund/go-cache"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const baseURL = "https://registry.hub.docker.com/"
const notFound = "Sorry we couldn't find this repo"
const unexpectedError = "Sorry but it seems something went wrong"

var port *string = flag.String("port", "7777", "Port to listen on")
var quiet *bool = flag.Bool("quiet", false, "Silence is golden")
var c *cache.Cache = cache.New(5*time.Minute, 1*time.Minute)

type ResponseBody struct {
	Stars int `json:"stars"`
	Pulls int `json:"pulls"`
}

func stringToInt(s string) (int, error) {
	i, err := strconv.Atoi(s)
	if err != nil {
		return -1, err
	}
	return i, nil
}

func logger(color string, code string, s string, i interface{}) {
	if *quiet == true {
		return
	}

	var colorCode string
	switch color {
	case "red":
		colorCode = "31"
	case "yellow":
		colorCode = "33"
	case "green":
		colorCode = "32"
	default:
		colorCode = "34"
	}

	prefix := "\033[" + colorCode + "m[" + code + "]\033[0m "
	if i != nil {
		log.Printf(prefix+s, i)
	} else {
		log.Printf(prefix + s)
	}
}

func cacheManager(repo string) interface{} {
	result, found := c.Get(repo)
	if found {
		return result
	} else {
		return false
	}
}

func getPulls(repo string) (int, error) {
	var path string

	firstChars := repo[:2]
	if firstChars != "_/" {
		path = "u/" + repo
	} else {
		path = repo
	}

	dockerhub := baseURL + path
	doc, err := goquery.NewDocument(dockerhub)
	if err != nil {
		logger("red", "502", "Could not get website at %s\n", dockerhub)
	}

	count, err := stringToInt(doc.Find(".downloads").Text())
	if err != nil {
		return -1, err
	}
	return count, nil
}

func getStars(repo string) (int, error) {
	var path string

	firstChars := repo[:2]
	if firstChars == "_/" {
		path = "library" + strings.TrimPrefix(repo, "_")
	} else {
		path = repo
	}

	dockerhub := baseURL + "v2/repositories/" + path + "/stars/count/"
	resp, err := http.Get(dockerhub)
	if err != nil {
		logger("red", "502", "Could not get API at %s\n", dockerhub)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		logger("red", "502", "Could not get response body from %s\n", dockerhub)
	}

	count, err := stringToInt(string(body))
	if err != nil {
		return -1, err
	}
	return count, nil
}

func serveDocker(w http.ResponseWriter, r *http.Request) {
	repo := r.URL.Path[1:]
	regex, err := regexp.Compile(`^.+\/.+$`)
	if err != nil {
		log.Println(err)
		logger("red", "500", "Unexpected error for %s\n", repo)
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(unexpectedError))
		return
	}

	if len(repo) < 4 || regex.MatchString(repo) == false {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Path required '/:owner/:repo' or '/_/:repo' for official repos."))
		logger("yellow", "400", "Did not attempt to get metrics for %s\n", repo)
		return
	}

	inCache := cacheManager(repo)
	if str, ok := inCache.(string); ok {
		if str == notFound {
			logger("yellow", "404", repo+" was not found (cache-hit)", nil)
			w.WriteHeader(http.StatusNotFound)
		} else {
			w.WriteHeader(http.StatusOK)
			logger("green", "200", repo+" was successful (cache-hit)", nil)
		}
		w.Write([]byte(str))
		return
	}

	pulls, pullsErr := getPulls(repo)
	stars, starsErr := getStars(repo)

	if pullsErr != nil || starsErr != nil {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte(notFound))
		c.Set(repo, string(notFound), cache.DefaultExpiration)
		logger("yellow", "404", "Could not get metrics for %s\n", repo)
		return
	}

	resBody := &ResponseBody{
		Stars: stars,
		Pulls: pulls,
	}

	body, err := json.Marshal(resBody)
	if err != nil {
		logger("red", "500", "Unexpected error for %s\n", repo)
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(unexpectedError))
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte(string(body)))
	c.Set(repo, string(body), cache.DefaultExpiration)
	logger("green", "200", repo+" was successful (cache-miss)", nil)
}

func main() {
	flag.Parse()
	logger("green", "GO", "Server starting on port "+*port, nil)
	http.HandleFunc("/", serveDocker)
	http.ListenAndServe(":"+*port, nil)
}