<commit_before>// Package drivers talks to various database backends and retrieves table,
// column, type, and foreign key information
package drivers

import (
	"sort"

	"github.com/friendsofgo/errors"
	"github.com/volatiletech/sqlboiler/v4/importers"
	"github.com/volatiletech/strmangle"
)

// These constants are used in the config map passed into the driver
const (
	ConfigBlacklist      = "blacklist"
	ConfigWhitelist      = "whitelist"
	ConfigSchema         = "schema"
	ConfigAddEnumTypes   = "add-enum-types"
	ConfigEnumNullPrefix = "enum-null-prefix"

	ConfigUser    = "user"
	ConfigPass    = "pass"
	ConfigHost    = "host"
	ConfigPort    = "port"
	ConfigDBName  = "dbname"
	ConfigSSLMode = "sslmode"
)

// Interface abstracts either a side-effect imported driver or a binary
// that is called in order to produce the data required for generation.
type Interface interface {
	// Assemble the database information into a nice struct
	Assemble(config Config) (*DBInfo, error)
	// Templates to add/replace for generation
	Templates() (map[string]string, error)
	// Imports to merge for generation
	Imports() (importers.Collection, error)
}

// DBInfo is the database's table data and dialect.
type DBInfo struct {
	Schema  string  `json:"schema"`
	Tables  []Table `json:"tables"`
	Dialect Dialect `json:"dialect"`
}

// Dialect describes the databases requirements in terms of which features
// it speaks and what kind of quoting mechanisms it uses.
//
// WARNING: When updating this struct there is a copy of it inside
// the boil_queries template that is used for users to create queries
// without having to figure out what their dialect is.
type Dialect struct {
	LQ rune `json:"lq"`
	RQ rune `json:"rq"`

	UseIndexPlaceholders bool `json:"use_index_placeholders"`
	UseLastInsertID      bool `json:"use_last_insert_id"`
	UseSchema            bool `json:"use_schema"`
	UseDefaultKeyword    bool `json:"use_default_keyword"`

	// The following is mostly for T-SQL/MSSQL, what a show
	UseTopClause            bool `json:"use_top_clause"`
	UseOutputClause         bool `json:"use_output_clause"`
	UseCaseWhenExistsClause bool `json:"use_case_when_exists_clause"`

	// No longer used, left for backwards compatibility
	// should be removed in v5
	UseAutoColumns bool `json:"use_auto_columns"`
}

// Constructor breaks down the functionality required to implement a driver
// such that the drivers.Tables method can be used to reduce duplication in driver
// implementations.
type Constructor interface {
	TableNames(schema string, whitelist, blacklist []string) ([]string, error)
	Columns(schema, tableName string, whitelist, blacklist []string) ([]Column, error)
	PrimaryKeyInfo(schema, tableName string) (*PrimaryKey, error)
	ForeignKeyInfo(schema, tableName string) ([]ForeignKey, error)

	// TranslateColumnType takes a Database column type and returns a go column type.
	TranslateColumnType(Column) Column
}

// Constructor breaks down the functionality required to implement a driver
// such that the drivers.Views method can be used to reduce duplication in driver
// implementations.
type ViewConstructor interface {
	ViewNames(schema string, whitelist, blacklist []string) ([]string, error)
	ViewCapabilities(schema, viewName string) (ViewCapabilities, error)
	ViewColumns(schema, tableName string, whitelist, blacklist []string) ([]Column, error)

	// TranslateColumnType takes a Database column type and returns a go column type.
	TranslateColumnType(Column) Column
}

type TableColumnTypeTranslator interface {
	// TranslateTableColumnType takes a Database column type and table name and returns a go column type.
	TranslateTableColumnType(c Column, tableName string) Column
}

// Tables returns the metadata for all tables, minus the tables
// specified in the blacklist.
func Tables(c Constructor, schema string, whitelist, blacklist []string) ([]Table, error) {
	var err error

	names, err := c.TableNames(schema, whitelist, blacklist)
	if err != nil {
		return nil, errors.Wrap(err, "unable to get table names")
	}

	sort.Strings(names)

	var tables []Table
	for _, name := range names {
		t := Table{
			Name: name,
		}

		if t.Columns, err = c.Columns(schema, name, whitelist, blacklist); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch table column info (%s)", name)
		}

		tr, ok := c.(TableColumnTypeTranslator)
		if ok {
			for i, col := range t.Columns {
				t.Columns[i] = tr.TranslateTableColumnType(col, name)
			}
		} else {
			for i, col := range t.Columns {
				t.Columns[i] = c.TranslateColumnType(col)
			}
		}

		if t.PKey, err = c.PrimaryKeyInfo(schema, name); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch table pkey info (%s)", name)
		}

		if t.FKeys, err = c.ForeignKeyInfo(schema, name); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch table fkey info (%s)", name)
		}

		filterForeignKeys(&t, whitelist, blacklist)

		setIsJoinTable(&t)

		tables = append(tables, t)
	}

	// Relationships have a dependency on foreign key nullability.
	for i := range tables {
		tbl := &tables[i]
		setForeignKeyConstraints(tbl, tables)
	}
	for i := range tables {
		tbl := &tables[i]
		setRelationships(tbl, tables)
	}

	if vc, ok := c.(ViewConstructor); ok {
		viewTables, err := views(vc, schema, whitelist, blacklist)
		if err != nil {
			return nil, errors.Wrap(err, "unable to load views")
		}

		tables = append(tables, viewTables...)
	}

	return tables, nil
}

// views returns the metadata for all views, minus the views
// specified in the blacklist.
func views(c ViewConstructor, schema string, whitelist, blacklist []string) ([]Table, error) {
	var err error

	names, err := c.ViewNames(schema, whitelist, blacklist)
	if err != nil {
		return nil, errors.Wrap(err, "unable to get view names")
	}

	sort.Strings(names)

	var views []Table
	for _, name := range names {
		t := Table{
			IsView: true,
			Name:   name,
		}

		if t.ViewCapabilities, err = c.ViewCapabilities(schema, name); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch view capabilities info (%s)", name)
		}

		if t.Columns, err = c.ViewColumns(schema, name, whitelist, blacklist); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch view column info (%s)", name)
		}

		tr, ok := c.(TableColumnTypeTranslator)
		if ok {
			for i, col := range t.Columns {
				t.Columns[i] = tr.TranslateTableColumnType(col, name)
			}
		} else {
			for i, col := range t.Columns {
				t.Columns[i] = c.TranslateColumnType(col)
			}
		}

		views = append(views, t)
	}

	return views, nil
}

func knownColumn(table string, column string, whitelist, blacklist []string) bool {
	return (len(whitelist) == 0 ||
		strmangle.SetInclude(table, whitelist) ||
		strmangle.SetInclude(table+"."+column, whitelist) ||
		strmangle.SetInclude("*."+column, whitelist)) &&

		(len(blacklist) == 0 || (!strmangle.SetInclude(table, blacklist) &&
			!strmangle.SetInclude(table+"."+column, blacklist) &&
			!strmangle.SetInclude("*."+column, blacklist)))
}

// filterForeignKeys filter FK whose ForeignTable is not in whitelist or in blacklist
func filterForeignKeys(t *Table, whitelist, blacklist []string) {
	var fkeys []ForeignKey

	for _, fkey := range t.FKeys {
		if knownColumn(fkey.ForeignTable, fkey.ForeignColumn, whitelist, blacklist) &&
			knownColumn(fkey.Table, fkey.Column, whitelist, blacklist) {
			fkeys = append(fkeys, fkey)
		}
	}
	t.FKeys = fkeys
}

// setIsJoinTable if there are:
// A composite primary key involving two columns
// Both primary key columns are also foreign keys
func setIsJoinTable(t *Table) {
	if t.PKey == nil || len(t.PKey.Columns) != 2 || len(t.FKeys) < 2 || len(t.Columns) > 2 {
		return
	}

	for _, c := range t.PKey.Columns {
		found := false
		for _, f := range t.FKeys {
			if c == f.Column {
				found = true
				break
			}
		}
		if !found {
			return
		}
	}

	t.IsJoinTable = true
}

func setForeignKeyConstraints(t *Table, tables []Table) {
	for i, fkey := range t.FKeys {
		localColumn := t.GetColumn(fkey.Column)
		foreignTable := GetTable(tables, fkey.ForeignTable)
		foreignColumn := foreignTable.GetColumn(fkey.ForeignColumn)

		t.FKeys[i].Nullable = localColumn.Nullable
		t.FKeys[i].Unique = localColumn.Unique
		t.FKeys[i].ForeignColumnNullable = foreignColumn.Nullable
		t.FKeys[i].ForeignColumnUnique = foreignColumn.Unique
	}
}

func setRelationships(t *Table, tables []Table) {
	t.ToOneRelationships = toOneRelationships(*t, tables)
	t.ToManyRelationships = toManyRelationships(*t, tables)
}
<commit_msg>Fix column ignore for primary key<commit_after>// Package drivers talks to various database backends and retrieves table,
// column, type, and foreign key information
package drivers

import (
	"sort"

	"github.com/friendsofgo/errors"
	"github.com/volatiletech/sqlboiler/v4/importers"
	"github.com/volatiletech/strmangle"
)

// These constants are used in the config map passed into the driver
const (
	ConfigBlacklist      = "blacklist"
	ConfigWhitelist      = "whitelist"
	ConfigSchema         = "schema"
	ConfigAddEnumTypes   = "add-enum-types"
	ConfigEnumNullPrefix = "enum-null-prefix"

	ConfigUser    = "user"
	ConfigPass    = "pass"
	ConfigHost    = "host"
	ConfigPort    = "port"
	ConfigDBName  = "dbname"
	ConfigSSLMode = "sslmode"
)

// Interface abstracts either a side-effect imported driver or a binary
// that is called in order to produce the data required for generation.
type Interface interface {
	// Assemble the database information into a nice struct
	Assemble(config Config) (*DBInfo, error)
	// Templates to add/replace for generation
	Templates() (map[string]string, error)
	// Imports to merge for generation
	Imports() (importers.Collection, error)
}

// DBInfo is the database's table data and dialect.
type DBInfo struct {
	Schema  string  `json:"schema"`
	Tables  []Table `json:"tables"`
	Dialect Dialect `json:"dialect"`
}

// Dialect describes the databases requirements in terms of which features
// it speaks and what kind of quoting mechanisms it uses.
//
// WARNING: When updating this struct there is a copy of it inside
// the boil_queries template that is used for users to create queries
// without having to figure out what their dialect is.
type Dialect struct {
	LQ rune `json:"lq"`
	RQ rune `json:"rq"`

	UseIndexPlaceholders bool `json:"use_index_placeholders"`
	UseLastInsertID      bool `json:"use_last_insert_id"`
	UseSchema            bool `json:"use_schema"`
	UseDefaultKeyword    bool `json:"use_default_keyword"`

	// The following is mostly for T-SQL/MSSQL, what a show
	UseTopClause            bool `json:"use_top_clause"`
	UseOutputClause         bool `json:"use_output_clause"`
	UseCaseWhenExistsClause bool `json:"use_case_when_exists_clause"`

	// No longer used, left for backwards compatibility
	// should be removed in v5
	UseAutoColumns bool `json:"use_auto_columns"`
}

// Constructor breaks down the functionality required to implement a driver
// such that the drivers.Tables method can be used to reduce duplication in driver
// implementations.
type Constructor interface {
	TableNames(schema string, whitelist, blacklist []string) ([]string, error)
	Columns(schema, tableName string, whitelist, blacklist []string) ([]Column, error)
	PrimaryKeyInfo(schema, tableName string) (*PrimaryKey, error)
	ForeignKeyInfo(schema, tableName string) ([]ForeignKey, error)

	// TranslateColumnType takes a Database column type and returns a go column type.
	TranslateColumnType(Column) Column
}

// Constructor breaks down the functionality required to implement a driver
// such that the drivers.Views method can be used to reduce duplication in driver
// implementations.
type ViewConstructor interface {
	ViewNames(schema string, whitelist, blacklist []string) ([]string, error)
	ViewCapabilities(schema, viewName string) (ViewCapabilities, error)
	ViewColumns(schema, tableName string, whitelist, blacklist []string) ([]Column, error)

	// TranslateColumnType takes a Database column type and returns a go column type.
	TranslateColumnType(Column) Column
}

type TableColumnTypeTranslator interface {
	// TranslateTableColumnType takes a Database column type and table name and returns a go column type.
	TranslateTableColumnType(c Column, tableName string) Column
}

// Tables returns the metadata for all tables, minus the tables
// specified in the blacklist.
func Tables(c Constructor, schema string, whitelist, blacklist []string) ([]Table, error) {
	var err error

	names, err := c.TableNames(schema, whitelist, blacklist)
	if err != nil {
		return nil, errors.Wrap(err, "unable to get table names")
	}

	sort.Strings(names)

	var tables []Table
	for _, name := range names {
		t := Table{
			Name: name,
		}

		if t.Columns, err = c.Columns(schema, name, whitelist, blacklist); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch table column info (%s)", name)
		}

		tr, ok := c.(TableColumnTypeTranslator)
		if ok {
			for i, col := range t.Columns {
				t.Columns[i] = tr.TranslateTableColumnType(col, name)
			}
		} else {
			for i, col := range t.Columns {
				t.Columns[i] = c.TranslateColumnType(col)
			}
		}

		if t.PKey, err = c.PrimaryKeyInfo(schema, name); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch table pkey info (%s)", name)
		}

		if t.FKeys, err = c.ForeignKeyInfo(schema, name); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch table fkey info (%s)", name)
		}

		filterPrimaryKey(&t, whitelist, blacklist)
		filterForeignKeys(&t, whitelist, blacklist)

		setIsJoinTable(&t)

		tables = append(tables, t)
	}

	// Relationships have a dependency on foreign key nullability.
	for i := range tables {
		tbl := &tables[i]
		setForeignKeyConstraints(tbl, tables)
	}
	for i := range tables {
		tbl := &tables[i]
		setRelationships(tbl, tables)
	}

	if vc, ok := c.(ViewConstructor); ok {
		viewTables, err := views(vc, schema, whitelist, blacklist)
		if err != nil {
			return nil, errors.Wrap(err, "unable to load views")
		}

		tables = append(tables, viewTables...)
	}

	return tables, nil
}

// views returns the metadata for all views, minus the views
// specified in the blacklist.
func views(c ViewConstructor, schema string, whitelist, blacklist []string) ([]Table, error) {
	var err error

	names, err := c.ViewNames(schema, whitelist, blacklist)
	if err != nil {
		return nil, errors.Wrap(err, "unable to get view names")
	}

	sort.Strings(names)

	var views []Table
	for _, name := range names {
		t := Table{
			IsView: true,
			Name:   name,
		}

		if t.ViewCapabilities, err = c.ViewCapabilities(schema, name); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch view capabilities info (%s)", name)
		}

		if t.Columns, err = c.ViewColumns(schema, name, whitelist, blacklist); err != nil {
			return nil, errors.Wrapf(err, "unable to fetch view column info (%s)", name)
		}

		tr, ok := c.(TableColumnTypeTranslator)
		if ok {
			for i, col := range t.Columns {
				t.Columns[i] = tr.TranslateTableColumnType(col, name)
			}
		} else {
			for i, col := range t.Columns {
				t.Columns[i] = c.TranslateColumnType(col)
			}
		}

		views = append(views, t)
	}

	return views, nil
}

func knownColumn(table string, column string, whitelist, blacklist []string) bool {
	return (len(whitelist) == 0 ||
		strmangle.SetInclude(table, whitelist) ||
		strmangle.SetInclude(table+"."+column, whitelist) ||
		strmangle.SetInclude("*."+column, whitelist)) &&

		(len(blacklist) == 0 || (!strmangle.SetInclude(table, blacklist) &&
			!strmangle.SetInclude(table+"."+column, blacklist) &&
			!strmangle.SetInclude("*."+column, blacklist)))
}

// filterPrimaryKey filter columns from the primary key that are not in whitelist or in blacklist
func filterPrimaryKey(t *Table, whitelist, blacklist []string) {
	pkeyColumns := make([]string, 0, len(t.PKey.Columns))
	for _, c := range t.PKey.Columns {
		if knownColumn(t.Name, c, whitelist, blacklist) {
			pkeyColumns = append(pkeyColumns, c)
		}
	}
	t.PKey.Columns = pkeyColumns
}

// filterForeignKeys filter FK whose ForeignTable is not in whitelist or in blacklist
func filterForeignKeys(t *Table, whitelist, blacklist []string) {
	var fkeys []ForeignKey

	for _, fkey := range t.FKeys {
		if knownColumn(fkey.ForeignTable, fkey.ForeignColumn, whitelist, blacklist) &&
			knownColumn(fkey.Table, fkey.Column, whitelist, blacklist) {
			fkeys = append(fkeys, fkey)
		}
	}
	t.FKeys = fkeys
}

// setIsJoinTable if there are:
// A composite primary key involving two columns
// Both primary key columns are also foreign keys
func setIsJoinTable(t *Table) {
	if t.PKey == nil || len(t.PKey.Columns) != 2 || len(t.FKeys) < 2 || len(t.Columns) > 2 {
		return
	}

	for _, c := range t.PKey.Columns {
		found := false
		for _, f := range t.FKeys {
			if c == f.Column {
				found = true
				break
			}
		}
		if !found {
			return
		}
	}

	t.IsJoinTable = true
}

func setForeignKeyConstraints(t *Table, tables []Table) {
	for i, fkey := range t.FKeys {
		localColumn := t.GetColumn(fkey.Column)
		foreignTable := GetTable(tables, fkey.ForeignTable)
		foreignColumn := foreignTable.GetColumn(fkey.ForeignColumn)

		t.FKeys[i].Nullable = localColumn.Nullable
		t.FKeys[i].Unique = localColumn.Unique
		t.FKeys[i].ForeignColumnNullable = foreignColumn.Nullable
		t.FKeys[i].ForeignColumnUnique = foreignColumn.Unique
	}
}

func setRelationships(t *Table, tables []Table) {
	t.ToOneRelationships = toOneRelationships(*t, tables)
	t.ToManyRelationships = toManyRelationships(*t, tables)
}
<|endoftext|>


<commit_before>// +build !unit

package e2e

import (
	"github.com/stretchr/testify/assert"
	"testing"
	"io/ioutil"
	"fmt"
	"os"
	"github.com/stelligent/mu/common"
	"github.com/termie/go-shutil"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/codecommit"
	"github.com/stelligent/mu/provider/aws"
	"gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/config"
	"gopkg.in/src-d/go-git.v4/plumbing/object"
	"path/filepath"
	"time"
	"gopkg.in/src-d/go-git.v4/plumbing/transport/http"
	"bytes"
	"net/url"
	"regexp"
	"crypto/sha256"
	hm "crypto/hmac"

)

var contexts []*common.Context

func TestMain(m *testing.M) {
	sessOptions := session.Options{SharedConfigState: session.SharedConfigEnable}
	sess, err := session.NewSessionWithOptions(sessOptions)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	dir, err := ioutil.TempDir("", "mu-e2e-repos")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	var retCode int
	err = setupContexts(sess, dir)
	if err != nil {
		fmt.Println(err)
		retCode = 1
	} else {
		retCode = m.Run()
	}

	//teardownContexts(sess)
	os.RemoveAll(dir)
	os.Exit(retCode)
}

func TestPipelineE2E(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping e2e test")
	}
	assert := assert.New(t)

	assert.Fail("force fail")
}

func runScenario(name string) {
	fmt.Printf("scenario: %s", name)
	// - pipeline up
	// - wait...
	// - pipeline term
	// - env term
}

func setupContexts(sess *session.Session, basedir string) error {
	contexts = make([]*common.Context, 0)
	files, err := ioutil.ReadDir(".")
	if err != nil {
		return err
	}
	for _, f := range files {
		if f.IsDir() {
			// create temp dir for repo
			dst := fmt.Sprintf("%s/%s",basedir,f.Name())
			err = shutil.CopyTree(f.Name(), dst, nil)
			if err != nil {
				fmt.Printf("Failure copying repo '%s'\n", f.Name())
				return err
			}

			// create a codecommit repo
			err = setupRepo(sess, f.Name(), dst)
			if err != nil {
				return err
			}

			// load mu.yml
			ctx, err := initContext(fmt.Sprintf("%s/mu.yml", dst))
			if err != nil {
				return err
			}

			contexts = append(contexts, ctx)
		}
	}

	return nil
}

func initContext(muConfigFile string) (*common.Context, error) {
	// initialize context
	context := common.NewContext()
	err := context.InitializeContext()
	if err != nil {
		return nil, err
	}

	err = aws.InitializeContext(context, "", "", false)
	if err != nil {
		return nil, err
	}

	err = context.InitializeConfigFromFile(muConfigFile)
	if err != nil {
		return nil, err
	}

	return context, nil
}

func teardownContexts(sess *session.Session) {
	for _, ctx := range contexts {
		err := teardownRepo(sess, ctx)
		if err != nil {
			fmt.Printf("Unable to teardown repo: %s", err)
		}
	}
}
func setupRepo(sess *session.Session, name string, dir string) error {
	// create codecommit repo
	ccAPI := codecommit.New(sess)
	ccAPI.DeleteRepository(&codecommit.DeleteRepositoryInput{
		RepositoryName: &name,
	})
	resp, err := ccAPI.CreateRepository(&codecommit.CreateRepositoryInput{
		RepositoryName: &name,
	})
	if err != nil {
		fmt.Printf("Failure creating codecommit repo '%s'\n", name)
		return err
	}

	// git init
	repo, err := git.PlainInit(dir, false)
	if err != nil {
		fmt.Printf("Failure init repo '%s'\n", name)
		return err
	}

	// git commit
	wt, err := repo.Worktree()
	if err != nil {
		return err
	}
	err = gitAddAll(dir, ".", wt)
	if err != nil {
		fmt.Printf("Failure add '%s'\n", name)
		return err
	}
	_, err = wt.Commit("initial commit", &git.CommitOptions{
		Author:  &object.Signature{
			Name: "mu-e2e",
			Email: "mu@stelligent.com",
			When: time.Now(),
		},
	})
	if err != nil {
		fmt.Printf("Failure commit '%s'\n", name)
		return err
	}

	// git push
	remote, err := repo.CreateRemote(&config.RemoteConfig{
		Name: "origin",
		URL: common.StringValue(resp.RepositoryMetadata.CloneUrlHttp),
	})
	if err != nil {
		fmt.Printf("Failure create remote '%s'\n", name)
		return err
	}

	username, password := credentialHelper(sess, common.StringValue(resp.RepositoryMetadata.CloneUrlHttp))

	auth := http.NewBasicAuth(username, password)
	err = remote.Push(&git.PushOptions{
		RemoteName: "origin",
		Auth:  auth,
	})
	if err != nil {
		fmt.Printf("Failure pushing '%s'\n", name)
		return err
	}

	return nil

}

func gitAddAll(gitdir string, subdir string, wt *git.Worktree) error {
	files, err := ioutil.ReadDir(fmt.Sprintf("%s/%s", gitdir, subdir))
	if err != nil {
		return err
	}
	for _, f := range files {
		if f.Name() == ".git" {
			continue
		}
		prefixedName := filepath.Clean(fmt.Sprintf("%s/%s", subdir, f.Name()))
		if f.IsDir() {
			err = gitAddAll(gitdir, prefixedName, wt)
			if err != nil {
				return err
			}
		} else {
			_, err =  wt.Add(prefixedName)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func teardownRepo(sess *session.Session, ctx *common.Context) error {
	// delete codecommit repo
	ccAPI := codecommit.New(sess)
	params := &codecommit.DeleteRepositoryInput{
		RepositoryName: &ctx.Config.Repo.Name,
	}
	_, err := ccAPI.DeleteRepository(params)
	if err != nil {
		return err
	}

	return nil
}

var regexpCC = regexp.MustCompile(`git-codecommit\.([^.]+)\.amazonaws\.com.*`)
func credentialHelper(sess *session.Session, giturl string) (string,string) {
	u, _ :=  url.Parse(giturl)

	t := time.Now().UTC()
	//t, _ := time.Parse("20060102T150405", "20170824T072504") // reference time the comments were generated with

	creds, _ := sess.Config.Credentials.Get()
	username := creds.AccessKeyID
	if creds.SessionToken != "" {
		username = fmt.Sprintf("%s%s%s", username, "%", creds.SessionToken)
	}

	region := common.StringValue(sess.Config.Region)

	if match := regexpCC.FindStringSubmatch(u.Host); len(match) > 1 {
		region = match[1]
	}

	// Build canonical request
	// 'GIT\n/v1/repos/httpd\n\nhost:git-codecommit.us-east-1.amazonaws.com\n\nhost\n'
	cr := new(bytes.Buffer)
	fmt.Fprintf(cr, "%s\n", "GIT")         // HTTPRequestMethod
	fmt.Fprintf(cr, "%s\n", u.Path)        // CanonicalURI
	fmt.Fprintf(cr, "%s\n", "")            // CanonicalQueryString
	fmt.Fprintf(cr, "host:%s\n\n", u.Host) // CanonicalHeaders
	fmt.Fprintf(cr, "%s\n", "host")        // SignedHeaders
	fmt.Fprintf(cr, "%s", "")              // HexEncode(Hash(Payload))

	// Build string to sign
	// 'AWS4-HMAC-SHA256\n20160227T172057\n20160227/us-east-1/codecommit/aws4_request\n650b9e2de2abce7c30f6ad51c4a84b361e1f8aaaa3152e93d35509450db2d869'
	sts := new(bytes.Buffer)
	fmt.Fprint(sts, "AWS4-HMAC-SHA256\n")                                                   // Algorithm
	fmt.Fprintf(sts, "%s\n", t.Format("20060102T150405"))                                   // RequestDate
	fmt.Fprintf(sts, "%s/%s/%s/aws4_request\n", t.Format("20060102"), region, "codecommit") // CredentialScope
	fmt.Fprintf(sts, "%s", hash(cr.String()))                                               // HexEncode(Hash(CanonicalRequest))

	// Calculate the AWS Signature Version 4
	// '56238a5ac75c8bd36bba91737377aca46c867f584a6695f0486f3e3bba9b4ed5'
	dsk := hmac([]byte("AWS4"+creds.SecretAccessKey), []byte(t.Format("20060102")))
	dsk = hmac(dsk, []byte(region))
	dsk = hmac(dsk, []byte("codecommit"))
	dsk = hmac(dsk, []byte("aws4_request"))
	h := hmac(dsk, []byte(sts.String()))
	sig := fmt.Sprintf("%x", h) // HexEncode(HMAC(derived-signing-key, string-to-sign))

	// codecommmit smart http password to use with AWS_ACCESS
	password := fmt.Sprintf("%sZ%s", t.Format("20060102T150405"), sig)
	//fmt.Printf("url:%s username:%s password:%s\n", u, username, password)

	return username, password
}

// hash method calculates the sha256 hash for a given string
func hash(in string) string {
	h := sha256.New()
	fmt.Fprintf(h, "%s", in)
	return fmt.Sprintf("%x", h.Sum(nil))
}

// hmac method calculates the sha256 hmac for a given slice of bytes
func hmac(key, data []byte) []byte {
	h := hm.New(sha256.New, key)
	h.Write(data)
	return h.Sum(nil)
}<commit_msg>stub out context loop for e2e tests<commit_after>// +build !unit

package e2e

import (
	"bytes"
	hm "crypto/hmac"
	"crypto/sha256"
	"fmt"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/codecommit"
	"github.com/stelligent/mu/common"
	"github.com/stelligent/mu/provider/aws"
	"github.com/stretchr/testify/assert"
	"github.com/termie/go-shutil"
	"gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/config"
	"gopkg.in/src-d/go-git.v4/plumbing/object"
	"gopkg.in/src-d/go-git.v4/plumbing/transport/http"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"testing"
	"time"
)

var contexts []*common.Context

func TestMain(m *testing.M) {
	sessOptions := session.Options{SharedConfigState: session.SharedConfigEnable}
	sess, err := session.NewSessionWithOptions(sessOptions)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	dir, err := ioutil.TempDir("", "mu-e2e-repos")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	var retCode int
	err = setupContexts(sess, dir)
	if err != nil {
		fmt.Println(err)
		retCode = 1
	} else {
		retCode = m.Run()
	}

	teardownContexts(sess)
	os.RemoveAll(dir)
	os.Exit(retCode)
}

func TestPipelineE2E(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping e2e test")
	}
	assert := assert.New(t)

	for _, ctx := range contexts {
		fmt.Printf("scenario: %s\n", ctx.Config.Repo.Name)

		// - pipeline up
		// - wait...
		// - pipeline term
		// - env term
	}
	assert.Fail("force fail")
}

func setupContexts(sess *session.Session, basedir string) error {
	contexts = make([]*common.Context, 0)
	files, err := ioutil.ReadDir(".")
	if err != nil {
		return err
	}
	for _, f := range files {
		if f.IsDir() {
			// create temp dir for repo
			dst := fmt.Sprintf("%s/%s", basedir, f.Name())
			err = shutil.CopyTree(f.Name(), dst, nil)
			if err != nil {
				fmt.Printf("Failure copying repo '%s'\n", f.Name())
				return err
			}

			// create a codecommit repo
			err = setupRepo(sess, f.Name(), dst)
			if err != nil {
				return err
			}

			// load mu.yml
			ctx, err := initContext(fmt.Sprintf("%s/mu.yml", dst))
			if err != nil {
				return err
			}

			contexts = append(contexts, ctx)
		}
	}

	return nil
}

func initContext(muConfigFile string) (*common.Context, error) {
	// initialize context
	context := common.NewContext()
	err := context.InitializeContext()
	if err != nil {
		return nil, err
	}

	err = aws.InitializeContext(context, "", "", false)
	if err != nil {
		return nil, err
	}

	err = context.InitializeConfigFromFile(muConfigFile)
	if err != nil {
		return nil, err
	}

	return context, nil
}

func teardownContexts(sess *session.Session) {
	for _, ctx := range contexts {
		err := teardownRepo(sess, ctx)
		if err != nil {
			fmt.Printf("Unable to teardown repo: %s", err)
		}
	}
}
func setupRepo(sess *session.Session, name string, dir string) error {
	// create codecommit repo
	ccAPI := codecommit.New(sess)
	ccAPI.DeleteRepository(&codecommit.DeleteRepositoryInput{
		RepositoryName: &name,
	})
	resp, err := ccAPI.CreateRepository(&codecommit.CreateRepositoryInput{
		RepositoryName: &name,
	})
	if err != nil {
		fmt.Printf("Failure creating codecommit repo '%s'\n", name)
		return err
	}

	// git init
	repo, err := git.PlainInit(dir, false)
	if err != nil {
		fmt.Printf("Failure init repo '%s'\n", name)
		return err
	}

	// git commit
	wt, err := repo.Worktree()
	if err != nil {
		return err
	}
	err = gitAddAll(dir, ".", wt)
	if err != nil {
		fmt.Printf("Failure add '%s'\n", name)
		return err
	}
	_, err = wt.Commit("initial commit", &git.CommitOptions{
		Author: &object.Signature{
			Name:  "mu-e2e",
			Email: "mu@stelligent.com",
			When:  time.Now(),
		},
	})
	if err != nil {
		fmt.Printf("Failure commit '%s'\n", name)
		return err
	}

	// git push
	remote, err := repo.CreateRemote(&config.RemoteConfig{
		Name: "origin",
		URL:  common.StringValue(resp.RepositoryMetadata.CloneUrlHttp),
	})
	if err != nil {
		fmt.Printf("Failure create remote '%s'\n", name)
		return err
	}

	username, password := credentialHelper(sess, common.StringValue(resp.RepositoryMetadata.CloneUrlHttp))

	auth := http.NewBasicAuth(username, password)
	err = remote.Push(&git.PushOptions{
		RemoteName: "origin",
		Auth:       auth,
	})
	if err != nil {
		fmt.Printf("Failure pushing '%s'\n", name)
		return err
	}

	return nil

}

func gitAddAll(gitdir string, subdir string, wt *git.Worktree) error {
	files, err := ioutil.ReadDir(fmt.Sprintf("%s/%s", gitdir, subdir))
	if err != nil {
		return err
	}
	for _, f := range files {
		if f.Name() == ".git" {
			continue
		}
		prefixedName := filepath.Clean(fmt.Sprintf("%s/%s", subdir, f.Name()))
		if f.IsDir() {
			err = gitAddAll(gitdir, prefixedName, wt)
			if err != nil {
				return err
			}
		} else {
			_, err = wt.Add(prefixedName)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func teardownRepo(sess *session.Session, ctx *common.Context) error {
	// delete codecommit repo
	ccAPI := codecommit.New(sess)
	params := &codecommit.DeleteRepositoryInput{
		RepositoryName: &ctx.Config.Repo.Name,
	}
	_, err := ccAPI.DeleteRepository(params)
	if err != nil {
		return err
	}

	return nil
}

var regexpCC = regexp.MustCompile(`git-codecommit\.([^.]+)\.amazonaws\.com.*`)

func credentialHelper(sess *session.Session, giturl string) (string, string) {
	u, _ := url.Parse(giturl)

	t := time.Now().UTC()
	//t, _ := time.Parse("20060102T150405", "20170824T072504") // reference time the comments were generated with

	creds, _ := sess.Config.Credentials.Get()
	username := creds.AccessKeyID
	if creds.SessionToken != "" {
		username = fmt.Sprintf("%s%s%s", username, "%", creds.SessionToken)
	}

	region := common.StringValue(sess.Config.Region)

	if match := regexpCC.FindStringSubmatch(u.Host); len(match) > 1 {
		region = match[1]
	}

	// Build canonical request
	// 'GIT\n/v1/repos/httpd\n\nhost:git-codecommit.us-east-1.amazonaws.com\n\nhost\n'
	cr := new(bytes.Buffer)
	fmt.Fprintf(cr, "%s\n", "GIT")         // HTTPRequestMethod
	fmt.Fprintf(cr, "%s\n", u.Path)        // CanonicalURI
	fmt.Fprintf(cr, "%s\n", "")            // CanonicalQueryString
	fmt.Fprintf(cr, "host:%s\n\n", u.Host) // CanonicalHeaders
	fmt.Fprintf(cr, "%s\n", "host")        // SignedHeaders
	fmt.Fprintf(cr, "%s", "")              // HexEncode(Hash(Payload))

	// Build string to sign
	// 'AWS4-HMAC-SHA256\n20160227T172057\n20160227/us-east-1/codecommit/aws4_request\n650b9e2de2abce7c30f6ad51c4a84b361e1f8aaaa3152e93d35509450db2d869'
	sts := new(bytes.Buffer)
	fmt.Fprint(sts, "AWS4-HMAC-SHA256\n")                                                   // Algorithm
	fmt.Fprintf(sts, "%s\n", t.Format("20060102T150405"))                                   // RequestDate
	fmt.Fprintf(sts, "%s/%s/%s/aws4_request\n", t.Format("20060102"), region, "codecommit") // CredentialScope
	fmt.Fprintf(sts, "%s", hash(cr.String()))                                               // HexEncode(Hash(CanonicalRequest))

	// Calculate the AWS Signature Version 4
	// '56238a5ac75c8bd36bba91737377aca46c867f584a6695f0486f3e3bba9b4ed5'
	dsk := hmac([]byte("AWS4"+creds.SecretAccessKey), []byte(t.Format("20060102")))
	dsk = hmac(dsk, []byte(region))
	dsk = hmac(dsk, []byte("codecommit"))
	dsk = hmac(dsk, []byte("aws4_request"))
	h := hmac(dsk, []byte(sts.String()))
	sig := fmt.Sprintf("%x", h) // HexEncode(HMAC(derived-signing-key, string-to-sign))

	// codecommmit smart http password to use with AWS_ACCESS
	password := fmt.Sprintf("%sZ%s", t.Format("20060102T150405"), sig)
	//fmt.Printf("url:%s username:%s password:%s\n", u, username, password)

	return username, password
}

// hash method calculates the sha256 hash for a given string
func hash(in string) string {
	h := sha256.New()
	fmt.Fprintf(h, "%s", in)
	return fmt.Sprintf("%x", h.Sum(nil))
}

// hmac method calculates the sha256 hmac for a given slice of bytes
func hmac(key, data []byte) []byte {
	h := hm.New(sha256.New, key)
	h.Write(data)
	return h.Sum(nil)
}
<|endoftext|>


<commit_before>package ethutil

import (
	"bytes"
	"testing"
)

func TestParseDataString(t *testing.T) {
	data := ParseData("hello", "world", "0x0106")
	exp := "68656c6c6f000000000000000000000000000000000000000000000000000000776f726c640000000000000000000000000000000000000000000000000000000106000000000000000000000000000000000000000000000000000000000000"
	if bytes.Compare(data, Hex2Bytes(exp)) != 0 {
		t.Error("Error parsing data")
	}
}

func TestParseDataBytes(t *testing.T) {
	data := []byte{232, 212, 165, 16, 0}
	exp := []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 232, 212, 165, 16, 0}

	res := ParseData(data)
	if bytes.Compare(res, exp) != 0 {
		t.Errorf("Expected %x got %x", exp, res)
	}
}

func TestLeftPadBytes(t *testing.T) {
	val := []byte{1, 2, 3, 4}
	exp := []byte{0, 0, 0, 0, 1, 2, 3, 4}

	resstd := LeftPadBytes(val, 8)
	if bytes.Compare(resstd, exp) != 0 {
		t.Errorf("Expected % x Got % x", exp, resstd)
	}

	resshrt := LeftPadBytes(val, 2)
	if bytes.Compare(resshrt, val) != 0 {
		t.Errorf("Expected % x Got % x", exp, resshrt)
	}
}

func TestRightPadBytes(t *testing.T) {
	val := []byte{1, 2, 3, 4}
	exp := []byte{1, 2, 3, 4, 0, 0, 0, 0}

	resstd := RightPadBytes(val, 8)
	if bytes.Compare(resstd, exp) != 0 {
		t.Errorf("Expected % x Got % x", exp, resstd)
	}

	resshrt := RightPadBytes(val, 2)
	if bytes.Compare(resshrt, val) != 0 {
		t.Errorf("Expected % x Got % x", exp, resshrt)
	}
}

func TestLeftPadString(t *testing.T) {
	val := "test"

	resstd := LeftPadString(val, 8)

	if resstd != "\x30\x30\x30\x30"+val {
		t.Errorf("Expected % x Got % x", val, resstd)
	}

	resshrt := LeftPadString(val, 2)

	if resshrt != val {
		t.Errorf("Expected % x Got % x", val, resshrt)
	}
}

func TestRightPadString(t *testing.T) {
	val := "test"

	resstd := RightPadString(val, 8)
	if resstd != val+"\x30\x30\x30\x30" {
		t.Errorf("Expected % x Got % x", val, resstd)
	}

	resshrt := RightPadString(val, 2)
	if resshrt != val {
		t.Errorf("Expected % x Got % x", val, resshrt)
	}
}

func TestReadVarInt(t *testing.T) {
	data8 := []byte{1, 2, 3, 4, 5, 6, 7, 8}
	data4 := []byte{1, 2, 3, 4}
	data2 := []byte{1, 2}
	data1 := []byte{1}

	exp8 := uint64(72623859790382856)
	exp4 := uint64(16909060)
	exp2 := uint64(258)
	exp1 := uint64(1)

	res8 := ReadVarInt(data8)
	res4 := ReadVarInt(data4)
	res2 := ReadVarInt(data2)
	res1 := ReadVarInt(data1)

	if res8 != exp8 {
		t.Errorf("Expected %d | Got %d", exp8, res8)
	}

	if res4 != exp4 {
		t.Errorf("Expected %d | Got %d", exp4, res4)
	}

	if res2 != exp2 {
		t.Errorf("Expected %d | Got %d", exp2, res2)
	}

	if res1 != exp1 {
		t.Errorf("Expected %d | Got %d", exp1, res1)
	}
}
<commit_msg>added test for BinaryLength<commit_after>package ethutil

import (
	"bytes"
	"testing"
)

func TestParseDataString(t *testing.T) {
	data := ParseData("hello", "world", "0x0106")
	exp := "68656c6c6f000000000000000000000000000000000000000000000000000000776f726c640000000000000000000000000000000000000000000000000000000106000000000000000000000000000000000000000000000000000000000000"
	if bytes.Compare(data, Hex2Bytes(exp)) != 0 {
		t.Error("Error parsing data")
	}
}

func TestParseDataBytes(t *testing.T) {
	data := []byte{232, 212, 165, 16, 0}
	exp := []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 232, 212, 165, 16, 0}

	res := ParseData(data)
	if bytes.Compare(res, exp) != 0 {
		t.Errorf("Expected %x got %x", exp, res)
	}
}

func TestLeftPadBytes(t *testing.T) {
	val := []byte{1, 2, 3, 4}
	exp := []byte{0, 0, 0, 0, 1, 2, 3, 4}

	resstd := LeftPadBytes(val, 8)
	if bytes.Compare(resstd, exp) != 0 {
		t.Errorf("Expected % x Got % x", exp, resstd)
	}

	resshrt := LeftPadBytes(val, 2)
	if bytes.Compare(resshrt, val) != 0 {
		t.Errorf("Expected % x Got % x", exp, resshrt)
	}
}

func TestRightPadBytes(t *testing.T) {
	val := []byte{1, 2, 3, 4}
	exp := []byte{1, 2, 3, 4, 0, 0, 0, 0}

	resstd := RightPadBytes(val, 8)
	if bytes.Compare(resstd, exp) != 0 {
		t.Errorf("Expected % x Got % x", exp, resstd)
	}

	resshrt := RightPadBytes(val, 2)
	if bytes.Compare(resshrt, val) != 0 {
		t.Errorf("Expected % x Got % x", exp, resshrt)
	}
}

func TestLeftPadString(t *testing.T) {
	val := "test"

	resstd := LeftPadString(val, 8)

	if resstd != "\x30\x30\x30\x30"+val {
		t.Errorf("Expected % x Got % x", val, resstd)
	}

	resshrt := LeftPadString(val, 2)

	if resshrt != val {
		t.Errorf("Expected % x Got % x", val, resshrt)
	}
}

func TestRightPadString(t *testing.T) {
	val := "test"

	resstd := RightPadString(val, 8)
	if resstd != val+"\x30\x30\x30\x30" {
		t.Errorf("Expected % x Got % x", val, resstd)
	}

	resshrt := RightPadString(val, 2)
	if resshrt != val {
		t.Errorf("Expected % x Got % x", val, resshrt)
	}
}

func TestReadVarInt(t *testing.T) {
	data8 := []byte{1, 2, 3, 4, 5, 6, 7, 8}
	data4 := []byte{1, 2, 3, 4}
	data2 := []byte{1, 2}
	data1 := []byte{1}

	exp8 := uint64(72623859790382856)
	exp4 := uint64(16909060)
	exp2 := uint64(258)
	exp1 := uint64(1)

	res8 := ReadVarInt(data8)
	res4 := ReadVarInt(data4)
	res2 := ReadVarInt(data2)
	res1 := ReadVarInt(data1)

	if res8 != exp8 {
		t.Errorf("Expected %d | Got %d", exp8, res8)
	}

	if res4 != exp4 {
		t.Errorf("Expected %d | Got %d", exp4, res4)
	}

	if res2 != exp2 {
		t.Errorf("Expected %d | Got %d", exp2, res2)
	}

	if res1 != exp1 {
		t.Errorf("Expected %d | Got %d", exp1, res1)
	}
}

func TestBinaryLength(t *testing.T) {
	data1 := 0
	data2 := 920987656789

	exp1 := 0
	exp2 := 5

	res1 := BinaryLength(data1)
	res2 := BinaryLength(data2)

	if res1 != exp1 {
		t.Errorf("Expected %d got %d", exp1, res1)
	}

	if res2 != exp2 {
		t.Errorf("Expected %d got %d", exp2, res2)
	}
}
<|endoftext|>


<commit_before>package stan_output

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/stan.go"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/reflect/protoreflect"
)

const (
	stanDefaultTimeout      = 10
	stanDefaultPingInterval = 5
	stanDefaultPingRetry    = 2

	defaultSubjectName = "gnmic-telemetry"

	defaultFormat = "json"
)

func init() {
	outputs.Register("stan", func() outputs.Output {
		return &StanOutput{
			Cfg: &Config{},
		}
	})
}

// StanOutput //
type StanOutput struct {
	Cfg     *Config
	conn    stan.Conn
	metrics []prometheus.Collector
	logger  *log.Logger
	mo      *collector.MarshalOptions
}

// Config //
type Config struct {
	Name          string `mapstructure:"name,omitempty"`
	Address       string `mapstructure:"address,omitempty"`
	SubjectPrefix string `mapstructure:"subject-prefix,omitempty"`
	Subject       string `mapstructure:"subject,omitempty"`
	Username      string `mapstructure:"username,omitempty"`
	Password      string `mapstructure:"password,omitempty"`
	ClusterName   string `mapstructure:"cluster-name,omitempty"`
	Timeout       int    `mapstructure:"timeout,omitempty"`
	PingInterval  int    `mapstructure:"ping-interval,omitempty"`
	PingRetry     int    `mapstructure:"ping-retry,omitempty"`
	Format        string `mapstructure:"format,omitempty"`
}

func (s *StanOutput) String() string {
	b, err := json.Marshal(s)
	if err != nil {
		return ""
	}
	return string(b)
}

// Init //
func (s *StanOutput) Init(cfg map[string]interface{}, logger *log.Logger) error {
	err := mapstructure.Decode(cfg, s.Cfg)
	if err != nil {
		return err
	}
	if s.Cfg.Name == "" {
		s.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if s.Cfg.ClusterName == "" {
		return fmt.Errorf("clusterName is mandatory")
	}
	if s.Cfg.Subject == "" && s.Cfg.SubjectPrefix == "" {
		s.Cfg.Subject = defaultSubjectName
	}
	s.conn, err = s.createSTANConn(s.Cfg)
	if err != nil {
		return err
	}
	s.logger = log.New(os.Stderr, "stan_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		s.logger.SetOutput(logger.Writer())
		s.logger.SetFlags(logger.Flags())
	}
	if s.Cfg.Format == "" {
		s.Cfg.Format = defaultFormat
	}
	if !(s.Cfg.Format == "event" || s.Cfg.Format == "protojson" || s.Cfg.Format == "proto" || s.Cfg.Format == "json") {
		return fmt.Errorf("unsupported output format: '%s' for output type STAN", s.Cfg.Format)
	}
	s.mo = &collector.MarshalOptions{Format: s.Cfg.Format}
	s.logger.Printf("initialized stan producer: %s", s.String())
	return nil
}

// Write //
func (s *StanOutput) Write(rsp protoreflect.ProtoMessage, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	ssb := strings.Builder{}
	ssb.WriteString(s.Cfg.SubjectPrefix)
	if s.Cfg.SubjectPrefix != "" {
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(s, ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(subname)
		}
	} else if s.Cfg.Subject != "" {
		ssb.WriteString(s.Cfg.Subject)
	}
	subject := strings.ReplaceAll(ssb.String(), " ", "_")
	b, err := s.mo.Marshal(rsp, meta)
	if err != nil {
		s.logger.Printf("failed marshaling proto msg: %v", err)
		return
	}
	err = s.conn.Publish(subject, b)
	if err != nil {
		log.Printf("failed to write to stan subject '%s': %v", subject, err)
		return
	}
	//s.logger.Printf("wrote %d bytes to stan_subject=%s", len(b), s.Cfg.Subject)
}

// Metrics //
func (s *StanOutput) Metrics() []prometheus.Collector { return s.metrics }

// Close //
func (s *StanOutput) Close() error {
	return s.conn.Close()
}

func (s *StanOutput) createSTANConn(c *Config) (stan.Conn, error) {
	if c.Timeout == 0 {
		c.Timeout = stanDefaultTimeout
	}
	if c.PingInterval == 0 {
		c.PingInterval = stanDefaultPingInterval
	}
	if c.PingRetry == 0 {
		c.PingRetry = stanDefaultPingRetry
	}
	opts := []nats.Option{
		nats.Name(c.Name),
		// nats.Timeout(time.Duration(c.Timeout) * time.Second),
		// nats.PingInterval(time.Duration(c.PingInterval) * time.Second),
		// nats.MaxPingsOutstanding(c.PingRetry),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	sc, err := stan.Connect(c.ClusterName, c.Name,
		stan.NatsConn(nc),
		stan.Pings(c.PingInterval, c.PingRetry),
		stan.SetConnectionLostHandler(func(_ stan.Conn, err error) {
			log.Fatalf("STAN connection lost, reason: %v", err)
		}),
	)
	if err != nil {
		return nil, fmt.Errorf("cannot connect to %s: %v", c.Address, err)
	}
	return sc, nil
}
<commit_msg>added ctx to stan output<commit_after>package stan_output

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/stan.go"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/reflect/protoreflect"
)

const (
	stanDefaultTimeout      = 10
	stanDefaultPingInterval = 5
	stanDefaultPingRetry    = 2

	defaultSubjectName = "gnmic-telemetry"

	defaultFormat = "json"
)

func init() {
	outputs.Register("stan", func() outputs.Output {
		return &StanOutput{
			Cfg: &Config{},
		}
	})
}

// StanOutput //
type StanOutput struct {
	Cfg     *Config
	conn    stan.Conn
	metrics []prometheus.Collector
	logger  *log.Logger
	mo      *collector.MarshalOptions
}

// Config //
type Config struct {
	Name          string `mapstructure:"name,omitempty"`
	Address       string `mapstructure:"address,omitempty"`
	SubjectPrefix string `mapstructure:"subject-prefix,omitempty"`
	Subject       string `mapstructure:"subject,omitempty"`
	Username      string `mapstructure:"username,omitempty"`
	Password      string `mapstructure:"password,omitempty"`
	ClusterName   string `mapstructure:"cluster-name,omitempty"`
	Timeout       int    `mapstructure:"timeout,omitempty"`
	PingInterval  int    `mapstructure:"ping-interval,omitempty"`
	PingRetry     int    `mapstructure:"ping-retry,omitempty"`
	Format        string `mapstructure:"format,omitempty"`
}

func (s *StanOutput) String() string {
	b, err := json.Marshal(s)
	if err != nil {
		return ""
	}
	return string(b)
}

// Init //
func (s *StanOutput) Init(ctx context.Context, cfg map[string]interface{}, logger *log.Logger) error {
	err := mapstructure.Decode(cfg, s.Cfg)
	if err != nil {
		return err
	}
	if s.Cfg.Name == "" {
		s.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if s.Cfg.ClusterName == "" {
		return fmt.Errorf("clusterName is mandatory")
	}
	if s.Cfg.Subject == "" && s.Cfg.SubjectPrefix == "" {
		s.Cfg.Subject = defaultSubjectName
	}
	s.conn, err = s.createSTANConn(s.Cfg)
	if err != nil {
		return err
	}
	s.logger = log.New(os.Stderr, "stan_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		s.logger.SetOutput(logger.Writer())
		s.logger.SetFlags(logger.Flags())
	}
	if s.Cfg.Format == "" {
		s.Cfg.Format = defaultFormat
	}
	if !(s.Cfg.Format == "event" || s.Cfg.Format == "protojson" || s.Cfg.Format == "proto" || s.Cfg.Format == "json") {
		return fmt.Errorf("unsupported output format: '%s' for output type STAN", s.Cfg.Format)
	}
	s.mo = &collector.MarshalOptions{Format: s.Cfg.Format}
	s.logger.Printf("initialized stan producer: %s", s.String())
	go func() {
		<-ctx.Done()
		s.Close()
	}()
	return nil
}

// Write //
func (s *StanOutput) Write(_ context.Context, rsp protoreflect.ProtoMessage, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	ssb := strings.Builder{}
	ssb.WriteString(s.Cfg.SubjectPrefix)
	if s.Cfg.SubjectPrefix != "" {
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(s, ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(subname)
		}
	} else if s.Cfg.Subject != "" {
		ssb.WriteString(s.Cfg.Subject)
	}
	subject := strings.ReplaceAll(ssb.String(), " ", "_")
	b, err := s.mo.Marshal(rsp, meta)
	if err != nil {
		s.logger.Printf("failed marshaling proto msg: %v", err)
		return
	}
	err = s.conn.Publish(subject, b)
	if err != nil {
		log.Printf("failed to write to stan subject '%s': %v", subject, err)
		return
	}
	//s.logger.Printf("wrote %d bytes to stan_subject=%s", len(b), s.Cfg.Subject)
}

// Metrics //
func (s *StanOutput) Metrics() []prometheus.Collector { return s.metrics }

// Close //
func (s *StanOutput) Close() error {
	return s.conn.Close()
}

func (s *StanOutput) createSTANConn(c *Config) (stan.Conn, error) {
	if c.Timeout == 0 {
		c.Timeout = stanDefaultTimeout
	}
	if c.PingInterval == 0 {
		c.PingInterval = stanDefaultPingInterval
	}
	if c.PingRetry == 0 {
		c.PingRetry = stanDefaultPingRetry
	}
	opts := []nats.Option{
		nats.Name(c.Name),
		// nats.Timeout(time.Duration(c.Timeout) * time.Second),
		// nats.PingInterval(time.Duration(c.PingInterval) * time.Second),
		// nats.MaxPingsOutstanding(c.PingRetry),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	sc, err := stan.Connect(c.ClusterName, c.Name,
		stan.NatsConn(nc),
		stan.Pings(c.PingInterval, c.PingRetry),
		stan.SetConnectionLostHandler(func(_ stan.Conn, err error) {
			log.Fatalf("STAN connection lost, reason: %v", err)
		}),
	)
	if err != nil {
		return nil, fmt.Errorf("cannot connect to %s: %v", c.Address, err)
	}
	return sc, nil
}
<|endoftext|>


<commit_before>package twitch

import (
	"bufio"
	"fmt"
	"net"
	"net/textproto"
	"strings"
	"time"
)

const (
	// ircTwitch constant for irc chat address
	ircTwitch = "irc.chat.twitch.tv:6667"
)

// User data you receive from tmi
type User struct {
	Username    string
	DisplayName string
	UserType    string
	Color       string
	Badges      map[string]int
}

// Message data you receive from tmi
type Message struct {
	Type   msgType
	Time   time.Time
	Action bool
	Emotes []*Emote
	Tags   map[string]string
	Text   string
}

// Client client to control your connection and attach callbacks
type Client struct {
	IrcAddress            string
	ircUser               string
	ircToken              string
	connection            *net.Conn
	connActive            bool
	closingConn           bool
	onNewMessage          MessageCallback
	onNewRoomstateMessage MessageCallback
	onNewClearchatMessage MessageCallback
}

// MessageCallback is called when a new message is received
type MessageCallback func(channel string, user User, message Message)

// NewClient to create a new client
func NewClient(username, oauth string) *Client {
	return &Client{
		ircUser:    username,
		ircToken:   oauth,
		IrcAddress: ircTwitch,
	}
}

// OnNewMessage attach callback to new standard chat messages
func (c *Client) OnNewMessage(callback MessageCallback) {
	c.onNewMessage = callback
}

// OnNewRoomstateMessage attach callback to new messages such as submode enabled
func (c *Client) OnNewRoomstateMessage(callback MessageCallback) {
	c.onNewRoomstateMessage = callback
}

// OnNewClearchatMessage attach callback to new messages such as timeouts
func (c *Client) OnNewClearchatMessage(callback MessageCallback) {
	c.onNewClearchatMessage = callback
}

// Say write something in a chat
func (c *Client) Say(channel, text string) {
	c.send(fmt.Sprintf("PRIVMSG #%s :%s", channel, text))
}

// Join enter a twitch channel to read more messages
func (c *Client) Join(channel string) {
	c.send(fmt.Sprintf("JOIN #%s", channel))
}

// Connect connect the client to the irc server
func (c *Client) Connect() error {
	for {
		conn, err := net.Dial("tcp", c.IrcAddress)
		c.connection = &conn
		if err != nil {
			return err
		}

		go c.setupConnection()

		err = c.readConnection(conn)
		if err != nil {
			time.Sleep(time.Millisecond * 200)
			continue
		}
	}
}

func (c *Client) readConnection(conn net.Conn) error {
	reader := bufio.NewReader(conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()
		if err != nil {
			return err
		}
		messages := strings.Split(line, "\r\n")
		for _, msg := range messages {
			if !c.connActive && strings.HasPrefix(msg, ":tmi.twitch.tv 001") {
				c.connActive = true
			}
			c.handleLine(msg)
		}
	}
}

func (c *Client) setupConnection() {
	fmt.Fprintf(*c.connection, "PASS %s\r\n", c.ircToken)
	fmt.Fprintf(*c.connection, "NICK %s\r\n", c.ircUser)
	fmt.Fprint(*c.connection, "CAP REQ :twitch.tv/tags\r\n")
	fmt.Fprint(*c.connection, "CAP REQ :twitch.tv/commands\r\n")
}

func (c *Client) send(line string) {
	for i := 0; i < 1000; i++ {
		if !c.connActive {
			time.Sleep(time.Millisecond * 2)
			continue
		}
		fmt.Fprint(*c.connection, line+"\r\n")
		return
	}
}

func (c *Client) handleLine(line string) {
	if strings.HasPrefix(line, "PING") {
		c.send(strings.Replace(line, "PING", "PONG", 1))
	}
	if strings.HasPrefix(line, "@") {
		message := parseMessage(line)

		Channel := message.Channel

		User := &User{
			Username:    message.Username,
			DisplayName: message.DisplayName,
			UserType:    message.UserType,
			Color:       message.Color,
			Badges:      message.Badges,
		}

		clientMessage := &Message{
			Type:   message.Type,
			Time:   message.Time,
			Action: message.Action,
			Emotes: message.Emotes,
			Tags:   message.Tags,
			Text:   message.Text,
		}

		switch message.Type {
		case PRIVMSG:
			if c.onNewMessage != nil {
				c.onNewMessage(Channel, *User, *clientMessage)
			}
		case ROOMSTATE:
			if c.onNewRoomstateMessage != nil {
				c.onNewRoomstateMessage(Channel, *User, *clientMessage)
			}
		case CLEARCHAT:
			if c.onNewClearchatMessage != nil {
				c.onNewClearchatMessage(Channel, *User, *clientMessage)
			}
		}
	}
}
<commit_msg>removed some unused code<commit_after>package twitch

import (
	"bufio"
	"fmt"
	"net"
	"net/textproto"
	"strings"
	"time"
)

const (
	// ircTwitch constant for twitch irc chat address
	ircTwitch = "irc.chat.twitch.tv:6667"
)

// User data you receive from tmi
type User struct {
	Username    string
	DisplayName string
	UserType    string
	Color       string
	Badges      map[string]int
}

// Message data you receive from tmi
type Message struct {
	Type   msgType
	Time   time.Time
	Action bool
	Emotes []*Emote
	Tags   map[string]string
	Text   string
}

// Client client to control your connection and attach callbacks
type Client struct {
	IrcAddress            string
	ircUser               string
	ircToken              string
	connection            *net.Conn
	connActive            bool
	onNewMessage          func(channel string, user User, message Message)
	onNewRoomstateMessage func(channel string, user User, message Message)
	onNewClearchatMessage func(channel string, user User, message Message)
}

// NewClient to create a new client
func NewClient(username, oauth string) *Client {
	return &Client{
		ircUser:    username,
		ircToken:   oauth,
		IrcAddress: ircTwitch,
	}
}

// OnNewMessage attach callback to new standard chat messages
func (c *Client) OnNewMessage(callback func(channel string, user User, message Message)) {
	c.onNewMessage = callback
}

// OnNewRoomstateMessage attach callback to new messages such as submode enabled
func (c *Client) OnNewRoomstateMessage(callback func(channel string, user User, message Message)) {
	c.onNewRoomstateMessage = callback
}

// OnNewClearchatMessage attach callback to new messages such as timeouts
func (c *Client) OnNewClearchatMessage(callback func(channel string, user User, message Message)) {
	c.onNewClearchatMessage = callback
}

// Say write something in a chat
func (c *Client) Say(channel, text string) {
	c.send(fmt.Sprintf("PRIVMSG #%s :%s", channel, text))
}

// Join enter a twitch channel to read more messages
func (c *Client) Join(channel string) {
	c.send(fmt.Sprintf("JOIN #%s", channel))
}

// Connect connect the client to the irc server
func (c *Client) Connect() error {
	for {
		conn, err := net.Dial("tcp", c.IrcAddress)
		c.connection = &conn
		if err != nil {
			return err
		}

		go c.setupConnection()

		err = c.readConnection(conn)
		if err != nil {
			time.Sleep(time.Millisecond * 200)
			continue
		}
	}
}

func (c *Client) readConnection(conn net.Conn) error {
	reader := bufio.NewReader(conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()
		if err != nil {
			return err
		}
		messages := strings.Split(line, "\r\n")
		for _, msg := range messages {
			if !c.connActive && strings.HasPrefix(msg, ":tmi.twitch.tv 001") {
				c.connActive = true
			}
			c.handleLine(msg)
		}
	}
}

func (c *Client) setupConnection() {
	fmt.Fprintf(*c.connection, "PASS %s\r\n", c.ircToken)
	fmt.Fprintf(*c.connection, "NICK %s\r\n", c.ircUser)
	fmt.Fprint(*c.connection, "CAP REQ :twitch.tv/tags\r\n")
	fmt.Fprint(*c.connection, "CAP REQ :twitch.tv/commands\r\n")
}

func (c *Client) send(line string) {
	for i := 0; i < 1000; i++ {
		if !c.connActive {
			time.Sleep(time.Millisecond * 2)
			continue
		}
		fmt.Fprint(*c.connection, line+"\r\n")
		return
	}
}

func (c *Client) handleLine(line string) {
	if strings.HasPrefix(line, "PING") {
		c.send(strings.Replace(line, "PING", "PONG", 1))
	}
	if strings.HasPrefix(line, "@") {
		message := parseMessage(line)

		Channel := message.Channel

		User := &User{
			Username:    message.Username,
			DisplayName: message.DisplayName,
			UserType:    message.UserType,
			Color:       message.Color,
			Badges:      message.Badges,
		}

		clientMessage := &Message{
			Type:   message.Type,
			Time:   message.Time,
			Action: message.Action,
			Emotes: message.Emotes,
			Tags:   message.Tags,
			Text:   message.Text,
		}

		switch message.Type {
		case PRIVMSG:
			if c.onNewMessage != nil {
				c.onNewMessage(Channel, *User, *clientMessage)
			}
		case ROOMSTATE:
			if c.onNewRoomstateMessage != nil {
				c.onNewRoomstateMessage(Channel, *User, *clientMessage)
			}
		case CLEARCHAT:
			if c.onNewClearchatMessage != nil {
				c.onNewClearchatMessage(Channel, *User, *clientMessage)
			}
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>client: correction in func composeStartupConfig()<commit_after><|endoftext|>


<commit_before>/**
 * Author:        Tony.Shao
 * Email:         xiocode@gmail.com
 * Github:        github.com/xiocode
 * File:          client.go
 * Description:   httplib client
 */

package httplib

import (
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"sync"
)

type status int

type Client struct {
	sync.Mutex
	conn    *http.Client
	session http.Header
}

func NewClient() *Client {
	return &Client{
		conn: new(http.Client),
	}
}

func NewProxyClient(proxy string) *Client {
	proxyURL, _ := url.Parse(proxy)
	transport := &http.Transport{Proxy: http.ProxyURL(proxyURL)}
	return &Client{
		conn: &http.Client{
			Transport: transport,
		},
	}
}

func NewSession() *Client {
	jar, _ := cookiejar.New(nil)
	return &Client{
		conn: &http.Client{
			Jar: jar,
		},
		session: make(http.Header),
	}
}

func NewProxySession(proxy string) *Client {
	proxyURL, _ := url.Parse(proxy)
	transport := &http.Transport{Proxy: http.ProxyURL(proxyURL)}
	jar, _ := cookiejar.New(nil)
	return &Client{
		conn: &http.Client{
			Jar:       jar,
			Transport: transport,
		},
		session: make(http.Header),
	}
}

func (c *Client) Do(method, url string, headers map[string][]string, body io.Reader) (*Response, error) {
	req, err := http.NewRequest(method, url, body)
	if checkError(err) {
		return nil, err
	}
	c.Lock()
	defer c.Unlock()

	if c.session == nil {
		if headers != nil {
			for key, v := range headers {
				for _, val := range v {
					c.session.Set(key, val)
				}
			}
		}

		for key, val := range c.session {
			for _, val := range v {
				req.Header.Set(key, val)
			}
		}

	} else {
		if headers != nil {
			for key, v := range headers {
				for _, val := range v {
					req.Header.Set(key, val)
				}
			}
		}
	}

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

	return &Response{resp}, nil
}

func (c *Client) GET(url string, headers map[string][]string) (*Response, error) {
	return c.Do("GET", url, headers, nil)
}

func (c *Client) POST(url string, headers map[string][]string, body io.Reader) (*Response, error) {
	if headers == nil {
		headers = make(map[string][]string)
	}
	if _, ok := headers["Content-Type"]; !ok {
		headers["Content-Type"] = []string{"application/x-www-form-urlencoded; param=value"}
	}
	return c.Do("POST", url, headers, body)
}
<commit_msg>update<commit_after>/**
 * Author:        Tony.Shao
 * Email:         xiocode@gmail.com
 * Github:        github.com/xiocode
 * File:          client.go
 * Description:   httplib client
 */

package httplib

import (
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"sync"
)

type status int

type Client struct {
	sync.Mutex
	conn    *http.Client
	session http.Header
}

func NewClient() *Client {
	return &Client{
		conn: new(http.Client),
	}
}

func NewProxyClient(proxy string) *Client {
	proxyURL, _ := url.Parse(proxy)
	transport := &http.Transport{Proxy: http.ProxyURL(proxyURL)}
	return &Client{
		conn: &http.Client{
			Transport: transport,
		},
	}
}

func NewSession() *Client {
	jar, _ := cookiejar.New(nil)
	return &Client{
		conn: &http.Client{
			Jar: jar,
		},
		session: make(http.Header),
	}
}

func NewProxySession(proxy string) *Client {
	proxyURL, _ := url.Parse(proxy)
	transport := &http.Transport{Proxy: http.ProxyURL(proxyURL)}
	jar, _ := cookiejar.New(nil)
	return &Client{
		conn: &http.Client{
			Jar:       jar,
			Transport: transport,
		},
		session: make(http.Header),
	}
}

func (c *Client) Do(method, url string, headers map[string][]string, body io.Reader) (*Response, error) {
	req, err := http.NewRequest(method, url, body)
	if checkError(err) {
		return nil, err
	}
	c.Lock()
	defer c.Unlock()

	if c.session == nil {
		if headers != nil {
			for key, v := range headers {
				for _, val := range v {
					c.session.Set(key, val)
				}
			}
		}

		for key, v := range c.session {
			for _, val := range v {
				req.Header.Set(key, val)
			}
		}

	} else {
		if headers != nil {
			for key, v := range headers {
				for _, val := range v {
					req.Header.Set(key, val)
				}
			}
		}
	}

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

	return &Response{resp}, nil
}

func (c *Client) GET(url string, headers map[string][]string) (*Response, error) {
	return c.Do("GET", url, headers, nil)
}

func (c *Client) POST(url string, headers map[string][]string, body io.Reader) (*Response, error) {
	if headers == nil {
		headers = make(map[string][]string)
	}
	if _, ok := headers["Content-Type"]; !ok {
		headers["Content-Type"] = []string{"application/x-www-form-urlencoded; param=value"}
	}
	return c.Do("POST", url, headers, body)
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"path"
	"strings"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
)

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

//  Will return the latest commit hash, waiting on https://github.com/google/go-github/pull/49
func latestGitHubCommit(c *github.Client, repoName string) *github.Repository {
	repo, _, err := c.Repositories.Get("gengo", repoName)
	if err != nil {
		log.Panic(err)
	}
	return repo
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	projects, deployUser := parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, _ := template.ParseFiles("templates/index.html")
	t.Execute(w, map[string]interface{}{"Projects": projects})
}

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name     string
	RepoPath string
	Hosts    []Host
	Branch   string
}

type Project struct {
	Name         string
	GitHubURL    string
	Environments []Environment
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = v.(yaml.Map)["branch"].(yaml.Scalar).String()
		e.RepoPath = v.(yaml.Map)["repo_path"].(yaml.Scalar).String()
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for k, v := range p.(yaml.Map) {
			proj := Project{Name: k, GitHubURL: v.(yaml.Map)["github_url"].(yaml.Scalar).String()}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	for i, project := range projects {
		for j, environment := range project.Environments {
			for k, host := range environment.Hosts {
				host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, environment))
				host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
				projects[i].Environments[j].Hosts[k] = host
			}
			projects[i].Environments[j] = environment
		}
	}
	return projects
}

func main() {
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	fmt.Println(client)

	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	fmt.Println("Running on localhost:" + port)
	http.ListenAndServe(":"+port, r)
}
<commit_msg>Use goroutines for ssh command<commit_after>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"path"
	"strings"
	"sync"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
)

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

//  Will return the latest commit hash, waiting on https://github.com/google/go-github/pull/49
func latestGitHubCommit(c *github.Client, repoName string) *github.Repository {
	repo, _, err := c.Repositories.Get("gengo", repoName)
	if err != nil {
		log.Panic(err)
	}
	return repo
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	projects, deployUser := parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, _ := template.ParseFiles("templates/index.html")
	t.Execute(w, map[string]interface{}{"Projects": projects})
}

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name     string
	RepoPath string
	Hosts    []Host
	Branch   string
}

type Project struct {
	Name         string
	GitHubURL    string
	Environments []Environment
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = v.(yaml.Map)["branch"].(yaml.Scalar).String()
		e.RepoPath = v.(yaml.Map)["repo_path"].(yaml.Scalar).String()
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for k, v := range p.(yaml.Map) {
			proj := Project{Name: k, GitHubURL: v.(yaml.Map)["github_url"].(yaml.Scalar).String()}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func doStuff(wg *sync.WaitGroup, projects []Project, env Environment, host Host, deployUser string, i, j, k int) {
	defer wg.Done()
	host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, env))
	host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
	projects[i].Environments[j].Hosts[k] = host
	fmt.Println("Retrieved commit for", host)
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	// define a wait group to wait for all goroutines to finish
	var wg sync.WaitGroup
	fmt.Println("Retrieving commits...")
	for i, project := range projects {
		fmt.Println("Project", i)
		for j, environment := range project.Environments {
			fmt.Println("Environment", j)
			for k, host := range environment.Hosts {
				fmt.Println("Host", k)
				// start a goroutine for SSH-ing on to the machine
				wg.Add(1)
				go doStuff(&wg, projects, environment, host, deployUser, i, j, k)
			}
			projects[i].Environments[j] = environment
		}
	}
	fmt.Println("All done! Waiting...")
	// wait for goroutines to finish
	wg.Wait()

	fmt.Println("Done!!")
	return projects
}

func main() {
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	fmt.Println(client)

	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	fmt.Println("Running on localhost:" + port)
	http.ListenAndServe(":"+port, r)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2020 The usbtmc developers. All rights reserved.
// Project site: https://github.com/gotmc/usbtmc
// Use of this source code is governed by a MIT-style license that
// can be found in the LICENSE.txt file for the project.

package usbtmc

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"strings"

	"github.com/gotmc/usbtmc/driver"
)

// Device models a USBTMC device, which includes a USB device and the required
// USBTMC attributes and methods.
type Device struct {
	usbDevice       driver.USBDevice
	bTag            byte
	termChar        byte
	termCharEnabled bool
}

// Write creates the appropriate USBMTC header, writes the header and data on
// the bulk out endpoint, and returns the number of bytes written and any
// errors.
func (d *Device) Write(p []byte) (n int, err error) {
	// FIXME(mdr): I need to change this so that I look at the size of the buf
	// being written to see if it can truly fit into one transfer, and if not
	// split it into multiple transfers.
	maxTransferSize := 512
	for pos := 0; pos < len(p); {
		d.bTag = nextbTag(d.bTag)
		thisLen := len(p[pos:])
		if thisLen > maxTransferSize - bulkOutHeaderSize {
			thisLen = maxTransferSize - bulkOutHeaderSize
		}
		header := encodeBulkOutHeader(d.bTag, uint32(thisLen), true)
		data := append(header[:], p[pos:thisLen]...)
		if moduloFour := len(data) % 4; moduloFour > 0 {
			numAlignment := 4 - moduloFour
			alignment := bytes.Repeat([]byte{0x00}, numAlignment)
			data = append(data, alignment...)
		}
		_, err := d.usbDevice.Write(data)
		if err != nil {
			return pos, err
		}
		pos += thisLen
	}
	return len(p), nil
}

// Read creates and sends the header on the bulk out endpoint and then reads
// from the bulk in endpoint per USBTMC standard.
func (d *Device) Read(p []byte) (n int, err error) {
	d.bTag = nextbTag(d.bTag)
	header := encodeMsgInBulkOutHeader(d.bTag, uint32(len(p)), d.termCharEnabled, d.termChar)
	if _, err = d.usbDevice.Write(header[:]); err != nil {
		return 0, err
	}
	pos := 0
	var transfer int
	for ; pos < len(p); {
		var resp int
		var err error
		if pos == 0 {
			resp, transfer, err = d.readRemoveHeader(p[pos:])
		} else {
			resp, err = d.readKeepHeader(p[pos:])
		}
		if err != nil {
			return pos, err
		}
		if resp == 0 {
			break
		}
		pos += resp
		if pos >= transfer {
			break
		}
	}
	return pos, nil
}

func (d *Device) readRemoveHeader(p []byte) (n int, transfer int, err error) {
	// FIXME(mdr): Seems like I shouldn't use 512 as a magic number or as a hard
	// size limit. I should grab the max size of the bulk in endpoint.
	usbtmcHeaderLen := 12
	temp := make([]byte, 512)
	n, err = d.usbDevice.Read(temp)
	// Remove the USBMTC Bulk-IN Header from the data and the number of bytes
	if n < usbtmcHeaderLen {
		return 0, 0, err
	}
	t32 := binary.LittleEndian.Uint32(temp[4:8])
	transfer = int(t32)
	reader := bytes.NewReader(temp)
	_, err = reader.ReadAt(p, int64(usbtmcHeaderLen))

	if err != nil && err != io.EOF {
		return n - usbtmcHeaderLen, transfer, err
	}
	return n - usbtmcHeaderLen, transfer, nil
}

func (d *Device) readKeepHeader(p []byte) (n int, err error) {
	return d.usbDevice.Read(p)
}

// Close closes the underlying USB device.
func (d *Device) Close() error {
	return d.usbDevice.Close()
}

// WriteString writes a string using the underlying USB device. A newline
// terminator is not automatically added.
func (d *Device) WriteString(s string) (n int, err error) {
	return d.Write([]byte(s))
}

// Command sends the SCPI/ASCII command to the underlying USB device. A newline
// character is automatically added to the end of the string.
func (d *Device) Command(format string, a ...interface{}) error {
	cmd := format
	if a != nil {
		cmd = fmt.Sprintf(format, a...)
	}
	_, err := d.WriteString(strings.TrimSpace(cmd) + "\n")
	return err
}

// Query writes the given string to the USBTMC device and returns the returned
// value as a string. A newline character is automatically added to the query
// command sent to the instrument.
func (d *Device) Query(s string) (string, error) {
	err := d.Command(s)
	if err != nil {
		return "", err
	}
	p := make([]byte, 512)
	n, err := d.Read(p)
	if err != nil {
		return "", err
	}
	s = fmt.Sprintf("%s", p[:n])
	p = nil
	return s, nil
}
<commit_msg>Formatting & len fixup<commit_after>// Copyright (c) 2015-2020 The usbtmc developers. All rights reserved.
// Project site: https://github.com/gotmc/usbtmc
// Use of this source code is governed by a MIT-style license that
// can be found in the LICENSE.txt file for the project.

package usbtmc

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"strings"

	"github.com/gotmc/usbtmc/driver"
)

// Device models a USBTMC device, which includes a USB device and the required
// USBTMC attributes and methods.
type Device struct {
	usbDevice       driver.USBDevice
	bTag            byte
	termChar        byte
	termCharEnabled bool
}

// Write creates the appropriate USBMTC header, writes the header and data on
// the bulk out endpoint, and returns the number of bytes written and any
// errors.
func (d *Device) Write(p []byte) (n int, err error) {
	// FIXME(mdr): I need to change this so that I look at the size of the buf
	// being written to see if it can truly fit into one transfer, and if not
	// split it into multiple transfers.
	maxTransferSize := 512
	for pos := 0; pos < len(p); {
		d.bTag = nextbTag(d.bTag)
		thisLen := len(p[pos:])
		if thisLen > maxTransferSize-bulkOutHeaderSize {
			thisLen = maxTransferSize - bulkOutHeaderSize
		}
		header := encodeBulkOutHeader(d.bTag, uint32(thisLen), true)
		data := append(header[:], p[pos:pos+thisLen]...)
		if moduloFour := len(data) % 4; moduloFour > 0 {
			numAlignment := 4 - moduloFour
			alignment := bytes.Repeat([]byte{0x00}, numAlignment)
			data = append(data, alignment...)
		}
		_, err := d.usbDevice.Write(data)
		if err != nil {
			return pos, err
		}
		pos += thisLen
	}
	return len(p), nil
}

// Read creates and sends the header on the bulk out endpoint and then reads
// from the bulk in endpoint per USBTMC standard.
func (d *Device) Read(p []byte) (n int, err error) {
	d.bTag = nextbTag(d.bTag)
	header := encodeMsgInBulkOutHeader(d.bTag, uint32(len(p)), d.termCharEnabled, d.termChar)
	if _, err = d.usbDevice.Write(header[:]); err != nil {
		return 0, err
	}
	pos := 0
	var transfer int
	for pos < len(p) {
		var resp int
		var err error
		if pos == 0 {
			resp, transfer, err = d.readRemoveHeader(p[pos:])
		} else {
			resp, err = d.readKeepHeader(p[pos:])
		}
		if err != nil {
			return pos, err
		}
		if resp == 0 {
			break
		}
		pos += resp
		if pos >= transfer {
			break
		}
	}
	return pos, nil
}

func (d *Device) readRemoveHeader(p []byte) (n int, transfer int, err error) {
	// FIXME(mdr): Seems like I shouldn't use 512 as a magic number or as a hard
	// size limit. I should grab the max size of the bulk in endpoint.
	usbtmcHeaderLen := 12
	temp := make([]byte, 512)
	n, err = d.usbDevice.Read(temp)
	// Remove the USBMTC Bulk-IN Header from the data and the number of bytes
	if n < usbtmcHeaderLen {
		return 0, 0, err
	}
	t32 := binary.LittleEndian.Uint32(temp[4:8])
	transfer = int(t32)
	reader := bytes.NewReader(temp)
	_, err = reader.ReadAt(p, int64(usbtmcHeaderLen))

	if err != nil && err != io.EOF {
		return n - usbtmcHeaderLen, transfer, err
	}
	return n - usbtmcHeaderLen, transfer, nil
}

func (d *Device) readKeepHeader(p []byte) (n int, err error) {
	return d.usbDevice.Read(p)
}

// Close closes the underlying USB device.
func (d *Device) Close() error {
	return d.usbDevice.Close()
}

// WriteString writes a string using the underlying USB device. A newline
// terminator is not automatically added.
func (d *Device) WriteString(s string) (n int, err error) {
	return d.Write([]byte(s))
}

// Command sends the SCPI/ASCII command to the underlying USB device. A newline
// character is automatically added to the end of the string.
func (d *Device) Command(format string, a ...interface{}) error {
	cmd := format
	if a != nil {
		cmd = fmt.Sprintf(format, a...)
	}
	_, err := d.WriteString(strings.TrimSpace(cmd) + "\n")
	return err
}

// Query writes the given string to the USBTMC device and returns the returned
// value as a string. A newline character is automatically added to the query
// command sent to the instrument.
func (d *Device) Query(s string) (string, error) {
	err := d.Command(s)
	if err != nil {
		return "", err
	}
	p := make([]byte, 512)
	n, err := d.Read(p)
	if err != nil {
		return "", err
	}
	s = fmt.Sprintf("%s", p[:n])
	p = nil
	return s, nil
}
<|endoftext|>


<commit_before>package validator

import (
	"errors"
	"net"
	"regexp"
	"strings"
)

var domainRegexp = regexp.MustCompile(`^(?i)[a-z0-9-]+(\.[a-z0-9-]+)+\.?$`)

// IsValidDomain returns true if the domain is valid.
//
// It uses a simple regular expression to check the domain validity.
func IsValidDomain(domain string) bool {
	return domainRegexp.MatchString(domain)
}

var ErrInvalidDomain = errors.New("invalid domain")

// ValidateDomainByResolvingIt queries DNS for the given domain name,
// and returns nil if the the name resolves, or an error.
func ValidateDomainByResolvingIt(domain string) error {
	if !IsValidDomain(domain) {
		return ErrInvalidDomain
	}
	addr, err := net.LookupHost(domain)
	if err != nil {
		return err
	}
	if len(addr) == 0 {
		return ErrInvalidDomain
	}
	return nil
}

// NormalizeEmail returns a normalized domain.
// It returns an empty string if the domain is not valid.
func NormalizeDomain(domain string) string {
	// Trim whitespace.
	domain = strings.TrimSpace(domain)
	// Check validity.
	if !IsValidDomain(domain) {
		return ""
	}
	// Remove trailing dot.
	domain = strings.TrimRight(domain, ".")
	// Convert to lower case.
	domain = strings.ToLower(domain)
	return domain
}
<commit_msg>Fix typo in comment<commit_after>package validator

import (
	"errors"
	"net"
	"regexp"
	"strings"
)

var domainRegexp = regexp.MustCompile(`^(?i)[a-z0-9-]+(\.[a-z0-9-]+)+\.?$`)

// IsValidDomain returns true if the domain is valid.
//
// It uses a simple regular expression to check the domain validity.
func IsValidDomain(domain string) bool {
	return domainRegexp.MatchString(domain)
}

var ErrInvalidDomain = errors.New("invalid domain")

// ValidateDomainByResolvingIt queries DNS for the given domain name,
// and returns nil if the the name resolves, or an error.
func ValidateDomainByResolvingIt(domain string) error {
	if !IsValidDomain(domain) {
		return ErrInvalidDomain
	}
	addr, err := net.LookupHost(domain)
	if err != nil {
		return err
	}
	if len(addr) == 0 {
		return ErrInvalidDomain
	}
	return nil
}

// NormalizeDomain returns a normalized domain.
// It returns an empty string if the domain is not valid.
func NormalizeDomain(domain string) string {
	// Trim whitespace.
	domain = strings.TrimSpace(domain)
	// Check validity.
	if !IsValidDomain(domain) {
		return ""
	}
	// Remove trailing dot.
	domain = strings.TrimRight(domain, ".")
	// Convert to lower case.
	domain = strings.ToLower(domain)
	return domain
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"io/ioutil"
	"log"
	"strconv"
	"time"

	"github.com/brotherlogic/goserver"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	pbs "github.com/brotherlogic/gobuildslave/proto"
)

const (
	intentWait = 1000
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	config  *pb.Config
	serving bool
}

type mainChecker struct{}

func getIP(servertype, servername string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	r, err := registry.ListAllServices(context.Background(), &pbd.Empty{})
	if err != nil {
		return "", -1
	}
	for _, s := range r.Services {
		log.Printf("Does %v = %v and %v = %v?", s.Name, servertype, s.Identifier, servername)
		if s.Name == servertype && s.Identifier == servername {
			return s.Ip, int(s.Port)
		}
	}

	return "", -1
}

func (t *mainChecker) assess(server string) (*pbs.JobList, *pbs.Config) {
	list := &pbs.JobList{}
	conf := &pbs.Config{}

	log.Printf("Assessing server %v", server)
	ip, port := getIP("gobuildslave", server)
	conn, _ := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
	defer conn.Close()

	slave := pbs.NewGoBuildSlaveClient(conn)
	r, err := slave.List(context.Background(), &pbs.Empty{})
	if err != nil {
		log.Printf("Err %v", err)
		return list, conf
	}

	r2, err := slave.GetConfig(context.Background(), &pbs.Empty{})
	if err != nil {
		log.Printf("Err %v", err)
		return list, conf
	}

	return r, r2
}

func runJob(job *pbs.JobSpec, server string) {
	log.Printf("RUNNING: %v on %v", job, server)
	if server != "" {
		ip, port := getIP("gobuildslave", server)
		conn, _ := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
		defer conn.Close()

		slave := pbs.NewGoBuildSlaveClient(conn)
		slave.Run(context.Background(), job)
		log.Printf("RUN COMMAND SENT %v", job)
	}
}

func (t *mainChecker) discover() *pbd.ServiceList {
	ret := &pbd.ServiceList{}

	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	r, err := registry.ListAllServices(context.Background(), &pbd.Empty{})
	log.Printf("DISCOVERED: %v", r)
	if err == nil {
		for _, s := range r.Services {
			ret.Services = append(ret.Services, s)
		}
	}

	return ret
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	// Do nothing
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

func getConfig(c checker) *pb.Config {
	list, _ := getFleetStatus(c)
	config := &pb.Config{}

	for _, jlist := range list {
		for _, job := range jlist.Details {
			found := false
			for _, ij := range config.Intents {
				if job.Spec.Name == ij.Spec.Name {
					ij.Count++
					found = true
				}
			}

			if !found {
				config.Intents = append(config.Intents, &pb.Intent{Spec: &pbs.JobSpec{Name: job.Spec.Name}, Count: 1})
			}
		}
	}

	return config
}

// MatchIntent tries to match the intent with the state of production
func (s Server) MatchIntent() {
	for s.serving {
		time.Sleep(intentWait)

		state := getConfig(&mainChecker{})
		diff := configDiff(s.config, state)
		joblist := runJobs(diff)
		log.Printf("FOUND %v from %v and %v", joblist, state, s.config)
		for _, job := range joblist {
			runJob(job, chooseServer(job, &mainChecker{}))
		}
	}
}

func main() {
	config, err := loadConfig("config.pb")
	if err != nil {
		log.Fatalf("Fatal loading of config: %v", err)
	}

	var sync = flag.Bool("once", false, "One pass intent match")
	s := Server{&goserver.GoServer{}, config, true}

	var quiet = flag.Bool("quiet", true, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	if *sync {
		s.MatchIntent()
	} else {
		s.Register = s
		s.PrepServer()
		s.RegisterServer("gobuildmaster", false)
		s.RegisterServingTask(s.MatchIntent)
		s.Serve()
	}
}
<commit_msg>Added mote comman<commit_after>package main

import (
	"context"
	"flag"
	"io/ioutil"
	"log"
	"strconv"
	"time"

	"github.com/brotherlogic/goserver"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	pbs "github.com/brotherlogic/gobuildslave/proto"
)

const (
	intentWait = 1000
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	config  *pb.Config
	serving bool
}

type mainChecker struct{}

func getIP(servertype, servername string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	r, err := registry.ListAllServices(context.Background(), &pbd.Empty{})
	if err != nil {
		return "", -1
	}
	for _, s := range r.Services {
		log.Printf("Does %v = %v and %v = %v?", s.Name, servertype, s.Identifier, servername)
		if s.Name == servertype && s.Identifier == servername {
			return s.Ip, int(s.Port)
		}
	}

	return "", -1
}

func (t *mainChecker) assess(server string) (*pbs.JobList, *pbs.Config) {
	list := &pbs.JobList{}
	conf := &pbs.Config{}

	log.Printf("Assessing server %v", server)
	ip, port := getIP("gobuildslave", server)
	conn, _ := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
	defer conn.Close()

	slave := pbs.NewGoBuildSlaveClient(conn)
	r, err := slave.List(context.Background(), &pbs.Empty{})
	if err != nil {
		log.Printf("Err %v", err)
		return list, conf
	}

	r2, err := slave.GetConfig(context.Background(), &pbs.Empty{})
	if err != nil {
		log.Printf("Err %v", err)
		return list, conf
	}

	return r, r2
}

func runJob(job *pbs.JobSpec, server string) {
	log.Printf("RUNNING: %v on %v", job, server)
	if server != "" {
		ip, port := getIP("gobuildslave", server)
		conn, _ := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
		defer conn.Close()

		slave := pbs.NewGoBuildSlaveClient(conn)
		slave.Run(context.Background(), job)
		log.Printf("RUN COMMAND SENT %v", job)
	}
}

func (t *mainChecker) discover() *pbd.ServiceList {
	ret := &pbd.ServiceList{}

	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	r, err := registry.ListAllServices(context.Background(), &pbd.Empty{})
	log.Printf("DISCOVERED: %v", r)
	if err == nil {
		for _, s := range r.Services {
			ret.Services = append(ret.Services, s)
		}
	}

	return ret
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	// Do nothing
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

// Mote promotes/demotes this server
func (s Server) Mote(master bool) error {
	return nil
}

func getConfig(c checker) *pb.Config {
	list, _ := getFleetStatus(c)
	config := &pb.Config{}

	for _, jlist := range list {
		for _, job := range jlist.Details {
			found := false
			for _, ij := range config.Intents {
				if job.Spec.Name == ij.Spec.Name {
					ij.Count++
					found = true
				}
			}

			if !found {
				config.Intents = append(config.Intents, &pb.Intent{Spec: &pbs.JobSpec{Name: job.Spec.Name}, Count: 1})
			}
		}
	}

	return config
}

// MatchIntent tries to match the intent with the state of production
func (s Server) MatchIntent() {
	for s.serving {
		time.Sleep(intentWait)

		state := getConfig(&mainChecker{})
		diff := configDiff(s.config, state)
		joblist := runJobs(diff)
		log.Printf("FOUND %v from %v and %v", joblist, state, s.config)
		for _, job := range joblist {
			runJob(job, chooseServer(job, &mainChecker{}))
		}
	}
}

func main() {
	config, err := loadConfig("config.pb")
	if err != nil {
		log.Fatalf("Fatal loading of config: %v", err)
	}

	var sync = flag.Bool("once", false, "One pass intent match")
	s := Server{&goserver.GoServer{}, config, true}

	var quiet = flag.Bool("quiet", true, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	if *sync {
		s.MatchIntent()
	} else {
		s.Register = s
		s.PrepServer()
		s.RegisterServer("gobuildmaster", false)
		s.RegisterServingTask(s.MatchIntent)
		s.Serve()
	}
}
<|endoftext|>