<commit_before>package plugins

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/plugins/backend"
	"github.com/grafana/grafana/pkg/tsdb"
	shared "github.com/grafana/grafana/pkg/tsdb/models/proxy"
	plugin "github.com/hashicorp/go-plugin"
)

type DataSourcePlugin struct {
	FrontendPluginBase
	Annotations  bool              `json:"annotations"`
	Metrics      bool              `json:"metrics"`
	Alerting     bool              `json:"alerting"`
	QueryOptions map[string]bool   `json:"queryOptions,omitempty"`
	BuiltIn      bool              `json:"builtIn,omitempty"`
	Mixed        bool              `json:"mixed,omitempty"`
	HasQueryHelp bool              `json:"hasQueryHelp,omitempty"`
	Routes       []*AppPluginRoute `json:"routes"`

	Backend    bool   `json:"backend,omitempty"`
	Executable string `json:"executable,omitempty"`

	log    log.Logger
	client *plugin.Client
}

func (p *DataSourcePlugin) Load(decoder *json.Decoder, pluginDir string) error {
	if err := decoder.Decode(&p); err != nil {
		return err
	}

	if err := p.registerPlugin(pluginDir); err != nil {
		return err
	}

	// look for help markdown
	helpPath := filepath.Join(p.PluginDir, "QUERY_HELP.md")
	if _, err := os.Stat(helpPath); os.IsNotExist(err) {
		helpPath = filepath.Join(p.PluginDir, "query_help.md")
	}
	if _, err := os.Stat(helpPath); err == nil {
		p.HasQueryHelp = true
	}

	DataSources[p.Id] = p
	return nil
}

var handshakeConfig = plugin.HandshakeConfig{
	ProtocolVersion:  1,
	MagicCookieKey:   "GRAFANA_BACKEND_DATASOURCE",
	MagicCookieValue: "55d2200a-6492-493a-9353-73b728d468aa",
}

func buildExecutablePath(pluginDir, executable, os, arch string) string {
	return path.Join(pluginDir, fmt.Sprintf("%s_%s_%s", executable, strings.ToLower(os), strings.ToLower(arch)))
}

func (p *DataSourcePlugin) initBackendPlugin(ctx context.Context, log log.Logger) error {
	p.log = log.New("plugin-id", p.Id)

	err := p.spawnSubProcess()
	if err == nil {
		go p.restartKilledProcess(ctx)
	}

	return err
}

func (p *DataSourcePlugin) spawnSubProcess() error {
	cmd := buildExecutablePath(p.PluginDir, p.Executable, runtime.GOOS, runtime.GOARCH)

	p.client = plugin.NewClient(&plugin.ClientConfig{
		HandshakeConfig:  handshakeConfig,
		Plugins:          map[string]plugin.Plugin{p.Id: &shared.TsdbPluginImpl{}},
		Cmd:              exec.Command(cmd),
		AllowedProtocols: []plugin.Protocol{plugin.ProtocolGRPC},
		Logger:           backend.LogWrapper{Logger: p.log},
	})

	rpcClient, err := p.client.Client()
	if err != nil {
		return err
	}

	raw, err := rpcClient.Dispense(p.Id)
	if err != nil {
		return err
	}

	plugin := raw.(shared.TsdbPlugin)

	tsdb.RegisterTsdbQueryEndpoint(p.Id, func(dsInfo *models.DataSource) (tsdb.TsdbQueryEndpoint, error) {
		return &shared.DatasourcePluginWrapper{TsdbPlugin: plugin}, nil
	})

	return nil
}

func (p *DataSourcePlugin) restartKilledProcess(ctx context.Context) error {
	ticker := time.NewTicker(time.Second * 1)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			if p.client.Exited() {
				err := p.spawnSubProcess()
				p.log.Debug("Spawning new sub process", "name", p.Name, "id", p.Id)
				if err != nil {
					p.log.Error("Failed to spawn subprocess")
				}
			}
		}
	}
}

func (p *DataSourcePlugin) Kill() {
	if p.client != nil {
		p.log.Debug("Killing subprocess ", "name", p.Name)
		p.client.Kill()
	}
}
<commit_msg>makes datasource handshake more explicit<commit_after>package plugins

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/plugins/backend"
	"github.com/grafana/grafana/pkg/tsdb"
	shared "github.com/grafana/grafana/pkg/tsdb/models/proxy"
	plugin "github.com/hashicorp/go-plugin"
)

type DataSourcePlugin struct {
	FrontendPluginBase
	Annotations  bool              `json:"annotations"`
	Metrics      bool              `json:"metrics"`
	Alerting     bool              `json:"alerting"`
	QueryOptions map[string]bool   `json:"queryOptions,omitempty"`
	BuiltIn      bool              `json:"builtIn,omitempty"`
	Mixed        bool              `json:"mixed,omitempty"`
	HasQueryHelp bool              `json:"hasQueryHelp,omitempty"`
	Routes       []*AppPluginRoute `json:"routes"`

	Backend    bool   `json:"backend,omitempty"`
	Executable string `json:"executable,omitempty"`

	log    log.Logger
	client *plugin.Client
}

func (p *DataSourcePlugin) Load(decoder *json.Decoder, pluginDir string) error {
	if err := decoder.Decode(&p); err != nil {
		return err
	}

	if err := p.registerPlugin(pluginDir); err != nil {
		return err
	}

	// look for help markdown
	helpPath := filepath.Join(p.PluginDir, "QUERY_HELP.md")
	if _, err := os.Stat(helpPath); os.IsNotExist(err) {
		helpPath = filepath.Join(p.PluginDir, "query_help.md")
	}
	if _, err := os.Stat(helpPath); err == nil {
		p.HasQueryHelp = true
	}

	DataSources[p.Id] = p
	return nil
}

var handshakeConfig = plugin.HandshakeConfig{
	ProtocolVersion:  1,
	MagicCookieKey:   "grafana_plugin_type",
	MagicCookieValue: "datasource",
}

func buildExecutablePath(pluginDir, executable, os, arch string) string {
	return path.Join(pluginDir, fmt.Sprintf("%s_%s_%s", executable, strings.ToLower(os), strings.ToLower(arch)))
}

func (p *DataSourcePlugin) initBackendPlugin(ctx context.Context, log log.Logger) error {
	p.log = log.New("plugin-id", p.Id)

	err := p.spawnSubProcess()
	if err == nil {
		go p.restartKilledProcess(ctx)
	}

	return err
}

func (p *DataSourcePlugin) spawnSubProcess() error {
	cmd := buildExecutablePath(p.PluginDir, p.Executable, runtime.GOOS, runtime.GOARCH)

	p.client = plugin.NewClient(&plugin.ClientConfig{
		HandshakeConfig:  handshakeConfig,
		Plugins:          map[string]plugin.Plugin{p.Id: &shared.TsdbPluginImpl{}},
		Cmd:              exec.Command(cmd),
		AllowedProtocols: []plugin.Protocol{plugin.ProtocolGRPC},
		Logger:           backend.LogWrapper{Logger: p.log},
	})

	rpcClient, err := p.client.Client()
	if err != nil {
		return err
	}

	raw, err := rpcClient.Dispense(p.Id)
	if err != nil {
		return err
	}

	plugin := raw.(shared.TsdbPlugin)

	tsdb.RegisterTsdbQueryEndpoint(p.Id, func(dsInfo *models.DataSource) (tsdb.TsdbQueryEndpoint, error) {
		return &shared.DatasourcePluginWrapper{TsdbPlugin: plugin}, nil
	})

	return nil
}

func (p *DataSourcePlugin) restartKilledProcess(ctx context.Context) error {
	ticker := time.NewTicker(time.Second * 1)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			if p.client.Exited() {
				err := p.spawnSubProcess()
				p.log.Debug("Spawning new sub process", "name", p.Name, "id", p.Id)
				if err != nil {
					p.log.Error("Failed to spawn subprocess")
				}
			}
		}
	}
}

func (p *DataSourcePlugin) Kill() {
	if p.client != nil {
		p.log.Debug("Killing subprocess ", "name", p.Name)
		p.client.Kill()
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package cachepool

import (
	"sync"
	"time"

	"github.com/wandoulabs/codis/pkg/proxy/redispool"

	"github.com/juju/errors"
)

type LivePool struct {
	pool *redispool.ConnectionPool
}

type CachePool struct {
	mu    sync.RWMutex
	pools map[string]*LivePool
}

func NewCachePool() *CachePool {
	return &CachePool{
		pools: make(map[string]*LivePool),
	}
}

func (cp *CachePool) GetConn(key string) (redispool.PoolConnection, error) {
	cp.mu.RLock()

	pool, ok := cp.pools[key]
	if !ok {
		cp.mu.RUnlock()
		return nil, errors.Errorf("pool %s not exist", key)
	}

	c, err := pool.pool.Get()

	cp.mu.RUnlock()
	return c, err
}

func (cp *CachePool) ReleaseConn(pc redispool.PoolConnection) {
	pc.Recycle()
}

func (cp *CachePool) AddPool(key string) error {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	pool, ok := cp.pools[key]
	if ok {
		return nil
	}
	pool = &LivePool{
		pool: redispool.NewConnectionPool("redis conn pool", 50, 120*time.Second),
	}

	pool.pool.Open(redispool.ConnectionCreator(key))

	cp.pools[key] = pool

	return nil
}

func (cp *CachePool) RemovePool(key string) error {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	pool, ok := cp.pools[key]
	if !ok {
		return errors.Errorf("pool %s already exist", key)
	}

	pool.pool.Close() //todo:async close
	delete(cp.pools, key)
	return nil
}
<commit_msg>fix error message<commit_after>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package cachepool

import (
	"sync"
	"time"

	"github.com/wandoulabs/codis/pkg/proxy/redispool"

	"github.com/juju/errors"
)

type LivePool struct {
	pool *redispool.ConnectionPool
}

type CachePool struct {
	mu    sync.RWMutex
	pools map[string]*LivePool
}

func NewCachePool() *CachePool {
	return &CachePool{
		pools: make(map[string]*LivePool),
	}
}

func (cp *CachePool) GetConn(key string) (redispool.PoolConnection, error) {
	cp.mu.RLock()

	pool, ok := cp.pools[key]
	if !ok {
		cp.mu.RUnlock()
		return nil, errors.Errorf("pool %s not exist", key)
	}

	c, err := pool.pool.Get()

	cp.mu.RUnlock()
	return c, err
}

func (cp *CachePool) ReleaseConn(pc redispool.PoolConnection) {
	pc.Recycle()
}

func (cp *CachePool) AddPool(key string) error {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	pool, ok := cp.pools[key]
	if ok {
		return nil
	}
	pool = &LivePool{
		pool: redispool.NewConnectionPool("redis conn pool", 50, 120*time.Second),
	}

	pool.pool.Open(redispool.ConnectionCreator(key))

	cp.pools[key] = pool

	return nil
}

func (cp *CachePool) RemovePool(key string) error {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	pool, ok := cp.pools[key]
	if !ok {
		return errors.Errorf("pool %s not exist", key)
	}

	pool.pool.Close() //todo:async close
	delete(cp.pools, key)
	return nil
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"io/ioutil"
	"strconv"
	"strings"
	"time"

	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
)

type Pps struct {
	Ui  cli.Ui
	Cmd string

	interval time.Duration
}

func (this *Pps) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("pps", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.DurationVar(&this.interval, "sleep", time.Second*2, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if len(args) == 0 {
		this.Ui.Error("missing nic")
		return 2
	}

	nic := args[len(args)-1]
	this.showPps(nic)

	return
}

func (this *Pps) showPps(nic string) {
	tx := fmt.Sprintf("/sys/class/net/%s/statistics/tx_packets", nic)
	rx := fmt.Sprintf("/sys/class/net/%s/statistics/rx_packets", nic)

	var lastTx, lastRx int64
	s := int64(this.interval.Seconds())
	for {
		brx, err := ioutil.ReadFile(rx)
		swallow(err)
		btx, err := ioutil.ReadFile(tx)
		swallow(err)

		rxN, err := strconv.ParseInt(string(brx), 10, 64)
		swallow(err)
		txN, err := strconv.ParseInt(string(btx), 10, 64)
		swallow(err)

		if lastRx != 0 && lastRx != 0 {
			rxPps := (rxN - lastRx) / s
			txPps := (txN - lastTx) / s
			sumPps := rxPps + txPps
			this.Ui.Output(fmt.Sprintf("%10s rx:%-8s tx:%-8s sum:%-8s",
				nic, gofmt.Comma(rxPps), gofmt.Comma(txPps), gofmt.Comma(sumPps)))
		}

		lastRx = rxN
		lastRx = txN

		time.Sleep(this.interval)
	}
}

func (*Pps) Synopsis() string {
	return "Show PacketPerSecond of a NIC on Linux"
}

func (this *Pps) Help() string {
	help := fmt.Sprintf(`
Usage: %s pps [option] nic

    %s

Options:

    -sleep duration

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

import (
	"flag"
	"fmt"
	"io/ioutil"
	"strconv"
	"strings"
	"time"

	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
)

type Pps struct {
	Ui  cli.Ui
	Cmd string

	interval time.Duration
}

func (this *Pps) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("pps", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.DurationVar(&this.interval, "sleep", time.Second*2, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if len(args) == 0 {
		this.Ui.Error("missing nic")
		return 2
	}

	nic := args[len(args)-1]
	this.showPps(nic)

	return
}

func (this *Pps) showPps(nic string) {
	tx := fmt.Sprintf("/sys/class/net/%s/statistics/tx_packets", nic)
	rx := fmt.Sprintf("/sys/class/net/%s/statistics/rx_packets", nic)

	var lastTx, lastRx int64
	s := int64(this.interval.Seconds())
	for {
		brx, err := ioutil.ReadFile(rx)
		swallow(err)
		btx, err := ioutil.ReadFile(tx)
		swallow(err)

		rxN, err := strconv.ParseInt(strings.TrimSpace(string(brx)), 10, 64)
		swallow(err)
		txN, err := strconv.ParseInt(strings.TrimSpace(string(btx)), 10, 64)
		swallow(err)

		if lastRx != 0 && lastTx != 0 {
			rxPps := (rxN - lastRx) / s
			txPps := (txN - lastTx) / s
			sumPps := rxPps + txPps

			this.Ui.Output(fmt.Sprintf("%10s rx:%-8s tx:%-8s sum:%-8s",
				nic, gofmt.Comma(rxPps), gofmt.Comma(txPps), gofmt.Comma(sumPps)))
		}

		lastRx = rxN
		lastTx = txN

		time.Sleep(this.interval)
	}
}

func (*Pps) Synopsis() string {
	return "Show PacketPerSecond of a NIC on Linux"
}

func (this *Pps) Help() string {
	help := fmt.Sprintf(`
Usage: %s pps [option] nic

    %s

Options:

    -sleep duration

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


<commit_before>/*
Copyright 2019 The Skaffold Authors

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 initializer

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/GoogleContainerTools/skaffold/cmd/skaffold/app/tips"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/analyze"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/config"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/deploy"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/prompt"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/yaml"
)

// DoInit executes the `skaffold init` flow.
func DoInit(ctx context.Context, out io.Writer, c config.Config) error {
	if c.ComposeFile != "" {
		if err := runKompose(ctx, c.ComposeFile); err != nil {
			return err
		}
	}

	a := analyze.NewAnalyzer(c)
	if err := a.Analyze("."); err != nil {
		return err
	}

	// helm projects can't currently be bootstrapped automatically by skaffold, so we fail fast and link to our docs instead.
	if len(a.ChartPaths()) > 0 {
		//nolint
		return errors.New(`Projects set up to deploy with helm must be configured with skaffold manually.

See https://skaffold.dev/docs/pipeline-stages/deployers/helm/ for a detailed guide on setting your project up with skaffold.`)
	}

	deployInitializer := deploy.NewInitializer(a.Manifests(), a.KustomizeBases(), a.KustomizePaths(), c)
	images := deployInitializer.GetImages()

	buildInitializer := build.NewInitializer(a.Builders(), c)
	if err := buildInitializer.ProcessImages(images); err != nil {
		return err
	}

	if c.Analyze {
		return buildInitializer.PrintAnalysis(out)
	}

	var generatedManifests map[string][]byte
	if c.EnableManifestGeneration {
		generatedManifestPairs, err := buildInitializer.GenerateManifests()
		if err != nil {
			return err
		}
		generatedManifests = make(map[string][]byte, len(generatedManifestPairs))
		for pair, manifest := range generatedManifestPairs {
			deployInitializer.AddManifestForImage(pair.ManifestPath, pair.ImageName)
			generatedManifests[pair.ManifestPath] = manifest
		}
	}

	if err := deployInitializer.Validate(); err != nil {
		return err
	}

	pipeline, err := yaml.Marshal(generateSkaffoldConfig(buildInitializer, deployInitializer))
	if err != nil {
		return err
	}
	if c.Opts.ConfigurationFile == "-" {
		out.Write(pipeline)
		return nil
	}

	if !c.Force {
		if done, err := prompt.WriteSkaffoldConfig(out, pipeline, generatedManifests, c.Opts.ConfigurationFile); done {
			return err
		}
	}

	for path, manifest := range generatedManifests {
		if err = ioutil.WriteFile(path, manifest, 0644); err != nil {
			return fmt.Errorf("writing k8s manifest to file: %w", err)
		}
		fmt.Fprintf(out, "Generated manifest %s was written\n", path)
	}

	if err = ioutil.WriteFile(c.Opts.ConfigurationFile, pipeline, 0644); err != nil {
		return fmt.Errorf("writing config to file: %w", err)
	}

	fmt.Fprintf(out, "Configuration %s was written\n", c.Opts.ConfigurationFile)
	tips.PrintForInit(out, c.Opts)

	return nil
}
<commit_msg>Update pkg/skaffold/initializer/init.go<commit_after>/*
Copyright 2019 The Skaffold Authors

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 initializer

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/GoogleContainerTools/skaffold/cmd/skaffold/app/tips"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/analyze"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/config"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/deploy"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/initializer/prompt"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/yaml"
)

// DoInit executes the `skaffold init` flow.
func DoInit(ctx context.Context, out io.Writer, c config.Config) error {
	if c.ComposeFile != "" {
		if err := runKompose(ctx, c.ComposeFile); err != nil {
			return err
		}
	}

	a := analyze.NewAnalyzer(c)
	if err := a.Analyze("."); err != nil {
		return err
	}

	// helm projects can't currently be bootstrapped automatically by skaffold, so we fail fast and link to our docs instead.
	if len(a.ChartPaths()) > 0 {
		//nolint
		return errors.New(`Projects set up to deploy with helm must be manually configured.

See https://skaffold.dev/docs/pipeline-stages/deployers/helm/ for a detailed guide on setting your project up with skaffold.`)
	}

	deployInitializer := deploy.NewInitializer(a.Manifests(), a.KustomizeBases(), a.KustomizePaths(), c)
	images := deployInitializer.GetImages()

	buildInitializer := build.NewInitializer(a.Builders(), c)
	if err := buildInitializer.ProcessImages(images); err != nil {
		return err
	}

	if c.Analyze {
		return buildInitializer.PrintAnalysis(out)
	}

	var generatedManifests map[string][]byte
	if c.EnableManifestGeneration {
		generatedManifestPairs, err := buildInitializer.GenerateManifests()
		if err != nil {
			return err
		}
		generatedManifests = make(map[string][]byte, len(generatedManifestPairs))
		for pair, manifest := range generatedManifestPairs {
			deployInitializer.AddManifestForImage(pair.ManifestPath, pair.ImageName)
			generatedManifests[pair.ManifestPath] = manifest
		}
	}

	if err := deployInitializer.Validate(); err != nil {
		return err
	}

	pipeline, err := yaml.Marshal(generateSkaffoldConfig(buildInitializer, deployInitializer))
	if err != nil {
		return err
	}
	if c.Opts.ConfigurationFile == "-" {
		out.Write(pipeline)
		return nil
	}

	if !c.Force {
		if done, err := prompt.WriteSkaffoldConfig(out, pipeline, generatedManifests, c.Opts.ConfigurationFile); done {
			return err
		}
	}

	for path, manifest := range generatedManifests {
		if err = ioutil.WriteFile(path, manifest, 0644); err != nil {
			return fmt.Errorf("writing k8s manifest to file: %w", err)
		}
		fmt.Fprintf(out, "Generated manifest %s was written\n", path)
	}

	if err = ioutil.WriteFile(c.Opts.ConfigurationFile, pipeline, 0644); err != nil {
		return fmt.Errorf("writing config to file: %w", err)
	}

	fmt.Fprintf(out, "Configuration %s was written\n", c.Opts.ConfigurationFile)
	tips.PrintForInit(out, c.Opts)

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/brutella/hc/accessory"
	"github.com/hemtjanst/hemtjanst/device"
	"github.com/hemtjanst/hemtjanst/homekit"
	"github.com/hemtjanst/hemtjanst/homekit/bridge"
	"github.com/hemtjanst/hemtjanst/messaging"
	"github.com/hemtjanst/hemtjanst/messaging/flagmqtt"
	"github.com/hemtjanst/hemtjanst/web"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	name     = flag.String("name", "hemtjanst", "Name of bridge instance")
	addr     = flag.String("address", "127.0.0.1", "IP or hostname for Hemtjänst to bind on")
	port     = flag.String("port", "12345", "Port for Hemtjänst to bind on")
	pin      = flag.String("pin", "01020304", "Pairing pin for the HomeKit bridge")
	startWeb = flag.Bool("web.ui", false, "Start the built-in web UI")
	wAddr    = flag.String("web.addr", ":8080", "IP/host:port to bind the webinterface to")
	dbPath   = flag.String("db.path", "./db", "Path to store the database with HomeKit key pairs etc.")
	hVersion = flag.Bool("version", false, "Print the version")

	version = "master"
)

const (
	announceTopicPrefix = "announce/"
	leaveTopic          = "leave"
	discoverTopic       = "discover"
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Parameters:\n\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}
	flag.Parse()
	if *hVersion {
		fmt.Println(version)
		os.Exit(0)
	}

	log.Print("Initialing Hemtjänst")
	quit := make(chan os.Signal)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	announce := make(chan messaging.Message)
	leave := make(chan messaging.Message)
	handlerInit := make(chan bool)
	managerInit := make(chan bool)

	bridgeConfig := bridge.Config{
		Pin:         *pin,
		Port:        *port,
		IP:          *addr,
		StoragePath: *dbPath,
	}
	bridgeInfo := accessory.Info{
		Name:         *name,
		SerialNumber: "12345",
		Manufacturer: "BEDS Inc.",
		Model:        "v0.1",
	}

	log.Print("Attempting to connect to MQTT broker")
	handler := &messaging.Handler{
		Ann:           announce,
		Leave:         leave,
		AnnounceTopic: announceTopicPrefix + "#",
		LeaveTopic:    leaveTopic,
		DiscoverTopic: discoverTopic,
		// Wait 5 seconds before sending discover to allow
		// persistent announcements to be fetched
		DiscoverDelay: 5 * time.Second,
		DiscoverStart: handlerInit,
	}
	cID := flagmqtt.NewUniqueIdentifier()
	conf := flagmqtt.ClientConfig{
		ClientID:                fmt.Sprintf("hemtjanst-%s", cID),
		OnConnectHandler:        handler.OnConnect,
		OnConnectionLostHandler: handler.OnConnectionLost,
		WillTopic:               "leave",
		WillPayload:             fmt.Sprintf("hemtjanst-%s", cID),
		WillRetain:              false,
		WillQoS:                 0,
	}

	c, err := flagmqtt.NewPersistentMqtt(conf)
	if err != nil {
		log.Fatal("Could not configure the MQTT client: ", err)
	}

	go func() {
		if token := c.Connect(); token.Wait() && token.Error() != nil {
			log.Fatal("Failed to establish connection with broker: ", token.Error())
		}
	}()

	hkBridge, err := bridge.NewBridge(bridgeConfig, bridgeInfo)
	if err != nil {
		log.Fatal("Could not start HomeKit bridge: ", err)
	}

	manager := device.NewManager(messaging.NewMQTTMessenger(c), managerInit)
	log.Print("Started device manager")

	hk := homekit.NewHomekit(hkBridge, manager)

	go func() {
		// Wait for handler to have sent its discover
		<-handlerInit

		// Tell manager that we're initialised
		managerInit <- true

		// Wait a few more seconds before starting bridge
		<-time.After(5 * time.Second)
		log.Print("Starting HomeKit bridge")
		manager.AddHandler(hk)
		hkBridge.Start()
	}()

	if *startWeb {
		go func() {
			<-time.After(5 * time.Second)
			web.Serve(manager, *wAddr)
		}()
		log.Print("Started web interface")
	}

loop:
	for {
		select {
		case sig := <-quit:
			log.Printf("Received signal: %s, proceeding to shutdown", sig)
			break loop
		case msg := <-announce:
			newReg := msg.Topic()
			log.Print("New announcement: ", newReg)
			if len(newReg) <= len(announceTopicPrefix) || newReg[0:len(announceTopicPrefix)] != announceTopicPrefix {
				// Announcement doesn't start with prefix
				log.Printf("Ignoring: Announcement doesn't start with " + announceTopicPrefix)
				continue
			}
			newReg = newReg[len(announceTopicPrefix):]
			if !strings.Contains(newReg, "/") {
				// We expect topics we care about to contain at least 1 /
				log.Printf("Ignoring: Malformed topic %s", newReg)
				continue
			}
			if len(msg.Payload()) == 0 {
				// Empty payload, remove device
				go manager.Remove(newReg)
				continue
			}

			go manager.Add(newReg, msg.Payload())
		case msg := <-leave:
			go manager.Leave(string(msg.Payload()))
		}
	}

	// When the MQTT lib is connecting but hasn't establish a conneciton yet
	// the IsConnected() method returns true. As a consequence, b/c it believes
	// it is connected the call to .Disconnect() will panic if we terminate
	// before we've managed to establish a connection to the broker, as it
	// tries to close one of its own channels that are currently still nil.
	//
	// To avoid getting an ugly panic printed for what is arguably a bug in the
	// library defer a recover that does nothing and exit normally.
	defer func() {
		recover()
	}()

	c.Disconnect(250)
	hkBridge.Stop()
	log.Print("Disconnected from broker. Bye!")
	os.Exit(0)
}
<commit_msg>:pencil: Register devices with HK before starting<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/brutella/hc/accessory"
	"github.com/hemtjanst/hemtjanst/device"
	"github.com/hemtjanst/hemtjanst/homekit"
	"github.com/hemtjanst/hemtjanst/homekit/bridge"
	"github.com/hemtjanst/hemtjanst/messaging"
	"github.com/hemtjanst/hemtjanst/messaging/flagmqtt"
	"github.com/hemtjanst/hemtjanst/web"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	name     = flag.String("name", "hemtjanst", "Name of bridge instance")
	addr     = flag.String("address", "127.0.0.1", "IP or hostname for Hemtjänst to bind on")
	port     = flag.String("port", "12345", "Port for Hemtjänst to bind on")
	pin      = flag.String("pin", "01020304", "Pairing pin for the HomeKit bridge")
	startWeb = flag.Bool("web.ui", false, "Start the built-in web UI")
	wAddr    = flag.String("web.addr", ":8080", "IP/host:port to bind the webinterface to")
	dbPath   = flag.String("db.path", "./db", "Path to store the database with HomeKit key pairs etc.")
	hVersion = flag.Bool("version", false, "Print the version")

	version = "master"
)

const (
	announceTopicPrefix = "announce/"
	leaveTopic          = "leave"
	discoverTopic       = "discover"
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Parameters:\n\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}
	flag.Parse()
	if *hVersion {
		fmt.Println(version)
		os.Exit(0)
	}

	log.Print("Initialing Hemtjänst")
	quit := make(chan os.Signal)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	announce := make(chan messaging.Message)
	leave := make(chan messaging.Message)
	handlerInit := make(chan bool)
	managerInit := make(chan bool)

	bridgeConfig := bridge.Config{
		Pin:         *pin,
		Port:        *port,
		IP:          *addr,
		StoragePath: *dbPath,
	}
	bridgeInfo := accessory.Info{
		Name:         *name,
		SerialNumber: "12345",
		Manufacturer: "BEDS Inc.",
		Model:        "v0.1",
	}

	log.Print("Attempting to connect to MQTT broker")
	handler := &messaging.Handler{
		Ann:           announce,
		Leave:         leave,
		AnnounceTopic: announceTopicPrefix + "#",
		LeaveTopic:    leaveTopic,
		DiscoverTopic: discoverTopic,
		// Wait 5 seconds before sending discover to allow
		// persistent announcements to be fetched
		DiscoverDelay: 5 * time.Second,
		DiscoverStart: handlerInit,
	}
	cID := flagmqtt.NewUniqueIdentifier()
	conf := flagmqtt.ClientConfig{
		ClientID:                fmt.Sprintf("hemtjanst-%s", cID),
		OnConnectHandler:        handler.OnConnect,
		OnConnectionLostHandler: handler.OnConnectionLost,
		WillTopic:               "leave",
		WillPayload:             fmt.Sprintf("hemtjanst-%s", cID),
		WillRetain:              false,
		WillQoS:                 0,
	}

	c, err := flagmqtt.NewPersistentMqtt(conf)
	if err != nil {
		log.Fatal("Could not configure the MQTT client: ", err)
	}

	go func() {
		if token := c.Connect(); token.Wait() && token.Error() != nil {
			log.Fatal("Failed to establish connection with broker: ", token.Error())
		}
	}()

	hkBridge, err := bridge.NewBridge(bridgeConfig, bridgeInfo)
	if err != nil {
		log.Fatal("Could not start HomeKit bridge: ", err)
	}

	manager := device.NewManager(messaging.NewMQTTMessenger(c), managerInit)
	log.Print("Started device manager")

	hk := homekit.NewHomekit(hkBridge, manager)
	manager.AddHandler(hk)

	go func() {
		// Wait for handler to have sent its discover
		<-handlerInit

		// Tell manager that we're initialised
		managerInit <- true

		// Wait a few more seconds before starting bridge
		<-time.After(5 * time.Second)
		log.Print("Starting HomeKit bridge")
		hkBridge.Start()
	}()

	if *startWeb {
		go func() {
			<-time.After(5 * time.Second)
			web.Serve(manager, *wAddr)
		}()
		log.Print("Started web interface")
	}

loop:
	for {
		select {
		case sig := <-quit:
			log.Printf("Received signal: %s, proceeding to shutdown", sig)
			break loop
		case msg := <-announce:
			newReg := msg.Topic()
			log.Print("New announcement: ", newReg)
			if len(newReg) <= len(announceTopicPrefix) || newReg[0:len(announceTopicPrefix)] != announceTopicPrefix {
				// Announcement doesn't start with prefix
				log.Printf("Ignoring: Announcement doesn't start with " + announceTopicPrefix)
				continue
			}
			newReg = newReg[len(announceTopicPrefix):]
			if !strings.Contains(newReg, "/") {
				// We expect topics we care about to contain at least 1 /
				log.Printf("Ignoring: Malformed topic %s", newReg)
				continue
			}
			if len(msg.Payload()) == 0 {
				// Empty payload, remove device
				go manager.Remove(newReg)
				continue
			}

			go manager.Add(newReg, msg.Payload())
		case msg := <-leave:
			go manager.Leave(string(msg.Payload()))
		}
	}

	// When the MQTT lib is connecting but hasn't establish a conneciton yet
	// the IsConnected() method returns true. As a consequence, b/c it believes
	// it is connected the call to .Disconnect() will panic if we terminate
	// before we've managed to establish a connection to the broker, as it
	// tries to close one of its own channels that are currently still nil.
	//
	// To avoid getting an ugly panic printed for what is arguably a bug in the
	// library defer a recover that does nothing and exit normally.
	defer func() {
		recover()
	}()

	c.Disconnect(250)
	hkBridge.Stop()
	log.Print("Disconnected from broker. Bye!")
	os.Exit(0)
}
<|endoftext|>


<commit_before>package dynamodb

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"log"
	"math"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	godynamodb "github.com/aws/aws-sdk-go/service/dynamodb"
	godynamodbiface "github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/pkg/config"
	"github.com/yuuki/diamondb/pkg/model"
	"github.com/yuuki/diamondb/pkg/storage/util"
	"github.com/yuuki/diamondb/pkg/timeparser"
)

//go:generate mockgen -source ../../../vendor/github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface/interface.go -destination dynamodb_mock.go -package dynamodb

// ReadWriter defines the interface for DynamoDB reader and writer.
type ReadWriter interface {
	Ping() error
	Client() godynamodbiface.DynamoDBAPI
	CreateTable(*CreateTableParam) error
	Fetch(string, time.Time, time.Time) (model.SeriesMap, error)
	batchGet(q *query) (model.SeriesMap, error)
	Put(string, string, string, int64, map[int64]float64) error
}

// DynamoDB provides a dynamodb client.
type DynamoDB struct {
	svc godynamodbiface.DynamoDBAPI
}

type timeSlot struct {
	itemEpoch int64
	step      int
}

type query struct {
	names []string
	start time.Time
	end   time.Time
	slot  *timeSlot
	// context
}

const (
	pingTimeout        = time.Duration(5) * time.Second
	dynamodbBatchLimit = 100

	oneYear time.Duration = time.Duration(24*360) * time.Hour
	oneWeek time.Duration = time.Duration(24*7) * time.Hour
	oneDay  time.Duration = time.Duration(24*1) * time.Hour
)

var (
	oneYearSeconds = int(oneYear.Seconds())
	oneWeekSeconds = int(oneWeek.Seconds())
	oneDaySeconds  = int(oneDay.Seconds())
)

var _ ReadWriter = &DynamoDB{}

// New creates a new DynamoDB.
func New() (*DynamoDB, error) {
	awsConf := aws.NewConfig().WithRegion(config.Config.DynamoDBRegion)
	if config.Config.DynamoDBEndpoint != "" {
		// For dynamodb-local configuration
		awsConf.WithEndpoint(config.Config.DynamoDBEndpoint)
		awsConf.WithCredentials(credentials.NewStaticCredentials("dummy", "dummy", "dummy"))
	}
	awsConf.WithHTTPClient(&http.Client{
		Timeout:   10 * time.Second,
		Transport: http.DefaultTransport,
	})
	sess, err := session.NewSession(awsConf)
	if err != nil {
		return nil, errors.Wrapf(err,
			"failed to create session for dynamodb (%s,%s)",
			config.Config.DynamoDBRegion,
			config.Config.DynamoDBEndpoint,
		)
	}
	return &DynamoDB{
		svc: godynamodb.New(sess),
	}, nil
}

// Client returns the DynamoDB client.
func (d *DynamoDB) Client() godynamodbiface.DynamoDBAPI {
	return d.svc
}

// Ping pings DynamoDB endpoint.
func (d *DynamoDB) Ping() error {
	ctx, cancel := context.WithTimeout(context.TODO(), pingTimeout)
	defer cancel()
	var opt request.Option = func(r *request.Request) {}
	_, err := d.svc.DescribeTableWithContext(ctx, &godynamodb.DescribeTableInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
	}, opt)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == request.CanceledErrorCode {
			// If the SDK can determine the request or retry delay was canceled
			// by a context the CanceledErrorCode error code will be returned.
			return errors.Wrap(err, "failed to ping dynamodb due to timeout")
		}
		return errors.Wrapf(err, "failed to ping dynamodb")
	}
	return nil
}

// CreateTableParam is parameter set of CreateTable.
type CreateTableParam struct {
	Name string
	RCU  int64 // ReadCapacityUnits
	WCU  int64 // WriteCapacityUnits
}

// CreateTable creates a dynamodb table to store time series data.
// Skip creating table if the table already exists.
func (d *DynamoDB) CreateTable(param *CreateTableParam) error {
	_, err := d.svc.CreateTable(&godynamodb.CreateTableInput{
		TableName: aws.String(param.Name),
		AttributeDefinitions: []*godynamodb.AttributeDefinition{
			{
				AttributeName: aws.String("Name"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
			{
				AttributeName: aws.String("Timestamp"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
		},
		KeySchema: []*godynamodb.KeySchemaElement{
			{
				AttributeName: aws.String("Name"),
				KeyType:       aws.String("HASH"),
			},
			{
				AttributeName: aws.String("Timestamp"),
				KeyType:       aws.String("RANGE"),
			},
		},
		ProvisionedThroughput: &godynamodb.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(param.RCU),
			WriteCapacityUnits: aws.Int64(param.WCU),
		},
		// TODO StreamSpecification to export to s3
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceInUseException" {
				// Skip if the table already exists
				log.Printf("Skip creating DynamoDB table because %s already exists\n", param.Name)
				return nil
			}
		}
		return errors.Wrapf(err, "failed to create dynamodb table (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	log.Printf("Creating DynamoDB table (name:%s, rcu:%d, wcu:%d) ...\n",
		param.Name, param.RCU, param.WCU)

	err = d.svc.WaitUntilTableExists(&godynamodb.DescribeTableInput{
		TableName: aws.String(param.Name),
	})
	if err != nil {
		return errors.Wrapf(err, "failed to wait until table exists (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	if config.Config.DynamoDBTTL {
		_, err = d.svc.UpdateTimeToLive(&godynamodb.UpdateTimeToLiveInput{
			TableName: aws.String(param.Name),
			TimeToLiveSpecification: &godynamodb.TimeToLiveSpecification{
				AttributeName: aws.String("TTL"),
				Enabled:       aws.Bool(true),
			},
		})
		if err != nil {
			return errors.Wrapf(err, "failed to set TTL to (%s,%d,%d)",
				param.Name, param.RCU, param.WCU)
		}
	}
	return nil
}

// Fetch fetches datapoints by name from start until end.
func (d *DynamoDB) Fetch(name string, start, end time.Time) (model.SeriesMap, error) {
	slots := selectTimeSlots(start, end)
	nameGroups := util.GroupNames(util.SplitName(name), dynamodbBatchLimit)
	numQueries := len(slots) * len(nameGroups)

	type result struct {
		value model.SeriesMap
		err   error
	}
	c := make(chan *result, numQueries)
	for _, slot := range slots {
		for _, names := range nameGroups {
			q := &query{
				names: names,
				start: start,
				end:   end,
				slot:  slot,
			}
			go func(q *query) {
				sm, err := d.batchGet(q)
				c <- &result{value: sm, err: err}
			}(q)
		}
	}
	sm := make(model.SeriesMap, len(nameGroups))
	for i := 0; i < numQueries; i++ {
		ret := <-c
		if ret.err != nil {
			return nil, ret.err
		}
		sm.MergePointsToMap(ret.value)
	}
	return sm, nil
}

func batchGetResultToMap(resp *godynamodb.BatchGetItemOutput, q *query) model.SeriesMap {
	sm := make(model.SeriesMap, len(resp.Responses))
	for _, xs := range resp.Responses {
		for _, x := range xs {
			name := (*x["Name"].S)
			points := make(model.DataPoints, 0, len(x["Values"].BS))
			for _, y := range x["Values"].BS {
				t := int64(binary.BigEndian.Uint64(y[0:8]))
				v := math.Float64frombits(binary.BigEndian.Uint64(y[8:]))
				// Trim datapoints out of [start, end]
				if t < q.start.Unix() || q.end.Unix() < t {
					continue
				}
				points = append(points, model.NewDataPoint(t, v))
			}
			sm[name] = model.NewSeriesPoint(name, points, q.slot.step)
		}
	}
	return sm
}

func (d *DynamoDB) batchGet(q *query) (model.SeriesMap, error) {
	var keys []map[string]*godynamodb.AttributeValue
	for _, name := range q.names {
		keys = append(keys, map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", q.slot.itemEpoch, q.slot.step))},
		})
	}
	items := make(map[string]*godynamodb.KeysAndAttributes)
	items[config.Config.DynamoDBTableName] = &godynamodb.KeysAndAttributes{Keys: keys}
	params := &godynamodb.BatchGetItemInput{
		RequestItems:           items,
		ReturnConsumedCapacity: aws.String("NONE"),
	}
	resp, err := d.svc.BatchGetItem(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				// Don't handle ResourceNotFoundException as error
				// bacause diamondb web return length 0 series as 200.
				return model.SeriesMap{}, nil
			}
		}
		return nil, errors.Wrapf(err,
			"failed to call dynamodb API batchGetItem (%s,%d,%d)",
			config.Config.DynamoDBTableName, q.slot.itemEpoch, q.slot.step,
		)
	}
	return batchGetResultToMap(resp, q), nil
}

// Put writes the datapoints into DynamoDB. It creates item
// if item doesn't exist and updates item if it exists.
func (d *DynamoDB) Put(name, slot, history string, itemEpoch int64, tv map[int64]float64) error {
	stepDuration, err := timeparser.ParseTimeOffset(slot)
	if err != nil {
		return err
	}
	historyDuration, err := timeparser.ParseTimeOffset(history)
	if err != nil {
		return err
	}
	ttl := itemEpoch + int64(historyDuration.Seconds())

	vals := make([][]byte, 0, len(tv))
	for timestamp, value := range tv {
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, timestamp)
		binary.Write(buf, binary.BigEndian, math.Float64bits(value))
		vals = append(vals, buf.Bytes())
	}

	params := &godynamodb.UpdateItemInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
		Key: map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", itemEpoch, int64(stepDuration.Seconds())))},
		},
		UpdateExpression: aws.String(`
			SET #ttl = :new_ttl
			ADD #values_set :new_values
		`),
		ExpressionAttributeNames: map[string]*string{
			"#ttl":        aws.String("TTL"),
			"#values_set": aws.String("Values"),
		},
		ExpressionAttributeValues: map[string]*godynamodb.AttributeValue{
			":new_ttl":    {N: aws.String(fmt.Sprintf("%d", ttl))},
			":new_values": {BS: vals},
		},
		ReturnValues: aws.String("NONE"),
	}
	if _, err := d.svc.UpdateItem(params); err != nil {
		return errors.Wrapf(err, "failed to call dynamodb API putItem (%s,%s,%d)",
			config.Config.DynamoDBTableName, name, itemEpoch)
	}
	return nil
}

func selectTimeSlots(startTime, endTime time.Time) []*timeSlot {
	var (
		step          int
		itemEpochStep int
	)
	diff := endTime.Sub(startTime)
	switch {
	case oneYear <= diff:
		itemEpochStep = oneYearSeconds
		step = 60 * 60 * 24
	case oneWeek <= diff:
		itemEpochStep = oneWeekSeconds
		step = 60 * 60
	case oneDay <= diff:
		itemEpochStep = oneDaySeconds
		step = 5 * 60
	default:
		itemEpochStep = 60 * 60
		step = 60
	}

	var slots []*timeSlot
	startItemEpoch := startTime.Unix() - startTime.Unix()%int64(itemEpochStep)
	endItemEpoch := endTime.Unix()
	for epoch := startItemEpoch; epoch < endItemEpoch; epoch += int64(itemEpochStep) {
		slots = append(slots, &timeSlot{itemEpoch: epoch, step: step})
	}

	return slots
}
<commit_msg>Revert "Fix dynamodbBatchLimit constant"<commit_after>package dynamodb

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"log"
	"math"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	godynamodb "github.com/aws/aws-sdk-go/service/dynamodb"
	godynamodbiface "github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/pkg/config"
	"github.com/yuuki/diamondb/pkg/model"
	"github.com/yuuki/diamondb/pkg/storage/util"
	"github.com/yuuki/diamondb/pkg/timeparser"
)

//go:generate mockgen -source ../../../vendor/github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface/interface.go -destination dynamodb_mock.go -package dynamodb

// ReadWriter defines the interface for DynamoDB reader and writer.
type ReadWriter interface {
	Ping() error
	Client() godynamodbiface.DynamoDBAPI
	CreateTable(*CreateTableParam) error
	Fetch(string, time.Time, time.Time) (model.SeriesMap, error)
	batchGet(q *query) (model.SeriesMap, error)
	Put(string, string, string, int64, map[int64]float64) error
}

// DynamoDB provides a dynamodb client.
type DynamoDB struct {
	svc godynamodbiface.DynamoDBAPI
}

type timeSlot struct {
	itemEpoch int64
	step      int
}

type query struct {
	names []string
	start time.Time
	end   time.Time
	slot  *timeSlot
	// context
}

const (
	pingTimeout = time.Duration(5) * time.Second

	oneYear time.Duration = time.Duration(24*360) * time.Hour
	oneWeek time.Duration = time.Duration(24*7) * time.Hour
	oneDay  time.Duration = time.Duration(24*1) * time.Hour
)

var (
	dynamodbBatchLimit = 100

	oneYearSeconds = int(oneYear.Seconds())
	oneWeekSeconds = int(oneWeek.Seconds())
	oneDaySeconds  = int(oneDay.Seconds())
)

var _ ReadWriter = &DynamoDB{}

// New creates a new DynamoDB.
func New() (*DynamoDB, error) {
	awsConf := aws.NewConfig().WithRegion(config.Config.DynamoDBRegion)
	if config.Config.DynamoDBEndpoint != "" {
		// For dynamodb-local configuration
		awsConf.WithEndpoint(config.Config.DynamoDBEndpoint)
		awsConf.WithCredentials(credentials.NewStaticCredentials("dummy", "dummy", "dummy"))
	}
	awsConf.WithHTTPClient(&http.Client{
		Timeout:   10 * time.Second,
		Transport: http.DefaultTransport,
	})
	sess, err := session.NewSession(awsConf)
	if err != nil {
		return nil, errors.Wrapf(err,
			"failed to create session for dynamodb (%s,%s)",
			config.Config.DynamoDBRegion,
			config.Config.DynamoDBEndpoint,
		)
	}
	return &DynamoDB{
		svc: godynamodb.New(sess),
	}, nil
}

// Client returns the DynamoDB client.
func (d *DynamoDB) Client() godynamodbiface.DynamoDBAPI {
	return d.svc
}

// Ping pings DynamoDB endpoint.
func (d *DynamoDB) Ping() error {
	ctx, cancel := context.WithTimeout(context.TODO(), pingTimeout)
	defer cancel()
	var opt request.Option = func(r *request.Request) {}
	_, err := d.svc.DescribeTableWithContext(ctx, &godynamodb.DescribeTableInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
	}, opt)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == request.CanceledErrorCode {
			// If the SDK can determine the request or retry delay was canceled
			// by a context the CanceledErrorCode error code will be returned.
			return errors.Wrap(err, "failed to ping dynamodb due to timeout")
		}
		return errors.Wrapf(err, "failed to ping dynamodb")
	}
	return nil
}

// CreateTableParam is parameter set of CreateTable.
type CreateTableParam struct {
	Name string
	RCU  int64 // ReadCapacityUnits
	WCU  int64 // WriteCapacityUnits
}

// CreateTable creates a dynamodb table to store time series data.
// Skip creating table if the table already exists.
func (d *DynamoDB) CreateTable(param *CreateTableParam) error {
	_, err := d.svc.CreateTable(&godynamodb.CreateTableInput{
		TableName: aws.String(param.Name),
		AttributeDefinitions: []*godynamodb.AttributeDefinition{
			{
				AttributeName: aws.String("Name"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
			{
				AttributeName: aws.String("Timestamp"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
		},
		KeySchema: []*godynamodb.KeySchemaElement{
			{
				AttributeName: aws.String("Name"),
				KeyType:       aws.String("HASH"),
			},
			{
				AttributeName: aws.String("Timestamp"),
				KeyType:       aws.String("RANGE"),
			},
		},
		ProvisionedThroughput: &godynamodb.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(param.RCU),
			WriteCapacityUnits: aws.Int64(param.WCU),
		},
		// TODO StreamSpecification to export to s3
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceInUseException" {
				// Skip if the table already exists
				log.Printf("Skip creating DynamoDB table because %s already exists\n", param.Name)
				return nil
			}
		}
		return errors.Wrapf(err, "failed to create dynamodb table (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	log.Printf("Creating DynamoDB table (name:%s, rcu:%d, wcu:%d) ...\n",
		param.Name, param.RCU, param.WCU)

	err = d.svc.WaitUntilTableExists(&godynamodb.DescribeTableInput{
		TableName: aws.String(param.Name),
	})
	if err != nil {
		return errors.Wrapf(err, "failed to wait until table exists (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	if config.Config.DynamoDBTTL {
		_, err = d.svc.UpdateTimeToLive(&godynamodb.UpdateTimeToLiveInput{
			TableName: aws.String(param.Name),
			TimeToLiveSpecification: &godynamodb.TimeToLiveSpecification{
				AttributeName: aws.String("TTL"),
				Enabled:       aws.Bool(true),
			},
		})
		if err != nil {
			return errors.Wrapf(err, "failed to set TTL to (%s,%d,%d)",
				param.Name, param.RCU, param.WCU)
		}
	}
	return nil
}

// Fetch fetches datapoints by name from start until end.
func (d *DynamoDB) Fetch(name string, start, end time.Time) (model.SeriesMap, error) {
	slots := selectTimeSlots(start, end)
	nameGroups := util.GroupNames(util.SplitName(name), dynamodbBatchLimit)
	numQueries := len(slots) * len(nameGroups)

	type result struct {
		value model.SeriesMap
		err   error
	}
	c := make(chan *result, numQueries)
	for _, slot := range slots {
		for _, names := range nameGroups {
			q := &query{
				names: names,
				start: start,
				end:   end,
				slot:  slot,
			}
			go func(q *query) {
				sm, err := d.batchGet(q)
				c <- &result{value: sm, err: err}
			}(q)
		}
	}
	sm := make(model.SeriesMap, len(nameGroups))
	for i := 0; i < numQueries; i++ {
		ret := <-c
		if ret.err != nil {
			return nil, ret.err
		}
		sm.MergePointsToMap(ret.value)
	}
	return sm, nil
}

func batchGetResultToMap(resp *godynamodb.BatchGetItemOutput, q *query) model.SeriesMap {
	sm := make(model.SeriesMap, len(resp.Responses))
	for _, xs := range resp.Responses {
		for _, x := range xs {
			name := (*x["Name"].S)
			points := make(model.DataPoints, 0, len(x["Values"].BS))
			for _, y := range x["Values"].BS {
				t := int64(binary.BigEndian.Uint64(y[0:8]))
				v := math.Float64frombits(binary.BigEndian.Uint64(y[8:]))
				// Trim datapoints out of [start, end]
				if t < q.start.Unix() || q.end.Unix() < t {
					continue
				}
				points = append(points, model.NewDataPoint(t, v))
			}
			sm[name] = model.NewSeriesPoint(name, points, q.slot.step)
		}
	}
	return sm
}

func (d *DynamoDB) batchGet(q *query) (model.SeriesMap, error) {
	var keys []map[string]*godynamodb.AttributeValue
	for _, name := range q.names {
		keys = append(keys, map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", q.slot.itemEpoch, q.slot.step))},
		})
	}
	items := make(map[string]*godynamodb.KeysAndAttributes)
	items[config.Config.DynamoDBTableName] = &godynamodb.KeysAndAttributes{Keys: keys}
	params := &godynamodb.BatchGetItemInput{
		RequestItems:           items,
		ReturnConsumedCapacity: aws.String("NONE"),
	}
	resp, err := d.svc.BatchGetItem(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				// Don't handle ResourceNotFoundException as error
				// bacause diamondb web return length 0 series as 200.
				return model.SeriesMap{}, nil
			}
		}
		return nil, errors.Wrapf(err,
			"failed to call dynamodb API batchGetItem (%s,%d,%d)",
			config.Config.DynamoDBTableName, q.slot.itemEpoch, q.slot.step,
		)
	}
	return batchGetResultToMap(resp, q), nil
}

// Put writes the datapoints into DynamoDB. It creates item
// if item doesn't exist and updates item if it exists.
func (d *DynamoDB) Put(name, slot, history string, itemEpoch int64, tv map[int64]float64) error {
	stepDuration, err := timeparser.ParseTimeOffset(slot)
	if err != nil {
		return err
	}
	historyDuration, err := timeparser.ParseTimeOffset(history)
	if err != nil {
		return err
	}
	ttl := itemEpoch + int64(historyDuration.Seconds())

	vals := make([][]byte, 0, len(tv))
	for timestamp, value := range tv {
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, timestamp)
		binary.Write(buf, binary.BigEndian, math.Float64bits(value))
		vals = append(vals, buf.Bytes())
	}

	params := &godynamodb.UpdateItemInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
		Key: map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", itemEpoch, int64(stepDuration.Seconds())))},
		},
		UpdateExpression: aws.String(`
			SET #ttl = :new_ttl
			ADD #values_set :new_values
		`),
		ExpressionAttributeNames: map[string]*string{
			"#ttl":        aws.String("TTL"),
			"#values_set": aws.String("Values"),
		},
		ExpressionAttributeValues: map[string]*godynamodb.AttributeValue{
			":new_ttl":    {N: aws.String(fmt.Sprintf("%d", ttl))},
			":new_values": {BS: vals},
		},
		ReturnValues: aws.String("NONE"),
	}
	if _, err := d.svc.UpdateItem(params); err != nil {
		return errors.Wrapf(err, "failed to call dynamodb API putItem (%s,%s,%d)",
			config.Config.DynamoDBTableName, name, itemEpoch)
	}
	return nil
}

func selectTimeSlots(startTime, endTime time.Time) []*timeSlot {
	var (
		step          int
		itemEpochStep int
	)
	diff := endTime.Sub(startTime)
	switch {
	case oneYear <= diff:
		itemEpochStep = oneYearSeconds
		step = 60 * 60 * 24
	case oneWeek <= diff:
		itemEpochStep = oneWeekSeconds
		step = 60 * 60
	case oneDay <= diff:
		itemEpochStep = oneDaySeconds
		step = 5 * 60
	default:
		itemEpochStep = 60 * 60
		step = 60
	}

	var slots []*timeSlot
	startItemEpoch := startTime.Unix() - startTime.Unix()%int64(itemEpochStep)
	endItemEpoch := endTime.Unix()
	for epoch := startItemEpoch; epoch < endItemEpoch; epoch += int64(itemEpochStep) {
		slots = append(slots, &timeSlot{itemEpoch: epoch, step: step})
	}

	return slots
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"launchpad.net/gnuflag"
	"launchpad.net/loggo"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/filestorage"
	"launchpad.net/juju-core/environs/sync"
	"launchpad.net/juju-core/version"
)

var syncTools = sync.SyncTools

// SyncToolsCommand copies all the tools from the us-east-1 bucket to the local
// bucket.
type SyncToolsCommand struct {
	cmd.EnvCommandBase
	allVersions  bool
	versionStr   string
	majorVersion int
	minorVersion int
	dryRun       bool
	dev          bool
	source       string
	destination  string
}

var _ cmd.Command = (*SyncToolsCommand)(nil)

func (c *SyncToolsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "sync-tools",
		Purpose: "copy tools from the official bucket into a local environment",
		Doc: `
This copies the Juju tools tarball from the official bucket into
your environment. This is generally done when you want Juju to be able
to run without having to access Amazon. Alternatively you can specify
a local directory as source.

Sometimes this is because the environment does not have public access,
and sometimes you just want to avoid having to access data outside of
the local cloud.
`,
	}
}

func (c *SyncToolsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.BoolVar(&c.allVersions, "all", false, "copy all versions, not just the latest")
	f.StringVar(&c.versionStr, "version", "", "copy a specific major[.minor] version")
	f.BoolVar(&c.dryRun, "dry-run", false, "don't copy, just print what would be copied")
	f.BoolVar(&c.dev, "dev", false, "consider development versions as well as released ones")
	f.StringVar(&c.source, "source", "", "chose a location on the file system as source")
	f.StringVar(&c.destination, "destination", "", "chose a location on the file system as destination")

	// BUG(lp:1163164)  jam 2013-04-2 we would like to add a "source"
	// location, rather than only copying from us-east-1
}

func (c *SyncToolsCommand) Init(args []string) error {
	if c.versionStr != "" {
		var err error
		if c.majorVersion, c.minorVersion, err = version.ParseMajorMinor(c.versionStr); err != nil {
			return err
		}
	}
	return cmd.CheckEmpty(args)
}

func (c *SyncToolsCommand) Run(ctx *cmd.Context) error {
	// Register writer for output on screen.
	loggo.RegisterWriter("synctools", cmd.NewCommandLogWriter("juju.environs.sync", ctx.Stdout, ctx.Stderr), loggo.INFO)
	defer loggo.RemoveWriter("synctools")
	environ, err := environs.PrepareFromName(c.EnvName)
	if err != nil {
		return err
	}

	target := environ.Storage()
	if c.destination != "" {
		target, err = filestorage.NewFileStorageWriter(c.destination)
		if err != nil {
			return err
		}
	}

	// Prepare syncing.
	sctx := &sync.SyncContext{
		Target:       target,
		AllVersions:  c.allVersions,
		MajorVersion: c.majorVersion,
		MinorVersion: c.minorVersion,
		DryRun:       c.dryRun,
		Dev:          c.dev,
		Source:       c.source,
	}
	return syncTools(sctx)
}
<commit_msg>fix a couple tiny typos<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"launchpad.net/gnuflag"
	"launchpad.net/loggo"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/filestorage"
	"launchpad.net/juju-core/environs/sync"
	"launchpad.net/juju-core/version"
)

var syncTools = sync.SyncTools

// SyncToolsCommand copies all the tools from the us-east-1 bucket to the local
// bucket.
type SyncToolsCommand struct {
	cmd.EnvCommandBase
	allVersions  bool
	versionStr   string
	majorVersion int
	minorVersion int
	dryRun       bool
	dev          bool
	source       string
	destination  string
}

var _ cmd.Command = (*SyncToolsCommand)(nil)

func (c *SyncToolsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "sync-tools",
		Purpose: "copy tools from the official bucket into a local environment",
		Doc: `
This copies the Juju tools tarball from the official bucket into
your environment. This is generally done when you want Juju to be able
to run without having to access Amazon. Alternatively you can specify
a local directory as source.

Sometimes this is because the environment does not have public access,
and sometimes you just want to avoid having to access data outside of
the local cloud.
`,
	}
}

func (c *SyncToolsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.BoolVar(&c.allVersions, "all", false, "copy all versions, not just the latest")
	f.StringVar(&c.versionStr, "version", "", "copy a specific major[.minor] version")
	f.BoolVar(&c.dryRun, "dry-run", false, "don't copy, just print what would be copied")
	f.BoolVar(&c.dev, "dev", false, "consider development versions as well as released ones")
	f.StringVar(&c.source, "source", "", "choose a location on the file system as source")
	f.StringVar(&c.destination, "destination", "", "choose a location on the file system as destination")

	// BUG(lp:1163164)  jam 2013-04-2 we would like to add a "source"
	// location, rather than only copying from us-east-1
}

func (c *SyncToolsCommand) Init(args []string) error {
	if c.versionStr != "" {
		var err error
		if c.majorVersion, c.minorVersion, err = version.ParseMajorMinor(c.versionStr); err != nil {
			return err
		}
	}
	return cmd.CheckEmpty(args)
}

func (c *SyncToolsCommand) Run(ctx *cmd.Context) error {
	// Register writer for output on screen.
	loggo.RegisterWriter("synctools", cmd.NewCommandLogWriter("juju.environs.sync", ctx.Stdout, ctx.Stderr), loggo.INFO)
	defer loggo.RemoveWriter("synctools")
	environ, err := environs.PrepareFromName(c.EnvName)
	if err != nil {
		return err
	}

	target := environ.Storage()
	if c.destination != "" {
		target, err = filestorage.NewFileStorageWriter(c.destination)
		if err != nil {
			return err
		}
	}

	// Prepare syncing.
	sctx := &sync.SyncContext{
		Target:       target,
		AllVersions:  c.allVersions,
		MajorVersion: c.majorVersion,
		MinorVersion: c.minorVersion,
		DryRun:       c.dryRun,
		Dev:          c.dev,
		Source:       c.source,
	}
	return syncTools(sctx)
}
<|endoftext|>


<commit_before>package rulesets

import (
	"encoding/json"
	"time"

	"github.com/ion-channel/ionic/rules"
)

type RuleSet struct {
	ID          string       `json:"id"`
	TeamID      string       `json:"team_id"`
	Name        string       `json:"name"`
	Description string       `json:"description"`
	RuleIDs     []string     `json:"rule_ids"`
	CreatedAt   time.Time    `json:"created_at"`
	UpdatedAt   time.Time    `json:"updated_at"`
	Rules       []rules.Rule `json:"rules"`
}

type AppliedRulesetSummary struct {
	ProjectID             string `json:"project_id"`
	TeamID                string `json:"team_id"`
	AnalysisID            string `json:"analysis_id"`
	RuleEvaluationSummary struct {
		Summary     string `json:"summary"`
		Ruleresults []struct {
			ID          string          `json:"id"`
			AnalysisID  string          `json:"analysis_id"`
			TeamID      string          `json:"team_id"`
			ProjectID   string          `json:"project_id"`
			Description string          `json:"description"`
			Name        string          `json:"name"`
			Summary     string          `json:"summary"`
			CreatedAt   time.Time       `json:"created_at"`
			UpdatedAt   time.Time       `json:"updated_at"`
			Results     json.RawMessage `json:"results"`
			Duration    float64         `json:"duration"`
			Passed      bool            `json:"passed"`
			Risk        string          `json:"risk"`
			Type        string          `json:"type"`
		} `json:"ruleresults"`
	} `json:"rule_evaluation_summary"`
	RuleEvalCreatedAt time.Time `json:"rule_eval_created_at"`
	CreatedAt         time.Time `json:"created_at"`
	UpdatedAt         time.Time `json:"updated_at"`
}
<commit_msg>adding ruleset name to ruleset summary<commit_after>package rulesets

import (
	"encoding/json"
	"time"

	"github.com/ion-channel/ionic/rules"
)

type RuleSet struct {
	ID          string       `json:"id"`
	TeamID      string       `json:"team_id"`
	Name        string       `json:"name"`
	Description string       `json:"description"`
	RuleIDs     []string     `json:"rule_ids"`
	CreatedAt   time.Time    `json:"created_at"`
	UpdatedAt   time.Time    `json:"updated_at"`
	Rules       []rules.Rule `json:"rules"`
}

type AppliedRulesetSummary struct {
	ProjectID             string `json:"project_id"`
	TeamID                string `json:"team_id"`
	AnalysisID            string `json:"analysis_id"`
	RulesetName           string `json:"ruleset_name"`
	RuleEvaluationSummary struct {
		Summary     string `json:"summary"`
		Ruleresults []struct {
			ID          string          `json:"id"`
			AnalysisID  string          `json:"analysis_id"`
			TeamID      string          `json:"team_id"`
			ProjectID   string          `json:"project_id"`
			Description string          `json:"description"`
			Name        string          `json:"name"`
			Summary     string          `json:"summary"`
			CreatedAt   time.Time       `json:"created_at"`
			UpdatedAt   time.Time       `json:"updated_at"`
			Results     json.RawMessage `json:"results"`
			Duration    float64         `json:"duration"`
			Passed      bool            `json:"passed"`
			Risk        string          `json:"risk"`
			Type        string          `json:"type"`
		} `json:"ruleresults"`
	} `json:"rule_evaluation_summary"`
	RuleEvalCreatedAt time.Time `json:"rule_eval_created_at"`
	CreatedAt         time.Time `json:"created_at"`
	UpdatedAt         time.Time `json:"updated_at"`
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/x509"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/log/cmdlogger"
	"github.com/Symantec/keymaster/lib/client/config"
	"github.com/Symantec/keymaster/lib/client/twofa"
	"github.com/Symantec/keymaster/lib/client/twofa/u2f"
	"github.com/Symantec/keymaster/lib/client/util"
)

const DefaultKeysLocation = "/.ssh/"
const FilePrefix = "keymaster"

var (
	// Must be a global variable in the data segment so that the build
	// process can inject the version number on the fly when building the
	// binary. Use only from the Usage() function.
	Version = "No version provided"
)

var (
	configFilename = flag.String("config", filepath.Join(os.Getenv("HOME"), ".keymaster", "client_config.yml"), "The filename of the configuration")
	rootCAFilename = flag.String("rootCAFilename", "", "(optional) name for using non OS root CA to verify TLS connections")
	configHost     = flag.String("configHost", "", "Get a bootstrap config from this host")
	cliUsername    = flag.String("username", "", "username for keymaster")
	checkDevices   = flag.Bool("checkDevices", false, "CheckU2F devices in your system")
)

func maybeGetRootCas(logger log.Logger) *x509.CertPool {
	var rootCAs *x509.CertPool
	if len(*rootCAFilename) > 1 {
		caData, err := ioutil.ReadFile(*rootCAFilename)
		if err != nil {
			logger.Printf("Failed to read caFilename")
			logger.Fatal(err)
		}
		rootCAs = x509.NewCertPool()
		if !rootCAs.AppendCertsFromPEM(caData) {
			logger.Fatal("cannot append file data")
		}

	}
	return rootCAs
}

func getUserNameAndHomeDir(logger log.Logger) (userName, homeDir string) {
	usr, err := user.Current()
	if err != nil {
		logger.Printf("cannot get current user info")
		logger.Fatal(err)
	}
	userName = usr.Username

	homeDir, err = util.GetUserHomeDir(usr)
	if err != nil {
		logger.Fatal(err)
	}
	return
}

func loadConfigFile(rootCAs *x509.CertPool, logger log.Logger) (
	configContents config.AppConfigFile) {
	configPath, _ := filepath.Split(*configFilename)

	err := os.MkdirAll(configPath, 0755)
	if err != nil {
		logger.Fatal(err)
	}

	if len(*configHost) > 1 {
		err = config.GetConfigFromHost(*configFilename, *configHost, rootCAs, logger)
		if err != nil {
			logger.Fatal(err)
		}
	} else if len(defaultConfigHost) > 1 { // if there is a configHost AND there is NO config file, create one
		if _, err := os.Stat(*configFilename); os.IsNotExist(err) {
			err = config.GetConfigFromHost(
				*configFilename, defaultConfigHost, rootCAs, logger)
			if err != nil {
				logger.Fatal(err)
			}
		}
	}

	configContents, err = config.LoadVerifyConfigFile(*configFilename)
	if err != nil {
		logger.Fatal(err)
	}
	return
}

func setupCerts(
	rootCAs *x509.CertPool,
	userName,
	homeDir string,
	configContents config.AppConfigFile,
	logger log.DebugLogger) {
	// create dirs
	privateKeyPath := filepath.Join(homeDir, DefaultKeysLocation, FilePrefix)
	sshConfigPath, _ := filepath.Split(privateKeyPath)
	err := os.MkdirAll(sshConfigPath, 0700)
	if err != nil {
		logger.Fatal(err)
	}
	// get signer
	tempPrivateKeyPath := filepath.Join(homeDir, DefaultKeysLocation, "keymaster-temp")
	signer, tempPublicKeyPath, err := util.GenKeyPair(
		tempPrivateKeyPath, userName+"@keymaster", logger)
	if err != nil {
		logger.Fatal(err)
	}
	defer os.Remove(tempPrivateKeyPath)
	defer os.Remove(tempPublicKeyPath)
	// Get user creds
	password, err := util.GetUserCreds(userName)
	if err != nil {
		logger.Fatal(err)
	}

	// Get the certs
	sshCert, x509Cert, kubernetesCert, err := twofa.GetCertFromTargetUrls(
		signer,
		userName,
		password,
		strings.Split(configContents.Base.Gen_Cert_URLS, ","),
		rootCAs,
		false,
		logger)
	if err != nil {
		logger.Fatal(err)
	}
	if sshCert == nil || x509Cert == nil {
		err := errors.New("Could not get cert from any url")
		logger.Fatal(err)
	}
	logger.Debugf(0, "Got Certs from server")
	//..
	if _, ok := os.LookupEnv("SSH_AUTH_SOCK"); ok {
		// TODO(rgooch): Parse certificate to get actual lifetime.
		cmd := exec.Command("ssh-add", "-d", privateKeyPath)
		cmd.Run()
	}

	//rename files to expected paths
	err = os.Rename(tempPrivateKeyPath, privateKeyPath)
	if err != nil {
		err := errors.New("Could not rename private Key")
		logger.Fatal(err)
	}

	err = os.Rename(tempPublicKeyPath, privateKeyPath+".pub")
	if err != nil {
		err := errors.New("Could not rename public Key")
		logger.Fatal(err)
	}

	// now we write the cert file...
	sshCertPath := privateKeyPath + "-cert.pub"
	err = ioutil.WriteFile(sshCertPath, sshCert, 0644)
	if err != nil {
		err := errors.New("Could not write ssh cert")
		logger.Fatal(err)
	}
	x509CertPath := privateKeyPath + "-x509Cert.pem"
	err = ioutil.WriteFile(x509CertPath, x509Cert, 0644)
	if err != nil {
		err := errors.New("Could not write ssh cert")
		logger.Fatal(err)
	}
	if kubernetesCert != nil {
		kubernetesCertPath := privateKeyPath + "-kubernetesCert.pem"
		err = ioutil.WriteFile(kubernetesCertPath, kubernetesCert, 0644)
		if err != nil {
			err := errors.New("Could not write ssh cert")
			logger.Fatal(err)
		}
	}

	logger.Printf("Success")
	if _, ok := os.LookupEnv("SSH_AUTH_SOCK"); ok {
		// TODO(rgooch): Parse certificate to get actual lifetime.
		lifetime := fmt.Sprintf("%ds", uint64((*twofa.Duration).Seconds()))
		cmd := exec.Command("ssh-add", "-t", lifetime, privateKeyPath)
		cmd.Run()
	}
}

func Usage() {
	fmt.Fprintf(
		os.Stderr, "Usage of %s (version %s):\n", os.Args[0], Version)
	flag.PrintDefaults()
}

func main() {
	flag.Usage = Usage
	flag.Parse()
	logger := cmdlogger.New()

	if *checkDevices {
		u2f.CheckU2FDevices(logger)
		return
	}

	rootCAs := maybeGetRootCas(logger)
	userName, homeDir := getUserNameAndHomeDir(logger)
	config := loadConfigFile(rootCAs, logger)

	// Adjust user name
	if len(config.Base.Username) > 0 {
		userName = config.Base.Username
	}
	// command line always wins over pref or config
	if *cliUsername != "" {
		userName = *cliUsername
	}
	setupCerts(rootCAs, userName, homeDir, config, logger)
}
<commit_msg>added base functionality tested on osx<commit_after>package main

import (
	"crypto/x509"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/log/cmdlogger"
	"github.com/Symantec/keymaster/lib/client/config"
	"github.com/Symantec/keymaster/lib/client/twofa"
	"github.com/Symantec/keymaster/lib/client/twofa/u2f"
	"github.com/Symantec/keymaster/lib/client/util"
)

const DefaultSSHKeysLocation = "/.ssh/"
const DefaultTLSKeysLocation = "/.ssl/"

const FilePrefix = "keymaster"

var (
	// Must be a global variable in the data segment so that the build
	// process can inject the version number on the fly when building the
	// binary. Use only from the Usage() function.
	Version = "No version provided"
)

var (
	configFilename = flag.String("config", filepath.Join(os.Getenv("HOME"), ".keymaster", "client_config.yml"), "The filename of the configuration")
	rootCAFilename = flag.String("rootCAFilename", "", "(optional) name for using non OS root CA to verify TLS connections")
	configHost     = flag.String("configHost", "", "Get a bootstrap config from this host")
	cliUsername    = flag.String("username", "", "username for keymaster")
	checkDevices   = flag.Bool("checkDevices", false, "CheckU2F devices in your system")
)

func maybeGetRootCas(logger log.Logger) *x509.CertPool {
	var rootCAs *x509.CertPool
	if len(*rootCAFilename) > 1 {
		caData, err := ioutil.ReadFile(*rootCAFilename)
		if err != nil {
			logger.Printf("Failed to read caFilename")
			logger.Fatal(err)
		}
		rootCAs = x509.NewCertPool()
		if !rootCAs.AppendCertsFromPEM(caData) {
			logger.Fatal("cannot append file data")
		}

	}
	return rootCAs
}

func getUserNameAndHomeDir(logger log.Logger) (userName, homeDir string) {
	usr, err := user.Current()
	if err != nil {
		logger.Printf("cannot get current user info")
		logger.Fatal(err)
	}
	userName = usr.Username

	homeDir, err = util.GetUserHomeDir(usr)
	if err != nil {
		logger.Fatal(err)
	}
	return
}

func loadConfigFile(rootCAs *x509.CertPool, logger log.Logger) (
	configContents config.AppConfigFile) {
	configPath, _ := filepath.Split(*configFilename)

	err := os.MkdirAll(configPath, 0755)
	if err != nil {
		logger.Fatal(err)
	}

	if len(*configHost) > 1 {
		err = config.GetConfigFromHost(*configFilename, *configHost, rootCAs, logger)
		if err != nil {
			logger.Fatal(err)
		}
	} else if len(defaultConfigHost) > 1 { // if there is a configHost AND there is NO config file, create one
		if _, err := os.Stat(*configFilename); os.IsNotExist(err) {
			err = config.GetConfigFromHost(
				*configFilename, defaultConfigHost, rootCAs, logger)
			if err != nil {
				logger.Fatal(err)
			}
		}
	}

	configContents, err = config.LoadVerifyConfigFile(*configFilename)
	if err != nil {
		logger.Fatal(err)
	}
	return
}

func setupCerts(
	rootCAs *x509.CertPool,
	userName,
	homeDir string,
	configContents config.AppConfigFile,
	logger log.DebugLogger) {
	// create dirs
	privateSSHKeyPath := filepath.Join(homeDir, DefaultSSHKeysLocation, FilePrefix)
	sshConfigPath, _ := filepath.Split(privateSSHKeyPath)
	err := os.MkdirAll(sshConfigPath, 0700)
	if err != nil {
		logger.Fatal(err)
	}
	privateTLSKeyPath := filepath.Join(homeDir, DefaultTLSKeysLocation, FilePrefix)
	tlsConfigPath, _ := filepath.Split(privateTLSKeyPath)
	err = os.MkdirAll(tlsConfigPath, 0700)
	if err != nil {
		logger.Fatal(err)
	}

	// get signer
	tempPrivateKeyPath := filepath.Join(homeDir, DefaultSSHKeysLocation, "keymaster-temp")
	signer, tempPublicKeyPath, err := util.GenKeyPair(
		tempPrivateKeyPath, userName+"@keymaster", logger)
	if err != nil {
		logger.Fatal(err)
	}
	defer os.Remove(tempPrivateKeyPath)
	defer os.Remove(tempPublicKeyPath)
	// Get user creds
	password, err := util.GetUserCreds(userName)
	if err != nil {
		logger.Fatal(err)
	}

	// Get the certs
	sshCert, x509Cert, kubernetesCert, err := twofa.GetCertFromTargetUrls(
		signer,
		userName,
		password,
		strings.Split(configContents.Base.Gen_Cert_URLS, ","),
		rootCAs,
		false,
		logger)
	if err != nil {
		logger.Fatal(err)
	}
	if sshCert == nil || x509Cert == nil {
		err := errors.New("Could not get cert from any url")
		logger.Fatal(err)
	}
	logger.Debugf(0, "Got Certs from server")
	//..
	if _, ok := os.LookupEnv("SSH_AUTH_SOCK"); ok {
		// TODO(rgooch): Parse certificate to get actual lifetime.
		cmd := exec.Command("ssh-add", "-d", privateSSHKeyPath)
		cmd.Run()
	}

	//rename files to expected paths
	err = os.Rename(tempPrivateKeyPath, privateSSHKeyPath)
	if err != nil {
		err := errors.New("Could not rename private Key")
		logger.Fatal(err)
	}

	err = os.Rename(tempPublicKeyPath, privateSSHKeyPath+".pub")
	if err != nil {
		err := errors.New("Could not rename public Key")
		logger.Fatal(err)
	}
	// Now handle the key in the tls directory
	tlsPrivateKeyName := filepath.Join(homeDir, DefaultTLSKeysLocation, "keymaster-key.pem")
	os.Remove(tlsPrivateKeyName)
	err = os.Symlink(privateSSHKeyPath, tlsPrivateKeyName)
	if err != nil {
		err := errors.New("Could not create new symlink for TLS key")
		logger.Fatal(err)
	}

	// now we write the cert file...
	sshCertPath := privateSSHKeyPath + "-cert.pub"
	err = ioutil.WriteFile(sshCertPath, sshCert, 0644)
	if err != nil {
		err := errors.New("Could not write ssh cert")
		logger.Fatal(err)
	}
	x509CertPath := privateTLSKeyPath + "-x509Cert.pem"
	err = ioutil.WriteFile(x509CertPath, x509Cert, 0644)
	if err != nil {
		err := errors.New("Could not write ssh cert")
		logger.Fatal(err)
	}
	if kubernetesCert != nil {
		kubernetesCertPath := privateTLSKeyPath + "-kubernetesCert.pem"
		err = ioutil.WriteFile(kubernetesCertPath, kubernetesCert, 0644)
		if err != nil {
			err := errors.New("Could not write ssh cert")
			logger.Fatal(err)
		}
	}

	logger.Printf("Success")
	if _, ok := os.LookupEnv("SSH_AUTH_SOCK"); ok {
		// TODO(rgooch): Parse certificate to get actual lifetime.
		lifetime := fmt.Sprintf("%ds", uint64((*twofa.Duration).Seconds()))
		cmd := exec.Command("ssh-add", "-t", lifetime, privateSSHKeyPath)
		cmd.Run()
	}
}

func Usage() {
	fmt.Fprintf(
		os.Stderr, "Usage of %s (version %s):\n", os.Args[0], Version)
	flag.PrintDefaults()
}

func main() {
	flag.Usage = Usage
	flag.Parse()
	logger := cmdlogger.New()

	if *checkDevices {
		u2f.CheckU2FDevices(logger)
		return
	}

	rootCAs := maybeGetRootCas(logger)
	userName, homeDir := getUserNameAndHomeDir(logger)
	config := loadConfigFile(rootCAs, logger)

	// Adjust user name
	if len(config.Base.Username) > 0 {
		userName = config.Base.Username
	}
	// command line always wins over pref or config
	if *cliUsername != "" {
		userName = *cliUsername
	}
	setupCerts(rootCAs, userName, homeDir, config, logger)
}
<|endoftext|>


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

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/naming"
	"v.io/v23/options"
	"v.io/v23/security/access"
	"v.io/v23/verror"
	"v.io/x/lib/cmdline"
	"v.io/x/lib/vlog"
)

var (
	flagLongGlob            bool
	flagInsecureResolve     bool
	flagInsecureResolveToMT bool
)

var cmdGlob = &cmdline.Command{
	Run:      runGlob,
	Name:     "glob",
	Short:    "Returns all matching entries from the namespace",
	Long:     "Returns all matching entries from the namespace.",
	ArgsName: "<pattern>",
	ArgsLong: `
<pattern> is a glob pattern that is matched against all the names below the
specified mount name.
`,
}

func runGlob(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("glob: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	pattern := args[0]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)

	c, err := ns.Glob(ctx, pattern)
	if err != nil {
		vlog.Infof("ns.Glob(%q) failed: %v", pattern, err)
		return err
	}
	if flagLongGlob {
		// Show all the information we received.
		for res := range c {
			switch v := res.(type) {
			case *naming.GlobReplyEntry:
				fmt.Fprint(cmd.Stdout(), v.Value.Name)
				for _, s := range v.Value.Servers {
					delta := s.Deadline.Time.Sub(time.Now())
					fmt.Fprintf(cmd.Stdout(), " %s (Expires in %d sec)", s.Server, int(delta.Seconds()))
				}
				fmt.Fprintln(cmd.Stdout())
			case *naming.GlobReplyError:
				fmt.Fprintf(cmd.Stderr(), "Error: %s: %v\n", v.Value.Name, v.Value.Error)
			}
		}
		return nil
	}
	// Show a sorted list of unique names, and any errors.
	resultSet := make(map[string]struct{})
	errors := []*naming.GlobError{}
	for res := range c {
		switch v := res.(type) {
		case *naming.GlobReplyEntry:
			if v.Value.Name != "" {
				resultSet[v.Value.Name] = struct{}{}
			}
		case *naming.GlobReplyError:
			errors = append(errors, &v.Value)
		}
	}
	results := []string{}
	for r := range resultSet {
		results = append(results, r)
	}
	sort.Strings(results)
	for _, result := range results {
		fmt.Fprintln(cmd.Stdout(), result)
	}
	for _, err := range errors {
		fmt.Fprintf(cmd.Stderr(), "Error: %s: %v\n", err.Name, err.Error)
	}
	return nil
}

var cmdMount = &cmdline.Command{
	Run:      runMount,
	Name:     "mount",
	Short:    "Adds a server to the namespace",
	Long:     "Adds server <server> to the namespace with name <name>.",
	ArgsName: "<name> <server> <ttl>",
	ArgsLong: `
<name> is the name to add to the namespace.
<server> is the object address of the server to add.
<ttl> is the TTL of the new entry. It is a decimal number followed by a unit
suffix (s, m, h). A value of 0s represents an infinite duration.
`,
}

func runMount(cmd *cmdline.Command, args []string) error {
	if expected, got := 3, len(args); expected != got {
		return cmd.UsageErrorf("mount: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	server := args[1]
	ttlArg := args[2]

	ttl, err := time.ParseDuration(ttlArg)
	if err != nil {
		return fmt.Errorf("TTL parse error: %v", err)
	}

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)
	if err = ns.Mount(ctx, name, server, ttl); err != nil {
		vlog.Infof("ns.Mount(%q, %q, %s) failed: %v", name, server, ttl, err)
		return err
	}
	fmt.Fprintln(cmd.Stdout(), "Server mounted successfully.")
	return nil
}

var cmdUnmount = &cmdline.Command{
	Run:      runUnmount,
	Name:     "unmount",
	Short:    "Removes a server from the namespace",
	Long:     "Removes server <server> with name <name> from the namespace.",
	ArgsName: "<name> <server>",
	ArgsLong: `
<name> is the name to remove from the namespace.
<server> is the object address of the server to remove.
`,
}

func runUnmount(cmd *cmdline.Command, args []string) error {
	if expected, got := 2, len(args); expected != got {
		return cmd.UsageErrorf("unmount: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	server := args[1]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)

	if err := ns.Unmount(ctx, name, server); err != nil {
		vlog.Infof("ns.Unmount(%q, %q) failed: %v", name, server, err)
		return err
	}
	fmt.Fprintln(cmd.Stdout(), "Server unmounted successfully.")
	return nil
}

var cmdResolve = &cmdline.Command{
	Run:      runResolve,
	Name:     "resolve",
	Short:    "Translates a object name to its object address(es)",
	Long:     "Translates a object name to its object address(es).",
	ArgsName: "<name>",
	ArgsLong: "<name> is the name to resolve.",
}

func runResolve(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("resolve: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)

	var opts []naming.NamespaceOpt
	if flagInsecureResolve {
		opts = append(opts, options.SkipServerEndpointAuthorization{})
	}
	me, err := ns.Resolve(ctx, name, opts...)
	if err != nil {
		vlog.Infof("ns.Resolve(%q) failed: %v", name, err)
		return err
	}
	for _, n := range me.Names() {
		fmt.Fprintln(cmd.Stdout(), n)
	}
	return nil
}

var cmdResolveToMT = &cmdline.Command{
	Run:      runResolveToMT,
	Name:     "resolvetomt",
	Short:    "Finds the address of the mounttable that holds an object name",
	Long:     "Finds the address of the mounttable that holds an object name.",
	ArgsName: "<name>",
	ArgsLong: "<name> is the name to resolve.",
}

func runResolveToMT(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("resolvetomt: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)
	var opts []naming.NamespaceOpt
	if flagInsecureResolveToMT {
		opts = append(opts, options.SkipServerEndpointAuthorization{})
	}
	e, err := ns.ResolveToMountTable(ctx, name, opts...)
	if err != nil {
		vlog.Infof("ns.ResolveToMountTable(%q) failed: %v", name, err)
		return err
	}
	for _, s := range e.Servers {
		fmt.Fprintln(cmd.Stdout(), naming.JoinAddressName(s.Server, e.Name))
	}
	return nil
}

var cmdPermissions = &cmdline.Command{
	Name:  "permissions",
	Short: "Manipulates permissions on an entry in the namespace",
	Long: `
Commands to get and set the permissions on a name - controlling the blessing
names required to resolve the name.

The permissions are provided as an JSON-encoded version of the Permissions type
defined in v.io/v23/security/access/types.vdl.
`,
	Children: []*cmdline.Command{cmdPermissionsGet, cmdPermissionsSet},
}

var cmdPermissionsSet = &cmdline.Command{
	Run:   runPermissionsSet,
	Name:  "set",
	Short: "Sets permissions on a mount name",
	Long: `
Set replaces the permissions controlling usage of a mount name.
`,
	ArgsName: "<name> <permissions>",
	ArgsLong: `
<name> is the name on which permissions are to be set.

<permissions> is the path to a file containing a JSON-encoded Permissions
object (defined in v.io/v23/security/access/types.vdl), or "-" for STDIN.
`,
}

func runPermissionsSet(cmd *cmdline.Command, args []string) error {
	if expected, got := 2, len(args); expected != got {
		return cmd.UsageErrorf("set: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	var perms access.Permissions
	file := os.Stdin
	if args[1] != "-" {
		var err error
		if file, err = os.Open(args[1]); err != nil {
			return err
		}
		defer file.Close()
	}
	if err := json.NewDecoder(file).Decode(&perms); err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()
	ns := v23.GetNamespace(ctx)
	for {
		_, etag, err := ns.GetPermissions(ctx, name)
		if err != nil {
			return err
		}
		if err = ns.SetPermissions(ctx, name, perms, etag); verror.ErrorID(err) == verror.ErrBadVersion.ID {
			vlog.Infof("SetPermissions(%q, %q) failed: %v, retrying...", name, etag, err)
			continue
		}
		return err
	}
}

var cmdPermissionsGet = &cmdline.Command{
	Run:      runPermissionsGet,
	Name:     "get",
	Short:    "Gets permissions on a mount name",
	ArgsName: "<name>",
	ArgsLong: `
<name> is a name in the namespace.
`,
	Long: `
Get retrieves the permissions on the usage of a name.

The output is a JSON-encoded Permissions object (defined in
v.io/v23/security/access/types.vdl).
`,
}

func runPermissionsGet(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("get: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	perms, _, err := v23.GetNamespace(ctx).GetPermissions(ctx, name)
	if err != nil {
		return err
	}
	return json.NewEncoder(cmd.Stdout()).Encode(perms)
}

func root() *cmdline.Command {
	cmdGlob.Flags.BoolVar(&flagLongGlob, "l", false, "Long listing format.")
	cmdResolve.Flags.BoolVar(&flagInsecureResolve, "insecure", false, "Insecure mode: May return results from untrusted servers and invoke Resolve on untrusted mounttables")
	cmdResolveToMT.Flags.BoolVar(&flagInsecureResolveToMT, "insecure", false, "Insecure mode: May return results from untrusted servers and invoke Resolve on untrusted mounttables")
	return &cmdline.Command{
		Name:  "namespace",
		Short: "resolves and manages names in the Vanadium namespace",
		Long: `
Command namespace resolves and manages names in the Vanadium namespace.

The namespace roots are set from the command line via --v23.namespace.root
command line option or from environment variables that have a name starting
with V23_NAMESPACE, e.g.  V23_NAMESPACE, V23_NAMESPACE_2, V23_NAMESPACE_GOOGLE,
etc.  The command line options override the environment.
`,
		Children: []*cmdline.Command{cmdGlob, cmdMount, cmdUnmount, cmdResolve, cmdResolveToMT, cmdPermissions},
	}
}
<commit_msg>cmd/namespace: Allow setting permissions on name that does not exist yet.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/naming"
	"v.io/v23/options"
	"v.io/v23/security/access"
	"v.io/v23/verror"
	"v.io/x/lib/cmdline"
	"v.io/x/lib/vlog"
)

var (
	flagLongGlob            bool
	flagInsecureResolve     bool
	flagInsecureResolveToMT bool
)

var cmdGlob = &cmdline.Command{
	Run:      runGlob,
	Name:     "glob",
	Short:    "Returns all matching entries from the namespace",
	Long:     "Returns all matching entries from the namespace.",
	ArgsName: "<pattern>",
	ArgsLong: `
<pattern> is a glob pattern that is matched against all the names below the
specified mount name.
`,
}

func runGlob(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("glob: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	pattern := args[0]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)

	c, err := ns.Glob(ctx, pattern)
	if err != nil {
		vlog.Infof("ns.Glob(%q) failed: %v", pattern, err)
		return err
	}
	if flagLongGlob {
		// Show all the information we received.
		for res := range c {
			switch v := res.(type) {
			case *naming.GlobReplyEntry:
				fmt.Fprint(cmd.Stdout(), v.Value.Name)
				for _, s := range v.Value.Servers {
					delta := s.Deadline.Time.Sub(time.Now())
					fmt.Fprintf(cmd.Stdout(), " %s (Expires in %d sec)", s.Server, int(delta.Seconds()))
				}
				fmt.Fprintln(cmd.Stdout())
			case *naming.GlobReplyError:
				fmt.Fprintf(cmd.Stderr(), "Error: %s: %v\n", v.Value.Name, v.Value.Error)
			}
		}
		return nil
	}
	// Show a sorted list of unique names, and any errors.
	resultSet := make(map[string]struct{})
	errors := []*naming.GlobError{}
	for res := range c {
		switch v := res.(type) {
		case *naming.GlobReplyEntry:
			if v.Value.Name != "" {
				resultSet[v.Value.Name] = struct{}{}
			}
		case *naming.GlobReplyError:
			errors = append(errors, &v.Value)
		}
	}
	results := []string{}
	for r := range resultSet {
		results = append(results, r)
	}
	sort.Strings(results)
	for _, result := range results {
		fmt.Fprintln(cmd.Stdout(), result)
	}
	for _, err := range errors {
		fmt.Fprintf(cmd.Stderr(), "Error: %s: %v\n", err.Name, err.Error)
	}
	return nil
}

var cmdMount = &cmdline.Command{
	Run:      runMount,
	Name:     "mount",
	Short:    "Adds a server to the namespace",
	Long:     "Adds server <server> to the namespace with name <name>.",
	ArgsName: "<name> <server> <ttl>",
	ArgsLong: `
<name> is the name to add to the namespace.
<server> is the object address of the server to add.
<ttl> is the TTL of the new entry. It is a decimal number followed by a unit
suffix (s, m, h). A value of 0s represents an infinite duration.
`,
}

func runMount(cmd *cmdline.Command, args []string) error {
	if expected, got := 3, len(args); expected != got {
		return cmd.UsageErrorf("mount: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	server := args[1]
	ttlArg := args[2]

	ttl, err := time.ParseDuration(ttlArg)
	if err != nil {
		return fmt.Errorf("TTL parse error: %v", err)
	}

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)
	if err = ns.Mount(ctx, name, server, ttl); err != nil {
		vlog.Infof("ns.Mount(%q, %q, %s) failed: %v", name, server, ttl, err)
		return err
	}
	fmt.Fprintln(cmd.Stdout(), "Server mounted successfully.")
	return nil
}

var cmdUnmount = &cmdline.Command{
	Run:      runUnmount,
	Name:     "unmount",
	Short:    "Removes a server from the namespace",
	Long:     "Removes server <server> with name <name> from the namespace.",
	ArgsName: "<name> <server>",
	ArgsLong: `
<name> is the name to remove from the namespace.
<server> is the object address of the server to remove.
`,
}

func runUnmount(cmd *cmdline.Command, args []string) error {
	if expected, got := 2, len(args); expected != got {
		return cmd.UsageErrorf("unmount: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	server := args[1]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)

	if err := ns.Unmount(ctx, name, server); err != nil {
		vlog.Infof("ns.Unmount(%q, %q) failed: %v", name, server, err)
		return err
	}
	fmt.Fprintln(cmd.Stdout(), "Server unmounted successfully.")
	return nil
}

var cmdResolve = &cmdline.Command{
	Run:      runResolve,
	Name:     "resolve",
	Short:    "Translates a object name to its object address(es)",
	Long:     "Translates a object name to its object address(es).",
	ArgsName: "<name>",
	ArgsLong: "<name> is the name to resolve.",
}

func runResolve(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("resolve: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)

	var opts []naming.NamespaceOpt
	if flagInsecureResolve {
		opts = append(opts, options.SkipServerEndpointAuthorization{})
	}
	me, err := ns.Resolve(ctx, name, opts...)
	if err != nil {
		vlog.Infof("ns.Resolve(%q) failed: %v", name, err)
		return err
	}
	for _, n := range me.Names() {
		fmt.Fprintln(cmd.Stdout(), n)
	}
	return nil
}

var cmdResolveToMT = &cmdline.Command{
	Run:      runResolveToMT,
	Name:     "resolvetomt",
	Short:    "Finds the address of the mounttable that holds an object name",
	Long:     "Finds the address of the mounttable that holds an object name.",
	ArgsName: "<name>",
	ArgsLong: "<name> is the name to resolve.",
}

func runResolveToMT(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("resolvetomt: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]

	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	ns := v23.GetNamespace(ctx)
	var opts []naming.NamespaceOpt
	if flagInsecureResolveToMT {
		opts = append(opts, options.SkipServerEndpointAuthorization{})
	}
	e, err := ns.ResolveToMountTable(ctx, name, opts...)
	if err != nil {
		vlog.Infof("ns.ResolveToMountTable(%q) failed: %v", name, err)
		return err
	}
	for _, s := range e.Servers {
		fmt.Fprintln(cmd.Stdout(), naming.JoinAddressName(s.Server, e.Name))
	}
	return nil
}

var cmdPermissions = &cmdline.Command{
	Name:  "permissions",
	Short: "Manipulates permissions on an entry in the namespace",
	Long: `
Commands to get and set the permissions on a name - controlling the blessing
names required to resolve the name.

The permissions are provided as an JSON-encoded version of the Permissions type
defined in v.io/v23/security/access/types.vdl.
`,
	Children: []*cmdline.Command{cmdPermissionsGet, cmdPermissionsSet},
}

var cmdPermissionsSet = &cmdline.Command{
	Run:   runPermissionsSet,
	Name:  "set",
	Short: "Sets permissions on a mount name",
	Long: `
Set replaces the permissions controlling usage of a mount name.
`,
	ArgsName: "<name> <permissions>",
	ArgsLong: `
<name> is the name on which permissions are to be set.

<permissions> is the path to a file containing a JSON-encoded Permissions
object (defined in v.io/v23/security/access/types.vdl), or "-" for STDIN.
`,
}

func runPermissionsSet(cmd *cmdline.Command, args []string) error {
	if expected, got := 2, len(args); expected != got {
		return cmd.UsageErrorf("set: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	var perms access.Permissions
	file := os.Stdin
	if args[1] != "-" {
		var err error
		if file, err = os.Open(args[1]); err != nil {
			return err
		}
		defer file.Close()
	}
	if err := json.NewDecoder(file).Decode(&perms); err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()
	ns := v23.GetNamespace(ctx)
	for {
		_, etag, err := ns.GetPermissions(ctx, name)
		if err != nil && verror.ErrorID(err) != naming.ErrNoSuchName.ID {
			return err
		}
		if err = ns.SetPermissions(ctx, name, perms, etag); verror.ErrorID(err) == verror.ErrBadVersion.ID {
			vlog.Infof("SetPermissions(%q, %q) failed: %v, retrying...", name, etag, err)
			continue
		}
		return err
	}
}

var cmdPermissionsGet = &cmdline.Command{
	Run:      runPermissionsGet,
	Name:     "get",
	Short:    "Gets permissions on a mount name",
	ArgsName: "<name>",
	ArgsLong: `
<name> is a name in the namespace.
`,
	Long: `
Get retrieves the permissions on the usage of a name.

The output is a JSON-encoded Permissions object (defined in
v.io/v23/security/access/types.vdl).
`,
}

func runPermissionsGet(cmd *cmdline.Command, args []string) error {
	if expected, got := 1, len(args); expected != got {
		return cmd.UsageErrorf("get: incorrect number of arguments, expected %d, got %d", expected, got)
	}
	name := args[0]
	ctx, cancel := context.WithTimeout(gctx, time.Minute)
	defer cancel()

	perms, _, err := v23.GetNamespace(ctx).GetPermissions(ctx, name)
	if err != nil {
		return err
	}
	return json.NewEncoder(cmd.Stdout()).Encode(perms)
}

func root() *cmdline.Command {
	cmdGlob.Flags.BoolVar(&flagLongGlob, "l", false, "Long listing format.")
	cmdResolve.Flags.BoolVar(&flagInsecureResolve, "insecure", false, "Insecure mode: May return results from untrusted servers and invoke Resolve on untrusted mounttables")
	cmdResolveToMT.Flags.BoolVar(&flagInsecureResolveToMT, "insecure", false, "Insecure mode: May return results from untrusted servers and invoke Resolve on untrusted mounttables")
	return &cmdline.Command{
		Name:  "namespace",
		Short: "resolves and manages names in the Vanadium namespace",
		Long: `
Command namespace resolves and manages names in the Vanadium namespace.

The namespace roots are set from the command line via --v23.namespace.root
command line option or from environment variables that have a name starting
with V23_NAMESPACE, e.g.  V23_NAMESPACE, V23_NAMESPACE_2, V23_NAMESPACE_GOOGLE,
etc.  The command line options override the environment.
`,
		Children: []*cmdline.Command{cmdGlob, cmdMount, cmdUnmount, cmdResolve, cmdResolveToMT, cmdPermissions},
	}
}
<|endoftext|>