<commit_before>package fly

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"github.com/EngineerBetter/concourse-up/config"
	"github.com/EngineerBetter/concourse-up/util"
)

// ConcourseUpVersion is a compile-time variable set with -ldflags
var ConcourseUpVersion = "COMPILE_TIME_VARIABLE_fly_concourse_up_version"

// IClient represents an interface for a client
type IClient interface {
	CanConnect() (bool, error)
	SetDefaultPipeline(deployArgs *config.DeployArgs, config *config.Config, allowFlyVersionDiscrepancy bool) error
	Cleanup() error
}

// Client represents a low-level wrapper for fly
type Client struct {
	tempDir     *util.TempDir
	creds       Credentials
	stdout      io.Writer
	stderr      io.Writer
	versionFile []byte
}

// Credentials represents credentials needed to connect to concourse
type Credentials struct {
	Target   string
	API      string
	Username string
	Password string
	CACert   string
}

// New returns a new fly client
func New(creds Credentials, stdout, stderr io.Writer, versionFile []byte) (IClient, error) {
	tempDir, err := util.NewTempDir()
	if err != nil {
		return nil, err
	}

	fileHandler, err := os.Create(tempDir.Path("fly"))
	if err != nil {
		return nil, err
	}
	defer fileHandler.Close()

	url, err := getFlyURL(versionFile)
	if err != nil {
		return nil, err
	}

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if _, err := io.Copy(fileHandler, resp.Body); err != nil {
		return nil, err
	}

	if err := fileHandler.Sync(); err != nil {
		return nil, err
	}

	if err := os.Chmod(fileHandler.Name(), 0700); err != nil {
		return nil, err
	}

	return &Client{
		tempDir,
		creds,
		stdout,
		stderr,
		versionFile,
	}, nil
}

// CanConnect returns true if it can connect to the concourse
func (client *Client) CanConnect() (bool, error) {
	cmd := exec.Command(
		client.tempDir.Path("fly"),
		"--target",
		client.creds.Target,
		"login",
		"--insecure",
		"--concourse-url",
		client.creds.API,
		"--username",
		client.creds.Username,
		"--password",
		client.creds.Password,
	)

	stderr := bytes.NewBuffer(nil)
	cmd.Stdout = client.stdout
	cmd.Stderr = stderr

	runErr := cmd.Run()
	if runErr == nil {
		return true, nil
	}

	stderrBytes, err := ioutil.ReadAll(stderr)
	if err != nil {
		return false, err
	}

	if strings.Contains(string(stderrBytes), "could not reach the Concourse server") {
		return false, nil
	}

	// if there is a legitimate error, copy it to stderr for debugging
	if _, err := client.stderr.Write(stderrBytes); err != nil {
		return false, err
	}

	return false, runErr
}

// SetDefaultPipeline sets the default pipeline against a given concourse
func (client *Client) SetDefaultPipeline(deployArgs *config.DeployArgs, config *config.Config, allowFlyVersionDiscrepancy bool) error {
	if err := client.login(); err != nil {
		return err
	}

	if allowFlyVersionDiscrepancy {
		if err := client.sync(); err != nil {
			return err
		}
		if err := client.login(); err != nil {
			return err
		}
	}

	pipelinePath := client.tempDir.Path("default-pipeline.yml")
	pipelineName := "concourse-up-self-update"

	if err := client.writePipelineConfig(pipelinePath, deployArgs, config); err != nil {
		return err
	}

	if err := client.run("set-pipeline", "--pipeline", pipelineName, "--config", pipelinePath, "--non-interactive"); err != nil {
		return err
	}

	if err := os.Remove(pipelinePath); err != nil {
		return err
	}

	if err := client.run("pause-job", "--job", pipelineName+"/self-update"); err != nil {
		return err
	}

	return client.run("unpause-pipeline", "--pipeline", pipelineName)
}

func (client *Client) writePipelineConfig(pipelinePath string, deployArgs *config.DeployArgs, config *config.Config) error {
	fileHandler, err := os.Create(pipelinePath)
	if err != nil {
		return err
	}
	defer fileHandler.Close()

	params, err := client.buildDefaultPipelineParams(deployArgs, config)
	if err != nil {
		return err
	}

	pipelineConfig, err := util.RenderTemplate(defaultPipelineTemplate, params)
	if err != nil {
		return err
	}

	if _, err := fileHandler.Write(pipelineConfig); err != nil {
		return err
	}

	if err := fileHandler.Sync(); err != nil {
		return err
	}

	return nil
}

// Cleanup removes tempfiles
func (client *Client) Cleanup() error {
	return client.tempDir.Cleanup()
}

func (client *Client) login() error {
	attempts := 50
	secondsBetweenAttempts := 4

	if _, err := client.stdout.Write([]byte("Waiting for Concourse ATC to start... \n")); err != nil {
		return err
	}

	for i := 0; i < attempts; i++ {
		canConnect, err := client.CanConnect()
		if err != nil {
			return err
		}
		if canConnect {
			return nil
		}

		time.Sleep(time.Second * time.Duration(secondsBetweenAttempts))
	}

	return fmt.Errorf("failed to log in to %s after %d seconds", client.creds.API, attempts*secondsBetweenAttempts)
}

func (client *Client) sync() error {
	return client.run("sync")
}

func (client *Client) run(args ...string) error {
	args = append([]string{"--target", client.creds.Target}, args...)
	cmd := exec.Command(client.tempDir.Path("fly"), args...)
	cmd.Stdout = client.stdout
	cmd.Stderr = client.stderr
	return cmd.Run()
}

func getFlyURL(versionFile []byte) (string, error) {
	var x map[string]map[string]string
	err := json.Unmarshal(versionFile, &x)
	if err != nil {
		return "", err
	}
	switch runtime.GOOS {
	case "darwin":
		return x["fly"]["mac"], nil
	case "linux":
		return x["fly"]["linux"], nil
	case "windows":
		return x["fly"]["windows"], nil
	default:
		return "", fmt.Errorf("unknown os: `%s`", runtime.GOOS)
	}
}

func (client *Client) buildDefaultPipelineParams(deployArgs *config.DeployArgs, config *config.Config) (*defaultPipelineParams, error) {
	awsAccessKeyID := os.Getenv("AWS_ACCESS_KEY_ID")
	if awsAccessKeyID == "" {
		return nil, errors.New("env var AWS_ACCESS_KEY_ID not found")
	}

	awsSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
	if awsSecretAccessKey == "" {
		return nil, errors.New("env var AWS_SECRET_ACCESS_KEY not found")
	}

	return &defaultPipelineParams{
		AWSAccessKeyID:     awsAccessKeyID,
		AWSSecretAccessKey: awsSecretAccessKey,
		Deployment:         strings.TrimPrefix(config.Deployment, "concourse-up-"),
		FlagAWSRegion:      deployArgs.AWSRegion,
		FlagDomain:         deployArgs.Domain,
		FlagTLSCert:        deployArgs.TLSCert,
		FlagTLSKey:         deployArgs.TLSKey,
		FlagWebSize:        deployArgs.WebSize,
		FlagWorkerSize:     deployArgs.WorkerSize,
		FlagWorkers:        deployArgs.WorkerCount,
		ConcourseUpVersion: ConcourseUpVersion,
	}, nil
}

type defaultPipelineParams struct {
	AWSAccessKeyID     string
	AWSDefaultRegion   string
	AWSSecretAccessKey string
	Deployment         string
	FlagAWSRegion      string
	FlagDomain         string
	FlagTLSCert        string
	FlagTLSKey         string
	FlagWebSize        string
	FlagWorkerSize     string
	FlagWorkers        int
	ConcourseUpVersion string
}

// Indent is a helper function to indent the field a given number of spaces
func (params defaultPipelineParams) Indent(countStr, field string) string {
	return util.Indent(countStr, field)
}

const defaultPipelineTemplate = `
---
resources:
- name: concourse-up-release
  type: github-release
  source:
    user: engineerbetter
    repository: concourse-up
    pre_release: true
- name: every-month
  type: time
  source: {interval: 730h}

jobs:
- name: self-update
  serial_groups: [cup]
  serial: true
  plan:
  - get: concourse-up-release
    trigger: true
  - task: update
    params:
      AWS_REGION: "<% .FlagAWSRegion %>"
      DOMAIN: "<% .FlagDomain %>"
      TLS_CERT: |-
        <% .Indent "8" .FlagTLSCert %>
      TLS_KEY: |-
        <% .Indent "8" .FlagTLSKey %>
      WORKERS: "<% .FlagWorkers %>"
      WORKER_SIZE: "<% .FlagWorkerSize %>"
      WEB_SIZE: "<% .FlagWebSize %>"
      DEPLOYMENT: "<% .Deployment %>"
      AWS_ACCESS_KEY_ID: "<% .AWSAccessKeyID %>"
      AWS_SECRET_ACCESS_KEY: "<% .AWSSecretAccessKey %>"
      SELF_UPDATE: true
    config:
      platform: linux
      image_resource:
        type: docker-image
        source:
          repository: engineerbetter/cup-image
      inputs:
      - name: concourse-up-release
      run:
        path: bash
        args:
        - -c
        - |
          set -eux

          cd concourse-up-release
          chmod +x concourse-up-linux-amd64
          ./concourse-up-linux-amd64 deploy $DEPLOYMENT
- name: renew-cert
  serial_groups: [cup]
  serial: true
  plan:
  - get: concourse-up-release
    version: {tag: <% .ConcourseUpVersion %> }
  - get: every-month
    trigger: true
  - task: update
    params:
      AWS_REGION: "<% .FlagAWSRegion %>"
      DOMAIN: "<% .FlagDomain %>"
      TLS_CERT: |-
        <% .Indent "8" .FlagTLSCert %>
      TLS_KEY: |-
        <% .Indent "8" .FlagTLSKey %>
      WORKERS: "<% .FlagWorkers %>"
      WORKER_SIZE: "<% .FlagWorkerSize %>"
      WEB_SIZE: "<% .FlagWebSize %>"
      DEPLOYMENT: "<% .Deployment %>"
      AWS_ACCESS_KEY_ID: "<% .AWSAccessKeyID %>"
      AWS_SECRET_ACCESS_KEY: "<% .AWSSecretAccessKey %>"
      SELF_UPDATE: true
    config:
      platform: linux
      image_resource:
        type: docker-image
        source:
          repository: engineerbetter/cup-image
      inputs:
      - name: concourse-up-release
      run:
        path: bash
        args:
        - -c
        - |
          set -eux

          cd concourse-up-release
          chmod +x concourse-up-linux-amd64
          ./concourse-up-linux-amd64 deploy $DEPLOYMENT
`
<commit_msg>use the correct image in the self-update pipeline<commit_after>package fly

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"github.com/EngineerBetter/concourse-up/config"
	"github.com/EngineerBetter/concourse-up/util"
)

// ConcourseUpVersion is a compile-time variable set with -ldflags
var ConcourseUpVersion = "COMPILE_TIME_VARIABLE_fly_concourse_up_version"

// IClient represents an interface for a client
type IClient interface {
	CanConnect() (bool, error)
	SetDefaultPipeline(deployArgs *config.DeployArgs, config *config.Config, allowFlyVersionDiscrepancy bool) error
	Cleanup() error
}

// Client represents a low-level wrapper for fly
type Client struct {
	tempDir     *util.TempDir
	creds       Credentials
	stdout      io.Writer
	stderr      io.Writer
	versionFile []byte
}

// Credentials represents credentials needed to connect to concourse
type Credentials struct {
	Target   string
	API      string
	Username string
	Password string
	CACert   string
}

// New returns a new fly client
func New(creds Credentials, stdout, stderr io.Writer, versionFile []byte) (IClient, error) {
	tempDir, err := util.NewTempDir()
	if err != nil {
		return nil, err
	}

	fileHandler, err := os.Create(tempDir.Path("fly"))
	if err != nil {
		return nil, err
	}
	defer fileHandler.Close()

	url, err := getFlyURL(versionFile)
	if err != nil {
		return nil, err
	}

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if _, err := io.Copy(fileHandler, resp.Body); err != nil {
		return nil, err
	}

	if err := fileHandler.Sync(); err != nil {
		return nil, err
	}

	if err := os.Chmod(fileHandler.Name(), 0700); err != nil {
		return nil, err
	}

	return &Client{
		tempDir,
		creds,
		stdout,
		stderr,
		versionFile,
	}, nil
}

// CanConnect returns true if it can connect to the concourse
func (client *Client) CanConnect() (bool, error) {
	cmd := exec.Command(
		client.tempDir.Path("fly"),
		"--target",
		client.creds.Target,
		"login",
		"--insecure",
		"--concourse-url",
		client.creds.API,
		"--username",
		client.creds.Username,
		"--password",
		client.creds.Password,
	)

	stderr := bytes.NewBuffer(nil)
	cmd.Stdout = client.stdout
	cmd.Stderr = stderr

	runErr := cmd.Run()
	if runErr == nil {
		return true, nil
	}

	stderrBytes, err := ioutil.ReadAll(stderr)
	if err != nil {
		return false, err
	}

	if strings.Contains(string(stderrBytes), "could not reach the Concourse server") {
		return false, nil
	}

	// if there is a legitimate error, copy it to stderr for debugging
	if _, err := client.stderr.Write(stderrBytes); err != nil {
		return false, err
	}

	return false, runErr
}

// SetDefaultPipeline sets the default pipeline against a given concourse
func (client *Client) SetDefaultPipeline(deployArgs *config.DeployArgs, config *config.Config, allowFlyVersionDiscrepancy bool) error {
	if err := client.login(); err != nil {
		return err
	}

	if allowFlyVersionDiscrepancy {
		if err := client.sync(); err != nil {
			return err
		}
		if err := client.login(); err != nil {
			return err
		}
	}

	pipelinePath := client.tempDir.Path("default-pipeline.yml")
	pipelineName := "concourse-up-self-update"

	if err := client.writePipelineConfig(pipelinePath, deployArgs, config); err != nil {
		return err
	}

	if err := client.run("set-pipeline", "--pipeline", pipelineName, "--config", pipelinePath, "--non-interactive"); err != nil {
		return err
	}

	if err := os.Remove(pipelinePath); err != nil {
		return err
	}

	if err := client.run("pause-job", "--job", pipelineName+"/self-update"); err != nil {
		return err
	}

	return client.run("unpause-pipeline", "--pipeline", pipelineName)
}

func (client *Client) writePipelineConfig(pipelinePath string, deployArgs *config.DeployArgs, config *config.Config) error {
	fileHandler, err := os.Create(pipelinePath)
	if err != nil {
		return err
	}
	defer fileHandler.Close()

	params, err := client.buildDefaultPipelineParams(deployArgs, config)
	if err != nil {
		return err
	}

	pipelineConfig, err := util.RenderTemplate(defaultPipelineTemplate, params)
	if err != nil {
		return err
	}

	if _, err := fileHandler.Write(pipelineConfig); err != nil {
		return err
	}

	if err := fileHandler.Sync(); err != nil {
		return err
	}

	return nil
}

// Cleanup removes tempfiles
func (client *Client) Cleanup() error {
	return client.tempDir.Cleanup()
}

func (client *Client) login() error {
	attempts := 50
	secondsBetweenAttempts := 4

	if _, err := client.stdout.Write([]byte("Waiting for Concourse ATC to start... \n")); err != nil {
		return err
	}

	for i := 0; i < attempts; i++ {
		canConnect, err := client.CanConnect()
		if err != nil {
			return err
		}
		if canConnect {
			return nil
		}

		time.Sleep(time.Second * time.Duration(secondsBetweenAttempts))
	}

	return fmt.Errorf("failed to log in to %s after %d seconds", client.creds.API, attempts*secondsBetweenAttempts)
}

func (client *Client) sync() error {
	return client.run("sync")
}

func (client *Client) run(args ...string) error {
	args = append([]string{"--target", client.creds.Target}, args...)
	cmd := exec.Command(client.tempDir.Path("fly"), args...)
	cmd.Stdout = client.stdout
	cmd.Stderr = client.stderr
	return cmd.Run()
}

func getFlyURL(versionFile []byte) (string, error) {
	var x map[string]map[string]string
	err := json.Unmarshal(versionFile, &x)
	if err != nil {
		return "", err
	}
	switch runtime.GOOS {
	case "darwin":
		return x["fly"]["mac"], nil
	case "linux":
		return x["fly"]["linux"], nil
	case "windows":
		return x["fly"]["windows"], nil
	default:
		return "", fmt.Errorf("unknown os: `%s`", runtime.GOOS)
	}
}

func (client *Client) buildDefaultPipelineParams(deployArgs *config.DeployArgs, config *config.Config) (*defaultPipelineParams, error) {
	awsAccessKeyID := os.Getenv("AWS_ACCESS_KEY_ID")
	if awsAccessKeyID == "" {
		return nil, errors.New("env var AWS_ACCESS_KEY_ID not found")
	}

	awsSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
	if awsSecretAccessKey == "" {
		return nil, errors.New("env var AWS_SECRET_ACCESS_KEY not found")
	}

	return &defaultPipelineParams{
		AWSAccessKeyID:     awsAccessKeyID,
		AWSSecretAccessKey: awsSecretAccessKey,
		Deployment:         strings.TrimPrefix(config.Deployment, "concourse-up-"),
		FlagAWSRegion:      deployArgs.AWSRegion,
		FlagDomain:         deployArgs.Domain,
		FlagTLSCert:        deployArgs.TLSCert,
		FlagTLSKey:         deployArgs.TLSKey,
		FlagWebSize:        deployArgs.WebSize,
		FlagWorkerSize:     deployArgs.WorkerSize,
		FlagWorkers:        deployArgs.WorkerCount,
		ConcourseUpVersion: ConcourseUpVersion,
	}, nil
}

type defaultPipelineParams struct {
	AWSAccessKeyID     string
	AWSDefaultRegion   string
	AWSSecretAccessKey string
	Deployment         string
	FlagAWSRegion      string
	FlagDomain         string
	FlagTLSCert        string
	FlagTLSKey         string
	FlagWebSize        string
	FlagWorkerSize     string
	FlagWorkers        int
	ConcourseUpVersion string
}

// Indent is a helper function to indent the field a given number of spaces
func (params defaultPipelineParams) Indent(countStr, field string) string {
	return util.Indent(countStr, field)
}

const defaultPipelineTemplate = `
---
resources:
- name: concourse-up-release
  type: github-release
  source:
    user: engineerbetter
    repository: concourse-up
    pre_release: true
- name: every-month
  type: time
  source: {interval: 730h}

jobs:
- name: self-update
  serial_groups: [cup]
  serial: true
  plan:
  - get: concourse-up-release
    trigger: true
  - task: update
    params:
      AWS_REGION: "<% .FlagAWSRegion %>"
      DOMAIN: "<% .FlagDomain %>"
      TLS_CERT: |-
        <% .Indent "8" .FlagTLSCert %>
      TLS_KEY: |-
        <% .Indent "8" .FlagTLSKey %>
      WORKERS: "<% .FlagWorkers %>"
      WORKER_SIZE: "<% .FlagWorkerSize %>"
      WEB_SIZE: "<% .FlagWebSize %>"
      DEPLOYMENT: "<% .Deployment %>"
      AWS_ACCESS_KEY_ID: "<% .AWSAccessKeyID %>"
      AWS_SECRET_ACCESS_KEY: "<% .AWSSecretAccessKey %>"
      SELF_UPDATE: true
    config:
      platform: linux
      image_resource:
        type: docker-image
        source:
          repository: engineerbetter/pcf-ops
      inputs:
      - name: concourse-up-release
      run:
        path: bash
        args:
        - -c
        - |
          set -eux

          cd concourse-up-release
          chmod +x concourse-up-linux-amd64
          ./concourse-up-linux-amd64 deploy $DEPLOYMENT
- name: renew-cert
  serial_groups: [cup]
  serial: true
  plan:
  - get: concourse-up-release
    version: {tag: <% .ConcourseUpVersion %> }
  - get: every-month
    trigger: true
  - task: update
    params:
      AWS_REGION: "<% .FlagAWSRegion %>"
      DOMAIN: "<% .FlagDomain %>"
      TLS_CERT: |-
        <% .Indent "8" .FlagTLSCert %>
      TLS_KEY: |-
        <% .Indent "8" .FlagTLSKey %>
      WORKERS: "<% .FlagWorkers %>"
      WORKER_SIZE: "<% .FlagWorkerSize %>"
      WEB_SIZE: "<% .FlagWebSize %>"
      DEPLOYMENT: "<% .Deployment %>"
      AWS_ACCESS_KEY_ID: "<% .AWSAccessKeyID %>"
      AWS_SECRET_ACCESS_KEY: "<% .AWSSecretAccessKey %>"
      SELF_UPDATE: true
    config:
      platform: linux
      image_resource:
        type: docker-image
        source:
          repository: engineerbetter/pcf-ops
      inputs:
      - name: concourse-up-release
      run:
        path: bash
        args:
        - -c
        - |
          set -eux

          cd concourse-up-release
          chmod +x concourse-up-linux-amd64
          ./concourse-up-linux-amd64 deploy $DEPLOYMENT
`
<|endoftext|>


<commit_before>// Package cli provides all methods to control command line functions
package cli

import (
	"encoding/json"
	"fmt"
	"github.com/mehrdadrad/mylg/banner"
	"gopkg.in/readline.v1"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"strings"
)

const usage = `Usage:
	The myLG tool developed to troubleshoot networking situations.
	The vi/emacs mode,almost all basic features is supported. press tab to see what options are available.

	connect <provider name>     connects to external looking glass, press tab to see the menu
	node <city/country name>    connects to specific node at current looking glass, press tab to see the available nodes
	local                       back to local
	lg                          change mode to extenal looking glass
	ns                          change mode to name server looking up
	ping                        ping ip address or domain name
	dig                         name server looking up
	whois                       resolve AS number/IP/CIDR to holder (provides by ripe ncc)
	hping                       Ping through HTTP/HTTPS w/ GET/HEAD methods
	scan                        scan tcp ports (you can provide range >scan host minport maxport)
	peering                     peering information (provides by peeringdb.com)
	web                         web dashboard - opens dashboard at your default browser
	`

// Readline structure
type Readline struct {
	instance  *readline.Instance
	completer *readline.PrefixCompleter
	prompt    string
	next      chan struct{}
}

// Init set readline imain items
func Init(prompt, version string) *Readline {
	var (
		r         Readline
		err       error
		completer = readline.NewPrefixCompleter(
			readline.PcItem("ping"),
			readline.PcItem("trace"),
			readline.PcItem("bgp"),
			readline.PcItem("hping"),
			readline.PcItem("connect"),
			readline.PcItem("node"),
			readline.PcItem("local"),
			readline.PcItem("lg"),
			readline.PcItem("ns"),
			readline.PcItem("dig"),
			readline.PcItem("whois"),
			readline.PcItem("scan"),
			readline.PcItem("peering"),
			readline.PcItem("help"),
			readline.PcItem("web"),
			readline.PcItem("exit"),
		)
	)
	r.completer = completer
	r.instance, err = readline.NewEx(&readline.Config{
		Prompt:          prompt + "> ",
		HistoryFile:     "/tmp/myping",
		InterruptPrompt: "^C",
		EOFPrompt:       "exit",
		AutoComplete:    completer,
	})
	if err != nil {
		panic(err)
	}
	banner.Println(version) // print banner
	checkUpdate(version)    // check update version
	r.prompt = prompt       // init local prompt
	return &r
}

// RemoveItemCompleter removes subitem(s) from a specific main item
func (r *Readline) RemoveItemCompleter(pcItem string) {
	child := []readline.PrefixCompleterInterface{}
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) != pcItem {
			child = append(child, p)
		}
	}

	r.completer.Children = child

}

// AddCompleter updates subitem(s) from a specific main item
func (r *Readline) AddCompleter(pcItem string, pcSubItems []string) {
	var pc readline.PrefixCompleter
	c := []readline.PrefixCompleterInterface{}
	for _, item := range pcSubItems {
		c = append(c, readline.PcItem(item))
	}
	pc.Name = []rune(pcItem + " ")
	pc.Children = c
	r.completer.Children = append(r.completer.Children, &pc)
}

// UpdateCompleter updates subitem(s) from a specific main item
func (r *Readline) UpdateCompleter(pcItem string, pcSubItems []string) {
	child := []readline.PrefixCompleterInterface{}
	var pc readline.PrefixCompleter
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) == pcItem {
			c := []readline.PrefixCompleterInterface{}
			for _, item := range pcSubItems {
				c = append(c, readline.PcItem(item))
			}
			pc.Name = []rune(pcItem + " ")
			pc.Children = c
			child = append(child, &pc)
		} else {
			child = append(child, p)
		}
	}
	if len(pc.Name) < 1 {
		// todo adding new
	}
	r.completer.Children = child
}

// SetPrompt set readline prompt and store it
func (r *Readline) SetPrompt(p string) {
	p = strings.ToLower(p)
	r.prompt = p
	r.instance.SetPrompt(p + "> ")
}

// UpdatePromptN appends readline prompt
func (r *Readline) UpdatePromptN(p string, n int) {
	var parts []string
	p = strings.ToLower(p)
	parts = strings.SplitAfterN(r.prompt, "/", n)
	if n <= len(parts) && n > -1 {
		parts[n-1] = p
		r.prompt = strings.Join(parts, "")
	} else {
		r.prompt += "/" + p
	}
	r.instance.SetPrompt(r.prompt + "> ")
}

// GetPrompt returns the current prompt string
func (r *Readline) GetPrompt() string {
	return r.prompt
}

// Refresh prompt
func (r *Readline) Refresh() {
	r.instance.Refresh()
}

// SetVim set mode to vim
func (r *Readline) SetVim() {
	if !r.instance.IsVimMode() {
		r.instance.SetVimMode(true)
		println("mode changed to vim")
	} else {
		println("mode already is vim")
	}
}

// SetEmacs set mode to emacs
func (r *Readline) SetEmacs() {
	if r.instance.IsVimMode() {
		r.instance.SetVimMode(false)
		println("mode changed to emacs")
	} else {
		println("mode already is emacs")
	}
}

// Next trigers to read next line
func (r *Readline) Next() {
	r.next <- struct{}{}
}

// Run the main loop
func (r *Readline) Run(cmd chan<- string, next chan struct{}) {
	r.next = next
	func() {
		for {
			line, err := r.instance.Readline()
			if err != nil { // io.EOF, readline.ErrInterrupt
				break
			}
			cmd <- line
			if _, ok := <-next; !ok {
				break
			}
		}
	}()
}

// Close the readline instance
func (r *Readline) Close(next chan struct{}) {
	r.instance.Close()
}

// Help print out the main help
func (r *Readline) Help() {
	fmt.Println(usage)
}

// checkUpdate checks if any new version is available
func checkUpdate(version string) {
	type mylg struct {
		Version string
	}
	var appCtl mylg

	if version == "test" {
		return
	}

	resp, err := http.Get("http://mylg.io/appctl/mylg")
	if err != nil {
		println("error: check update has been failed ")
		return
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		println("error: check update has been failed (2)" + err.Error())
		return
	}
	err = json.Unmarshal(body, &appCtl)
	if err != nil {
		println(err.Error())
		return
	}
	if version != appCtl.Version {
		fmt.Printf("New version is available (v%s) mylg.io/download\n", appCtl.Version)
	}
}

//Flag parses the command arguments syntax:
// -flag=x
// -flag x
// help
func Flag(args string) (string, map[string]interface{}) {
	var (
		r   = make(map[string]interface{}, 10)
		err error
	)
	args = strings.TrimSpace(args)
	re := regexp.MustCompile(`(?i)-([a-z]+)={0,1}\s{0,1}([0-9|a-z|\-|'"{}:]+)`)
	f := re.FindAllStringSubmatch(args, -1)
	for _, kv := range f {
		if len(kv) > 1 {
			// trim extra characters (' and ") from value
			kv[2] = strings.Trim(kv[2], "'")
			kv[2] = strings.Trim(kv[2], `"`)
			r[kv[1]], err = strconv.Atoi(kv[2])
			if err != nil {
				r[kv[1]] = kv[2]
			}
			args = strings.Replace(args, kv[0], "", -1)
		}
	}
	if m, _ := regexp.MatchString(`(?i)help$`, args); m {
		r["help"] = true
	}
	args = strings.TrimSpace(args)
	return args, r
}

// SetFlag returns command option(s)
func SetFlag(flag map[string]interface{}, option string, v interface{}) interface{} {
	if sValue, ok := flag[option]; ok {
		switch v.(type) {
		case int:
			return sValue.(int)
		default:
			return sValue.(string)
		}
	} else {
		return v
	}
}
<commit_msg>add gorouting for checking new version<commit_after>// Package cli provides all methods to control command line functions
package cli

import (
	"encoding/json"
	"fmt"
	"github.com/mehrdadrad/mylg/banner"
	"gopkg.in/readline.v1"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"strings"
)

const usage = `Usage:
	The myLG tool developed to troubleshoot networking situations.
	The vi/emacs mode,almost all basic features is supported. press tab to see what options are available.

	connect <provider name>     connects to external looking glass, press tab to see the menu
	node <city/country name>    connects to specific node at current looking glass, press tab to see the available nodes
	local                       back to local
	lg                          change mode to extenal looking glass
	ns                          change mode to name server looking up
	ping                        ping ip address or domain name
	dig                         name server looking up
	whois                       resolve AS number/IP/CIDR to holder (provides by ripe ncc)
	hping                       Ping through HTTP/HTTPS w/ GET/HEAD methods
	scan                        scan tcp ports (you can provide range >scan host minport maxport)
	peering                     peering information (provides by peeringdb.com)
	web                         web dashboard - opens dashboard at your default browser
	`

// Readline structure
type Readline struct {
	instance  *readline.Instance
	completer *readline.PrefixCompleter
	prompt    string
	next      chan struct{}
}

// Init set readline imain items
func Init(prompt, version string) *Readline {
	var (
		r         Readline
		err       error
		completer = readline.NewPrefixCompleter(
			readline.PcItem("ping"),
			readline.PcItem("trace"),
			readline.PcItem("bgp"),
			readline.PcItem("hping"),
			readline.PcItem("connect"),
			readline.PcItem("node"),
			readline.PcItem("local"),
			readline.PcItem("lg"),
			readline.PcItem("ns"),
			readline.PcItem("dig"),
			readline.PcItem("whois"),
			readline.PcItem("scan"),
			readline.PcItem("peering"),
			readline.PcItem("help"),
			readline.PcItem("web"),
			readline.PcItem("exit"),
		)
	)
	r.completer = completer
	r.instance, err = readline.NewEx(&readline.Config{
		Prompt:          prompt + "> ",
		HistoryFile:     "/tmp/myping",
		InterruptPrompt: "^C",
		EOFPrompt:       "exit",
		AutoComplete:    completer,
	})
	if err != nil {
		panic(err)
	}
	banner.Println(version) // print banner
	go checkUpdate(version) // check update version
	r.prompt = prompt       // init local prompt
	return &r
}

// RemoveItemCompleter removes subitem(s) from a specific main item
func (r *Readline) RemoveItemCompleter(pcItem string) {
	child := []readline.PrefixCompleterInterface{}
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) != pcItem {
			child = append(child, p)
		}
	}

	r.completer.Children = child

}

// AddCompleter updates subitem(s) from a specific main item
func (r *Readline) AddCompleter(pcItem string, pcSubItems []string) {
	var pc readline.PrefixCompleter
	c := []readline.PrefixCompleterInterface{}
	for _, item := range pcSubItems {
		c = append(c, readline.PcItem(item))
	}
	pc.Name = []rune(pcItem + " ")
	pc.Children = c
	r.completer.Children = append(r.completer.Children, &pc)
}

// UpdateCompleter updates subitem(s) from a specific main item
func (r *Readline) UpdateCompleter(pcItem string, pcSubItems []string) {
	child := []readline.PrefixCompleterInterface{}
	var pc readline.PrefixCompleter
	for _, p := range r.completer.Children {
		if strings.TrimSpace(string(p.GetName())) == pcItem {
			c := []readline.PrefixCompleterInterface{}
			for _, item := range pcSubItems {
				c = append(c, readline.PcItem(item))
			}
			pc.Name = []rune(pcItem + " ")
			pc.Children = c
			child = append(child, &pc)
		} else {
			child = append(child, p)
		}
	}
	if len(pc.Name) < 1 {
		// todo adding new
	}
	r.completer.Children = child
}

// SetPrompt set readline prompt and store it
func (r *Readline) SetPrompt(p string) {
	p = strings.ToLower(p)
	r.prompt = p
	r.instance.SetPrompt(p + "> ")
}

// UpdatePromptN appends readline prompt
func (r *Readline) UpdatePromptN(p string, n int) {
	var parts []string
	p = strings.ToLower(p)
	parts = strings.SplitAfterN(r.prompt, "/", n)
	if n <= len(parts) && n > -1 {
		parts[n-1] = p
		r.prompt = strings.Join(parts, "")
	} else {
		r.prompt += "/" + p
	}
	r.instance.SetPrompt(r.prompt + "> ")
}

// GetPrompt returns the current prompt string
func (r *Readline) GetPrompt() string {
	return r.prompt
}

// Refresh prompt
func (r *Readline) Refresh() {
	r.instance.Refresh()
}

// SetVim set mode to vim
func (r *Readline) SetVim() {
	if !r.instance.IsVimMode() {
		r.instance.SetVimMode(true)
		println("mode changed to vim")
	} else {
		println("mode already is vim")
	}
}

// SetEmacs set mode to emacs
func (r *Readline) SetEmacs() {
	if r.instance.IsVimMode() {
		r.instance.SetVimMode(false)
		println("mode changed to emacs")
	} else {
		println("mode already is emacs")
	}
}

// Next trigers to read next line
func (r *Readline) Next() {
	r.next <- struct{}{}
}

// Run the main loop
func (r *Readline) Run(cmd chan<- string, next chan struct{}) {
	r.next = next
	func() {
		for {
			line, err := r.instance.Readline()
			if err != nil { // io.EOF, readline.ErrInterrupt
				break
			}
			cmd <- line
			if _, ok := <-next; !ok {
				break
			}
		}
	}()
}

// Close the readline instance
func (r *Readline) Close(next chan struct{}) {
	r.instance.Close()
}

// Help print out the main help
func (r *Readline) Help() {
	fmt.Println(usage)
}

// checkUpdate checks if any new version is available
func checkUpdate(version string) {
	type mylg struct {
		Version string
	}
	var appCtl mylg

	if version == "test" {
		return
	}

	resp, err := http.Get("http://mylg.io/appctl/mylg")
	if err != nil {
		println("error: check update has been failed ")
		return
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		println("error: check update has been failed (2)" + err.Error())
		return
	}
	err = json.Unmarshal(body, &appCtl)
	if err != nil {
		println(err.Error())
		return
	}
	if version != appCtl.Version {
		fmt.Printf("New version is available (v%s) mylg.io/download\n", appCtl.Version)
	}
}

//Flag parses the command arguments syntax:
// -flag=x
// -flag x
// help
func Flag(args string) (string, map[string]interface{}) {
	var (
		r   = make(map[string]interface{}, 10)
		err error
	)
	args = strings.TrimSpace(args)
	re := regexp.MustCompile(`(?i)-([a-z]+)={0,1}\s{0,1}([0-9|a-z|\-|'"{}:]+)`)
	f := re.FindAllStringSubmatch(args, -1)
	for _, kv := range f {
		if len(kv) > 1 {
			// trim extra characters (' and ") from value
			kv[2] = strings.Trim(kv[2], "'")
			kv[2] = strings.Trim(kv[2], `"`)
			r[kv[1]], err = strconv.Atoi(kv[2])
			if err != nil {
				r[kv[1]] = kv[2]
			}
			args = strings.Replace(args, kv[0], "", -1)
		}
	}
	if m, _ := regexp.MatchString(`(?i)help$`, args); m {
		r["help"] = true
	}
	args = strings.TrimSpace(args)
	return args, r
}

// SetFlag returns command option(s)
func SetFlag(flag map[string]interface{}, option string, v interface{}) interface{} {
	if sValue, ok := flag[option]; ok {
		switch v.(type) {
		case int:
			return sValue.(int)
		default:
			return sValue.(string)
		}
	} else {
		return v
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"os"

	"pkg.re/essentialkaos/ek.v9/fmtc"
	"pkg.re/essentialkaos/ek.v9/fmtutil/table"
	"pkg.re/essentialkaos/ek.v9/knf"
	"pkg.re/essentialkaos/ek.v9/options"
	"pkg.re/essentialkaos/ek.v9/terminal"
	"pkg.re/essentialkaos/ek.v9/usage"

	"github.com/gongled/vgrepo/prefs"
	"github.com/gongled/vgrepo/repository"
	"github.com/gongled/vgrepo/storage"
	"github.com/gongled/vgrepo/index"
)

// ////////////////////////////////////////////////////////////////////////////////// //

const (
	APP  = "vgrepo"
	VER  = "2.0.0"
	DESC = "Simple CLI tool for managing Vagrant repositories"
)

const (
	CMD_ADD    = "add"
	CMD_DELETE = "delete"
	CMD_LIST   = "list"
	CMD_INFO   = "info"
	CMD_RENDER = "render"
	CMD_HELP   = "help"

	CMD_ADD_SHORTCUT    = "a"
	CMD_DELETE_SHORTCUT = "d"
	CMD_LIST_SHORTCUT   = "l"
	CMD_INFO_SHORTCUT   = "i"
	CMD_RENDER_SHORTCUT = "r"
)

const (
	KNF_STORAGE_URL  = "storage:url"
	KNF_STORAGE_PATH = "storage:path"
)

const (
	ARG_NO_COLOR = "nc:no-color"
	ARG_HELP     = "h:help"
	ARG_VER      = "v:version"
)

const (
	ERROR_UNSUPPORTED      = 1
	ERROR_INVALID_SETTINGS = 2
)

const CONFIG_FILE = "/etc/vgrepo/vgrepo.knf"

// ////////////////////////////////////////////////////////////////////////////////// //

var optionsMap = options.Map{
	ARG_NO_COLOR: {Type: options.BOOL},
	ARG_HELP:     {Type: options.BOOL, Alias: "u:usage"},
	ARG_VER:      {Type: options.BOOL, Alias: "ver"},
}

var preferences *prefs.Preferences

// ////////////////////////////////////////////////////////////////////////////////// //

func Init() {
	opts, errs := options.Parse(optionsMap)

	if len(errs) != 0 {
		fmtc.Println("Arguments parsing errors:")

		for _, err := range errs {
			fmtc.Printf("  %s\n", err.Error())
		}

		os.Exit(1)
	}

	if options.GetB(ARG_NO_COLOR) {
		fmtc.DisableColors = true
	}

	if options.GetB(ARG_VER) {
		showAbout()
		return
	}

	if options.GetB(ARG_HELP) || len(opts) == 0 {
		showUsage()
		return
	}

	switch len(opts) {
	case 0:
		showUsage()
		return
	case 1:
		processCommand(opts[0], nil)
	default:
		processCommand(opts[0], opts[1:])
	}
}

// ////////////////////////////////////////////////////////////////////////////////// //

func prepare() {
	err := knf.Global(CONFIG_FILE)

	preferences = prefs.NewPreferences(
		knf.GetS(KNF_STORAGE_PATH),
		knf.GetS(KNF_STORAGE_URL),
	)

	if err != nil {
		terminal.PrintErrorMessage(err.Error())
		os.Exit(ERROR_INVALID_SETTINGS)
	}
}

func processCommand(cmd string, args []string) {
	prepare()

	switch cmd {
	case CMD_ADD, CMD_ADD_SHORTCUT:
		addCommand(args)
	case CMD_DELETE, CMD_DELETE_SHORTCUT:
		deleteCommand(args)
	case CMD_LIST, CMD_LIST_SHORTCUT:
		listCommand()
	case CMD_INFO, CMD_INFO_SHORTCUT:
		infoCommand(args)
	case CMD_RENDER, CMD_RENDER_SHORTCUT:
		renderCommand(args)
	case CMD_HELP:
		showUsage()
	default:
		terminal.PrintErrorMessage("Error: unknown command %s", cmd)
		os.Exit(ERROR_UNSUPPORTED)
	}
}

// ////////////////////////////////////////////////////////////////////////////////// //

func addCommand(args []string) {
	if len(args) < 4 {
		terminal.PrintErrorMessage(
			"Error: unable to handle %v arguments",
			len(args),
		)
		os.Exit(1)
	}

	var (
		src      = args[0]
		name     = args[1]
		version  = args[2]
		provider = args[3]
	)

	r := repository.NewRepository(preferences, name)

	terminal.PrintActionMessage("Importing package")
	err := r.AddPackage(src, repository.NewPackage(name, version, provider))

	if err != nil {
		terminal.PrintActionStatus(1)
		terminal.PrintErrorMessage("Error: %s", err.Error())
		os.Exit(1)
	} else {
		terminal.PrintActionStatus(0)
	}
}

func deleteCommand(args []string) {
	if len(args) < 3 {
		terminal.PrintErrorMessage("Error: name, version and provider must be set")
		os.Exit(1)
	}

	var (
		name     = args[0]
		version  = args[1]
		provider = args[2]
	)

	r := repository.NewRepository(preferences, name)

	terminal.PrintActionMessage("Removing package")
	err := r.RemovePackage(repository.NewPackage(name, version, provider))

	if err != nil {
		terminal.PrintActionStatus(1)
		terminal.PrintErrorMessage("Error: %s", err.Error())
		os.Exit(1)
	} else {
		terminal.PrintActionStatus(0)
	}
}

func listCommand() {
	s := storage.NewStorage(preferences)

	listTableRender(s.Repositories())
}

func infoCommand(args []string) {
	if len(args) < 1 {
		terminal.PrintErrorMessage("Error: name must be set")
		os.Exit(1)
	}

	name := args[0]

	infoTableRender(repository.NewRepository(preferences, name))
}

func renderCommand(args []string) {
	if len(args) < 1 {
		terminal.PrintErrorMessage("Error: template must be set")
		os.Exit(1)
	}

	output := args[0]
	template := "/etc/vgrepo/templates/default.tpl"

	if len(args) >= 2 {
		template = args[1]
	}

	terminal.PrintActionMessage("Rendering template")
	err := index.ExportIndex(storage.NewStorage(preferences), template, output)

	if err != nil {
		terminal.PrintActionStatus(1)
		terminal.PrintErrorMessage("Error: ", err.Error())
		os.Exit(1)
	}

	terminal.PrintActionStatus(0)
}

// ////////////////////////////////////////////////////////////////////////////////// //

func infoTableRender(r *repository.VRepository) {
	t := table.NewTable("Name", "Provider", "Version", "Checksum")
	table.HeaderCapitalize = true

	t.SetAlignments(table.ALIGN_LEFT, table.ALIGN_LEFT, table.ALIGN_RIGHT, table.ALIGN_LEFT)

	for _, v := range r.Versions {
		for _, p := range v.Providers {
			t.Add(r.Name, p.Name, v.Version, p.Checksum)
		}
	}

	if t.HasData() {
		t.Render()
	} else {
		terminal.PrintWarnMessage("Repository does not exist")
	}
}

func listTableRender(repos repository.VRepositoryList) {
	t := table.NewTable("Name", "Latest", "Metadata URL")
	table.HeaderCapitalize = true

	t.SetAlignments(table.ALIGN_LEFT, table.ALIGN_RIGHT, table.ALIGN_LEFT)
	for _, r := range repos {
		if r.HasMeta() {
			t.Add(r.Name, r.LatestVersion().Version, r.MetaURL())
		}
	}

	if t.HasData() {
		t.Render()
	} else {
		terminal.PrintWarnMessage("No repositories yet")
	}
}

// ////////////////////////////////////////////////////////////////////////////////// //

func setUsageCommands(info *usage.Info) {
	info.AddCommand(
		CMD_ADD,
		"Add image to the Vagrant repository",
		"source",
		"name",
		"version",
		"provider",
	)
	info.AddCommand(
		CMD_LIST,
		"Show the list of available images",
	)
	info.AddCommand(
		CMD_DELETE,
		"Delete the image from the repository",
		"name",
		"version",
		"provider",
	)
	info.AddCommand(
		CMD_INFO,
		"Display info of the particular repository",
		"name",
	)
	info.AddCommand(
		CMD_RENDER,
		"Create index by given template file",
		"template",
		"?output",
	)
	info.AddCommand(
		CMD_HELP,
		"Display the current help message",
	)
}

func setUsageOptions(info *usage.Info) {
	info.AddOption(ARG_NO_COLOR, "Disable colors in output")
	info.AddOption(ARG_HELP, "Show this help message")
	info.AddOption(ARG_VER, "Show version")
}

func setUsageExamples(info *usage.Info) {
	info.AddExample(
		"add $HOME/powerbox-1.0.0.box powerbox 1.1.0 virtualbox",
		"Add image to the Vagrant repository",
	)
	info.AddExample(
		"list",
		"Show the list of available repositories",
	)
	info.AddExample(
		"delete powerbox 1.1.0",
		"Remove the image from the repository",
	)
	info.AddExample(
		"info powerbox",
		"Show detailed info about the repository",
	)
	info.AddExample(
		"render index.html /etc/vgrepo/templates/default.tpl",
		"Create index file by given template with output index.html",
	)
}

func showUsage() {
	info := usage.NewInfo(APP)

	setUsageCommands(info)
	setUsageOptions(info)
	setUsageExamples(info)

	info.Render()
}

func showAbout() {
	about := &usage.About{
		App:     APP,
		Version: VER,
		Desc:    DESC,
		Year:    2014,
		Owner:   "Gleb E Goncharov",
		License: "MIT License",
	}

	about.Render()
}

// ////////////////////////////////////////////////////////////////////////////////// //
<commit_msg>Fix CLI helper<commit_after>package cli

import (
	"os"

	"pkg.re/essentialkaos/ek.v9/fmtc"
	"pkg.re/essentialkaos/ek.v9/fmtutil/table"
	"pkg.re/essentialkaos/ek.v9/knf"
	"pkg.re/essentialkaos/ek.v9/options"
	"pkg.re/essentialkaos/ek.v9/terminal"
	"pkg.re/essentialkaos/ek.v9/usage"

	"github.com/gongled/vgrepo/prefs"
	"github.com/gongled/vgrepo/repository"
	"github.com/gongled/vgrepo/storage"
	"github.com/gongled/vgrepo/index"
)

// ////////////////////////////////////////////////////////////////////////////////// //

const (
	APP  = "vgrepo"
	VER  = "2.0.0"
	DESC = "Simple CLI tool for managing Vagrant repositories"
)

const (
	CMD_ADD    = "add"
	CMD_DELETE = "delete"
	CMD_LIST   = "list"
	CMD_INFO   = "info"
	CMD_RENDER = "render"
	CMD_HELP   = "help"

	CMD_ADD_SHORTCUT    = "a"
	CMD_DELETE_SHORTCUT = "d"
	CMD_LIST_SHORTCUT   = "l"
	CMD_INFO_SHORTCUT   = "i"
	CMD_RENDER_SHORTCUT = "r"
)

const (
	KNF_STORAGE_URL  = "storage:url"
	KNF_STORAGE_PATH = "storage:path"
)

const (
	ARG_NO_COLOR = "nc:no-color"
	ARG_HELP     = "h:help"
	ARG_VER      = "v:version"
)

const (
	ERROR_UNSUPPORTED      = 1
	ERROR_INVALID_SETTINGS = 2
)

const CONFIG_FILE = "/etc/vgrepo/vgrepo.knf"

// ////////////////////////////////////////////////////////////////////////////////// //

var optionsMap = options.Map{
	ARG_NO_COLOR: {Type: options.BOOL},
	ARG_HELP:     {Type: options.BOOL, Alias: "u:usage"},
	ARG_VER:      {Type: options.BOOL, Alias: "ver"},
}

var preferences *prefs.Preferences

// ////////////////////////////////////////////////////////////////////////////////// //

func Init() {
	opts, errs := options.Parse(optionsMap)

	if len(errs) != 0 {
		fmtc.Println("Arguments parsing errors:")

		for _, err := range errs {
			fmtc.Printf("  %s\n", err.Error())
		}

		os.Exit(1)
	}

	if options.GetB(ARG_NO_COLOR) {
		fmtc.DisableColors = true
	}

	if options.GetB(ARG_VER) {
		showAbout()
		return
	}

	if options.GetB(ARG_HELP) || len(opts) == 0 {
		showUsage()
		return
	}

	switch len(opts) {
	case 0:
		showUsage()
		return
	case 1:
		processCommand(opts[0], nil)
	default:
		processCommand(opts[0], opts[1:])
	}
}

// ////////////////////////////////////////////////////////////////////////////////// //

func prepare() {
	err := knf.Global(CONFIG_FILE)

	preferences = prefs.NewPreferences(
		knf.GetS(KNF_STORAGE_PATH),
		knf.GetS(KNF_STORAGE_URL),
	)

	if err != nil {
		terminal.PrintErrorMessage(err.Error())
		os.Exit(ERROR_INVALID_SETTINGS)
	}
}

func processCommand(cmd string, args []string) {
	prepare()

	switch cmd {
	case CMD_ADD, CMD_ADD_SHORTCUT:
		addCommand(args)
	case CMD_DELETE, CMD_DELETE_SHORTCUT:
		deleteCommand(args)
	case CMD_LIST, CMD_LIST_SHORTCUT:
		listCommand()
	case CMD_INFO, CMD_INFO_SHORTCUT:
		infoCommand(args)
	case CMD_RENDER, CMD_RENDER_SHORTCUT:
		renderCommand(args)
	case CMD_HELP:
		showUsage()
	default:
		terminal.PrintErrorMessage("Error: unknown command %s", cmd)
		os.Exit(ERROR_UNSUPPORTED)
	}
}

// ////////////////////////////////////////////////////////////////////////////////// //

func addCommand(args []string) {
	if len(args) < 4 {
		terminal.PrintErrorMessage(
			"Error: unable to handle %v arguments",
			len(args),
		)
		os.Exit(1)
	}

	var (
		src      = args[0]
		name     = args[1]
		version  = args[2]
		provider = args[3]
	)

	r := repository.NewRepository(preferences, name)

	terminal.PrintActionMessage("Importing package")
	err := r.AddPackage(src, repository.NewPackage(name, version, provider))

	if err != nil {
		terminal.PrintActionStatus(1)
		terminal.PrintErrorMessage("Error: %s", err.Error())
		os.Exit(1)
	} else {
		terminal.PrintActionStatus(0)
	}
}

func deleteCommand(args []string) {
	if len(args) < 3 {
		terminal.PrintErrorMessage("Error: name, version and provider must be set")
		os.Exit(1)
	}

	var (
		name     = args[0]
		version  = args[1]
		provider = args[2]
	)

	r := repository.NewRepository(preferences, name)

	terminal.PrintActionMessage("Removing package")
	err := r.RemovePackage(repository.NewPackage(name, version, provider))

	if err != nil {
		terminal.PrintActionStatus(1)
		terminal.PrintErrorMessage("Error: %s", err.Error())
		os.Exit(1)
	} else {
		terminal.PrintActionStatus(0)
	}
}

func listCommand() {
	s := storage.NewStorage(preferences)

	listTableRender(s.Repositories())
}

func infoCommand(args []string) {
	if len(args) < 1 {
		terminal.PrintErrorMessage("Error: name must be set")
		os.Exit(1)
	}

	name := args[0]

	infoTableRender(repository.NewRepository(preferences, name))
}

func renderCommand(args []string) {
	if len(args) < 1 {
		terminal.PrintErrorMessage("Error: template must be set")
		os.Exit(1)
	}

	output := args[0]
	template := "/etc/vgrepo/templates/default.tpl"

	if len(args) >= 2 {
		template = args[1]
	}

	terminal.PrintActionMessage("Rendering template")
	err := index.ExportIndex(storage.NewStorage(preferences), template, output)

	if err != nil {
		terminal.PrintActionStatus(1)
		terminal.PrintErrorMessage("Error: ", err.Error())
		os.Exit(1)
	}

	terminal.PrintActionStatus(0)
}

// ////////////////////////////////////////////////////////////////////////////////// //

func infoTableRender(r *repository.VRepository) {
	t := table.NewTable("Name", "Provider", "Version", "Checksum")
	table.HeaderCapitalize = true

	t.SetAlignments(table.ALIGN_LEFT, table.ALIGN_LEFT, table.ALIGN_RIGHT, table.ALIGN_LEFT)

	for _, v := range r.Versions {
		for _, p := range v.Providers {
			t.Add(r.Name, p.Name, v.Version, p.Checksum)
		}
	}

	if t.HasData() {
		t.Render()
	} else {
		terminal.PrintWarnMessage("Repository does not exist")
	}
}

func listTableRender(repos repository.VRepositoryList) {
	t := table.NewTable("Name", "Latest", "Metadata URL")
	table.HeaderCapitalize = true

	t.SetAlignments(table.ALIGN_LEFT, table.ALIGN_RIGHT, table.ALIGN_LEFT)
	for _, r := range repos {
		if r.HasMeta() {
			t.Add(r.Name, r.LatestVersion().Version, r.MetaURL())
		}
	}

	if t.HasData() {
		t.Render()
	} else {
		terminal.PrintWarnMessage("No repositories yet")
	}
}

// ////////////////////////////////////////////////////////////////////////////////// //

func setUsageCommands(info *usage.Info) {
	info.AddCommand(
		CMD_ADD,
		"Add image to the Vagrant repository",
		"source",
		"name",
		"version",
		"provider",
	)
	info.AddCommand(
		CMD_LIST,
		"Show the list of available images",
	)
	info.AddCommand(
		CMD_DELETE,
		"Delete the image from the repository",
		"name",
		"version",
		"provider",
	)
	info.AddCommand(
		CMD_INFO,
		"Display info of the particular repository",
		"name",
	)
	info.AddCommand(
		CMD_RENDER,
		"Create index by given template file",
		"output",
		"?template",
	)
	info.AddCommand(
		CMD_HELP,
		"Display the current help message",
	)
}

func setUsageOptions(info *usage.Info) {
	info.AddOption(ARG_NO_COLOR, "Disable colors in output")
	info.AddOption(ARG_HELP, "Show this help message")
	info.AddOption(ARG_VER, "Show version")
}

func setUsageExamples(info *usage.Info) {
	info.AddExample(
		"add $HOME/powerbox-1.0.0.box powerbox 1.1.0 virtualbox",
		"Add image to the Vagrant repository",
	)
	info.AddExample(
		"list",
		"Show the list of available repositories",
	)
	info.AddExample(
		"delete powerbox 1.1.0",
		"Remove the image from the repository",
	)
	info.AddExample(
		"info powerbox",
		"Show detailed info about the repository",
	)
	info.AddExample(
		"render index.html /etc/vgrepo/templates/default.tpl",
		"Create index file by given template with output index.html",
	)
}

func showUsage() {
	info := usage.NewInfo(APP)

	setUsageCommands(info)
	setUsageOptions(info)
	setUsageExamples(info)

	info.Render()
}

func showAbout() {
	about := &usage.About{
		App:     APP,
		Version: VER,
		Desc:    DESC,
		Year:    2014,
		Owner:   "Gleb E Goncharov",
		License: "MIT License",
	}

	about.Render()
}

// ////////////////////////////////////////////////////////////////////////////////// //
<|endoftext|>


<commit_before>package main

import (
    "os"
    "fmt"
    "flag"
    "time"
    "github.com/garyburd/redigo/redis"
)

func main() {

    // Vars
    var src, dest string
    var sp, dp, tc, tr, tw int

    // Flag declaration
    // For source
    flag.StringVar(&src, "src", "127.0.0.1", "Source redis server hostname")
    flag.IntVar(&sp, "sp", 6379, "Source redis server port")

    // For destination
    flag.StringVar(&dest, "dest", "none", "Destination redis server hostname")
    flag.IntVar(&dp, "dp", 6379, "Destination redis server port")

    // Timeouts
    flag.IntVar(&tc, "tc", 10, "Connection timeout in seconds")
    flag.IntVar(&tr, "tr", 5, "Read timeout in seconds")
    flag.IntVar(&tw, "tw", 5, "Write timeout in seconds")

    // Do the flag parsing
    flag.Usage = usage
    flag.Parse()

    // Check if there's a -dest flag
    if dest == "none" {
        usage()
    }

    // Format addresses
    srcaddr := fmt.Sprintf("%s:%d", src, sp)
    destaddr := fmt.Sprintf("%s:%d", dest, dp)

    // Set timeouts
    connt := time.Duration(tc) * time.Second
    readt := time.Duration(tr) * time.Second
    writet := time.Duration(tw) * time.Second

    // Connect to Redis Source
    srcc, err := redis.DialTimeout("tcp", srcaddr, connt, readt, writet)

    // Exit if connection failed (source)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Cannot connect to source redis server: %v\n", err)
        os.Exit(-1)
    }

    // Connect to Redis Destination
    destc, err := redis.DialTimeout("tcp", destaddr, connt, readt, writet)

    // Exit if connection failed (destination)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Cannot connect to destination redis server: %v\n", err)
        os.Exit(-1)
    }

    // Go through each keys in source
    srckeys, _ := redis.Strings(srcc.Do("KEYS", "*"))

    for _, key := range srckeys {

        // Get the object type of each key
        rtype, _ := redis.String(srcc.Do("TYPE", key))

        fmt.Println("Migrating", key, "of type", rtype)

        // Map with the correct object data type
        switch rtype {
        case "string":

            // Raw key-value
            val, _ := redis.String(srcc.Do("GET", key))
            destc.Do("SET", key, val)
            break
        case "list":

            // Lists
            val, _ := redis.Strings(srcc.Do("LRANGE", key, 0, -1))

            for _, v := range val {
                destc.Do("RPUSH", key, v)
            }
            break
        case "set":

            // Set
            val, _ := redis.Strings(srcc.Do("SMEMBERS", key))

            for _, v := range val {
                destc.Do("SADD", key, v)
            }
            break
        case "zset":

            // Sorted Set
            val, _ := redis.Strings(srcc.Do("ZRANGE", key, 0, -1, "WITHSCORES"))

            // Since this zset will include scores we need to adapt
            zval := ""
            for i, v := range val {
                if (i + 1) % 2 == 1 {

                    // This is a value
                    zval = v
                } else {

                    // This is a score
                    // And here we do our thing
                    destc.Do("ZADD", key, v, zval)
                }
            }
            break
        case "hash":

            // Hashes, finally
            val, _ := redis.Strings(srcc.Do("HGETALL", key))

            // Do the same like in zset
            hfield := ""
            for i, v := range val {
                if (i + 1) % 2 == 1 {

                    // Value
                    hfield = v
                } else {

                    // And field
                    destc.Do("HSET", key, hfield, v)
                }
            }
            break
        default:

            // What?
            fmt.Println("Unknown type for", key)
        }

        // Check TTL
        checkttl(key, srcc, destc);
    }

    // Done
    fmt.Println("Done migrating keys from", srcaddr, "to", destaddr)

    // Cleanup
    srcc.Close()
    destc.Close()
}

// Check for key expiration and set expiry. This is a problematic one though, and may be inaccurate for a second at most
// (but I'm pretty confident it'll hardly reach that, since one command will just take a couple of ms)
func checkttl(key string, src, dst redis.Conn) {
    ttl, _ := redis.Int(src.Do("TTL", key))

    // Expire the key if it's greater than 0
    if ttl > 0 {
        fmt.Println("Setting expiration of", key)
        dst.Do("EXPIRE", key, ttl)
    }
}

// For usage
func usage() {
    fmt.Fprintf(os.Stderr, "Usage: redback [-src <source-host>] [-sp <source-port>] -dest <destination-host> [-dp <destination-port>]\n")
    flag.PrintDefaults()
    os.Exit(2)
}<commit_msg>Added DB options<commit_after>package main

import (
    "os"
    "fmt"
    "flag"
    "time"
    "github.com/garyburd/redigo/redis"
)

func main() {

    // Vars
    var src, dest string
    var sdb, ddb, sp, dp, tc, tr, tw int

    // Flag declaration
    // For source
    flag.StringVar(&src, "src", "127.0.0.1", "Source redis server hostname")
    flag.IntVar(&sp, "sp", 6379, "Source redis server port")
    flag.IntVar(&sdb, "sdb", 0, "Source redis db number")

    // For destination
    flag.StringVar(&dest, "dest", "none", "Destination redis server hostname")
    flag.IntVar(&dp, "dp", 6379, "Destination redis server port")
    flag.IntVar(&ddb, "ddb", 0, "Destination redis db number")

    // Timeouts
    flag.IntVar(&tc, "tc", 10, "Connection timeout in seconds")
    flag.IntVar(&tr, "tr", 5, "Read timeout in seconds")
    flag.IntVar(&tw, "tw", 5, "Write timeout in seconds")

    // Do the flag parsing
    flag.Usage = usage
    flag.Parse()

    // Check if there's a -dest flag
    if dest == "none" {
        usage()
    }

    // Format addresses
    srcaddr := fmt.Sprintf("%s:%d", src, sp)
    destaddr := fmt.Sprintf("%s:%d", dest, dp)

    // Set timeouts
    connt := time.Duration(tc) * time.Second
    readt := time.Duration(tr) * time.Second
    writet := time.Duration(tw) * time.Second

    // Connect to Redis Source
    srcc, err := redis.DialTimeout("tcp", srcaddr, connt, readt, writet)

    // Exit if connection failed (source)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Cannot connect to source redis server: %v\n", err)
        os.Exit(-1)
    }

    // Connect to Redis Destination
    destc, err := redis.DialTimeout("tcp", destaddr, connt, readt, writet)

    // Exit if connection failed (destination)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Cannot connect to destination redis server: %v\n", err)
        os.Exit(-1)
    }

    // Select DBs
    srcc.Do("SELECT", sdb)
    destc.Do("SELECT", ddb)

    // Go through each keys in source
    srckeys, _ := redis.Strings(srcc.Do("KEYS", "*"))

    for _, key := range srckeys {

        // Get the object type of each key
        rtype, _ := redis.String(srcc.Do("TYPE", key))

        fmt.Println("Migrating", key, "of type", rtype)

        // Map with the correct object data type
        switch rtype {
        case "string":

            // Raw key-value
            val, _ := redis.String(srcc.Do("GET", key))
            destc.Do("SET", key, val)
            break
        case "list":

            // Lists
            val, _ := redis.Strings(srcc.Do("LRANGE", key, 0, -1))

            for _, v := range val {
                destc.Do("RPUSH", key, v)
            }
            break
        case "set":

            // Set
            val, _ := redis.Strings(srcc.Do("SMEMBERS", key))

            for _, v := range val {
                destc.Do("SADD", key, v)
            }
            break
        case "zset":

            // Sorted Set
            val, _ := redis.Strings(srcc.Do("ZRANGE", key, 0, -1, "WITHSCORES"))

            // Since this zset will include scores we need to adapt
            zval := ""
            for i, v := range val {
                if (i + 1) % 2 == 1 {

                    // This is a value
                    zval = v
                } else {

                    // This is a score
                    // And here we do our thing
                    destc.Do("ZADD", key, v, zval)
                }
            }
            break
        case "hash":

            // Hashes, finally
            val, _ := redis.Strings(srcc.Do("HGETALL", key))

            // Do the same like in zset
            hfield := ""
            for i, v := range val {
                if (i + 1) % 2 == 1 {

                    // Value
                    hfield = v
                } else {

                    // And field
                    destc.Do("HSET", key, hfield, v)
                }
            }
            break
        default:

            // What?
            fmt.Println("Unknown type for", key)
        }

        // Check TTL
        checkttl(key, srcc, destc);
    }

    // Done
    fmt.Println("Done migrating keys from", srcaddr, "to", destaddr)

    // Cleanup
    srcc.Close()
    destc.Close()
}

// Check for key expiration and set expiry. This is a problematic one though, and may be inaccurate for a second at most
// (but I'm pretty confident it'll hardly reach that, since one command will just take a couple of ms)
func checkttl(key string, src, dst redis.Conn) {
    ttl, _ := redis.Int(src.Do("TTL", key))

    // Expire the key if it's greater than 0
    if ttl > 0 {
        fmt.Println("Setting expiration of", key)
        dst.Do("EXPIRE", key, ttl)
    }
}

// For usage
func usage() {
    fmt.Fprintf(os.Stderr, "Usage: redback [-src <source-host>] [-sp <source-port>] [-sdb <source-db>] -dest <destination-host> [-dp <destination-port>] [-ddb <destination-db>]\n")
    flag.PrintDefaults()
    os.Exit(2)
}
<|endoftext|>


<commit_before>// Redwood is an internet content-filtering program.
// It is designed to replace and improve on DansGuardian
// as the core of the Security Appliance internet filter.
package main

import (
	"flag"
	"go-icap.googlecode.com/hg"
	"log"
	"os"
	"runtime/pprof"
)

var configFile = flag.String("c", "/etc/redwood/redwood.conf", "configuration file path")
var testURL = flag.String("test", "", "URL to test instead of running ICAP server")
var cpuProfile = flag.String("cpuprofile", "", "write cpu profile to file")

func main() {
	flag.Parse()

	if *cpuProfile != "" {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	loadConfiguration()

	if *testURL != "" {
		runURLTest(*testURL)
		return
	}

	icap.HandleFunc("/reqmod", handleRequest)
	icap.HandleFunc("/respmod", handleResponse)
	icap.ListenAndServe(":1344", nil)
}
<commit_msg>Handle signals.<commit_after>// Redwood is an internet content-filtering program.
// It is designed to replace and improve on DansGuardian
// as the core of the Security Appliance internet filter.
package main

import (
	"flag"
	"go-icap.googlecode.com/hg"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"syscall"
)

var configFile = flag.String("c", "/etc/redwood/redwood.conf", "configuration file path")
var testURL = flag.String("test", "", "URL to test instead of running ICAP server")
var cpuProfile = flag.String("cpuprofile", "", "write cpu profile to file")

func main() {
	flag.Parse()

	if *cpuProfile != "" {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	loadConfiguration()

	if *testURL != "" {
		runURLTest(*testURL)
		return
	}

	go func() {
		for {
			select {
			case sig := <-signal.Incoming:
				switch sig.(os.UnixSignal) {
				case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT:
					log.Println("Terminating on signal", sig)
					if *cpuProfile != "" {
						pprof.StopCPUProfile()
					}
					os.Exit(0)
				}
			}
		}
	}()

	icap.HandleFunc("/reqmod", handleRequest)
	icap.HandleFunc("/respmod", handleResponse)
	icap.ListenAndServe(":1344", nil)
}
<|endoftext|>


<commit_before>//wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/md5"
	"fmt"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	ID string
	//Rooms the client has been added to
	rooms     []string
	roomsLock *sync.RWMutex

	conn     *ws.Conn
	connLock *sync.RWMutex
}

//Server
type Server struct {
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func([]byte, int) string
	//Called when the websocket connection is called. The only argument is
	//the disconnected client's session ID
	OnDisconnect func(string)

	handlers     map[string]func([]byte, int) ([]byte, int)
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genId(r *http.Request) string {
	hash := fmt.Sprintf("%s%d", r.RemoteAddr, time.Now().UnixNano())
	return fmt.Sprintf("%x", md5.Sum([]byte(hash)))
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	var arr []string
	client := &Client{
		ID:        genId(r),
		rooms:     arr,
		conn:      conn,
		connLock:  new(sync.RWMutex),
		roomsLock: new(sync.RWMutex)}
	s.newClient <- client

	return client, nil
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data []byte, messageType int) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	return c.conn.WriteMessage(messageType, data)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:        make(map[string]([]*Client)),
		roomsLock:    new(sync.RWMutex),
		handlers:     make(map[string](func([]byte, int) ([]byte, int))),
		handlersLock: new(sync.RWMutex),
		newClient:    make(chan *Client),
	}

	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	c.roomsLock.Lock()
	defer c.roomsLock.Unlock()
	c.rooms = append(c.rooms, r)
}

//Sends all clients in room data with type messageType
func (s *Server) Broadcast(room string, data []byte, messageType int) {
	wg := new(sync.WaitGroup)

	for _, client := range s.rooms[room] {
		go func(c *Client) {
			wg.Add(1)
			defer wg.Done()
			c.Emit(data, messageType)
		}(client)
	}

	wg.Wait()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()
	c.roomsLock.RLock()
	defer c.roomsLock.RUnlock()

	for _, room := range c.rooms {
		s.roomsLock.Lock()
		delete(s.rooms, room)
		s.roomsLock.Unlock()
	}

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.ID)
	}
}

func (s *Server) Listener() {
	for {
		c := <-s.newClient
		go func(c *Client) {
			for {
				messageType, data, err := c.conn.ReadMessage()
				if err != nil {
					c.cleanup(s)
					return
				}

				callName := s.Extractor(data, messageType)

				s.handlersLock.RLock()
				f, ok := s.handlers[callName]
				s.handlersLock.RUnlock()

				if !ok {
					continue
				}
				c.Emit(f(data, messageType))
			}
		}(c)
	}
}

//Registers a callback for the event string. The callback must take two arguments,
//a byte array it's type, and return a byte array and it's type.
func (s *Server) On(event string, f func([]byte, int) ([]byte, int)) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<commit_msg>Store joined rooms per client in server<commit_after>//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/md5"
	"fmt"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
}

//Server
type Server struct {
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs tothe list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func([]byte, int) string
	//Called when the websocket connection is called. The only argument is
	//the disconnected client's session ID
	OnDisconnect func(string)

	handlers     map[string]func([]byte, int) ([]byte, int)
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID(r *http.Request) string {
	hash := fmt.Sprintf("%s%d", r.RemoteAddr, time.Now().UnixNano())
	return fmt.Sprintf("%x", md5.Sum([]byte(hash)))
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       genID(r),
		conn:     conn,
		connLock: new(sync.RWMutex),
	}
	s.newClient <- client

	return client, nil
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data []byte, messageType int) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	return c.conn.WriteMessage(messageType, data)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func([]byte, int) ([]byte, int))),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
}

//Sends all clients in room data with type messageType
func (s *Server) Broadcast(room string, data []byte, messageType int) {
	wg := new(sync.WaitGroup)

	for _, client := range s.rooms[room] {
		go func(c *Client) {
			wg.Add(1)
			defer wg.Done()
			c.Emit(data, messageType)
		}(client)
	}

	wg.Wait()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	for _, room := range s.joinedRooms[c.id] {
		s.roomsLock.Lock()
		delete(s.rooms, room)
		s.roomsLock.Unlock()
	}

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Starts listening for events on added sockets. Needs to be called only once.
func (s *Server) Listener() {
	for {
		c := <-s.newClient
		go func(c *Client) {
			for {
				messageType, data, err := c.conn.ReadMessage()
				if err != nil {
					c.cleanup(s)
					return
				}

				callName := s.Extractor(data, messageType)

				s.handlersLock.RLock()
				f, ok := s.handlers[callName]
				s.handlersLock.RUnlock()

				if !ok {
					continue
				}
				c.Emit(f(data, messageType))
			}
		}(c)
	}
}

//Registers a callback for the event string. The callback must take two arguments,
//a byte array it's type, and return a byte array and it's type.
func (s *Server) On(event string, f func([]byte, int) ([]byte, int)) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<|endoftext|>


<commit_before>//Copyright 2015 Vibhav Pant. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)

	handlers     map[string]func(*Server, *Client, string) string
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID(r *http.Request) string {
	hash := fmt.Sprintf("%s%d", r.RemoteAddr, time.Now().UnixNano())
	return fmt.Sprintf("%x", sha1.Sum([]byte(hash)))
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       genID(r),
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := struct {
		Id   int         `json:"id"`
		Data interface{} `json:"data"`
	}{-1, v}

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, string) string)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, clientID := range s.joinedRooms[c.id] {
		if clientID == c.id {
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.Lock()

	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
		}
	}
	if index == -1 {
		s.roomsLock.Unlock()
		return
	}

	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	if index == -1 {
		return
	}

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data with type messageType
func (s *Server) Broadcast(room string, data string) {
	wg := new(sync.WaitGroup)

	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			wg.Add(1)
			defer wg.Done()
			c.Emit(data)
		}(client)
	}

	wg.Wait()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	wg := new(sync.WaitGroup)

	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		wg.Add(1)
		go func(c *Client) {
			defer wg.Done()
			c.EmitJSON(v)
		}(client)
	}

	wg.Wait()

}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.roomsLock.Lock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}
		if index == -1 {
			s.roomsLock.Unlock()
			return
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
	}
	s.roomsLock.Unlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	var rooms []string
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

//Starts listening for events on added sockets. Needs to be called only once.
func (s *Server) Listener() {
	for {
		c := <-s.newClient
		go func(c *Client) {
			for {
				mtype, data, err := c.conn.ReadMessage()
				if err != nil {
					c.cleanup(s)
					return
				}

				var js struct {
					Id   string
					Data json.RawMessage
				}
				err = json.Unmarshal(data, &js)

				if err != nil || mtype != ws.TextMessage {
					log.Println(err)
					continue
				}

				callName := s.Extractor(string(js.Data))

				s.handlersLock.RLock()
				f, ok := s.handlers[callName]
				s.handlersLock.RUnlock()

				if !ok {
					continue
				}

				rtrn := f(s, c, string(js.Data))
				reply := struct {
					Id   string `json:"id"`
					Data string `json:"data,string"`
				}{js.Id, rtrn}

				bytes, _ := json.Marshal(reply)
				c.Emit(string(bytes))
			}
		}(c)
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, string) string) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<commit_msg>Add more loggin messages.<commit_after>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)

	handlers     map[string]func(*Server, *Client, string) string
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID(r *http.Request) string {
	hash := fmt.Sprintf("%s%d", r.RemoteAddr, time.Now().UnixNano())
	return fmt.Sprintf("%x", sha1.Sum([]byte(hash)))
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       genID(r),
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := struct {
		Id   int         `json:"id"`
		Data interface{} `json:"data"`
	}{-1, v}

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, string) string)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, clientID := range s.joinedRooms[c.id] {
		if clientID == c.id {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.Lock()

	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		s.roomsLock.Unlock()
		return
	}

	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	if index == -1 {
		return
	}

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data with type messageType
func (s *Server) Broadcast(room string, data string) {
	wg := new(sync.WaitGroup)

	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			wg.Add(1)
			defer wg.Done()
			c.Emit(data)
		}(client)
	}

	wg.Wait()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	wg := new(sync.WaitGroup)

	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		wg.Add(1)
		go func(c *Client) {
			defer wg.Done()
			c.EmitJSON(v)
		}(client)
	}

	wg.Wait()

}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.roomsLock.Lock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}
		if index == -1 {
			s.roomsLock.Unlock()
			return
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
	}
	s.roomsLock.Unlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	var rooms []string
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

//Starts listening for events on added sockets. Needs to be called only once.
func (s *Server) Listener() {
	for {
		c := <-s.newClient
		go func(c *Client) {
			for {
				mtype, data, err := c.conn.ReadMessage()
				if err != nil {
					c.cleanup(s)
					return
				}

				var js struct {
					Id   string
					Data json.RawMessage
				}
				err = json.Unmarshal(data, &js)

				if err != nil || mtype != ws.TextMessage {
					log.Println(err)
					continue
				}

				callName := s.Extractor(string(js.Data))

				s.handlersLock.RLock()
				f, ok := s.handlers[callName]
				s.handlersLock.RUnlock()

				if !ok {
					continue
				}

				rtrn := f(s, c, string(js.Data))
				reply := struct {
					Id   string `json:"id"`
					Data string `json:"data,string"`
				}{js.Id, rtrn}

				bytes, _ := json.Marshal(reply)
				c.Emit(string(bytes))
			}
		}(c)
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, string) string) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<|endoftext|>