<|start_filename|>zabbix/protocol_test.go<|end_filename|>
package zabbix_test

import (
	"bytes"
	"testing"

	"github.com/fujiwara/go-zabbix-get/zabbix"
)

func TestPacket(t *testing.T) {
	source := []byte("test.data")
	packet := zabbix.Data2Packet(source)
	compare := []byte{
		90, 66, 88, 68, // "ZBXD"
		1,                      // \x01
		9, 0, 0, 0, 0, 0, 0, 0, // binary int64(9) == len("test.data")
		116, 101, 115, 116, // "test"
		46,               // "."
		100, 97, 116, 97, // "data"
	}
	if !bytes.Equal(packet, compare) {
		t.Errorf("invalid packet %v <=> %v", packet, compare)
	}
	restored, err := zabbix.Packet2Data(packet)
	if err != nil {
		t.Errorf("must be no error: %v", err)
	}
	if !bytes.Equal(restored, source) {
		t.Errorf("invalid data %v <=> %v", restored, source)
	}
}

func TestCorruptedPacket(t *testing.T) {
	packet := []byte{
		90, 66, 88, 68, // "ZBXD"
		0,                      // CORRUPTED!!
		9, 0, 0, 0, 0, 0, 0, 0, // binary int64(9) == len("test.data")
		116, 101, 115, 116, // "test"
		46,               // "."
		100, 97, 116, 97, // "data"
	}
	restored, err := zabbix.Packet2Data(packet)
	if err == nil {
		t.Errorf("must be error")
	}
	t.Logf("restored: %v", restored)
}

func TestStreamBinary(t *testing.T) {
	source := []byte("test.data")
	packet := zabbix.Data2Packet(source)
	buf := bytes.NewReader(packet)
	restored, err := zabbix.Stream2Data(buf)
	if err != nil {
		t.Errorf("must be no error: %v", err)
	}
	if !bytes.Equal(restored, source) {
		t.Errorf("invalid data %v <=> %v", restored, source)
	}
}

func TestStreamText(t *testing.T) {
	source := []byte("test.data\n")
	reader := bytes.NewReader(source)
	restored, err := zabbix.Stream2Data(reader)
	if err != nil {
		t.Errorf("must be no error: %v", err)
	}
	if !bytes.Equal(restored, source) {
		t.Errorf("invalid data %v <=> %v", restored, source)
	}
}


<|start_filename|>zabbix/client_test.go<|end_filename|>
package zabbix_test

import (
	"testing"

	"github.com/fujiwara/go-zabbix-get/zabbix"
)

func TestIncludesPortNumber(t *testing.T) {
	var addrWithPort = []string{
		"example.com:10050",
		"192.0.2.1:10050",
		"[2001:fd00:c2b6:b24b:be67:2827:688d:e6a1:6a3b:beef]:10050",
	}
	for _, addr := range addrWithPort {
		if zabbix.FillDefaultPort(addr, zabbix.AgentDefaultPort) != addr {
			t.Errorf("%s includes port number", addr)
		}
	}
}

func TestNotIncludesPortNumber(t *testing.T) {
	var addrWithoutPort = []string{
		"example.com",
		"192.0.2.1",
		"[2001:fd00:c2b6:b24b:be67:2827:688d:e6a1:6a3b:beef]",
	}
	for _, addr := range addrWithoutPort {
		if zabbix.FillDefaultPort(addr, zabbix.AgentDefaultPort) != addr+":10050" {
			t.Errorf("%s not includes port number", addr)
		}
	}
}


<|start_filename|>zabbix/client.go<|end_filename|>
package zabbix

import (
	"fmt"
	"net"
	"time"
)

const (
	AgentDefaultPort  = 10050
	ServerDefaultPort = 10051
)

func FillDefaultPort(addr string, port int) string {
	_, _, err := net.SplitHostPort(addr)
	if err != nil {
		return fmt.Sprintf("%s:%d", addr, port)
	}
	return addr
}

func Get(addr string, key string, timeout time.Duration) (value string, err error) {
	addr = FillDefaultPort(addr, AgentDefaultPort)
	conn, err := net.DialTimeout("tcp", addr, timeout)
	if err != nil {
		return
	}
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(timeout))

	msg := Data2Packet([]byte(key))
	_, err = conn.Write(msg)
	if err != nil {
		return
	}
	_value, err := Stream2Data(conn)
	return string(_value), err
}


<|start_filename|>main.go<|end_filename|>
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/fujiwara/go-zabbix-get/zabbix"
)

var (
	Version string
)

func main() {
	var (
		port        int
		server      string
		key         string
		timeout     int
		showVersion bool
		format      string
		outputKey   string
	)
	flag.IntVar(&port, "p", 10050, "port")
	flag.StringVar(&server, "s", "127.0.0.1", "hostname or IP")
	flag.StringVar(&key, "k", "", "key")
	flag.IntVar(&timeout, "t", 30, "timeout")
	flag.BoolVar(&showVersion, "V", false, "show Version")
	flag.StringVar(&format, "f", "zabbix", "output format (zabbix or sensu)")
	flag.StringVar(&outputKey, "o", "", "output key string (format=sensu only)")
	flag.Parse()

	if showVersion {
		fmt.Printf("go-zabbix-get version %s\n", Version)
		os.Exit(255)
	}

	if key == "" {
		flag.PrintDefaults()
		os.Exit(255)
	}

	address := fmt.Sprintf("%s:%d", server, port)
	value, err := zabbix.Get(address, key, time.Duration(timeout)*time.Second)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	switch format {
	case "sensu":
		if outputKey == "" {
			printFormatSensu(key, value)
		} else {
			printFormatSensu(outputKey, value)
		}
	default:
		printFormatZabbix(value)
	}
}

func printFormatZabbix(value string) {
	fmt.Printf("%s\n", value)
}

func printFormatSensu(key string, value string) {
	fmt.Printf("%s\t%s\t%d\n", key, value, time.Now().Unix())
}


<|start_filename|>zabbix/trapper.go<|end_filename|>
package zabbix

import (
	"encoding/json"
	"fmt"
	"log"
	"net"
	"time"
)

const (
	TrapperResponseSuccess = "success"
	TrapperResponseInfoFmt = "processed: %d; failed: %d; total: %d; seconds spent: %f"
	TrapperRequestString   = "sender data"
)

type TrapperData struct {
	Host  string `json:"host"`
	Key   string `json:"key"`
	Value string `json:"value"`
	Clock int64  `json:"clock"`
}

type TrapperRequest struct {
	Request string        `json:"request"`
	Clock   int64         `json:"clock"`
	Data    []TrapperData `json:"data"`
}

func (r TrapperRequest) ToPacket() []byte {
	data, _ := json.Marshal(r)
	return Data2Packet(data)
}

type TrapperResponse struct {
	Response  string `json:"response"`
	Info      string `json:"info"`
	Proceeded int    `json:""`
	Failed    int    `json:""`
	Total     int    `json:""`
}

func (r TrapperResponse) ToPacket() []byte {
	data, _ := json.Marshal(r)
	return Data2Packet(data)
}

func SendBulk(addr string, req TrapperRequest, timeout time.Duration) (res TrapperResponse, err error) {
	packet := req.ToPacket()

	addr = FillDefaultPort(addr, ServerDefaultPort)
	conn, err := net.DialTimeout("tcp", addr, timeout)
	if err != nil {
		return
	}
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(timeout))

	_, err = conn.Write(packet)
	if err != nil {
		return
	}

	data, err := Stream2Data(conn)
	if err != nil {
		return
	}
	err = json.Unmarshal(data, &res)
	if err != nil {
		return
	}
	return res, nil
}

func Send(addr string, data TrapperData, timeout time.Duration) (TrapperResponse, error) {
	req := TrapperRequest{
		Request: TrapperRequestString,
		Data:    []TrapperData{data},
	}
	return SendBulk(addr, req, timeout)
}

func RunTrapper(addr string, callback func(TrapperRequest) (TrapperResponse, error)) error {
	addr = FillDefaultPort(addr, ServerDefaultPort)
	log.Println("Starting trapper on", addr)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("Can't lesten tcp", addr, err)
	}
	log.Println("Ready for connection")
	var conn net.Conn
	for {
		conn, err = listener.Accept()
		if err != nil {
			log.Println("Error accept:", err)
			continue
		}
		// log.Println("Accept connection from", conn.RemoteAddr())
		go handleTrapperConn(conn, callback)
	}
	return nil
}

func handleTrapperConn(conn net.Conn, callback func(TrapperRequest) (TrapperResponse, error)) {
	defer conn.Close()
	start := time.Now()
	var request TrapperRequest
	input, err := Stream2Data(conn)
	if err != nil {
		log.Println("request error:", err)
		return
	}
	err = json.Unmarshal(input, &request)
	if err != nil {
		log.Println("decode request error:", err)
		return
	}
	if request.Request != TrapperRequestString {
		log.Println("invalid request.Request", request.Request)
		return
	}

	res, err := callback(request)
	if err != nil {
		log.Println("process callback error", err)
		return
	}
	res.Total = res.Proceeded + res.Failed
	if res.Response == "" {
		res.Response = TrapperResponseSuccess
	}
	if res.Info == "" {
		res.Info = fmt.Sprintf(
			TrapperResponseInfoFmt,
			res.Proceeded,
			res.Failed,
			res.Total,
			time.Now().Sub(start).Seconds(),
		)
	}
	conn.Write(res.ToPacket())
}


<|start_filename|>zabbix/agent_test.go<|end_filename|>
package zabbix_test

import (
	"fmt"
	"log"
	"net"
	"testing"
	"time"

	"github.com/fujiwara/go-zabbix-get/zabbix"
)

var (
	done    chan (bool)
	timeout = 3 * time.Second
)

func TestAgent(t *testing.T) {
	done = make(chan bool)
	go zabbix.RunAgent("localhost", func(key string) (string, error) {
		switch key {
		case "agent.ping":
			log.Println("key", key)
			return "1", nil
		case "agent.uptime":
			log.Println("key", key)
			return "123", nil
		case "timeout":
			log.Println("key", key, "sleeping...")
			time.Sleep(timeout + time.Second)
			log.Println("wake up. response ok!")
			return "ok", nil
		case "shutdown":
			done <- true
			return "", nil
		default:
			return "", fmt.Errorf("not supported")
		}
	})
	time.Sleep(1 * time.Second)
}

func TestAgentCannotConnect(t *testing.T) {
	value, err := zabbix.Get("localhost:10049", "agent.ping", timeout)
	if err == nil {
		t.Errorf("agent is not runnig, but not error value:", value)
	}
}

func TestAgentGetPing(t *testing.T) {
	value, err := zabbix.Get("localhost", "agent.ping", timeout)
	if err != nil {
		t.Error("get agent.ping failed", err)
	}
	if value != "1" {
		t.Error("agent.ping value expected: 1, got:", value)
	}
}

func TestAgentGetUptime(t *testing.T) {
	value, err := zabbix.Get("localhost", "agent.uptime", timeout)
	if err != nil {
		t.Error("get agent.ping failed", err)
	}
	if value != "123" {
		t.Error("agent.uptime value expected: 123, got:", value)
	}
}

func TestAgentGetError(t *testing.T) {
	value, err := zabbix.Get("localhost", "xxx", timeout)
	if err != nil {
		t.Error("xxx failed", err)
	}
	if value != zabbix.ErrorMessage {
		t.Error("xxx value expected: ", zabbix.ErrorMessage, "got:", value)
	}
}

func TestAgentGetTimeout(t *testing.T) {
	_, err := zabbix.Get("localhost", "timeout", timeout)
	if err == nil {
		t.Error("timeout must be timeouted.", err)
	}
	if _err := err.(*net.OpError); !_err.Timeout() {
		t.Error("err expected i/o timeout. got:", err)
	}
	log.Println("client timeout")
}

func TestAgentShutdown(t *testing.T) {
	zabbix.Get("localhost", "shutdown", timeout)
	<-done
}


<|start_filename|>zabbix/protocol.go<|end_filename|>
package zabbix

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
)

const (
	HeaderString     = "ZBXD"
	HeaderLength     = len(HeaderString)
	HeaderVersion    = uint8(1)
	DataLengthOffset = int64(HeaderLength + 1)
	DataLengthSize   = int64(8)
	DataOffset       = int64(DataLengthOffset + DataLengthSize)
	ErrorMessage     = "ZBX_NOTSUPPORTED"
)

var (
	ErrorMessageBytes = []byte(ErrorMessage)
	Terminator        = []byte("\n")
	HeaderBytes       = []byte(HeaderString)
)

func Data2Packet(data []byte) []byte {
	buf := new(bytes.Buffer)
	Data2Stream(data, buf)
	return buf.Bytes()
}

func Data2Stream(data []byte, conn io.Writer) (int, error) {
	conn.Write(HeaderBytes)
	binary.Write(conn, binary.LittleEndian, HeaderVersion)
	binary.Write(conn, binary.LittleEndian, int64(len(data)))
	return conn.Write(data)
}

func Packet2Data(packet []byte) (data []byte, err error) {
	var dataLength int64
	if len(packet) < int(DataOffset) {
		err = errors.New("zabbix protocol packet too short")
		return
	}

	// read header
	headBuf := bytes.NewReader(packet[0:DataLengthOffset])
	head := make([]byte, DataLengthOffset)
	_, err = headBuf.Read(head)
	if !bytes.Equal(head[0:HeaderLength], HeaderBytes) || head[HeaderLength] != byte(HeaderVersion) {
		err = errors.New("invalid packet header")
		return
	}

	// read data
	buf := bytes.NewReader(packet[DataLengthOffset:DataOffset])
	err = binary.Read(buf, binary.LittleEndian, &dataLength)
	if err != nil {
		return
	}
	data = packet[DataOffset : DataOffset+dataLength]
	return
}

func Stream2Data(conn io.Reader) (rdata []byte, err error) {
	// read header "ZBXD\x01"
	head := make([]byte, DataLengthOffset)
	_, err = conn.Read(head)
	if err != nil {
		return
	}
	if bytes.Equal(head[0:HeaderLength], HeaderBytes) && head[HeaderLength] == byte(HeaderVersion) {
		rdata, err = parseBinary(conn)
	} else {
		rdata, err = parseText(conn, head)
	}
	return
}

func parseBinary(conn io.Reader) (rdata []byte, err error) {
	// read data length
	var dataLength int64
	err = binary.Read(conn, binary.LittleEndian, &dataLength)
	if err != nil {
		return
	}
	// read data body
	buf := make([]byte, 1024)
	data := new(bytes.Buffer)
	total := 0
	size := 0
	for total < int(dataLength) {
		size, err = conn.Read(buf)
		if err != nil {
			return
		}
		if size == 0 {
			break
		}
		total = total + size
		data.Write(buf[0:size])
	}
	rdata = data.Bytes()
	return
}

func parseText(conn io.Reader, head []byte) (rdata []byte, err error) {
	data := new(bytes.Buffer)
	data.Write(head)
	buf := make([]byte, 1024)
	size := 0
	for {
		// read data while "\n" found
		size, err = conn.Read(buf)
		if err != nil {
			return
		}
		if size == 0 {
			break
		}
		i := bytes.Index(buf[0:size], Terminator)
		if i == -1 {
			// terminator not found
			data.Write(buf[0:size])
			continue
		}
		// terminator found
		data.Write(buf[0 : i+1])
		break
	}
	rdata = data.Bytes()
	return
}


<|start_filename|>zabbix/trapper_test.go<|end_filename|>
package zabbix_test

import (
	"log"
	"net"
	"testing"
	"time"

	"github.com/fujiwara/go-zabbix-get/zabbix"
)

func TestTrapper(t *testing.T) {
	done = make(chan bool)
	go zabbix.RunTrapper("localhost", func(req zabbix.TrapperRequest) (res zabbix.TrapperResponse, err error) {
		switch req.Request {
		case "timeout":
			log.Println("request timeout sleeping...")
			time.Sleep(timeout + time.Second)
			log.Println("wake up")
		case "shutdown":
			done <- true
		}
		for _, data := range req.Data {
			log.Println(data)
		}
		res.Proceeded = len(req.Data)
		return res, nil
	})
	time.Sleep(1 * time.Second)
}

func TestTrapperCannotConnect(t *testing.T) {
	value, err := zabbix.Send(
		"localhost:10049",
		zabbix.TrapperData{Host: "localhost", Key: "foo", Value: "bar"},
		timeout,
	)
	if err == nil {
		t.Errorf("trapper is not runnig, but not error value:", value)
	}
}

func TestTrapperSend(t *testing.T) {
	res, err := zabbix.Send(
		"localhost",
		zabbix.TrapperData{Host: "localhost", Key: "foo", Value: "bar"},
		timeout,
	)
	if err != nil {
		t.Errorf("send failed", err)
	}
	if res.Proceeded != 1 {
		t.Errorf("proceeded expected 1 got", res.Proceeded)
	}
	if res.Failed != 0 {
		t.Errorf("failed expected 0 got", res.Failed)
	}
	if res.Total != 1 {
		t.Errorf("total expected 1 got", res.Total)
	}
}

func TestTrapperSendBulk(t *testing.T) {
	res, err := zabbix.SendBulk(
		"localhost",
		zabbix.TrapperRequest{
			Data: []zabbix.TrapperData{
				zabbix.TrapperData{Host: "localhost", Key: "foo", Value: "bar"},
				zabbix.TrapperData{Host: "localhost", Key: "xxx", Value: "yyy"},
			},
		},
		timeout,
	)
	if err != nil {
		t.Errorf("send failed", err)
	}
	if res.Proceeded != 2 {
		t.Errorf("proceeded expected 2 got", res.Proceeded)
	}
	if res.Failed != 0 {
		t.Errorf("failed expected 0 got", res.Failed)
	}
	if res.Total != 2 {
		t.Errorf("total expected 2 got", res.Total)
	}
}

func TestTrapperSendTimeout(t *testing.T) {
	res, err := zabbix.SendBulk(
		"localhost",
		zabbix.TrapperRequest{
			Request: "timeout",
			Data: []zabbix.TrapperData{
				zabbix.TrapperData{Host: "localhost", Key: "foo", Value: "bar"},
			},
		},
		timeout,
	)
	if err == nil {
		t.Error("timeout must be timeouted.", err)
	}
	if _err := err.(*net.OpError); !_err.Timeout() {
		t.Error("err expected i/o timeout. got:", err)
	}
	log.Println("client timeout", res)
}

func TestTrapperShutdown(t *testing.T) {
	zabbix.SendBulk(
		"localhost",
		zabbix.TrapperRequest{Request: "shutdown"},
		timeout,
	)
	<-done
}


<|start_filename|>Makefile<|end_filename|>
GIT_VER := $(shell git describe --tags)

all:
	go get github.com/fujiwara/go-zabbix-get

binary:
	gox -os="linux darwin" -arch="amd64" -output "pkg/{{.Dir}}-${GIT_VER}-{{.OS}}-{{.Arch}}" -ldflags "-X main.Version ${GIT_VER}"
	gox -os="linux" -arch="arm" -output "pkg/{{.Dir}}-${GIT_VER}-{{.OS}}-{{.Arch}}" -ldflags "-X main.Version ${GIT_VER}"
	cd pkg && find . -name "*${GIT_VER}*" -type f -exec zip {}.zip {} \;

test:
	cd zabbix && go test

clean:
	rm -f pkg/*


<|start_filename|>zabbix/agent.go<|end_filename|>
package zabbix

import (
	"log"
	"net"
	"strings"
)

func RunAgent(addr string, callback func(string) (string, error)) error {
	addr = FillDefaultPort(addr, AgentDefaultPort)
	log.Println("Starting agent on", addr)

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("Can't lesten", addr, err)
	}
	log.Println("Ready for connection")
	var conn net.Conn
	for {
		conn, err = listener.Accept()
		if err != nil {
			log.Println("Error accept:", err)
			continue
		}
		// log.Println("Accept connection from", conn.RemoteAddr())
		go handleAgentConn(conn, callback)
	}
	return nil
}

func handleAgentConn(conn net.Conn, callback func(string) (string, error)) {
	var value string
	defer conn.Close()

	key, err := Stream2Data(conn)
	if err != nil {
		log.Println("request error:", err)
		Data2Stream(ErrorMessageBytes, conn)
		return
	}
	keyStr := strings.TrimRight(string(key), "\n")
	value, err = callback(keyStr)
	if err != nil {
		log.Println("process callback error:", err)
		value = ErrorMessage
		Data2Stream(ErrorMessageBytes, conn)
		return
	}
	Data2Stream([]byte(value), conn)
}
