<commit_before>// Copyright (c) 2015, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

// Package xurls extracts urls from plain text using regular expressions.
package xurls

import "regexp"

//go:generate go run generate/tldsgen/main.go
//go:generate go run generate/regexgen/main.go

const (
	letter    = `\p{L}`
	number    = `\p{N}`
	iriChar   = letter + number
	pathChar  = iriChar + `/\-+_&~*%=@|#.,:;'?!` + `\p{Sc}\p{Sk}\p{So}`
	endChar   = iriChar + `/\-+_&~*%=` + `\p{Sc}`
	wellParen = `\([` + pathChar + `]*(\([` + pathChar + `]*\)[` + pathChar + `]*)*\)`
	wellBrack = `\[[` + pathChar + `]*(\[[` + pathChar + `]*\][` + pathChar + `]*)*\]`
	wellAll   = wellParen + `|` + wellBrack
	pathCont  = `([` + pathChar + `]*(` + wellAll + `|[` + endChar + `])+)+`
	comScheme = `[a-zA-Z][a-zA-Z.\-+]*://`
	scheme    = `(` + comScheme + `|` + otherScheme + `)`

	iri      = `[` + iriChar + `]([` + iriChar + `\-]*[` + iriChar + `])?`
	domain   = `(` + iri + `\.)+`
	octet    = `(25[0-5]|2[0-4][0-9]|1[0-9]{2}|[1-9][0-9]|[0-9])`
	ipv4Addr = `\b` + octet + `\.` + octet + `\.` + octet + `\.` + octet + `\b`
	ipv6Addr = `([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:[0-9a-fA-F]{0,4}|:[0-9a-fA-F]{1,4})?|(:[0-9a-fA-F]{1,4}){0,2})|(:[0-9a-fA-F]{1,4}){0,3})|(:[0-9a-fA-F]{1,4}){0,4})|:(:[0-9a-fA-F]{1,4}){0,5})((:[0-9a-fA-F]{1,4}){2}|:(25[0-5]|(2[0-4]|1[0-9]|[1-9])?[0-9])(\.(25[0-5]|(2[0-4]|1[0-9]|[1-9])?[0-9])){3})|(([0-9a-fA-F]{1,4}:){1,6}|:):[0-9a-fA-F]{1,4}|([0-9a-fA-F]{1,4}:){7}:`
	ipAddr   = `(` + ipv4Addr + `|` + ipv6Addr + `)`
	hostName = `(` + domain + gtld + `|` + ipAddr + `)`
	port     = `(:[0-9]*)?`
	path     = `(/|/` + pathCont + `?|\b|$)`
	webURL   = hostName + port + path
	email    = `[a-zA-Z0-9._%\-+]+@` + hostName

	strict  = `\b` + scheme + pathCont
	relaxed = strict + `|` + webURL + `|` + email
)

var (
	// Relaxed matches all the urls it can find.
	Relaxed = regexp.MustCompile(relaxed)
	// Strict only matches urls with a scheme to avoid false positives.
	Strict = regexp.MustCompile(strict)
)

func init() {
	Relaxed.Longest()
	Strict.Longest()
}

// StrictMatchingScheme produces a regexp that matches urls like Strict but
// whose scheme matches the given regular expression.
func StrictMatchingScheme(exp string) (*regexp.Regexp, error) {
	strictMatching := `\b(?i)(` + exp + `)(?-i)` + pathCont
	re, err := regexp.Compile(strictMatching)
	if err != nil {
		return nil, err
	}
	re.Longest()
	return re, nil
}
<commit_msg>Wrap exposed expressions in parenthesis<commit_after>// Copyright (c) 2015, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

// Package xurls extracts urls from plain text using regular expressions.
package xurls

import "regexp"

//go:generate go run generate/tldsgen/main.go
//go:generate go run generate/regexgen/main.go

const (
	letter    = `\p{L}`
	number    = `\p{N}`
	iriChar   = letter + number
	pathChar  = iriChar + `/\-+_&~*%=@|#.,:;'?!` + `\p{Sc}\p{Sk}\p{So}`
	endChar   = iriChar + `/\-+_&~*%=` + `\p{Sc}`
	wellParen = `\([` + pathChar + `]*(\([` + pathChar + `]*\)[` + pathChar + `]*)*\)`
	wellBrack = `\[[` + pathChar + `]*(\[[` + pathChar + `]*\][` + pathChar + `]*)*\]`
	wellAll   = wellParen + `|` + wellBrack
	pathCont  = `([` + pathChar + `]*(` + wellAll + `|[` + endChar + `])+)+`
	comScheme = `[a-zA-Z][a-zA-Z.\-+]*://`
	scheme    = `(` + comScheme + `|` + otherScheme + `)`

	iri      = `[` + iriChar + `]([` + iriChar + `\-]*[` + iriChar + `])?`
	domain   = `(` + iri + `\.)+`
	octet    = `(25[0-5]|2[0-4][0-9]|1[0-9]{2}|[1-9][0-9]|[0-9])`
	ipv4Addr = `\b` + octet + `\.` + octet + `\.` + octet + `\.` + octet + `\b`
	ipv6Addr = `([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:([0-9a-fA-F]{1,4}:[0-9a-fA-F]{0,4}|:[0-9a-fA-F]{1,4})?|(:[0-9a-fA-F]{1,4}){0,2})|(:[0-9a-fA-F]{1,4}){0,3})|(:[0-9a-fA-F]{1,4}){0,4})|:(:[0-9a-fA-F]{1,4}){0,5})((:[0-9a-fA-F]{1,4}){2}|:(25[0-5]|(2[0-4]|1[0-9]|[1-9])?[0-9])(\.(25[0-5]|(2[0-4]|1[0-9]|[1-9])?[0-9])){3})|(([0-9a-fA-F]{1,4}:){1,6}|:):[0-9a-fA-F]{1,4}|([0-9a-fA-F]{1,4}:){7}:`
	ipAddr   = `(` + ipv4Addr + `|` + ipv6Addr + `)`
	hostName = `(` + domain + gtld + `|` + ipAddr + `)`
	port     = `(:[0-9]*)?`
	path     = `(/|/` + pathCont + `?|\b|$)`
	webURL   = hostName + port + path
	email    = `[a-zA-Z0-9._%\-+]+@` + hostName

	strict  = `(\b` + scheme + pathCont + `)`
	relaxed = `(` + strict + `|` + webURL + `|` + email + `)`
)

var (
	// Relaxed matches all the urls it can find.
	Relaxed = regexp.MustCompile(relaxed)
	// Strict only matches urls with a scheme to avoid false positives.
	Strict = regexp.MustCompile(strict)
)

func init() {
	Relaxed.Longest()
	Strict.Longest()
}

// StrictMatchingScheme produces a regexp that matches urls like Strict but
// whose scheme matches the given regular expression.
func StrictMatchingScheme(exp string) (*regexp.Regexp, error) {
	strictMatching := `(\b(?i)(` + exp + `)(?-i)` + pathCont + `)`
	re, err := regexp.Compile(strictMatching)
	if err != nil {
		return nil, err
	}
	re.Longest()
	return re, nil
}
<|endoftext|>


<commit_before>package lib

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"time"

	"github.com/metral/goutils"
)

var K8S_API_VERSION string = "v1beta1"
var K8S_API_PORT string = "8080"

type PreregisteredMinion struct {
	Kind       string `json:"kind,omitempty"`
	ID         string `json:"id,omitempty"`
	HostIP     string `json:"hostIP,omitempty"`
	APIVersion string `json:"apiVersion,omitempty"`
}

type MinionsResult struct {
	Kind              string  `json:"kind,omitempty"`
	CreationTimestamp string  `json:"creationTimestamp,omitempty"`
	SelfLink          string  `json:"selfLink,omitempty"`
	APIVersion        string  `json:"apiVersion,omitempty"`
	Minions           Minions `json:"minions,omitempty"`
}

type Minions []Minion
type Minion struct {
	ID                string `json:"id,omitempty"`
	UID               string `json:"uid,omitempty"`
	CreationTimestamp string `json:"creationTimestamp,omitempty"`
	SelfLink          string `json:"selfLink,omitempty"`
	ResourceVersion   int    `json:"resourceVersion,omitempty"`
	HostIP            string `json:"hostIP,omitempty"`
	resources         map[interface{}]interface{}
}

func isK8sMachine(
	fleetMachine,
	master *FleetMachine,
	minions *FleetMachines) bool {

	role := fleetMachine.Metadata["kubernetes_role"]

	switch role {
	case "master":
		*master = *fleetMachine
		return true
	case "minion":
		*minions = append(*minions, *fleetMachine)
		return true
	}
	return false
}

func registerMinions(master *FleetMachine, minions *FleetMachines) {

	// Get registered minions, if any
	endpoint := fmt.Sprintf("http://%s:%s", master.PublicIP, K8S_API_PORT)
	masterAPIurl := fmt.Sprintf("%s/api/%s/minions", endpoint, K8S_API_VERSION)
	jsonResponse := goutils.HttpGetRequest(masterAPIurl)
	m := *minions

	var minionsResult MinionsResult
	err := json.Unmarshal(jsonResponse, &minionsResult)
	goutils.CheckForErrors(goutils.ErrorParams{Err: err, CallerNum: 1})

	// See if minions discovered have been registered. If not, register
	for _, minion := range m {
		registered := false
		for _, registeredMinion := range minionsResult.Minions {
			if registeredMinion.HostIP == minion.PublicIP {
				registered = true
			}
		}

		if !registered {
			register(endpoint, minion.PublicIP)
			time.Sleep(500 * time.Millisecond)
		}
	}
}

func register(endpoint, addr string) error {
	m := &PreregisteredMinion{
		Kind:       "Minion",
		APIVersion: K8S_API_VERSION,
		ID:         addr,
		HostIP:     addr,
	}
	data, err := json.Marshal(m)
	goutils.CheckForErrors(goutils.ErrorParams{Err: err, CallerNum: 1})

	url := fmt.Sprintf("%s/api/%s/minions", endpoint, K8S_API_VERSION)

	headers := map[string]string{
		"Content-Type": "application/json",
	}

	h := goutils.HttpRequestParams{
		HttpRequestType: "POST",
		Url:             url,
		Data:            data,
		Headers:         headers,
	}
	statusCode := goutils.HttpCreateRequest(h)

	switch statusCode {
	case 200, 202:
		log.Printf("------------------------------------------------")
		log.Printf("Registered machine with the master: %s\n", addr)
		return nil
	case 409:
		//log.Printf("Machine has already registered with master: %s\n", addr)
		return nil
	}
	data, err = ioutil.ReadAll(resp.Body)
	log.Printf("Response: %#v", resp)
	log.Printf("Response Body:\n%s", string(data))
	return errors.New("error registering: " + addr)
}
<commit_msg>Hotfix: remove usage of resp<commit_after>package lib

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/metral/goutils"
)

var K8S_API_VERSION string = "v1beta1"
var K8S_API_PORT string = "8080"

type PreregisteredMinion struct {
	Kind       string `json:"kind,omitempty"`
	ID         string `json:"id,omitempty"`
	HostIP     string `json:"hostIP,omitempty"`
	APIVersion string `json:"apiVersion,omitempty"`
}

type MinionsResult struct {
	Kind              string  `json:"kind,omitempty"`
	CreationTimestamp string  `json:"creationTimestamp,omitempty"`
	SelfLink          string  `json:"selfLink,omitempty"`
	APIVersion        string  `json:"apiVersion,omitempty"`
	Minions           Minions `json:"minions,omitempty"`
}

type Minions []Minion
type Minion struct {
	ID                string `json:"id,omitempty"`
	UID               string `json:"uid,omitempty"`
	CreationTimestamp string `json:"creationTimestamp,omitempty"`
	SelfLink          string `json:"selfLink,omitempty"`
	ResourceVersion   int    `json:"resourceVersion,omitempty"`
	HostIP            string `json:"hostIP,omitempty"`
	resources         map[interface{}]interface{}
}

func isK8sMachine(
	fleetMachine,
	master *FleetMachine,
	minions *FleetMachines) bool {

	role := fleetMachine.Metadata["kubernetes_role"]

	switch role {
	case "master":
		*master = *fleetMachine
		return true
	case "minion":
		*minions = append(*minions, *fleetMachine)
		return true
	}
	return false
}

func registerMinions(master *FleetMachine, minions *FleetMachines) {

	// Get registered minions, if any
	endpoint := fmt.Sprintf("http://%s:%s", master.PublicIP, K8S_API_PORT)
	masterAPIurl := fmt.Sprintf("%s/api/%s/minions", endpoint, K8S_API_VERSION)
	jsonResponse := goutils.HttpGetRequest(masterAPIurl)
	m := *minions

	var minionsResult MinionsResult
	err := json.Unmarshal(jsonResponse, &minionsResult)
	goutils.CheckForErrors(goutils.ErrorParams{Err: err, CallerNum: 1})

	// See if minions discovered have been registered. If not, register
	for _, minion := range m {
		registered := false
		for _, registeredMinion := range minionsResult.Minions {
			if registeredMinion.HostIP == minion.PublicIP {
				registered = true
			}
		}

		if !registered {
			register(endpoint, minion.PublicIP)
			time.Sleep(500 * time.Millisecond)
		}
	}
}

func register(endpoint, addr string) error {
	m := &PreregisteredMinion{
		Kind:       "Minion",
		APIVersion: K8S_API_VERSION,
		ID:         addr,
		HostIP:     addr,
	}
	data, err := json.Marshal(m)
	goutils.CheckForErrors(goutils.ErrorParams{Err: err, CallerNum: 1})

	url := fmt.Sprintf("%s/api/%s/minions", endpoint, K8S_API_VERSION)

	headers := map[string]string{
		"Content-Type": "application/json",
	}

	h := goutils.HttpRequestParams{
		HttpRequestType: "POST",
		Url:             url,
		Data:            data,
		Headers:         headers,
	}
	statusCode := goutils.HttpCreateRequest(h)

	switch statusCode {
	case 200, 202:
		log.Printf("------------------------------------------------")
		log.Printf("Registered machine with the master: %s\n", addr)
		return nil
	case 409:
		return nil
	}
}
<|endoftext|>


<commit_before>package mount

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/dnephin/dobi/config"
	"github.com/dnephin/dobi/logging"
	"github.com/dnephin/dobi/tasks/context"
	"github.com/dnephin/dobi/tasks/task"
	"github.com/dnephin/dobi/tasks/types"
	docker "github.com/fsouza/go-dockerclient"
)

// Task is a mount task
type Task struct {
	types.NoStop
	name   task.Name
	config *config.MountConfig
	run    func(*Task, *context.ExecuteContext) (bool, error)
}

// Name returns the name of the task
func (t *Task) Name() task.Name {
	return t.name
}

// Repr formats the task for logging
func (t *Task) Repr() string {
	return fmt.Sprintf("%s %s:%s", t.name.Format("mount"), t.config.Bind, t.config.Path)
}

// Run performs the task action
func (t *Task) Run(ctx *context.ExecuteContext, _ bool) (bool, error) {
	return t.run(t, ctx)
}

type createAction struct {
	task *Task
}

// Run creates the host path if it doesn't already exist
func runCreate(task *Task, ctx *context.ExecuteContext) (bool, error) {
	c := createAction{task: task}
	return c.run(ctx)
}

func (t *createAction) run(ctx *context.ExecuteContext) (bool, error) {
	logger := logging.ForTask(t.task)

	if t.exists(ctx) {
		logger.Debug("is fresh")
		return false, nil
	}

	var err error
	switch {
	case t.task.config.IsBind():
		err = t.createBind(ctx)
	default:
		err = t.createNamed(ctx)
	}
	if err != nil {
		return false, err
	}
	logger.Info("Created")
	return true, nil
}

func (t *createAction) createBind(ctx *context.ExecuteContext) error {
	path := AbsBindPath(t.task.config, ctx.WorkingDir)
	mode := os.FileMode(t.task.config.Mode)

	switch t.task.config.File {
	case true:
		return ioutil.WriteFile(path, []byte{}, mode)
	default:
		return os.MkdirAll(path, mode)
	}
}

func (t *createAction) createNamed(ctx *context.ExecuteContext) error {
	_, err := ctx.Client.CreateVolume(docker.CreateVolumeOptions{
		Name: t.task.config.Name,
	})
	return err
}

func (t *createAction) exists(ctx *context.ExecuteContext) bool {
	_, err := os.Stat(AbsBindPath(t.task.config, ctx.WorkingDir))
	return err == nil
}

func remove(task *Task, ctx *context.ExecuteContext) (bool, error) {
	if task.config.Name == "" {
		logging.ForTask(task).Warn("Bind mounts are not removable")
		return false, nil
	}

	return true, ctx.Client.RemoveVolume(task.config.Name)
}
<commit_msg>Catch and log mount removal failure as `Warning` instead of propogating an error.<commit_after>package mount

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/dnephin/dobi/config"
	"github.com/dnephin/dobi/logging"
	"github.com/dnephin/dobi/tasks/context"
	"github.com/dnephin/dobi/tasks/task"
	"github.com/dnephin/dobi/tasks/types"
	docker "github.com/fsouza/go-dockerclient"
	log "github.com/sirupsen/logrus"
)

// Task is a mount task
type Task struct {
	types.NoStop
	name   task.Name
	config *config.MountConfig
	run    func(*Task, *context.ExecuteContext) (bool, error)
}

// Name returns the name of the task
func (t *Task) Name() task.Name {
	return t.name
}

func (t *Task) logger() *log.Entry {
	return logging.ForTask(t)
}

// Repr formats the task for logging
func (t *Task) Repr() string {
	return fmt.Sprintf("%s %s:%s", t.name.Format("mount"), t.config.Bind, t.config.Path)
}

// Run performs the task action
func (t *Task) Run(ctx *context.ExecuteContext, _ bool) (bool, error) {
	return t.run(t, ctx)
}

type createAction struct {
	task *Task
}

// Run creates the host path if it doesn't already exist
func runCreate(task *Task, ctx *context.ExecuteContext) (bool, error) {
	c := createAction{task: task}
	return c.run(ctx)
}

func (t *createAction) run(ctx *context.ExecuteContext) (bool, error) {
	logger := logging.ForTask(t.task)

	if t.exists(ctx) {
		logger.Debug("is fresh")
		return false, nil
	}

	var err error
	switch {
	case t.task.config.IsBind():
		err = t.createBind(ctx)
	default:
		err = t.createNamed(ctx)
	}
	if err != nil {
		return false, err
	}
	logger.Info("Created")
	return true, nil
}

func (t *createAction) createBind(ctx *context.ExecuteContext) error {
	path := AbsBindPath(t.task.config, ctx.WorkingDir)
	mode := os.FileMode(t.task.config.Mode)

	switch t.task.config.File {
	case true:
		return ioutil.WriteFile(path, []byte{}, mode)
	default:
		return os.MkdirAll(path, mode)
	}
}

func (t *createAction) createNamed(ctx *context.ExecuteContext) error {
	_, err := ctx.Client.CreateVolume(docker.CreateVolumeOptions{
		Name: t.task.config.Name,
	})
	return err
}

func (t *createAction) exists(ctx *context.ExecuteContext) bool {
	_, err := os.Stat(AbsBindPath(t.task.config, ctx.WorkingDir))
	return err == nil
}

func remove(task *Task, ctx *context.ExecuteContext) (bool, error) {
	if task.config.Name == "" {
		logging.ForTask(task).Warn("Bind mounts are not removable")
		return false, nil
	}

	if err := ctx.Client.RemoveVolume(task.config.Name); err != nil {
		task.logger().Warnf("failed to remove %q: %s", task.config.Name, err)
	}

	return true, nil
}
<|endoftext|>


<commit_before>package migration_test

import (
	"database/sql"
	"fmt"
	"regexp"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/nicday/turtle/db"
	. "github.com/nicday/turtle/migration"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("migration", func() {
	// Setup a mock DB connection
	mockDB, err := sqlmock.New()
	if err != nil {
		panic(err)
	}

	db.Conn = mockDB

	// Setup a mock filesystem
	mockFS := NewMockFS()
	mockFS.AddFiles(
		"",
		NewMockFile("migrations", []byte(""),
			NewMockFile("20150703234300001_first_up.sql", []byte("CREATE TABLE first")),
			NewMockFile("20150703234300002_second_up.sql", []byte("CREATE TABLE second")),
			NewMockFile("20150703234300003_third_up.sql", []byte("CREATE TABLE third")),
			NewMockFile("20150703234300001_first_down.sql", []byte("DROP TABLE first")),
			NewMockFile("20150703234300002_second_down.sql", []byte("DROP TABLE second")),
			NewMockFile("20150703234300003_third_down.sql", []byte("DROP TABLE third")),
		),
	)

	FS = mockFS

	db.MigrationsTableName = "migrations"

	Describe("#AddPath", func() {
		Context("with an up migration path", func() {
			It("sets UpPath to the path", func() {
				p := "20150703234300001_first_up.sql"
				m := Migration{}
				m.AddPath(p)

				Expect(m.UpPath).To(Equal(p))
				Expect(m.DownPath).To(Equal(""))
			})
		})

		Context("with an down migration path", func() {
			It("sets DownPath to the path", func() {
				p := "20150703234300001_first_down.sql"
				m := Migration{}
				m.AddPath(p)

				Expect(m.DownPath).To(Equal(p))
				Expect(m.UpPath).To(Equal(""))
			})
		})
	})

	Describe("#Apply", func() {

	})

	Describe("#Revert", func() {

	})

	Describe(".ApplyAll", func() {
		Context("with no active migrations", func() {
			It("applies all migrations", func() {
				ExpectedMigrationActiveQuery("20150703234300001_first", false)
				ExpectedMigration("CREATE TABLE first")
				ExpectedMigrationLogInsert("20150703234300001_first")

				ExpectedMigrationActiveQuery("20150703234300002_second", false)
				ExpectedMigration("CREATE TABLE second")
				ExpectedMigrationLogInsert("20150703234300002_second")

				ExpectedMigrationActiveQuery("20150703234300003_third", false)
				ExpectedMigration("CREATE TABLE third")
				ExpectedMigrationLogInsert("20150703234300003_third")

				err := ApplyAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with some active migrations", func() {
			It("applies all inactive migrations", func() {
				ExpectedMigrationActiveQuery("20150703234300001_first", true)

				ExpectedMigrationActiveQuery("20150703234300002_second", false)
				ExpectedMigration("CREATE TABLE second")
				ExpectedMigrationLogInsert("20150703234300002_second")

				ExpectedMigrationActiveQuery("20150703234300003_third", false)
				ExpectedMigration("CREATE TABLE third")
				ExpectedMigrationLogInsert("20150703234300003_third")

				err := ApplyAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with all active migrations", func() {
			It("doesn't apply any migrations", func() {
				ExpectedMigrationActiveQuery("20150703234300001_first", true)

				ExpectedMigrationActiveQuery("20150703234300002_second", true)

				ExpectedMigrationActiveQuery("20150703234300003_third", true)

				err := ApplyAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

	})

	Describe(".RevertAll", func() {
		Context("with all active migrations", func() {
			It("reverts all migrations", func() {
				ExpectedMigrationActiveQuery("20150703234300003_third", true)
				ExpectedMigration("DROP TABLE third")
				ExpectedMigrationLogDelete("20150703234300003_third")

				ExpectedMigrationActiveQuery("20150703234300002_second", true)
				ExpectedMigration("DROP TABLE second")
				ExpectedMigrationLogDelete("20150703234300002_second")

				ExpectedMigrationActiveQuery("20150703234300001_first", true)
				ExpectedMigration("DROP TABLE first")
				ExpectedMigrationLogDelete("20150703234300001_first")

				err := RevertAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with some active migrations", func() {
			It("reverts all active migrations", func() {
				ExpectedMigrationActiveQuery("20150703234300003_third", false)

				ExpectedMigrationActiveQuery("20150703234300002_second", false)

				ExpectedMigrationActiveQuery("20150703234300001_first", true)
				ExpectedMigration("DROP TABLE first")
				ExpectedMigrationLogDelete("20150703234300001_first")

				err := RevertAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with all migrations inactive", func() {
			It("doesn't revert any migrations", func() {
				ExpectedMigrationActiveQuery("20150703234300003_third", false)

				ExpectedMigrationActiveQuery("20150703234300002_second", false)

				ExpectedMigrationActiveQuery("20150703234300001_first", false)

				err := RevertAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})
	})
})

func ExpectedMigration(sql string) {
	// Migration transaction
	sqlmock.ExpectBegin()
	sqlmock.ExpectExec(regexp.QuoteMeta(sql)).
		WillReturnResult(sqlmock.NewResult(0, 0))
	sqlmock.ExpectCommit()
}

func ExpectedMigrationLogInsert(id string) {
	expectedMigrationLogSQL := fmt.Sprintf(
		"INSERT INTO %s (migration_id) VALUES (?)",
		db.MigrationsTableName,
	)
	sqlmock.ExpectExec(regexp.QuoteMeta(expectedMigrationLogSQL)).
		WithArgs(id).
		WillReturnResult(sqlmock.NewResult(0, 1))
}

func ExpectedMigrationLogDelete(id string) {
	expectedMigrationLogSQL := fmt.Sprintf(
		"DELETE FROM %s WHERE migration_id=?",
		db.MigrationsTableName,
	)
	sqlmock.ExpectExec(regexp.QuoteMeta(expectedMigrationLogSQL)).
		WithArgs(id).
		WillReturnResult(sqlmock.NewResult(0, 1))
}

func ExpectedMigrationActiveQuery(id string, active bool) {
	expectedSQL := fmt.Sprintf(
		"SELECT id FROM %s WHERE migration_id=?",
		db.MigrationsTableName,
	)
	query := sqlmock.ExpectQuery(regexp.QuoteMeta(expectedSQL)).
		WithArgs(id)

	if active {
		query.WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
	} else {
		query.WillReturnError(sql.ErrNoRows)
	}

func expectMigrationsTablePresenceQuery() {
	expectedSQL := fmt.Sprintf(
		"SELECT 1 FROM %s LIMIT 1",
		db.MigrationsTableName,
	)
	sqlmock.ExpectQuery(regexp.QuoteMeta(expectedSQL)).
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
}
<commit_msg>Unexport test helpers<commit_after>package migration_test

import (
	"database/sql"
	"fmt"
	"regexp"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/nicday/turtle/db"
	. "github.com/nicday/turtle/migration"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("migration", func() {
	// Setup a mock DB connection
	mockDB, err := sqlmock.New()
	if err != nil {
		panic(err)
	}

	db.Conn = mockDB

	// Setup a mock filesystem
	mockFS := NewMockFS()
	mockFS.AddFiles(
		"",
		NewMockFile("migrations", []byte(""),
			NewMockFile("20150703234300001_first_up.sql", []byte("CREATE TABLE first")),
			NewMockFile("20150703234300002_second_up.sql", []byte("CREATE TABLE second")),
			NewMockFile("20150703234300003_third_up.sql", []byte("CREATE TABLE third")),
			NewMockFile("20150703234300001_first_down.sql", []byte("DROP TABLE first")),
			NewMockFile("20150703234300002_second_down.sql", []byte("DROP TABLE second")),
			NewMockFile("20150703234300003_third_down.sql", []byte("DROP TABLE third")),
		),
	)

	FS = mockFS

	db.MigrationsTableName = "migrations"

	Describe("#AddPath", func() {
		Context("with an up migration path", func() {
			It("sets UpPath to the path", func() {
				p := "20150703234300001_first_up.sql"
				m := Migration{}
				m.AddPath(p)

				Expect(m.UpPath).To(Equal(p))
				Expect(m.DownPath).To(Equal(""))
			})
		})

		Context("with an down migration path", func() {
			It("sets DownPath to the path", func() {
				p := "20150703234300001_first_down.sql"
				m := Migration{}
				m.AddPath(p)

				Expect(m.DownPath).To(Equal(p))
				Expect(m.UpPath).To(Equal(""))
			})
		})
	})

	Describe("#Apply", func() {

	})

	Describe("#Revert", func() {

	})

	Describe(".ApplyAll", func() {
		Context("with no active migrations", func() {
			It("applies all migrations", func() {
				expectMigrationsTablePresenceQuery()

				expectedMigrationActiveQuery("20150703234300001_first", false)
				expectedMigration("CREATE TABLE first")
				expectedMigrationLogInsert("20150703234300001_first")

				expectedMigrationActiveQuery("20150703234300002_second", false)
				expectedMigration("CREATE TABLE second")
				expectedMigrationLogInsert("20150703234300002_second")

				expectedMigrationActiveQuery("20150703234300003_third", false)
				expectedMigration("CREATE TABLE third")
				expectedMigrationLogInsert("20150703234300003_third")

				err := ApplyAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with some active migrations", func() {
			It("applies all inactive migrations", func() {
				expectMigrationsTablePresenceQuery()

				expectedMigrationActiveQuery("20150703234300001_first", true)

				expectedMigrationActiveQuery("20150703234300002_second", false)
				expectedMigration("CREATE TABLE second")
				expectedMigrationLogInsert("20150703234300002_second")

				expectedMigrationActiveQuery("20150703234300003_third", false)
				expectedMigration("CREATE TABLE third")
				expectedMigrationLogInsert("20150703234300003_third")

				err := ApplyAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with all active migrations", func() {
			It("doesn't apply any migrations", func() {
				expectMigrationsTablePresenceQuery()

				expectedMigrationActiveQuery("20150703234300001_first", true)

				expectedMigrationActiveQuery("20150703234300002_second", true)

				expectedMigrationActiveQuery("20150703234300003_third", true)

				err := ApplyAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

	})

	Describe(".RevertAll", func() {
		Context("with all active migrations", func() {
			It("reverts all migrations", func() {
				expectMigrationsTablePresenceQuery()

				expectedMigrationActiveQuery("20150703234300003_third", true)
				expectedMigration("DROP TABLE third")
				expectedMigrationLogDelete("20150703234300003_third")

				expectedMigrationActiveQuery("20150703234300002_second", true)
				expectedMigration("DROP TABLE second")
				expectedMigrationLogDelete("20150703234300002_second")

				expectedMigrationActiveQuery("20150703234300001_first", true)
				expectedMigration("DROP TABLE first")
				expectedMigrationLogDelete("20150703234300001_first")

				err := RevertAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with some active migrations", func() {
			It("reverts all active migrations", func() {
				expectMigrationsTablePresenceQuery()

				expectedMigrationActiveQuery("20150703234300003_third", false)

				expectedMigrationActiveQuery("20150703234300002_second", false)

				expectedMigrationActiveQuery("20150703234300001_first", true)
				expectedMigration("DROP TABLE first")
				expectedMigrationLogDelete("20150703234300001_first")

				err := RevertAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})

		Context("with all migrations inactive", func() {
			It("doesn't revert any migrations", func() {
				expectMigrationsTablePresenceQuery()

				expectedMigrationActiveQuery("20150703234300003_third", false)

				expectedMigrationActiveQuery("20150703234300002_second", false)

				expectedMigrationActiveQuery("20150703234300001_first", false)

				err := RevertAll()

				Expect(err).NotTo(HaveOccurred())
			})
		})
	})
})

func expectedMigration(sql string) {
	// Migration transaction
	sqlmock.ExpectBegin()
	sqlmock.ExpectExec(regexp.QuoteMeta(sql)).
		WillReturnResult(sqlmock.NewResult(0, 0))
	sqlmock.ExpectCommit()
}

func expectedMigrationLogInsert(id string) {
	expectedSQL := fmt.Sprintf(
		"INSERT INTO %s (migration_id) VALUES (?)",
		db.MigrationsTableName,
	)
	sqlmock.ExpectExec(regexp.QuoteMeta(expectedSQL)).
		WithArgs(id).
		WillReturnResult(sqlmock.NewResult(0, 1))
}

func expectedMigrationLogDelete(id string) {
	expectedSQL := fmt.Sprintf(
		"DELETE FROM %s WHERE migration_id=?",
		db.MigrationsTableName,
	)
	sqlmock.ExpectExec(regexp.QuoteMeta(expectedSQL)).
		WithArgs(id).
		WillReturnResult(sqlmock.NewResult(0, 1))
}

func expectedMigrationActiveQuery(id string, active bool) {
	expectedSQL := fmt.Sprintf(
		"SELECT id FROM %s WHERE migration_id=?",
		db.MigrationsTableName,
	)
	query := sqlmock.ExpectQuery(regexp.QuoteMeta(expectedSQL)).
		WithArgs(id)

	if active {
		query.WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
	} else {
		query.WillReturnError(sql.ErrNoRows)
	}
}

func expectMigrationsTablePresenceQuery() {
	expectedSQL := fmt.Sprintf(
		"SELECT 1 FROM %s LIMIT 1",
		db.MigrationsTableName,
	)
	sqlmock.ExpectQuery(regexp.QuoteMeta(expectedSQL)).
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
}
<|endoftext|>


<commit_before>// Package telegram implements Telegram handlers for github.com/go-chat-bot/bot
package telegram

import (
	"log"
	"strconv"

	"github.com/Syfaro/telegram-bot-api"
	"github.com/go-chat-bot/bot"
)

var (
	tg *tgbotapi.BotAPI
)

func responseHandler(target string, message string, sender *bot.User) {
	id, err := strconv.Atoi(target)
	if err != nil {
		log.Println(err)
		return
	}
	msg := tgbotapi.NewMessage(id, message)
	msg.ReplyToMessageID = msg.ReplyToMessageID

	tg.Send(msg)
}

// Run executes the bot and connects to Telegram using the provided token. Use the debug flag if you wish to see all trafic logged
func Run(token string, debug bool) {
	var err error
	tg, err = tgbotapi.NewBotAPI(token)
	if err != nil {
		log.Fatal(err)
	}

	tg.Debug = debug

	log.Printf("Authorized on account %s", tg.Self.UserName)

	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60

	updates, err := tg.GetUpdatesChan(u)
	if err != nil {
		log.Fatal(err)
	}

	b := bot.New(&bot.Handlers{
		Response: responseHandler,
	})
	b.Disable([]string{"url"})

	for update := range updates {
		target := strconv.Itoa(update.Message.Chat.ID)
		sender := strconv.Itoa(update.Message.From.ID)
		b.MessageReceived(target, update.Message.Text, &bot.User{Nick: sender})
	}
}
<commit_msg>Ensure Chat IDs are valid for Telegram supergroups.<commit_after>// Package telegram implements Telegram handlers for github.com/go-chat-bot/bot
package telegram

import (
	"log"
	"strconv"

	"github.com/go-chat-bot/bot"
	"gopkg.in/telegram-bot-api.v3"
)

var (
	tg *tgbotapi.BotAPI
)

func responseHandler(target string, message string, sender *bot.User) {
	id, err := strconv.ParseInt(target, 10, 64)
	if err != nil {
		log.Println(err)
		return
	}
	msg := tgbotapi.NewMessage(id, message)
	msg.ReplyToMessageID = msg.ReplyToMessageID

	tg.Send(msg)
}

// Run executes the bot and connects to Telegram using the provided token. Use the debug flag if you wish to see all trafic logged
func Run(token string, debug bool) {
	var err error
	tg, err = tgbotapi.NewBotAPI(token)
	if err != nil {
		log.Fatal(err)
	}

	tg.Debug = debug

	log.Printf("Authorized on account %s", tg.Self.UserName)

	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60

	updates, err := tg.GetUpdatesChan(u)
	if err != nil {
		log.Fatal(err)
	}

	b := bot.New(&bot.Handlers{
		Response: responseHandler,
	})
	b.Disable([]string{"url"})

	for update := range updates {
		target := strconv.FormatInt(update.Message.Chat.ID, 10)
		sender := strconv.Itoa(update.Message.From.ID)
		b.MessageReceived(target, update.Message.Text, &bot.User{Nick: sender})
	}
}
<|endoftext|>


<commit_before>package registry

import (
	"net"
	"net/url"
	"time"

	"github.com/pivotal-golang/bytefmt"
	"github.com/snagles/docker-registry-manager/models/client"
	"github.com/snagles/docker-registry-manager/utilities"
)

// Registries contains a map of all active registries identified by their name
var Registries map[string]*Registry

func init() {
	Registries = make(map[string]*Registry, 0)
}

func AddRegistry(uri string) error {
	r, err := ParseRegistry(uri)
	if err != nil {
		return err
	}

	repos, err := client.GetRepositories(uri)
	if err != nil {
		return err
	}

	// Use a layer map to de-duplicate shared layers
	layerMap := make(map[string]int64, 0)

	// Get the tags and image information
	for _, repoName := range repos {
		r.Repositories = append(r.Repositories, Repository{Name: repoName})

		tagList, err := client.GetTags(uri, repoName)
		if err == nil {
			r.TagCount = len(tagList)
			for _, tagName := range tagList {
				img, _ := client.GetImage(uri, repoName, tagName)
				for _, layer := range img.FsLayers {
					layerMap[layer.BlobSum] = layer.Size
				}
			}
		}
	}

	// Total the size of the layers
	for _, size := range layerMap {
		r.RepoTotalSize += size
	}
	r.RepoTotalSizeStr = bytefmt.ByteSize(uint64(r.RepoTotalSize))

	Registries[r.Name] = &r

	return nil
}

// Registry contains all identifying information for communicating with a registry
type Registry struct {
	Name             string
	IP               string
	Scheme           string
	Port             string
	Version          string
	Repositories     []Repository
	TagCount         int
	Status           string
	RepoTotalSize    int64
	RepoTotalSizeStr string
}

// URI returns the full url path for communicating with this registry
func (r *Registry) URI() string {
	return r.Scheme + "://" + r.Name + ":" + r.Port + "/v2"
}

type Repository struct {
	Name string
	Tags []Tag
}

type Tag struct {
	Image client.Image

	ID              string
	Name            string
	UpdatedTime     time.Time
	UpdatedTimeUnix int64
	TimeAgo         string
	Layers          int
	Size            string
	SizeInt         int64
}

// ParseRegistry takes in a registry URI string and converts it into a registry object
func ParseRegistry(registry string) (Registry, error) {

	// only supports v2 currently
	r := Registry{
		Version: "v2",
	}

	// Parse the URL and get the scheme
	// e.g https, http, etc.
	u, err := url.Parse(registry)
	if err != nil {
		utils.Log.Error(err)
		return r, err
	}

	// Set scheme
	r.Scheme = u.Scheme

	// Get the host and port
	// e.g test.domain.com and 5000, etc.
	host, port, err := net.SplitHostPort(u.Host)
	if err != nil {
		utils.Log.Error(err)
		return r, err
	}

	// Set name and port
	r.Name = host
	r.Port = port

	// Lookup the ip for the passed host
	ip, err := net.LookupHost(host)
	if err != nil {
		utils.Log.Error(err)
		return r, err
	}
	if len(ip) > 0 {
		r.IP = ip[0]
	}

	return r, err
}
<commit_msg>Changed registry 'building' to a method to allow for future refreshes at set intervals<commit_after>package registry

import (
	"net"
	"net/url"
	"time"

	"github.com/pivotal-golang/bytefmt"
	"github.com/snagles/docker-registry-manager/models/client"
	"github.com/snagles/docker-registry-manager/utilities"
)

// Registries contains a map of all active registries identified by their name
var Registries map[string]*Registry

func init() {
	Registries = make(map[string]*Registry, 0)
}

func AddRegistry(uri string) error {
	r, err := ParseRegistry(uri)
	if err != nil {
		return err
	}

	r.Refresh()

	return nil
}

// Registry contains all identifying information for communicating with a registry
type Registry struct {
	Name             string
	IP               string
	Scheme           string
	Port             string
	Version          string
	Repositories     []Repository
	TagCount         int
	Status           string
	RepoTotalSize    int64
	RepoTotalSizeStr string
}

// URI returns the full url path for communicating with this registry
func (r *Registry) URI() string {
	return r.Scheme + "://" + r.Name + ":" + r.Port + "/v2"
}

func (r *Registry) Refresh() {

	// Get the lsit of repositories on this registry
	repoList, _ := client.GetRepositories(r.URI())

	// Use a layer map to de-duplicate shared layers across the registry
	registryLayerMap := make(map[string]int64, 0)

	// Get the repository information
	for _, repoName := range repoList {

		// Build a repository object
		repo := Repository{Name: repoName}

		tagList, _ := client.GetTags(r.URI(), repoName)
		for _, tagName := range tagList {
			tag := Tag{Name: tagName}
			tag.Image, _ = client.GetImage(r.URI(), repoName, tagName)
			for _, layer := range tag.Image.FsLayers {
				registryLayerMap[layer.BlobSum] = layer.Size
			}

			// Add the tag to the repository
			repo.Tags = append(repo.Tags, tag)
		}
		r.Repositories = append(r.Repositories, repo)
		r.TagCount += len(tagList)
	}

	// Total the size of the layers
	for _, size := range registryLayerMap {
		r.RepoTotalSize += size
	}
	r.RepoTotalSizeStr = bytefmt.ByteSize(uint64(r.RepoTotalSize))

	Registries[r.Name] = r
}

type Repository struct {
	Name string
	Tags []Tag
}

type Tag struct {
	Image client.Image

	ID              string
	Name            string
	UpdatedTime     time.Time
	UpdatedTimeUnix int64
	TimeAgo         string
	Layers          int
	Size            string
	SizeInt         int64
}

// ParseRegistry takes in a registry URI string and converts it into a registry object
func ParseRegistry(registry string) (Registry, error) {

	// only supports v2 currently
	r := Registry{
		Version: "v2",
	}

	// Parse the URL and get the scheme
	// e.g https, http, etc.
	u, err := url.Parse(registry)
	if err != nil {
		utils.Log.Error(err)
		return r, err
	}

	// Set scheme
	r.Scheme = u.Scheme

	// Get the host and port
	// e.g test.domain.com and 5000, etc.
	host, port, err := net.SplitHostPort(u.Host)
	if err != nil {
		utils.Log.Error(err)
		return r, err
	}

	// Set name and port
	r.Name = host
	r.Port = port

	// Lookup the ip for the passed host
	ip, err := net.LookupHost(host)
	if err != nil {
		utils.Log.Error(err)
		return r, err
	}
	if len(ip) > 0 {
		r.IP = ip[0]
	}

	return r, err
}
<|endoftext|>


<commit_before>package plugins

import (
    "github.com/bwmarrin/discordgo"
    "strings"
    "github.com/Seklfreak/Robyul2/helpers"
    "fmt"
    "github.com/getsentry/raven-go"
)

type AutoRoles struct{}

func (a *AutoRoles) Commands() []string {
    return []string{
        "autorole",
        "autoroles",
    }
}

func (a *AutoRoles) Init(session *discordgo.Session) {
}

func (a *AutoRoles) Action(command string, content string, msg *discordgo.Message, session *discordgo.Session) {
    args := strings.Fields(content)
    if len(args) >= 1 {
        switch args[0] {
        case "add":
            session.ChannelTyping(msg.ChannelID)
            helpers.RequireAdmin(msg, func() {
                if len(args) < 2 {
                    _, err := session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
                    helpers.Relax(err)
                    return
                }
                channel, err := helpers.GetChannel(msg.ChannelID)
                helpers.Relax(err)

                serverRoles, err := session.GuildRoles(channel.GuildID)
                helpers.Relax(err)

                roleNameToMatch := strings.TrimSpace(strings.Replace(content, strings.Join(args[:1], " "), "", 1))

                var targetRole *discordgo.Role
                for _, role := range serverRoles {
                    if strings.ToLower(role.Name) == strings.ToLower(roleNameToMatch) || role.ID == roleNameToMatch {
                        targetRole = role
                    }
                }
                if targetRole == nil || targetRole.ID == "" {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
                    helpers.Relax(err)
                    return
                }

                settings := helpers.GuildSettingsGetCached(channel.GuildID)

                for _, role := range settings.AutoRoleIDs {
                    if role == targetRole.ID {
                        _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-add-error-duplicate"))
                        helpers.Relax(err)
                        return
                    }
                }

                settings.AutoRoleIDs = append(settings.AutoRoleIDs, targetRole.ID)

                err = helpers.GuildSettingsSet(channel.GuildID, settings)
                helpers.Relax(err)

                _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetTextF("plugins.autorole.role-add-success",
                    targetRole.Name))
                helpers.Relax(err)
                return
            })
            return
        case "list":
            session.ChannelTyping(msg.ChannelID)
            channel, err := helpers.GetChannel(msg.ChannelID)
            helpers.Relax(err)
            settings := helpers.GuildSettingsGetCached(channel.GuildID)

            if len(settings.AutoRoleIDs) <= 0 {
                _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-list-none"))
                helpers.Relax(err)
                return
            }

            result := "AutoRoles on this server: "

            for _, roleID := range settings.AutoRoleIDs {
                role, err := session.State.Role(channel.GuildID, roleID)
                if err == nil {
                    result += fmt.Sprintf("`%s (#%s)` ", role.Name, role.ID)
                } else {
                    result += fmt.Sprintf("`N/A (#%s)` ", roleID)
                }
            }

            result += fmt.Sprintf("(%d role(s))", len(settings.AutoRoleIDs))

            _, err = session.ChannelMessageSend(msg.ChannelID, result)
            helpers.Relax(err)
            return
        case "delete", "remove":
            session.ChannelTyping(msg.ChannelID)
            helpers.RequireAdmin(msg, func() {
                if len(args) < 2 {
                    _, err := session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
                    helpers.Relax(err)
                    return
                }
                channel, err := helpers.GetChannel(msg.ChannelID)
                helpers.Relax(err)

                serverRoles, err := session.GuildRoles(channel.GuildID)
                helpers.Relax(err)

                roleNameToMatch := strings.TrimSpace(strings.Replace(content, strings.Join(args[:1], " "), "", 1))

                var targetRole *discordgo.Role
                for _, role := range serverRoles {
                    if strings.ToLower(role.Name) == strings.ToLower(roleNameToMatch) || role.ID == roleNameToMatch {
                        targetRole = role
                    }
                }
                if targetRole == nil || targetRole.ID == "" {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
                    helpers.Relax(err)
                    return
                }

                settings := helpers.GuildSettingsGetCached(channel.GuildID)

                roleWasInList := false
                newRoleIDs := make([]string, 0)

                for _, role := range settings.AutoRoleIDs {
                    if role == targetRole.ID {
                        roleWasInList = true
                    } else {
                        newRoleIDs = append(newRoleIDs, role)
                    }
                }

                if roleWasInList == false {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-remove-error-not-found"))
                    helpers.Relax(err)
                    return
                }

                settings.AutoRoleIDs = newRoleIDs

                err = helpers.GuildSettingsSet(channel.GuildID, settings)
                helpers.Relax(err)

                _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-remove-success"))
                helpers.Relax(err)
                return
            })
            return
        case "apply":
            session.ChannelTyping(msg.ChannelID)
            helpers.RequireAdmin(msg, func() {
                if len(args) < 2 {
                    _, err := session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
                    helpers.Relax(err)
                    return
                }
                channel, err := helpers.GetChannel(msg.ChannelID)
                helpers.Relax(err)

                serverRoles, err := session.GuildRoles(channel.GuildID)
                helpers.Relax(err)

                roleNameToMatch := strings.TrimSpace(strings.Replace(content, strings.Join(args[:1], " "), "", 1))

                var targetRole *discordgo.Role
                for _, role := range serverRoles {
                    if strings.ToLower(role.Name) == strings.ToLower(roleNameToMatch) || role.ID == roleNameToMatch {
                        targetRole = role
                    }
                }
                if targetRole == nil || targetRole.ID == "" {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
                    helpers.Relax(err)
                    return
                }

                users := make([]string, 0)
                lastAfterMemberId := ""
                for {
                    members, err := session.GuildMembers(channel.GuildID, lastAfterMemberId, 1000)
                    helpers.Relax(err)
                    if len(members) <= 0 {
                        break
                    }

                    lastAfterMemberId = members[len(members)-1].User.ID
                    for _, u := range members {
                        users = append(users, u.User.ID)
                    }
                }

                if helpers.ConfirmEmbed(msg.ChannelID, msg.Author, helpers.GetTextF("plugins.autorole.apply-confirm",
                    targetRole.Name, targetRole.ID, len(users)),"✅", "🚫") {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.apply-started"))
                    helpers.Relax(err)

                    addedSuccess := 0
                    addedError := 0

                    for _, userID := range users {
                        err := session.GuildMemberRoleAdd(channel.GuildID, userID, targetRole.ID)
                        if err != nil {
                            addedError += 1
                        } else {
                            addedSuccess += 1
                        }
                    }

                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetTextF("plugins.autorole.apply-done",
                    msg.Author.ID, addedSuccess, addedError))
                    helpers.Relax(err)
                    return
                }
                return
            })
            return
        }
    }
}

func (a *AutoRoles) OnMessage(content string, msg *discordgo.Message, session *discordgo.Session) {

}

func (a *AutoRoles) OnGuildMemberAdd(member *discordgo.Member, session *discordgo.Session) {
    go func() {
        settings := helpers.GuildSettingsGetCached(member.GuildID)
        for _, roleID := range settings.AutoRoleIDs {
            err := session.GuildMemberRoleAdd(member.GuildID, member.User.ID, roleID)
            if err != nil {
                raven.CaptureError(fmt.Errorf("%#v", err), map[string]string{})
                continue
            }
        }
    }()
}

func (a *AutoRoles) OnGuildMemberRemove(member *discordgo.Member, session *discordgo.Session) {

}

func (a *AutoRoles) OnReactionAdd(reaction *discordgo.MessageReactionAdd, session *discordgo.Session) {

}

func (a *AutoRoles) OnReactionRemove(reaction *discordgo.MessageReactionRemove, session *discordgo.Session) {

}

func (a *AutoRoles) OnGuildBanAdd(user *discordgo.GuildBanAdd, session *discordgo.Session) {

}

func (a *AutoRoles) OnGuildBanRemove(user *discordgo.GuildBanRemove, session *discordgo.Session) {

}
<commit_msg>[AutoRole] better error handling<commit_after>package plugins

import (
    "github.com/bwmarrin/discordgo"
    "strings"
    "github.com/Seklfreak/Robyul2/helpers"
    "fmt"
    "github.com/getsentry/raven-go"
)

type AutoRoles struct{}

func (a *AutoRoles) Commands() []string {
    return []string{
        "autorole",
        "autoroles",
    }
}

func (a *AutoRoles) Init(session *discordgo.Session) {
}

func (a *AutoRoles) Action(command string, content string, msg *discordgo.Message, session *discordgo.Session) {
    args := strings.Fields(content)
    if len(args) >= 1 {
        switch args[0] {
        case "add":
            session.ChannelTyping(msg.ChannelID)
            helpers.RequireAdmin(msg, func() {
                if len(args) < 2 {
                    _, err := session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
                    helpers.Relax(err)
                    return
                }
                channel, err := helpers.GetChannel(msg.ChannelID)
                helpers.Relax(err)

                serverRoles, err := session.GuildRoles(channel.GuildID)
                if err != nil {
                    if errD := err.(*discordgo.RESTError); errD != nil {
                        if errD.Message.Code == 50013 {
                            _, err = session.ChannelMessageSend(msg.ChannelID, "Please give me the `Manage Roles` permission to use this feature.")
                            helpers.Relax(err)
                            return
                        } else {
                            helpers.Relax(err)
                        }
                    } else {
                        helpers.Relax(err)
                    }
                }

                roleNameToMatch := strings.TrimSpace(strings.Replace(content, strings.Join(args[:1], " "), "", 1))

                var targetRole *discordgo.Role
                for _, role := range serverRoles {
                    if strings.ToLower(role.Name) == strings.ToLower(roleNameToMatch) || role.ID == roleNameToMatch {
                        targetRole = role
                    }
                }
                if targetRole == nil || targetRole.ID == "" {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
                    helpers.Relax(err)
                    return
                }

                settings := helpers.GuildSettingsGetCached(channel.GuildID)

                for _, role := range settings.AutoRoleIDs {
                    if role == targetRole.ID {
                        _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-add-error-duplicate"))
                        helpers.Relax(err)
                        return
                    }
                }

                settings.AutoRoleIDs = append(settings.AutoRoleIDs, targetRole.ID)

                err = helpers.GuildSettingsSet(channel.GuildID, settings)
                helpers.Relax(err)

                _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetTextF("plugins.autorole.role-add-success",
                    targetRole.Name))
                helpers.Relax(err)
                return
            })
            return
        case "list":
            session.ChannelTyping(msg.ChannelID)
            channel, err := helpers.GetChannel(msg.ChannelID)
            helpers.Relax(err)
            settings := helpers.GuildSettingsGetCached(channel.GuildID)

            if len(settings.AutoRoleIDs) <= 0 {
                _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-list-none"))
                helpers.Relax(err)
                return
            }

            result := "AutoRoles on this server: "

            for _, roleID := range settings.AutoRoleIDs {
                role, err := session.State.Role(channel.GuildID, roleID)
                if err == nil {
                    result += fmt.Sprintf("`%s (#%s)` ", role.Name, role.ID)
                } else {
                    result += fmt.Sprintf("`N/A (#%s)` ", roleID)
                }
            }

            result += fmt.Sprintf("(%d role(s))", len(settings.AutoRoleIDs))

            _, err = session.ChannelMessageSend(msg.ChannelID, result)
            helpers.Relax(err)
            return
        case "delete", "remove":
            session.ChannelTyping(msg.ChannelID)
            helpers.RequireAdmin(msg, func() {
                if len(args) < 2 {
                    _, err := session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
                    helpers.Relax(err)
                    return
                }
                channel, err := helpers.GetChannel(msg.ChannelID)
                helpers.Relax(err)

                serverRoles, err := session.GuildRoles(channel.GuildID)
                if err != nil {
                    if errD := err.(*discordgo.RESTError); errD != nil {
                        if errD.Message.Code == 50013 {
                            _, err = session.ChannelMessageSend(msg.ChannelID, "Please give me the `Manage Roles` permission to use this feature.")
                            helpers.Relax(err)
                            return
                        } else {
                            helpers.Relax(err)
                        }
                    } else {
                        helpers.Relax(err)
                    }
                }

                roleNameToMatch := strings.TrimSpace(strings.Replace(content, strings.Join(args[:1], " "), "", 1))

                var targetRole *discordgo.Role
                for _, role := range serverRoles {
                    if strings.ToLower(role.Name) == strings.ToLower(roleNameToMatch) || role.ID == roleNameToMatch {
                        targetRole = role
                    }
                }
                if targetRole == nil || targetRole.ID == "" {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
                    helpers.Relax(err)
                    return
                }

                settings := helpers.GuildSettingsGetCached(channel.GuildID)

                roleWasInList := false
                newRoleIDs := make([]string, 0)

                for _, role := range settings.AutoRoleIDs {
                    if role == targetRole.ID {
                        roleWasInList = true
                    } else {
                        newRoleIDs = append(newRoleIDs, role)
                    }
                }

                if roleWasInList == false {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-remove-error-not-found"))
                    helpers.Relax(err)
                    return
                }

                settings.AutoRoleIDs = newRoleIDs

                err = helpers.GuildSettingsSet(channel.GuildID, settings)
                helpers.Relax(err)

                _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.role-remove-success"))
                helpers.Relax(err)
                return
            })
            return
        case "apply":
            session.ChannelTyping(msg.ChannelID)
            helpers.RequireAdmin(msg, func() {
                if len(args) < 2 {
                    _, err := session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
                    helpers.Relax(err)
                    return
                }
                channel, err := helpers.GetChannel(msg.ChannelID)
                helpers.Relax(err)

                serverRoles, err := session.GuildRoles(channel.GuildID)
                if err != nil {
                    if errD := err.(*discordgo.RESTError); errD != nil {
                        if errD.Message.Code == 50013 {
                            _, err = session.ChannelMessageSend(msg.ChannelID, "Please give me the `Manage Roles` permission to use this feature.")
                            helpers.Relax(err)
                            return
                        } else {
                            helpers.Relax(err)
                        }
                    } else {
                        helpers.Relax(err)
                    }
                }

                roleNameToMatch := strings.TrimSpace(strings.Replace(content, strings.Join(args[:1], " "), "", 1))

                var targetRole *discordgo.Role
                for _, role := range serverRoles {
                    if strings.ToLower(role.Name) == strings.ToLower(roleNameToMatch) || role.ID == roleNameToMatch {
                        targetRole = role
                    }
                }
                if targetRole == nil || targetRole.ID == "" {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
                    helpers.Relax(err)
                    return
                }

                users := make([]string, 0)
                lastAfterMemberId := ""
                for {
                    members, err := session.GuildMembers(channel.GuildID, lastAfterMemberId, 1000)
                    helpers.Relax(err)
                    if len(members) <= 0 {
                        break
                    }

                    lastAfterMemberId = members[len(members)-1].User.ID
                    for _, u := range members {
                        users = append(users, u.User.ID)
                    }
                }

                if helpers.ConfirmEmbed(msg.ChannelID, msg.Author, helpers.GetTextF("plugins.autorole.apply-confirm",
                    targetRole.Name, targetRole.ID, len(users)),"✅", "🚫") {
                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetText("plugins.autorole.apply-started"))
                    helpers.Relax(err)

                    addedSuccess := 0
                    addedError := 0

                    for _, userID := range users {
                        err := session.GuildMemberRoleAdd(channel.GuildID, userID, targetRole.ID)
                        if err != nil {
                            addedError += 1
                        } else {
                            addedSuccess += 1
                        }
                    }

                    _, err = session.ChannelMessageSend(msg.ChannelID, helpers.GetTextF("plugins.autorole.apply-done",
                    msg.Author.ID, addedSuccess, addedError))
                    helpers.Relax(err)
                    return
                }
                return
            })
            return
        }
    }
}

func (a *AutoRoles) OnMessage(content string, msg *discordgo.Message, session *discordgo.Session) {

}

func (a *AutoRoles) OnGuildMemberAdd(member *discordgo.Member, session *discordgo.Session) {
    go func() {
        settings := helpers.GuildSettingsGetCached(member.GuildID)
        for _, roleID := range settings.AutoRoleIDs {
            err := session.GuildMemberRoleAdd(member.GuildID, member.User.ID, roleID)
            if err != nil {
                raven.CaptureError(fmt.Errorf("%#v", err), map[string]string{})
                continue
            }
        }
    }()
}

func (a *AutoRoles) OnGuildMemberRemove(member *discordgo.Member, session *discordgo.Session) {

}

func (a *AutoRoles) OnReactionAdd(reaction *discordgo.MessageReactionAdd, session *discordgo.Session) {

}

func (a *AutoRoles) OnReactionRemove(reaction *discordgo.MessageReactionRemove, session *discordgo.Session) {

}

func (a *AutoRoles) OnGuildBanAdd(user *discordgo.GuildBanAdd, session *discordgo.Session) {

}

func (a *AutoRoles) OnGuildBanRemove(user *discordgo.GuildBanRemove, session *discordgo.Session) {

}
<|endoftext|>


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

package helpers

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

	"github.com/cilium/cilium/test/config"

	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
)

var (
	//SSHMetaLogs is a buffer where all commands sent over ssh are saved.
	SSHMetaLogs = NewWriter(new(bytes.Buffer))
)

// SSHMeta contains metadata to SSH into a remote location to run tests
type SSHMeta struct {
	sshClient *SSHClient
	env       []string
	rawConfig []byte
	nodeName  string
	logger    *logrus.Entry
}

// CreateSSHMeta returns an SSHMeta with the specified host, port, and user, as
// well as an according SSHClient.
func CreateSSHMeta(host string, port int, user string) *SSHMeta {
	return &SSHMeta{
		sshClient: GetSSHClient(host, port, user),
	}
}

func (s *SSHMeta) String() string {
	return fmt.Sprintf("environment: %s, SSHClient: %s", s.env, s.sshClient.String())

}

// GetVagrantSSHMeta returns a SSHMeta initialized based on the provided
// SSH-config target.
func GetVagrantSSHMeta(vmName string) *SSHMeta {
	config, err := GetVagrantSSHMetadata(vmName)
	if err != nil {
		return nil
	}

	log.Debugf("generated SSHConfig for node %s", vmName)
	nodes, err := ImportSSHconfig(config)
	if err != nil {
		log.WithError(err).Error("Error importing ssh config")
		return nil
	}
	log.Debugf("done importing ssh config")
	node := nodes[vmName]
	if node == nil {
		log.Errorf("Node %s not found in ssh config", vmName)
		return nil
	}
	sshMeta := &SSHMeta{
		sshClient: node.GetSSHClient(),
		rawConfig: config,
		nodeName:  vmName,
	}

	sshMeta.setBasePath()
	return sshMeta
}

// setBasePath if the SSHConfig is defined we set the BasePath to the GOPATH,
// from golang 1.8 GOPATH is by default $HOME/go so we also check that.
func (s *SSHMeta) setBasePath() {
	if config.CiliumTestConfig.SSHConfig == "" {
		return
	}

	gopath := s.Exec("echo $GOPATH").SingleOut()
	if gopath != "" {
		BasePath = filepath.Join(gopath, CiliumPath)
		return
	}

	home := s.Exec("echo $HOME").SingleOut()
	if home == "" {
		return
	}

	BasePath = filepath.Join(home, "go", CiliumPath)
	return
}

// Execute executes cmd on the provided node and stores the stdout / stderr of
// the command in the provided buffers. Returns false if the command failed
// during its execution.
func (s *SSHMeta) Execute(cmd string, stdout io.Writer, stderr io.Writer) error {
	if stdout == nil {
		stdout = os.Stdout
	}

	if stderr == nil {
		stderr = os.Stderr
	}
	fmt.Fprintln(SSHMetaLogs, cmd)
	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}
	err := s.sshClient.RunCommand(command)
	return err
}

// ExecWithSudo returns the result of executing the provided cmd via SSH using
// sudo.
func (s *SSHMeta) ExecWithSudo(cmd string) *CmdRes {
	command := fmt.Sprintf("sudo %s", cmd)
	return s.Exec(command)
}

// Exec returns the results of executing the provided cmd via SSH.
func (s *SSHMeta) Exec(cmd string) *CmdRes {
	log.Debugf("running command: %s", cmd)
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)
	exit := true
	err := s.Execute(cmd, stdout, stderr)
	if err != nil {
		exit = false
	}
	res := CmdRes{
		cmd:     cmd,
		stdout:  stdout,
		stderr:  stderr,
		success: exit,
	}
	// Set the exitCode if the error is an ExitError
	if exiterr, ok := err.(*ssh.ExitError); ok {
		res.exitcode = exiterr.Waitmsg.ExitStatus()
	}

	res.SendToLog()
	return &res
}

// GetCopy returns a copy of SSHMeta, useful for parallel requests
func (s *SSHMeta) GetCopy() *SSHMeta {
	nodes, err := ImportSSHconfig(s.rawConfig)
	if err != nil {
		log.WithError(err).Error("while importing ssh config for meta copy")
		return nil
	}

	config := nodes[s.nodeName]
	if config == nil {
		log.Errorf("no node %s in imported config", s.nodeName)
		return nil
	}

	copy := &SSHMeta{
		sshClient: config.GetSSHClient(),
		rawConfig: s.rawConfig,
		nodeName:  s.nodeName,
	}

	return copy
}

// ExecContext returns the results of running cmd via SSH in the specified
// context.
func (s *SSHMeta) ExecContext(ctx context.Context, cmd string) *CmdRes {
	if ctx == nil {
		panic("no context provided")
	}

	fmt.Fprintln(SSHMetaLogs, cmd)
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}

	res := CmdRes{
		cmd:     cmd,
		stdout:  stdout,
		stderr:  stderr,
		success: false,
	}

	go func() {
		s.sshClient.RunCommandContext(ctx, command)
		res.SendToLog()
	}()

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

package helpers

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/cilium/cilium/test/config"

	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
)

var (
	//SSHMetaLogs is a buffer where all commands sent over ssh are saved.
	SSHMetaLogs = NewWriter(new(bytes.Buffer))
)

// SSHMeta contains metadata to SSH into a remote location to run tests
type SSHMeta struct {
	sshClient *SSHClient
	env       []string
	rawConfig []byte
	nodeName  string
	logger    *logrus.Entry
}

// CreateSSHMeta returns an SSHMeta with the specified host, port, and user, as
// well as an according SSHClient.
func CreateSSHMeta(host string, port int, user string) *SSHMeta {
	return &SSHMeta{
		sshClient: GetSSHClient(host, port, user),
	}
}

func (s *SSHMeta) String() string {
	return fmt.Sprintf("environment: %s, SSHClient: %s", s.env, s.sshClient.String())

}

// GetVagrantSSHMeta returns a SSHMeta initialized based on the provided
// SSH-config target.
func GetVagrantSSHMeta(vmName string) *SSHMeta {
	config, err := GetVagrantSSHMetadata(vmName)
	if err != nil {
		return nil
	}

	log.Debugf("generated SSHConfig for node %s", vmName)
	nodes, err := ImportSSHconfig(config)
	if err != nil {
		log.WithError(err).Error("Error importing ssh config")
		return nil
	}
	var node *SSHConfig
	log.Debugf("done importing ssh config")
	for name := range nodes {
		if strings.HasPrefix(name, vmName) {
			node = nodes[name]
			break
		}
	}
	if node == nil {
		log.Errorf("Node %s not found in ssh config", vmName)
		return nil
	}
	sshMeta := &SSHMeta{
		sshClient: node.GetSSHClient(),
		rawConfig: config,
		nodeName:  vmName,
	}

	sshMeta.setBasePath()
	return sshMeta
}

// setBasePath if the SSHConfig is defined we set the BasePath to the GOPATH,
// from golang 1.8 GOPATH is by default $HOME/go so we also check that.
func (s *SSHMeta) setBasePath() {
	if config.CiliumTestConfig.SSHConfig == "" {
		return
	}

	gopath := s.Exec("echo $GOPATH").SingleOut()
	if gopath != "" {
		BasePath = filepath.Join(gopath, CiliumPath)
		return
	}

	home := s.Exec("echo $HOME").SingleOut()
	if home == "" {
		return
	}

	BasePath = filepath.Join(home, "go", CiliumPath)
	return
}

// Execute executes cmd on the provided node and stores the stdout / stderr of
// the command in the provided buffers. Returns false if the command failed
// during its execution.
func (s *SSHMeta) Execute(cmd string, stdout io.Writer, stderr io.Writer) error {
	if stdout == nil {
		stdout = os.Stdout
	}

	if stderr == nil {
		stderr = os.Stderr
	}
	fmt.Fprintln(SSHMetaLogs, cmd)
	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}
	err := s.sshClient.RunCommand(command)
	return err
}

// ExecWithSudo returns the result of executing the provided cmd via SSH using
// sudo.
func (s *SSHMeta) ExecWithSudo(cmd string) *CmdRes {
	command := fmt.Sprintf("sudo %s", cmd)
	return s.Exec(command)
}

// Exec returns the results of executing the provided cmd via SSH.
func (s *SSHMeta) Exec(cmd string) *CmdRes {
	log.Debugf("running command: %s", cmd)
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)
	exit := true
	err := s.Execute(cmd, stdout, stderr)
	if err != nil {
		exit = false
	}
	res := CmdRes{
		cmd:     cmd,
		stdout:  stdout,
		stderr:  stderr,
		success: exit,
	}
	// Set the exitCode if the error is an ExitError
	if exiterr, ok := err.(*ssh.ExitError); ok {
		res.exitcode = exiterr.Waitmsg.ExitStatus()
	}

	res.SendToLog()
	return &res
}

// GetCopy returns a copy of SSHMeta, useful for parallel requests
func (s *SSHMeta) GetCopy() *SSHMeta {
	nodes, err := ImportSSHconfig(s.rawConfig)
	if err != nil {
		log.WithError(err).Error("while importing ssh config for meta copy")
		return nil
	}

	config := nodes[s.nodeName]
	if config == nil {
		log.Errorf("no node %s in imported config", s.nodeName)
		return nil
	}

	copy := &SSHMeta{
		sshClient: config.GetSSHClient(),
		rawConfig: s.rawConfig,
		nodeName:  s.nodeName,
	}

	return copy
}

// ExecContext returns the results of running cmd via SSH in the specified
// context.
func (s *SSHMeta) ExecContext(ctx context.Context, cmd string) *CmdRes {
	if ctx == nil {
		panic("no context provided")
	}

	fmt.Fprintln(SSHMetaLogs, cmd)
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}

	res := CmdRes{
		cmd:     cmd,
		stdout:  stdout,
		stderr:  stderr,
		success: false,
	}

	go func() {
		s.sshClient.RunCommandContext(ctx, command)
		res.SendToLog()
	}()

	return &res
}
<|endoftext|>


<commit_before>package main

import (
	j "encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/html"

	. "github.com/WhiteBlue/bilibili-service/lib"
)

func main() {
	sr := os.Args[1]
	sn := os.Args[2]
	r, _ := strconv.Atoi(sr)
	n, _ := strconv.Atoi(sn)

	chr := make(chan int)      //任务
	csvch := make(chan string) //结果 作为csv的一行
	quitch := make(chan int)
	go Deal(chr, csvch)
	go Save(csvch, quitch, n+1)
	for i := 0; i < n; i++ {
		aid := r - i
		chr <- aid
	}

	<-quitch
}

func Deal(ch chan int, ch2 chan string) {
	client := NewBiliClient()
	var b []byte
	ch2 <- "play,comments,danmu,favorites,coins"
	for i := 0; i < 1; i++ {
		go func() {
			for {
				aid := <-ch
				p := map[string]string{}
				json, err := client.GetVideoInfo2(strconv.Itoa(aid))
				time.Sleep(1 * time.Second)
				if err != nil {
					fmt.Println(err)
					ch2 <- ""
					continue
				}
				fmt.Println(json)
				b, _ = j.Marshal(json.Get("coins"))
				p["coin"] = strings.Trim(string(b), "\"")
				b, _ = j.Marshal(json.Get("favorites"))
				p["fav"] = strings.Trim(string(b), "\"")
				b, _ = j.Marshal(json.Get("play"))
				p["play"] = strings.Trim(string(b), "\"")
				b, _ = j.Marshal(json.Get("review"))
				p["comm"] = strings.Trim(string(b), "\"")
				b, _ = j.Marshal(json.Get("video_review"))
				p["danmu"] = strings.Trim(string(b), "\"")
				//url := fmt.Sprintf("http://www.bilibili.tv/video/av%d/index.html", aid)
				//fmt.Println("get task:" + url)
				//s, _ := fetchBody(url)
				////fmt.Println(s)
				//result := parseAv(s)
				//fmt.Println(result)
				//sl, _ := getComments(url)
				//ch2 <- "save: " + strconv.Itoa(len(sl))
				fmt.Println(p)
				line := p["play"] + "," + p["comm"] + "," + p["danmu"] + "," + p["fav"] + "," + p["coin"]
				ch2 <- line
			}
		}()
	}
}

func Save(ch chan string, quitch chan int, n int) {
	f, err := os.Create("data.csv")
	check(err)
	defer f.Close()
	for {
		line := <-ch
		if line != "" {
			fmt.Println(line)
			_, err := f.Write([]byte(line + "\n"))
			check(err)
		}

		n--
		if n == 0 {
			quitch <- -1
		}
	}
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func fetchBody(url string) (string, error) {
	fmt.Println("fetch: " + url)
	resp, err := http.Get(url)
	defer resp.Body.Close()
	if err != nil {
		return "", err
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("getting %s: %s", url, resp.Status)
	}

	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return string(contents), nil
}

func parseAv(s string) map[string]string {
	var pattern string
	var reg *regexp.Regexp
	var ss [][]string

	// play,comm,coin,fav
	//pattern = `(?s)<div\sclass="v-title-info">(.*?)<div\sclass="upinfo">`
	//reg = regexp.MustCompile(pattern)
	//ss = reg.FindAllStringSubmatch(s, -1)
	//fmt.Printf("%q\n", ss[0][0])

	//pattern = `<span\sid="dianji">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//play := ss1[0][1]

	//pattern = `<span\sid="dm_count">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//comm := ss1[0][1]

	//pattern = `<span\sid="v_ctimes">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//coin := ss1[0][1]

	//pattern = `<span\sid="stow_count">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//fav := ss1[0][1]

	// comment url
	pattern = `(?s)<div\sclass="scontent"\sid="bofqi">(.*?)</div>`
	reg = regexp.MustCompile(pattern)
	ss = reg.FindAllStringSubmatch(s, -1)
	fmt.Printf("%q\n", ss[0][0])
	pattern = `[^a]id=(\d+)`
	reg = regexp.MustCompile(pattern)
	ss = reg.FindAllStringSubmatch(ss[0][0], -1)
	fmt.Printf("%q\n", ss)
	id := ss[0][1]
	commurl := fmt.Sprintf("http://comment.bilibili.tv/%s.xml", id)

	var result = map[string]string{}
	result["url"] = commurl
	//result["play"] = play
	//result["comm"] = comm
	//result["coin"] = coin
	//result["fav"] = fav
	return result
}

func getComments(url string) ([]string, error) {
	fmt.Println("fetch: " + url)
	resp, err := http.Get(url)
	defer resp.Body.Close()
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("getting %s: %s", url, resp.Status)
	}

	doc, err := html.Parse(resp.Body)
	if err != nil {
		return nil, err
	}

	var comments []string
	visitNode := func(n *html.Node) {
		fmt.Println(n)
		if n.Type == html.ElementNode && n.Data == "d" {
			fmt.Println(n.Namespace)
			comments = append(comments, n.Namespace)
		}
	}

	//fmt.Println(doc)
	forEachNode(doc, visitNode, nil)
	fmt.Println(comments)
	return comments, nil
}

// Copied from gopl.io/ch5/outline2.
func forEachNode(n *html.Node, pre, post func(n *html.Node)) {
	if pre != nil {
		pre(n)
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		forEachNode(c, pre, post)
	}
	if post != nil {
		post(n)
	}
}
<commit_msg>use new lib<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/html"

	"github.com/chrisju/bilibili-service/lib"
)

type MyBClient struct {
	*lib.BClient
}

func NewClient() *MyBClient {
	return &MyBClient{lib.NewClient()}
}

//取得视频详细信息
func (this *MyBClient) GetVideoInfo2(aid int) (*lib.JSON, error) {
	params := map[string][]string{
		"appver":     {"2310"},
		"build":      {"2310"},
		"batch":      {"1"},
		"check_area": {"1"},
		"id":         {strconv.Itoa(aid)},
		"platform":   {"ios"},
		"type":       {"json"},
	}
	json, err := this.Get("http://api.bilibili.com/view", params)
	if err != nil {
		return nil, err
	}
	return json, nil
}

func main() {
	sr := os.Args[1]
	sn := os.Args[2]
	r, _ := strconv.Atoi(sr)
	n, _ := strconv.Atoi(sn)

	chr := make(chan int)      //任务
	csvch := make(chan string) //结果 作为csv的一行
	quitch := make(chan int)
	go Deal(chr, csvch)
	go Save(csvch, quitch, n+1)
	for i := 0; i < n; i++ {
		aid := r - i
		chr <- aid
	}

	<-quitch
}
func Deal(ch chan int, ch2 chan string) {
	client := NewClient()
	var b string
	ch2 <- "play,comments,danmu,favorites,coins"
	for i := 0; i < 1; i++ {
		go func() {
			for {
				aid := <-ch
				p := map[string]string{}
				json, err := client.GetVideoInfo2(aid)
				time.Sleep(1 * time.Second)
				if err != nil {
					fmt.Println(err)
					ch2 <- ""
					continue
				}
				fmt.Println(json)
				b, _ = json.Get("coins").String()
				p["coin"] = strings.Trim(string(b), "\"")
				b, _ = json.Get("favorites").String()
				p["fav"] = strings.Trim(string(b), "\"")
				b, _ = json.Get("play").String()
				p["play"] = strings.Trim(string(b), "\"")
				b, _ = json.Get("review").String()
				p["comm"] = strings.Trim(string(b), "\"")
				b, _ = json.Get("video_review").String()
				p["danmu"] = strings.Trim(string(b), "\"")
				//url := fmt.Sprintf("http://www.bilibili.tv/video/av%d/index.html", aid)
				//fmt.Println("get task:" + url)
				//s, _ := fetchBody(url)
				////fmt.Println(s)
				//result := parseAv(s)
				//fmt.Println(result)
				//sl, _ := getComments(url)
				//ch2 <- "save: " + strconv.Itoa(len(sl))
				fmt.Println(p)
				line := p["play"] + "," + p["comm"] + "," + p["danmu"] + "," + p["fav"] + "," + p["coin"]
				ch2 <- line
			}
		}()
	}
}

func Save(ch chan string, quitch chan int, n int) {
	f, err := os.Create("data.csv")
	check(err)
	defer f.Close()
	for {
		line := <-ch
		if line != "" {
			fmt.Println(line)
			_, err := f.Write([]byte(line + "\n"))
			check(err)
		}

		n--
		if n == 0 {
			quitch <- -1
		}
	}
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func fetchBody(url string) (string, error) {
	fmt.Println("fetch: " + url)
	resp, err := http.Get(url)
	defer resp.Body.Close()
	if err != nil {
		return "", err
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("getting %s: %s", url, resp.Status)
	}

	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return string(contents), nil
}

func parseAv(s string) map[string]string {
	var pattern string
	var reg *regexp.Regexp
	var ss [][]string

	// play,comm,coin,fav
	//pattern = `(?s)<div\sclass="v-title-info">(.*?)<div\sclass="upinfo">`
	//reg = regexp.MustCompile(pattern)
	//ss = reg.FindAllStringSubmatch(s, -1)
	//fmt.Printf("%q\n", ss[0][0])

	//pattern = `<span\sid="dianji">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//play := ss1[0][1]

	//pattern = `<span\sid="dm_count">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//comm := ss1[0][1]

	//pattern = `<span\sid="v_ctimes">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//coin := ss1[0][1]

	//pattern = `<span\sid="stow_count">(\d*)</span>`
	//reg = regexp.MustCompile(pattern)
	//ss1 = reg.FindAllStringSubmatch(ss[0][0], -1)
	//fmt.Printf("%q\n", ss1)
	//fav := ss1[0][1]

	// comment url
	pattern = `(?s)<div\sclass="scontent"\sid="bofqi">(.*?)</div>`
	reg = regexp.MustCompile(pattern)
	ss = reg.FindAllStringSubmatch(s, -1)
	fmt.Printf("%q\n", ss[0][0])
	pattern = `[^a]id=(\d+)`
	reg = regexp.MustCompile(pattern)
	ss = reg.FindAllStringSubmatch(ss[0][0], -1)
	fmt.Printf("%q\n", ss)
	id := ss[0][1]
	commurl := fmt.Sprintf("http://comment.bilibili.tv/%s.xml", id)

	var result = map[string]string{}
	result["url"] = commurl
	//result["play"] = play
	//result["comm"] = comm
	//result["coin"] = coin
	//result["fav"] = fav
	return result
}

func getComments(url string) ([]string, error) {
	fmt.Println("fetch: " + url)
	resp, err := http.Get(url)
	defer resp.Body.Close()
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("getting %s: %s", url, resp.Status)
	}

	doc, err := html.Parse(resp.Body)
	if err != nil {
		return nil, err
	}

	var comments []string
	visitNode := func(n *html.Node) {
		fmt.Println(n)
		if n.Type == html.ElementNode && n.Data == "d" {
			fmt.Println(n.Namespace)
			comments = append(comments, n.Namespace)
		}
	}

	//fmt.Println(doc)
	forEachNode(doc, visitNode, nil)
	fmt.Println(comments)
	return comments, nil
}

// Copied from gopl.io/ch5/outline2.
func forEachNode(n *html.Node, pre, post func(n *html.Node)) {
	if pre != nil {
		pre(n)
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		forEachNode(c, pre, post)
	}
	if post != nil {
		post(n)
	}
}
<|endoftext|>


<commit_before>package azuresink

import (
	"bytes"
	"context"
	"fmt"
	"net/url"

	"github.com/Azure/azure-storage-blob-go/2016-05-31/azblob"
	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/replication/sink"
	"github.com/chrislusf/seaweedfs/weed/replication/source"
	"github.com/chrislusf/seaweedfs/weed/util"
)

type AzureSink struct {
	containerURL azblob.ContainerURL
	container    string
	dir          string
	filerSource  *source.FilerSource
}

func init() {
	sink.Sinks = append(sink.Sinks, &AzureSink{})
}

func (g *AzureSink) GetName() string {
	return "azure"
}

func (g *AzureSink) GetSinkToDirectory() string {
	return g.dir
}

func (g *AzureSink) Initialize(configuration util.Configuration) error {
	return g.initialize(
		configuration.GetString("account_name"),
		configuration.GetString("account_key"),
		configuration.GetString("container"),
		configuration.GetString("directory"),
	)
}

func (g *AzureSink) SetSourceFiler(s *source.FilerSource) {
	g.filerSource = s
}

func (g *AzureSink) initialize(accountName, accountKey, container, dir string) error {
	g.container = container
	g.dir = dir

	// Use your Storage account's name and key to create a credential object.
	credential := azblob.NewSharedKeyCredential(accountName, accountKey)

	// Create a request pipeline that is used to process HTTP(S) requests and responses.
	p := azblob.NewPipeline(credential, azblob.PipelineOptions{})

	// Create an ServiceURL object that wraps the service URL and a request pipeline.
	u, _ := url.Parse(fmt.Sprintf("https://%s.blob.core.windows.net", accountName))
	serviceURL := azblob.NewServiceURL(*u, p)

	g.containerURL = serviceURL.NewContainerURL(g.container)

	return nil
}

func (g *AzureSink) DeleteEntry(key string, isDirectory, deleteIncludeChunks bool) error {

	if isDirectory {
		key = key + "/"
	}

	ctx := context.Background()

	if _, err := g.containerURL.NewBlobURL(key).Delete(ctx,
		azblob.DeleteSnapshotsOptionInclude, azblob.BlobAccessConditions{}); err != nil {
		return fmt.Errorf("azure delete %s/%s: %v", g.container, key, err)
	}

	return nil

}

func (g *AzureSink) CreateEntry(key string, entry *filer_pb.Entry) error {

	if entry.IsDirectory {
		return nil
	}

	totalSize := filer2.TotalSize(entry.Chunks)
	chunkViews := filer2.ViewFromChunks(entry.Chunks, 0, int(totalSize))

	ctx := context.Background()

	// Create a URL that references a to-be-created blob in your
	// Azure Storage account's container.
	appendBlobURL := g.containerURL.NewAppendBlobURL(key)

	_, err := appendBlobURL.Create(ctx, azblob.BlobHTTPHeaders{}, azblob.Metadata{}, azblob.BlobAccessConditions{})
	if err != nil {
		return err
	}

	for _, chunk := range chunkViews {

		fileUrl, err := g.filerSource.LookupFileId(chunk.FileId)
		if err != nil {
			return err
		}

		var writeErr error
		_, readErr := util.ReadUrlAsStream(fileUrl, chunk.Offset, int(chunk.Size), func(data []byte) {
			_, writeErr = appendBlobURL.AppendBlock(ctx, bytes.NewReader(data), azblob.BlobAccessConditions{})
		})

		if readErr != nil {
			return readErr
		}
		if writeErr != nil {
			return writeErr
		}

	}

	return nil

}

func (g *AzureSink) UpdateEntry(key string, oldEntry, newEntry *filer_pb.Entry, deleteIncludeChunks bool) (foundExistingEntry bool, err error) {
	// TODO improve efficiency
	return false, nil
}
<commit_msg>use released version of azblob<commit_after>package azuresink

import (
	"bytes"
	"context"
	"fmt"
	"net/url"

	"github.com/Azure/azure-storage-blob-go/azblob"
	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/replication/sink"
	"github.com/chrislusf/seaweedfs/weed/replication/source"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/glog"
)

type AzureSink struct {
	containerURL azblob.ContainerURL
	container    string
	dir          string
	filerSource  *source.FilerSource
}

func init() {
	sink.Sinks = append(sink.Sinks, &AzureSink{})
}

func (g *AzureSink) GetName() string {
	return "azure"
}

func (g *AzureSink) GetSinkToDirectory() string {
	return g.dir
}

func (g *AzureSink) Initialize(configuration util.Configuration) error {
	return g.initialize(
		configuration.GetString("account_name"),
		configuration.GetString("account_key"),
		configuration.GetString("container"),
		configuration.GetString("directory"),
	)
}

func (g *AzureSink) SetSourceFiler(s *source.FilerSource) {
	g.filerSource = s
}

func (g *AzureSink) initialize(accountName, accountKey, container, dir string) error {
	g.container = container
	g.dir = dir

	// Use your Storage account's name and key to create a credential object.
	credential, err := azblob.NewSharedKeyCredential(accountName, accountKey)
	if err != nil {
		glog.Fatalf("failed to create Azure credential with account name:%s key:%s", accountName, accountKey)
	}

	// Create a request pipeline that is used to process HTTP(S) requests and responses.
	p := azblob.NewPipeline(credential, azblob.PipelineOptions{})

	// Create an ServiceURL object that wraps the service URL and a request pipeline.
	u, _ := url.Parse(fmt.Sprintf("https://%s.blob.core.windows.net", accountName))
	serviceURL := azblob.NewServiceURL(*u, p)

	g.containerURL = serviceURL.NewContainerURL(g.container)

	return nil
}

func (g *AzureSink) DeleteEntry(key string, isDirectory, deleteIncludeChunks bool) error {

	if isDirectory {
		key = key + "/"
	}

	ctx := context.Background()

	if _, err := g.containerURL.NewBlobURL(key).Delete(ctx,
		azblob.DeleteSnapshotsOptionInclude, azblob.BlobAccessConditions{}); err != nil {
		return fmt.Errorf("azure delete %s/%s: %v", g.container, key, err)
	}

	return nil

}

func (g *AzureSink) CreateEntry(key string, entry *filer_pb.Entry) error {

	if entry.IsDirectory {
		return nil
	}

	totalSize := filer2.TotalSize(entry.Chunks)
	chunkViews := filer2.ViewFromChunks(entry.Chunks, 0, int(totalSize))

	ctx := context.Background()

	// Create a URL that references a to-be-created blob in your
	// Azure Storage account's container.
	appendBlobURL := g.containerURL.NewAppendBlobURL(key)

	_, err := appendBlobURL.Create(ctx, azblob.BlobHTTPHeaders{}, azblob.Metadata{}, azblob.BlobAccessConditions{})
	if err != nil {
		return err
	}

	for _, chunk := range chunkViews {

		fileUrl, err := g.filerSource.LookupFileId(chunk.FileId)
		if err != nil {
			return err
		}

		var writeErr error
		_, readErr := util.ReadUrlAsStream(fileUrl, chunk.Offset, int(chunk.Size), func(data []byte) {
			_, writeErr = appendBlobURL.AppendBlock(ctx, bytes.NewReader(data), azblob.AppendBlobAccessConditions{}, nil)
		})

		if readErr != nil {
			return readErr
		}
		if writeErr != nil {
			return writeErr
		}

	}

	return nil

}

func (g *AzureSink) UpdateEntry(key string, oldEntry, newEntry *filer_pb.Entry, deleteIncludeChunks bool) (foundExistingEntry bool, err error) {
	// TODO improve efficiency
	return false, nil
}
<|endoftext|>


<commit_before>package ping

import (
	"bytes"
	"time"

	"github.com/hsheth2/gonet/icmp"

	"github.com/hsheth2/gonet/ipv4"

	"github.com/hsheth2/logs"
	"sync/atomic"
)

const DATA_56_BYTES = "abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcd"

func (pm *Ping_Manager) ping_response_dealer() {
	for {
		ping := <-pm.reply
		identNum := uint16(ping.Header.Opt >> 16)
		if _, ok := pm.identifiers[identNum]; !ok {
			/*logs*/ logs.Info.Println("Dropped something from response dealer, identnum=", identNum, "options=", ping.Header.Opt)
			continue
		}
		pm.identifiers[identNum] <- ping
	}
}

func sendSinglePing(writer ipv4.Writer, id uint16, seq uint16, timeout time.Duration, reply chan *icmp.Packet) {
	// prepare packet
	packet := &icmp.Header{
		Tp:   icmp.EchoRequest,
		Code: PING_ICMP_CODE,
		Opt:  uint32(id)<<16 | uint32(seq),
		Data: []byte(DATA_56_BYTES), // TODO make legit by putting the timestamp in the data
	}

	// make data
	icmp.SendPacket(writer, packet)

	time1 := time.Now()
	timeoutTimer := time.NewTimer(timeout)
	go func(seqChan chan *icmp.Packet, header *icmp.Header, time1 *time.Time, timer *time.Timer) {
		defer close(seqChan)
		for {
			select {
			case pingResponse := <-seqChan:
				if !bytes.Equal(pingResponse.Header.Data, header.Data) {
					/*logs*/ logs.Info.Println("Dropped packet because header data not equal to ping sent")
					continue
				}
				time2 := time.Now()
				logs.Info.Printf("%d bytes from %v: icmp_seq=%d time=%f ms",
					len(header.Data)+icmp.HeaderMinSize,
					pingResponse.RIP.IP,
					uint16(header.Opt),
					float32(time2.Sub(*time1).Nanoseconds())/1000000) // put ttl
				return
			case <-timer.C:
				logs.Info.Println("Seq num of", uint16(header.Opt), "timed out")
				return
			}
		}
	}(reply, packet, &time1, timeoutTimer)
}

func (pm *Ping_Manager) initIdentifier(terminate chan bool, numPings uint16, pingCount *uint32) (uint16, [](chan *icmp.Packet), error) {
	// get identifier
	pm.currentIdentifier++
	id := pm.currentIdentifier

	// setup sequence number dealer
	pm.identifiers[id] = make(chan *icmp.Packet)
	seqChannel := make([](chan *icmp.Packet), numPings)

	// create go routine function to deal packets
	go sequenceDealer(pm.identifiers[id], seqChannel, terminate, pingCount)

	return id, seqChannel, nil
}

func sequenceDealer(idInput chan *icmp.Packet, seqChan [](chan *icmp.Packet), terminate chan bool, pingCount *uint32) {
	// TODO verify IPs
	for {
		select {
		case <-terminate:
		//			/*logs*/logs.Info.Println("Terminating seq dealer")
			return
		case packet := <-idInput:
		// /*logs*/logs.Info.Println("icmp in =", packet.Header.Opt)
			seqNum := packet.Header.Opt << 16 >> 16
			pingC := atomic.LoadUint32(pingCount)
			if seqNum < pingC-uint32(1) {
				seqChan[seqNum-1] <- packet
			} else {
				/*logs*/ logs.Info.Println("Dropping bad seq num packet with existing identifier, seqNum-id:", seqNum, "pingcount:", pingC)
			}
		}
	}
}

const FLOOD_INTERVAL = 0 // TODO add a flood option

func (pm *Ping_Manager) SendPing(ip *ipv4.Address, interval, timeout time.Duration, numPings uint16) error {
	terminate := make(chan bool)
	pingCount := uint32(1)
	id, seqChannel, err := pm.initIdentifier(terminate, numPings, &pingCount)
	if err != nil {
		logs.Error.Println(err)
		return err
	}
	defer func() { terminate <- true }()

	// get ip writer
	writer, err := pm.getIP_Writer(ip)
	if err != nil {
		logs.Error.Println(err)
		return err
	}

	for pingCount <= uint32(numPings) {
		seqChannel[pingCount-1] = make(chan *icmp.Packet)

		sendSinglePing(writer, id, uint16(pingCount), timeout, seqChannel[pingCount-1]) // function is non-blocking

		atomic.AddUint32(&pingCount, 1)
		// not last
		if uint16(pingCount) != numPings {
			time.Sleep(interval)
		}
	}

	time.Sleep(timeout)
	return nil
}
<commit_msg>Incremented ping count just before sending<commit_after>package ping

import (
	"bytes"
	"time"

	"github.com/hsheth2/gonet/icmp"

	"github.com/hsheth2/gonet/ipv4"

	"github.com/hsheth2/logs"
	"sync/atomic"
)

const DATA_56_BYTES = "abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcd"

func (pm *Ping_Manager) ping_response_dealer() {
	for {
		ping := <-pm.reply
		identNum := uint16(ping.Header.Opt >> 16)
		if _, ok := pm.identifiers[identNum]; !ok {
			/*logs*/ logs.Info.Println("Dropped something from response dealer, identnum=", identNum, "options=", ping.Header.Opt)
			continue
		}
		pm.identifiers[identNum] <- ping
	}
}

func sendSinglePing(writer ipv4.Writer, id uint16, seq uint16, timeout time.Duration, reply chan *icmp.Packet) {
	// prepare packet
	packet := &icmp.Header{
		Tp:   icmp.EchoRequest,
		Code: PING_ICMP_CODE,
		Opt:  uint32(id)<<16 | uint32(seq),
		Data: []byte(DATA_56_BYTES), // TODO make legit by putting the timestamp in the data
	}

	// make data
	icmp.SendPacket(writer, packet)

	time1 := time.Now()
	timeoutTimer := time.NewTimer(timeout)
	go func(seqChan chan *icmp.Packet, header *icmp.Header, time1 *time.Time, timer *time.Timer) {
		defer close(seqChan)
		for {
			select {
			case pingResponse := <-seqChan:
				if !bytes.Equal(pingResponse.Header.Data, header.Data) {
					/*logs*/ logs.Info.Println("Dropped packet because header data not equal to ping sent")
					continue
				}
				time2 := time.Now()
				logs.Info.Printf("%d bytes from %v: icmp_seq=%d time=%f ms",
					len(header.Data)+icmp.HeaderMinSize,
					pingResponse.RIP.IP,
					uint16(header.Opt),
					float32(time2.Sub(*time1).Nanoseconds())/1000000) // put ttl
				return
			case <-timer.C:
				logs.Info.Println("Seq num of", uint16(header.Opt), "timed out")
				return
			}
		}
	}(reply, packet, &time1, timeoutTimer)
}

func (pm *Ping_Manager) initIdentifier(terminate chan bool, numPings uint16, pingCount *uint32) (uint16, [](chan *icmp.Packet), error) {
	// get identifier
	pm.currentIdentifier++
	id := pm.currentIdentifier

	// setup sequence number dealer
	pm.identifiers[id] = make(chan *icmp.Packet)
	seqChannel := make([](chan *icmp.Packet), numPings)

	// create go routine function to deal packets
	go sequenceDealer(pm.identifiers[id], seqChannel, terminate, pingCount)

	return id, seqChannel, nil
}

func sequenceDealer(idInput chan *icmp.Packet, seqChan [](chan *icmp.Packet), terminate chan bool, pingCount *uint32) {
	// TODO verify IPs
	for {
		select {
		case <-terminate:
		//			/*logs*/logs.Info.Println("Terminating seq dealer")
			return
		case packet := <-idInput:
		// /*logs*/logs.Info.Println("icmp in =", packet.Header.Opt)
			seqNum := packet.Header.Opt << 16 >> 16
			pingC := atomic.LoadUint32(pingCount)
			if seqNum < pingC {
				seqChan[seqNum-1] <- packet
			} else {
				/*logs*/ logs.Info.Println("Dropping bad seq num packet with existing identifier, seqNum-id:", seqNum, "pingcount:", pingC)
			}
		}
	}
}

const FLOOD_INTERVAL = 0 // TODO add a flood option

func (pm *Ping_Manager) SendPing(ip *ipv4.Address, interval, timeout time.Duration, numPings uint16) error {
	terminate := make(chan bool)
	pingCount := uint32(0)
	id, seqChannel, err := pm.initIdentifier(terminate, numPings, &pingCount)
	if err != nil {
		logs.Error.Println(err)
		return err
	}
	defer func() { terminate <- true }()

	// get ip writer
	writer, err := pm.getIP_Writer(ip)
	if err != nil {
		logs.Error.Println(err)
		return err
	}

	for pingCount <= uint32(numPings) {
		atomic.AddUint32(&pingCount, 1)
		seqChannel[pingCount-1] = make(chan *icmp.Packet)

		sendSinglePing(writer, id, uint16(pingCount), timeout, seqChannel[pingCount-1]) // function is non-blocking

		// not last
		if uint16(pingCount) != numPings {
			time.Sleep(interval)
		}
	}

	time.Sleep(timeout)
	return nil
}
<|endoftext|>