<commit_before>// (c) Copyright 2015-2017 JONNALAGADDA Srinivas
//
// 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 flow

import (
	"database/sql"
	"errors"
	"fmt"
)

// NodeID is the type of unique identifiers of nodes.
type NodeID int64

// NodeFunc defines the type of functions that generate notification
// messages in workflows.
//
// These functions are triggered by appropriate nodes, when document
// events are applied to documents to possibly transform them.
// Invocation of a `NodeFunc` should result in a message that can then
// be dispatched to applicable mailboxes.
//
// Error should be returned only when an impossible situation arises,
// and processing needs to abort.  Note that returning an error stops
// the workflow.  Manual intervention will be needed to move the
// document further.
//
// N. B. NodeFunc instances must be referentially transparent --
// stateless and not capture their environment in any manner.
// Unexpected bad things could happen otherwise!
type NodeFunc func(*Document, *DocEvent) *Message

// defNodeFunc prepares a simple message that can be posted to
// applicable mailboces.
func defNodeFunc(d *Document, event *DocEvent) *Message {
	return &Message{
		dtype: d.dtype.id,
		docID: d.id,
		event: event.id,
		title: d.title,
		data:  event.text,
	}
}

// Node represents a specific logical unit of processing and routing
// in a workflow.
type Node struct {
	id    NodeID     // Unique identifier of this node
	dtype DocTypeID  // Document type which this node's workflow manages
	state DocStateID // A document arriving at this node must be in this state
	wflow WorkflowID // Containing flow of this node
	name  string     // Unique within its workflow
	ntype NodeType   // Topology type of this node
	nfunc NodeFunc   // Processing function of this node
}

// ID answers the unique identifier of this workflow node.
func (n *Node) ID() NodeID {
	return n.id
}

// Workflow answers this node definition's containing workflow.
func (n *Node) Workflow() WorkflowID {
	return n.wflow
}

// State answers the state that any document arriving at this node
// must be in.
func (n *Node) State() DocStateID {
	return n.state
}

// Name answers the descriptive title of this node.
func (n *Node) Name() string {
	return n.name
}

// Type answers the type of this node.
func (n *Node) Type() NodeType {
	return n.ntype
}

// Transitions answers the possible document states into which a
// document currently in the given state can transition.
func (n *Node) Transitions() (map[DocActionID]DocStateID, error) {
	q := `
	SELECT docaction_id, to_state_id
	FROM wf_docstate_transitions
	WHERE doctype_id = ?
	AND from_state_id = ?
	`
	rows, err := db.Query(q, n.dtype, n.state)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	hash := make(map[DocActionID]DocStateID)
	for rows.Next() {
		var da DocActionID
		var ds DocStateID
		err := rows.Scan(&da, &ds)
		if err != nil {
			return nil, err
		}
		hash[da] = ds
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return hash, nil
}

// SetFunc registers the given node function with this node.
//
// If `nil` is given, a default node function is registered instead.
// This default function sets the document title as the message
// subject, and the event's data as the message body.
func (n *Node) SetFunc(fn NodeFunc) error {
	if fn == nil {
		n.nfunc = defNodeFunc
		return nil
	}

	n.nfunc = fn
	return nil
}

// Func answers the processing function registered in this node
// definition.
func (n *Node) Func() NodeFunc {
	return n.nfunc
}

// applyEvent checks to see if the given event can be applied
// successfully.  Accordingly, it prepares a message by utilising the
// registered node function, and posts it to applicable mailboxes.
func (n *Node) applyEvent(otx *sql.Tx, event *DocEvent, recipients []GroupID) (DocStateID, error) {
	ts, err := n.Transitions()
	nstate, ok := ts[event.action]
	if !ok {
		return 0, fmt.Errorf("given event's action '%d' cannot be performed on a document in the state '%d'", event.action, n.state)
	}

	doc, err := _documents.Get(event.dtype, event.docID)
	if err != nil {
		return 0, err
	}
	if doc.state.id != event.state && doc.state.id != nstate {
		return 0, fmt.Errorf("document state is : %d, but event is targeting state : %d", doc.state.id, event.state)
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	// Process state transition according to the target node type.

	tnode, err := _nodes.GetByState(n.dtype, nstate)
	if err != nil {
		return 0, err
	}

	switch tnode.ntype {
	case NodeTypeJoinAny:
		// Multiple 'in's, but any one suffices.

		// Document has already transitioned; nothing to do.
		if doc.state.id == nstate {
			return nstate, nil
		}

		// This is the first event; process it.
		fallthrough

	case NodeTypeBegin, NodeTypeEnd, NodeTypeLinear, NodeTypeBranch:
		// Any node type having a single 'in'.

		// Update the document to transition the state.
		tbl := _doctypes.docStorName(event.dtype)
		q := `UPDATE ` + tbl + ` SET state = ? WHERE doc_id = ?`
		_, err = tx.Exec(q, nstate, event.docID)
		if err != nil {
			return 0, err
		}

		// Record event application.
		err = n.recordEvent(tx, event, nstate)
		if err != nil {
			return 0, err
		}

		// Post messages.
		msg := n.nfunc(doc, event)
		err = n.postMessage(tx, msg, recipients)
		if err != nil {
			return 0, err
		}

	case NodeTypeJoinAll:
		// Multiple 'in's, and all are required.

		// TODO(js)

	default:
		return 0, fmt.Errorf("unknown node type encountered : %s", tnode.ntype)
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return nstate, nil
}

// recordEvent writes a record stating that the given event has
// successfully been applied to effect a document state transition.
func (n *Node) recordEvent(otx *sql.Tx, event *DocEvent, nstate DocStateID) error {
	q := `
	INSERT INTO wf_docevent_application(doctype_id, doc_id, from_state_id, docevent_id, to_state_id)
	VALUES(?, ?, ?, ?, ?)
	`
	_, err := otx.Exec(q, event.dtype, event.docID, event.state, event.action, nstate)
	if err != nil {
		return err
	}

	return nil
}

// postMessage posts the given message into the mailboxes of the
// specified recipients.
func (n *Node) postMessage(otx *sql.Tx, msg *Message, recipients []GroupID) error {
	// Record the message.

	q := `
	INSERT INTO wf_messages(doctype_id, doc_id, docevent_id, title, data)
	VALUES(?, ?, ?, ?, ?)
	`
	res, err := otx.Exec(q, msg.dtype, msg.docID, msg.event, msg.title, msg.data)
	if err != nil {
		return err
	}
	var msgid int64
	if msgid, err = res.LastInsertId(); err != nil {
		return err
	}

	// Post it into applicable mailboxes.

	q = `
	INSERT INTO wf_mailboxes(group_id, message_id)
	VALUES(?, ?)
	`
	for _, gid := range recipients {
		res, err = otx.Exec(q, gid, msgid)
		if err != nil {
			return err
		}
	}

	return nil
}

// Unexported type, only for convenience methods.
type _Nodes struct{}

var _nodes *_Nodes

func init() {
	_nodes = &_Nodes{}
}

// Nodes provides a resource-like interface to the nodes defined in
// this system.
func Nodes() *_Nodes {
	return _nodes
}

// Nodes answers a list of the nodes comprising the given workflow.
func (ns *_Nodes) Nodes(id WorkflowID) ([]*Node, error) {
	q := `
	SELECT id, doctype_id, docstate_id, workflow_id, name, type
	FROM wf_workflow_nodes
	WHERE workflow_id = ?
	`
	rows, err := db.Query(q, id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ary := make([]*Node, 0, 5)
	for rows.Next() {
		var elem Node
		err = rows.Scan(&elem.id, &elem.dtype, &elem.state, &elem.wflow, &elem.name, &elem.ntype)
		if err != nil {
			return nil, err
		}
		elem.nfunc = defNodeFunc
		ary = append(ary, &elem)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return ary, nil
}

// Get retrieves the requested node from the database.
func (ns *_Nodes) Get(id NodeID) (*Node, error) {
	if id <= 0 {
		return nil, errors.New("node ID must be a positive integer")
	}

	var elem Node
	q := `
	SELECT id, doctype_id, docstate_id, workflow_id, name, type
	FROM wf_workflow_nodes
	WHERE id = ?
	`
	row := db.QueryRow(q, id)
	err := row.Scan(&elem.id, &elem.dtype, &elem.state, &elem.wflow, &elem.name, &elem.ntype)
	if err != nil {
		return nil, err
	}

	elem.nfunc = defNodeFunc
	return &elem, nil
}

// GetByState retrieves the requested node from the database, as per
// the document state specification.
func (ns *_Nodes) GetByState(dtype DocTypeID, state DocStateID) (*Node, error) {
	var elem Node
	q := `
	SELECT id, doctype_id, docstate_id, workflow_id, name, type
	FROM wf_workflow_nodes
	WHERE doctype_id = ?
	AND docstate_id = ?
	`
	row := db.QueryRow(q, dtype, state)
	err := row.Scan(&elem.id, &elem.dtype, &elem.state, &elem.wflow, &elem.name, &elem.ntype)
	if err != nil {
		return nil, err
	}

	elem.nfunc = defNodeFunc
	return &elem, nil
}
<commit_msg>Update document event status to 'A' upon a successful application<commit_after>// (c) Copyright 2015-2017 JONNALAGADDA Srinivas
//
// 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 flow

import (
	"database/sql"
	"errors"
	"fmt"
)

// NodeID is the type of unique identifiers of nodes.
type NodeID int64

// NodeFunc defines the type of functions that generate notification
// messages in workflows.
//
// These functions are triggered by appropriate nodes, when document
// events are applied to documents to possibly transform them.
// Invocation of a `NodeFunc` should result in a message that can then
// be dispatched to applicable mailboxes.
//
// Error should be returned only when an impossible situation arises,
// and processing needs to abort.  Note that returning an error stops
// the workflow.  Manual intervention will be needed to move the
// document further.
//
// N. B. NodeFunc instances must be referentially transparent --
// stateless and not capture their environment in any manner.
// Unexpected bad things could happen otherwise!
type NodeFunc func(*Document, *DocEvent) *Message

// defNodeFunc prepares a simple message that can be posted to
// applicable mailboces.
func defNodeFunc(d *Document, event *DocEvent) *Message {
	return &Message{
		dtype: d.dtype.id,
		docID: d.id,
		event: event.id,
		title: d.title,
		data:  event.text,
	}
}

// Node represents a specific logical unit of processing and routing
// in a workflow.
type Node struct {
	id    NodeID     // Unique identifier of this node
	dtype DocTypeID  // Document type which this node's workflow manages
	state DocStateID // A document arriving at this node must be in this state
	wflow WorkflowID // Containing flow of this node
	name  string     // Unique within its workflow
	ntype NodeType   // Topology type of this node
	nfunc NodeFunc   // Processing function of this node
}

// ID answers the unique identifier of this workflow node.
func (n *Node) ID() NodeID {
	return n.id
}

// Workflow answers this node definition's containing workflow.
func (n *Node) Workflow() WorkflowID {
	return n.wflow
}

// State answers the state that any document arriving at this node
// must be in.
func (n *Node) State() DocStateID {
	return n.state
}

// Name answers the descriptive title of this node.
func (n *Node) Name() string {
	return n.name
}

// Type answers the type of this node.
func (n *Node) Type() NodeType {
	return n.ntype
}

// Transitions answers the possible document states into which a
// document currently in the given state can transition.
func (n *Node) Transitions() (map[DocActionID]DocStateID, error) {
	q := `
	SELECT docaction_id, to_state_id
	FROM wf_docstate_transitions
	WHERE doctype_id = ?
	AND from_state_id = ?
	`
	rows, err := db.Query(q, n.dtype, n.state)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	hash := make(map[DocActionID]DocStateID)
	for rows.Next() {
		var da DocActionID
		var ds DocStateID
		err := rows.Scan(&da, &ds)
		if err != nil {
			return nil, err
		}
		hash[da] = ds
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return hash, nil
}

// SetFunc registers the given node function with this node.
//
// If `nil` is given, a default node function is registered instead.
// This default function sets the document title as the message
// subject, and the event's data as the message body.
func (n *Node) SetFunc(fn NodeFunc) error {
	if fn == nil {
		n.nfunc = defNodeFunc
		return nil
	}

	n.nfunc = fn
	return nil
}

// Func answers the processing function registered in this node
// definition.
func (n *Node) Func() NodeFunc {
	return n.nfunc
}

// applyEvent checks to see if the given event can be applied
// successfully.  Accordingly, it prepares a message by utilising the
// registered node function, and posts it to applicable mailboxes.
func (n *Node) applyEvent(otx *sql.Tx, event *DocEvent, recipients []GroupID) (DocStateID, error) {
	ts, err := n.Transitions()
	nstate, ok := ts[event.action]
	if !ok {
		return 0, fmt.Errorf("given event's action '%d' cannot be performed on a document in the state '%d'", event.action, n.state)
	}

	doc, err := _documents.Get(event.dtype, event.docID)
	if err != nil {
		return 0, err
	}
	if doc.state.id != event.state && doc.state.id != nstate {
		return 0, fmt.Errorf("document state is : %d, but event is targeting state : %d", doc.state.id, event.state)
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	// Process state transition according to the target node type.

	tnode, err := _nodes.GetByState(n.dtype, nstate)
	if err != nil {
		return 0, err
	}

	switch tnode.ntype {
	case NodeTypeJoinAny:
		// Multiple 'in's, but any one suffices.

		// Document has already transitioned; nothing to do.
		if doc.state.id == nstate {
			return nstate, nil
		}

		// This is the first event; process it.
		fallthrough

	case NodeTypeBegin, NodeTypeEnd, NodeTypeLinear, NodeTypeBranch:
		// Any node type having a single 'in'.

		// Update the document to transition the state.
		tbl := _doctypes.docStorName(event.dtype)
		q := `UPDATE ` + tbl + ` SET state = ? WHERE doc_id = ?`
		_, err = tx.Exec(q, nstate, event.docID)
		if err != nil {
			return 0, err
		}

		// Record event application.
		err = n.recordEvent(tx, event, nstate)
		if err != nil {
			return 0, err
		}

		// Post messages.
		msg := n.nfunc(doc, event)
		err = n.postMessage(tx, msg, recipients)
		if err != nil {
			return 0, err
		}

	case NodeTypeJoinAll:
		// Multiple 'in's, and all are required.

		// TODO(js)

	default:
		return 0, fmt.Errorf("unknown node type encountered : %s", tnode.ntype)
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return nstate, nil
}

// recordEvent writes a record stating that the given event has
// successfully been applied to effect a document state transition.
func (n *Node) recordEvent(otx *sql.Tx, event *DocEvent, nstate DocStateID) error {
	q := `
	INSERT INTO wf_docevent_application(doctype_id, doc_id, from_state_id, docevent_id, to_state_id)
	VALUES(?, ?, ?, ?, ?)
	`
	_, err := otx.Exec(q, event.dtype, event.docID, event.state, event.id, nstate)
	if err != nil {
		return err
	}

	q = `UPDATE wf_docevents SET status = 'A' WHERE id = ?`
	_, err = otx.Exec(q, event.id)
	if err != nil {
		return err
	}

	return nil
}

// postMessage posts the given message into the mailboxes of the
// specified recipients.
func (n *Node) postMessage(otx *sql.Tx, msg *Message, recipients []GroupID) error {
	// Record the message.

	q := `
	INSERT INTO wf_messages(doctype_id, doc_id, docevent_id, title, data)
	VALUES(?, ?, ?, ?, ?)
	`
	res, err := otx.Exec(q, msg.dtype, msg.docID, msg.event, msg.title, msg.data)
	if err != nil {
		return err
	}
	var msgid int64
	if msgid, err = res.LastInsertId(); err != nil {
		return err
	}

	// Post it into applicable mailboxes.

	q = `
	INSERT INTO wf_mailboxes(group_id, message_id)
	VALUES(?, ?)
	`
	for _, gid := range recipients {
		res, err = otx.Exec(q, gid, msgid)
		if err != nil {
			return err
		}
	}

	return nil
}

// Unexported type, only for convenience methods.
type _Nodes struct{}

var _nodes *_Nodes

func init() {
	_nodes = &_Nodes{}
}

// Nodes provides a resource-like interface to the nodes defined in
// this system.
func Nodes() *_Nodes {
	return _nodes
}

// Nodes answers a list of the nodes comprising the given workflow.
func (ns *_Nodes) Nodes(id WorkflowID) ([]*Node, error) {
	q := `
	SELECT id, doctype_id, docstate_id, workflow_id, name, type
	FROM wf_workflow_nodes
	WHERE workflow_id = ?
	`
	rows, err := db.Query(q, id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ary := make([]*Node, 0, 5)
	for rows.Next() {
		var elem Node
		err = rows.Scan(&elem.id, &elem.dtype, &elem.state, &elem.wflow, &elem.name, &elem.ntype)
		if err != nil {
			return nil, err
		}
		elem.nfunc = defNodeFunc
		ary = append(ary, &elem)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return ary, nil
}

// Get retrieves the requested node from the database.
func (ns *_Nodes) Get(id NodeID) (*Node, error) {
	if id <= 0 {
		return nil, errors.New("node ID must be a positive integer")
	}

	var elem Node
	q := `
	SELECT id, doctype_id, docstate_id, workflow_id, name, type
	FROM wf_workflow_nodes
	WHERE id = ?
	`
	row := db.QueryRow(q, id)
	err := row.Scan(&elem.id, &elem.dtype, &elem.state, &elem.wflow, &elem.name, &elem.ntype)
	if err != nil {
		return nil, err
	}

	elem.nfunc = defNodeFunc
	return &elem, nil
}

// GetByState retrieves the requested node from the database, as per
// the document state specification.
func (ns *_Nodes) GetByState(dtype DocTypeID, state DocStateID) (*Node, error) {
	var elem Node
	q := `
	SELECT id, doctype_id, docstate_id, workflow_id, name, type
	FROM wf_workflow_nodes
	WHERE doctype_id = ?
	AND docstate_id = ?
	`
	row := db.QueryRow(q, dtype, state)
	err := row.Scan(&elem.id, &elem.dtype, &elem.state, &elem.wflow, &elem.name, &elem.ntype)
	if err != nil {
		return nil, err
	}

	elem.nfunc = defNodeFunc
	return &elem, nil
}
<|endoftext|>


<commit_before>package openzwave

// #cgo LDFLAGS: -lopenzwave -Lgo/src/github.com/ninjasphere/go-openzwave/openzwave
// #cgo CPPFLAGS: -Iopenzwave/cpp/src/platform -Iopenzwave/cpp/src -Iopenzwave/cpp/src/value_classes
//
// #include "api.h"
import "C"

import (
	"fmt"
	"unsafe"

	"github.com/ninjasphere/go-openzwave/NT"
)

type state int

const (
	STATE_INIT  state = iota
	STATE_READY       = iota
)

//export newGoNode
func newGoNode(cRef *C.Node) unsafe.Pointer {
	goRef := &node{cRef, make(map[uint8]*valueClass), STATE_INIT, nil}
	cRef.goRef = unsafe.Pointer(goRef)
	return cRef.goRef
}

// A device is an abstraction used by higher-level parts of the system for the underlying ZWave node.
type Device interface{}

type Node interface {
	GetHomeId() uint32
	GetId() uint8

	GetDevice() Device
	SetDevice(device Device)

	GetProductId() *ProductId
	GetProductDescription() *ProductDescription
	GetNodeName() string
}

type ProductId struct {
	ManufacturerId string
	ProductId      string
}

type ProductDescription struct {
     ManufacturerName string
     ProductName      string
     ProductType      string
}

type node struct {
	cRef    *C.Node
	classes map[uint8]*valueClass
	state   state
	device  Device
}

type valueClass struct {
	commandClass uint8
	instances    map[uint8]*valueInstance
}

type valueInstance struct {
	instance uint8
	values   map[uint8]*value
}

func (self *node) String() string {
	cRef := self.cRef

	return fmt.Sprintf(
		"Node["+
			"homeId=0x%08x, "+
			"nodeId=%03d, "+
			"basicType=%02x, "+
			"genericType=%02x, "+
			"specificType=%02x, "+
			"nodeType='%s', "+
			"manufacturerName='%s', "+
			"productName='%s', "+
			"location='%s', "+
			"manufacturerId=%s, "+
			"productType=%s, "+
			"productId=%s]",
		uint32(cRef.nodeId.homeId),
		uint8(cRef.nodeId.nodeId),
		uint8(cRef.basicType),
		uint8(cRef.genericType),
		uint8(cRef.specificType),
		C.GoString(cRef.nodeType),
		C.GoString(cRef.manufacturerName),
		C.GoString(cRef.productName),
		C.GoString(cRef.location),
		C.GoString(cRef.manufacturerId),
		C.GoString(cRef.productType),
		C.GoString(cRef.productId))
}

// convert a reference from the C Node to the Go Node
func asNode(cRef *C.Node) Node {
	return Node((*node)(cRef.goRef))
}

func (self *node) GetHomeId() uint32 {
	return uint32(self.cRef.nodeId.homeId)
}

func (self *node) GetId() uint8 {
	return uint8(self.cRef.nodeId.nodeId)
}

func (self *node) notify(api *api, nt *notification) {

	var event Event

	notificationType := nt.cRef.notificationType
	switch notificationType {
	case NT.NODE_REMOVED:
		api.notifyEvent(&NodeUnavailable{nodeEvent{self}})
		break

	case NT.VALUE_REMOVED:
		self.removeValue(nt)
		break

	case NT.ESSENTIAL_NODE_QUERIES_COMPLETE:
	case NT.NODE_QUERIES_COMPLETE:
		// move the node into the initialized state
		// begin admission processing for the node

		switch self.state {
		case STATE_INIT:
			self.state = STATE_READY
			event = &NodeAvailable{nodeEvent{self}}
		default:
			event = &NodeChanged{nodeEvent{self}}
		}
		api.notifyEvent(event)
		break

	case NT.VALUE_ADDED:
	case NT.VALUE_CHANGED:
	case NT.VALUE_REFRESHED:
		self.takeValue(nt)
		break

	case NT.NODE_NAMING:
	case NT.NODE_PROTOCOL_INFO:
		// log the related information for diagnostics purposes

	}
}

// take the value structure from the notification
func (self *node) takeValue(nt *notification) *value {
	commandClassId := (uint8)(nt.cRef.value.valueId.commandClassId)
	instanceId := (uint8)(nt.cRef.value.valueId.instance)
	index := (uint8)(nt.cRef.value.valueId.index)

	instance := self.createOrGetInstance(commandClassId, instanceId)
	v, ok := instance.values[index]
	if !ok {
		v = nt.swapValueImpl(nil)
		instance.values[index] = v
	} else {
		nt.swapValueImpl(v)
	}
	return v
}

func (self *node) createOrGetInstance(commandClassId uint8, instanceId uint8) *valueInstance {
	class, ok := self.classes[commandClassId]
	if !ok {
		class = &valueClass{commandClassId, make(map[uint8]*valueInstance)}
		self.classes[commandClassId] = class
	}
	instance, ok := class.instances[instanceId]
	if !ok {
		instance = &valueInstance{instanceId, make(map[uint8]*value)}
		class.instances[instanceId] = instance
	}
	return instance
}

func (self *node) removeValue(nt *notification) {
	commandClassId := (uint8)(nt.cRef.value.valueId.commandClassId)
	instanceId := (uint8)(nt.cRef.value.valueId.instance)
	index := (uint8)(nt.cRef.value.valueId.index)

	class, ok := self.classes[commandClassId]
	if !ok {
		return
	}

	instance, ok := class.instances[instanceId]
	if !ok {
		return
	}

	value, ok := instance.values[index]
	_ = value

	if !ok {
		return
	} else {
		delete(instance.values, index)
		if len(instance.values) == 0 {
			delete(class.instances, instanceId)
			if len(class.instances) == 0 {
				delete(self.classes, commandClassId)
			}
		}
	}

}

func (self *node) GetDevice() Device {
	return self.device
}

func (self *node) SetDevice(device Device) {
	self.device = device
}

func (self *node) GetProductId() *ProductId {
	return &ProductId{C.GoString(self.cRef.manufacturerId), C.GoString(self.cRef.productId)}
}

func (self *node) GetProductDescription() *ProductDescription {
	return &ProductDescription{
	       C.GoString(self.cRef.manufacturerName),
	       C.GoString(self.cRef.productName),
	       C.GoString(self.cRef.productType)}
}

func (self *node) GetNodeName() string {
     return C.GoString(self.cRef.nodeName)
}<commit_msg>Reformatting.<commit_after>package openzwave

// #cgo LDFLAGS: -lopenzwave -Lgo/src/github.com/ninjasphere/go-openzwave/openzwave
// #cgo CPPFLAGS: -Iopenzwave/cpp/src/platform -Iopenzwave/cpp/src -Iopenzwave/cpp/src/value_classes
//
// #include "api.h"
import "C"

import (
	"fmt"
	"unsafe"

	"github.com/ninjasphere/go-openzwave/NT"
)

type state int

const (
	STATE_INIT  state = iota
	STATE_READY       = iota
)

//export newGoNode
func newGoNode(cRef *C.Node) unsafe.Pointer {
	goRef := &node{cRef, make(map[uint8]*valueClass), STATE_INIT, nil}
	cRef.goRef = unsafe.Pointer(goRef)
	return cRef.goRef
}

// A device is an abstraction used by higher-level parts of the system for the underlying ZWave node.
type Device interface{}

type Node interface {
	GetHomeId() uint32
	GetId() uint8

	GetDevice() Device
	SetDevice(device Device)

	GetProductId() *ProductId
	GetProductDescription() *ProductDescription
	GetNodeName() string
}

type ProductId struct {
	ManufacturerId string
	ProductId      string
}

type ProductDescription struct {
	ManufacturerName string
	ProductName      string
	ProductType      string
}

type node struct {
	cRef    *C.Node
	classes map[uint8]*valueClass
	state   state
	device  Device
}

type valueClass struct {
	commandClass uint8
	instances    map[uint8]*valueInstance
}

type valueInstance struct {
	instance uint8
	values   map[uint8]*value
}

func (self *node) String() string {
	cRef := self.cRef

	return fmt.Sprintf(
		"Node["+
			"homeId=0x%08x, "+
			"nodeId=%03d, "+
			"basicType=%02x, "+
			"genericType=%02x, "+
			"specificType=%02x, "+
			"nodeType='%s', "+
			"manufacturerName='%s', "+
			"productName='%s', "+
			"location='%s', "+
			"manufacturerId=%s, "+
			"productType=%s, "+
			"productId=%s]",
		uint32(cRef.nodeId.homeId),
		uint8(cRef.nodeId.nodeId),
		uint8(cRef.basicType),
		uint8(cRef.genericType),
		uint8(cRef.specificType),
		C.GoString(cRef.nodeType),
		C.GoString(cRef.manufacturerName),
		C.GoString(cRef.productName),
		C.GoString(cRef.location),
		C.GoString(cRef.manufacturerId),
		C.GoString(cRef.productType),
		C.GoString(cRef.productId))
}

// convert a reference from the C Node to the Go Node
func asNode(cRef *C.Node) Node {
	return Node((*node)(cRef.goRef))
}

func (self *node) GetHomeId() uint32 {
	return uint32(self.cRef.nodeId.homeId)
}

func (self *node) GetId() uint8 {
	return uint8(self.cRef.nodeId.nodeId)
}

func (self *node) notify(api *api, nt *notification) {

	var event Event

	notificationType := nt.cRef.notificationType
	switch notificationType {
	case NT.NODE_REMOVED:
		api.notifyEvent(&NodeUnavailable{nodeEvent{self}})
		break

	case NT.VALUE_REMOVED:
		self.removeValue(nt)
		break

	case NT.ESSENTIAL_NODE_QUERIES_COMPLETE:
	case NT.NODE_QUERIES_COMPLETE:
		// move the node into the initialized state
		// begin admission processing for the node

		switch self.state {
		case STATE_INIT:
			self.state = STATE_READY
			event = &NodeAvailable{nodeEvent{self}}
		default:
			event = &NodeChanged{nodeEvent{self}}
		}
		api.notifyEvent(event)
		break

	case NT.VALUE_ADDED:
	case NT.VALUE_CHANGED:
	case NT.VALUE_REFRESHED:
		self.takeValue(nt)
		break

	case NT.NODE_NAMING:
	case NT.NODE_PROTOCOL_INFO:
		// log the related information for diagnostics purposes

	}
}

// take the value structure from the notification
func (self *node) takeValue(nt *notification) *value {
	commandClassId := (uint8)(nt.cRef.value.valueId.commandClassId)
	instanceId := (uint8)(nt.cRef.value.valueId.instance)
	index := (uint8)(nt.cRef.value.valueId.index)

	instance := self.createOrGetInstance(commandClassId, instanceId)
	v, ok := instance.values[index]
	if !ok {
		v = nt.swapValueImpl(nil)
		instance.values[index] = v
	} else {
		nt.swapValueImpl(v)
	}
	return v
}

func (self *node) createOrGetInstance(commandClassId uint8, instanceId uint8) *valueInstance {
	class, ok := self.classes[commandClassId]
	if !ok {
		class = &valueClass{commandClassId, make(map[uint8]*valueInstance)}
		self.classes[commandClassId] = class
	}
	instance, ok := class.instances[instanceId]
	if !ok {
		instance = &valueInstance{instanceId, make(map[uint8]*value)}
		class.instances[instanceId] = instance
	}
	return instance
}

func (self *node) removeValue(nt *notification) {
	commandClassId := (uint8)(nt.cRef.value.valueId.commandClassId)
	instanceId := (uint8)(nt.cRef.value.valueId.instance)
	index := (uint8)(nt.cRef.value.valueId.index)

	class, ok := self.classes[commandClassId]
	if !ok {
		return
	}

	instance, ok := class.instances[instanceId]
	if !ok {
		return
	}

	value, ok := instance.values[index]
	_ = value

	if !ok {
		return
	} else {
		delete(instance.values, index)
		if len(instance.values) == 0 {
			delete(class.instances, instanceId)
			if len(class.instances) == 0 {
				delete(self.classes, commandClassId)
			}
		}
	}

}

func (self *node) GetDevice() Device {
	return self.device
}

func (self *node) SetDevice(device Device) {
	self.device = device
}

func (self *node) GetProductId() *ProductId {
	return &ProductId{C.GoString(self.cRef.manufacturerId), C.GoString(self.cRef.productId)}
}

func (self *node) GetProductDescription() *ProductDescription {
	return &ProductDescription{
		C.GoString(self.cRef.manufacturerName),
		C.GoString(self.cRef.productName),
		C.GoString(self.cRef.productType)}
}

func (self *node) GetNodeName() string {
	return C.GoString(self.cRef.nodeName)
}

<|endoftext|>


<commit_before>// Copyright 2015 Prometheus Team
// 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 main

import (
	"database/sql"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"

	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"

	"github.com/prometheus/alertmanager/config"
	"github.com/prometheus/alertmanager/notify"
	"github.com/prometheus/alertmanager/provider"
	"github.com/prometheus/alertmanager/types"
)

var (
	configFile    = flag.String("config.file", "config.yml", "The configuration file")
	dataDir       = flag.String("data.dir", "data/", "The data directory")
	listenAddress = flag.String("web.listen-address", ":9093", "Address to listen on for the web interface and API.")
)

func main() {
	flag.Parse()

	db, err := sql.Open("ql", filepath.Join(*dataDir, "am.db"))
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	alerts, err := provider.NewSQLAlerts(db)
	if err != nil {
		log.Fatal(err)
	}
	notifies, err := provider.NewSQLNotifyInfo(db)
	if err != nil {
		log.Fatal(err)
	}
	silences, err := provider.NewSQLSilences(db)
	if err != nil {
		log.Fatal(err)
	}

	inhibitor := &Inhibitor{alerts: alerts}

	routedNotifier := notify.NewRoutedNotifier(func(confs []*config.NotificationConfig) map[string]notify.Notifier {
		res := notify.Build(confs)
		for name, n := range res {
			n = &notify.LogNotifier{
				Log:      log.With("notifier", fmt.Sprintf("%T", n)),
				Notifier: n,
			}
			res[name] = &notify.LogNotifier{
				Log:      log.With("notifier", "dedup"),
				Notifier: notify.NewDedupingNotifier(notifies, n),
			}
		}
		return res
	})

	var notifier notify.Notifier
	notifier = &notify.LogNotifier{
		Log:      log.With("notifier", "routed"),
		Notifier: routedNotifier,
	}
	notifier = &notify.MutingNotifier{
		Notifier: notifier,
		Muter:    inhibitor,
	}
	notifier = &notify.LogNotifier{
		Log:      log.With("notifier", "inhibit"),
		Notifier: notifier,
	}
	notifier = &notify.MutingNotifier{
		Notifier: notifier,
		Muter:    silences,
	}
	notifier = &notify.LogNotifier{
		Log:      log.With("notifier", "silencer"),
		Notifier: notifier,
	}

	disp := NewDispatcher(alerts, notifier)

	if !reloadConfig(*configFile, disp, routedNotifier, inhibitor) {
		os.Exit(1)
	}

	go disp.Run()
	defer disp.Stop()

	router := route.New()

	NewAPI(router.WithPrefix("/api/v1"), alerts, silences)

	go http.ListenAndServe(*listenAddress, router)

	var (
		hup  = make(chan os.Signal)
		term = make(chan os.Signal)
	)
	signal.Notify(hup, syscall.SIGHUP)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)

	go func() {
		for range hup {
			reloadConfig(*configFile, disp, routedNotifier, inhibitor)
		}
	}()

	<-term

	log.Infoln("Received SIGTERM, exiting gracefully...")
	os.Exit(0)
}

func reloadConfig(filename string, rls ...types.Reloadable) (success bool) {
	log.Infof("Loading configuration file %s", filename)

	conf, err := config.LoadFile(filename)
	if err != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", filename, err)
		return false
	}
	success = true

	for _, rl := range rls {
		success = success && rl.ApplyConfig(conf)
	}
	return success
}
<commit_msg>Change notifier pipeline order.<commit_after>// Copyright 2015 Prometheus Team
// 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 main

import (
	"database/sql"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"

	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"

	"github.com/prometheus/alertmanager/config"
	"github.com/prometheus/alertmanager/notify"
	"github.com/prometheus/alertmanager/provider"
	"github.com/prometheus/alertmanager/types"
)

var (
	configFile    = flag.String("config.file", "config.yml", "The configuration file")
	dataDir       = flag.String("data.dir", "data/", "The data directory")
	listenAddress = flag.String("web.listen-address", ":9093", "Address to listen on for the web interface and API.")
)

func main() {
	flag.Parse()

	db, err := sql.Open("ql", filepath.Join(*dataDir, "am.db"))
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	alerts, err := provider.NewSQLAlerts(db)
	if err != nil {
		log.Fatal(err)
	}
	notifies, err := provider.NewSQLNotifyInfo(db)
	if err != nil {
		log.Fatal(err)
	}
	silences, err := provider.NewSQLSilences(db)
	if err != nil {
		log.Fatal(err)
	}

	inhibitor := &Inhibitor{alerts: alerts}

	routedNotifier := notify.NewRoutedNotifier(func(confs []*config.NotificationConfig) map[string]notify.Notifier {
		res := notify.Build(confs)
		for name, n := range res {
			res[name] = &notify.LogNotifier{
				Log:      log.With("notifier", fmt.Sprintf("%T", n)),
				Notifier: n,
			}
		}
		return res
	})

	// Connect the pipeline of notifiers. Notifications will be sent
	// through them in inverted order.
	var notifier notify.Notifier
	notifier = &notify.LogNotifier{
		Log:      log.With("notifier", "routed"),
		Notifier: routedNotifier,
	}

	notifier = notify.NewDedupingNotifier(notifies, notifier)
	notifier = &notify.LogNotifier{
		Log:      log.With("notifier", "dedup"),
		Notifier: notifier,
	}

	notifier = &notify.MutingNotifier{
		Notifier: notifier,
		Muter:    inhibitor,
	}
	notifier = &notify.LogNotifier{
		Log:      log.With("notifier", "inhibit"),
		Notifier: notifier,
	}

	notifier = &notify.MutingNotifier{
		Notifier: notifier,
		Muter:    silences,
	}
	notifier = &notify.LogNotifier{
		Log:      log.With("notifier", "silencer"),
		Notifier: notifier,
	}

	disp := NewDispatcher(alerts, notifier)

	if !reloadConfig(*configFile, disp, routedNotifier, inhibitor) {
		os.Exit(1)
	}

	go disp.Run()
	defer disp.Stop()

	router := route.New()

	NewAPI(router.WithPrefix("/api/v1"), alerts, silences)

	go http.ListenAndServe(*listenAddress, router)

	var (
		hup  = make(chan os.Signal)
		term = make(chan os.Signal)
	)
	signal.Notify(hup, syscall.SIGHUP)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)

	go func() {
		for range hup {
			reloadConfig(*configFile, disp, routedNotifier, inhibitor)
		}
	}()

	<-term

	log.Infoln("Received SIGTERM, exiting gracefully...")
	os.Exit(0)
}

func reloadConfig(filename string, rls ...types.Reloadable) (success bool) {
	log.Infof("Loading configuration file %s", filename)

	conf, err := config.LoadFile(filename)
	if err != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", filename, err)
		return false
	}
	success = true

	for _, rl := range rls {
		success = success && rl.ApplyConfig(conf)
	}
	return success
}
<|endoftext|>


<commit_before>// Copyright 2015 Joel Wu
// Copyright 2012 Gary Burd
//
// 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 redis

import (
    "fmt"
    "net"
    "sync"
    "time"
    "bufio"
    "errors"
    "strconv"
    "container/list"
)

type redisConn struct {
    c	net.Conn
    t	time.Time

    br	*bufio.Reader
    bw	*bufio.Writer

    readTimeout	    time.Duration
    writeTimeout    time.Duration

    // Pending replies to be read in redis pipeling.
    pending	int

    // Scratch space for formatting argument length.
    lenScratch [32]byte

    // Scratch space for formatting integer and float.
    numScratch [40]byte
}

type redisNode struct {
    address	    string

    conns	    list.List
    keepAlive	    int
    aliveTime	    time.Duration

    connTimeout	    time.Duration
    readTimeout	    time.Duration
    writeTimeout    time.Duration

    mutex	    sync.Mutex

    updateTime	    time.Time

    closed	    bool
}

func (node *redisNode) getConn() (*redisConn, error) {
    node.mutex.Lock()

    if node.closed {
	node.mutex.Unlock()
	return nil, fmt.Errorf("getConn: connection has been closed")
    }

    // remove stale connections
    if node.connTimeout > 0 {
	for {
	    elem := node.conns.Back()
	    if elem == nil {
		break
	    }
	    conn := elem.Value.(*redisConn)
	    if conn.t.Add(node.aliveTime).After(time.Now()) {
		break
	    }
	    node.conns.Remove(elem)
	}
    }

    if node.conns.Len() <= 0 {
	node.mutex.Unlock()

	c, err := net.DialTimeout("tcp", node.address, node.connTimeout)
	if err != nil {
	    return nil, err
	}

	conn := &redisConn{
	    c: c,
	    br: bufio.NewReader(c),
	    bw: bufio.NewWriter(c),
	    readTimeout: node.readTimeout,
	    writeTimeout: node.writeTimeout,
	}

	return conn, nil
    }

    elem := node.conns.Back()
    node.conns.Remove(elem)
    node.mutex.Unlock()

    return elem.Value.(*redisConn), nil
}

func (node *redisNode) releaseConn(conn *redisConn) {
    node.mutex.Lock()
    defer node.mutex.Unlock()

    // Connection still has pending replies, just close it.
    if conn.pending > 0 || node.closed {
	conn.shutdown()
	return
    }

    if node.conns.Len() >= node.keepAlive || node.aliveTime <= 0 {
	conn.shutdown()
	return
    }

    if conn.c.RemoteAddr().String() != node.address {
	panic("unreachable")
    }

    conn.t = time.Now()
    node.conns.PushFront(conn)
}

func (conn *redisConn) shutdown() {
    conn.c.Close()
}

func (node *redisNode) shutdown() {
    node.mutex.Lock()
    defer node.mutex.Unlock()

    for {
	elem := node.conns.Back()
	if elem == nil {
	    break
	}

	conn := elem.Value.(*redisConn)
	conn.c.Close()
	node.conns.Remove(elem)
    }

    node.closed = true
}

func (conn *redisConn) send(cmd string, args ...interface{}) error {
    conn.pending += 1

    if conn.writeTimeout > 0 {
	conn.c.SetWriteDeadline(time.Now().Add(conn.writeTimeout))
    }

    if err := conn.writeCommand(cmd, args); err != nil {
	return err
    }

    return nil
}

func (conn *redisConn) flush() error {
    if conn.writeTimeout > 0 {
	conn.c.SetWriteDeadline(time.Now().Add(conn.writeTimeout))
    }

    if err := conn.bw.Flush(); err != nil {
	return err
    }

    return nil
}

func (conn *redisConn) receive() (interface{}, error) {
    if conn.readTimeout > 0 {
	conn.c.SetWriteDeadline(time.Now().Add(conn.readTimeout))
    }

    if conn.pending <= 0 {
	return nil, errors.New("no more pending reply")
    }

    conn.pending -= 1

    return conn.readReply()
}

func (node *redisNode) do(cmd string, args ...interface{}) (interface{}, error) {
    conn, err := node.getConn()
    if err != nil {
	return redisError("ECONNTIMEOUT"), nil
    }

    if err = conn.send(cmd, args...); err != nil {
	conn.shutdown()
	return nil, err
    }

    if err = conn.flush(); err != nil {
	conn.shutdown()
	return nil, err
    }

    reply, err := conn.receive()
    if err != nil {
	conn.shutdown()
	return nil, err
    }

    node.releaseConn(conn)

    return reply, err
}

func (conn *redisConn) writeLen(prefix byte, n int) error {
    conn.lenScratch[len(conn.lenScratch) - 1] = '\n'
    conn.lenScratch[len(conn.lenScratch) - 2] = '\r'
    i := len(conn.lenScratch) - 3

    for {
        conn.lenScratch[i] = byte('0' + n % 10)
        i -= 1
        n = n / 10
        if n == 0 {
	    break
        }
    }

    conn.lenScratch[i] = prefix
    _, err := conn.bw.Write(conn.lenScratch[i:])

    return err
}

func (conn *redisConn) writeString(s string) error {
    conn.writeLen('$', len(s))
    conn.bw.WriteString(s)
    _, err := conn.bw.WriteString("\r\n")

    return err
}

func (conn *redisConn) writeBytes(p []byte) error {
    conn.writeLen('$', len(p))
    conn.bw.Write(p)
    _, err := conn.bw.WriteString("\r\n")

    return err
}

func (conn *redisConn) writeInt64(n int64) error {
    return conn.writeBytes(strconv.AppendInt(conn.numScratch[:0], n, 10))
}

func (conn *redisConn) writeFloat64(n float64) error {
    return conn.writeBytes(strconv.AppendFloat(conn.numScratch[:0], n, 'g', -1, 64))
}

// Args must be int64, float64, string, []byte, other types are not supported for safe reason.
func (conn *redisConn) writeCommand(cmd string, args []interface{}) error {
    conn.writeLen('*', len(args) + 1)
    err := conn.writeString(cmd)

    for _, arg := range args {
	if err != nil {
		break
	}
	switch arg := arg.(type) {
	case int:
	    err = conn.writeInt64(int64(arg))
	case int64:
	    err = conn.writeInt64(arg)
	case float64:
	    err = conn.writeFloat64(arg)
	case string:
	    err = conn.writeString(arg)
	case []byte:
	    err = conn.writeBytes(arg)
	default:
	    err = fmt.Errorf("unknown type %T", arg)
	}
    }

    return err
}

// readLine read a single line terminated with CRLF.
func (conn *redisConn) readLine() ([]byte, error) {
    var line []byte
    for {
	p, err := conn.br.ReadBytes('\n')
	if err != nil {
	    return nil, err
	}

	n := len(p) - 2
	if n < 0 {
	    return nil, errors.New("invalid response")
	}

	// bulk string may contain '\n', such as CLUSTER NODES
	if p[n] != '\r' {
	    if line != nil {
		line = append(line, p[:]...)
	    } else {
		line = p
	    }
	    continue
	}

	if line != nil {
	    return append(line, p[:n]...), nil
	} else {
	    return p[:n], nil
	}
    }
}

// parseLen parses bulk string and array length.
func parseLen(p []byte) (int, error) {
    if len(p) == 0 {
	return -1, errors.New("invalid response")
    }

    // null element.
    if p[0] == '-' && len(p) == 2 && p[1] == '1' {
	return -1, nil
    }

    var n int
    for _, b := range p {
	n *= 10
	if b < '0' || b > '9' {
		return -1, errors.New("invalid response")
	}
	n += int(b - '0')
    }

    return n, nil
}

// parseInt parses an integer reply.
func parseInt(p []byte) (int64, error) {
    if len(p) == 0 {
	return 0, errors.New("invalid response")
    }

    var negate bool
    if p[0] == '-' {
	negate = true
	p = p[1:]
	if len(p) == 0 {
	    return 0, errors.New("invalid response")
	}
    }

    var n int64
    for _, b := range p {
	n *= 10
	if b < '0' || b > '9' {
		return 0, errors.New("invalid response")
	}
	n += int64(b - '0')
    }

    if negate {
	n = -n
    }

    return n, nil
}

var (
    okReply   interface{} = "OK"
    pongReply interface{} = "PONG"
)

type redisError string
func (err redisError) Error() string { return string(err) }

func (conn *redisConn) readReply() (interface{}, error) {
    line, err := conn.readLine()
    if err != nil {
	return nil, err
    }
    if len(line) == 0 {
	return nil, errors.New("invalid reponse")
    }

    switch line[0] {
    case '+':
	switch {
	case len(line) == 3 && line[1] == 'O' && line[2] == 'K':
	    // Avoid allocation for frequent "+OK" response.
	    return okReply, nil
	case len(line) == 5 && line[1] == 'P' && line[2] == 'O' && line[3] == 'N' && line[4] == 'G':
	    // Avoid allocation in PING command benchmarks :)
	    return pongReply, nil
	default:
	    return string(line[1:]), nil
	}
    case '-':
	return redisError(string(line[1:])), nil
    case ':':
	return parseInt(line[1:])
    case '$':
	n, err := parseLen(line[1:])
	if n < 0 || err != nil {
	    return nil, err
	}

	line, err = conn.readLine()
	if err != nil {
	    return nil, err
	}
	if len(line) != n {
	    return nil, errors.New("invalid response")
	}

	return line, nil
    case '*':
	n, err := parseLen(line[1:])
	if n < 0 || err != nil {
	    return nil, err
	}

	r := make([]interface{}, n)
	for i := range r {
	    r[i], err = conn.readReply()
	    if err != nil {
		return nil, err
	    }
	}

	return r, nil
    }

    return nil, errors.New("invalid response")
}
<commit_msg>bugfix: avoid panic when initializing cluster with domain name<commit_after>// Copyright 2015 Joel Wu
// Copyright 2012 Gary Burd
//
// 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 redis

import (
    "fmt"
    "net"
    "sync"
    "time"
    "bufio"
    "errors"
    "strconv"
    "container/list"
)

type redisConn struct {
    c	net.Conn
    t	time.Time

    br	*bufio.Reader
    bw	*bufio.Writer

    readTimeout	    time.Duration
    writeTimeout    time.Duration

    // Pending replies to be read in redis pipeling.
    pending	int

    // Scratch space for formatting argument length.
    lenScratch [32]byte

    // Scratch space for formatting integer and float.
    numScratch [40]byte
}

type redisNode struct {
    address	    string

    conns	    list.List
    keepAlive	    int
    aliveTime	    time.Duration

    connTimeout	    time.Duration
    readTimeout	    time.Duration
    writeTimeout    time.Duration

    mutex	    sync.Mutex

    updateTime	    time.Time

    closed	    bool
}

func (node *redisNode) getConn() (*redisConn, error) {
    node.mutex.Lock()

    if node.closed {
	node.mutex.Unlock()
	return nil, fmt.Errorf("getConn: connection has been closed")
    }

    // remove stale connections
    if node.connTimeout > 0 {
	for {
	    elem := node.conns.Back()
	    if elem == nil {
		break
	    }
	    conn := elem.Value.(*redisConn)
	    if conn.t.Add(node.aliveTime).After(time.Now()) {
		break
	    }
	    node.conns.Remove(elem)
	}
    }

    if node.conns.Len() <= 0 {
	node.mutex.Unlock()

	c, err := net.DialTimeout("tcp", node.address, node.connTimeout)
	if err != nil {
	    return nil, err
	}

	conn := &redisConn{
	    c: c,
	    br: bufio.NewReader(c),
	    bw: bufio.NewWriter(c),
	    readTimeout: node.readTimeout,
	    writeTimeout: node.writeTimeout,
	}

	return conn, nil
    }

    elem := node.conns.Back()
    node.conns.Remove(elem)
    node.mutex.Unlock()

    return elem.Value.(*redisConn), nil
}

func (node *redisNode) releaseConn(conn *redisConn) {
    node.mutex.Lock()
    defer node.mutex.Unlock()

    // Connection still has pending replies, just close it.
    if conn.pending > 0 || node.closed {
	conn.shutdown()
	return
    }

    if node.conns.Len() >= node.keepAlive || node.aliveTime <= 0 {
	conn.shutdown()
	return
    }

    conn.t = time.Now()
    node.conns.PushFront(conn)
}

func (conn *redisConn) shutdown() {
    conn.c.Close()
}

func (node *redisNode) shutdown() {
    node.mutex.Lock()
    defer node.mutex.Unlock()

    for {
	elem := node.conns.Back()
	if elem == nil {
	    break
	}

	conn := elem.Value.(*redisConn)
	conn.c.Close()
	node.conns.Remove(elem)
    }

    node.closed = true
}

func (conn *redisConn) send(cmd string, args ...interface{}) error {
    conn.pending += 1

    if conn.writeTimeout > 0 {
	conn.c.SetWriteDeadline(time.Now().Add(conn.writeTimeout))
    }

    if err := conn.writeCommand(cmd, args); err != nil {
	return err
    }

    return nil
}

func (conn *redisConn) flush() error {
    if conn.writeTimeout > 0 {
	conn.c.SetWriteDeadline(time.Now().Add(conn.writeTimeout))
    }

    if err := conn.bw.Flush(); err != nil {
	return err
    }

    return nil
}

func (conn *redisConn) receive() (interface{}, error) {
    if conn.readTimeout > 0 {
	conn.c.SetWriteDeadline(time.Now().Add(conn.readTimeout))
    }

    if conn.pending <= 0 {
	return nil, errors.New("no more pending reply")
    }

    conn.pending -= 1

    return conn.readReply()
}

func (node *redisNode) do(cmd string, args ...interface{}) (interface{}, error) {
    conn, err := node.getConn()
    if err != nil {
	return redisError("ECONNTIMEOUT"), nil
    }

    if err = conn.send(cmd, args...); err != nil {
	conn.shutdown()
	return nil, err
    }

    if err = conn.flush(); err != nil {
	conn.shutdown()
	return nil, err
    }

    reply, err := conn.receive()
    if err != nil {
	conn.shutdown()
	return nil, err
    }

    node.releaseConn(conn)

    return reply, err
}

func (conn *redisConn) writeLen(prefix byte, n int) error {
    conn.lenScratch[len(conn.lenScratch) - 1] = '\n'
    conn.lenScratch[len(conn.lenScratch) - 2] = '\r'
    i := len(conn.lenScratch) - 3

    for {
        conn.lenScratch[i] = byte('0' + n % 10)
        i -= 1
        n = n / 10
        if n == 0 {
	    break
        }
    }

    conn.lenScratch[i] = prefix
    _, err := conn.bw.Write(conn.lenScratch[i:])

    return err
}

func (conn *redisConn) writeString(s string) error {
    conn.writeLen('$', len(s))
    conn.bw.WriteString(s)
    _, err := conn.bw.WriteString("\r\n")

    return err
}

func (conn *redisConn) writeBytes(p []byte) error {
    conn.writeLen('$', len(p))
    conn.bw.Write(p)
    _, err := conn.bw.WriteString("\r\n")

    return err
}

func (conn *redisConn) writeInt64(n int64) error {
    return conn.writeBytes(strconv.AppendInt(conn.numScratch[:0], n, 10))
}

func (conn *redisConn) writeFloat64(n float64) error {
    return conn.writeBytes(strconv.AppendFloat(conn.numScratch[:0], n, 'g', -1, 64))
}

// Args must be int64, float64, string, []byte, other types are not supported for safe reason.
func (conn *redisConn) writeCommand(cmd string, args []interface{}) error {
    conn.writeLen('*', len(args) + 1)
    err := conn.writeString(cmd)

    for _, arg := range args {
	if err != nil {
		break
	}
	switch arg := arg.(type) {
	case int:
	    err = conn.writeInt64(int64(arg))
	case int64:
	    err = conn.writeInt64(arg)
	case float64:
	    err = conn.writeFloat64(arg)
	case string:
	    err = conn.writeString(arg)
	case []byte:
	    err = conn.writeBytes(arg)
	default:
	    err = fmt.Errorf("unknown type %T", arg)
	}
    }

    return err
}

// readLine read a single line terminated with CRLF.
func (conn *redisConn) readLine() ([]byte, error) {
    var line []byte
    for {
	p, err := conn.br.ReadBytes('\n')
	if err != nil {
	    return nil, err
	}

	n := len(p) - 2
	if n < 0 {
	    return nil, errors.New("invalid response")
	}

	// bulk string may contain '\n', such as CLUSTER NODES
	if p[n] != '\r' {
	    if line != nil {
		line = append(line, p[:]...)
	    } else {
		line = p
	    }
	    continue
	}

	if line != nil {
	    return append(line, p[:n]...), nil
	} else {
	    return p[:n], nil
	}
    }
}

// parseLen parses bulk string and array length.
func parseLen(p []byte) (int, error) {
    if len(p) == 0 {
	return -1, errors.New("invalid response")
    }

    // null element.
    if p[0] == '-' && len(p) == 2 && p[1] == '1' {
	return -1, nil
    }

    var n int
    for _, b := range p {
	n *= 10
	if b < '0' || b > '9' {
		return -1, errors.New("invalid response")
	}
	n += int(b - '0')
    }

    return n, nil
}

// parseInt parses an integer reply.
func parseInt(p []byte) (int64, error) {
    if len(p) == 0 {
	return 0, errors.New("invalid response")
    }

    var negate bool
    if p[0] == '-' {
	negate = true
	p = p[1:]
	if len(p) == 0 {
	    return 0, errors.New("invalid response")
	}
    }

    var n int64
    for _, b := range p {
	n *= 10
	if b < '0' || b > '9' {
		return 0, errors.New("invalid response")
	}
	n += int64(b - '0')
    }

    if negate {
	n = -n
    }

    return n, nil
}

var (
    okReply   interface{} = "OK"
    pongReply interface{} = "PONG"
)

type redisError string
func (err redisError) Error() string { return string(err) }

func (conn *redisConn) readReply() (interface{}, error) {
    line, err := conn.readLine()
    if err != nil {
	return nil, err
    }
    if len(line) == 0 {
	return nil, errors.New("invalid reponse")
    }

    switch line[0] {
    case '+':
	switch {
	case len(line) == 3 && line[1] == 'O' && line[2] == 'K':
	    // Avoid allocation for frequent "+OK" response.
	    return okReply, nil
	case len(line) == 5 && line[1] == 'P' && line[2] == 'O' && line[3] == 'N' && line[4] == 'G':
	    // Avoid allocation in PING command benchmarks :)
	    return pongReply, nil
	default:
	    return string(line[1:]), nil
	}
    case '-':
	return redisError(string(line[1:])), nil
    case ':':
	return parseInt(line[1:])
    case '$':
	n, err := parseLen(line[1:])
	if n < 0 || err != nil {
	    return nil, err
	}

	line, err = conn.readLine()
	if err != nil {
	    return nil, err
	}
	if len(line) != n {
	    return nil, errors.New("invalid response")
	}

	return line, nil
    case '*':
	n, err := parseLen(line[1:])
	if n < 0 || err != nil {
	    return nil, err
	}

	r := make([]interface{}, n)
	for i := range r {
	    r[i], err = conn.readReply()
	    if err != nil {
		return nil, err
	    }
	}

	return r, nil
    }

    return nil, errors.New("invalid response")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"path"
	"time"

	docker "github.com/fsouza/go-dockerclient"
)

func main() {
	config := GetConfig()
	fmt.Println(config)

	// Init the client
	client, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		log.Fatal(err.Error())
	}

	toPause, err := getRunningContainers(client)
	if err != nil {
		log.Fatal(err.Error())
	}
	err = pauseContainers(client, toPause)
	if err != nil {
		log.Fatal(err.Error())
	}
	log.Printf("Paused %d cointainers.\n", len(toPause))

	volumes := getDataVolumes(client)

	// set destination for all backups
	dst := "/tmp/"
	t := time.Now().UTC()
	datetime := fmt.Sprintf("%d-%02d-%02d_%02d-%02d-%02d", t.Year(), t.Month(), t.Day(), t.Hour(), t.Minute(), t.Second())
	dst += datetime + "/"
	err = os.Mkdir(dst, 0700)
	if err != nil {
		log.Fatalf("Could not create backup destination: %s", dst)
	}

	log.Printf("%d data volumes to backup in %s .", len(volumes), dst)

	for _, vol := range volumes {
		file, err := backupVolume(client, dst, vol)
		if err != nil {
			log.Fatal(err.Error())
		}
		log.Printf(" * %s", path.Base(file))
	}

	err = unpauseContainers(client, toPause)
	if err != nil {
		log.Fatal(err.Error())
	}
	log.Printf("Unpaused %d cointainers.\n", len(toPause))
}

func getRunningContainers(cli *docker.Client) ([]string, error) {
	containers, err := cli.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		return nil, err
	}

	ids := make([]string, 0)
	for _, cont := range containers {
		if isRunning(cont) {
			ids = append(ids, cont.ID)
		}
	}

	return ids, nil
}

func getDataVolumes(cli *docker.Client) []Volume {
	containers, err := cli.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Fatal(err)
	}

	vols := make([]Volume, 0, 0)

	for _, cont := range containers {
		c := cont
		if isVolumeContainer(&c) {
			moreVols := getVolumes(cli, &c)
			vols = append(vols, moreVols...)
		}
	}

	return vols
}
<commit_msg>wipe archive dir when done<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"path"
	"time"

	docker "github.com/fsouza/go-dockerclient"
)

func main() {
	config := GetConfig()
	fmt.Println(config)

	// Init the client
	client, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		log.Fatal(err.Error())
	}

	toPause, err := getRunningContainers(client)
	if err != nil {
		log.Fatal(err.Error())
	}
	err = pauseContainers(client, toPause)
	if err != nil {
		log.Fatal(err.Error())
	}
	log.Printf("Paused %d cointainers.\n", len(toPause))

	volumes := getDataVolumes(client)

	// set destination for all backups
	dst := "/tmp/"
	t := time.Now().UTC()
	datetime := fmt.Sprintf("%d-%02d-%02d_%02d-%02d-%02d", t.Year(), t.Month(), t.Day(), t.Hour(), t.Minute(), t.Second())
	dst += datetime + "/"
	err = os.Mkdir(dst, 0700)
	if err != nil {
		log.Fatalf("Could not create backup destination: %s", dst)
	}

	wipetmp := func() {
		err := os.RemoveAll(dst)
		if err != nil {
			log.Fatalf("Failed to remove archives at %s", dst)
		}
		log.Printf("Removed archives at %s", dst)
	}
	defer wipetmp()

	log.Printf("%d data volumes to backup in %s .", len(volumes), dst)

	for _, vol := range volumes {
		file, err := backupVolume(client, dst, vol)
		if err != nil {
			wipetmp()
			log.Fatal(err.Error())
		}
		log.Printf(" * %s", path.Base(file))
	}

	err = unpauseContainers(client, toPause)
	if err != nil {
		wipetmp()
		log.Fatal(err.Error())
	}
	log.Printf("Unpaused %d cointainers.\n", len(toPause))
}

func getRunningContainers(cli *docker.Client) ([]string, error) {
	containers, err := cli.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		return nil, err
	}

	ids := make([]string, 0)
	for _, cont := range containers {
		if isRunning(cont) {
			ids = append(ids, cont.ID)
		}
	}

	return ids, nil
}

func getDataVolumes(cli *docker.Client) []Volume {
	containers, err := cli.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Fatal(err)
	}

	vols := make([]Volume, 0, 0)

	for _, cont := range containers {
		c := cont
		if isVolumeContainer(&c) {
			moreVols := getVolumes(cli, &c)
			vols = append(vols, moreVols...)
		}
	}

	return vols
}
<|endoftext|>


<commit_before>package main

import (
	ps "github.com/mitchellh/go-ps"
	"log"
	"os/exec"
	"regexp"
	"strconv"
	"github.com/andlabs/ui"
	"time"
	"net/http"
	"os"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"syscall"
	"strings"
)

// Telegram and program settings (config.yml)
// Guide: http://sweetohm.net/article/go-yaml-parsers.en.html
type Config struct {
	Token string
	Botid string
	Chatid string
	Message string
	StayAlive bool
	Process string
	TimeBetweenChecksInS int
	KillOnDC bool
	ShutdownOnDC bool
	KillCoherentUI bool
}

// Variables
var STATUS bool = false
var CONNECTION bool = false
var PID int

func main() {

	////  SETTINGS
	//--------------------------------------------------------------------------------------------------------------
	// YAML PARSING
	ex, err := os.Executable()
	if err != nil {
		panic(err)
	}
	newex := strings.Replace(ex, "BDO-Watchdog.exe", "config.yml", -1)
	// This is necessary for dynamic builds in Jetbrains Gogland IDE
	newex = strings.Replace(ex, "Application.exe", "config.yml", -1)

	var config Config
	source, err := ioutil.ReadFile(newex)

	if err != nil {
		// in theory, using yml.Marshal() would be more elegant, but we want to preserve the yaml comments
		// as well as set some default values/hints
		defconf :=
				"## Telegram Bot Settings\r\n" +
				"token: \r\n" +
				"botid: \r\n" +
				"chatid: \r\n" +
				"message: BDO disconnected \r\n" +
				"\r\n" +
				"## Program Settings\r\n" +
				"stayalive: false\r\n" +
				"process: BlackDesert64.exe\r\n" +
				"timebetweenchecksins: 60\r\n" +
				"\r\n" +
				"# These settings require the .exe to be run with admin rights! \r\n" +
				"killondc: true\r\n" +
				"shutdownondc: false\r\n" +
				"killcoherentui: false"
		ioutil.WriteFile("config.yml", []byte(defconf), os.FileMode(int(0777)))
		panic(err)
	}
	err = yaml.Unmarshal(source, &config)
	if err != nil {
		panic(err)
	}

	//// GUI
	//--------------------------------------------------------------------------------------------------------------
	ui := ui.Main(func() {
		window := ui.NewWindow("BDO Watchdog", 300, 80, false)

		label_Process := ui.NewLabel("  Process: " + config.Process)
		label_Status := ui.NewLabel("  Initializing...")
		label_PID := ui.NewLabel("-")
		label_Connection := ui.NewLabel("-")
		label_Update := ui.NewLabel("")

		box := ui.NewVerticalBox()
		sep := ui.NewHorizontalSeparator()
		pb := ui.NewProgressBar()

		// Append all UI elements to the box container
		box.Append(label_Process, false)
		box.Append(label_Status, false)
		box.Append(label_PID, false)
		box.Append(label_Connection, false)
		box.Append(label_Update, false)
		box.Append(sep, false)
		box.Append(pb, true)

		// If this tool requires admin rights add a text label to inform the user
		if config.KillOnDC || config.KillCoherentUI {
			box.Append(ui.NewLabel("Make sure to have started this program with admin rights!"), true)
		}

		window.SetChild(box)
		window.OnClosing(func(*ui.Window) bool {
			ui.Quit()
			return true
		})
		window.Show()
		go observer(config, label_Status, label_PID, label_Connection, label_Update, pb)
	})
	if ui != nil {
		panic(ui)
	}
}

//--------------------------------------------------------------------------------------------------------------
// PROCESS
//--------------------------------------------------------------------------------------------------------------
func observer(
	config Config,
	label_Status *ui.Label,
	label_PID *ui.Label,
	label_Connection *ui.Label,
	label_Update *ui.Label,
	pb *ui.ProgressBar) {

	// KILL CoherentUI_Host.exe
	if config.KillCoherentUI {

		// Find process(es)
		chp, err := ps.Processes()
		if err != nil {
			log.Fatal(err)
		}

		// Find PID and kill
		for _, v := range chp {
			if v.Executable() == "CoherentUI_Host.exe" {
				proc, err := os.FindProcess(v.Pid())

				if err != nil {
					log.Println(err)
				}
				// Kill the process
				defer proc.Kill()
			}
		}
	}

	// INFINITE MAIN LOOP
	for {
		label_Update.SetText("")

		//// EXIT CONDITION
		//-----------------
		// If the process is running, but no longer connected we trigger the following actions
		if STATUS && !CONNECTION {

			// Use the Telegram API to send a message
			send_TelegramMessage(config)

			// Optional: shutdown the computer if the monitored process is disconnected
			if config.ShutdownOnDC {
				exec.Command("cmd", "/C", "shutdown", "/s").Run()
			}

			// Optional: kill the monitored process if it is disconnected
			// requires elevated rights --> start .exe as administrator
			if config.KillOnDC {

				proc, err := os.FindProcess(PID)

				if err != nil {
					log.Println(err)
				}
				// Kill the process
				defer proc.Kill()

				time.Sleep(5 * time.Second)
			}

			// Optional (YAML file, default: false): keep this program open even if
			// the process is disconnected
			if !config.StayAlive {
				os.Exit(1)
			}
		}

		//// PROCESS
		//----------
		p, err := ps.Processes()
		if err != nil {
			log.Fatal(err)
		}

		//// PID
		//------
		for _, v := range p {
			if v.Executable() == config.Process {
				PID = v.Pid()
			}
		}
		if (PID == 0) {
			ui.QueueMain(func () {
				STATUS = false
				label_Status.SetText("  Status: not running")
				label_PID.SetText("  PID: -")
				label_Connection.SetText("  Connection: -" )
			})

			wait(config, label_Update, pb)
			continue
		} else {

			ui.QueueMain(func () {
				STATUS = true
				label_Status.SetText("  Status: running")
				label_PID.SetText("  PID: " + strconv.Itoa(PID))
			})
		}

		//// CONNECTION STATUS
		//--------------------
		// NETSTAT
		// the syscall.SysProcAttr trick found here:
		// https://www.reddit.com/r/golang/comments/2c1g3x/build_golang_app_reverse_shell_to_run_in_windows/
		cmd := exec.Command("cmd.exe", "/C netstat -aon")
		cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
		out, err := cmd.Output()
		if err != nil {
			log.Fatal(err)
		}

		// RegEx matching; try to find the PID in the netstat output
		re := regexp.MustCompile(strconv.Itoa(PID))
		byteIndex := re.FindIndex([]byte(out))

		if (len(byteIndex) == 0) {
			ui.QueueMain(func () {
				CONNECTION = false
				label_Connection.SetText("  Connection: Offline" )
			})
		} else {
			// Update labels
			ui.QueueMain(func () {
				CONNECTION = true
				label_Connection.SetText("  Connection: online")
			})
		}

		// Wait x seconds before next iteration
		wait(config, label_Update, pb)
	}
}

// ---------------------------------------------------------------------------------------------------------------------
// A wrapper for time.Sleep() that also updates the UI label and progressbar
func wait(config Config, label_Update *ui.Label, pb *ui.ProgressBar) {
	tstep := config.TimeBetweenChecksInS
	if tstep <= 0 {
		tstep = 1
	} // otherwise division by 0
	for i := 0; i <= tstep; i++ {
		pb.SetValue(int(100/float32(tstep) * float32(i + 1)))
		label_Update.SetText("  Next update in... " + strconv.Itoa(tstep - i) + " s")
		time.Sleep(1 * time.Second)
	}
	pb.SetValue(0)
}

// ---------------------------------------------------------------------------------------------------------------------
// Send a telegram message using a query URL
func send_TelegramMessage(config Config) {
	// Learn how to setup a telegram bot: https://core.telegram.org/bots
	resp, _ := http.Get("https://api.telegram.org/bot" + config.Botid +
		":" + config.Token +
		"/sendMessage?chat_id=" + config.Chatid +
		"&text=" + config.Message)
	defer resp.Body.Close()
}
<commit_msg>Uncomment a line only used for local building<commit_after>package main

import (
	ps "github.com/mitchellh/go-ps"
	"log"
	"os/exec"
	"regexp"
	"strconv"
	"github.com/andlabs/ui"
	"time"
	"net/http"
	"os"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"syscall"
	"strings"
)

// Telegram and program settings (config.yml)
// Guide: http://sweetohm.net/article/go-yaml-parsers.en.html
type Config struct {
	Token string
	Botid string
	Chatid string
	Message string
	StayAlive bool
	Process string
	TimeBetweenChecksInS int
	KillOnDC bool
	ShutdownOnDC bool
	KillCoherentUI bool
}

// Variables
var STATUS bool = false
var CONNECTION bool = false
var PID int

func main() {

	////  SETTINGS
	//--------------------------------------------------------------------------------------------------------------
	// YAML PARSING
	ex, err := os.Executable()
	if err != nil {
		panic(err)
	}
	newex := strings.Replace(ex, "BDO-Watchdog.exe", "config.yml", -1)
	// This is necessary for dynamic builds in Jetbrains Gogland IDE
	//newex = strings.Replace(ex, "Application.exe", "config.yml", -1)

	var config Config
	source, err := ioutil.ReadFile(newex)

	if err != nil {
		// in theory, using yml.Marshal() would be more elegant, but we want to preserve the yaml comments
		// as well as set some default values/hints
		defconf :=
				"## Telegram Bot Settings\r\n" +
				"token: \r\n" +
				"botid: \r\n" +
				"chatid: \r\n" +
				"message: BDO disconnected \r\n" +
				"\r\n" +
				"## Program Settings\r\n" +
				"stayalive: false\r\n" +
				"process: BlackDesert64.exe\r\n" +
				"timebetweenchecksins: 60\r\n" +
				"\r\n" +
				"# These settings require the .exe to be run with admin rights! \r\n" +
				"killondc: true\r\n" +
				"shutdownondc: false\r\n" +
				"killcoherentui: false"
		ioutil.WriteFile("config.yml", []byte(defconf), os.FileMode(int(0777)))
		panic(err)
	}
	err = yaml.Unmarshal(source, &config)
	if err != nil {
		panic(err)
	}

	//// GUI
	//--------------------------------------------------------------------------------------------------------------
	ui := ui.Main(func() {
		window := ui.NewWindow("BDO Watchdog", 300, 80, false)

		label_Process := ui.NewLabel("  Process: " + config.Process)
		label_Status := ui.NewLabel("  Initializing...")
		label_PID := ui.NewLabel("-")
		label_Connection := ui.NewLabel("-")
		label_Update := ui.NewLabel("")

		box := ui.NewVerticalBox()
		sep := ui.NewHorizontalSeparator()
		pb := ui.NewProgressBar()

		// Append all UI elements to the box container
		box.Append(label_Process, false)
		box.Append(label_Status, false)
		box.Append(label_PID, false)
		box.Append(label_Connection, false)
		box.Append(label_Update, false)
		box.Append(sep, false)
		box.Append(pb, true)

		// If this tool requires admin rights add a text label to inform the user
		if config.KillOnDC || config.KillCoherentUI {
			box.Append(ui.NewLabel("Make sure to have started this program with admin rights!"), true)
		}

		window.SetChild(box)
		window.OnClosing(func(*ui.Window) bool {
			ui.Quit()
			return true
		})
		window.Show()
		go observer(config, label_Status, label_PID, label_Connection, label_Update, pb)
	})
	if ui != nil {
		panic(ui)
	}
}

//--------------------------------------------------------------------------------------------------------------
// PROCESS
//--------------------------------------------------------------------------------------------------------------
func observer(
	config Config,
	label_Status *ui.Label,
	label_PID *ui.Label,
	label_Connection *ui.Label,
	label_Update *ui.Label,
	pb *ui.ProgressBar) {

	// KILL CoherentUI_Host.exe
	if config.KillCoherentUI {

		// Find process(es)
		chp, err := ps.Processes()
		if err != nil {
			log.Fatal(err)
		}

		// Find PID and kill
		for _, v := range chp {
			if v.Executable() == "CoherentUI_Host.exe" {
				proc, err := os.FindProcess(v.Pid())

				if err != nil {
					log.Println(err)
				}
				// Kill the process
				defer proc.Kill()
			}
		}
	}

	// INFINITE MAIN LOOP
	for {
		label_Update.SetText("")

		//// EXIT CONDITION
		//-----------------
		// If the process is running, but no longer connected we trigger the following actions
		if STATUS && !CONNECTION {

			// Use the Telegram API to send a message
			send_TelegramMessage(config)

			// Optional: shutdown the computer if the monitored process is disconnected
			if config.ShutdownOnDC {
				exec.Command("cmd", "/C", "shutdown", "/s").Run()
			}

			// Optional: kill the monitored process if it is disconnected
			// requires elevated rights --> start .exe as administrator
			if config.KillOnDC {

				proc, err := os.FindProcess(PID)

				if err != nil {
					log.Println(err)
				}
				// Kill the process
				defer proc.Kill()

				time.Sleep(5 * time.Second)
			}

			// Optional (YAML file, default: false): keep this program open even if
			// the process is disconnected
			if !config.StayAlive {
				os.Exit(1)
			}
		}

		//// PROCESS
		//----------
		p, err := ps.Processes()
		if err != nil {
			log.Fatal(err)
		}

		//// PID
		//------
		for _, v := range p {
			if v.Executable() == config.Process {
				PID = v.Pid()
			}
		}
		if (PID == 0) {
			ui.QueueMain(func () {
				STATUS = false
				label_Status.SetText("  Status: not running")
				label_PID.SetText("  PID: -")
				label_Connection.SetText("  Connection: -" )
			})

			wait(config, label_Update, pb)
			continue
		} else {

			ui.QueueMain(func () {
				STATUS = true
				label_Status.SetText("  Status: running")
				label_PID.SetText("  PID: " + strconv.Itoa(PID))
			})
		}

		//// CONNECTION STATUS
		//--------------------
		// NETSTAT
		// the syscall.SysProcAttr trick found here:
		// https://www.reddit.com/r/golang/comments/2c1g3x/build_golang_app_reverse_shell_to_run_in_windows/
		cmd := exec.Command("cmd.exe", "/C netstat -aon")
		cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
		out, err := cmd.Output()
		if err != nil {
			log.Fatal(err)
		}

		// RegEx matching; try to find the PID in the netstat output
		re := regexp.MustCompile(strconv.Itoa(PID))
		byteIndex := re.FindIndex([]byte(out))

		if (len(byteIndex) == 0) {
			ui.QueueMain(func () {
				CONNECTION = false
				label_Connection.SetText("  Connection: Offline" )
			})
		} else {
			// Update labels
			ui.QueueMain(func () {
				CONNECTION = true
				label_Connection.SetText("  Connection: online")
			})
		}

		// Wait x seconds before next iteration
		wait(config, label_Update, pb)
	}
}

// ---------------------------------------------------------------------------------------------------------------------
// A wrapper for time.Sleep() that also updates the UI label and progressbar
func wait(config Config, label_Update *ui.Label, pb *ui.ProgressBar) {
	tstep := config.TimeBetweenChecksInS
	if tstep <= 0 {
		tstep = 1
	} // otherwise division by 0
	for i := 0; i <= tstep; i++ {
		pb.SetValue(int(100/float32(tstep) * float32(i + 1)))
		label_Update.SetText("  Next update in... " + strconv.Itoa(tstep - i) + " s")
		time.Sleep(1 * time.Second)
	}
	pb.SetValue(0)
}

// ---------------------------------------------------------------------------------------------------------------------
// Send a telegram message using a query URL
func send_TelegramMessage(config Config) {
	// Learn how to setup a telegram bot: https://core.telegram.org/bots
	resp, _ := http.Get("https://api.telegram.org/bot" + config.Botid +
		":" + config.Token +
		"/sendMessage?chat_id=" + config.Chatid +
		"&text=" + config.Message)
	defer resp.Body.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"

	flag "github.com/dotcloud/docker/pkg/mflag"

	"github.com/reiki4040/cstore"
	"github.com/reiki4040/peco"
)

const (
	Usage = `rnssh - easy ssh login to EC2.

usage:

  rnssh [-f] [-p] [-s] [user@]query strings ...
  rnssh --init

options:
  -f: reload ec2 instances infomaion. connect to AWS.
      you have to specify this option after you modified instances.
	  (launch, start, stop, modify name, etc...)

  -P: use Public IP address. this is default ssh host type.
  -p: use Private IP address. for VPN/Direct connect.
  -n: use Name tag.
      this option for ssh config that Host named by ec2 Name tag.

  -r: target region. you can set default by --init (~/.rnssh/config)

  -s: show ssh command string that would be run. (debug)

  --init: start wizard for default setting AWS region and rnssh host type.
          and save to config file (~/.rnssh/config)

options for ssh:
  -l: ssh user.
  -i: identity file path.
  --port: ssh port.
  --strict-host-key-checking-no: suppress host key checking.
                                 1: using 'ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
                                 0: OFF
                                -1: Default(OFF)

options for help:
  -h: show this usage.
  -v: show version.

args:
  query string...: filtering ec2 instances list.
`

	ENV_AWS_REGION      = "AWS_REGION"
	ENV_RNSSH_HOST_TYPE = "RNSSH_HOST_TYPE"

	ENV_HOME = "HOME"

	RNSSH_DIR_NAME = ".rnssh"
)

type CommandOption struct {
	Reload                  bool
	Region                  string
	PrivateIP               bool
	PublicIP                bool
	NameTag                 bool
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
	UseSshConfig            bool
	UseEC2                  bool
}

func (o *CommandOption) Validate() error {
	if err := duplicateHostTypeOption(o.PublicIP, o.PrivateIP, o.NameTag); err != nil {
		return err
	}

	if err := IdentityFileCheck(o.IdentityFile); err != nil {
		return err
	}

	if err := StrictHostKeyCheckingNoCheck(o.StrictHostKeyCheckingNo); err != nil {
		return err
	}

	if o.UseSshConfig && o.UseEC2 {
		return fmt.Errorf("can not specify both --use-ssh-config and --use-ec2")
	}

	return nil
}

func duplicateHostTypeOption(public, private, name bool) error {
	if public && private || private && name || public && name {
		return fmt.Errorf("duplicate specify option -P/-p/-n. please spcify only one.")
	}

	return nil
}

type RnsshOption struct {
	Reload                  bool
	Region                  string
	HostType                string
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
	UseSshConfig            bool
}

var (
	version   string
	hash      string
	builddate string

	show_version bool
	show_usage   bool
	initWizard   bool
	showCommand  bool

	// command option
	opt = &CommandOption{}
)

func init() {
	flag.BoolVar(&show_version, []string{"v", "-version"}, false, "show version.")
	flag.BoolVar(&show_usage, []string{"h", "-help"}, false, "show this usage.")
	flag.BoolVar(&initWizard, []string{"-init"}, false, "run initial configuration wizard.")

	flag.BoolVar(&opt.Reload, []string{"f", "-force"}, false, "reload ec2 (force connect to AWS)")
	flag.BoolVar(&opt.PublicIP, []string{"P", "-public-ip"}, false, "ssh with EC2 Public IP")
	flag.BoolVar(&opt.PrivateIP, []string{"p", "-private-ip"}, false, "ssh with EC2 Private IP")
	flag.BoolVar(&opt.NameTag, []string{"n", "-name-tag"}, false, "ssh with EC2 Name tag")
	flag.BoolVar(&showCommand, []string{"s", "-show-command"}, false, "show ssh command that will do (debug)")

	flag.StringVar(&opt.Region, []string{"r", "-region"}, "", "specify region")
	flag.StringVar(&opt.SshUser, []string{"l", "-user"}, "", "specify ssh user")
	flag.StringVar(&opt.IdentityFile, []string{"i", "-identity-file"}, "", "specify ssh identity file")
	flag.IntVar(&opt.Port, []string{"-port"}, 0, "specify ssh port")
	flag.IntVar(&opt.StrictHostKeyCheckingNo, []string{"-strict-host-key-checking-no"}, -1, "suppress host key checking. 1: ON, 0: OFF, -1: default(OFF)")

	flag.BoolVar(&opt.UseSshConfig, []string{"-use-ssh-config"}, false, "load from ssh config")
	flag.BoolVar(&opt.UseEC2, []string{"-use-ec2"}, false, "load from ec2")

	flag.Parse()
}

func showVersion() {
	fmt.Printf("%s (%s) built:%s\n", version, hash, builddate)
}

func usage() {
	fmt.Printf("%s\n", Usage)
}

func main() {
	if show_usage {
		usage()
		os.Exit(0)
	}

	if show_version {
		showVersion()
		os.Exit(0)
	}

	err := opt.Validate()
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	m, err := cstore.NewManager("rnssh", getRnsshDir())
	if err != nil {
		fmt.Printf("can not create rnssh dir: %s\n", err.Error())
		os.Exit(1)
	}

	cs, err := m.New("config", cstore.TOML)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if initWizard {
		if err := DoConfigWizard(cs); err != nil {
			fmt.Println(err)
			os.Exit(1)
		} else {
			fmt.Println("saved rnssh config.")
			os.Exit(0)
		}
	}

	conf := Config{}
	err = cs.Get(&conf)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	rOpt := mergeConfig(&conf.Default, *opt)
	if !rOpt.UseSshConfig && rOpt.Region == "" {
		fmt.Println("region is empty. please specify by region option (-r) or AWS_REGION envirnment variable.")
		os.Exit(1)
	}

	sshArgs, err := chooseAndGenSshArgs(rOpt, flag.Args(), m)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if showCommand {
		fmt.Printf("%s %s\n", "ssh", strings.Join(sshArgs, " "))
		os.Exit(0)
	}

	cmd := exec.Command("ssh", sshArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}

// merge option, config, ENV
// priority [high] option > config > ENV [low]
func mergeConfig(conf *RnsshConfig, opt CommandOption) *RnsshOption {

	region := os.Getenv(ENV_AWS_REGION)
	if opt.Region != "" {
		region = opt.Region
	} else {
		if conf.AWSRegion != "" {
			region = conf.AWSRegion
		}
	}

	hostType := os.Getenv(ENV_RNSSH_HOST_TYPE)
	optHostType := getSshTargetType(opt.PublicIP, opt.PrivateIP, opt.NameTag)
	if optHostType != "" {
		hostType = optHostType
	} else {
		if conf.HostType != "" {
			hostType = conf.HostType
		}
	}

	sshUser := conf.SshUser
	if opt.SshUser != "" {
		sshUser = opt.SshUser
	}

	identityFile := conf.SshIdentityFile
	if opt.IdentityFile != "" {
		identityFile = opt.IdentityFile
	}

	port := conf.SshPort
	if opt.Port > 0 {
		port = opt.Port
	}

	strictHostKeyCheckingNo := conf.SshStrictHostKeyCheckingNo
	if opt.StrictHostKeyCheckingNo != -1 {
		strictHostKeyCheckingNo = opt.StrictHostKeyCheckingNo
	}

	useSshConfig := conf.UseSshConfig
	if opt.UseSshConfig {
		useSshConfig = true
	}

	if opt.UseEC2 {
		useSshConfig = false
	}

	return &RnsshOption{
		Reload:       opt.Reload,
		Region:       region,
		HostType:     hostType,
		SshUser:      sshUser,
		IdentityFile: identityFile,
		Port:         port,
		StrictHostKeyCheckingNo: strictHostKeyCheckingNo,
		UseSshConfig:            useSshConfig,
	}
}

func chooseAndGenSshArgs(rOpt *RnsshOption, cmdArgs []string, manager *cstore.Manager) ([]string, error) {

	// support user@host format
	sshUser, hostname, err := getSshUserAndHostname(strings.Join(cmdArgs, " "))
	if err != nil {
		return nil, fmt.Errorf("%s\n", err.Error())
	}

	hostType := HOST_TYPE_PUBLIC_IP
	if rOpt.HostType != "" {
		hostType = rOpt.HostType
	}

	var choosableList []peco.Choosable
	if rOpt.UseSshConfig {
		var err error
		choosableList, err = LoadSshConfigChoosableList()
		if err != nil {
			fmt.Printf("%s\n", err.Error())
			os.Exit(1)
		}

		if len(choosableList) == 0 {
			fmt.Println("ssh config does not have host settings")
			os.Exit(1)
		}
	} else {
		var err error
		handler := NewEC2Handler(manager)
		choosableList, err = handler.LoadTargetHost(hostType, rOpt.Region, rOpt.Reload)
		if err != nil {
			fmt.Printf("%s\n", err.Error())
			os.Exit(1)
		}

		if len(choosableList) == 0 {
			fmt.Printf("there is no instance. not running %s\n", rOpt.Region)
			os.Exit(1)
		}
	}

	// show ec2 instances and choose intactive
	targetHosts, err := peco.Choose("server", "which servers connect with ssh?", hostname, choosableList)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	l := len(targetHosts) - 1
	targetHost := targetHosts[l]
	sshHost := targetHost.Value()
	sshArgs := genSshArgs(rOpt.SshUser, rOpt.IdentityFile, rOpt.Port, rOpt.StrictHostKeyCheckingNo, sshUser, sshHost)

	return sshArgs, nil
}

func getSshUserAndHostname(sshTarget string) (string, string, error) {
	// support user@host format
	idx := strings.Index(sshTarget, "@")

	// not include(-1) or first char(0)
	if idx > 0 {
		splited := strings.SplitN(sshTarget, "@", 2)
		return splited[0], splited[1], nil
	}

	return "", sshTarget, nil
}

func getSshTargetType(publicIP, privateIP, nameTag bool) string {

	// overwrite by option
	if publicIP {
		return HOST_TYPE_PUBLIC_IP
	}

	if privateIP {
		return HOST_TYPE_PRIVATE_IP
	}

	if nameTag {
		return HOST_TYPE_NAME_TAG
	}

	return ""
}

func genSshArgs(optSshUser, optIdentityFile string, optPort, optStrictHostKeyCheckingNo int, sshUser, sshHost string) []string {
	args := make([]string, 0)
	if optSshUser != "" {
		args = append(args, "-l"+optSshUser)
	}

	if optIdentityFile != "" {
		args = append(args, "-i"+optIdentityFile)
	}

	if optPort > 0 {
		args = append(args, "-p"+strconv.Itoa(optPort))
	}

	if optStrictHostKeyCheckingNo == 1 {
		args = append(args, "-oStrictHostKeyChecking=no")
		args = append(args, "-oUserKnownHostsFile=/dev/null")
	}

	if sshUser != "" {
		sshHost = sshUser + "@" + sshHost
	}

	args = append(args, sshHost)

	return args
}

func getRnsshDir() string {
	rnsshDir := os.Getenv(ENV_HOME) + string(os.PathSeparator) + RNSSH_DIR_NAME
	return rnsshDir
}
<commit_msg>little fix error message.<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"

	flag "github.com/dotcloud/docker/pkg/mflag"

	"github.com/reiki4040/cstore"
	"github.com/reiki4040/peco"
)

const (
	Usage = `rnssh - easy ssh login to EC2.

usage:

  rnssh [-f] [-p] [-s] [user@]query strings ...
  rnssh --init

options:
  -f: reload ec2 instances infomaion. connect to AWS.
      you have to specify this option after you modified instances.
	  (launch, start, stop, modify name, etc...)

  -P: use Public IP address. this is default ssh host type.
  -p: use Private IP address. for VPN/Direct connect.
  -n: use Name tag.
      this option for ssh config that Host named by ec2 Name tag.

  -r: target region. you can set default by --init (~/.rnssh/config)

  -s: show ssh command string that would be run. (debug)

  --init: start wizard for default setting AWS region and rnssh host type.
          and save to config file (~/.rnssh/config)

options for ssh:
  -l: ssh user.
  -i: identity file path.
  --port: ssh port.
  --strict-host-key-checking-no: suppress host key checking.
                                 1: using 'ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
                                 0: OFF
                                -1: Default(OFF)

options for help:
  -h: show this usage.
  -v: show version.

args:
  query string...: filtering ec2 instances list.
`

	ENV_AWS_REGION      = "AWS_REGION"
	ENV_RNSSH_HOST_TYPE = "RNSSH_HOST_TYPE"

	ENV_HOME = "HOME"

	RNSSH_DIR_NAME = ".rnssh"
)

type CommandOption struct {
	Reload                  bool
	Region                  string
	PrivateIP               bool
	PublicIP                bool
	NameTag                 bool
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
	UseSshConfig            bool
	UseEC2                  bool
}

func (o *CommandOption) Validate() error {
	if err := duplicateHostTypeOption(o.PublicIP, o.PrivateIP, o.NameTag); err != nil {
		return err
	}

	if err := IdentityFileCheck(o.IdentityFile); err != nil {
		return err
	}

	if err := StrictHostKeyCheckingNoCheck(o.StrictHostKeyCheckingNo); err != nil {
		return err
	}

	if o.UseSshConfig && o.UseEC2 {
		return fmt.Errorf("can not specify both --use-ssh-config and --use-ec2")
	}

	return nil
}

func duplicateHostTypeOption(public, private, name bool) error {
	if public && private || private && name || public && name {
		return fmt.Errorf("duplicate specify option -P/-p/-n. please spcify only one.")
	}

	return nil
}

type RnsshOption struct {
	Reload                  bool
	Region                  string
	HostType                string
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
	UseSshConfig            bool
}

var (
	version   string
	hash      string
	builddate string

	show_version bool
	show_usage   bool
	initWizard   bool
	showCommand  bool

	// command option
	opt = &CommandOption{}
)

func init() {
	flag.BoolVar(&show_version, []string{"v", "-version"}, false, "show version.")
	flag.BoolVar(&show_usage, []string{"h", "-help"}, false, "show this usage.")
	flag.BoolVar(&initWizard, []string{"-init"}, false, "run initial configuration wizard.")

	flag.BoolVar(&opt.Reload, []string{"f", "-force"}, false, "reload ec2 (force connect to AWS)")
	flag.BoolVar(&opt.PublicIP, []string{"P", "-public-ip"}, false, "ssh with EC2 Public IP")
	flag.BoolVar(&opt.PrivateIP, []string{"p", "-private-ip"}, false, "ssh with EC2 Private IP")
	flag.BoolVar(&opt.NameTag, []string{"n", "-name-tag"}, false, "ssh with EC2 Name tag")
	flag.BoolVar(&showCommand, []string{"s", "-show-command"}, false, "show ssh command that will do (debug)")

	flag.StringVar(&opt.Region, []string{"r", "-region"}, "", "specify region")
	flag.StringVar(&opt.SshUser, []string{"l", "-user"}, "", "specify ssh user")
	flag.StringVar(&opt.IdentityFile, []string{"i", "-identity-file"}, "", "specify ssh identity file")
	flag.IntVar(&opt.Port, []string{"-port"}, 0, "specify ssh port")
	flag.IntVar(&opt.StrictHostKeyCheckingNo, []string{"-strict-host-key-checking-no"}, -1, "suppress host key checking. 1: ON, 0: OFF, -1: default(OFF)")

	flag.BoolVar(&opt.UseSshConfig, []string{"-use-ssh-config"}, false, "load from ssh config")
	flag.BoolVar(&opt.UseEC2, []string{"-use-ec2"}, false, "load from ec2")

	flag.Parse()
}

func showVersion() {
	fmt.Printf("%s (%s) built:%s\n", version, hash, builddate)
}

func usage() {
	fmt.Printf("%s\n", Usage)
}

func main() {
	if show_usage {
		usage()
		os.Exit(0)
	}

	if show_version {
		showVersion()
		os.Exit(0)
	}

	err := opt.Validate()
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	m, err := cstore.NewManager("rnssh", getRnsshDir())
	if err != nil {
		fmt.Printf("can not create rnssh dir: %s\n", err.Error())
		os.Exit(1)
	}

	cs, err := m.New("config", cstore.TOML)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if initWizard {
		if err := DoConfigWizard(cs); err != nil {
			fmt.Println(err)
			os.Exit(1)
		} else {
			fmt.Println("saved rnssh config.")
			os.Exit(0)
		}
	}

	conf := Config{}
	err = cs.Get(&conf)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	rOpt := mergeConfig(&conf.Default, *opt)
	if !rOpt.UseSshConfig && rOpt.Region == "" {
		fmt.Println("region is empty. please specify by region option (-r) or set default region with --init option")
		os.Exit(1)
	}

	sshArgs, err := chooseAndGenSshArgs(rOpt, flag.Args(), m)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if showCommand {
		fmt.Printf("%s %s\n", "ssh", strings.Join(sshArgs, " "))
		os.Exit(0)
	}

	cmd := exec.Command("ssh", sshArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}

// merge option, config, ENV
// priority [high] option > config > ENV [low]
func mergeConfig(conf *RnsshConfig, opt CommandOption) *RnsshOption {

	region := os.Getenv(ENV_AWS_REGION)
	if opt.Region != "" {
		region = opt.Region
	} else {
		if conf.AWSRegion != "" {
			region = conf.AWSRegion
		}
	}

	hostType := os.Getenv(ENV_RNSSH_HOST_TYPE)
	optHostType := getSshTargetType(opt.PublicIP, opt.PrivateIP, opt.NameTag)
	if optHostType != "" {
		hostType = optHostType
	} else {
		if conf.HostType != "" {
			hostType = conf.HostType
		}
	}

	sshUser := conf.SshUser
	if opt.SshUser != "" {
		sshUser = opt.SshUser
	}

	identityFile := conf.SshIdentityFile
	if opt.IdentityFile != "" {
		identityFile = opt.IdentityFile
	}

	port := conf.SshPort
	if opt.Port > 0 {
		port = opt.Port
	}

	strictHostKeyCheckingNo := conf.SshStrictHostKeyCheckingNo
	if opt.StrictHostKeyCheckingNo != -1 {
		strictHostKeyCheckingNo = opt.StrictHostKeyCheckingNo
	}

	useSshConfig := conf.UseSshConfig
	if opt.UseSshConfig {
		useSshConfig = true
	}

	if opt.UseEC2 {
		useSshConfig = false
	}

	return &RnsshOption{
		Reload:       opt.Reload,
		Region:       region,
		HostType:     hostType,
		SshUser:      sshUser,
		IdentityFile: identityFile,
		Port:         port,
		StrictHostKeyCheckingNo: strictHostKeyCheckingNo,
		UseSshConfig:            useSshConfig,
	}
}

func chooseAndGenSshArgs(rOpt *RnsshOption, cmdArgs []string, manager *cstore.Manager) ([]string, error) {

	// support user@host format
	sshUser, hostname, err := getSshUserAndHostname(strings.Join(cmdArgs, " "))
	if err != nil {
		return nil, fmt.Errorf("%s\n", err.Error())
	}

	hostType := HOST_TYPE_PUBLIC_IP
	if rOpt.HostType != "" {
		hostType = rOpt.HostType
	}

	var choosableList []peco.Choosable
	if rOpt.UseSshConfig {
		var err error
		choosableList, err = LoadSshConfigChoosableList()
		if err != nil {
			fmt.Printf("%s\n", err.Error())
			os.Exit(1)
		}

		if len(choosableList) == 0 {
			fmt.Println("ssh config does not have host settings")
			os.Exit(1)
		}
	} else {
		var err error
		handler := NewEC2Handler(manager)
		choosableList, err = handler.LoadTargetHost(hostType, rOpt.Region, rOpt.Reload)
		if err != nil {
			fmt.Printf("%s\n", err.Error())
			os.Exit(1)
		}

		if len(choosableList) == 0 {
			fmt.Printf("there is no instance. not running %s\n", rOpt.Region)
			os.Exit(1)
		}
	}

	// show ec2 instances and choose intactive
	targetHosts, err := peco.Choose("server", "which servers connect with ssh?", hostname, choosableList)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	l := len(targetHosts) - 1
	targetHost := targetHosts[l]
	sshHost := targetHost.Value()
	sshArgs := genSshArgs(rOpt.SshUser, rOpt.IdentityFile, rOpt.Port, rOpt.StrictHostKeyCheckingNo, sshUser, sshHost)

	return sshArgs, nil
}

func getSshUserAndHostname(sshTarget string) (string, string, error) {
	// support user@host format
	idx := strings.Index(sshTarget, "@")

	// not include(-1) or first char(0)
	if idx > 0 {
		splited := strings.SplitN(sshTarget, "@", 2)
		return splited[0], splited[1], nil
	}

	return "", sshTarget, nil
}

func getSshTargetType(publicIP, privateIP, nameTag bool) string {

	// overwrite by option
	if publicIP {
		return HOST_TYPE_PUBLIC_IP
	}

	if privateIP {
		return HOST_TYPE_PRIVATE_IP
	}

	if nameTag {
		return HOST_TYPE_NAME_TAG
	}

	return ""
}

func genSshArgs(optSshUser, optIdentityFile string, optPort, optStrictHostKeyCheckingNo int, sshUser, sshHost string) []string {
	args := make([]string, 0)
	if optSshUser != "" {
		args = append(args, "-l"+optSshUser)
	}

	if optIdentityFile != "" {
		args = append(args, "-i"+optIdentityFile)
	}

	if optPort > 0 {
		args = append(args, "-p"+strconv.Itoa(optPort))
	}

	if optStrictHostKeyCheckingNo == 1 {
		args = append(args, "-oStrictHostKeyChecking=no")
		args = append(args, "-oUserKnownHostsFile=/dev/null")
	}

	if sshUser != "" {
		sshHost = sshUser + "@" + sshHost
	}

	args = append(args, sshHost)

	return args
}

func getRnsshDir() string {
	rnsshDir := os.Getenv(ENV_HOME) + string(os.PathSeparator) + RNSSH_DIR_NAME
	return rnsshDir
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/hashicorp/terraform/plugin"
	"github.com/jayhding/terraform-provider-pingdom/pingdom"
)

func main() {
	plugin.Serve(&plugin.ServeOpts{
		ProviderFunc: pingdom.Provider,
	})
}
<commit_msg>Change upstream repo back for pull request<commit_after>package main

import (
	"github.com/hashicorp/terraform/plugin"
	"github.com/russellcardullo/terraform-provider-pingdom/pingdom"
)

func main() {
	plugin.Serve(&plugin.ServeOpts{
		ProviderFunc: pingdom.Provider,
	})
}
<|endoftext|>