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

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/gorhill/cronexpr"
	"github.com/mozilla/tls-observatory/certificate"
	"github.com/mozilla/tls-observatory/database"

	"gopkg.in/yaml.v2"
)

type Configuration struct {
	Runs []Run
	Smtp struct {
		Host string
		Port int
		From string
		Auth struct {
			User, Pass string
		}
	}
}
type Run struct {
	Targets       []string
	Assertions    []Assertion
	Cron          string
	Notifications NotificationsConf
}

type Assertion struct {
	Certificate struct {
		Validity struct {
			NotBefore string
			NotAfter  string
		}
	}
	Analysis struct {
		Analyzer string
		Result   string `json:"result"`
	}
}

type NotificationsConf struct {
	Irc struct {
		Channels []string
	}
	Email struct {
		Recipients []string
	}
}

var (
	cfgFile     string
	observatory string
	debug       bool
	conf        Configuration
)

func main() {
	flag.StringVar(&observatory, "observatory", "https://tls-observatory.services.mozilla.com", "URL of the observatory")
	flag.StringVar(&cfgFile, "c", "/etc/tls-observatory/runner.yaml", "YAML configuration file")
	flag.BoolVar(&debug, "debug", false, "Set debug logging")
	flag.Parse()
	conf = getConf(cfgFile)
	exit := make(chan bool)
	for i, run := range conf.Runs {
		go run.start(i)
	}
	<-exit
}

func (r Run) start(id int) {
	for {
		cexpr, err := cronexpr.Parse(r.Cron)
		if err != nil {
			panic(err)
		}
		// sleep until the next run is scheduled to happen
		nrun := cexpr.Next(time.Now())
		waitduration := nrun.Sub(time.Now())
		log.Printf("[info] run %d will start at %v (in %v)", id, nrun, waitduration)
		time.Sleep(waitduration)

		notifchan := make(chan Notification)
		done := make(chan bool)
		go processNotifications(notifchan, done)
		var wg sync.WaitGroup
		for _, target := range r.Targets {
			log.Printf("[info] run %d starting scan of target %q", id, target)
			id, err := r.scan(target)
			debugprint("got scan id %s", id)
			if err != nil {
				log.Printf("[error] failed to launch against %q: %v", target, err)
				continue
			}
			wg.Add(1)
			go r.evaluate(id, notifchan, &wg)
		}
		wg.Wait()
		close(notifchan)
		<-done
	}
}

type scan struct {
	ID string `json:"scan_id"`
}

func (r Run) scan(target string) (id string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("scan(target=%q) -> %v", e)
		}
	}()
	resp, err := http.Post(observatory+"/api/v1/scan?target="+target, "application/json", nil)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	var s scan
	err = json.Unmarshal(body, &s)
	if err != nil {
		panic(err)
	}
	if s.ID == "" {
		panic("failed to launch scan on target " + target)
	}
	id = s.ID
	return
}

func (r Run) evaluate(id string, notifchan chan Notification, wg *sync.WaitGroup) {
	defer func() {
		if e := recover(); e != nil {
			log.Printf("[error] evaluate(id=%q) -> %v", id, e)
		}
		wg.Done()
	}()
	var (
		results database.Scan
		cert    certificate.Certificate
		err     error
	)
	for {
		resp, err := http.Get(observatory + "/api/v1/results?id=" + id)
		if err != nil {
			panic(err)
		}
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			panic(err)
		}
		err = json.Unmarshal(body, &results)
		if err != nil {
			panic(err)
		}
		if results.Complperc >= 100 {
			debugprint("scan id %s completed", id)
			break
		}
		time.Sleep(5 * time.Second)
	}
	debugprint("getting certificate id %d", results.Cert_id)
	if !results.Has_tls && results.Cert_id < 1 {
		log.Printf("[info] target %q is not TLS enabled", results.Target)
		return
	}
	cert, err = getCert(results.Cert_id)
	if err != nil {
		panic(err)
	}
	for _, a := range r.Assertions {
		r.AssertNotBefore(a, results.Target, cert.Validity.NotBefore, notifchan)
		r.AssertNotAfter(a, results.Target, cert.Validity.NotAfter, notifchan)
		r.AssertAnalysis(a, results, cert, notifchan)
	}
	return
}

func getCert(id int64) (cert certificate.Certificate, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getCert(id=%q) -> %v", e)
		}
	}()
	resp, err := http.Get(fmt.Sprintf("%s/api/v1/certificate?id=%d", observatory, id))
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	err = json.Unmarshal(body, &cert)
	if err != nil {
		panic(err)
	}
	return
}

func debugprint(format string, a ...interface{}) {
	if debug {
		log.Printf("[debug] "+format, a...)
	}
}

// getConf first read the configuration from a local YAML file then overrides it
// with the content of the TLSOBS_RUNNER_CONF var (which much contain a full yaml file
// encoded in base64), and then overrides the SMTP settings with various SMTP env var
func getConf(cfg string) (c Configuration) {
	// load the local configuration file
	fd, err := ioutil.ReadFile(cfg)
	if err != nil {
		log.Fatal(err)
	}
	err = yaml.Unmarshal(fd, &c)
	if err != nil {
		log.Fatalf("error: %v", err)
	}
	// iterate over notifications in targets and unbase64 the values
	for i, run := range c.Runs {
		for j, rcpt := range run.Notifications.Email.Recipients {
			if len(rcpt) < 5 || rcpt[0:4] != "b64:" {
				continue
			}
			data, err := base64.StdEncoding.DecodeString(rcpt[4:])
			if err != nil {
				log.Fatalf("error while decoding b64 recipient: %v", err)
			}
			c.Runs[i].Notifications.Email.Recipients[j] = fmt.Sprintf("%s", bytes.TrimRight(data, "\n"))
		}
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_HOST") != "" {
		c.Smtp.Host = os.Getenv("TLSOBS_RUNNER_SMTP_HOST")
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_PORT") != "" {
		var err error
		c.Smtp.Port, err = strconv.Atoi(os.Getenv("TLSOBS_RUNNER_SMTP_PORT"))
		if err != nil {
			log.Printf("[error] failed to read smtp port from env variable: %v", err)
		}
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_FROM") != "" {
		c.Smtp.From = os.Getenv("TLSOBS_RUNNER_SMTP_FROM")
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_USER") != "" {
		c.Smtp.Auth.User = os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_USER")
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_PASS") != "" {
		c.Smtp.Auth.Pass = os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_PASS")
	}
	return c
}
<commit_msg>Fix handling of scanid in tlsobs-runner<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
   License, v. 2.0. If a copy of the MPL was not distributed with this
   file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/gorhill/cronexpr"
	"github.com/mozilla/tls-observatory/certificate"
	"github.com/mozilla/tls-observatory/database"

	"gopkg.in/yaml.v2"
)

type Configuration struct {
	Runs []Run
	Smtp struct {
		Host string
		Port int
		From string
		Auth struct {
			User, Pass string
		}
	}
}
type Run struct {
	Targets       []string
	Assertions    []Assertion
	Cron          string
	Notifications NotificationsConf
}

type Assertion struct {
	Certificate struct {
		Validity struct {
			NotBefore string
			NotAfter  string
		}
	}
	Analysis struct {
		Analyzer string
		Result   string `json:"result"`
	}
}

type NotificationsConf struct {
	Irc struct {
		Channels []string
	}
	Email struct {
		Recipients []string
	}
}

var (
	cfgFile     string
	observatory string
	debug       bool
	conf        Configuration
)

func main() {
	flag.StringVar(&observatory, "observatory", "https://tls-observatory.services.mozilla.com", "URL of the observatory")
	flag.StringVar(&cfgFile, "c", "/etc/tls-observatory/runner.yaml", "YAML configuration file")
	flag.BoolVar(&debug, "debug", false, "Set debug logging")
	flag.Parse()
	conf = getConf(cfgFile)
	exit := make(chan bool)
	for i, run := range conf.Runs {
		go run.start(i)
	}
	<-exit
}

func (r Run) start(id int) {
	for {
		cexpr, err := cronexpr.Parse(r.Cron)
		if err != nil {
			panic(err)
		}
		// sleep until the next run is scheduled to happen
		nrun := cexpr.Next(time.Now())
		waitduration := nrun.Sub(time.Now())
		log.Printf("[info] run %d will start at %v (in %v)", id, nrun, waitduration)
		time.Sleep(waitduration)

		notifchan := make(chan Notification)
		done := make(chan bool)
		go processNotifications(notifchan, done)
		var wg sync.WaitGroup
		for _, target := range r.Targets {
			log.Printf("[info] run %d starting scan of target %q", id, target)
			id, err := r.scan(target)
			debugprint("got scan id %d", id)
			if err != nil {
				log.Printf("[error] failed to launch against %q: %v", target, err)
				continue
			}
			wg.Add(1)
			go r.evaluate(id, notifchan, &wg)
		}
		wg.Wait()
		close(notifchan)
		<-done
	}
}

type scan struct {
	ID int64 `json:"scan_id"`
}

func (r Run) scan(target string) (id int64, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("scan(target=%q) -> %v", e)
		}
	}()
	resp, err := http.Post(observatory+"/api/v1/scan?target="+target, "application/json", nil)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	var s scan
	err = json.Unmarshal(body, &s)
	if err != nil {
		panic(err)
	}
	if s.ID < 1 {
		panic("failed to launch scan on target " + target)
	}
	id = s.ID
	return
}

func (r Run) evaluate(id int64, notifchan chan Notification, wg *sync.WaitGroup) {
	defer func() {
		if e := recover(); e != nil {
			log.Printf("[error] evaluate(id=%q) -> %v", id, e)
		}
		wg.Done()
	}()
	var (
		results database.Scan
		cert    certificate.Certificate
		err     error
	)
	for {
		resp, err := http.Get(fmt.Sprintf("%s/api/v1/results?id=%d", observatory, id))
		if err != nil {
			panic(err)
		}
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			panic(err)
		}
		err = json.Unmarshal(body, &results)
		if err != nil {
			panic(err)
		}
		if results.Complperc >= 100 {
			debugprint("scan id %s completed", id)
			break
		}
		time.Sleep(5 * time.Second)
	}
	debugprint("getting certificate id %d", results.Cert_id)
	if !results.Has_tls && results.Cert_id < 1 {
		log.Printf("[info] target %q is not TLS enabled", results.Target)
		return
	}
	cert, err = getCert(results.Cert_id)
	if err != nil {
		panic(err)
	}
	for _, a := range r.Assertions {
		r.AssertNotBefore(a, results.Target, cert.Validity.NotBefore, notifchan)
		r.AssertNotAfter(a, results.Target, cert.Validity.NotAfter, notifchan)
		r.AssertAnalysis(a, results, cert, notifchan)
	}
	return
}

func getCert(id int64) (cert certificate.Certificate, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getCert(id=%q) -> %v", e)
		}
	}()
	resp, err := http.Get(fmt.Sprintf("%s/api/v1/certificate?id=%d", observatory, id))
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	err = json.Unmarshal(body, &cert)
	if err != nil {
		panic(err)
	}
	return
}

func debugprint(format string, a ...interface{}) {
	if debug {
		log.Printf("[debug] "+format, a...)
	}
}

// getConf first read the configuration from a local YAML file then overrides it
// with the content of the TLSOBS_RUNNER_CONF var (which much contain a full yaml file
// encoded in base64), and then overrides the SMTP settings with various SMTP env var
func getConf(cfg string) (c Configuration) {
	// load the local configuration file
	fd, err := ioutil.ReadFile(cfg)
	if err != nil {
		log.Fatal(err)
	}
	err = yaml.Unmarshal(fd, &c)
	if err != nil {
		log.Fatalf("error: %v", err)
	}
	// iterate over notifications in targets and unbase64 the values
	for i, run := range c.Runs {
		for j, rcpt := range run.Notifications.Email.Recipients {
			if len(rcpt) < 5 || rcpt[0:4] != "b64:" {
				continue
			}
			data, err := base64.StdEncoding.DecodeString(rcpt[4:])
			if err != nil {
				log.Fatalf("error while decoding b64 recipient: %v", err)
			}
			c.Runs[i].Notifications.Email.Recipients[j] = fmt.Sprintf("%s", bytes.TrimRight(data, "\n"))
		}
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_HOST") != "" {
		c.Smtp.Host = os.Getenv("TLSOBS_RUNNER_SMTP_HOST")
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_PORT") != "" {
		var err error
		c.Smtp.Port, err = strconv.Atoi(os.Getenv("TLSOBS_RUNNER_SMTP_PORT"))
		if err != nil {
			log.Printf("[error] failed to read smtp port from env variable: %v", err)
		}
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_FROM") != "" {
		c.Smtp.From = os.Getenv("TLSOBS_RUNNER_SMTP_FROM")
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_USER") != "" {
		c.Smtp.Auth.User = os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_USER")
	}
	if os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_PASS") != "" {
		c.Smtp.Auth.Pass = os.Getenv("TLSOBS_RUNNER_SMTP_AUTH_PASS")
	}
	return c
}
<|endoftext|>


<commit_before>package app_files

import (
	"crypto/sha1"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"

	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/gofileutils/fileutils"
)

type AppFiles interface {
	AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error)
	CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) (err error)
	CountFiles(directory string) int64
	WalkAppFiles(dir string, onEachFile func(string, string) error) (err error)
}

type ApplicationFiles struct{}

func (appfiles ApplicationFiles) AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error) {
	dir, err = filepath.Abs(dir)
	if err != nil {
		return
	}

	err = appfiles.WalkAppFiles(dir, func(fileName string, fullPath string) error {
		fileInfo, err := os.Lstat(fullPath)
		if err != nil {
			return err
		}

		appFile := models.AppFileFields{
			Path: filepath.ToSlash(fileName),
			Size: fileInfo.Size(),
		}

		if fileInfo.IsDir() {
			appFile.Sha1 = "0"
			appFile.Size = 0
		} else {
			hash := sha1.New()
			file, err := os.Open(fullPath)
			if err != nil {
				return err
			}
			defer file.Close()

			_, err = io.Copy(hash, file)
			if err != nil {
				return err
			}

			appFile.Sha1 = fmt.Sprintf("%x", hash.Sum(nil))
		}

		appFiles = append(appFiles, appFile)

		return nil
	})

	return
}

func (appfiles ApplicationFiles) CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) error {
	for _, file := range appFiles {
		err := func() error {
			fromPath := filepath.Join(fromDir, file.Path)
			srcFileInfo, err := os.Stat(fromPath)
			if err != nil {
				return err
			}

			toPath := filepath.Join(toDir, file.Path)

			if srcFileInfo.IsDir() {
				err = os.MkdirAll(toPath, srcFileInfo.Mode())
				if err != nil {
					return err
				}
				return nil
			}

			var dst *os.File
			dst, err = fileutils.Create(toPath)
			if err != nil {
				return err
			}
			defer dst.Close()

			dst.Chmod(srcFileInfo.Mode())

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

			_, err = io.Copy(dst, src)
			if err != nil {
				return err
			}

			return nil
		}()

		if err != nil {
			return err
		}
	}

	return nil
}

func (appfiles ApplicationFiles) CountFiles(directory string) int64 {
	var count int64
	appfiles.WalkAppFiles(directory, func(_, _ string) error {
		count++
		return nil
	})
	return count
}

func (appfiles ApplicationFiles) WalkAppFiles(dir string, onEachFile func(string, string) error) error {
	cfIgnore := loadIgnoreFile(dir)
	walkFunc := func(fullPath string, f os.FileInfo, err error) error {
		fileRelativePath, _ := filepath.Rel(dir, fullPath)
		fileRelativeUnixPath := filepath.ToSlash(fileRelativePath)

		if cfIgnore.FileShouldBeIgnored(fileRelativeUnixPath) {
			if err == nil {
				if f.IsDir() {
					return filepath.SkipDir
				}
			}

			if runtime.GOOS == "windows" {
				fi, statErr := os.Lstat(`\\?\` + fullPath)
				if statErr != nil {
					return statErr
				}

				if fi.IsDir() {
					return filepath.SkipDir
				}
			}

			return err
		}

		if err != nil {
			return err
		}

		if fullPath == dir {
			return nil
		}

		if !f.Mode().IsRegular() && !f.IsDir() {
			return nil
		}

		return onEachFile(fileRelativePath, fullPath)
	}

	return filepath.Walk(dir, walkFunc)
}

func loadIgnoreFile(dir string) CfIgnore {
	fileContents, err := ioutil.ReadFile(filepath.Join(dir, ".cfignore"))
	if err != nil {
		return NewCfIgnore("")
	}

	return NewCfIgnore(string(fileContents))
}
<commit_msg>WalkAppFiles handles paths that exceed MAX_PATH<commit_after>package app_files

import (
	"crypto/sha1"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"

	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/gofileutils/fileutils"
)

type AppFiles interface {
	AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error)
	CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) (err error)
	CountFiles(directory string) int64
	WalkAppFiles(dir string, onEachFile func(string, string) error) (err error)
}

type ApplicationFiles struct{}

func (appfiles ApplicationFiles) AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error) {
	dir, err = filepath.Abs(dir)
	if err != nil {
		return
	}

	err = appfiles.WalkAppFiles(dir, func(fileName string, fullPath string) error {
		fileInfo, err := os.Lstat(fullPath)
		if err != nil {
			return err
		}

		appFile := models.AppFileFields{
			Path: filepath.ToSlash(fileName),
			Size: fileInfo.Size(),
		}

		if fileInfo.IsDir() {
			appFile.Sha1 = "0"
			appFile.Size = 0
		} else {
			hash := sha1.New()
			file, err := os.Open(fullPath)
			if err != nil {
				return err
			}
			defer file.Close()

			_, err = io.Copy(hash, file)
			if err != nil {
				return err
			}

			appFile.Sha1 = fmt.Sprintf("%x", hash.Sum(nil))
		}

		appFiles = append(appFiles, appFile)

		return nil
	})

	return
}

func (appfiles ApplicationFiles) CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) error {
	for _, file := range appFiles {
		err := func() error {
			fromPath := filepath.Join(fromDir, file.Path)
			srcFileInfo, err := os.Stat(fromPath)
			if err != nil {
				return err
			}

			toPath := filepath.Join(toDir, file.Path)

			if srcFileInfo.IsDir() {
				err = os.MkdirAll(toPath, srcFileInfo.Mode())
				if err != nil {
					return err
				}
				return nil
			}

			var dst *os.File
			dst, err = fileutils.Create(toPath)
			if err != nil {
				return err
			}
			defer dst.Close()

			dst.Chmod(srcFileInfo.Mode())

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

			_, err = io.Copy(dst, src)
			if err != nil {
				return err
			}

			return nil
		}()

		if err != nil {
			return err
		}
	}

	return nil
}

func (appfiles ApplicationFiles) CountFiles(directory string) int64 {
	var count int64
	appfiles.WalkAppFiles(directory, func(_, _ string) error {
		count++
		return nil
	})
	return count
}

func (appfiles ApplicationFiles) WalkAppFiles(dir string, onEachFile func(string, string) error) error {
	cfIgnore := loadIgnoreFile(dir)
	walkFunc := func(fullPath string, f os.FileInfo, err error) error {
		fileRelativePath, _ := filepath.Rel(dir, fullPath)
		fileRelativeUnixPath := filepath.ToSlash(fileRelativePath)

		if err != nil && runtime.GOOS == "windows" {
			f, err = os.Lstat(`\\?\` + fullPath)
		}

		if cfIgnore.FileShouldBeIgnored(fileRelativeUnixPath) {
			if err == nil && f.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}

		if err != nil {
			return err
		}

		if fullPath == dir {
			return nil
		}

		if !f.Mode().IsRegular() && !f.IsDir() {
			return nil
		}

		return onEachFile(fileRelativePath, fullPath)
	}

	return filepath.Walk(dir, walkFunc)
}

func loadIgnoreFile(dir string) CfIgnore {
	fileContents, err := ioutil.ReadFile(filepath.Join(dir, ".cfignore"))
	if err != nil {
		return NewCfIgnore("")
	}

	return NewCfIgnore(string(fileContents))
}
<|endoftext|>


<commit_before>package lib

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/golang/glog"
	"github.com/hashicorp/go-version"
	"io/ioutil"
	"net/http"
	"strings"
)

var (
	insecure = flag.Bool("insecure", true, "Ignore server certificate if using https")
)

type BasicAuth struct {
	Username string
	Password string
}

type CouchdbClient struct {
	baseUri   string
	basicAuth BasicAuth
	databases []string
	client    *http.Client
}

type RootResponse struct {
	Couchdb string `json:"couchdb"`
	Version string `json:"version"`
}

type MembershipResponse struct {
	AllNodes     []string `json:"all_nodes"`
	ClusterNodes []string `json:"cluster_nodes"`
}

func (c *CouchdbClient) getServerVersion() (string, error) {
	data, err := c.request("GET", fmt.Sprintf("%s/", c.baseUri))
	if err != nil {
		return "", err
	}
	var root RootResponse
	err = json.Unmarshal(data, &root)
	if err != nil {
		return "", err
	}
	return root.Version, nil
}

func (c *CouchdbClient) isCouchDbV2() (bool, error) {
	clusteredCouch, err := version.NewConstraint(">= 2.0")
	if err != nil {
		return false, err
	}

	serverVersion, err := c.getServerVersion()
	if err != nil {
		return false, err
	}

	couchDbVersion, err := version.NewVersion(serverVersion)
	if err != nil {
		return false, err
	}

	glog.Infof("relaxing on couch@%s", couchDbVersion)
	//fmt.Printf("relaxing on couch@%s\n", couchDbVersion)
	return clusteredCouch.Check(couchDbVersion), nil
}

func (c *CouchdbClient) getNodeNames() ([]string, error) {
	data, err := c.request("GET", fmt.Sprintf("%s/_membership", c.baseUri))
	if err != nil {
		return nil, err
	}
	var membership MembershipResponse
	err = json.Unmarshal(data, &membership)
	if err != nil {
		return nil, err
	}
	for i, name := range membership.ClusterNodes {
		glog.Infof("node[%d]: %s\n", i, name)
	}
	return membership.ClusterNodes, nil
}

func (c *CouchdbClient) getNodeBaseUrisByNodeName(baseUri string) (map[string]string, error) {
	names, err := c.getNodeNames()
	if err != nil {
		return nil, err
	}
	urisByNodeName := make(map[string]string)
	for _, name := range names {
		urisByNodeName[name] = fmt.Sprintf("%s/_node/%s", baseUri, name)
	}
	return urisByNodeName, nil
}

func (c *CouchdbClient) getStatsByNodeName(urisByNodeName map[string]string) (map[string]StatsResponse, error) {
	statsByNodeName := make(map[string]StatsResponse)
	for name, uri := range urisByNodeName {
		data, err := c.request("GET", fmt.Sprintf("%s/_stats", uri))
		if err != nil {
			err = fmt.Errorf("error reading couchdb stats: %v", err)
			if !strings.Contains(err.Error(), "\"error\":\"nodedown\"") {
				return nil, err
			}

			glog.Error(fmt.Errorf("continuing despite error: %v", err))
			continue
		}

		var stats StatsResponse
		err = json.Unmarshal(data, &stats)
		if err != nil {
			return nil, fmt.Errorf("error unmarshalling stats: %v", err)
		}
		statsByNodeName[name] = stats
	}
	return statsByNodeName, nil
}

func (c *CouchdbClient) getStats() (Stats, error) {
	isCouchDbV2, err := c.isCouchDbV2()
	if err != nil {
		return Stats{}, err
	}
	if isCouchDbV2 {
		urisByNode, err := c.getNodeBaseUrisByNodeName(c.baseUri)
		if err != nil {
			return Stats{}, err
		}
		nodeStats, err := c.getStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		databaseStats, err := c.getDatabasesStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		activeTasks, err := c.getActiveTasks()
		if err != nil {
			return Stats{}, err
		}
		return Stats{
			StatsByNodeName:       nodeStats,
			DatabaseStatsByDbName: databaseStats,
			ActiveTasksResponse:   activeTasks,
			ApiVersion:            "2"}, nil
	} else {
		urisByNode := map[string]string{
			"master": c.baseUri,
		}
		nodeStats, err := c.getStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		databaseStats, err := c.getDatabasesStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		activeTasks, err := c.getActiveTasks()
		if err != nil {
			return Stats{}, err
		}
		return Stats{
			StatsByNodeName:       nodeStats,
			DatabaseStatsByDbName: databaseStats,
			ActiveTasksResponse:   activeTasks,
			ApiVersion:            "1"}, nil
	}
}

func (c *CouchdbClient) getDatabasesStatsByNodeName(urisByNodeName map[string]string) (map[string]DatabaseStats, error) {
	dbStatsByDbName := make(map[string]DatabaseStats)
	for _, dbName := range c.databases {
		data, err := c.request("GET", fmt.Sprintf("%s/%s", c.baseUri, dbName))
		if err != nil {
			return nil, fmt.Errorf("error reading database '%s' stats: %v", dbName, err)
		}

		var dbStats DatabaseStats
		err = json.Unmarshal(data, &dbStats)
		if err != nil {
			return nil, fmt.Errorf("error unmarshalling database '%s' stats: %v", dbName, err)
		}
		dbStats.DiskSizeOverhead = dbStats.DiskSize - dbStats.DataSize
		dbStatsByDbName[dbName] = dbStats
	}
	return dbStatsByDbName, nil
}

func (c *CouchdbClient) getActiveTasks() (ActiveTasksResponse, error) {
	data, err := c.request("GET", fmt.Sprintf("%s/_active_tasks", c.baseUri))
	if err != nil {
		return ActiveTasksResponse{}, fmt.Errorf("error reading active tasks: %v", err)
	}

	var activeTasks ActiveTasksResponse
	err = json.Unmarshal(data, &activeTasks)
	if err != nil {
		return ActiveTasksResponse{}, fmt.Errorf("error unmarshalling active tasks: %v", err)
	}
	for _, activeTask := range activeTasks {
		// CouchDB 1.x doesn't know anything about nodes.
		if activeTask.Node == "" {
			activeTask.Node = "master"
		}
	}
	return activeTasks, nil
}

func (c *CouchdbClient) request(method string, uri string) (respData []byte, err error) {
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, err
	}
	if len(c.basicAuth.Username) > 0 {
		req.SetBasicAuth(c.basicAuth.Username, c.basicAuth.Password)
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	respData, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		if err != nil {
			respData = []byte(err.Error())
		}
		return nil, fmt.Errorf("status %s (%d): %s", resp.Status, resp.StatusCode, respData)
	}

	return respData, nil
}

func NewCouchdbClient(uri string, basicAuth BasicAuth, databases []string) *CouchdbClient {
	httpClient := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: *insecure},
		},
	}

	return &CouchdbClient{
		baseUri:   uri,
		basicAuth: basicAuth,
		databases: databases,
		client:    httpClient,
	}
}
<commit_msg>Handle case where all nodes are down<commit_after>package lib

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/golang/glog"
	"github.com/hashicorp/go-version"
	"io/ioutil"
	"net/http"
	"strings"
)

var (
	insecure = flag.Bool("insecure", true, "Ignore server certificate if using https")
)

type BasicAuth struct {
	Username string
	Password string
}

type CouchdbClient struct {
	baseUri   string
	basicAuth BasicAuth
	databases []string
	client    *http.Client
}

type RootResponse struct {
	Couchdb string `json:"couchdb"`
	Version string `json:"version"`
}

type MembershipResponse struct {
	AllNodes     []string `json:"all_nodes"`
	ClusterNodes []string `json:"cluster_nodes"`
}

func (c *CouchdbClient) getServerVersion() (string, error) {
	data, err := c.request("GET", fmt.Sprintf("%s/", c.baseUri))
	if err != nil {
		return "", err
	}
	var root RootResponse
	err = json.Unmarshal(data, &root)
	if err != nil {
		return "", err
	}
	return root.Version, nil
}

func (c *CouchdbClient) isCouchDbV2() (bool, error) {
	clusteredCouch, err := version.NewConstraint(">= 2.0")
	if err != nil {
		return false, err
	}

	serverVersion, err := c.getServerVersion()
	if err != nil {
		return false, err
	}

	couchDbVersion, err := version.NewVersion(serverVersion)
	if err != nil {
		return false, err
	}

	glog.Infof("relaxing on couch@%s", couchDbVersion)
	//fmt.Printf("relaxing on couch@%s\n", couchDbVersion)
	return clusteredCouch.Check(couchDbVersion), nil
}

func (c *CouchdbClient) getNodeNames() ([]string, error) {
	data, err := c.request("GET", fmt.Sprintf("%s/_membership", c.baseUri))
	if err != nil {
		return nil, err
	}
	var membership MembershipResponse
	err = json.Unmarshal(data, &membership)
	if err != nil {
		return nil, err
	}
	for i, name := range membership.ClusterNodes {
		glog.Infof("node[%d]: %s\n", i, name)
	}
	return membership.ClusterNodes, nil
}

func (c *CouchdbClient) getNodeBaseUrisByNodeName(baseUri string) (map[string]string, error) {
	names, err := c.getNodeNames()
	if err != nil {
		return nil, err
	}
	urisByNodeName := make(map[string]string)
	for _, name := range names {
		urisByNodeName[name] = fmt.Sprintf("%s/_node/%s", baseUri, name)
	}
	return urisByNodeName, nil
}

func (c *CouchdbClient) getStatsByNodeName(urisByNodeName map[string]string) (map[string]StatsResponse, error) {
	statsByNodeName := make(map[string]StatsResponse)
	for name, uri := range urisByNodeName {
		data, err := c.request("GET", fmt.Sprintf("%s/_stats", uri))
		if err != nil {
			err = fmt.Errorf("error reading couchdb stats: %v", err)
			if !strings.Contains(err.Error(), "\"error\":\"nodedown\"") {
				return nil, err
			}

			delete(urisByNodeName, name)
			glog.Error(fmt.Errorf("continuing despite error: %v", err))
			continue
		}

		var stats StatsResponse
		err = json.Unmarshal(data, &stats)
		if err != nil {
			return nil, fmt.Errorf("error unmarshalling stats: %v", err)
		}
		statsByNodeName[name] = stats
	}

	if len(urisByNodeName) == 0 {
		return nil, fmt.Errorf("all nodes down")
	}

	return statsByNodeName, nil
}

func (c *CouchdbClient) getStats() (Stats, error) {
	isCouchDbV2, err := c.isCouchDbV2()
	if err != nil {
		return Stats{}, err
	}
	if isCouchDbV2 {
		urisByNode, err := c.getNodeBaseUrisByNodeName(c.baseUri)
		if err != nil {
			return Stats{}, err
		}
		nodeStats, err := c.getStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		databaseStats, err := c.getDatabasesStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		activeTasks, err := c.getActiveTasks()
		if err != nil {
			return Stats{}, err
		}
		return Stats{
			StatsByNodeName:       nodeStats,
			DatabaseStatsByDbName: databaseStats,
			ActiveTasksResponse:   activeTasks,
			ApiVersion:            "2"}, nil
	} else {
		urisByNode := map[string]string{
			"master": c.baseUri,
		}
		nodeStats, err := c.getStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		databaseStats, err := c.getDatabasesStatsByNodeName(urisByNode)
		if err != nil {
			return Stats{}, err
		}
		activeTasks, err := c.getActiveTasks()
		if err != nil {
			return Stats{}, err
		}
		return Stats{
			StatsByNodeName:       nodeStats,
			DatabaseStatsByDbName: databaseStats,
			ActiveTasksResponse:   activeTasks,
			ApiVersion:            "1"}, nil
	}
}

func (c *CouchdbClient) getDatabasesStatsByNodeName(urisByNodeName map[string]string) (map[string]DatabaseStats, error) {
	dbStatsByDbName := make(map[string]DatabaseStats)
	for _, dbName := range c.databases {
		data, err := c.request("GET", fmt.Sprintf("%s/%s", c.baseUri, dbName))
		if err != nil {
			return nil, fmt.Errorf("error reading database '%s' stats: %v", dbName, err)
		}

		var dbStats DatabaseStats
		err = json.Unmarshal(data, &dbStats)
		if err != nil {
			return nil, fmt.Errorf("error unmarshalling database '%s' stats: %v", dbName, err)
		}
		dbStats.DiskSizeOverhead = dbStats.DiskSize - dbStats.DataSize
		dbStatsByDbName[dbName] = dbStats
	}
	return dbStatsByDbName, nil
}

func (c *CouchdbClient) getActiveTasks() (ActiveTasksResponse, error) {
	data, err := c.request("GET", fmt.Sprintf("%s/_active_tasks", c.baseUri))
	if err != nil {
		return ActiveTasksResponse{}, fmt.Errorf("error reading active tasks: %v", err)
	}

	var activeTasks ActiveTasksResponse
	err = json.Unmarshal(data, &activeTasks)
	if err != nil {
		return ActiveTasksResponse{}, fmt.Errorf("error unmarshalling active tasks: %v", err)
	}
	for _, activeTask := range activeTasks {
		// CouchDB 1.x doesn't know anything about nodes.
		if activeTask.Node == "" {
			activeTask.Node = "master"
		}
	}
	return activeTasks, nil
}

func (c *CouchdbClient) request(method string, uri string) (respData []byte, err error) {
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, err
	}
	if len(c.basicAuth.Username) > 0 {
		req.SetBasicAuth(c.basicAuth.Username, c.basicAuth.Password)
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	respData, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		if err != nil {
			respData = []byte(err.Error())
		}
		return nil, fmt.Errorf("status %s (%d): %s", resp.Status, resp.StatusCode, respData)
	}

	return respData, nil
}

func NewCouchdbClient(uri string, basicAuth BasicAuth, databases []string) *CouchdbClient {
	httpClient := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: *insecure},
		},
	}

	return &CouchdbClient{
		baseUri:   uri,
		basicAuth: basicAuth,
		databases: databases,
		client:    httpClient,
	}
}
<|endoftext|>


<commit_before>package goose_test

import (
    "testing"
    "github.com/jkrobles/goose"
)

func TestSplit(t *testing.T) {
    t.Log("Starting Split Test")
    g := new(goose.Goose)
    strs := g.Split("we are good", " ")
    test := g.Split("we don't like you"," ")
    if  test[0] == strs[0] {
        t.Log("Correct answer")
    } else {
        t.Log("Arrays don't match ", strs, test)
        t.Fail()
    }
}
<commit_msg>Update client test<commit_after>package goose_test

import (
    "testing"
    "github.com/jkrobles/goose"
)

func TestSplit(t *testing.T) {
    t.Log("Starting Split Test")
    g := new(goose.Goose)
    strs := g.Split("we are good", " ")
    var test []string = []string{"we","are","good"}
    if  test[0] == strs[0] && test[1] == strs[1] {
        t.Log("Correct answer")
    } else {
        t.Log("Arrays don't match ", strs, test)
        t.Fail()
    }
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/github/git-lfs/vendor/_nuts/github.com/cheggaaa/pb"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
	contentaddressable "github.com/github/git-lfs/vendor/_nuts/github.com/technoweenie/go-contentaddressable"
)

func PointerSmudgeToFile(filename string, ptr *Pointer, cb CopyCallback) error {
	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("Could not create working directory file: %v", err)
	}
	defer file.Close()
	if err := PointerSmudge(file, ptr, filename, cb); err != nil {
		return fmt.Errorf("Could not write working directory file: %v", err)
	}
	return nil
}
func PointerSmudge(writer io.Writer, ptr *Pointer, workingfile string, cb CopyCallback) error {
	mediafile, err := LocalMediaPath(ptr.Oid)
	if err != nil {
		return err
	}

	stat, statErr := os.Stat(mediafile)
	if statErr == nil && stat != nil {
		fileSize := stat.Size()
		if fileSize == 0 || fileSize != ptr.Size {
			tracerx.Printf("Removing %s, size %d is invalid", mediafile, fileSize)
			os.RemoveAll(mediafile)
			stat = nil
		}
	}

	var wErr *WrappedError
	if statErr != nil || stat == nil {
		wErr = downloadFile(writer, ptr, workingfile, mediafile, cb)
	} else {
		sendApiEvent(apiEventSuccess)
		wErr = readLocalFile(writer, ptr, mediafile, cb)
	}

	if wErr != nil {
		return &SmudgeError{ptr.Oid, mediafile, wErr}
	}

	return nil
}

// PointerSmudgeObject uses a Pointer and objectResource to download the object to the
// media directory. It does not write the file to the working directory.
func PointerSmudgeObject(ptr *Pointer, obj *objectResource, cb CopyCallback) error {
	mediafile, err := LocalMediaPath(obj.Oid)
	if err != nil {
		return err
	}

	stat, statErr := os.Stat(mediafile)
	if statErr == nil && stat != nil {
		fileSize := stat.Size()
		if fileSize == 0 || fileSize != obj.Size {
			tracerx.Printf("Removing %s, size %d is invalid", mediafile, fileSize)
			os.RemoveAll(mediafile)
			stat = nil
		}
	}

	if statErr != nil || stat == nil {
		wErr := downloadObject(ptr, obj, mediafile, cb)

		if wErr != nil {
			sendApiEvent(apiEventFail)
			return &SmudgeError{obj.Oid, mediafile, wErr}
		}
	}

	sendApiEvent(apiEventSuccess)
	return nil
}

func downloadObject(ptr *Pointer, obj *objectResource, mediafile string, cb CopyCallback) *WrappedError {
	reader, size, wErr := DownloadObject(obj)
	if reader != nil {
		defer reader.Close()
	}

	// TODO this can be unified with the same code in downloadFile
	if wErr != nil {
		wErr.Errorf("Error downloading %s.", mediafile)
		return wErr
	}

	if ptr.Size == 0 {
		ptr.Size = size
	}

	mediaFile, err := contentaddressable.NewFile(mediafile)
	if err != nil {
		return Errorf(err, "Error opening media file buffer.")
	}

	_, err = CopyWithCallback(mediaFile, reader, ptr.Size, cb)
	if err == nil {
		err = mediaFile.Accept()
	}
	mediaFile.Close()

	if err != nil {
		return Errorf(err, "Error buffering media file.")
	}

	return nil
}

func downloadFile(writer io.Writer, ptr *Pointer, workingfile, mediafile string, cb CopyCallback) *WrappedError {
	fmt.Fprintf(os.Stderr, "Downloading %s (%s)\n", workingfile, pb.FormatBytes(ptr.Size))
	reader, size, wErr := Download(filepath.Base(mediafile))
	if reader != nil {
		defer reader.Close()
	}

	if wErr != nil {
		wErr.Errorf("Error downloading %s.", mediafile)
		return wErr
	}

	if ptr.Size == 0 {
		ptr.Size = size
	}

	mediaFile, err := contentaddressable.NewFile(mediafile)
	if err != nil {
		return Errorf(err, "Error opening media file buffer.")
	}

	_, err = CopyWithCallback(mediaFile, reader, ptr.Size, cb)
	if err == nil {
		err = mediaFile.Accept()
	}
	mediaFile.Close()

	if err != nil {
		return Errorf(err, "Error buffering media file.")
	}

	return readLocalFile(writer, ptr, mediafile, nil)
}

func readLocalFile(writer io.Writer, ptr *Pointer, mediafile string, cb CopyCallback) *WrappedError {
	reader, err := os.Open(mediafile)
	if err != nil {
		return Errorf(err, "Error opening media file.")
	}
	defer reader.Close()

	if ptr.Size == 0 {
		if stat, _ := os.Stat(mediafile); stat != nil {
			ptr.Size = stat.Size()
		}
	}

	_, err = CopyWithCallback(writer, reader, ptr.Size, cb)
	return Errorf(err, "Error reading from media file.")
}

type SmudgeError struct {
	Oid      string
	Filename string
	*WrappedError
}
<commit_msg>アー アアアア アーアー<commit_after>package lfs

import (
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/github/git-lfs/vendor/_nuts/github.com/cheggaaa/pb"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
	contentaddressable "github.com/github/git-lfs/vendor/_nuts/github.com/technoweenie/go-contentaddressable"
)

func PointerSmudgeToFile(filename string, ptr *Pointer, cb CopyCallback) error {
	os.MkdirAll(filepath.Dir(filename), 0755)
	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("Could not create working directory file: %v", err)
	}
	defer file.Close()
	if err := PointerSmudge(file, ptr, filename, cb); err != nil {
		return fmt.Errorf("Could not write working directory file: %v", err)
	}
	return nil
}
func PointerSmudge(writer io.Writer, ptr *Pointer, workingfile string, cb CopyCallback) error {
	mediafile, err := LocalMediaPath(ptr.Oid)
	if err != nil {
		return err
	}

	stat, statErr := os.Stat(mediafile)
	if statErr == nil && stat != nil {
		fileSize := stat.Size()
		if fileSize == 0 || fileSize != ptr.Size {
			tracerx.Printf("Removing %s, size %d is invalid", mediafile, fileSize)
			os.RemoveAll(mediafile)
			stat = nil
		}
	}

	var wErr *WrappedError
	if statErr != nil || stat == nil {
		wErr = downloadFile(writer, ptr, workingfile, mediafile, cb)
	} else {
		sendApiEvent(apiEventSuccess)
		wErr = readLocalFile(writer, ptr, mediafile, cb)
	}

	if wErr != nil {
		return &SmudgeError{ptr.Oid, mediafile, wErr}
	}

	return nil
}

// PointerSmudgeObject uses a Pointer and objectResource to download the object to the
// media directory. It does not write the file to the working directory.
func PointerSmudgeObject(ptr *Pointer, obj *objectResource, cb CopyCallback) error {
	mediafile, err := LocalMediaPath(obj.Oid)
	if err != nil {
		return err
	}

	stat, statErr := os.Stat(mediafile)
	if statErr == nil && stat != nil {
		fileSize := stat.Size()
		if fileSize == 0 || fileSize != obj.Size {
			tracerx.Printf("Removing %s, size %d is invalid", mediafile, fileSize)
			os.RemoveAll(mediafile)
			stat = nil
		}
	}

	if statErr != nil || stat == nil {
		wErr := downloadObject(ptr, obj, mediafile, cb)

		if wErr != nil {
			sendApiEvent(apiEventFail)
			return &SmudgeError{obj.Oid, mediafile, wErr}
		}
	}

	sendApiEvent(apiEventSuccess)
	return nil
}

func downloadObject(ptr *Pointer, obj *objectResource, mediafile string, cb CopyCallback) *WrappedError {
	reader, size, wErr := DownloadObject(obj)
	if reader != nil {
		defer reader.Close()
	}

	// TODO this can be unified with the same code in downloadFile
	if wErr != nil {
		wErr.Errorf("Error downloading %s.", mediafile)
		return wErr
	}

	if ptr.Size == 0 {
		ptr.Size = size
	}

	mediaFile, err := contentaddressable.NewFile(mediafile)
	if err != nil {
		return Errorf(err, "Error opening media file buffer.")
	}

	_, err = CopyWithCallback(mediaFile, reader, ptr.Size, cb)
	if err == nil {
		err = mediaFile.Accept()
	}
	mediaFile.Close()

	if err != nil {
		return Errorf(err, "Error buffering media file.")
	}

	return nil
}

func downloadFile(writer io.Writer, ptr *Pointer, workingfile, mediafile string, cb CopyCallback) *WrappedError {
	fmt.Fprintf(os.Stderr, "Downloading %s (%s)\n", workingfile, pb.FormatBytes(ptr.Size))
	reader, size, wErr := Download(filepath.Base(mediafile))
	if reader != nil {
		defer reader.Close()
	}

	if wErr != nil {
		wErr.Errorf("Error downloading %s.", mediafile)
		return wErr
	}

	if ptr.Size == 0 {
		ptr.Size = size
	}

	mediaFile, err := contentaddressable.NewFile(mediafile)
	if err != nil {
		return Errorf(err, "Error opening media file buffer.")
	}

	_, err = CopyWithCallback(mediaFile, reader, ptr.Size, cb)
	if err == nil {
		err = mediaFile.Accept()
	}
	mediaFile.Close()

	if err != nil {
		return Errorf(err, "Error buffering media file.")
	}

	return readLocalFile(writer, ptr, mediafile, nil)
}

func readLocalFile(writer io.Writer, ptr *Pointer, mediafile string, cb CopyCallback) *WrappedError {
	reader, err := os.Open(mediafile)
	if err != nil {
		return Errorf(err, "Error opening media file.")
	}
	defer reader.Close()

	if ptr.Size == 0 {
		if stat, _ := os.Stat(mediafile); stat != nil {
			ptr.Size = stat.Size()
		}
	}

	_, err = CopyWithCallback(writer, reader, ptr.Size, cb)
	return Errorf(err, "Error reading from media file.")
}

type SmudgeError struct {
	Oid      string
	Filename string
	*WrappedError
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/cheggaaa/pb"
	"github.com/rubyist/tracerx"
)

const (
	batchSize = 100
)

type Transferable interface {
	Check() (*objectResource, *WrappedError)
	Transfer(CopyCallback) *WrappedError
	Object() *objectResource
	Oid() string
	Size() int64
	Name() string
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	filesAdded    int32 // Number of files added to the queue
	filesFinished int32 // Number of files that have finished transfering
	transferIdx   int32 // Used to track transfer number for progress monitor
	workers       int   // Number of transfer workers to spawn
	transferKind  string
	errors        []*WrappedError
	transferables map[string]Transferable
	bar           *pb.ProgressBar
	batcher       *Batcher
	apic          chan Transferable  // Channel for processing individual API requests
	transferc     chan Transferable  // Channel for processing transfers
	errorc        chan *WrappedError // Channel for processing errors
	progressc     chan string        // Channel for GIT_LFS_PROGRESS monitor
	watchers      []chan string
	wait          sync.WaitGroup
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(workers int) *TransferQueue {
	q := &TransferQueue{
		apic:          make(chan Transferable, batchSize),
		transferc:     make(chan Transferable, batchSize),
		errorc:        make(chan *WrappedError),
		watchers:      make([]chan string, 0),
		progressc:     make(chan string, batchSize),
		workers:       workers,
		transferables: make(map[string]Transferable),
	}

	q.run()

	return q
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	atomic.AddInt32(&q.filesAdded, 1)

	// Sneak in and update the progress bar's total size
	atomic.AddInt64(&q.bar.Total, t.Size())

	q.wait.Add(1)
	q.transferables[t.Oid()] = t

	if q.batcher != nil {
		q.batcher.Add(t)
		return
	}

	q.apic <- t
}

// Wait waits for the queue to finish processing all transfers
func (q *TransferQueue) Wait() {

	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.filesFinished, q.filesAdded))

	if q.batcher != nil {
		q.batcher.Exit()
	}
	q.wait.Wait()
	close(q.apic)
	close(q.transferc)
	close(q.errorc)
	for _, watcher := range q.watchers {
		close(watcher)
	}
	close(q.progressc)
	q.bar.Finish()
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, batchSize)
	q.watchers = append(q.watchers, c)
	return c
}

// individualApiRoutine processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) individualApiRoutine(apiWaiter chan interface{}) {
	for t := range q.apic {
		obj, err := t.Check()
		if err != nil {
			q.wait.Done()
			q.errorc <- err
			continue
		}

		if apiWaiter != nil { // Signal to launch more individual api workers
			select {
			case apiWaiter <- 1:
			default:
			}
		}

		if obj != nil {
			t.SetObject(obj)
			q.transferc <- t
		}
	}
}

// legacyFallback is used when a batch request is made to a server that does
// not support the batch endpoint. When this happens, the Transferables are
// feed from the batcher into apic to be processed individually.
func (q *TransferQueue) legacyFallback(failedBatch []Transferable) {
	tracerx.Printf("tq: batch api not implemented, falling back to individual")

	q.launchIndividualApiRoutines()

	for _, t := range failedBatch {
		q.apic <- t
	}

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		for _, t := range batch {
			q.apic <- t
		}
	}
}

// batchApiRoutine processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) batchApiRoutine() {
	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		transfers := make([]*objectResource, 0, len(batch))
		for _, t := range batch {
			transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
		}

		objects, err := Batch(transfers, q.transferKind)
		if err != nil {
			if isNotImplError(err) {
				configFile := filepath.Join(LocalGitDir, "config")
				git.Config.SetLocal(configFile, "lfs.batch", "false")

				go q.legacyFallback(batch)
				return
			}
		}

		for _, o := range objects {
			if _, ok := o.Links[q.transferKind]; ok {
				// This object needs to be transfered
				if transfer, ok := q.transferables[o.Oid]; ok {
					transfer.SetObject(o)
					q.transferc <- transfer
				}
			}
		}
	}
}

// This goroutine collects errors returned from transfers
func (q *TransferQueue) errorCollector() {
	for err := range q.errorc {
		q.errors = append(q.errors, err)
	}
}

func (q *TransferQueue) progressMonitor() {
	output, err := newProgressLogger()
	if err != nil {
		q.errorc <- Error(err)
	}

	for l := range q.progressc {
		if err := output.Write([]byte(l)); err != nil {
			q.errorc <- Error(err)
			output.Shutdown()
		}
	}

	output.Close()
}

func (q *TransferQueue) transferWorker() {
	direction := "push"
	if q.transferKind == "download" {
		direction = "pull"
	}

	for transfer := range q.transferc {
		c := atomic.AddInt32(&q.transferIdx, 1)
		cb := func(total, read int64, current int) error {
			q.progressc <- fmt.Sprintf("%s %d/%d %d/%d %s\n", direction, c, q.filesAdded, read, total, transfer.Name())
			q.bar.Add(current)
			return nil
		}

		if err := transfer.Transfer(cb); err != nil {
			q.errorc <- err
		} else {
			oid := transfer.Oid()
			for _, c := range q.watchers {
				c <- oid
			}
		}

		f := atomic.AddInt32(&q.filesFinished, 1)
		q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", f, q.filesAdded))
		q.wait.Done()
	}
}

// launchIndividualApiRoutines first launches a single api worker. When it
// receives the first successful api request it launches workers - 1 more
// workers. This prevents being prompted for credentials multiple times at once
// when they're needed.
func (q *TransferQueue) launchIndividualApiRoutines() {
	go func() {
		apiWaiter := make(chan interface{})
		go q.individualApiRoutine(apiWaiter)

		<-apiWaiter

		for i := 0; i < q.workers-1; i++ {
			go q.individualApiRoutine(nil)
		}
	}()
}

// run starts the transfer queue and displays a progress bar. Process will
// do individual or batch transfers depending on the Config.BatchTransfer() value.
// Process will transfer files sequentially or concurrently depending on the
// Concig.ConcurrentTransfers() value.
func (q *TransferQueue) run() {
	// Set up the pb progress bar. The total size will be updated as files are
	// added to the queue.
	q.bar = pb.New64(0)
	q.bar.SetUnits(pb.U_BYTES)
	q.bar.ShowBar = false
	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.filesFinished, q.filesAdded))
	q.bar.Start()

	go q.errorCollector()
	go q.progressMonitor()

	for i := 0; i < q.workers; i++ {
		go q.transferWorker()
	}

	if Config.BatchTransfer() {
		q.batcher = NewBatcher(batchSize)
		go q.batchApiRoutine()
	} else {
		q.launchIndividualApiRoutines()
	}
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []*WrappedError {
	return q.errors
}

// progressLogger provides a wrapper around an os.File that can either
// write to the file or ignore all writes completely.
type progressLogger struct {
	writeData bool
	log       *os.File
}

// Write will write to the file and perform a Sync() if writing succeeds.
func (l *progressLogger) Write(b []byte) error {
	if l.writeData {
		if _, err := l.log.Write(b); err != nil {
			return err
		}
		return l.log.Sync()
	}
	return nil
}

// Close will call Close() on the underlying file
func (l *progressLogger) Close() error {
	if l.log != nil {
		return l.log.Close()
	}
	return nil
}

// Shutdown will cause the logger to ignore any further writes. It should
// be used when writing causes an error.
func (l *progressLogger) Shutdown() {
	l.writeData = false
}

// newProgressLogger creates a progressLogger based on the presence of
// the GIT_LFS_PROGRESS environment variable. If it is present and a log file
// is able to be created, the logger will write to the file. If it is absent,
// or there is an err creating the file, the logger will ignore all writes.
func newProgressLogger() (*progressLogger, error) {
	logPath := Config.Getenv("GIT_LFS_PROGRESS")

	if len(logPath) == 0 {
		return &progressLogger{}, nil
	}
	if !filepath.IsAbs(logPath) {
		return &progressLogger{}, fmt.Errorf("GIT_LFS_PROGRESS must be an absolute path")
	}

	cbDir := filepath.Dir(logPath)
	if err := os.MkdirAll(cbDir, 0755); err != nil {
		return &progressLogger{}, err
	}

	file, err := os.OpenFile(logPath, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		return &progressLogger{}, err
	}

	return &progressLogger{true, file}, nil
}
<commit_msg>アー アアアア アーアー<commit_after>package lfs

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/cheggaaa/pb"
	"github.com/rubyist/tracerx"
)

const (
	batchSize = 100
)

type Transferable interface {
	Check() (*objectResource, *WrappedError)
	Transfer(CopyCallback) *WrappedError
	Object() *objectResource
	Oid() string
	Size() int64
	Name() string
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	filesAdded    int32 // Number of files added to the queue
	filesFinished int32 // Number of files that have finished transfering
	transferIdx   int32 // Used to track transfer number for progress monitor
	workers       int   // Number of transfer workers to spawn
	transferKind  string
	errors        []*WrappedError
	transferables map[string]Transferable
	bar           *pb.ProgressBar
	batcher       *Batcher
	apic          chan Transferable  // Channel for processing individual API requests
	transferc     chan Transferable  // Channel for processing transfers
	errorc        chan *WrappedError // Channel for processing errors
	progressc     chan string        // Channel for GIT_LFS_PROGRESS monitor
	watchers      []chan string
	wait          sync.WaitGroup
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(workers int) *TransferQueue {
	q := &TransferQueue{
		apic:          make(chan Transferable, batchSize),
		transferc:     make(chan Transferable, batchSize),
		errorc:        make(chan *WrappedError),
		watchers:      make([]chan string, 0),
		progressc:     make(chan string, batchSize),
		workers:       workers,
		transferables: make(map[string]Transferable),
	}

	q.run()

	return q
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	atomic.AddInt32(&q.filesAdded, 1)

	// Sneak in and update the progress bar's total size
	atomic.AddInt64(&q.bar.Total, t.Size())

	q.wait.Add(1)
	q.transferables[t.Oid()] = t

	if q.batcher != nil {
		q.batcher.Add(t)
		return
	}

	q.apic <- t
}

// Wait waits for the queue to finish processing all transfers
func (q *TransferQueue) Wait() {

	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.filesFinished, q.filesAdded))

	if q.batcher != nil {
		q.batcher.Exit()
	}
	q.wait.Wait()
	close(q.apic)
	close(q.transferc)
	close(q.errorc)
	for _, watcher := range q.watchers {
		close(watcher)
	}
	close(q.progressc)
	q.bar.Finish()
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, batchSize)
	q.watchers = append(q.watchers, c)
	return c
}

// individualApiRoutine processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) individualApiRoutine(apiWaiter chan interface{}) {
	for t := range q.apic {
		obj, err := t.Check()
		if err != nil {
			q.wait.Done()
			q.errorc <- err
			continue
		}

		if apiWaiter != nil { // Signal to launch more individual api workers
			select {
			case apiWaiter <- 1:
			default:
			}
		}

		if obj != nil {
			t.SetObject(obj)
			q.transferc <- t
		}
	}
}

// legacyFallback is used when a batch request is made to a server that does
// not support the batch endpoint. When this happens, the Transferables are
// feed from the batcher into apic to be processed individually.
func (q *TransferQueue) legacyFallback(failedBatch []Transferable) {
	tracerx.Printf("tq: batch api not implemented, falling back to individual")

	q.launchIndividualApiRoutines()

	for _, t := range failedBatch {
		q.apic <- t
	}

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		for _, t := range batch {
			q.apic <- t
		}
	}
}

// batchApiRoutine processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) batchApiRoutine() {
	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		tracerx.Printf("tq: sending batch of size %d", len(batch))

		transfers := make([]*objectResource, 0, len(batch))
		for _, t := range batch {
			transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
		}

		objects, err := Batch(transfers, q.transferKind)
		if err != nil {
			if isNotImplError(err) {
				configFile := filepath.Join(LocalGitDir, "config")
				git.Config.SetLocal(configFile, "lfs.batch", "false")

				go q.legacyFallback(batch)
				return
			}
		}

		for _, o := range objects {
			if _, ok := o.Links[q.transferKind]; ok {
				// This object needs to be transfered
				if transfer, ok := q.transferables[o.Oid]; ok {
					transfer.SetObject(o)
					q.transferc <- transfer
				}
			}
		}
	}
}

// This goroutine collects errors returned from transfers
func (q *TransferQueue) errorCollector() {
	for err := range q.errorc {
		q.errors = append(q.errors, err)
	}
}

func (q *TransferQueue) progressMonitor() {
	output, err := newProgressLogger()
	if err != nil {
		q.errorc <- Error(err)
	}

	for l := range q.progressc {
		if err := output.Write([]byte(l)); err != nil {
			q.errorc <- Error(err)
			output.Shutdown()
		}
	}

	output.Close()
}

func (q *TransferQueue) transferWorker() {
	direction := "push"
	if q.transferKind == "download" {
		direction = "pull"
	}

	for transfer := range q.transferc {
		c := atomic.AddInt32(&q.transferIdx, 1)
		cb := func(total, read int64, current int) error {
			q.progressc <- fmt.Sprintf("%s %d/%d %d/%d %s\n", direction, c, q.filesAdded, read, total, transfer.Name())
			q.bar.Add(current)
			return nil
		}

		if err := transfer.Transfer(cb); err != nil {
			q.errorc <- err
		} else {
			oid := transfer.Oid()
			for _, c := range q.watchers {
				c <- oid
			}
		}

		f := atomic.AddInt32(&q.filesFinished, 1)
		q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", f, q.filesAdded))
		q.wait.Done()
	}
}

// launchIndividualApiRoutines first launches a single api worker. When it
// receives the first successful api request it launches workers - 1 more
// workers. This prevents being prompted for credentials multiple times at once
// when they're needed.
func (q *TransferQueue) launchIndividualApiRoutines() {
	go func() {
		apiWaiter := make(chan interface{})
		go q.individualApiRoutine(apiWaiter)

		<-apiWaiter

		for i := 0; i < q.workers-1; i++ {
			go q.individualApiRoutine(nil)
		}
	}()
}

// run starts the transfer queue and displays a progress bar. Process will
// do individual or batch transfers depending on the Config.BatchTransfer() value.
// Process will transfer files sequentially or concurrently depending on the
// Concig.ConcurrentTransfers() value.
func (q *TransferQueue) run() {
	// Set up the pb progress bar. The total size will be updated as files are
	// added to the queue.
	q.bar = pb.New64(0)
	q.bar.SetUnits(pb.U_BYTES)
	q.bar.ShowBar = false
	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.filesFinished, q.filesAdded))
	q.bar.Start()

	go q.errorCollector()
	go q.progressMonitor()

	tracerx.Printf("tq: starting %d transfer workers", q.workers)
	for i := 0; i < q.workers; i++ {
		go q.transferWorker()
	}

	if Config.BatchTransfer() {
		tracerx.Printf("tq: running as batched queue, batch size of %d", batchSize)
		q.batcher = NewBatcher(batchSize)
		go q.batchApiRoutine()
	} else {
		tracerx.Printf("tq: running as individual queue")
		q.launchIndividualApiRoutines()
	}
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []*WrappedError {
	return q.errors
}

// progressLogger provides a wrapper around an os.File that can either
// write to the file or ignore all writes completely.
type progressLogger struct {
	writeData bool
	log       *os.File
}

// Write will write to the file and perform a Sync() if writing succeeds.
func (l *progressLogger) Write(b []byte) error {
	if l.writeData {
		if _, err := l.log.Write(b); err != nil {
			return err
		}
		return l.log.Sync()
	}
	return nil
}

// Close will call Close() on the underlying file
func (l *progressLogger) Close() error {
	if l.log != nil {
		return l.log.Close()
	}
	return nil
}

// Shutdown will cause the logger to ignore any further writes. It should
// be used when writing causes an error.
func (l *progressLogger) Shutdown() {
	l.writeData = false
}

// newProgressLogger creates a progressLogger based on the presence of
// the GIT_LFS_PROGRESS environment variable. If it is present and a log file
// is able to be created, the logger will write to the file. If it is absent,
// or there is an err creating the file, the logger will ignore all writes.
func newProgressLogger() (*progressLogger, error) {
	logPath := Config.Getenv("GIT_LFS_PROGRESS")

	if len(logPath) == 0 {
		return &progressLogger{}, nil
	}
	if !filepath.IsAbs(logPath) {
		return &progressLogger{}, fmt.Errorf("GIT_LFS_PROGRESS must be an absolute path")
	}

	cbDir := filepath.Dir(logPath)
	if err := os.MkdirAll(cbDir, 0755); err != nil {
		return &progressLogger{}, err
	}

	file, err := os.OpenFile(logPath, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		return &progressLogger{}, err
	}

	return &progressLogger{true, file}, nil
}
<|endoftext|>


<commit_before>/*
 * MinIO Client (C) 2019 MinIO, Inc.
 *
 * 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 cmd

import (
	"errors"
	"strings"

	"github.com/fatih/color"
	"github.com/minio/cli"
	"github.com/minio/mc/pkg/probe"
	"github.com/minio/minio/pkg/console"
)

var adminPolicySetCmd = cli.Command{
	Name:         "set",
	Usage:        "set IAM policy on a user or group",
	Action:       mainAdminPolicySet,
	OnUsageError: onUsageError,
	Before:       setGlobalsFromContext,
	Flags:        globalFlags,
	CustomHelpTemplate: `NAME:
  {{.HelpName}} - {{.Usage}}

USAGE:
  {{.HelpName}} TARGET POLICYNAME [ user=username1 | group=groupname1 ]

POLICYNAME:
  Name of the policy on the MinIO server.

FLAGS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
EXAMPLES:
  1. Set the "readwrite" policy for user "james".
     {{.Prompt}} {{.HelpName}} myminio readwrite user=james

  2. Set the "readonly" policy for group "auditors".
     {{.Prompt}} {{.HelpName}} myminio readonly group=auditors
`,
}

var (
	errBadUserGroupArg = errors.New("Last argument must be of the form user=xx or group=xx")
)

func checkAdminPolicySetSyntax(ctx *cli.Context) {
	if len(ctx.Args()) != 3 {
		cli.ShowCommandHelpAndExit(ctx, "set", 1) // last argument is exit code
	}
}

func parseEntityArg(arg string) (userOrGroup string, isGroup bool, err error) {
	parts := strings.Split(arg, "=")
	switch {
	case len(parts) != 2:
		fallthrough
	case parts[1] == "":
		err = errBadUserGroupArg
	case strings.ToLower(parts[0]) == "user":
		userOrGroup = parts[1]
		isGroup = false
	case strings.ToLower(parts[0]) == "group":
		userOrGroup = parts[1]
		isGroup = true
	default:
		err = errBadUserGroupArg

	}
	return
}

// mainAdminPolicySet is the handler for "mc admin policy set" command.
func mainAdminPolicySet(ctx *cli.Context) error {
	checkAdminPolicySetSyntax(ctx)

	console.SetColor("PolicyMessage", color.New(color.FgGreen))
	console.SetColor("Policy", color.New(color.FgBlue))

	// Get the alias parameter from cli
	args := ctx.Args()
	aliasedURL := args.Get(0)
	policyName := args.Get(1)
	entityArg := args.Get(2)

	userOrGroup, isGroup, e1 := parseEntityArg(entityArg)
	fatalIf(probe.NewError(e1).Trace(args...), "Bad last argument")

	// Create a new MinIO Admin Client
	client, err := newAdminClient(aliasedURL)
	fatalIf(err, "Unable to initialize admin connection.")

	e := client.SetPolicy(globalContext, policyName, userOrGroup, isGroup)

	if e == nil {
		printMsg(userPolicyMessage{
			op:          "set",
			Policy:      policyName,
			UserOrGroup: userOrGroup,
			IsGroup:     isGroup,
		})
	} else {
		fatalIf(probe.NewError(e), "Unable to set the policy")
	}
	return nil
}
<commit_msg>Fix user/group argument parsing (#3602)<commit_after>/*
 * MinIO Client (C) 2019 MinIO, Inc.
 *
 * 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 cmd

import (
	"errors"
	"strings"

	"github.com/fatih/color"
	"github.com/minio/cli"
	"github.com/minio/mc/pkg/probe"
	"github.com/minio/minio/pkg/console"
)

var adminPolicySetCmd = cli.Command{
	Name:         "set",
	Usage:        "set IAM policy on a user or group",
	Action:       mainAdminPolicySet,
	OnUsageError: onUsageError,
	Before:       setGlobalsFromContext,
	Flags:        globalFlags,
	CustomHelpTemplate: `NAME:
  {{.HelpName}} - {{.Usage}}

USAGE:
  {{.HelpName}} TARGET POLICYNAME [ user=username1 | group=groupname1 ]

POLICYNAME:
  Name of the policy on the MinIO server.

FLAGS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
EXAMPLES:
  1. Set the "readwrite" policy for user "james".
     {{.Prompt}} {{.HelpName}} myminio readwrite user=james

  2. Set the "readonly" policy for group "auditors".
     {{.Prompt}} {{.HelpName}} myminio readonly group=auditors
`,
}

var (
	errBadUserGroupArg = errors.New("Last argument must be of the form user=xx or group=xx")
)

func checkAdminPolicySetSyntax(ctx *cli.Context) {
	if len(ctx.Args()) != 3 {
		cli.ShowCommandHelpAndExit(ctx, "set", 1) // last argument is exit code
	}
}

func parseEntityArg(arg string) (userOrGroup string, isGroup bool, err error) {
	parts := strings.SplitN(arg, "=", 2)
	switch {
	case parts[1] == "":
		err = errBadUserGroupArg
	case strings.ToLower(parts[0]) == "user":
		userOrGroup = parts[1]
		isGroup = false
	case strings.ToLower(parts[0]) == "group":
		userOrGroup = parts[1]
		isGroup = true
	default:
		err = errBadUserGroupArg

	}
	return
}

// mainAdminPolicySet is the handler for "mc admin policy set" command.
func mainAdminPolicySet(ctx *cli.Context) error {
	checkAdminPolicySetSyntax(ctx)

	console.SetColor("PolicyMessage", color.New(color.FgGreen))
	console.SetColor("Policy", color.New(color.FgBlue))

	// Get the alias parameter from cli
	args := ctx.Args()
	aliasedURL := args.Get(0)
	policyName := args.Get(1)
	entityArg := args.Get(2)

	userOrGroup, isGroup, e1 := parseEntityArg(entityArg)
	fatalIf(probe.NewError(e1).Trace(args...), "Bad last argument")

	// Create a new MinIO Admin Client
	client, err := newAdminClient(aliasedURL)
	fatalIf(err, "Unable to initialize admin connection.")

	e := client.SetPolicy(globalContext, policyName, userOrGroup, isGroup)

	if e == nil {
		printMsg(userPolicyMessage{
			op:          "set",
			Policy:      policyName,
			UserOrGroup: userOrGroup,
			IsGroup:     isGroup,
		})
	} else {
		fatalIf(probe.NewError(e), "Unable to set the policy")
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/util"
	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/row"
	"github.com/olekukonko/tablewriter"
	"github.com/urfave/cli"
	"net"
	"reflect"
	"sort"
	"strings"
)

// Context is a wrapper around urfave/cli.Context which provides easy access to
// the next unused argument and can have various bytemarky types attached to it
// in order to keep code DRY
type Context struct {
	Context        innerContext
	Account        *lib.Account
	Authed         bool
	Definitions    *lib.Definitions
	Disc           *brain.Disc
	Group          *brain.Group
	Privilege      brain.Privilege
	User           *brain.User
	VirtualMachine *brain.VirtualMachine

	currentArgIndex int
}

// Reset replaces the Context with a blank one (keeping the cli.Context)
func (c *Context) Reset() {
	*c = Context{
		Context: c.Context,
	}
}

// App returns the cli.App that this context is part of. Usually this will be the same as global.App, but it's nice to depend less on globals.
func (c *Context) App() *cli.App {
	return c.App()
}

// args returns all the args that were passed to the Context (i.e. all the args passed to this (sub)command)
func (c *Context) args() cli.Args {
	return c.Context.Args()
}

// Args returns all the unused arguments
func (c *Context) Args() []string {
	return c.args()[c.currentArgIndex:]
}

// Command returns the cli.Command this context is for
func (c *Context) Command() cli.Command {
	return c.Context.Command()
}

// NextArg returns the next unused argument, and marks it as used.
func (c *Context) NextArg() (string, error) {
	if len(c.args()) <= c.currentArgIndex {
		return "", c.Help("not enough arguments were specified")
	}
	arg := c.args()[c.currentArgIndex]
	c.currentArgIndex++
	return arg, nil
}

// Help creates a UsageDisplayedError that will output the issue and a message to consult the documentation
func (c *Context) Help(whatsyourproblem string) (err error) {
	return util.UsageDisplayedError{TheProblem: whatsyourproblem, Command: c.Command().FullName()}
}

// flags below

// Bool returns the value of the named flag as a bool
func (c *Context) Bool(flagname string) bool {
	return c.Context.Bool(flagname)
}

// Discs returns the discs passed along as the named flag.
// I can't imagine why I'd ever name a disc flag anything other than --disc, but the flexibility is there just in case.
func (c *Context) Discs(flagname string) []brain.Disc {
	disc, ok := c.Context.Generic(flagname).(*util.DiscSpecFlag)
	if ok {
		return []brain.Disc(*disc)
	}
	return []brain.Disc{}
}

// FileName returns the name of the file given by the named flag
func (c *Context) FileName(flagname string) string {
	file, ok := c.Context.Generic(flagname).(*util.FileFlag)
	if ok {
		return file.FileName
	}
	return ""
}

// FileContents returns the contents of the file given by the named flag.
func (c *Context) FileContents(flagname string) string {
	file, ok := c.Context.Generic(flagname).(*util.FileFlag)
	if ok {
		return file.Value
	}
	return ""
}

// GroupName returns the named flag as a lib.GroupName
func (c *Context) GroupName(flagname string) lib.GroupName {
	gpNameFlag, ok := c.Context.Generic(flagname).(*GroupNameFlag)
	if !ok {
		return lib.GroupName{}
	}
	return lib.GroupName(*gpNameFlag)
}

// Int returns the value of the named flag as an int
func (c *Context) Int(flagname string) int {
	return c.Context.Int(flagname)
}

// Int64 returns the value of the named flag as an int64
func (c *Context) Int64(flagname string) int64 {
	return c.Context.Int64(flagname)
}

// IPs returns the ips passed along as the named flag.
func (c *Context) IPs(flagname string) []net.IP {
	ips, ok := c.Context.Generic(flagname).(*util.IPFlag)
	if ok {
		return []net.IP(*ips)
	}
	return []net.IP{}
}

// PrivilegeFlag returns the named flag as a PrivilegeFlag
func (c *Context) PrivilegeFlag(flagname string) PrivilegeFlag {
	priv, ok := c.Context.Generic(flagname).(*PrivilegeFlag)
	if ok {
		return *priv
	}
	return PrivilegeFlag{}
}

// String returns the value of the named flag as a string
func (c *Context) String(flagname string) string {
	if c.Context.IsSet(flagname) {
		return c.Context.String(flagname)
	}
	return c.Context.GlobalString(flagname)
}

// Size returns the value of the named SizeSpecFlag as an int in megabytes
func (c *Context) Size(flagname string) int {
	size, ok := c.Context.Generic(flagname).(*util.SizeSpecFlag)
	if ok {
		return int(*size)
	}
	return 0
}

// ResizeFlag returns the named ResizeFlag
func (c *Context) ResizeFlag(flagname string) ResizeFlag {
	size, ok := c.Context.Generic(flagname).(*ResizeFlag)
	if ok {
		return *size
	}
	return ResizeFlag{}
}

// VirtualMachineName returns the named flag as a lib.VirtualMachineName
func (c *Context) VirtualMachineName(flagname string) lib.VirtualMachineName {
	vmNameFlag, ok := c.Context.Generic(flagname).(*VirtualMachineNameFlag)
	if !ok {
		return *global.Config.GetVirtualMachine()
	}
	return lib.VirtualMachineName(*vmNameFlag)
}

// OutputJSON outputs a nicely-indented JSON object that represents obj
func (c *Context) OutputJSON(obj interface{}) error {
	js, err := json.MarshalIndent(obj, "", "    ")
	if err != nil {
		return err
	}
	fmt.Fprintf(global.App.Writer, string(js))
	return nil
}

// OutputTable creates a table for the given object. This makes
// most sense when it's an array, but a regular struct-y object works fine too.
func (c *Context) OutputTable(obj interface{}, fields []string) error {
	table := tablewriter.NewWriter(global.App.Writer)
	// don't autowrap because fields that are slices output one element per line
	// and autowrap
	table.SetAutoWrapText(false)
	// lines between rows!
	table.SetRowLine(true)
	// don't autoformat the headers - autoformat makes them ALLCAPS which makes
	// it hard to figure out what to set --table-fields to.
	// with autoformat off, --table-fields can be set by copying and pasting
	// from the table header.
	table.SetAutoFormatHeaders(false)

	table.SetHeader(fields)
	v := reflect.ValueOf(obj)

	// indirect pointers so we can switch on Kind()
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	// output a single table row for a struct, or several for a slice / array
	switch v.Kind() {
	case reflect.Struct:
		r, err := row.From(obj, fields)
		if err != nil {
			return err
		}
		table.Append(r)
	case reflect.Slice, reflect.Array:
		length := v.Len()
		for i := 0; i < length; i++ {
			el := v.Index(i)
			r, err := row.From(el.Interface(), fields)
			if err != nil {
				return err
			}
			table.Append(r)
		}
	default:
		return fmt.Errorf("%T is not a struct or slice type - please file a bug report", obj)
	}

	table.Render()
	return nil
}

const (
	DefaultAccountTableFields        = "BillingID, Name, Suspended, Groups"
	DefaultBackupTableFields         = "ID, Manual, Label, StorageGrade, Size, BackupCount, BackupSchedules"
	DefaultBackupScheduleTableFields = "ID, StartDate, Interval"
	DefaultDiscTableFields           = "ID, Label, StorageGrade, Size, BackupCount, BackupSchedules"
	DefaultGroupTableFields          = "ID, Name, VirtualMachines"
	DefaultPrivilegeTableFields      = "ID, Username, Level, Target, YubikeyRequired"
	DefaultServerTableFields         = "ID, HostName, ManagementIP, Memory, Cores, Discs, CdromURL, Autoreboot, PowerOn, Deleted"

	DefaultHeadTableFields        = "ID, Label, Online?, UsageStrategy, UUID, CCAddress, VirtualMachinesCount, MemoryFree, UsedCores, Memory, Note, Architecture, Models, ZoneName"
	DefaultTailTableFields        = "ID, Label, Online?, UUID, CCAddress, StoragePools, ZoneName"
	DefaultStoragePoolTableFields = "Label, Discs, Name, Size, FreeSpace, StorageGrade"
	DefaultIPRangeTableFields     = "ID, Spec, VLANNum, Available, Zones"
	DefaultVLANTableFields        = "ID, Num, UsageType, IPRanges"

	DefaultDefinitionTableFields = ""
)

// OutputFlags creates some cli.Flags for when you wanna use OutputInDesiredForm
// thing should be like "server", "servers", "group", "groups"
// jsonType should be "array" or "object"
func OutputFlags(thing string, jsonType string, defaultTableFields string) []cli.Flag {
	return []cli.Flag{
		cli.BoolFlag{
			Name:  "json",
			Usage: fmt.Sprintf("Output the %s as a JSON %s", thing, jsonType),
		},
		cli.BoolFlag{
			Name:  "table",
			Usage: fmt.Sprintf("Output the %s as a table", thing),
		},
		cli.StringFlag{
			Name:  "table-fields",
			Usage: fmt.Sprintf("The fields of the %s to output in the table, comma separated. set to 'help' for a list of fields for this command", thing),
			Value: defaultTableFields,
		},
	}
}

// OutputInDesiredForm outputs obj as a JSON object if --json is set,
// or as a table / table row if --table is set
// otherwise calls humanOutputFn (which should output it in a very human form - PrettyPrint or such
// defaultFormat is an optional string stating that the default format should be
func (c *Context) OutputInDesiredForm(obj interface{}, humanOutputFn func() error, defaultFormat ...string) error {
	format, err := global.Config.GetV("output-format")
	if err != nil {
		return err
	}
	if len(defaultFormat) > 0 && format.Source == "CODE" {
		format.Value = defaultFormat[0]
	}

	if c.Bool("json") {
		format.Value = "json"
	} else if c.Bool("table") || c.String("table-fields") != "" {
		format.Value = "table"
	}

	switch format.Value {
	case "debug":
		fmt.Printf("%#v", obj)
		return nil
	case "json":
		return c.OutputJSON(obj)
	case "table":
		// this is so horrible
		fields := strings.Split(c.String("table-fields"), ",")
		for i, f := range fields {
			fields[i] = strings.TrimSpace(f)
		}
		fieldsList := row.FieldsFrom(obj)
		sort.Strings(fieldsList)
		if len(fields) > 0 && fields[0] == "help" {
			fmt.Fprintf(global.App.Writer, "Table fields available for this command: \r\n  %s\r\n\r\n", strings.Join(fieldsList, "\r\n  "))
			return nil
		} else if len(fields) > 0 && fields[0] != "" {
			return c.OutputTable(obj, fields)
		} else {
			return c.OutputTable(obj, fieldsList)
		}
	}
	return humanOutputFn()
}
<commit_msg>Fix Context.String() incorrectly trying GlobalString when a nonzero default Value has been set<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/util"
	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/row"
	"github.com/olekukonko/tablewriter"
	"github.com/urfave/cli"
	"net"
	"reflect"
	"sort"
	"strings"
)

// Context is a wrapper around urfave/cli.Context which provides easy access to
// the next unused argument and can have various bytemarky types attached to it
// in order to keep code DRY
type Context struct {
	Context        innerContext
	Account        *lib.Account
	Authed         bool
	Definitions    *lib.Definitions
	Disc           *brain.Disc
	Group          *brain.Group
	Privilege      brain.Privilege
	User           *brain.User
	VirtualMachine *brain.VirtualMachine

	currentArgIndex int
}

// Reset replaces the Context with a blank one (keeping the cli.Context)
func (c *Context) Reset() {
	*c = Context{
		Context: c.Context,
	}
}

// App returns the cli.App that this context is part of. Usually this will be the same as global.App, but it's nice to depend less on globals.
func (c *Context) App() *cli.App {
	return c.App()
}

// args returns all the args that were passed to the Context (i.e. all the args passed to this (sub)command)
func (c *Context) args() cli.Args {
	return c.Context.Args()
}

// Args returns all the unused arguments
func (c *Context) Args() []string {
	return c.args()[c.currentArgIndex:]
}

// Command returns the cli.Command this context is for
func (c *Context) Command() cli.Command {
	return c.Context.Command()
}

// NextArg returns the next unused argument, and marks it as used.
func (c *Context) NextArg() (string, error) {
	if len(c.args()) <= c.currentArgIndex {
		return "", c.Help("not enough arguments were specified")
	}
	arg := c.args()[c.currentArgIndex]
	c.currentArgIndex++
	return arg, nil
}

// Help creates a UsageDisplayedError that will output the issue and a message to consult the documentation
func (c *Context) Help(whatsyourproblem string) (err error) {
	return util.UsageDisplayedError{TheProblem: whatsyourproblem, Command: c.Command().FullName()}
}

// flags below

// Bool returns the value of the named flag as a bool
func (c *Context) Bool(flagname string) bool {
	return c.Context.Bool(flagname)
}

// Discs returns the discs passed along as the named flag.
// I can't imagine why I'd ever name a disc flag anything other than --disc, but the flexibility is there just in case.
func (c *Context) Discs(flagname string) []brain.Disc {
	disc, ok := c.Context.Generic(flagname).(*util.DiscSpecFlag)
	if ok {
		return []brain.Disc(*disc)
	}
	return []brain.Disc{}
}

// FileName returns the name of the file given by the named flag
func (c *Context) FileName(flagname string) string {
	file, ok := c.Context.Generic(flagname).(*util.FileFlag)
	if ok {
		return file.FileName
	}
	return ""
}

// FileContents returns the contents of the file given by the named flag.
func (c *Context) FileContents(flagname string) string {
	file, ok := c.Context.Generic(flagname).(*util.FileFlag)
	if ok {
		return file.Value
	}
	return ""
}

// GroupName returns the named flag as a lib.GroupName
func (c *Context) GroupName(flagname string) lib.GroupName {
	gpNameFlag, ok := c.Context.Generic(flagname).(*GroupNameFlag)
	if !ok {
		return lib.GroupName{}
	}
	return lib.GroupName(*gpNameFlag)
}

// Int returns the value of the named flag as an int
func (c *Context) Int(flagname string) int {
	return c.Context.Int(flagname)
}

// Int64 returns the value of the named flag as an int64
func (c *Context) Int64(flagname string) int64 {
	return c.Context.Int64(flagname)
}

// IPs returns the ips passed along as the named flag.
func (c *Context) IPs(flagname string) []net.IP {
	ips, ok := c.Context.Generic(flagname).(*util.IPFlag)
	if ok {
		return []net.IP(*ips)
	}
	return []net.IP{}
}

// PrivilegeFlag returns the named flag as a PrivilegeFlag
func (c *Context) PrivilegeFlag(flagname string) PrivilegeFlag {
	priv, ok := c.Context.Generic(flagname).(*PrivilegeFlag)
	if ok {
		return *priv
	}
	return PrivilegeFlag{}
}

// String returns the value of the named flag as a string
func (c *Context) String(flagname string) string {
	if c.Context.IsSet(flagname) || c.Context.String(flagname) != "" {
		return c.Context.String(flagname)
	}
	return c.Context.GlobalString(flagname)
}

// Size returns the value of the named SizeSpecFlag as an int in megabytes
func (c *Context) Size(flagname string) int {
	size, ok := c.Context.Generic(flagname).(*util.SizeSpecFlag)
	if ok {
		return int(*size)
	}
	return 0
}

// ResizeFlag returns the named ResizeFlag
func (c *Context) ResizeFlag(flagname string) ResizeFlag {
	size, ok := c.Context.Generic(flagname).(*ResizeFlag)
	if ok {
		return *size
	}
	return ResizeFlag{}
}

// VirtualMachineName returns the named flag as a lib.VirtualMachineName
func (c *Context) VirtualMachineName(flagname string) lib.VirtualMachineName {
	vmNameFlag, ok := c.Context.Generic(flagname).(*VirtualMachineNameFlag)
	if !ok {
		return *global.Config.GetVirtualMachine()
	}
	return lib.VirtualMachineName(*vmNameFlag)
}

// OutputJSON outputs a nicely-indented JSON object that represents obj
func (c *Context) OutputJSON(obj interface{}) error {
	js, err := json.MarshalIndent(obj, "", "    ")
	if err != nil {
		return err
	}
	fmt.Fprintf(global.App.Writer, string(js))
	return nil
}

// OutputTable creates a table for the given object. This makes
// most sense when it's an array, but a regular struct-y object works fine too.
func (c *Context) OutputTable(obj interface{}, fields []string) error {
	table := tablewriter.NewWriter(global.App.Writer)
	// don't autowrap because fields that are slices output one element per line
	// and autowrap
	table.SetAutoWrapText(false)
	// lines between rows!
	table.SetRowLine(true)
	// don't autoformat the headers - autoformat makes them ALLCAPS which makes
	// it hard to figure out what to set --table-fields to.
	// with autoformat off, --table-fields can be set by copying and pasting
	// from the table header.
	table.SetAutoFormatHeaders(false)

	table.SetHeader(fields)
	v := reflect.ValueOf(obj)

	// indirect pointers so we can switch on Kind()
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	// output a single table row for a struct, or several for a slice / array
	switch v.Kind() {
	case reflect.Struct:
		r, err := row.From(obj, fields)
		if err != nil {
			return err
		}
		table.Append(r)
	case reflect.Slice, reflect.Array:
		length := v.Len()
		for i := 0; i < length; i++ {
			el := v.Index(i)
			r, err := row.From(el.Interface(), fields)
			if err != nil {
				return err
			}
			table.Append(r)
		}
	default:
		return fmt.Errorf("%T is not a struct or slice type - please file a bug report", obj)
	}

	table.Render()
	return nil
}

const (
	// DefaultAccountTableFields is the default for --table-fields for lib.Account
	DefaultAccountTableFields = "BillingID, Name, Suspended, Groups"
	// DefaultBackupTableFields is the default for --table-fields for brain.Backup
	DefaultBackupTableFields = "ID, Manual, Label, StorageGrade, Size, BackupCount, BackupSchedules"
	// DefaultBackupScheduleTableFields is the default for --table-fields for brain.BackupSchedule
	DefaultBackupScheduleTableFields = "ID, StartDate, Interval"
	// DefaultDiscTableFields is the default for --table-fields for brain.Disc
	DefaultDiscTableFields = "ID, Label, StorageGrade, Size, BackupCount, BackupSchedules"
	// DefaultGroupTableFields is the default for --table-fields for brain.Group
	DefaultGroupTableFields = "ID, Name, VirtualMachines"
	// DefaultPrivilegeTableFields is the default for --table-fields for brain.Privilege
	DefaultPrivilegeTableFields = "ID, Username, Level, Target, YubikeyRequired"
	// DefaultServerTableFields is the default for --table-fields for brain.VirtualMachine
	DefaultServerTableFields = "ID, Hostname, ManagementAddress, Memory, Cores, Discs, CdromURL, Autoreboot, PowerOn, Deleted"

	// DefaultHeadTableFields is the default for --table-fields for brain.Head
	DefaultHeadTableFields = "ID, Label, Online?, UsageStrategy, UUID, CCAddress, VirtualMachinesCount, MemoryFree, UsedCores, Memory, Note, Architecture, Models, ZoneName"
	// DefaultTailTableFields is the default for --table-fields for brain.Tail
	DefaultTailTableFields = "ID, Label, Online?, UUID, CCAddress, StoragePools, ZoneName"
	// DefaultStoragePoolTableFields is the default for --table-fields for brain.StoragePool
	DefaultStoragePoolTableFields = "Label, Discs, Name, Size, FreeSpace, StorageGrade"
	// DefaultIPRangeTableFields is the default for --table-fields for brain.IPRange
	DefaultIPRangeTableFields = "ID, Spec, VLANNum, Available, Zones"
	// DefaultVLANTableFields is the default for --table-fields for brain.VLAN
	DefaultVLANTableFields = "ID, Num, UsageType, IPRanges"

	// DefaultDefinitionTableFields is the default for --table-fields for the *Definition types, because they're all the same at the moment.
	DefaultDefinitionTableFields = "Name, Description"
)

// OutputFlags creates some cli.Flags for when you wanna use OutputInDesiredForm
// thing should be like "server", "servers", "group", "groups"
// jsonType should be "array" or "object"
func OutputFlags(thing string, jsonType string, defaultTableFields string) []cli.Flag {
	return []cli.Flag{
		cli.BoolFlag{
			Name:  "json",
			Usage: fmt.Sprintf("Output the %s as a JSON %s", thing, jsonType),
		},
		cli.BoolFlag{
			Name:  "table",
			Usage: fmt.Sprintf("Output the %s as a table", thing),
		},
		cli.StringFlag{
			Name:  "table-fields",
			Usage: fmt.Sprintf("The fields of the %s to output in the table, comma separated. set to 'help' for a list of fields for this command", thing),
			Value: defaultTableFields,
		},
	}
}

// OutputInDesiredForm outputs obj as a JSON object if --json is set,
// or as a table / table row if --table is set
// otherwise calls humanOutputFn (which should output it in a very human form - PrettyPrint or such
// defaultFormat is an optional string stating that the default format should be
func (c *Context) OutputInDesiredForm(obj interface{}, humanOutputFn func() error, defaultFormat ...string) error {
	format, err := global.Config.GetV("output-format")
	if err != nil {
		return err
	}
	if len(defaultFormat) > 0 && format.Source == "CODE" {
		format.Value = defaultFormat[0]
	}

	if c.Bool("json") {
		format.Value = "json"
	} else if c.Bool("table") || c.String("table-fields") != "" {
		format.Value = "table"
	}

	switch format.Value {
	case "debug":
		fmt.Printf("%#v", obj)
		return nil
	case "json":
		return c.OutputJSON(obj)
	case "table":
		// this is so horrible
		fields := strings.Split(c.String("table-fields"), ",")
		for i, f := range fields {
			fields[i] = strings.TrimSpace(f)
		}
		fieldsList := row.FieldsFrom(obj)
		sort.Strings(fieldsList)
		if len(fields) > 0 && fields[0] == "help" {
			fmt.Fprintf(global.App.Writer, "Table fields available for this command: \r\n  %s\r\n\r\n", strings.Join(fieldsList, "\r\n  "))
			return nil
		} else if len(fields) > 0 && fields[0] != "" {
			return c.OutputTable(obj, fields)
		} else {
			return c.OutputTable(obj, fieldsList)
		}
	}
	return humanOutputFn()
}
<|endoftext|>