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

package github

import (
	"fmt"
	"log"
	"strings"

	"github.com/google/go-github/github"
	"go.mozilla.org/userplex/modules"
	"golang.org/x/oauth2"
)

func init() {
	modules.Register("github", new(module))
}

type module struct{}

func (m *module) NewRun(c modules.Configuration) modules.Runner {
	r := new(run)
	r.Conf = c
	return r
}

type run struct {
	Conf         modules.Configuration
	p            parameters
	c            credentials
	ghclient     *github.Client
	githubToLdap map[string]string
}

type organization struct {
	Name  string
	Teams []string
}

type parameters struct {
	Organization     organization
	UserplexTeamName string
	Enforce2FA       bool
}

type credentials struct {
	OAuthToken string `yaml:"oauthtoken"`
}

func (r *run) Run() (err error) {
	var resp *github.Response

	err = r.Conf.GetParameters(&r.p)
	if err != nil {
		return
	}
	err = r.Conf.GetCredentials(&r.c)
	if err != nil {
		return
	}

	if r.p.UserplexTeamName == "" {
		return fmt.Errorf("[error] github: UserplexTeamName is not set")
	}

	ts := oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: r.c.OAuthToken},
	)
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	r.ghclient = github.NewClient(tc)

	// initialize mapping of github usernames -> ldap uids
	r.githubToLdap = make(map[string]string)
	ldapers := r.getLdapers()

	// get all members for the organization
	// name -> bool
	membersMap := r.getOrgMembersMap(r.p.Organization, "all")
	if r.Conf.Debug {
		log.Printf("[debug] github: found %d users for organization %s", len(membersMap), r.p.Organization.Name)
	}

	// get all teams for the organization
	// name -> team
	teamsMap := r.getOrgTeamsMap(r.p.Organization)
	if r.Conf.Debug {
		log.Printf("[debug] github: found %d teams for organization %s", len(teamsMap), r.p.Organization.Name)
	}

	teamMembersMap := make(map[string]map[string]bool)
	for _, team := range teamsMap {
		teamMembersMap[*team.Name] = make(map[string]bool)
		teamMembersMap[*team.Name] = r.getTeamMembersMap(team)
	}

	if _, ok := teamsMap[r.p.UserplexTeamName]; !ok {
		return fmt.Errorf("[error] github: could not find UserplexTeam \"%s\" for organization %s", r.p.UserplexTeamName, r.p.Organization.Name)
	}
	userplexedUsers := teamMembersMap[r.p.UserplexTeamName]

	var no2fa map[string]bool
	if r.p.Enforce2FA {
		no2fa = r.getOrgMembersMap(r.p.Organization, "2fa_disabled")
		log.Printf("[info] github: organization %s has %d total members and %d with 2fa disabled. %.2f%% have 2fa enabled.",
			r.p.Organization.Name, len(membersMap), len(no2fa), 100-100*float64(len(no2fa))/float64(len(membersMap)))
	}

	// member or admin
	membershipType := "member"

	countAdded := 0
	for user := range ldapers {
		// set to true to indicate that user in github has ldap match
		membersMap[user] = true

		userWasAddedToTeam := false
		// teams in config
		for _, teamName := range r.p.Organization.Teams {
			// if the team in config doesn't exist on github
			if team, ok := teamsMap[teamName]; !ok {
				return fmt.Errorf("[error] github: could not find team %s for organization %s", team, r.p.Organization.Name)
			} else {
				// if the user is already in the team, skip adding them
				if _, ok := teamMembersMap[teamName][user]; ok {
					continue
				}
				// user not in team, add them
				if r.Conf.ApplyChanges && r.Conf.Create {
					// add user to team
					_, resp, err = r.ghclient.Organizations.AddTeamMembership(*team.ID, user, &github.OrganizationAddTeamMembershipOptions{
						Role: membershipType,
					})
					if err != nil || resp.StatusCode != 200 {
						return fmt.Errorf("[error] github: could not add user %s to %s: %s, error: %v with status %s", user, r.p.Organization.Name, *team.Name, err, resp.Status)
					}
				}
				if r.Conf.Create {
					userWasAddedToTeam = true
				}
			}
		}
		if userWasAddedToTeam {
			if !r.Conf.ApplyChanges {
				log.Printf("[dryrun] github: would have added %s to GitHub organization %s and teams %v", user, r.p.Organization.Name, r.p.Organization.Teams)
			}
			countAdded++
			r.notify(user, fmt.Sprintf("Userplex added %s to GitHub organization %s and teams %v", user, r.p.Organization.Name, r.p.Organization.Teams))
		}
	}

	countRemoved := 0
	countSkipped := 0
	for member := range membersMap {
		var ldapUsername, ldapUsernameString string
		member = strings.ToLower(member)
		if _, ok := r.githubToLdap[member]; ok {
			ldapUsername = r.githubToLdap[member]
			ldapUsernameString = ldapUsername + " / "
		}

		var userTeams []string
		// icky iterating over all these teams
		for _, team := range teamsMap {
			if _, ok := teamMembersMap[*team.Name][member]; ok {
				userTeams = append(userTeams, *team.Name)
			}
		}

		// if the member does not have 2fa
		_, no2fa := no2fa[member]

		// if the user is in ldap
		inLdap, _ := membersMap[member]

		// if the member is not in the userplex team
		_, isUserplexed := userplexedUsers[member]

		shouldDelete := false
		if !inLdap {
			if r.Conf.Debug {
				log.Printf("[debug] github: user %s%s is not in ldap groups %s but is a member of github organization %s and teams %v", ldapUsernameString, member, r.Conf.LdapGroups, r.p.Organization.Name, userTeams)
			}
			shouldDelete = true
		}

		if r.p.Enforce2FA && no2fa {
			log.Printf("[info] github: user %s%s does not have 2FA enabled and is a member of github organization %s and teams %v", ldapUsernameString, member, r.p.Organization.Name, userTeams)
			shouldDelete = true
		}

		if shouldDelete && r.Conf.Delete {
			// not in UserplexTeam -> skip
			if !isUserplexed {
				log.Printf("[info] github: would have removed member %s in organization %s, but skipped because they are not in UserplexTeam %q", member, r.p.Organization.Name, r.p.UserplexTeamName)
				countSkipped++
				continue
			}

			if !r.Conf.ApplyChanges {
				log.Printf("[dryrun] github: Userplex would have removed %s%s from GitHub organization %s", ldapUsernameString, member, r.p.Organization.Name)
			} else {
				// applying changes, user is userplexed -> remove them
				resp, err = r.ghclient.Organizations.RemoveOrgMembership(r.p.Organization.Name, member)
				if err != nil || resp.StatusCode != 200 {
					log.Printf("[error] github: could not remove user %s from %s, error: %v with status %s", member, r.p.Organization.Name, err, resp.Status)
				}
			}

			// update count and send notification here regardless of ApplyChanges
			countRemoved++
			r.notify(member, fmt.Sprintf("Userplex removed %s to GitHub organization %s", member, r.p.Organization.Name))
		}
	}

	log.Printf("[info] github %q: summary added=%d, removed=%d, skipped=%d",
		r.p.Organization.Name, countAdded, countRemoved, countSkipped)

	return nil
}

func (r *run) getOrgMembersMap(org organization, filter string) (membersMap map[string]bool) {
	membersMap = make(map[string]bool)
	opt := &github.ListMembersOptions{
		Filter:      filter,
		ListOptions: github.ListOptions{PerPage: 100},
	}
	for {
		members, resp, err := r.ghclient.Organizations.ListMembers(org.Name, opt)
		if err != nil || resp.StatusCode != 200 {
			log.Printf("[error] github: could not list members for organization %s, error: %v with status %s", org, err, resp.Status)
			return
		}
		for _, member := range members {
			memberName := strings.ToLower(*member.Login)
			membersMap[memberName] = false
		}
		if resp.NextPage == 0 {
			break
		}
		opt.ListOptions.Page = resp.NextPage
	}
	return membersMap
}

func (r *run) getTeamMembersMap(team *github.Team) (membersMap map[string]bool) {
	membersMap = make(map[string]bool)
	opt := &github.OrganizationListTeamMembersOptions{
		ListOptions: github.ListOptions{
			PerPage: 100,
		},
	}
	for {
		members, resp, err := r.ghclient.Organizations.ListTeamMembers(*team.ID, opt)
		if err != nil || resp.StatusCode != 200 {
			log.Printf("[error] github: could not list members for organization %s, error: %v with status %s", team, err, resp.Status)
			return
		}
		for _, member := range members {
			memberName := strings.ToLower(*member.Login)
			membersMap[memberName] = false
		}
		if resp.NextPage == 0 {
			break
		}
		opt.ListOptions.Page = resp.NextPage
	}
	return membersMap
}

func (r *run) getOrgTeamsMap(org organization) (teamsMap map[string]*github.Team) {
	teamsMap = make(map[string]*github.Team)
	opt := &github.ListOptions{
		PerPage: 100,
	}
	for {
		teams, resp, err := r.ghclient.Organizations.ListTeams(org.Name, opt)
		if err != nil || resp.StatusCode != 200 {
			log.Printf("[error] github: could not list teams for organization %s, error: %v", org.Name, err)
			return
		}
		for _, team := range teams {
			teamName := strings.ToLower(*team.Name)
			teamsMap[teamName] = team
		}
		if resp.NextPage == 0 {
			break
		}
		opt.Page = resp.NextPage
	}
	return teamsMap
}

func (r *run) notify(user string, body string) (err error) {
	rcpt := r.Conf.Notify.Recipient
	if rcpt == "{ldap:mail}" {
		if _, ok := r.githubToLdap[user]; ok {
			user = r.githubToLdap[user]
		}
		rcpt, err = r.Conf.LdapCli.GetUserEmailByUid(user)
		if err != nil {
			return fmt.Errorf("[error] github: couldn't find email of user %q in ldap, notification not sent: %v", user, err)
		}
	}
	r.Conf.Notify.Channel <- modules.Notification{
		Module:      "github",
		Recipient:   rcpt,
		Mode:        r.Conf.Notify.Mode,
		MustEncrypt: false,
		Body:        []byte(body),
	}
	return
}

func (r *run) getLdapers() map[string]bool {
	ldapers := make(map[string]bool)
	users, err := r.Conf.LdapCli.GetEnabledUsersInGroups(r.Conf.LdapGroups)
	if err != nil {
		return ldapers
	}
	for _, user := range users {
		shortdn := strings.Split(user, ",")[0]
		uid, err := r.Conf.LdapCli.GetUserId(shortdn)
		if err != nil {
			log.Printf("[error] github: ldap query GetUserId(%q) failed with error %v", uid, err)
			continue
		}

		github, getGithubAccountErr := r.Conf.LdapCli.GetUserGithubByUID(uid)
		if getGithubAccountErr != nil {
			// this error can happen if the user does not have a githubProfile in ldap
			if r.Conf.Debug {
				log.Printf("[debug] github: ldap query GetUserGithubByUID(%q) failed with error %q", uid, getGithubAccountErr.Error())
			}
		} else {
			// has a githubProfile in LDAP
			r.githubToLdap[github] = uid
			uid = github
		}
		uid = strings.ToLower(uid)
		ldapers[uid] = false
	}
	return ldapers
}
<commit_msg>github module: ensure getting user works, skip errored users<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Miles Crabill <mcrabill@mozilla.com>

package github

import (
	"fmt"
	"log"
	"strings"

	"github.com/google/go-github/github"
	"go.mozilla.org/userplex/modules"
	"golang.org/x/oauth2"
)

func init() {
	modules.Register("github", new(module))
}

type module struct{}

func (m *module) NewRun(c modules.Configuration) modules.Runner {
	r := new(run)
	r.Conf = c
	return r
}

type run struct {
	Conf         modules.Configuration
	p            parameters
	c            credentials
	ghclient     *github.Client
	githubToLdap map[string]string
}

type organization struct {
	Name  string
	Teams []string
}

type parameters struct {
	Organization     organization
	UserplexTeamName string
	Enforce2FA       bool
}

type credentials struct {
	OAuthToken string `yaml:"oauthtoken"`
}

func (r *run) Run() (err error) {
	var resp *github.Response

	err = r.Conf.GetParameters(&r.p)
	if err != nil {
		return
	}
	err = r.Conf.GetCredentials(&r.c)
	if err != nil {
		return
	}

	if r.p.UserplexTeamName == "" {
		return fmt.Errorf("[error] github: UserplexTeamName is not set")
	}

	ts := oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: r.c.OAuthToken},
	)
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	r.ghclient = github.NewClient(tc)

	// initialize mapping of github usernames -> ldap uids
	r.githubToLdap = make(map[string]string)
	ldapers := r.getLdapers()

	// get all members for the organization
	// name -> bool
	membersMap := r.getOrgMembersMap(r.p.Organization, "all")
	if r.Conf.Debug {
		log.Printf("[debug] github: found %d users for organization %s", len(membersMap), r.p.Organization.Name)
	}

	// get all teams for the organization
	// name -> team
	teamsMap := r.getOrgTeamsMap(r.p.Organization)
	if r.Conf.Debug {
		log.Printf("[debug] github: found %d teams for organization %s", len(teamsMap), r.p.Organization.Name)
	}

	teamMembersMap := make(map[string]map[string]bool)
	for _, team := range teamsMap {
		teamMembersMap[*team.Name] = make(map[string]bool)
		teamMembersMap[*team.Name] = r.getTeamMembersMap(team)
	}

	if _, ok := teamsMap[r.p.UserplexTeamName]; !ok {
		return fmt.Errorf("[error] github: could not find UserplexTeam \"%s\" for organization %s", r.p.UserplexTeamName, r.p.Organization.Name)
	}
	userplexedUsers := teamMembersMap[r.p.UserplexTeamName]

	var no2fa map[string]bool
	if r.p.Enforce2FA {
		no2fa = r.getOrgMembersMap(r.p.Organization, "2fa_disabled")
		log.Printf("[info] github: organization %s has %d total members and %d with 2fa disabled. %.2f%% have 2fa enabled.",
			r.p.Organization.Name, len(membersMap), len(no2fa), 100-100*float64(len(no2fa))/float64(len(membersMap)))
	}

	// member or admin
	membershipType := "member"

	countAdded := 0

userloop:
	for user := range ldapers {
		_, resp, err = r.ghclient.Users.Get(user)
		if err != nil || resp.StatusCode != 200 {
			log.Printf("[error] github: could not get user %s: %s, error: %v with status %s", user, r.p.Organization.Name, err, resp.Status)
			continue
		}

		// set to true to indicate that user in github has ldap match
		membersMap[user] = true

		userWasAddedToTeam := false
		// teams in config
		for _, teamName := range r.p.Organization.Teams {
			// if the team in config doesn't exist on github
			if team, ok := teamsMap[teamName]; !ok {
				return fmt.Errorf("[error] github: could not find team %s for organization %s", team, r.p.Organization.Name)
			} else {
				// if the user is already in the team, skip adding them
				if _, ok := teamMembersMap[teamName][user]; ok {
					continue
				}
				// user not in team, add them
				if r.Conf.ApplyChanges && r.Conf.Create {
					// add user to team
					_, resp, err = r.ghclient.Organizations.AddTeamMembership(*team.ID, user, &github.OrganizationAddTeamMembershipOptions{
						Role: membershipType,
					})
					if err != nil || resp.StatusCode != 200 {
						log.Printf("[error] github: could not add user %s to %s: %s, error: %v with status %s", user, r.p.Organization.Name, *team.Name, err, resp.Status)
						continue userloop
					}
				}
				if r.Conf.Create {
					userWasAddedToTeam = true
				}
			}
		}
		if userWasAddedToTeam {
			if !r.Conf.ApplyChanges {
				log.Printf("[dryrun] github: would have added %s to GitHub organization %s and teams %v", user, r.p.Organization.Name, r.p.Organization.Teams)
			}
			countAdded++
			r.notify(user, fmt.Sprintf("Userplex added %s to GitHub organization %s and teams %v", user, r.p.Organization.Name, r.p.Organization.Teams))
		}
	}

	countRemoved := 0
	countSkipped := 0
	for member := range membersMap {
		var ldapUsername, ldapUsernameString string
		member = strings.ToLower(member)
		if _, ok := r.githubToLdap[member]; ok {
			ldapUsername = r.githubToLdap[member]
			ldapUsernameString = ldapUsername + " / "
		}

		var userTeams []string
		// icky iterating over all these teams
		for _, team := range teamsMap {
			if _, ok := teamMembersMap[*team.Name][member]; ok {
				userTeams = append(userTeams, *team.Name)
			}
		}

		// if the member does not have 2fa
		_, no2fa := no2fa[member]

		// if the user is in ldap
		inLdap, _ := membersMap[member]

		// if the member is not in the userplex team
		_, isUserplexed := userplexedUsers[member]

		shouldDelete := false
		if !inLdap {
			if r.Conf.Debug {
				log.Printf("[debug] github: user %s%s is not in ldap groups %s but is a member of github organization %s and teams %v", ldapUsernameString, member, r.Conf.LdapGroups, r.p.Organization.Name, userTeams)
			}
			shouldDelete = true
		}

		if r.p.Enforce2FA && no2fa {
			log.Printf("[info] github: user %s%s does not have 2FA enabled and is a member of github organization %s and teams %v", ldapUsernameString, member, r.p.Organization.Name, userTeams)
			shouldDelete = true
		}

		if shouldDelete && r.Conf.Delete {
			// not in UserplexTeam -> skip
			if !isUserplexed {
				log.Printf("[info] github: would have removed member %s in organization %s, but skipped because they are not in UserplexTeam %q", member, r.p.Organization.Name, r.p.UserplexTeamName)
				countSkipped++
				continue
			}

			if !r.Conf.ApplyChanges {
				log.Printf("[dryrun] github: Userplex would have removed %s%s from GitHub organization %s", ldapUsernameString, member, r.p.Organization.Name)
			} else {
				// applying changes, user is userplexed -> remove them
				resp, err = r.ghclient.Organizations.RemoveOrgMembership(r.p.Organization.Name, member)
				if err != nil || resp.StatusCode != 200 {
					log.Printf("[error] github: could not remove user %s from %s, error: %v with status %s", member, r.p.Organization.Name, err, resp.Status)
				}
			}

			// update count and send notification here regardless of ApplyChanges
			countRemoved++
			r.notify(member, fmt.Sprintf("Userplex removed %s to GitHub organization %s", member, r.p.Organization.Name))
		}
	}

	log.Printf("[info] github %q: summary added=%d, removed=%d, skipped=%d",
		r.p.Organization.Name, countAdded, countRemoved, countSkipped)

	return nil
}

func (r *run) getOrgMembersMap(org organization, filter string) (membersMap map[string]bool) {
	membersMap = make(map[string]bool)
	opt := &github.ListMembersOptions{
		Filter:      filter,
		ListOptions: github.ListOptions{PerPage: 100},
	}
	for {
		members, resp, err := r.ghclient.Organizations.ListMembers(org.Name, opt)
		if err != nil || resp.StatusCode != 200 {
			log.Printf("[error] github: could not list members for organization %s, error: %v with status %s", org, err, resp.Status)
			return
		}
		for _, member := range members {
			memberName := strings.ToLower(*member.Login)
			membersMap[memberName] = false
		}
		if resp.NextPage == 0 {
			break
		}
		opt.ListOptions.Page = resp.NextPage
	}
	return membersMap
}

func (r *run) getTeamMembersMap(team *github.Team) (membersMap map[string]bool) {
	membersMap = make(map[string]bool)
	opt := &github.OrganizationListTeamMembersOptions{
		ListOptions: github.ListOptions{
			PerPage: 100,
		},
	}
	for {
		members, resp, err := r.ghclient.Organizations.ListTeamMembers(*team.ID, opt)
		if err != nil || resp.StatusCode != 200 {
			log.Printf("[error] github: could not list members for organization %s, error: %v with status %s", team, err, resp.Status)
			return
		}
		for _, member := range members {
			memberName := strings.ToLower(*member.Login)
			membersMap[memberName] = false
		}
		if resp.NextPage == 0 {
			break
		}
		opt.ListOptions.Page = resp.NextPage
	}
	return membersMap
}

func (r *run) getOrgTeamsMap(org organization) (teamsMap map[string]*github.Team) {
	teamsMap = make(map[string]*github.Team)
	opt := &github.ListOptions{
		PerPage: 100,
	}
	for {
		teams, resp, err := r.ghclient.Organizations.ListTeams(org.Name, opt)
		if err != nil || resp.StatusCode != 200 {
			log.Printf("[error] github: could not list teams for organization %s, error: %v", org.Name, err)
			return
		}
		for _, team := range teams {
			teamName := strings.ToLower(*team.Name)
			teamsMap[teamName] = team
		}
		if resp.NextPage == 0 {
			break
		}
		opt.Page = resp.NextPage
	}
	return teamsMap
}

func (r *run) notify(user string, body string) (err error) {
	rcpt := r.Conf.Notify.Recipient
	if rcpt == "{ldap:mail}" {
		if _, ok := r.githubToLdap[user]; ok {
			user = r.githubToLdap[user]
		}
		rcpt, err = r.Conf.LdapCli.GetUserEmailByUid(user)
		if err != nil {
			return fmt.Errorf("[error] github: couldn't find email of user %q in ldap, notification not sent: %v", user, err)
		}
	}
	r.Conf.Notify.Channel <- modules.Notification{
		Module:      "github",
		Recipient:   rcpt,
		Mode:        r.Conf.Notify.Mode,
		MustEncrypt: false,
		Body:        []byte(body),
	}
	return
}

func (r *run) getLdapers() map[string]bool {
	ldapers := make(map[string]bool)
	users, err := r.Conf.LdapCli.GetEnabledUsersInGroups(r.Conf.LdapGroups)
	if err != nil {
		return ldapers
	}
	for _, user := range users {
		shortdn := strings.Split(user, ",")[0]
		uid, err := r.Conf.LdapCli.GetUserId(shortdn)
		if err != nil {
			log.Printf("[error] github: ldap query GetUserId(%q) failed with error %v", uid, err)
			continue
		}

		github, getGithubAccountErr := r.Conf.LdapCli.GetUserGithubByUID(uid)
		if getGithubAccountErr != nil {
			// this error can happen if the user does not have a githubProfile in ldap
			if r.Conf.Debug {
				log.Printf("[debug] github: ldap query GetUserGithubByUID(%q) failed with error %q", uid, getGithubAccountErr.Error())
			}
		} else {
			// has a githubProfile in LDAP
			r.githubToLdap[github] = uid
			uid = github
		}
		uid = strings.ToLower(uid)
		ldapers[uid] = false
	}
	return ldapers
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package mailer

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/mail"
	"net/smtp"
	"os"
	"strings"

	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/setting"
)

type loginAuth struct {
	username, password string
}

// SMTP AUTH LOGIN Auth Handler
func LoginAuth(username, password string) smtp.Auth {
	return &loginAuth{username, password}
}

func (a *loginAuth) Start(server *smtp.ServerInfo) (string, []byte, error) {
	return "LOGIN", []byte{}, nil
}

func (a *loginAuth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		switch string(fromServer) {
		case "Username:":
			return []byte(a.username), nil
		case "Password:":
			return []byte(a.password), nil
		default:
			return nil, fmt.Errorf("unknwon fromServer: %s", string(fromServer))
		}
	}
	return nil, nil
}

type Message struct {
	To      []string
	From    string
	Subject string
	Body    string
	Type    string
	Massive bool
	Info    string
}

// create mail content
func (m Message) Content() string {
	// set mail type
	contentType := "text/plain; charset=UTF-8"
	if m.Type == "html" {
		contentType = "text/html; charset=UTF-8"
	}

	// create mail content
	content := "From: " + m.From + "\r\nSubject: " + m.Subject + "\r\nContent-Type: " + contentType + "\r\n\r\n" + m.Body
	return content
}

var mailQueue chan *Message

func NewMailerContext() {
	mailQueue = make(chan *Message, setting.Cfg.Section("mailer").Key("SEND_BUFFER_LEN").MustInt(10))
	go processMailQueue()
}

func processMailQueue() {
	for {
		select {
		case msg := <-mailQueue:
			num, err := Send(msg)
			tos := strings.Join(msg.To, "; ")
			info := ""
			if err != nil {
				if len(msg.Info) > 0 {
					info = ", info: " + msg.Info
				}
				log.Error(4, fmt.Sprintf("Async sent email %d succeed, not send emails: %s%s err: %s", num, tos, info, err))
			} else {
				log.Trace(fmt.Sprintf("Async sent email %d succeed, sent emails: %s%s", num, tos, info))
			}
		}
	}
}

// sendMail allows mail with self-signed certificates.
func sendMail(settings *setting.Mailer, recipients []string, msgContent []byte) error {
	host, port, err := net.SplitHostPort(settings.Host)
	if err != nil {
		return err
	}

	tlsconfig := &tls.Config{
		InsecureSkipVerify: settings.SkipVerify,
		ServerName:         host,
	}

	if settings.UseCertificate {
		cert, err := tls.LoadX509KeyPair(settings.CertFile, settings.KeyFile)
		if err != nil {
			return err
		}
		tlsconfig.Certificates = []tls.Certificate{cert}
	}

	conn, err := net.Dial("tcp", net.JoinHostPort(host, port))
	if err != nil {
		return err
	}
	defer conn.Close()

	isSecureConn := false
	// Start TLS directly if the port ends with 465 (SMTPS protocol)
	if strings.HasSuffix(port, "465") {
		conn = tls.Client(conn, tlsconfig)
		isSecureConn = true
	}

	client, err := smtp.NewClient(conn, host)
	if err != nil {
		return err
	}

	if !setting.MailService.DisableHelo {
		hostname := setting.MailService.HeloHostname
		if len(hostname) == 0 {
			hostname, err = os.Hostname()
			if err != nil {
				return err
			}
		}

		if err = client.Hello(hostname); err != nil {
			return err
		}
	}

	// If not using SMTPS, alway use STARTTLS if available
	hasStartTLS, _ := client.Extension("STARTTLS")
	if !isSecureConn && hasStartTLS {
		if err = client.StartTLS(tlsconfig); err != nil {
			return err
		}
	}

	canAuth, options := client.Extension("AUTH")

	if canAuth && len(settings.User) > 0 {
		var auth smtp.Auth

		if strings.Contains(options, "CRAM-MD5") {
			auth = smtp.CRAMMD5Auth(settings.User, settings.Passwd)
		} else if strings.Contains(options, "PLAIN") {
			auth = smtp.PlainAuth("", settings.User, settings.Passwd, host)
		} else if strings.Contains(options, "LOGIN") {
			// Patch for AUTH LOGIN
			auth = LoginAuth(settings.User, settings.Passwd)
		}

		if auth != nil {
			if err = client.Auth(auth); err != nil {
				return err
			}
		}
	}

	if fromAddress, err := mail.ParseAddress(settings.From); err != nil {
		return err
	} else {
		if err = client.Mail(fromAddress.Address); err != nil {
			return err
		}
	}

	for _, rec := range recipients {
		if err = client.Rcpt(rec); err != nil {
			return err
		}
	}

	w, err := client.Data()
	if err != nil {
		return err
	}
	if _, err = w.Write([]byte(msgContent)); err != nil {
		return err
	}

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

	return client.Quit()
}

// Direct Send mail message
func Send(msg *Message) (int, error) {
	log.Trace("Sending mails to: %s", strings.Join(msg.To, "; "))

	// get message body
	content := msg.Content()

	if len(msg.To) == 0 {
		return 0, fmt.Errorf("empty receive emails")
	} else if len(msg.Body) == 0 {
		return 0, fmt.Errorf("empty email body")
	}

	if msg.Massive {
		// send mail to multiple emails one by one
		num := 0
		for _, to := range msg.To {
			body := []byte("To: " + to + "\r\n" + content)
			err := sendMail(setting.MailService, []string{to}, body)
			if err != nil {
				return num, err
			}
			num++
		}
		return num, nil
	} else {
		body := []byte("To: " + strings.Join(msg.To, ";") + "\r\n" + content)

		// send to multiple emails in one message
		err := sendMail(setting.MailService, msg.To, body)
		if err != nil {
			return 0, err
		} else {
			return 1, nil
		}
	}
}

// Async Send mail message
func SendAsync(msg *Message) {
	go func() {
		mailQueue <- msg
	}()
}

// Create html mail message
func NewHtmlMessage(To []string, From, Subject, Body string) Message {
	return Message{
		To:      To,
		From:    From,
		Subject: Subject,
		Body:    Body,
		Type:    "html",
	}
}
<commit_msg>implemented #1528<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package mailer

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/mail"
	"net/smtp"
	"os"
	"strings"

	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/setting"
)

type loginAuth struct {
	username, password string
}

// SMTP AUTH LOGIN Auth Handler
func LoginAuth(username, password string) smtp.Auth {
	return &loginAuth{username, password}
}

func (a *loginAuth) Start(server *smtp.ServerInfo) (string, []byte, error) {
	return "LOGIN", []byte{}, nil
}

func (a *loginAuth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		switch string(fromServer) {
		case "Username:":
			return []byte(a.username), nil
		case "Password:":
			return []byte(a.password), nil
		default:
			return nil, fmt.Errorf("unknwon fromServer: %s", string(fromServer))
		}
	}
	return nil, nil
}

type Message struct {
	To      []string
	From    string
	Subject string
	Reply-To string
	Body    string
	Type    string
	Massive bool
	Info    string
}

// create mail content
func (m Message) Content() string {
	// set mail type
	contentType := "text/plain; charset=UTF-8"
	if m.Type == "html" {
		contentType = "text/html; charset=UTF-8"
	}

	// create mail content
	content := "From: " + m.From + "\r\nReply-To: " + m.Reply-To + "\r\nSubject: " + m.Subject + "\r\nContent-Type: " + contentType + "\r\n\r\n" + m.Body
	return content
}

var mailQueue chan *Message

func NewMailerContext() {
	mailQueue = make(chan *Message, setting.Cfg.Section("mailer").Key("SEND_BUFFER_LEN").MustInt(10))
	go processMailQueue()
}

func processMailQueue() {
	for {
		select {
		case msg := <-mailQueue:
			num, err := Send(msg)
			tos := strings.Join(msg.To, "; ")
			info := ""
			if err != nil {
				if len(msg.Info) > 0 {
					info = ", info: " + msg.Info
				}
				log.Error(4, fmt.Sprintf("Async sent email %d succeed, not send emails: %s%s err: %s", num, tos, info, err))
			} else {
				log.Trace(fmt.Sprintf("Async sent email %d succeed, sent emails: %s%s", num, tos, info))
			}
		}
	}
}

// sendMail allows mail with self-signed certificates.
func sendMail(settings *setting.Mailer, recipients []string, msgContent []byte) error {
	host, port, err := net.SplitHostPort(settings.Host)
	if err != nil {
		return err
	}

	tlsconfig := &tls.Config{
		InsecureSkipVerify: settings.SkipVerify,
		ServerName:         host,
	}

	if settings.UseCertificate {
		cert, err := tls.LoadX509KeyPair(settings.CertFile, settings.KeyFile)
		if err != nil {
			return err
		}
		tlsconfig.Certificates = []tls.Certificate{cert}
	}

	conn, err := net.Dial("tcp", net.JoinHostPort(host, port))
	if err != nil {
		return err
	}
	defer conn.Close()

	isSecureConn := false
	// Start TLS directly if the port ends with 465 (SMTPS protocol)
	if strings.HasSuffix(port, "465") {
		conn = tls.Client(conn, tlsconfig)
		isSecureConn = true
	}

	client, err := smtp.NewClient(conn, host)
	if err != nil {
		return err
	}

	if !setting.MailService.DisableHelo {
		hostname := setting.MailService.HeloHostname
		if len(hostname) == 0 {
			hostname, err = os.Hostname()
			if err != nil {
				return err
			}
		}

		if err = client.Hello(hostname); err != nil {
			return err
		}
	}

	// If not using SMTPS, alway use STARTTLS if available
	hasStartTLS, _ := client.Extension("STARTTLS")
	if !isSecureConn && hasStartTLS {
		if err = client.StartTLS(tlsconfig); err != nil {
			return err
		}
	}

	canAuth, options := client.Extension("AUTH")

	if canAuth && len(settings.User) > 0 {
		var auth smtp.Auth

		if strings.Contains(options, "CRAM-MD5") {
			auth = smtp.CRAMMD5Auth(settings.User, settings.Passwd)
		} else if strings.Contains(options, "PLAIN") {
			auth = smtp.PlainAuth("", settings.User, settings.Passwd, host)
		} else if strings.Contains(options, "LOGIN") {
			// Patch for AUTH LOGIN
			auth = LoginAuth(settings.User, settings.Passwd)
		}

		if auth != nil {
			if err = client.Auth(auth); err != nil {
				return err
			}
		}
	}

	if fromAddress, err := mail.ParseAddress(settings.From); err != nil {
		return err
	} else {
		if err = client.Mail(fromAddress.Address); err != nil {
			return err
		}
	}

	for _, rec := range recipients {
		if err = client.Rcpt(rec); err != nil {
			return err
		}
	}

	w, err := client.Data()
	if err != nil {
		return err
	}
	if _, err = w.Write([]byte(msgContent)); err != nil {
		return err
	}

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

	return client.Quit()
}

// Direct Send mail message
func Send(msg *Message) (int, error) {
	log.Trace("Sending mails to: %s", strings.Join(msg.To, "; "))

	// get message body
	content := msg.Content()

	if len(msg.To) == 0 {
		return 0, fmt.Errorf("empty receive emails")
	} else if len(msg.Body) == 0 {
		return 0, fmt.Errorf("empty email body")
	}

	if msg.Massive {
		// send mail to multiple emails one by one
		num := 0
		for _, to := range msg.To {
			body := []byte("To: " + to + "\r\n" + content)
			err := sendMail(setting.MailService, []string{to}, body)
			if err != nil {
				return num, err
			}
			num++
		}
		return num, nil
	} else {
		body := []byte("To: " + strings.Join(msg.To, ";") + "\r\n" + content)

		// send to multiple emails in one message
		err := sendMail(setting.MailService, msg.To, body)
		if err != nil {
			return 0, err
		} else {
			return 1, nil
		}
	}
}

// Async Send mail message
func SendAsync(msg *Message) {
	go func() {
		mailQueue <- msg
	}()
}

// Create html mail message
func NewHtmlMessage(To []string, From, Subject, Body string) Message {
	return Message{
		To:      To,
		From:    setting.MailService.From,
		Reply-To: From,
		Subject: Subject,
		Body:    Body,
		Type:    "html",
	}
}
<|endoftext|>


<commit_before>package main

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

	xj "github.com/basgys/goxml2json"
	"github.com/gorilla/websocket"
)

//VERSION of the program
var version = "undefined-autogenerated"

//gobal variable to store all websocket connections for the update
var cons = make(map[*websocket.Conn]bool)

//mutex to lock the NMAP JSON information
var lock = sync.RWMutex{}

var nmapJSON *bytes.Buffer

//just for development
type msg struct {
	Command string
	Num     int
}

//Config data struct to read the config file
type Config struct {
	NMAPRange    string
	NMAPPorts    string //comma separated
	HTTPPort     int
	ScanInterval int //seconds
}

func echo(conn *websocket.Conn) {
	for {
		m := msg{}
		err := conn.ReadJSON(&m)
		if err != nil {
			log.Println("Error reading json.", err)
			break
		}
		lock.RLock()
		if nmapJSON != nil {
			err := conn.WriteMessage(1, nmapJSON.Bytes())
			if err != nil {
				log.Println(err)
				break
			}
		}
		lock.RUnlock()
	}
	delete(cons, conn)
	conn.Close()
}

//ReadConfig reads the config file
func readConfig(configfile string) Config {
	var config Config
	data, err := ioutil.ReadFile(configfile)
	if err != nil {
		log.Println("Error opening for JSON parsing - reverting to defaults")
	}

	err = json.Unmarshal(data, &config)
	if err != nil {
		log.Println("Error parsing json config:", err)
		log.Println("Continuing with defaults", err)
	}

	return config
}

func callNMAP(conf Config) {
	log.Println("Starting nmap caller")
	var counter = 1
	var scanResultsFileName = "scan.xml"

	for {
		cmd := exec.Command("nmap", "-p", conf.NMAPPorts, "-oX", scanResultsFileName, conf.NMAPRange)
		log.Println("Init NMAP scan no:", counter)
		var out bytes.Buffer
		cmd.Stdout = &out
		err := cmd.Run()
		if err != nil {
			log.Println(err)
		}
		log.Println("Scan no.", counter, "complete")
		counter = counter + 1

		//read the xml file, convert to json and send via the websockets
		xmlfile, err := os.Open(scanResultsFileName)
		if err != nil {
			log.Println(err)
		}
		lock.Lock()
		nmapJSON, err = xj.Convert(xmlfile)
		lock.Unlock()
		if err != nil {
			log.Println(err)
		}
		log.Println("Writing to websockets")
		for c := range cons {
			c.WriteMessage(1, nmapJSON.Bytes())
		}
		<-time.After(time.Duration(conf.ScanInterval) * time.Second)
	}
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("Ws request from:", r.Header.Get("Origin"))
	// log.Println("accepted from: http://" + r.Host)
	// if r.Header.Get("Origin") != "http://"+r.Host {
	// 	http.Error(w, "Origin not allowed", 403)
	// 	log.Println("Origin not allowed")
	// 	return
	// }
	conn, err := websocket.Upgrade(w, r, w.Header(), 0, 0)
	if err != nil {
		http.Error(w, "Could not open websocket connection", http.StatusBadRequest)
		log.Println("Upgrade error")
	}

	//add the connection to the pool for broadcast
	cons[conn] = true
	go echo(conn)
}

func main() {
	log.Println("Starting lan-monitor-server ver: " + version)

	//process the config
	//1st the config file is read and set parameters applied
	//2nd the command line parameters are interpreted,
	//if they are set they will overrule the config file
	//3rd if none of the above is applied the program reverts to the hardcoded defaults

	//defaults
	var config Config
	defaultConfigFileLocation := "/etc/lan-monitor.json"
	config.HTTPPort = 8080
	config.NMAPRange = "192.168.0.1/24"
	config.NMAPPorts = "22,80"
	config.ScanInterval = 120 //seconds

	displayVersion := flag.Bool("version", false, "Prints the version number")
	createExampleConfig := flag.Bool("config", false, "Writes the lan-monitor.json to local dir as example")
	cmdlineHTTPPort := flag.Int("port", config.HTTPPort, "HTTP port for the webserver")
	cmdlineNMAPScanRange := flag.String("range", config.NMAPRange, "The range NMAP should scan e.g. 192.168.1.1/24 it has to be nmap compatible")
	cmdlineScanInterval := flag.Int("scan-rate", config.ScanInterval, "The interval of the scans in seconds")
	configFileLocation := flag.String("config-file", defaultConfigFileLocation, "Location of the config file")
	cmdlinePorts := flag.String("scan-ports", config.NMAPPorts, "The ports that will be scanned")
	flag.Parse()

	//try to read the configfile
	_, err := os.Stat(*configFileLocation)
	if err == nil {
		config = readConfig(*configFileLocation)
	} else {
		log.Println("Config file is missing - looked at:", *configFileLocation)
		log.Println("Reverting to commandline/defaults")
	}

	//if no range is defined in the config file
	if config.NMAPRange == "" {
		config.NMAPRange = *cmdlineNMAPScanRange
	}

	//if no port is defined in the config file
	if config.HTTPPort == 0 {
		config.HTTPPort = *cmdlineHTTPPort
	}

	//if no scan interval is defined in the config file
	if config.ScanInterval == 0 {
		config.ScanInterval = *cmdlineScanInterval
	}

	//if no ports to be scanned are defined
	if config.NMAPPorts == "" {
		config.NMAPPorts = *cmdlinePorts
	}

	log.Println("Config - range:", config.NMAPRange, "webserver port:", config.HTTPPort, "interval:", config.ScanInterval, "sec", "scan-ports", config.NMAPPorts)

	if *displayVersion == true {
		fmt.Println("Version: " + version)
		return
	}

	if *createExampleConfig == true {
		workingDir, _ := os.Getwd()
		fmt.Println("Config in JSON format written to:", workingDir)
		configJSON, _ := json.Marshal(config)
		err := ioutil.WriteFile("lan-monitor.json", configJSON, 0644)
		if err != nil {
			fmt.Println("Error writing demo config file")
		}
		return
	}

	workingDir, _ := os.Getwd()
	log.Println("Dir:" + workingDir)

	//init the scanning routine
	go callNMAP(config)

	//starting and configuring the webserver
	fs := http.FileServer(http.Dir(workingDir))
	http.HandleFunc("/ws", wsHandler)
	http.Handle("/", fs)
	listenAddress := ":" + strconv.Itoa(config.HTTPPort)
	log.Fatal(http.ListenAndServe(listenAddress, nil))
}
<commit_msg>correction new path<commit_after>package main

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

	xj "github.com/basgys/goxml2json"
	"github.com/gorilla/websocket"
)

//VERSION of the program
var version = "undefined-autogenerated"

//gobal variable to store all websocket connections for the update
var cons = make(map[*websocket.Conn]bool)

//mutex to lock the NMAP JSON information
var lock = sync.RWMutex{}

var nmapJSON *bytes.Buffer

//just for development
type msg struct {
	Command string
	Num     int
}

//Config data struct to read the config file
type Config struct {
	NMAPRange    string
	NMAPPorts    string //comma separated
	HTTPPort     int
	ScanInterval int //seconds
}

func echo(conn *websocket.Conn) {
	for {
		m := msg{}
		err := conn.ReadJSON(&m)
		if err != nil {
			log.Println("Error reading json.", err)
			break
		}
		lock.RLock()
		if nmapJSON != nil {
			err := conn.WriteMessage(1, nmapJSON.Bytes())
			if err != nil {
				log.Println(err)
				break
			}
		}
		lock.RUnlock()
	}
	delete(cons, conn)
	conn.Close()
}

//ReadConfig reads the config file
func readConfig(configfile string) Config {
	var config Config
	data, err := ioutil.ReadFile(configfile)
	if err != nil {
		log.Println("Error opening for JSON parsing - reverting to defaults")
	}

	err = json.Unmarshal(data, &config)
	if err != nil {
		log.Println("Error parsing json config:", err)
		log.Println("Continuing with defaults", err)
	}

	return config
}

func callNMAP(conf Config) {
	log.Println("Starting nmap caller")
	var counter = 1
	var scanResultsFileName = "scan.xml"

	for {
		cmd := exec.Command("nmap", "-p", conf.NMAPPorts, "-oX", scanResultsFileName, conf.NMAPRange)
		log.Println("Init NMAP scan no:", counter)
		var out bytes.Buffer
		cmd.Stdout = &out
		err := cmd.Run()
		if err != nil {
			log.Println(err)
		}
		log.Println("Scan no.", counter, "complete")
		counter = counter + 1

		//read the xml file, convert to json and send via the websockets
		xmlfile, err := os.Open(scanResultsFileName)
		if err != nil {
			log.Println(err)
		}
		lock.Lock()
		nmapJSON, err = xj.Convert(xmlfile)
		lock.Unlock()
		if err != nil {
			log.Println(err)
		}
		log.Println("Writing to websockets")
		for c := range cons {
			c.WriteMessage(1, nmapJSON.Bytes())
		}
		<-time.After(time.Duration(conf.ScanInterval) * time.Second)
	}
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("Ws request from:", r.Header.Get("Origin"))
	// log.Println("accepted from: http://" + r.Host)
	// if r.Header.Get("Origin") != "http://"+r.Host {
	// 	http.Error(w, "Origin not allowed", 403)
	// 	log.Println("Origin not allowed")
	// 	return
	// }
	conn, err := websocket.Upgrade(w, r, w.Header(), 0, 0)
	if err != nil {
		http.Error(w, "Could not open websocket connection", http.StatusBadRequest)
		log.Println("Upgrade error")
	}

	//add the connection to the pool for broadcast
	cons[conn] = true
	go echo(conn)
}

func main() {
	log.Println("Starting lan-monitor-server ver: " + version)

	//process the config
	//1st the config file is read and set parameters applied
	//2nd the command line parameters are interpreted,
	//if they are set they will overrule the config file
	//3rd if none of the above is applied the program reverts to the hardcoded defaults

	//defaults
	var config Config
	defaultConfigFileLocation := "/etc/lan-monitor.json"
	config.HTTPPort = 8080
	config.NMAPRange = "192.168.0.1/24"
	config.NMAPPorts = "22,80"
	config.ScanInterval = 120 //seconds

	displayVersion := flag.Bool("version", false, "Prints the version number")
	createExampleConfig := flag.Bool("config", false, "Writes the lan-monitor.json to local dir as example")
	cmdlineHTTPPort := flag.Int("port", config.HTTPPort, "HTTP port for the webserver")
	cmdlineNMAPScanRange := flag.String("range", config.NMAPRange, "The range NMAP should scan e.g. 192.168.1.1/24 it has to be nmap compatible")
	cmdlineScanInterval := flag.Int("scan-rate", config.ScanInterval, "The interval of the scans in seconds")
	configFileLocation := flag.String("config-file", defaultConfigFileLocation, "Location of the config file")
	cmdlinePorts := flag.String("scan-ports", config.NMAPPorts, "The ports that will be scanned")
	flag.Parse()

	//try to read the configfile
	_, err := os.Stat(*configFileLocation)
	if err == nil {
		config = readConfig(*configFileLocation)
	} else {
		log.Println("Config file is missing - looked at:", *configFileLocation)
		log.Println("Reverting to commandline/defaults")
	}

	//if no range is defined in the config file
	if config.NMAPRange == "" {
		config.NMAPRange = *cmdlineNMAPScanRange
	}

	//if no port is defined in the config file
	if config.HTTPPort == 0 {
		config.HTTPPort = *cmdlineHTTPPort
	}

	//if no scan interval is defined in the config file
	if config.ScanInterval == 0 {
		config.ScanInterval = *cmdlineScanInterval
	}

	//if no ports to be scanned are defined
	if config.NMAPPorts == "" {
		config.NMAPPorts = *cmdlinePorts
	}

	log.Println("Config - range:", config.NMAPRange, "webserver port:", config.HTTPPort, "interval:", config.ScanInterval, "sec", "scan-ports", config.NMAPPorts)

	if *displayVersion == true {
		fmt.Println("Version: " + version)
		return
	}

	if *createExampleConfig == true {
		workingDir, _ := os.Getwd()
		fmt.Println("Config in JSON format written to:", workingDir)
		configJSON, _ := json.Marshal(config)
		err := ioutil.WriteFile("lan-monitor.json", configJSON, 0644)
		if err != nil {
			fmt.Println("Error writing demo config file")
		}
		return
	}

	workingDir, _ := os.Getwd()
	log.Println("Dir:" + workingDir)

	//init the scanning routine
	go callNMAP(config)

	//starting and configuring the webserver
	fs := http.FileServer(http.Dir(filepath.Dir(workingDir) + "/www"))
	http.HandleFunc("/ws", wsHandler)
	http.Handle("/", fs)
	listenAddress := ":" + strconv.Itoa(config.HTTPPort)
	log.Fatal(http.ListenAndServe(listenAddress, nil))
}
<|endoftext|>


<commit_before>package chainsync

import (
	"fmt"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/wire"
)

// ErrCheckpointMismatch is returned if given filter headers don't pass our
// control check.
var ErrCheckpointMismatch = fmt.Errorf("checkpoint doesn't match")

// filterHeaderCheckpoints holds a mapping from heights to filter headers for
// various heights. We use them to check whether peers are serving us the
// expected filter headers.
var filterHeaderCheckpoints = map[wire.BitcoinNet]map[uint32]*chainhash.Hash{
	// Mainnet filter header checkpoints.
	chaincfg.MainNetParams.Net: map[uint32]*chainhash.Hash{},

	// Testnet filter header checkpoints.
	chaincfg.TestNet3Params.Net: map[uint32]*chainhash.Hash{},
}

// ControlCFHeader controls the given filter header against our list of
// checkpoints. It returns ErrCheckpointMismatch if we have a checkpoint at the
// given height, and it doesn't match.
func ControlCFHeader(params chaincfg.Params, fType wire.FilterType,
	height uint32, filterHeader *chainhash.Hash) error {

	if fType != wire.GCSFilterRegular {
		return fmt.Errorf("unsupported filter type %v", fType)
	}

	control, ok := filterHeaderCheckpoints[params.Net]
	if !ok {
		return nil
	}

	hash, ok := control[height]
	if !ok {
		return nil
	}

	if *filterHeader != *hash {
		return ErrCheckpointMismatch
	}

	return nil
}

func hashFromStr(hexStr string) *chainhash.Hash {
	hash, _ := chainhash.NewHashFromStr(hexStr)
	return hash
}
<commit_msg>chainsync: populate testnet+mainnet filter header meta checkpoints<commit_after>package chainsync

import (
	"fmt"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/wire"
)

// ErrCheckpointMismatch is returned if given filter headers don't pass our
// control check.
var ErrCheckpointMismatch = fmt.Errorf("checkpoint doesn't match")

// filterHeaderCheckpoints holds a mapping from heights to filter headers for
// various heights. We use them to check whether peers are serving us the
// expected filter headers.
var filterHeaderCheckpoints = map[wire.BitcoinNet]map[uint32]*chainhash.Hash{
	// Mainnet filter header checkpoints.
	chaincfg.MainNetParams.Net: map[uint32]*chainhash.Hash{
		100000: hashFromStr("f28cbc1ab369eb01b7b5fe8bf59763abb73a31471fe404a26a06be4153aa7fa5"),
		200000: hashFromStr("e5031471732f4fbfe7a25f6a03acc1413300d5c56ae8e06b95046b8e4c0f32b3"),
		300000: hashFromStr("1bd50220fcdde929ca3143c91d2dd9a9bfedb38c452ba98dbb51e719bff8aa5b"),
		400000: hashFromStr("5d973ab1f1c569c70deec1c1a8fb2e317a260f1656edb3b262c65f78ef192e3a"),
		500000: hashFromStr("5d16ca293c9bdc0a9bc279b63f99fb661be38b095a59a44200a807caaa631a3c"),
		540000: hashFromStr("bbabb3b757ff0776971c5719e58a4fdc7f2a8159c9028be62896bc17ba14dda1"),
	},

	// Testnet filter header checkpoints.
	chaincfg.TestNet3Params.Net: map[uint32]*chainhash.Hash{
		100000:  hashFromStr("97c0633f14625627fcd133250ad8cc525937e776b5f3fd272b06d02c58b65a1c"),
		200000:  hashFromStr("51aa817e5abe3acdcf103616b1a5736caf84bc3773a7286e9081108ecc38cc87"),
		400000:  hashFromStr("4aab9b3d4312cd85cfcd48a08b36c4402bfdc1e8395dcf4236c3029dfa837c48"),
		600000:  hashFromStr("713d9c9198e2dba0739e85aab6875cb951c36297b95a2d51131aa6919753b55d"),
		800000:  hashFromStr("0dafdff27269a70293c120b14b1f5e9a72a5e8688098cfc6140b9d64f8325b99"),
		1000000: hashFromStr("c2043fa2f6eb5f8f8d2c5584f743187f36302ed86b62c302e31155f378da9c5f"),
		1390000: hashFromStr("ec71c508c02f59b2af2f34b64dfd79ffba55d9ef7d00589b0a2c3178da89e4c0"),
		1400000: hashFromStr("f9ae1750483d4c8ce82512616b1ded932886af46decb8d3e575907930542d9b3"),
		1500000: hashFromStr("dc0cfa13daf09df9b8dbe7532f75ebdb4255860b295016b2ca4b789394bc5090"),
	},
}

// ControlCFHeader controls the given filter header against our list of
// checkpoints. It returns ErrCheckpointMismatch if we have a checkpoint at the
// given height, and it doesn't match.
func ControlCFHeader(params chaincfg.Params, fType wire.FilterType,
	height uint32, filterHeader *chainhash.Hash) error {

	if fType != wire.GCSFilterRegular {
		return fmt.Errorf("unsupported filter type %v", fType)
	}

	control, ok := filterHeaderCheckpoints[params.Net]
	if !ok {
		return nil
	}

	hash, ok := control[height]
	if !ok {
		return nil
	}

	if *filterHeader != *hash {
		return ErrCheckpointMismatch
	}

	return nil
}

// hashFromStr makes a chainhash.Hash from a valid hex string. If the string is
// invalid, a nil pointer will be returned.
func hashFromStr(hexStr string) *chainhash.Hash {
	hash, _ := chainhash.NewHashFromStr(hexStr)
	return hash
}
<|endoftext|>


<commit_before>// All names loose the CK_ prefix
// All names loose the hungarian notation
//
package pkcs11

/*
#define CK_PTR *
#ifndef NULL_PTR
#define NULL_PTR 0
#endif
#define CK_DEFINE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION_POINTER(returnType, name) returnType (* name)
#define CK_CALLBACK_FUNCTION(returnType, name) returnType (* name)

#include <stdlib.h>
#include "pkcs11.h"
*/
import "C"

type Attribute struct {
	Type  C.CK_ATTRIBUTE_TYPE
	Value C.CK_VOID_PTR
	Len   C.CK_ULONG
}
<commit_msg>Add the most basic structures<commit_after>// All names loose the CK_ prefix
// All names loose the hungarian notation
// All the defines are kept from the C package so: C.CKM_RSA_X_509
// All struct's get a Go variant
//
package pkcs11

/*
#define CK_PTR *
#ifndef NULL_PTR
#define NULL_PTR 0
#endif
#define CK_DEFINE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION_POINTER(returnType, name) returnType (* name)
#define CK_CALLBACK_FUNCTION(returnType, name) returnType (* name)

#include <stdlib.h>
#include "pkcs11.h"
*/
import "C"

// A Void is used a lot in the PKCS#11 library, it always consists of a *void and a length.
// We use []byte as a subsitite in Go.
type Void []byte

type Version struct {
	Major byte
	Minor byte
}

type Info struct {
	// TODO
}

type SlotInfo struct {
	SlotDescription [64]byte
	ManufacturerID  [32]byte
	Flags           C.CK_FLAGS
	HardwareVersion Version
	FirmwareVersion Version
}

type TokenInfo struct {
	Label              [32]byte
	ManufacturerID     [32]byte
	Model              [16]byte
	SerialNumber       [16]byte
	Flags              C.CK_FLAGS
	MaxSessionCount    C.CK_ULONG
	SessionCount       C.CK_ULONG
	MaxRwSessionCount  C.CK_ULONG
	RwSessionCount     C.CK_ULONG
	MaxPinLen          C.CK_ULONG
	MinPinLen          C.CK_ULONG
	TotalPublicMemory  C.CK_ULONG
	FreePublicMemory   C.CK_ULONG
	TotalPrivateMemory C.CK_ULONG
	FreePrivateMemory  C.CK_ULONG
	hardwareVersion    Version
	firmwareVersion    Version
	UTCTime            [16]byte
}

type SessionInfo struct {
	SlotID      C.CK_SLOT_ID
	Sate        C.CK_STATE
	Flags       C.CK_FLAGS
	DeviceError C.CK_ULONG
}

type Attribute struct {
	Type  C.CK_ATTRIBUTE_TYPE
	Value Void
}

type Date struct {
	// TODO
}

type Mechanism struct {
	Type      C.CK_MECHANISM_TYPE
	Parameter Void
}

type MechanismInfo struct {
	MinKeySize C.CK_ULONG
	MaxKeySize C.CK_ULONG
	Flags      C.CK_FLAGS
}

// stopped after this one
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/codegangsta/cli"
	"github.com/gin-gonic/gin"
	"github.com/johnwilson/bytengine"
)

var (
	EngineRequestChan chan *EngineRequest
	Configuration     *Config
)

const (
	Version = "0.2.2"
)

type EngineRequest struct {
	Token        string
	Script       string
	Command      *bytengine.Command
	ResponseChan chan EngineResponse
}

type EngineResponse struct {
	Response interface{}
	Error    error
}

func errorResponse(err error) []byte {
	val := map[string]string{
		"status": "error",
		"msg":    err.Error(),
	}

	b, e := json.Marshal(val)
	if e != nil {
		return []byte{}
	}
	return b
}

func okResponse(data interface{}) []byte {
	val := map[string]interface{}{
		"status": "ok",
		"data":   data,
	}

	b, e := json.Marshal(val)
	if e != nil {
		return []byte{}
	}
	return b
}

func Worker(config *Config, requests chan *EngineRequest) {
	// create engine and start
	engine := bytengine.NewEngine()
	err := engine.Start(config.Bytengine)
	if err != nil {
		fmt.Println("Error: ", err)
		os.Exit(1)
	}

	for req := range requests {
		// check if script request or command request
		if req.Command != nil {
			r, err := engine.ExecuteCommand(req.Token, *req.Command)
			rep := EngineResponse{r, err}
			req.ResponseChan <- rep
		} else {
			r, err := engine.ExecuteScript(req.Token, req.Script)
			rep := EngineResponse{r, err}
			req.ResponseChan <- rep
		}
	}
}

func WorkerPool(workers int64, config *Config) chan *EngineRequest {
	requests := make(chan *EngineRequest)
	for i := int64(0); i < workers; i++ {
		go Worker(config, requests)
	}

	return requests
}

func welcomeHandler(ctx *gin.Context) {
	msg := fmt.Sprintf(`{"bytengine":"welcome","version":"%s"}`, Version)
	ctx.Data(200, "application/json", []byte(msg))
}

func runScriptHandler(ctx *gin.Context) {
	var form struct {
		Token string `form:"token" binding:"required"`
		Query string `form:"query" binding:"required"`
	}

	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	req := EngineRequest{
		Token:        form.Token,
		Script:       form.Query,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func getTokenHandler(ctx *gin.Context) {
	var form struct {
		Username string `form:"username" binding:"required"`
		Password string `form:"password" binding:"required"`
	}
	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "login",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Args["username"] = form.Username
	cmd.Args["password"] = form.Password

	duration := Configuration.Timeout.AuthToken // in minutes
	cmd.Args["duration"] = duration

	req := EngineRequest{
		Token:        "",
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func getUploadTicketHandler(ctx *gin.Context) {
	var form struct {
		Token    string `form:"token" binding:"required"`
		Database string `form:"database" binding:"required"`
		Path     string `form:"path" binding:"required"`
	}
	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "uploadticket",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Database = form.Database
	cmd.Args["path"] = form.Path

	duration := Configuration.Timeout.UploadTicket // in minutes
	cmd.Args["duration"] = duration

	req := EngineRequest{
		Token:        form.Token,
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func uploadFileHelper(max int, ctx *gin.Context) (string, int, error) {
	total := 0             // total bytes read/written
	maxbytes := 1024 * max // maximum upload size in bytes

	tmpfile, err := ioutil.TempFile("", "bytengine_upload_") // upload file name from header
	if err != nil {
		return "", 0, err
	}
	defer tmpfile.Close()

	// create read buffer
	var bsize int64 = 16 * 1024 // 16 kb
	buffer := make([]byte, bsize)

	// get stream
	mr, err := ctx.Request.MultipartReader()
	if err != nil {
		return "", 0, err
	}
	in_f, err := mr.NextPart()
	if err != nil {
		return "", 0, err
	}
	defer in_f.Close()

	// start reading/writing

	for {
		// read
		n, err := in_f.Read(buffer)
		if n == 0 {
			break
		}
		if err != nil {
			return "", 0, err
		}
		// update total bytes
		total += n
		if total > maxbytes {
			return "", 0, fmt.Errorf("exceeded maximum file size of %d bytes", maxbytes)
		}
		// write
		n, err = tmpfile.Write(buffer[:n])
		if err != nil {
			return "", 0, err
		}
	}

	return tmpfile.Name(), total, nil
}

func uploadFileHandler(ctx *gin.Context) {
	ticket := ctx.Params.ByName("ticket")
	filename, _, err := uploadFileHelper(300, ctx)
	if err != nil {
		data := errorResponse(fmt.Errorf("upload failed: %s", err.Error()))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "writebytes",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Args["ticket"] = ticket
	cmd.Args["tmpfile"] = filename

	req := EngineRequest{
		Token:        "",
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func downloadFileHandler(ctx *gin.Context) {
	var form struct {
		Token    string `form:"token" binding:"required"`
		Database string `form:"database" binding:"required"`
		Path     string `form:"path" binding:"required"`
	}
	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "readbytes",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Database = form.Database
	cmd.Args["path"] = form.Path
	cmd.Args["writer"] = ctx.Writer

	ctx.Writer.Header().Set("Content-Type", "application/octet-stream")
	req := EngineRequest{
		Token:        form.Token,
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.String(400, string(data))
		return
	}
}

func directaccessHandler(ctx *gin.Context) {
	db := ctx.Params.ByName("database")
	path := ctx.Params.ByName("path")
	layer := ctx.Params.ByName("layer")

	cmd := bytengine.Command{
		Name:    "directaccess",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Args["database"] = db
	cmd.Args["path"] = path
	cmd.Args["layer"] = layer
	cmd.Args["writer"] = ctx.Writer

	if layer == "json" {
		ctx.Writer.Header().Set("Content-Type", "application/json")
	} else {
		ctx.Writer.Header().Set("Content-Type", "application/octet-stream")
	}

	req := EngineRequest{
		Token:        "",
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.String(404, string(data))
		return
	}
}

func writeConfig(b []byte) error {
	Configuration = &Config{}
	return json.Unmarshal(b, Configuration)
}

func main() {
	app := cli.NewApp()

	createadminCmd := cli.Command{
		Name: "createadmin",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "u", Value: "", Usage: "username"},
			cli.StringFlag{Name: "p", Value: "", Usage: "password"},
			cli.StringFlag{Name: "c", Value: "config.json"},
		},
		Action: func(c *cli.Context) {
			usr := c.String("u")
			pw := c.String("p")

			b, err := configToJSON()
			if err != nil {
				fmt.Println("Error: ", err)
			}
			err = writeConfig(b)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}

			// create and start bytengine
			engine := bytengine.NewEngine()
			err = engine.Start(Configuration.Bytengine)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}

			err = engine.CreateAdminUser(usr, pw)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}
			fmt.Println("...done")
		},
	}

	run := cli.Command{
		Name: "run",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "c", Value: "config.json"},
		},
		Action: func(c *cli.Context) {
			b, err := configToJSON()
			if err != nil {
				fmt.Println("Error: ", err)
			}
			err = writeConfig(b)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}

			addr := Configuration.Address
			port := Configuration.Port
			workers := Configuration.Workers

			// start worker pool
			EngineRequestChan = WorkerPool(workers, Configuration)

			// setup routes
			router := gin.Default()
			router.GET("/", welcomeHandler)
			router.POST("/bfs/query", runScriptHandler)
			router.POST("/bfs/token", getTokenHandler)
			router.POST("/bfs/uploadticket", getUploadTicketHandler)
			router.POST("/bfs/writebytes/:ticket", uploadFileHandler)
			router.POST("/bfs/readbytes", downloadFileHandler)
			router.GET("/bfs/direct/:layer/:database/*path", directaccessHandler)

			router.Run(fmt.Sprintf("%s:%d", addr, port))
		},
	}
	app.Commands = []cli.Command{createadminCmd, run}
	app.Run(os.Args)
}
<commit_msg>updated version<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/codegangsta/cli"
	"github.com/gin-gonic/gin"
	"github.com/johnwilson/bytengine"
)

var (
	EngineRequestChan chan *EngineRequest
	Configuration     *Config
)

const (
	Version = "0.3.0"
)

type EngineRequest struct {
	Token        string
	Script       string
	Command      *bytengine.Command
	ResponseChan chan EngineResponse
}

type EngineResponse struct {
	Response interface{}
	Error    error
}

func errorResponse(err error) []byte {
	val := map[string]string{
		"status": "error",
		"msg":    err.Error(),
	}

	b, e := json.Marshal(val)
	if e != nil {
		return []byte{}
	}
	return b
}

func okResponse(data interface{}) []byte {
	val := map[string]interface{}{
		"status": "ok",
		"data":   data,
	}

	b, e := json.Marshal(val)
	if e != nil {
		return []byte{}
	}
	return b
}

func Worker(config *Config, requests chan *EngineRequest) {
	// create engine and start
	engine := bytengine.NewEngine()
	err := engine.Start(config.Bytengine)
	if err != nil {
		fmt.Println("Error: ", err)
		os.Exit(1)
	}

	for req := range requests {
		// check if script request or command request
		if req.Command != nil {
			r, err := engine.ExecuteCommand(req.Token, *req.Command)
			rep := EngineResponse{r, err}
			req.ResponseChan <- rep
		} else {
			r, err := engine.ExecuteScript(req.Token, req.Script)
			rep := EngineResponse{r, err}
			req.ResponseChan <- rep
		}
	}
}

func WorkerPool(workers int64, config *Config) chan *EngineRequest {
	requests := make(chan *EngineRequest)
	for i := int64(0); i < workers; i++ {
		go Worker(config, requests)
	}

	return requests
}

func welcomeHandler(ctx *gin.Context) {
	msg := fmt.Sprintf(`{"bytengine":"welcome","version":"%s"}`, Version)
	ctx.Data(200, "application/json", []byte(msg))
}

func runScriptHandler(ctx *gin.Context) {
	var form struct {
		Token string `form:"token" binding:"required"`
		Query string `form:"query" binding:"required"`
	}

	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	req := EngineRequest{
		Token:        form.Token,
		Script:       form.Query,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func getTokenHandler(ctx *gin.Context) {
	var form struct {
		Username string `form:"username" binding:"required"`
		Password string `form:"password" binding:"required"`
	}
	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "login",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Args["username"] = form.Username
	cmd.Args["password"] = form.Password

	duration := Configuration.Timeout.AuthToken // in minutes
	cmd.Args["duration"] = duration

	req := EngineRequest{
		Token:        "",
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func getUploadTicketHandler(ctx *gin.Context) {
	var form struct {
		Token    string `form:"token" binding:"required"`
		Database string `form:"database" binding:"required"`
		Path     string `form:"path" binding:"required"`
	}
	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "uploadticket",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Database = form.Database
	cmd.Args["path"] = form.Path

	duration := Configuration.Timeout.UploadTicket // in minutes
	cmd.Args["duration"] = duration

	req := EngineRequest{
		Token:        form.Token,
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func uploadFileHelper(max int, ctx *gin.Context) (string, int, error) {
	total := 0             // total bytes read/written
	maxbytes := 1024 * max // maximum upload size in bytes

	tmpfile, err := ioutil.TempFile("", "bytengine_upload_") // upload file name from header
	if err != nil {
		return "", 0, err
	}
	defer tmpfile.Close()

	// create read buffer
	var bsize int64 = 16 * 1024 // 16 kb
	buffer := make([]byte, bsize)

	// get stream
	mr, err := ctx.Request.MultipartReader()
	if err != nil {
		return "", 0, err
	}
	in_f, err := mr.NextPart()
	if err != nil {
		return "", 0, err
	}
	defer in_f.Close()

	// start reading/writing

	for {
		// read
		n, err := in_f.Read(buffer)
		if n == 0 {
			break
		}
		if err != nil {
			return "", 0, err
		}
		// update total bytes
		total += n
		if total > maxbytes {
			return "", 0, fmt.Errorf("exceeded maximum file size of %d bytes", maxbytes)
		}
		// write
		n, err = tmpfile.Write(buffer[:n])
		if err != nil {
			return "", 0, err
		}
	}

	return tmpfile.Name(), total, nil
}

func uploadFileHandler(ctx *gin.Context) {
	ticket := ctx.Params.ByName("ticket")
	filename, _, err := uploadFileHelper(300, ctx)
	if err != nil {
		data := errorResponse(fmt.Errorf("upload failed: %s", err.Error()))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "writebytes",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Args["ticket"] = ticket
	cmd.Args["tmpfile"] = filename

	req := EngineRequest{
		Token:        "",
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.Data(400, "application/json", data)
		return
	}

	ctx.Data(200, "application/json", okResponse(rep.Response))
}

func downloadFileHandler(ctx *gin.Context) {
	var form struct {
		Token    string `form:"token" binding:"required"`
		Database string `form:"database" binding:"required"`
		Path     string `form:"path" binding:"required"`
	}
	ok := ctx.Bind(&form)
	if !ok {
		data := errorResponse(errors.New("Missing parameters"))
		ctx.Data(400, "application/json", data)
		return
	}

	cmd := bytengine.Command{
		Name:    "readbytes",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Database = form.Database
	cmd.Args["path"] = form.Path
	cmd.Args["writer"] = ctx.Writer

	ctx.Writer.Header().Set("Content-Type", "application/octet-stream")
	req := EngineRequest{
		Token:        form.Token,
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.String(400, string(data))
		return
	}
}

func directaccessHandler(ctx *gin.Context) {
	db := ctx.Params.ByName("database")
	path := ctx.Params.ByName("path")
	layer := ctx.Params.ByName("layer")

	cmd := bytengine.Command{
		Name:    "directaccess",
		IsAdmin: false,
		Args:    make(map[string]interface{}),
		Options: make(map[string]interface{}),
	}
	cmd.Args["database"] = db
	cmd.Args["path"] = path
	cmd.Args["layer"] = layer
	cmd.Args["writer"] = ctx.Writer

	if layer == "json" {
		ctx.Writer.Header().Set("Content-Type", "application/json")
	} else {
		ctx.Writer.Header().Set("Content-Type", "application/octet-stream")
	}

	req := EngineRequest{
		Token:        "",
		Command:      &cmd,
		ResponseChan: make(chan EngineResponse),
	}
	EngineRequestChan <- &req
	rep := <-req.ResponseChan
	if rep.Error != nil {
		data := errorResponse(rep.Error)
		ctx.String(404, string(data))
		return
	}
}

func writeConfig(b []byte) error {
	Configuration = &Config{}
	return json.Unmarshal(b, Configuration)
}

func main() {
	app := cli.NewApp()

	createadminCmd := cli.Command{
		Name: "createadmin",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "u", Value: "", Usage: "username"},
			cli.StringFlag{Name: "p", Value: "", Usage: "password"},
			cli.StringFlag{Name: "c", Value: "config.json"},
		},
		Action: func(c *cli.Context) {
			usr := c.String("u")
			pw := c.String("p")

			b, err := configToJSON()
			if err != nil {
				fmt.Println("Error: ", err)
			}
			err = writeConfig(b)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}

			// create and start bytengine
			engine := bytengine.NewEngine()
			err = engine.Start(Configuration.Bytengine)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}

			err = engine.CreateAdminUser(usr, pw)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}
			fmt.Println("...done")
		},
	}

	run := cli.Command{
		Name: "run",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "c", Value: "config.json"},
		},
		Action: func(c *cli.Context) {
			b, err := configToJSON()
			if err != nil {
				fmt.Println("Error: ", err)
			}
			err = writeConfig(b)
			if err != nil {
				fmt.Println("Error: ", err)
				os.Exit(1)
			}

			addr := Configuration.Address
			port := Configuration.Port
			workers := Configuration.Workers

			// start worker pool
			EngineRequestChan = WorkerPool(workers, Configuration)

			// setup routes
			router := gin.Default()
			router.GET("/", welcomeHandler)
			router.POST("/bfs/query", runScriptHandler)
			router.POST("/bfs/token", getTokenHandler)
			router.POST("/bfs/uploadticket", getUploadTicketHandler)
			router.POST("/bfs/writebytes/:ticket", uploadFileHandler)
			router.POST("/bfs/readbytes", downloadFileHandler)
			router.GET("/bfs/direct/:layer/:database/*path", directaccessHandler)

			router.Run(fmt.Sprintf("%s:%d", addr, port))
		},
	}
	app.Commands = []cli.Command{createadminCmd, run}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"strings"

	"github.com/funkygao/columnize"
	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/gocli"

	// bootstrap plugins
	_ "github.com/funkygao/dbus/plugins/filter"
	_ "github.com/funkygao/dbus/plugins/input"
	_ "github.com/funkygao/dbus/plugins/output"
)

type Plugins struct {
	Ui  cli.Ui
	Cmd string

	fn string
}

func (this *Plugins) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("binlog", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.fn, "c", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	e := engine.New(nil).LoadFrom(this.fn)
	lines := []string{"Plugin|Name|Class"}
	for _, ir := range e.InputRunners {
		lines = append(lines, fmt.Sprintf("Input|%s|%s", ir.Name(), ir.Class()))
	}
	for _, fr := range e.FilterRunners {
		lines = append(lines, fmt.Sprintf("Filter|%s|%s", fr.Name(), fr.Class()))
	}
	for _, or := range e.OutputRunners {
		lines = append(lines, fmt.Sprintf("Output|%s|%s", or.Name(), or.Class()))
	}

	if len(lines) > 1 {
		this.Ui.Output(columnize.SimpleFormat(lines))
	}

	return
}

func (*Plugins) Synopsis() string {
	return "Display mysql binlog related info"
}

func (this *Plugins) Help() string {
	help := fmt.Sprintf(`
Usage: %s plugins [options]

    %s

Options:

    -c config location
      If empty, load from zookeeper
      zk location example:
      localhost:2181/dbus/conf

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>fix typo<commit_after>package command

import (
	"flag"
	"fmt"
	"strings"

	"github.com/funkygao/columnize"
	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/gocli"

	// bootstrap plugins
	_ "github.com/funkygao/dbus/plugins/filter"
	_ "github.com/funkygao/dbus/plugins/input"
	_ "github.com/funkygao/dbus/plugins/output"
)

type Plugins struct {
	Ui  cli.Ui
	Cmd string

	fn string
}

func (this *Plugins) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("binlog", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.fn, "c", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	e := engine.New(nil).LoadFrom(this.fn)
	lines := []string{"Plugin|Name|Class"}
	for _, ir := range e.InputRunners {
		lines = append(lines, fmt.Sprintf("Input|%s|%s", ir.Name(), ir.Class()))
	}
	for _, fr := range e.FilterRunners {
		lines = append(lines, fmt.Sprintf("Filter|%s|%s", fr.Name(), fr.Class()))
	}
	for _, or := range e.OutputRunners {
		lines = append(lines, fmt.Sprintf("Output|%s|%s", or.Name(), or.Class()))
	}

	if len(lines) > 1 {
		this.Ui.Output(columnize.SimpleFormat(lines))
	}

	return
}

func (*Plugins) Synopsis() string {
	return "List all plugins"
}

func (this *Plugins) Help() string {
	help := fmt.Sprintf(`
Usage: %s plugins [options]

    %s

Options:

    -c config location
      If empty, load from zookeeper
      zk location example:
      localhost:2181/dbus/conf

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"io"
	"time"

	"github.com/spf13/cobra"

	"github.com/weaveworks/flux"
	"github.com/weaveworks/flux/api/v11"
	"github.com/weaveworks/flux/api/v6"
	"github.com/weaveworks/flux/cluster"
	"github.com/weaveworks/flux/job"
	"github.com/weaveworks/flux/update"
)

type workloadReleaseOpts struct {
	*rootOpts
	namespace    string
	workloads    []string
	allWorkloads bool
	image        string
	allImages    bool
	exclude      []string
	dryRun       bool
	interactive  bool
	force        bool
	watch        bool
	outputOpts
	cause update.Cause

	// Deprecated
	controllers []string
}

func newWorkloadRelease(parent *rootOpts) *workloadReleaseOpts {
	return &workloadReleaseOpts{rootOpts: parent}
}

func (opts *workloadReleaseOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "release",
		Short: "Release a new version of a controller.",
		Example: makeExample(
			"fluxctl release -n default --workload=deployment/foo --update-image=library/hello:v2",
			"fluxctl release --all --update-image=library/hello:v2",
			"fluxctl release --workload=default:deployment/foo --update-all-images",
		),
		RunE: opts.RunE,
	}

	AddOutputFlags(cmd, &opts.outputOpts)
	AddCauseFlags(cmd, &opts.cause)
	cmd.Flags().StringVarP(&opts.namespace, "namespace", "n", "default", "Workload namespace")
	// Note: we cannot define a shorthand for --workload since it clashes with the shorthand of --watch
	cmd.Flags().StringSliceVarP(&opts.workloads, "workload", "", []string{}, "List of workloads to release <namespace>:<kind>/<name>")
	cmd.Flags().BoolVar(&opts.allWorkloads, "all", false, "Release all controllers")
	cmd.Flags().StringVarP(&opts.image, "update-image", "i", "", "Update a specific image")
	cmd.Flags().BoolVar(&opts.allImages, "update-all-images", false, "Update all images to latest versions")
	cmd.Flags().StringSliceVar(&opts.exclude, "exclude", []string{}, "List of controllers to exclude")
	cmd.Flags().BoolVar(&opts.dryRun, "dry-run", false, "Do not release anything; just report back what would have been done")
	cmd.Flags().BoolVar(&opts.interactive, "interactive", false, "Select interactively which containers to update")
	cmd.Flags().BoolVarP(&opts.force, "force", "f", false, "Disregard locks and container image filters (has no effect when used with --all or --update-all-images)")
	cmd.Flags().BoolVarP(&opts.watch, "watch", "w", false, "Watch rollout progress during release")

	// Deprecated
	cmd.Flags().StringSliceVarP(&opts.controllers, "controller", "c", []string{}, "List of controllers to release <namespace>:<kind>/<name>")
	cmd.Flags().MarkDeprecated("controller", "changed to --workspace, use that instead")

	return cmd
}

func (opts *workloadReleaseOpts) RunE(cmd *cobra.Command, args []string) error {
	if len(args) != 0 {
		return errorWantedNoArgs
	}

	if err := checkExactlyOne("--update-image=<image> or --update-all-images", opts.image != "", opts.allImages); err != nil {
		return err
	}

	// Backwards compatibility with --controller until we remove it
	opts.workloads = append(opts.workloads, opts.controllers...)

	switch {
	case len(opts.workloads) <= 0 && !opts.allWorkloads:
		return newUsageError("please supply either --all, or at least one --workload=<workload>")
	case opts.watch && opts.dryRun:
		return newUsageError("cannot use --watch with --dry-run")
	case opts.force && opts.allWorkloads && opts.allImages:
		return newUsageError("--force has no effect when used with --all and --update-all-images")
	case opts.force && opts.allWorkloads:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore locked workloads when used with --all\n")
	case opts.force && opts.allImages:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore container image tags when used with --update-all-images\n")
	}

	var workloads []update.ResourceSpec

	if opts.allWorkloads {
		workloads = []update.ResourceSpec{update.ResourceSpecAll}
	} else {
		for _, workload := range opts.workloads {
			id, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, workload)
			if err != nil {
				return err
			}
			workloads = append(workloads, update.MakeResourceSpec(id))
		}
	}

	var (
		image update.ImageSpec
		err   error
	)
	switch {
	case opts.image != "":
		image, err = update.ParseImageSpec(opts.image)
		if err != nil {
			return err
		}
	case opts.allImages:
		image = update.ImageSpecLatest
	}

	var kind update.ReleaseKind = update.ReleaseKindExecute
	if opts.dryRun || opts.interactive {
		kind = update.ReleaseKindPlan
	}

	var excludes []flux.ResourceID
	for _, exclude := range opts.exclude {
		s, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, exclude)
		if err != nil {
			return err
		}
		excludes = append(excludes, s)
	}

	if kind == update.ReleaseKindPlan {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting dry-run release ...\n")
	} else {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting release ...\n")
	}

	ctx := context.Background()
	spec := update.ReleaseImageSpec{
		ServiceSpecs: workloads,
		ImageSpec:    image,
		Kind:         kind,
		Excludes:     excludes,
		Force:        opts.force,
	}
	jobID, err := opts.API.UpdateManifests(ctx, update.Spec{
		Type:  update.Images,
		Cause: opts.cause,
		Spec:  spec,
	})
	if err != nil {
		return err
	}

	result, err := awaitJob(ctx, opts.API, jobID)
	if err != nil {
		return err
	}
	if opts.interactive {
		spec, err := promptSpec(cmd.OutOrStdout(), result, opts.verbosity)
		spec.Force = opts.force
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		fmt.Fprintf(cmd.OutOrStderr(), "Submitting selected release ...\n")
		jobID, err = opts.API.UpdateManifests(ctx, update.Spec{
			Type:  update.Containers,
			Cause: opts.cause,
			Spec:  spec,
		})
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		opts.dryRun = false
	}

	err = await(ctx, cmd.OutOrStdout(), cmd.OutOrStderr(), opts.API, jobID, !opts.dryRun, opts.verbosity)
	if !opts.watch || err != nil {
		return err
	}

	fmt.Fprintf(cmd.OutOrStderr(), "Monitoring rollout ...\n")
	for {
		completed := 0
		workloads, err := opts.API.ListServicesWithOptions(ctx, v11.ListServicesOptions{Services: result.Result.AffectedResources()})
		if err != nil {
			return err
		}

		for _, workload := range workloads {
			writeRolloutStatus(workload, opts.verbosity)

			if workload.Status == cluster.StatusReady {
				completed++
			}

			if workload.Rollout.Messages != nil {
				fmt.Fprintf(cmd.OutOrStderr(), "There was a problem releasing %s:\n", workload.ID)
				for _, msg := range workload.Rollout.Messages {
					fmt.Fprintf(cmd.OutOrStderr(), "%s\n", msg)
				}
				return nil
			}
		}

		if completed == len(workloads) {
			fmt.Fprintf(cmd.OutOrStderr(), "All workloads ready.\n")
			return nil
		}

		time.Sleep(2000 * time.Millisecond)
	}
}

func writeRolloutStatus(workload v6.ControllerStatus, verbosity int) {
	w := newTabwriter()
	fmt.Fprintf(w, "WORKLOAD\tCONTAINER\tIMAGE\tRELEASE\tREPLICAS\n")

	if len(workload.Containers) > 0 {
		c := workload.Containers[0]
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d/%d", workload.ID, c.Name, c.Current.ID, workload.Status, workload.Rollout.Updated, workload.Rollout.Desired)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d ready)", workload.Rollout.Outdated, workload.Rollout.Ready)
		}
		fmt.Fprintf(w, "\n")
		for _, c := range workload.Containers[1:] {
			fmt.Fprintf(w, "\t%s\t%s\t\t\n", c.Name, c.Current.ID)
		}
	} else {
		fmt.Fprintf(w, "%s\t\t\t%s\t%d/%d", workload.ID, workload.Status, workload.Rollout.Updated, workload.Rollout.Desired)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d ready)", workload.Rollout.Outdated, workload.Rollout.Ready)
		}
		fmt.Fprintf(w, "\n")
	}
	fmt.Fprintln(w)
	w.Flush()
}

func promptSpec(out io.Writer, result job.Result, verbosity int) (update.ReleaseContainersSpec, error) {
	menu := update.NewMenu(out, result.Result, verbosity)
	containerSpecs, err := menu.Run()
	return update.ReleaseContainersSpec{
		Kind:           update.ReleaseKindExecute,
		ContainerSpecs: containerSpecs,
		SkipMismatches: false,
	}, err
}
<commit_msg>Fix workspace to workload typo<commit_after>package main

import (
	"context"
	"fmt"
	"io"
	"time"

	"github.com/spf13/cobra"

	"github.com/weaveworks/flux"
	"github.com/weaveworks/flux/api/v11"
	"github.com/weaveworks/flux/api/v6"
	"github.com/weaveworks/flux/cluster"
	"github.com/weaveworks/flux/job"
	"github.com/weaveworks/flux/update"
)

type workloadReleaseOpts struct {
	*rootOpts
	namespace    string
	workloads    []string
	allWorkloads bool
	image        string
	allImages    bool
	exclude      []string
	dryRun       bool
	interactive  bool
	force        bool
	watch        bool
	outputOpts
	cause update.Cause

	// Deprecated
	controllers []string
}

func newWorkloadRelease(parent *rootOpts) *workloadReleaseOpts {
	return &workloadReleaseOpts{rootOpts: parent}
}

func (opts *workloadReleaseOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "release",
		Short: "Release a new version of a controller.",
		Example: makeExample(
			"fluxctl release -n default --workload=deployment/foo --update-image=library/hello:v2",
			"fluxctl release --all --update-image=library/hello:v2",
			"fluxctl release --workload=default:deployment/foo --update-all-images",
		),
		RunE: opts.RunE,
	}

	AddOutputFlags(cmd, &opts.outputOpts)
	AddCauseFlags(cmd, &opts.cause)
	cmd.Flags().StringVarP(&opts.namespace, "namespace", "n", "default", "Workload namespace")
	// Note: we cannot define a shorthand for --workload since it clashes with the shorthand of --watch
	cmd.Flags().StringSliceVarP(&opts.workloads, "workload", "", []string{}, "List of workloads to release <namespace>:<kind>/<name>")
	cmd.Flags().BoolVar(&opts.allWorkloads, "all", false, "Release all controllers")
	cmd.Flags().StringVarP(&opts.image, "update-image", "i", "", "Update a specific image")
	cmd.Flags().BoolVar(&opts.allImages, "update-all-images", false, "Update all images to latest versions")
	cmd.Flags().StringSliceVar(&opts.exclude, "exclude", []string{}, "List of controllers to exclude")
	cmd.Flags().BoolVar(&opts.dryRun, "dry-run", false, "Do not release anything; just report back what would have been done")
	cmd.Flags().BoolVar(&opts.interactive, "interactive", false, "Select interactively which containers to update")
	cmd.Flags().BoolVarP(&opts.force, "force", "f", false, "Disregard locks and container image filters (has no effect when used with --all or --update-all-images)")
	cmd.Flags().BoolVarP(&opts.watch, "watch", "w", false, "Watch rollout progress during release")

	// Deprecated
	cmd.Flags().StringSliceVarP(&opts.controllers, "controller", "c", []string{}, "List of controllers to release <namespace>:<kind>/<name>")
	cmd.Flags().MarkDeprecated("controller", "changed to --workload, use that instead")

	return cmd
}

func (opts *workloadReleaseOpts) RunE(cmd *cobra.Command, args []string) error {
	if len(args) != 0 {
		return errorWantedNoArgs
	}

	if err := checkExactlyOne("--update-image=<image> or --update-all-images", opts.image != "", opts.allImages); err != nil {
		return err
	}

	// Backwards compatibility with --controller until we remove it
	opts.workloads = append(opts.workloads, opts.controllers...)

	switch {
	case len(opts.workloads) <= 0 && !opts.allWorkloads:
		return newUsageError("please supply either --all, or at least one --workload=<workload>")
	case opts.watch && opts.dryRun:
		return newUsageError("cannot use --watch with --dry-run")
	case opts.force && opts.allWorkloads && opts.allImages:
		return newUsageError("--force has no effect when used with --all and --update-all-images")
	case opts.force && opts.allWorkloads:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore locked workloads when used with --all\n")
	case opts.force && opts.allImages:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore container image tags when used with --update-all-images\n")
	}

	var workloads []update.ResourceSpec

	if opts.allWorkloads {
		workloads = []update.ResourceSpec{update.ResourceSpecAll}
	} else {
		for _, workload := range opts.workloads {
			id, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, workload)
			if err != nil {
				return err
			}
			workloads = append(workloads, update.MakeResourceSpec(id))
		}
	}

	var (
		image update.ImageSpec
		err   error
	)
	switch {
	case opts.image != "":
		image, err = update.ParseImageSpec(opts.image)
		if err != nil {
			return err
		}
	case opts.allImages:
		image = update.ImageSpecLatest
	}

	var kind update.ReleaseKind = update.ReleaseKindExecute
	if opts.dryRun || opts.interactive {
		kind = update.ReleaseKindPlan
	}

	var excludes []flux.ResourceID
	for _, exclude := range opts.exclude {
		s, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, exclude)
		if err != nil {
			return err
		}
		excludes = append(excludes, s)
	}

	if kind == update.ReleaseKindPlan {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting dry-run release ...\n")
	} else {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting release ...\n")
	}

	ctx := context.Background()
	spec := update.ReleaseImageSpec{
		ServiceSpecs: workloads,
		ImageSpec:    image,
		Kind:         kind,
		Excludes:     excludes,
		Force:        opts.force,
	}
	jobID, err := opts.API.UpdateManifests(ctx, update.Spec{
		Type:  update.Images,
		Cause: opts.cause,
		Spec:  spec,
	})
	if err != nil {
		return err
	}

	result, err := awaitJob(ctx, opts.API, jobID)
	if err != nil {
		return err
	}
	if opts.interactive {
		spec, err := promptSpec(cmd.OutOrStdout(), result, opts.verbosity)
		spec.Force = opts.force
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		fmt.Fprintf(cmd.OutOrStderr(), "Submitting selected release ...\n")
		jobID, err = opts.API.UpdateManifests(ctx, update.Spec{
			Type:  update.Containers,
			Cause: opts.cause,
			Spec:  spec,
		})
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		opts.dryRun = false
	}

	err = await(ctx, cmd.OutOrStdout(), cmd.OutOrStderr(), opts.API, jobID, !opts.dryRun, opts.verbosity)
	if !opts.watch || err != nil {
		return err
	}

	fmt.Fprintf(cmd.OutOrStderr(), "Monitoring rollout ...\n")
	for {
		completed := 0
		workloads, err := opts.API.ListServicesWithOptions(ctx, v11.ListServicesOptions{Services: result.Result.AffectedResources()})
		if err != nil {
			return err
		}

		for _, workload := range workloads {
			writeRolloutStatus(workload, opts.verbosity)

			if workload.Status == cluster.StatusReady {
				completed++
			}

			if workload.Rollout.Messages != nil {
				fmt.Fprintf(cmd.OutOrStderr(), "There was a problem releasing %s:\n", workload.ID)
				for _, msg := range workload.Rollout.Messages {
					fmt.Fprintf(cmd.OutOrStderr(), "%s\n", msg)
				}
				return nil
			}
		}

		if completed == len(workloads) {
			fmt.Fprintf(cmd.OutOrStderr(), "All workloads ready.\n")
			return nil
		}

		time.Sleep(2000 * time.Millisecond)
	}
}

func writeRolloutStatus(workload v6.ControllerStatus, verbosity int) {
	w := newTabwriter()
	fmt.Fprintf(w, "WORKLOAD\tCONTAINER\tIMAGE\tRELEASE\tREPLICAS\n")

	if len(workload.Containers) > 0 {
		c := workload.Containers[0]
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d/%d", workload.ID, c.Name, c.Current.ID, workload.Status, workload.Rollout.Updated, workload.Rollout.Desired)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d ready)", workload.Rollout.Outdated, workload.Rollout.Ready)
		}
		fmt.Fprintf(w, "\n")
		for _, c := range workload.Containers[1:] {
			fmt.Fprintf(w, "\t%s\t%s\t\t\n", c.Name, c.Current.ID)
		}
	} else {
		fmt.Fprintf(w, "%s\t\t\t%s\t%d/%d", workload.ID, workload.Status, workload.Rollout.Updated, workload.Rollout.Desired)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d ready)", workload.Rollout.Outdated, workload.Rollout.Ready)
		}
		fmt.Fprintf(w, "\n")
	}
	fmt.Fprintln(w)
	w.Flush()
}

func promptSpec(out io.Writer, result job.Result, verbosity int) (update.ReleaseContainersSpec, error) {
	menu := update.NewMenu(out, result.Result, verbosity)
	containerSpecs, err := menu.Run()
	return update.ReleaseContainersSpec{
		Kind:           update.ReleaseKindExecute,
		ContainerSpecs: containerSpecs,
		SkipMismatches: false,
	}, err
}
<|endoftext|>