<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"

	lz4 "github.com/cloudflare/golz4"
	"github.com/MemeLabs/overrustlelogs/common"
)

var commands = map[string]command{
	"compress":      compress,
	"uncompress":    uncompress,
	"read":          read,
	"readnicks":     readNicks,
	"nicks":         nicks,
	"migrate":       migrate,
	"namechange":    namechange,
	"cleanup":       cleanup,
	"convert":       convertToZSTD,
	"createtoplist": createTopList,
}

func main() {
	log.SetFlags(log.Lshortfile | log.LstdFlags)
	if len(os.Args) < 2 {
		os.Exit(1)
	}
	if c, ok := commands[os.Args[1]]; ok {
		if err := c(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	} else {
		fmt.Println("invalid command")
		os.Exit(1)
	}
	os.Exit(0)
}

type command func() error

func compress() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	_, err := common.CompressFile(path)
	return err
}

func uncompress() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	_, err := common.UncompressFile(path)
	return err
}

func nicks() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	var data []byte
	data, err := common.ReadCompressedFile(path)
	if os.IsNotExist(err) {
		f, err := os.Open(path)
		if err != nil {
			return err
		}
		data, err = ioutil.ReadAll(f)
		if err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	r := bufio.NewReaderSize(bytes.NewReader(data), len(data))
	nick := regexp.MustCompile("^\\[[^\\]]+\\]\\s*([a-zA-Z0-9\\_\\-]+):")
	nicks := common.NickList{}
	for {
		line, err := r.ReadSlice('\n')
		if err != nil {
			if err != io.EOF {
				return err
			}
			break
		}
		if ok := nick.Match(line); ok {
			match := nick.FindSubmatch(line)
			nicks.Add(string(match[1]))
		}
	}
	return nicks.WriteTo(regexp.MustCompile("\\.txt(\\.gz)?$").ReplaceAllString(path, ".nicks"))
}

func read() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	if regexp.MustCompile("\\.txt\\.gz$").MatchString(path) {
		buf, err := common.ReadCompressedFile(path)
		if err != nil {
			return err
		}
		os.Stdout.Write(buf)
	} else {
		return errors.New("invalid file")
	}
	return nil
}

func readNicks() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	if regexp.MustCompile("\\.nicks\\.gz$").MatchString(path) {
		nicks := common.NickList{}
		if err := common.ReadNickList(nicks, path); err != nil {
			return err
		}
		for nick := range nicks {
			fmt.Println(nick)
		}
	} else {
		return errors.New("invalid file")
	}
	return nil
}

func namechange() error {
	if len(os.Args) < 5 {
		return errors.New("not enough args")
	}
	validNick := regexp.MustCompile("^[a-zA-Z0-9_]+$")
	log := os.Args[2]
	oldName := os.Args[3]
	if !validNick.Match([]byte(oldName)) {
		return errors.New("the old name is not a valid nick")
	}
	newName := os.Args[4]

	replacer := strings.NewReplacer(
		"] "+oldName+":", "] "+newName+":",
		" "+oldName+" ", " "+newName+" ",
		" "+oldName+"\n", " "+newName+"\n",
	)

	log = strings.Replace(log, "txt", "nicks", 1)

	if strings.Contains(log, time.Now().UTC().Format("2006-01-02")) {
		return errors.New("can't modify todays log file")
	}
	fmt.Println(log)

	n := common.NickList{}
	err := common.ReadNickList(n, log)
	if err != nil {
		fmt.Println(err)
		return err
	}

	if _, ok := n[newName]; ok {
		return errors.New("nick already used, choose another one")
	}
	if _, ok := n[oldName]; !ok {
		return errors.New("nick not found")
	}
	n.Remove(oldName)
	n.Add(newName)
	err = n.WriteTo(log[:len(log)-4])
	if err != nil {
		fmt.Println(err)
		return err
	}

	log = strings.Replace(log, "nicks", "txt", 1)

	d, err := common.ReadCompressedFile(log)
	if err != nil {
		fmt.Println(err)
		return err
	}

	newData := []byte(replacer.Replace(string(d)))
	f, err := common.WriteCompressedFile(log, newData)
	if err != nil {
		fmt.Println(err)
		return err
	}
	fmt.Println("replaced nicks in", f.Name())
	f.Close()
	return nil
}

func cleanup() error {
	now := time.Now()

	logsPath := os.Args[2]

	filepaths, err := filepath.Glob(filepath.Join(logsPath, "/*/*/*"))
	if err != nil {
		log.Printf("error getting filepaths: %v", err)
		return err
	}
	log.Printf("found %d files, starting cleanup...", len(filepaths))

	r := regexp.MustCompile(`\.gz$`)

	for _, fp := range filepaths {
		if r.MatchString(fp) || strings.Contains(fp, now.Format("2006-01-02")) {
			continue
		}
		_, err := common.CompressFile(fp)
		if err != nil {
			log.Panicf("error writing compressed file: %v", err)
		}
		log.Println("compressed", fp)
	}
	return nil
}

func convertToZSTD() error {
	logsPath := os.Args[2]

	filepaths, err := filepath.Glob(filepath.Join(logsPath, "/*/*/*"))
	if err != nil {
		log.Printf("error getting filepaths: %v", err)
		return err
	}
	log.Printf("found %d files, starting cleanup...", len(filepaths))
	// now := time.Now().UTC()
	for _, fp := range filepaths {
		if strings.HasSuffix(fp, ".lz4") { //!strings.Contains(fp, now.Format("2006-01-02")) {
			data, err := UncompressFile(fp)
			if err != nil {
				log.Printf("error reading compressed file: %v", err)
			}
			data.Close()
			fp = fp[:len(fp)-4]
		}
		if strings.HasSuffix(fp, ".txt") || strings.HasSuffix(fp, ".nicks") {
			_, err = common.CompressFile(fp)
			if err != nil {
				log.Println(err)
			}
			// log.Println("compressed", fp)
			continue
		}
		// log.Println("error:", fp)
	}
	return nil
}

// ReadCompressedFile read compressed file
func ReadCompressedFile(path string) ([]byte, error) {
	f, err := os.Open(lz4Path(path))
	if err != nil {
		return nil, err
	}
	defer f.Close()
	c, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	size := uint32(0)
	size |= uint32(c[0]) << 24
	size |= uint32(c[1]) << 16
	size |= uint32(c[2]) << 8
	size |= uint32(c[3])
	data := make([]byte, size)
	if err := lz4.Uncompress(c[4:], data); err != nil {
		return nil, err
	}
	return data, nil
}

// UncompressFile uncompress an existing file
func UncompressFile(path string) (*os.File, error) {
	d, err := ReadCompressedFile(path)
	if err != nil {
		return nil, err
	}
	f, err := os.OpenFile(strings.Replace(path, ".lz4", "", -1), os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	if _, err := f.Write(d); err != nil {
		return nil, err
	}
	if err := os.Remove(lz4Path(path)); err != nil {
		return nil, err
	}
	return f, nil
}

func lz4Path(path string) string {
	if path[len(path)-4:] != ".lz4" {
		path += ".lz4"
	}
	return path
}

// ./tool createToplist /path/to/logs/ "September *"
func createTopList() error {

	month := os.Args[3]
	logsPath := os.Args[2]

	filepaths, err := filepath.Glob(filepath.Join(logsPath, "/*", month))
	if err != nil {
		log.Printf("error getting filepaths: %v", err)
		return err
	}

	for _, mpath := range filepaths {
		fmt.Printf("creating toplist for %s\n", mpath)
		toplist := make(map[string]*user)

		files, err := ioutil.ReadDir(mpath)
		if err != nil {
			fmt.Printf("error reading folder: %s with: %v", mpath, err)
			continue
		}

		for _, file := range files {
			if !strings.HasSuffix(file.Name(), ".txt.gz") {
				continue
			}
			b, err := common.ReadCompressedFile(filepath.Join(mpath, file.Name()))
			if err != nil {
				fmt.Printf("error: %v reading file %s", err, file.Name())
				continue
			}

			buf := bytes.NewBuffer(b)
			scanner := bufio.NewScanner(buf)

			for scanner.Scan() {
				line := scanner.Bytes()
				// some 2016 files are borked
				if !strings.HasPrefix(scanner.Text(), "[") {
					continue
				}

				endofdate := len("[2017-08-27 01:57:59 UTC] ")
				endofnick := bytes.Index(line[endofdate:], []byte(":"))

				nick := line[endofdate : endofnick+endofdate]
				date, err := time.Parse("[2006-01-02 15:04:05 MST] ", string(line[:endofdate]))
				if err != nil {
					fmt.Println(err)
				}

				if _, ok := toplist[string(nick)]; !ok {
					toplist[string(nick)] = &user{
						Lines:    1,
						Bytes:    len(line[endofnick+endofdate:]),
						Username: string(nick),
						Seen:     date.Unix(),
					}
					continue
				}

				toplist[string(nick)].Lines++
				toplist[string(nick)].Bytes += len(line[endofnick+endofdate:])
				if toplist[string(nick)].Seen < date.Unix() {
					toplist[string(nick)].Seen = date.Unix()
				}
			}

			if err := scanner.Err(); err != nil {
				fmt.Fprintln(os.Stderr, "reading standard input:", err)
			}
		}

		users := []*user{}
		for _, u := range toplist {
			users = append(users, u)
		}
		sort.Sort(ByLines(users))

		var buf bytes.Buffer
		err = gob.NewEncoder(&buf).Encode(users)
		if err != nil {
			fmt.Printf("error encoding users: %v", err)
			continue
		}

		_, err = common.WriteCompressedFile(filepath.Join(mpath, "toplist.json.gz"), buf.Bytes())
		if err != nil {
			fmt.Printf("error writing toplist file: %v", err)
		}
	}

	return nil
}

type user struct {
	Username string
	Lines    int
	Bytes    int
	Seen     int64
}

type ByLines []*user

func (a ByLines) Len() int           { return len(a) }
func (a ByLines) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByLines) Less(i, j int) bool { return a[i].Lines > a[j].Lines }
<commit_msg>uncompress all function<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/MemeLabs/overrustlelogs/common"
	lz4 "github.com/cloudflare/golz4"
	pb "gopkg.in/cheggaaa/pb.v1"
)

var commands = map[string]command{
	"compress":      compress,
	"uncompress":    uncompress,
	"uncompressAll": uncompressAll,
	"read":          read,
	"readnicks":     readNicks,
	"nicks":         nicks,
	"migrate":       migrate,
	"namechange":    namechange,
	"cleanup":       cleanup,
	"convert":       convertToZSTD,
	"createtoplist": createTopList,
}

func main() {
	log.SetFlags(log.Lshortfile | log.LstdFlags)
	if len(os.Args) < 2 {
		os.Exit(1)
	}
	if c, ok := commands[os.Args[1]]; ok {
		if err := c(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	} else {
		fmt.Println("invalid command")
		os.Exit(1)
	}
	os.Exit(0)
}

type command func() error

func compress() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	_, err := common.CompressFile(path)
	return err
}

func uncompress() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	_, err := common.UncompressFile(path)
	return err
}

func nicks() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	var data []byte
	data, err := common.ReadCompressedFile(path)
	if os.IsNotExist(err) {
		f, err := os.Open(path)
		if err != nil {
			return err
		}
		data, err = ioutil.ReadAll(f)
		if err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	r := bufio.NewReaderSize(bytes.NewReader(data), len(data))
	nick := regexp.MustCompile("^\\[[^\\]]+\\]\\s*([a-zA-Z0-9\\_\\-]+):")
	nicks := common.NickList{}
	for {
		line, err := r.ReadSlice('\n')
		if err != nil {
			if err != io.EOF {
				return err
			}
			break
		}
		if ok := nick.Match(line); ok {
			match := nick.FindSubmatch(line)
			nicks.Add(string(match[1]))
		}
	}
	return nicks.WriteTo(regexp.MustCompile("\\.txt(\\.gz)?$").ReplaceAllString(path, ".nicks"))
}

func read() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	if regexp.MustCompile("\\.txt\\.gz$").MatchString(path) {
		buf, err := common.ReadCompressedFile(path)
		if err != nil {
			return err
		}
		os.Stdout.Write(buf)
	} else {
		return errors.New("invalid file")
	}
	return nil
}

func readNicks() error {
	if len(os.Args) < 3 {
		return errors.New("not enough args")
	}
	path := os.Args[2]
	if regexp.MustCompile("\\.nicks\\.gz$").MatchString(path) {
		nicks := common.NickList{}
		if err := common.ReadNickList(nicks, path); err != nil {
			return err
		}
		for nick := range nicks {
			fmt.Println(nick)
		}
	} else {
		return errors.New("invalid file")
	}
	return nil
}

func namechange() error {
	if len(os.Args) < 5 {
		return errors.New("not enough args")
	}
	validNick := regexp.MustCompile("^[a-zA-Z0-9_]+$")
	log := os.Args[2]
	oldName := os.Args[3]
	if !validNick.Match([]byte(oldName)) {
		return errors.New("the old name is not a valid nick")
	}
	newName := os.Args[4]

	replacer := strings.NewReplacer(
		"] "+oldName+":", "] "+newName+":",
		" "+oldName+" ", " "+newName+" ",
		" "+oldName+"\n", " "+newName+"\n",
	)

	log = strings.Replace(log, "txt", "nicks", 1)

	if strings.Contains(log, time.Now().UTC().Format("2006-01-02")) {
		return errors.New("can't modify todays log file")
	}
	fmt.Println(log)

	n := common.NickList{}
	err := common.ReadNickList(n, log)
	if err != nil {
		fmt.Println(err)
		return err
	}

	if _, ok := n[newName]; ok {
		return errors.New("nick already used, choose another one")
	}
	if _, ok := n[oldName]; !ok {
		return errors.New("nick not found")
	}
	n.Remove(oldName)
	n.Add(newName)
	err = n.WriteTo(log[:len(log)-4])
	if err != nil {
		fmt.Println(err)
		return err
	}

	log = strings.Replace(log, "nicks", "txt", 1)

	d, err := common.ReadCompressedFile(log)
	if err != nil {
		fmt.Println(err)
		return err
	}

	newData := []byte(replacer.Replace(string(d)))
	f, err := common.WriteCompressedFile(log, newData)
	if err != nil {
		fmt.Println(err)
		return err
	}
	fmt.Println("replaced nicks in", f.Name())
	f.Close()
	return nil
}

func cleanup() error {
	now := time.Now()

	logsPath := os.Args[2]

	filepaths, err := filepath.Glob(filepath.Join(logsPath, "/*/*/*"))
	if err != nil {
		log.Printf("error getting filepaths: %v", err)
		return err
	}
	log.Printf("found %d files, starting cleanup...", len(filepaths))

	r := regexp.MustCompile(`\.gz$`)

	for _, fp := range filepaths {
		if r.MatchString(fp) || strings.Contains(fp, now.Format("2006-01-02")) {
			continue
		}
		_, err := common.CompressFile(fp)
		if err != nil {
			log.Panicf("error writing compressed file: %v", err)
		}
		log.Println("compressed", fp)
	}
	return nil
}

func convertToZSTD() error {
	logsPath := os.Args[2]

	filepaths, err := filepath.Glob(filepath.Join(logsPath, "/*/*/*"))
	if err != nil {
		log.Printf("error getting filepaths: %v", err)
		return err
	}
	log.Printf("found %d files, starting cleanup...", len(filepaths))
	// now := time.Now().UTC()
	for _, fp := range filepaths {
		if strings.HasSuffix(fp, ".lz4") { //!strings.Contains(fp, now.Format("2006-01-02")) {
			data, err := UncompressFile(fp)
			if err != nil {
				log.Printf("error reading compressed file: %v", err)
			}
			data.Close()
			fp = fp[:len(fp)-4]
		}
		if strings.HasSuffix(fp, ".txt") || strings.HasSuffix(fp, ".nicks") {
			_, err = common.CompressFile(fp)
			if err != nil {
				log.Println(err)
			}
			// log.Println("compressed", fp)
			continue
		}
		// log.Println("error:", fp)
	}
	return nil
}

// ReadCompressedFile read compressed file
func ReadCompressedFile(path string) ([]byte, error) {
	f, err := os.Open(lz4Path(path))
	if err != nil {
		return nil, err
	}
	defer f.Close()
	c, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	size := uint32(0)
	size |= uint32(c[0]) << 24
	size |= uint32(c[1]) << 16
	size |= uint32(c[2]) << 8
	size |= uint32(c[3])
	data := make([]byte, size)
	if err := lz4.Uncompress(c[4:], data); err != nil {
		return nil, err
	}
	return data, nil
}

// UncompressFile uncompress an existing file
func UncompressFile(path string) (*os.File, error) {
	d, err := ReadCompressedFile(path)
	if err != nil {
		return nil, err
	}
	f, err := os.OpenFile(strings.Replace(path, ".lz4", "", -1), os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	if _, err := f.Write(d); err != nil {
		return nil, err
	}
	if err := os.Remove(lz4Path(path)); err != nil {
		return nil, err
	}
	return f, nil
}

func lz4Path(path string) string {
	if path[len(path)-4:] != ".lz4" {
		path += ".lz4"
	}
	return path
}

// ./tool createToplist /path/to/logs/ "September *"
func createTopList() error {

	month := os.Args[3]
	logsPath := os.Args[2]

	filepaths, err := filepath.Glob(filepath.Join(logsPath, "/*", month))
	if err != nil {
		log.Printf("error getting filepaths: %v", err)
		return err
	}

	for _, mpath := range filepaths {
		fmt.Printf("creating toplist for %s\n", mpath)
		toplist := make(map[string]*user)

		files, err := ioutil.ReadDir(mpath)
		if err != nil {
			fmt.Printf("error reading folder: %s with: %v", mpath, err)
			continue
		}

		for _, file := range files {
			if !strings.HasSuffix(file.Name(), ".txt.gz") {
				continue
			}
			b, err := common.ReadCompressedFile(filepath.Join(mpath, file.Name()))
			if err != nil {
				fmt.Printf("error: %v reading file %s", err, file.Name())
				continue
			}

			buf := bytes.NewBuffer(b)
			scanner := bufio.NewScanner(buf)

			for scanner.Scan() {
				line := scanner.Bytes()
				// some 2016 files are borked
				if !strings.HasPrefix(scanner.Text(), "[") {
					continue
				}

				endofdate := len("[2017-08-27 01:57:59 UTC] ")
				endofnick := bytes.Index(line[endofdate:], []byte(":"))

				nick := line[endofdate : endofnick+endofdate]
				date, err := time.Parse("[2006-01-02 15:04:05 MST] ", string(line[:endofdate]))
				if err != nil {
					fmt.Println(err)
				}

				if _, ok := toplist[string(nick)]; !ok {
					toplist[string(nick)] = &user{
						Lines:    1,
						Bytes:    len(line[endofnick+endofdate:]),
						Username: string(nick),
						Seen:     date.Unix(),
					}
					continue
				}

				toplist[string(nick)].Lines++
				toplist[string(nick)].Bytes += len(line[endofnick+endofdate:])
				if toplist[string(nick)].Seen < date.Unix() {
					toplist[string(nick)].Seen = date.Unix()
				}
			}

			if err := scanner.Err(); err != nil {
				fmt.Fprintln(os.Stderr, "reading standard input:", err)
			}
		}

		users := []*user{}
		for _, u := range toplist {
			users = append(users, u)
		}
		sort.Sort(ByLines(users))

		var buf bytes.Buffer
		err = gob.NewEncoder(&buf).Encode(users)
		if err != nil {
			fmt.Printf("error encoding users: %v", err)
			continue
		}

		_, err = common.WriteCompressedFile(filepath.Join(mpath, "toplist.json.gz"), buf.Bytes())
		if err != nil {
			fmt.Printf("error writing toplist file: %v", err)
		}
	}

	return nil
}

type user struct {
	Username string
	Lines    int
	Bytes    int
	Seen     int64
}

type ByLines []*user

func (a ByLines) Len() int           { return len(a) }
func (a ByLines) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByLines) Less(i, j int) bool { return a[i].Lines > a[j].Lines }

func uncompressAll() error {
	logsPath := os.Args[2]
	if logsPath == "" {
		return fmt.Errorf("didn't provide a path to the logs")
	}

	MAXWORKERS := runtime.NumCPU()

	files, err := filepath.Glob(filepath.Join(logsPath, "/*/*/*.gz"))
	if err != nil || len(files) < 1 {
		log.Println("couldn't find any compressed log files in", logsPath)
		return err
	}

	bar := pb.StartNew(len(files))

	wg := sync.WaitGroup{}
	wg.Add(MAXWORKERS)

	queue := make(chan string, len(files))
	for i := 0; i < MAXWORKERS; i++ {
		go func(queue <-chan string, i int, cb *pb.ProgressBar) {
			for file := range queue {
				f, err := common.UncompressFile(file)
				if err != nil {
					log.Println(err, file)
					continue
				}
				f.Close()
				cb.Increment()
			}
			wg.Done()
		}(queue, i, bar)
	}

	for _, file := range files {
		queue <- file
	}

	close(queue)
	wg.Wait()
	bar.Finish()
	return nil
}
<|endoftext|>


<commit_before>// Package magick is a hacked up port of the minimal functionality we need
// to satisfy the img.Decoder interface.  Code is based in part on
// github.com/quirkey/magick
package main

/*
#cgo pkg-config: MagickCore
#include <magick/MagickCore.h>
*/
import "C"
import (
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"rais/src/img"
	"rais/src/plugins"
	"strings"
	"unsafe"

	"github.com/uoregon-libraries/gopkg/logger"
)

var l *logger.Logger

// SetLogger is called by the RAIS server's plugin manager to let plugins use
// the central logger
func SetLogger(raisLogger *logger.Logger) {
	l = raisLogger
}

// Initialize sets up the MagickCore stuff and registers the TIFF, PNG, JPG,
// and GIF decoders
func Initialize() {
	path, _ := os.Getwd()
	cPath := C.CString(path)
	defer C.free(unsafe.Pointer(cPath))
	C.MagickCoreGenesis(cPath, C.MagickFalse)
	img.RegisterDecodeHandler(decodeCommonFile)
}

func makeError(exception *C.ExceptionInfo) error {
	return fmt.Errorf("%v: %v - %v", exception.severity, exception.reason, exception.description)
}

var validExtensions = []string{".tif", ".tiff", ".png", ".jpg", ".jpeg", ".gif"}

func validExt(u *url.URL) bool {
	var ext = strings.ToLower(filepath.Ext(u.Path))
	for _, validExt := range validExtensions {
		if ext == validExt {
			return true
		}
	}

	return false
}

func validScheme(u *url.URL) bool {
	return u.Scheme == "file"
}

func decodeCommonFile(s img.Streamer) (img.DecodeFunc, error) {
	var u = s.Location()
	if !validExt(u) {
		l.Infof("plugins/imagick-decoder: skipping unsupported image extension %q (must be one of %s)",
			s.Location(), strings.Join(validExtensions, ", "))
		return nil, plugins.ErrSkipped
	}

	// This is sorta of overly "loud" (warning), but generally speaking, a
	// decoder shouldn't be requiring local files, so we want people to be made
	// aware this plugin's not great....
	if !validScheme(u) {
		l.Warnf("plugins/imagick-decoder: skipping unsupported URL scheme %q (must be file)", u.Scheme)
		return nil, plugins.ErrSkipped
	}

	return func() (img.Decoder, error) { return NewImage(u.Path) }, nil
}
<commit_msg>plugins/imagick-decoder: SILENCE YOUR LOGS!<commit_after>// Package magick is a hacked up port of the minimal functionality we need
// to satisfy the img.Decoder interface.  Code is based in part on
// github.com/quirkey/magick
package main

/*
#cgo pkg-config: MagickCore
#include <magick/MagickCore.h>
*/
import "C"
import (
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"rais/src/img"
	"rais/src/plugins"
	"strings"
	"unsafe"

	"github.com/uoregon-libraries/gopkg/logger"
)

var l *logger.Logger

// SetLogger is called by the RAIS server's plugin manager to let plugins use
// the central logger
func SetLogger(raisLogger *logger.Logger) {
	l = raisLogger
}

// Initialize sets up the MagickCore stuff and registers the TIFF, PNG, JPG,
// and GIF decoders
func Initialize() {
	path, _ := os.Getwd()
	cPath := C.CString(path)
	defer C.free(unsafe.Pointer(cPath))
	C.MagickCoreGenesis(cPath, C.MagickFalse)
	img.RegisterDecodeHandler(decodeCommonFile)
}

func makeError(exception *C.ExceptionInfo) error {
	return fmt.Errorf("%v: %v - %v", exception.severity, exception.reason, exception.description)
}

var validExtensions = []string{".tif", ".tiff", ".png", ".jpg", ".jpeg", ".gif"}

func validExt(u *url.URL) bool {
	var ext = strings.ToLower(filepath.Ext(u.Path))
	for _, validExt := range validExtensions {
		if ext == validExt {
			return true
		}
	}

	return false
}

func validScheme(u *url.URL) bool {
	return u.Scheme == "file"
}

func decodeCommonFile(s img.Streamer) (img.DecodeFunc, error) {
	var u = s.Location()
	if !validExt(u) {
		l.Debugf("plugins/imagick-decoder: skipping unsupported image extension %q (must be one of %s)",
			s.Location(), strings.Join(validExtensions, ", "))
		return nil, plugins.ErrSkipped
	}

	if !validScheme(u) {
		l.Debugf("plugins/imagick-decoder: skipping unsupported URL scheme %q (must be file)", u.Scheme)
		return nil, plugins.ErrSkipped
	}

	return func() (img.Decoder, error) { return NewImage(u.Path) }, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Google Inc. All Rights Reserved.
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 remove handles the removal of packages.
package remove

import (
	"context"
	"fmt"
	"os"
	"sort"

	"github.com/google/googet/client"
	"github.com/google/googet/download"
	"github.com/google/googet/goolib"
	"github.com/google/googet/oswrap"
	"github.com/google/googet/system"
	"github.com/google/logger"
)

func uninstallPkg(ctx context.Context, pi goolib.PackageInfo, state *client.GooGetState, dbOnly bool, proxyServer string) error {
	logger.Infof("Executing removal of package %q", pi.Name)
	ps, err := state.GetPackageState(pi)
	if err != nil {
		return fmt.Errorf("package not found in state file: %v", err)
	}
	// Fix for package install by older versions of GooGet.
	if ps.LocalPath == "" && ps.UnpackDir != "" {
		ps.LocalPath = ps.UnpackDir + ".goo"
	}
	if ps.LocalPath == "" {
		return fmt.Errorf("no local path available for package %q", pi.Name)
	}

	if !dbOnly {
		f, err := os.Open(ps.LocalPath)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		var rd bool
		if os.IsNotExist(err) {
			logger.Infof("Local package does not exist for %s.%s.%s, redownloading...", pi.Name, pi.Arch, pi.Ver)
			rd = true
		}
		// Force redownload if checksum does not match.
		// If checksum is empty this was a local install so ignore.
		if !rd && ps.Checksum != "" && goolib.Checksum(f) != ps.Checksum {
			logger.Info("Local package checksum does not match, redownloading...")
			rd = true
		}
		f.Close()

		if rd {
			if ps.DownloadURL == "" {
				return fmt.Errorf("can not redownload %s.%s.%s, DownloadURL not saved", pi.Name, pi.Arch, pi.Ver)
			}
			if err := download.Package(ctx, ps.DownloadURL, ps.LocalPath, ps.Checksum, proxyServer); err != nil {
				return fmt.Errorf("error redownloading %s.%s.%s, package may no longer exist in the repo, you can use the '-db_only' flag to remove it form the database: %v", pi.Name, pi.Arch, pi.Ver, err)
			}
		}

		eDir, err := download.ExtractPkg(ps.LocalPath)
		if err != nil {
			return err
		}

		if err := system.Uninstall(eDir, ps.PackageSpec); err != nil {
			return err
		}

		if err := oswrap.RemoveAll(eDir); err != nil {
			logger.Error(err)
		}

		if len(ps.InstalledFiles) > 0 {
			var dirs []string
			for file, chksum := range ps.InstalledFiles {
				if chksum == "" {
					dirs = append(dirs, file)
					continue
				}
				logger.Infof("Removing %q", file)
				if _, err := client.RemoveOrRename(file); err != nil {
					logger.Error(err)
				}
			}
			sort.Sort(sort.Reverse(sort.StringSlice(dirs)))
			for _, dir := range dirs {
				logger.Infof("Removing %q", dir)
				if _, err := client.RemoveOrRename(dir); err != nil {
					logger.Info(err)
				}
			}
		}
	}

	if err := oswrap.RemoveAll(ps.LocalPath); err != nil {
		logger.Errorf("error removing package data from cache directory: %v", err)
	}
	return state.Remove(pi)
}

// DepMap is a map of packages to dependant packages.
type DepMap map[string][]string

func (deps DepMap) remove(name string) {
	for dep, s := range deps {
		for i, d := range s {
			if d == name {
				s[i] = s[len(s)-1]
				s = s[:len(s)-1]
				deps[dep] = s
				break
			}
		}
	}
	delete(deps, name)
}

func (deps DepMap) build(name, arch string, state client.GooGetState) {
	logger.Infof("Building dependency map for %q", name)
	deps[name+"."+arch] = nil
	for _, p := range state {
		if p.PackageSpec.Name == name && p.PackageSpec.Arch == arch {
			continue
		}
		for d := range p.PackageSpec.PkgDependencies {
			di := goolib.PkgNameSplit(d)
			if di.Name == name && (di.Arch == arch || di.Arch == "") {
				n, a := p.PackageSpec.Name, p.PackageSpec.Arch
				deps[name+"."+arch] = append(deps[name+"."+arch], n+"."+a)
				deps.build(n, a, state)
			}
		}
	}
}

// EnumerateDeps returns a DepMap and list of dependencies for a package.
func EnumerateDeps(pi goolib.PackageInfo, state client.GooGetState) (DepMap, []string) {
	dm := make(DepMap)
	dm.build(pi.Name, pi.Arch, state)
	var dl []string
	for k := range dm {
		di := goolib.PkgNameSplit(k)
		ps, err := state.GetPackageState(di)
		if err != nil {
			logger.Fatalf("error finding package in state file, even though the dependancy map was just built: %v", err)
		}
		dl = append(dl, k+" "+ps.PackageSpec.Version)
	}
	return dm, dl
}

// All removes a package and all dependant packages. Packages with no dependant packages
// will be removed first.
func All(ctx context.Context, pi goolib.PackageInfo, deps DepMap, state *client.GooGetState, dbOnly bool, proxyServer string) error {
	for len(deps) > 1 {
		for dep := range deps {
			if len(deps[dep]) == 0 {
				di := goolib.PkgNameSplit(dep)
				if err := uninstallPkg(ctx, di, state, dbOnly, proxyServer); err != nil {
					return err
				}
				deps.remove(dep)
			}
		}
	}
	return uninstallPkg(ctx, pi, state, dbOnly, proxyServer)
}
<commit_msg>Update remove.go (#56)<commit_after>/*
Copyright 2016 Google Inc. All Rights Reserved.
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 remove handles the removal of packages.
package remove

import (
	"context"
	"fmt"
	"os"
	"sort"

	"github.com/google/googet/client"
	"github.com/google/googet/download"
	"github.com/google/googet/goolib"
	"github.com/google/googet/oswrap"
	"github.com/google/googet/system"
	"github.com/google/logger"
)

func uninstallPkg(ctx context.Context, pi goolib.PackageInfo, state *client.GooGetState, dbOnly bool, proxyServer string) error {
	logger.Infof("Executing removal of package %q", pi.Name)
	ps, err := state.GetPackageState(pi)
	if err != nil {
		return fmt.Errorf("package not found in state file: %v", err)
	}

	if !dbOnly {
		// Fix for package install by older versions of GooGet.
		if ps.LocalPath == "" && ps.UnpackDir != "" {
			ps.LocalPath = ps.UnpackDir + ".goo"
		}
		if ps.LocalPath == "" {
			return fmt.Errorf("no local path available for package %q", pi.Name)
		}

		f, err := os.Open(ps.LocalPath)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		var rd bool
		if os.IsNotExist(err) {
			logger.Infof("Local package does not exist for %s.%s.%s, redownloading...", pi.Name, pi.Arch, pi.Ver)
			rd = true
		}
		// Force redownload if checksum does not match.
		// If checksum is empty this was a local install so ignore.
		if !rd && ps.Checksum != "" && goolib.Checksum(f) != ps.Checksum {
			logger.Info("Local package checksum does not match, redownloading...")
			rd = true
		}
		f.Close()

		if rd {
			if ps.DownloadURL == "" {
				return fmt.Errorf("can not redownload %s.%s.%s, DownloadURL not saved", pi.Name, pi.Arch, pi.Ver)
			}
			if err := download.Package(ctx, ps.DownloadURL, ps.LocalPath, ps.Checksum, proxyServer); err != nil {
				return fmt.Errorf("error redownloading %s.%s.%s, package may no longer exist in the repo, you can use the '-db_only' flag to remove it form the database: %v", pi.Name, pi.Arch, pi.Ver, err)
			}
		}

		eDir, err := download.ExtractPkg(ps.LocalPath)
		if err != nil {
			return err
		}

		if err := system.Uninstall(eDir, ps.PackageSpec); err != nil {
			return err
		}

		if err := oswrap.RemoveAll(eDir); err != nil {
			logger.Error(err)
		}

		if len(ps.InstalledFiles) > 0 {
			var dirs []string
			for file, chksum := range ps.InstalledFiles {
				if chksum == "" {
					dirs = append(dirs, file)
					continue
				}
				logger.Infof("Removing %q", file)
				if _, err := client.RemoveOrRename(file); err != nil {
					logger.Error(err)
				}
			}
			sort.Sort(sort.Reverse(sort.StringSlice(dirs)))
			for _, dir := range dirs {
				logger.Infof("Removing %q", dir)
				if _, err := client.RemoveOrRename(dir); err != nil {
					logger.Info(err)
				}
			}
		}
		if err := oswrap.RemoveAll(ps.LocalPath); err != nil {
			logger.Errorf("error removing package data from cache directory: %v", err)
		}
	}
	return state.Remove(pi)
}

// DepMap is a map of packages to dependant packages.
type DepMap map[string][]string

func (deps DepMap) remove(name string) {
	for dep, s := range deps {
		for i, d := range s {
			if d == name {
				s[i] = s[len(s)-1]
				s = s[:len(s)-1]
				deps[dep] = s
				break
			}
		}
	}
	delete(deps, name)
}

func (deps DepMap) build(name, arch string, state client.GooGetState) {
	logger.Infof("Building dependency map for %q", name)
	deps[name+"."+arch] = nil
	for _, p := range state {
		if p.PackageSpec.Name == name && p.PackageSpec.Arch == arch {
			continue
		}
		for d := range p.PackageSpec.PkgDependencies {
			di := goolib.PkgNameSplit(d)
			if di.Name == name && (di.Arch == arch || di.Arch == "") {
				n, a := p.PackageSpec.Name, p.PackageSpec.Arch
				deps[name+"."+arch] = append(deps[name+"."+arch], n+"."+a)
				deps.build(n, a, state)
			}
		}
	}
}

// EnumerateDeps returns a DepMap and list of dependencies for a package.
func EnumerateDeps(pi goolib.PackageInfo, state client.GooGetState) (DepMap, []string) {
	dm := make(DepMap)
	dm.build(pi.Name, pi.Arch, state)
	var dl []string
	for k := range dm {
		di := goolib.PkgNameSplit(k)
		ps, err := state.GetPackageState(di)
		if err != nil {
			logger.Fatalf("error finding package in state file, even though the dependancy map was just built: %v", err)
		}
		dl = append(dl, k+" "+ps.PackageSpec.Version)
	}
	return dm, dl
}

// All removes a package and all dependant packages. Packages with no dependant packages
// will be removed first.
func All(ctx context.Context, pi goolib.PackageInfo, deps DepMap, state *client.GooGetState, dbOnly bool, proxyServer string) error {
	for len(deps) > 1 {
		for dep := range deps {
			if len(deps[dep]) == 0 {
				di := goolib.PkgNameSplit(dep)
				if err := uninstallPkg(ctx, di, state, dbOnly, proxyServer); err != nil {
					return err
				}
				deps.remove(dep)
			}
		}
	}
	return uninstallPkg(ctx, pi, state, dbOnly, proxyServer)
}
<|endoftext|>


<commit_before>package tunnel

import (
	"fmt"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
	"github.com/ondevice/ondevice-cli/rest"
)

// WSListener -- WebSocket listener
type WSListener interface {
	OnMessage(msgType int, data []byte)
}

// Connection -- WebSocket connection
type Connection struct {
	ws        *websocket.Conn
	onMessage func(int, []byte)
	done      chan struct{}
}

// open -- Open a websocket connection
func open(c *Connection, endpoint string, params map[string]string, onMessage func(int, []byte), auths ...rest.Authentication) error {
	hdr := http.Header{}

	var auth rest.Authentication
	if len(auths) == 0 {
		var err error
		if auth, err = rest.CreateClientAuth(); err != nil {
			return err
		}
	}

	hdr.Add("Authorization", auth.GetAuthHeader())

	ws, resp, err := websocket.DefaultDialer.Dial(auth.GetURL(endpoint+"/websocket", params, "wss"), hdr)
	if resp.StatusCode == 401 {
		return fmt.Errorf("API server authentication failed")
	} else if err != nil {
		log.Fatalf("error opening websocket (response code: %s): %s", resp.Status, err)
	}

	c.ws = ws
	c.onMessage = onMessage
	c.done = make(chan struct{})

	go c.receive()

	return nil
}

// Close -- Close the underlying WebSocket connection
func (c Connection) Close() {
	c.ws.Close()
}

// OnMessage -- pass incoming WebSocket messages on to the listener function
func (c Connection) OnMessage(msgType int, msg []byte) {
	c.onMessage(msgType, msg)
}

func (c Connection) receive() {
	defer c.ws.Close()
	defer close(c.done)

	for {
		msgType, msg, err := c.ws.ReadMessage()
		if err != nil {
			log.Println("read error:", err)
			return
		}
		c.onMessage(msgType, msg)
	}
}

// SendBinary -- Send binary WebSocket message
func (c Connection) SendBinary(data []byte) {
	c.ws.WriteMessage(websocket.BinaryMessage, data)
}

// SendJSON -- Send a JSON text message to the WebSocket
func (c Connection) SendJSON(value interface{}) {
	c.ws.WriteJSON(value)
}

// SendText -- send a raw text websocket messge (use SendJson instead where possible)
func (c Connection) SendText(msg string) {
	c.ws.WriteMessage(websocket.TextMessage, []byte(msg))
}

// Wait -- Wait for the connection to close
func (c Connection) Wait() {
	<-c.done
}
<commit_msg>tunnel.open(): fixed issue with missing auth propagation<commit_after>package tunnel

import (
	"fmt"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
	"github.com/ondevice/ondevice-cli/rest"
)

// WSListener -- WebSocket listener
type WSListener interface {
	OnMessage(msgType int, data []byte)
}

// Connection -- WebSocket connection
type Connection struct {
	ws        *websocket.Conn
	onMessage func(int, []byte)
	done      chan struct{}
}

// open -- Open a websocket connection
func open(c *Connection, endpoint string, params map[string]string, onMessage func(int, []byte), auths ...rest.Authentication) error {
	hdr := http.Header{}

	var auth rest.Authentication
	if len(auths) == 0 {
		var err error
		if auth, err = rest.CreateClientAuth(); err != nil {
			return err
		}
	} else {
		auth = auths[0]
	}

	hdr.Add("Authorization", auth.GetAuthHeader())

	url := auth.GetURL(endpoint+"/websocket", params, "wss")
	//log.Printf("Opening websocket connection to '%s' (auth: '%s')", url, auth.GetAuthHeader())

	ws, resp, err := websocket.DefaultDialer.Dial(url, hdr)
	if resp.StatusCode == 401 {
		return fmt.Errorf("API server authentication failed")
	} else if err != nil {
		log.Fatalf("error opening websocket (response code: %s): %s", resp.Status, err)
	}

	c.ws = ws
	c.onMessage = onMessage
	c.done = make(chan struct{})

	go c.receive()

	return nil
}

// Close -- Close the underlying WebSocket connection
func (c Connection) Close() {
	c.ws.Close()
}

// OnMessage -- pass incoming WebSocket messages on to the listener function
func (c Connection) OnMessage(msgType int, msg []byte) {
	c.onMessage(msgType, msg)
}

func (c Connection) receive() {
	defer c.ws.Close()
	defer close(c.done)

	for {
		msgType, msg, err := c.ws.ReadMessage()
		if err != nil {
			log.Println("read error:", err)
			return
		}
		c.onMessage(msgType, msg)
	}
}

// SendBinary -- Send binary WebSocket message
func (c Connection) SendBinary(data []byte) {
	c.ws.WriteMessage(websocket.BinaryMessage, data)
}

// SendJSON -- Send a JSON text message to the WebSocket
func (c Connection) SendJSON(value interface{}) {
	c.ws.WriteJSON(value)
}

// SendText -- send a raw text websocket messge (use SendJson instead where possible)
func (c Connection) SendText(msg string) {
	c.ws.WriteMessage(websocket.TextMessage, []byte(msg))
}

// Wait -- Wait for the connection to close
func (c Connection) Wait() {
	<-c.done
}
<|endoftext|>


<commit_before>package repository

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

	"github.com/naoty/todo/todo"
)

// FS represents a repository backed by file system.
type FS struct {
	root string
}

// NewFS returns a new FS.
func NewFS(root string) *FS {
	return &FS{root: root}
}

// List implements Repository interface.
func (repo *FS) List() ([]*todo.Todo, error) {
	todos := []*todo.Todo{}

	err := filepath.Walk(repo.root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		id, err := parseID(path)
		if err != nil {
			return err
		}

		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()

		content, err := ioutil.ReadAll(file)
		if err != nil {
			return err
		}

		td, err := todo.Parse(string(content))
		if err != nil {
			return err
		}

		td.SetID(id)
		todos = append(todos, td)

		return nil
	})

	if err != nil {
		return todos, fmt.Errorf("failed to get todos from %s: %w", repo.root, err)
	}

	return todos, nil
}

// Add implements Repository interface.
func (repo *FS) Add(title string) error {
	todos, err := repo.List()
	if err != nil {
		return fmt.Errorf("failed to get next id: %w", err)
	}

	lastID := 0
	for _, td := range todos {
		id := td.ID()
		if id > lastID {
			lastID = id
		}
	}

	nextID := lastID + 1
	filename := fmt.Sprintf("%d.md", nextID)
	path := filepath.Join(repo.root, filename)

	if _, err := os.Stat(path); err == nil {
		return fmt.Errorf("already exist: %s", path)
	}

	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return fmt.Errorf("failed to add TODO: %w", err)
	}
	defer file.Close()

	content := newContent(title)
	_, err = file.WriteString(content)
	if err != nil {
		return fmt.Errorf("failed to add TODO: %w", err)
	}

	return nil
}

// Open implements Repository interface.
func (repo *FS) Open(id int) error {
	filename := fmt.Sprintf("%d.md", id)
	path := filepath.Join(repo.root, filename)

	cmd := exec.Command("open", path)
	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("failed to open %s: %w", path, err)
	}

	return nil
}

func parseID(path string) (int, error) {
	text := strings.TrimRight(filepath.Base(path), filepath.Ext(path))
	return strconv.Atoi(text)
}

func newContent(title string) string {
	return strings.TrimLeft(fmt.Sprintf(`
---
title: %s
---

`, title), "\n")
}
<commit_msg>Trim line breaks in TODO template<commit_after>package repository

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

	"github.com/naoty/todo/todo"
)

// FS represents a repository backed by file system.
type FS struct {
	root string
}

// NewFS returns a new FS.
func NewFS(root string) *FS {
	return &FS{root: root}
}

// List implements Repository interface.
func (repo *FS) List() ([]*todo.Todo, error) {
	todos := []*todo.Todo{}

	err := filepath.Walk(repo.root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		id, err := parseID(path)
		if err != nil {
			return err
		}

		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()

		content, err := ioutil.ReadAll(file)
		if err != nil {
			return err
		}

		td, err := todo.Parse(string(content))
		if err != nil {
			return err
		}

		td.SetID(id)
		todos = append(todos, td)

		return nil
	})

	if err != nil {
		return todos, fmt.Errorf("failed to get todos from %s: %w", repo.root, err)
	}

	return todos, nil
}

// Add implements Repository interface.
func (repo *FS) Add(title string) error {
	todos, err := repo.List()
	if err != nil {
		return fmt.Errorf("failed to get next id: %w", err)
	}

	lastID := 0
	for _, td := range todos {
		id := td.ID()
		if id > lastID {
			lastID = id
		}
	}

	nextID := lastID + 1
	filename := fmt.Sprintf("%d.md", nextID)
	path := filepath.Join(repo.root, filename)

	if _, err := os.Stat(path); err == nil {
		return fmt.Errorf("already exist: %s", path)
	}

	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return fmt.Errorf("failed to add TODO: %w", err)
	}
	defer file.Close()

	content := newContent(title)
	_, err = file.WriteString(content)
	if err != nil {
		return fmt.Errorf("failed to add TODO: %w", err)
	}

	return nil
}

// Open implements Repository interface.
func (repo *FS) Open(id int) error {
	filename := fmt.Sprintf("%d.md", id)
	path := filepath.Join(repo.root, filename)

	cmd := exec.Command("open", path)
	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("failed to open %s: %w", path, err)
	}

	return nil
}

func parseID(path string) (int, error) {
	text := strings.TrimRight(filepath.Base(path), filepath.Ext(path))
	return strconv.Atoi(text)
}

func newContent(title string) string {
	return strings.Trim(fmt.Sprintf(`
---
title: %s
---
`, title), "\n")
}
<|endoftext|>


<commit_before>package requests

import (
	//"fmt"
	"net/http"
	"testing"
	"bytes"
	//"errors"
	"io"
	"net/url"
	"io/ioutil"
	"os"
	"reflect"
	"time"
	//"runtime"
)

// Stub `requests` namespace
type stubRequests struct {
	Get      func(string, string, map[string]string) (*stubResponse, error)
	GetAsync func(string, string, map[string]string, int) (chan interface{}, error)
	//Post     func(string, string, map[string]interface{}) (*http.Response, error)
}

// Stub connection
type stubConnection struct {
	channel chan interface{}
	done chan bool
	latency time.Duration
}

func newStubConnection(latency time.Duration) *stubConnection {
	channel := make(chan interface{}, 1)
	done := make(chan bool, 1)
	connection := &stubConnection{
		channel: channel,
		done: done,
		latency: latency,
	}
	return connection
}

func (c *stubConnection) Close() {
	close(c.channel)
	close(c.done)
}

// Stub `http.Request`
type stubRequest struct { http.Request }

func newStubRequest(method, rawurl string, body io.ReadCloser) (*stubRequest, error) {
	uri, err := url.ParseRequestURI(rawurl)
	if err != nil {
		panic("Something's wrong with your URI")
	}
	request := &stubRequest{
		Request: http.Request{
			Method: method,
			URL: uri,
			Body: body,
		},
	}
	return request, nil
}

// Stub `http.Response`
type stubResponse struct { http.Response }

func newStubResponse(status string, code int, header http.Header, body io.ReadCloser) *stubResponse {
	response := &stubResponse{
		Response: http.Response{
			Status: status,
			StatusCode: code,
			Proto: "HTTP/1.0",
			Header: header,
			Body: body,
		},
	}
	return response
}

// Stub `http.Server`
type stubServer struct {
	http.Server
	response *stubResponse
	latency  time.Duration
}

func newStubServer(addr string, res *stubResponse, lat time.Duration) *stubServer {
	server := &stubServer{
		Server: http.Server{ Addr: addr },
		response: res,
		latency: lat,
	}
	return server
}

func (s *stubServer) Reply(code statusCode) *stubResponse {
	// Block for server's latency
	//<-time.Tick(s.latency)
	time.Sleep(s.latency)
	// Create status code and return the response
	s.response.StatusCode = (int)(code)
	// TODO: Assign Status
	return s.response
}

// Stub the `http.Client`
type stubClient struct { http.Client }

func newStubClient(timeout time.Duration) *stubClient {
	client := &stubClient{
		Client: http.Client{ Timeout: timeout },
	}
	return client
}

// Inject `*stubConnection` and `*stubServer` to simulate a server call
func (c *stubClient) Do(req *stubRequest, conn *stubConnection, server *stubServer) (*stubResponse, error) {
	// Block for the duration of `conn.latency` + `server.latency`
	// to simulaate real-world latencies and test timeoutn
	code := (statusCode)(server.response.StatusCode)
	conn.channel <- req
	go func() {
		//request := <-conn
		//<-time.Tick(conn.latency)
		time.Sleep(conn.latency)
		<-conn.channel
		// TODO: Do something with the receive request
		conn.channel <- server.Reply(code)
		//<-time.Tick(server.latency)
		time.Sleep(server.latency)
		conn.done <- true
	}()

	<-conn.done
	res := <-conn.channel
	return res.(*stubResponse), nil
}

func (c *stubClient) DoAsync(req *stubRequest, conn *stubConnection, server *stubServer) (chan interface{}, error) {
	code := (statusCode)(server.response.StatusCode)

	// Network latency
	time.Sleep(conn.latency)
	conn.channel <- req

	go func(conn *stubConnection) {
		<-conn.channel
		// TODO: Do something with the receive request
		conn.channel <- server.Reply(code)
		//<-time.Tick(server.latency)
	}(conn)
	return conn.channel, nil
}

var (
	// Setup connection
	networkLatency = time.Duration(100) * time.Millisecond
	conn = newStubConnection(networkLatency)

	// Setup client
	timeout = time.Duration(3) * time.Second
	client = newStubClient(timeout)

	// Setup server
	res = newStubResponse("200 OK", 200, header, body)
	endpoint = "http://jochasinga.io"
	serverLatency = time.Duration(100) * time.Millisecond
	server = newStubServer(endpoint, res, serverLatency)

	// Setup request
	header = http.Header{}
	jsonStr = `{"foo": ["bar", "baz"]}`
	body = ioutil.NopCloser(bytes.NewBuffer([]byte(jsonStr)))
	auth = map[string]string{"user": "pass"}
	requests = &stubRequests{
		Get: func(url, body string, auth map[string]string) (*stubResponse, error) {
			// Convert body from string to io.ReadCloser
			bodyReadCloser := ioutil.NopCloser(bytes.NewBuffer([]byte(body)))
			//req, err := http.NewRequest("GET", url, bodyReadCloser)
			req, err := newStubRequest("GET", url, bodyReadCloser)
			if err != nil {
				panic(err)
			}
			// TODO: include basic auth
			/*
			if len(auth) > 0 {
				for user, password := range auth {
					req.SetBasicAuth(user, password)
				}
			}
                        */
			res, err := client.Do(req, conn, server)
			if err != nil {
				panic(err)
			}
			return res, nil
		},
		GetAsync: func(url, body string, auth map[string]string, timeout int) (chan interface{}, error) {
			waitUntil := time.Duration(timeout) * time.Second
			data := ioutil.NopCloser(bytes.NewBuffer([]byte(body)))
			client.Timeout = waitUntil

			req, err := newStubRequest("GET", url, data)
			if err != nil {
				panic(err)
			}
			// TODO: include basic auth
			/*
			if len(auth) > 0 {
				for user, password := range auth {
					req.SetBasicAuth(user, password)
				}
			}
                        */
			channel, err := client.DoAsync(req, conn, server)
			if err != nil {
				panic(err)
			}

			return channel, nil

			/*
			p := NewPromise(func() (*http.Response, error) {
				re := make(chan *http.Response)
				er := make(chan error)
				go func() {
					res, err := client.Do(req, server)
					if err != nil {
						er <- err
					}
					re <- res
					return
				}()
				defer close(re)
				defer close(er)
				return (*http.Response)(nil), errors.New("Time out")
			})
			return p, nil
                        */
		},
	}
)

// Get response back as `*stubResponse`
func TestGetResponseType(t *testing.T) {
	resp, err := requests.Get(endpoint, jsonStr, auth)
	if err != nil {
		t.Error(err)
	}
	returnType := reflect.TypeOf(resp)
	responseType := reflect.TypeOf((*stubResponse)(nil))
	if returnType != responseType {
		t.Errorf("Expected return type of `*stubResponse`, but it was %v instead.", returnType)
	}
}

// Get response back with status 200
func TestGetResponseStatus(t *testing.T) {
	resp, err := requests.Get(endpoint, jsonStr, auth)
	if err != nil {
		t.Error(err)
	}
	if resp.StatusCode != 200 {
		t.Errorf("Expected StatusCode `200`, but it was %s instead.", resp.Status)
	}
}

func TestGetAsyncResponseType(t *testing.T) {
	// 3 seconds timeout
	timeout := 1
	resultChan, err := requests.GetAsync(endpoint, jsonStr, auth, timeout)
	if err != nil {
		t.Error(err)
	}
	returnType := reflect.TypeOf(resultChan)
	responseType := reflect.TypeOf(chan interface{}(nil))
	if returnType != responseType {
		t.Errorf("Expected return type of `chan interface{}`, but it was %v instead.", returnType)
	}
}

/*
func TestGetAsyncResponseStatus(t *testing.T) {
	timeout := 0
	p, err := requests.GetAsync("http://example.com", htmlStr, auth, timeout)
	if err != nil {
		t.Error(err)
	}

	result := p.Then(func() *http.Response { return p.res }, func() error { return p.err })

	// Result should be either `*http.Response` or `error`
	switch result := result.(type) {
	default:
		t.Error()
	case *http.Response:
    		if result.Status != "200 OK" {
			t.Errorf("Expected Status `200 OK`, but it was %s instead.", result.Status)
		}
		break
	case error:
		if result.Error() == "" {
			t.Errorf("Expected `error.`")
		}
		break
	}
}
*/

func TestMain(m *testing.M) {
	v := m.Run()
	defer conn.Close()
	if v == 0 {
		os.Exit(1)
	}
	os.Exit(v)
}
<commit_msg>TestGetAsyncResponseStatus PASS<commit_after>package requests

import (
	//"fmt"
	"net/http"
	"testing"
	"bytes"
	//"errors"
	"io"
	"net/url"
	"io/ioutil"
	"os"
	"reflect"
	"time"
	//"runtime"
)

// Stub `requests` namespace
type stubRequests struct {
	Get      func(string, string, map[string]string) (*stubResponse, error)
	GetAsync func(string, string, map[string]string, int) (chan *stubResponse, error)
	//Post     func(string, string, map[string]interface{}) (*http.Response, error)
}

// Stub connection
type stubConnection struct {
	channel chan interface{}
	done chan bool
	latency time.Duration
}

func newStubConnection(latency time.Duration) *stubConnection {
	channel := make(chan interface{}, 1)
	done := make(chan bool, 1)
	connection := &stubConnection{
		channel: channel,
		done: done,
		latency: latency,
	}
	return connection
}

func (c *stubConnection) Close() {
	close(c.channel)
	close(c.done)
}

// Stub `http.Request`
type stubRequest struct { http.Request }

func newStubRequest(method, rawurl string, body io.ReadCloser) (*stubRequest, error) {
	uri, err := url.ParseRequestURI(rawurl)
	if err != nil {
		panic("Something's wrong with your URI")
	}
	request := &stubRequest{
		Request: http.Request{
			Method: method,
			URL: uri,
			Body: body,
		},
	}
	return request, nil
}

// Stub `http.Response`
type stubResponse struct { http.Response }

func newStubResponse(status string, code int, header http.Header, body io.ReadCloser) *stubResponse {
	response := &stubResponse{
		Response: http.Response{
			Status: status,
			StatusCode: code,
			Proto: "HTTP/1.0",
			Header: header,
			Body: body,
		},
	}
	return response
}

// Stub `http.Server`
type stubServer struct {
	http.Server
	response *stubResponse
	latency  time.Duration
}

func newStubServer(addr string, res *stubResponse, lat time.Duration) *stubServer {
	server := &stubServer{
		Server: http.Server{ Addr: addr },
		response: res,
		latency: lat,
	}
	return server
}

func (s *stubServer) Reply(code statusCode) *stubResponse {
	// Block for server's latency
	<-time.Tick(s.latency)
	//time.Sleep(s.latency)
	// Create status code and return the response
	s.response.StatusCode = (int)(code)
	// TODO: Assign Status
	return s.response
}

// Stub the `http.Client`
type stubClient struct { http.Client }

func newStubClient(timeout time.Duration) *stubClient {
	client := &stubClient{
		Client: http.Client{ Timeout: timeout },
	}
	return client
}

// Inject `*stubConnection` and `*stubServer` to simulate a server call
func (c *stubClient) Do(req *stubRequest, conn *stubConnection, server *stubServer) (*stubResponse, error) {
	// Block for the duration of `conn.latency` + `server.latency`
	// to simulaate real-world latencies and test timeoutn
	code := (statusCode)(server.response.StatusCode)

	conn.channel <- req

	go func(conn *stubConnection) {
		<-time.Tick(conn.latency)
		<-conn.channel
		// TODO: Do something with the receive request
		conn.channel <- server.Reply(code)
		<-time.Tick(server.latency)
		conn.done <- true
	}(conn)
	<-conn.done
	res := <-conn.channel
	return res.(*stubResponse), nil
}

var (
	// Setup connection
	networkLatency = time.Duration(100) * time.Millisecond
	conn = newStubConnection(networkLatency)

	// Setup client
	timeout = time.Duration(3) * time.Second
	client = newStubClient(timeout)

	// Setup server
	res = newStubResponse("200 OK", 200, header, body)
	endpoint = "http://jochasinga.io"
	serverLatency = time.Duration(100) * time.Millisecond
	server = newStubServer(endpoint, res, serverLatency)

	// Setup request
	header = http.Header{}
	jsonStr = `{"foo": ["bar", "baz"]}`
	body = ioutil.NopCloser(bytes.NewBuffer([]byte(jsonStr)))
	auth = map[string]string{"user": "pass"}
	requests = &stubRequests{
		Get: func(url, body string, auth map[string]string) (*stubResponse, error) {
			// Convert body from string to io.ReadCloser
			bodyReadCloser := ioutil.NopCloser(bytes.NewBuffer([]byte(body)))
			//req, err := http.NewRequest("GET", url, bodyReadCloser)
			req, err := newStubRequest("GET", url, bodyReadCloser)
			if err != nil {
				panic(err)
			}
			// TODO: include basic auth
			/*
			if len(auth) > 0 {
				for user, password := range auth {
					req.SetBasicAuth(user, password)
				}
			}
                        */
			res, err := client.Do(req, conn, server)
			if err != nil {
				panic(err)
			}
			return res, nil
		},
		GetAsync: func(url, body string, auth map[string]string, timeout int) (chan *stubResponse, error) {
			data := ioutil.NopCloser(bytes.NewBuffer([]byte(body)))

			req, err := newStubRequest("GET", url, data)
			if err != nil {
				panic(err)
			}

			temp := make(chan *stubResponse, 1)

			go func(t chan *stubResponse) {
				res, err := client.Do(req, conn, server)
				if err != nil {
					panic(err)
				}
				t <- res
			}(temp)
			return temp, nil
		},
	}
)

func TestGetResponseType(t *testing.T) {
	resp, err := requests.Get(endpoint, jsonStr, auth)
	if err != nil {
		t.Error(err)
	}
	returnType := reflect.TypeOf(resp)
	responseType := reflect.TypeOf((*stubResponse)(nil))
	if returnType != responseType {
		t.Errorf("Expected return type of `*stubResponse`, but it was %v instead.", returnType)
	}
}

func TestGetResponseStatus(t *testing.T) {
	resp, err := requests.Get(endpoint, jsonStr, auth)
	if err != nil {
		t.Error(err)
	}
	if resp.StatusCode != 200 {
		t.Errorf("Expected StatusCode `200`, but it was %s instead.", resp.Status)
	}
}

func TestGetAsyncResponseType(t *testing.T) {
	timeout := 1
	resultChan, err := requests.GetAsync(endpoint, jsonStr, auth, timeout)
	if err != nil {
		t.Error(err)
	}
	returnType := reflect.TypeOf(resultChan)
	responseType := reflect.TypeOf((chan *stubResponse)(nil))
	if returnType != responseType {
		t.Errorf("Expected return type of `chan *stubResponse`, but it was %v instead.", returnType)
	}
}

func TestGetAsyncResponseStatus(t *testing.T) {
	timeout := 1
	resultChan, err := requests.GetAsync(endpoint, jsonStr, auth, timeout)
	if err != nil {
		t.Error(err)
	}

	select {
	case result := <-resultChan:
		if result.StatusCode != 200 {
			t.Errorf("Expected Status of `200 OK`, but it was `%s` instead.", res.Status)
		}
		break
	// TODO: Fix this timeout
	case <-time.Tick(time.Duration(timeout) * time.Second):
		t.Log("time out!")
	}
}

func TestMain(m *testing.M) {
	v := m.Run()
	defer conn.Close()
	if v == 0 {
		os.Exit(1)
	}
	os.Exit(v)
}
<|endoftext|>


<commit_before>package state

import "fmt"

// ProgressCallback is called periodically to announce the degree of completeness of an operation
type ProgressCallback func(percent float64)

// ProgressLabelCallback is called when the progress label should be changed
type ProgressLabelCallback func(label string)

// MessageCallback is called when a log message has to be printed
type MessageCallback func(level, msg string)

type VoidCallback func()

// Consumer holds callbacks for the various state changes one
// might want to consume (show progress to the user, store messages
// in a text file, etc.)
type Consumer struct {
	OnProgress       ProgressCallback
	OnPauseProgress  VoidCallback
	OnResumeProgress VoidCallback
	OnProgressLabel  ProgressLabelCallback
	OnMessage        MessageCallback
}

// Progress announces the degree of completion of a task, in the [0,1] interval
func (c *Consumer) Progress(progress float64) {
	if c.OnProgress != nil {
		c.OnProgress(progress)
	}
}

func (c *Consumer) PauseProgress() {
	if c.OnPauseProgress != nil {
		c.OnPauseProgress()
	}
}

func (c *Consumer) ResumeProgress() {
	if c.OnResumeProgress != nil {
		c.OnResumeProgress()
	}
}

// ProgressLabel gives extra info about which task is currently being executed
func (c *Consumer) ProgressLabel(label string) {
	if c.OnProgressLabel != nil {
		c.OnProgressLabel(label)
	}
}

// Debug logs debug-level messages
func (c *Consumer) Debug(msg string) {
	if c.OnMessage != nil {
		c.OnMessage("debug", msg)
	}
}

// Debugf is a formatted variant of Debug
func (c *Consumer) Debugf(msg string, args ...interface{}) {
	if c.OnMessage != nil {
		c.OnMessage("debug", fmt.Sprintf(msg, args...))
	}
}

// Info logs info-level messages
func (c *Consumer) Info(msg string) {
	if c.OnMessage != nil {
		c.OnMessage("info", msg)
	}
}

// Infof is a formatted variant of Info
func (c *Consumer) Infof(msg string, args ...interface{}) {
	if c.OnMessage != nil {
		c.OnMessage("info", fmt.Sprintf(msg, args...))
	}
}

// Warn logs warning-level messages
func (c *Consumer) Warn(msg string) {
	if c.OnMessage != nil {
		c.OnMessage("warning", msg)
	}
}

// Warnf is a formatted version of Warn
func (c *Consumer) Warnf(msg string, args ...interface{}) {
	if c.OnMessage != nil {
		c.OnMessage("warning", fmt.Sprintf(msg, args...))
	}
}
<commit_msg>state: Introduce Consumer.CountCallback<commit_after>package state

import (
	"fmt"

	"github.com/itchio/wharf/counter"
)

// ProgressCallback is called periodically to announce the degree of completeness of an operation
type ProgressCallback func(percent float64)

// ProgressLabelCallback is called when the progress label should be changed
type ProgressLabelCallback func(label string)

// MessageCallback is called when a log message has to be printed
type MessageCallback func(level, msg string)

type VoidCallback func()

// Consumer holds callbacks for the various state changes one
// might want to consume (show progress to the user, store messages
// in a text file, etc.)
type Consumer struct {
	OnProgress       ProgressCallback
	OnPauseProgress  VoidCallback
	OnResumeProgress VoidCallback
	OnProgressLabel  ProgressLabelCallback
	OnMessage        MessageCallback
}

// Progress announces the degree of completion of a task, in the [0,1] interval
func (c *Consumer) Progress(progress float64) {
	if c.OnProgress != nil {
		c.OnProgress(progress)
	}
}

func (c *Consumer) PauseProgress() {
	if c.OnPauseProgress != nil {
		c.OnPauseProgress()
	}
}

func (c *Consumer) ResumeProgress() {
	if c.OnResumeProgress != nil {
		c.OnResumeProgress()
	}
}

// ProgressLabel gives extra info about which task is currently being executed
func (c *Consumer) ProgressLabel(label string) {
	if c.OnProgressLabel != nil {
		c.OnProgressLabel(label)
	}
}

// Debug logs debug-level messages
func (c *Consumer) Debug(msg string) {
	if c.OnMessage != nil {
		c.OnMessage("debug", msg)
	}
}

// Debugf is a formatted variant of Debug
func (c *Consumer) Debugf(msg string, args ...interface{}) {
	if c.OnMessage != nil {
		c.OnMessage("debug", fmt.Sprintf(msg, args...))
	}
}

// Info logs info-level messages
func (c *Consumer) Info(msg string) {
	if c.OnMessage != nil {
		c.OnMessage("info", msg)
	}
}

// Infof is a formatted variant of Info
func (c *Consumer) Infof(msg string, args ...interface{}) {
	if c.OnMessage != nil {
		c.OnMessage("info", fmt.Sprintf(msg, args...))
	}
}

// Warn logs warning-level messages
func (c *Consumer) Warn(msg string) {
	if c.OnMessage != nil {
		c.OnMessage("warning", msg)
	}
}

// Warnf is a formatted version of Warn
func (c *Consumer) Warnf(msg string, args ...interface{}) {
	if c.OnMessage != nil {
		c.OnMessage("warning", fmt.Sprintf(msg, args...))
	}
}

// CountCallback returns a function suitable for counter.NewWriterCallback
// or counter.NewReaderCallback
func (c *Consumer) CountCallback(totalSize int64) counter.CountCallback {
	return func(count int64) {
		c.Progress(float64(count) / float64(totalSize))
	}
}
<|endoftext|>


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

package txtdirect

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strconv"
	"strings"
	"testing"

	"github.com/miekg/dns"
)

// Testing TXT records
var txts = map[string]string{
	// type=host
	"_redirect.e2e.txtdirect.":       "v=txtv0;to=https://e2e.txtdirect.org;type=host",
	"_redirect.test.path.txtdirect.": "v=txtv0;to=https://path.e2e.txtdirect.org;type=host",
	// type=path
	"_redirect.path.txtdirect.": "v=txtv0;type=path",
	// type=""
	"_redirect.about.txtdirect.": "v=txtv0;to=https://about.txtdirect.org",
	"_redirect.pkg.txtdirect.":   "v=txtv0;to=https://pkg.txtdirect.org;type=gometa",
}

// Testing DNS server port
const port = 6000

// Initialize dns server instance
var server = &dns.Server{Addr: ":" + strconv.Itoa(port), Net: "udp"}

func TestMain(m *testing.M) {
	go RunDNSServer()
	os.Exit(m.Run())
}

func TestParse(t *testing.T) {
	tests := []struct {
		txtRecord string
		expected  record
		err       error
	}{
		{
			"v=txtv0;to=https://example.com/;code=302",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    301,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=302",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=302;vcs=hg;type=gometa",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Vcs:     "hg",
				Type:    "gometa",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=302;type=gometa;vcs=git",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Vcs:     "git",
				Type:    "gometa",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=test",
			record{},
			fmt.Errorf("could not parse status code"),
		},
		{
			"v=txtv1;to=https://example.com/;code=test",
			record{},
			fmt.Errorf("unhandled version 'txtv1'"),
		},
		{
			"v=txtv0;https://example.com/",
			record{},
			fmt.Errorf("arbitrary data not allowed"),
		},
		{
			"v=txtv0;to=https://example.com/caddy;type=path;code=302",
			record{
				Version: "txtv0",
				To:      "https://example.com/caddy",
				Type:    "path",
				Code:    302,
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;key=value",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    301,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to={?url}",
			record{
				Version: "txtv0",
				To:      "https://example.com/testing",
				Code:    301,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to={?url};from={method}",
			record{
				Version: "txtv0",
				To:      "https://example.com/testing",
				Code:    301,
				Type:    "host",
				From:    "GET",
			},
			nil,
		},
	}

	for i, test := range tests {
		r := record{}
		req, _ := http.NewRequest("GET", "http://example.com?url=https://example.com/testing", nil)
		err := r.Parse(test.txtRecord, req)

		if err != nil {
			if test.err == nil || !strings.HasPrefix(err.Error(), test.err.Error()) {
				t.Errorf("Test %d: Unexpected error: %s", i, err)
			}
			continue
		}
		if err == nil && test.err != nil {
			t.Errorf("Test %d: Expected error, got nil", i)
			continue
		}

		if got, want := r.Version, test.expected.Version; got != want {
			t.Errorf("Test %d: Expected Version to be '%s', got '%s'", i, want, got)
		}
		if got, want := r.To, test.expected.To; got != want {
			t.Errorf("Test %d: Expected To to be '%s', got '%s'", i, want, got)
		}
		if got, want := r.Code, test.expected.Code; got != want {
			t.Errorf("Test %d: Expected Code to be '%d', got '%d'", i, want, got)
		}
		if got, want := r.Type, test.expected.Type; got != want {
			t.Errorf("Test %d: Expected Type to be '%s', got '%s'", i, want, got)
		}
		if got, want := r.Vcs, test.expected.Vcs; got != want {
			t.Errorf("Test %d: Expected Vcs to be '%s', got '%s'", i, want, got)
		}
	}
}

/*
DNS TXT records currently registered at _td.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestRedirectDefault(t *testing.T) {
	testURL := "https://%d._td.test.txtdirect.org"
	dnsURL := "_redirect.%d._td.test.txtdirect.org"

	config := Config{
		Enable: []string{"host"},
	}

	for i := 0; ; i++ {
		_, err := net.LookupTXT(fmt.Sprintf(dnsURL, i))
		if err != nil {
			break
		}
		req, _ := http.NewRequest("GET", fmt.Sprintf(testURL, i), nil)
		rec := httptest.NewRecorder()
		err = Redirect(rec, req, config)
		if err != nil {
			t.Errorf("test %d: Unexpected error: %s", i, err)
		}
	}
}

/*
DNS TXT records currently registered at _ths.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestRedirectSuccess(t *testing.T) {
	testURL := "https://%d._ths.test.txtdirect.org"
	dnsURL := "_redirect.%d._ths.test.txtdirect.org"

	config := Config{
		Enable: []string{"host", "gometa"},
	}

	for i := 0; ; i++ {
		_, err := net.LookupTXT(fmt.Sprintf(dnsURL, i))
		if err != nil {
			break
		}
		req, _ := http.NewRequest("GET", fmt.Sprintf(testURL, i), nil)
		rec := httptest.NewRecorder()
		err = Redirect(rec, req, config)
		if err != nil {
			t.Errorf("test %d: Unexpected error: %s", i, err)
		}
	}
}

/*
DNS TXT records currently registered at _thf.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestRedirectFailure(t *testing.T) {
	testURL := "https://%d._thf.test.txtdirect.org"
	dnsURL := "_redirect.%d._thf.test.txtdirect.org"

	config := Config{
		Enable: []string{"host"},
	}

	for i := 0; ; i++ {
		_, err := net.LookupTXT(fmt.Sprintf(dnsURL, i))
		if err != nil {
			break
		}
		req, _ := http.NewRequest("GET", fmt.Sprintf(testURL, i), nil)
		rec := httptest.NewRecorder()
		err = Redirect(rec, req, config)
		if err == nil {
			t.Errorf("test %d: Expected error, got nil", i)
		}
	}
}

/*
DNS TXT records currently registered at _ths.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestPathBasedRoutingRedirect(t *testing.T) {
	config := Config{
		Enable: []string{"path"},
	}
	req := httptest.NewRequest("GET", "https://pkg.txtdirect.com/caddy/v1/", nil)
	w := httptest.NewRecorder()

	err := Redirect(w, req, config)
	if err != nil {
		t.Errorf("Unexpected error: %s", err)
	}
}

func TestRedirectBlacklist(t *testing.T) {
	config := Config{
		Enable: []string{"path"},
	}
	req := httptest.NewRequest("GET", "https://txtdirect.com/favicon.ico", nil)
	w := httptest.NewRecorder()

	err := Redirect(w, req, config)
	if err != nil {
		t.Errorf("Unexpected error: %s", err)
	}
}

func TestParsePlaceholders(t *testing.T) {
	tests := []struct {
		url         string
		placeholder string
		expected    string
	}{
		{
			"example.com{uri}",
			"/?test=test",
			"example.com/?test=test",
		},
		{
			"example.com{uri}/{~test}",
			"/?test=test",
			"example.com/?test=test/test",
		},
		{
			"example.com{uri}/{>Test}",
			"/?test=test",
			"example.com/?test=test/test-header",
		},
		{
			"example.com{uri}/{?test}",
			"/?test=test",
			"example.com/?test=test/test",
		},
	}
	for _, test := range tests {
		req := httptest.NewRequest("GET", "https://example.com"+test.placeholder, nil)
		req.AddCookie(&http.Cookie{Name: "test", Value: "test"})
		req.Header.Add("Test", "test-header")
		result := parsePlaceholders(test.url, req)
		if result != test.expected {
			t.Errorf("Expected %s, got %s", test.expected, result)
		}
	}
}

func Test_query(t *testing.T) {
	tests := []struct {
		zone string
		txt  string
	}{
		{
			"_redirect.about.txtdirect.",
			txts["_redirect.about.txtdirect."],
		},
		{
			"_redirect.pkg.txtdirect.",
			txts["_redirect.pkg.txtdirect."],
		},
	}
	for _, test := range tests {
		ctx := context.Background()
		c := Config{
			Resolver: "127.0.0.1:" + strconv.Itoa(port),
		}
		resp, err := query(test.zone, ctx, c)
		if err != nil {
			t.Fatal(err)
		}
		if resp[0] != txts[test.zone] {
			t.Fatalf("Expected %s, got %s", txts[test.zone], resp[0])
		}
	}
}

func parseDNSQuery(m *dns.Msg) {
	for _, q := range m.Question {
		switch q.Qtype {
		case dns.TypeTXT:
			log.Printf("Query for %s\n", q.Name)
			m.Answer = append(m.Answer, &dns.TXT{
				Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 60},
				Txt: []string{txts[q.Name]},
			})
		}
	}
}

func handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false

	switch r.Opcode {
	case dns.OpcodeQuery:
		parseDNSQuery(m)
	}

	w.WriteMsg(m)
}

func RunDNSServer() {
	dns.HandleFunc("txtdirect.", handleDNSRequest)
	err := server.ListenAndServe()
	defer server.Shutdown()
	if err != nil {
		log.Fatalf("Failed to start server: %s\n ", err.Error())
	}
}
<commit_msg>Add e2e tests for Redirect using mock DNS Server<commit_after>/*
Copyright 2017 - The TXTdirect Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package txtdirect

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strconv"
	"strings"
	"testing"

	"github.com/miekg/dns"
)

// Testing TXT records
var txts = map[string]string{
	// type=host
	"_redirect.e2e.txtdirect.":       "v=txtv0;to=https://e2e.txtdirect.org;type=host",
	"_redirect.test.path.txtdirect.": "v=txtv0;to=https://path.e2e.txtdirect.org;type=host",
	// type=path
	"_redirect.path.txtdirect.": "v=txtv0;type=path",
	// type=gometa
	"_redirect.gometa.txtdirect.": "v=txtv0;to=https://pkg.txtdirect.org;type=gometa",
	// type=""
	"_redirect.about.txtdirect.": "v=txtv0;to=https://about.txtdirect.org",
	"_redirect.pkg.txtdirect.":   "v=txtv0;to=https://pkg.txtdirect.org;type=gometa",
}

// Testing DNS server port
const port = 6000

// Initialize dns server instance
var server = &dns.Server{Addr: ":" + strconv.Itoa(port), Net: "udp"}

func TestMain(m *testing.M) {
	go RunDNSServer()
	os.Exit(m.Run())
}

func TestParse(t *testing.T) {
	tests := []struct {
		txtRecord string
		expected  record
		err       error
	}{
		{
			"v=txtv0;to=https://example.com/;code=302",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    301,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=302",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=302;vcs=hg;type=gometa",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Vcs:     "hg",
				Type:    "gometa",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=302;type=gometa;vcs=git",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    302,
				Vcs:     "git",
				Type:    "gometa",
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;code=test",
			record{},
			fmt.Errorf("could not parse status code"),
		},
		{
			"v=txtv1;to=https://example.com/;code=test",
			record{},
			fmt.Errorf("unhandled version 'txtv1'"),
		},
		{
			"v=txtv0;https://example.com/",
			record{},
			fmt.Errorf("arbitrary data not allowed"),
		},
		{
			"v=txtv0;to=https://example.com/caddy;type=path;code=302",
			record{
				Version: "txtv0",
				To:      "https://example.com/caddy",
				Type:    "path",
				Code:    302,
			},
			nil,
		},
		{
			"v=txtv0;to=https://example.com/;key=value",
			record{
				Version: "txtv0",
				To:      "https://example.com/",
				Code:    301,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to={?url}",
			record{
				Version: "txtv0",
				To:      "https://example.com/testing",
				Code:    301,
				Type:    "host",
			},
			nil,
		},
		{
			"v=txtv0;to={?url};from={method}",
			record{
				Version: "txtv0",
				To:      "https://example.com/testing",
				Code:    301,
				Type:    "host",
				From:    "GET",
			},
			nil,
		},
	}

	for i, test := range tests {
		r := record{}
		req, _ := http.NewRequest("GET", "http://example.com?url=https://example.com/testing", nil)
		err := r.Parse(test.txtRecord, req)

		if err != nil {
			if test.err == nil || !strings.HasPrefix(err.Error(), test.err.Error()) {
				t.Errorf("Test %d: Unexpected error: %s", i, err)
			}
			continue
		}
		if err == nil && test.err != nil {
			t.Errorf("Test %d: Expected error, got nil", i)
			continue
		}

		if got, want := r.Version, test.expected.Version; got != want {
			t.Errorf("Test %d: Expected Version to be '%s', got '%s'", i, want, got)
		}
		if got, want := r.To, test.expected.To; got != want {
			t.Errorf("Test %d: Expected To to be '%s', got '%s'", i, want, got)
		}
		if got, want := r.Code, test.expected.Code; got != want {
			t.Errorf("Test %d: Expected Code to be '%d', got '%d'", i, want, got)
		}
		if got, want := r.Type, test.expected.Type; got != want {
			t.Errorf("Test %d: Expected Type to be '%s', got '%s'", i, want, got)
		}
		if got, want := r.Vcs, test.expected.Vcs; got != want {
			t.Errorf("Test %d: Expected Vcs to be '%s', got '%s'", i, want, got)
		}
	}
}

/*
DNS TXT records currently registered at _td.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestRedirectDefault(t *testing.T) {
	testURL := "https://%d._td.test.txtdirect.org"
	dnsURL := "_redirect.%d._td.test.txtdirect.org"

	config := Config{
		Enable: []string{"host"},
	}

	for i := 0; ; i++ {
		_, err := net.LookupTXT(fmt.Sprintf(dnsURL, i))
		if err != nil {
			break
		}
		req, _ := http.NewRequest("GET", fmt.Sprintf(testURL, i), nil)
		rec := httptest.NewRecorder()
		err = Redirect(rec, req, config)
		if err != nil {
			t.Errorf("test %d: Unexpected error: %s", i, err)
		}
	}
}

/*
DNS TXT records currently registered at _ths.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestRedirectSuccess(t *testing.T) {
	testURL := "https://%d._ths.test.txtdirect.org"
	dnsURL := "_redirect.%d._ths.test.txtdirect.org"

	config := Config{
		Enable: []string{"host", "gometa"},
	}

	for i := 0; ; i++ {
		_, err := net.LookupTXT(fmt.Sprintf(dnsURL, i))
		if err != nil {
			break
		}
		req, _ := http.NewRequest("GET", fmt.Sprintf(testURL, i), nil)
		rec := httptest.NewRecorder()
		err = Redirect(rec, req, config)
		if err != nil {
			t.Errorf("test %d: Unexpected error: %s", i, err)
		}
	}
}

/*
DNS TXT records currently registered at _thf.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestRedirectFailure(t *testing.T) {
	testURL := "https://%d._thf.test.txtdirect.org"
	dnsURL := "_redirect.%d._thf.test.txtdirect.org"

	config := Config{
		Enable: []string{"host"},
	}

	for i := 0; ; i++ {
		_, err := net.LookupTXT(fmt.Sprintf(dnsURL, i))
		if err != nil {
			break
		}
		req, _ := http.NewRequest("GET", fmt.Sprintf(testURL, i), nil)
		rec := httptest.NewRecorder()
		err = Redirect(rec, req, config)
		if err == nil {
			t.Errorf("test %d: Expected error, got nil", i)
		}
	}
}

/*
DNS TXT records currently registered at _ths.test.txtdirect.org available in:
https://raw.githubusercontent.com/txtdirect/_test-records/master/test.txtdirect.org
*/
func TestPathBasedRoutingRedirect(t *testing.T) {
	config := Config{
		Enable: []string{"path"},
	}
	req := httptest.NewRequest("GET", "https://pkg.txtdirect.com/caddy/v1/", nil)
	w := httptest.NewRecorder()

	err := Redirect(w, req, config)
	if err != nil {
		t.Errorf("Unexpected error: %s", err)
	}
}

func TestRedirectBlacklist(t *testing.T) {
	config := Config{
		Enable: []string{"path"},
	}
	req := httptest.NewRequest("GET", "https://txtdirect.com/favicon.ico", nil)
	w := httptest.NewRecorder()

	err := Redirect(w, req, config)
	if err != nil {
		t.Errorf("Unexpected error: %s", err)
	}
}

func TestParsePlaceholders(t *testing.T) {
	tests := []struct {
		url         string
		placeholder string
		expected    string
	}{
		{
			"example.com{uri}",
			"/?test=test",
			"example.com/?test=test",
		},
		{
			"example.com{uri}/{~test}",
			"/?test=test",
			"example.com/?test=test/test",
		},
		{
			"example.com{uri}/{>Test}",
			"/?test=test",
			"example.com/?test=test/test-header",
		},
		{
			"example.com{uri}/{?test}",
			"/?test=test",
			"example.com/?test=test/test",
		},
	}
	for _, test := range tests {
		req := httptest.NewRequest("GET", "https://example.com"+test.placeholder, nil)
		req.AddCookie(&http.Cookie{Name: "test", Value: "test"})
		req.Header.Add("Test", "test-header")
		result := parsePlaceholders(test.url, req)
		if result != test.expected {
			t.Errorf("Expected %s, got %s", test.expected, result)
		}
	}
}

func Test_query(t *testing.T) {
	tests := []struct {
		zone string
		txt  string
	}{
		{
			"_redirect.about.txtdirect.",
			txts["_redirect.about.txtdirect."],
		},
		{
			"_redirect.pkg.txtdirect.",
			txts["_redirect.pkg.txtdirect."],
		},
	}
	for _, test := range tests {
		ctx := context.Background()
		c := Config{
			Resolver: "127.0.0.1:" + strconv.Itoa(port),
		}
		resp, err := query(test.zone, ctx, c)
		if err != nil {
			t.Fatal(err)
		}
		if resp[0] != txts[test.zone] {
			t.Fatalf("Expected %s, got %s", txts[test.zone], resp[0])
		}
	}
}

func parseDNSQuery(m *dns.Msg) {
	for _, q := range m.Question {
		switch q.Qtype {
		case dns.TypeTXT:
			log.Printf("Query for %s\n", q.Name)
			m.Answer = append(m.Answer, &dns.TXT{
				Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 60},
				Txt: []string{txts[q.Name]},
			})
		}
	}
}

func handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false

	switch r.Opcode {
	case dns.OpcodeQuery:
		parseDNSQuery(m)
	}

	w.WriteMsg(m)
}

func RunDNSServer() {
	dns.HandleFunc("txtdirect.", handleDNSRequest)
	err := server.ListenAndServe()
	defer server.Shutdown()
	if err != nil {
		log.Printf("Failed to start server: %s\n ", err.Error())
	}
}

func TestRedirectTypeHost(t *testing.T) {
	tests := []struct {
		url      string
		txt      string
		expected string
		enable   []string
	}{
		{
			"https://e2e.txtdirect",
			txts["_redirect.path.txtdirect."],
			"https://e2e.txtdirect.org",
			[]string{"host"},
		},
		{
			"https://path.txtdirect/test",
			txts["_redirect.path.e2e.txtdirect."],
			"https://path.e2e.txtdirect.org",
			[]string{"path", "host"},
		},
		{
			"https://gometa.txtdirect",
			txts["_redirect.gometa.txtdirect."],
			"https://pkg.txtdirect.org",
			[]string{"gometa"},
		},
	}
	for _, test := range tests {
		req := httptest.NewRequest("GET", test.url, nil)
		resp := httptest.NewRecorder()
		c := Config{
			Resolver: "127.0.0.1:" + strconv.Itoa(port),
			Enable:   test.enable,
		}
		if err := Redirect(resp, req, c); err != nil {
			t.Fatalf("Unexpected error occured: %s", err.Error())
		}
		if !strings.Contains(resp.Body.String(), test.expected) {
			t.Fatalf("Expected %s to be in \"%s\"", test.expected, resp.Body.String())
		}
	}
}
<|endoftext|>


<commit_before>package udata

//-----------------------------------------------------------------------------
// Import statements:
//-----------------------------------------------------------------------------

import "sort"

//-----------------------------------------------------------------------------
// Typedefs:
//-----------------------------------------------------------------------------

type service struct {
	name   string
	roles  []string
	groups []string
	ports  []port
}

type port struct {
	num      int
	protocol string
	ingress  string
}

type serviceMap map[string]service

//-----------------------------------------------------------------------------
// func: findOne
//-----------------------------------------------------------------------------

func findOne(src, dst []string) bool {
	for _, i := range src {
		for _, j := range dst {
			if i == j {
				return true
			}
		}
	}
	return false
}

//-----------------------------------------------------------------------------
// func: listUnits
//-----------------------------------------------------------------------------

func (s *serviceMap) listUnits() (list []string) {

	// Map as set:
	m := map[string]struct{}{}
	for _, service := range *s {
		m[service.name] = struct{}{}
	}

	// Set to slice:
	for k := range m {
		list = append(list, k)
	}

	// Sort and return:
	sort.Strings(list)
	return
}

//-----------------------------------------------------------------------------
// func: listPorts
//-----------------------------------------------------------------------------

func (s *serviceMap) listPorts(protocol string) (list []int) {

	// Default ports:
	m := map[int]struct{}{}
	if protocol == "tcp" {
		m[22] = struct{}{}
	}

	// Map as set:
	for _, service := range *s {
		for _, port := range service.ports {
			if port.protocol == protocol {
				m[port.num] = struct{}{}
			}
		}
	}

	// Set to slice:
	for k := range m {
		list = append(list, k)
	}

	// Sort and return:
	sort.Ints(list)
	return
}

//-----------------------------------------------------------------------------
// func: load
//-----------------------------------------------------------------------------

func (s *serviceMap) load(roles, groups []string) {

	// Map roles to services:
	roleServices := map[string][]string{

		"quorum": {
			"docker", "rexray", "etchost", "zookeeper", "etcd-master", "rkt-api",
			"cadvisor", "node-exporter", "zookeeper-exporter"},

		"master": {
			"docker", "rexray", "etchost", "etcd-proxy", "calico", "mesos-dns",
			"mesos-master", "marathon", "rkt-api", "cadvisor", "node-exporter",
			"mesos-master-exporter", "confd", "alertmanager", "prometheus"},

		"worker": {
			"docker", "rexray", "etchost", "etcd-proxy", "calico", "go-dnsmasq",
			"marathon-lb", "mesos-agent", "rkt-api", "cadvisor", "node-exporter",
			"mesos-agent-exporter", "haproxy-exporter"},

		"border": {
			"docker", "rexray", "etchost", "etcd-proxy", "calico", "mongodb",
			"pritunl", "rkt-api", "cadvisor", "node-exporter"},
	}

	// Map services to config:
	serviceConfig := map[string]service{

		"docker": {
			name:   "docker.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2375, protocol: "tcp", ingress: ""},
			},
		},

		"rexray": {
			name:   "rexray.service",
			groups: []string{"base"},
			ports: []port{
				{num: 7979, protocol: "tcp", ingress: ""},
			},
		},

		"etchost": {
			name:   "etchost.timer",
			groups: []string{"base"},
		},

		"etcd-proxy": {
			name:   "etcd2.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2379, protocol: "tcp", ingress: ""},
			},
		},

		"calico": {
			name:   "calico.service",
			groups: []string{"base"},
		},

		"zookeeper": {
			name:   "zookeeper.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2181, protocol: "tcp", ingress: ""},
				{num: 2888, protocol: "tcp", ingress: ""},
				{num: 3888, protocol: "tcp", ingress: ""},
			},
		},

		"etcd-master": {
			name:   "etcd2.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2379, protocol: "tcp", ingress: ""},
				{num: 2380, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-dns": {
			name:   "mesos-dns.service",
			groups: []string{"base"},
			ports: []port{
				{num: 53, protocol: "tcp", ingress: ""},
				{num: 54, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-master": {
			name:   "mesos-master.service",
			groups: []string{"base"},
			ports: []port{
				{num: 5050, protocol: "tcp", ingress: ""},
			},
		},

		"marathon": {
			name:   "marathon.service",
			groups: []string{"base"},
			ports: []port{
				{num: 8080, protocol: "tcp", ingress: ""},
				{num: 9292, protocol: "tcp", ingress: ""},
			},
		},

		"go-dnsmasq": {
			name:   "go-dnsmasq.service",
			groups: []string{"base"},
			ports: []port{
				{num: 53, protocol: "tcp", ingress: ""},
			},
		},

		"marathon-lb": {
			name:   "marathon-lb.service",
			groups: []string{"base"},
			ports: []port{
				{num: 80, protocol: "tcp", ingress: ""},
				{num: 443, protocol: "tcp", ingress: ""},
				{num: 9090, protocol: "tcp", ingress: ""},
				{num: 9091, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-agent": {
			name:   "mesos-agent.service",
			groups: []string{"base"},
			ports: []port{
				{num: 5051, protocol: "tcp", ingress: ""},
			},
		},

		"mongodb": {
			name:   "mongodb.service",
			groups: []string{"base"},
			ports: []port{
				{num: 27017, protocol: "tcp", ingress: ""},
			},
		},

		"pritunl": {
			name:   "pritunl.service",
			groups: []string{"base"},
			ports: []port{
				{num: 80, protocol: "tcp", ingress: ""},
				{num: 443, protocol: "tcp", ingress: ""},
				{num: 9756, protocol: "tcp", ingress: ""},
				{num: 18443, protocol: "udp", ingress: ""},
			},
		},

		"rkt-api": {
			name:   "rkt-api.service",
			groups: []string{"insight"},
		},

		"cadvisor": {
			name:   "cadvisor.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 4194, protocol: "tcp", ingress: ""},
			},
		},

		"node-exporter": {
			name:   "node-exporter.service",
			groups: []string{"insight"},
		},

		"zookeeper-exporter": {
			name:   "zookeeper-exporter.service",
			groups: []string{"insight"},
		},

		"mesos-master-exporter": {
			name:   "mesos-master-exporter.service",
			groups: []string{"insight"},
		},

		"mesos-agent-exporter": {
			name:   "mesos-agent-exporter.service",
			groups: []string{"insight"},
		},

		"haproxy-exporter": {
			name:   "haproxy-exporter.service",
			groups: []string{"insight"},
		},

		"confd": {
			name:   "confd.service",
			groups: []string{"insight"},
		},

		"alertmanager": {
			name:   "alertmanager.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9093, protocol: "tcp", ingress: ""},
			},
		},

		"prometheus": {
			name:   "prometheus.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9191, protocol: "tcp", ingress: ""},
			},
		},
	}

	// Filter my services:
	*s = serviceMap{}
	for _, role := range roles {
		for _, service := range roleServices[role] {
			if findOne(serviceConfig[service].groups, groups) {
				(*s)[service] = serviceConfig[service]
			}
		}
	}
}
<commit_msg>Added missing service ports<commit_after>package udata

//-----------------------------------------------------------------------------
// Import statements:
//-----------------------------------------------------------------------------

import "sort"

//-----------------------------------------------------------------------------
// Typedefs:
//-----------------------------------------------------------------------------

type service struct {
	name   string
	roles  []string
	groups []string
	ports  []port
}

type port struct {
	num      int
	protocol string
	ingress  string
}

type serviceMap map[string]service

//-----------------------------------------------------------------------------
// func: findOne
//-----------------------------------------------------------------------------

func findOne(src, dst []string) bool {
	for _, i := range src {
		for _, j := range dst {
			if i == j {
				return true
			}
		}
	}
	return false
}

//-----------------------------------------------------------------------------
// func: listUnits
//-----------------------------------------------------------------------------

func (s *serviceMap) listUnits() (list []string) {

	// Map as set:
	m := map[string]struct{}{}
	for _, service := range *s {
		m[service.name] = struct{}{}
	}

	// Set to slice:
	for k := range m {
		list = append(list, k)
	}

	// Sort and return:
	sort.Strings(list)
	return
}

//-----------------------------------------------------------------------------
// func: listPorts
//-----------------------------------------------------------------------------

func (s *serviceMap) listPorts(protocol string) (list []int) {

	// Default ports:
	m := map[int]struct{}{}
	if protocol == "tcp" {
		m[22] = struct{}{}
	}

	// Map as set:
	for _, service := range *s {
		for _, port := range service.ports {
			if port.protocol == protocol {
				m[port.num] = struct{}{}
			}
		}
	}

	// Set to slice:
	for k := range m {
		list = append(list, k)
	}

	// Sort and return:
	sort.Ints(list)
	return
}

//-----------------------------------------------------------------------------
// func: load
//-----------------------------------------------------------------------------

func (s *serviceMap) load(roles, groups []string) {

	// Map roles to services:
	roleServices := map[string][]string{

		"quorum": {
			"docker", "rexray", "etchost", "zookeeper", "etcd-master", "rkt-api",
			"cadvisor", "node-exporter", "zookeeper-exporter"},

		"master": {
			"docker", "rexray", "etchost", "etcd-proxy", "calico", "mesos-dns",
			"mesos-master", "marathon", "rkt-api", "cadvisor", "node-exporter",
			"mesos-master-exporter", "confd", "alertmanager", "prometheus"},

		"worker": {
			"docker", "rexray", "etchost", "etcd-proxy", "calico", "go-dnsmasq",
			"marathon-lb", "mesos-agent", "rkt-api", "cadvisor", "node-exporter",
			"mesos-agent-exporter", "haproxy-exporter"},

		"border": {
			"docker", "rexray", "etchost", "etcd-proxy", "calico", "mongodb",
			"pritunl", "rkt-api", "cadvisor", "node-exporter"},
	}

	// Map services to config:
	serviceConfig := map[string]service{

		"docker": {
			name:   "docker.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2375, protocol: "tcp", ingress: ""},
			},
		},

		"rexray": {
			name:   "rexray.service",
			groups: []string{"base"},
			ports: []port{
				{num: 7979, protocol: "tcp", ingress: ""},
			},
		},

		"etchost": {
			name:   "etchost.timer",
			groups: []string{"base"},
		},

		"etcd-proxy": {
			name:   "etcd2.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2379, protocol: "tcp", ingress: ""},
			},
		},

		"calico": {
			name:   "calico.service",
			groups: []string{"base"},
		},

		"zookeeper": {
			name:   "zookeeper.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2181, protocol: "tcp", ingress: ""},
				{num: 2888, protocol: "tcp", ingress: ""},
				{num: 3888, protocol: "tcp", ingress: ""},
			},
		},

		"etcd-master": {
			name:   "etcd2.service",
			groups: []string{"base"},
			ports: []port{
				{num: 2379, protocol: "tcp", ingress: ""},
				{num: 2380, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-dns": {
			name:   "mesos-dns.service",
			groups: []string{"base"},
			ports: []port{
				{num: 53, protocol: "tcp", ingress: ""},
				{num: 54, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-master": {
			name:   "mesos-master.service",
			groups: []string{"base"},
			ports: []port{
				{num: 5050, protocol: "tcp", ingress: ""},
			},
		},

		"marathon": {
			name:   "marathon.service",
			groups: []string{"base"},
			ports: []port{
				{num: 8080, protocol: "tcp", ingress: ""},
				{num: 9292, protocol: "tcp", ingress: ""},
			},
		},

		"go-dnsmasq": {
			name:   "go-dnsmasq.service",
			groups: []string{"base"},
			ports: []port{
				{num: 53, protocol: "tcp", ingress: ""},
			},
		},

		"marathon-lb": {
			name:   "marathon-lb.service",
			groups: []string{"base"},
			ports: []port{
				{num: 80, protocol: "tcp", ingress: ""},
				{num: 443, protocol: "tcp", ingress: ""},
				{num: 9090, protocol: "tcp", ingress: ""},
				{num: 9091, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-agent": {
			name:   "mesos-agent.service",
			groups: []string{"base"},
			ports: []port{
				{num: 5051, protocol: "tcp", ingress: ""},
			},
		},

		"mongodb": {
			name:   "mongodb.service",
			groups: []string{"base"},
			ports: []port{
				{num: 27017, protocol: "tcp", ingress: ""},
			},
		},

		"pritunl": {
			name:   "pritunl.service",
			groups: []string{"base"},
			ports: []port{
				{num: 80, protocol: "tcp", ingress: ""},
				{num: 443, protocol: "tcp", ingress: ""},
				{num: 9756, protocol: "tcp", ingress: ""},
				{num: 18443, protocol: "udp", ingress: ""},
			},
		},

		"rkt-api": {
			name:   "rkt-api.service",
			groups: []string{"insight"},
		},

		"cadvisor": {
			name:   "cadvisor.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 4194, protocol: "tcp", ingress: ""},
			},
		},

		"node-exporter": {
			name:   "node-exporter.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9101, protocol: "tcp", ingress: ""},
			},
		},

		"zookeeper-exporter": {
			name:   "zookeeper-exporter.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9103, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-master-exporter": {
			name:   "mesos-master-exporter.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9104, protocol: "tcp", ingress: ""},
			},
		},

		"mesos-agent-exporter": {
			name:   "mesos-agent-exporter.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9105, protocol: "tcp", ingress: ""},
			},
		},

		"haproxy-exporter": {
			name:   "haproxy-exporter.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9102, protocol: "tcp", ingress: ""},
			},
		},

		"confd": {
			name:   "confd.service",
			groups: []string{"insight"},
		},

		"alertmanager": {
			name:   "alertmanager.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9093, protocol: "tcp", ingress: ""},
			},
		},

		"prometheus": {
			name:   "prometheus.service",
			groups: []string{"insight"},
			ports: []port{
				{num: 9191, protocol: "tcp", ingress: ""},
			},
		},
	}

	// Filter my services:
	*s = serviceMap{}
	for _, role := range roles {
		for _, service := range roleServices[role] {
			if findOne(serviceConfig[service].groups, groups) {
				(*s)[service] = serviceConfig[service]
			}
		}
	}
}
<|endoftext|>