<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/golib/stress"
	"github.com/funkygao/redigo/redis"
)

var (
	addr          string
	mode          string
	topic         string
	loops         int
	suppressError bool
	pubkey        string
	appid         string
	sleep         time.Duration
	sz            int
	async         bool
)

func main() {
	flag.StringVar(&addr, "addr", "http://localhost:9191/", "kateway pub addr, must start with http or https")
	flag.IntVar(&loops, "loops", 1000, "loops in each thread")
	flag.IntVar(&sz, "size", 200, "each pub message size")
	flag.StringVar(&topic, "topic", "foobar", "pub topic")
	flag.StringVar(&mode, "mode", "gw", "<gw|kafka|http|redis>")
	flag.StringVar(&appid, "appid", "app1", "appid of pub")
	flag.StringVar(&pubkey, "pubkey", "mypubkey", "pubkey")
	flag.BoolVar(&async, "async", false, "async pub")
	flag.DurationVar(&sleep, "sleep", 0, "sleep between pub")
	flag.BoolVar(&suppressError, "noerr", true, "suppress error output")
	flag.Parse()

	switch mode {
	case "gw":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(pubGatewayLoop)

	case "kafka":
		if async {
			stress.RunStress(pubKafkaAsyncLoop)
		} else {
			stress.RunStress(pubKafkaLoop)
		}

	case "redis":
		stress.RunStress(redisLoop)

	case "http":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(getHttpLoop)
	}

}

func getHttpLoop(seq int) {
	client := createHttpClient()
	req, _ := http.NewRequest("GET", "http://localhost:9090/", nil)
	for i := 0; i < loops; i++ {
		response, err := client.Do(req)
		if err == nil {
			ioutil.ReadAll(response.Body)
			response.Body.Close() // reuse the connection

			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
			//log.Println(err)
		}
	}
}

func pubKafkaLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewSyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, _, err := producer.SendMessage(&sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		})
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
		}
	}

}

func redisLoop(seq int) {
	conn, err := redis.DialTimeout("tcp", ":6379", 0, 1*time.Second, 1*time.Second)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer conn.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, err := conn.Do("SET", "key", msg)
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			if !suppressError {
				log.Println(err)
			}
			stress.IncCounter("fail", 1)
		}
	}

}

func pubKafkaAsyncLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.Flush.Frequency = time.Second * 10
	cf.Producer.Flush.Messages = 1000
	cf.Producer.Flush.MaxMessages = 1000
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewAsyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		producer.Input() <- &sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		}
		stress.IncCounter("ok", 1)
	}

}

func createHttpClient() *http.Client {
	timeout := 3 * time.Second
	httpClient := &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   timeout,
				KeepAlive: 60 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: timeout,
		},
	}

	return httpClient
}

func pubGatewayLoop(seq int) {
	httpClient := createHttpClient()
	url := fmt.Sprintf("%s/topics/%s/v1?", addr, topic)
	if async {
		url += "async=1"
	}
	for n := 0; n < loops; n++ {
		req, err := http.NewRequest("POST", url,
			bytes.NewBuffer([]byte(strings.Repeat("X", sz))))
		if err != nil {
			log.Fatalf("Error Occured. %+v", err)
			stress.IncCounter("fail", 1)
			return
		}

		req.Header.Set("Appid", appid)
		req.Header.Set("Pubkey", pubkey)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		// use httpClient to send request
		response, err := httpClient.Do(req)
		if err != nil && response == nil {
			stress.IncCounter("fail", 1)
			if !suppressError {
				log.Printf("Error sending request to API endpoint. %+v", err)
			}
		} else {
			if response.StatusCode != http.StatusOK {
				stress.IncCounter("fail", 1)
				if !suppressError {
					log.Printf("Error sending request to API endpoint. %+v", response.Status)
				}
				return
			}

			// Let's check if the work actually is done
			// We have seen inconsistencies even when we get 200 OK response
			body, err := ioutil.ReadAll(response.Body)
			if err != nil {
				log.Fatalf("Couldn't parse response body. %+v", err)
			}

			// Close the connection to reuse it
			response.Body.Close()

			stress.IncCounter("ok", 1)

			if false {
				log.Println("Response Body:", string(body))
			}
		}

		if sleep > 0 {
			time.Sleep(sleep)
		}
	}
}
<commit_msg>bench will output err by default<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/golib/stress"
	"github.com/funkygao/redigo/redis"
)

var (
	addr          string
	mode          string
	topic         string
	loops         int
	suppressError bool
	pubkey        string
	appid         string
	sleep         time.Duration
	sz            int
	async         bool
)

func main() {
	flag.StringVar(&addr, "addr", "http://localhost:9191/", "kateway pub addr, must start with http or https")
	flag.IntVar(&loops, "loops", 1000, "loops in each thread")
	flag.IntVar(&sz, "size", 200, "each pub message size")
	flag.StringVar(&topic, "topic", "foobar", "pub topic")
	flag.StringVar(&mode, "mode", "gw", "<gw|kafka|http|redis>")
	flag.StringVar(&appid, "appid", "app1", "appid of pub")
	flag.StringVar(&pubkey, "pubkey", "mypubkey", "pubkey")
	flag.BoolVar(&async, "async", false, "async pub")
	flag.DurationVar(&sleep, "sleep", 0, "sleep between pub")
	flag.BoolVar(&suppressError, "noerr", false, "suppress error output")
	flag.Parse()

	switch mode {
	case "gw":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(pubGatewayLoop)

	case "kafka":
		if async {
			stress.RunStress(pubKafkaAsyncLoop)
		} else {
			stress.RunStress(pubKafkaLoop)
		}

	case "redis":
		stress.RunStress(redisLoop)

	case "http":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(getHttpLoop)
	}

}

func getHttpLoop(seq int) {
	client := createHttpClient()
	req, _ := http.NewRequest("GET", "http://localhost:9090/", nil)
	for i := 0; i < loops; i++ {
		response, err := client.Do(req)
		if err == nil {
			ioutil.ReadAll(response.Body)
			response.Body.Close() // reuse the connection

			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
			//log.Println(err)
		}
	}
}

func pubKafkaLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewSyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, _, err := producer.SendMessage(&sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		})
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
		}
	}

}

func redisLoop(seq int) {
	conn, err := redis.DialTimeout("tcp", ":6379", 0, 1*time.Second, 1*time.Second)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer conn.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, err := conn.Do("SET", "key", msg)
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			if !suppressError {
				log.Println(err)
			}
			stress.IncCounter("fail", 1)
		}
	}

}

func pubKafkaAsyncLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.Flush.Frequency = time.Second * 10
	cf.Producer.Flush.Messages = 1000
	cf.Producer.Flush.MaxMessages = 1000
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewAsyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		producer.Input() <- &sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		}
		stress.IncCounter("ok", 1)
	}

}

func createHttpClient() *http.Client {
	timeout := 3 * time.Second
	httpClient := &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   timeout,
				KeepAlive: 60 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: timeout,
		},
	}

	return httpClient
}

func pubGatewayLoop(seq int) {
	httpClient := createHttpClient()
	url := fmt.Sprintf("%s/topics/%s/v1?", addr, topic)
	if async {
		url += "async=1"
	}
	for n := 0; n < loops; n++ {
		req, err := http.NewRequest("POST", url,
			bytes.NewBuffer([]byte(strings.Repeat("X", sz))))
		if err != nil {
			log.Fatalf("Error Occured. %+v", err)
			stress.IncCounter("fail", 1)
			return
		}

		req.Header.Set("Appid", appid)
		req.Header.Set("Pubkey", pubkey)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		// use httpClient to send request
		response, err := httpClient.Do(req)
		if err != nil && response == nil {
			stress.IncCounter("fail", 1)
			if !suppressError {
				log.Printf("Error sending request to API endpoint. %+v", err)
			}
		} else {
			if response.StatusCode != http.StatusOK {
				stress.IncCounter("fail", 1)
				if !suppressError {
					log.Printf("Error sending request to API endpoint. %+v", response.Status)
				}
				return
			}

			// Let's check if the work actually is done
			// We have seen inconsistencies even when we get 200 OK response
			body, err := ioutil.ReadAll(response.Body)
			if err != nil {
				log.Fatalf("Couldn't parse response body. %+v", err)
			}

			// Close the connection to reuse it
			response.Body.Close()

			stress.IncCounter("ok", 1)

			if false {
				log.Println("Response Body:", string(body))
			}
		}

		if sleep > 0 {
			time.Sleep(sleep)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

type PackageIndex struct {
	URI string
}

type Requirement struct {
	Name string
}

var allPkgRegexp = regexp.MustCompile(`<a href='([A-Za-z0-9\._\-]+)'>([A-Za-z0-9\._\-]+)</a><br/>`)
var pkgFilesRegexp = regexp.MustCompile(`<a href="([/A-Za-z0-9\._\-]+)#md5=[0-9a-z]+"[^>]*>([A-Za-z0-9\._\-]+)</a><br/>`)
var tarRegexp = regexp.MustCompile(`[/A-Za-z0-9\._\-]+\.tar\.(?:gz|bz2)`)
var zipRegexp = regexp.MustCompile(`[/A-Za-z0-9\._\-]+\.zip`)
var eggRegexp = regexp.MustCompile(`[/A-Za-z0-9\._\-]+\.egg`)
var requirementRegexp = regexp.MustCompile(`(?P<package>[A-Za-z0-9\._\-]+)(?:\[([A-Za-z0-9\._\-]+)\])?\s*(?:(?P<constraint>==|>=|>|<|<=)\s*(?P<version>[A-Za-z0-9\._\-]+)(?:\s*,\s*[<>=!]+\s*[a-z0-9\.]+)?)?`)
var reqHeaderRegexp = regexp.MustCompile(`\[[A-Za-z0-9\._\-]+\]`)

func (p *PackageIndex) AllPackages() ([]string, error) {
	pkgs := make([]string, 0)

	resp, err := http.Get(fmt.Sprintf("%s/simple", p.URI))
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	matches := allPkgRegexp.FindAllStringSubmatch(string(body), -1)
	for _, match := range matches {
		if len(match) != 3 {
			return nil, fmt.Errorf("Unexpected number of submatches: %d, %v", len(match), match)
		} else if match[1] != match[2] {
			return nil, fmt.Errorf("Names do not match %s != %s", match[1], match[2])
		} else {
			pkgs = append(pkgs, match[1])
		}
	}

	return pkgs, nil
}

func (p *PackageIndex) PackageRequirements(pkg string) ([]*Requirement, error) {
	files, err := p.pkgFiles(pkg)
	if err != nil {
		return nil, err
	} else if len(files) == 0 {
		// os.Stderr.WriteString(fmt.Sprintf("[no-files] no files found for pkg %s\n", pkg))
		return nil, nil
	}

	if path := lastTar(files); path != "" {
		return p.fetchRequiresTar(path)
	} else if path := lastEgg(files); path != "" {
		return p.fetchRequiresZip(path, true)
	} else if path := lastZip(files); path != "" {
		return p.fetchRequiresZip(path, false)
	} else {
		os.Stderr.WriteString(fmt.Sprintf("[tar/zip] no tar or zip found in %+v for pkg %s\n", files, pkg))
		return nil, nil
	}
}

func (p *PackageIndex) pkgFiles(pkg string) ([]string, error) {
	files := make([]string, 0)

	uriPath := fmt.Sprintf("/simple/%s", pkg)
	uri := fmt.Sprintf("%s%s", p.URI, uriPath)
	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	matches := pkgFilesRegexp.FindAllStringSubmatch(string(body), -1)
	for _, match := range matches {
		if len(match) != 3 {
			return nil, fmt.Errorf("Unexpected number of submatches: %d, %v", len(match), match)
		} else {
			files = append(files, filepath.Clean(filepath.Join(uriPath, match[1])))
		}
	}

	return files, nil
}

func (p *PackageIndex) fetchRequiresZip(path string, isEgg bool) ([]*Requirement, error) {
	f, err := ioutil.TempFile("", "pypigraph_zip")
	if err != nil {
		return nil, err
	}
	defer os.Remove(f.Name())
	if err := f.Close(); err != nil {
		return nil, err
	}

	uri := fmt.Sprintf("%s%s", p.URI, path)
	wget := exec.Command("wget", uri, "-O", f.Name())
	var eggInfoFilePattern string
	if isEgg {
		eggInfoFilePattern = "EGG-INFO/requires.txt"
	} else {
		eggInfoFilePattern = "**/*.egg-info/requires.txt"
	}
	unzip := exec.Command("unzip", "-cq", f.Name(), eggInfoFilePattern)

	err = wget.Run()
	if err != nil {
		return nil, fmt.Errorf("Error running wget: %s", err)
	}

	unzipOut, err := unzip.StdoutPipe()
	if err != nil {
		return nil, err
	}
	unzipErr, err := unzip.StderrPipe()
	if err != nil {
		return nil, err
	}

	err = unzip.Start()
	if err != nil {
		return nil, err
	}
	var unzipErrput, unzipOutput []byte
	go func() {
		unzipErrput, _ = ioutil.ReadAll(unzipErr)
	}()
	go func() {
		unzipOutput, _ = ioutil.ReadAll(unzipOut)
	}()
	err = unzip.Wait()
	if err != nil {
		if strings.Contains(string(unzipErrput), "filename not matched:") {
			// os.Stderr.WriteString(fmt.Sprintf("[requires.txt] no requires.txt found in %s\n", uri))
			return nil, nil
		} else {
			return nil, fmt.Errorf("Error running unzip on file %s: %s, [%s]", f.Name(), err, string(unzipErrput))
		}
	}

	rawReqs := strings.TrimSpace(string(unzipOutput))
	return parseRequirements(rawReqs)
}

func (p *PackageIndex) fetchRequiresTar(path string) ([]*Requirement, error) {
	uri := fmt.Sprintf("%s%s", p.URI, path)
	curl := exec.Command("curl", uri)
	tar := exec.Command("tar", "-xvO", "--include", "**/*.egg-info/requires.txt")

	curlOut, err := curl.StdoutPipe()
	if err != nil {
		return nil, err
	}
	tarIn, err := tar.StdinPipe()
	if err != nil {
		return nil, err
	}
	tarOut, err := tar.StdoutPipe()
	if err != nil {
		return nil, err
	}

	go func() {
		io.Copy(tarIn, curlOut)
		tarIn.Close()
	}()

	curl.Start()
	tar.Start()

	tarOutput, err := ioutil.ReadAll(tarOut)
	if err != nil {
		return nil, err
	}

	curl.Wait()
	tar.Wait()

	rawReqs := strings.TrimSpace(string(tarOutput))
	if rawReqs == "" {
		// os.Stderr.WriteString(fmt.Sprintf("[requires.txt] no requires.txt found in %s\n", uri))
		return nil, nil
	}

	return parseRequirements(rawReqs)
}

func lastTar(files []string) string {
	for f := len(files) - 1; f >= 0; f-- {
		if tarRegexp.MatchString(files[f]) {
			return files[f]
		}
	}
	return ""
}

func lastEgg(files []string) string {
	for f := len(files) - 1; f >= 0; f-- {
		if eggRegexp.MatchString(files[f]) {
			return files[f]
		}
	}
	return ""
}

func lastZip(files []string) string {
	for f := len(files) - 1; f >= 0; f-- {
		if zipRegexp.MatchString(files[f]) {
			return files[f]
		}
	}
	return ""
}

func parseRequirements(rawReqs string) ([]*Requirement, error) {
	rawReqs = strings.TrimSpace(rawReqs)

	reqStrs := strings.Split(rawReqs, "\n")
	reqs := make([]*Requirement, 0)
	for _, reqStr := range reqStrs {
		if reqStr == "" {
			continue
		}

		if req, err := parseRequirement(reqStr); err == nil {
			reqs = append(reqs, req)
		} else if reqHeaderRegexp.MatchString(reqStr) {
			// do nothing
		} else {
			os.Stderr.WriteString(fmt.Sprintf("[req] Could not parse requirement: %s\n", err))
		}
	}
	return reqs, nil
}

func parseRequirement(reqStr string) (*Requirement, error) {
	reqStr = strings.TrimSpace(reqStr)
	match := requirementRegexp.FindStringSubmatch(reqStr)
	if len(match) != 5 {
		return nil, fmt.Errorf("Expected match of length 5, but got %+v from '%s'", match, reqStr)
	} else if match[0] != reqStr {
		return nil, fmt.Errorf("Unable to parse requirement from string: '%s'", reqStr)
	}

	return &Requirement{Name: match[1]}, nil
}

func main() {
	p := &PackageIndex{URI: "https://pypi.python.org"}
	pkgs, err := p.AllPackages()
	if err != nil {
		os.Stderr.WriteString(fmt.Sprintf("[FATAL] %s\n", err))
		os.Exit(1)
	}

	for _, pkg := range pkgs {
		reqs, err := p.PackageRequirements(pkg)
		if err != nil {
			os.Stderr.WriteString(fmt.Sprintf("[ERROR] unable to parse pkg %s due to error: %s\n", pkg, err))
		} else {
			fmt.Println(pkg)
			for _, req := range reqs {
				fmt.Printf("  %+v\n", req.Name)
			}
		}
	}
}
<commit_msg>handle .tgz extension<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

type PackageIndex struct {
	URI string
}

type Requirement struct {
	Name string
}

var allPkgRegexp = regexp.MustCompile(`<a href='([A-Za-z0-9\._\-]+)'>([A-Za-z0-9\._\-]+)</a><br/>`)
var pkgFilesRegexp = regexp.MustCompile(`<a href="([/A-Za-z0-9\._\-]+)#md5=[0-9a-z]+"[^>]*>([A-Za-z0-9\._\-]+)</a><br/>`)
var tarRegexp = regexp.MustCompile(`[/A-Za-z0-9\._\-]+\.(?:tar\.(?:gz|bz2)|tgz)`)
var zipRegexp = regexp.MustCompile(`[/A-Za-z0-9\._\-]+\.zip`)
var eggRegexp = regexp.MustCompile(`[/A-Za-z0-9\._\-]+\.egg`)
var requirementRegexp = regexp.MustCompile(`(?P<package>[A-Za-z0-9\._\-]+)(?:\[([A-Za-z0-9\._\-]+)\])?\s*(?:(?P<constraint>==|>=|>|<|<=)\s*(?P<version>[A-Za-z0-9\._\-]+)(?:\s*,\s*[<>=!]+\s*[a-z0-9\.]+)?)?`)
var reqHeaderRegexp = regexp.MustCompile(`\[[A-Za-z0-9\._\-]+\]`)

func (p *PackageIndex) AllPackages() ([]string, error) {
	pkgs := make([]string, 0)

	resp, err := http.Get(fmt.Sprintf("%s/simple", p.URI))
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	matches := allPkgRegexp.FindAllStringSubmatch(string(body), -1)
	for _, match := range matches {
		if len(match) != 3 {
			return nil, fmt.Errorf("Unexpected number of submatches: %d, %v", len(match), match)
		} else if match[1] != match[2] {
			return nil, fmt.Errorf("Names do not match %s != %s", match[1], match[2])
		} else {
			pkgs = append(pkgs, match[1])
		}
	}

	return pkgs, nil
}

func (p *PackageIndex) PackageRequirements(pkg string) ([]*Requirement, error) {
	files, err := p.pkgFiles(pkg)
	if err != nil {
		return nil, err
	} else if len(files) == 0 {
		// os.Stderr.WriteString(fmt.Sprintf("[no-files] no files found for pkg %s\n", pkg))
		return nil, nil
	}

	if path := lastTar(files); path != "" {
		return p.fetchRequiresTar(path)
	} else if path := lastEgg(files); path != "" {
		return p.fetchRequiresZip(path, true)
	} else if path := lastZip(files); path != "" {
		return p.fetchRequiresZip(path, false)
	} else {
		os.Stderr.WriteString(fmt.Sprintf("[tar/zip] no tar or zip found in %+v for pkg %s\n", files, pkg))
		return nil, nil
	}
}

func (p *PackageIndex) pkgFiles(pkg string) ([]string, error) {
	files := make([]string, 0)

	uriPath := fmt.Sprintf("/simple/%s", pkg)
	uri := fmt.Sprintf("%s%s", p.URI, uriPath)
	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	matches := pkgFilesRegexp.FindAllStringSubmatch(string(body), -1)
	for _, match := range matches {
		if len(match) != 3 {
			return nil, fmt.Errorf("Unexpected number of submatches: %d, %v", len(match), match)
		} else {
			files = append(files, filepath.Clean(filepath.Join(uriPath, match[1])))
		}
	}

	return files, nil
}

func (p *PackageIndex) fetchRequiresZip(path string, isEgg bool) ([]*Requirement, error) {
	f, err := ioutil.TempFile("", "pypigraph_zip")
	if err != nil {
		return nil, err
	}
	defer os.Remove(f.Name())
	if err := f.Close(); err != nil {
		return nil, err
	}

	uri := fmt.Sprintf("%s%s", p.URI, path)
	wget := exec.Command("wget", uri, "-O", f.Name())
	var eggInfoFilePattern string
	if isEgg {
		eggInfoFilePattern = "EGG-INFO/requires.txt"
	} else {
		eggInfoFilePattern = "**/*.egg-info/requires.txt"
	}
	unzip := exec.Command("unzip", "-cq", f.Name(), eggInfoFilePattern)

	err = wget.Run()
	if err != nil {
		return nil, fmt.Errorf("Error running wget: %s", err)
	}

	unzipOut, err := unzip.StdoutPipe()
	if err != nil {
		return nil, err
	}
	unzipErr, err := unzip.StderrPipe()
	if err != nil {
		return nil, err
	}

	err = unzip.Start()
	if err != nil {
		return nil, err
	}
	var unzipErrput, unzipOutput []byte
	go func() {
		unzipErrput, _ = ioutil.ReadAll(unzipErr)
	}()
	go func() {
		unzipOutput, _ = ioutil.ReadAll(unzipOut)
	}()
	err = unzip.Wait()
	if err != nil {
		if strings.Contains(string(unzipErrput), "filename not matched:") {
			// os.Stderr.WriteString(fmt.Sprintf("[requires.txt] no requires.txt found in %s\n", uri))
			return nil, nil
		} else {
			return nil, fmt.Errorf("Error running unzip on file %s: %s, [%s]", f.Name(), err, string(unzipErrput))
		}
	}

	rawReqs := strings.TrimSpace(string(unzipOutput))
	return parseRequirements(rawReqs)
}

func (p *PackageIndex) fetchRequiresTar(path string) ([]*Requirement, error) {
	uri := fmt.Sprintf("%s%s", p.URI, path)
	curl := exec.Command("curl", uri)
	tar := exec.Command("tar", "-xvO", "--include", "**/*.egg-info/requires.txt")

	curlOut, err := curl.StdoutPipe()
	if err != nil {
		return nil, err
	}
	tarIn, err := tar.StdinPipe()
	if err != nil {
		return nil, err
	}
	tarOut, err := tar.StdoutPipe()
	if err != nil {
		return nil, err
	}

	go func() {
		io.Copy(tarIn, curlOut)
		tarIn.Close()
	}()

	curl.Start()
	tar.Start()

	tarOutput, err := ioutil.ReadAll(tarOut)
	if err != nil {
		return nil, err
	}

	curl.Wait()
	tar.Wait()

	rawReqs := strings.TrimSpace(string(tarOutput))
	if rawReqs == "" {
		// os.Stderr.WriteString(fmt.Sprintf("[requires.txt] no requires.txt found in %s\n", uri))
		return nil, nil
	}

	return parseRequirements(rawReqs)
}

func lastTar(files []string) string {
	for f := len(files) - 1; f >= 0; f-- {
		if tarRegexp.MatchString(files[f]) {
			return files[f]
		}
	}
	return ""
}

func lastEgg(files []string) string {
	for f := len(files) - 1; f >= 0; f-- {
		if eggRegexp.MatchString(files[f]) {
			return files[f]
		}
	}
	return ""
}

func lastZip(files []string) string {
	for f := len(files) - 1; f >= 0; f-- {
		if zipRegexp.MatchString(files[f]) {
			return files[f]
		}
	}
	return ""
}

func parseRequirements(rawReqs string) ([]*Requirement, error) {
	rawReqs = strings.TrimSpace(rawReqs)

	reqStrs := strings.Split(rawReqs, "\n")
	reqs := make([]*Requirement, 0)
	for _, reqStr := range reqStrs {
		if reqStr == "" {
			continue
		}

		if req, err := parseRequirement(reqStr); err == nil {
			reqs = append(reqs, req)
		} else if reqHeaderRegexp.MatchString(reqStr) {
			// do nothing
		} else {
			os.Stderr.WriteString(fmt.Sprintf("[req] Could not parse requirement: %s\n", err))
		}
	}
	return reqs, nil
}

func parseRequirement(reqStr string) (*Requirement, error) {
	reqStr = strings.TrimSpace(reqStr)
	match := requirementRegexp.FindStringSubmatch(reqStr)
	if len(match) != 5 {
		return nil, fmt.Errorf("Expected match of length 5, but got %+v from '%s'", match, reqStr)
	} else if match[0] != reqStr {
		return nil, fmt.Errorf("Unable to parse requirement from string: '%s'", reqStr)
	}

	return &Requirement{Name: match[1]}, nil
}

func main() {
	p := &PackageIndex{URI: "https://pypi.python.org"}
	pkgs, err := p.AllPackages()
	if err != nil {
		os.Stderr.WriteString(fmt.Sprintf("[FATAL] %s\n", err))
		os.Exit(1)
	}

	for _, pkg := range pkgs {
		reqs, err := p.PackageRequirements(pkg)
		if err != nil {
			os.Stderr.WriteString(fmt.Sprintf("[ERROR] unable to parse pkg %s due to error: %s\n", pkg, err))
		} else {
			fmt.Println(pkg)
			for _, req := range reqs {
				fmt.Printf("  %+v\n", req.Name)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"

	"github.com/squaremo/ambergreen/common/data"
	"github.com/squaremo/ambergreen/common/store"
	"github.com/squaremo/ambergreen/common/store/etcdstore"

	"github.com/gorilla/mux"
)

func main() {
	prom := os.Getenv("PROM_ADDRESS")
	if prom == "" {
		prom = "http://localhost:9090"
	}

	store := etcdstore.NewFromEnv()
	if err := store.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Printf("Connected to backend\n")
	api := &api{store, prom}

	router := mux.NewRouter()

	router.HandleFunc("/", homePage)
	router.HandleFunc("/index.html", homePage)
	router.PathPrefix("/res/").HandlerFunc(handleResource)

	router.HandleFunc("/api/{service}/", api.listInstances)
	router.HandleFunc("/api/", api.listServices)

	router.PathPrefix("/stats/").HandlerFunc(api.proxyStats)

	http.ListenAndServe("0.0.0.0:7070", router)
}

func handleResource(w http.ResponseWriter, r *http.Request) {
	file := r.URL.Path[1:]
	http.ServeFile(w, r, file)
}

func homePage(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}

//=== API handlers

type api struct {
	store   store.Store
	promURL string
}

type serviceDetails struct {
	Name    string       `json:"name"`
	Details data.Service `json:"details"`
}

type service struct {
	Name     string            `json:"name"`
	Children []instanceDetails `json:"children"`
	Details  data.Service      `json:"details"`
}

type instanceDetails struct {
	Name    string        `json:"name"`
	Details data.Instance `json:"details"`
}

func (api *api) listServices(w http.ResponseWriter, r *http.Request) {
	var currentService serviceDetails
	services := []serviceDetails{}

	api.store.ForeachServiceInstance(func(name string, details data.Service) {
		currentService = serviceDetails{
			Name:    name,
			Details: details,
		}
		services = append(services, currentService)
	}, nil)
	json.NewEncoder(w).Encode(services)
}

func (api *api) listInstances(w http.ResponseWriter, r *http.Request) {
	args := mux.Vars(r)
	serviceName := args["service"]
	details, err := api.store.GetServiceDetails(serviceName)
	if err != nil {
		http.NotFound(w, r)
	}
	children := []instanceDetails{}
	api.store.ForeachInstance(serviceName, func(name string, details data.Instance) {
		instance := instanceDetails{
			Name:    name,
			Details: details,
		}
		children = append(children, instance)
	})
	service := service{
		Name:     serviceName,
		Details:  details,
		Children: children,
	}
	json.NewEncoder(w).Encode(service)
}

/* Proxy for prometheus, as a stop-gap */

func (api *api) proxyStats(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path[len("/stats"):] + "?" + r.URL.RawQuery
	log.Println(path)
	resp, err := http.Get(api.promURL + path)
	if err != nil {
		http.Error(w, "Error contacting prometheus server: "+err.Error(), 500)
		return
	}
	defer resp.Body.Close()
	for k, vs := range resp.Header {
		w.Header()[k] = vs
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}
<commit_msg>Missing return<commit_after>package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"

	"github.com/squaremo/ambergreen/common/data"
	"github.com/squaremo/ambergreen/common/store"
	"github.com/squaremo/ambergreen/common/store/etcdstore"

	"github.com/gorilla/mux"
)

func main() {
	prom := os.Getenv("PROM_ADDRESS")
	if prom == "" {
		prom = "http://localhost:9090"
	}

	store := etcdstore.NewFromEnv()
	if err := store.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Printf("Connected to backend\n")
	api := &api{store, prom}

	router := mux.NewRouter()

	router.HandleFunc("/", homePage)
	router.HandleFunc("/index.html", homePage)
	router.PathPrefix("/res/").HandlerFunc(handleResource)

	router.HandleFunc("/api/{service}/", api.listInstances)
	router.HandleFunc("/api/", api.listServices)

	router.PathPrefix("/stats/").HandlerFunc(api.proxyStats)

	http.ListenAndServe("0.0.0.0:7070", router)
}

func handleResource(w http.ResponseWriter, r *http.Request) {
	file := r.URL.Path[1:]
	http.ServeFile(w, r, file)
}

func homePage(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}

//=== API handlers

type api struct {
	store   store.Store
	promURL string
}

type serviceDetails struct {
	Name    string       `json:"name"`
	Details data.Service `json:"details"`
}

type service struct {
	Name     string            `json:"name"`
	Children []instanceDetails `json:"children"`
	Details  data.Service      `json:"details"`
}

type instanceDetails struct {
	Name    string        `json:"name"`
	Details data.Instance `json:"details"`
}

func (api *api) listServices(w http.ResponseWriter, r *http.Request) {
	var currentService serviceDetails
	services := []serviceDetails{}

	api.store.ForeachServiceInstance(func(name string, details data.Service) {
		currentService = serviceDetails{
			Name:    name,
			Details: details,
		}
		services = append(services, currentService)
	}, nil)
	json.NewEncoder(w).Encode(services)
}

func (api *api) listInstances(w http.ResponseWriter, r *http.Request) {
	args := mux.Vars(r)
	serviceName := args["service"]
	details, err := api.store.GetServiceDetails(serviceName)
	if err != nil {
		http.NotFound(w, r)
		return
	}
	children := []instanceDetails{}
	api.store.ForeachInstance(serviceName, func(name string, details data.Instance) {
		instance := instanceDetails{
			Name:    name,
			Details: details,
		}
		children = append(children, instance)
	})
	service := service{
		Name:     serviceName,
		Details:  details,
		Children: children,
	}
	json.NewEncoder(w).Encode(service)
}

/* Proxy for prometheus, as a stop-gap */

func (api *api) proxyStats(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path[len("/stats"):] + "?" + r.URL.RawQuery
	log.Println(path)
	resp, err := http.Get(api.promURL + path)
	if err != nil {
		http.Error(w, "Error contacting prometheus server: "+err.Error(), 500)
		return
	}
	defer resp.Body.Close()
	for k, vs := range resp.Header {
		w.Header()[k] = vs
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}
<|endoftext|>


<commit_before>package caddymain

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	"gopkg.in/natefinch/lumberjack.v2"

	"github.com/xenolf/lego/acme"

	"github.com/mholt/caddy"
	// plug in the HTTP server type
	_ "github.com/mholt/caddy/caddyhttp"

	"github.com/mholt/caddy/caddytls"
	// This is where other plugins get plugged in (imported)
)

func init() {
	caddy.TrapSignals()
	setVersion()

	flag.BoolVar(&caddytls.Agreed, "agree", false, "Agree to the CA's Subscriber Agreement")
	flag.StringVar(&caddytls.DefaultCAUrl, "ca", "https://acme-v01.api.letsencrypt.org/directory", "URL to certificate authority's ACME server directory")
	flag.StringVar(&conf, "conf", "", "Caddyfile to load (default \""+caddy.DefaultConfigFile+"\")")
	flag.StringVar(&cpu, "cpu", "100%", "CPU cap")
	flag.BoolVar(&plugins, "plugins", false, "List installed plugins")
	flag.StringVar(&caddytls.DefaultEmail, "email", "", "Default ACME CA account email address")
	flag.StringVar(&logfile, "log", "", "Process log file")
	flag.StringVar(&caddy.PidFile, "pidfile", "", "Path to write pid file")
	flag.BoolVar(&caddy.Quiet, "quiet", false, "Quiet mode (no initialization output)")
	flag.StringVar(&revoke, "revoke", "", "Hostname for which to revoke the certificate")
	flag.StringVar(&serverType, "type", "http", "Type of server to run")
	flag.BoolVar(&version, "version", false, "Show version")

	caddy.RegisterCaddyfileLoader("flag", caddy.LoaderFunc(confLoader))
	caddy.SetDefaultCaddyfileLoader("default", caddy.LoaderFunc(defaultLoader))
}

// Run is Caddy's main() function.
func Run() {
	flag.Parse()
	moveStorage() // TODO: This is temporary for the 0.9 release, or until most users upgrade to 0.9+

	caddy.AppName = appName
	caddy.AppVersion = appVersion
	acme.UserAgent = appName + "/" + appVersion

	// Set up process log before anything bad happens
	switch logfile {
	case "stdout":
		log.SetOutput(os.Stdout)
	case "stderr":
		log.SetOutput(os.Stderr)
	case "":
		log.SetOutput(ioutil.Discard)
	default:
		log.SetOutput(&lumberjack.Logger{
			Filename:   logfile,
			MaxSize:    100,
			MaxAge:     14,
			MaxBackups: 10,
		})
	}

	// Check for one-time actions
	if revoke != "" {
		err := caddytls.Revoke(revoke)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("Revoked certificate for %s\n", revoke)
		os.Exit(0)
	}
	if version {
		fmt.Printf("%s %s\n", appName, appVersion)
		if devBuild && gitShortStat != "" {
			fmt.Printf("%s\n%s\n", gitShortStat, gitFilesModified)
		}
		os.Exit(0)
	}
	if plugins {
		fmt.Println(caddy.DescribePlugins())
		os.Exit(0)
	}

	// Set CPU cap
	err := setCPU(cpu)
	if err != nil {
		mustLogFatal(err)
	}

	// Get Caddyfile input
	caddyfile, err := caddy.LoadCaddyfile(serverType)
	if err != nil {
		mustLogFatal(err)
	}

	// Start your engines
	instance, err := caddy.Start(caddyfile)
	if err != nil {
		mustLogFatal(err)
	}

	// Twiddle your thumbs
	instance.Wait()
}

// mustLogFatal wraps log.Fatal() in a way that ensures the
// output is always printed to stderr so the user can see it
// if the user is still there, even if the process log was not
// enabled. If this process is an upgrade, however, and the user
// might not be there anymore, this just logs to the process
// log and exits.
func mustLogFatal(args ...interface{}) {
	if !caddy.IsUpgrade() {
		log.SetOutput(os.Stderr)
	}
	log.Fatal(args...)
}

// confLoader loads the Caddyfile using the -conf flag.
func confLoader(serverType string) (caddy.Input, error) {
	if conf == "" {
		return nil, nil
	}

	if conf == "stdin" {
		return caddy.CaddyfileFromPipe(os.Stdin)
	}

	contents, err := ioutil.ReadFile(conf)
	if err != nil {
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       conf,
		ServerTypeName: serverType,
	}, nil
}

// defaultLoader loads the Caddyfile from the current working directory.
func defaultLoader(serverType string) (caddy.Input, error) {
	contents, err := ioutil.ReadFile(caddy.DefaultConfigFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       caddy.DefaultConfigFile,
		ServerTypeName: serverType,
	}, nil
}

// moveStorage moves the old certificate storage location by
// renaming the "letsencrypt" folder to the hostname of the
// CA URL. This is TEMPORARY until most users have upgraded to 0.9+.
func moveStorage() {
	oldPath := filepath.Join(caddy.AssetsPath(), "letsencrypt")
	_, err := os.Stat(oldPath)
	if os.IsNotExist(err) {
		return
	}
	// Just use a default config to get default (file) storage
	fileStorage, err := new(caddytls.Config).StorageFor(caddytls.DefaultCAUrl)
	if err != nil {
		log.Fatalf("[ERROR] Unable to get new path for certificate storage: %v", err)
	}
	newPath := string(fileStorage.(caddytls.FileStorage))
	err = os.MkdirAll(string(newPath), 0700)
	if err != nil {
		log.Fatalf("[ERROR] Unable to make new certificate storage path: %v\n\nPlease follow instructions at:\nhttps://github.com/mholt/caddy/issues/902#issuecomment-228876011", err)
	}
	err = os.Rename(oldPath, string(newPath))
	if err != nil {
		log.Fatalf("[ERROR] Unable to migrate certificate storage: %v\n\nPlease follow instructions at:\nhttps://github.com/mholt/caddy/issues/902#issuecomment-228876011", err)
	}
	// convert mixed case folder and file names to lowercase
	filepath.Walk(string(newPath), func(path string, info os.FileInfo, err error) error {
		// must be careful to only lowercase the base of the path, not the whole thing!!
		base := filepath.Base(path)
		if lowerBase := strings.ToLower(base); base != lowerBase {
			lowerPath := filepath.Join(filepath.Dir(path), lowerBase)
			err = os.Rename(path, lowerPath)
			if err != nil {
				log.Fatalf("[ERROR] Unable to lower-case: %v\n\nPlease follow instructions at:\nhttps://github.com/mholt/caddy/issues/902#issuecomment-228876011", err)
			}
		}
		return nil
	})
}

// setVersion figures out the version information
// based on variables set by -ldflags.
func setVersion() {
	// A development build is one that's not at a tag or has uncommitted changes
	devBuild = gitTag == "" || gitShortStat != ""

	// Only set the appVersion if -ldflags was used
	if gitNearestTag != "" || gitTag != "" {
		if devBuild && gitNearestTag != "" {
			appVersion = fmt.Sprintf("%s (+%s %s)",
				strings.TrimPrefix(gitNearestTag, "v"), gitCommit, buildDate)
		} else if gitTag != "" {
			appVersion = strings.TrimPrefix(gitTag, "v")
		}
	}
}

// setCPU parses string cpu and sets GOMAXPROCS
// according to its value. It accepts either
// a number (e.g. 3) or a percent (e.g. 50%).
func setCPU(cpu string) error {
	var numCPU int

	availCPU := runtime.NumCPU()

	if strings.HasSuffix(cpu, "%") {
		// Percent
		var percent float32
		pctStr := cpu[:len(cpu)-1]
		pctInt, err := strconv.Atoi(pctStr)
		if err != nil || pctInt < 1 || pctInt > 100 {
			return errors.New("invalid CPU value: percentage must be between 1-100")
		}
		percent = float32(pctInt) / 100
		numCPU = int(float32(availCPU) * percent)
	} else {
		// Number
		num, err := strconv.Atoi(cpu)
		if err != nil || num < 1 {
			return errors.New("invalid CPU value: provide a number or percent greater than 0")
		}
		numCPU = num
	}

	if numCPU > availCPU {
		numCPU = availCPU
	}

	runtime.GOMAXPROCS(numCPU)
	return nil
}

const appName = "Caddy"

// Flags that control program flow or startup
var (
	serverType string
	conf       string
	cpu        string
	logfile    string
	revoke     string
	version    bool
	plugins    bool
)

// Build information obtained with the help of -ldflags
var (
	appVersion = "(untracked dev build)" // inferred at startup
	devBuild   = true                    // inferred at startup

	buildDate        string // date -u
	gitTag           string // git describe --exact-match HEAD 2> /dev/null
	gitNearestTag    string // git describe --abbrev=0 --tags HEAD
	gitCommit        string // git rev-parse HEAD
	gitShortStat     string // git diff-index --shortstat
	gitFilesModified string // git diff-index --name-only HEAD
)
<commit_msg>Fix ACME asset migration when renaming folders<commit_after>package caddymain

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	"gopkg.in/natefinch/lumberjack.v2"

	"github.com/xenolf/lego/acme"

	"github.com/mholt/caddy"
	// plug in the HTTP server type
	_ "github.com/mholt/caddy/caddyhttp"

	"github.com/mholt/caddy/caddytls"
	// This is where other plugins get plugged in (imported)
)

func init() {
	caddy.TrapSignals()
	setVersion()

	flag.BoolVar(&caddytls.Agreed, "agree", false, "Agree to the CA's Subscriber Agreement")
	flag.StringVar(&caddytls.DefaultCAUrl, "ca", "https://acme-v01.api.letsencrypt.org/directory", "URL to certificate authority's ACME server directory")
	flag.StringVar(&conf, "conf", "", "Caddyfile to load (default \""+caddy.DefaultConfigFile+"\")")
	flag.StringVar(&cpu, "cpu", "100%", "CPU cap")
	flag.BoolVar(&plugins, "plugins", false, "List installed plugins")
	flag.StringVar(&caddytls.DefaultEmail, "email", "", "Default ACME CA account email address")
	flag.StringVar(&logfile, "log", "", "Process log file")
	flag.StringVar(&caddy.PidFile, "pidfile", "", "Path to write pid file")
	flag.BoolVar(&caddy.Quiet, "quiet", false, "Quiet mode (no initialization output)")
	flag.StringVar(&revoke, "revoke", "", "Hostname for which to revoke the certificate")
	flag.StringVar(&serverType, "type", "http", "Type of server to run")
	flag.BoolVar(&version, "version", false, "Show version")

	caddy.RegisterCaddyfileLoader("flag", caddy.LoaderFunc(confLoader))
	caddy.SetDefaultCaddyfileLoader("default", caddy.LoaderFunc(defaultLoader))
}

// Run is Caddy's main() function.
func Run() {
	flag.Parse()
	moveStorage() // TODO: This is temporary for the 0.9 release, or until most users upgrade to 0.9+

	caddy.AppName = appName
	caddy.AppVersion = appVersion
	acme.UserAgent = appName + "/" + appVersion

	// Set up process log before anything bad happens
	switch logfile {
	case "stdout":
		log.SetOutput(os.Stdout)
	case "stderr":
		log.SetOutput(os.Stderr)
	case "":
		log.SetOutput(ioutil.Discard)
	default:
		log.SetOutput(&lumberjack.Logger{
			Filename:   logfile,
			MaxSize:    100,
			MaxAge:     14,
			MaxBackups: 10,
		})
	}

	// Check for one-time actions
	if revoke != "" {
		err := caddytls.Revoke(revoke)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("Revoked certificate for %s\n", revoke)
		os.Exit(0)
	}
	if version {
		fmt.Printf("%s %s\n", appName, appVersion)
		if devBuild && gitShortStat != "" {
			fmt.Printf("%s\n%s\n", gitShortStat, gitFilesModified)
		}
		os.Exit(0)
	}
	if plugins {
		fmt.Println(caddy.DescribePlugins())
		os.Exit(0)
	}

	// Set CPU cap
	err := setCPU(cpu)
	if err != nil {
		mustLogFatal(err)
	}

	// Get Caddyfile input
	caddyfile, err := caddy.LoadCaddyfile(serverType)
	if err != nil {
		mustLogFatal(err)
	}

	// Start your engines
	instance, err := caddy.Start(caddyfile)
	if err != nil {
		mustLogFatal(err)
	}

	// Twiddle your thumbs
	instance.Wait()
}

// mustLogFatal wraps log.Fatal() in a way that ensures the
// output is always printed to stderr so the user can see it
// if the user is still there, even if the process log was not
// enabled. If this process is an upgrade, however, and the user
// might not be there anymore, this just logs to the process
// log and exits.
func mustLogFatal(args ...interface{}) {
	if !caddy.IsUpgrade() {
		log.SetOutput(os.Stderr)
	}
	log.Fatal(args...)
}

// confLoader loads the Caddyfile using the -conf flag.
func confLoader(serverType string) (caddy.Input, error) {
	if conf == "" {
		return nil, nil
	}

	if conf == "stdin" {
		return caddy.CaddyfileFromPipe(os.Stdin)
	}

	contents, err := ioutil.ReadFile(conf)
	if err != nil {
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       conf,
		ServerTypeName: serverType,
	}, nil
}

// defaultLoader loads the Caddyfile from the current working directory.
func defaultLoader(serverType string) (caddy.Input, error) {
	contents, err := ioutil.ReadFile(caddy.DefaultConfigFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       caddy.DefaultConfigFile,
		ServerTypeName: serverType,
	}, nil
}

// moveStorage moves the old certificate storage location by
// renaming the "letsencrypt" folder to the hostname of the
// CA URL. This is TEMPORARY until most users have upgraded to 0.9+.
func moveStorage() {
	oldPath := filepath.Join(caddy.AssetsPath(), "letsencrypt")
	_, err := os.Stat(oldPath)
	if os.IsNotExist(err) {
		return
	}
	// Just use a default config to get default (file) storage
	fileStorage, err := new(caddytls.Config).StorageFor(caddytls.DefaultCAUrl)
	if err != nil {
		log.Fatalf("[ERROR] Unable to get new path for certificate storage: %v", err)
	}
	newPath := string(fileStorage.(caddytls.FileStorage))
	err = os.MkdirAll(string(newPath), 0700)
	if err != nil {
		log.Fatalf("[ERROR] Unable to make new certificate storage path: %v\n\nPlease follow instructions at:\nhttps://github.com/mholt/caddy/issues/902#issuecomment-228876011", err)
	}
	err = os.Rename(oldPath, string(newPath))
	if err != nil {
		log.Fatalf("[ERROR] Unable to migrate certificate storage: %v\n\nPlease follow instructions at:\nhttps://github.com/mholt/caddy/issues/902#issuecomment-228876011", err)
	}
	// convert mixed case folder and file names to lowercase
	var done bool // walking is recursive and preloads the file names, so we must restart walk after a change until no changes
	for !done {
		done = true
		filepath.Walk(string(newPath), func(path string, info os.FileInfo, err error) error {
			// must be careful to only lowercase the base of the path, not the whole thing!!
			base := filepath.Base(path)
			if lowerBase := strings.ToLower(base); base != lowerBase {
				lowerPath := filepath.Join(filepath.Dir(path), lowerBase)
				err = os.Rename(path, lowerPath)
				if err != nil {
					log.Fatalf("[ERROR] Unable to lower-case: %v\n\nPlease follow instructions at:\nhttps://github.com/mholt/caddy/issues/902#issuecomment-228876011", err)
				}
				// terminate traversal and restart since Walk needs the updated file list with new file names
				done = false
				return errors.New("start over")
			}
			return nil
		})
	}
}

// setVersion figures out the version information
// based on variables set by -ldflags.
func setVersion() {
	// A development build is one that's not at a tag or has uncommitted changes
	devBuild = gitTag == "" || gitShortStat != ""

	// Only set the appVersion if -ldflags was used
	if gitNearestTag != "" || gitTag != "" {
		if devBuild && gitNearestTag != "" {
			appVersion = fmt.Sprintf("%s (+%s %s)",
				strings.TrimPrefix(gitNearestTag, "v"), gitCommit, buildDate)
		} else if gitTag != "" {
			appVersion = strings.TrimPrefix(gitTag, "v")
		}
	}
}

// setCPU parses string cpu and sets GOMAXPROCS
// according to its value. It accepts either
// a number (e.g. 3) or a percent (e.g. 50%).
func setCPU(cpu string) error {
	var numCPU int

	availCPU := runtime.NumCPU()

	if strings.HasSuffix(cpu, "%") {
		// Percent
		var percent float32
		pctStr := cpu[:len(cpu)-1]
		pctInt, err := strconv.Atoi(pctStr)
		if err != nil || pctInt < 1 || pctInt > 100 {
			return errors.New("invalid CPU value: percentage must be between 1-100")
		}
		percent = float32(pctInt) / 100
		numCPU = int(float32(availCPU) * percent)
	} else {
		// Number
		num, err := strconv.Atoi(cpu)
		if err != nil || num < 1 {
			return errors.New("invalid CPU value: provide a number or percent greater than 0")
		}
		numCPU = num
	}

	if numCPU > availCPU {
		numCPU = availCPU
	}

	runtime.GOMAXPROCS(numCPU)
	return nil
}

const appName = "Caddy"

// Flags that control program flow or startup
var (
	serverType string
	conf       string
	cpu        string
	logfile    string
	revoke     string
	version    bool
	plugins    bool
)

// Build information obtained with the help of -ldflags
var (
	appVersion = "(untracked dev build)" // inferred at startup
	devBuild   = true                    // inferred at startup

	buildDate        string // date -u
	gitTag           string // git describe --exact-match HEAD 2> /dev/null
	gitNearestTag    string // git describe --abbrev=0 --tags HEAD
	gitCommit        string // git rev-parse HEAD
	gitShortStat     string // git diff-index --shortstat
	gitFilesModified string // git diff-index --name-only HEAD
)
<|endoftext|>


<commit_before>// Copyright 2017 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package k8sTest

import (
	"fmt"
	"sync"
	"time"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
	"github.com/sirupsen/logrus"
)

var _ = Describe("K8sValidatedTunnelTest", func() {

	var kubectl *helpers.Kubectl
	var demoDSPath string
	var once sync.Once
	var logger *logrus.Entry

	initialize := func() {
		logger = log.WithFields(logrus.Fields{"testName": "K8sTunnelTest"})
		logger.Info("Starting")

		kubectl = helpers.CreateKubectl(helpers.K8s1VMName(), logger)
		demoDSPath = helpers.ManifestGet("demo_ds.yaml")
		kubectl.Exec("kubectl -n kube-system delete ds cilium")
		// Expect(res.Correct()).Should(BeTrue())

		waitToDeleteCilium(kubectl, logger)
	}

	BeforeEach(func() {
		once.Do(initialize)
		kubectl.NodeCleanMetadata()
		kubectl.Apply(demoDSPath)
	}, 600)

	AfterFailed(func() {
		kubectl.CiliumReport(helpers.KubeSystemNamespace,
			"cilium bpf tunnel list",
			"cilium endpoint list")
	})

	JustAfterEach(func() {
		kubectl.ValidateNoErrorsOnLogs(CurrentGinkgoTestDescription().Duration)
	})

	AfterEach(func() {
		kubectl.Delete(demoDSPath)
		ExpectAllPodsTerminated(kubectl)
	})

	Context("VXLan", func() {

		BeforeEach(func() {
			err := kubectl.CiliumInstall(helpers.CiliumDSPath)
			Expect(err).To(BeNil(), "Cilium cannot be installed")
		})

		AfterEach(func() {
			// Do not assert on success in AfterEach intentionally to avoid
			// incomplete teardown.
			_ = kubectl.DeleteResource(
				"ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))
			waitToDeleteCilium(kubectl, logger)
		})

		It("Check VXLAN mode", func() {

			ExpectCiliumReady(kubectl)

			ciliumPod, err := kubectl.GetCiliumPodOnNode(helpers.KubeSystemNamespace, helpers.K8s1)
			Expect(err).Should(BeNil())

			By("Making sure all endpoints are in ready state")
			err = kubectl.CiliumEndpointWaitReady()
			Expect(err).To(BeNil(), "Endpoints are not ready after timeout")

			_, err = kubectl.CiliumNodesWait()
			Expect(err).Should(BeNil())

			By("Checking that BPF tunnels are in place")
			status := kubectl.CiliumExec(ciliumPod, "cilium bpf tunnel list | wc -l")
			status.ExpectSuccess()
			Expect(status.IntOutput()).Should(Equal(5))

			By("Checking that BPF tunnels are working correctly")
			tunnStatus := isNodeNetworkingWorking(kubectl, "zgroup=testDS")
			Expect(tunnStatus).Should(BeTrue())
		}, 600)
	})

	Context("Geneve", func() {

		BeforeEach(func() {
			err := kubectl.CiliumInstall("cilium_ds_geneve.jsonnet")
			Expect(err).To(BeNil(), "Cilium cannot be installed")
		})

		AfterEach(func() {
			// Do not assert on success in AfterEach intentionally to avoid
			// incomplete teardown.
			_ = kubectl.DeleteResource(
				"ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))
			waitToDeleteCilium(kubectl, logger)
		})

		It("Check Geneve mode", func() {

			ExpectCiliumReady(kubectl)

			ciliumPod, err := kubectl.GetCiliumPodOnNode(helpers.KubeSystemNamespace, helpers.K8s1)
			Expect(err).Should(BeNil())

			_, err = kubectl.CiliumNodesWait()
			Expect(err).Should(BeNil())

			By("Making sure all endpoints are in ready state")
			err = kubectl.CiliumEndpointWaitReady()
			Expect(err).To(BeNil(), "Endpoints are not ready after timeout")

			//Check that cilium detects a
			By("Checking that BPF tunnels are in place")
			status := kubectl.CiliumExec(ciliumPod, "cilium bpf tunnel list | wc -l")
			status.ExpectSuccess()
			Expect(status.IntOutput()).Should(Equal(5))

			By("Checking that BPF tunnels are working correctly")
			tunnStatus := isNodeNetworkingWorking(kubectl, "zgroup=testDS")
			Expect(tunnStatus).Should(BeTrue())
			//FIXME: Maybe added here a cilium bpf tunnel status?
		}, 600)

	})

})

func isNodeNetworkingWorking(kubectl *helpers.Kubectl, filter string) bool {
	err := kubectl.WaitforPods(helpers.DefaultNamespace, fmt.Sprintf("-l %s", filter), 3000)
	Expect(err).Should(BeNil())
	pods, err := kubectl.GetPodNames(helpers.DefaultNamespace, filter)
	Expect(err).Should(BeNil())
	podIP, err := kubectl.Get(
		helpers.DefaultNamespace,
		fmt.Sprintf("pod %s -o json", pods[1])).Filter("{.status.podIP}")
	Expect(err).Should(BeNil())
	res := kubectl.ExecPodCmd(helpers.DefaultNamespace, pods[0], helpers.Ping(podIP.String()))
	return res.WasSuccessful()
}

func waitToDeleteCilium(kubectl *helpers.Kubectl, logger *logrus.Entry) {
	status := 1
	for status > 0 {
		pods, err := kubectl.GetCiliumPods(helpers.KubeSystemNamespace)
		status := len(pods)
		logger.Infof("Cilium pods terminating '%d' err='%v' pods='%v'", status, err, pods)
		if status == 0 {
			return
		}
		time.Sleep(1 * time.Second)
	}
}
<commit_msg>Test: K8st Tunnels delete services before delete Cilium<commit_after>// Copyright 2017 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package k8sTest

import (
	"fmt"
	"time"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
	"github.com/sirupsen/logrus"
)

var _ = Describe("K8sValidatedTunnelTest", func() {

	var kubectl *helpers.Kubectl
	var demoDSPath string
	var logger *logrus.Entry

	BeforeAll(func() {
		logger = log.WithFields(logrus.Fields{"testName": "K8sTunnelTest"})
		logger.Info("Starting")

		kubectl = helpers.CreateKubectl(helpers.K8s1VMName(), logger)
		demoDSPath = helpers.ManifestGet("demo_ds.yaml")

		kubectl.Exec("kubectl -n kube-system delete ds cilium")
		// Expect(res.Correct()).Should(BeTrue())

		waitToDeleteCilium(kubectl, logger)
	})

	BeforeEach(func() {
		kubectl.Apply(demoDSPath).ExpectSuccess("cannot install Demo application")
		kubectl.NodeCleanMetadata()
	})

	AfterEach(func() {
		_ = kubectl.Delete(demoDSPath)
	})

	AfterAll(func() {
		ExpectAllPodsTerminated(kubectl)
	})

	AfterFailed(func() {
		kubectl.CiliumReport(helpers.KubeSystemNamespace,
			"cilium bpf tunnel list",
			"cilium endpoint list")
	})

	JustAfterEach(func() {
		kubectl.ValidateNoErrorsOnLogs(CurrentGinkgoTestDescription().Duration)
	})

	cleanService := func() {
		// To avoid hit GH-4384
		kubectl.DeleteResource("service", "test-nodeport testds-service").ExpectSuccess(
			"Service is deleted")
	}

	Context("VXLan", func() {
		AfterEach(func() {
			// Do not assert on success in AfterEach intentionally to avoid
			// incomplete teardown.
			_ = kubectl.DeleteResource(
				"ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))
			waitToDeleteCilium(kubectl, logger)
		})

		It("Check VXLAN mode", func() {

			err := kubectl.CiliumInstall(helpers.CiliumDSPath)
			Expect(err).To(BeNil(), "Cilium cannot be installed")

			ExpectCiliumReady(kubectl)

			err = kubectl.WaitforPods(helpers.DefaultNamespace, "", 300)
			Expect(err).Should(BeNil(), "Pods are not ready after timeout")

			ciliumPod, err := kubectl.GetCiliumPodOnNode(helpers.KubeSystemNamespace, helpers.K8s1)
			Expect(err).Should(BeNil())

			By("Making sure all endpoints are in ready state")
			err = kubectl.CiliumEndpointWaitReady()
			Expect(err).To(BeNil(), "Endpoints are not ready after timeout")

			_, err = kubectl.CiliumNodesWait()
			Expect(err).Should(BeNil())

			By("Checking that BPF tunnels are in place")
			status := kubectl.CiliumExec(ciliumPod, "cilium bpf tunnel list | wc -l")
			status.ExpectSuccess()
			Expect(status.IntOutput()).Should(Equal(5))

			By("Checking that BPF tunnels are working correctly")
			tunnStatus := isNodeNetworkingWorking(kubectl, "zgroup=testDS")
			Expect(tunnStatus).Should(BeTrue())

			// FIXME GH-4456
			cleanService()
		}, 600)
	})

	Context("Geneve", func() {

		AfterEach(func() {
			// Do not assert on success in AfterEach intentionally to avoid
			// incomplete teardown.
			_ = kubectl.DeleteResource(
				"ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))
			waitToDeleteCilium(kubectl, logger)
		})

		It("Check Geneve mode", func() {

			err := kubectl.CiliumInstall("cilium_ds_geneve.jsonnet")
			Expect(err).To(BeNil(), "Cilium cannot be installed")

			ExpectCiliumReady(kubectl)

			err = kubectl.WaitforPods(helpers.DefaultNamespace, "", 300)
			Expect(err).Should(BeNil(), "Pods are not ready after timeout")

			ciliumPod, err := kubectl.GetCiliumPodOnNode(helpers.KubeSystemNamespace, helpers.K8s1)
			Expect(err).Should(BeNil())

			_, err = kubectl.CiliumNodesWait()
			Expect(err).Should(BeNil())

			By("Making sure all endpoints are in ready state")
			err = kubectl.CiliumEndpointWaitReady()
			Expect(err).To(BeNil(), "Endpoints are not ready after timeout")

			//Check that cilium detects a
			By("Checking that BPF tunnels are in place")
			status := kubectl.CiliumExec(ciliumPod, "cilium bpf tunnel list | wc -l")
			status.ExpectSuccess()
			Expect(status.IntOutput()).Should(Equal(5))

			By("Checking that BPF tunnels are working correctly")
			tunnStatus := isNodeNetworkingWorking(kubectl, "zgroup=testDS")
			Expect(tunnStatus).Should(BeTrue())

			// FIXME GH-4456
			cleanService()
		}, 600)

	})

})

func isNodeNetworkingWorking(kubectl *helpers.Kubectl, filter string) bool {
	err := kubectl.WaitforPods(helpers.DefaultNamespace, fmt.Sprintf("-l %s", filter), 3000)
	Expect(err).Should(BeNil())
	pods, err := kubectl.GetPodNames(helpers.DefaultNamespace, filter)
	Expect(err).Should(BeNil())
	podIP, err := kubectl.Get(
		helpers.DefaultNamespace,
		fmt.Sprintf("pod %s -o json", pods[1])).Filter("{.status.podIP}")
	Expect(err).Should(BeNil())
	res := kubectl.ExecPodCmd(helpers.DefaultNamespace, pods[0], helpers.Ping(podIP.String()))
	return res.WasSuccessful()
}

func waitToDeleteCilium(kubectl *helpers.Kubectl, logger *logrus.Entry) {
	status := 1
	for status > 0 {
		pods, err := kubectl.GetCiliumPods(helpers.KubeSystemNamespace)
		status := len(pods)
		logger.Infof("Cilium pods terminating '%d' err='%v' pods='%v'", status, err, pods)
		if status == 0 {
			return
		}
		time.Sleep(1 * time.Second)
	}
}
<|endoftext|>


<commit_before>package vnet

import (
	"net"
	"sync"
	"time"
)

// UDPProxy is a proxy between real server(net.UDPConn) and vnet.UDPConn.
//
// High level design:
//                           ..............................................
//                           :         Virtual Network (vnet)             :
//                           :                                            :
//   +-------+ *         1 +----+         +--------+                      :
//   | :App  |------------>|:Net|--o<-----|:Router |          .............................
//   +-------+             +----+         |        |          :        UDPProxy           :
//                           :            |        |       +----+     +---------+     +---------+     +--------+
//                           :            |        |--->o--|:Net|-->o-| vnet.   |-->o-|  net.   |--->-| :Real  |
//                           :            |        |       +----+     | UDPConn |     | UDPConn |     | Server |
//                           :            |        |          :       +---------+     +---------+     +--------+
//                           :            |        |          ............................:
//                           :            +--------+                       :
//                           ...............................................
type UDPProxy struct {
	// The router bind to.
	router *Router

	// Each vnet source, bind to a real socket to server.
	// key is real server addr, which is net.Addr
	// value is *aUDPProxyWorker
	workers sync.Map

	// For each endpoint, we never know when to start and stop proxy,
	// so we stop the endpoint when timeout.
	timeout time.Duration

	// For utest, to mock the target real server.
	// Optional, use the address of received client packet.
	mockRealServerAddr *net.UDPAddr
}

// NewProxy create a proxy, the router for this proxy belongs/bind to. If need to proxy for
// please create a new proxy for each router. For all addresses we proxy, we will create a
// vnet.Net in this router and proxy all packets.
func NewProxy(router *Router) (*UDPProxy, error) {
	v := &UDPProxy{router: router, timeout: 2 * time.Minute}
	return v, nil
}

// Close the proxy, stop all workers.
func (v *UDPProxy) Close() error {
	// nolint:godox // TODO: FIXME: Do cleanup.
	return nil
}

// Proxy starts a worker for server, ignore if already started.
func (v *UDPProxy) Proxy(client *Net, server *net.UDPAddr) error {
	// Note that even if the worker exists, it's also ok to create a same worker,
	// because the router will use the last one, and the real server will see a address
	// change event after we switch to the next worker.
	if _, ok := v.workers.Load(server.String()); ok {
		// nolint:godox // TODO: Need to restart the stopped worker?
		return nil
	}

	// Not exists, create a new one.
	worker := &aUDPProxyWorker{
		router: v.router, mockRealServerAddr: v.mockRealServerAddr,
	}
	v.workers.Store(server.String(), worker)

	return worker.Proxy(client, server)
}

// A proxy worker for a specified proxy server.
type aUDPProxyWorker struct {
	router             *Router
	mockRealServerAddr *net.UDPAddr

	// Each vnet source, bind to a real socket to server.
	// key is vnet client addr, which is net.Addr
	// value is *net.UDPConn
	endpoints sync.Map
}

func (v *aUDPProxyWorker) Proxy(client *Net, serverAddr *net.UDPAddr) error { // nolint:gocognit
	// Create vnet for real server by serverAddr.
	nw := NewNet(&NetConfig{
		StaticIP: serverAddr.IP.String(),
	})
	if err := v.router.AddNet(nw); err != nil {
		return err
	}

	// We must create a "same" vnet.UDPConn as the net.UDPConn,
	// which has the same ip:port, to copy packets between them.
	vnetSocket, err := nw.ListenUDP("udp4", serverAddr)
	if err != nil {
		return err
	}

	// Got new vnet client, start a new endpoint.
	findEndpointBy := func(addr net.Addr) (*net.UDPConn, error) {
		// Exists binding.
		if value, ok := v.endpoints.Load(addr.String()); ok {
			// Exists endpoint, reuse it.
			return value.(*net.UDPConn), nil
		}

		// The real server we proxy to, for utest to mock it.
		realAddr := serverAddr
		if v.mockRealServerAddr != nil {
			realAddr = v.mockRealServerAddr
		}

		// Got new vnet client, create new endpoint.
		realSocket, err := net.DialUDP("udp4", nil, realAddr)
		if err != nil {
			return nil, err
		}

		// Bind address.
		v.endpoints.Store(addr.String(), realSocket)

		// Got packet from real serverAddr, we should proxy it to vnet.
		// nolint:godox // TODO: FIXME: Do cleanup.
		go func(vnetClientAddr net.Addr) {
			buf := make([]byte, 1500)
			for {
				n, _, err := realSocket.ReadFrom(buf)
				if err != nil {
					return
				}

				if n <= 0 {
					continue // Drop packet
				}

				if _, err := vnetSocket.WriteTo(buf[:n], vnetClientAddr); err != nil {
					return
				}
			}
		}(addr)

		return realSocket, nil
	}

	// Start a proxy goroutine.
	// nolint:godox // TODO: FIXME: Do cleanup.
	go func() {
		buf := make([]byte, 1500)

		for {
			n, addr, err := vnetSocket.ReadFrom(buf)
			if err != nil {
				return
			}

			if n <= 0 || addr == nil {
				continue // Drop packet
			}

			realSocket, err := findEndpointBy(addr)
			if err != nil {
				continue // Drop packet.
			}

			if _, err := realSocket.Write(buf[:n]); err != nil {
				return
			}
		}
	}()

	return nil
}
<commit_msg>Do cleanup when close UDPProxy<commit_after>package vnet

import (
	"context"
	"net"
	"sync"
	"time"
)

// UDPProxy is a proxy between real server(net.UDPConn) and vnet.UDPConn.
//
// High level design:
//                           ..............................................
//                           :         Virtual Network (vnet)             :
//                           :                                            :
//   +-------+ *         1 +----+         +--------+                      :
//   | :App  |------------>|:Net|--o<-----|:Router |          .............................
//   +-------+             +----+         |        |          :        UDPProxy           :
//                           :            |        |       +----+     +---------+     +---------+     +--------+
//                           :            |        |--->o--|:Net|-->o-| vnet.   |-->o-|  net.   |--->-| :Real  |
//                           :            |        |       +----+     | UDPConn |     | UDPConn |     | Server |
//                           :            |        |          :       +---------+     +---------+     +--------+
//                           :            |        |          ............................:
//                           :            +--------+                       :
//                           ...............................................
type UDPProxy struct {
	// The router bind to.
	router *Router

	// Each vnet source, bind to a real socket to server.
	// key is real server addr, which is net.Addr
	// value is *aUDPProxyWorker
	workers sync.Map

	// For each endpoint, we never know when to start and stop proxy,
	// so we stop the endpoint when timeout.
	timeout time.Duration

	// For utest, to mock the target real server.
	// Optional, use the address of received client packet.
	mockRealServerAddr *net.UDPAddr
}

// NewProxy create a proxy, the router for this proxy belongs/bind to. If need to proxy for
// please create a new proxy for each router. For all addresses we proxy, we will create a
// vnet.Net in this router and proxy all packets.
func NewProxy(router *Router) (*UDPProxy, error) {
	v := &UDPProxy{router: router, timeout: 2 * time.Minute}
	return v, nil
}

// Close the proxy, stop all workers.
func (v *UDPProxy) Close() error {
	v.workers.Range(func(key, value interface{}) bool {
		_ = value.(*aUDPProxyWorker).Close()
		return true
	})
	return nil
}

// Proxy starts a worker for server, ignore if already started.
func (v *UDPProxy) Proxy(client *Net, server *net.UDPAddr) error {
	// Note that even if the worker exists, it's also ok to create a same worker,
	// because the router will use the last one, and the real server will see a address
	// change event after we switch to the next worker.
	if _, ok := v.workers.Load(server.String()); ok {
		// nolint:godox // TODO: Need to restart the stopped worker?
		return nil
	}

	// Not exists, create a new one.
	worker := &aUDPProxyWorker{
		router: v.router, mockRealServerAddr: v.mockRealServerAddr,
	}

	// Create context for cleanup.
	var ctx context.Context
	ctx, worker.ctxDisposeCancel = context.WithCancel(context.Background())

	v.workers.Store(server.String(), worker)

	return worker.Proxy(ctx, client, server)
}

// A proxy worker for a specified proxy server.
type aUDPProxyWorker struct {
	router             *Router
	mockRealServerAddr *net.UDPAddr

	// Each vnet source, bind to a real socket to server.
	// key is vnet client addr, which is net.Addr
	// value is *net.UDPConn
	endpoints sync.Map

	// For cleanup.
	ctxDisposeCancel context.CancelFunc
	wg               sync.WaitGroup
}

func (v *aUDPProxyWorker) Close() error {
	// Notify all goroutines to dispose.
	v.ctxDisposeCancel()

	// Wait for all goroutines quit.
	v.wg.Wait()

	return nil
}

func (v *aUDPProxyWorker) Proxy(ctx context.Context, client *Net, serverAddr *net.UDPAddr) error { // nolint:gocognit
	// Create vnet for real server by serverAddr.
	nw := NewNet(&NetConfig{
		StaticIP: serverAddr.IP.String(),
	})
	if err := v.router.AddNet(nw); err != nil {
		return err
	}

	// We must create a "same" vnet.UDPConn as the net.UDPConn,
	// which has the same ip:port, to copy packets between them.
	vnetSocket, err := nw.ListenUDP("udp4", serverAddr)
	if err != nil {
		return err
	}

	// User stop proxy, we should close the socket.
	go func() {
		<-ctx.Done()
		_ = vnetSocket.Close()
	}()

	// Got new vnet client, start a new endpoint.
	findEndpointBy := func(addr net.Addr) (*net.UDPConn, error) {
		// Exists binding.
		if value, ok := v.endpoints.Load(addr.String()); ok {
			// Exists endpoint, reuse it.
			return value.(*net.UDPConn), nil
		}

		// The real server we proxy to, for utest to mock it.
		realAddr := serverAddr
		if v.mockRealServerAddr != nil {
			realAddr = v.mockRealServerAddr
		}

		// Got new vnet client, create new endpoint.
		realSocket, err := net.DialUDP("udp4", nil, realAddr)
		if err != nil {
			return nil, err
		}

		// User stop proxy, we should close the socket.
		go func() {
			<-ctx.Done()
			_ = realSocket.Close()
		}()

		// Bind address.
		v.endpoints.Store(addr.String(), realSocket)

		// Got packet from real serverAddr, we should proxy it to vnet.
		v.wg.Add(1)
		go func(vnetClientAddr net.Addr) {
			defer v.wg.Done()

			buf := make([]byte, 1500)
			for {
				n, _, err := realSocket.ReadFrom(buf)
				if err != nil {
					return
				}

				if n <= 0 {
					continue // Drop packet
				}

				if _, err := vnetSocket.WriteTo(buf[:n], vnetClientAddr); err != nil {
					return
				}
			}
		}(addr)

		return realSocket, nil
	}

	// Start a proxy goroutine.
	v.wg.Add(1)
	go func() {
		defer v.wg.Done()

		buf := make([]byte, 1500)

		for {
			n, addr, err := vnetSocket.ReadFrom(buf)
			if err != nil {
				return
			}

			if n <= 0 || addr == nil {
				continue // Drop packet
			}

			realSocket, err := findEndpointBy(addr)
			if err != nil {
				continue // Drop packet.
			}

			if _, err := realSocket.Write(buf[:n]); err != nil {
				return
			}
		}
	}()

	return nil
}
<|endoftext|>


<commit_before>package vulcan

import (
	"encoding/json"
	"fmt"

	"github.com/coreos/go-etcd/etcd"
)

const (
	etcdMachine   = "http://127.0.0.1:4001"
	frontendKey   = "%s/frontends/%s.%s/frontend"
	middlewareKey = "%s/frontends/%s.%s/middlewares/%s"
	backendKey    = "%s/backends/%s/backend"
	serverKey     = "%s/backends/%s/servers/%s"
)

type Client struct {
	Key  string
	etcd *etcd.Client
}

func NewClient(key string) *Client {
	etcd := etcd.NewClient([]string{etcdMachine})

	return &Client{Key: key, etcd: etcd}
}

func (c *Client) CreateServer(endpoint *Endpoint, ttl uint64) error {
	key := fmt.Sprintf(serverKey, c.Key, endpoint.Name, endpoint.ID)
	server, err := endpoint.ServerSpec()
	if err != nil {
		return nil
	}

	_, err = c.etcd.Create(key, server, ttl)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) UpdateServer(endpoint *Endpoint, ttl uint64) error {
	key := fmt.Sprintf(serverKey, c.Key, endpoint.Name, endpoint.ID)
	server, err := endpoint.ServerSpec()
	if err != nil {
		return nil
	}

	_, err = c.etcd.CompareAndSwap(key, server, ttl, server, 0)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) UpsertServer(endpoint *Endpoint, ttl uint64) error {
	key := fmt.Sprintf(serverKey, c.Key, endpoint.Name, endpoint.ID)
	server, err := endpoint.ServerSpec()
	if err != nil {
		return nil
	}

	_, err = c.etcd.Set(key, server, ttl)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) RegisterBackend(endpoint *Endpoint) error {
	key := fmt.Sprintf(backendKey, c.Key, endpoint.Name)
	backend, err := endpoint.BackendSpec()
	if err != nil {
		return err
	}

	_, err = c.etcd.Set(key, backend, 0)
	if err != nil {
		return err
	}

	return err
}

func (c *Client) RegisterFrontend(location *Location) error {
	key := fmt.Sprintf(frontendKey, c.Key, location.Host, location.ID)
	frontend, err := location.Spec()
	if err != nil {
		return err
	}

	_, err = c.etcd.Set(key, frontend, 0)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) RegisterMiddleware(location *Location) error {
	for i, m := range location.Middlewares {
		m.Priority = i

		key := fmt.Sprintf(middlewareKey, c.Key, location.Host, location.ID, m.ID)
		middleware, err := json.Marshal(m)
		if err != nil {
			return err
		}

		_, err = c.etcd.Set(key, string(middleware), 0)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Sync etcd cluster when creating a vulcan client<commit_after>package vulcan

import (
	"encoding/json"
	"fmt"

	"github.com/coreos/go-etcd/etcd"
)

const (
	etcdMachine   = "http://127.0.0.1:4001"
	frontendKey   = "%s/frontends/%s.%s/frontend"
	middlewareKey = "%s/frontends/%s.%s/middlewares/%s"
	backendKey    = "%s/backends/%s/backend"
	serverKey     = "%s/backends/%s/servers/%s"
)

type Client struct {
	Key  string
	etcd *etcd.Client
}

func NewClient(key string) *Client {
	etcd := etcd.NewClient([]string{etcdMachine})

	etcd.SyncCluster()

	return &Client{Key: key, etcd: etcd}
}

func (c *Client) CreateServer(endpoint *Endpoint, ttl uint64) error {
	key := fmt.Sprintf(serverKey, c.Key, endpoint.Name, endpoint.ID)
	server, err := endpoint.ServerSpec()
	if err != nil {
		return nil
	}

	_, err = c.etcd.Create(key, server, ttl)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) UpdateServer(endpoint *Endpoint, ttl uint64) error {
	key := fmt.Sprintf(serverKey, c.Key, endpoint.Name, endpoint.ID)
	server, err := endpoint.ServerSpec()
	if err != nil {
		return nil
	}

	_, err = c.etcd.CompareAndSwap(key, server, ttl, server, 0)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) UpsertServer(endpoint *Endpoint, ttl uint64) error {
	key := fmt.Sprintf(serverKey, c.Key, endpoint.Name, endpoint.ID)
	server, err := endpoint.ServerSpec()
	if err != nil {
		return nil
	}

	_, err = c.etcd.Set(key, server, ttl)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) RegisterBackend(endpoint *Endpoint) error {
	key := fmt.Sprintf(backendKey, c.Key, endpoint.Name)
	backend, err := endpoint.BackendSpec()
	if err != nil {
		return err
	}

	_, err = c.etcd.Set(key, backend, 0)
	if err != nil {
		return err
	}

	return err
}

func (c *Client) RegisterFrontend(location *Location) error {
	key := fmt.Sprintf(frontendKey, c.Key, location.Host, location.ID)
	frontend, err := location.Spec()
	if err != nil {
		return err
	}

	_, err = c.etcd.Set(key, frontend, 0)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) RegisterMiddleware(location *Location) error {
	for i, m := range location.Middlewares {
		m.Priority = i

		key := fmt.Sprintf(middlewareKey, c.Key, location.Host, location.ID, m.ID)
		middleware, err := json.Marshal(m)
		if err != nil {
			return err
		}

		_, err = c.etcd.Set(key, string(middleware), 0)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]

// runner-scribe is a mig-runner plugin that processes results coming from automated
// actions and forwards the results as vulnerability events to MozDef
package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"

	"github.com/mozilla/gozdef"
	"github.com/mozilla/mig"
	scribemod "github.com/mozilla/mig/modules/scribe"
	"gopkg.in/gcfg.v1"
)

const sourceName = "runner-scribe"

// config represents the configuration used by runner-scribe, and is read in on
// initialization
//
// URL is mandatory
type config struct {
	MozDef struct {
		URL      string // URL to post events to MozDef
		UseProxy bool   // A switch to enable/disable the use of a system-configured proxy
	}
}

const configPath string = "/etc/mig/runner-scribe.conf"

var conf config

func main() {
	defer func() {
		if e := recover(); e != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", e)
			os.Exit(1)
		}
	}()

	var (
		err     error
		results mig.RunnerResult
	)

	err = gcfg.ReadFileInto(&conf, configPath)
	if err != nil {
		panic(err)
	}

	buf, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		panic(err)
	}
	err = json.Unmarshal(buf, &results)
	if err != nil {
		panic(err)
	}
	var items []gozdef.VulnEvent
	for _, x := range results.Commands {
		// Process the incoming commands, under normal circumstances we will have one
		// returned command per host. However, this function can handle cases where
		// more than one command applies to a given host. If data for a host already
		// exists in items, makeVulnerability should attempt to append this data to
		// the host rather than add a new item.
		var err error
		items, err = makeVulnerability(items, x)
		if err != nil {
			panic(err)
		}
	}
	for _, y := range items {
		y.SourceName = sourceName
		err = sendVulnerability(y)
		if err != nil {
			panic(err)
		}
	}
}

func sendVulnerability(item gozdef.VulnEvent) (err error) {
	ac := gozdef.APIConf{
		URL:      conf.MozDef.URL,
		UseProxy: conf.MozDef.UseProxy,
	}
	pub, err := gozdef.InitAPI(ac)
	if err != nil {
		return
	}
	err = pub.Send(item)
	return
}

func makeVulnerability(initems []gozdef.VulnEvent, cmd mig.Command) (items []gozdef.VulnEvent, err error) {
	var (
		itemptr                       *gozdef.VulnEvent
		assethostname, assetipaddress string
		insertNew                     bool
	)
	items = initems

	assethostname = cmd.Agent.Name
	for _, x := range cmd.Agent.Env.Addresses {
		if !strings.Contains(x, ".") {
			continue
		}
		ipt, _, err := net.ParseCIDR(x)
		if err != nil {
			continue
		}
		assetipaddress = ipt.String()
		break
	}

	// First, see if we can locate a preexisting item for this asset
	for i := range items {
		if items[i].Asset.Hostname == assethostname &&
			items[i].Asset.IPAddress == assetipaddress {
			itemptr = &items[i]
			break
		}
	}
	if itemptr == nil {
		// Initialize a new event we will insert later
		newevent, err := gozdef.NewVulnEvent()
		if err != nil {
			return items, err
		}
		newevent.Description = "MIG vulnerability identification"
		newevent.Zone = "mig"
		newevent.Asset.Hostname = assethostname
		newevent.Asset.IPAddress = assetipaddress
		newevent.Asset.OS = cmd.Agent.Env.OS
		if len(cmd.Agent.Tags) != 0 {
			if _, ok := cmd.Agent.Tags["operator"]; ok {
				newevent.Asset.Owner.Operator = cmd.Agent.Tags["operator"]
			}
		}
		// Apply a v2bkey to the event. This should be set using integration
		// with service-map, but here for now we just apply it based on the operator
		// and team values which may be present in the event.
		if newevent.Asset.Owner.V2Bkey == "" {
			if newevent.Asset.Owner.Operator != "" {
				newevent.Asset.Owner.V2Bkey = newevent.Asset.Owner.Operator
			}
			if newevent.Asset.Owner.Team != "" {
				newevent.Asset.Owner.V2Bkey += "-" + newevent.Asset.Owner.Team
			}
		}
		// Always set credentialed checks here
		newevent.CredentialedChecks = true
		insertNew = true
		itemptr = &newevent
	}

	for _, result := range cmd.Results {
		var el scribemod.ScribeElements
		err = result.GetElements(&el)
		if err != nil {
			return items, err
		}
		for _, x := range el.Results {
			if !x.MasterResult {
				// Result was false (vulnerability did not match)
				continue
			}
			newve := gozdef.VulnVuln{}
			newve.Name = x.TestName
			for _, y := range x.Tags {
				if y.Key == "severity" {
					newve.Risk = y.Value
				} else if y.Key == "link" {
					newve.Link = y.Value
				}
			}
			// If no risk value is set on the vulnerability, we just treat this as
			// informational and ignore it. This will apply to things like the result
			// from platform dependency checks associated with real vulnerability checks.
			if newve.Risk == "" {
				continue
			}
			newve.Risk = normalizeRisk(newve.Risk)
			newve.LikelihoodIndicator = likelihoodFromRisk(newve.Risk)
			if newve.CVSS == "" {
				newve.CVSS = cvssFromRisk(newve.Risk)
			}
			// Use the identifier for each true subresult in the
			// test as a proof section
			for _, y := range x.Results {
				if y.Result {
					newve.Packages = append(newve.Packages, y.Identifier)
				}
			}
			itemptr.Vuln = append(itemptr.Vuln, newve)
		}
	}
	if insertNew {
		items = append(items, *itemptr)
	}
	return
}

// cvssFromRisk returns a synthesized CVSS score as a string given a risk label
func cvssFromRisk(risk string) string {
	switch risk {
	case "critical":
		return "10.0"
	case "high":
		return "8.0"
	case "medium":
		return "5.0"
	case "low":
		return "2.5"
	}
	return "0.0"
}

// likelihoodFromRisk returns a likelihood indicator value given a risk label
func likelihoodFromRisk(risk string) string {
	switch risk {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "maximum"
	}
	return "unknown"
}

// normalizeRisk converts known risk labels into a standardized form, if we can't identify
// the value we just return it as is
func normalizeRisk(in string) string {
	switch strings.ToLower(in) {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "critical"
	}
	return in
}
<commit_msg>Add ability to request an Auth0 token, query ServiceAPI, and lookup results<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]
// Contributor: Tristan Weir tweir@mozilla.com [:weir]

// runner-scribe is a mig-runner plugin that processes results coming from automated
// actions and forwards the results as vulnerability events to MozDef
package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strings"

	"github.com/mozilla/gozdef"
	"github.com/mozilla/mig"
	scribemod "github.com/mozilla/mig/modules/scribe"
	"gopkg.in/gcfg.v1"
)

const sourceName = "runner-scribe"

// config represents the configuration used by runner-scribe, and is read in on
// initialization
//
// URL is mandatory
type config struct {
	MozDef struct {
		URL      string // URL to post events to MozDef
		UseProxy bool   // A switch to enable/disable the use of a system-configured proxy
	}
	api 		ServiceApi
}

type ServiceApiAsset struct {
	Id 				string
	Asset_type 		string
	Asset_identifier string
	Team 			string
	Operator 		string
	Zone 			string
	Timestamp 		string
	Description 	string
	Score 			int
}

type ServiceApi struct {
	URL				string
	AuthEndpoint 	string
	ClientID 		string
	ClientSecret	string
	Token 			string // ephemeral token we generate to connect to ServiceAPI
}

type Auth0Token struct {
	Access_token	string
	Scope			string
	Expires_in		int
	Token_type		string
}

const configPath string = "/etc/mig/runner-scribe.conf"

var conf config

func main() {
	defer func() {
		if e := recover(); e != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", e)
			os.Exit(1)
		}
	}()

	var (
		err     error
		results mig.RunnerResult
	)

	err = gcfg.ReadFileInto(&conf, configPath)
	if err != nil {
		panic(err)
	}

	buf, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		panic(err)
	}
	err = json.Unmarshal(buf, &results)
	if err != nil {
		panic(err)
	}
	var items []gozdef.VulnEvent
	for _, x := range results.Commands {
		// Process the incoming commands, under normal circumstances we will have one
		// returned command per host. However, this function can handle cases where
		// more than one command applies to a given host. If data for a host already
		// exists in items, makeVulnerability should attempt to append this data to
		// the host rather than add a new item.
		var err error
		items, err = makeVulnerability(items, x)
		if err != nil {
			panic(err)
		}
	}
	for _, y := range items {
		y.SourceName = sourceName
		err = sendVulnerability(y)
		if err != nil {
			panic(err)
		}
	}
}

func sendVulnerability(item gozdef.VulnEvent) (err error) {
	ac := gozdef.APIConf{
		URL:      conf.MozDef.URL,
		UseProxy: conf.MozDef.UseProxy,
	}
	pub, err := gozdef.InitAPI(ac)
	if err != nil {
		return
	}
	err = pub.Send(item)
	return
}

func makeVulnerability(initems []gozdef.VulnEvent, cmd mig.Command) (items []gozdef.VulnEvent, err error) {
	var (
		itemptr                       *gozdef.VulnEvent
		assethostname, assetipaddress string
		insertNew                     bool
	)
	items = initems

	assethostname = cmd.Agent.Name
	for _, x := range cmd.Agent.Env.Addresses {
		if !strings.Contains(x, ".") {
			continue
		}
		ipt, _, err := net.ParseCIDR(x)
		if err != nil {
			continue
		}
		assetipaddress = ipt.String()
		break
	}

	// First, see if we can locate a preexisting item for this asset
	for i := range items {
		if items[i].Asset.Hostname == assethostname &&
			items[i].Asset.IPAddress == assetipaddress {
			itemptr = &items[i]
			break
		}
	}
	if itemptr == nil {
		// Initialize a new event we will insert later
		newevent, err := gozdef.NewVulnEvent()
		if err != nil {
			return items, err
		}
		newevent.Description = "MIG vulnerability identification"
		newevent.Zone = "mig"
		newevent.Asset.Hostname = assethostname
		newevent.Asset.IPAddress = assetipaddress
		newevent.Asset.OS = cmd.Agent.Env.OS
		if len(cmd.Agent.Tags) != 0 {
			if _, ok := cmd.Agent.Tags["operator"]; ok {
				newevent.Asset.Owner.Operator = cmd.Agent.Tags["operator"]
			}
		}
		// Apply a v2bkey to the event. This should be set using integration
		// with service-map, but here for now we just apply it based on the operator
		// and team values which may be present in the event.
		if newevent.Asset.Owner.V2Bkey == "" {
			if newevent.Asset.Owner.Operator != "" {
				newevent.Asset.Owner.V2Bkey = newevent.Asset.Owner.Operator
			}
			if newevent.Asset.Owner.Team != "" {
				newevent.Asset.Owner.V2Bkey += "-" + newevent.Asset.Owner.Team
			}
		}
		// Always set credentialed checks here
		newevent.CredentialedChecks = true
		insertNew = true
		itemptr = &newevent
	}

	for _, result := range cmd.Results {
		var el scribemod.ScribeElements
		err = result.GetElements(&el)
		if err != nil {
			return items, err
		}
		for _, x := range el.Results {
			if !x.MasterResult {
				// Result was false (vulnerability did not match)
				continue
			}
			newve := gozdef.VulnVuln{}
			newve.Name = x.TestName
			for _, y := range x.Tags {
				if y.Key == "severity" {
					newve.Risk = y.Value
				} else if y.Key == "link" {
					newve.Link = y.Value
				}
			}
			// If no risk value is set on the vulnerability, we just treat this as
			// informational and ignore it. This will apply to things like the result
			// from platform dependency checks associated with real vulnerability checks.
			if newve.Risk == "" {
				continue
			}
			newve.Risk = normalizeRisk(newve.Risk)
			newve.LikelihoodIndicator = likelihoodFromRisk(newve.Risk)
			if newve.CVSS == "" {
				newve.CVSS = cvssFromRisk(newve.Risk)
			}
			// Use the identifier for each true subresult in the
			// test as a proof section
			for _, y := range x.Results {
				if y.Result {
					newve.Packages = append(newve.Packages, y.Identifier)
				}
			}
			itemptr.Vuln = append(itemptr.Vuln, newve)
		}
	}
	if insertNew {
		items = append(items, *itemptr)
	}
	return
}

// given config for an API behind Auth0 (including client ID and Secret),
// return an Auth0 access token beginning with "Bearer "
// pattern from https://auth0.com/docs/api-auth/tutorials/client-credentials
func GetAuthToken(api ServiceApi) (authToken string) {
	payload := strings.NewReader("{\"grant_type\":\"client_credentials\",\"client_id\": \"" + api.ClientID + "\",\"client_secret\": \"" + api.ClientSecret + "\",\"audience\": \"" + api.URL + "\"}")
	req, _ := http.NewRequest("POST", api.AuthEndpoint, payload)
	req.Header.Add("content-type", "application/json")

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		panic(err)
	}

	defer res.Body.Close()
	bodyJSON, _ := ioutil.ReadAll(res.Body)

	// unpack the JSON into an Auth0 token struct
	var body Auth0Token
	err = json.Unmarshal(bodyJSON, &body)
	if err != nil {
		panic(err)
	}

	// serviceAPI expects the Access token in the form of "Bearer <token>"
	authToken = "Bearer " + body.Access_token
	return
}

// query a ServiceAPI instance for the set of all assets
// load them into a searchable map, keyed to asset hostname
// the ServiceAPI object must already be loaded with a Bearer token
func GetAssets(m map[string]ServiceApiAsset, api ServiceApi) (err error){

	// get json array of assets from serviceapi
	requestURL := api.URL + "api/v1/assets/"
	req, err := http.NewRequest(http.MethodGet, requestURL, nil)
	if err != nil {
		return err
	}
	req.Header.Add("accept", "application/json")
	req.Header.Add("Authorization", api.Token)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	// unpack the HTTP request response
	defer res.Body.Close()
	body, readErr := ioutil.ReadAll(res.Body)
	if readErr != nil {
		return readErr
	}

	// because of the way that ServiceAPI returns the JSON content,
	// we need to Unmarshal it twice
	var allAssetsJson string
	err  = json.Unmarshal(body, &allAssetsJson)
	if err != nil {
		return err
	}

	// convert json into array of ServiceApiAsset objects
	var allAssets []ServiceApiAsset
	err  = json.Unmarshal([]byte(allAssetsJson), &allAssets)
	if err != nil {
		return err
	}

	// build a searchable map, keyed on Asset_identifier (which is usually hostname)
	for _, tempAsset := range allAssets {
		permanentAsset := tempAsset				//not sure if this is needed
		m[tempAsset.Asset_identifier] = permanentAsset
	}

	return
}

// return the operator and team for a given hostname, provided they are in the map of
// ServiceApiAssets. If they are not in the map or if the values are not present,
// operator and/or team will return as an empty string ""
func LookupOperatorTeam(hostname string, m map[string]ServiceApiAsset) (operator string, team string) {

	operator = m[hostname].Operator
	team = m[hostname].Team

	return
}

// cvssFromRisk returns a synthesized CVSS score as a string given a risk label
func cvssFromRisk(risk string) string {
	switch risk {
	case "critical":
		return "10.0"
	case "high":
		return "8.0"
	case "medium":
		return "5.0"
	case "low":
		return "2.5"
	}
	return "0.0"
}

// likelihoodFromRisk returns a likelihood indicator value given a risk label
func likelihoodFromRisk(risk string) string {
	switch risk {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "maximum"
	}
	return "unknown"
}

// normalizeRisk converts known risk labels into a standardized form, if we can't identify
// the value we just return it as is
func normalizeRisk(in string) string {
	switch strings.ToLower(in) {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "critical"
	}
	return in
}
<|endoftext|>