<commit_before>package mpb_test

import (
	"bytes"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb"
)

func TestBarSetWidth(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	// overwrite default width 80
	customWidth := 60
	bar := p.AddBar(100).SetWidth(customWidth).
		TrimLeftSpace().TrimRightSpace()
	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}
	p.Stop()

	gotWidth := len(buf.Bytes())
	if gotWidth != customWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", customWidth, gotWidth)
	}
}

func TestBarSetInvalidWidth(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	bar := p.AddBar(100).SetWidth(1).
		TrimLeftSpace().TrimRightSpace()
	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}
	p.Stop()

	wantWidth := 80
	gotWidth := len(buf.Bytes())
	if gotWidth != wantWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", wantWidth, gotWidth)
	}
}

func TestBarFormat(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	p := mpb.New().WithCancel(cancel).SetOut(&buf)
	customFormat := "(#>_)"
	bar := p.AddBar(100).Format(customFormat).
		TrimLeftSpace().TrimRightSpace()

	go func() {
		for i := 0; i < 100; i++ {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	// removing new line
	bytes := removeLastRune(buf.Bytes())

	seen := make(map[rune]bool)
	for _, r := range string(bytes) {
		if !seen[r] {
			seen[r] = true
		}
	}
	for _, r := range customFormat {
		if !seen[r] {
			t.Errorf("Rune %#U not found in bar\n", r)
		}
	}
}

func TestBarInvalidFormat(t *testing.T) {
	var buf bytes.Buffer
	customWidth := 60
	p := mpb.New().SetWidth(customWidth).SetOut(&buf)
	customFormat := "(#>=_)"
	bar := p.AddBar(100).Format(customFormat).
		TrimLeftSpace().TrimRightSpace()

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	got := buf.String()
	want := fmt.Sprintf("[%s]", strings.Repeat("=", customWidth-2))
	if !strings.Contains(got, want) {
		t.Errorf("Expected format: %s, got %s\n", want, got)
	}
}

func TestBarInProgress(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	p := mpb.New().WithCancel(cancel).SetOut(&buf)
	bar := p.AddBar(100).TrimLeftSpace().TrimRightSpace()

	stopped := make(chan struct{})

	go func() {
		defer close(stopped)
		for bar.InProgress() {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	select {
	case <-stopped:
	case <-time.After(300 * time.Millisecond):
		t.Error("bar.InProgress returns true after cancel")
	}
}

func TestGetSpinner(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	bar := p.AddBar(0).TrimLeftSpace().TrimRightSpace()

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	spinnerChars := []byte(`-\|/`)
	seen := make(map[byte]bool)
	for _, b := range buf.Bytes() {
		if !seen[b] {
			seen[b] = true
		}
	}
	for _, b := range spinnerChars {
		if !seen[b] {
			t.Errorf("Char %#U not found in bar's output\n", b)
		}
	}
}

func TestBarGetID(t *testing.T) {
	var wg sync.WaitGroup
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)

	numBars := 3
	wg.Add(numBars)

	bars := make([]*mpb.Bar, numBars)
	for i := 0; i < numBars; i++ {
		bars[i] = p.AddBarWithID(i, 100)

		go func(bar *mpb.Bar) {
			defer wg.Done()
			for i := 0; i < 100; i++ {
				time.Sleep(10 * time.Millisecond)
				bar.Incr(1)
			}
		}(bars[i])
	}

	for wantID, bar := range bars {
		gotID := bar.GetID()
		if gotID != wantID {
			t.Errorf("Expected bar id: %d, got %d\n", wantID, gotID)
		}
	}

	wg.Wait()
	p.Stop()
}

func TestBarIncrWithReFill(t *testing.T) {
	var buf bytes.Buffer

	width := 100
	p := mpb.New().SetWidth(width).SetOut(&buf)

	total := 100
	refill := 30
	delta := total - refill
	refillChar := '+'
	bar := p.AddBar(int64(total)).TrimLeftSpace().TrimRightSpace()

	bar.IncrWithReFill(refill, &mpb.Refill{Char: refillChar})

	for i := 0; i < delta; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	bytes := removeLastRune(buf.Bytes())

	gotBar := string(bytes[len(bytes)-width:])
	wantBar := fmt.Sprintf("[%s%s]",
		strings.Repeat(string(refillChar), refill-1),
		strings.Repeat("=", delta-1))
	if gotBar != wantBar {
		t.Errorf("Want bar: %s, got bar: %s\n", wantBar, gotBar)
	}
}

func TestBarPanics(t *testing.T) {
	var wg sync.WaitGroup
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)

	wantPanic := "Upps!!!"
	numBars := 3
	wg.Add(numBars)

	for i := 0; i < numBars; i++ {
		name := fmt.Sprintf("b#%02d:", i)
		bar := p.AddBarWithID(i, 100).
			PrependFunc(func(s *mpb.Statistics, _ chan<- int, _ <-chan int) string {
				if s.ID == 2 && s.Current >= 42 {
					panic(wantPanic)
				}
				return name
			})

		go func() {
			defer wg.Done()
			for i := 0; i < 100; i++ {
				time.Sleep(10 * time.Millisecond)
				bar.Incr(1)
			}
		}()
	}

	wg.Wait()
	p.Stop()

	bytes := removeLastRune(buf.Bytes())
	out := strings.Split(string(bytes), "\n")
	gotPanic := out[len(out)-1]
	if gotPanic != wantPanic {
		t.Errorf("Want panic: %s, got panic: %s\n", wantPanic, gotPanic)
	}
}

func removeLastRune(bytes []byte) []byte {
	_, size := utf8.DecodeLastRune(bytes)
	return bytes[:len(bytes)-size]
}
<commit_msg>update bar_test<commit_after>package mpb_test

import (
	"bytes"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb"
	"github.com/vbauerster/mpb/decor"
)

func TestBarSetWidth(t *testing.T) {
	var buf bytes.Buffer
	// overwrite default width 80
	customWidth := 60
	p := mpb.New(mpb.Output(&buf), mpb.WithWidth(customWidth))
	bar := p.AddBar(100, mpb.BarTrim())

	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}

	p.Stop()

	gotWidth := len(buf.Bytes())
	if gotWidth != customWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", customWidth, gotWidth)
	}
}

func TestBarSetInvalidWidth(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New(mpb.Output(&buf), mpb.WithWidth(1))
	bar := p.AddBar(100, mpb.BarTrim())

	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}

	p.Stop()

	wantWidth := 80
	gotWidth := len(buf.Bytes())
	if gotWidth != wantWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", wantWidth, gotWidth)
	}
}

func TestBarFormat(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	customFormat := "(#>_)"
	p := mpb.New(
		mpb.Output(&buf),
		mpb.WithCancel(cancel),
		mpb.WithFormat(customFormat),
	)
	bar := p.AddBar(100, mpb.BarTrim())

	go func() {
		for i := 0; i < 100; i++ {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	barAsStr := strings.Trim(buf.String(), "\n")

	seen := make(map[rune]bool)
	for _, r := range barAsStr {
		if !seen[r] {
			seen[r] = true
		}
	}
	for _, r := range customFormat {
		if !seen[r] {
			t.Errorf("Rune %#U not found in bar\n", r)
		}
	}
}

func TestBarInvalidFormat(t *testing.T) {
	var buf bytes.Buffer
	customWidth := 60
	customFormat := "(#>=_)"
	p := mpb.New(
		mpb.Output(&buf),
		mpb.WithWidth(customWidth),
		mpb.WithFormat(customFormat),
	)
	bar := p.AddBar(100, mpb.BarTrim())

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	got := buf.String()
	want := fmt.Sprintf("[%s]", strings.Repeat("=", customWidth-2))
	if !strings.Contains(got, want) {
		t.Errorf("Expected format: %s, got %s\n", want, got)
	}
}

func TestBarInProgress(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	p := mpb.New(
		mpb.Output(&buf),
		mpb.WithCancel(cancel),
	)
	bar := p.AddBar(100, mpb.BarTrim())

	stopped := make(chan struct{})

	go func() {
		defer close(stopped)
		for bar.InProgress() {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	select {
	case <-stopped:
	case <-time.After(300 * time.Millisecond):
		t.Error("bar.InProgress returns true after cancel")
	}
}

func TestGetSpinner(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New(mpb.Output(&buf))
	bar := p.AddBar(0, mpb.BarTrim())

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	spinnerChars := []byte(`-\|/`)
	seen := make(map[byte]bool)
	for _, b := range bytes.Trim(buf.Bytes(), "\n") {
		if !seen[b] {
			seen[b] = true
		}
	}
	for _, b := range spinnerChars {
		if !seen[b] {
			t.Errorf("Char %#U not found in bar's output\n", b)
		}
	}
}

func TestBarGetID(t *testing.T) {
	var wg sync.WaitGroup
	var buf bytes.Buffer
	p := mpb.New(mpb.Output(&buf))

	numBars := 3
	wg.Add(numBars)

	bars := make([]*mpb.Bar, numBars)
	for i := 0; i < numBars; i++ {
		bars[i] = p.AddBar(100, mpb.BarID(i))

		go func(bar *mpb.Bar) {
			defer wg.Done()
			for i := 0; i < 100; i++ {
				time.Sleep(10 * time.Millisecond)
				bar.Incr(1)
			}
		}(bars[i])
	}

	for wantID, bar := range bars {
		gotID := bar.ID()
		if gotID != wantID {
			t.Errorf("Expected bar id: %d, got %d\n", wantID, gotID)
		}
	}

	wg.Wait()
	p.Stop()
}

func TestBarIncrWithReFill(t *testing.T) {
	var buf bytes.Buffer

	width := 100
	p := mpb.New(
		mpb.Output(&buf),
		mpb.WithWidth(width),
	)

	total := 100
	till := 30
	refillChar := '+'

	bar := p.AddBar(100, mpb.BarTrim())

	bar.ResumeFill(refillChar, int64(till))

	for i := 0; i < total; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	bytes := removeLastRune(buf.Bytes())

	gotBar := string(bytes[len(bytes)-width:])
	wantBar := fmt.Sprintf("[%s%s]",
		strings.Repeat(string(refillChar), till-1),
		strings.Repeat("=", total-till-1))
	if gotBar != wantBar {
		t.Errorf("Want bar: %s, got bar: %s\n", wantBar, gotBar)
	}
}

func TestBarPanics(t *testing.T) {
	var wg sync.WaitGroup
	var buf bytes.Buffer
	p := mpb.New(mpb.Output(&buf))

	wantPanic := "Upps!!!"
	numBars := 3
	wg.Add(numBars)

	for i := 0; i < numBars; i++ {
		name := fmt.Sprintf("b#%02d:", i)
		bar := p.AddBar(100, mpb.BarID(i), mpb.PrependDecorators(
			func(s *decor.Statistics, _ chan<- int, _ <-chan int) string {
				if s.ID == 2 && s.Current >= 42 {
					panic(wantPanic)
				}
				return name
			},
		))

		go func() {
			defer wg.Done()
			for i := 0; i < 100; i++ {
				time.Sleep(10 * time.Millisecond)
				bar.Incr(1)
			}
		}()
	}

	wg.Wait()
	p.Stop()

	bytes := removeLastRune(buf.Bytes())
	out := strings.Split(string(bytes), "\n")
	gotPanic := out[len(out)-1]
	if gotPanic != wantPanic {
		t.Errorf("Want panic: %s, got panic: %s\n", wantPanic, gotPanic)
	}
}

func removeLastRune(bytes []byte) []byte {
	_, size := utf8.DecodeLastRune(bytes)
	return bytes[:len(bytes)-size]
}
<|endoftext|>


<commit_before>package pelicantun

import (
	"fmt"
	"net/http"
	"testing"

	//cv "github.com/glycerine/goconvey/convey"
)

func TestFullRoundtripSocksProxyTalksToReverseProxy002(t *testing.T) {

	// setup a mock web server that replies to ping with pong.
	mux := http.NewServeMux()

	// ping allows our test machinery to function
	mux.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) {
		r.Body.Close()
		fmt.Fprintf(w, "pong")
	})

	web := NewWebServer(WebServerConfig{}, mux)
	web.Start()
	defer web.Stop()

	if !PortIsBound(web.Cfg.Listen.IpPort) {
		panic("web server did not come up")
	}

	// start a reverse proxy
	rev := NewReverseProxy(ReverseProxyConfig{Dest: web.Cfg.Listen})
	rev.Start()
	defer rev.Stop()

	if !PortIsBound(rev.Cfg.Listen.IpPort) {
		panic("rev proxy not up")
	}

	// start the forward proxy, talks to the reverse proxy.
	fwd := NewPelicanSocksProxy(PelicanSocksProxyConfig{
		Dest: rev.Cfg.Listen,
	})
	fwd.Start()
	//defer fwd.Stop()
	if !PortIsBound(fwd.Cfg.Listen.IpPort) {
		panic("fwd proxy not up")
	}
	fwd.Stop()
	/*
		cv.Convey("Given a ForwardProxy and a ReverseProxy, they should communicate over http", t, func() {

			po("\n fetching url from %v\n", fwd.Cfg.Listen.IpPort)

			by, err := FetchUrl("http://" + fwd.Cfg.Listen.IpPort + "/ping")
			cv.So(err, cv.ShouldEqual, nil)
			//fmt.Printf("by:'%s'\n", string(by))
			cv.So(string(by), cv.ShouldEqual, "pong")
		})
	*/
	fmt.Printf("\n done with TestSocksProxyTalksToReverseProxy002()\n")
}
<commit_msg>restore full 002 test; hangs<commit_after>package pelicantun

import (
	"fmt"
	"net/http"
	"testing"

	cv "github.com/glycerine/goconvey/convey"
)

func TestFullRoundtripSocksProxyTalksToReverseProxy002(t *testing.T) {

	// setup a mock web server that replies to ping with pong.
	mux := http.NewServeMux()

	// ping allows our test machinery to function
	mux.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) {
		r.Body.Close()
		fmt.Fprintf(w, "pong")
	})

	web := NewWebServer(WebServerConfig{}, mux)
	web.Start()
	defer web.Stop()

	if !PortIsBound(web.Cfg.Listen.IpPort) {
		panic("web server did not come up")
	}

	// start a reverse proxy
	rev := NewReverseProxy(ReverseProxyConfig{Dest: web.Cfg.Listen})
	rev.Start()
	defer rev.Stop()

	if !PortIsBound(rev.Cfg.Listen.IpPort) {
		panic("rev proxy not up")
	}

	// start the forward proxy, talks to the reverse proxy.
	fwd := NewPelicanSocksProxy(PelicanSocksProxyConfig{
		Dest: rev.Cfg.Listen,
	})
	fwd.Start()
	defer fwd.Stop()
	if !PortIsBound(fwd.Cfg.Listen.IpPort) {
		panic("fwd proxy not up")
	}

	cv.Convey("Given a ForwardProxy and a ReverseProxy, they should communicate over http", t, func() {

		po("\n fetching url from %v\n", fwd.Cfg.Listen.IpPort)

		by, err := FetchUrl("http://" + fwd.Cfg.Listen.IpPort + "/ping")
		cv.So(err, cv.ShouldEqual, nil)
		//fmt.Printf("by:'%s'\n", string(by))
		cv.So(string(by), cv.ShouldEqual, "pong")
	})
	fmt.Printf("\n done with TestSocksProxyTalksToReverseProxy002()\n")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/cactus/go-statsd-client/statsd"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"
)

func dieIfError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s\n", err.Error())
		os.Exit(1)
	}
}

func doEvery(d time.Duration, f func(*statsd.Client), s *statsd.Client) {
	f(s)
	for _ = range time.Tick(d) {
		f(s)
	}
}

func process_targets(s *statsd.Client) {
	content, err := ioutil.ReadFile("targets")
	if err != nil {
		fmt.Println("couldn't open targets file")
		return
	}
	targets := strings.Split(string(content), "\n")
	for _, target := range targets {
		if len(target) < 1 {
			continue
		}
		go test(target, s)
	}
}

func test(target string, s *statsd.Client) {
	tuple := strings.Split(target, ":")
	host := tuple[0]
	port := tuple[1]
	subhost := strings.Replace(host, ".", "_", -1)

	pre := time.Now()
	conn, err := net.Dial("tcp", target)
	if err != nil {
		fmt.Println("connect error", target)
		s.Inc(fmt.Sprintf("%s.%s.dial_failed", subhost, port), 1, 1)
		return
	}
	duration := time.Since(pre)
	ms := int64(duration / time.Millisecond)
	fmt.Printf("%s.%s.duration %d\n", subhost, port, ms)
	s.Timing(fmt.Sprintf("%s.%s", subhost, port), ms, 1)
	conn.Close()
}

func main() {
	if len(os.Args) < 3 {
		fmt.Println("Usage: smoketcp <statsd_host>:<statsd_port> <bucket_prefix>")
		fmt.Println("\nEx: smoketcp statsd.example.com:8125 Location.for.smokeping.values")
		os.Exit(1)
	}

	s, err := statsd.Dial(os.Args[1], fmt.Sprintf("%s", os.Args[2]))
	dieIfError(err)
	defer s.Close()
	doEvery(time.Second, process_targets, s)
}
<commit_msg>use flags on the commandline, statsd_host, statsd_port, and bucket<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/cactus/go-statsd-client/statsd"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"
)

func dieIfError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s\n", err.Error())
		os.Exit(1)
	}
}

func doEvery(d time.Duration, f func(*statsd.Client), s *statsd.Client) {
	f(s)
	for _ = range time.Tick(d) {
		f(s)
	}
}

func process_targets(s *statsd.Client) {
	content, err := ioutil.ReadFile("targets")
	if err != nil {
		fmt.Println("couldn't open targets file")
		return
	}
	targets := strings.Split(string(content), "\n")
	for _, target := range targets {
		if len(target) < 1 {
			continue
		}
		go test(target, s)
	}
}

func test(target string, s *statsd.Client) {
	tuple := strings.Split(target, ":")
	host := tuple[0]
	port := tuple[1]
	subhost := strings.Replace(host, ".", "_", -1)

	pre := time.Now()
	conn, err := net.Dial("tcp", target)
	if err != nil {
		fmt.Println("connect error", target)
		s.Inc(fmt.Sprintf("%s.%s.dial_failed", subhost, port), 1, 1)
		return
	}
	duration := time.Since(pre)
	ms := int64(duration / time.Millisecond)
	fmt.Printf("%s.%s.duration %d\n", subhost, port, ms)
	s.Timing(fmt.Sprintf("%s.%s", subhost, port), ms, 1)
	conn.Close()
}

func main() {
  var statsd_host = flag.String("statsd_host", "localhost", "Statsd Hostname")
  var statsd_port = flag.String("statsd_port", "8125", "Statsd port")
  var bucket = flag.String("bucket", "smoketcp", "Graphite bucket prefix")
  flag.Parse()

	s, err := statsd.Dial(fmt.Sprintf("%s:%s", *statsd_host, *statsd_port), fmt.Sprintf("%s", *bucket))
	dieIfError(err)
	defer s.Close()
	doEvery(time.Second, process_targets, s)
}
<|endoftext|>


<commit_before>package junos

import (
	"encoding/xml"
	"fmt"
)

// SoftwarePackages contains a list of software packages managed by Junos Space.
type SoftwarePackages struct {
	Packages []SoftwarePackage `xml:"package"`
}

// A SoftwarePackage contains information about each individual software package.
type SoftwarePackage struct {
	ID       int    `xml:"key,attr"`
	Name     string `xml:"fileName"`
	Version  string `xml:"version"`
	Platform string `xml:"platformType"`
}

// SoftwareUpgrade consists of options available to use before issuing a software upgrade.
type SoftwareUpgrade struct {
	UseDownloaded bool // Use an image already staged on the device.
	Validate      bool // Check/don't check compatibility with current configuration.
	Reboot        bool // Reboot system after adding package.
	RebootAfter   int  // Reboot the system after "x" minutes.
	Cleanup       bool // Remove any pre-existing packages on the device.
	RemoveAfter   bool // Remove the package after successful installation.
}

// deployXML is XML we send (POST) for image deployment.
var deployXML = `
<exec-deploy>
    <devices>
        <device href= "/api/space/device-management/devices/%d"/>
    </devices>
    <deployOptions>
        <useAlreadyDownloaded>%t</useAlreadyDownloaded>
        <validate>%t</validate>
        <bestEffortLoad>false</bestEffortLoad>
        <snapShotRequired>false</snapShotRequired>
        <rebootDevice>%t</rebootDevice>
        <rebootAfterXMinutes>%d</rebootAfterXMinutes>
        <cleanUpExistingOnDevice>%t</cleanUpExistingOnDevice>
        <removePkgAfterInstallation>%t</removePkgAfterInstallation>
    </deployOptions>
</exec-deploy>
`

// removeStagedXML is XML we send (POST) for removing a staged image.
var removeStagedXML = `
<exec-remove>
    <devices>
        <device href="/api/space/device-management/devices/%d"/>
    </devices>
</exec-remove>
`

// stageXML is XML we send (POST) for staging an image on a device.
var stageXML = `
<exec-stage>
    <devices>
        <device href="/api/space/device-management/devices/%d"/>
    </devices>
    <stageOptions>
        <cleanUpExistingOnDevice>%t</cleanUpExistingOnDevice>
    </stageOptions>
</exec-stage>
`

// getSoftwareID returns the ID of the software package.
func (s *JunosSpace) getSoftwareID(image string) (int, error) {
	var err error
	var softwareID int
	images, err := s.Software()
	if err != nil {
		return 0, err
	}

	for _, sw := range images.Packages {
		if sw.Name == image {
			softwareID = sw.ID
		}
	}

	return softwareID, nil
}

// DeploySoftware starts the upgrade process on the device, using the given image along
// with the options specified.
func (s *JunosSpace) DeploySoftware(device, image string, options *SoftwareUpgrade) (int, error) {
	var job jobID
	deviceID, _ := s.getDeviceID(device, false)
	softwareID, _ := s.getSoftwareID(image)
	deploy := fmt.Sprintf(deployXML, deviceID, options.UseDownloaded, options.Validate, options.Reboot, options.RebootAfter, options.Cleanup, options.RemoveAfter)
	req := &APIRequest{
		Method:      "post",
		URL:         fmt.Sprintf("/api/space/software-management/packages/%d/exec-deploy", softwareID),
		Body:        deploy,
		ContentType: contentExecDeploy,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}

// RemoveStagedSoftware will delete the staged software image on the device.
func (s *JunosSpace) RemoveStagedSoftware(device, image string) (int, error) {
	var job jobID
	deviceID, _ := s.getDeviceID(device, false)
	softwareID, _ := s.getSoftwareID(image)
	remove := fmt.Sprintf(removeStagedXML, deviceID)
	req := &APIRequest{
		Method:      "post",
		URL:         fmt.Sprintf("/api/space/software-management/packages/%d/exec-remove", softwareID),
		Body:        remove,
		ContentType: contentExecRemove,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}

// Software queries the Junos Space server and returns all of the information
// about each software image that Space manages.
func (s *JunosSpace) Software() (*SoftwarePackages, error) {
	var software SoftwarePackages
	req := &APIRequest{
		Method: "get",
		URL:    "/api/space/software-management/packages",
	}
	data, err := s.APICall(req)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(data, &software)
	if err != nil {
		return nil, err
	}

	return &software, nil
}

// StageSoftware loads the given software image onto the device but does not
// upgrade it. The package is placed in the /var/tmp directory.
func (s *JunosSpace) StageSoftware(device, image string, cleanup bool) (int, error) {
	var job jobID
	deviceID, _ := s.getDeviceID(device, false)
	softwareID, _ := s.getSoftwareID(image)
	stage := fmt.Sprintf(stageXML, deviceID, cleanup)
	req := &APIRequest{
		Method:      "post",
		URL:         fmt.Sprintf("/api/space/software-management/packages/%d/exec-stage", softwareID),
		Body:        stage,
		ContentType: contentExecStage,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}
<commit_msg>Removed unused parameter in getDeviceID<commit_after>package junos

import (
	"encoding/xml"
	"fmt"
)

// SoftwarePackages contains a list of software packages managed by Junos Space.
type SoftwarePackages struct {
	Packages []SoftwarePackage `xml:"package"`
}

// A SoftwarePackage contains information about each individual software package.
type SoftwarePackage struct {
	ID       int    `xml:"key,attr"`
	Name     string `xml:"fileName"`
	Version  string `xml:"version"`
	Platform string `xml:"platformType"`
}

// SoftwareUpgrade consists of options available to use before issuing a software upgrade.
type SoftwareUpgrade struct {
	UseDownloaded bool // Use an image already staged on the device.
	Validate      bool // Check/don't check compatibility with current configuration.
	Reboot        bool // Reboot system after adding package.
	RebootAfter   int  // Reboot the system after "x" minutes.
	Cleanup       bool // Remove any pre-existing packages on the device.
	RemoveAfter   bool // Remove the package after successful installation.
}

// deployXML is XML we send (POST) for image deployment.
var deployXML = `
<exec-deploy>
    <devices>
        <device href= "/api/space/device-management/devices/%d"/>
    </devices>
    <deployOptions>
        <useAlreadyDownloaded>%t</useAlreadyDownloaded>
        <validate>%t</validate>
        <bestEffortLoad>false</bestEffortLoad>
        <snapShotRequired>false</snapShotRequired>
        <rebootDevice>%t</rebootDevice>
        <rebootAfterXMinutes>%d</rebootAfterXMinutes>
        <cleanUpExistingOnDevice>%t</cleanUpExistingOnDevice>
        <removePkgAfterInstallation>%t</removePkgAfterInstallation>
    </deployOptions>
</exec-deploy>
`

// removeStagedXML is XML we send (POST) for removing a staged image.
var removeStagedXML = `
<exec-remove>
    <devices>
        <device href="/api/space/device-management/devices/%d"/>
    </devices>
</exec-remove>
`

// stageXML is XML we send (POST) for staging an image on a device.
var stageXML = `
<exec-stage>
    <devices>
        <device href="/api/space/device-management/devices/%d"/>
    </devices>
    <stageOptions>
        <cleanUpExistingOnDevice>%t</cleanUpExistingOnDevice>
    </stageOptions>
</exec-stage>
`

// getSoftwareID returns the ID of the software package.
func (s *JunosSpace) getSoftwareID(image string) (int, error) {
	var err error
	var softwareID int
	images, err := s.Software()
	if err != nil {
		return 0, err
	}

	for _, sw := range images.Packages {
		if sw.Name == image {
			softwareID = sw.ID
		}
	}

	return softwareID, nil
}

// DeploySoftware starts the upgrade process on the device, using the given image along
// with the options specified.
func (s *JunosSpace) DeploySoftware(device, image string, options *SoftwareUpgrade) (int, error) {
	var job jobID
	deviceID, _ := s.getDeviceID(device)
	softwareID, _ := s.getSoftwareID(image)
	deploy := fmt.Sprintf(deployXML, deviceID, options.UseDownloaded, options.Validate, options.Reboot, options.RebootAfter, options.Cleanup, options.RemoveAfter)
	req := &APIRequest{
		Method:      "post",
		URL:         fmt.Sprintf("/api/space/software-management/packages/%d/exec-deploy", softwareID),
		Body:        deploy,
		ContentType: contentExecDeploy,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}

// RemoveStagedSoftware will delete the staged software image on the device.
func (s *JunosSpace) RemoveStagedSoftware(device, image string) (int, error) {
	var job jobID
	deviceID, _ := s.getDeviceID(device)
	softwareID, _ := s.getSoftwareID(image)
	remove := fmt.Sprintf(removeStagedXML, deviceID)
	req := &APIRequest{
		Method:      "post",
		URL:         fmt.Sprintf("/api/space/software-management/packages/%d/exec-remove", softwareID),
		Body:        remove,
		ContentType: contentExecRemove,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}

// Software queries the Junos Space server and returns all of the information
// about each software image that Space manages.
func (s *JunosSpace) Software() (*SoftwarePackages, error) {
	var software SoftwarePackages
	req := &APIRequest{
		Method: "get",
		URL:    "/api/space/software-management/packages",
	}
	data, err := s.APICall(req)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(data, &software)
	if err != nil {
		return nil, err
	}

	return &software, nil
}

// StageSoftware loads the given software image onto the device but does not
// upgrade it. The package is placed in the /var/tmp directory.
func (s *JunosSpace) StageSoftware(device, image string, cleanup bool) (int, error) {
	var job jobID
	deviceID, _ := s.getDeviceID(device)
	softwareID, _ := s.getSoftwareID(image)
	stage := fmt.Sprintf(stageXML, deviceID, cleanup)
	req := &APIRequest{
		Method:      "post",
		URL:         fmt.Sprintf("/api/space/software-management/packages/%d/exec-stage", softwareID),
		Body:        stage,
		ContentType: contentExecStage,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}
<|endoftext|>


<commit_before>package auth

import (
	"time"

	"golang.org/x/net/context"
)

// AuthenticationProvider provides helper methods to convert tokens to sessions
// using our own internal authorization services
type AuthenticationProvider interface {
	// RecoverSession from a given access token, converting this into a set of credentials
	RecoverCredentials(ctx context.Context, accessToken string) (Credentials, error)
}

// Credentials
type Credentials interface {
	AccessToken() string
	RefreshToken() string
	Expiry() time.Time
	Scopes() []string
}

// Authorizer provides an interface to validate authorization credentials
// for access to resources, eg. oauth scopes, or other access control
type Authorizer func(ctx context.Context, creds Credentials) error
<commit_msg>Add User and Client to the credential interface<commit_after>package auth

import (
	"time"

	"golang.org/x/net/context"
)

// AuthenticationProvider provides helper methods to convert tokens to sessions
// using our own internal authorization services
type AuthenticationProvider interface {
	// RecoverSession from a given access token, converting this into a set of credentials
	RecoverCredentials(ctx context.Context, accessToken string) (Credentials, error)
}

// Credentials
type Credentials interface {
	AccessToken() string
	RefreshToken() string
	Expiry() time.Time
	Scopes() []string // aggregated scope information from a combination of the user and client scopes
	User() User
	Client() Client
}

// Authorizer provides an interface to validate authorization credentials
// for access to resources, eg. oauth scopes, or other access control
type Authorizer func(ctx context.Context, creds Credentials) error

// User represents the resource owner ie. an end-user of the application
type User interface {
	ID() string
	Scopes() []string
}

// Client represents the application making a request on behalf of a User
type Client interface {
	ID() string
	Scopes() []string
}
<|endoftext|>


<commit_before>package auth

import (
	"crypto/md5"
	"crypto/sha256"
	"fmt"
	"math/rand"
	"net/http"
	"strings"
	"time"

	"github.com/subutai-io/agent/log"

	"github.com/subutai-io/gorjun/db"
	"github.com/subutai-io/gorjun/pgp"
)

func Register(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		r.ParseMultipartForm(32 << 20)
		if strings.Split(r.RemoteAddr, ":")[0] == "127.0.0.1" && len(r.MultipartForm.Value["name"]) > 0 && len(r.MultipartForm.Value["key"]) > 0 {
			name := r.MultipartForm.Value["name"][0]
			key := r.MultipartForm.Value["key"][0]

			w.Write([]byte("Name: " + name + "\n"))
			w.Write([]byte("PGP key: " + key + "\n"))

			db.RegisterUser([]byte(name), []byte(key))
			return
		} else if len(r.MultipartForm.Value["key"]) > 0 {
			key := pgp.Verify("Hub", r.MultipartForm.Value["key"][0])
			if len(key) == 0 {
				w.Write([]byte("Signature check failed"))
				w.WriteHeader(http.StatusForbidden)
				return
			}
			fingerprint := pgp.Fingerprint(key)
			if len(fingerprint) == 0 {
				w.Write([]byte("Filed to get key fingerprint"))
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			db.RegisterUser([]byte(fmt.Sprintf("%x", fingerprint)), []byte(key))
			return
		}
	}
	w.WriteHeader(http.StatusUnauthorized)
	w.Write([]byte("Not allowed"))
}

func Token(w http.ResponseWriter, r *http.Request) {
	rand.Seed(time.Now().UnixNano())
	if r.Method == http.MethodGet {
		name := r.URL.Query().Get("user")
		if len(name) != 0 {
			hash := md5.New()
			hash.Write([]byte(fmt.Sprint(time.Now().String(), name, rand.Float64())))
			authID := fmt.Sprintf("%x", hash.Sum(nil))
			db.SaveAuthID(name, authID)
			w.Write([]byte(authID))
		}
	} else if r.Method == http.MethodPost {
		name := r.FormValue("user")
		message := r.FormValue("message")
		if len(name) == 0 || len(message) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Please specify user name and auth message"))
			log.Warn(r.RemoteAddr + " - empty user name or message filed")
			return
		}
		authid := pgp.Verify(name, message)
		if db.CheckAuthID(authid) == name {
			token := fmt.Sprintf("%x", sha256.Sum256([]byte(fmt.Sprint(time.Now().String(), name, rand.Float64()))))
			db.SaveToken(name, fmt.Sprintf("%x", sha256.Sum256([]byte(token))))
			w.Write([]byte(token))
		} else {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte("Signature verification failed"))
		}
	}
}

func Validate(w http.ResponseWriter, r *http.Request) {
	token := r.URL.Query().Get("token")
	if len(token) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty token"))
		return
	}
	if len(db.CheckToken(token)) == 0 {
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte("Forbidden"))
		return
	}
	w.Write([]byte("Success"))
}

func Key(w http.ResponseWriter, r *http.Request) {
	user := r.URL.Query().Get("user")
	if len(user) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty user"))
		return
	}
	key := db.UserKey(user)
	if len(key) == 0 {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("User key not found"))
		return
	}
	w.Write([]byte(key))
}

func Sign(w http.ResponseWriter, r *http.Request) {
	r.ParseMultipartForm(32 << 20)
	if len(r.MultipartForm.Value["token"]) == 0 || len(db.CheckToken(r.MultipartForm.Value["token"][0])) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Not authorized"))
		log.Warn(r.RemoteAddr + " - rejecting unauthorized sign request")
		return
	}
	owner := db.CheckToken(r.MultipartForm.Value["token"][0])
	if len(r.MultipartForm.Value["signature"]) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty signature"))
		log.Warn("auth.Sign received empty signature")
		return
	}
	signature := r.MultipartForm.Value["signature"][0]
	hash := pgp.Verify(owner, signature)
	if len(hash) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Failed to verify signature with user key"))
		log.Warn("Failed to verify signature with user key")
		return
	}
	info := db.Info(hash)
	repo := strings.Split(r.URL.EscapedPath(), "/")
	if len(repo) < 4 {
		log.Warn(r.URL.EscapedPath() + " - bad share request")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Bad request"))
		return
	}
	if db.CheckRepo(owner, repo[3], hash) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("File and signature have different owner"))
		log.Warn("File and signature have different owner")
		return
	}
	db.Write(owner, hash, signature)
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("File has been signed"))
	log.Info("File " + info["Name"] + "(" + hash + ")" + " has been signed by " + owner)
	return
}
<commit_msg>File signing fixed<commit_after>package auth

import (
	"crypto/md5"
	"crypto/sha256"
	"fmt"
	"math/rand"
	"net/http"
	"strings"
	"time"

	"github.com/subutai-io/agent/log"

	"github.com/subutai-io/gorjun/db"
	"github.com/subutai-io/gorjun/pgp"
)

func Register(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		r.ParseMultipartForm(32 << 20)
		if strings.Split(r.RemoteAddr, ":")[0] == "127.0.0.1" && len(r.MultipartForm.Value["name"]) > 0 && len(r.MultipartForm.Value["key"]) > 0 {
			name := r.MultipartForm.Value["name"][0]
			key := r.MultipartForm.Value["key"][0]

			w.Write([]byte("Name: " + name + "\n"))
			w.Write([]byte("PGP key: " + key + "\n"))

			db.RegisterUser([]byte(name), []byte(key))
			return
		} else if len(r.MultipartForm.Value["key"]) > 0 {
			key := pgp.Verify("Hub", r.MultipartForm.Value["key"][0])
			if len(key) == 0 {
				w.Write([]byte("Signature check failed"))
				w.WriteHeader(http.StatusForbidden)
				return
			}
			fingerprint := pgp.Fingerprint(key)
			if len(fingerprint) == 0 {
				w.Write([]byte("Filed to get key fingerprint"))
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			db.RegisterUser([]byte(fmt.Sprintf("%x", fingerprint)), []byte(key))
			return
		}
	}
	w.WriteHeader(http.StatusUnauthorized)
	w.Write([]byte("Not allowed"))
}

func Token(w http.ResponseWriter, r *http.Request) {
	rand.Seed(time.Now().UnixNano())
	if r.Method == http.MethodGet {
		name := r.URL.Query().Get("user")
		if len(name) != 0 {
			hash := md5.New()
			hash.Write([]byte(fmt.Sprint(time.Now().String(), name, rand.Float64())))
			authID := fmt.Sprintf("%x", hash.Sum(nil))
			db.SaveAuthID(name, authID)
			w.Write([]byte(authID))
		}
	} else if r.Method == http.MethodPost {
		name := r.FormValue("user")
		message := r.FormValue("message")
		if len(name) == 0 || len(message) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Please specify user name and auth message"))
			log.Warn(r.RemoteAddr + " - empty user name or message filed")
			return
		}
		authid := pgp.Verify(name, message)
		if db.CheckAuthID(authid) == name {
			token := fmt.Sprintf("%x", sha256.Sum256([]byte(fmt.Sprint(time.Now().String(), name, rand.Float64()))))
			db.SaveToken(name, fmt.Sprintf("%x", sha256.Sum256([]byte(token))))
			w.Write([]byte(token))
		} else {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte("Signature verification failed"))
		}
	}
}

func Validate(w http.ResponseWriter, r *http.Request) {
	token := r.URL.Query().Get("token")
	if len(token) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty token"))
		return
	}
	if len(db.CheckToken(token)) == 0 {
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte("Forbidden"))
		return
	}
	w.Write([]byte("Success"))
}

func Key(w http.ResponseWriter, r *http.Request) {
	user := r.URL.Query().Get("user")
	if len(user) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty user"))
		return
	}
	key := db.UserKey(user)
	if len(key) == 0 {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("User key not found"))
		return
	}
	w.Write([]byte(key))
}

func Sign(w http.ResponseWriter, r *http.Request) {
	r.ParseMultipartForm(32 << 20)
	if len(r.MultipartForm.Value["token"]) == 0 || len(db.CheckToken(r.MultipartForm.Value["token"][0])) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Not authorized"))
		log.Warn(r.RemoteAddr + " - rejecting unauthorized sign request")
		return
	}
	owner := db.CheckToken(r.MultipartForm.Value["token"][0])
	if len(r.MultipartForm.Value["signature"]) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty signature"))
		log.Warn("auth.Sign received empty signature")
		return
	}
	signature := r.MultipartForm.Value["signature"][0]
	hash := pgp.Verify(owner, signature)
	if len(hash) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Failed to verify signature with user key"))
		log.Warn("Failed to verify signature with user key")
		return
	}
	if db.CheckRepo(owner, "", hash) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("File and signature have different owner"))
		log.Warn("File and signature have different owner")
		return
	}
	db.Write(owner, hash, signature)
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("File has been signed"))
	log.Info("File " + hash + " has been signed by " + owner)
	return
}
<|endoftext|>


<commit_before>package ofutils

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/smtp"
)

func SendQQMail(email, password, host string, port int, from, toEmail, subject, body string) error {
	header := make(map[string]string)
	header["From"] = from + " " + "<" + email + ">"
	header["To"] = toEmail
	header["Subject"] = subject
	header["Content-Type"] = "text/html; charset=UTF-8"

	message := ""
	for k, v := range header {
		message += fmt.Sprintf("%s: %s\r\n", k, v)
	}
	message += "\r\n" + body

	auth := smtp.PlainAuth(
		"",
		email,
		password,
		host,
	)

	return SendMailUsingTLS(
		fmt.Sprintf("%s:%d", host, port),
		auth,
		email,
		[]string{toEmail},
		[]byte(message),
	)
}

//return a smtp client
func Dial(addr string) (*smtp.Client, error) {
	conn, err := tls.Dial("tcp", addr, nil)
	if err != nil {
		log.Println("Dialing Error:", err)
		return nil, err
	}
	//分解主机端口字符串
	host, _, _ := net.SplitHostPort(addr)
	return smtp.NewClient(conn, host)
}

//参考net/smtp的func SendMail()
//使用net.Dial连接tls(ssl)端口时,smtp.NewClient()会卡住且不提示err
//len(to)>1时,to[1]开始提示是密送
func SendMailUsingTLS(addr string, auth smtp.Auth, from string,
	to []string, msg []byte) (err error) {

	//create smtp client
	c, err := Dial(addr)
	if err != nil {
		log.Println("Create smpt client error:", err)
		return err
	}
	defer c.Close()

	if auth != nil {
		if ok, _ := c.Extension("AUTH"); ok {
			if err = c.Auth(auth); err != nil {
				log.Println("Error during AUTH", err)
				return err
			}
		}
	}

	if err = c.Mail(from); err != nil {
		return err
	}

	for _, addr := range to {
		if err = c.Rcpt(addr); err != nil {
			return err
		}
	}

	w, err := c.Data()
	if err != nil {
		return err
	}

	_, err = w.Write(msg)
	if err != nil {
		return err
	}

	err = w.Close()
	if err != nil {
		return err
	}

	return c.Quit()
}
<commit_msg>update<commit_after>package ofutils

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/smtp"
)

func SendQQMail(email, password, host string, port int, from, toEmail, subject, body string) error {
	header := make(map[string]string)
	header["From"] = from + " " + "<" + email + ">"
	header["To"] = toEmail
	header["Subject"] = subject
	header["Content-Type"] = "text/html; charset=GBK"

	message := ""
	for k, v := range header {
		message += fmt.Sprintf("%s: %s\r\n", k, v)
	}
	message += "\r\n" + body

	auth := smtp.PlainAuth(
		"",
		email,
		password,
		host,
	)

	return SendMailUsingTLS(
		fmt.Sprintf("%s:%d", host, port),
		auth,
		email,
		[]string{toEmail},
		[]byte(message),
	)
}

//return a smtp client
func Dial(addr string) (*smtp.Client, error) {
	conn, err := tls.Dial("tcp", addr, nil)
	if err != nil {
		log.Println("Dialing Error:", err)
		return nil, err
	}
	//分解主机端口字符串
	host, _, _ := net.SplitHostPort(addr)
	return smtp.NewClient(conn, host)
}

//参考net/smtp的func SendMail()
//使用net.Dial连接tls(ssl)端口时,smtp.NewClient()会卡住且不提示err
//len(to)>1时,to[1]开始提示是密送
func SendMailUsingTLS(addr string, auth smtp.Auth, from string,
	to []string, msg []byte) (err error) {

	//create smtp client
	c, err := Dial(addr)
	if err != nil {
		log.Println("Create smpt client error:", err)
		return err
	}
	defer c.Close()

	if auth != nil {
		if ok, _ := c.Extension("AUTH"); ok {
			if err = c.Auth(auth); err != nil {
				log.Println("Error during AUTH", err)
				return err
			}
		}
	}

	if err = c.Mail(from); err != nil {
		return err
	}

	for _, addr := range to {
		if err = c.Rcpt(addr); err != nil {
			return err
		}
	}

	w, err := c.Data()
	if err != nil {
		return err
	}

	_, err = w.Write(msg)
	if err != nil {
		return err
	}

	err = w.Close()
	if err != nil {
		return err
	}

	return c.Quit()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"flag"
	"log"
	"os"

	// Load the common drivers
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"

	// Load sqlx over database/sql
	"github.com/jmoiron/sqlx"

	"github.com/StabbyCutyou/sqltocsv/converters"
)

type config struct {
	dbAdapter       string
	connString      string
	sqlQuery        string
	outputFile      string
	delimeter       string
	obfuscateFields string
	quoteFields     string
	quoteType       string
}

var delimeters = map[string]rune{
	"tab": rune('	'), //That's a tab in there, yo
	"comma": rune(','),
}

func main() {
	cfg := getConfig()
	db, err := sqlx.Open(cfg.dbAdapter, cfg.connString)
	if err != nil {
		log.Fatal(err)
	}

	results, err := db.Queryx(cfg.sqlQuery)
	if err != nil {
		log.Fatal(err)
	}

	csvWriter := csv.NewWriter(os.Stdout)
	if comma, ok := delimeters[cfg.delimeter]; ok {
		csvWriter.Comma = comma
	} else {
		log.Printf("Warning: No known delimeter for %s, defaulting to Comma", cfg.delimeter)
	}

	converter := converters.GetConverter(cfg.dbAdapter)

	count := 0
	for results.Next() {
		row, err := results.SliceScan()
		if err != nil {
			log.Fatal(err)
		}
		// Only do this for the first line, aka the headers
		if count == 0 {
			cols, err := results.Columns()
			if err != nil {
				log.Fatal(err)
			}
			csvWriter.Write(cols)
		}

		rowStrings := make([]string, len(row))
		// It seems for mysql, the case is always []byte of a string?
		for i, col := range row {
			val, err := converter.ColumnToString(col)
			if err != nil {
				log.Fatal(err)
			}
			// Inject quoting, obfuscating here
			rowStrings[i] = val
		}
		csvWriter.Write(rowStrings)
		count++
	}

	csvWriter.Flush()
	log.Printf("\nFinished processing %d lines\n", count)
}

func getConfig() *config {
	d := flag.String("d", "mysql", "The (d)atabase adapter to use")
	c := flag.String("c", "", "The (c)onnection string to use")
	q := flag.String("q", "", "The (q)uery to use")
	m := flag.String("m", "comma", "The deli(m)eter to use: 'comma' or 'tab'. Defaults to 'comma'")
	o := flag.String("o", "", "The fields to (o)bfuscate")
	w := flag.String("w", "", "The fields to (w)rap in quotes")
	t := flag.String("t", "double", "The (t)ype of quote to use with -w: 'single' or 'double'. Defaults to 'double'")

	flag.Parse()

	if *q == "" {
		log.Fatal("You must provide query via -q")
	}
	if *c == "" {
		log.Fatal("You must provide a connection string via -c")
	}

	return &config{
		dbAdapter:       *d,
		connString:      *c,
		sqlQuery:        *q,
		obfuscateFields: *o,
		delimeter:       *m,
		quoteFields:     *w,
		quoteType:       *t,
	}
}

//SELECT * FROM users WHERE created_at >= '2015-01-01 00:00:00' AND created_at < '2015-02-01 00:00:00'
<commit_msg>Spelling error lol<commit_after>package main

import (
	"encoding/csv"
	"flag"
	"log"
	"os"

	// Load the common drivers
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"

	// Load sqlx over database/sql
	"github.com/jmoiron/sqlx"

	"github.com/StabbyCutyou/sqltocsv/converters"
)

type config struct {
	dbAdapter       string
	connString      string
	sqlQuery        string
	outputFile      string
	delimiter       string
	obfuscateFields string
	quoteFields     string
	quoteType       string
}

var delimiters = map[string]rune{
	"tab": rune('	'), //That's a tab in there, yo
	"comma": rune(','),
}

func main() {
	cfg := getConfig()
	db, err := sqlx.Open(cfg.dbAdapter, cfg.connString)
	if err != nil {
		log.Fatal(err)
	}

	results, err := db.Queryx(cfg.sqlQuery)
	if err != nil {
		log.Fatal(err)
	}

	csvWriter := csv.NewWriter(os.Stdout)
	if comma, ok := delimiters[cfg.delimiter]; ok {
		csvWriter.Comma = comma
	} else {
		log.Printf("Warning: No known delimiter for %s, defaulting to Comma", cfg.delimiter)
	}

	converter := converters.GetConverter(cfg.dbAdapter)

	count := 0
	for results.Next() {
		row, err := results.SliceScan()
		if err != nil {
			log.Fatal(err)
		}
		// Only do this for the first line, aka the headers
		if count == 0 {
			cols, err := results.Columns()
			if err != nil {
				log.Fatal(err)
			}
			csvWriter.Write(cols)
		}

		rowStrings := make([]string, len(row))
		// It seems for mysql, the case is always []byte of a string?
		for i, col := range row {
			val, err := converter.ColumnToString(col)
			if err != nil {
				log.Fatal(err)
			}
			// Inject quoting, obfuscating here
			rowStrings[i] = val
		}
		csvWriter.Write(rowStrings)
		count++
	}

	csvWriter.Flush()
	log.Printf("\nFinished processing %d lines\n", count)
}

func getConfig() *config {
	d := flag.String("d", "mysql", "The (d)atabase adapter to use")
	c := flag.String("c", "", "The (c)onnection string to use")
	q := flag.String("q", "", "The (q)uery to use")
	m := flag.String("m", "comma", "The deli(m)iter to use: 'comma' or 'tab'. Defaults to 'comma'")
	o := flag.String("o", "", "The fields to (o)bfuscate")
	w := flag.String("w", "", "The fields to (w)rap in quotes")
	t := flag.String("t", "double", "The (t)ype of quote to use with -w: 'single' or 'double'. Defaults to 'double'")

	flag.Parse()

	if *q == "" {
		log.Fatal("You must provide query via -q")
	}
	if *c == "" {
		log.Fatal("You must provide a connection string via -c")
	}

	return &config{
		dbAdapter:       *d,
		connString:      *c,
		sqlQuery:        *q,
		obfuscateFields: *o,
		delimiter:       *m,
		quoteFields:     *w,
		quoteType:       *t,
	}
}

//SELECT * FROM users WHERE created_at >= '2015-01-01 00:00:00' AND created_at < '2015-02-01 00:00:00'
<|endoftext|>


<commit_before>// dedupe - gets rid of identical files remotes which can have duplicate file names (drive, mega)

package operations

import (
	"context"
	"fmt"
	"log"
	"path"
	"sort"
	"strings"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/config"
	"github.com/rclone/rclone/fs/hash"
	"github.com/rclone/rclone/fs/walk"
	"github.com/spf13/pflag"
)

// dedupeRename renames the objs slice to different names
func dedupeRename(ctx context.Context, f fs.Fs, remote string, objs []fs.Object) {
	doMove := f.Features().Move
	if doMove == nil {
		log.Fatalf("Fs %v doesn't support Move", f)
	}
	ext := path.Ext(remote)
	base := remote[:len(remote)-len(ext)]

outer:
	for i, o := range objs {
		suffix := 1
		newName := fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
		_, err := f.NewObject(ctx, newName)
		for ; err != fs.ErrorObjectNotFound; suffix++ {
			if err != nil {
				err = fs.CountError(err)
				fs.Errorf(o, "Failed to check for existing object: %v", err)
				continue outer
			}
			if suffix > 100 {
				fs.Errorf(o, "Could not find an available new name")
				continue outer
			}
			newName = fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
			_, err = f.NewObject(ctx, newName)
		}
		if !fs.Config.DryRun {
			newObj, err := doMove(ctx, o, newName)
			if err != nil {
				err = fs.CountError(err)
				fs.Errorf(o, "Failed to rename: %v", err)
				continue
			}
			fs.Infof(newObj, "renamed from: %v", o)
		} else {
			fs.Logf(remote, "Not renaming to %q as --dry-run", newName)
		}
	}
}

// dedupeDeleteAllButOne deletes all but the one in keep
func dedupeDeleteAllButOne(ctx context.Context, keep int, remote string, objs []fs.Object) {
	count := 0
	for i, o := range objs {
		if i == keep {
			continue
		}
		err := DeleteFile(ctx, o)
		if err == nil {
			count++
		}
	}
	if count > 0 {
		fs.Logf(remote, "Deleted %d extra copies", count)
	}
}

// dedupeDeleteIdentical deletes all but one of identical (by hash) copies
func dedupeDeleteIdentical(ctx context.Context, ht hash.Type, remote string, objs []fs.Object) (remainingObjs []fs.Object) {
	// See how many of these duplicates are identical
	byHash := make(map[string][]fs.Object, len(objs))
	for _, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil || md5sum == "" {
			remainingObjs = append(remainingObjs, o)
		} else {
			byHash[md5sum] = append(byHash[md5sum], o)
		}
	}

	// Delete identical duplicates, filling remainingObjs with the ones remaining
	for md5sum, hashObjs := range byHash {
		remainingObjs = append(remainingObjs, hashObjs[0])
		if len(hashObjs) > 1 {
			fs.Logf(remote, "Deleting %d/%d identical duplicates (%v %q)", len(hashObjs)-1, len(hashObjs), ht, md5sum)
			for _, o := range hashObjs[1:] {
				err := DeleteFile(ctx, o)
				if err != nil {
					remainingObjs = append(remainingObjs, o)
				}
			}
		}
	}

	return remainingObjs
}

// dedupeInteractive interactively dedupes the slice of objects
func dedupeInteractive(ctx context.Context, f fs.Fs, ht hash.Type, remote string, objs []fs.Object) {
	fmt.Printf("%s: %d duplicates remain\n", remote, len(objs))
	for i, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil {
			md5sum = err.Error()
		}
		fmt.Printf("  %d: %12d bytes, %s, %v %32s\n", i+1, o.Size(), o.ModTime(ctx).Local().Format("2006-01-02 15:04:05.000000000"), ht, md5sum)
	}
	switch config.Command([]string{"sSkip and do nothing", "kKeep just one (choose which in next step)", "rRename all to be different (by changing file.jpg to file-1.jpg)"}) {
	case 's':
	case 'k':
		keep := config.ChooseNumber("Enter the number of the file to keep", 1, len(objs))
		dedupeDeleteAllButOne(ctx, keep-1, remote, objs)
	case 'r':
		dedupeRename(ctx, f, remote, objs)
	}
}

type objectsSortedByModTime []fs.Object

func (objs objectsSortedByModTime) Len() int      { return len(objs) }
func (objs objectsSortedByModTime) Swap(i, j int) { objs[i], objs[j] = objs[j], objs[i] }
func (objs objectsSortedByModTime) Less(i, j int) bool {
	return objs[i].ModTime(context.TODO()).Before(objs[j].ModTime(context.TODO()))
}

// DeduplicateMode is how the dedupe command chooses what to do
type DeduplicateMode int

// Deduplicate modes
const (
	DeduplicateInteractive DeduplicateMode = iota // interactively ask the user
	DeduplicateSkip                               // skip all conflicts
	DeduplicateFirst                              // choose the first object
	DeduplicateNewest                             // choose the newest object
	DeduplicateOldest                             // choose the oldest object
	DeduplicateRename                             // rename the objects
	DeduplicateLargest                            // choose the largest object
)

func (x DeduplicateMode) String() string {
	switch x {
	case DeduplicateInteractive:
		return "interactive"
	case DeduplicateSkip:
		return "skip"
	case DeduplicateFirst:
		return "first"
	case DeduplicateNewest:
		return "newest"
	case DeduplicateOldest:
		return "oldest"
	case DeduplicateRename:
		return "rename"
	case DeduplicateLargest:
		return "largest"
	}
	return "unknown"
}

// Set a DeduplicateMode from a string
func (x *DeduplicateMode) Set(s string) error {
	switch strings.ToLower(s) {
	case "interactive":
		*x = DeduplicateInteractive
	case "skip":
		*x = DeduplicateSkip
	case "first":
		*x = DeduplicateFirst
	case "newest":
		*x = DeduplicateNewest
	case "oldest":
		*x = DeduplicateOldest
	case "rename":
		*x = DeduplicateRename
	case "largest":
		*x = DeduplicateLargest
	default:
		return errors.Errorf("Unknown mode for dedupe %q.", s)
	}
	return nil
}

// Type of the value
func (x *DeduplicateMode) Type() string {
	return "string"
}

// Check it satisfies the interface
var _ pflag.Value = (*DeduplicateMode)(nil)

// dedupeFindDuplicateDirs scans f for duplicate directories
func dedupeFindDuplicateDirs(ctx context.Context, f fs.Fs) ([][]fs.Directory, error) {
	dirs := map[string][]fs.Directory{}
	err := walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListDirs, func(entries fs.DirEntries) error {
		entries.ForDir(func(d fs.Directory) {
			dirs[d.Remote()] = append(dirs[d.Remote()], d)
		})
		return nil
	})
	if err != nil {
		return nil, errors.Wrap(err, "find duplicate dirs")
	}
	duplicateDirs := [][]fs.Directory{}
	for _, ds := range dirs {
		if len(ds) > 1 {
			duplicateDirs = append(duplicateDirs, ds)
		}
	}
	return duplicateDirs, nil
}

// dedupeMergeDuplicateDirs merges all the duplicate directories found
func dedupeMergeDuplicateDirs(ctx context.Context, f fs.Fs, duplicateDirs [][]fs.Directory) error {
	mergeDirs := f.Features().MergeDirs
	if mergeDirs == nil {
		return errors.Errorf("%v: can't merge directories", f)
	}
	dirCacheFlush := f.Features().DirCacheFlush
	if dirCacheFlush == nil {
		return errors.Errorf("%v: can't flush dir cache", f)
	}
	for _, dirs := range duplicateDirs {
		if !fs.Config.DryRun {
			fs.Infof(dirs[0], "Merging contents of duplicate directories")
			err := mergeDirs(ctx, dirs)
			if err != nil {
				return errors.Wrap(err, "merge duplicate dirs")
			}
		} else {
			fs.Infof(dirs[0], "NOT Merging contents of duplicate directories as --dry-run")
		}
	}
	dirCacheFlush()
	return nil
}

// Deduplicate interactively finds duplicate files and offers to
// delete all but one or rename them to be different. Only useful with
// Google Drive which can have duplicate file names.
func Deduplicate(ctx context.Context, f fs.Fs, mode DeduplicateMode) error {
	fs.Infof(f, "Looking for duplicates using %v mode.", mode)

	// Find duplicate directories first and fix them - repeat
	// until all fixed
	for {
		duplicateDirs, err := dedupeFindDuplicateDirs(ctx, f)
		if err != nil {
			return err
		}
		if len(duplicateDirs) == 0 {
			break
		}
		err = dedupeMergeDuplicateDirs(ctx, f, duplicateDirs)
		if err != nil {
			return err
		}
		if fs.Config.DryRun {
			break
		}
	}

	// find a hash to use
	ht := f.Hashes().GetOne()

	// Now find duplicate files
	files := map[string][]fs.Object{}
	err := walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListObjects, func(entries fs.DirEntries) error {
		entries.ForObject(func(o fs.Object) {
			remote := o.Remote()
			files[remote] = append(files[remote], o)
		})
		return nil
	})
	if err != nil {
		return err
	}

	for remote, objs := range files {
		if len(objs) > 1 {
			fs.Logf(remote, "Found %d duplicates - deleting identical copies", len(objs))
			objs = dedupeDeleteIdentical(ctx, ht, remote, objs)
			if len(objs) <= 1 {
				fs.Logf(remote, "All duplicates removed")
				continue
			}
			switch mode {
			case DeduplicateInteractive:
				dedupeInteractive(ctx, f, ht, remote, objs)
			case DeduplicateFirst:
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateNewest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, len(objs)-1, remote, objs)
			case DeduplicateOldest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateRename:
				dedupeRename(ctx, f, remote, objs)
			case DeduplicateLargest:
				largest, largestIndex := int64(-1), -1
				for i, obj := range objs {
					size := obj.Size()
					if size > largest {
						largest, largestIndex = size, i
					}
				}
				if largestIndex > -1 {
					dedupeDeleteAllButOne(ctx, largestIndex, remote, objs)
				}
			case DeduplicateSkip:
				// skip
			default:
				//skip
			}
		}
	}
	return nil
}
<commit_msg>operations: fix dedupe continuing on errors like insufficientFilePermisson - fixes #3470<commit_after>// dedupe - gets rid of identical files remotes which can have duplicate file names (drive, mega)

package operations

import (
	"context"
	"fmt"
	"log"
	"path"
	"sort"
	"strings"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/config"
	"github.com/rclone/rclone/fs/hash"
	"github.com/rclone/rclone/fs/walk"
	"github.com/spf13/pflag"
)

// dedupeRename renames the objs slice to different names
func dedupeRename(ctx context.Context, f fs.Fs, remote string, objs []fs.Object) {
	doMove := f.Features().Move
	if doMove == nil {
		log.Fatalf("Fs %v doesn't support Move", f)
	}
	ext := path.Ext(remote)
	base := remote[:len(remote)-len(ext)]

outer:
	for i, o := range objs {
		suffix := 1
		newName := fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
		_, err := f.NewObject(ctx, newName)
		for ; err != fs.ErrorObjectNotFound; suffix++ {
			if err != nil {
				err = fs.CountError(err)
				fs.Errorf(o, "Failed to check for existing object: %v", err)
				continue outer
			}
			if suffix > 100 {
				fs.Errorf(o, "Could not find an available new name")
				continue outer
			}
			newName = fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
			_, err = f.NewObject(ctx, newName)
		}
		if !fs.Config.DryRun {
			newObj, err := doMove(ctx, o, newName)
			if err != nil {
				err = fs.CountError(err)
				fs.Errorf(o, "Failed to rename: %v", err)
				continue
			}
			fs.Infof(newObj, "renamed from: %v", o)
		} else {
			fs.Logf(remote, "Not renaming to %q as --dry-run", newName)
		}
	}
}

// dedupeDeleteAllButOne deletes all but the one in keep
func dedupeDeleteAllButOne(ctx context.Context, keep int, remote string, objs []fs.Object) {
	count := 0
	for i, o := range objs {
		if i == keep {
			continue
		}
		err := DeleteFile(ctx, o)
		if err == nil {
			count++
		}
	}
	if count > 0 {
		fs.Logf(remote, "Deleted %d extra copies", count)
	}
}

// dedupeDeleteIdentical deletes all but one of identical (by hash) copies
func dedupeDeleteIdentical(ctx context.Context, ht hash.Type, remote string, objs []fs.Object) (remainingObjs []fs.Object) {
	// See how many of these duplicates are identical
	byHash := make(map[string][]fs.Object, len(objs))
	for _, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil || md5sum == "" {
			remainingObjs = append(remainingObjs, o)
		} else {
			byHash[md5sum] = append(byHash[md5sum], o)
		}
	}

	// Delete identical duplicates, filling remainingObjs with the ones remaining
	for md5sum, hashObjs := range byHash {
		remainingObjs = append(remainingObjs, hashObjs[0])
		if len(hashObjs) > 1 {
			fs.Logf(remote, "Deleting %d/%d identical duplicates (%v %q)", len(hashObjs)-1, len(hashObjs), ht, md5sum)
			for _, o := range hashObjs[1:] {
				err := DeleteFile(ctx, o)
				if err != nil {
					remainingObjs = append(remainingObjs, o)
				}
			}
		}
	}

	return remainingObjs
}

// dedupeInteractive interactively dedupes the slice of objects
func dedupeInteractive(ctx context.Context, f fs.Fs, ht hash.Type, remote string, objs []fs.Object) {
	fmt.Printf("%s: %d duplicates remain\n", remote, len(objs))
	for i, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil {
			md5sum = err.Error()
		}
		fmt.Printf("  %d: %12d bytes, %s, %v %32s\n", i+1, o.Size(), o.ModTime(ctx).Local().Format("2006-01-02 15:04:05.000000000"), ht, md5sum)
	}
	switch config.Command([]string{"sSkip and do nothing", "kKeep just one (choose which in next step)", "rRename all to be different (by changing file.jpg to file-1.jpg)"}) {
	case 's':
	case 'k':
		keep := config.ChooseNumber("Enter the number of the file to keep", 1, len(objs))
		dedupeDeleteAllButOne(ctx, keep-1, remote, objs)
	case 'r':
		dedupeRename(ctx, f, remote, objs)
	}
}

type objectsSortedByModTime []fs.Object

func (objs objectsSortedByModTime) Len() int      { return len(objs) }
func (objs objectsSortedByModTime) Swap(i, j int) { objs[i], objs[j] = objs[j], objs[i] }
func (objs objectsSortedByModTime) Less(i, j int) bool {
	return objs[i].ModTime(context.TODO()).Before(objs[j].ModTime(context.TODO()))
}

// DeduplicateMode is how the dedupe command chooses what to do
type DeduplicateMode int

// Deduplicate modes
const (
	DeduplicateInteractive DeduplicateMode = iota // interactively ask the user
	DeduplicateSkip                               // skip all conflicts
	DeduplicateFirst                              // choose the first object
	DeduplicateNewest                             // choose the newest object
	DeduplicateOldest                             // choose the oldest object
	DeduplicateRename                             // rename the objects
	DeduplicateLargest                            // choose the largest object
)

func (x DeduplicateMode) String() string {
	switch x {
	case DeduplicateInteractive:
		return "interactive"
	case DeduplicateSkip:
		return "skip"
	case DeduplicateFirst:
		return "first"
	case DeduplicateNewest:
		return "newest"
	case DeduplicateOldest:
		return "oldest"
	case DeduplicateRename:
		return "rename"
	case DeduplicateLargest:
		return "largest"
	}
	return "unknown"
}

// Set a DeduplicateMode from a string
func (x *DeduplicateMode) Set(s string) error {
	switch strings.ToLower(s) {
	case "interactive":
		*x = DeduplicateInteractive
	case "skip":
		*x = DeduplicateSkip
	case "first":
		*x = DeduplicateFirst
	case "newest":
		*x = DeduplicateNewest
	case "oldest":
		*x = DeduplicateOldest
	case "rename":
		*x = DeduplicateRename
	case "largest":
		*x = DeduplicateLargest
	default:
		return errors.Errorf("Unknown mode for dedupe %q.", s)
	}
	return nil
}

// Type of the value
func (x *DeduplicateMode) Type() string {
	return "string"
}

// Check it satisfies the interface
var _ pflag.Value = (*DeduplicateMode)(nil)

// dedupeFindDuplicateDirs scans f for duplicate directories
func dedupeFindDuplicateDirs(ctx context.Context, f fs.Fs) ([][]fs.Directory, error) {
	dirs := map[string][]fs.Directory{}
	err := walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListDirs, func(entries fs.DirEntries) error {
		entries.ForDir(func(d fs.Directory) {
			dirs[d.Remote()] = append(dirs[d.Remote()], d)
		})
		return nil
	})
	if err != nil {
		return nil, errors.Wrap(err, "find duplicate dirs")
	}
	// make sure parents are before children
	duplicateNames := []string{}
	for name, ds := range dirs {
		if len(ds) > 1 {
			duplicateNames = append(duplicateNames, name)
		}
	}
	sort.Strings(duplicateNames)
	duplicateDirs := [][]fs.Directory{}
	for _, name := range duplicateNames {
		duplicateDirs = append(duplicateDirs, dirs[name])
	}
	return duplicateDirs, nil
}

// dedupeMergeDuplicateDirs merges all the duplicate directories found
func dedupeMergeDuplicateDirs(ctx context.Context, f fs.Fs, duplicateDirs [][]fs.Directory) error {
	mergeDirs := f.Features().MergeDirs
	if mergeDirs == nil {
		return errors.Errorf("%v: can't merge directories", f)
	}
	dirCacheFlush := f.Features().DirCacheFlush
	if dirCacheFlush == nil {
		return errors.Errorf("%v: can't flush dir cache", f)
	}
	for _, dirs := range duplicateDirs {
		if !fs.Config.DryRun {
			fs.Infof(dirs[0], "Merging contents of duplicate directories")
			err := mergeDirs(ctx, dirs)
			if err != nil {
				err = fs.CountError(err)
				fs.Errorf(nil, "merge duplicate dirs: %v", err)
			}
		} else {
			fs.Infof(dirs[0], "NOT Merging contents of duplicate directories as --dry-run")
		}
	}
	dirCacheFlush()
	return nil
}

// Deduplicate interactively finds duplicate files and offers to
// delete all but one or rename them to be different. Only useful with
// Google Drive which can have duplicate file names.
func Deduplicate(ctx context.Context, f fs.Fs, mode DeduplicateMode) error {
	fs.Infof(f, "Looking for duplicates using %v mode.", mode)

	// Find duplicate directories first and fix them
	duplicateDirs, err := dedupeFindDuplicateDirs(ctx, f)
	if err != nil {
		return err
	}
	if len(duplicateDirs) != 0 {
		err = dedupeMergeDuplicateDirs(ctx, f, duplicateDirs)
		if err != nil {
			return err
		}
	}

	// find a hash to use
	ht := f.Hashes().GetOne()

	// Now find duplicate files
	files := map[string][]fs.Object{}
	err = walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListObjects, func(entries fs.DirEntries) error {
		entries.ForObject(func(o fs.Object) {
			remote := o.Remote()
			files[remote] = append(files[remote], o)
		})
		return nil
	})
	if err != nil {
		return err
	}

	for remote, objs := range files {
		if len(objs) > 1 {
			fs.Logf(remote, "Found %d duplicates - deleting identical copies", len(objs))
			objs = dedupeDeleteIdentical(ctx, ht, remote, objs)
			if len(objs) <= 1 {
				fs.Logf(remote, "All duplicates removed")
				continue
			}
			switch mode {
			case DeduplicateInteractive:
				dedupeInteractive(ctx, f, ht, remote, objs)
			case DeduplicateFirst:
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateNewest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, len(objs)-1, remote, objs)
			case DeduplicateOldest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateRename:
				dedupeRename(ctx, f, remote, objs)
			case DeduplicateLargest:
				largest, largestIndex := int64(-1), -1
				for i, obj := range objs {
					size := obj.Size()
					if size > largest {
						largest, largestIndex = size, i
					}
				}
				if largestIndex > -1 {
					dedupeDeleteAllButOne(ctx, largestIndex, remote, objs)
				}
			case DeduplicateSkip:
				// skip
			default:
				//skip
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.

// XML
package gxml

import (
    "github.com/clbanning/mxj"
    "encoding/xml"
    "io"
	"gitee.com/wenzi1/gf/g/encoding/gcharset"
	"gitee.com/johng/gf/g/util/gregx"
)

// 将XML内容解析为map变量
func Decode(xmlbyte []byte) (map[string]interface{}, error) {
    Prepare(xmlbyte)
    return mxj.NewMapXml(xmlbyte)
}

// 将map变量解析为XML格式内容
func Encode(v map[string]interface{}, rootTag...string) ([]byte, error) {
    return mxj.Map(v).Xml(rootTag...)
}

func EncodeWithIndent(v map[string]interface{}, rootTag...string) ([]byte, error) {
    return mxj.Map(v).XmlIndent("", "\t", rootTag...)
}

// XML格式内容直接转换为JSON格式内容
func ToJson(xmlbyte []byte) ([]byte, error) {
    Prepare(xmlbyte)
	mv, err := mxj.NewMapXml(xmlbyte)
	if err == nil {
        return mv.Json()
    } else {
        return nil, err
    }
}

//XML字符集预处理
//@author wenzi1
//@date 20180604
func Prepare(xmlbyte []byte) error {
	patten := "<\\?xml\\s+version\\s*=.*?\\s+encoding\\s*=\\s*[\\'|\"](.*?)[\\'|\"]\\s*\\?\\s*>"
	charsetReader := func(charset string, input io.Reader) (io.Reader, error) {
		reader, err := gcharset.GetCharset(charset)
		if err != nil {
			return nil, err
		}
		return reader.NewDecoder().NewReader(input), nil
	}

	matchStr, err := gregx.MatchString(patten, string(xmlbyte))
	if err != nil {
		return err
	}

	charset, err := gcharset.GetCharset(matchStr[1])
	if err != nil {
		return err
	}

	if charset.Name != "UTF-8" {
		mxj.CustomDecoder = &xml.Decoder{Strict:false,CharsetReader:charsetReader}
	}
	return nil
}<commit_msg>修复gxml字符集转换问题<commit_after>// Copyright 2017 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.

// XML
package gxml

import (
    "github.com/clbanning/mxj"
    "encoding/xml"
    "io"
	"gitee.com/johng/gf/g/encoding/gcharset"
	"gitee.com/johng/gf/g/util/gregx"
)

// 将XML内容解析为map变量
func Decode(xmlbyte []byte) (map[string]interface{}, error) {
    Prepare(xmlbyte)
    return mxj.NewMapXml(xmlbyte)
}

// 将map变量解析为XML格式内容
func Encode(v map[string]interface{}, rootTag...string) ([]byte, error) {
    return mxj.Map(v).Xml(rootTag...)
}

func EncodeWithIndent(v map[string]interface{}, rootTag...string) ([]byte, error) {
    return mxj.Map(v).XmlIndent("", "\t", rootTag...)
}

// XML格式内容直接转换为JSON格式内容
func ToJson(xmlbyte []byte) ([]byte, error) {
    Prepare(xmlbyte)
	mv, err := mxj.NewMapXml(xmlbyte)
	if err == nil {
        return mv.Json()
    } else {
        return nil, err
    }
}

//XML字符集预处理
//@author wenzi1
//@date 20180604
func Prepare(xmlbyte []byte) error {
	patten := "<\\?xml\\s+version\\s*=.*?\\s+encoding\\s*=\\s*[\\'|\"](.*?)[\\'|\"]\\s*\\?\\s*>"
	charsetReader := func(charset string, input io.Reader) (io.Reader, error) {
		reader, err := gcharset.GetCharset(charset)
		if err != nil {
			return nil, err
		}
		return reader.NewDecoder().NewReader(input), nil
	}

	matchStr, err := gregx.MatchString(patten, string(xmlbyte))
	if err != nil {
		return err
	}

	charset, err := gcharset.GetCharset(matchStr[1])
	if err != nil {
		return err
	}

	if charset.Name != "UTF-8" {
		mxj.CustomDecoder = &xml.Decoder{Strict:false,CharsetReader:charsetReader}
	}
	return nil
}<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// Package gxml provides accessing and converting for XML content.
//
// XML数据格式解析。
package gxml

import (
    "github.com/gogf/gf/third/github.com/clbanning/mxj"
    "encoding/xml"
    "io"
	"github.com/gogf/gf/g/text/gregex"
	"github.com/gogf/gf/third/github.com/axgle/mahonia"
	"errors"
	"fmt"
	"strings"
)

// 将XML内容解析为map变量
func Decode(content []byte) (map[string]interface{}, error) {
    prepare(content)
    return mxj.NewMapXml(content)
}

// 将map变量解析为XML格式内容
func Encode(v map[string]interface{}, rootTag...string) ([]byte, error) {
    return mxj.Map(v).Xml(rootTag...)
}

func EncodeWithIndent(v map[string]interface{}, rootTag...string) ([]byte, error) {
    return mxj.Map(v).XmlIndent("", "\t", rootTag...)
}

// XML格式内容直接转换为JSON格式内容
func ToJson(content []byte) ([]byte, error) {
    prepare(content)
	mv, err := mxj.NewMapXml(content)
	if err == nil {
        return mv.Json()
    } else {
        return nil, err
    }
}

// XML字符集预处理
// @author wenzi1
// @date 20180604
func prepare(xmlbyte []byte) error {
	patten := `<\?xml.*encoding\s*=\s*['|"](.*?)['|"].*\?>`
	charsetReader := func(charset string, input io.Reader) (io.Reader, error) {
		reader := mahonia.GetCharset(charset)
		if reader == nil {
			return nil, errors.New(fmt.Sprintf("not support charset:%s", charset))
		}
		return reader.NewDecoder().NewReader(input), nil
	}

	matchStr, err := gregex.MatchString(patten, string(xmlbyte))
	if err != nil {
		return err
	}

	xmlEncode := "UTF-8"
	if len(matchStr) == 2 {
		xmlEncode = matchStr[1]
	}

	charset := mahonia.GetCharset(xmlEncode)
	if charset == nil {
		return errors.New(fmt.Sprintf("not support charset:%s", xmlEncode))
	}

	if !strings.EqualFold(charset.Name, "UTF-8") {
		mxj.CustomDecoder = &xml.Decoder{Strict : false, CharsetReader : charsetReader}
	}
	return nil
}<commit_msg>修复并发安全问题,改为如果非UTF8字符集则先做字符集转换<commit_after>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// Package gxml provides accessing and converting for XML content.
//
// XML数据格式解析。
package gxml

import (
	"fmt"
	"github.com/gogf/gf/g/text/gregex"
	"github.com/gogf/gf/third/github.com/axgle/mahonia"
	"github.com/gogf/gf/third/github.com/clbanning/mxj"
	"strings"
)

// 将XML内容解析为map变量
func Decode(content []byte) (map[string]interface{}, error) {
	res, err := convert(content)
	if err != nil {
		return nil, err
	}
	return mxj.NewMapXml(res)
}

// 将map变量解析为XML格式内容
func Encode(v map[string]interface{}, rootTag ...string) ([]byte, error) {
	return mxj.Map(v).Xml(rootTag...)
}

func EncodeWithIndent(v map[string]interface{}, rootTag ...string) ([]byte, error) {
	return mxj.Map(v).XmlIndent("", "\t", rootTag...)
}

// XML格式内容直接转换为JSON格式内容
func ToJson(content []byte) ([]byte, error) {
	res, err := convert(content)
	if err != nil {
		fmt.Println("convert error. ", err)
		return nil, err
	}

	mv, err := mxj.NewMapXml(res)
	if err == nil {
		return mv.Json()
	} else {
		return nil, err
	}
}

// XML字符集预处理
// @author wenzi1
// @date 20180604  修复并发安全问题,改为如果非UTF8字符集则先做字符集转换
func convert(xmlbyte []byte) (res []byte, err error) {
	patten := `<\?xml.*encoding\s*=\s*['|"](.*?)['|"].*\?>`
	matchStr, err := gregex.MatchString(patten, string(xmlbyte))
	if err != nil {
		return nil, err
	}

	xmlEncode := "UTF-8"
	if len(matchStr) == 2 {
		xmlEncode = matchStr[1]
	}

	s := mahonia.GetCharset(xmlEncode)
	if s == nil {
		return nil, fmt.Errorf("not support charset:%s\n", xmlEncode)
	}
	fmt.Println(s.Name, xmlEncode)
	res, err = gregex.Replace(patten, []byte(""), []byte(xmlbyte))
	if err != nil {
		return nil, err
	}

	if !strings.EqualFold(s.Name, "UTF-8") {
		res = []byte(s.NewDecoder().ConvertString(string(res)))
	}

	return res, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package main

import (
	"context"
	"fmt"
	"strings"
	"testing"
	"time"

	computebeta "google.golang.org/api/compute/v0.beta"

	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud/meta"

	"k8s.io/ingress-gce/pkg/annotations"
	"k8s.io/ingress-gce/pkg/e2e"
	"k8s.io/ingress-gce/pkg/fuzz"
	"k8s.io/ingress-gce/pkg/fuzz/features"
	"k8s.io/ingress-gce/pkg/utils"
)

const (
	policyUpdateInterval = 15 * time.Second
	policyUpdateTimeout  = 3 * time.Minute
)

func buildPolicyAllowAll(name string) *computebeta.SecurityPolicy {
	return &computebeta.SecurityPolicy{
		Name: name,
	}
}

func buildPolicyDisallowAll(name string) *computebeta.SecurityPolicy {
	return &computebeta.SecurityPolicy{
		Name: name,
		Rules: []*computebeta.SecurityPolicyRule{
			&computebeta.SecurityPolicyRule{
				Action: "deny(403)",
				Match: &computebeta.SecurityPolicyRuleMatcher{
					Config: &computebeta.SecurityPolicyRuleMatcherConfig{
						SrcIpRanges: []string{"*"},
					},
					VersionedExpr: "SRC_IPS_V1",
				},
				Priority: 2147483647,
			},
		},
	}
}

func TestSecurityPolicyEnable(t *testing.T) {
	ctx := context.Background()
	t.Parallel()

	Framework.RunWithSandbox("Security Policy Enable", t, func(t *testing.T, s *e2e.Sandbox) {
		policies := []*computebeta.SecurityPolicy{
			buildPolicyAllowAll(fmt.Sprintf("enable-test-allow-all-%s", s.Namespace)),
		}
		defer func() {
			if err := cleanupSecurityPolicies(t, ctx, Framework.Cloud, policies); err != nil {
				t.Errorf("cleanupSecurityPolicies(...) =  %v, want nil", err)
			}
		}()
		policies, err := createSecurityPolicies(t, ctx, Framework.Cloud, policies)
		if err != nil {
			t.Fatalf("createSecurityPolicies(...) = _, %v, want _, nil", err)
		}
		// Re-assign to get the populated self-link.
		testSecurityPolicy := policies[0]

		testBackendConfigAnnotation := map[string]string{
			annotations.BackendConfigKey: `{"default":"backendconfig-1"}`,
		}
		_, testSvc, err := e2e.CreateEchoService(s, "service-1", testBackendConfigAnnotation)
		if err != nil {
			t.Fatalf("e2e.CreateEchoService(s, service-1, %q) = _, _, %v, want _, _, nil", testBackendConfigAnnotation, err)
		}

		testBackendConfig := fuzz.NewBackendConfigBuilder("", "backendconfig-1").SetSecurityPolicy(testSecurityPolicy.Name).Build()
		testBackendConfig, err = Framework.BackendConfigClient.CloudV1beta1().BackendConfigs(s.Namespace).Create(testBackendConfig)
		if err != nil {
			t.Fatalf("Error creating test backend config: %v", err)
		}
		t.Logf("Backend config %s/%s created", s.Namespace, testBackendConfig.Name)

		port80 := intstr.FromInt(80)
		testIng := fuzz.NewIngressBuilder("", "ingress-1", "").DefaultBackend("service-1", port80).AddPath("test.com", "/", "service-1", port80).Build()
		testIng, err = Framework.Clientset.Extensions().Ingresses(s.Namespace).Create(testIng)
		if err != nil {
			t.Fatalf("error creating Ingress spec: %v", err)
		}
		t.Logf("Ingress %s/%s created", s.Namespace, testIng.Name)

		t.Logf("Checking on relevant backend service whether security policy is properly attached")

		testIng, err = e2e.WaitForIngress(s, testIng)
		if err != nil {
			t.Fatalf("e2e.WaitForIngress(s, %q) = _, %v; want _, nil", testIng.Name, err)
		}
		if len(testIng.Status.LoadBalancer.Ingress) < 1 {
			t.Fatalf("Ingress does not have an IP: %+v", testIng.Status)
		}

		vip := testIng.Status.LoadBalancer.Ingress[0].IP
		gclb, err := fuzz.GCLBForVIP(ctx, Framework.Cloud, vip, fuzz.FeatureValidators([]fuzz.Feature{features.SecurityPolicy}))
		if err != nil {
			t.Fatalf("fuzz.GCLBForVIP(..., %q, %q) = _, %v; want _, nil", vip, features.SecurityPolicy, err)
		}

		if err := verifySecurityPolicy(t, gclb, s.Namespace, testSvc.Name, testSecurityPolicy.SelfLink); err != nil {
			t.Errorf("verifySecurityPolicy(..., %q, %q, %q) = %v, want nil", s.Namespace, testSvc.Name, testSecurityPolicy.SelfLink, err)
		}

		t.Logf("Cleaning up test")

		if err := e2e.WaitForIngressDeletion(ctx, gclb, s, testIng, nil); err != nil {
			t.Errorf("e2e.WaitForIngressDeletion(..., %q, nil) = %v, want nil", testIng.Name, err)
		}
	})
}

func TestSecurityPolicyTransition(t *testing.T) {
	ctx := context.Background()
	t.Parallel()

	Framework.RunWithSandbox("Security Policy Transition", t, func(t *testing.T, s *e2e.Sandbox) {
		policies := []*computebeta.SecurityPolicy{
			buildPolicyAllowAll(fmt.Sprintf("transition-test-allow-all-%s", s.Namespace))
			buildPolicyDisallowAll(fmt.Sprintf("transition-test-disallow-all-%s", s.Namespace))
		}
		defer func() {
			if err := cleanupSecurityPolicies(t, ctx, Framework.Cloud, policies); err != nil {
				t.Errorf("cleanupSecurityPolicies(...) = %v, want nil", err)
			}
		}()
		policies, err := createSecurityPolicies(t, ctx, Framework.Cloud, policies)
		if err != nil {
			t.Fatalf("createSecurityPolicies(...) = _, %v, want _, nil", err)
		}
		// Re-assign to get the populated self-link.
		testSecurityPolicyAllow, testSecurityPolicyDisallow := policies[0], policies[1]

		testBackendConfigAnnotation := map[string]string{
			annotations.BackendConfigKey: `{"default":"backendconfig-1"}`,
		}
		_, testSvc, err := e2e.CreateEchoService(s, "service-1", testBackendConfigAnnotation)
		if err != nil {
			t.Fatalf("e2e.CreateEchoService(s, service-1, %q) = _, _, %v, want _, _, nil", testBackendConfigAnnotation, err)
		}

		testBackendConfig := fuzz.NewBackendConfigBuilder("", "backendconfig-1").SetSecurityPolicy(testSecurityPolicyAllow.Name).Build()
		testBackendConfig, err = Framework.BackendConfigClient.CloudV1beta1().BackendConfigs(s.Namespace).Create(testBackendConfig)
		if err != nil {
			t.Fatalf("Error creating test backend config: %v", err)
		}
		t.Logf("Backend config %s/%s created", s.Namespace, testBackendConfig.Name)

		port80 := intstr.FromInt(80)
		testIng := fuzz.NewIngressBuilder("", "ingress-1", "").DefaultBackend("service-1", port80).AddPath("test.com", "/", "service-1", port80).Build()
		testIng, err = Framework.Clientset.Extensions().Ingresses(s.Namespace).Create(testIng)
		if err != nil {
			t.Fatalf("error creating Ingress spec: %v", err)
		}
		t.Logf("Ingress %s/%s created", s.Namespace, testIng.Name)

		ing, err := e2e.WaitForIngress(s, testIng)
		if err != nil {
			t.Fatalf("e2e.WaitForIngress(s, %q) = _, %v; want _, nil", testIng.Name, err)
		}
		if len(ing.Status.LoadBalancer.Ingress) < 1 {
			t.Fatalf("Ingress does not have an IP: %+v", ing.Status)
		}

		vip := ing.Status.LoadBalancer.Ingress[0].IP
		var gclb *fuzz.GCLB

		steps := []struct {
			desc                string
			securityPolicyToSet string
			expectedpolicyLink  string
		}{
			{
				desc:                "update to use policy that disallows all",
				securityPolicyToSet: testSecurityPolicyDisallow.Name,
				expectedpolicyLink:  testSecurityPolicyDisallow.SelfLink,
			},
			{
				desc:                "detach policy",
				securityPolicyToSet: "",
				expectedpolicyLink:  "",
			},
		}

		for _, step := range steps {
			testBackendConfig.Spec.SecurityPolicy.Name = step.securityPolicyToSet
			testBackendConfig, err = Framework.BackendConfigClient.CloudV1beta1().BackendConfigs(s.Namespace).Update(testBackendConfig)
			if err != nil {
				t.Fatalf("Error updating test backend config: %v", err)
			}
			t.Logf("Backend config %s/%s updated", testBackendConfig.Name, s.Namespace)

			t.Logf("Checking on relevant backend service whether security policy is properly updated")

			if err := wait.Poll(policyUpdateInterval, policyUpdateTimeout, func() (bool, error) {
				gclb, err = fuzz.GCLBForVIP(ctx, Framework.Cloud, vip, fuzz.FeatureValidators([]fuzz.Feature{features.SecurityPolicy}))
				if err != nil {
					t.Fatalf("fuzz.GCLBForVIP(..., %q, %q) = _, %v; want _, nil", vip, features.SecurityPolicy, err)
				}

				if err := verifySecurityPolicy(t, gclb, s.Namespace, testSvc.Name, step.expectedpolicyLink); err != nil {
					t.Logf("verifySecurityPolicy(..., %q, %q, %q) = %v, want nil", s.Namespace, testSvc.Name, step.expectedpolicyLink, err)
					return false, nil
				}
				return true, nil
			}); err != nil {
				t.Errorf("Failed to wait for security policy updated: %v", err)
			}
		}

		t.Logf("Cleaning up test")

		if err := e2e.WaitForIngressDeletion(ctx, gclb, s, ing, nil); err != nil {
			t.Errorf("e2e.WaitForIngressDeletion(..., %q, nil) = %v, want nil", ing.Name, err)
		}
	})
}

func createSecurityPolicies(t *testing.T, ctx context.Context, c cloud.Cloud, policies []*computebeta.SecurityPolicy) ([]*computebeta.SecurityPolicy, error) {
	t.Logf("Creating security policies...")
	createdPolicies := []*computebeta.SecurityPolicy{}
	for _, policy := range policies {
		if err := c.BetaSecurityPolicies().Insert(ctx, meta.GlobalKey(policy.Name), policy); err != nil {
			return nil, fmt.Errorf("error creating security policy %q: %v", policy.Name, err)
		}
		t.Logf("Security policy %q created", policy.Name)
		policy, err := c.BetaSecurityPolicies().Get(ctx, meta.GlobalKey(policy.Name))
		if err != nil {
			return nil, fmt.Errorf("error getting security policy %q: %v", policy.Name, err)
		}
		createdPolicies = append(createdPolicies, policy)
	}
	return createdPolicies, nil
}

func cleanupSecurityPolicies(t *testing.T, ctx context.Context, c cloud.Cloud, policies []*computebeta.SecurityPolicy) error {
	t.Logf("Deleting security policies...")
	var errs []string
	for _, policy := range policies {
		if err := c.BetaSecurityPolicies().Delete(ctx, meta.GlobalKey(policy.Name)); err != nil {
			errs = append(errs, err.Error())
		}
		t.Logf("Security policy %q deleted", policy.Name)
	}
	if len(errs) != 0 {
		return fmt.Errorf("failed to delete security policies: %s", strings.Join(errs, "\n"))
	}
	return nil
}

func verifySecurityPolicy(t *testing.T, gclb *fuzz.GCLB, svcNamespace, svcName, policyLink string) error {
	numBsWithPolicy := 0
	for _, bs := range gclb.BackendService {
		// Check on relevant backend services.
		desc := utils.DescriptionFromString(bs.GA.Description)
		if desc.ServiceName != fmt.Sprintf("%s/%s", svcNamespace, svcName) {
			continue
		}

		if bs.Beta == nil {
			return fmt.Errorf("Beta BackendService resource not found: %v", bs)
		}
		if bs.Beta.SecurityPolicy != policyLink {
			return fmt.Errorf("backend service %q has security policy %q, want %q", bs.Beta.Name, bs.Beta.SecurityPolicy, policyLink)
		}
		t.Logf("Backend service %q has the expected security policy %q attached", bs.Beta.Name, bs.Beta.SecurityPolicy)
		numBsWithPolicy = numBsWithPolicy + 1
	}
	if numBsWithPolicy != 1 {
		return fmt.Errorf("unexpected number of backend service has security policy attached: got %d, want 1", numBsWithPolicy)
	}
	return nil
}
<commit_msg>Update security_policy_test.go<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package main

import (
	"context"
	"fmt"
	"strings"
	"testing"
	"time"

	computebeta "google.golang.org/api/compute/v0.beta"

	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud/meta"

	"k8s.io/ingress-gce/pkg/annotations"
	"k8s.io/ingress-gce/pkg/e2e"
	"k8s.io/ingress-gce/pkg/fuzz"
	"k8s.io/ingress-gce/pkg/fuzz/features"
	"k8s.io/ingress-gce/pkg/utils"
)

const (
	policyUpdateInterval = 15 * time.Second
	policyUpdateTimeout  = 3 * time.Minute
)

func buildPolicyAllowAll(name string) *computebeta.SecurityPolicy {
	return &computebeta.SecurityPolicy{
		Name: name,
	}
}

func buildPolicyDisallowAll(name string) *computebeta.SecurityPolicy {
	return &computebeta.SecurityPolicy{
		Name: name,
		Rules: []*computebeta.SecurityPolicyRule{
			&computebeta.SecurityPolicyRule{
				Action: "deny(403)",
				Match: &computebeta.SecurityPolicyRuleMatcher{
					Config: &computebeta.SecurityPolicyRuleMatcherConfig{
						SrcIpRanges: []string{"*"},
					},
					VersionedExpr: "SRC_IPS_V1",
				},
				Priority: 2147483647,
			},
		},
	}
}

func TestSecurityPolicyEnable(t *testing.T) {
	ctx := context.Background()
	t.Parallel()

	Framework.RunWithSandbox("Security Policy Enable", t, func(t *testing.T, s *e2e.Sandbox) {
		policies := []*computebeta.SecurityPolicy{
			buildPolicyAllowAll(fmt.Sprintf("enable-test-allow-all-%s", s.Namespace)),
		}
		defer func() {
			if err := cleanupSecurityPolicies(t, ctx, Framework.Cloud, policies); err != nil {
				t.Errorf("cleanupSecurityPolicies(...) =  %v, want nil", err)
			}
		}()
		policies, err := createSecurityPolicies(t, ctx, Framework.Cloud, policies)
		if err != nil {
			t.Fatalf("createSecurityPolicies(...) = _, %v, want _, nil", err)
		}
		// Re-assign to get the populated self-link.
		testSecurityPolicy := policies[0]

		testBackendConfigAnnotation := map[string]string{
			annotations.BackendConfigKey: `{"default":"backendconfig-1"}`,
		}
		_, testSvc, err := e2e.CreateEchoService(s, "service-1", testBackendConfigAnnotation)
		if err != nil {
			t.Fatalf("e2e.CreateEchoService(s, service-1, %q) = _, _, %v, want _, _, nil", testBackendConfigAnnotation, err)
		}

		testBackendConfig := fuzz.NewBackendConfigBuilder("", "backendconfig-1").SetSecurityPolicy(testSecurityPolicy.Name).Build()
		testBackendConfig, err = Framework.BackendConfigClient.CloudV1beta1().BackendConfigs(s.Namespace).Create(testBackendConfig)
		if err != nil {
			t.Fatalf("Error creating test backend config: %v", err)
		}
		t.Logf("Backend config %s/%s created", s.Namespace, testBackendConfig.Name)

		port80 := intstr.FromInt(80)
		testIng := fuzz.NewIngressBuilder("", "ingress-1", "").DefaultBackend("service-1", port80).AddPath("test.com", "/", "service-1", port80).Build()
		testIng, err = Framework.Clientset.Extensions().Ingresses(s.Namespace).Create(testIng)
		if err != nil {
			t.Fatalf("error creating Ingress spec: %v", err)
		}
		t.Logf("Ingress %s/%s created", s.Namespace, testIng.Name)

		t.Logf("Checking on relevant backend service whether security policy is properly attached")

		testIng, err = e2e.WaitForIngress(s, testIng)
		if err != nil {
			t.Fatalf("e2e.WaitForIngress(s, %q) = _, %v; want _, nil", testIng.Name, err)
		}
		if len(testIng.Status.LoadBalancer.Ingress) < 1 {
			t.Fatalf("Ingress does not have an IP: %+v", testIng.Status)
		}

		vip := testIng.Status.LoadBalancer.Ingress[0].IP
		gclb, err := fuzz.GCLBForVIP(ctx, Framework.Cloud, vip, fuzz.FeatureValidators([]fuzz.Feature{features.SecurityPolicy}))
		if err != nil {
			t.Fatalf("fuzz.GCLBForVIP(..., %q, %q) = _, %v; want _, nil", vip, features.SecurityPolicy, err)
		}

		if err := verifySecurityPolicy(t, gclb, s.Namespace, testSvc.Name, testSecurityPolicy.SelfLink); err != nil {
			t.Errorf("verifySecurityPolicy(..., %q, %q, %q) = %v, want nil", s.Namespace, testSvc.Name, testSecurityPolicy.SelfLink, err)
		}

		t.Logf("Cleaning up test")

		if err := e2e.WaitForIngressDeletion(ctx, gclb, s, testIng, nil); err != nil {
			t.Errorf("e2e.WaitForIngressDeletion(..., %q, nil) = %v, want nil", testIng.Name, err)
		}
	})
}

func TestSecurityPolicyTransition(t *testing.T) {
	ctx := context.Background()
	t.Parallel()

	Framework.RunWithSandbox("Security Policy Transition", t, func(t *testing.T, s *e2e.Sandbox) {
		policies := []*computebeta.SecurityPolicy{
			buildPolicyAllowAll(fmt.Sprintf("transition-test-allow-all-%s", s.Namespace)),
			buildPolicyDisallowAll(fmt.Sprintf("transition-test-disallow-all-%s", s.Namespace)),
		}
		defer func() {
			if err := cleanupSecurityPolicies(t, ctx, Framework.Cloud, policies); err != nil {
				t.Errorf("cleanupSecurityPolicies(...) = %v, want nil", err)
			}
		}()
		policies, err := createSecurityPolicies(t, ctx, Framework.Cloud, policies)
		if err != nil {
			t.Fatalf("createSecurityPolicies(...) = _, %v, want _, nil", err)
		}
		// Re-assign to get the populated self-link.
		testSecurityPolicyAllow, testSecurityPolicyDisallow := policies[0], policies[1]

		testBackendConfigAnnotation := map[string]string{
			annotations.BackendConfigKey: `{"default":"backendconfig-1"}`,
		}
		_, testSvc, err := e2e.CreateEchoService(s, "service-1", testBackendConfigAnnotation)
		if err != nil {
			t.Fatalf("e2e.CreateEchoService(s, service-1, %q) = _, _, %v, want _, _, nil", testBackendConfigAnnotation, err)
		}

		testBackendConfig := fuzz.NewBackendConfigBuilder("", "backendconfig-1").SetSecurityPolicy(testSecurityPolicyAllow.Name).Build()
		testBackendConfig, err = Framework.BackendConfigClient.CloudV1beta1().BackendConfigs(s.Namespace).Create(testBackendConfig)
		if err != nil {
			t.Fatalf("Error creating test backend config: %v", err)
		}
		t.Logf("Backend config %s/%s created", s.Namespace, testBackendConfig.Name)

		port80 := intstr.FromInt(80)
		testIng := fuzz.NewIngressBuilder("", "ingress-1", "").DefaultBackend("service-1", port80).AddPath("test.com", "/", "service-1", port80).Build()
		testIng, err = Framework.Clientset.Extensions().Ingresses(s.Namespace).Create(testIng)
		if err != nil {
			t.Fatalf("error creating Ingress spec: %v", err)
		}
		t.Logf("Ingress %s/%s created", s.Namespace, testIng.Name)

		ing, err := e2e.WaitForIngress(s, testIng)
		if err != nil {
			t.Fatalf("e2e.WaitForIngress(s, %q) = _, %v; want _, nil", testIng.Name, err)
		}
		if len(ing.Status.LoadBalancer.Ingress) < 1 {
			t.Fatalf("Ingress does not have an IP: %+v", ing.Status)
		}

		vip := ing.Status.LoadBalancer.Ingress[0].IP
		var gclb *fuzz.GCLB

		steps := []struct {
			desc                string
			securityPolicyToSet string
			expectedpolicyLink  string
		}{
			{
				desc:                "update to use policy that disallows all",
				securityPolicyToSet: testSecurityPolicyDisallow.Name,
				expectedpolicyLink:  testSecurityPolicyDisallow.SelfLink,
			},
			{
				desc:                "detach policy",
				securityPolicyToSet: "",
				expectedpolicyLink:  "",
			},
		}

		for _, step := range steps {
			testBackendConfig.Spec.SecurityPolicy.Name = step.securityPolicyToSet
			testBackendConfig, err = Framework.BackendConfigClient.CloudV1beta1().BackendConfigs(s.Namespace).Update(testBackendConfig)
			if err != nil {
				t.Fatalf("Error updating test backend config: %v", err)
			}
			t.Logf("Backend config %s/%s updated", testBackendConfig.Name, s.Namespace)

			t.Logf("Checking on relevant backend service whether security policy is properly updated")

			if err := wait.Poll(policyUpdateInterval, policyUpdateTimeout, func() (bool, error) {
				gclb, err = fuzz.GCLBForVIP(ctx, Framework.Cloud, vip, fuzz.FeatureValidators([]fuzz.Feature{features.SecurityPolicy}))
				if err != nil {
					t.Fatalf("fuzz.GCLBForVIP(..., %q, %q) = _, %v; want _, nil", vip, features.SecurityPolicy, err)
				}

				if err := verifySecurityPolicy(t, gclb, s.Namespace, testSvc.Name, step.expectedpolicyLink); err != nil {
					t.Logf("verifySecurityPolicy(..., %q, %q, %q) = %v, want nil", s.Namespace, testSvc.Name, step.expectedpolicyLink, err)
					return false, nil
				}
				return true, nil
			}); err != nil {
				t.Errorf("Failed to wait for security policy updated: %v", err)
			}
		}

		t.Logf("Cleaning up test")

		if err := e2e.WaitForIngressDeletion(ctx, gclb, s, ing, nil); err != nil {
			t.Errorf("e2e.WaitForIngressDeletion(..., %q, nil) = %v, want nil", ing.Name, err)
		}
	})
}

func createSecurityPolicies(t *testing.T, ctx context.Context, c cloud.Cloud, policies []*computebeta.SecurityPolicy) ([]*computebeta.SecurityPolicy, error) {
	t.Logf("Creating security policies...")
	createdPolicies := []*computebeta.SecurityPolicy{}
	for _, policy := range policies {
		if err := c.BetaSecurityPolicies().Insert(ctx, meta.GlobalKey(policy.Name), policy); err != nil {
			return nil, fmt.Errorf("error creating security policy %q: %v", policy.Name, err)
		}
		t.Logf("Security policy %q created", policy.Name)
		policy, err := c.BetaSecurityPolicies().Get(ctx, meta.GlobalKey(policy.Name))
		if err != nil {
			return nil, fmt.Errorf("error getting security policy %q: %v", policy.Name, err)
		}
		createdPolicies = append(createdPolicies, policy)
	}
	return createdPolicies, nil
}

func cleanupSecurityPolicies(t *testing.T, ctx context.Context, c cloud.Cloud, policies []*computebeta.SecurityPolicy) error {
	t.Logf("Deleting security policies...")
	var errs []string
	for _, policy := range policies {
		if err := c.BetaSecurityPolicies().Delete(ctx, meta.GlobalKey(policy.Name)); err != nil {
			errs = append(errs, err.Error())
		}
		t.Logf("Security policy %q deleted", policy.Name)
	}
	if len(errs) != 0 {
		return fmt.Errorf("failed to delete security policies: %s", strings.Join(errs, "\n"))
	}
	return nil
}

func verifySecurityPolicy(t *testing.T, gclb *fuzz.GCLB, svcNamespace, svcName, policyLink string) error {
	numBsWithPolicy := 0
	for _, bs := range gclb.BackendService {
		// Check on relevant backend services.
		desc := utils.DescriptionFromString(bs.GA.Description)
		if desc.ServiceName != fmt.Sprintf("%s/%s", svcNamespace, svcName) {
			continue
		}

		if bs.Beta == nil {
			return fmt.Errorf("Beta BackendService resource not found: %v", bs)
		}
		if bs.Beta.SecurityPolicy != policyLink {
			return fmt.Errorf("backend service %q has security policy %q, want %q", bs.Beta.Name, bs.Beta.SecurityPolicy, policyLink)
		}
		t.Logf("Backend service %q has the expected security policy %q attached", bs.Beta.Name, bs.Beta.SecurityPolicy)
		numBsWithPolicy = numBsWithPolicy + 1
	}
	if numBsWithPolicy != 1 {
		return fmt.Errorf("unexpected number of backend service has security policy attached: got %d, want 1", numBsWithPolicy)
	}
	return nil
}
<|endoftext|>