<commit_before>// This package implements a provisioner for Packer that executes
// Puppet on the remote machine, configured to apply a local manifest
// versus connecting to a Puppet master.
package puppetmasterless

import (
	"fmt"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"os"
	"path/filepath"
	"strings"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	tpl                 *packer.ConfigTemplate

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// Path to a hiera configuration file to upload and use.
	HieraConfigPath string `mapstructure:"hiera_config_path"`

	// An array of local paths of modules to upload.
	ModulePaths []string `mapstructure:"module_paths"`

	// The main manifest file to apply to kick off the entire thing.
	ManifestFile string `mapstructure:"manifest_file"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_directory"`
}

type Provisioner struct {
	config Config
}

type ExecuteTemplate struct {
	FacterVars         string
	HasHieraConfigPath bool
	HieraConfigPath    string
	ModulePath         string
	ManifestFile       string
	Sudo               bool
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	md, err := common.DecodeConfig(&p.config, raws...)
	if err != nil {
		return err
	}

	p.config.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	p.config.tpl.UserVars = p.config.PackerUserVars

	// Accumulate any errors
	errs := common.CheckUnusedConfig(md)

	// Set some defaults
	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "{{.FacterVars}} {{if .Sudo}} sudo -E {{end}}" +
			"puppet apply --verbose --modulepath='{{.ModulePath}}' " +
			"{{if .HasHieraConfigPath}}--hiera_config='{{.HieraConfigPath}}' {{end}}" +
			"--detailed-exitcodes " +
			"{{.ManifestFile}}"
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = "/tmp/packer-puppet-masterless"
	}

	// Templates
	templates := map[string]*string{
		"staging_dir": &p.config.StagingDir,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.config.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"module_paths": p.config.ModulePaths,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.config.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	validates := map[string]*string{
		"execute_command": &p.config.ExecuteCommand,
	}

	for n, ptr := range validates {
		if err := p.config.tpl.Validate(*ptr); err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error parsing %s: %s", n, err))
		}
	}

	newFacts := make(map[string]string)
	for k, v := range p.config.Facter {
		k, err := p.config.tpl.Process(k, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Error processing facter key %s: %s", k, err))
			continue
		}

		v, err := p.config.tpl.Process(v, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Error processing facter value '%s': %s", v, err))
			continue
		}

		newFacts[k] = v
	}

	p.config.Facter = newFacts

	// Validation
	if p.config.HieraConfigPath != "" {
		info, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path must point to a file"))
		}
	}

	if p.config.ManifestFile == "" {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("A manifest_file must be specified."))
	} else {
		info, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file must point to a file"))
		}
	}

	for i, path := range p.config.ModulePaths {
		info, err := os.Stat(path)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] is invalid: %s", i, err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] must point to a directory"))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload hiera config if set
	remoteHieraConfigPath := ""
	if p.config.HieraConfigPath != "" {
		var err error
		remoteHieraConfigPath, err = p.uploadHieraConfig(ui, comm)
		if err != nil {
			return fmt.Errorf("Error uploading hiera config: %s", err)
		}
	}

	// Upload all modules
	modulePaths := make([]string, 0, len(p.config.ModulePaths))
	for i, path := range p.config.ModulePaths {
		ui.Message(fmt.Sprintf("Uploading local modules from: %s", path))
		targetPath := fmt.Sprintf("%s/module-%d", p.config.StagingDir, i)
		if err := p.uploadDirectory(ui, comm, targetPath, path); err != nil {
			return fmt.Errorf("Error uploading modules: %s", err)
		}

		modulePaths = append(modulePaths, targetPath)
	}

	// Upload manifests
	ui.Message("Uploading manifests...")
	remoteManifestFile, err := p.uploadManifests(ui, comm)
	if err != nil {
		return fmt.Errorf("Error uploading manifests: %s", err)
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf("FACTER_%s='%s'", k, v))
	}

	// Execute Puppet
	command, err := p.config.tpl.Process(p.config.ExecuteCommand, &ExecuteTemplate{
		FacterVars:         strings.Join(facterVars, " "),
		HasHieraConfigPath: remoteHieraConfigPath != "",
		HieraConfigPath:    remoteHieraConfigPath,
		ManifestFile:       remoteManifestFile,
		ModulePath:         strings.Join(modulePaths, ":"),
		Sudo:               !p.config.PreventSudo,
	})
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) uploadHieraConfig(ui packer.Ui, comm packer.Communicator) (string, error) {
	ui.Message("Uploading hiera configuration...")
	f, err := os.Open(p.config.HieraConfigPath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	path := fmt.Sprintf("%s/hiera.yaml", p.config.StagingDir)
	if err := comm.Upload(path, f); err != nil {
		return "", err
	}

	return path, nil
}

func (p *Provisioner) uploadManifests(ui packer.Ui, comm packer.Communicator) (string, error) {
	// Create the remote manifests directory...
	ui.Message("Uploading manifests...")
	remoteManifestsPath := fmt.Sprintf("%s/manifests", p.config.StagingDir)
	if err := p.createDir(ui, comm, remoteManifestsPath); err != nil {
		return "", fmt.Errorf("Error creating manifests directory: %s", err)
	}

	// Upload the main manifest
	f, err := os.Open(p.config.ManifestFile)
	if err != nil {
		return "", err
	}
	defer f.Close()

	manifestFilename := filepath.Base(p.config.ManifestFile)
	remoteManifestFile := fmt.Sprintf("%s/%s", remoteManifestsPath, manifestFilename)
	if err := comm.Upload(remoteManifestFile, f); err != nil {
		return "", err
	}

	return remoteManifestFile, nil
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("mkdir -p '%s'", dir),
	}

	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<commit_msg>provisioner/puppet-masterless: remove duplicate manifest upload message<commit_after>// This package implements a provisioner for Packer that executes
// Puppet on the remote machine, configured to apply a local manifest
// versus connecting to a Puppet master.
package puppetmasterless

import (
	"fmt"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"os"
	"path/filepath"
	"strings"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	tpl                 *packer.ConfigTemplate

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// Path to a hiera configuration file to upload and use.
	HieraConfigPath string `mapstructure:"hiera_config_path"`

	// An array of local paths of modules to upload.
	ModulePaths []string `mapstructure:"module_paths"`

	// The main manifest file to apply to kick off the entire thing.
	ManifestFile string `mapstructure:"manifest_file"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_directory"`
}

type Provisioner struct {
	config Config
}

type ExecuteTemplate struct {
	FacterVars         string
	HasHieraConfigPath bool
	HieraConfigPath    string
	ModulePath         string
	ManifestFile       string
	Sudo               bool
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	md, err := common.DecodeConfig(&p.config, raws...)
	if err != nil {
		return err
	}

	p.config.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	p.config.tpl.UserVars = p.config.PackerUserVars

	// Accumulate any errors
	errs := common.CheckUnusedConfig(md)

	// Set some defaults
	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "{{.FacterVars}} {{if .Sudo}} sudo -E {{end}}" +
			"puppet apply --verbose --modulepath='{{.ModulePath}}' " +
			"{{if .HasHieraConfigPath}}--hiera_config='{{.HieraConfigPath}}' {{end}}" +
			"--detailed-exitcodes " +
			"{{.ManifestFile}}"
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = "/tmp/packer-puppet-masterless"
	}

	// Templates
	templates := map[string]*string{
		"staging_dir": &p.config.StagingDir,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.config.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"module_paths": p.config.ModulePaths,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.config.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	validates := map[string]*string{
		"execute_command": &p.config.ExecuteCommand,
	}

	for n, ptr := range validates {
		if err := p.config.tpl.Validate(*ptr); err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error parsing %s: %s", n, err))
		}
	}

	newFacts := make(map[string]string)
	for k, v := range p.config.Facter {
		k, err := p.config.tpl.Process(k, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Error processing facter key %s: %s", k, err))
			continue
		}

		v, err := p.config.tpl.Process(v, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Error processing facter value '%s': %s", v, err))
			continue
		}

		newFacts[k] = v
	}

	p.config.Facter = newFacts

	// Validation
	if p.config.HieraConfigPath != "" {
		info, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path must point to a file"))
		}
	}

	if p.config.ManifestFile == "" {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("A manifest_file must be specified."))
	} else {
		info, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file must point to a file"))
		}
	}

	for i, path := range p.config.ModulePaths {
		info, err := os.Stat(path)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] is invalid: %s", i, err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] must point to a directory"))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload hiera config if set
	remoteHieraConfigPath := ""
	if p.config.HieraConfigPath != "" {
		var err error
		remoteHieraConfigPath, err = p.uploadHieraConfig(ui, comm)
		if err != nil {
			return fmt.Errorf("Error uploading hiera config: %s", err)
		}
	}

	// Upload all modules
	modulePaths := make([]string, 0, len(p.config.ModulePaths))
	for i, path := range p.config.ModulePaths {
		ui.Message(fmt.Sprintf("Uploading local modules from: %s", path))
		targetPath := fmt.Sprintf("%s/module-%d", p.config.StagingDir, i)
		if err := p.uploadDirectory(ui, comm, targetPath, path); err != nil {
			return fmt.Errorf("Error uploading modules: %s", err)
		}

		modulePaths = append(modulePaths, targetPath)
	}

	// Upload manifests
	remoteManifestFile, err := p.uploadManifests(ui, comm)
	if err != nil {
		return fmt.Errorf("Error uploading manifests: %s", err)
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf("FACTER_%s='%s'", k, v))
	}

	// Execute Puppet
	command, err := p.config.tpl.Process(p.config.ExecuteCommand, &ExecuteTemplate{
		FacterVars:         strings.Join(facterVars, " "),
		HasHieraConfigPath: remoteHieraConfigPath != "",
		HieraConfigPath:    remoteHieraConfigPath,
		ManifestFile:       remoteManifestFile,
		ModulePath:         strings.Join(modulePaths, ":"),
		Sudo:               !p.config.PreventSudo,
	})
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) uploadHieraConfig(ui packer.Ui, comm packer.Communicator) (string, error) {
	ui.Message("Uploading hiera configuration...")
	f, err := os.Open(p.config.HieraConfigPath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	path := fmt.Sprintf("%s/hiera.yaml", p.config.StagingDir)
	if err := comm.Upload(path, f); err != nil {
		return "", err
	}

	return path, nil
}

func (p *Provisioner) uploadManifests(ui packer.Ui, comm packer.Communicator) (string, error) {
	// Create the remote manifests directory...
	ui.Message("Uploading manifests...")
	remoteManifestsPath := fmt.Sprintf("%s/manifests", p.config.StagingDir)
	if err := p.createDir(ui, comm, remoteManifestsPath); err != nil {
		return "", fmt.Errorf("Error creating manifests directory: %s", err)
	}

	// Upload the main manifest
	f, err := os.Open(p.config.ManifestFile)
	if err != nil {
		return "", err
	}
	defer f.Close()

	manifestFilename := filepath.Base(p.config.ManifestFile)
	remoteManifestFile := fmt.Sprintf("%s/%s", remoteManifestsPath, manifestFilename)
	if err := comm.Upload(remoteManifestFile, f); err != nil {
		return "", err
	}

	return remoteManifestFile, nil
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("mkdir -p '%s'", dir),
	}

	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<|endoftext|>


<commit_before>package messenger

import (
	"fmt"
	"net/http"
	"regexp"

	"encoding/json"

	"github.com/julienschmidt/httprouter"
	"github.com/meifamily/ptt-alertor/command"
	"github.com/meifamily/ptt-alertor/myutil"
	log "github.com/meifamily/logrus"
)

const (
	SendAPIURL = "https://graph.facebook.com/v2.6/me/messages?access_token="
	ProfileURL = "https://graph.facebook.com/v2.6/me/messenger_profile?access_token="
)

type Messenger struct {
	VerifyToken string
	AccessToken string
}

var config map[string]string

func init() {
	config = myutil.Config("messenger")
}

func New() Messenger {
	return Messenger{
		VerifyToken: config["verifyToken"],
		AccessToken: config["accessToken"],
	}
}

func (m *Messenger) Verify(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if r.FormValue("hub.mode") == "subscribe" && r.FormValue("hub.verify_token") == m.VerifyToken {
		log.Info("Validating webhook")
		resStr := r.FormValue("hub.challenge")
		fmt.Fprintln(w, resStr)
	} else {
		log.Info("Failed validation. Make sure the validation tokens match.")
		http.Error(w, http.StatusText(http.StatusForbidden), http.StatusForbidden)
	}
}

func (m *Messenger) Received(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	data := Webhook{}
	json.NewDecoder(r.Body).Decode(&data)
	if data.Object == "page" {
		for _, entry := range data.Entry {
			for _, messaging := range entry.Messaging {
				id := messaging.Sender.ID
				if messaging.Message != nil {
					text := messaging.Message.Text
					if text != "" {
						if match, _ := regexp.MatchString("^(刪除|刪除作者)+\\s.*\\*+", text); match {
							m.SendConfirmation(id, text)
							return
						}
						responseText := command.HandleCommand(text, id)
						m.SendTextMessage(id, responseText)
					}
				} else if messaging.Postback != nil {
					payload := messaging.Postback.Payload
					log.WithField("payload", payload).Info("Messenger Postback")
					m.handlePostback(id, payload)
				}
			}
		}
	}
}

func (m *Messenger) handlePostback(id string, payload string) {
	var responseText string
	switch payload {
	case "GET_STARTED_PAYLOAD":
		err := command.HandleMessengerFollow(id)
		if err != nil {
			log.WithError(err).Error("Messenger Follow Error")
		}
		responseText = "歡迎使用 Ptt Alertor\n輸入「指令」查看相關功能。\n\n觀看Demo:\nhttps://media.giphy.com/media/3ohzdF6vidM6I49lQs/giphy.gif"
	case "COMMANDS_PAYLOAD":
		// responseText = command.HandleCommand("指令", id)
		var str string
		commands := make(map[string]string)
		for cat, cmds := range command.Commands {
			for cmd, doc := range cmds {
				str += cmd
				if doc != "" {
					str += "：" + doc
				}
				str += "\n"
			}
			commands[cat] = str
			str = ""
		}
		m.SendListMessage(id, commands)
	case "SUBSCRIPTIONS_PAYLOAD":
		responseText = command.HandleCommand("清單", id)
	case "CANCEL":
		responseText = "取消"
	default:
		responseText = command.HandleCommand(payload, id)
	}
	m.SendTextMessage(id, responseText)
}

func (m *Messenger) SendTextMessage(id string, message string) {
	body := Request{
		Recipient{id},
		Message{Text: message},
	}
	m.callSendAPI(body)
}

func (m *Messenger) SendConfirmation(id string, cmd string) {
	attachment := Attachment{
		Type: "template",
		Payload: ButtonPayload{
			TemplateType: "button",
			Text:         "確認" + cmd,
			Buttons: Buttons{
				Button{"postback", "是", cmd},
				Button{"postback", "否", "CANCEL"},
			},
		},
	}
	body := Request{}
	body.Recipient.ID = id
	body.Message.Attachment = &attachment
	m.callSendAPI(body)
}

func (m *Messenger) SendQuickReplies(id string, payload string) {
	qrs := QuickReplies{
		QuickReply{"text", "是", payload},
		QuickReply{"text", "否", "CANCEL"},
	}
	body := Request{
		Recipient{id},
		Message{Text: "確定" + payload, QuickReplies: &qrs},
	}
	log.WithField("ID", id).Info("Messenger Confirmation Sent")
	m.callSendAPI(body)
}

func (m *Messenger) SendListMessage(id string, StringMap map[string]string) {
	elements := []Element{}
	for key, str := range StringMap {
		elements = append(elements, Element{
			Title:    key,
			Subtitle: str,
		})
	}
	attachment := Attachment{
		Type: "template",
		Payload: ListPayload{
			TemplateType:    "list",
			TopElementStyle: "compact",
			Elements:        elements,
		},
	}
	body := Request{}
	body.Recipient.ID = id
	body.Message.Attachment = &attachment
	m.callSendAPI(body)
}

func (m *Messenger) callSendAPI(body Request) {
	url := SendAPIURL + m.AccessToken
	err := callAPI(url, body)
	if err != nil {
		log.WithError(err).Error("Call Send API Error")
	}
}
<commit_msg>add messenger menu postback<commit_after>package messenger

import (
	"fmt"
	"net/http"
	"regexp"

	"encoding/json"

	"github.com/julienschmidt/httprouter"
	log "github.com/meifamily/logrus"
	"github.com/meifamily/ptt-alertor/command"
	"github.com/meifamily/ptt-alertor/myutil"
)

const (
	SendAPIURL = "https://graph.facebook.com/v2.6/me/messages?access_token="
	ProfileURL = "https://graph.facebook.com/v2.6/me/messenger_profile?access_token="
)

type Messenger struct {
	VerifyToken string
	AccessToken string
}

var config map[string]string

func init() {
	config = myutil.Config("messenger")
}

func New() Messenger {
	return Messenger{
		VerifyToken: config["verifyToken"],
		AccessToken: config["accessToken"],
	}
}

func (m *Messenger) Verify(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if r.FormValue("hub.mode") == "subscribe" && r.FormValue("hub.verify_token") == m.VerifyToken {
		log.Info("Validating webhook")
		resStr := r.FormValue("hub.challenge")
		fmt.Fprintln(w, resStr)
	} else {
		log.Info("Failed validation. Make sure the validation tokens match.")
		http.Error(w, http.StatusText(http.StatusForbidden), http.StatusForbidden)
	}
}

func (m *Messenger) Received(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	data := Webhook{}
	json.NewDecoder(r.Body).Decode(&data)
	if data.Object == "page" {
		for _, entry := range data.Entry {
			for _, messaging := range entry.Messaging {
				id := messaging.Sender.ID
				if messaging.Message != nil {
					text := messaging.Message.Text
					if text != "" {
						if match, _ := regexp.MatchString("^(刪除|刪除作者)+\\s.*\\*+", text); match {
							m.SendConfirmation(id, text)
							return
						}
						responseText := command.HandleCommand(text, id)
						m.SendTextMessage(id, responseText)
					}
				} else if messaging.Postback != nil {
					payload := messaging.Postback.Payload
					log.WithField("payload", payload).Info("Messenger Postback")
					m.handlePostback(id, payload)
				}
			}
		}
	}
}

func (m *Messenger) handlePostback(id string, payload string) {
	var responseText string
	switch payload {
	case "GET_STARTED_PAYLOAD":
		err := command.HandleMessengerFollow(id)
		if err != nil {
			log.WithError(err).Error("Messenger Follow Error")
		}
		responseText = "歡迎使用 Ptt Alertor\n輸入「指令」查看相關功能。\n\n觀看Demo:\nhttps://media.giphy.com/media/3ohzdF6vidM6I49lQs/giphy.gif"
	case "COMMANDS_PAYLOAD":
		// responseText = command.HandleCommand("指令", id)
		var str string
		commands := make(map[string]string)
		for cat, cmds := range command.Commands {
			for cmd, doc := range cmds {
				str += cmd
				if doc != "" {
					str += "：" + doc
				}
				str += "\n"
			}
			commands[cat] = str
			str = ""
		}
		m.SendListMessage(id, commands)
	case "SUBSCRIPTIONS_PAYLOAD":
		responseText = command.HandleCommand("清單", id)
	case "TOP_PAYLOAD":
		responseText = command.HandleCommand("排行", id)
	case "PUSHLIST_PAYLOAD":
		responseText = command.HandleCommand("推文清單", id)
	case "CANCEL":
		responseText = "取消"
	default:
		responseText = command.HandleCommand(payload, id)
	}
	m.SendTextMessage(id, responseText)
}

func (m *Messenger) SendTextMessage(id string, message string) {
	body := Request{
		Recipient{id},
		Message{Text: message},
	}
	m.callSendAPI(body)
}

func (m *Messenger) SendConfirmation(id string, cmd string) {
	attachment := Attachment{
		Type: "template",
		Payload: ButtonPayload{
			TemplateType: "button",
			Text:         "確認" + cmd,
			Buttons: Buttons{
				Button{"postback", "是", cmd},
				Button{"postback", "否", "CANCEL"},
			},
		},
	}
	body := Request{}
	body.Recipient.ID = id
	body.Message.Attachment = &attachment
	m.callSendAPI(body)
}

func (m *Messenger) SendQuickReplies(id string, payload string) {
	qrs := QuickReplies{
		QuickReply{"text", "是", payload},
		QuickReply{"text", "否", "CANCEL"},
	}
	body := Request{
		Recipient{id},
		Message{Text: "確定" + payload, QuickReplies: &qrs},
	}
	log.WithField("ID", id).Info("Messenger Confirmation Sent")
	m.callSendAPI(body)
}

func (m *Messenger) SendListMessage(id string, StringMap map[string]string) {
	elements := []Element{}
	for key, str := range StringMap {
		elements = append(elements, Element{
			Title:    key,
			Subtitle: str,
		})
	}
	attachment := Attachment{
		Type: "template",
		Payload: ListPayload{
			TemplateType:    "list",
			TopElementStyle: "compact",
			Elements:        elements,
		},
	}
	body := Request{}
	body.Recipient.ID = id
	body.Message.Attachment = &attachment
	m.callSendAPI(body)
}

func (m *Messenger) callSendAPI(body Request) {
	url := SendAPIURL + m.AccessToken
	err := callAPI(url, body)
	if err != nil {
		log.WithError(err).Error("Call Send API Error")
	}
}
<|endoftext|>


<commit_before>package models

import (
	"strings"
	"testing"

	"github.com/PuerkitoBio/goquery"
	"github.com/gophish/gophish/config"
	"gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { check.TestingT(t) }

type ModelsSuite struct{}

var _ = check.Suite(&ModelsSuite{})

func (s *ModelsSuite) SetUpSuite(c *check.C) {
	config.Conf.DBPath = ":memory:"
	config.Conf.MigrationsPath = "../db/migrations/"
	err := Setup()
	if err != nil {
		c.Fatalf("Failed creating database: %v", err)
	}
}

func (s *ModelsSuite) TearDownTest(c *check.C) {
	// Clear database tables between each test. If new tables are
	// used in this test suite they will need to be cleaned up here.
	db.Delete(Group{})
	db.Delete(Target{})
	db.Delete(GroupTarget{})
	db.Delete(SMTP{})
	db.Delete(Page{})

	// Reset users table to default state.
	db.Not("id", 1).Delete(User{})
	db.Model(User{}).Update("username", "admin")
}

func (s *ModelsSuite) TestGetUser(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(u.Username, check.Equals, "admin")
}

func (s *ModelsSuite) TestGeneratedAPIKey(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(u.ApiKey, check.Not(check.Equals), "12345678901234567890123456789012")
}

func (s *ModelsSuite) TestPostGroup(c *check.C) {
	g := Group{Name: "Test Group"}
	g.Targets = []Target{Target{Email: "test@example.com"}}
	g.UserId = 1
	err := PostGroup(&g)
	c.Assert(err, check.Equals, nil)
	c.Assert(g.Name, check.Equals, "Test Group")
	c.Assert(g.Targets[0].Email, check.Equals, "test@example.com")
}

func (s *ModelsSuite) TestPostGroupNoName(c *check.C) {
	g := Group{Name: ""}
	g.Targets = []Target{Target{Email: "test@example.com"}}
	g.UserId = 1
	err := PostGroup(&g)
	c.Assert(err, check.Equals, ErrGroupNameNotSpecified)
}

func (s *ModelsSuite) TestPostGroupNoTargets(c *check.C) {
	g := Group{Name: "No Target Group"}
	g.Targets = []Target{}
	g.UserId = 1
	err := PostGroup(&g)
	c.Assert(err, check.Equals, ErrNoTargetsSpecified)
}

func (s *ModelsSuite) TestPutGroup(c *check.C) {
	// Add test group.
	group := Group{Name: "Put Test Group"}
	group.Targets = []Target{
		Target{Email: "test1@example.com", FirstName: "First", LastName: "Example"},
		Target{Email: "test2@example.com", FirstName: "Second", LastName: "Example"},
	}
	group.UserId = 1
	PostGroup(&group)

	// Update one of group's targets.
	group.Targets[0].FirstName = "New"
	err := PutGroup(&group)
	c.Assert(err, check.Equals, nil)

	// Verify updated target information.
	targets, _ := GetTargets(group.Id)
	c.Assert(targets[0].Email, check.Equals, "test1@example.com")
	c.Assert(targets[0].FirstName, check.Equals, "New")
	c.Assert(targets[0].LastName, check.Equals, "Example")
	c.Assert(targets[1].Email, check.Equals, "test2@example.com")
	c.Assert(targets[1].FirstName, check.Equals, "Second")
	c.Assert(targets[1].LastName, check.Equals, "Example")
}

func (s *ModelsSuite) TestPutGroupEmptyAttribute(c *check.C) {
	// Add test group.
	group := Group{Name: "Put Empty Attribute Test Group"}
	group.Targets = []Target{
		Target{Email: "test3@example.com", FirstName: "Third", LastName: "Example"},
		Target{Email: "test4@example.com", FirstName: "Fourth", LastName: "Example"},
	}
	group.UserId = 1
	PostGroup(&group)

	// Update one of group's targets.
	group.Targets[0].FirstName = ""
	err := PutGroup(&group)
	c.Assert(err, check.Equals, nil)

	// Verify updated empty attribute was saved.
	targets, _ := GetTargets(group.Id)
	c.Assert(targets[0].Email, check.Equals, "test3@example.com")
	c.Assert(targets[0].FirstName, check.Equals, "")
	c.Assert(targets[0].LastName, check.Equals, "Example")
	c.Assert(targets[1].Email, check.Equals, "test4@example.com")
	c.Assert(targets[1].FirstName, check.Equals, "Fourth")
	c.Assert(targets[1].LastName, check.Equals, "Example")
}

func (s *ModelsSuite) TestPostSMTP(c *check.C) {
	smtp := SMTP{
		Name:        "Test SMTP",
		Host:        "1.1.1.1:25",
		FromAddress: "Foo Bar <foo@example.com>",
		UserId:      1,
	}
	err = PostSMTP(&smtp)
	c.Assert(err, check.Equals, nil)
	ss, err := GetSMTPs(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(len(ss), check.Equals, 1)
}

func (s *ModelsSuite) TestPostSMTPNoHost(c *check.C) {
	smtp := SMTP{
		Name:        "Test SMTP",
		FromAddress: "Foo Bar <foo@example.com>",
		UserId:      1,
	}
	err = PostSMTP(&smtp)
	c.Assert(err, check.Equals, ErrHostNotSpecified)
}

func (s *ModelsSuite) TestPostSMTPNoFrom(c *check.C) {
	smtp := SMTP{
		Name:   "Test SMTP",
		UserId: 1,
		Host:   "1.1.1.1:25",
	}
	err = PostSMTP(&smtp)
	c.Assert(err, check.Equals, ErrFromAddressNotSpecified)
}

func (s *ModelsSuite) TestPostPage(c *check.C) {
	html := `

			<form action="example.com">
				<input name="username"/>
				<input name="password" type="password"/>
			</form>
		  `
	p := Page{
		Name:        "Test Page",
		HTML:        html,
		RedirectURL: "http://example.com",
	}
	// Check the capturing credentials and passwords
	p.CaptureCredentials = true
	p.CapturePasswords = true
	err := PostPage(&p)
	c.Assert(err, check.Equals, nil)
	c.Assert(p.RedirectURL, check.Equals, "http://example.com")
	d, err := goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
	c.Assert(err, check.Equals, nil)
	forms := d.Find("form")
	forms.Each(func(i int, f *goquery.Selection) {
		// Check the action has been set
		a, _ := f.Attr("action")
		c.Assert(a, check.Equals, "")
		// Check the password still has a name
		_, ok := f.Find("input[type=\"password\"]").Attr("name")
		c.Assert(ok, check.Equals, true)
		// Check the username is still correct
		u, ok := f.Find("input").Attr("name")
		c.Assert(ok, check.Equals, true)
		c.Assert(u, check.Equals, "username")
	})
	// Check what happens when we don't capture passwords
	p.CapturePasswords = false
	p.HTML = html
	p.RedirectURL = ""
	err = PutPage(&p)
	c.Assert(err, check.Equals, nil)
	c.Assert(p.RedirectURL, check.Equals, "")
	d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
	c.Assert(err, check.Equals, nil)
	forms = d.Find("form")
	forms.Each(func(i int, f *goquery.Selection) {
		// Check the action has been set
		a, _ := f.Attr("action")
		c.Assert(a, check.Equals, "")
		// Check the password still has a name
		_, ok := f.Find("input[type=\"password\"]").Attr("name")
		c.Assert(ok, check.Equals, false)
		// Check the username is still correct
		u, ok := f.Find("input").Attr("name")
		c.Assert(ok, check.Equals, true)
		c.Assert(u, check.Equals, "username")
	})
	// Finally, check when we don't capture credentials
	p.CaptureCredentials = false
	p.HTML = html
	err = PutPage(&p)
	c.Assert(err, check.Equals, nil)
	d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
	c.Assert(err, check.Equals, nil)
	forms = d.Find("form")
	forms.Each(func(i int, f *goquery.Selection) {
		// Check the action has been set
		a, _ := f.Attr("action")
		c.Assert(a, check.Equals, "")
		// Check the password still has a name
		_, ok := f.Find("input[type=\"password\"]").Attr("name")
		c.Assert(ok, check.Equals, false)
		// Check the username is still correct
		_, ok = f.Find("input").Attr("name")
		c.Assert(ok, check.Equals, false)
	})
}

func (s *ModelsSuite) TestPutUser(c *check.C) {
	u, err := GetUser(1)
	u.Username = "admin_changed"
	err = PutUser(&u)
	c.Assert(err, check.Equals, nil)
	u, err = GetUser(1)
	c.Assert(u.Username, check.Equals, "admin_changed")
}
<commit_msg>Put User tests together<commit_after>package models

import (
	"strings"
	"testing"

	"github.com/PuerkitoBio/goquery"
	"github.com/gophish/gophish/config"
	"gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { check.TestingT(t) }

type ModelsSuite struct{}

var _ = check.Suite(&ModelsSuite{})

func (s *ModelsSuite) SetUpSuite(c *check.C) {
	config.Conf.DBPath = ":memory:"
	config.Conf.MigrationsPath = "../db/migrations/"
	err := Setup()
	if err != nil {
		c.Fatalf("Failed creating database: %v", err)
	}
}

func (s *ModelsSuite) TearDownTest(c *check.C) {
	// Clear database tables between each test. If new tables are
	// used in this test suite they will need to be cleaned up here.
	db.Delete(Group{})
	db.Delete(Target{})
	db.Delete(GroupTarget{})
	db.Delete(SMTP{})
	db.Delete(Page{})

	// Reset users table to default state.
	db.Not("id", 1).Delete(User{})
	db.Model(User{}).Update("username", "admin")
}

func (s *ModelsSuite) TestGetUser(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(u.Username, check.Equals, "admin")
}

func (s *ModelsSuite) TestPutUser(c *check.C) {
	u, err := GetUser(1)
	u.Username = "admin_changed"
	err = PutUser(&u)
	c.Assert(err, check.Equals, nil)
	u, err = GetUser(1)
	c.Assert(u.Username, check.Equals, "admin_changed")
}

func (s *ModelsSuite) TestGeneratedAPIKey(c *check.C) {
	u, err := GetUser(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(u.ApiKey, check.Not(check.Equals), "12345678901234567890123456789012")
}

func (s *ModelsSuite) TestPostGroup(c *check.C) {
	g := Group{Name: "Test Group"}
	g.Targets = []Target{Target{Email: "test@example.com"}}
	g.UserId = 1
	err := PostGroup(&g)
	c.Assert(err, check.Equals, nil)
	c.Assert(g.Name, check.Equals, "Test Group")
	c.Assert(g.Targets[0].Email, check.Equals, "test@example.com")
}

func (s *ModelsSuite) TestPostGroupNoName(c *check.C) {
	g := Group{Name: ""}
	g.Targets = []Target{Target{Email: "test@example.com"}}
	g.UserId = 1
	err := PostGroup(&g)
	c.Assert(err, check.Equals, ErrGroupNameNotSpecified)
}

func (s *ModelsSuite) TestPostGroupNoTargets(c *check.C) {
	g := Group{Name: "No Target Group"}
	g.Targets = []Target{}
	g.UserId = 1
	err := PostGroup(&g)
	c.Assert(err, check.Equals, ErrNoTargetsSpecified)
}

func (s *ModelsSuite) TestPutGroup(c *check.C) {
	// Add test group.
	group := Group{Name: "Put Test Group"}
	group.Targets = []Target{
		Target{Email: "test1@example.com", FirstName: "First", LastName: "Example"},
		Target{Email: "test2@example.com", FirstName: "Second", LastName: "Example"},
	}
	group.UserId = 1
	PostGroup(&group)

	// Update one of group's targets.
	group.Targets[0].FirstName = "New"
	err := PutGroup(&group)
	c.Assert(err, check.Equals, nil)

	// Verify updated target information.
	targets, _ := GetTargets(group.Id)
	c.Assert(targets[0].Email, check.Equals, "test1@example.com")
	c.Assert(targets[0].FirstName, check.Equals, "New")
	c.Assert(targets[0].LastName, check.Equals, "Example")
	c.Assert(targets[1].Email, check.Equals, "test2@example.com")
	c.Assert(targets[1].FirstName, check.Equals, "Second")
	c.Assert(targets[1].LastName, check.Equals, "Example")
}

func (s *ModelsSuite) TestPutGroupEmptyAttribute(c *check.C) {
	// Add test group.
	group := Group{Name: "Put Empty Attribute Test Group"}
	group.Targets = []Target{
		Target{Email: "test3@example.com", FirstName: "Third", LastName: "Example"},
		Target{Email: "test4@example.com", FirstName: "Fourth", LastName: "Example"},
	}
	group.UserId = 1
	PostGroup(&group)

	// Update one of group's targets.
	group.Targets[0].FirstName = ""
	err := PutGroup(&group)
	c.Assert(err, check.Equals, nil)

	// Verify updated empty attribute was saved.
	targets, _ := GetTargets(group.Id)
	c.Assert(targets[0].Email, check.Equals, "test3@example.com")
	c.Assert(targets[0].FirstName, check.Equals, "")
	c.Assert(targets[0].LastName, check.Equals, "Example")
	c.Assert(targets[1].Email, check.Equals, "test4@example.com")
	c.Assert(targets[1].FirstName, check.Equals, "Fourth")
	c.Assert(targets[1].LastName, check.Equals, "Example")
}

func (s *ModelsSuite) TestPostSMTP(c *check.C) {
	smtp := SMTP{
		Name:        "Test SMTP",
		Host:        "1.1.1.1:25",
		FromAddress: "Foo Bar <foo@example.com>",
		UserId:      1,
	}
	err = PostSMTP(&smtp)
	c.Assert(err, check.Equals, nil)
	ss, err := GetSMTPs(1)
	c.Assert(err, check.Equals, nil)
	c.Assert(len(ss), check.Equals, 1)
}

func (s *ModelsSuite) TestPostSMTPNoHost(c *check.C) {
	smtp := SMTP{
		Name:        "Test SMTP",
		FromAddress: "Foo Bar <foo@example.com>",
		UserId:      1,
	}
	err = PostSMTP(&smtp)
	c.Assert(err, check.Equals, ErrHostNotSpecified)
}

func (s *ModelsSuite) TestPostSMTPNoFrom(c *check.C) {
	smtp := SMTP{
		Name:   "Test SMTP",
		UserId: 1,
		Host:   "1.1.1.1:25",
	}
	err = PostSMTP(&smtp)
	c.Assert(err, check.Equals, ErrFromAddressNotSpecified)
}

func (s *ModelsSuite) TestPostPage(c *check.C) {
	html := `

			<form action="example.com">
				<input name="username"/>
				<input name="password" type="password"/>
			</form>
		  `
	p := Page{
		Name:        "Test Page",
		HTML:        html,
		RedirectURL: "http://example.com",
	}
	// Check the capturing credentials and passwords
	p.CaptureCredentials = true
	p.CapturePasswords = true
	err := PostPage(&p)
	c.Assert(err, check.Equals, nil)
	c.Assert(p.RedirectURL, check.Equals, "http://example.com")
	d, err := goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
	c.Assert(err, check.Equals, nil)
	forms := d.Find("form")
	forms.Each(func(i int, f *goquery.Selection) {
		// Check the action has been set
		a, _ := f.Attr("action")
		c.Assert(a, check.Equals, "")
		// Check the password still has a name
		_, ok := f.Find("input[type=\"password\"]").Attr("name")
		c.Assert(ok, check.Equals, true)
		// Check the username is still correct
		u, ok := f.Find("input").Attr("name")
		c.Assert(ok, check.Equals, true)
		c.Assert(u, check.Equals, "username")
	})
	// Check what happens when we don't capture passwords
	p.CapturePasswords = false
	p.HTML = html
	p.RedirectURL = ""
	err = PutPage(&p)
	c.Assert(err, check.Equals, nil)
	c.Assert(p.RedirectURL, check.Equals, "")
	d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
	c.Assert(err, check.Equals, nil)
	forms = d.Find("form")
	forms.Each(func(i int, f *goquery.Selection) {
		// Check the action has been set
		a, _ := f.Attr("action")
		c.Assert(a, check.Equals, "")
		// Check the password still has a name
		_, ok := f.Find("input[type=\"password\"]").Attr("name")
		c.Assert(ok, check.Equals, false)
		// Check the username is still correct
		u, ok := f.Find("input").Attr("name")
		c.Assert(ok, check.Equals, true)
		c.Assert(u, check.Equals, "username")
	})
	// Finally, check when we don't capture credentials
	p.CaptureCredentials = false
	p.HTML = html
	err = PutPage(&p)
	c.Assert(err, check.Equals, nil)
	d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
	c.Assert(err, check.Equals, nil)
	forms = d.Find("form")
	forms.Each(func(i int, f *goquery.Selection) {
		// Check the action has been set
		a, _ := f.Attr("action")
		c.Assert(a, check.Equals, "")
		// Check the password still has a name
		_, ok := f.Find("input[type=\"password\"]").Attr("name")
		c.Assert(ok, check.Equals, false)
		// Check the username is still correct
		_, ok = f.Find("input").Attr("name")
		c.Assert(ok, check.Equals, false)
	})
}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/tcnksm/gcli/skeleton"
)

// ApplyCommand is a Command that generates a new cli project
type ApplyCommand struct {
	Meta
}

// Run generates a new cli project. It returns exit code
func (c *ApplyCommand) Run(args []string) int {

	var (
		frameworkStr string
		skipTest     bool
		verbose      bool
		owner        string
		name         string
	)

	uflag := flag.NewFlagSet("apply", flag.ContinueOnError)
	uflag.Usage = func() { c.UI.Error(c.Help()) }

	uflag.StringVar(&frameworkStr, "framework", "", "framework")
	uflag.StringVar(&frameworkStr, "F", "", "framework (short)")

	uflag.BoolVar(&skipTest, "skip-test", false, "skip-test")
	uflag.BoolVar(&skipTest, "T", false, "skip-test (short)")

	uflag.BoolVar(&verbose, "verbose", false, "verbose")
	uflag.BoolVar(&verbose, "V", false, "verbose (short)")

	// These flags are supposed only to use in test
	uflag.StringVar(&owner, "owner", "", "owner (Should only for test)")
	uflag.StringVar(&name, "name", "", "name (Should only for test)")

	errR, errW := io.Pipe()
	errScanner := bufio.NewScanner(errR)
	uflag.SetOutput(errW)

	go func() {
		for errScanner.Scan() {
			c.UI.Error(errScanner.Text())
		}
	}()

	if err := uflag.Parse(args); err != nil {
		return 1
	}

	parsedArgs := uflag.Args()
	if len(parsedArgs) != 1 {
		c.UI.Error("Invalid argument: Usage glic apply [options] FILE")
		return 1
	}

	planFile := parsedArgs[0]
	c.UI.Info(fmt.Sprintf(
		"Use plan file %q for generating new cli project", planFile))

	// Check file is exist or not
	if _, err := os.Stat(planFile); os.IsNotExist(err) {
		c.UI.Error(fmt.Sprintf(
			"Plan file does not exsit"))
		return 1
	}

	// Decode plan file as skeleton.Executable
	executable := skeleton.NewExecutable()
	if _, err := toml.DecodeFile(planFile, executable); err != nil {
		c.UI.Error(fmt.Sprintf(
			"Failed to decode plan file %q: %s", planFile, err))
		return 1
	}

	output := executable.Name
	if _, err := os.Stat(output); !os.IsNotExist(err) {
		msg := fmt.Sprintf("Cannot create directory %s: file exists", output)
		c.UI.Error(msg)
		return 1
	}

	// Check option input first and if it's specified use it
	if len(frameworkStr) == 0 {
		if len(executable.FrameworkStr) != 0 {
			// If FrameworStr is specified from design file use it
			frameworkStr = executable.FrameworkStr
		} else {
			frameworkStr = defaultFrameworkString
		}
	}

	fmt.Println(frameworkStr)
	framework, err := skeleton.FrameworkByName(frameworkStr)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Failed to generate %q: %s", executable.Name, err.Error()))
		return 1
	}

	if len(name) != 0 {
		executable.Name = name
		output = name
	}

	if len(owner) != 0 {
		executable.Owner = owner
	}

	// Channels to receive artifact path (result) and error
	artifactCh, errCh := make(chan string), make(chan error)

	// Define Skeleton
	skeleton := &skeleton.Skeleton{
		Path:       output,
		Framework:  framework,
		SkipTest:   skipTest,
		Executable: executable,
		ArtifactCh: artifactCh,
		ErrCh:      errCh,
		Verbose:    verbose,
		LogWriter:  os.Stdout,
	}

	// Create project directory
	doneCh := skeleton.Generate()

	for {
		select {
		case artifact := <-artifactCh:
			c.UI.Output(fmt.Sprintf("  Created %s", artifact))
		case err := <-errCh:
			c.UI.Error(fmt.Sprintf("Failed to generate %s: %s", output, err.Error()))

			// If some file are created before error happend
			// Should be cleanuped
			if _, err := os.Stat(output); !os.IsNotExist(err) {
				c.UI.Output(fmt.Sprintf("Cleanup %s", output))
				os.RemoveAll(output)
			}
			return ExitCodeFailed
		case <-doneCh:
			c.UI.Info(fmt.Sprintf("====> Successfully generated %s", executable.Name))
			return ExitCodeOK
		}
	}
}

// Synopsis is a one-line, short synopsis of the command.
func (c *ApplyCommand) Synopsis() string {
	return "Apply design template file for generating cli project"
}

// Help is a long-form help text that includes the command-line
// usage, a brief few sentences explaining the function of the command,
// and the complete list of flags the command accepts.
func (c *ApplyCommand) Help() string {
	helpText := `
Usage: gcli apply [option] FILE

  Apply design template file for generating cli project. You can generate
  design template file via 'gcli design' command. If framework name is not
  specified gcli use codegangsta/cli. You can set framework name via '-F'
  option. To check cli framework you can use, run 'gcli list'.

Options:

   -framework=name, -F        Cli framework name. By default, gcli use "codegangsta/cli"
                              To check cli framework you can use, run 'gcli list'.
                              If you set invalid framework, it will be failed.

   -skip-test, -T             Skip generating *_test.go file. By default, gcli generates
                              test file If you specify this flag, gcli will not generate
                              test files.
`
	return strings.TrimSpace(helpText)
}
<commit_msg>Fix output message<commit_after>package command

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/tcnksm/gcli/skeleton"
)

// ApplyCommand is a Command that generates a new cli project
type ApplyCommand struct {
	Meta
}

// Run generates a new cli project. It returns exit code
func (c *ApplyCommand) Run(args []string) int {

	var (
		frameworkStr string
		skipTest     bool
		verbose      bool
		owner        string
		name         string
	)

	uflag := flag.NewFlagSet("apply", flag.ContinueOnError)
	uflag.Usage = func() { c.UI.Error(c.Help()) }

	uflag.StringVar(&frameworkStr, "framework", "", "framework")
	uflag.StringVar(&frameworkStr, "F", "", "framework (short)")

	uflag.BoolVar(&skipTest, "skip-test", false, "skip-test")
	uflag.BoolVar(&skipTest, "T", false, "skip-test (short)")

	uflag.BoolVar(&verbose, "verbose", false, "verbose")
	uflag.BoolVar(&verbose, "V", false, "verbose (short)")

	// These flags are supposed only to use in test
	uflag.StringVar(&owner, "owner", "", "owner (Should only for test)")
	uflag.StringVar(&name, "name", "", "name (Should only for test)")

	errR, errW := io.Pipe()
	errScanner := bufio.NewScanner(errR)
	uflag.SetOutput(errW)

	go func() {
		for errScanner.Scan() {
			c.UI.Error(errScanner.Text())
		}
	}()

	if err := uflag.Parse(args); err != nil {
		return 1
	}

	parsedArgs := uflag.Args()
	if len(parsedArgs) != 1 {
		c.UI.Error("Invalid argument: Usage glic apply [options] FILE")
		return 1
	}

	designFile := parsedArgs[0]
	c.UI.Info(fmt.Sprintf(
		"Use design template %q for generating new cli project", designFile))

	// Check file is exist or not
	if _, err := os.Stat(designFile); os.IsNotExist(err) {
		c.UI.Error(fmt.Sprintf(
			"Design file does not exsit"))
		return 1
	}

	// Decode design file as skeleton.Executable
	executable := skeleton.NewExecutable()
	if _, err := toml.DecodeFile(designFile, executable); err != nil {
		c.UI.Error(fmt.Sprintf(
			"Failed to decode design file %q: %s", designFile, err))
		return 1
	}

	output := executable.Name
	if _, err := os.Stat(output); !os.IsNotExist(err) {
		msg := fmt.Sprintf("Cannot create directory %s: file exists", output)
		c.UI.Error(msg)
		return 1
	}

	// Check option input first and if it's specified use it
	if len(frameworkStr) == 0 {
		if len(executable.FrameworkStr) != 0 {
			// If FrameworStr is specified from design file use it
			frameworkStr = executable.FrameworkStr
		} else {
			frameworkStr = defaultFrameworkString
		}
	}

	fmt.Println(frameworkStr)
	framework, err := skeleton.FrameworkByName(frameworkStr)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Failed to generate %q: %s", executable.Name, err.Error()))
		return 1
	}

	if len(name) != 0 {
		executable.Name = name
		output = name
	}

	if len(owner) != 0 {
		executable.Owner = owner
	}

	// Channels to receive artifact path (result) and error
	artifactCh, errCh := make(chan string), make(chan error)

	// Define Skeleton
	skeleton := &skeleton.Skeleton{
		Path:       output,
		Framework:  framework,
		SkipTest:   skipTest,
		Executable: executable,
		ArtifactCh: artifactCh,
		ErrCh:      errCh,
		Verbose:    verbose,
		LogWriter:  os.Stdout,
	}

	// Create project directory
	doneCh := skeleton.Generate()

	for {
		select {
		case artifact := <-artifactCh:
			c.UI.Output(fmt.Sprintf("  Created %s", artifact))
		case err := <-errCh:
			c.UI.Error(fmt.Sprintf("Failed to generate %s: %s", output, err.Error()))

			// If some file are created before error happend
			// Should be cleanuped
			if _, err := os.Stat(output); !os.IsNotExist(err) {
				c.UI.Output(fmt.Sprintf("Cleanup %s", output))
				os.RemoveAll(output)
			}
			return ExitCodeFailed
		case <-doneCh:
			c.UI.Info(fmt.Sprintf("====> Successfully generated %s", executable.Name))
			return ExitCodeOK
		}
	}
}

// Synopsis is a one-line, short synopsis of the command.
func (c *ApplyCommand) Synopsis() string {
	return "Apply design template file for generating cli project"
}

// Help is a long-form help text that includes the command-line
// usage, a brief few sentences explaining the function of the command,
// and the complete list of flags the command accepts.
func (c *ApplyCommand) Help() string {
	helpText := `
Usage: gcli apply [option] FILE

  Apply design template file for generating cli project. You can generate
  design template file via 'gcli design' command. If framework name is not
  specified gcli use codegangsta/cli. You can set framework name via '-F'
  option. To check cli framework you can use, run 'gcli list'.

Options:

   -framework=name, -F        Cli framework name. By default, gcli use "codegangsta/cli"
                              To check cli framework you can use, run 'gcli list'.
                              If you set invalid framework, it will be failed.

   -skip-test, -T             Skip generating *_test.go file. By default, gcli generates
                              test file If you specify this flag, gcli will not generate
                              test files.
`
	return strings.TrimSpace(helpText)
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"github.com/hashicorp/serf/client"
	"github.com/hashicorp/serf/command/agent"
	"github.com/mitchellh/cli"
	"strings"
	"time"
)

// QueryCommand is a Command implementation that is used to trigger a new
// query and wait for responses and acks
type QueryCommand struct {
	ShutdownCh <-chan struct{}
	Ui         cli.Ui
}

func (c *QueryCommand) Help() string {
	helpText := `
Usage: serf query [options] name payload

  Dispatches a query to the Serf cluster.

Options:

  -node=NAME                This flag can be provided multiple times to filter
                            responses to only named nodes.

  -tag key=regexp           This flag can be provided multiple times to filter
                            responses to only nodes matching the tags

  -timeout="15s"            Providing a timeout overrides the default timeout

  -no-ack                   Setting this prevents nodes from sending an acknowledgement
                            of the query

  -rpc-addr=127.0.0.1:7373  RPC address of the Serf agent.
`
	return strings.TrimSpace(helpText)
}

func (c *QueryCommand) Run(args []string) int {
	var noAck bool
	var nodes []string
	var tags []string
	var timeout time.Duration
	cmdFlags := flag.NewFlagSet("event", flag.ContinueOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }
	cmdFlags.Var((*agent.AppendSliceValue)(&nodes), "node", "node filter")
	cmdFlags.Var((*agent.AppendSliceValue)(&tags), "tag", "tag filter")
	cmdFlags.DurationVar(&timeout, "timeout", 0, "query timeout")
	cmdFlags.BoolVar(&noAck, "no-ack", false, "no-ack")
	rpcAddr := RPCAddrFlag(cmdFlags)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// Setup the filter tags
	filterTags := make(map[string]string)
	for _, tag := range tags {
		parts := strings.SplitN(tag, "=", 2)
		if len(parts) != 2 {
			c.Ui.Error(fmt.Sprintf("Invalid tag '%s' provided", tag))
			return 1
		}
		filterTags[parts[0]] = parts[1]
	}

	args = cmdFlags.Args()
	if len(args) < 1 {
		c.Ui.Error("A query name must be specified.")
		c.Ui.Error("")
		c.Ui.Error(c.Help())
		return 1
	} else if len(args) > 2 {
		c.Ui.Error("Too many command line arguments. Only a name and payload must be specified.")
		c.Ui.Error("")
		c.Ui.Error(c.Help())
		return 1
	}

	name := args[0]
	var payload []byte
	if len(args) == 2 {
		payload = []byte(args[1])
	}

	cl, err := RPCClient(*rpcAddr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error connecting to Serf agent: %s", err))
		return 1
	}
	defer cl.Close()

	ackCh := make(chan string, 128)
	respCh := make(chan client.NodeResponse, 128)

	params := client.QueryParam{
		FilterNodes: nodes,
		FilterTags:  filterTags,
		RequestAck:  !noAck,
		Timeout:     timeout,
		Name:        name,
		Payload:     payload,
		AckCh:       ackCh,
		RespCh:      respCh,
	}
	if err := cl.Query(&params); err != nil {
		c.Ui.Error(fmt.Sprintf("Error sending query: %s", err))
		return 1
	}
	c.Ui.Output(fmt.Sprintf("Query '%s' dispatched", name))

	// Track responses and acknowledgements
	numAcks := 0
	numResp := 0

OUTER:
	for {
		select {
		case a := <-ackCh:
			if a == "" {
				break OUTER
			}
			numAcks++
			c.Ui.Info(fmt.Sprintf("Ack from '%s'", a))

		case r := <-respCh:
			if r.From == "" {
				break OUTER
			}
			numResp++
			c.Ui.Info(fmt.Sprintf("Response from '%s':", r.From, r.Payload))

		case <-c.ShutdownCh:
			return 1
		}
	}

	if !noAck {
		c.Ui.Output(fmt.Sprintf("Total Acks: %d", numAcks))
	}
	c.Ui.Output(fmt.Sprintf("Total Responses: %d", numResp))
	return 0
}

func (c *QueryCommand) Synopsis() string {
	return "Send a query to the Serf cluster"
}
<commit_msg>command/query: remove the trailing newline from a response<commit_after>package command

import (
	"flag"
	"fmt"
	"github.com/hashicorp/serf/client"
	"github.com/hashicorp/serf/command/agent"
	"github.com/mitchellh/cli"
	"strings"
	"time"
)

// QueryCommand is a Command implementation that is used to trigger a new
// query and wait for responses and acks
type QueryCommand struct {
	ShutdownCh <-chan struct{}
	Ui         cli.Ui
}

func (c *QueryCommand) Help() string {
	helpText := `
Usage: serf query [options] name payload

  Dispatches a query to the Serf cluster.

Options:

  -node=NAME                This flag can be provided multiple times to filter
                            responses to only named nodes.

  -tag key=regexp           This flag can be provided multiple times to filter
                            responses to only nodes matching the tags

  -timeout="15s"            Providing a timeout overrides the default timeout

  -no-ack                   Setting this prevents nodes from sending an acknowledgement
                            of the query

  -rpc-addr=127.0.0.1:7373  RPC address of the Serf agent.
`
	return strings.TrimSpace(helpText)
}

func (c *QueryCommand) Run(args []string) int {
	var noAck bool
	var nodes []string
	var tags []string
	var timeout time.Duration
	cmdFlags := flag.NewFlagSet("event", flag.ContinueOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }
	cmdFlags.Var((*agent.AppendSliceValue)(&nodes), "node", "node filter")
	cmdFlags.Var((*agent.AppendSliceValue)(&tags), "tag", "tag filter")
	cmdFlags.DurationVar(&timeout, "timeout", 0, "query timeout")
	cmdFlags.BoolVar(&noAck, "no-ack", false, "no-ack")
	rpcAddr := RPCAddrFlag(cmdFlags)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// Setup the filter tags
	filterTags := make(map[string]string)
	for _, tag := range tags {
		parts := strings.SplitN(tag, "=", 2)
		if len(parts) != 2 {
			c.Ui.Error(fmt.Sprintf("Invalid tag '%s' provided", tag))
			return 1
		}
		filterTags[parts[0]] = parts[1]
	}

	args = cmdFlags.Args()
	if len(args) < 1 {
		c.Ui.Error("A query name must be specified.")
		c.Ui.Error("")
		c.Ui.Error(c.Help())
		return 1
	} else if len(args) > 2 {
		c.Ui.Error("Too many command line arguments. Only a name and payload must be specified.")
		c.Ui.Error("")
		c.Ui.Error(c.Help())
		return 1
	}

	name := args[0]
	var payload []byte
	if len(args) == 2 {
		payload = []byte(args[1])
	}

	cl, err := RPCClient(*rpcAddr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error connecting to Serf agent: %s", err))
		return 1
	}
	defer cl.Close()

	ackCh := make(chan string, 128)
	respCh := make(chan client.NodeResponse, 128)

	params := client.QueryParam{
		FilterNodes: nodes,
		FilterTags:  filterTags,
		RequestAck:  !noAck,
		Timeout:     timeout,
		Name:        name,
		Payload:     payload,
		AckCh:       ackCh,
		RespCh:      respCh,
	}
	if err := cl.Query(&params); err != nil {
		c.Ui.Error(fmt.Sprintf("Error sending query: %s", err))
		return 1
	}
	c.Ui.Output(fmt.Sprintf("Query '%s' dispatched", name))

	// Track responses and acknowledgements
	numAcks := 0
	numResp := 0

OUTER:
	for {
		select {
		case a := <-ackCh:
			if a == "" {
				break OUTER
			}
			numAcks++
			c.Ui.Info(fmt.Sprintf("Ack from '%s'", a))

		case r := <-respCh:
			if r.From == "" {
				break OUTER
			}
			numResp++

			// Remove the trailing newline if there is one
			payload := r.Payload
			if n := len(payload); n > 0 && payload[n-1] == '\n' {
				payload = payload[:n-1]
			}

			c.Ui.Info(fmt.Sprintf("Response from '%s': %s", r.From, payload))

		case <-c.ShutdownCh:
			return 1
		}
	}

	if !noAck {
		c.Ui.Output(fmt.Sprintf("Total Acks: %d", numAcks))
	}
	c.Ui.Output(fmt.Sprintf("Total Responses: %d", numResp))
	return 0
}

func (c *QueryCommand) Synopsis() string {
	return "Send a query to the Serf cluster"
}
<|endoftext|>


<commit_before>package file

import (
	"fmt"
	"log"

	"github.com/miekg/coredns/middleware"

	"github.com/miekg/dns"
	"golang.org/x/net/context"
)

type (
	Xfr struct {
		*Zone
	}
)

// Serve an AXFR (or maybe later an IXFR) as well.
func (x Xfr) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
	state := middleware.State{W: w, Req: r}
	if !x.TransferAllowed(state) {
		return dns.RcodeServerFailure, nil
	}
	if state.QType() != dns.TypeAXFR {
		return 0, fmt.Errorf("xfr called with non transfer type: %d", state.QType())
	}

	records := x.All()
	if len(records) == 0 {
		return dns.RcodeServerFailure, nil
	}

	ch := make(chan *dns.Envelope)
	defer close(ch)
	tr := new(dns.Transfer)
	go tr.Out(w, r, ch)

	j, l := 0, 0
	records = append(records, records[0]) // add closing SOA to the end
	log.Printf("[INFO] Outgoing transfer of %d records of zone %s to %s started", len(records), x.name, state.IP())
	for i, r := range records {
		l += dns.Len(r)
		if l > transferLength {
			ch <- &dns.Envelope{RR: records[j:i]}
			l = 0
			j = i
		}
	}
	if j < len(records) {
		ch <- &dns.Envelope{RR: records[j:]}
	}

	w.Hijack()
	// w.Close() // Client closes connection
	return dns.RcodeSuccess, nil
}

const transferLength = 1000 // Start a new envelop after message reaches this size in bytes. Intentionally small to test multi envelope parsing.
<commit_msg>Allow IXFR as well (#97)<commit_after>package file

import (
	"fmt"
	"log"

	"github.com/miekg/coredns/middleware"

	"github.com/miekg/dns"
	"golang.org/x/net/context"
)

type (
	Xfr struct {
		*Zone
	}
)

// Serve an AXFR (and fallback of IXFR) as well.
func (x Xfr) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
	state := middleware.State{W: w, Req: r}
	if !x.TransferAllowed(state) {
		return dns.RcodeServerFailure, nil
	}
	if state.QType() != dns.TypeAXFR && state.QType() != dns.TypeIXFR {
		return 0, fmt.Errorf("xfr called with non transfer type: %d", state.QType())
	}

	records := x.All()
	if len(records) == 0 {
		return dns.RcodeServerFailure, nil
	}

	ch := make(chan *dns.Envelope)
	defer close(ch)
	tr := new(dns.Transfer)
	go tr.Out(w, r, ch)

	j, l := 0, 0
	records = append(records, records[0]) // add closing SOA to the end
	log.Printf("[INFO] Outgoing transfer of %d records of zone %s to %s started", len(records), x.name, state.IP())
	for i, r := range records {
		l += dns.Len(r)
		if l > transferLength {
			ch <- &dns.Envelope{RR: records[j:i]}
			l = 0
			j = i
		}
	}
	if j < len(records) {
		ch <- &dns.Envelope{RR: records[j:]}
	}

	w.Hijack()
	// w.Close() // Client closes connection
	return dns.RcodeSuccess, nil
}

const transferLength = 1000 // Start a new envelop after message reaches this size in bytes. Intentionally small to test multi envelope parsing.
<|endoftext|>


<commit_before>package middleware

import (
	"net/http"
	"time"
)

// responseRecorder is a type of ResponseWriter that captures
// the status code written to it and also the size of the body
// written in the response. A status code does not have
// to be written, however, in which case 200 must be assumed.
// It is best to have the constructor initialize this type
// with that default status code.
type responseRecorder struct {
	http.ResponseWriter
	status int
	size   int
	start  time.Time
}

// NewResponseRecorder makes and returns a new responseRecorder,
// which captures the HTTP Status code from the ResponseWriter
// and also the length of the response body written through it.
// Because a status is not set unless WriteHeader is called
// explicitly, this constructor initializes with a status code
// of 200 to cover the default case.
func NewResponseRecorder(w http.ResponseWriter) *responseRecorder {
	return &responseRecorder{
		ResponseWriter: w,
		status:         http.StatusOK,
		start:          time.Now(),
	}
}

// WriteHeader records the status code and calls the
// underlying ResponseWriter's WriteHeader method.
func (r *responseRecorder) WriteHeader(status int) {
	r.status = status
	r.ResponseWriter.WriteHeader(status)
}

// Write is a wrapper that records the size of the body
// that gets written.
func (r *responseRecorder) Write(buf []byte) (int, error) {
	n, err := r.ResponseWriter.Write(buf)
	if err == nil {
		r.size += n
	}
	return n, err
}
<commit_msg>let middleware.responseRecorder be a http.Hijacker<commit_after>package middleware

import (
	"bufio"
	"errors"
	"net"
	"net/http"
	"time"
)

// responseRecorder is a type of ResponseWriter that captures
// the status code written to it and also the size of the body
// written in the response. A status code does not have
// to be written, however, in which case 200 must be assumed.
// It is best to have the constructor initialize this type
// with that default status code.
type responseRecorder struct {
	http.ResponseWriter
	status int
	size   int
	start  time.Time
}

// NewResponseRecorder makes and returns a new responseRecorder,
// which captures the HTTP Status code from the ResponseWriter
// and also the length of the response body written through it.
// Because a status is not set unless WriteHeader is called
// explicitly, this constructor initializes with a status code
// of 200 to cover the default case.
func NewResponseRecorder(w http.ResponseWriter) *responseRecorder {
	return &responseRecorder{
		ResponseWriter: w,
		status:         http.StatusOK,
		start:          time.Now(),
	}
}

// WriteHeader records the status code and calls the
// underlying ResponseWriter's WriteHeader method.
func (r *responseRecorder) WriteHeader(status int) {
	r.status = status
	r.ResponseWriter.WriteHeader(status)
}

// Write is a wrapper that records the size of the body
// that gets written.
func (r *responseRecorder) Write(buf []byte) (int, error) {
	n, err := r.ResponseWriter.Write(buf)
	if err == nil {
		r.size += n
	}
	return n, err
}

// Hijacker is a wrapper of http.Hijacker underearth if any,
// otherwise it just returns an error.
func (r *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := r.ResponseWriter.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, errors.New("I'm not a Hijacker")
}
<|endoftext|>


<commit_before>package common

import (
	"fmt"
	"runtime"
)

const VerBinary = "0.0.8"

func VerString(app string) string {
	return fmt.Sprintf("%s v%s (built w/%s)", app, VerBinary, runtime.Version())
}
<commit_msg>bump version<commit_after>package common

import (
	"fmt"
	"runtime"
)

const VerBinary = "0.1.0"

func VerString(app string) string {
	return fmt.Sprintf("%s v%s (built w/%s)", app, VerBinary, runtime.Version())
}
<|endoftext|>


<commit_before>// Copyright 2017 Gruppe 12 IS-105. All rights reserved.
package functions

import (
	"fmt"
	"math"
	"os"
	"strconv"
)

func LogaritmeFunc() {
	base := os.Args[1]
	tall := os.Args[2]
	tallFloat, _ := strconv.ParseFloat(tall, 64)
	baseFloat, _ := strconv.Atoi(base)

	if baseFloat == 2 {
		fmt.Println("Logaritmen av ",tall, " blir: ", math.Log2(tallFloat), ", beregnet med base 2.")
	} else if baseFloat == 10 {
		fmt.Println("Logaritmen av ",tall, " blir: ", math.Log10(tallFloat), ", beregnet med base 10.")
	} else {
		fmt.Println("Første argumentet må være 2 eller 10")
	}
}
<commit_msg>Update log.go<commit_after>// Copyright 2017 Gruppe 12 IS-105. All rights reserved.
package functions

import (
	"fmt"
	"math"
	"os"
	"strconv"
)

func LogaritmeFunc() {
	base := os.Args[1]
	tall := os.Args[2]
	tallFloat, _ := strconv.ParseFloat(tall, 64)
	baseFloat, _ := strconv.Atoi(base)

	if baseFloat == 2 {
		fmt.Println("Logaritmen av ",tall, " blir: ", math.Log2(tallFloat), ", beregnet med base 2.")
	} else if baseFloat == 10 {
		fmt.Println("Logaritmen av ",tall, " blir: ", math.Log10(tallFloat), ", beregnet med base 10.")
	}
		else {
		fmt.Println("Første argumentet må være 2 eller 10")
	}

}
<|endoftext|>


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

package heartbeat

import (
	"sync"
	"time"

	"github.com/cloudfoundry-community/stackdriver-tools/src/stackdriver-nozzle/messages"
	"github.com/cloudfoundry-community/stackdriver-tools/src/stackdriver-nozzle/stackdriver"
	"github.com/cloudfoundry/lager"
)

type metricHandler struct {
	start      time.Time
	logger     lager.Logger
	ma         stackdriver.MetricAdapter
	nozzleId   string
	nozzleName string
	nozzleZone string

	counterMu *sync.Mutex // Guards counter
	counter   map[string]uint
}

func NewMetricHandler(ma stackdriver.MetricAdapter, logger lager.Logger, nozzleId, nozzleName, nozzleZone string) *metricHandler {
	return &metricHandler{
		logger:     logger,
		ma:         ma,
		nozzleId:   nozzleId,
		nozzleName: nozzleName,
		nozzleZone: nozzleZone,
		start:      time.Now(),
		counterMu:  &sync.Mutex{},
		counter:    map[string]uint{},
	}
}

func (h *metricHandler) Name() string {
	return "metricHandler"
}

func (h *metricHandler) Handle(event string, count uint) {
	h.counterMu.Lock()
	defer h.counterMu.Unlock()
	h.counter[event] += count
	return
}

func (h *metricHandler) Flush() error {
	h.counterMu.Lock()
	defer h.counterMu.Unlock()

	metrics := []*messages.Metric{}
	t := time.Now()
	labels := map[string]string{
		"instance": h.nozzleName,
		"zone":     h.nozzleZone,
	}

	for k, v := range h.counter {
		metrics = append(metrics, &messages.Metric{
			Name:      "heartbeat." + k,
			Value:     float64(v),
			EventTime: t,
		})
	}
	h.counter = map[string]uint{}
	return h.ma.PostMetricEvents([]*messages.MetricEvent{{Labels: labels, Metrics: metrics}})
}
<commit_msg>heartbeat/metric_handler: don't block incomming during flush<commit_after>/*
 * Copyright 2017 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package heartbeat

import (
	"sync"
	"time"

	"github.com/cloudfoundry-community/stackdriver-tools/src/stackdriver-nozzle/messages"
	"github.com/cloudfoundry-community/stackdriver-tools/src/stackdriver-nozzle/stackdriver"
	"github.com/cloudfoundry/lager"
)

type metricHandler struct {
	start      time.Time
	logger     lager.Logger
	ma         stackdriver.MetricAdapter
	nozzleId   string
	nozzleName string
	nozzleZone string

	counterMu *sync.Mutex // Guards counter
	counter   map[string]uint
}

func NewMetricHandler(ma stackdriver.MetricAdapter, logger lager.Logger, nozzleId, nozzleName, nozzleZone string) *metricHandler {
	return &metricHandler{
		logger:     logger,
		ma:         ma,
		nozzleId:   nozzleId,
		nozzleName: nozzleName,
		nozzleZone: nozzleZone,
		start:      time.Now(),
		counterMu:  &sync.Mutex{},
		counter:    map[string]uint{},
	}
}

func (h *metricHandler) Name() string {
	return "metricHandler"
}

func (h *metricHandler) Handle(event string, count uint) {
	h.counterMu.Lock()
	defer h.counterMu.Unlock()
	h.counter[event] += count
	return
}

func (h *metricHandler) Flush() error {
	counter := h.flushInternal()

	metrics := []*messages.Metric{}
	t := time.Now()
	labels := map[string]string{
		"instance": h.nozzleName,
		"zone":     h.nozzleZone,
	}

	for k, v := range counter {
		metrics = append(metrics, &messages.Metric{
			Name:      "heartbeat." + k,
			Value:     float64(v),
			EventTime: t,
		})
	}
	return h.ma.PostMetricEvents([]*messages.MetricEvent{{Labels: labels, Metrics: metrics}})
}

func (h *metricHandler) flushInternal() map[string]uint {
	h.counterMu.Lock()
	defer h.counterMu.Unlock()

	counter := h.counter
	h.counter = map[string]uint{}
	return counter
}
<|endoftext|>


<commit_before>// Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

package eventhandler

import (
	"container/list"
	"sync"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/utils"
	"github.com/cihub/seelog"
)

// Maximum number of tasks that may be handled at once by the TaskHandler
const concurrentEventCalls = 3

type eventList struct {
	// events is a list of *sendableEvents
	events *list.List
	// sending will check whether the list is already being handlerd
	sending bool
	//eventsListLock locks both the list and sending bool
	eventListLock sync.Mutex
}

// TaskHandler encapsulates the the map of a task arn to task and container events
// associated with said task
type TaskHandler struct {
	// submitSemaphore for the number of tasks that may be handled at once
	submitSemaphore utils.Semaphore
	// taskToEvents is arn:*eventList map so events may be serialized per task
	tasksToEvents map[string]*eventList
	// tasksToEventsLock for locking the map
	tasksToEventsLock sync.RWMutex
}

// NewTaskHandler returns a pointer to TaskHandler
func NewTaskHandler() *TaskHandler {
	return &TaskHandler{
		tasksToEvents:   make(map[string]*eventList),
		submitSemaphore: utils.NewSemaphore(concurrentEventCalls),
	}
}

// AddTaskEvent queues up a state change for sending using the given client.
func (handler *TaskHandler) AddTaskEvent(change api.TaskStateChange, client api.ECSClient) {
	handler.addEvent(newSendableTaskEvent(change), client)
}

// AddContainerEvent queues up a state change for sending using the given client.
func (handler *TaskHandler) AddContainerEvent(change api.ContainerStateChange, client api.ECSClient) {
	handler.addEvent(newSendableContainerEvent(change), client)
}

// Prepares a given event to be sent by adding it to the handler's appropriate
// eventList
func (handler *TaskHandler) addEvent(change *sendableEvent, client api.ECSClient) {

	seelog.Info("TaskHandler, Adding event: ", change)

	eventsToSubmit := handler.getTaskEventList(change)

	eventsToSubmit.eventListLock.Lock()
	defer eventsToSubmit.eventListLock.Unlock()

	// Update taskEvent
	eventsToSubmit.events.PushBack(change)

	if !eventsToSubmit.sending {
		eventsToSubmit.sending = true
		go handler.SubmitTaskEvents(eventsToSubmit, client)
	}
}

// getTaskEventList gets the eventList from taskToEvent map, and reduces the
// scope of the taskToEventsLock to just this function
func (handler *TaskHandler) getTaskEventList(change *sendableEvent) (eventsToSubmit *eventList) {

	handler.tasksToEventsLock.Lock()
	defer handler.tasksToEventsLock.Unlock()

	eventsToSubmit, ok := handler.tasksToEvents[change.taskArn()]

	if !ok {
		seelog.Debug("TaskHandler, New event: ", change)
		eventsToSubmit = &eventList{events: list.New(), sending: false}
		handler.tasksToEvents[change.taskArn()] = eventsToSubmit
	}

	return eventsToSubmit
}

// Continuously retries sending an event until it succeeds, sleeping between each
// attempt
func (handler *TaskHandler) SubmitTaskEvents(eventsToSubmit *eventList, client api.ECSClient) {
	backoff := utils.NewSimpleBackoff(1*time.Second, 30*time.Second, 0.20, 1.3)

	// Mirror events.sending, but without the need to lock since this is local
	// to our goroutine
	done := false

	for !done {
		// If we looped back up here, we successfully submitted an event, but
		// we haven't emptied the list so we should keep submitting
		backoff.Reset()
		utils.RetryWithBackoff(backoff, func() error {
			// Lock and unlock within this function, allowing the list to be added
			// to while we're not actively sending an event
			seelog.Debug("TaskHandler, Waiting on semaphore to send...")
			handler.submitSemaphore.Wait()
			defer handler.submitSemaphore.Post()

			seelog.Debug("TaskHandler, Aquiring lock for sending event...")
			eventsToSubmit.eventListLock.Lock()
			defer eventsToSubmit.eventListLock.Unlock()
			seelog.Debug("TaskHandler, Aquired lock!")

			var err error

			if eventsToSubmit.events.Len() == 0 {
				seelog.Debug("TaskHandler, No events left; not retrying more")

				eventsToSubmit.sending = false
				done = true
				return nil
			}

			eventToSubmit := eventsToSubmit.events.Front()
			event := eventToSubmit.Value.(*sendableEvent)

			if event.containerShouldBeSent() {
				seelog.Info("TaskHandler, Sending container change: ", event)
				err = client.SubmitContainerStateChange(event.containerChange)
				if err == nil {
					// submitted; ensure we don't retry it
					event.containerSent = true
					if event.containerChange.Container != nil {
						event.containerChange.Container.SetSentStatus(event.containerChange.Status)
					}
					statesaver.Save()
					seelog.Debug("TaskHandler, Submitted container state change")
					backoff.Reset()
					eventsToSubmit.events.Remove(eventToSubmit)
				} else {
					seelog.Error("TaskHandler, Unretriable error submitting container state change ", err)
				}
			} else if event.taskShouldBeSent() {
				seelog.Info("TaskHandler, Sending task change: ", event)
				err = client.SubmitTaskStateChange(event.taskChange)
				if err == nil {
					// submitted or can't be retried; ensure we don't retry it
					event.taskSent = true
					if event.taskChange.Task != nil {
						event.taskChange.Task.SetSentStatus(event.taskChange.Status)
					}
					statesaver.Save()
					seelog.Debug("TaskHandler, Submitted task state change")
					backoff.Reset()
					eventsToSubmit.events.Remove(eventToSubmit)
				} else {
					seelog.Error("TaskHandler, Unretriable error submitting container state change: ", err)
				}
			} else {
				// Shouldn't be sent as either a task or container change event; must have been already sent
				seelog.Info("TaskHandler, Not submitting redundant event; just removing")
				eventsToSubmit.events.Remove(eventToSubmit)
			}

			if eventsToSubmit.events.Len() == 0 {
				seelog.Debug("TaskHandler, Removed the last element, no longer sending")
				eventsToSubmit.sending = false
				done = true
				return nil
			}

			return err
		})
	}
}
<commit_msg>eventhandler: address cr, formatting, naming<commit_after>// Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

package eventhandler

import (
	"container/list"
	"sync"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/utils"
	"github.com/cihub/seelog"
)

// Maximum number of tasks that may be handled at once by the TaskHandler
const concurrentEventCalls = 3

type eventList struct {
	// events is a list of *sendableEvents
	events *list.List
	// sending will check whether the list is already being handlerd
	sending bool
	//eventsListLock locks both the list and sending bool
	eventListLock sync.Mutex
}

// TaskHandler encapsulates the the map of a task arn to task and container events
// associated with said task
type TaskHandler struct {
	// submitSemaphore for the number of tasks that may be handled at once
	submitSemaphore utils.Semaphore
	// taskToEvents is arn:*eventList map so events may be serialized per task
	//TODO: fix leak, currently items are never removed from this map
	tasksToEvents map[string]*eventList
	// tasksToEventsLock for locking the map
	tasksToEventsLock sync.RWMutex
}

// NewTaskHandler returns a pointer to TaskHandler
func NewTaskHandler() *TaskHandler {
	return &TaskHandler{
		tasksToEvents:   make(map[string]*eventList),
		submitSemaphore: utils.NewSemaphore(concurrentEventCalls),
	}
}

// AddTaskEvent queues up a state change for sending using the given client.
func (handler *TaskHandler) AddTaskEvent(change api.TaskStateChange, client api.ECSClient) {
	handler.addEvent(newSendableTaskEvent(change), client)
}

// AddContainerEvent queues up a state change for sending using the given client.
func (handler *TaskHandler) AddContainerEvent(change api.ContainerStateChange, client api.ECSClient) {
	handler.addEvent(newSendableContainerEvent(change), client)
}

// Prepares a given event to be sent by adding it to the handler's appropriate
// eventList
func (handler *TaskHandler) addEvent(change *sendableEvent, client api.ECSClient) {
	seelog.Info("TaskHandler, Adding event: ", change)

	taskEvents := handler.getTaskEventList(change)

	taskEvents.eventListLock.Lock()
	defer taskEvents.eventListLock.Unlock()

	// Update taskEvent
	taskEvents.events.PushBack(change)

	if !taskEvents.sending {
		taskEvents.sending = true
		go handler.SubmitTaskEvents(taskEvents, client)
	}
}

// getTaskEventList gets the eventList from taskToEvent map, and reduces the
// scope of the taskToEventsLock to just this function
func (handler *TaskHandler) getTaskEventList(change *sendableEvent) (taskEvents *eventList) {
	handler.tasksToEventsLock.Lock()
	defer handler.tasksToEventsLock.Unlock()

	taskEvents, ok := handler.tasksToEvents[change.taskArn()]
	if !ok {
		seelog.Debug("TaskHandler, collecting events for new task ", change)
		taskEvents = &eventList{events: list.New(), sending: false}
		handler.tasksToEvents[change.taskArn()] = taskEvents
	}

	return taskEvents
}

// Continuously retries sending an event until it succeeds, sleeping between each
// attempt
func (handler *TaskHandler) SubmitTaskEvents(taskEvents *eventList, client api.ECSClient) {
	backoff := utils.NewSimpleBackoff(1*time.Second, 30*time.Second, 0.20, 1.3)

	// Mirror events.sending, but without the need to lock since this is local
	// to our goroutine
	done := false

	for !done {
		// If we looped back up here, we successfully submitted an event, but
		// we haven't emptied the list so we should keep submitting
		backoff.Reset()
		utils.RetryWithBackoff(backoff, func() error {
			// Lock and unlock within this function, allowing the list to be added
			// to while we're not actively sending an event
			seelog.Debug("TaskHandler, Waiting on semaphore to send...")
			handler.submitSemaphore.Wait()
			defer handler.submitSemaphore.Post()

			seelog.Debug("TaskHandler, Aquiring lock for sending event...")
			taskEvents.eventListLock.Lock()
			defer taskEvents.eventListLock.Unlock()
			seelog.Debug("TaskHandler, Aquired lock!")

			var err error

			if taskEvents.events.Len() == 0 {
				seelog.Debug("TaskHandler, No events left; not retrying more")

				taskEvents.sending = false
				done = true
				return nil
			}

			eventToSubmit := taskEvents.events.Front()
			event := eventToSubmit.Value.(*sendableEvent)

			if event.containerShouldBeSent() {
				seelog.Info("TaskHandler, Sending container change: ", event)
				err = client.SubmitContainerStateChange(event.containerChange)
				if err == nil {
					// submitted; ensure we don't retry it
					event.containerSent = true
					if event.containerChange.Container != nil {
						event.containerChange.Container.SetSentStatus(event.containerChange.Status)
					}
					statesaver.Save()
					seelog.Debug("TaskHandler, Submitted container state change")
					backoff.Reset()
					taskEvents.events.Remove(eventToSubmit)
				} else {
					seelog.Error("TaskHandler, Unretriable error submitting container state change ", err)
				}
			} else if event.taskShouldBeSent() {
				seelog.Info("TaskHandler, Sending task change: ", event)
				err = client.SubmitTaskStateChange(event.taskChange)
				if err == nil {
					// submitted or can't be retried; ensure we don't retry it
					event.taskSent = true
					if event.taskChange.Task != nil {
						event.taskChange.Task.SetSentStatus(event.taskChange.Status)
					}
					statesaver.Save()
					seelog.Debug("TaskHandler, Submitted task state change")
					backoff.Reset()
					taskEvents.events.Remove(eventToSubmit)
				} else {
					seelog.Error("TaskHandler, Unretriable error submitting container state change: ", err)
				}
			} else {
				// Shouldn't be sent as either a task or container change event; must have been already sent
				seelog.Info("TaskHandler, Not submitting redundant event; just removing")
				taskEvents.events.Remove(eventToSubmit)
			}

			if taskEvents.events.Len() == 0 {
				seelog.Debug("TaskHandler, Removed the last element, no longer sending")
				taskEvents.sending = false
				done = true
				return nil
			}

			return err
		})
	}
}
<|endoftext|>


<commit_before>package common

import (
	"fmt"
	"image/color"
	"sort"

	"engo.io/ecs"
	"engo.io/engo"
	"engo.io/gl"
	"sync"
)

const (
	RenderSystemPriority = -1000
)

type renderChangeMessage struct{}

func (renderChangeMessage) Type() string {
	return "renderChangeMessage"
}

type Drawable interface {
	Texture() *gl.Texture
	Width() float32
	Height() float32
	View() (float32, float32, float32, float32)
	Close()
}

type TextureRepeating uint8

const (
	ClampToEdge TextureRepeating = iota
	ClampToBorder
	Repeat
	MirroredRepeat
)

type RenderComponent struct {
	// Hidden is used to prevent drawing by OpenGL
	Hidden bool
	// Scale is the scale at which to render, in the X and Y axis. Not defining Scale, will default to engo.Point{1, 1}
	Scale engo.Point
	// Color defines how much of the color-components of the texture get used
	Color color.Color
	// Drawable refers to the Texture that should be drawn
	Drawable Drawable
	// Repeat defines how to repeat the Texture if the viewport of the texture is larger than the texture itself
	Repeat TextureRepeating

	shader Shader
	zIndex float32

	buffer        *gl.Buffer
	bufferContent []float32
}

func (r *RenderComponent) SetShader(s Shader) {
	r.shader = s
	engo.Mailbox.Dispatch(&renderChangeMessage{})
}

func (r *RenderComponent) SetZIndex(index float32) {
	r.zIndex = index
	engo.Mailbox.Dispatch(&renderChangeMessage{})
}

type renderEntity struct {
	*ecs.BasicEntity
	*RenderComponent
	*SpaceComponent
}

type renderEntityList []renderEntity

func (r renderEntityList) Len() int {
	return len(r)
}

func (r renderEntityList) Less(i, j int) bool {
	// Sort by shader-pointer if they have the same zIndex
	if r[i].RenderComponent.zIndex == r[j].RenderComponent.zIndex {
		// TODO: optimize this for performance
		return fmt.Sprintf("%p", r[i].RenderComponent.shader) < fmt.Sprintf("%p", r[j].RenderComponent.shader)
	}

	return r[i].RenderComponent.zIndex < r[j].RenderComponent.zIndex
}

func (r renderEntityList) Swap(i, j int) {
	r[i], r[j] = r[j], r[i]
}

type RenderSystem struct {
	entities renderEntityList
	world    *ecs.World

	sortingNeeded bool
	currentShader Shader
}

func (*RenderSystem) Priority() int { return RenderSystemPriority }

func (rs *RenderSystem) New(w *ecs.World) {
	rs.world = w

	addCameraSystemOnce(w)

	if !engo.Headless() {
		initShaders(w)
		engo.Gl.Enable(engo.Gl.MULTISAMPLE)
	}

	engo.Mailbox.Listen("renderChangeMessage", func(engo.Message) {
		rs.sortingNeeded = true
	})
}

var cameraInitMutex sync.Mutex
func addCameraSystemOnce(w *ecs.World) {
	cameraInitMutex.Lock()
	defer cameraInitMutex.Unlock()

	camSystemAlreadyAdded := false
	for _, system := range w.Systems() {
		switch system.(type) {
		case *CameraSystem:
			camSystemAlreadyAdded = true
		}
	}
	if !camSystemAlreadyAdded {
		w.AddSystem(&CameraSystem{})
	}
}

func (rs *RenderSystem) Add(basic *ecs.BasicEntity, render *RenderComponent, space *SpaceComponent) {
	if rs.EntityExists(basic) == -1 {
		// Setting default shader
		if render.shader == nil {
			switch render.Drawable.(type) {
			case Triangle:
				render.shader = LegacyShader
			case Circle:
				render.shader = LegacyShader
			case Rectangle:
				render.shader = LegacyShader
			case ComplexTriangles:
				render.shader = LegacyShader
			case Text:
				render.shader = TextShader
			default:
				render.shader = DefaultShader
			}
		}

		// This is to prevent users from using the wrong one
		if render.shader == HUDShader {
			switch render.Drawable.(type) {
			case Triangle:
				render.shader = LegacyHUDShader
			case Circle:
				render.shader = LegacyHUDShader
			case Rectangle:
				render.shader = LegacyHUDShader
			case ComplexTriangles:
				render.shader = LegacyHUDShader
			case Text:
				render.shader = TextHUDShader
			default:
				render.shader = HUDShader
			}
		}

		rs.entities = append(rs.entities, renderEntity{basic, render, space})
		rs.sortingNeeded = true
	}
}

// EntityExists looks if the entity is already into the System's entities. It will return the index >= 0 of the object into de rs.entities or -1 if it could not be found.
func (rs *RenderSystem) EntityExists(basic *ecs.BasicEntity) int {
	for index, entity := range rs.entities {
		if entity.ID() == basic.ID() {
			return index
			break
		}
	}

	return -1
}

// AddByInterface adds any Renderable to the render system. Any Entity containing a BasicEntity,RenderComponent, and SpaceComponent anonymously does this automatically
func (rs *RenderSystem) AddByInterface(o Renderable) {
	rs.Add(o.GetBasicEntity(), o.GetRenderComponent(), o.GetSpaceComponent())
}

func (rs *RenderSystem) Remove(basic ecs.BasicEntity) {
	var delete int = rs.EntityExists(&basic)
	if delete >= 0 {
		rs.entities = append(rs.entities[:delete], rs.entities[delete+1:]...)
		rs.sortingNeeded = true
	}
}

func (rs *RenderSystem) Update(dt float32) {
	if engo.Headless() {
		return
	}

	if rs.sortingNeeded {
		sort.Sort(rs.entities)
		rs.sortingNeeded = false
	}

	engo.Gl.Clear(engo.Gl.COLOR_BUFFER_BIT)

	// TODO: it's linear for now, but that might very well be a bad idea
	for _, e := range rs.entities {
		if e.RenderComponent.Hidden {
			continue // with other entities
		}

		// Retrieve a shader, may be the default one -- then use it if we aren't already using it
		shader := e.RenderComponent.shader

		// Change Shader if we have to
		if shader != rs.currentShader {
			if rs.currentShader != nil {
				rs.currentShader.Post()
			}
			shader.Pre()
			rs.currentShader = shader
		}

		// Setting default scale to 1
		if e.RenderComponent.Scale.X == 0 && e.RenderComponent.Scale.Y == 0 {
			e.RenderComponent.Scale = engo.Point{1, 1}
		}

		// Setting default to white
		if e.RenderComponent.Color == nil {
			e.RenderComponent.Color = color.White
		}

		rs.currentShader.Draw(e.RenderComponent, e.SpaceComponent)
	}

	if rs.currentShader != nil {
		rs.currentShader.Post()
		rs.currentShader = nil
	}
}

func SetBackground(c color.Color) {
	if !engo.Headless() {
		r, g, b, a := c.RGBA()

		engo.Gl.ClearColor(float32(r)/0xffff, float32(g)/0xffff, float32(b)/0xffff, float32(a)/0xffff)
	}
}
<commit_msg>Update render.go<commit_after>package common

import (
	"fmt"
	"image/color"
	"sort"

	"engo.io/ecs"
	"engo.io/engo"
	"engo.io/gl"
	"sync"
)

const (
	RenderSystemPriority = -1000
)

type renderChangeMessage struct{}

func (renderChangeMessage) Type() string {
	return "renderChangeMessage"
}

type Drawable interface {
	Texture() *gl.Texture
	Width() float32
	Height() float32
	View() (float32, float32, float32, float32)
	Close()
}

type TextureRepeating uint8

const (
	ClampToEdge TextureRepeating = iota
	ClampToBorder
	Repeat
	MirroredRepeat
)

type RenderComponent struct {
	// Hidden is used to prevent drawing by OpenGL
	Hidden bool
	// Scale is the scale at which to render, in the X and Y axis. Not defining Scale, will default to engo.Point{1, 1}
	Scale engo.Point
	// Color defines how much of the color-components of the texture get used
	Color color.Color
	// Drawable refers to the Texture that should be drawn
	Drawable Drawable
	// Repeat defines how to repeat the Texture if the viewport of the texture is larger than the texture itself
	Repeat TextureRepeating

	shader Shader
	zIndex float32

	buffer        *gl.Buffer
	bufferContent []float32
}

func (r *RenderComponent) SetShader(s Shader) {
	r.shader = s
	engo.Mailbox.Dispatch(&renderChangeMessage{})
}

func (r *RenderComponent) SetZIndex(index float32) {
	r.zIndex = index
	engo.Mailbox.Dispatch(&renderChangeMessage{})
}

type renderEntity struct {
	*ecs.BasicEntity
	*RenderComponent
	*SpaceComponent
}

type renderEntityList []renderEntity

func (r renderEntityList) Len() int {
	return len(r)
}

func (r renderEntityList) Less(i, j int) bool {
	// Sort by shader-pointer if they have the same zIndex
	if r[i].RenderComponent.zIndex == r[j].RenderComponent.zIndex {
		// TODO: optimize this for performance
		return fmt.Sprintf("%p", r[i].RenderComponent.shader) < fmt.Sprintf("%p", r[j].RenderComponent.shader)
	}

	return r[i].RenderComponent.zIndex < r[j].RenderComponent.zIndex
}

func (r renderEntityList) Swap(i, j int) {
	r[i], r[j] = r[j], r[i]
}

type RenderSystem struct {
	entities renderEntityList
	world    *ecs.World

	sortingNeeded bool
	currentShader Shader
}

func (*RenderSystem) Priority() int { return RenderSystemPriority }

func (rs *RenderSystem) New(w *ecs.World) {
	rs.world = w

	addCameraSystemOnce(w)

	if !engo.Headless() {
		initShaders(w)
		engo.Gl.Enable(engo.Gl.MULTISAMPLE)
	}

	engo.Mailbox.Listen("renderChangeMessage", func(engo.Message) {
		rs.sortingNeeded = true
	})
}

var cameraInitMutex sync.Mutex
func addCameraSystemOnce(w *ecs.World) {
	cameraInitMutex.Lock()
	defer cameraInitMutex.Unlock()

	camSystemAlreadyAdded := false
	for _, system := range w.Systems() {
		switch system.(type) {
		case *CameraSystem:
			camSystemAlreadyAdded = true
		}
	}
	if !camSystemAlreadyAdded {
		w.AddSystem(&CameraSystem{})
	}
}

func (rs *RenderSystem) Add(basic *ecs.BasicEntity, render *RenderComponent, space *SpaceComponent) {
	// Do nothing if entity already exists
	if rs.EntityExists(basic) >= 0 {
		return
	}

	// Setting default shader
	if render.shader == nil {
		switch render.Drawable.(type) {
		case Triangle:
			render.shader = LegacyShader
		case Circle:
			render.shader = LegacyShader
		case Rectangle:
			render.shader = LegacyShader
		case ComplexTriangles:
			render.shader = LegacyShader
		case Text:
			render.shader = TextShader
		default:
			render.shader = DefaultShader
		}
	}

	// This is to prevent users from using the wrong one
	if render.shader == HUDShader {
		switch render.Drawable.(type) {
		case Triangle:
			render.shader = LegacyHUDShader
		case Circle:
			render.shader = LegacyHUDShader
		case Rectangle:
			render.shader = LegacyHUDShader
		case ComplexTriangles:
			render.shader = LegacyHUDShader
		case Text:
			render.shader = TextHUDShader
		default:
			render.shader = HUDShader
		}
	}

	rs.entities = append(rs.entities, renderEntity{basic, render, space})
	rs.sortingNeeded = true
}

// EntityExists looks if the entity is already into the System's entities. It will return the index >= 0 of the object into de rs.entities or -1 if it could not be found.
func (rs *RenderSystem) EntityExists(basic *ecs.BasicEntity) int {
	for index, entity := range rs.entities {
		if entity.ID() == basic.ID() {
			return index
		}
	}

	return -1
}

// AddByInterface adds any Renderable to the render system. Any Entity containing a BasicEntity,RenderComponent, and SpaceComponent anonymously does this automatically
func (rs *RenderSystem) AddByInterface(o Renderable) {
	rs.Add(o.GetBasicEntity(), o.GetRenderComponent(), o.GetSpaceComponent())
}

func (rs *RenderSystem) Remove(basic ecs.BasicEntity) {
	var delete int = rs.EntityExists(&basic)
	if delete >= 0 {
		rs.entities = append(rs.entities[:delete], rs.entities[delete+1:]...)
		rs.sortingNeeded = true
	}
}

func (rs *RenderSystem) Update(dt float32) {
	if engo.Headless() {
		return
	}

	if rs.sortingNeeded {
		sort.Sort(rs.entities)
		rs.sortingNeeded = false
	}

	engo.Gl.Clear(engo.Gl.COLOR_BUFFER_BIT)

	// TODO: it's linear for now, but that might very well be a bad idea
	for _, e := range rs.entities {
		if e.RenderComponent.Hidden {
			continue // with other entities
		}

		// Retrieve a shader, may be the default one -- then use it if we aren't already using it
		shader := e.RenderComponent.shader

		// Change Shader if we have to
		if shader != rs.currentShader {
			if rs.currentShader != nil {
				rs.currentShader.Post()
			}
			shader.Pre()
			rs.currentShader = shader
		}

		// Setting default scale to 1
		if e.RenderComponent.Scale.X == 0 && e.RenderComponent.Scale.Y == 0 {
			e.RenderComponent.Scale = engo.Point{1, 1}
		}

		// Setting default to white
		if e.RenderComponent.Color == nil {
			e.RenderComponent.Color = color.White
		}

		rs.currentShader.Draw(e.RenderComponent, e.SpaceComponent)
	}

	if rs.currentShader != nil {
		rs.currentShader.Post()
		rs.currentShader = nil
	}
}

func SetBackground(c color.Color) {
	if !engo.Headless() {
		r, g, b, a := c.RGBA()

		engo.Gl.ClearColor(float32(r)/0xffff, float32(g)/0xffff, float32(b)/0xffff, float32(a)/0xffff)
	}
}
<|endoftext|>