main.go

package main

import (
        "bytes"
        "encoding/json"
        "flag"
        "fmt"
        "io/ioutil"
        "log"
        "net/http"
        "net/url"
        "sort"
        "time"

        memdb "github.com/hashicorp/go-memdb"
        "github.com/julienschmidt/httprouter"
)

var (
        configFile   = flag.String("c", "conf.json", "config file location")
        parsedconfig = conf{}
        db           *memdb.MemDB
)

type conf struct {
        Subscription string `json:"subscription"`
        ClientID     string `json:"clientid"`
        ClientSecret string `json:"clientsecret"`
        TenantName   string `json:"tenantname"`
}

func Index(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
        fmt.Fprint(w, "Welcome!\n")
}

func loadMeters() {
        authResp, err := getAuthToken(parsedconfig.ClientID, parsedconfig.ClientSecret, parsedconfig.TenantName)
        if err != nil {
                log.Println("error getting auth token: ", err)
        }

        meters, err := getMeters(authResp.AccessToken)
        if err != nil {
                log.Println("error getting meters: ", err)
        }

        txn := db.Txn(true)

        for _, v := range meters {
                if err := txn.Insert("meters", v); err != nil {
                        panic(err)
                }
        }
        txn.Commit()
}

func GetVMs(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
        //TODO: should break this out to a function
        //TODO: should filter in query, not iterator
        txn := db.Txn(false)
        defer txn.Abort()
        results, err := txn.Get("meters", "id")
        if err != nil {
                log.Println("error getting meters from db: ", err)
        }

        vms := make([]Meter, 0)

        for i := results.Next(); i != nil; i = results.Next() {
                tempMeter := i.(Meter)
                if tempMeter.MeterName == "Compute Hours" && tempMeter.MeterCategory == "Virtual Machines" && tempMeter.MeterStatus == "Active" {
                        vms = append(vms, tempMeter)
                }
        }

        sort.Sort(BySubCat(vms))
        buffer, err := json.MarshalIndent(vms, "", "  ")
        if err != nil {
                log.Println("error doing indent: ", err)
        }
        _, err = w.Write(buffer)
}

func GetStorage(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
        txn := db.Txn(false)
        defer txn.Abort()
        results, err := txn.Get("meters", "id")
        if err != nil {
                log.Println("error getting meters from db: ", err)
        }

        storage := make([]Meter, 0)

        for i := results.Next(); i != nil; i = results.Next() {
                tempMeter := i.(Meter)
                if tempMeter.MeterCategory == "Storage" && tempMeter.MeterStatus == "Active" {
                        storage = append(storage, tempMeter)
                }
        }

        sort.Sort(BySubCat(storage))
        buffer, err := json.MarshalIndent(storage, "", "  ")
        if err != nil {
                log.Println("error doing indent: ", err)
        }
        _, err = w.Write(buffer)
}

func GetIPs(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
        txn := db.Txn(false)
        defer txn.Abort()
        results, err := txn.Get("meters", "id")
        if err != nil {
                log.Println("error getting meters from db: ", err)
        }

        ips := make([]Meter, 0)

        for i := results.Next(); i != nil; i = results.Next() {
                tempMeter := i.(Meter)
                if tempMeter.MeterCategory == "Networking" && tempMeter.MeterStatus == "Active" && (tempMeter.MeterSubCategory == "Reserved IP" || tempMeter.MeterSubCategory == "Public IP Addresses") {
                        ips = append(ips, tempMeter)
                }
        }

        sort.Sort(BySubCat(ips))
        buffer, err := json.MarshalIndent(ips, "", "  ")
        if err != nil {
                log.Println("error doing indent: ", err)
        }
        _, err = w.Write(buffer)
}

func main() {
        flag.Parse()

        file, err := ioutil.ReadFile(*configFile)
        if err != nil {
                log.Fatal("unable to read config file, exiting...")
        }
        if err := json.Unmarshal(file, &parsedconfig); err != nil {
                log.Fatal("unable to marshal config file, exiting...")
        }

        // Create the DB schema
        schema := &memdb.DBSchema{
                Tables: map[string]*memdb.TableSchema{
                        "meters": &memdb.TableSchema{
                                Name: "meters",
                                Indexes: map[string]*memdb.IndexSchema{
                                        "id": &memdb.IndexSchema{
                                                Name:    "id",
                                                Unique:  true,
                                                Indexer: &memdb.StringFieldIndex{Field: "MeterID"},
                                        },
                                },
                        },
                },
        }

        // Create a new data base
        db, err = memdb.NewMemDB(schema)
        if err != nil {
                panic(err)
        }

        log.Println("loading meters...")
        loadMeters()
        log.Println("all meters loaded, continuing...")
        router := httprouter.New()
        router.GET("/vm", GetVMs)
        router.GET("/storage", GetStorage)
        router.GET("/ips", GetIPs)
        router.NotFound = http.FileServer(http.Dir("public"))
        log.Println("starting server...")
        log.Fatal(http.ListenAndServe(":8080", router))
}

func httpReq(authToken string, method string, URL string, postData []byte, isJson bool) (*http.Response, error) {
        client := &http.Client{
                Timeout: 20 * time.Second,
        }
        req, err := http.NewRequest(method, URL, bytes.NewBuffer(postData))
        if err != nil {
                log.Println("Error with building request for "+URL+": ", err)
                return &http.Response{}, err
        }
        if authToken != "" {
                req.Header.Add("Authorization", "Bearer "+authToken)
        }
        if isJson {
                req.Header.Add("Content-Type", "application/json")
        }
        resp, err := client.Do(req)
        if err != nil {
                log.Println("Error with request for "+URL+": ", err)
                return &http.Response{}, err
        }
        return resp, nil
}

func getAuthToken(clientID, clientSecret, tenantName string) (AuthTokenResp, error) {
        data := url.Values{}
        data.Set("grant_type", "client_credentials")
        data.Add("client_id", parsedconfig.ClientID)
        data.Add("client_secret", parsedconfig.ClientSecret)
        data.Add("resource", "https://management.azure.com/")

        r, _ := httpReq("", "POST", "https://login.microsoftonline.com/"+parsedconfig.TenantName+"/oauth2/token", []byte(data.Encode()), false)
        var authResp AuthTokenResp
        defer r.Body.Close()
        body, err := ioutil.ReadAll(r.Body)
        if err != nil {
                return authResp, err
        }
        err = json.Unmarshal(body, &authResp)
        if err != nil {
                return authResp, err
        }
        return authResp, nil
}

func getMeters(authToken string) ([]Meter, error) {
        data := url.Values{}
        data.Set("api-version", "2016-08-31-preview")
        data.Add("$filter", "OfferDurableId eq 'MS-AZR-0003p' and Currency eq 'USD' and Locale eq 'en-US' and RegionInfo eq 'US'")

        s, _ := httpReq(authToken, "GET", "https://management.azure.com/subscriptions/"+parsedconfig.Subscription+"/providers/Microsoft.Commerce/RateCard?"+data.Encode(), nil, true)
        var meterList MeterResponse
        body2, err := ioutil.ReadAll(s.Body)

        if err != nil {
                return nil, err
        }

        err = json.Unmarshal(body2, &meterList)
        if err != nil {
                return nil, err
        }

        return meterList.Meters, nil
}