<commit_before>package main

import (
	"bytes"
	"code.google.com/p/go.crypto/openpgp"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/user"
)

const jaegerJSONGPGDBExtension = ".jgrdb"
const jaegerDBDescription = "JaegerDB - Jaeger database management program\n\nJaeger is a JSON encoded GPG encrypted key value store. It is useful for separating development with operations and keeping configuration files secure."
const jaegerQuote = "\"Stacker Pentecost: Haven't you heard Mr. Beckett? The world is coming to an end. So where would you rather die? Here? Or in a Jaeger!\" - Pacific Rim"
const jaegerDBRecommendedUsage = "RECOMMENDED:\n    jaegerdb -j file.txt.jgrdb -a \"Field1\" -v \"Secret value\"\n\nThis will run JaegerDB with the default options and assume the following:\n    Keyring file: ~/.gnupg/jaeger_pubring.gpg"

var debug debugging = false

type debugging bool

func (d debugging) Printf(format string, args ...interface{}) {
	// From: https://groups.google.com/forum/#!msg/golang-nuts/gU7oQGoCkmg/BNIl-TqB-4wJ
	if d {
		log.Printf(format, args...)
	}
}

type Data struct {
	Properties []Property
}

type Property struct {
	Name           string `json:"Name"`
	EncryptedValue string `json:"EncryptedValue"`
}

func main() {
	// Define flags
	// TODO: View individual property and unencrypted value. 'get'
	var (
		addKey       = flag.String("a", "", "Add property")
		changeKey    = flag.String("c", "", "Change property")
		debugFlag    = flag.Bool("d", false, "Enable Debug")
		deleteKey    = flag.String("delete", "", "Delete property")
		inializeFlag = flag.Bool("init", false, "Create an initial blank JSON GPG database file")
		jsonGPGDB    = flag.String("j", "", "JSON GPG database file. eg. file.txt.jgrdb")
		keyringFile  = flag.String("k", "", "Keyring file. Public key in ASCII armored format. eg. pubring.asc")
		value        = flag.String("v", "", "Value for property to use")
	)

	flag.Usage = func() {
		fmt.Printf("%s\n%s\n\n%s\n\n", jaegerDBDescription, jaegerQuote, jaegerDBRecommendedUsage)
		fmt.Fprintf(os.Stderr, "OPTIONS:\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	if *debugFlag {
		debug = true
	}

	if *jsonGPGDB == "" {
		flag.Usage()
		log.Fatalf("\n\nError: No JSON GPG database file specified")
		return
	}

	if *inializeFlag {
		initializeJSONGPGDB(jsonGPGDB)
	}

	if *deleteKey != "" {
		deleteKeyJaegerDB(deleteKey, jsonGPGDB)
	}

	var entity *openpgp.Entity
	var entitylist openpgp.EntityList

	if *keyringFile == "" {
		entity, entitylist = processPublicKeyRing()
	} else {
		entity, entitylist = processArmoredKeyRingFile(keyringFile)
	}

	if *addKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for add key operation specified")
		}
		addKeyJaegerDB(addKey, value, jsonGPGDB, entitylist)
	}

	if *changeKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for change key operation specified")
		}
		changeKeyJaegerDB(changeKey, value, jsonGPGDB, entitylist)
	}

	debug.Printf("End - Delete this line", *jsonGPGDB, *keyringFile, entity, entitylist)
}

func initializeJSONGPGDB(jsonGPGDB *string) {
	if _, err := os.Stat(*jsonGPGDB); err == nil {
		log.Fatalf("ERR: File already exists: %v", *jsonGPGDB)
	}

	var newP []Property

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote new file:", *jsonGPGDB)
	}

}

func encodeBase64EncryptedMessage(s string, entitylist openpgp.EntityList) string {
	// Encrypt message using public key and then encode with base64
	debug.Printf("entitylist: #%v", entitylist)
	buf := new(bytes.Buffer)
	w, err := openpgp.Encrypt(buf, entitylist, nil, nil, nil)
	if err != nil {
		log.Fatalln("ERR: Error encrypting message - ", err)
	}

	_, err = w.Write([]byte(s))
	if err != nil {
	}
	err = w.Close()
	if err != nil {
	}

	// Output as base64 encoded string
	bytes, err := ioutil.ReadAll(buf)
	str := base64.StdEncoding.EncodeToString(bytes)

	debug.Printf("Public key encrypted message (base64 encoded): %v", str)

	return str
}

func processPublicKeyRing() (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	// TODO: Handle a specified recipient
	// Get default public keyring location
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}

	jaegerPublicKeyRing := fmt.Sprintf("%v/.gnupg/jaeger_pubring.gpg", usr.HomeDir)
	publicKeyRing := ""

	if _, err := os.Stat(jaegerPublicKeyRing); err == nil {
		publicKeyRing = jaegerPublicKeyRing
	} else {
		publicKeyRing = fmt.Sprintf("%v/.gnupg/pubring.gpg", usr.HomeDir)
	}

	debug.Printf("publicKeyRing file:", publicKeyRing)
	publicKeyRingBuffer, err := os.Open(publicKeyRing)
	if err != nil {
		panic(err)
	}
	entitylist, err = openpgp.ReadKeyRing(publicKeyRingBuffer)
	if err != nil {
		log.Fatal(err)
	}

	entity = entitylist[0]
	debug.Printf("Public key default keyring:", entity.Identities)

	return entity, entitylist
}

func processArmoredKeyRingFile(keyringFile *string) (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	keyringFileBuffer, err := os.Open(*keyringFile)
	if err != nil {
		log.Fatalln("ERROR: Unable to read keyring file")
	}
	entitylist, err = openpgp.ReadArmoredKeyRing(keyringFileBuffer)
	if err != nil {
		log.Fatal(err)
	}
	entity = entitylist[0]
	debug.Printf("Public key from ASCII armored string:", entity.Identities)

	return entity, entitylist
}

func addKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		log.Fatalln("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		panic(err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	var newP []Property

	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search
	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if found {
		log.Fatalf("\n\nError: Property '%s' already exists.", *key)
	}

	newP = append(j.Properties, p)

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote file:", *jsonGPGDB)
	}

}

func changeKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		log.Fatalln("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		panic(err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	// New property to replace the old
	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search and replace
	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if !found {
		log.Fatalf("\n\nError: Property '%s' not found.", *key)
	}

	bytes, err := json.MarshalIndent(j, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote file:", *jsonGPGDB)
	}

}

func deleteKeyJaegerDB(key *string, jsonGPGDB *string) {
	debug.Printf("deleteKeyJaegerDB key: %v", *key)

	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		log.Fatalln("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		panic(err)
	}
	debug.Printf("json unmarshal: %v", j)

	var newP []Property
	found := false

	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			// https://code.google.com/p/go-wiki/wiki/SliceTricks
			newP = j.Properties[:i+copy(j.Properties[i:], j.Properties[i+1:])]
			found = true
			break
		}
	}

	if !found {
		log.Fatalf("\n\nError: Property '%s' not found.", *key)
	}

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote file:", *jsonGPGDB)
	}

}
<commit_msg>Mention if no JSON GPG database operations specified. Removed unused entity variable.<commit_after>package main

import (
	"bytes"
	"code.google.com/p/go.crypto/openpgp"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/user"
)

const jaegerJSONGPGDBExtension = ".jgrdb"
const jaegerDBDescription = "JaegerDB - Jaeger database management program\n\nJaeger is a JSON encoded GPG encrypted key value store. It is useful for separating development with operations and keeping configuration files secure."
const jaegerQuote = "\"Stacker Pentecost: Haven't you heard Mr. Beckett? The world is coming to an end. So where would you rather die? Here? Or in a Jaeger!\" - Pacific Rim"
const jaegerDBRecommendedUsage = "RECOMMENDED:\n    jaegerdb -j file.txt.jgrdb -a \"Field1\" -v \"Secret value\"\n\nThis will run JaegerDB with the default options and assume the following:\n    Keyring file: ~/.gnupg/jaeger_pubring.gpg"

var debug debugging = false

type debugging bool

func (d debugging) Printf(format string, args ...interface{}) {
	// From: https://groups.google.com/forum/#!msg/golang-nuts/gU7oQGoCkmg/BNIl-TqB-4wJ
	if d {
		log.Printf(format, args...)
	}
}

type Data struct {
	Properties []Property
}

type Property struct {
	Name           string `json:"Name"`
	EncryptedValue string `json:"EncryptedValue"`
}

func main() {
	// Define flags
	// TODO: View individual property and unencrypted value. 'get'
	var (
		addKey       = flag.String("a", "", "Add property")
		changeKey    = flag.String("c", "", "Change property")
		debugFlag    = flag.Bool("d", false, "Enable Debug")
		deleteKey    = flag.String("delete", "", "Delete property")
		inializeFlag = flag.Bool("init", false, "Create an initial blank JSON GPG database file")
		jsonGPGDB    = flag.String("j", "", "JSON GPG database file. eg. file.txt.jgrdb")
		keyringFile  = flag.String("k", "", "Keyring file. Public key in ASCII armored format. eg. pubring.asc")
		value        = flag.String("v", "", "Value for property to use")
	)

	flag.Usage = func() {
		fmt.Printf("%s\n%s\n\n%s\n\n", jaegerDBDescription, jaegerQuote, jaegerDBRecommendedUsage)
		fmt.Fprintf(os.Stderr, "OPTIONS:\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	if *debugFlag {
		debug = true
	}

	if *jsonGPGDB == "" {
		flag.Usage()
		log.Fatalf("\n\nError: No JSON GPG database file specified")
		return
	}

	if *inializeFlag {
		initializeJSONGPGDB(jsonGPGDB)
	}

	if *deleteKey != "" {
		deleteKeyJaegerDB(deleteKey, jsonGPGDB)
	}

	var entitylist openpgp.EntityList

	if *keyringFile == "" {
		_, entitylist = processPublicKeyRing()
	} else {
		_, entitylist = processArmoredKeyRingFile(keyringFile)
	}

	if *addKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for add key operation specified")
		}
		addKeyJaegerDB(addKey, value, jsonGPGDB, entitylist)
	}

	if *changeKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for change key operation specified")
		}
		changeKeyJaegerDB(changeKey, value, jsonGPGDB, entitylist)
	}

	if *deleteKey == "" && *addKey == "" && *changeKey == "" {
		log.Fatalf("\n\nError: No JSON GPG database operations specified")
	}

}

func initializeJSONGPGDB(jsonGPGDB *string) {
	if _, err := os.Stat(*jsonGPGDB); err == nil {
		log.Fatalf("ERR: File already exists: %v", *jsonGPGDB)
	}

	var newP []Property

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote new file:", *jsonGPGDB)
	}

}

func encodeBase64EncryptedMessage(s string, entitylist openpgp.EntityList) string {
	// Encrypt message using public key and then encode with base64
	debug.Printf("entitylist: #%v", entitylist)
	buf := new(bytes.Buffer)
	w, err := openpgp.Encrypt(buf, entitylist, nil, nil, nil)
	if err != nil {
		log.Fatalln("ERR: Error encrypting message - ", err)
	}

	_, err = w.Write([]byte(s))
	if err != nil {
	}
	err = w.Close()
	if err != nil {
	}

	// Output as base64 encoded string
	bytes, err := ioutil.ReadAll(buf)
	str := base64.StdEncoding.EncodeToString(bytes)

	debug.Printf("Public key encrypted message (base64 encoded): %v", str)

	return str
}

func processPublicKeyRing() (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	// TODO: Handle a specified recipient
	// Get default public keyring location
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}

	jaegerPublicKeyRing := fmt.Sprintf("%v/.gnupg/jaeger_pubring.gpg", usr.HomeDir)
	publicKeyRing := ""

	if _, err := os.Stat(jaegerPublicKeyRing); err == nil {
		publicKeyRing = jaegerPublicKeyRing
	} else {
		publicKeyRing = fmt.Sprintf("%v/.gnupg/pubring.gpg", usr.HomeDir)
	}

	debug.Printf("publicKeyRing file:", publicKeyRing)
	publicKeyRingBuffer, err := os.Open(publicKeyRing)
	if err != nil {
		panic(err)
	}
	entitylist, err = openpgp.ReadKeyRing(publicKeyRingBuffer)
	if err != nil {
		log.Fatal(err)
	}

	entity = entitylist[0]
	debug.Printf("Public key default keyring:", entity.Identities)

	return entity, entitylist
}

func processArmoredKeyRingFile(keyringFile *string) (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	keyringFileBuffer, err := os.Open(*keyringFile)
	if err != nil {
		log.Fatalln("ERROR: Unable to read keyring file")
	}
	entitylist, err = openpgp.ReadArmoredKeyRing(keyringFileBuffer)
	if err != nil {
		log.Fatal(err)
	}
	entity = entitylist[0]
	debug.Printf("Public key from ASCII armored string:", entity.Identities)

	return entity, entitylist
}

func addKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		log.Fatalln("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		panic(err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	var newP []Property

	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search
	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if found {
		log.Fatalf("\n\nError: Property '%s' already exists.", *key)
	}

	newP = append(j.Properties, p)

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote file:", *jsonGPGDB)
	}

}

func changeKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		log.Fatalln("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		panic(err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	// New property to replace the old
	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search and replace
	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if !found {
		log.Fatalf("\n\nError: Property '%s' not found.", *key)
	}

	bytes, err := json.MarshalIndent(j, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote file:", *jsonGPGDB)
	}

}

func deleteKeyJaegerDB(key *string, jsonGPGDB *string) {
	debug.Printf("deleteKeyJaegerDB key: %v", *key)

	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		log.Fatalln("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		panic(err)
	}
	debug.Printf("json unmarshal: %v", j)

	var newP []Property
	found := false

	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			// https://code.google.com/p/go-wiki/wiki/SliceTricks
			newP = j.Properties[:i+copy(j.Properties[i:], j.Properties[i+1:])]
			found = true
			break
		}
	}

	if !found {
		log.Fatalf("\n\nError: Property '%s' not found.", *key)
	}

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		fmt.Println("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		panic(err)
	} else {
		log.Fatalln("Wrote file:", *jsonGPGDB)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var index map[string](int) = map[string](int){
	"NAME":       0,
	"STATE":      1,
	"CPU_SEC":    2,
	"CPU_PER":    3,
	"MEM_K":      4,
	"MEM_PER":    5,
	"MAXMEM_K":   6,
	"MAXMEM_PER": 7,
	"VCPUS":      8,
	"NETS":       9,
	"NETTX":      10,
	"NETRX":      11,
	"VBDS":       12,
	"VBD_OO":     13,
	"VBD_RD":     14,
	"VBD_WR":     15,
	"VBD_RSECT":  16,
	"VBD_WSECT":  17,
	"SSID":       18,
}

// GraphDefinitionで動的に定義するので，これは多分必要なくなる
var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"xentop.cpu": mp.Graphs{
		Label:   "Xentop CPU",
		Unit:    "float",
		Metrics: [](mp.Metrics){},
	},
	"xentop.memory": mp.Graphs{
		Label:   "Xentop Memory",
		Unit:    "float",
		Metrics: [](mp.Metrics){},
	},
	"xentop.network": mp.Graphs{
		Label:   "Xentop Network",
		Unit:    "float",
		Metrics: [](mp.Metrics){},
	},
	"xentop.io": mp.Graphs{},
}

type XentopMetrics struct {
	HostName string
	Metrics  mp.Metrics
}

type XentopPlugin struct {
	GraphName          string
	GraphUnit          string
	XentopMetricsSlice []XentopMetrics
}

func parseXentop() error {
	cmd := exec.Command("/bin/sh", "-c", "sudo xentop --batch -i 1 -f")
	s, err := cmd.Run()
	if err != nil {
		return nil, err
	}
	return nil
}

func (m XentopPlugin) FetchMetrics() (map[string]float64, error) {
	stat := make(map[string]float64)
	cmd := exec.Command("/bin/sh", "-c", "sudo xentop --batch -i 1 -f")
	stdout, err := cmd.StdoutPipe()

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	cmd.Start()

	scanner := bufio.NewScanner(stdout)
	scanner.Scan()
	for scanner.Scan() {
		sf := strings.Fields(string(scanner.Text()))
		name := sf[index["NAME"]]
		// TODO: みにくいのでなんとかしたい
		stat[fmt.Sprintf("cpu_%s", name)] = strconv.Atof(sf[index["CPU_PER"]])
		stat[fmt.Sprintf("memory_%s", name)] = strconv.Atof(sf[index["MEM_PER"]])
		stat[fmt.Sprintf("vbd_read_%s", name)] = strconv.Atof(sf[index["VBD_RD"]])
		stat[fmt.Sprintf("vbd_write_%s", name)] = strconv.Atof(sf[index["VBD_WR"]])
		fmt.Println(sf[index["NAME"]])
	}
	if err := scanner.Err(); err != nil {
		panic(err)
	}

}

// ここでグラフを定義する
func (m XentopPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := []mp.Metrics{}
}

func main() {
	// flagの取得

	var xentop XentopPlugin

	helper := mp.NewMackerelPlugin(xentop)

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>メトリクスの返却<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var index map[string](int) = map[string](int){
	"NAME":       0,
	"STATE":      1,
	"CPU_SEC":    2,
	"CPU_PER":    3,
	"MEM_K":      4,
	"MEM_PER":    5,
	"MAXMEM_K":   6,
	"MAXMEM_PER": 7,
	"VCPUS":      8,
	"NETS":       9,
	"NETTX":      10,
	"NETRX":      11,
	"VBDS":       12,
	"VBD_OO":     13,
	"VBD_RD":     14,
	"VBD_WR":     15,
	"VBD_RSECT":  16,
	"VBD_WSECT":  17,
	"SSID":       18,
}

// GraphDefinitionで動的に定義するので，これは多分必要なくなる
var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"xentop.cpu": mp.Graphs{
		Label:   "Xentop CPU",
		Unit:    "float",
		Metrics: [](mp.Metrics){},
	},
	"xentop.memory": mp.Graphs{
		Label:   "Xentop Memory",
		Unit:    "float",
		Metrics: [](mp.Metrics){},
	},
	"xentop.network": mp.Graphs{
		Label:   "Xentop Network",
		Unit:    "float",
		Metrics: [](mp.Metrics){},
	},
	"xentop.io": mp.Graphs{},
}

type XentopMetrics struct {
	HostName string
	Metrics  mp.Metrics
}

type XentopPlugin struct {
	GraphName          string
	GraphUnit          string
	XentopMetricsSlice []XentopMetrics
}

func parseXentop() error {
	cmd := exec.Command("/bin/sh", "-c", "sudo xentop --batch -i 1 -f")
	s, err := cmd.Run()
	if err != nil {
		return nil, err
	}
	return nil
}

func (m XentopPlugin) FetchMetrics() (map[string]float64, error) {
	stat := make(map[string]float64)
	cmd := exec.Command("/bin/sh", "-c", "sudo xentop --batch -i 1 -f")
	stdout, err := cmd.StdoutPipe()

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	cmd.Start()

	scanner := bufio.NewScanner(stdout)
	scanner.Scan()
	for scanner.Scan() {
		sf := strings.Fields(string(scanner.Text()))
		name := sf[index["NAME"]]
		// TODO: みにくいのでなんとかしたい
		stat[fmt.Sprintf("cpu_%s", name)] = strconv.Atof(sf[index["CPU_PER"]])
		stat[fmt.Sprintf("memory_%s", name)] = strconv.Atof(sf[index["MEM_PER"]])
		stat[fmt.Sprintf("vbd_read_%s", name)] = strconv.Atof(sf[index["VBD_RD"]])
		stat[fmt.Sprintf("vbd_write_%s", name)] = strconv.Atof(sf[index["VBD_WR"]])
		fmt.Println(sf[index["NAME"]])
	}
	if err := scanner.Err(); err != nil {
		panic(err)
	}

	return stat
}

// ここでグラフを定義する
func (m XentopPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := []mp.Metrics{}
}

func main() {
	// flagの取得

	var xentop XentopPlugin

	helper := mp.NewMackerelPlugin(xentop)

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>


<commit_before>// +build !debug

package web

import (
	"bytes"
	"io"

	"github.com/go-martini/martini"
	"github.com/martini-contrib/staticbin"
	"github.com/tyler-sommer/squircy2/web/generated"
	"github.com/tyler-sommer/squircy2/web/generated/manage"
	"github.com/tyler-sommer/squircy2/web/generated/repl"
	"github.com/tyler-sommer/squircy2/web/generated/script"
	"github.com/tyler-sommer/squircy2/web/generated/webhook"
	"github.com/tyler-sommer/stick"
)

var templateMapping = map[string]generatedTemplate{
	"index.html.twig":         generated.TemplateIndexHtmlTwig,
	"webhook/edit.html.twig":  webhook.TemplateWebhookEditHtmlTwig,
	"webhook/index.html.twig": webhook.TemplateWebhookIndexHtmlTwig,
	"webhook/new.html.twig":   webhook.TemplateWebhookNewHtmlTwig,
	"script/edit.html.twig":   script.TemplateScriptEditHtmlTwig,
	"script/index.html.twig":  script.TemplateScriptIndexHtmlTwig,
	"script/new.html.twig":    script.TemplateScriptNewHtmlTwig,
	"repl/index.html.twig":    repl.TemplateReplIndexHtmlTwig,
	"manage/edit.html.twig":   manage.TemplateManageEditHtmlTwig,
}

type stringTemplate struct {
	name     string
	contents string
}

func (t *stringTemplate) Name() string {
	return t.name
}

func (t *stringTemplate) Contents() io.Reader {
	return bytes.NewBufferString(t.contents)
}

func newTemplateLoader() stick.Loader {
	return &assetLoader{}
}

type assetLoader struct{}

func (l *assetLoader) Load(name string) (stick.Template, error) {
	res, err := Asset("web/views/" + name)
	if err != nil {
		return nil, err
	}
	return &stringTemplate{name, string(res)}, nil
}

func newStaticHandler() martini.Handler {
	return staticbin.Static("web/public", Asset)
}
<commit_msg>Fixed issue with static asset loader.<commit_after>// +build !debug

package web

import (
	"bytes"
	"io"

	"github.com/go-martini/martini"
	"github.com/martini-contrib/staticbin"
	"github.com/tyler-sommer/squircy2/web/generated"
	"github.com/tyler-sommer/squircy2/web/generated/manage"
	"github.com/tyler-sommer/squircy2/web/generated/repl"
	"github.com/tyler-sommer/squircy2/web/generated/script"
	"github.com/tyler-sommer/squircy2/web/generated/webhook"
	"github.com/tyler-sommer/stick"
)

var templateMapping = map[string]generatedTemplate{
	"index.html.twig":         generated.TemplateIndexHtmlTwig,
	"webhook/edit.html.twig":  webhook.TemplateWebhookEditHtmlTwig,
	"webhook/index.html.twig": webhook.TemplateWebhookIndexHtmlTwig,
	"webhook/new.html.twig":   webhook.TemplateWebhookNewHtmlTwig,
	"script/edit.html.twig":   script.TemplateScriptEditHtmlTwig,
	"script/index.html.twig":  script.TemplateScriptIndexHtmlTwig,
	"script/new.html.twig":    script.TemplateScriptNewHtmlTwig,
	"repl/index.html.twig":    repl.TemplateReplIndexHtmlTwig,
	"manage/edit.html.twig":   manage.TemplateManageEditHtmlTwig,
}

type stringTemplate struct {
	name     string
	contents string
}

func (t *stringTemplate) Name() string {
	return t.name
}

func (t *stringTemplate) Contents() io.Reader {
	return bytes.NewBufferString(t.contents)
}

func newTemplateLoader() stick.Loader {
	return &assetLoader{}
}

type assetLoader struct{}

func (l *assetLoader) Load(name string) (stick.Template, error) {
	res, err := Asset("web/views/" + name)
	if err != nil {
		return nil, err
	}
	return &stringTemplate{name, string(res)}, nil
}

func newStaticHandler() martini.Handler {
	return staticbin.Static("public", Asset)
}
<|endoftext|>


<commit_before>package web

import (
	"github.com/masterminds/cookoo"
	//cookoo "../"
	"net/http"
	"fmt"
)

// Create a new Cookoo web server.
//
// Important details:
//
// - A URIPathResolver is used for resolving request names.
// - The following datasources are added to the Context:
//   * url: A URLDatasource (Provides access to parts of the URL)
//   * path: A PathDatasource (Provides access to parts of a path. E.g. "/foo/bar")
//   * query: A QueryParameterDatasource (Provides access to URL query parameters.)
//   * post: A FormValuesDatasource (Provides access to form data or the body of a request.)
// - The following context variables are set:
//   * http.Request: A pointer to the http.Request object
//   * http.ResponseWriter: The response writer.
//   * server.Address: The server's address and port (NOT ALWAYS PRESENT)
// - The handler includes logic to redirect "not found" errors to a path named "@404" if present.
//
// Context Params:
//
// - server.Address: If this key exists in the context, it will be used to determine the host/port the
//   server runes on. EXPERIMENTAL. Default is ":8080".
//
// Example:
//
//    package main
//
//    import (
//      //This is the path to Cookoo
//      "github.com/masterminds/cookoo/src/cookoo"
//      "github.com/masterminds/cookoo/src/cookoo/web"
//      "fmt"
//    )
//
//    func main() {
//      // Build a new Cookoo app.
//      handler, registry, router, context := cookoo.Cookoo()
//
//      // Fill the registry.
//      registry.Route("GET /", "The index").Does(web.Flush, "example").
//      	Using("content").WithDefault("Hello World")
//      	Using("writer").From(
//
//    	// Create a server
//    	web.Serve(reg, router, cxt)
//    }
//
func Serve(reg *cookoo.Registry, router *cookoo.Router, cxt cookoo.Context) {

	addr := cxt.Get("server.Address", ":8080").(string)

	handler := NewCookooHandler(reg, router, cxt)
	http.Handle("/", handler)
	http.ListenAndServe(addr, nil)
}

// The handler for Cookoo.
// You way use this handler in your own web apps, or you can use
// the Serve() function to create and manage a handler for you.
type CookooHandler struct {
	Registry *cookoo.Registry
	Router *cookoo.Router
	BaseContext cookoo.Context
}

// Create a new Cookoo HTTP handler.
//
// This will create an HTTP hanlder, but will not automatically attach it to a server. Implementors
// can take the handler and attach it to an existing HTTP server wiht http.HandleFunc() or
// http.ListenAndServe().
//
// For simple web servers, using this package's Serve() function may be the easier route.
//
// Important details:
//
// - A URIPathResolver is used for resolving request names.
// - The following datasources are added to the Context:
//   * url: A URLDatasource (Provides access to parts of the URL)
//   * path: A PathDatasource (Provides access to parts of a path. E.g. "/foo/bar")
//   * query: A QueryParameterDatasource (Provides access to URL query parameters.)
//   * post: A FormValuesDatasource (Provides access to form data or the body of a request.)
// - The following context variables are set:
//   * http.Request: A pointer to the http.Request object
//   * http.ResponseWriter: The response writer.
//   * server.Address: The server's address and port (NOT ALWAYS PRESENT)
func NewCookooHandler(reg *cookoo.Registry, router *cookoo.Router, cxt cookoo.Context) *CookooHandler {
	handler := new(CookooHandler)
	handler.Registry = reg
	handler.Router = router
	handler.BaseContext = cxt

	// Use the URI oriented request resolver in this package.
	resolver := new(URIPathResolver)
	resolver.Init(reg)
	router.SetRequestResolver(resolver)

	return handler
}

// Adds the built-in HTTP-specific datasources.
func (h *CookooHandler) addDatasources(cxt cookoo.Context, req *http.Request) {
	parsedURL := req.URL
	urlDS := new(URLDatasource).Init(parsedURL)
	queryDS := new(QueryParameterDatasource).Init(parsedURL.Query())
	formDS := new(FormValuesDatasource).Init(req)
	pathDS := new(PathDatasource).Init(parsedURL.Path)

	cxt.AddDatasource("url", urlDS)
	cxt.AddDatasource("query", queryDS)
	// cxt.AddDatasource("q", queryDS)
	cxt.AddDatasource("post", formDS)
	cxt.AddDatasource("path", pathDS)
}

// The Cookoo request handling function.
func (h *CookooHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	// Trap panics and make them 500 errors:
	defer func() {
		// fmt.Printf("Deferred function executed for path %s\n", req.URL.Path)
		if err := recover(); err != nil {
			fmt.Printf("FOUND ERROR: %v", err)
			http.Error(res, "An internal error occurred.", http.StatusInternalServerError)
		}
	}()
	// First we need to clone the context so we have a mutable copy.
	cxt := h.BaseContext.Copy()

	cxt.Add("http.Request", req)
	cxt.Add("http.ResponseWriter", res)

	// Next, we add the datasources for URL and Query params.
	h.addDatasources(cxt, req)

	// Find the route
	path := req.Method + " " + req.URL.Path;

	fmt.Printf("Handling request for %s\n", path)

	// If a route matches, run it.
	if h.Router.HasRoute(path) {
		h.Router.HandleRequest(path, cxt, true)

	// Else if there is a custom 404 handler, run it.
	} else if h.Router.HasRoute("@404") {
		// Taint mode is false for error paths.
		h.Router.HandleRequest("@404", cxt, false)

	// Else run the default 404 handler.
	} else {
		http.NotFound(res, req)
	}
}

<commit_msg>Web server now logs errors.<commit_after>package web

import (
	"github.com/masterminds/cookoo"
	//cookoo "../"
	"net/http"
	"fmt"
	"log"
)

// Create a new Cookoo web server.
//
// Important details:
//
// - A URIPathResolver is used for resolving request names.
// - The following datasources are added to the Context:
//   * url: A URLDatasource (Provides access to parts of the URL)
//   * path: A PathDatasource (Provides access to parts of a path. E.g. "/foo/bar")
//   * query: A QueryParameterDatasource (Provides access to URL query parameters.)
//   * post: A FormValuesDatasource (Provides access to form data or the body of a request.)
// - The following context variables are set:
//   * http.Request: A pointer to the http.Request object
//   * http.ResponseWriter: The response writer.
//   * server.Address: The server's address and port (NOT ALWAYS PRESENT)
// - The handler includes logic to redirect "not found" errors to a path named "@404" if present.
//
// Context Params:
//
// - server.Address: If this key exists in the context, it will be used to determine the host/port the
//   server runes on. EXPERIMENTAL. Default is ":8080".
//
// Example:
//
//    package main
//
//    import (
//      //This is the path to Cookoo
//      "github.com/masterminds/cookoo/src/cookoo"
//      "github.com/masterminds/cookoo/src/cookoo/web"
//      "fmt"
//    )
//
//    func main() {
//      // Build a new Cookoo app.
//      handler, registry, router, context := cookoo.Cookoo()
//
//      // Fill the registry.
//      registry.Route("GET /", "The index").Does(web.Flush, "example").
//      	Using("content").WithDefault("Hello World")
//      	Using("writer").From(
//
//    	// Create a server
//    	web.Serve(reg, router, cxt)
//    }
//
func Serve(reg *cookoo.Registry, router *cookoo.Router, cxt cookoo.Context) {

	addr := cxt.Get("server.Address", ":8080").(string)

	handler := NewCookooHandler(reg, router, cxt)
	http.Handle("/", handler)
	http.ListenAndServe(addr, nil)
}

// The handler for Cookoo.
// You way use this handler in your own web apps, or you can use
// the Serve() function to create and manage a handler for you.
type CookooHandler struct {
	Registry *cookoo.Registry
	Router *cookoo.Router
	BaseContext cookoo.Context
}

// Create a new Cookoo HTTP handler.
//
// This will create an HTTP hanlder, but will not automatically attach it to a server. Implementors
// can take the handler and attach it to an existing HTTP server wiht http.HandleFunc() or
// http.ListenAndServe().
//
// For simple web servers, using this package's Serve() function may be the easier route.
//
// Important details:
//
// - A URIPathResolver is used for resolving request names.
// - The following datasources are added to the Context:
//   * url: A URLDatasource (Provides access to parts of the URL)
//   * path: A PathDatasource (Provides access to parts of a path. E.g. "/foo/bar")
//   * query: A QueryParameterDatasource (Provides access to URL query parameters.)
//   * post: A FormValuesDatasource (Provides access to form data or the body of a request.)
// - The following context variables are set:
//   * http.Request: A pointer to the http.Request object
//   * http.ResponseWriter: The response writer.
//   * server.Address: The server's address and port (NOT ALWAYS PRESENT)
func NewCookooHandler(reg *cookoo.Registry, router *cookoo.Router, cxt cookoo.Context) *CookooHandler {
	handler := new(CookooHandler)
	handler.Registry = reg
	handler.Router = router
	handler.BaseContext = cxt

	// Use the URI oriented request resolver in this package.
	resolver := new(URIPathResolver)
	resolver.Init(reg)
	router.SetRequestResolver(resolver)

	return handler
}

// Adds the built-in HTTP-specific datasources.
func (h *CookooHandler) addDatasources(cxt cookoo.Context, req *http.Request) {
	parsedURL := req.URL
	urlDS := new(URLDatasource).Init(parsedURL)
	queryDS := new(QueryParameterDatasource).Init(parsedURL.Query())
	formDS := new(FormValuesDatasource).Init(req)
	pathDS := new(PathDatasource).Init(parsedURL.Path)

	cxt.AddDatasource("url", urlDS)
	cxt.AddDatasource("query", queryDS)
	// cxt.AddDatasource("q", queryDS)
	cxt.AddDatasource("post", formDS)
	cxt.AddDatasource("path", pathDS)
}

// The Cookoo request handling function.
func (h *CookooHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	// Trap panics and make them 500 errors:
	defer func() {
		// fmt.Printf("Deferred function executed for path %s\n", req.URL.Path)
		if err := recover(); err != nil {
			fmt.Printf("FOUND ERROR: %v", err)
			http.Error(res, "An internal error occurred.", http.StatusInternalServerError)
		}
	}()
	// First we need to clone the context so we have a mutable copy.
	cxt := h.BaseContext.Copy()

	cxt.Add("http.Request", req)
	cxt.Add("http.ResponseWriter", res)

	// Next, we add the datasources for URL and Query params.
	h.addDatasources(cxt, req)

	// Find the route
	path := req.Method + " " + req.URL.Path;

	fmt.Printf("Handling request for %s\n", path)

	// If a route matches, run it.
	if h.Router.HasRoute(path) {
		err := h.Router.HandleRequest(path, cxt, true)

		if err != nil {
			fatal, ok := err.(*cookoo.FatalError)
			if !ok {
				log.Printf("Unknown error: %v %T", err, err)
			} else {
				log.Printf("Fatal Error: %s", fatal)
			}
		}

	// Else if there is a custom 404 handler, run it.
	} else if h.Router.HasRoute("@404") {
		// Taint mode is false for error paths.
		h.Router.HandleRequest("@404", cxt, false)

	// Else run the default 404 handler.
	} else {
		http.NotFound(res, req)
	}
}

<|endoftext|>


<commit_before>package jpath

import (
	"fmt"
	. "github.com/pelletier/go-toml"
	"testing"
)

func assertQuery(t *testing.T, toml, query string, ref []interface{}) {
	tree, err := Load(toml)
	if err != nil {
		t.Errorf("Non-nil toml parse error: %v", err)
		return
	}
	_, flow := lex(query)
	if err != nil {
		t.Errorf("Non-nil query lex error: %v", err)
		return
	}
	path := parse(flow)
	result := processPath(path, tree)
	assertValue(t, result, ref, "")
}

func assertValue(t *testing.T, result, ref interface{}, location string) {
	switch node := ref.(type) {
	case []interface{}:
		if resultNode, ok := result.([]interface{}); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
      if len(node) != len(resultNode) {
        t.Errorf("{%s} lengths do not match: %v vs %v",
          location, node, resultNode)
      } else {
        for i, v := range node {
          assertValue(t, resultNode[i], v, fmt.Sprintf("%s[%d]", location, i))
        }
      }
		}
	case map[string]interface{}:
		if resultNode, ok := result.(*TomlTree); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
			for k, v := range node {
				assertValue(t, resultNode.GetPath([]string{k}), v, location+"."+k)
			}
		}
	case int64:
		if resultNode, ok := result.(int64); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
			if node != resultNode {
				t.Errorf("{%s} result value does not match", location)
			}
		}
	case string:
		if resultNode, ok := result.(string); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
			if node != resultNode {
				t.Errorf("{%s} result value does not match", location)
			}
		}
	default:
		if fmt.Sprintf("%v", node) != fmt.Sprintf("%v", ref) {
			t.Errorf("{%s} result value does not match: %v != %v",
				location, node, ref)
		}
	}
}

func TestQueryRoot(t *testing.T) {
	assertQuery(t,
		"a = 42",
		"$",
		[]interface{}{
			map[string]interface{}{
				"a": int64(42),
			},
		})
}

/*
// NOTE: from the BurntSushi test suite
// NOTE: this test is pure evil due to the embedded '.'
func TestSpecialKV(t *testing.T) {
	tree, err := Load("~!@#$^&*()_+-`1234567890[]\\|/?><.,;: = 1")
	assertTree(t, tree, err, map[string]interface{}{
		"~!@#$^&*()_+-`1234567890[]\\|/?><.,;:": int64(1),
	})
}

func TestSimpleNumbers(t *testing.T) {
	tree, err := Load("a = +42\nb = -21\nc = +4.2\nd = -2.1")
	assertTree(t, tree, err, map[string]interface{}{
		"a": int64(42),
		"b": int64(-21),
		"c": float64(4.2),
		"d": float64(-2.1),
	})
}

func TestSimpleDate(t *testing.T) {
	tree, err := Load("a = 1979-05-27T07:32:00Z")
	assertTree(t, tree, err, map[string]interface{}{
		"a": time.Date(1979, time.May, 27, 7, 32, 0, 0, time.UTC),
	})
}

func TestSimpleString(t *testing.T) {
	tree, err := Load("a = \"hello world\"")
	assertTree(t, tree, err, map[string]interface{}{
		"a": "hello world",
	})
}

func TestStringEscapables(t *testing.T) {
	tree, err := Load("a = \"a \\n b\"")
	assertTree(t, tree, err, map[string]interface{}{
		"a": "a \n b",
	})

	tree, err = Load("a = \"a \\t b\"")
	assertTree(t, tree, err, map[string]interface{}{
		"a": "a \t b",
	})

	tree, err = Load("a = \"a \\r b\"")
	assertTree(t, tree, err, map[string]interface{}{
		"a": "a \r b",
	})

	tree, err = Load("a = \"a \\\\ b\"")
	assertTree(t, tree, err, map[string]interface{}{
		"a": "a \\ b",
	})
}

func TestBools(t *testing.T) {
	tree, err := Load("a = true\nb = false")
	assertTree(t, tree, err, map[string]interface{}{
		"a": true,
		"b": false,
	})
}

func TestNestedKeys(t *testing.T) {
	tree, err := Load("[a.b.c]\nd = 42")
	assertTree(t, tree, err, map[string]interface{}{
		"a": map[string]interface{}{
			"b": map[string]interface{}{
				"c": map[string]interface{}{
					"d": int64(42),
				},
			},
		},
	})
}

func TestArrayOne(t *testing.T) {
	tree, err := Load("a = [1]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []int64{int64(1)},
	})
}

func TestArrayZero(t *testing.T) {
	tree, err := Load("a = []")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []interface{}{},
	})
}

func TestArraySimple(t *testing.T) {
	tree, err := Load("a = [42, 21, 10]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []int64{int64(42), int64(21), int64(10)},
	})

	tree, _ = Load("a = [42, 21, 10,]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []int64{int64(42), int64(21), int64(10)},
	})
}

func TestArrayMultiline(t *testing.T) {
	tree, err := Load("a = [42,\n21, 10,]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []int64{int64(42), int64(21), int64(10)},
	})
}

func TestArrayNested(t *testing.T) {
	tree, err := Load("a = [[42, 21], [10]]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": [][]int64{[]int64{int64(42), int64(21)}, []int64{int64(10)}},
	})
}

func TestNestedEmptyArrays(t *testing.T) {
	tree, err := Load("a = [[[]]]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": [][][]interface{}{[][]interface{}{[]interface{}{}}},
	})
}

func TestArrayMixedTypes(t *testing.T) {
	_, err := Load("a = [42, 16.0]")
	if err.Error() != "(1, 10): mixed types in array" {
		t.Error("Bad error message:", err.Error())
	}

	_, err = Load("a = [42, \"hello\"]")
	if err.Error() != "(1, 11): mixed types in array" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestArrayNestedStrings(t *testing.T) {
	tree, err := Load("data = [ [\"gamma\", \"delta\"], [\"Foo\"] ]")
	assertTree(t, tree, err, map[string]interface{}{
		"data": [][]string{[]string{"gamma", "delta"}, []string{"Foo"}},
	})
}

func TestMissingValue(t *testing.T) {
	_, err := Load("a = ")
	if err.Error() != "(1, 4): expecting a value" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestUnterminatedArray(t *testing.T) {
	_, err := Load("a = [1,")
	if err.Error() != "(1, 8): unterminated array" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestNewlinesInArrays(t *testing.T) {
	tree, err := Load("a = [1,\n2,\n3]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []int64{int64(1), int64(2), int64(3)},
	})
}

func TestArrayWithExtraComma(t *testing.T) {
	tree, err := Load("a = [1,\n2,\n3,\n]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []int64{int64(1), int64(2), int64(3)},
	})
}

func TestArrayWithExtraCommaComment(t *testing.T) {
	tree, err := Load("a = [1, # wow\n2, # such items\n3, # so array\n]")
	assertTree(t, tree, err, map[string]interface{}{
		"a": []int64{int64(1), int64(2), int64(3)},
	})
}

func TestDuplicateGroups(t *testing.T) {
	_, err := Load("[foo]\na=2\n[foo]b=3")
	if err.Error() != "(3, 2): duplicated tables" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestDuplicateKeys(t *testing.T) {
	_, err := Load("foo = 2\nfoo = 3")
	if err.Error() != "(2, 1): The following key was defined twice: foo" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestEmptyIntermediateTable(t *testing.T) {
	_, err := Load("[foo..bar]")
	if err.Error() != "(1, 2): empty intermediate table" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestImplicitDeclarationBefore(t *testing.T) {
	tree, err := Load("[a.b.c]\nanswer = 42\n[a]\nbetter = 43")
	assertTree(t, tree, err, map[string]interface{}{
		"a": map[string]interface{}{
			"b": map[string]interface{}{
				"c": map[string]interface{}{
					"answer": int64(42),
				},
			},
			"better": int64(43),
		},
	})
}

func TestFloatsWithoutLeadingZeros(t *testing.T) {
	_, err := Load("a = .42")
	if err.Error() != "(1, 4): cannot start float with a dot" {
		t.Error("Bad error message:", err.Error())
	}

	_, err = Load("a = -.42")
	if err.Error() != "(1, 5): cannot start float with a dot" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestMissingFile(t *testing.T) {
	_, err := LoadFile("foo.toml")
	if err.Error() != "open foo.toml: no such file or directory" {
		t.Error("Bad error message:", err.Error())
	}
}

func TestParseFile(t *testing.T) {
	tree, err := LoadFile("example.toml")

	assertTree(t, tree, err, map[string]interface{}{
		"title": "TOML Example",
		"owner": map[string]interface{}{
			"name":         "Tom Preston-Werner",
			"organization": "GitHub",
			"bio":          "GitHub Cofounder & CEO\nLikes tater tots and beer.",
			"dob":          time.Date(1979, time.May, 27, 7, 32, 0, 0, time.UTC),
		},
		"database": map[string]interface{}{
			"server":         "192.168.1.1",
			"ports":          []int64{8001, 8001, 8002},
			"connection_max": 5000,
			"enabled":        true,
		},
		"servers": map[string]interface{}{
			"alpha": map[string]interface{}{
				"ip": "10.0.0.1",
				"dc": "eqdc10",
			},
			"beta": map[string]interface{}{
				"ip": "10.0.0.2",
				"dc": "eqdc10",
			},
		},
		"clients": map[string]interface{}{
			"data": []interface{}{
				[]string{"gamma", "delta"},
				[]int64{1, 2},
			},
		},
	})
}

func TestParseKeyGroupArray(t *testing.T) {
	tree, err := Load("[[foo.bar]] a = 42\n[[foo.bar]] a = 69")
	assertTree(t, tree, err, map[string]interface{}{
		"foo": map[string]interface{}{
			"bar": []map[string]interface{}{
				{"a": int64(42)},
				{"a": int64(69)},
			},
		},
	})
}

func TestParseKeyGroupArraySpec(t *testing.T) {
	tree, err := Load("[[fruit]]\n name=\"apple\"\n [fruit.physical]\n color=\"red\"\n shape=\"round\"\n [[fruit]]\n name=\"banana\"")
	assertTree(t, tree, err, map[string]interface{}{
		"fruit": []map[string]interface{}{
			{"name": "apple", "physical": map[string]interface{}{"color": "red", "shape": "round"}},
			{"name": "banana"},
		},
	})
}

func TestToTomlValue(t *testing.T) {
	for idx, item := range []struct {
		Value  interface{}
		Expect string
	}{
		{int64(12345), "12345"},
		{float64(123.45), "123.45"},
		{bool(true), "true"},
		{"hello world", "\"hello world\""},
		{"\b\t\n\f\r\"\\", "\"\\b\\t\\n\\f\\r\\\"\\\\\""},
		{"\x05", "\"\\u0005\""},
		{time.Date(1979, time.May, 27, 7, 32, 0, 0, time.UTC),
			"1979-05-27T07:32:00Z"},
		{[]interface{}{"gamma", "delta"},
			"[\n  \"gamma\",\n  \"delta\",\n]"},
	} {
		result := toTomlValue(item.Value, 0)
		if result != item.Expect {
			t.Errorf("Test %d - got '%s', expected '%s'", idx, result, item.Expect)
		}
	}
}

func TestToString(t *testing.T) {
	tree, err := Load("[foo]\n\n[[foo.bar]]\na = 42\n\n[[foo.bar]]\na = 69\n")
	if err != nil {
		t.Errorf("Test failed to parse: %v", err)
		return
	}
	result := tree.ToString()
	expected := "\n[foo]\n\n  [[foo.bar]]\n    a = 42\n\n  [[foo.bar]]\n    a = 69\n"
	if result != expected {
		t.Errorf("Expected got '%s', expected '%s'", result, expected)
	}
}

func assertPosition(t *testing.T, text string, ref map[string]Position) {
	tree, err := Load(text)
	if err != nil {
		t.Errorf("Error loading document text: `%v`", text)
		t.Errorf("Error: %v", err)
	}
	for path, pos := range ref {
		testPos := tree.GetPosition(path)
		if testPos.Invalid() {
			t.Errorf("Failed to query tree path: %s", path)
		} else if pos != testPos {
			t.Errorf("Expected position %v, got %v instead", pos, testPos)
		}
	}
}

func TestDocumentPositions(t *testing.T) {
	assertPosition(t,
		"[foo]\nbar=42\nbaz=69",
		map[string]Position{
			"foo":     Position{1, 1},
			"foo.bar": Position{2, 1},
			"foo.baz": Position{3, 1},
		})
}

func TestDocumentPositionsWithSpaces(t *testing.T) {
	assertPosition(t,
		"  [foo]\n  bar=42\n  baz=69",
		map[string]Position{
			"foo":     Position{1, 3},
			"foo.bar": Position{2, 3},
			"foo.baz": Position{3, 3},
		})
}

func TestDocumentPositionsWithGroupArray(t *testing.T) {
	assertPosition(t,
		"[[foo]]\nbar=42\nbaz=69",
		map[string]Position{
			"foo":     Position{1, 1},
			"foo.bar": Position{2, 1},
			"foo.baz": Position{3, 1},
		})
}

func TestDocumentPositionsEmptyPath(t *testing.T) {
	text := "[foo]\nbar=42\nbaz=69"
	tree, err := Load(text)
	if err != nil {
		t.Errorf("Error loading document text: `%v`", text)
		t.Errorf("Error: %v", err)
	}
	if pos := tree.GetPosition(""); !pos.Invalid() {
		t.Errorf("Valid position was returned for empty path")
	}
}
*/
<commit_msg>Added additional parser tests<commit_after>package jpath

import (
	"fmt"
	. "github.com/pelletier/go-toml"
	"testing"
)

func assertQuery(t *testing.T, toml, query string, ref []interface{}) {
	tree, err := Load(toml)
	if err != nil {
		t.Errorf("Non-nil toml parse error: %v", err)
		return
	}
	_, flow := lex(query)
	if err != nil {
		t.Errorf("Non-nil query lex error: %v", err)
		return
	}
	path := parse(flow)
	result := processPath(path, tree)
	assertValue(t, result, ref, "((" + query + ")) -> ")
}

func assertValue(t *testing.T, result, ref interface{}, location string) {
	switch node := ref.(type) {
	case []interface{}:
		if resultNode, ok := result.([]interface{}); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
      if len(node) != len(resultNode) {
        t.Errorf("{%s} lengths do not match: %v vs %v",
          location, node, resultNode)
      } else {
        for i, v := range node {
          assertValue(t, resultNode[i], v, fmt.Sprintf("%s[%d]", location, i))
        }
      }
		}
	case map[string]interface{}:
		if resultNode, ok := result.(*TomlTree); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
			for k, v := range node {
				assertValue(t, resultNode.GetPath([]string{k}), v, location+"."+k)
			}
		}
	case int64:
		if resultNode, ok := result.(int64); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
			if node != resultNode {
				t.Errorf("{%s} result value does not match", location)
			}
		}
	case string:
		if resultNode, ok := result.(string); !ok {
			t.Errorf("{%s} result value not of type %T: %T",
				location, node, resultNode)
		} else {
			if node != resultNode {
				t.Errorf("{%s} result value does not match", location)
			}
		}
	default:
		if fmt.Sprintf("%v", node) != fmt.Sprintf("%v", ref) {
			t.Errorf("{%s} result value does not match: %v != %v",
				location, node, ref)
		}
	}
}

func TestQueryRoot(t *testing.T) {
	assertQuery(t,
		"a = 42",
		"$",
		[]interface{}{
			map[string]interface{}{
				"a": int64(42),
			},
		})
}

func TestQueryKey(t *testing.T) {
  assertQuery(t,
    "[foo]\na = 42",
    "$.foo.a",
		[]interface{}{
		  int64(42),
	  })
}

func TestQueryKeyString(t *testing.T) {
  assertQuery(t,
    "[foo]\na = 42",
    "$.foo['a']",
		[]interface{}{
		  int64(42),
	  })
}

func TestQueryIndex(t *testing.T) {
  assertQuery(t,
    "[foo]\na = [1,2,3,4,5,6,7,8,9,0]",
    "$.foo.a[0]",
		[]interface{}{
		  int64(1),
	  })
}

func TestQuerySliceRange(t *testing.T) {
  assertQuery(t,
    "[foo]\na = [1,2,3,4,5,6,7,8,9,0]",
    "$.foo.a[0:5]",
		[]interface{}{
		  int64(1),
		  int64(2),
		  int64(3),
		  int64(4),
		  int64(5),
	  })
}

func TestQuerySliceStep(t *testing.T) {
  assertQuery(t,
    "[foo]\na = [1,2,3,4,5,6,7,8,9,0]",
    "$.foo.a[0:5:2]",
		[]interface{}{
		  int64(1),
		  int64(3),
		  int64(5),
	  })
}

func TestQueryAny(t *testing.T) {
  assertQuery(t,
    "[foo.bar]\na=1\nb=2\n[foo.baz]\na=3\nb=4",
    "$.foo.*",
		[]interface{}{
		  map[string]interface{}{
		    "a": int64(1),
		    "b": int64(2),
      },
		  map[string]interface{}{
		    "a": int64(3),
		    "b": int64(4),
      },
	  })
}
func TestQueryUnionSimple(t *testing.T) {
  assertQuery(t,
    "[foo.bar]\na=1\nb=2\n[baz.foo]\na=3\nb=4\n[gorf.foo]\na=5\nb=6",
    "$.*[bar,foo]",
		[]interface{}{
		  map[string]interface{}{
		    "a": int64(1),
		    "b": int64(2),
      },
		  map[string]interface{}{
		    "a": int64(3),
		    "b": int64(4),
      },
		  map[string]interface{}{
		    "a": int64(5),
		    "b": int64(6),
      },
	  })
}

func TestQueryRecursionAll(t *testing.T) {
  assertQuery(t,
    "[foo.bar]\na=1\nb=2\n[baz.foo]\na=3\nb=4\n[gorf.foo]\na=5\nb=6",
    "$..*",
		[]interface{}{
		  map[string]interface{}{
		    "bar": map[string]interface{}{
          "a": int64(1),
          "b": int64(2),
        },
      },
		  map[string]interface{}{
		    "a": int64(1),
		    "b": int64(2),
      },
		  int64(1),
		  int64(2),
		  map[string]interface{}{
        "foo": map[string]interface{}{
          "a": int64(3),
          "b": int64(4),
        },
      },
		  map[string]interface{}{
		    "a": int64(3),
		    "b": int64(4),
      },
		  int64(3),
		  int64(4),
		  map[string]interface{}{
        "foo": map[string]interface{}{
          "a": int64(5),
          "b": int64(6),
        },
      },
		  map[string]interface{}{
		    "a": int64(5),
		    "b": int64(6),
      },
		  int64(5),
		  int64(6),
	  })
}

func TestQueryRecursionUnionSimple(t *testing.T) {
  assertQuery(t,
    "[foo.bar]\na=1\nb=2\n[baz.foo]\na=3\nb=4\n[gorf.foo]\na=5\nb=6",
    "$..['foo','bar']",
		[]interface{}{
		  map[string]interface{}{
		    "a": int64(1),
		    "b": int64(2),
      },
		  map[string]interface{}{
		    "a": int64(3),
		    "b": int64(4),
      },
		  map[string]interface{}{
		    "a": int64(5),
		    "b": int64(6),
      },
	  })
}
<|endoftext|>


<commit_before>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"testing"

	"k8s.io/registry.k8s.io/pkg/net/cidrs/aws"
)

func TestRegionToAWSRegionToS3URL(t *testing.T) {
	// ensure all known regions return a configured bucket
	regions := aws.Regions()
	for region := range regions {
		bucket := awsRegionToS3URL(region)
		if bucket == "" {
			t.Fatalf("received empty string for known region %q bucket", region)
		}
	}
}
<commit_msg>enhance test coverage back to 100% for app package<commit_after>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"testing"

	"k8s.io/registry.k8s.io/pkg/net/cidrs/aws"
)

func TestRegionToAWSRegionToS3URL(t *testing.T) {
	// ensure all known regions return a configured bucket
	regions := aws.Regions()
	for region := range regions {
		url := awsRegionToS3URL(region)
		if url == "" {
			t.Fatalf("received empty string for known region %q url", region)
		}
	}
	// ensure bogus region would return "" so we know above test is valid
	if url := awsRegionToS3URL("nonsensical-region"); url != "" {
		t.Fatalf("received non-empty URL string for made up region \"nonsensical-region\": %q", url)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"container/heap"
	"context"
	"encoding/json"
	"expvar"
	"flag"
	"fmt"
	"github.com/coreos/go-systemd/daemon"
	"github.com/gorilla/websocket"
	"github.com/omegaup/quark/broadcaster"
	"github.com/omegaup/quark/common"
	"github.com/prometheus/client_golang/prometheus"
	"math"
	"net/http"
	_ "net/http/pprof"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

var (
	version    = flag.Bool("version", false, "Print the version and exit")
	insecure   = flag.Bool("insecure", false, "Do not use TLS")
	configPath = flag.String(
		"config",
		"/etc/omegaup/broadcaster/config.json",
		"Grader configuration file",
	)
	globalContext atomic.Value
	upgrader      = websocket.Upgrader{
		Subprotocols: []string{"com.omegaup.events"},
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	// ProgramVersion is the version of the code from which the binary was built from.
	ProgramVersion string
)

func loadContext() error {
	f, err := os.Open(*configPath)
	if err != nil {
		return err
	}
	defer f.Close()

	ctx, err := common.NewContextFromReader(f, "broadcaster")
	if err != nil {
		return err
	}
	globalContext.Store(ctx)
	return nil
}

func broadcasterContext() *common.Context {
	return globalContext.Load().(*common.Context)
}

func mustParseURL(rawurl string, relative ...string) *url.URL {
	parsed, err := url.Parse(rawurl)
	if err != nil {
		panic(err)
	}
	for _, rel := range relative {
		parsed, err = parsed.Parse(rel)
		if err != nil {
			panic(err)
		}
	}
	return parsed
}

type updateScoreboardEvent struct {
	contestAlias string
	deadline     time.Time
}

type updateScoreboardEventHeap []*updateScoreboardEvent

func (h updateScoreboardEventHeap) Len() int      { return len(h) }
func (h updateScoreboardEventHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h updateScoreboardEventHeap) Less(i, j int) bool {
	return h[i].deadline.Before(h[j].deadline)
}

func (h *updateScoreboardEventHeap) Push(x interface{}) {
	*h = append(*h, x.(*updateScoreboardEvent))
}

func (h *updateScoreboardEventHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func updateScoreboardForContest(
	ctx *common.Context,
	client *http.Client,
	updateScoreboardURL *url.URL,
	contestAlias string,
) {
	ctx.Log.Info("Requesting scoreboard update", "contest", contestAlias)
	resp, err := client.PostForm(
		updateScoreboardURL.String(),
		url.Values{
			"token": {ctx.Config.Broadcaster.ScoreboardUpdateSecret},
			"alias": {contestAlias},
		},
	)
	if err != nil {
		ctx.Log.Error(
			"Error requesting scoreboard update",
			"contest", contestAlias,
			"err", err,
		)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		ctx.Log.Error(
			"Failed to request scoreboard update",
			"contest", contestAlias,
			"status code", resp.StatusCode,
		)
	}
}

func updateScoreboardLoop(
	ctx *common.Context,
	client *http.Client,
	updateScoreboardURL *url.URL,
	contestChan <-chan string,
) {
	const infinity = time.Duration(math.MaxInt64)
	timer := time.NewTimer(infinity)
	events := updateScoreboardEventHeap{}
	eventSet := make(map[string]bool)

	for {
		select {
		case contestAlias := <-contestChan:
			if _, ok := eventSet[contestAlias]; ok {
				eventSet[contestAlias] = true
				continue
			}

			eventSet[contestAlias] = false
			heap.Push(&events, &updateScoreboardEvent{
				contestAlias: contestAlias,
				deadline:     time.Now().Add(time.Duration(ctx.Config.Broadcaster.ScoreboardUpdateTimeout)),
			})
			if len(events) == 1 {
				if !timer.Stop() {
					<-timer.C
				}
				timer.Reset(time.Duration(ctx.Config.Broadcaster.ScoreboardUpdateTimeout))
			}

			updateScoreboardForContest(
				ctx,
				client,
				updateScoreboardURL,
				contestAlias,
			)
		case <-timer.C:
			if len(events) == 0 {
				timer.Reset(infinity)
				continue
			}
			event := heap.Pop(&events).(*updateScoreboardEvent)
			if len(events) == 0 {
				timer.Reset(infinity)
			} else {
				timer.Reset(events[0].deadline.Sub(time.Now()))
			}

			if eventSet[event.contestAlias] {
				updateScoreboardForContest(
					ctx,
					client,
					updateScoreboardURL,
					event.contestAlias,
				)
			}
			delete(eventSet, event.contestAlias)
		}
	}
}

func main() {
	flag.Parse()

	if *version {
		fmt.Printf("omegaup-broadcaster %s\n", ProgramVersion)
		return
	}

	stopChan := make(chan os.Signal)
	signal.Notify(stopChan, syscall.SIGINT, syscall.SIGTERM)

	if err := loadContext(); err != nil {
		panic(err)
	}

	ctx := broadcasterContext()
	expvar.Publish("config", &ctx.Config)

	b := broadcaster.NewBroadcaster(ctx, &PrometheusMetrics{})
	contestChan := make(chan string, 1)

	client := http.Client{}

	metricsMux := http.NewServeMux()
	metricsMux.Handle("/metrics", prometheus.Handler())
	go func() {
		addr := fmt.Sprintf(":%d", ctx.Config.Metrics.Port)
		ctx.Log.Error(
			"http listen and serve",
			"err", http.ListenAndServe(addr, metricsMux),
		)
	}()

	http.HandleFunc("/deauthenticate/", func(w http.ResponseWriter, r *http.Request) {
		pathComponents := strings.Split(r.URL.Path, "/")
		if len(pathComponents) < 3 {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		b.Deauthenticate(pathComponents[2])
		w.WriteHeader(http.StatusOK)
	})
	http.HandleFunc("/broadcast/", func(w http.ResponseWriter, r *http.Request) {
		if *insecure {
			w.Header().Set("Access-Control-Methods", "POST")
			w.Header().Set("Access-Control-Allow-Origin", "*")
		}
		if r.Method != "POST" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		defer r.Body.Close()
		var message broadcaster.Message
		decoder := json.NewDecoder(r.Body)
		if err := decoder.Decode(&message); err != nil {
			ctx.Log.Error("Error decoding broadcast message", "err", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		ctx.Log.Debug("/broadcast/", "message", message)
		if !b.Broadcast(&message) {
			ctx.Log.Error("Error sending message, queue too large")
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		// TODO(lhchavez): Figure out a better way of checking this.
		if len(message.Contest) > 0 && strings.Contains(message.Message, "\"message\":\"/run/update/\"") {
			contestChan <- message.Contest
		}
		w.WriteHeader(http.StatusOK)
		w.Header().Set("Content-Type", "text/json; charset=utf-8")
		w.Write([]byte("{\"status\":\"ok\"}"))
	})

	eventsMux := http.NewServeMux()
	eventsMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		ctx := broadcasterContext()

		authToken := ""
		if ouat, _ := r.Cookie("ouat"); ouat != nil {
			authToken = ouat.Value
		}

		var transport broadcaster.Transport

		if common.AcceptsMimeType(r, "text/event-stream") {
			transport = broadcaster.NewSSETransport(w)
		} else {
			conn, err := upgrader.Upgrade(w, r, nil)
			if err != nil {
				ctx.Log.Error("Failed to upgrade connection", "err", err)
				return
			}
			defer conn.Close()

			transport = broadcaster.NewWebSocketTransport(
				conn,
				time.Duration(ctx.Config.Broadcaster.WriteDeadline),
			)
		}

		subscriber, err := broadcaster.NewSubscriber(
			ctx,
			&client,
			mustParseURL(
				ctx.Config.Broadcaster.FrontendURL,
				"api/user/validateFilter/",
			),
			authToken,
			strings.Join(r.URL.Query()["filter"], ","),
			transport,
		)
		if err != nil {
			ctx.Log.Error("Failed to create subscriber", "err", err)
			if upstream, ok := err.(*broadcaster.UpstreamError); ok {
				w.WriteHeader(upstream.HTTPStatusCode)
				w.Write(upstream.Contents)
			}
			return
		}
		if !b.Subscribe(subscriber) {
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		defer b.Unsubscribe(subscriber)

		subscriber.Run()
	})
	go b.Run()

	eventsHost := ""
	if ctx.Config.Broadcaster.Proxied {
		eventsHost = "localhost"
	}
	var servers []*http.Server
	var wg sync.WaitGroup
	servers = append(
		servers,
		common.RunServer(
			&ctx.Config.Broadcaster.TLS,
			eventsMux,
			&wg,
			fmt.Sprintf("%s:%d", eventsHost, ctx.Config.Broadcaster.EventsPort),
			ctx.Config.Broadcaster.Proxied,
		),
	)
	go updateScoreboardLoop(
		ctx,
		&client,
		mustParseURL(
			ctx.Config.Broadcaster.FrontendURL,
			"api/scoreboard/refresh/",
		),
		contestChan,
	)
	servers = append(
		servers,
		common.RunServer(
			&ctx.Config.TLS,
			nil,
			&wg,
			fmt.Sprintf(":%d", ctx.Config.Broadcaster.Port),
			*insecure,
		),
	)

	ctx.Log.Info(
		"omegaUp broadcaster started",
		"broadcaster port", ctx.Config.Broadcaster.Port,
		"events port", ctx.Config.Broadcaster.EventsPort,
	)
	daemon.SdNotify(false, "READY=1")

	<-stopChan

	ctx.Log.Info("Shutting down server...")
	cancelCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	for _, server := range servers {
		server.Shutdown(cancelCtx)
	}

	cancel()
	wg.Wait()

	ctx.Log.Info("Server gracefully stopped.")
}
<commit_msg>Let the broadcaster bind to the any interface<commit_after>package main

import (
	"container/heap"
	"context"
	"encoding/json"
	"expvar"
	"flag"
	"fmt"
	"github.com/coreos/go-systemd/daemon"
	"github.com/gorilla/websocket"
	"github.com/omegaup/quark/broadcaster"
	"github.com/omegaup/quark/common"
	"github.com/prometheus/client_golang/prometheus"
	"math"
	"net/http"
	_ "net/http/pprof"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

var (
	version    = flag.Bool("version", false, "Print the version and exit")
	insecure   = flag.Bool("insecure", false, "Do not use TLS")
	configPath = flag.String(
		"config",
		"/etc/omegaup/broadcaster/config.json",
		"Grader configuration file",
	)
	globalContext atomic.Value
	upgrader      = websocket.Upgrader{
		Subprotocols: []string{"com.omegaup.events"},
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	// ProgramVersion is the version of the code from which the binary was built from.
	ProgramVersion string
)

func loadContext() error {
	f, err := os.Open(*configPath)
	if err != nil {
		return err
	}
	defer f.Close()

	ctx, err := common.NewContextFromReader(f, "broadcaster")
	if err != nil {
		return err
	}
	globalContext.Store(ctx)
	return nil
}

func broadcasterContext() *common.Context {
	return globalContext.Load().(*common.Context)
}

func mustParseURL(rawurl string, relative ...string) *url.URL {
	parsed, err := url.Parse(rawurl)
	if err != nil {
		panic(err)
	}
	for _, rel := range relative {
		parsed, err = parsed.Parse(rel)
		if err != nil {
			panic(err)
		}
	}
	return parsed
}

type updateScoreboardEvent struct {
	contestAlias string
	deadline     time.Time
}

type updateScoreboardEventHeap []*updateScoreboardEvent

func (h updateScoreboardEventHeap) Len() int      { return len(h) }
func (h updateScoreboardEventHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h updateScoreboardEventHeap) Less(i, j int) bool {
	return h[i].deadline.Before(h[j].deadline)
}

func (h *updateScoreboardEventHeap) Push(x interface{}) {
	*h = append(*h, x.(*updateScoreboardEvent))
}

func (h *updateScoreboardEventHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func updateScoreboardForContest(
	ctx *common.Context,
	client *http.Client,
	updateScoreboardURL *url.URL,
	contestAlias string,
) {
	ctx.Log.Info("Requesting scoreboard update", "contest", contestAlias)
	resp, err := client.PostForm(
		updateScoreboardURL.String(),
		url.Values{
			"token": {ctx.Config.Broadcaster.ScoreboardUpdateSecret},
			"alias": {contestAlias},
		},
	)
	if err != nil {
		ctx.Log.Error(
			"Error requesting scoreboard update",
			"contest", contestAlias,
			"err", err,
		)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		ctx.Log.Error(
			"Failed to request scoreboard update",
			"contest", contestAlias,
			"status code", resp.StatusCode,
		)
	}
}

func updateScoreboardLoop(
	ctx *common.Context,
	client *http.Client,
	updateScoreboardURL *url.URL,
	contestChan <-chan string,
) {
	const infinity = time.Duration(math.MaxInt64)
	timer := time.NewTimer(infinity)
	events := updateScoreboardEventHeap{}
	eventSet := make(map[string]bool)

	for {
		select {
		case contestAlias := <-contestChan:
			if _, ok := eventSet[contestAlias]; ok {
				eventSet[contestAlias] = true
				continue
			}

			eventSet[contestAlias] = false
			heap.Push(&events, &updateScoreboardEvent{
				contestAlias: contestAlias,
				deadline:     time.Now().Add(time.Duration(ctx.Config.Broadcaster.ScoreboardUpdateTimeout)),
			})
			if len(events) == 1 {
				if !timer.Stop() {
					<-timer.C
				}
				timer.Reset(time.Duration(ctx.Config.Broadcaster.ScoreboardUpdateTimeout))
			}

			updateScoreboardForContest(
				ctx,
				client,
				updateScoreboardURL,
				contestAlias,
			)
		case <-timer.C:
			if len(events) == 0 {
				timer.Reset(infinity)
				continue
			}
			event := heap.Pop(&events).(*updateScoreboardEvent)
			if len(events) == 0 {
				timer.Reset(infinity)
			} else {
				timer.Reset(events[0].deadline.Sub(time.Now()))
			}

			if eventSet[event.contestAlias] {
				updateScoreboardForContest(
					ctx,
					client,
					updateScoreboardURL,
					event.contestAlias,
				)
			}
			delete(eventSet, event.contestAlias)
		}
	}
}

func main() {
	flag.Parse()

	if *version {
		fmt.Printf("omegaup-broadcaster %s\n", ProgramVersion)
		return
	}

	stopChan := make(chan os.Signal)
	signal.Notify(stopChan, syscall.SIGINT, syscall.SIGTERM)

	if err := loadContext(); err != nil {
		panic(err)
	}

	ctx := broadcasterContext()
	expvar.Publish("config", &ctx.Config)

	b := broadcaster.NewBroadcaster(ctx, &PrometheusMetrics{})
	contestChan := make(chan string, 1)

	client := http.Client{}

	metricsMux := http.NewServeMux()
	metricsMux.Handle("/metrics", prometheus.Handler())
	go func() {
		addr := fmt.Sprintf(":%d", ctx.Config.Metrics.Port)
		ctx.Log.Error(
			"http listen and serve",
			"err", http.ListenAndServe(addr, metricsMux),
		)
	}()

	http.HandleFunc("/deauthenticate/", func(w http.ResponseWriter, r *http.Request) {
		pathComponents := strings.Split(r.URL.Path, "/")
		if len(pathComponents) < 3 {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		b.Deauthenticate(pathComponents[2])
		w.WriteHeader(http.StatusOK)
	})
	http.HandleFunc("/broadcast/", func(w http.ResponseWriter, r *http.Request) {
		if *insecure {
			w.Header().Set("Access-Control-Methods", "POST")
			w.Header().Set("Access-Control-Allow-Origin", "*")
		}
		if r.Method != "POST" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		defer r.Body.Close()
		var message broadcaster.Message
		decoder := json.NewDecoder(r.Body)
		if err := decoder.Decode(&message); err != nil {
			ctx.Log.Error("Error decoding broadcast message", "err", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		ctx.Log.Debug("/broadcast/", "message", message)
		if !b.Broadcast(&message) {
			ctx.Log.Error("Error sending message, queue too large")
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		// TODO(lhchavez): Figure out a better way of checking this.
		if len(message.Contest) > 0 && strings.Contains(message.Message, "\"message\":\"/run/update/\"") {
			contestChan <- message.Contest
		}
		w.WriteHeader(http.StatusOK)
		w.Header().Set("Content-Type", "text/json; charset=utf-8")
		w.Write([]byte("{\"status\":\"ok\"}"))
	})

	eventsMux := http.NewServeMux()
	eventsMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		ctx := broadcasterContext()

		authToken := ""
		if ouat, _ := r.Cookie("ouat"); ouat != nil {
			authToken = ouat.Value
		}

		var transport broadcaster.Transport

		if common.AcceptsMimeType(r, "text/event-stream") {
			transport = broadcaster.NewSSETransport(w)
		} else {
			conn, err := upgrader.Upgrade(w, r, nil)
			if err != nil {
				ctx.Log.Error("Failed to upgrade connection", "err", err)
				return
			}
			defer conn.Close()

			transport = broadcaster.NewWebSocketTransport(
				conn,
				time.Duration(ctx.Config.Broadcaster.WriteDeadline),
			)
		}

		subscriber, err := broadcaster.NewSubscriber(
			ctx,
			&client,
			mustParseURL(
				ctx.Config.Broadcaster.FrontendURL,
				"api/user/validateFilter/",
			),
			authToken,
			strings.Join(r.URL.Query()["filter"], ","),
			transport,
		)
		if err != nil {
			ctx.Log.Error("Failed to create subscriber", "err", err)
			if upstream, ok := err.(*broadcaster.UpstreamError); ok {
				w.WriteHeader(upstream.HTTPStatusCode)
				w.Write(upstream.Contents)
			}
			return
		}
		if !b.Subscribe(subscriber) {
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		defer b.Unsubscribe(subscriber)

		subscriber.Run()
	})
	go b.Run()

	var servers []*http.Server
	var wg sync.WaitGroup
	servers = append(
		servers,
		common.RunServer(
			&ctx.Config.Broadcaster.TLS,
			eventsMux,
			&wg,
			fmt.Sprintf(":%d", ctx.Config.Broadcaster.EventsPort),
			ctx.Config.Broadcaster.Proxied,
		),
	)
	go updateScoreboardLoop(
		ctx,
		&client,
		mustParseURL(
			ctx.Config.Broadcaster.FrontendURL,
			"api/scoreboard/refresh/",
		),
		contestChan,
	)
	servers = append(
		servers,
		common.RunServer(
			&ctx.Config.TLS,
			nil,
			&wg,
			fmt.Sprintf(":%d", ctx.Config.Broadcaster.Port),
			*insecure,
		),
	)

	ctx.Log.Info(
		"omegaUp broadcaster started",
		"broadcaster port", ctx.Config.Broadcaster.Port,
		"events port", ctx.Config.Broadcaster.EventsPort,
	)
	daemon.SdNotify(false, "READY=1")

	<-stopChan

	ctx.Log.Info("Shutting down server...")
	cancelCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	for _, server := range servers {
		server.Shutdown(cancelCtx)
	}

	cancel()
	wg.Wait()

	ctx.Log.Info("Server gracefully stopped.")
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/golang/glog"
	gogit "gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/plumbing"

	"k8s.io/publishing-bot/cmd/publishing-bot/config"
)

// PublisherMunger publishes content from one repository to another one.
type PublisherMunger struct {
	reposRules config.RepositoryRules
	config     *config.Config
	// plog duplicates the logs at glog and a file
	plog *plog
	// absolute path to the repos.
	baseRepoPath string
}

// New will create a new munger.
func New(cfg *config.Config, baseRepoPath string) *PublisherMunger {
	// create munger
	return &PublisherMunger{
		baseRepoPath: baseRepoPath,
		config:       cfg,
	}
}

// update the local checkout of the source repository. It returns the branch heads.
func (p *PublisherMunger) updateSourceRepo() (map[string]plumbing.Hash, error) {
	repoDir := filepath.Join(p.baseRepoPath, p.config.SourceRepo)

	// fetch origin
	glog.Infof("Fetching origin at %s.", repoDir)
	r, err := gogit.PlainOpen(repoDir)
	if err != nil {
		return nil, fmt.Errorf("failed to open repo at %s: %v", repoDir, err)
	}
	if err := r.Fetch(&gogit.FetchOptions{Progress: os.Stdout}); err != nil && err != gogit.NoErrAlreadyUpToDate {
		return nil, fmt.Errorf("failed to fetch at %s: %v", repoDir, err)
	}

	// checkout head
	glog.Infof("Checking out HEAD at %s.", repoDir)
	w, err := r.Worktree()
	if err != nil {
		return nil, fmt.Errorf("failed to open worktree at %s: %v", repoDir, err)
	}
	head, err := r.Head()
	if err != nil {
		return nil, fmt.Errorf("failed to get head at %s: %v", repoDir, err)
	}
	if err := w.Checkout(&gogit.CheckoutOptions{Hash: head.Hash()}); err != nil {
		return nil, fmt.Errorf("failed to checkout HEAD at %s: %v", repoDir, err)
	}

	// create/update local branch for all origin branches. Those are fetches into the destination repos later (as upstream/<branch>).
	refs, err := r.Storer.IterReferences()
	if err != nil {
		return nil, fmt.Errorf("failed to get branches: %v", err)
	}
	glog.Infof("Updating local branches at %s.", repoDir)
	heads := map[string]plumbing.Hash{}
	if err = refs.ForEach(func(ref *plumbing.Reference) error {
		name := ref.Name().String()

		originPrefix := "refs/remotes/origin/"
		if !strings.Contains(name, originPrefix) || ref.Type() != plumbing.HashReference {
			return nil
		}

		shortName := strings.TrimPrefix(name, originPrefix)
		localBranch := plumbing.NewHashReference(plumbing.ReferenceName("refs/heads/"+shortName), ref.Hash())
		if err := r.Storer.SetReference(localBranch); err != nil {
			return fmt.Errorf("failed to create reference %s pointing to %s", localBranch.Name(), localBranch.Hash().String())
		}

		heads[localBranch.Name().String()] = localBranch.Hash()

		return nil
	}); err != nil {
		return nil, fmt.Errorf("failed to process branches: %v", err)
	}

	return heads, nil
}

// update the active rules
func (p *PublisherMunger) updateRules() error {
	repoDir := filepath.Join(p.baseRepoPath, p.config.SourceRepo)

	glog.Infof("Checking out master at %s.", repoDir)
	cmd := exec.Command("git", "checkout", "master")
	cmd.Dir = repoDir
	if _, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("failed to checkout master: %v", err)
	}

	rules, err := config.LoadRules(p.config.RulesFile)
	if err != nil {
		return err
	}
	p.reposRules = *rules
	glog.Infof("Loaded %d repository rules from %s.", len(p.reposRules.Rules), p.config.RulesFile)
	return nil
}

func (p *PublisherMunger) skippedBranch(b string) bool {
	for _, skipped := range p.reposRules.SkippedSourceBranches {
		if b == skipped {
			return true
		}
	}
	return false
}

// git clone dstURL to dst if dst doesn't exist yet.
func (p *PublisherMunger) ensureCloned(dst string, dstURL string) error {
	if _, err := os.Stat(dst); err == nil {
		return nil
	}

	cmd := exec.Command("mkdir", "-p", dst)
	if err := p.plog.Run(cmd); err != nil {
		return err
	}
	cmd = exec.Command("git", "clone", dstURL, dst)
	if err := p.plog.Run(cmd); err != nil {
		return err
	}
	cmd = exec.Command("/bin/bash", "-c", "git tag -l | xargs git tag -d")
	cmd.Dir = dst
	return p.plog.Run(cmd)
}

// constructs all the repos, but does not push the changes to remotes.
func (p *PublisherMunger) construct() error {
	sourceRemote := filepath.Join(p.baseRepoPath, p.config.SourceRepo, ".git")
	for _, repoRule := range p.reposRules.Rules {
		if repoRule.Skip {
			continue
		}

		// clone the destination repo
		dstDir := filepath.Join(p.baseRepoPath, repoRule.DestinationRepository, "")
		dstURL := fmt.Sprintf("https://%s/%s/%s.git", p.config.GithubHost, p.config.TargetOrg, repoRule.DestinationRepository)
		if err := p.ensureCloned(dstDir, dstURL); err != nil {
			p.plog.Errorf("%v", err)
			return err
		}
		p.plog.Infof("Successfully ensured %s exists", dstDir)
		if err := os.Chdir(dstDir); err != nil {
			return err
		}

		// delete tags
		cmd := exec.Command("/bin/bash", "-c", "git tag | xargs git tag -d >/dev/null")
		if err := p.plog.Run(cmd); err != nil {
			return err
		}

		formatDeps := func(deps []config.Dependency) string {
			var depStrings []string
			for _, dep := range deps {
				depStrings = append(depStrings, fmt.Sprintf("%s:%s", dep.Repository, dep.Branch))
			}
			return strings.Join(depStrings, ",")
		}

		// construct branches
		for _, branchRule := range repoRule.Branches {
			if p.skippedBranch(branchRule.Source.Branch) {
				continue
			}
			if len(branchRule.Source.Dir) == 0 {
				branchRule.Source.Dir = "."
				p.plog.Infof("%v: 'dir' cannot be empty, defaulting to '.'", branchRule)
			}

			// get old HEAD. Ignore errors as the branch might be non-existent
			oldHead, _ := exec.Command("git", "rev-parse", fmt.Sprintf("origin/%s", branchRule.Name)).Output()

			goPath := os.Getenv("GOPATH")
			branchEnv := append([]string(nil), os.Environ()...) // make mutable
			if branchRule.GoVersion != "" {
				goRoot := filepath.Join(goPath, "go-"+branchRule.GoVersion)
				branchEnv = append(branchEnv, "GOROOT="+goRoot)
				goBin := filepath.Join(goRoot, "bin")
				branchEnv = updateEnv(branchEnv, "PATH", prependPath(goBin), goBin)
			}

			skipTags := ""
			if p.reposRules.SkipTags {
				skipTags = "true"
				p.plog.Infof("synchronizing tags is disabled")
			}

			// get old published hash to eventually skip cherry picking
			var lastPublishedUpstreamHash string
			bs, err := ioutil.ReadFile(path.Join(p.baseRepoPath, publishedFileName(repoRule.DestinationRepository, branchRule.Name)))
			if err != nil && !os.IsNotExist(err) {
				return err
			}
			if err == nil {
				lastPublishedUpstreamHash = string(bs)
			}

			// TODO: Refactor this to use environment variables instead
			repoPublishScriptPath := filepath.Join(p.config.BasePublishScriptPath, "construct.sh")
			cmd := exec.Command(repoPublishScriptPath,
				repoRule.DestinationRepository,
				branchRule.Source.Branch,
				branchRule.Name,
				formatDeps(branchRule.Dependencies),
				strings.Join(branchRule.RequiredPackages, ":"),
				sourceRemote,
				branchRule.Source.Dir,
				p.config.SourceRepo,
				p.config.SourceRepo,
				p.config.BasePackage,
				fmt.Sprintf("%v", repoRule.Library),
				strings.Join(p.reposRules.RecursiveDeletePatterns, " "),
				skipTags,
				lastPublishedUpstreamHash,
			)
			cmd.Env = append([]string(nil), branchEnv...) // make mutable
			if p.reposRules.SkipGodeps {
				cmd.Env = append(cmd.Env, "PUBLISHER_BOT_SKIP_GODEPS=true")
			}
			if err := p.plog.Run(cmd); err != nil {
				return err
			}

			newHead, _ := exec.Command("git", "rev-parse", fmt.Sprintf("origin/%s", branchRule.Name)).Output()
			if len(repoRule.SmokeTest) > 0 && string(oldHead) != string(newHead) {
				p.plog.Infof("Running smoke tests for branch %s", branchRule.Name)
				cmd := exec.Command("/bin/bash", "-xec", repoRule.SmokeTest)
				cmd.Env = append([]string(nil), branchEnv...) // make mutable
				if err := p.plog.Run(cmd); err != nil {
					// do not clean up to allow debugging with kubectl-exec.
					return err
				}
				exec.Command("git", "reset", "--hard").Run()
				exec.Command("git", "clean", "-f", "-f", "-d").Run()
			}

			p.plog.Infof("Successfully constructed %s", branchRule.Name)
		}
	}
	return nil
}

func updateEnv(env []string, key string, change func(string) string, val string) []string {
	for i := range env {
		if strings.HasPrefix(env[i], key+"=") {
			ss := strings.SplitN(env[i], "=", 2)
			env[i] = fmt.Sprintf("%s=%s", key, change(ss[1]))
			return env
		}
	}
	return append(env, fmt.Sprintf("%s=%s", key, val))
}

func prependPath(p string) func(string) string {
	return func(s string) string {
		if s == "" {
			return p
		}
		return p + ":" + s
	}
}

// publish to remotes.
func (p *PublisherMunger) publish(newUpstreamHeads map[string]plumbing.Hash) error {
	if p.config.DryRun {
		p.plog.Infof("Skipping push in dry-run mode")
		return nil
	}

	if p.config.TokenFile == "" {
		return fmt.Errorf("token cannot be empty in non-dry-run mode")
	}

	// NOTE: because some repos depend on each other, e.g., client-go depends on
	// apimachinery, they should be published atomically, but it's not supported
	// by github.
	for _, repoRules := range p.reposRules.Rules {
		if repoRules.Skip {
			continue
		}

		dstDir := filepath.Join(p.baseRepoPath, repoRules.DestinationRepository, "")
		if err := os.Chdir(dstDir); err != nil {
			return err
		}
		for _, branchRule := range repoRules.Branches {
			if p.skippedBranch(branchRule.Source.Branch) {
				continue
			}

			cmd := exec.Command(p.config.BasePublishScriptPath+"/push.sh", p.config.TokenFile, branchRule.Name)
			if err := p.plog.Run(cmd); err != nil {
				return err
			}

			upstreamBranchHead, ok := newUpstreamHeads[branchRule.Source.Branch]
			if !ok {
				return fmt.Errorf("no upstream branch %q found", branchRule.Source.Branch)
			}
			if err := ioutil.WriteFile(path.Join(path.Dir(dstDir), publishedFileName(repoRules.DestinationRepository, branchRule.Name)), []byte(upstreamBranchHead.String()), 0644); err != nil {
				return err
			}
		}
	}
	return nil
}

func publishedFileName(repo, branch string) string {
	return fmt.Sprintf("published-%s-%s", repo, branch)
}

// Run constructs the repos and pushes them. It returns logs and the last master hash.
func (p *PublisherMunger) Run() (string, string, error) {
	buf := bytes.NewBuffer(nil)
	var err error
	if p.plog, err = NewPublisherLog(buf, path.Join(p.baseRepoPath, "run.log")); err != nil {
		return "", "", err
	}

	newUpstreamHeads, err := p.updateSourceRepo()
	if err != nil {
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}
	if err := p.updateRules(); err != nil { // this comes after the source update because we might fetch the rules from there.
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}
	if err := p.construct(); err != nil {
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}
	if err := p.publish(newUpstreamHeads); err != nil {
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}

	var masterHead string
	if h, ok := newUpstreamHeads["master"]; ok {
		masterHead = h.String()
	}

	return p.plog.Logs(), masterHead, nil
}
<commit_msg>Fix branch-name-hash map to be used on publish<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/golang/glog"
	gogit "gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/plumbing"

	"k8s.io/publishing-bot/cmd/publishing-bot/config"
)

// PublisherMunger publishes content from one repository to another one.
type PublisherMunger struct {
	reposRules config.RepositoryRules
	config     *config.Config
	// plog duplicates the logs at glog and a file
	plog *plog
	// absolute path to the repos.
	baseRepoPath string
}

// New will create a new munger.
func New(cfg *config.Config, baseRepoPath string) *PublisherMunger {
	// create munger
	return &PublisherMunger{
		baseRepoPath: baseRepoPath,
		config:       cfg,
	}
}

// update the local checkout of the source repository. It returns the branch heads.
func (p *PublisherMunger) updateSourceRepo() (map[string]plumbing.Hash, error) {
	repoDir := filepath.Join(p.baseRepoPath, p.config.SourceRepo)

	// fetch origin
	glog.Infof("Fetching origin at %s.", repoDir)
	r, err := gogit.PlainOpen(repoDir)
	if err != nil {
		return nil, fmt.Errorf("failed to open repo at %s: %v", repoDir, err)
	}
	if err := r.Fetch(&gogit.FetchOptions{Progress: os.Stdout}); err != nil && err != gogit.NoErrAlreadyUpToDate {
		return nil, fmt.Errorf("failed to fetch at %s: %v", repoDir, err)
	}

	// checkout head
	glog.Infof("Checking out HEAD at %s.", repoDir)
	w, err := r.Worktree()
	if err != nil {
		return nil, fmt.Errorf("failed to open worktree at %s: %v", repoDir, err)
	}
	head, err := r.Head()
	if err != nil {
		return nil, fmt.Errorf("failed to get head at %s: %v", repoDir, err)
	}
	if err := w.Checkout(&gogit.CheckoutOptions{Hash: head.Hash()}); err != nil {
		return nil, fmt.Errorf("failed to checkout HEAD at %s: %v", repoDir, err)
	}

	// create/update local branch for all origin branches. Those are fetches into the destination repos later (as upstream/<branch>).
	refs, err := r.Storer.IterReferences()
	if err != nil {
		return nil, fmt.Errorf("failed to get branches: %v", err)
	}
	glog.Infof("Updating local branches at %s.", repoDir)
	heads := map[string]plumbing.Hash{}
	if err = refs.ForEach(func(ref *plumbing.Reference) error {
		name := ref.Name().String()

		originPrefix := "refs/remotes/origin/"
		if !strings.Contains(name, originPrefix) || ref.Type() != plumbing.HashReference {
			return nil
		}

		shortName := strings.TrimPrefix(name, originPrefix)
		localBranch := plumbing.NewHashReference(plumbing.ReferenceName("refs/heads/"+shortName), ref.Hash())
		if err := r.Storer.SetReference(localBranch); err != nil {
			return fmt.Errorf("failed to create reference %s pointing to %s", localBranch.Name(), localBranch.Hash().String())
		}

		heads[shortName] = localBranch.Hash()

		return nil
	}); err != nil {
		return nil, fmt.Errorf("failed to process branches: %v", err)
	}

	return heads, nil
}

// update the active rules
func (p *PublisherMunger) updateRules() error {
	repoDir := filepath.Join(p.baseRepoPath, p.config.SourceRepo)

	glog.Infof("Checking out master at %s.", repoDir)
	cmd := exec.Command("git", "checkout", "master")
	cmd.Dir = repoDir
	if _, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("failed to checkout master: %v", err)
	}

	rules, err := config.LoadRules(p.config.RulesFile)
	if err != nil {
		return err
	}
	p.reposRules = *rules
	glog.Infof("Loaded %d repository rules from %s.", len(p.reposRules.Rules), p.config.RulesFile)
	return nil
}

func (p *PublisherMunger) skippedBranch(b string) bool {
	for _, skipped := range p.reposRules.SkippedSourceBranches {
		if b == skipped {
			return true
		}
	}
	return false
}

// git clone dstURL to dst if dst doesn't exist yet.
func (p *PublisherMunger) ensureCloned(dst string, dstURL string) error {
	if _, err := os.Stat(dst); err == nil {
		return nil
	}

	cmd := exec.Command("mkdir", "-p", dst)
	if err := p.plog.Run(cmd); err != nil {
		return err
	}
	cmd = exec.Command("git", "clone", dstURL, dst)
	if err := p.plog.Run(cmd); err != nil {
		return err
	}
	cmd = exec.Command("/bin/bash", "-c", "git tag -l | xargs git tag -d")
	cmd.Dir = dst
	return p.plog.Run(cmd)
}

// constructs all the repos, but does not push the changes to remotes.
func (p *PublisherMunger) construct() error {
	sourceRemote := filepath.Join(p.baseRepoPath, p.config.SourceRepo, ".git")
	for _, repoRule := range p.reposRules.Rules {
		if repoRule.Skip {
			continue
		}

		// clone the destination repo
		dstDir := filepath.Join(p.baseRepoPath, repoRule.DestinationRepository, "")
		dstURL := fmt.Sprintf("https://%s/%s/%s.git", p.config.GithubHost, p.config.TargetOrg, repoRule.DestinationRepository)
		if err := p.ensureCloned(dstDir, dstURL); err != nil {
			p.plog.Errorf("%v", err)
			return err
		}
		p.plog.Infof("Successfully ensured %s exists", dstDir)
		if err := os.Chdir(dstDir); err != nil {
			return err
		}

		// delete tags
		cmd := exec.Command("/bin/bash", "-c", "git tag | xargs git tag -d >/dev/null")
		if err := p.plog.Run(cmd); err != nil {
			return err
		}

		formatDeps := func(deps []config.Dependency) string {
			var depStrings []string
			for _, dep := range deps {
				depStrings = append(depStrings, fmt.Sprintf("%s:%s", dep.Repository, dep.Branch))
			}
			return strings.Join(depStrings, ",")
		}

		// construct branches
		for _, branchRule := range repoRule.Branches {
			if p.skippedBranch(branchRule.Source.Branch) {
				continue
			}
			if len(branchRule.Source.Dir) == 0 {
				branchRule.Source.Dir = "."
				p.plog.Infof("%v: 'dir' cannot be empty, defaulting to '.'", branchRule)
			}

			// get old HEAD. Ignore errors as the branch might be non-existent
			oldHead, _ := exec.Command("git", "rev-parse", fmt.Sprintf("origin/%s", branchRule.Name)).Output()

			goPath := os.Getenv("GOPATH")
			branchEnv := append([]string(nil), os.Environ()...) // make mutable
			if branchRule.GoVersion != "" {
				goRoot := filepath.Join(goPath, "go-"+branchRule.GoVersion)
				branchEnv = append(branchEnv, "GOROOT="+goRoot)
				goBin := filepath.Join(goRoot, "bin")
				branchEnv = updateEnv(branchEnv, "PATH", prependPath(goBin), goBin)
			}

			skipTags := ""
			if p.reposRules.SkipTags {
				skipTags = "true"
				p.plog.Infof("synchronizing tags is disabled")
			}

			// get old published hash to eventually skip cherry picking
			var lastPublishedUpstreamHash string
			bs, err := ioutil.ReadFile(path.Join(p.baseRepoPath, publishedFileName(repoRule.DestinationRepository, branchRule.Name)))
			if err != nil && !os.IsNotExist(err) {
				return err
			}
			if err == nil {
				lastPublishedUpstreamHash = string(bs)
			}

			// TODO: Refactor this to use environment variables instead
			repoPublishScriptPath := filepath.Join(p.config.BasePublishScriptPath, "construct.sh")
			cmd := exec.Command(repoPublishScriptPath,
				repoRule.DestinationRepository,
				branchRule.Source.Branch,
				branchRule.Name,
				formatDeps(branchRule.Dependencies),
				strings.Join(branchRule.RequiredPackages, ":"),
				sourceRemote,
				branchRule.Source.Dir,
				p.config.SourceRepo,
				p.config.SourceRepo,
				p.config.BasePackage,
				fmt.Sprintf("%v", repoRule.Library),
				strings.Join(p.reposRules.RecursiveDeletePatterns, " "),
				skipTags,
				lastPublishedUpstreamHash,
			)
			cmd.Env = append([]string(nil), branchEnv...) // make mutable
			if p.reposRules.SkipGodeps {
				cmd.Env = append(cmd.Env, "PUBLISHER_BOT_SKIP_GODEPS=true")
			}
			if err := p.plog.Run(cmd); err != nil {
				return err
			}

			newHead, _ := exec.Command("git", "rev-parse", fmt.Sprintf("origin/%s", branchRule.Name)).Output()
			if len(repoRule.SmokeTest) > 0 && string(oldHead) != string(newHead) {
				p.plog.Infof("Running smoke tests for branch %s", branchRule.Name)
				cmd := exec.Command("/bin/bash", "-xec", repoRule.SmokeTest)
				cmd.Env = append([]string(nil), branchEnv...) // make mutable
				if err := p.plog.Run(cmd); err != nil {
					// do not clean up to allow debugging with kubectl-exec.
					return err
				}
				exec.Command("git", "reset", "--hard").Run()
				exec.Command("git", "clean", "-f", "-f", "-d").Run()
			}

			p.plog.Infof("Successfully constructed %s", branchRule.Name)
		}
	}
	return nil
}

func updateEnv(env []string, key string, change func(string) string, val string) []string {
	for i := range env {
		if strings.HasPrefix(env[i], key+"=") {
			ss := strings.SplitN(env[i], "=", 2)
			env[i] = fmt.Sprintf("%s=%s", key, change(ss[1]))
			return env
		}
	}
	return append(env, fmt.Sprintf("%s=%s", key, val))
}

func prependPath(p string) func(string) string {
	return func(s string) string {
		if s == "" {
			return p
		}
		return p + ":" + s
	}
}

// publish to remotes.
func (p *PublisherMunger) publish(newUpstreamHeads map[string]plumbing.Hash) error {
	if p.config.DryRun {
		p.plog.Infof("Skipping push in dry-run mode")
		return nil
	}

	if p.config.TokenFile == "" {
		return fmt.Errorf("token cannot be empty in non-dry-run mode")
	}

	// NOTE: because some repos depend on each other, e.g., client-go depends on
	// apimachinery, they should be published atomically, but it's not supported
	// by github.
	for _, repoRules := range p.reposRules.Rules {
		if repoRules.Skip {
			continue
		}

		dstDir := filepath.Join(p.baseRepoPath, repoRules.DestinationRepository, "")
		if err := os.Chdir(dstDir); err != nil {
			return err
		}
		for _, branchRule := range repoRules.Branches {
			if p.skippedBranch(branchRule.Source.Branch) {
				continue
			}

			cmd := exec.Command(p.config.BasePublishScriptPath+"/push.sh", p.config.TokenFile, branchRule.Name)
			if err := p.plog.Run(cmd); err != nil {
				return err
			}

			upstreamBranchHead, ok := newUpstreamHeads[branchRule.Source.Branch]
			if !ok {
				return fmt.Errorf("no upstream branch %q found", branchRule.Source.Branch)
			}
			if err := ioutil.WriteFile(path.Join(path.Dir(dstDir), publishedFileName(repoRules.DestinationRepository, branchRule.Name)), []byte(upstreamBranchHead.String()), 0644); err != nil {
				return err
			}
		}
	}
	return nil
}

func publishedFileName(repo, branch string) string {
	return fmt.Sprintf("published-%s-%s", repo, branch)
}

// Run constructs the repos and pushes them. It returns logs and the last master hash.
func (p *PublisherMunger) Run() (string, string, error) {
	buf := bytes.NewBuffer(nil)
	var err error
	if p.plog, err = NewPublisherLog(buf, path.Join(p.baseRepoPath, "run.log")); err != nil {
		return "", "", err
	}

	newUpstreamHeads, err := p.updateSourceRepo()
	if err != nil {
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}
	if err := p.updateRules(); err != nil { // this comes after the source update because we might fetch the rules from there.
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}
	if err := p.construct(); err != nil {
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}
	if err := p.publish(newUpstreamHeads); err != nil {
		p.plog.Errorf("%v", err)
		p.plog.Flush()
		return p.plog.Logs(), "", err
	}

	var masterHead string
	if h, ok := newUpstreamHeads["master"]; ok {
		masterHead = h.String()
	}

	return p.plog.Logs(), masterHead, nil
}
<|endoftext|>