<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/lair-framework/api-server/client"
	"github.com/lair-framework/go-lair"
	"github.com/tomsteele/go-shodan"
)

const (
	version  = "1.0.0"
	tool     = "shodan"
	osWeight = 50
	usage    = `
Provided a file containing a new line delimited file containing cidr netblocks or ip
addresses, this drone uses shodan's 'net' and 'host' search to identify and import available
services into lair. Requests are made to shodan concurrently using a pool of 10 goroutines.

Usage:
  drone-shodan <id> <filename>
  export LAIR_ID=<id>; drone-shodan <filename>
Options:
  -v              show version and exit
  -h              show usage and exit
  -k              allow insecure SSL connections
  -force-ports    disable data protection in the API server for excessive ports
  -tags           a comma separated list of tags to add to every host that is imported
`
)

func removeDuplicates(in []string) []string {
	m := map[string]bool{}
	out := []string{}
	for _, i := range in {
		if i == "" {
			continue
		}
		if _, ok := m[i]; ok {
			continue
		}
		m[i] = true
		out = append(out, i)
	}
	return out
}

func shodanIPsFromShodanNetSearch(client *shodan.Client, netblock string) ([]string, error) {
	ips := []string{}
	result, err := client.HostSearch("net:"+netblock, []string{}, url.Values{})
	if err != nil {
		return ips, err
	}
	for _, m := range result.Matches {
		ips = append(ips, m.IPStr)
	}
	return ips, nil
}

func main() {
	showVersion := flag.Bool("v", false, "")
	insecureSSL := flag.Bool("k", false, "")
	forcePorts := flag.Bool("force-ports", false, "")
	tags := flag.String("tags", "", "")
	flag.Usage = func() {
		fmt.Println(usage)
	}
	flag.Parse()
	if *showVersion {
		log.Println(version)
		os.Exit(0)
	}
	lairURL := os.Getenv("LAIR_API_SERVER")
	if lairURL == "" {
		log.Fatal("Fatal: Missing LAIR_API_SERVER environment variable")
	}
	lairPID := os.Getenv("LAIR_ID")
	var filename string
	switch len(flag.Args()) {
	case 2:
		lairPID = flag.Arg(0)
		filename = flag.Arg(1)
	case 1:
		filename = flag.Arg(0)
	default:
		log.Fatal("Fatal: Missing required argument")
	}
	u, err := url.Parse(lairURL)
	if err != nil {
		log.Fatalf("Fatal: Error parsing LAIR_API_SERVER URL. Error %s", err.Error())
	}
	if u.User == nil {
		log.Fatal("Fatal: Missing username and/or password")
	}
	user := u.User.Username()
	pass, _ := u.User.Password()
	if user == "" || pass == "" {
		log.Fatal("Fatal: Missing username and/or password")
	}
	c, err := client.New(&client.COptions{
		User:               user,
		Password:           pass,
		Host:               u.Host,
		Scheme:             u.Scheme,
		InsecureSkipVerify: *insecureSSL,
	})
	if err != nil {
		log.Fatalf("Fatal: Error setting up client: Error %s", err.Error())
	}
	hostTags := []string{}
	if *tags != "" {
		hostTags = strings.Split(*tags, ",")
	}
	l := lair.Project{
		ID:   lairPID,
		Tool: tool,
		Commands: []lair.Command{lair.Command{
			Tool:    tool,
			Command: "",
		}},
	}

	shodanKey := os.Getenv("SHODAN_KEY")
	if shodanKey == "" {
		log.Fatal("Fatal: Missing SHODAN_KEY environment variable")
	}

	sclient := shodan.New(shodanKey)
	serviceMap, err := sclient.Services()
	if err != nil {
		log.Fatalf("Fatal: Error getting services from shodan. Error %s", err.Error())
	}

	ips := []string{}
	file, err := os.Open(filename)
	if err != nil {
		log.Fatalf("Fatal: Could not open file. Error %s", err.Error())
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		_, _, err := net.ParseCIDR(line)
		if err != nil {
			ip := net.ParseIP(line)
			if ip == nil {
				log.Fatalf("Fatal: %s in file is not an ip or cidr netblock", ip)
			}
			ips = append(ips, line)
		} else {
			if netIPs, err := shodanIPsFromShodanNetSearch(sclient, line); err != nil {
				log.Fatalf("Fatal: Error returned from shodan. Error %s", err.Error())
			} else {
				ips = append(ips, netIPs...)
			}
		}
	}

	lk := sync.Mutex{}
	wg := sync.WaitGroup{}
	wg.Add(10)
	ipChan := make(chan string, 10)

	for i := 0; i < 10; i++ {
		go func(s shodan.Client) {
			for ip := range ipChan {
				host, err := s.Host(ip, url.Values{})
				if err != nil {
					log.Printf("Error: Error returned from shodan for %s. Error %s", ip, err.Error())
					continue
				}
				h := lair.Host{
					Hostnames:      host.Hostnames,
					IPv4:           ip,
					LastModifiedBy: tool,
					Tags:           hostTags,
				}
				for _, d := range host.Data {
					service := lair.Service{
						Port:     d.Port,
						Protocol: "tcp",
						Service:  serviceMap[strconv.Itoa(d.Port)],
						Product:  d.Product,
						Notes: []lair.Note{lair.Note{
							Title:          "Shodan Banner",
							Content:        d.Data,
							LastModifiedBy: tool,
						}},
					}
					if fingerprint, ok := d.Os.(string); ok {
						h.OS = lair.OS{
							Fingerprint: fingerprint,
							Weight:      osWeight,
							Tool:        tool,
						}
					}
					h.Hostnames = removeDuplicates(append(h.Hostnames, d.Hostnames...))
					h.Services = append(h.Services, service)
				}
				lk.Lock()
				l.Hosts = append(l.Hosts, h)
				lk.Unlock()
			}
			wg.Done()
		}(*sclient)
	}

	for _, ip := range ips {
		ipChan <- ip
	}
	close(ipChan)
	wg.Wait()

	res, err := c.ImportProject(&client.DOptions{ForcePorts: *forcePorts}, &l)
	if err != nil {
		log.Fatalf("Fatal: Unable to import project. Error %s", err)
	}
	defer res.Body.Close()
	droneRes := &client.Response{}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.Fatalf("Fatal: Error %s", err.Error())
	}
	if err := json.Unmarshal(body, droneRes); err != nil {
		log.Fatalf("Fatal: Could not unmarshal JSON. Error %s", err.Error())
	}
	if droneRes.Status == "Error" {
		log.Fatalf("Fatal: Import failed. Error %s", droneRes.Message)
	}
	log.Println("Success: Operation completed successfully")
}
<commit_msg>updated help info<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/lair-framework/api-server/client"
	"github.com/lair-framework/go-lair"
	"github.com/tomsteele/go-shodan"
)

const (
	version  = "1.0.0"
	tool     = "shodan"
	osWeight = 50
	usage    = `
Provided a file containing a new line delimited file containing cidr netblocks or ip
addresses, this drone uses shodan's 'net' and 'host' search to identify and import available
services into lair. Requests are made to shodan concurrently using a pool of 10 goroutines.

Usage:
  drone-shodan [options] <id> <filename>
  export LAIR_ID=<id>; drone-shodan [options] <filename>
Options:
  -v              show version and exit
  -h              show usage and exit
  -k              allow insecure SSL connections
  -force-ports    disable data protection in the API server for excessive ports
  -tags           a comma separated list of tags to add to every host that is imported
`
)

func removeDuplicates(in []string) []string {
	m := map[string]bool{}
	out := []string{}
	for _, i := range in {
		if i == "" {
			continue
		}
		if _, ok := m[i]; ok {
			continue
		}
		m[i] = true
		out = append(out, i)
	}
	return out
}

func shodanIPsFromShodanNetSearch(client *shodan.Client, netblock string) ([]string, error) {
	ips := []string{}
	result, err := client.HostSearch("net:"+netblock, []string{}, url.Values{})
	if err != nil {
		return ips, err
	}
	for _, m := range result.Matches {
		ips = append(ips, m.IPStr)
	}
	return ips, nil
}

func main() {
	showVersion := flag.Bool("v", false, "")
	insecureSSL := flag.Bool("k", false, "")
	forcePorts := flag.Bool("force-ports", false, "")
	tags := flag.String("tags", "", "")
	flag.Usage = func() {
		fmt.Println(usage)
	}
	flag.Parse()
	if *showVersion {
		log.Println(version)
		os.Exit(0)
	}
	lairURL := os.Getenv("LAIR_API_SERVER")
	if lairURL == "" {
		log.Fatal("Fatal: Missing LAIR_API_SERVER environment variable")
	}
	lairPID := os.Getenv("LAIR_ID")
	var filename string
	switch len(flag.Args()) {
	case 2:
		lairPID = flag.Arg(0)
		filename = flag.Arg(1)
	case 1:
		filename = flag.Arg(0)
	default:
		log.Fatal("Fatal: Missing required argument")
	}
	u, err := url.Parse(lairURL)
	if err != nil {
		log.Fatalf("Fatal: Error parsing LAIR_API_SERVER URL. Error %s", err.Error())
	}
	if u.User == nil {
		log.Fatal("Fatal: Missing username and/or password")
	}
	user := u.User.Username()
	pass, _ := u.User.Password()
	if user == "" || pass == "" {
		log.Fatal("Fatal: Missing username and/or password")
	}
	c, err := client.New(&client.COptions{
		User:               user,
		Password:           pass,
		Host:               u.Host,
		Scheme:             u.Scheme,
		InsecureSkipVerify: *insecureSSL,
	})
	if err != nil {
		log.Fatalf("Fatal: Error setting up client: Error %s", err.Error())
	}
	hostTags := []string{}
	if *tags != "" {
		hostTags = strings.Split(*tags, ",")
	}
	l := lair.Project{
		ID:   lairPID,
		Tool: tool,
		Commands: []lair.Command{lair.Command{
			Tool:    tool,
			Command: "",
		}},
	}

	shodanKey := os.Getenv("SHODAN_KEY")
	if shodanKey == "" {
		log.Fatal("Fatal: Missing SHODAN_KEY environment variable")
	}

	sclient := shodan.New(shodanKey)
	serviceMap, err := sclient.Services()
	if err != nil {
		log.Fatalf("Fatal: Error getting services from shodan. Error %s", err.Error())
	}

	ips := []string{}
	file, err := os.Open(filename)
	if err != nil {
		log.Fatalf("Fatal: Could not open file. Error %s", err.Error())
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		_, _, err := net.ParseCIDR(line)
		if err != nil {
			ip := net.ParseIP(line)
			if ip == nil {
				log.Fatalf("Fatal: %s in file is not an ip or cidr netblock", ip)
			}
			ips = append(ips, line)
		} else {
			if netIPs, err := shodanIPsFromShodanNetSearch(sclient, line); err != nil {
				log.Fatalf("Fatal: Error returned from shodan. Error %s", err.Error())
			} else {
				ips = append(ips, netIPs...)
			}
		}
	}

	lk := sync.Mutex{}
	wg := sync.WaitGroup{}
	wg.Add(10)
	ipChan := make(chan string, 10)

	for i := 0; i < 10; i++ {
		go func(s shodan.Client) {
			for ip := range ipChan {
				host, err := s.Host(ip, url.Values{})
				if err != nil {
					log.Printf("Error: Error returned from shodan for %s. Error %s", ip, err.Error())
					continue
				}
				h := lair.Host{
					Hostnames:      host.Hostnames,
					IPv4:           ip,
					LastModifiedBy: tool,
					Tags:           hostTags,
				}
				for _, d := range host.Data {
					service := lair.Service{
						Port:     d.Port,
						Protocol: "tcp",
						Service:  serviceMap[strconv.Itoa(d.Port)],
						Product:  d.Product,
						Notes: []lair.Note{lair.Note{
							Title:          "Shodan Banner",
							Content:        d.Data,
							LastModifiedBy: tool,
						}},
					}
					if fingerprint, ok := d.Os.(string); ok {
						h.OS = lair.OS{
							Fingerprint: fingerprint,
							Weight:      osWeight,
							Tool:        tool,
						}
					}
					h.Hostnames = removeDuplicates(append(h.Hostnames, d.Hostnames...))
					h.Services = append(h.Services, service)
				}
				lk.Lock()
				l.Hosts = append(l.Hosts, h)
				lk.Unlock()
			}
			wg.Done()
		}(*sclient)
	}

	for _, ip := range ips {
		ipChan <- ip
	}
	close(ipChan)
	wg.Wait()

	res, err := c.ImportProject(&client.DOptions{ForcePorts: *forcePorts}, &l)
	if err != nil {
		log.Fatalf("Fatal: Unable to import project. Error %s", err)
	}
	defer res.Body.Close()
	droneRes := &client.Response{}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.Fatalf("Fatal: Error %s", err.Error())
	}
	if err := json.Unmarshal(body, droneRes); err != nil {
		log.Fatalf("Fatal: Could not unmarshal JSON. Error %s", err.Error())
	}
	if droneRes.Status == "Error" {
		log.Fatalf("Fatal: Import failed. Error %s", droneRes.Message)
	}
	log.Println("Success: Operation completed successfully")
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"net/http"

	"github.com/twanies/weavebox-docker-example/Godeps/_workspace/src/github.com/twanies/weavebox"
)

func main() {
	var listen = flag.Int("listen", 3000, "target listen address of the application")
	flag.Parse()

	app := weavebox.New()
	app.Get("/", helloHandler)
	log.Fatal(app.Serve(*listen))
}

func helloHandler(ctx *weavebox.Context) error {
	return ctx.Text(http.StatusOK, "Hello from weavebox in a docker container")
}
<commit_msg>updated main<commit_after>package main

import (
	"flag"
	"log"
	"net/http"

	"github.com/twanies/weavebox-docker-example/Godeps/_workspace/src/github.com/twanies/weavebox"
)

func main() {
	var listen = flag.Int("listen", 3000, "target listen address of the application")
	flag.Parse()

	app := weavebox.New()
	app.Get("/", helloHandler)
	log.Fatal(app.Serve(*listen))
}

func helloHandler(ctx *weavebox.Context) error {
	return ctx.Text(http.StatusOK, "Hello from weavebox in a fancy docker container")
}
<|endoftext|>


<commit_before>// Copyright 2015 Stuart Glenn, OMRF. All rights reserved.
// Use of this code is governed by a 3 clause BSD style license
// Full license details in LICENSE file distributed with this software

package main

import (
	"encoding/csv"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/oklasoft/mmatcher/matcher"
	"gopkg.in/alecthomas/kingpin.v1"
)

func loadData(path string, skipHeader bool) matcher.Records {
	file, err := os.Open(path)
	if nil != err {
		log.Fatal(err)
	}
	defer file.Close()

	data, err := matcher.NewRecordsFromCSV(file, skipHeader)
	if nil != err {
		log.Fatal(err)
	}
	return data
}

func parseKeys(s string) ([]int, []matcher.Atter) {
	parts := strings.Split(s, ",")
	positions := make([]int, len(parts))
	ranges := make([]matcher.Atter, len(parts))
	for i, v := range parts {
		k := strings.Split(v, ":")
		p, err := strconv.ParseInt(k[0], 10, 32)
		if nil != err {
			log.Fatal(err)
		}
		positions[i] = int(p) - 1
		if len(k) > 1 {
			r, err := strconv.ParseFloat(k[1], 32)
			if nil != err {
				log.Fatal(err)
			}
			ranges[i] = matcher.NumericAtt{r}
		}
	}
	return positions, ranges
}

func version() string {
	return fmt.Sprintf("mmatcher - Multi Matcher 0.6.0 (20150403 %s)", build)
}

var (
	verbose      = kingpin.Flag("verbose", "Increase verbosity").Short('v').Bool()
	skipHeaders  = kingpin.Flag("skip-header", "Inputs have header line to be skipped, default is use everyline").Short('h').Bool()
	outFile      = kingpin.Flag("output", "Output file").Short('o').PlaceHolder("STDOUT").OpenFile(os.O_WRONLY|os.O_CREATE, 0660)
	key          = kingpin.Arg("keys", "Keys to compare. A comma separated list of columns starting a 1, with optional :# +/- window").Required().String()
	case_file    = kingpin.Arg("case", "CSV file representing the cases").Required().ExistingFile()
	control_file = kingpin.Arg("controls", "CSV file representing the controls").Required().ExistingFile()
	build        string
)

func main() {
	kingpin.Version(version())
	kingpin.Parse()

	if nil == *outFile {
		outFile = &os.Stdout
	}

	positions, ranges := parseKeys(*key)
	cases := loadData(*case_file, *skipHeaders)
	controls := loadData(*control_file, *skipHeaders)

	all_matches := matcher.NewMatchSet()

	for _, r := range cases {
		spots := r.Matches(controls, positions, ranges...)
		for _, i := range spots {
			all_matches.AddPair(matcher.NewPair(r.ID, controls[i].ID))
		}
	}

	opti := all_matches.QuantityOptimized()

	out := csv.NewWriter(*outFile)
	out.Write([]string{"case", "control"})

	for _, r := range cases {
		m := opti.MatchesFor(r.ID)
		if 0 == len(m) {
			out.Write([]string{r.ID})
			continue
		}
		line := []string{r.ID, m[0]}
		if *verbose {
			control := controls.Get(m[0])
			for _, p := range positions {
				line = append(line, r.Atts[p].String(), control.Atts[p].String())
			}
		}
		out.Write(line)
	}
	out.Flush()
	(*outFile).Close()

}
<commit_msg>Add --out-separator option<commit_after>// Copyright 2015 Stuart Glenn, OMRF. All rights reserved.
// Use of this code is governed by a 3 clause BSD style license
// Full license details in LICENSE file distributed with this software

package main

import (
	"encoding/csv"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/oklasoft/mmatcher/matcher"
	"gopkg.in/alecthomas/kingpin.v1"
)

func loadData(path string, skipHeader bool) matcher.Records {
	file, err := os.Open(path)
	if nil != err {
		log.Fatal(err)
	}
	defer file.Close()

	data, err := matcher.NewRecordsFromCSV(file, skipHeader)
	if nil != err {
		log.Fatal(err)
	}
	return data
}

func parseKeys(s string) ([]int, []matcher.Atter) {
	parts := strings.Split(s, ",")
	positions := make([]int, len(parts))
	ranges := make([]matcher.Atter, len(parts))
	for i, v := range parts {
		k := strings.Split(v, ":")
		p, err := strconv.ParseInt(k[0], 10, 32)
		if nil != err {
			log.Fatal(err)
		}
		positions[i] = int(p) - 1
		if len(k) > 1 {
			r, err := strconv.ParseFloat(k[1], 32)
			if nil != err {
				log.Fatal(err)
			}
			ranges[i] = matcher.NumericAtt{r}
		}
	}
	return positions, ranges
}

func version() string {
	return fmt.Sprintf("mmatcher - Multi Matcher 0.6.0 (20150403 %s)", build)
}

var (
	verbose      = kingpin.Flag("verbose", "Increase verbosity").Short('v').Bool()
	skipHeaders  = kingpin.Flag("skip-header", "Inputs have header line to be skipped, default is use everyline").Short('h').Bool()
	outFile      = kingpin.Flag("output", "Output file").Short('o').PlaceHolder("STDOUT").OpenFile(os.O_WRONLY|os.O_CREATE, 0660)
	outSep       = kingpin.Flag("out-separator", "Output field separator").Default(",").String()
	key          = kingpin.Arg("keys", "Keys to compare. A comma separated list of columns starting a 1, with optional :# +/- window").Required().String()
	case_file    = kingpin.Arg("case", "CSV file representing the cases").Required().ExistingFile()
	control_file = kingpin.Arg("controls", "CSV file representing the controls").Required().ExistingFile()
	build        string
)

func main() {
	kingpin.Version(version())
	kingpin.Parse()

	if nil == *outFile {
		outFile = &os.Stdout
	}

	positions, ranges := parseKeys(*key)
	cases := loadData(*case_file, *skipHeaders)
	controls := loadData(*control_file, *skipHeaders)

	all_matches := matcher.NewMatchSet()

	for _, r := range cases {
		spots := r.Matches(controls, positions, ranges...)
		for _, i := range spots {
			all_matches.AddPair(matcher.NewPair(r.ID, controls[i].ID))
		}
	}

	opti := all_matches.QuantityOptimized()

	out := csv.NewWriter(*outFile)
	sep, err := strconv.Unquote("'" + *outSep + "'")
	if nil != err {
		log.Fatal(err)
	}
	out.Comma = ([]rune(sep))[0]
	out.Write([]string{"case", "control"})

	for _, r := range cases {
		m := opti.MatchesFor(r.ID)
		if 0 == len(m) {
			out.Write([]string{r.ID})
			continue
		}
		line := []string{r.ID, m[0]}
		if *verbose {
			control := controls.Get(m[0])
			for _, p := range positions {
				line = append(line, r.Atts[p].String(), control.Atts[p].String())
			}
		}
		out.Write(line)
	}
	out.Flush()
	(*outFile).Close()

}
<|endoftext|>


<commit_before>package main

import (
"fmt"
//	"net"
	"os"
//	"errors"
//	"log"
	"io"
	"github.com/vaughan0/go-ini"
	"encoding/json"
	"net/http"
	"io/ioutil"
	"strconv"
	)

func main(){

	fmt.Println("Derpiboo.ru Downloader version 0.0.2 \nWorking")
//	fmt.Println("Working")

	config, err := ini.LoadFile("config.ini") // Loading default config file and checking for various errors.

	if os.IsNotExist(err) {
		panic("config.ini does not exist, create it")
		}

	if err != nil { panic(err) }

	key, ok := config.Get("main", "key") //Need to make things for key == nil

	if !ok {
		panic("'key' variable missing from 'main' section")
		}

	length := len(os.Args)
	if length == 1 {
		fmt.Println("Nothing to download, bye!")
		os.Exit(0)
	}

	imgid := os.Args[length - 1] //Last argument given presumed to be number of image on site. Temporally, because later would do with flags.

//	fmt.Println(key) //Just checking that I am not wrong

	fmt.Println("Processing image No " + imgid)

	imgdat := parse (imgid, key)

	dlimage(imgdat)

}

type Image struct {
	url		string
	filename	string
	hash		string
	}

func parse(imgid string, key string) (imgdata Image) {

	source := "http://derpiboo.ru/" + imgid + ".json?nofav=&nocomments=?key=" + key //correct way is to assemble all the different arguments and only then append them to source url. I can live with hardcoded source site. May be add check for derpiboo.ru or derpibooru.org ?
//	fmt.Println(source)

	resp, err := http.Get(source)	//Getting our nice http response. Needs checking for 404 and other responses that are... less expected
		if err != nil {
			panic(err)
		}

	defer resp.Body.Close()	//and not forgetting to close it when it's done

	var dat map[string]interface{}

	body, err := ioutil.ReadAll(resp.Body)	//stolen from official documentation
		if err != nil {
			panic(err)
		}

	if err := json.Unmarshal(body, &dat); //transforming json into native map

	err != nil {
	        panic(err)
		}

	imgdata.url = "http:" + dat["image"].(string)
	imgdata.hash = dat["sha512_hash"].(string)  //for the future and checking that we got file right
	imgdata.filename = strconv.FormatFloat(dat["id_number"].(float64), 'f', -1, 64) + "." + dat["file_name"].(string)

//	fmt.Println(strconv.FormatFloat(dat["id_number"].(float64), 'f', -1, 64))

//	fmt.Println(dat)

// for now and troubleshooting
//	fmt.Println(imgdata.url)
//	fmt.Println(imgdata.hash)
//	fmt.Println(imgdata.filename)

	return
	}

func dlimage(imgdata Image) {
	fmt.Println("Saving as ", imgdata.filename)
	output, err := os.Create(imgdata.filename)
	defer output.Close()

	response, err := http.Get(imgdata.url)
    if err != nil {
		fmt.Println("Error while downloading", imgdata.url, "-", err)
		return
    }
    defer response.Body.Close()

    io.Copy(output, response.Body)

	}
<commit_msg>science - in the channels<commit_after>package main

import (
"fmt"
//	"net"
	"os"
//	"errors"
//	"log"
	"io"
	"github.com/vaughan0/go-ini"
	"encoding/json"
	"net/http"
	"io/ioutil"
	"strconv"
	)

func main(){

	fmt.Println("Derpiboo.ru Downloader version 0.0.2 \nWorking")
//	fmt.Println("Working")

	config, err := ini.LoadFile("config.ini") // Loading default config file and checking for various errors.

	if os.IsNotExist(err) {
		panic("config.ini does not exist, create it")
		}

	if err != nil { panic(err) }

	key, ok := config.Get("main", "key") //Need to make things for key == nil

	if !ok {
		panic("'key' variable missing from 'main' section")
		}

	length := len(os.Args)
	if length == 1 {
		fmt.Println("Nothing to download, bye!")
		os.Exit(0)
	}

	imgid := os.Args[length - 1] //Last argument given presumed to be number of image on site. Temporally, because later would do with flags.

//	fmt.Println(key) //Just checking that I am not wrong

	fmt.Println("Processing image No " + imgid)

	imgdat := parseImg (imgid, key)

	dlimage(imgdat)

}

type Image struct {
	url		string
	filename	string
	hash		string
	}

func parseImg(imgid string, key string) (imgdata Image) {

	source := "http://derpiboo.ru/" + imgid + ".json?nofav=&nocomments=?key=" + key //correct way is to assemble all the different arguments and only then append them to source url. I can live with hardcoded source site. May be add check for derpiboo.ru or derpibooru.org ?
//	fmt.Println(source)

	resp, err := http.Get(source)	//Getting our nice http response. Needs checking for 404 and other responses that are... less expected
		if err != nil {
			panic(err)
		}

	defer resp.Body.Close()	//and not forgetting to close it when it's done

	var dat map[string]interface{}

	body, err := ioutil.ReadAll(resp.Body)	//stolen from official documentation
		if err != nil {
			panic(err)
		}

	if err := json.Unmarshal(body, &dat); //transforming json into native map

	err != nil {
	        panic(err)
		}

	imgdata.url = "http:" + dat["image"].(string)
	imgdata.hash = dat["sha512_hash"].(string)  //for the future and checking that we got file right
	imgdata.filename = strconv.FormatFloat(dat["id_number"].(float64), 'f', -1, 64) + "." + dat["file_name"].(string)

//	fmt.Println(strconv.FormatFloat(dat["id_number"].(float64), 'f', -1, 64))

//	fmt.Println(dat)

// for now and troubleshooting
//	fmt.Println(imgdata.url)
//	fmt.Println(imgdata.hash)
//	fmt.Println(imgdata.filename)

	return
	}

func dlimage(imgdata Image) {
	fmt.Println("Saving as ", imgdata.filename)
	output, err := os.Create(imgdata.filename)
	defer output.Close()

	response, err := http.Get(imgdata.url)
    if err != nil {
		fmt.Println("Error while downloading", imgdata.url, "-", err)
		return
    }
    defer response.Body.Close()

    io.Copy(output, response.Body)

	}

func parseTag ( tag string ) (imgdata chan-< Image) {}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strconv"

	"time"

	"strings"

	"syscall"

	"os/signal"

	"runtime"

	"github.com/claudetech/loggo"
	. "github.com/claudetech/loggo/default"
	"github.com/plexdrive/plexdrive/chunk"
	"github.com/plexdrive/plexdrive/config"
	"github.com/plexdrive/plexdrive/drive"
	"github.com/plexdrive/plexdrive/mount"
	flag "github.com/ogier/pflag"
	"golang.org/x/sys/unix"
)

func main() {
	// Find users home directory
	usr, err := user.Current()
	home := ""
	if err != nil {
		// Fall back to reading $HOME - work around user.Current() not
		// working for cross compiled binaries on OSX or freebsd.
		// https://github.com/golang/go/issues/6376
		home = os.Getenv("HOME")
		if home == "" {
			panic(fmt.Sprintf("Could not read users homedir and HOME is not set: %v\n", err))
		}
	} else {
		home = usr.HomeDir
	}

	// parse the command line arguments
	argLogLevel := flag.IntP("verbosity", "v", 0, "Set the log level (0 = error, 1 = warn, 2 = info, 3 = debug, 4 = trace)")
	argRootNodeID := flag.String("root-node-id", "root", "The ID of the root node to mount (use this for only mount a sub directory)")
	argDriveID := flag.String("drive-id", "", "The ID of the shared drive to mount (including team drives)")
	argConfigPath := flag.StringP("config", "c", filepath.Join(home, ".plexdrive"), "The path to the configuration directory")
	argCacheFile := flag.String("cache-file", filepath.Join(*argConfigPath, "cache.bolt"), "Path of the cache file")
	argChunkSize := flag.String("chunk-size", "10M", "The size of each chunk that is downloaded (units: B, K, M, G)")
	argChunkLoadThreads := flag.Int("chunk-load-threads", max(runtime.NumCPU()/2, 1), "The number of threads to use for downloading chunks")
	argChunkCheckThreads := flag.Int("chunk-check-threads", max(runtime.NumCPU()/2, 1), "The number of threads to use for checking chunk existence")
	argChunkLoadAhead := flag.Int("chunk-load-ahead", max(runtime.NumCPU()-1, 1), "The number of chunks that should be read ahead")
	argMaxChunks := flag.Int("max-chunks", runtime.NumCPU()*2, "The maximum number of chunks to be stored in memory")
	argRefreshInterval := flag.Duration("refresh-interval", 1*time.Minute, "The time to wait till checking for changes")
	argMountOptions := flag.StringP("fuse-options", "o", "", "Fuse mount options (e.g. --fuse-options allow_other,direct_io,...)")
	argVersion := flag.Bool("version", false, "Displays program's version information")
	argUID := flag.Int64("uid", -1, "Set the mounts UID (-1 = default permissions)")
	argGID := flag.Int64("gid", -1, "Set the mounts GID (-1 = default permissions)")
	argUmask := flag.Uint32("umask", 0, "Override the default file permissions")
	// argDownloadSpeedLimit := flag.String("speed-limit", "", "This value limits the download speed, e.g. 5M = 5MB/s per chunk (units: B, K, M, G)")
	flag.Parse()

	// display version information
	if *argVersion {
		fmt.Println("%VERSION%")
		return
	}

	argCommand := flag.Arg(0)

	if argCommand == "mount" {
		// check if mountpoint is specified
		argMountPoint := flag.Arg(1)
		if "" == argMountPoint {
			flag.Usage()
			fmt.Println()
			panic(fmt.Errorf("Mountpoint not specified"))
		}

		// calculate uid / gid
		uid := uint32(unix.Geteuid())
		gid := uint32(unix.Getegid())
		if *argUID > -1 {
			uid = uint32(*argUID)
		}
		if *argGID > -1 {
			gid = uint32(*argGID)
		}

		// parse filemode
		umask := os.FileMode(*argUmask)

		// parse the mount options
		var mountOptions []string
		if "" != *argMountOptions {
			mountOptions = strings.Split(*argMountOptions, ",")
		}

		// initialize the logger with the specific log level
		var logLevel loggo.Level
		switch *argLogLevel {
		case 0:
			logLevel = loggo.Error
		case 1:
			logLevel = loggo.Warning
		case 2:
			logLevel = loggo.Info
		case 3:
			logLevel = loggo.Debug
		case 4:
			logLevel = loggo.Trace
		default:
			logLevel = loggo.Warning
		}
		Log.SetLevel(logLevel)

		// debug all given parameters
		Log.Debugf("verbosity            : %v", logLevel)
		Log.Debugf("root-node-id         : %v", *argRootNodeID)
		Log.Debugf("drive-id             : %v", *argDriveID)
		Log.Debugf("config               : %v", *argConfigPath)
		Log.Debugf("cache-file           : %v", *argCacheFile)
		Log.Debugf("chunk-size           : %v", *argChunkSize)
		Log.Debugf("chunk-load-threads   : %v", *argChunkLoadThreads)
		Log.Debugf("chunk-check-threads  : %v", *argChunkCheckThreads)
		Log.Debugf("chunk-load-ahead     : %v", *argChunkLoadAhead)
		Log.Debugf("max-chunks           : %v", *argMaxChunks)
		Log.Debugf("refresh-interval     : %v", *argRefreshInterval)
		Log.Debugf("fuse-options         : %v", *argMountOptions)
		Log.Debugf("UID                  : %v", uid)
		Log.Debugf("GID                  : %v", gid)
		Log.Debugf("umask                : %v", umask)
		// Log.Debugf("speed-limit          : %v", *argDownloadSpeedLimit)
		// version missing here

		// create all directories
		if err := os.MkdirAll(*argConfigPath, 0766); nil != err {
			Log.Errorf("Could not create configuration directory")
			Log.Debugf("%v", err)
			os.Exit(1)
		}
		if err := os.MkdirAll(filepath.Dir(*argCacheFile), 0766); nil != err {
			Log.Errorf("Could not create cache file directory")
			Log.Debugf("%v", err)
			os.Exit(1)
		}

		// set the global buffer configuration
		chunkSize, err := parseSizeArg(*argChunkSize)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(2)
		}

		// read the configuration
		configPath := filepath.Join(*argConfigPath, "config.json")
		cfg, err := config.Read(configPath)
		if nil != err {
			cfg, err = config.Create(configPath)
			if nil != err {
				Log.Errorf("Could not read configuration")
				Log.Debugf("%v", err)
				os.Exit(3)
			}
		}

		cache, err := drive.NewCache(*argCacheFile, *argConfigPath, *argLogLevel > 3)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(4)
		}
		defer cache.Close()

		client, err := drive.NewClient(cfg, cache, *argRefreshInterval, *argRootNodeID, *argDriveID)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(4)
		}

		chunkManager, err := chunk.NewManager(
			chunkSize,
			*argChunkLoadAhead,
			*argChunkCheckThreads,
			*argChunkLoadThreads,
			client,
			*argMaxChunks)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(4)
		}

		// check os signals like SIGINT/TERM
		checkOsSignals(argMountPoint)
		if err := mount.Mount(client, chunkManager, argMountPoint, mountOptions, uid, gid, umask); nil != err {
			Log.Debugf("%v", err)
			os.Exit(5)
		}
	} else {
		Log.Errorf("Command %v not found", argCommand)
	}
}

func checkOsSignals(mountpoint string) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT)

	go func() {
		for sig := range signals {
			if sig == syscall.SIGINT {
				if err := mount.Unmount(mountpoint, false); nil != err {
					Log.Warningf("%v", err)
				}
			}
		}
	}()
}

func max(x, y int) int {
	if x > y {
		return x
	}
	return y
}

func parseSizeArg(input string) (int64, error) {
	if "" == input {
		return 0, nil
	}

	suffix := input[len(input)-1]
	suffixLen := 1
	var multiplier float64
	switch suffix {
	case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '.':
		suffixLen = 0
	case 'b', 'B':
		multiplier = 1
	case 'k', 'K':
		multiplier = 1024
	case 'm', 'M':
		multiplier = 1024 * 1024
	case 'g', 'G':
		multiplier = 1024 * 1024 * 1024
	default:
		return 0, fmt.Errorf("Invalid unit %v for %v", suffix, input)
	}
	input = input[:len(input)-suffixLen]
	value, err := strconv.ParseFloat(input, 64)
	if nil != err {
		Log.Debugf("%v", err)
		return 0, fmt.Errorf("Could not parse numeric value %v", input)
	}
	if value < 0 {
		return 0, fmt.Errorf("Numeric value must not be negative %v", input)
	}
	value *= multiplier
	return int64(value), nil
}
<commit_msg>Do graceful unmount on SIGTERM<commit_after>package main

import (
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strconv"

	"time"

	"strings"

	"syscall"

	"os/signal"

	"runtime"

	"github.com/claudetech/loggo"
	. "github.com/claudetech/loggo/default"
	"github.com/plexdrive/plexdrive/chunk"
	"github.com/plexdrive/plexdrive/config"
	"github.com/plexdrive/plexdrive/drive"
	"github.com/plexdrive/plexdrive/mount"
	flag "github.com/ogier/pflag"
	"golang.org/x/sys/unix"
)

func main() {
	// Find users home directory
	usr, err := user.Current()
	home := ""
	if err != nil {
		// Fall back to reading $HOME - work around user.Current() not
		// working for cross compiled binaries on OSX or freebsd.
		// https://github.com/golang/go/issues/6376
		home = os.Getenv("HOME")
		if home == "" {
			panic(fmt.Sprintf("Could not read users homedir and HOME is not set: %v\n", err))
		}
	} else {
		home = usr.HomeDir
	}

	// parse the command line arguments
	argLogLevel := flag.IntP("verbosity", "v", 0, "Set the log level (0 = error, 1 = warn, 2 = info, 3 = debug, 4 = trace)")
	argRootNodeID := flag.String("root-node-id", "root", "The ID of the root node to mount (use this for only mount a sub directory)")
	argDriveID := flag.String("drive-id", "", "The ID of the shared drive to mount (including team drives)")
	argConfigPath := flag.StringP("config", "c", filepath.Join(home, ".plexdrive"), "The path to the configuration directory")
	argCacheFile := flag.String("cache-file", filepath.Join(*argConfigPath, "cache.bolt"), "Path of the cache file")
	argChunkSize := flag.String("chunk-size", "10M", "The size of each chunk that is downloaded (units: B, K, M, G)")
	argChunkLoadThreads := flag.Int("chunk-load-threads", max(runtime.NumCPU()/2, 1), "The number of threads to use for downloading chunks")
	argChunkCheckThreads := flag.Int("chunk-check-threads", max(runtime.NumCPU()/2, 1), "The number of threads to use for checking chunk existence")
	argChunkLoadAhead := flag.Int("chunk-load-ahead", max(runtime.NumCPU()-1, 1), "The number of chunks that should be read ahead")
	argMaxChunks := flag.Int("max-chunks", runtime.NumCPU()*2, "The maximum number of chunks to be stored in memory")
	argRefreshInterval := flag.Duration("refresh-interval", 1*time.Minute, "The time to wait till checking for changes")
	argMountOptions := flag.StringP("fuse-options", "o", "", "Fuse mount options (e.g. --fuse-options allow_other,direct_io,...)")
	argVersion := flag.Bool("version", false, "Displays program's version information")
	argUID := flag.Int64("uid", -1, "Set the mounts UID (-1 = default permissions)")
	argGID := flag.Int64("gid", -1, "Set the mounts GID (-1 = default permissions)")
	argUmask := flag.Uint32("umask", 0, "Override the default file permissions")
	// argDownloadSpeedLimit := flag.String("speed-limit", "", "This value limits the download speed, e.g. 5M = 5MB/s per chunk (units: B, K, M, G)")
	flag.Parse()

	// display version information
	if *argVersion {
		fmt.Println("%VERSION%")
		return
	}

	argCommand := flag.Arg(0)

	if argCommand == "mount" {
		// check if mountpoint is specified
		argMountPoint := flag.Arg(1)
		if "" == argMountPoint {
			flag.Usage()
			fmt.Println()
			panic(fmt.Errorf("Mountpoint not specified"))
		}

		// calculate uid / gid
		uid := uint32(unix.Geteuid())
		gid := uint32(unix.Getegid())
		if *argUID > -1 {
			uid = uint32(*argUID)
		}
		if *argGID > -1 {
			gid = uint32(*argGID)
		}

		// parse filemode
		umask := os.FileMode(*argUmask)

		// parse the mount options
		var mountOptions []string
		if "" != *argMountOptions {
			mountOptions = strings.Split(*argMountOptions, ",")
		}

		// initialize the logger with the specific log level
		var logLevel loggo.Level
		switch *argLogLevel {
		case 0:
			logLevel = loggo.Error
		case 1:
			logLevel = loggo.Warning
		case 2:
			logLevel = loggo.Info
		case 3:
			logLevel = loggo.Debug
		case 4:
			logLevel = loggo.Trace
		default:
			logLevel = loggo.Warning
		}
		Log.SetLevel(logLevel)

		// debug all given parameters
		Log.Debugf("verbosity            : %v", logLevel)
		Log.Debugf("root-node-id         : %v", *argRootNodeID)
		Log.Debugf("drive-id             : %v", *argDriveID)
		Log.Debugf("config               : %v", *argConfigPath)
		Log.Debugf("cache-file           : %v", *argCacheFile)
		Log.Debugf("chunk-size           : %v", *argChunkSize)
		Log.Debugf("chunk-load-threads   : %v", *argChunkLoadThreads)
		Log.Debugf("chunk-check-threads  : %v", *argChunkCheckThreads)
		Log.Debugf("chunk-load-ahead     : %v", *argChunkLoadAhead)
		Log.Debugf("max-chunks           : %v", *argMaxChunks)
		Log.Debugf("refresh-interval     : %v", *argRefreshInterval)
		Log.Debugf("fuse-options         : %v", *argMountOptions)
		Log.Debugf("UID                  : %v", uid)
		Log.Debugf("GID                  : %v", gid)
		Log.Debugf("umask                : %v", umask)
		// Log.Debugf("speed-limit          : %v", *argDownloadSpeedLimit)
		// version missing here

		// create all directories
		if err := os.MkdirAll(*argConfigPath, 0766); nil != err {
			Log.Errorf("Could not create configuration directory")
			Log.Debugf("%v", err)
			os.Exit(1)
		}
		if err := os.MkdirAll(filepath.Dir(*argCacheFile), 0766); nil != err {
			Log.Errorf("Could not create cache file directory")
			Log.Debugf("%v", err)
			os.Exit(1)
		}

		// set the global buffer configuration
		chunkSize, err := parseSizeArg(*argChunkSize)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(2)
		}

		// read the configuration
		configPath := filepath.Join(*argConfigPath, "config.json")
		cfg, err := config.Read(configPath)
		if nil != err {
			cfg, err = config.Create(configPath)
			if nil != err {
				Log.Errorf("Could not read configuration")
				Log.Debugf("%v", err)
				os.Exit(3)
			}
		}

		cache, err := drive.NewCache(*argCacheFile, *argConfigPath, *argLogLevel > 3)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(4)
		}
		defer cache.Close()

		client, err := drive.NewClient(cfg, cache, *argRefreshInterval, *argRootNodeID, *argDriveID)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(4)
		}

		chunkManager, err := chunk.NewManager(
			chunkSize,
			*argChunkLoadAhead,
			*argChunkCheckThreads,
			*argChunkLoadThreads,
			client,
			*argMaxChunks)
		if nil != err {
			Log.Errorf("%v", err)
			os.Exit(4)
		}

		// check os signals like SIGINT/TERM
		checkOsSignals(argMountPoint)
		if err := mount.Mount(client, chunkManager, argMountPoint, mountOptions, uid, gid, umask); nil != err {
			Log.Debugf("%v", err)
			os.Exit(5)
		}
	} else {
		Log.Errorf("Command %v not found", argCommand)
	}
}

func checkOsSignals(mountpoint string) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		for sig := range signals {
			Log.Infof("Received signal %v, stopping mount", sig)
			if err := mount.Unmount(mountpoint, false); nil != err {
				Log.Warningf("%v", err)
			}
		}
	}()
}

func max(x, y int) int {
	if x > y {
		return x
	}
	return y
}

func parseSizeArg(input string) (int64, error) {
	if "" == input {
		return 0, nil
	}

	suffix := input[len(input)-1]
	suffixLen := 1
	var multiplier float64
	switch suffix {
	case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '.':
		suffixLen = 0
	case 'b', 'B':
		multiplier = 1
	case 'k', 'K':
		multiplier = 1024
	case 'm', 'M':
		multiplier = 1024 * 1024
	case 'g', 'G':
		multiplier = 1024 * 1024 * 1024
	default:
		return 0, fmt.Errorf("Invalid unit %v for %v", suffix, input)
	}
	input = input[:len(input)-suffixLen]
	value, err := strconv.ParseFloat(input, 64)
	if nil != err {
		Log.Debugf("%v", err)
		return 0, fmt.Errorf("Could not parse numeric value %v", input)
	}
	if value < 0 {
		return 0, fmt.Errorf("Numeric value must not be negative %v", input)
	}
	value *= multiplier
	return int64(value), nil
}
<|endoftext|>


<commit_before>/*
Quick and simple monitoring system

Usage:
 jsonmon path.to/config.yml
 jsonmon -v # Prints version to stdout and exits

Environment:
 HOST
  - defaults to localhost
  - the JSON API network interface
 PORT
  - defaults to 3000
  - the JSON API port
 GOMAXPROCS
  - defaults to CPU number + 1
  - number of threads to start
*/
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Application version
const Version = "0.9.7"

// Check details
type Check struct {
	Name   string      `json:"name,omitempty" yaml:"name"`
	Web    string      `json:"web,omitempty" yaml:"web"`
	Shell  string      `json:"shell,omitempty" yaml:"shell"`
	Notify interface{} `json:"-" yaml:"notify"`
	Tries  int         `json:"-" yaml:"tries"`
	Repeat int         `json:"-" yaml:"repeat"`
	Failed bool        `json:"failed" yaml:"-"`
	Since  string      `json:"since,omitempty" yaml:"-"`
}

// Global checks list. Need to share it with workers and Web UI.
var checks []Check

// The main loop.
func main() {
	// Parse CLI args.
	version := flag.Bool("v", false, "print version to stdout and exit")
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage:", path.Base(os.Args[0]), "path.to/config.yml")
		fmt.Fprintln(os.Stderr, "Docs:  https://github.com/chillum/jsonmon/wiki")
	}
	flag.Parse()
	// -v for version.
	if *version {
		fmt.Println("jsonmon", Version)
		fmt.Println(runtime.Version(), runtime.GOOS, runtime.GOARCH)
		os.Exit(0)
	}
	// Should supply a config file.
	args := flag.Args()
	if len(args) < 1 {
		flag.Usage()
	}
	// Tune concurrency.
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU() + 1)
	}
	// Read config file or exit with error.
	config, err := ioutil.ReadFile(args[0])
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
		os.Exit(3)
	}
	err = yaml.Unmarshal(config, &checks)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", config[0], err)
		os.Exit(3)
	}
	// Run checks.
	var wg sync.WaitGroup
	wg.Add(1)
	for i := range checks {
		go worker(&checks[i])
	}
	// Launch the JSON API.
	host := os.Getenv("HOST")
	port := os.Getenv("PORT")
	if host == "" {
		host = "localhost"
	}
	if port == "" {
		port = "3000"
	}
	http.HandleFunc("/", display)
	err = http.ListenAndServe(host+":"+port, nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
	}
	// Wait forever.
	wg.Wait()
}

// Background worker.
func worker(check *Check) {
	for {
		if check.Repeat == 0 { // Set default timeout.
			check.Repeat = 60
		}
		if check.Tries == 0 { // Default to 1 attempt.
			check.Tries = 1
		}
		if check.Web != "" {
			web(check)
		}
		if check.Shell != "" {
			shell(check)
		}
		time.Sleep(time.Second * time.Duration(check.Repeat))
	}
}

// Shell worker.
func shell(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Shell
	}
	// Execute with shell in N attemps.
	var out []byte
	var err error
	for i := 0; i < check.Tries; i++ {
		out, err = exec.Command("/bin/sh", "-c", check.Shell).CombinedOutput()
		if err == nil {
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Fixed: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Failed: "+name, strings.TrimSpace(string(out)))
		}
	}
}

// Web worker.
func web(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Web
	}
	// Get the URL in N attempts.
	var err error
	for i := 0; i < check.Tries; i++ {
		err = fetch(check.Web)
		if err == nil {
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FIXED: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FAILED: "+name, err.Error())
		}
	}
}

// The actual HTTP GET. Fails if HTTP status code is not 200.
func fetch(url string) error {
	resp, err := http.Get(url)
	if err == nil && resp.StatusCode != 200 {
		err = errors.New(url + " returned " + strconv.Itoa(resp.StatusCode))
	}
	if resp != nil {
		resp.Body.Close()
	}
	return err
}

// Logs and mail alerting.
func alert(mail interface{}, subject string, message string) {
	fmt.Println(subject)
	// Log the alerts.
	if message != "" {
		fmt.Println(message)
	}
	// Mail the alerts.
	if mail != nil {
		// Make the message.
		var rcpt string
		var ok bool
		if rcpt, ok = mail.(string); !ok {
			for i, v := range mail.([]interface{}) {
				if i != 0 {
					rcpt += ", "
				}
				rcpt += v.(string)
			}
		}
		msg := "To: " + rcpt + "\nSubject: " + subject + "\n\n"
		if message != "" {
			msg += message
		}
		msg += "\n.\n"
		// And send it.
		sendmail := exec.Command("/usr/sbin/sendmail", "-t")
		stdin, err := sendmail.StdinPipe()
		err = sendmail.Start()
		if err != nil {
			fmt.Fprintln(os.Stderr, "ERROR:", err)
		}
		io.WriteString(stdin, msg)
		sendmail.Wait()
	}
}

// Format JSON for output.
func display(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" { // Serve for root page only, 404 otherwise.
		http.NotFound(w, r)
		return
	}
	json, _ := json.MarshalIndent(&checks, "", "  ")
	w.Header().Set("Content-Type", "application/json")
	w.Write(json)
	fmt.Fprintln(w, "") // Trailing newline.
}
<commit_msg>increment version<commit_after>/*
Quick and simple monitoring system

Usage:
 jsonmon path.to/config.yml
 jsonmon -v # Prints version to stdout and exits

Environment:
 HOST
  - defaults to localhost
  - the JSON API network interface
 PORT
  - defaults to 3000
  - the JSON API port
 GOMAXPROCS
  - defaults to CPU number + 1
  - number of threads to start
*/
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Application version
const Version = "0.9.8"

// Check details
type Check struct {
	Name   string      `json:"name,omitempty" yaml:"name"`
	Web    string      `json:"web,omitempty" yaml:"web"`
	Shell  string      `json:"shell,omitempty" yaml:"shell"`
	Notify interface{} `json:"-" yaml:"notify"`
	Tries  int         `json:"-" yaml:"tries"`
	Repeat int         `json:"-" yaml:"repeat"`
	Failed bool        `json:"failed" yaml:"-"`
	Since  string      `json:"since,omitempty" yaml:"-"`
}

// Global checks list. Need to share it with workers and Web UI.
var checks []Check

// The main loop.
func main() {
	// Parse CLI args.
	version := flag.Bool("v", false, "print version to stdout and exit")
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage:", path.Base(os.Args[0]), "path.to/config.yml")
		fmt.Fprintln(os.Stderr, "Docs:  https://github.com/chillum/jsonmon/wiki")
	}
	flag.Parse()
	// -v for version.
	if *version {
		fmt.Println("jsonmon", Version)
		fmt.Println(runtime.Version(), runtime.GOOS, runtime.GOARCH)
		os.Exit(0)
	}
	// Should supply a config file.
	args := flag.Args()
	if len(args) < 1 {
		flag.Usage()
	}
	// Tune concurrency.
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU() + 1)
	}
	// Read config file or exit with error.
	config, err := ioutil.ReadFile(args[0])
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
		os.Exit(3)
	}
	err = yaml.Unmarshal(config, &checks)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", config[0], err)
		os.Exit(3)
	}
	// Run checks.
	var wg sync.WaitGroup
	wg.Add(1)
	for i := range checks {
		go worker(&checks[i])
	}
	// Launch the JSON API.
	host := os.Getenv("HOST")
	port := os.Getenv("PORT")
	if host == "" {
		host = "localhost"
	}
	if port == "" {
		port = "3000"
	}
	http.HandleFunc("/", display)
	err = http.ListenAndServe(host+":"+port, nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
	}
	// Wait forever.
	wg.Wait()
}

// Background worker.
func worker(check *Check) {
	for {
		if check.Repeat == 0 { // Set default timeout.
			check.Repeat = 60
		}
		if check.Tries == 0 { // Default to 1 attempt.
			check.Tries = 1
		}
		if check.Web != "" {
			web(check)
		}
		if check.Shell != "" {
			shell(check)
		}
		time.Sleep(time.Second * time.Duration(check.Repeat))
	}
}

// Shell worker.
func shell(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Shell
	}
	// Execute with shell in N attemps.
	var out []byte
	var err error
	for i := 0; i < check.Tries; i++ {
		out, err = exec.Command("/bin/sh", "-c", check.Shell).CombinedOutput()
		if err == nil {
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Fixed: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Failed: "+name, strings.TrimSpace(string(out)))
		}
	}
}

// Web worker.
func web(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Web
	}
	// Get the URL in N attempts.
	var err error
	for i := 0; i < check.Tries; i++ {
		err = fetch(check.Web)
		if err == nil {
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FIXED: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FAILED: "+name, err.Error())
		}
	}
}

// The actual HTTP GET. Fails if HTTP status code is not 200.
func fetch(url string) error {
	resp, err := http.Get(url)
	if err == nil && resp.StatusCode != 200 {
		err = errors.New(url + " returned " + strconv.Itoa(resp.StatusCode))
	}
	if resp != nil {
		resp.Body.Close()
	}
	return err
}

// Logs and mail alerting.
func alert(mail interface{}, subject string, message string) {
	fmt.Println(subject)
	// Log the alerts.
	if message != "" {
		fmt.Println(message)
	}
	// Mail the alerts.
	if mail != nil {
		// Make the message.
		var rcpt string
		var ok bool
		if rcpt, ok = mail.(string); !ok {
			for i, v := range mail.([]interface{}) {
				if i != 0 {
					rcpt += ", "
				}
				rcpt += v.(string)
			}
		}
		msg := "To: " + rcpt + "\nSubject: " + subject + "\n\n"
		if message != "" {
			msg += message
		}
		msg += "\n.\n"
		// And send it.
		sendmail := exec.Command("/usr/sbin/sendmail", "-t")
		stdin, err := sendmail.StdinPipe()
		err = sendmail.Start()
		if err != nil {
			fmt.Fprintln(os.Stderr, "ERROR:", err)
		}
		io.WriteString(stdin, msg)
		sendmail.Wait()
	}
}

// Format JSON for output.
func display(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" { // Serve for root page only, 404 otherwise.
		http.NotFound(w, r)
		return
	}
	json, _ := json.MarshalIndent(&checks, "", "  ")
	w.Header().Set("Content-Type", "application/json")
	w.Write(json)
	fmt.Fprintln(w, "") // Trailing newline.
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	numPollers     = 1
	pollInterval   = 60 * time.Second
	statusInterval = 1 * time.Minute
	errTimeout     = 10 * time.Second
)

var (
	threshold = flag.Uint64("threshold", uint64(1.074*10E11), "Threshold to start serving 500's over HTTP")
	addr      = flag.String("addr", ":8080", "Listen address for HTTP")
	path      = flag.String("path", "/mnt/storage", "The path to query for disk usage")
)

// State represents the last-known state of a path
// This is sent around between the Poller & StateMonitor's channels.
type State struct {
	path  string
	bytes uint64
}

// Disk status storage, with an RWMutex for safe read/write access
// across multiple goroutines
type diskStatus struct {
	state map[string]uint64
	sync.RWMutex
}

// StateMonitor maintains a map that stores the disk usage for paths
// being
// polled, and prints the current state every updateInterval
// nanoseconds.
// It returns a chan State to which resource state should be sent.
// It also serves a HTTP Handler Func for threshold checking. It's
// probably doing too many things! :D
func StateMonitor(updateInterval time.Duration) chan<- State {
	updates := make(chan State)
	diskStatus := &diskStatus{state: make(map[string]uint64)}
	ticker := time.NewTicker(updateInterval)
	go func() {
		for {
			select {
			case <-ticker.C:
				logState(diskStatus)
			case s := <-updates:
				// Write lock
				diskStatus.Lock()
				diskStatus.state[s.path] = s.bytes
				diskStatus.Unlock()
			}
		}
	}()
	go func() {
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			// Read lock
			diskStatus.RLock()
			bytes := diskStatus.state[*path]
			diskStatus.RUnlock()
			switch {
			case bytes == 0:
				http.Error(w, "Disk status not cached yet", http.StatusServiceUnavailable)
			case bytes > *threshold:
				err := fmt.Sprintf("ERROR: Bytes exceed threshold (%v/%v)", bytes, *threshold)
				http.Error(w, err, http.StatusInternalServerError)
			default:
				fmt.Fprintf(w, "OK: %v\n", bytes)
			}
		})
		if err := http.ListenAndServe(*addr, nil); err != nil {
			log.Fatal("ListenAndServe failed: ", err)
		}
	}()
	return updates
}

// logState prints a state map.
func logState(ds *diskStatus) {
	log.Println("Current state:")
	// Read Lock
	ds.RLock()
	for k, v := range ds.state {
		log.Printf(" %s %v", k, v)
	}
	ds.RUnlock()
}

// Path represents a filesystem directory to be polled with du and a
// count of errors when interacting with said path
type Path struct {
	path     string
	errCount int
}

// Poll executes du for a path
// and returns the disk usage in bytes or an error string
func (r *Path) Poll() (bytes uint64) {
	out, err := exec.Command("du", "-sbx", r.path).Output()
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	// Tidy up the line
	s := string(out)
	s = strings.TrimSpace(s)

	// Parse tabulation
	bytes_str := strings.Split(s, "\t")[0]

	// Parse uint64 from string
	bytes, err = strconv.ParseUint(bytes_str, 0, 64)
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	r.errCount = 0
	return bytes
}

// Sleep sleeps for an appropriate interval (dependent on error state)
// before sending the Path to done.
func (r *Path) Sleep(done chan<- *Path) {
	time.Sleep(pollInterval + errTimeout*time.Duration(r.errCount))
	done <- r
}

// Poller pulls paths off the input queue, runs Poll on the paths,
// sends the output along the status channel then sends the path to
// the complete channel
func Poller(in <-chan *Path, out chan<- *Path, status chan<- State) {
	for r := range in {
		bytes := r.Poll()
		status <- State{r.path, bytes}
		out <- r
	}
}

func main() {
	// parse CLI flags
	flag.Parse()

	// Create our input and output channels.
	pending := make(chan *Path)
	complete := make(chan *Path)

	// Launch the StateMonitor.
	status := StateMonitor(statusInterval)
	log.Println("State Monitor started")

	// Launch some Poller goroutines.
	for i := 0; i < numPollers; i++ {
		go Poller(pending, complete, status)
	}

	// Send the path flag to the pending queue
	go func() {
		pending <- &Path{path: *path}
	}()

	for r := range complete {
		go r.Sleep(pending)
	}

}
<commit_msg>go lint (underscores)<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	numPollers     = 1
	pollInterval   = 60 * time.Second
	statusInterval = 1 * time.Minute
	errTimeout     = 10 * time.Second
)

var (
	threshold = flag.Uint64("threshold", uint64(1.074*10E11), "Threshold to start serving 500's over HTTP")
	addr      = flag.String("addr", ":8080", "Listen address for HTTP")
	path      = flag.String("path", "/mnt/storage", "The path to query for disk usage")
)

// State represents the last-known state of a path
// This is sent around between the Poller & StateMonitor's channels.
type State struct {
	path  string
	bytes uint64
}

// Disk status storage, with an RWMutex for safe read/write access
// across multiple goroutines
type diskStatus struct {
	state map[string]uint64
	sync.RWMutex
}

// StateMonitor maintains a map that stores the disk usage for paths
// being
// polled, and prints the current state every updateInterval
// nanoseconds.
// It returns a chan State to which resource state should be sent.
// It also serves a HTTP Handler Func for threshold checking. It's
// probably doing too many things! :D
func StateMonitor(updateInterval time.Duration) chan<- State {
	updates := make(chan State)
	diskStatus := &diskStatus{state: make(map[string]uint64)}
	ticker := time.NewTicker(updateInterval)
	go func() {
		for {
			select {
			case <-ticker.C:
				logState(diskStatus)
			case s := <-updates:
				// Write lock
				diskStatus.Lock()
				diskStatus.state[s.path] = s.bytes
				diskStatus.Unlock()
			}
		}
	}()
	go func() {
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			// Read lock
			diskStatus.RLock()
			bytes := diskStatus.state[*path]
			diskStatus.RUnlock()
			switch {
			case bytes == 0:
				http.Error(w, "Disk status not cached yet", http.StatusServiceUnavailable)
			case bytes > *threshold:
				err := fmt.Sprintf("ERROR: Bytes exceed threshold (%v/%v)", bytes, *threshold)
				http.Error(w, err, http.StatusInternalServerError)
			default:
				fmt.Fprintf(w, "OK: %v\n", bytes)
			}
		})
		if err := http.ListenAndServe(*addr, nil); err != nil {
			log.Fatal("ListenAndServe failed: ", err)
		}
	}()
	return updates
}

// logState prints a state map.
func logState(ds *diskStatus) {
	log.Println("Current state:")
	// Read Lock
	ds.RLock()
	for k, v := range ds.state {
		log.Printf(" %s %v", k, v)
	}
	ds.RUnlock()
}

// Path represents a filesystem directory to be polled with du and a
// count of errors when interacting with said path
type Path struct {
	path     string
	errCount int
}

// Poll executes du for a path
// and returns the disk usage in bytes or an error string
func (r *Path) Poll() (bytes uint64) {
	out, err := exec.Command("du", "-sbx", r.path).Output()
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	// Tidy up the line
	s := string(out)
	s = strings.TrimSpace(s)

	// Parse tabulation
	bytesStr := strings.Split(s, "\t")[0]

	// Parse uint64 from string
	bytes, err = strconv.ParseUint(bytesStr, 0, 64)
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	r.errCount = 0
	return bytes
}

// Sleep sleeps for an appropriate interval (dependent on error state)
// before sending the Path to done.
func (r *Path) Sleep(done chan<- *Path) {
	time.Sleep(pollInterval + errTimeout*time.Duration(r.errCount))
	done <- r
}

// Poller pulls paths off the input queue, runs Poll on the paths,
// sends the output along the status channel then sends the path to
// the complete channel
func Poller(in <-chan *Path, out chan<- *Path, status chan<- State) {
	for r := range in {
		bytes := r.Poll()
		status <- State{r.path, bytes}
		out <- r
	}
}

func main() {
	// parse CLI flags
	flag.Parse()

	// Create our input and output channels.
	pending := make(chan *Path)
	complete := make(chan *Path)

	// Launch the StateMonitor.
	status := StateMonitor(statusInterval)
	log.Println("State Monitor started")

	// Launch some Poller goroutines.
	for i := 0; i < numPollers; i++ {
		go Poller(pending, complete, status)
	}

	// Send the path flag to the pending queue
	go func() {
		pending <- &Path{path: *path}
	}()

	for r := range complete {
		go r.Sleep(pending)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof" // Register pprof handlers with DefaultServeMux.
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"

	auth "github.com/abbot/go-http-auth"

	"github.com/buchgr/bazel-remote/cache/disk"

	"github.com/buchgr/bazel-remote/config"
	"github.com/buchgr/bazel-remote/server"
	"github.com/buchgr/bazel-remote/utils/flags"
	"github.com/buchgr/bazel-remote/utils/idle"
	"github.com/buchgr/bazel-remote/utils/rlimit"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	httpmetrics "github.com/slok/go-http-metrics/metrics/prometheus"
	middleware "github.com/slok/go-http-metrics/middleware"
	middlewarestd "github.com/slok/go-http-metrics/middleware/std"
	"github.com/urfave/cli/v2"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/semaphore"
)

// gitCommit is the version stamp for the server. The value of this var
// is set through linker options.
var gitCommit string

func main() {
	log.SetFlags(config.LogFlags)

	maybeGitCommitMsg := ""
	if len(gitCommit) > 0 && gitCommit != "{STABLE_GIT_COMMIT}" {
		maybeGitCommitMsg = fmt.Sprintf(" from git commit %s", gitCommit)
	}
	log.Printf("bazel-remote built with %s%s.",
		runtime.Version(), maybeGitCommitMsg)

	app := cli.NewApp()

	cli.AppHelpTemplate = flags.Template
	cli.HelpPrinterCustom = flags.HelpPrinter
	// Force the use of cli.HelpPrinterCustom.
	app.ExtraInfo = func() map[string]string { return map[string]string{} }

	app.Flags = flags.GetCliFlags()
	app.Action = run

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal("bazel-remote terminated:", err)
	}
}

func run(ctx *cli.Context) error {
	c, err := config.Get(ctx)
	if err != nil {
		fmt.Fprintf(ctx.App.Writer, "%v\n\n", err)
		_ = cli.ShowAppHelp(ctx)
		return cli.Exit(err.Error(), 1)
	}

	if ctx.NArg() > 0 {
		fmt.Fprintf(ctx.App.Writer,
			"Error: bazel-remote does not take positional aguments\n")
		for i := 0; i < ctx.NArg(); i++ {
			fmt.Fprintf(ctx.App.Writer, "arg: %s\n", ctx.Args().Get(i))
		}
		fmt.Fprintf(ctx.App.Writer, "\n")

		_ = cli.ShowAppHelp(ctx)
		os.Exit(1)
	}

	rlimit.Raise()

	grpcSem := semaphore.NewWeighted(1)
	var grpcServer *grpc.Server

	httpSem := semaphore.NewWeighted(1)
	var httpServer *http.Server

	idleTimeoutChan := make(chan struct{}, 1)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	go func() {
		select {
		case sig := <-sigChan:
			log.Printf("Received signal: %s, attempting graceful shutdown", sig)
		case <-idleTimeoutChan:
			log.Println("Idle timeout reached, attempting graceful shutdown")
		}

		go func() {
			if !grpcSem.TryAcquire(1) {
				if grpcServer != nil {
					log.Println("Stopping gRPC server")
					grpcServer.GracefulStop()
					log.Println("gRPC server stopped")
				}
			}
		}()

		go func() {
			if !httpSem.TryAcquire(1) {
				if httpServer != nil {
					log.Println("Stopping HTTP server")
					err := httpServer.Shutdown(context.Background())
					if err != nil {
						log.Println("Error occurred while stopping HTTP server:", err)
					} else {
						log.Println("HTTP server stopped")
					}
				}
			}
		}()
	}()

	log.Println("Storage mode:", c.StorageMode)
	if c.StorageMode == "zstd" {
		log.Println("Zstandard implementation:", c.ZstdImplementation)
	}

	opts := []disk.Option{
		disk.WithStorageMode(c.StorageMode),
		disk.WithZstdImplementation(c.ZstdImplementation),
		disk.WithMaxBlobSize(c.MaxBlobSize),
		disk.WithProxyMaxBlobSize(c.MaxProxyBlobSize),
		disk.WithAccessLogger(c.AccessLogger),
	}
	if c.ProxyBackend != nil {
		opts = append(opts, disk.WithProxyBackend(c.ProxyBackend))
	}
	if c.EnableEndpointMetrics {
		opts = append(opts, disk.WithEndpointMetrics())
	}

	diskCache, err := disk.New(c.Dir, int64(c.MaxSize)*1024*1024*1024, opts...)
	if err != nil {
		log.Fatal(err)
	}
	diskCache.RegisterMetrics()

	servers := new(errgroup.Group)

	var htpasswdSecrets auth.SecretProvider

	authMode := "disabled"
	if c.HtpasswdFile != "" {
		authMode = "basic"
		htpasswdSecrets = auth.HtpasswdFileProvider(c.HtpasswdFile)
	} else if c.TLSCaFile != "" {
		authMode = "mTLS"
	}
	log.Println("Authentication:", authMode)

	if authMode != "disabled" {
		if c.AllowUnauthenticatedReads {
			log.Println("Access mode: authentication required for writes, unauthenticated reads allowed")
		} else {
			log.Println("Access mode: authentication required")
		}
	}

	var idleTimer *idle.Timer
	if c.IdleTimeout > 0 {
		idleTimer = idle.NewTimer(c.IdleTimeout, idleTimeoutChan)
	}

	acKeyManglingStatus := "disabled"
	if c.EnableACKeyInstanceMangling {
		acKeyManglingStatus = "enabled"
	}
	log.Println("Mangling non-empty instance names with AC keys:", acKeyManglingStatus)

	servers.Go(func() error {
		return startHttpServer(c, &httpServer, htpasswdSecrets, idleTimer, httpSem, diskCache)
	})

	if c.GRPCAddress != "none" {
		servers.Go(func() error {
			return startGrpcServer(c, &grpcServer, htpasswdSecrets, idleTimer, grpcSem, diskCache)
		})
	}

	if c.ProfileAddress != "" {
		go func() {
			// Allow access to /debug/pprof/ URLs.
			log.Printf("Starting HTTP server for profiling on address %s",
				c.ProfileAddress)
			log.Fatal(`Failed to listen on address: "`, c.ProfileAddress,
				`": `, http.ListenAndServe(c.ProfileAddress, nil))
		}()
	}

	if idleTimer != nil {
		log.Printf("Starting idle timer with value %v", c.IdleTimeout)
		idleTimer.Start()
	}

	return servers.Wait()
}

func startHttpServer(c *config.Config, httpServer **http.Server,
	htpasswdSecrets auth.SecretProvider, idleTimer *idle.Timer,
	httpSem *semaphore.Weighted, diskCache disk.Cache) error {

	mux := http.NewServeMux()
	*httpServer = &http.Server{
		Handler:      mux,
		ReadTimeout:  c.HTTPReadTimeout,
		TLSConfig:    c.TLSConfig,
		WriteTimeout: c.HTTPWriteTimeout,
	}

	checkClientCertForWrites := c.TLSCaFile != "" && c.AllowUnauthenticatedReads
	validateAC := !c.DisableHTTPACValidation
	h := server.NewHTTPCache(diskCache, c.AccessLogger, c.ErrorLogger, validateAC,
		c.EnableACKeyInstanceMangling, checkClientCertForWrites, gitCommit)

	cacheHandler := h.CacheHandler
	if c.HtpasswdFile != "" {
		if c.AllowUnauthenticatedReads {
			cacheHandler = unauthenticatedReadWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
		} else {
			cacheHandler = authWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
		}
	}

	if c.IdleTimeout > 0 {
		cacheHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			idleTimer.ResetTimer()
			cacheHandler(w, r)
		})
	}

	if c.EnableEndpointMetrics {
		metricsMdlw := middleware.New(middleware.Config{
			Recorder: httpmetrics.NewRecorder(httpmetrics.Config{
				DurationBuckets: c.MetricsDurationBuckets,
			}),
		})
		mux.Handle("/metrics", middlewarestd.Handler("metrics", metricsMdlw,
			promhttp.Handler()))
		mux.Handle("/status", middlewarestd.Handler("status", metricsMdlw,
			http.HandlerFunc(h.StatusPageHandler)))
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			middlewarestd.Handler(r.Method, metricsMdlw, http.HandlerFunc(cacheHandler)).ServeHTTP(w, r)
		})
	} else {
		mux.HandleFunc("/status", h.StatusPageHandler)
		mux.HandleFunc("/", cacheHandler)
	}

	var ln net.Listener
	var err error
	if strings.HasPrefix(c.HTTPAddress, "unix://") {
		ln, err = net.Listen("unix", c.HTTPAddress[len("unix://"):])
	} else {
		ln, err = net.Listen("tcp", c.HTTPAddress)
	}
	if err != nil {
		log.Fatal(`Failed to listen on address: "`, c.HTTPAddress, `": `, err)
	}

	validateStatus := "disabled"
	if validateAC {
		validateStatus = "enabled"
	}
	log.Println("HTTP AC validation:", validateStatus)

	if len(c.TLSCertFile) > 0 && len(c.TLSKeyFile) > 0 {
		log.Printf("Starting HTTPS server on address %s", (*httpServer).Addr)

		if !httpSem.TryAcquire(1) {
			log.Println("bazel-remote is shutting down, not starting HTTPS server")
			return nil
		}

		err = (*httpServer).ServeTLS(ln, c.TLSCertFile, c.TLSKeyFile)
		if err == http.ErrServerClosed {
			log.Println("HTTPS server stopped")
			return nil
		}

		return err
	}

	log.Printf("Starting HTTP server on address %s", c.HTTPAddress)

	if !httpSem.TryAcquire(1) {
		log.Println("bazel-remote is shutting down, not starting HTTP server")
		return nil
	}

	err = (*httpServer).Serve(ln)
	if err == http.ErrServerClosed {
		return nil
	}

	return err
}

func startGrpcServer(c *config.Config, grpcServer **grpc.Server,
	htpasswdSecrets auth.SecretProvider, idleTimer *idle.Timer,
	grpcSem *semaphore.Weighted, diskCache disk.Cache) error {

	opts := []grpc.ServerOption{}
	streamInterceptors := []grpc.StreamServerInterceptor{}
	unaryInterceptors := []grpc.UnaryServerInterceptor{}

	if c.EnableEndpointMetrics {
		streamInterceptors = append(streamInterceptors, grpc_prometheus.StreamServerInterceptor)
		unaryInterceptors = append(unaryInterceptors, grpc_prometheus.UnaryServerInterceptor)
		grpc_prometheus.EnableHandlingTimeHistogram(grpc_prometheus.WithHistogramBuckets(c.MetricsDurationBuckets))
	}

	if c.TLSConfig != nil {
		opts = append(opts, grpc.Creds(credentials.NewTLS(c.TLSConfig)))

		if c.TLSCaFile != "" {
			streamInterceptors = append(streamInterceptors,
				server.GRPCmTLSStreamServerInterceptor(c.AllowUnauthenticatedReads))
			unaryInterceptors = append(unaryInterceptors,
				server.GRPCmTLSUnaryServerInterceptor(c.AllowUnauthenticatedReads))
		}
	}

	if htpasswdSecrets != nil {
		gba := server.NewGrpcBasicAuth(htpasswdSecrets, c.AllowUnauthenticatedReads)
		streamInterceptors = append(streamInterceptors, gba.StreamServerInterceptor)
		unaryInterceptors = append(unaryInterceptors, gba.UnaryServerInterceptor)
	}

	if idleTimer != nil {
		it := server.NewGrpcIdleTimer(idleTimer)
		streamInterceptors = append(streamInterceptors, it.StreamServerInterceptor)
		unaryInterceptors = append(unaryInterceptors, it.UnaryServerInterceptor)
	}

	opts = append(opts, grpc.ChainStreamInterceptor(streamInterceptors...))
	opts = append(opts, grpc.ChainUnaryInterceptor(unaryInterceptors...))

	validateAC := !c.DisableGRPCACDepsCheck
	validateStatus := "disabled"
	if validateAC {
		validateStatus = "enabled"
	}
	log.Println("gRPC AC dependency checks:", validateStatus)

	enableRemoteAssetAPI := c.ExperimentalRemoteAssetAPI
	remoteAssetStatus := "disabled"
	if enableRemoteAssetAPI {
		remoteAssetStatus = "enabled"
	}
	log.Println("experimental gRPC remote asset API:", remoteAssetStatus)

	network := "tcp"
	addr := c.GRPCAddress
	if strings.HasPrefix(c.GRPCAddress, "unix://") {
		network = "unix"
		addr = c.GRPCAddress[len("unix://"):]
	}

	*grpcServer = grpc.NewServer(opts...)

	if !grpcSem.TryAcquire(1) {
		log.Println("bazel-remote is shutting down, not starting gRPC server")
		return nil
	}

	log.Println("Starting gRPC server on address", addr)
	err := server.ListenAndServeGRPC(*grpcServer,
		network, addr,
		validateAC,
		c.EnableACKeyInstanceMangling,
		enableRemoteAssetAPI,
		diskCache, c.AccessLogger, c.ErrorLogger)
	if err != nil {
		log.Fatal("gRPC server returned fatal error:", err)
	}

	return nil
}

// A http.HandlerFunc wrapper which requires successful basic
// authentication for all requests.
func authWrapper(handler http.HandlerFunc, secrets auth.SecretProvider, addr string) http.HandlerFunc {
	authenticator := &auth.BasicAuth{Realm: addr, Secrets: secrets}
	return auth.JustCheck(authenticator, handler)
}

// A http.HandlerFunc wrapper which requires successful basic
// authentication for write requests, but allows unauthenticated
// read requests.
func unauthenticatedReadWrapper(handler http.HandlerFunc, secrets auth.SecretProvider, addr string) http.HandlerFunc {
	authenticator := &auth.BasicAuth{Realm: addr, Secrets: secrets}

	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet || r.Method == http.MethodHead {
			handler(w, r)
			return
		}

		if authenticator.CheckAuth(r) != "" {
			handler(w, r)
			return
		}

		http.Error(w, "Authorization required", http.StatusUnauthorized)
		// TODO: pass in a logger so we can log this event?
	}
}
<commit_msg>Make startHttpServer log messages consistent with startGrpcServer<commit_after>package main

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof" // Register pprof handlers with DefaultServeMux.
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"

	auth "github.com/abbot/go-http-auth"

	"github.com/buchgr/bazel-remote/cache/disk"

	"github.com/buchgr/bazel-remote/config"
	"github.com/buchgr/bazel-remote/server"
	"github.com/buchgr/bazel-remote/utils/flags"
	"github.com/buchgr/bazel-remote/utils/idle"
	"github.com/buchgr/bazel-remote/utils/rlimit"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	httpmetrics "github.com/slok/go-http-metrics/metrics/prometheus"
	middleware "github.com/slok/go-http-metrics/middleware"
	middlewarestd "github.com/slok/go-http-metrics/middleware/std"
	"github.com/urfave/cli/v2"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/semaphore"
)

// gitCommit is the version stamp for the server. The value of this var
// is set through linker options.
var gitCommit string

func main() {
	log.SetFlags(config.LogFlags)

	maybeGitCommitMsg := ""
	if len(gitCommit) > 0 && gitCommit != "{STABLE_GIT_COMMIT}" {
		maybeGitCommitMsg = fmt.Sprintf(" from git commit %s", gitCommit)
	}
	log.Printf("bazel-remote built with %s%s.",
		runtime.Version(), maybeGitCommitMsg)

	app := cli.NewApp()

	cli.AppHelpTemplate = flags.Template
	cli.HelpPrinterCustom = flags.HelpPrinter
	// Force the use of cli.HelpPrinterCustom.
	app.ExtraInfo = func() map[string]string { return map[string]string{} }

	app.Flags = flags.GetCliFlags()
	app.Action = run

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal("bazel-remote terminated:", err)
	}
}

func run(ctx *cli.Context) error {
	c, err := config.Get(ctx)
	if err != nil {
		fmt.Fprintf(ctx.App.Writer, "%v\n\n", err)
		_ = cli.ShowAppHelp(ctx)
		return cli.Exit(err.Error(), 1)
	}

	if ctx.NArg() > 0 {
		fmt.Fprintf(ctx.App.Writer,
			"Error: bazel-remote does not take positional aguments\n")
		for i := 0; i < ctx.NArg(); i++ {
			fmt.Fprintf(ctx.App.Writer, "arg: %s\n", ctx.Args().Get(i))
		}
		fmt.Fprintf(ctx.App.Writer, "\n")

		_ = cli.ShowAppHelp(ctx)
		os.Exit(1)
	}

	rlimit.Raise()

	grpcSem := semaphore.NewWeighted(1)
	var grpcServer *grpc.Server

	httpSem := semaphore.NewWeighted(1)
	var httpServer *http.Server

	idleTimeoutChan := make(chan struct{}, 1)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	go func() {
		select {
		case sig := <-sigChan:
			log.Printf("Received signal: %s, attempting graceful shutdown", sig)
		case <-idleTimeoutChan:
			log.Println("Idle timeout reached, attempting graceful shutdown")
		}

		go func() {
			if !grpcSem.TryAcquire(1) {
				if grpcServer != nil {
					log.Println("Stopping gRPC server")
					grpcServer.GracefulStop()
					log.Println("gRPC server stopped")
				}
			}
		}()

		go func() {
			if !httpSem.TryAcquire(1) {
				if httpServer != nil {
					log.Println("Stopping HTTP server")
					err := httpServer.Shutdown(context.Background())
					if err != nil {
						log.Println("Error occurred while stopping HTTP server:", err)
					} else {
						log.Println("HTTP server stopped")
					}
				}
			}
		}()
	}()

	log.Println("Storage mode:", c.StorageMode)
	if c.StorageMode == "zstd" {
		log.Println("Zstandard implementation:", c.ZstdImplementation)
	}

	opts := []disk.Option{
		disk.WithStorageMode(c.StorageMode),
		disk.WithZstdImplementation(c.ZstdImplementation),
		disk.WithMaxBlobSize(c.MaxBlobSize),
		disk.WithProxyMaxBlobSize(c.MaxProxyBlobSize),
		disk.WithAccessLogger(c.AccessLogger),
	}
	if c.ProxyBackend != nil {
		opts = append(opts, disk.WithProxyBackend(c.ProxyBackend))
	}
	if c.EnableEndpointMetrics {
		opts = append(opts, disk.WithEndpointMetrics())
	}

	diskCache, err := disk.New(c.Dir, int64(c.MaxSize)*1024*1024*1024, opts...)
	if err != nil {
		log.Fatal(err)
	}
	diskCache.RegisterMetrics()

	servers := new(errgroup.Group)

	var htpasswdSecrets auth.SecretProvider

	authMode := "disabled"
	if c.HtpasswdFile != "" {
		authMode = "basic"
		htpasswdSecrets = auth.HtpasswdFileProvider(c.HtpasswdFile)
	} else if c.TLSCaFile != "" {
		authMode = "mTLS"
	}
	log.Println("Authentication:", authMode)

	if authMode != "disabled" {
		if c.AllowUnauthenticatedReads {
			log.Println("Access mode: authentication required for writes, unauthenticated reads allowed")
		} else {
			log.Println("Access mode: authentication required")
		}
	}

	var idleTimer *idle.Timer
	if c.IdleTimeout > 0 {
		idleTimer = idle.NewTimer(c.IdleTimeout, idleTimeoutChan)
	}

	acKeyManglingStatus := "disabled"
	if c.EnableACKeyInstanceMangling {
		acKeyManglingStatus = "enabled"
	}
	log.Println("Mangling non-empty instance names with AC keys:", acKeyManglingStatus)

	servers.Go(func() error {
		return startHttpServer(c, &httpServer, htpasswdSecrets, idleTimer, httpSem, diskCache)
	})

	if c.GRPCAddress != "none" {
		servers.Go(func() error {
			return startGrpcServer(c, &grpcServer, htpasswdSecrets, idleTimer, grpcSem, diskCache)
		})
	}

	if c.ProfileAddress != "" {
		go func() {
			// Allow access to /debug/pprof/ URLs.
			log.Printf("Starting HTTP server for profiling on address %s",
				c.ProfileAddress)
			log.Fatal(`Failed to listen on address: "`, c.ProfileAddress,
				`": `, http.ListenAndServe(c.ProfileAddress, nil))
		}()
	}

	if idleTimer != nil {
		log.Printf("Starting idle timer with value %v", c.IdleTimeout)
		idleTimer.Start()
	}

	return servers.Wait()
}

func startHttpServer(c *config.Config, httpServer **http.Server,
	htpasswdSecrets auth.SecretProvider, idleTimer *idle.Timer,
	httpSem *semaphore.Weighted, diskCache disk.Cache) error {

	mux := http.NewServeMux()
	*httpServer = &http.Server{
		Handler:      mux,
		ReadTimeout:  c.HTTPReadTimeout,
		TLSConfig:    c.TLSConfig,
		WriteTimeout: c.HTTPWriteTimeout,
	}

	checkClientCertForWrites := c.TLSCaFile != "" && c.AllowUnauthenticatedReads
	validateAC := !c.DisableHTTPACValidation
	h := server.NewHTTPCache(diskCache, c.AccessLogger, c.ErrorLogger, validateAC,
		c.EnableACKeyInstanceMangling, checkClientCertForWrites, gitCommit)

	cacheHandler := h.CacheHandler
	if c.HtpasswdFile != "" {
		if c.AllowUnauthenticatedReads {
			cacheHandler = unauthenticatedReadWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
		} else {
			cacheHandler = authWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
		}
	}

	if c.IdleTimeout > 0 {
		cacheHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			idleTimer.ResetTimer()
			cacheHandler(w, r)
		})
	}

	if c.EnableEndpointMetrics {
		metricsMdlw := middleware.New(middleware.Config{
			Recorder: httpmetrics.NewRecorder(httpmetrics.Config{
				DurationBuckets: c.MetricsDurationBuckets,
			}),
		})
		mux.Handle("/metrics", middlewarestd.Handler("metrics", metricsMdlw,
			promhttp.Handler()))
		mux.Handle("/status", middlewarestd.Handler("status", metricsMdlw,
			http.HandlerFunc(h.StatusPageHandler)))
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			middlewarestd.Handler(r.Method, metricsMdlw, http.HandlerFunc(cacheHandler)).ServeHTTP(w, r)
		})
	} else {
		mux.HandleFunc("/status", h.StatusPageHandler)
		mux.HandleFunc("/", cacheHandler)
	}

	var ln net.Listener
	var err error
	if strings.HasPrefix(c.HTTPAddress, "unix://") {
		ln, err = net.Listen("unix", c.HTTPAddress[len("unix://"):])
	} else {
		ln, err = net.Listen("tcp", c.HTTPAddress)
	}
	if err != nil {
		log.Fatal(`Failed to listen on address: "`, c.HTTPAddress, `": `, err)
	}

	validateStatus := "disabled"
	if validateAC {
		validateStatus = "enabled"
	}
	log.Println("HTTP AC validation:", validateStatus)

	if len(c.TLSCertFile) > 0 && len(c.TLSKeyFile) > 0 {
		if !httpSem.TryAcquire(1) {
			log.Println("bazel-remote is shutting down, not starting HTTPS server")
			return nil
		}

		log.Printf("Starting HTTPS server on address %s", (*httpServer).Addr)
		err = (*httpServer).ServeTLS(ln, c.TLSCertFile, c.TLSKeyFile)
		if err == http.ErrServerClosed {
			log.Println("HTTPS server stopped")
			return nil
		}

		return err
	}

	if !httpSem.TryAcquire(1) {
		log.Println("bazel-remote is shutting down, not starting HTTP server")
		return nil
	}

	log.Printf("Starting HTTP server on address %s", c.HTTPAddress)
	err = (*httpServer).Serve(ln)
	if err == http.ErrServerClosed {
		return nil
	}

	return err
}

func startGrpcServer(c *config.Config, grpcServer **grpc.Server,
	htpasswdSecrets auth.SecretProvider, idleTimer *idle.Timer,
	grpcSem *semaphore.Weighted, diskCache disk.Cache) error {

	opts := []grpc.ServerOption{}
	streamInterceptors := []grpc.StreamServerInterceptor{}
	unaryInterceptors := []grpc.UnaryServerInterceptor{}

	if c.EnableEndpointMetrics {
		streamInterceptors = append(streamInterceptors, grpc_prometheus.StreamServerInterceptor)
		unaryInterceptors = append(unaryInterceptors, grpc_prometheus.UnaryServerInterceptor)
		grpc_prometheus.EnableHandlingTimeHistogram(grpc_prometheus.WithHistogramBuckets(c.MetricsDurationBuckets))
	}

	if c.TLSConfig != nil {
		opts = append(opts, grpc.Creds(credentials.NewTLS(c.TLSConfig)))

		if c.TLSCaFile != "" {
			streamInterceptors = append(streamInterceptors,
				server.GRPCmTLSStreamServerInterceptor(c.AllowUnauthenticatedReads))
			unaryInterceptors = append(unaryInterceptors,
				server.GRPCmTLSUnaryServerInterceptor(c.AllowUnauthenticatedReads))
		}
	}

	if htpasswdSecrets != nil {
		gba := server.NewGrpcBasicAuth(htpasswdSecrets, c.AllowUnauthenticatedReads)
		streamInterceptors = append(streamInterceptors, gba.StreamServerInterceptor)
		unaryInterceptors = append(unaryInterceptors, gba.UnaryServerInterceptor)
	}

	if idleTimer != nil {
		it := server.NewGrpcIdleTimer(idleTimer)
		streamInterceptors = append(streamInterceptors, it.StreamServerInterceptor)
		unaryInterceptors = append(unaryInterceptors, it.UnaryServerInterceptor)
	}

	opts = append(opts, grpc.ChainStreamInterceptor(streamInterceptors...))
	opts = append(opts, grpc.ChainUnaryInterceptor(unaryInterceptors...))

	validateAC := !c.DisableGRPCACDepsCheck
	validateStatus := "disabled"
	if validateAC {
		validateStatus = "enabled"
	}
	log.Println("gRPC AC dependency checks:", validateStatus)

	enableRemoteAssetAPI := c.ExperimentalRemoteAssetAPI
	remoteAssetStatus := "disabled"
	if enableRemoteAssetAPI {
		remoteAssetStatus = "enabled"
	}
	log.Println("experimental gRPC remote asset API:", remoteAssetStatus)

	network := "tcp"
	addr := c.GRPCAddress
	if strings.HasPrefix(c.GRPCAddress, "unix://") {
		network = "unix"
		addr = c.GRPCAddress[len("unix://"):]
	}

	*grpcServer = grpc.NewServer(opts...)

	if !grpcSem.TryAcquire(1) {
		log.Println("bazel-remote is shutting down, not starting gRPC server")
		return nil
	}

	log.Println("Starting gRPC server on address", addr)
	err := server.ListenAndServeGRPC(*grpcServer,
		network, addr,
		validateAC,
		c.EnableACKeyInstanceMangling,
		enableRemoteAssetAPI,
		diskCache, c.AccessLogger, c.ErrorLogger)
	if err != nil {
		log.Fatal("gRPC server returned fatal error:", err)
	}

	return nil
}

// A http.HandlerFunc wrapper which requires successful basic
// authentication for all requests.
func authWrapper(handler http.HandlerFunc, secrets auth.SecretProvider, addr string) http.HandlerFunc {
	authenticator := &auth.BasicAuth{Realm: addr, Secrets: secrets}
	return auth.JustCheck(authenticator, handler)
}

// A http.HandlerFunc wrapper which requires successful basic
// authentication for write requests, but allows unauthenticated
// read requests.
func unauthenticatedReadWrapper(handler http.HandlerFunc, secrets auth.SecretProvider, addr string) http.HandlerFunc {
	authenticator := &auth.BasicAuth{Realm: addr, Secrets: secrets}

	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet || r.Method == http.MethodHead {
			handler(w, r)
			return
		}

		if authenticator.CheckAuth(r) != "" {
			handler(w, r)
			return
		}

		http.Error(w, "Authorization required", http.StatusUnauthorized)
		// TODO: pass in a logger so we can log this event?
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/geofffranks/simpleyaml" // FIXME: switch back to smallfish/simpleyaml after https://github.com/smallfish/simpleyaml/pull/1 is merged
	"github.com/voxelbrain/goptions"
	"gopkg.in/yaml.v2"
	"regexp"
	"strings"
)

// Current version of spruce
var VERSION = "0.7.0" // SED MARKER FOR AUTO VERSION BUMPING

var printfStdOut = func(format string, args ...interface{}) {
	fmt.Fprintf(os.Stdout, format, args...)
}

var printfStdErr = func(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args...)
}

var getopts = func(o interface{}) {
	err := goptions.Parse(o)
	if err != nil {
		usage()
	}
}

var exit = func(code int) {
	os.Exit(code)
}

var usage = func() {
	goptions.PrintHelp()
	exit(1)
}

var debug bool
var handleConcourseQuoting bool

// DEBUG - Prints out a debug message
func DEBUG(format string, args ...interface{}) {
	if debug {
		content := fmt.Sprintf(format, args...)
		lines := strings.Split(content, "\n")
		for i, line := range lines {
			lines[i] = "DEBUG> " + line
		}
		content = strings.Join(lines, "\n")
		printfStdErr("%s\n", content)
	}
}

func main() {
	var options struct {
		Debug     bool `goptions:"-D, --debug, description='Enable debugging'"`
		Version   bool `goptions:"-v, --version, description='Display version information'"`
		Concourse bool `goptions:"--concourse, description='Pre/Post-process YAML for Concourse CI (handles {{ }} quoting)'"`
		Action    goptions.Verbs
		Merge     struct {
			Prune []string           `goptions:"--prune, description='Specify keys to prune from final output (may be specified more than once'"`
			Files goptions.Remainder `goptions:"description='Merges file2.yml through fileN.yml on top of file1.yml'"`
		} `goptions:"merge"`
	}
	getopts(&options)

	if os.Getenv("DEBUG") != "" && strings.ToLower(os.Getenv("DEBUG")) != "false" && os.Getenv("DEBUG") != "0" {
		debug = true
	}
	if options.Debug {
		debug = options.Debug
	}

	handleConcourseQuoting = options.Concourse

	if options.Version {
		printfStdErr("%s - Version %s\n", os.Args[0], VERSION)
		exit(0)
	} else {
		switch {
		case options.Action == "merge":
			if len(options.Merge.Files) >= 1 {
				root := make(map[interface{}]interface{})

				err := mergeAllDocs(root, options.Merge.Files)
				if err != nil {
					printfStdErr("%s\n", err.Error())
					exit(2)
				}

				m := &Merger{}
				m.Visit(root, &StaticIPGenerator{root: root})
				m.Visit(root, &DeReferencer{root: root})
				m.Visit(root, &Concatenator{root: root})

				for _, toPrune := range options.Merge.Prune {
					DEBUG("%s: Pruning", toPrune)
					keys := strings.Split(toPrune, ".")
					if len(keys) > 1 {
						parent := strings.Join(keys[0:(len(keys)-1)], ".")
						target := keys[(len(keys) - 1)]
						node, _ := resolveNode(parent, root)
						if _, ok := node.(map[interface{}]interface{}); ok {
							DEBUG("   PRUNED %q from %q", target, parent)
							delete(node.(map[interface{}]interface{}), target)
						}
					} else {
						DEBUG("   PRUNED %q from \"$\"", keys[0])
						delete(root, keys[0])
					}
				}

				m.Visit(root, &ParamChecker{})

				if err := m.Error(); err != nil {
					printfStdErr("%s\n", err.Error())
					exit(2)

				} else {
					DEBUG("Converting the following data back to YML:")
					DEBUG("%#v", root)
					merged, err := yaml.Marshal(root)
					if err != nil {
						printfStdErr("Unable to convert merged result back to YAML: %s\nData:\n%#v", err.Error(), root)
						exit(2)
					} else {
						var output string
						if handleConcourseQuoting {
							output = dequoteConcourse(merged)
						} else {
							output = string(merged)
						}
						printfStdOut("%s\n", output)
					}
				}
			} else {
				usage()
			}
		default:
			usage()
		}
	}
}

func parseYAML(data []byte) (map[interface{}]interface{}, error) {
	y, err := simpleyaml.NewYaml(data)
	if err != nil {
		return nil, err
	}

	doc, err := y.Map()
	if err != nil {
		return nil, fmt.Errorf("Root of YAML document is not a hash/map: %s\n", err.Error())
	}

	return doc, nil
}

func mergeAllDocs(root map[interface{}]interface{}, paths []string) error {
	m := &Merger{}

	for _, path := range paths {
		DEBUG("Processing file '%s'", path)
		data, err := readFile(path)
		if err != nil {
			return fmt.Errorf("Error reading file %s: %s\n", path, err.Error())
		}

		if handleConcourseQuoting {
			data = quoteConcourse(data)
		}

		doc, err := parseYAML(data)
		if err != nil {
			return fmt.Errorf("%s: %s\n", path, err.Error())
		}

		m.Merge(root, doc)

		tmpYaml, _ := yaml.Marshal(root) // we don't care about errors for debugging
		DEBUG("Current data after processing '%s':\n%s", path, tmpYaml)
	}

	return m.Error()
}

var concourseRegex = `\{\{([-\w\p{L}]+)\}\}`

func quoteConcourse(input []byte) []byte {
	re := regexp.MustCompile("(" + concourseRegex + ")")
	return re.ReplaceAll(input, []byte("\"$1\""))
}

func dequoteConcourse(input []byte) string {
	re := regexp.MustCompile("['\"](" + concourseRegex + ")[\"']")
	return re.ReplaceAllString(string(input), "$1")
}
<commit_msg>Bump version to 0.7.1 to troubleshoot build something-something<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/geofffranks/simpleyaml" // FIXME: switch back to smallfish/simpleyaml after https://github.com/smallfish/simpleyaml/pull/1 is merged
	"github.com/voxelbrain/goptions"
	"gopkg.in/yaml.v2"
	"regexp"
	"strings"
)

// Current version of spruce
var VERSION = "0.7.1" // SED MARKER FOR AUTO VERSION BUMPING

var printfStdOut = func(format string, args ...interface{}) {
	fmt.Fprintf(os.Stdout, format, args...)
}

var printfStdErr = func(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args...)
}

var getopts = func(o interface{}) {
	err := goptions.Parse(o)
	if err != nil {
		usage()
	}
}

var exit = func(code int) {
	os.Exit(code)
}

var usage = func() {
	goptions.PrintHelp()
	exit(1)
}

var debug bool
var handleConcourseQuoting bool

// DEBUG - Prints out a debug message
func DEBUG(format string, args ...interface{}) {
	if debug {
		content := fmt.Sprintf(format, args...)
		lines := strings.Split(content, "\n")
		for i, line := range lines {
			lines[i] = "DEBUG> " + line
		}
		content = strings.Join(lines, "\n")
		printfStdErr("%s\n", content)
	}
}

func main() {
	var options struct {
		Debug     bool `goptions:"-D, --debug, description='Enable debugging'"`
		Version   bool `goptions:"-v, --version, description='Display version information'"`
		Concourse bool `goptions:"--concourse, description='Pre/Post-process YAML for Concourse CI (handles {{ }} quoting)'"`
		Action    goptions.Verbs
		Merge     struct {
			Prune []string           `goptions:"--prune, description='Specify keys to prune from final output (may be specified more than once'"`
			Files goptions.Remainder `goptions:"description='Merges file2.yml through fileN.yml on top of file1.yml'"`
		} `goptions:"merge"`
	}
	getopts(&options)

	if os.Getenv("DEBUG") != "" && strings.ToLower(os.Getenv("DEBUG")) != "false" && os.Getenv("DEBUG") != "0" {
		debug = true
	}
	if options.Debug {
		debug = options.Debug
	}

	handleConcourseQuoting = options.Concourse

	if options.Version {
		printfStdErr("%s - Version %s\n", os.Args[0], VERSION)
		exit(0)
	} else {
		switch {
		case options.Action == "merge":
			if len(options.Merge.Files) >= 1 {
				root := make(map[interface{}]interface{})

				err := mergeAllDocs(root, options.Merge.Files)
				if err != nil {
					printfStdErr("%s\n", err.Error())
					exit(2)
				}

				m := &Merger{}
				m.Visit(root, &StaticIPGenerator{root: root})
				m.Visit(root, &DeReferencer{root: root})
				m.Visit(root, &Concatenator{root: root})

				for _, toPrune := range options.Merge.Prune {
					DEBUG("%s: Pruning", toPrune)
					keys := strings.Split(toPrune, ".")
					if len(keys) > 1 {
						parent := strings.Join(keys[0:(len(keys)-1)], ".")
						target := keys[(len(keys) - 1)]
						node, _ := resolveNode(parent, root)
						if _, ok := node.(map[interface{}]interface{}); ok {
							DEBUG("   PRUNED %q from %q", target, parent)
							delete(node.(map[interface{}]interface{}), target)
						}
					} else {
						DEBUG("   PRUNED %q from \"$\"", keys[0])
						delete(root, keys[0])
					}
				}

				m.Visit(root, &ParamChecker{})

				if err := m.Error(); err != nil {
					printfStdErr("%s\n", err.Error())
					exit(2)

				} else {
					DEBUG("Converting the following data back to YML:")
					DEBUG("%#v", root)
					merged, err := yaml.Marshal(root)
					if err != nil {
						printfStdErr("Unable to convert merged result back to YAML: %s\nData:\n%#v", err.Error(), root)
						exit(2)
					} else {
						var output string
						if handleConcourseQuoting {
							output = dequoteConcourse(merged)
						} else {
							output = string(merged)
						}
						printfStdOut("%s\n", output)
					}
				}
			} else {
				usage()
			}
		default:
			usage()
		}
	}
}

func parseYAML(data []byte) (map[interface{}]interface{}, error) {
	y, err := simpleyaml.NewYaml(data)
	if err != nil {
		return nil, err
	}

	doc, err := y.Map()
	if err != nil {
		return nil, fmt.Errorf("Root of YAML document is not a hash/map: %s\n", err.Error())
	}

	return doc, nil
}

func mergeAllDocs(root map[interface{}]interface{}, paths []string) error {
	m := &Merger{}

	for _, path := range paths {
		DEBUG("Processing file '%s'", path)
		data, err := readFile(path)
		if err != nil {
			return fmt.Errorf("Error reading file %s: %s\n", path, err.Error())
		}

		if handleConcourseQuoting {
			data = quoteConcourse(data)
		}

		doc, err := parseYAML(data)
		if err != nil {
			return fmt.Errorf("%s: %s\n", path, err.Error())
		}

		m.Merge(root, doc)

		tmpYaml, _ := yaml.Marshal(root) // we don't care about errors for debugging
		DEBUG("Current data after processing '%s':\n%s", path, tmpYaml)
	}

	return m.Error()
}

var concourseRegex = `\{\{([-\w\p{L}]+)\}\}`

func quoteConcourse(input []byte) []byte {
	re := regexp.MustCompile("(" + concourseRegex + ")")
	return re.ReplaceAll(input, []byte("\"$1\""))
}

func dequoteConcourse(input []byte) string {
	re := regexp.MustCompile("['\"](" + concourseRegex + ")[\"']")
	return re.ReplaceAllString(string(input), "$1")
}
<|endoftext|>