<commit_before>// Websocket Client
//
// Connects to the configured websocket server.
// Usage:
// ws := web.New(...)
// go ws.Connect() // Gracefull Reconnection
// defer ws.Close()
// i, err := svc.Write([]byte("hello"))
// for {
//     b, _ := ws.Read()
//     fmt.Println(string(b))
// }
//

package web

import (
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"player/event"
	"player/logger"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
)

type message struct {
	typ int
	msg []byte
	err error
}

// Websocket connection interface
type ReadWriteCloser interface {
	ReadMessage() (messageType int, p []byte, err error)
	WriteMessage(messageType int, data []byte) error
	Close() error
}

// Default dialer function
var dialer Dialer = &websocket.Dialer{}

// Implemented by websocket.Dialer
type Dialer interface {
	Dial(urlStr string, headers http.Header) (*websocket.Conn, *http.Response, error)
}

// Websocket client
type Client struct {
	// Exported Fields
	Config Configurer
	// Unexported Fields
	id string
	// Connection & state
	conn      ReadWriteCloser
	connected bool
	// Received messages
	messageC chan message
	// Orchestraion
	wg       *sync.WaitGroup
	closeC   chan bool
	connectC chan bool
}

// Constructs the connection url
func (c *Client) url() string {
	u := url.URL{Scheme: "ws", Host: c.Config.Host(), Path: "/"}
	return u.String()
}

// Returns headers tp use for connecting to the server
func (c *Client) headers() http.Header {
	headers := http.Header{}
	// Authorization header
	plain := fmt.Sprintf("%s:%s", c.Config.Username(), c.Config.Password())
	encoded := base64.StdEncoding.EncodeToString([]byte(plain))
	headers.Add("Authorization", fmt.Sprintf("Basic %s", encoded))
	// Topics we want to subscribe too
	// TODO: live in config?
	topics := []string{
		"player:play",
		"player:stop",
		"player:pause",
		"player:resume",
	}
	headers.Add("Topics", strings.Join(topics, ","))
	return headers
}

// Connect to server
func (c *Client) connect() error {
	logger.WithField("url", c.url()).Debug("connecting to websocket server")
	conn, _, err := dialer.Dial(c.url(), c.headers())
	if err != nil {
		return err
	}
	c.connected = true
	c.conn = conn
	return nil
}

// Reads messages from the websocket connection
func (c *Client) read() {
	c.wg.Add(1)
	defer c.wg.Done()
	logger.Debug("start websocket read loop")
	defer logger.Debug("exit websocket read loop")
	for c.connected {
		typ, msg, err := c.conn.ReadMessage()
		if err != nil {
			c.connected = false
			c.conn = nil
			event.Del(c) // Remove from event hub
			defer logger.WithError(err).Error("error reading websocket server")
			select {
			case <-c.closeC:
				// Don't connect if closing
				return
			default:
				go c.Connect() // Start the connection routine
				return
			}
		}
		c.messageC <- message{typ, msg, err} // Place message on channel
	}
}

// Returns instance ID
func (c *Client) ID() string {
	return c.id
}

// Connect to the websocket server
func (c *Client) Connect() {
	c.wg.Add(1)
	defer c.wg.Done()
	logger.Debug("start websocket connect loop")
	defer logger.Debug("exit websocket service connect loop")
	c.connectC <- true // connect immediately
	for {
		select {
		case <-c.closeC:
			return
		case <-c.connectC:
			break
		case <-time.After(c.Config.Retry()):
			break
		}
		if err := c.connect(); err != nil {
			logger.WithError(err).WithFields(logger.F{
				"retry": c.Config.Retry(),
				"url":   c.url(),
			}).Error("failed connecting to websocket server")
			continue
		}
		event.Add(c) // Add to event hub
		go c.read()  // Start a read routine
		return
	}
}

// Connected state
func (c *Client) Connected() bool {
	return c.connected
}

// Read messages from the websocket server
func (c *Client) Read() ([]byte, error) {
	select {
	case <-c.closeC:
		return nil, io.EOF
	case message := <-c.messageC:
		return message.msg, message.err
	}
}

// Writes messages to websocket server
func (c *Client) Write(b []byte) (int, error) {
	if c.connected && c.conn != nil {
		if err := c.conn.WriteMessage(websocket.TextMessage, b); err != nil {
			return 0, err
		}
		return len(b), nil
	}
	logger.Warn("unable to write to websocket server")
	return 0, nil
}

// Gracefully closes the websocket connection
func (c *Client) Close() error {
	logger.Debug("close websocket client")
	defer logger.Info("closed websocket client")
	// Close the closeC
	close(c.closeC)
	// Close the websocket connection
	if c.connected && c.conn != nil {
		err := c.conn.WriteMessage(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(
				websocket.CloseNormalClosure, ""))
		if err != nil {
			logger.WithError(err).Error("error closing connection")
		}
		if err := c.conn.Close(); err != nil {
			logger.WithError(err).Error("error closing connection")
		}
	}
	// Wait for routines to exit
	c.wg.Wait()
	return nil
}

// Constructs a new websocket Client
func New(c Configurer) *Client {
	return &Client{
		// Exported Fields
		Config: c,
		// ID
		id: xid.New().String(),
		// Read messages
		messageC: make(chan message),
		// Orechestration
		wg:       &sync.WaitGroup{},
		closeC:   make(chan bool, 1),
		connectC: make(chan bool, 1),
	}
}
<commit_msg>Improved connection retry logic<commit_after>// Websocket Client
//
// Connects to the configured websocket server.
// Usage:
// ws := web.New(...)
// go ws.Connect() // Gracefull Reconnection
// defer ws.Close()
// i, err := svc.Write([]byte("hello"))
// for {
//     b, _ := ws.Read()
//     fmt.Println(string(b))
// }
//

package web

import (
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"player/event"
	"player/logger"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
)

type message struct {
	typ int
	msg []byte
	err error
}

// Websocket connection interface
type ReadWriteCloser interface {
	ReadMessage() (messageType int, p []byte, err error)
	WriteMessage(messageType int, data []byte) error
	Close() error
}

// Default dialer function
var dialer Dialer = &websocket.Dialer{}

// Implemented by websocket.Dialer
type Dialer interface {
	Dial(urlStr string, headers http.Header) (*websocket.Conn, *http.Response, error)
}

// Websocket client
type Client struct {
	// Exported Fields
	Config Configurer
	// Unexported Fields
	id string
	// Connection & state
	conn      ReadWriteCloser
	connected bool
	// Received messages
	messageC chan message
	// Orchestraion
	wg       *sync.WaitGroup
	closeC   chan bool
	connectC chan bool
}

// Constructs the connection url
func (c *Client) url() string {
	u := url.URL{Scheme: "ws", Host: c.Config.Host(), Path: "/"}
	return u.String()
}

// Returns headers tp use for connecting to the server
func (c *Client) headers() http.Header {
	headers := http.Header{}
	// Authorization header
	plain := fmt.Sprintf("%s:%s", c.Config.Username(), c.Config.Password())
	encoded := base64.StdEncoding.EncodeToString([]byte(plain))
	headers.Add("Authorization", fmt.Sprintf("Basic %s", encoded))
	// Topics we want to subscribe too
	// TODO: live in config?
	topics := []string{
		"player:play",
		"player:stop",
		"player:pause",
		"player:resume",
	}
	headers.Add("Topics", strings.Join(topics, ","))
	return headers
}

// Connect to server
func (c *Client) connect() {
	logger.Debug("start websocket connect lopp")
	defer logger.Debug("exit websocket connect loop")
	for {
		select {
		case <-c.closeC:
			return
		case <-c.connectC:
			break
		case <-time.After(c.Config.Retry()):
			break
		}
		logger.WithField("url", c.url()).Debug("connecting to websocket server")
		conn, _, err := dialer.Dial(c.url(), c.headers())
		if err != nil {
			logger.WithError(err).Error("failed to connect to websocket server")
			continue
		}
		conn.SetPingHandler(c.ping)
		c.connected = true
		c.conn = conn
		event.Add(c) // Add to event hub
		go c.read()  // Start a read routine
		break
	}
}

// Ping handler, pongs back
func (c *Client) ping(string) error {
	logger.Debug("ping from websocket server")
	return c.conn.WriteMessage(websocket.PongMessage, []byte{})
}

// Reads messages from the websocket connection
func (c *Client) read() {
	c.wg.Add(1)
	defer c.wg.Done()
	logger.Debug("start websocket read loop")
	defer logger.Debug("exit websocket read loop")
	for c.connected {
		typ, msg, err := c.conn.ReadMessage()
		if err != nil {
			c.connected = false
			c.conn = nil
			event.Del(c) // Remove from event hub
			defer logger.WithError(err).Error("error reading websocket server")
			select {
			case <-c.closeC:
				// Don't connect if closing
				return
			default:
				go func() {
					c.wg.Add(1)
					c.connect()
					c.wg.Done()
				}()
				return
			}
		}
		logger.WithFields(logger.F{
			"type":    typ,
			"message": string(msg),
		}).Debug("received websocket message")
		if typ == websocket.TextMessage {
			c.messageC <- message{typ, msg, err} // Place message on channel
		}
	}
}

// Returns instance ID
func (c *Client) ID() string {
	return c.id
}

// Connect to the websocket server
func (c *Client) Connect() {
	c.connectC <- true // connect immediately
	go func() {
		c.wg.Add(1)
		c.connect()
		c.wg.Done()
	}()
}

// Connected state
func (c *Client) Connected() bool {
	return c.connected
}

// Read messages from the websocket server
func (c *Client) Read() ([]byte, error) {
	select {
	case <-c.closeC:
		return nil, io.EOF
	case message := <-c.messageC:
		return message.msg, message.err
	}
}

// Writes messages to websocket server
func (c *Client) Write(b []byte) (int, error) {
	if c.connected && c.conn != nil {
		if err := c.conn.WriteMessage(websocket.TextMessage, b); err != nil {
			return 0, err
		}
		return len(b), nil
	}
	logger.Warn("unable to write to websocket server")
	return 0, nil
}

// Gracefully closes the websocket connection
func (c *Client) Close() error {
	logger.Debug("close websocket client")
	defer logger.Info("closed websocket client")
	// Close the closeC
	close(c.closeC)
	// Close the websocket connection
	if c.connected && c.conn != nil {
		err := c.conn.WriteMessage(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(
				websocket.CloseNormalClosure, ""))
		if err != nil {
			logger.WithError(err).Error("error closing connection")
		}
		if err := c.conn.Close(); err != nil {
			logger.WithError(err).Error("error closing connection")
		}
	}
	// Wait for routines to exit
	c.wg.Wait()
	return nil
}

// Constructs a new websocket Client
func New(c Configurer) *Client {
	return &Client{
		// Exported Fields
		Config: c,
		// ID
		id: xid.New().String(),
		// Read messages
		messageC: make(chan message),
		// Orechestration
		wg:       &sync.WaitGroup{},
		closeC:   make(chan bool, 1),
		connectC: make(chan bool, 1),
	}
}
<|endoftext|>


<commit_before>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// 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 check

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// test:
// flag: OPTION
// set: (true|false)
// compare:
//   op: (eq|gt|gte|lt|lte|has)
//   value: val

type binOp string

const (
	and binOp = "and"
	or        = "or"
)

type testItem struct {
	Flag    string
	Output  string
	Value   string
	Set     bool
	Compare compare
}

type compare struct {
	Op    string
	Value string
}

func (t *testItem) execute(s string) (result bool) {
	result = false
	match := strings.Contains(s, t.Flag)

	if t.Set {
		var flagVal string
		isset := match

		if isset && t.Compare.Op != "" {
			// Expects flags in the form;
			// --flag=somevalue
			// --flag
			// somevalue
			pttn := `(` + t.Flag + `)(=)*([^\s,]*) *`
			flagRe := regexp.MustCompile(pttn)
			vals := flagRe.FindStringSubmatch(s)

			if len(vals) > 0 {
				if vals[3] != "" {
					flagVal = vals[3]
				} else {
					flagVal = vals[1]
				}
			} else {
				fmt.Fprintf(os.Stderr, "invalid flag in testitem definition")
				os.Exit(1)
			}

			switch t.Compare.Op {
			case "eq":
				result = flagVal == t.Compare.Value

			case "noteq":
				result = !(flagVal == t.Compare.Value)

			case "gt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a > b

			case "gte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a >= b

			case "lt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a < b

			case "lte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a <= b

			case "has":
				result = strings.Contains(flagVal, t.Compare.Value)

			case "nothave":
				result = !strings.Contains(flagVal, t.Compare.Value)
			}
		} else {
			result = isset
		}

	} else {
		notset := !match
		result = notset
	}

	return
}

type tests struct {
	TestItems []*testItem `yaml:"test_items"`
	BinOp     binOp       `yaml:"bin_op"`
}

func (ts *tests) execute(s string) (result bool) {
	res := make([]bool, len(ts.TestItems))

	for i, t := range ts.TestItems {
		res[i] = t.execute(s)
	}

	// If no binary operation is specified, default to AND
	switch ts.BinOp {
	default:
		fmt.Fprintf(os.Stderr, "unknown binary operator for tests %s\n", ts.BinOp)
		os.Exit(1)
	case and, "":
		result = true
		for i := range res {
			result = result && res[i]
		}
	case or:
		result = false
		for i := range res {
			result = result || res[i]
		}
	}

	return
}

func toNumeric(a, b string) (c, d int) {
	var err error
	c, err = strconv.Atoi(a)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", a, err)
		os.Exit(1)
	}
	d, err = strconv.Atoi(b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", b, err)
		os.Exit(1)
	}

	return c, d
}
<commit_msg>Fixed issue of The controls for master - admission control showing wrong status #49<commit_after>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// 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 check

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// test:
// flag: OPTION
// set: (true|false)
// compare:
//   op: (eq|gt|gte|lt|lte|has)
//   value: val

type binOp string

const (
	and binOp = "and"
	or        = "or"
)

type testItem struct {
	Flag    string
	Output  string
	Value   string
	Set     bool
	Compare compare
}

type compare struct {
	Op    string
	Value string
}

func (t *testItem) execute(s string) (result bool) {
	result = false
	match := strings.Contains(s, t.Flag)

	if t.Set {
		var flagVal string
		isset := match

		if isset && t.Compare.Op != "" {
			// Expects flags in the form;
			// --flag=somevalue
			// --flag
			// somevalue
			//pttn := `(` + t.Flag + `)(=)*([^\s,]*) *`
			pttn := `(` + t.Flag + `)(=)*([^\s]*) *`
			flagRe := regexp.MustCompile(pttn)
			vals := flagRe.FindStringSubmatch(s)

			if len(vals) > 0 {
				if vals[3] != "" {
					flagVal = vals[3]
				} else {
					flagVal = vals[1]
				}
			} else {
				fmt.Fprintf(os.Stderr, "invalid flag in testitem definition")
				os.Exit(1)
			}

			switch t.Compare.Op {
			case "eq":
				result = flagVal == t.Compare.Value

			case "noteq":
				result = !(flagVal == t.Compare.Value)

			case "gt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a > b

			case "gte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a >= b

			case "lt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a < b

			case "lte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a <= b

			case "has":
				result = strings.Contains(flagVal, t.Compare.Value)

			case "nothave":
				result = !strings.Contains(flagVal, t.Compare.Value)
			}
		} else {
			result = isset
		}

	} else {
		notset := !match
		result = notset
	}

	return
}

type tests struct {
	TestItems []*testItem `yaml:"test_items"`
	BinOp     binOp       `yaml:"bin_op"`
}

func (ts *tests) execute(s string) (result bool) {
	res := make([]bool, len(ts.TestItems))

	for i, t := range ts.TestItems {
		res[i] = t.execute(s)
	}

	// If no binary operation is specified, default to AND
	switch ts.BinOp {
	default:
		fmt.Fprintf(os.Stderr, "unknown binary operator for tests %s\n", ts.BinOp)
		os.Exit(1)
	case and, "":
		result = true
		for i := range res {
			result = result && res[i]
		}
	case or:
		result = false
		for i := range res {
			result = result || res[i]
		}
	}

	return
}

func toNumeric(a, b string) (c, d int) {
	var err error
	c, err = strconv.Atoi(a)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", a, err)
		os.Exit(1)
	}
	d, err = strconv.Atoi(b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", b, err)
		os.Exit(1)
	}

	return c, d
}
<|endoftext|>


<commit_before>package common

import (
	"github.com/dgrijalva/jwt-go"
)

type Consumer struct {
	Iss    string `json:"iss"`
	Secret string `json:"secret"`
}

func (c *Consumer) CreateToken() (string, error) {
	token := jwt.New(jwt.SigningMethodHS256)
	claims := token.Claims.(jwt.MapClaims)
	claims["iss"] = c.Iss

	tokenString, err := token.SignedString([]byte(c.Secret))
	if err != nil {
		return "", err
	}
	return tokenString, err
}
<commit_msg>add comment<commit_after>package common

import (
	"github.com/dgrijalva/jwt-go"
)

// Consumer defined consumer iss and secret use jwt token
type Consumer struct {
	Iss    string `json:"iss"`
	Secret string `json:"secret"`
}

// CreateToken with iss , secret
func (c *Consumer) CreateToken() (string, error) {
	token := jwt.New(jwt.SigningMethodHS256)
	claims := token.Claims.(jwt.MapClaims)
	claims["iss"] = c.Iss

	tokenString, err := token.SignedString([]byte(c.Secret))
	if err != nil {
		return "", err
	}
	return tokenString, err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"encoding/json"

	"github.com/grafana/grafana-plugin-sdk-go/backend"
	"github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
	"github.com/grafana/grafana-plugin-sdk-go/data"
	"golang.org/x/net/context"
)

func GetDatasourceServeOpts() datasource.ServeOpts {
	ds := &ClickHouseDatasource{
		im: datasource.NewInstanceManager(NewDatasourceSettings),
	}

	return datasource.ServeOpts{
		QueryDataHandler:   ds,
		CheckHealthHandler: ds,
	}
}

type ClickHouseDatasource struct {
	im instancemgmt.InstanceManager
}

func (ds *ClickHouseDatasource) getClient(ctx backend.PluginContext) (*ClickHouseClient, error) {
	im, err := ds.im.Get(ctx)
	if err != nil {
		return nil, err
	}

	return &ClickHouseClient{
		settings: im.(*DatasourceSettings),
	}, nil
}

func (ds *ClickHouseDatasource) query(ctx backend.PluginContext, query *Query) backend.DataResponse {

	onErr := func(err error) backend.DataResponse {
		backend.Logger.Error(fmt.Sprintf("Datasource query error: %w", err.Error()))
		return backend.DataResponse{Error: err}
	}

	client, err := ds.getClient(ctx)
	if err != nil {
		return onErr(err)
	}

	res, err := client.Query(query.Format())
	if err != nil {
		return onErr(err)
	}
	frame := res.toFrame(query.RefId, client.FetchTimeZone)
	response := backend.DataResponse{}

	if frame != nil {
		response.Frames = []*data.Frame{frame}
	}

	return response
}

func (ds *ClickHouseDatasource) QueryData(
	ctx context.Context,
	req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {

	onErr := func(err error) (*backend.QueryDataResponse, error) {
		backend.Logger.Error("QueryData error: " + err.Error())
		return nil, err
	}

	response := backend.NewQueryDataResponse()

	for _, query := range req.Queries {
		var q = Query{}
		err := json.Unmarshal(query.JSON, &q)
		if err != nil {
			return onErr(fmt.Errorf("Unable to parse json %s. Error: %w", query.JSON, err))
		}

		response.Responses[q.RefId] = ds.query(req.PluginContext, &q)
	}

	return response, nil
}

func (ds *ClickHouseDatasource) CheckHealth(
	ctx context.Context,
	req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {

	onErr := func(err error) (*backend.CheckHealthResult, error) {
		backend.Logger.Error(fmt.Sprintf("HealthCheck error: %w", err))
		return &backend.CheckHealthResult{
			Status:  backend.HealthStatusError,
			Message: err.Error(),
		}, err
	}

	client, err := ds.getClient(req.PluginContext)
	if err != nil {
		return onErr(err)
	}
	_, err = client.Query(DefaultQuery)
	if err != nil {
		return onErr(err)
	}

	return &backend.CheckHealthResult{
		Status:  backend.HealthStatusOk,
		Message: "OK",
	}, nil
}
<commit_msg>Fix error formatting.<commit_after>package main

import (
	"fmt"
	"encoding/json"

	"github.com/grafana/grafana-plugin-sdk-go/backend"
	"github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
	"github.com/grafana/grafana-plugin-sdk-go/data"
	"golang.org/x/net/context"
)

func GetDatasourceServeOpts() datasource.ServeOpts {
	ds := &ClickHouseDatasource{
		im: datasource.NewInstanceManager(NewDatasourceSettings),
	}

	return datasource.ServeOpts{
		QueryDataHandler:   ds,
		CheckHealthHandler: ds,
	}
}

type ClickHouseDatasource struct {
	im instancemgmt.InstanceManager
}

func (ds *ClickHouseDatasource) getClient(ctx backend.PluginContext) (*ClickHouseClient, error) {
	im, err := ds.im.Get(ctx)
	if err != nil {
		return nil, err
	}

	return &ClickHouseClient{
		settings: im.(*DatasourceSettings),
	}, nil
}

func (ds *ClickHouseDatasource) query(ctx backend.PluginContext, query *Query) backend.DataResponse {

	onErr := func(err error) backend.DataResponse {
		backend.Logger.Error(fmt.Sprintf("Datasource query error: %w", err))
		return backend.DataResponse{Error: err}
	}

	client, err := ds.getClient(ctx)
	if err != nil {
		return onErr(err)
	}

	res, err := client.Query(query.Format())
	if err != nil {
		return onErr(err)
	}
	frame := res.toFrame(query.RefId, client.FetchTimeZone)
	response := backend.DataResponse{}

	if frame != nil {
		response.Frames = []*data.Frame{frame}
	}

	return response
}

func (ds *ClickHouseDatasource) QueryData(
	ctx context.Context,
	req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {

	onErr := func(err error) (*backend.QueryDataResponse, error) {
		backend.Logger.Error(fmt.Sprintf("QueryData error: %w", err))
		return nil, err
	}

	response := backend.NewQueryDataResponse()

	for _, query := range req.Queries {
		var q = Query{}
		err := json.Unmarshal(query.JSON, &q)
		if err != nil {
			return onErr(fmt.Errorf("Unable to parse json %s. Error: %w", query.JSON, err))
		}

		response.Responses[q.RefId] = ds.query(req.PluginContext, &q)
	}

	return response, nil
}

func (ds *ClickHouseDatasource) CheckHealth(
	ctx context.Context,
	req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {

	onErr := func(err error) (*backend.CheckHealthResult, error) {
		backend.Logger.Error(fmt.Sprintf("HealthCheck error: %w", err))
		return &backend.CheckHealthResult{
			Status:  backend.HealthStatusError,
			Message: err.Error(),
		}, err
	}

	client, err := ds.getClient(req.PluginContext)
	if err != nil {
		return onErr(err)
	}
	_, err = client.Query(DefaultQuery)
	if err != nil {
		return onErr(err)
	}

	return &backend.CheckHealthResult{
		Status:  backend.HealthStatusOk,
		Message: "OK",
	}, nil
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"reflect"
	"strconv"
)

type option struct {
	field    string
	isFlag   bool
	desc     string
	short    string
	long     string
	required bool
	value    string
}

// Reflect the gathered information into the concrete action instance.
func (o *option) reflectTo(value reflect.Value) (e error) {
	if o.value == "" {
		if o.required {
			return fmt.Errorf("option %q is required but not set", o.field)
		}
		return nil
	}

	field := value.FieldByName(o.field)

	switch field.Kind() {
	case reflect.String:
		field.SetString(o.value)
	case reflect.Int:
		i, e := strconv.Atoi(o.value)
		if e != nil {
			return e
		}
		field.SetInt(int64(i))
	case reflect.Bool:
		field.SetBool(o.value == "true")
	default:
		return fmt.Errorf("invalid type %q", field.Type().String())
	}
	return nil
}

func (o *option) description() string {
	desc := "    "
	desc += o.shortDescription(" ")
	desc += fmt.Sprintf("%-*s", 30-len(desc), " ") + o.desc
	if o.value != "" {
		desc += " (default: " + o.value + ")"
	}
	return desc
}

func (o *option) shortDescription(sep string) (desc string) {
	if o.short != "" {
		desc += "-" + o.short
	}
	if o.long != "" {
		if o.short != "" {
			desc += sep
		}
		desc += "--" + o.long
	}

	if !o.isFlag {
		desc += " <" + o.field + ">"
	}

	return desc
}

func (a *action) createOption(field reflect.StructField, value reflect.Value, tagMap map[string]string) (e error) {
	if e := validateTagMap(tagMap, "type", "desc", "short", "long", "required", "default"); e != nil {
		return fmt.Errorf("[option:%s] %s", field.Name, e.Error())
	}
	opt := &option{field: field.Name}

	if field.Type.Kind() == reflect.Bool {
		opt.isFlag = true
	}

	opt.short, e = handleShortIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.long, e = handleLongIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.required, e = handleRequired(tagMap)
	if e != nil {
		return e
	}

	if opt.required && opt.isFlag {
		return fmt.Errorf(`field %q is a flag and required, that doesn't make much sense`, field.Name)
	}

	opt.value = handlePresetValue(field, value)

	// Do that for error checking (not necessary if value preset otherwise).
	defaultValue, e := handleDefault(field, tagMap)
	if e != nil {
		return fmt.Errorf(`wrong value for "default" tag: %s`, e.Error())
	}

	if opt.value == "" {
		opt.value = defaultValue
	} else { // With a preset value that value is not required any more.
		opt.required = false
	}

	opt.desc = handleDescription(tagMap)

	if opt.short == "" && opt.long == "" {
		return fmt.Errorf("option %q has neither long nor short accessor set", field.Name)
	}

	if opt.short != "" {
		if o, found := a.params[opt.short]; found {
			return fmt.Errorf("short option %q already used (option %q)", opt.short, o.field)
		}
		a.params[opt.short] = opt
	}
	if opt.long != "" {
		if o, found := a.params[opt.long]; found {
			return fmt.Errorf("long option %q already used (option %q)", opt.long, o.field)
		}
		a.params[opt.long] = opt
	}
	a.opts = append(a.opts, opt)
	return nil
}
<commit_msg>fixed whitespace in help error<commit_after>package cli

import (
	"fmt"
	"reflect"
	"strconv"
)

type option struct {
	field    string
	isFlag   bool
	desc     string
	short    string
	long     string
	required bool
	value    string
}

// Reflect the gathered information into the concrete action instance.
func (o *option) reflectTo(value reflect.Value) (e error) {
	if o.value == "" {
		if o.required {
			return fmt.Errorf("option %q is required but not set", o.field)
		}
		return nil
	}

	field := value.FieldByName(o.field)

	switch field.Kind() {
	case reflect.String:
		field.SetString(o.value)
	case reflect.Int:
		i, e := strconv.Atoi(o.value)
		if e != nil {
			return e
		}
		field.SetInt(int64(i))
	case reflect.Bool:
		field.SetBool(o.value == "true")
	default:
		return fmt.Errorf("invalid type %q", field.Type().String())
	}
	return nil
}

func (o *option) description() string {
	desc := "    "
	desc += o.shortDescription(" ")
	desc += fmt.Sprintf("%-*s", 30-len(desc), " ") + o.desc
	if o.value != "" {
		if o.desc != "" {
			desc += " "
		}
		desc += "(default: " + o.value + ")"
	}
	return desc
}

func (o *option) shortDescription(sep string) (desc string) {
	if o.short != "" {
		desc += "-" + o.short
	}
	if o.long != "" {
		if o.short != "" {
			desc += sep
		}
		desc += "--" + o.long
	}

	if !o.isFlag {
		desc += " <" + o.field + ">"
	}

	return desc
}

func (a *action) createOption(field reflect.StructField, value reflect.Value, tagMap map[string]string) (e error) {
	if e := validateTagMap(tagMap, "type", "desc", "short", "long", "required", "default"); e != nil {
		return fmt.Errorf("[option:%s] %s", field.Name, e.Error())
	}
	opt := &option{field: field.Name}

	if field.Type.Kind() == reflect.Bool {
		opt.isFlag = true
	}

	opt.short, e = handleShortIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.long, e = handleLongIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.required, e = handleRequired(tagMap)
	if e != nil {
		return e
	}

	if opt.required && opt.isFlag {
		return fmt.Errorf(`field %q is a flag and required, that doesn't make much sense`, field.Name)
	}

	opt.value = handlePresetValue(field, value)

	// Do that for error checking (not necessary if value preset otherwise).
	defaultValue, e := handleDefault(field, tagMap)
	if e != nil {
		return fmt.Errorf(`wrong value for "default" tag: %s`, e.Error())
	}

	if opt.value == "" {
		opt.value = defaultValue
	} else { // With a preset value that value is not required any more.
		opt.required = false
	}

	opt.desc = handleDescription(tagMap)

	if opt.short == "" && opt.long == "" {
		return fmt.Errorf("option %q has neither long nor short accessor set", field.Name)
	}

	if opt.short != "" {
		if o, found := a.params[opt.short]; found {
			return fmt.Errorf("short option %q already used (option %q)", opt.short, o.field)
		}
		a.params[opt.short] = opt
	}
	if opt.long != "" {
		if o, found := a.params[opt.long]; found {
			return fmt.Errorf("long option %q already used (option %q)", opt.long, o.field)
		}
		a.params[opt.long] = opt
	}
	a.opts = append(a.opts, opt)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os/exec"

	"github.com/flynn/flynn/controller/client"
	"github.com/flynn/go-docopt"
)

func init() {
	register("remote", runRemote, `
usage: flynn remote add [<remote>] [-y]

Create a git remote that allows deploying the application via git.
If a name for the remote is not provided 'flynn' will be used.

Options:
	-y, --yes              Skip the confirmation prompt if the git remote already exists.

Examples:

	$ flynn -a turkeys-stupefy-perry remote add
	Created remote flynn with url https://git.dev.localflynn.com/turkeys-stupefy-perry.git

	$ flynn -a turkeys-stupefy-perry remote add staging
	Created remote staging with url https://git.dev.localflynn.com/turkeys-stupefy-perry.git
`)
}

func runRemote(args *docopt.Args, client controller.Client) error {
	app, err := client.GetApp(mustApp())
	if err != nil {
		return err
	}

	remote := args.String["<remote>"]
	if remote == "" {
		remote = "flynn"
	}

	if !inGitRepo() {
		log.Print("Must be executed within a git repository.")
		return nil
	}

	if !args.Bool["--yes"] {
		update, err := promptReplaceRemote(remote)
		if err != nil {
			return err
		}
		if update == false {
			return nil
		}
	}

	// Register git remote
	url := gitURL(clusterConf, app.Name)
	exec.Command("git", "remote", "remove", remote).Run()
	exec.Command("git", "remote", "add", "--", remote, url).Run()

	log.Printf("Created remote %s with url %s.", remote, url)
	return nil
}
<commit_msg>cli: Document use of -a <app> in 'remote add'<commit_after>package main

import (
	"log"
	"os/exec"

	"github.com/flynn/flynn/controller/client"
	"github.com/flynn/go-docopt"
)

func init() {
	register("remote", runRemote, `
usage: flynn remote add [<remote>] [-y]

Create a git remote that allows deploying the application via git.
If a name for the remote is not provided 'flynn' will be used.

Note that the -a <app> option must be given so the remote to add is known.

Options:
	-y, --yes              Skip the confirmation prompt if the git remote already exists.

Examples:

	$ flynn -a turkeys-stupefy-perry remote add
	Created remote flynn with url https://git.dev.localflynn.com/turkeys-stupefy-perry.git

	$ flynn -a turkeys-stupefy-perry remote add staging
	Created remote staging with url https://git.dev.localflynn.com/turkeys-stupefy-perry.git
`)
}

func runRemote(args *docopt.Args, client controller.Client) error {
	app, err := client.GetApp(mustApp())
	if err != nil {
		return err
	}

	remote := args.String["<remote>"]
	if remote == "" {
		remote = "flynn"
	}

	if !inGitRepo() {
		log.Print("Must be executed within a git repository.")
		return nil
	}

	if !args.Bool["--yes"] {
		update, err := promptReplaceRemote(remote)
		if err != nil {
			return err
		}
		if update == false {
			return nil
		}
	}

	// Register git remote
	url := gitURL(clusterConf, app.Name)
	exec.Command("git", "remote", "remove", remote).Run()
	exec.Command("git", "remote", "add", "--", remote, url).Run()

	log.Printf("Created remote %s with url %s.", remote, url)
	return nil
}
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"fmt"
	"html/template"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/exercism/cli/config"
)

// Status represents the results of a CLI self test.
type Status struct {
	Censor          bool
	Version         versionStatus
	System          systemStatus
	Configuration   configurationStatus
	APIReachability apiReachabilityStatus
	cli             *CLI
	cfg             config.UserConfig
}

type versionStatus struct {
	Current  string
	Latest   string
	Status   string
	Error    error
	UpToDate bool
}

type systemStatus struct {
	OS           string
	Architecture string
	Build        string
}

type configurationStatus struct {
	Home      string
	Workspace string
	File      string
	Token     string
	TokenURL  string
}

type apiReachabilityStatus struct {
	Services []*apiPing
}

type apiPing struct {
	Service string
	URL     string
	Status  string
	Latency time.Duration
}

// NewStatus prepares a value to perform a diagnostic self-check.
func NewStatus(c *CLI, uc config.UserConfig) Status {
	status := Status{
		cli: c,
		cfg: uc,
	}
	return status
}

// Check runs the CLI's diagnostic self-check.
func (status *Status) Check() (string, error) {
	status.Version = newVersionStatus(status.cli)
	status.System = newSystemStatus()

	cs, err := newConfigurationStatus(status)
	if err != nil {
		return "", err
	}
	status.Configuration = cs

	ar, err := newAPIReachabilityStatus()
	if err != nil {
		return "", err
	}
	status.APIReachability = ar

	return status.compile()
}
func (status *Status) compile() (string, error) {
	t, err := template.New("self-test").Parse(tmplSelfTest)
	if err != nil {
		return "", err
	}

	var bb bytes.Buffer
	t.Execute(&bb, status)
	return bb.String(), nil
}

func newAPIReachabilityStatus() (apiReachabilityStatus, error) {
	apiCfg, err := config.NewAPIConfig()
	if err != nil {
		return apiReachabilityStatus{}, nil
	}
	apiCfg.SetDefaults()
	ar := apiReachabilityStatus{
		Services: []*apiPing{
			{Service: "GitHub", URL: "https://api.github.com"},
			{Service: "Exercism", URL: fmt.Sprintf("%s/ping", apiCfg.BaseURL)},
		},
	}
	var wg sync.WaitGroup
	wg.Add(len(ar.Services))
	for _, service := range ar.Services {
		go service.Call(&wg)
	}
	wg.Wait()
	return ar, nil
}

func newVersionStatus(cli *CLI) versionStatus {
	vs := versionStatus{
		Current: cli.Version,
	}
	ok, err := cli.IsUpToDate()
	if err == nil {
		vs.Latest = cli.LatestRelease.Version()
	} else {
		vs.Error = fmt.Errorf("Error: %s", err)
	}
	vs.UpToDate = ok
	return vs
}

func newSystemStatus() systemStatus {
	ss := systemStatus{
		OS:           runtime.GOOS,
		Architecture: runtime.GOARCH,
	}
	if BuildOS != "" && BuildARCH != "" {
		ss.Build = fmt.Sprintf("%s/%s", BuildOS, BuildARCH)
	}
	if BuildARM != "" {
		ss.Build = fmt.Sprintf("%s ARMv%s", ss.Build, BuildARM)
	}
	return ss
}

func newConfigurationStatus(status *Status) (configurationStatus, error) {
	apiCfg, err := config.NewAPIConfig()
	if err != nil {
		return configurationStatus{}, err
	}
	apiCfg.SetDefaults()
	cs := configurationStatus{
		Home:      status.cfg.Home,
		Workspace: status.cfg.Workspace,
		File:      status.cfg.File(),
		Token:     status.cfg.Token,
		TokenURL:  config.InferSiteURL(apiCfg.BaseURL) + "/my/settings",
	}
	if status.Censor && status.cfg.Token != "" {
		cs.Token = redactToken(status.cfg.Token)
	}
	return cs, nil
}

func (ping *apiPing) Call(wg *sync.WaitGroup) {
	defer wg.Done()

	now := time.Now()
	res, err := HTTPClient.Get(ping.URL)
	delta := time.Since(now)
	ping.Latency = delta
	if err != nil {
		ping.Status = err.Error()
		return
	}
	res.Body.Close()
	ping.Status = "connected"
}

func redactToken(token string) string {
	str := token[4 : len(token)-3]
	redaction := strings.Repeat("*", len(str))
	return string(token[:4]) + redaction + string(token[len(token)-3:])
}

const tmplSelfTest = `
Troubleshooting Information
===========================

Version
----------------
Current: {{ .Version.Current }}
Latest:  {{ with .Version.Latest }}{{ . }}{{ else }}<unknown>{{ end }}
{{ with .Version.Error }}
{{ . }}
{{ end -}}
{{ if not .Version.UpToDate }}
Call 'exercism upgrade' to get the latest version.
See the release notes at https://github.com/exercism/cli/releases/tag/{{ .Version.Latest }} for details.
{{ end }}

Operating System
----------------
OS:           {{ .System.OS }}
Architecture: {{ .System.Architecture }}
{{ with .System.Build }}
Build: {{ . }}
{{ end }}

Configuration
----------------
Home:      {{ .Configuration.Home }}
Workspace: {{ .Configuration.Workspace }}
Config:    {{ .Configuration.File }}
API key:   {{ with .Configuration.Token }}{{ . }}{{ else }}<not configured>
Find your API key at {{ .Configuration.TokenURL }}{{ end }}

API Reachability
----------------
{{ range .APIReachability.Services }}
{{ .Service }}:
    * {{ .URL }}
    * [{{ .Status }}]
    * {{ .Latency }}
{{ end }}

If you are having trouble please file a GitHub issue at
https://github.com/exercism/exercism.io/issues and include
this information.
{{ if not .Censor }}
Don't share your API key. Keep that private.
{{ end }}`
<commit_msg>Switch status check to only use user config<commit_after>package cli

import (
	"bytes"
	"fmt"
	"html/template"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/exercism/cli/config"
)

// Status represents the results of a CLI self test.
type Status struct {
	Censor          bool
	Version         versionStatus
	System          systemStatus
	Configuration   configurationStatus
	APIReachability apiReachabilityStatus
	cli             *CLI
	cfg             config.UserConfig
}

type versionStatus struct {
	Current  string
	Latest   string
	Status   string
	Error    error
	UpToDate bool
}

type systemStatus struct {
	OS           string
	Architecture string
	Build        string
}

type configurationStatus struct {
	Home      string
	Workspace string
	File      string
	Token     string
	TokenURL  string
}

type apiReachabilityStatus struct {
	Services []*apiPing
}

type apiPing struct {
	Service string
	URL     string
	Status  string
	Latency time.Duration
}

// NewStatus prepares a value to perform a diagnostic self-check.
func NewStatus(c *CLI, uc config.UserConfig) Status {
	status := Status{
		cli: c,
		cfg: uc,
	}
	return status
}

// Check runs the CLI's diagnostic self-check.
func (status *Status) Check() (string, error) {
	status.Version = newVersionStatus(status.cli)
	status.System = newSystemStatus()

	cs, err := newConfigurationStatus(status)
	if err != nil {
		return "", err
	}
	status.Configuration = cs

	ar, err := newAPIReachabilityStatus(status.cfg.APIBaseURL)
	if err != nil {
		return "", err
	}
	status.APIReachability = ar

	return status.compile()
}
func (status *Status) compile() (string, error) {
	t, err := template.New("self-test").Parse(tmplSelfTest)
	if err != nil {
		return "", err
	}

	var bb bytes.Buffer
	t.Execute(&bb, status)
	return bb.String(), nil
}

func newAPIReachabilityStatus(baseURL string) (apiReachabilityStatus, error) {
	ar := apiReachabilityStatus{
		Services: []*apiPing{
			{Service: "GitHub", URL: "https://api.github.com"},
			{Service: "Exercism", URL: fmt.Sprintf("%s/ping", baseURL)},
		},
	}
	var wg sync.WaitGroup
	wg.Add(len(ar.Services))
	for _, service := range ar.Services {
		go service.Call(&wg)
	}
	wg.Wait()
	return ar, nil
}

func newVersionStatus(cli *CLI) versionStatus {
	vs := versionStatus{
		Current: cli.Version,
	}
	ok, err := cli.IsUpToDate()
	if err == nil {
		vs.Latest = cli.LatestRelease.Version()
	} else {
		vs.Error = fmt.Errorf("Error: %s", err)
	}
	vs.UpToDate = ok
	return vs
}

func newSystemStatus() systemStatus {
	ss := systemStatus{
		OS:           runtime.GOOS,
		Architecture: runtime.GOARCH,
	}
	if BuildOS != "" && BuildARCH != "" {
		ss.Build = fmt.Sprintf("%s/%s", BuildOS, BuildARCH)
	}
	if BuildARM != "" {
		ss.Build = fmt.Sprintf("%s ARMv%s", ss.Build, BuildARM)
	}
	return ss
}

func newConfigurationStatus(status *Status) (configurationStatus, error) {
	cs := configurationStatus{
		Home:      status.cfg.Home,
		Workspace: status.cfg.Workspace,
		File:      status.cfg.File(),
		Token:     status.cfg.Token,
		TokenURL:  config.InferSiteURL(status.cfg.APIBaseURL) + "/my/settings",
	}
	if status.Censor && status.cfg.Token != "" {
		cs.Token = redactToken(status.cfg.Token)
	}
	return cs, nil
}

func (ping *apiPing) Call(wg *sync.WaitGroup) {
	defer wg.Done()

	now := time.Now()
	res, err := HTTPClient.Get(ping.URL)
	delta := time.Since(now)
	ping.Latency = delta
	if err != nil {
		ping.Status = err.Error()
		return
	}
	res.Body.Close()
	ping.Status = "connected"
}

func redactToken(token string) string {
	str := token[4 : len(token)-3]
	redaction := strings.Repeat("*", len(str))
	return string(token[:4]) + redaction + string(token[len(token)-3:])
}

const tmplSelfTest = `
Troubleshooting Information
===========================

Version
----------------
Current: {{ .Version.Current }}
Latest:  {{ with .Version.Latest }}{{ . }}{{ else }}<unknown>{{ end }}
{{ with .Version.Error }}
{{ . }}
{{ end -}}
{{ if not .Version.UpToDate }}
Call 'exercism upgrade' to get the latest version.
See the release notes at https://github.com/exercism/cli/releases/tag/{{ .Version.Latest }} for details.
{{ end }}

Operating System
----------------
OS:           {{ .System.OS }}
Architecture: {{ .System.Architecture }}
{{ with .System.Build }}
Build: {{ . }}
{{ end }}

Configuration
----------------
Home:      {{ .Configuration.Home }}
Workspace: {{ .Configuration.Workspace }}
Config:    {{ .Configuration.File }}
API key:   {{ with .Configuration.Token }}{{ . }}{{ else }}<not configured>
Find your API key at {{ .Configuration.TokenURL }}{{ end }}

API Reachability
----------------
{{ range .APIReachability.Services }}
{{ .Service }}:
    * {{ .URL }}
    * [{{ .Status }}]
    * {{ .Latency }}
{{ end }}

If you are having trouble please file a GitHub issue at
https://github.com/exercism/exercism.io/issues and include
this information.
{{ if not .Censor }}
Don't share your API key. Keep that private.
{{ end }}`
<|endoftext|>


<commit_before>package httputil

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os"
)

// GetWriteFile performs a HTTP GET request and saves the response body
// to the file path specified
func GetWriteFile(url string, filename string, perm os.FileMode) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	bytes, err := ResponseBody(resp)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename, bytes, perm)
	return err
}

// ResponseBody returns the body as a byte array
func ResponseBody(res *http.Response) ([]byte, error) {
	defer res.Body.Close()
	contents, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return []byte{}, err
	} else {
		return contents, nil
	}
}

// ResponseBodyJsonMapIndent returns the body as a generic JSON dictionary
func ResponseBodyJsonMapIndent(res *http.Response, prefix string, indent string) ([]byte, error) {
	body, err := ResponseBody(res)
	if err != nil {
		return body, err
	}
	any := map[string]interface{}{}
	json.Unmarshal(body, &any)
	return json.MarshalIndent(any, prefix, indent)
}
<commit_msg>update doc formatting<commit_after>package httputil

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os"
)

// GetWriteFile performs a HTTP GET request and saves the response body
// to the file path specified

func GetWriteFile(url string, filename string, perm os.FileMode) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	bytes, err := ResponseBody(resp)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename, bytes, perm)
	return err
}

// ResponseBody returns the body as a byte array

func ResponseBody(res *http.Response) ([]byte, error) {
	defer res.Body.Close()
	contents, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return []byte{}, err
	} else {
		return contents, nil
	}
}

// ResponseBodyJsonMapIndent returns the body as a generic JSON dictionary

func ResponseBodyJsonMapIndent(res *http.Response, prefix string, indent string) ([]byte, error) {
	body, err := ResponseBody(res)
	if err != nil {
		return body, err
	}
	any := map[string]interface{}{}
	json.Unmarshal(body, &any)
	return json.MarshalIndent(any, prefix, indent)
}
<|endoftext|>


<commit_before>package garagepi_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/robdimsdale/garagepi"
	fakes "github.com/robdimsdale/garagepi/fakes"
)

var _ = Describe("OsHelper", func() {
	var osHelper garagepi.OsHelper

	BeforeEach(func() {
		fakeLogger := &fakes.FakeLogger{}
		osHelper = garagepi.NewOsHelperImpl(fakeLogger)
	})

	Describe("Exec", func() {
		It("Executes 'echo hello world' successfully", func() {
			resp, err := osHelper.Exec("echo", "hello", "world")
			Expect(err).NotTo(HaveOccurred())
			Expect(resp).To(Equal("hello world\n"))
		})
	})
})
<commit_msg>Add osHelper sleep test.<commit_after>package garagepi_test

import (
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/robdimsdale/garagepi"
	fakes "github.com/robdimsdale/garagepi/fakes"
)

var _ = Describe("OsHelper", func() {
	var osHelper garagepi.OsHelper

	BeforeEach(func() {
		fakeLogger := &fakes.FakeLogger{}
		osHelper = garagepi.NewOsHelperImpl(fakeLogger)
	})

	Describe("Exec", func() {
		It("Executes 'echo hello world' successfully", func() {
			resp, err := osHelper.Exec("echo", "hello", "world")
			Expect(err).NotTo(HaveOccurred())
			Expect(resp).To(Equal("hello world\n"))
		})
	})

	Describe("Sleep", func() {
		It("Sleeps for at least the time provided", func() {
			start := time.Now()
			osHelper.Sleep(1 * time.Second)
			end := time.Now()
			Expect((end.Sub(start)).Seconds()).Should(BeNumerically(">", 1))
		})
	})
})
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"strconv"

	"github.com/bcrusu/kcm/cmd/create"
	"github.com/bcrusu/kcm/cmd/download"
	"github.com/bcrusu/kcm/cmd/start"
	"github.com/bcrusu/kcm/cmd/validate"
	"github.com/bcrusu/kcm/repository"
	"github.com/bcrusu/kcm/util"
	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

const DefaultKubernetesVersion = "1.6.4"
const DefaultCoreOSVersion = "1353.7.0"
const DefaultCoreOsChannel = "stable"

type createCmdState struct {
	KubernetesVersion  string
	CoreOSVersion      string
	CoreOSChannel      string
	LibvirtStoragePool string
	NodesCount         uint
	KubernetesNetwork  string
	SSHPublicKeyPath   string
	Start              bool
	IPv4CIDR           string
	MasterCPUs         uint
	MasterMemory       uint
	NondeCPUs          uint
	NodeMemory         uint
}

func newCreateCmd() *cobra.Command {
	var cmd = &cobra.Command{
		Use:          "create CLUSTER_NAME",
		Short:        "Create a new cluster",
		SilenceUsage: true,
	}

	state := &createCmdState{}

	cmd.PersistentFlags().StringVar(&state.KubernetesVersion, "kubernetes-version", DefaultKubernetesVersion, "Kubernetes version to use")
	cmd.PersistentFlags().StringVar(&state.CoreOSVersion, "coreos-version", DefaultCoreOSVersion, "CoreOS version to use")
	cmd.PersistentFlags().StringVar(&state.CoreOSChannel, "coreos-channel", DefaultCoreOsChannel, "CoreOS release channel: stable, beta, alpha")
	cmd.PersistentFlags().StringVar(&state.LibvirtStoragePool, "libvirt-pool", "default", "Libvirt storage pool")
	cmd.PersistentFlags().UintVar(&state.NodesCount, "node-count", 1, "Initial number of nodes in the cluster") //TODO: set value to 3
	cmd.PersistentFlags().StringVar(&state.KubernetesNetwork, "kubernetes-network", "flannel", "Networking mode to use. Only flannel is suppoted at the moment")
	cmd.PersistentFlags().StringVar(&state.SSHPublicKeyPath, "ssh-public-key", util.GetUserDefaultSSHPublicKeyPath(), "SSH public key to use")
	cmd.PersistentFlags().BoolVarP(&state.Start, "start", "s", false, "Start the cluster immediately")
	cmd.PersistentFlags().StringVar(&state.IPv4CIDR, "ipv4-cidr", "10.1.0.0/16", "Libvirt network IPv4 CIDR. Network 10.2.0.0/16 is reserved for pods/services network")
	cmd.PersistentFlags().UintVar(&state.MasterCPUs, "master-cpu", 1, "Master node allocated CPUs")
	cmd.PersistentFlags().UintVar(&state.MasterMemory, "master-memory", 512, "Master node memory (in MiB)")
	cmd.PersistentFlags().UintVar(&state.NondeCPUs, "node-cpu", 1, "Node allocated CPUs")
	cmd.PersistentFlags().UintVar(&state.NodeMemory, "node-memory", 512, "Node memory (in MiB)")

	cmd.RunE = state.runE
	return cmd
}

func (s *createCmdState) runE(cmd *cobra.Command, args []string) error {
	if len(args) != 1 {
		return errors.New("invalid command arguments")
	}

	cluster, err := s.createClusterDefinition(args[0])
	if err != nil {
		return err
	}

	// lightweight validation - valid names, empty strings, no nodes defined, etc.
	if err := cluster.Validate(); err != nil {
		return err
	}

	clusterRepository, err := newClusterRepository()
	if err != nil {
		return err
	}

	exists, err := clusterRepository.Exists(cluster.Name)
	if err != nil {
		return err
	}

	if exists {
		return errors.Errorf("failed to create cluster. Cluster '%s' already exists", cluster.Name)
	}

	sshPublicKey, err := readSSHPublicKey(s.SSHPublicKeyPath)
	if err != nil {
		return err
	}

	connection, err := connectLibvirt()
	if err != nil {
		return err
	}
	defer connection.Close()

	// check for name conflicts/missing libvirt objects
	if err := validate.LibvirtClusterObjects(connection, cluster); err != nil {
		return err
	}

	if err := download.DownloadPrerequisites(connection, cluster, kubernetesCacheDir()); err != nil {
		return err
	}

	//persist cluster definition before creating any artefacts (libvirt/files on disk/etc.)
	if err := clusterRepository.Save(cluster); err != nil {
		return err
	}

	clusterConfig, err := getClusterConfig(cluster)
	if err != nil {
		return err
	}

	if err := create.Cluster(connection, clusterConfig, cluster, sshPublicKey); err != nil {
		return err
	}

	s.setActiveCluster(clusterRepository, cluster.Name)

	if s.Start {
		return start.Cluster(connection, cluster)
	}

	return nil
}

func (s *createCmdState) createClusterDefinition(clusterName string) (repository.Cluster, error) {
	backingStorageVolume := coreOSStorageVolumeName(s.CoreOSVersion)

	caCertificateBytes, caKeyBytes, err := util.CreateCACertificate(clusterName + "-ca")
	if err != nil {
		return repository.Cluster{}, err
	}

	cluster := repository.Cluster{
		Name:                 clusterName,
		KubernetesVersion:    s.KubernetesVersion,
		CoreOSChannel:        s.CoreOSChannel,
		CoreOSVersion:        s.CoreOSVersion,
		StoragePool:          s.LibvirtStoragePool,
		BackingStorageVolume: backingStorageVolume,
		Network: repository.Network{
			Name:     libvirtNetworkName(clusterName),
			IPv4CIDR: s.IPv4CIDR,
		},
		Nodes:         make(map[string]repository.Node),
		CACertificate: caCertificateBytes,
		CAPrivateKey:  caKeyBytes,
		DNSDomain:     clusterName + ".local",
	}

	addNode := func(name string, isMaster bool) error {
		domainName := libvirtDomainName(clusterName, name)
		dnsName := nodeDNSName(name, cluster.DNSDomain)

		node := repository.Node{
			Name:                 name,
			IsMaster:             isMaster,
			Domain:               domainName,
			StoragePool:          s.LibvirtStoragePool,
			BackingStorageVolume: backingStorageVolume,
			StorageVolume:        libvirtStorageVolumeName(domainName),
			DNSName:              dnsName,
		}

		if isMaster {
			node.CPUs = s.MasterCPUs
			node.MemoryMiB = s.MasterMemory
		} else {
			node.CPUs = s.NondeCPUs
			node.MemoryMiB = s.NodeMemory
		}

		cluster.Nodes[name] = node
		return nil
	}

	masterName := MasterNodeNamePrefix
	if err := addNode(masterName, true); err != nil {
		return repository.Cluster{}, err
	}

	for i := uint(1); i <= s.NodesCount; i++ {
		name := NodeNamePrefix + strconv.FormatUint(uint64(i), 10)
		if err := addNode(name, false); err != nil {
			return repository.Cluster{}, err
		}
	}

	cluster.ServerURL = fmt.Sprintf("https://%s:6443", cluster.Nodes[masterName].DNSName)

	return cluster, nil
}

func (s *createCmdState) setActiveCluster(clusterRepository repository.ClusterRepository, name string) {
	currentActiveCluster, err := clusterRepository.Current()
	if err != nil {
		glog.Error(err)
		return
	}

	if currentActiveCluster != nil {
		// only set the active cluster if none is currently set
		return
	}

	if err := clusterRepository.SetCurrent(name); err != nil {
		glog.Errorf("failed to set current cluster. Error: %v", err)
	}
}
<commit_msg>updated coreos version<commit_after>package cmd

import (
	"fmt"
	"strconv"

	"github.com/bcrusu/kcm/cmd/create"
	"github.com/bcrusu/kcm/cmd/download"
	"github.com/bcrusu/kcm/cmd/start"
	"github.com/bcrusu/kcm/cmd/validate"
	"github.com/bcrusu/kcm/repository"
	"github.com/bcrusu/kcm/util"
	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

const DefaultKubernetesVersion = "1.6.4"
const DefaultCoreOSVersion = "1353.8.0"
const DefaultCoreOsChannel = "stable"

type createCmdState struct {
	KubernetesVersion  string
	CoreOSVersion      string
	CoreOSChannel      string
	LibvirtStoragePool string
	NodesCount         uint
	KubernetesNetwork  string
	SSHPublicKeyPath   string
	Start              bool
	IPv4CIDR           string
	MasterCPUs         uint
	MasterMemory       uint
	NondeCPUs          uint
	NodeMemory         uint
}

func newCreateCmd() *cobra.Command {
	var cmd = &cobra.Command{
		Use:          "create CLUSTER_NAME",
		Short:        "Create a new cluster",
		SilenceUsage: true,
	}

	state := &createCmdState{}

	cmd.PersistentFlags().StringVar(&state.KubernetesVersion, "kubernetes-version", DefaultKubernetesVersion, "Kubernetes version to use")
	cmd.PersistentFlags().StringVar(&state.CoreOSVersion, "coreos-version", DefaultCoreOSVersion, "CoreOS version to use")
	cmd.PersistentFlags().StringVar(&state.CoreOSChannel, "coreos-channel", DefaultCoreOsChannel, "CoreOS release channel: stable, beta, alpha")
	cmd.PersistentFlags().StringVar(&state.LibvirtStoragePool, "libvirt-pool", "default", "Libvirt storage pool")
	cmd.PersistentFlags().UintVar(&state.NodesCount, "node-count", 1, "Initial number of nodes in the cluster") //TODO: set value to 3
	cmd.PersistentFlags().StringVar(&state.KubernetesNetwork, "kubernetes-network", "flannel", "Networking mode to use. Only flannel is suppoted at the moment")
	cmd.PersistentFlags().StringVar(&state.SSHPublicKeyPath, "ssh-public-key", util.GetUserDefaultSSHPublicKeyPath(), "SSH public key to use")
	cmd.PersistentFlags().BoolVarP(&state.Start, "start", "s", false, "Start the cluster immediately")
	cmd.PersistentFlags().StringVar(&state.IPv4CIDR, "ipv4-cidr", "10.1.0.0/16", "Libvirt network IPv4 CIDR. Network 10.2.0.0/16 is reserved for pods/services network")
	cmd.PersistentFlags().UintVar(&state.MasterCPUs, "master-cpu", 1, "Master node allocated CPUs")
	cmd.PersistentFlags().UintVar(&state.MasterMemory, "master-memory", 512, "Master node memory (in MiB)")
	cmd.PersistentFlags().UintVar(&state.NondeCPUs, "node-cpu", 1, "Node allocated CPUs")
	cmd.PersistentFlags().UintVar(&state.NodeMemory, "node-memory", 512, "Node memory (in MiB)")

	cmd.RunE = state.runE
	return cmd
}

func (s *createCmdState) runE(cmd *cobra.Command, args []string) error {
	if len(args) != 1 {
		return errors.New("invalid command arguments")
	}

	cluster, err := s.createClusterDefinition(args[0])
	if err != nil {
		return err
	}

	// lightweight validation - valid names, empty strings, no nodes defined, etc.
	if err := cluster.Validate(); err != nil {
		return err
	}

	clusterRepository, err := newClusterRepository()
	if err != nil {
		return err
	}

	exists, err := clusterRepository.Exists(cluster.Name)
	if err != nil {
		return err
	}

	if exists {
		return errors.Errorf("failed to create cluster. Cluster '%s' already exists", cluster.Name)
	}

	sshPublicKey, err := readSSHPublicKey(s.SSHPublicKeyPath)
	if err != nil {
		return err
	}

	connection, err := connectLibvirt()
	if err != nil {
		return err
	}
	defer connection.Close()

	// check for name conflicts/missing libvirt objects
	if err := validate.LibvirtClusterObjects(connection, cluster); err != nil {
		return err
	}

	if err := download.DownloadPrerequisites(connection, cluster, kubernetesCacheDir()); err != nil {
		return err
	}

	//persist cluster definition before creating any artefacts (libvirt/files on disk/etc.)
	if err := clusterRepository.Save(cluster); err != nil {
		return err
	}

	clusterConfig, err := getClusterConfig(cluster)
	if err != nil {
		return err
	}

	if err := create.Cluster(connection, clusterConfig, cluster, sshPublicKey); err != nil {
		return err
	}

	s.setActiveCluster(clusterRepository, cluster.Name)

	if s.Start {
		return start.Cluster(connection, cluster)
	}

	return nil
}

func (s *createCmdState) createClusterDefinition(clusterName string) (repository.Cluster, error) {
	backingStorageVolume := coreOSStorageVolumeName(s.CoreOSVersion)

	caCertificateBytes, caKeyBytes, err := util.CreateCACertificate(clusterName + "-ca")
	if err != nil {
		return repository.Cluster{}, err
	}

	cluster := repository.Cluster{
		Name:                 clusterName,
		KubernetesVersion:    s.KubernetesVersion,
		CoreOSChannel:        s.CoreOSChannel,
		CoreOSVersion:        s.CoreOSVersion,
		StoragePool:          s.LibvirtStoragePool,
		BackingStorageVolume: backingStorageVolume,
		Network: repository.Network{
			Name:     libvirtNetworkName(clusterName),
			IPv4CIDR: s.IPv4CIDR,
		},
		Nodes:         make(map[string]repository.Node),
		CACertificate: caCertificateBytes,
		CAPrivateKey:  caKeyBytes,
		DNSDomain:     clusterName + ".local",
	}

	addNode := func(name string, isMaster bool) error {
		domainName := libvirtDomainName(clusterName, name)
		dnsName := nodeDNSName(name, cluster.DNSDomain)

		node := repository.Node{
			Name:                 name,
			IsMaster:             isMaster,
			Domain:               domainName,
			StoragePool:          s.LibvirtStoragePool,
			BackingStorageVolume: backingStorageVolume,
			StorageVolume:        libvirtStorageVolumeName(domainName),
			DNSName:              dnsName,
		}

		if isMaster {
			node.CPUs = s.MasterCPUs
			node.MemoryMiB = s.MasterMemory
		} else {
			node.CPUs = s.NondeCPUs
			node.MemoryMiB = s.NodeMemory
		}

		cluster.Nodes[name] = node
		return nil
	}

	masterName := MasterNodeNamePrefix
	if err := addNode(masterName, true); err != nil {
		return repository.Cluster{}, err
	}

	for i := uint(1); i <= s.NodesCount; i++ {
		name := NodeNamePrefix + strconv.FormatUint(uint64(i), 10)
		if err := addNode(name, false); err != nil {
			return repository.Cluster{}, err
		}
	}

	cluster.ServerURL = fmt.Sprintf("https://%s:6443", cluster.Nodes[masterName].DNSName)

	return cluster, nil
}

func (s *createCmdState) setActiveCluster(clusterRepository repository.ClusterRepository, name string) {
	currentActiveCluster, err := clusterRepository.Current()
	if err != nil {
		glog.Error(err)
		return
	}

	if currentActiveCluster != nil {
		// only set the active cluster if none is currently set
		return
	}

	if err := clusterRepository.SetCurrent(name); err != nil {
		glog.Errorf("failed to set current cluster. Error: %v", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/summerwind/h2spec"
	"os"
	"strings"
	"time"
)

const VERSION = "v0.0.9"

type sections []string

func (s *sections) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *sections) Set(v string) error {
	*s = append(*s, v)
	return nil
}

func main() {
	port := flag.Int("p", 0, "Target port.")
	host := flag.String("h", "127.0.0.1", "Target host.")
	useTls := flag.Bool("t", false, "Connect over TLS.")
	insecureSkipVerify := flag.Bool("k", false, "Don't verify server's certificate.")
	timeout := flag.Int("o", 2, "Maximum time allowed for test.")
	strict := flag.Bool("S", false, "Strict mode.")
	version := flag.Bool("version", false, "Display version information and exit.")

	var sectionFlag sections
	flag.Var(&sectionFlag, "s", "Section number on which to run the test")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS]\n\n", os.Args[0])
		fmt.Println("Options:")
		fmt.Println("  -p:        Target port. (Default: 80 or 443)")
		fmt.Println("  -h:        Target host. (Default: 127.0.0.1)")
		fmt.Println("  -t:        Connect over TLS. (Default: false)")
		fmt.Println("  -k:        Don't verify server's certificate. (Default: false)")
		fmt.Println("  -o:        Maximum time allowed for test. (Default: 2)")
		fmt.Println("  -s:        Section number on which to run the test. (Example: -s 6.1 -s 6.2)")
		fmt.Println("  -S:        Run the test cases marked as \"strict\".")
		fmt.Println("  --version: Display version information and exit.")
		fmt.Println("  --help:    Display this help and exit.")
		os.Exit(1)
	}

	flag.Parse()

	if *version {
		fmt.Fprintf(os.Stderr, "h2spec %s\n", VERSION)
		os.Exit(0)
	}

	if *port == 0 {
		if *useTls {
			*port = 443
		} else {
			*port = 80
		}
	}

	var ctx h2spec.Context
	ctx.Port = *port
	ctx.Host = *host
	ctx.Timeout = time.Duration(*timeout) * time.Second
	ctx.Strict = *strict
	ctx.Tls = *useTls
	ctx.TlsConfig = &tls.Config{
		InsecureSkipVerify: *insecureSkipVerify,
	}

	if len(sectionFlag) > 0 {
		ctx.Sections = map[string]bool{}
		for _, sec := range sectionFlag {
			splitedSec := strings.Split(sec, ".")
			lastIndex := len(splitedSec) - 1

			num := []string{}
			for i, sec := range splitedSec {
				num = append(num, sec)
				if i != 0 {
					key := strings.Join(num, ".")

					runAll := false
					if i == lastIndex || ctx.Sections[key] {
						runAll = true
					}

					ctx.Sections[key] = runAll
				}
			}
		}
	}

	h2spec.Run(&ctx)
}
<commit_msg>Marked as "v1.0.0"<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/summerwind/h2spec"
	"os"
	"strings"
	"time"
)

const VERSION = "v1.0.0"

type sections []string

func (s *sections) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *sections) Set(v string) error {
	*s = append(*s, v)
	return nil
}

func main() {
	port := flag.Int("p", 0, "Target port.")
	host := flag.String("h", "127.0.0.1", "Target host.")
	useTls := flag.Bool("t", false, "Connect over TLS.")
	insecureSkipVerify := flag.Bool("k", false, "Don't verify server's certificate.")
	timeout := flag.Int("o", 2, "Maximum time allowed for test.")
	strict := flag.Bool("S", false, "Strict mode.")
	version := flag.Bool("version", false, "Display version information and exit.")

	var sectionFlag sections
	flag.Var(&sectionFlag, "s", "Section number on which to run the test")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS]\n\n", os.Args[0])
		fmt.Println("Options:")
		fmt.Println("  -p:        Target port. (Default: 80 or 443)")
		fmt.Println("  -h:        Target host. (Default: 127.0.0.1)")
		fmt.Println("  -t:        Connect over TLS. (Default: false)")
		fmt.Println("  -k:        Don't verify server's certificate. (Default: false)")
		fmt.Println("  -o:        Maximum time allowed for test. (Default: 2)")
		fmt.Println("  -s:        Section number on which to run the test. (Example: -s 6.1 -s 6.2)")
		fmt.Println("  -S:        Run the test cases marked as \"strict\".")
		fmt.Println("  --version: Display version information and exit.")
		fmt.Println("  --help:    Display this help and exit.")
		os.Exit(1)
	}

	flag.Parse()

	if *version {
		fmt.Fprintf(os.Stderr, "h2spec %s\n", VERSION)
		os.Exit(0)
	}

	if *port == 0 {
		if *useTls {
			*port = 443
		} else {
			*port = 80
		}
	}

	var ctx h2spec.Context
	ctx.Port = *port
	ctx.Host = *host
	ctx.Timeout = time.Duration(*timeout) * time.Second
	ctx.Strict = *strict
	ctx.Tls = *useTls
	ctx.TlsConfig = &tls.Config{
		InsecureSkipVerify: *insecureSkipVerify,
	}

	if len(sectionFlag) > 0 {
		ctx.Sections = map[string]bool{}
		for _, sec := range sectionFlag {
			splitedSec := strings.Split(sec, ".")
			lastIndex := len(splitedSec) - 1

			num := []string{}
			for i, sec := range splitedSec {
				num = append(num, sec)
				if i != 0 {
					key := strings.Join(num, ".")

					runAll := false
					if i == lastIndex || ctx.Sections[key] {
						runAll = true
					}

					ctx.Sections[key] = runAll
				}
			}
		}
	}

	h2spec.Run(&ctx)
}
<|endoftext|>


<commit_before>// Copyright © 2017 Douglas Chimento <dchimento@gmail.com>
//
// 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 cmd

import (
	"database/sql"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"sync"
	"time"
)

var eventChan = make(chan *Event)
var geoCache *Cache

type mockGeoClient struct {
}

var mutex = &sync.Mutex{}
var geoClient = geoClientTransporter(defaultGeoClient())
var geoPool = sync.Pool{
	New: func() interface{} {
		return &Geo{}
	},
}

func insertGeo(geo *Geo, db *sql.DB) (int64, error) {
	var id int64
	r, err := db.Query(`INSERT INTO geo
	(ip, country_code, region_code, region_name, city, time_zone, latitude, longitude, metro_code, last_update)
	VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING ID`,
		geo.IP, geo.CountryCode, geo.RegionCode, geo.RegionName, geo.City, geo.TimeZone, geo.Latitude, geo.Longitude, geo.MetroCode, geo.LastUpdate)
	if err != nil {
		return 0, err
	}
	if !r.Next() {
		return 0, errors.New(fmt.Sprintf("Failed inserting %s", geo))
	}
	err = r.Scan(&id)
	return id, err
}

func (c *eventClient) resolveAddr(addr string) (int64, error) {
	if config.UseCache {
		cachedGeo, found := geoCache.Get(addr)
		if found {
			return cachedGeo, nil
		}
	}
	mutex.Lock()
	defer mutex.Unlock()
	geo := geoPool.Get().(*Geo)
	defer geoPool.Put(geo)
	expire := time.Now().AddDate(0, -1, 0)
	r := c.db.QueryRow(`SELECT
	id, ip, country_code, region_code, region_name, city, time_zone, latitude, longitude, metro_code, last_update
	FROM geo
	WHERE ip = $1
	ORDER BY last_update DESC LIMIT 1`, addr)
	err := r.Scan(&geo.ID, &geo.IP, &geo.CountryCode, &geo.RegionCode, &geo.RegionName, &geo.City, &geo.TimeZone, &geo.Latitude, &geo.Longitude, &geo.MetroCode, &geo.LastUpdate)
	if err != nil && err != sql.ErrNoRows {
		return 0, err
	}
	if err == sql.ErrNoRows {
		log.Infof("New addr found %s", addr)
		geo, err = c.geoClient.getLocationForAddr(addr)
		if err != nil {
			return 0, err
		}
		id, err := insertGeo(geo, c.db)
		if err != nil {
			return 0, err
		}
		geo.ID = id
	} else if geo.LastUpdate.Before(expire) {
		log.Infof("Found expired addr %s (%s) (%s)\n", addr, geo.LastUpdate, expire)
		var newGeo = &Geo{}
		newGeo, err = c.geoClient.getLocationForAddr(addr)
		if err != nil {
			return 0, err
		}
		if geo.equals(newGeo) {
			log.Infof("Updating last_update for id %d ", geo.ID)
			if _, err = c.db.Exec("UPDATE geo SET last_update = now() WHERE id  = $1", geo.ID); err != nil {
				return 0, err
			}
		} else {
			log.Infof("Inserting new record for id %d ", geo.ID)
			id, err := insertGeo(newGeo, c.db)
			if err != nil {
				return 0, err
			}
			geo.ID = id
		}
	}
	geoCache.Set(addr, geo.ID)
	return geo.ID, nil
}

func runLookup() {
	log.Infof("Initalize lookup channel")
	for {
		select {
		case event := <-eventChan:
			go defaultEventClient.resolveGeoEvent(*event)
		}
	}
}

func init() {
	for i := 0; i < 100; i++ {
		geoPool.Put(&Geo{})
	}
	geoCache = NewCache()
}
<commit_msg>closing insert<commit_after>// Copyright © 2017 Douglas Chimento <dchimento@gmail.com>
//
// 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 cmd

import (
	"database/sql"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"sync"
	"time"
)

var eventChan = make(chan *Event)
var geoCache *Cache

type mockGeoClient struct {
}

var mutex = &sync.Mutex{}
var geoClient = geoClientTransporter(defaultGeoClient())
var geoPool = sync.Pool{
	New: func() interface{} {
		return &Geo{}
	},
}

func insertGeo(geo *Geo, db *sql.DB) (int64, error) {
	var id int64
	r, err := db.Query(`INSERT INTO geo
	(ip, country_code, region_code, region_name, city, time_zone, latitude, longitude, metro_code, last_update)
	VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING ID`,
		geo.IP, geo.CountryCode, geo.RegionCode, geo.RegionName, geo.City, geo.TimeZone, geo.Latitude, geo.Longitude, geo.MetroCode, geo.LastUpdate)
	if err != nil {
		return 0, err
	}
	defer r.Close()
	if !r.Next() {
		return 0, errors.New(fmt.Sprintf("Failed inserting %s", geo))
	}
	err = r.Scan(&id)

	return id, err
}

func (c *eventClient) resolveAddr(addr string) (int64, error) {
	if config.UseCache {
		cachedGeo, found := geoCache.Get(addr)
		if found {
			return cachedGeo, nil
		}
	}
	mutex.Lock()
	defer mutex.Unlock()
	geo := geoPool.Get().(*Geo)
	defer geoPool.Put(geo)
	expire := time.Now().AddDate(0, -1, 0)
	r := c.db.QueryRow(`SELECT
	id, ip, country_code, region_code, region_name, city, time_zone, latitude, longitude, metro_code, last_update
	FROM geo
	WHERE ip = $1
	ORDER BY last_update DESC LIMIT 1`, addr)
	err := r.Scan(&geo.ID, &geo.IP, &geo.CountryCode, &geo.RegionCode, &geo.RegionName, &geo.City, &geo.TimeZone, &geo.Latitude, &geo.Longitude, &geo.MetroCode, &geo.LastUpdate)
	if err != nil && err != sql.ErrNoRows {
		return 0, err
	}
	if err == sql.ErrNoRows {
		log.Infof("New addr found %s", addr)
		geo, err = c.geoClient.getLocationForAddr(addr)
		if err != nil {
			return 0, err
		}
		id, err := insertGeo(geo, c.db)
		if err != nil {
			return 0, err
		}
		geo.ID = id
	} else if geo.LastUpdate.Before(expire) {
		log.Infof("Found expired addr %s (%s) (%s)", addr, geo.LastUpdate, expire)
		var newGeo = &Geo{}
		newGeo, err = c.geoClient.getLocationForAddr(addr)
		if err != nil {
			return 0, err
		}
		if geo.equals(newGeo) {
			log.Infof("Updating last_update for id %d ", geo.ID)
			if _, err = c.db.Exec("UPDATE geo SET last_update = now() WHERE id  = $1", geo.ID); err != nil {
				return 0, err
			}
		} else {
			log.Infof("Inserting new record for id %d ", geo.ID)
			id, err := insertGeo(newGeo, c.db)
			if err != nil {
				return 0, err
			}
			geo.ID = id
		}
	}
	geoCache.Set(addr, geo.ID)
	return geo.ID, nil
}

func runLookup() {
	log.Infof("Initalize lookup channel")
	for {
		select {
		case event := <-eventChan:
			go defaultEventClient.resolveGeoEvent(*event)
		}
	}
}

func init() {
	for i := 0; i < 100; i++ {
		geoPool.Put(&Geo{})
	}
	geoCache = NewCache()
}
<|endoftext|>


<commit_before>// Copyright © 2017 The VirusTotal CLI authors. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// 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 cmd

import (
	"encoding/hex"
	"fmt"
	"strings"

	"github.com/spf13/viper"

	"github.com/VirusTotal/vt-cli/utils"
	vt "github.com/VirusTotal/vt-go"
	"github.com/spf13/cobra"
)

func preRunSearchCmd(c *cobra.Command, args []string) error {

	if viper.GetBool("download") {
		for _, flag := range []string{"include", "identifiers-only"} {
			if c.Flag(flag).Changed {
				return fmt.Errorf("--%s can't be used with --download", flag)
			}
		}
	} else {
		for _, flag := range []string{"output", "threads"} {
			if c.Flag(flag).Changed {
				return fmt.Errorf("--%s must be used with --download", flag)
			}
		}
	}

	return nil
}

func runSearchCmd(cmd *cobra.Command, args []string) error {

	batchSize := 25
	if viper.GetInt("limit") < batchSize {
		batchSize = viper.GetInt("limit")
	}

	client, err := NewAPIClient()
	if err != nil {
		return err
	}

	it, err := client.Search(args[0],
		vt.WithLimit(viper.GetInt("limit")),
		vt.WithCursor(viper.GetString("cursor")),
		vt.WithBatchSize(batchSize),
		vt.WithDescriptorsOnly(
			viper.GetBool("identifiers-only") || viper.GetBool("download")))

	if err != nil {
		return err
	}

	if viper.GetBool("download") {
		ch := make(chan interface{})
		go func() {
			for it.Next() {
				obj := it.Get()
				ch <- obj.ID
			}
			close(ch)
		}()
		c := utils.NewCoordinator(viper.GetInt("threads"))
		c.DoWithItemsFromChannel(&downloader{client: client}, ch)
		return it.Error()
	}

	p, err := NewObjectPrinter(cmd)
	if err != nil {
		return err
	}
	return p.PrintIter(it)
}

var cmdSearchHelp = `Search for files using VirusTotal Intelligence's query language.`

var cmdSearchExample = `  vt search eicar
  vt search "foobar p:1+"`

// NewSearchCmd returns a new instance of the 'search' command.
func NewSearchCmd() *cobra.Command {

	cmd := &cobra.Command{
		Args:    cobra.ExactArgs(1),
		Use:     "search [query]",
		Short:   "Search for files in VirusTotal Intelligence",
		Long:    cmdSearchHelp,
		Example: cmdSearchExample,
		PreRunE: preRunSearchCmd,
		RunE:    runSearchCmd,
	}

	cmd.Flags().BoolP("download", "d", false, "download files")

	addIDOnlyFlag(cmd.Flags())
	addIncludeExcludeFlags(cmd.Flags())
	addThreadsFlag(cmd.Flags())
	addLimitFlag(cmd.Flags())
	addCursorFlag(cmd.Flags())
	addOutputFlag(cmd.Flags())

	cmd.AddCommand(NewContentSearchCmd())

	return cmd
}

type matchPrinter struct {
	client *utils.APIClient
	idOnly bool
}

func (m *matchPrinter) Do(fileObj interface{}, ds *utils.DoerState) string {
	f := fileObj.(*vt.Object)
	var line string
	if m.idOnly {
		line = f.ID
	} else {
		var s string
		confidence, _ := f.GetContextAttributeFloat64("confidence")
		snippetID, _ := f.GetContextAttributeString("snippet")
		inSubFile, _ := f.GetContextAttributeString("match_in_subfile")
		snippets := make([]string, 0)
		_, err := m.client.GetData(vt.URL("intelligence/search/snippets/%s", snippetID), &snippets)
		if err == nil {
			s = strings.Join(snippets, "\n\n")
			s = strings.Replace(s, "\x1c", "\033[1m", -1)
			s = strings.Replace(s, "\x1d", "\033[0m", -1)
		} else {
			s = "<no snippet available>"
		}
		line = fmt.Sprintf(
			"%s\n\nsha256  : %s\nscore   : %03.1f \nsubfile : %v\n\n%s\n",
			strings.Repeat("_", 76), f.ID, confidence, inSubFile, s)
	}
	return line
}

func getIgnoredSubstrings(meta map[string]interface{}) []string {
	if i, ok := meta["ignored_substrings"]; ok {
		ii := i.([]interface{})
		ss := make([]string, len(ii))
		for i := range ss {
			s := ii[i].(string)
			h, _ := hex.DecodeString(s)
			ss[i] = fmt.Sprintf("{% X} %q", h, h)
		}
		return ss
	}
	return nil
}

func runContentSearchCmd(cmd *cobra.Command, args []string) error {

	batchSize := 10
	if viper.GetInt("limit") < batchSize {
		batchSize = viper.GetInt("limit")
	}

	client, err := NewAPIClient()
	if err != nil {
		return err
	}

	terms := make([]string, len(args))
	for i, arg := range args {
		terms[i] = fmt.Sprintf("content:%s", arg)
	}

	it, err := client.Search(strings.Join(terms, " "),
		vt.WithLimit(viper.GetInt("limit")),
		vt.WithCursor(viper.GetString("cursor")),
		vt.WithBatchSize(batchSize),
		vt.WithDescriptorsOnly(
			viper.GetBool("identifiers-only") || viper.GetBool("download")))

	if err != nil {
		return err
	}

	c := utils.NewCoordinator(viper.GetInt("threads"))

	var doer utils.Doer
	if viper.GetBool("download") {
		doer = &downloader{client: client}
	} else {
		doer = &matchPrinter{client, viper.GetBool("identifiers-only")}
		c.EnableSpinner()
	}

	c.DoWithObjectsFromIterator(doer, it, batchSize)

	if ignored := getIgnoredSubstrings(it.Meta()); ignored != nil {
		colorScheme.CommentColor.Printf(
			"IGNORED SUBSTRINGS:\n%s\n",
			strings.Join(ignored, "\n"))
	}

	PrintCommandLineWithCursor(cmd, it)
	return it.Error()
}

var cmdContentSearchHelp = `Search for content within files in VirusTotal`

var cmdContentSearchExample = `  vt search content '{cafebabe}'
  vt search content '{70 6C 75 73 76 69 63 [1] 79 61 72 61}'
  vt search content '/virustotal(.org|.com)/'`

// NewContentSearchCmd returns a new instance of the 'search content' command.
func NewContentSearchCmd() *cobra.Command {

	cmd := &cobra.Command{
		Args:    cobra.MinimumNArgs(1),
		Use:     "content [query]",
		Short:   "Search for patterns within files in VirusTotal Intelligence",
		Long:    cmdContentSearchHelp,
		Example: cmdContentSearchExample,
		RunE:    runContentSearchCmd,
	}

	cmd.Flags().BoolP("download", "d", false, "download files")
	cmd.Flags().BoolP("exact-matches-only", "e", false, "exact matches only")

	addThreadsFlag(cmd.Flags())
	addIDOnlyFlag(cmd.Flags())
	addLimitFlag(cmd.Flags())
	addCursorFlag(cmd.Flags())

	return cmd
}
<commit_msg>Don't decode ignored substrings in content searches, as they now come ready to be printed.<commit_after>// Copyright © 2017 The VirusTotal CLI authors. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// 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 cmd

import (
	"fmt"
	"strings"

	"github.com/spf13/viper"

	"github.com/VirusTotal/vt-cli/utils"
	vt "github.com/VirusTotal/vt-go"
	"github.com/spf13/cobra"
)

func preRunSearchCmd(c *cobra.Command, args []string) error {

	if viper.GetBool("download") {
		for _, flag := range []string{"include", "identifiers-only"} {
			if c.Flag(flag).Changed {
				return fmt.Errorf("--%s can't be used with --download", flag)
			}
		}
	} else {
		for _, flag := range []string{"output", "threads"} {
			if c.Flag(flag).Changed {
				return fmt.Errorf("--%s must be used with --download", flag)
			}
		}
	}

	return nil
}

func runSearchCmd(cmd *cobra.Command, args []string) error {

	batchSize := 25
	if viper.GetInt("limit") < batchSize {
		batchSize = viper.GetInt("limit")
	}

	client, err := NewAPIClient()
	if err != nil {
		return err
	}

	it, err := client.Search(args[0],
		vt.WithLimit(viper.GetInt("limit")),
		vt.WithCursor(viper.GetString("cursor")),
		vt.WithBatchSize(batchSize),
		vt.WithDescriptorsOnly(
			viper.GetBool("identifiers-only") || viper.GetBool("download")))

	if err != nil {
		return err
	}

	if viper.GetBool("download") {
		ch := make(chan interface{})
		go func() {
			for it.Next() {
				obj := it.Get()
				ch <- obj.ID
			}
			close(ch)
		}()
		c := utils.NewCoordinator(viper.GetInt("threads"))
		c.DoWithItemsFromChannel(&downloader{client: client}, ch)
		return it.Error()
	}

	p, err := NewObjectPrinter(cmd)
	if err != nil {
		return err
	}
	return p.PrintIter(it)
}

var cmdSearchHelp = `Search for files using VirusTotal Intelligence's query language.`

var cmdSearchExample = `  vt search eicar
  vt search "foobar p:1+"`

// NewSearchCmd returns a new instance of the 'search' command.
func NewSearchCmd() *cobra.Command {

	cmd := &cobra.Command{
		Args:    cobra.ExactArgs(1),
		Use:     "search [query]",
		Short:   "Search for files in VirusTotal Intelligence",
		Long:    cmdSearchHelp,
		Example: cmdSearchExample,
		PreRunE: preRunSearchCmd,
		RunE:    runSearchCmd,
	}

	cmd.Flags().BoolP("download", "d", false, "download files")

	addIDOnlyFlag(cmd.Flags())
	addIncludeExcludeFlags(cmd.Flags())
	addThreadsFlag(cmd.Flags())
	addLimitFlag(cmd.Flags())
	addCursorFlag(cmd.Flags())
	addOutputFlag(cmd.Flags())

	cmd.AddCommand(NewContentSearchCmd())

	return cmd
}

type matchPrinter struct {
	client *utils.APIClient
	idOnly bool
}

func (m *matchPrinter) Do(fileObj interface{}, ds *utils.DoerState) string {
	f := fileObj.(*vt.Object)
	var line string
	if m.idOnly {
		line = f.ID
	} else {
		var s string
		confidence, _ := f.GetContextAttributeFloat64("confidence")
		snippetID, _ := f.GetContextAttributeString("snippet")
		inSubFile, _ := f.GetContextAttributeString("match_in_subfile")
		snippets := make([]string, 0)
		_, err := m.client.GetData(vt.URL("intelligence/search/snippets/%s", snippetID), &snippets)
		if err == nil {
			s = strings.Join(snippets, "\n\n")
			s = strings.Replace(s, "\x1c", "\033[1m", -1)
			s = strings.Replace(s, "\x1d", "\033[0m", -1)
		} else {
			s = "<no snippet available>"
		}
		line = fmt.Sprintf(
			"%s\n\nsha256  : %s\nscore   : %03.1f \nsubfile : %v\n\n%s\n",
			strings.Repeat("_", 76), f.ID, confidence, inSubFile, s)
	}
	return line
}

func getIgnoredSubstrings(meta map[string]interface{}) []string {
	if i, ok := meta["ignored_substrings"]; ok {
		ii := i.([]interface{})
		ss := make([]string, len(ii))
		for i := range ss {
			ss[i] = ii[i].(string)
		}
		return ss
	}
	return nil
}

func runContentSearchCmd(cmd *cobra.Command, args []string) error {

	batchSize := 10
	if viper.GetInt("limit") < batchSize {
		batchSize = viper.GetInt("limit")
	}

	client, err := NewAPIClient()
	if err != nil {
		return err
	}

	terms := make([]string, len(args))
	for i, arg := range args {
		terms[i] = fmt.Sprintf("content:%s", arg)
	}

	it, err := client.Search(strings.Join(terms, " "),
		vt.WithLimit(viper.GetInt("limit")),
		vt.WithCursor(viper.GetString("cursor")),
		vt.WithBatchSize(batchSize),
		vt.WithDescriptorsOnly(
			viper.GetBool("identifiers-only") || viper.GetBool("download")))

	if err != nil {
		return err
	}

	c := utils.NewCoordinator(viper.GetInt("threads"))

	var doer utils.Doer
	if viper.GetBool("download") {
		doer = &downloader{client: client}
	} else {
		doer = &matchPrinter{client, viper.GetBool("identifiers-only")}
		c.EnableSpinner()
	}

	c.DoWithObjectsFromIterator(doer, it, batchSize)

	if ignored := getIgnoredSubstrings(it.Meta()); ignored != nil {
		colorScheme.CommentColor.Printf(
			"IGNORED SUBSTRINGS:\n%s\n",
			strings.Join(ignored, "\n"))
	}

	PrintCommandLineWithCursor(cmd, it)
	return it.Error()
}

var cmdContentSearchHelp = `Search for content within files in VirusTotal`

var cmdContentSearchExample = `  vt search content '{cafebabe}'
  vt search content '{70 6C 75 73 76 69 63 [1] 79 61 72 61}'
  vt search content '/virustotal(.org|.com)/'`

// NewContentSearchCmd returns a new instance of the 'search content' command.
func NewContentSearchCmd() *cobra.Command {

	cmd := &cobra.Command{
		Args:    cobra.MinimumNArgs(1),
		Use:     "content [query]",
		Short:   "Search for patterns within files in VirusTotal Intelligence",
		Long:    cmdContentSearchHelp,
		Example: cmdContentSearchExample,
		RunE:    runContentSearchCmd,
	}

	cmd.Flags().BoolP("download", "d", false, "download files")
	cmd.Flags().BoolP("exact-matches-only", "e", false, "exact matches only")

	addThreadsFlag(cmd.Flags())
	addIDOnlyFlag(cmd.Flags())
	addLimitFlag(cmd.Flags())
	addCursorFlag(cmd.Flags())

	return cmd
}
<|endoftext|>


<commit_before>// Copyright © 2017 Douglas Chimento <dchimento@gmail.com>
//
// 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 cmd

import (
	"github.com/spf13/cobra"

	log "github.com/Sirupsen/logrus"
	"time"

	"bytes"
	"fmt"
	"github.com/dougEfresh/passwd-pot/api"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"net/http"
	"os"
	"strings"
)

// streamCmd represents the stream command
const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var (
	newline = []byte{'\n'}
	space   = []byte{' '}
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

var streamCmd = &cobra.Command{
	Use:   "stream",
	Short: "Stream events to websocket clients",
	Long:  "",
	Run: func(cmd *cobra.Command, args []string) {
		var err error
		setup(cmd, args)
		r := mux.NewRouter()
		r.HandleFunc(getHandler(api.StreamURL, streamEvents)).Methods("GET")
		r.HandleFunc(getHandler(api.StreamURL+"/random", streamEvents)).Methods("GET")

		srv := &http.Server{
			Handler:      r,
			Addr:         config.BindAddr,
			WriteTimeout: 10 * time.Second,
			ReadTimeout:  10 * time.Second,
		}

		log.Infof("Listing on %s", config.BindAddr)
		//websocket requests
		go hub.run()
		go randomDataHub.run()
		go startRandomHub(randomDataHub)
		err = srv.ListenAndServe()
		if err != nil {
			log.Errorf("Caught error %s", err)
			os.Exit(-1)
		}
	},
}

// Client is a middleman between the websocket connection and the hub.
type Client struct {
	hub *Hub

	// The websocket connection.
	conn *websocket.Conn

	// Buffered channel of outbound messages.
	send chan []byte
}

// readPump pumps messages from the websocket connection to the hub.
//
// The application runs readPump in a per-connection goroutine. The application
// ensures that there is at most one reader on a connection by executing all
// reads from this goroutine.
func (c *Client) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
	}()
	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Printf("error: %v", err)
			}
			break
		}
		message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
		c.hub.broadcast <- message
	}
}

// writePump pumps messages from the hub to the websocket connection.
//
// A goroutine running writePump is started for each connection. The
// application ensures that there is at most one writer to a connection by
// executing all writes from this goroutine.
func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				// The hub closed the channel.
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			// Add queued chat messages to the current websocket message.
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write(newline)
				w.Write(<-c.send)
			}

			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func streamEvents(w http.ResponseWriter, r *http.Request) {
	if strings.Contains(r.URL.Path, "random") {
		serveRandomWs(randomDataHub, w, r)
	} else {
		serveWs(hub, w, r)
	}
}

// serveWs handles websocket requests from the peer.
func serveRandomWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error(err)
		return
	}
	client := &Client{hub: hub, conn: conn, send: make(chan []byte, 512)}
	client.hub.register <- client
	go client.writePump()
	client.readPump()
}

// serveWs handles websocket requests from the peer.
func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error(err)
		return
	}
	client := &Client{hub: hub, conn: conn, send: make(chan []byte, 512)}
	client.hub.register <- client
	go client.writePump()
	client.readPump()
}

var lastRandomEvent *EventGeo

func startRandomHub(hub *Hub) {
	log.Info("Starting random hub")

	var id int64
	query := fmt.Sprintf("SELECT id FROM %s WHERE id != $1 and remote_latitude != $2 and remote_longitude != ? and id >= (SELECT max(id) * RANDOM() FROM %s) ORDER BY id LIMIT 1 ", eventGeoTable, eventTable)
	for {
		time.Sleep(1 * time.Second)
		if len(hub.clients) == 0 {
			continue
		}

		if lastRandomEvent == nil {
			r := defaultEventClient.db.QueryRow("SELECT max(id) FROM event_geo")
			err := r.Scan(&id)
			if err != nil {
				log.Error("Error getting max id")
				continue
			}
		} else {
			r := defaultEventClient.db.QueryRow(query, lastRandomEvent.ID, lastRandomEvent.RemoteLatitude, lastRandomEvent.RemoteLongitude)
			r.Scan(id)
		}
		if id == 0 {
			log.Error("Could not find an random event!")
		}
		if id != 0 {
			lastRandomEvent = defaultEventClient.broadcastEvent(id, hub)
		}
	}
}

func init() {
	RootCmd.AddCommand(streamCmd)
	streamCmd.PersistentFlags().StringVar(&config.Dsn, "dsn", "postgres://postgres:@172.17.0.1/?sslmode=disable", "DSN database url")
	streamCmd.PersistentFlags().StringVar(&config.BindAddr, "bind", "localhost:8080", "bind to this address:port")
	streamCmd.PersistentFlags().StringVar(&config.NewRelic, "new-relic", "", "new relic api key")
}
<commit_msg>fixing scan bug<commit_after>// Copyright © 2017 Douglas Chimento <dchimento@gmail.com>
//
// 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 cmd

import (
	"github.com/spf13/cobra"

	log "github.com/Sirupsen/logrus"
	"time"

	"bytes"
	"fmt"
	"github.com/dougEfresh/passwd-pot/api"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"net/http"
	"os"
	"strings"
)

// streamCmd represents the stream command
const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var (
	newline = []byte{'\n'}
	space   = []byte{' '}
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

var streamCmd = &cobra.Command{
	Use:   "stream",
	Short: "Stream events to websocket clients",
	Long:  "",
	Run: func(cmd *cobra.Command, args []string) {
		var err error
		setup(cmd, args)
		r := mux.NewRouter()
		r.HandleFunc(getHandler(api.StreamURL, streamEvents)).Methods("GET")
		r.HandleFunc(getHandler(api.StreamURL+"/random", streamEvents)).Methods("GET")

		srv := &http.Server{
			Handler:      r,
			Addr:         config.BindAddr,
			WriteTimeout: 10 * time.Second,
			ReadTimeout:  10 * time.Second,
		}

		log.Infof("Listing on %s", config.BindAddr)
		//websocket requests
		go hub.run()
		go randomDataHub.run()
		go startRandomHub(randomDataHub)
		err = srv.ListenAndServe()
		if err != nil {
			log.Errorf("Caught error %s", err)
			os.Exit(-1)
		}
	},
}

// Client is a middleman between the websocket connection and the hub.
type Client struct {
	hub *Hub

	// The websocket connection.
	conn *websocket.Conn

	// Buffered channel of outbound messages.
	send chan []byte
}

// readPump pumps messages from the websocket connection to the hub.
//
// The application runs readPump in a per-connection goroutine. The application
// ensures that there is at most one reader on a connection by executing all
// reads from this goroutine.
func (c *Client) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
	}()
	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Printf("error: %v", err)
			}
			break
		}
		message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
		c.hub.broadcast <- message
	}
}

// writePump pumps messages from the hub to the websocket connection.
//
// A goroutine running writePump is started for each connection. The
// application ensures that there is at most one writer to a connection by
// executing all writes from this goroutine.
func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				// The hub closed the channel.
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			// Add queued chat messages to the current websocket message.
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write(newline)
				w.Write(<-c.send)
			}

			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func streamEvents(w http.ResponseWriter, r *http.Request) {
	if strings.Contains(r.URL.Path, "random") {
		serveRandomWs(randomDataHub, w, r)
	} else {
		serveWs(hub, w, r)
	}
}

// serveWs handles websocket requests from the peer.
func serveRandomWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error(err)
		return
	}
	client := &Client{hub: hub, conn: conn, send: make(chan []byte, 512)}
	client.hub.register <- client
	go client.writePump()
	client.readPump()
}

// serveWs handles websocket requests from the peer.
func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error(err)
		return
	}
	client := &Client{hub: hub, conn: conn, send: make(chan []byte, 512)}
	client.hub.register <- client
	go client.writePump()
	client.readPump()
}

var lastRandomEvent *EventGeo

func startRandomHub(hub *Hub) {
	log.Info("Starting random hub")

	var id int64
	query := fmt.Sprintf("SELECT id FROM %s WHERE id != $1 and remote_latitude != $2 and remote_longitude != ? and id >= (SELECT max(id) * RANDOM() FROM %s) ORDER BY id LIMIT 1 ", eventGeoTable, eventTable)
	for {
		time.Sleep(1250 * time.Millisecond)
		if len(hub.clients) == 0 {
			continue
		}

		if lastRandomEvent == nil {
			r := defaultEventClient.db.QueryRow("SELECT max(id) FROM event_geo")
			err := r.Scan(&id)
			if err != nil {
				log.Error("Error getting max id")
				continue
			}
		} else {
			r := defaultEventClient.db.QueryRow(query, lastRandomEvent.ID, lastRandomEvent.RemoteLatitude, lastRandomEvent.RemoteLongitude)
			r.Scan(&id)
		}
		if id == 0 {
			log.Error("Could not find an random event!")
		}
		if id != 0 {
			lastRandomEvent = defaultEventClient.broadcastEvent(id, hub)
		}
	}
}

func init() {
	RootCmd.AddCommand(streamCmd)
	streamCmd.PersistentFlags().StringVar(&config.Dsn, "dsn", "postgres://postgres:@172.17.0.1/?sslmode=disable", "DSN database url")
	streamCmd.PersistentFlags().StringVar(&config.BindAddr, "bind", "localhost:8080", "bind to this address:port")
	streamCmd.PersistentFlags().StringVar(&config.NewRelic, "new-relic", "", "new relic api key")
}
<|endoftext|>