<commit_before>
// Simple chat server which uses connection properties from "connection.json" in same directory
// Simple chat client is intended to be used but a standard telnet connection can be used
// > telnet {host} {port}
// > /user {username}
// > /message {message}
//
// reference: https://parroty00.wordpress.com/2013/07/18/golang-tcp-server-example/
package main

import (
  "net"
  "fmt"
  "bufio"
  "strings"
  "regexp"
)

// Container for client username and connection details
type Client struct {
  Connection net.Conn
  Username string
}

// Close the client connection and clenup
func (client Client) Close() {
  client.Connection.Close();
  // FIXME need to clean the connection up from availableClients
}

// Register the connection and cache it
func (client *Client) Register() {
  numClients := len(clients)
  availableClients[numClients] = client;
  clients = availableClients[0:numClients+1]
}

// static client list
var availableClients [256]*Client
var clients []*Client

// program main
func main() {
  // start the server
  psock, err := net.Listen("tcp", ":5000")
  if err != nil {
    fmt.Printf("Can't create server %v\n", err)
    return
  }

  for {
    // accept connections
    conn, err := psock.Accept()
    if err != nil {
      fmt.Printf("Can't accept connections %v\n", err)
      return
    }

    // keep track of the client details
    client := Client{Connection: conn}
    client.Register();

    // allow non-blocking client request handling
    channel := make(chan string)
    go waitForInput(channel, &client)
    go handleInput(channel, &client)
  }
}

// wait for client input (buffered by newlines) and signal the channel
func waitForInput(out chan string, client *Client) {
  defer close(out)

  for {
    line, err := bufio.NewReader(client.Connection).ReadBytes('\n')
    if err != nil {
      // connection has been closed, remove the client
      client.Close();
      return
    }
    out <- string(line)
  }
}

// listen for channel updates for a client and handle the message
// messages must be in the format of /{action} {content} where content is optional depending on the action
// supported actions are "user", "chat", and "quit".  the "user" must be set before any chat messages are allowed
func handleInput(in <-chan string, client *Client) {
  for {
    message := <- in
    message = strings.TrimSpace(message)
    action, body := getAction(message)

    if (action != "") {
      switch action {
        case "message":
          sendMessage("message", body, client, false)
        case "user":
          client.Username = body
          sendMessage("enter", "", client, false)
        case "leave":
          sendMessage("leave", "", client, false)
          client.Close();
        default:
          sendMessage("unrecognized", action, client, true)
      }
    }
  }
}

// sent a message to all clients (except the sender)
func sendMessage(messageType string, message string, client *Client, thisClientOnly bool) {
  message = fmt.Sprintf("/%v [%v] %v\n", messageType, client.Username, message);

  for _, _client := range clients {
    // write the message to the client
    if ((thisClientOnly && _client.Username == client.Username) ||
        (!thisClientOnly && _client != client && _client.Username != "")) {
      // you won't hear any activity if you are anonymous unless thisClientOnly
      // when current client will *only* be messaged
      fmt.Fprintf(_client.Connection, message)
    }
  }
}

// parse out message contents (/{action} {message}) and return individual values
func getAction(message string) (string, string) {
  actionRegex, _ := regexp.Compile(`^\/([^\s]*)\s*(.*)$`)
  res := actionRegex.FindAllStringSubmatch(message, -1)
  if (len(res) == 1) {
    return res[0][1], res[0][2]
  }
  return "", ""
}
<commit_msg>clean up client cache<commit_after>
// Simple chat server which uses connection properties from "connection.json" in same directory
// Simple chat client is intended to be used but a standard telnet connection can be used
// > telnet {host} {port}
// > /user {username}
// > /message {message}
//
// reference: https://parroty00.wordpress.com/2013/07/18/golang-tcp-server-example/
package main

import (
  "net"
  "fmt"
  "bufio"
  "strings"
  "regexp"
)

// Container for client username and connection details
type Client struct {
  Connection net.Conn
  Username string
}

// Close the client connection and clenup
func (client *Client) Close(doSendMessage bool) {
  if (doSendMessage) {
    // if we send the close command, the connection will terminate causing another close
    // which will send the message
    sendMessage("leave", "", client, false)
  }
  client.Connection.Close();
  clients = removeEntry(client, clients);
}

// Register the connection and cache it
func (client *Client) Register() {
  clients = append(clients, client);
}

// static client list
var clients []*Client

// program main
func main() {
  // start the server
  psock, err := net.Listen("tcp", ":5000")
  if err != nil {
    fmt.Printf("Can't create server %v\n", err)
    return
  }

  for {
    // accept connections
    conn, err := psock.Accept()
    if err != nil {
      fmt.Printf("Can't accept connections %v\n", err)
      return
    }

    // keep track of the client details
    client := Client{Connection: conn}
    client.Register();

    // allow non-blocking client request handling
    channel := make(chan string)
    go waitForInput(channel, &client)
    go handleInput(channel, &client)
  }
}

// wait for client input (buffered by newlines) and signal the channel
func waitForInput(out chan string, client *Client) {
  defer close(out)

  for {
    line, err := bufio.NewReader(client.Connection).ReadBytes('\n')
    if err != nil {
      // connection has been closed, remove the client
      client.Close(true);
      return
    }
    out <- string(line)
  }
}

// listen for channel updates for a client and handle the message
// messages must be in the format of /{action} {content} where content is optional depending on the action
// supported actions are "user", "chat", and "quit".  the "user" must be set before any chat messages are allowed
func handleInput(in <-chan string, client *Client) {
  for {
    message := <- in
    message = strings.TrimSpace(message)
    action, body := getAction(message)

    if (action != "") {
      switch action {
        case "message":
          sendMessage("message", body, client, false)
        case "user":
          client.Username = body
          sendMessage("enter", "", client, false)
        case "leave":
          client.Close(false);
        default:
          sendMessage("unrecognized", action, client, true)
      }
    }
  }
}

// sent a message to all clients (except the sender)
func sendMessage(messageType string, message string, client *Client, thisClientOnly bool) {
  message = fmt.Sprintf("/%v [%v] %v\n", messageType, client.Username, message);

  for _, _client := range clients {
    // write the message to the client
    if ((thisClientOnly && _client.Username == client.Username) ||
        (!thisClientOnly && _client != client && _client.Username != "")) {
      // you won't hear any activity if you are anonymous unless thisClientOnly
      // when current client will *only* be messaged
      fmt.Fprintf(_client.Connection, message)
    }
  }
}

// parse out message contents (/{action} {message}) and return individual values
func getAction(message string) (string, string) {
  actionRegex, _ := regexp.Compile(`^\/([^\s]*)\s*(.*)$`)
  res := actionRegex.FindAllStringSubmatch(message, -1)
  if (len(res) == 1) {
    return res[0][1], res[0][2]
  }
  return "", ""
}

// remove client entry from stored clients
func removeEntry(client *Client, arr []*Client) []*Client {
  rtn := arr
  index := -1
  for i, value := range arr {
    if (value == client) {
      index = i;
      break;
    }
  }

  if (index >= 0) {
    // we have a match, create a new array without the match
    rtn = make([]*Client, len(arr)-1)
    copy(rtn, arr[:index])
    copy(rtn[index:], arr[index+1:])
    fmt.Printf("found entry and new arr is %v", rtn);
  }

  return rtn;
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/Shopify/sarama"
	"github.com/cenkalti/backoff"
	rfc3164 "github.com/jeromer/syslogparser/rfc3164"
	"github.com/pingles/go-metrics-riemann"
	"github.com/rcrowley/go-metrics"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	cfg         *config
	logger      = log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Lshortfile)
	port        = flag.Int("port", 514, "port on which to listen")
	verbose     = flag.Bool("verbose", false, "print all messages to stdout")
	publish     = flag.Bool("publish", false, "publish messages to kafka")
	topic       = flag.String("topic", "syslog", "kafka topic to publish on")
	zkstring    = flag.String("zkstring", "localhost:2181", "ZooKeeper broker connection string")
	riemannHost = flag.String("riemann", "localhost:5555", "Riemann TCP host:port")

	connectionsCounter = metrics.NewCounter()
	sendMeter          = metrics.NewMeter()
	sendErrorsMeter    = metrics.NewMeter()
	sendTimer          = metrics.NewTimer()
)

type config struct {
	port        int
	verbose     bool
	publish     bool
	topic       string
	zkstring    string
	riemannHost string
}

func (c *config) String() string {
	return fmt.Sprintf("port: %d, verbose: %t, publish: %t, topic: %s, zkstring: %s, riemannHost: %s",
		c.port, c.verbose, c.publish, c.topic, c.zkstring, c.riemannHost)
}

type openConnection struct {
	connection net.Conn
	done       chan bool
}

type server struct {
	listener    net.Listener
	connections []*openConnection
	client      *sarama.Client
	producer    *sarama.Producer
	shutdown    chan bool
}

func (s *server) process(line []byte) {
	p := rfc3164.NewParser(line)
	if err := p.Parse(); err != nil {
		logger.Println("failed to parse:", err)
		return
	}

	parts := p.Dump()

	jsonBytes, err := json.Marshal(parts)
	if err != nil {
		logger.Println("error marshaling message, skipping.", err)
		return
	}

	if cfg.verbose {
		logger.Println("enqueuing", string(jsonBytes))
	}

	if cfg.publish {

		put := func() error {
			var err error

			send := func() {
				err = s.producer.SendMessage(cfg.topic, nil, sarama.ByteEncoder(jsonBytes))
			}
			sendTimer.Time(send)

			if err != nil {
				logger.Println("error sending message, will retry.", err)
				sendErrorsMeter.Mark(1)
			} else {
				sendMeter.Mark(1)
			}
			return err
		}

		policy := backoff.NewExponentialBackOff()
		policy.MaxElapsedTime = time.Minute * 5
		err := backoff.Retry(put, policy)
		if err != nil {
			// retrying a bunch of times failed...
			logger.Println("failed sending message.", err)
		}

	}
}

func (s *server) handleConnection(conn *openConnection) {
	defer conn.connection.Close()

	logger.Println("got connection from:", conn.connection.RemoteAddr())

	incomingMessageMeter := metrics.NewMeter()
	metrics.Register(fmt.Sprintf("%s incomingMessages", conn.connection.RemoteAddr()), incomingMessageMeter)

	scanner := bufio.NewScanner(conn.connection)
	for scanner.Scan() {
		b := []byte(scanner.Text())
		incomingMessageMeter.Mark(1)
		s.process(b)
	}

	if err := scanner.Err(); err != nil {
		logger.Println("error reading from connection:", err)
	}

	logger.Println("exiting connection handler")
	conn.done <- true
	connectionsCounter.Dec(1)
}

func (s *server) stop() {
	s.shutdown <- true
	s.listener.Close()

	for _, conn := range s.connections {
		logger.Println("closing connection to", conn.connection.RemoteAddr())
		conn.connection.Close()
	}

	logger.Println("waiting for", len(s.connections), "connections to close")
	for _, conn := range s.connections {
		<-conn.done
	}
	logger.Println("all connections closed")

	if cfg.publish {
		s.producer.Close()
		s.client.Close()
		logger.Println("finished stopping producer")
	}
}

func (s *server) start() error {
	logger.Printf("starting to listen on %d\n", cfg.port)

	if cfg.verbose {
		logger.Println("verbose is enabled, all messages will be printed to stderr")
	}

	// listen for inbound syslog messages over tcp
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", cfg.port))
	if err != nil {
		return err
	}
	s.listener = listener

	if cfg.publish {
		// client for kafka
		brokers, err := LookupBrokers(cfg.zkstring)
		if err != nil {
			return err
		}

		brokerStr := make([]string, len(brokers))
		for i, b := range brokers {
			brokerStr[i] = fmt.Sprintf("%s:%d", b.Host, b.Port)
		}

		logger.Println("connecting to Kafka, using brokers from ZooKeeper:", brokerStr)
		client, err := sarama.NewClient("syslog", brokerStr, sarama.NewClientConfig())
		if err != nil {
			return err
		}
		s.client = client

		producer, err := sarama.NewProducer(client, sarama.NewProducerConfig())
		if err != nil {
			return err
		}
		s.producer = producer
	}

	connections := make(chan *openConnection)
	errors := make(chan error)

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				errors <- err
			} else {
				connection := &openConnection{
					connection: conn,
					done:       make(chan bool, 1),
				}
				connectionsCounter.Inc(1)
				connections <- connection
			}
		}
	}()

	logger.Println("waiting for syslog connection")
	for {
		select {
		case conn := <-connections:
			s.connections = append(s.connections, conn)
			go s.handleConnection(conn)
		case err := <-errors:
			logger.Println("failed to accept connection:", err)
		case <-s.shutdown:
			goto exit
		}
	}

exit:
	logger.Println("exiting listen loop")
	return nil
}

func handleInterrupt(s *server) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGUSR1)

	go func() {
		for sig := range signals {
			switch sig {
			case syscall.SIGINT:
				logger.Println("got kill signal, closing", len(s.connections), "connections")
				s.stop()
				os.Exit(0)
				break
			case syscall.SIGUSR1:
				logger.Println("refreshing kafka metadata")
				err := s.client.RefreshAllMetadata()
				if err != nil {
					logger.Println("error refreshing metadata.", err)
				}
				break
			}
		}
	}()
}

func registerMetrics() {
	metrics.Register("openConnections count", connectionsCounter)
	metrics.Register("sentMessages", sendMeter)
	metrics.Register("sendMessageErrors count", sendErrorsMeter)
	metrics.Register("sendMessage", sendTimer)
}

func main() {
	flag.Parse()

	cfg = &config{
		port:        *port,
		verbose:     *verbose,
		publish:     *publish,
		topic:       *topic,
		zkstring:    *zkstring,
		riemannHost: *riemannHost,
	}

	logger.Println("starting with config:", cfg)

	registerMetrics()

	go riemann.Report(metrics.DefaultRegistry, time.Second, cfg.riemannHost)

	server := &server{
		connections: []*openConnection{},
		shutdown:    make(chan bool),
	}

	handleInterrupt(server)

	if err := server.start(); err != nil {
		logger.Println("failed to listen:", err)
		os.Exit(1)
	}
}
<commit_msg>extract kafka producer function<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/Shopify/sarama"
	"github.com/cenkalti/backoff"
	rfc3164 "github.com/jeromer/syslogparser/rfc3164"
	"github.com/pingles/go-metrics-riemann"
	"github.com/rcrowley/go-metrics"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	cfg         *config
	logger      = log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Lshortfile)
	port        = flag.Int("port", 1514, "port on which to listen")
	verbose     = flag.Bool("verbose", false, "print all messages to stdout")
	publish     = flag.Bool("publish", false, "publish messages to kafka")
	topic       = flag.String("topic", "syslog", "kafka topic to publish on")
	zkstring    = flag.String("zkstring", "localhost:2181", "ZooKeeper broker connection string")
	riemannHost = flag.String("riemann", "localhost:5555", "Riemann TCP host:port")

	connectionsCounter = metrics.NewCounter()
	sendMeter          = metrics.NewMeter()
	sendErrorsMeter    = metrics.NewMeter()
	sendTimer          = metrics.NewTimer()
)

type config struct {
	port        int
	verbose     bool
	publish     bool
	topic       string
	zkstring    string
	riemannHost string
}

func (c *config) String() string {
	return fmt.Sprintf("port: %d, verbose: %t, publish: %t, topic: %s, zkstring: %s, riemannHost: %s",
		c.port, c.verbose, c.publish, c.topic, c.zkstring, c.riemannHost)
}

type openConnection struct {
	connection net.Conn
	done       chan bool
}

type server struct {
	listener    net.Listener
	connections []*openConnection
	client      *sarama.Client
	producer    *sarama.Producer
	shutdown    chan bool
}

func (s *server) process(line []byte) {
	p := rfc3164.NewParser(line)
	if err := p.Parse(); err != nil {
		logger.Println("failed to parse:", err)
		return
	}

	parts := p.Dump()

	jsonBytes, err := json.Marshal(parts)
	if err != nil {
		logger.Println("error marshaling message, skipping.", err)
		return
	}

	if cfg.verbose {
		logger.Println("enqueuing", string(jsonBytes))
	}

	if cfg.publish {
		put := func() error {
			var err error

			send := func() {
				err = s.producer.SendMessage(cfg.topic, nil, sarama.ByteEncoder(jsonBytes))
			}
			sendTimer.Time(send)

			if err != nil {
				logger.Println("error sending message, will retry.", err)
				sendErrorsMeter.Mark(1)
			} else {
				sendMeter.Mark(1)
			}
			return err
		}

		policy := backoff.NewExponentialBackOff()
		policy.MaxElapsedTime = time.Minute * 5
		err := backoff.Retry(put, policy)
		if err != nil {
			// retrying a bunch of times failed...
			logger.Println("failed sending message.", err)
		}
	}
}

func (s *server) handleConnection(conn *openConnection) {
	defer conn.connection.Close()

	logger.Println("got connection from:", conn.connection.RemoteAddr())

	incomingMessageMeter := metrics.NewMeter()
	metrics.Register(fmt.Sprintf("%s incomingMessages", conn.connection.RemoteAddr()), incomingMessageMeter)

	scanner := bufio.NewScanner(conn.connection)
	for scanner.Scan() {
		b := []byte(scanner.Text())
		incomingMessageMeter.Mark(1)
		s.process(b)
	}

	if err := scanner.Err(); err != nil {
		logger.Println("error reading from connection:", err)
	}

	logger.Println("exiting connection handler")
	conn.done <- true
	connectionsCounter.Dec(1)
}

func (s *server) stop() {
	s.shutdown <- true
	s.listener.Close()

	for _, conn := range s.connections {
		logger.Println("closing connection to", conn.connection.RemoteAddr())
		conn.connection.Close()
	}

	logger.Println("waiting for", len(s.connections), "connections to close")
	for _, conn := range s.connections {
		<-conn.done
	}
	logger.Println("all connections closed")

	if cfg.publish {
		s.producer.Close()
		s.client.Close()
		logger.Println("finished stopping producer")
	}
}

func newProducerFromZookeeper() (*sarama.Client, *sarama.Producer, error) {
	brokers, err := LookupBrokers(cfg.zkstring)
	if err != nil {
		return nil, nil, err
	}

	brokerStr := make([]string, len(brokers))
	for i, b := range brokers {
		brokerStr[i] = fmt.Sprintf("%s:%d", b.Host, b.Port)
	}

	logger.Println("connecting to Kafka, using brokers from ZooKeeper:", brokerStr)
	client, err := sarama.NewClient("syslog", brokerStr, sarama.NewClientConfig())
	if err != nil {
		return nil, nil, err
	}

	producer, err := sarama.NewProducer(client, sarama.NewProducerConfig())
	if err != nil {
		return nil, nil, err
	}

	return client, producer, nil
}

func (s *server) start() error {
	logger.Printf("starting to listen on %d\n", cfg.port)

	if cfg.verbose {
		logger.Println("verbose is enabled, all messages will be printed to stderr")
	}

	// listen for inbound syslog messages over tcp
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", cfg.port))
	if err != nil {
		return err
	}
	s.listener = listener

	if cfg.publish {
		client, producer, err := newProducerFromZookeeper()
		if err != nil {
			return err
		}
		s.client = client
		s.producer = producer
	}

	connections := make(chan *openConnection)
	errors := make(chan error)

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				errors <- err
			} else {
				connection := &openConnection{
					connection: conn,
					done:       make(chan bool, 1),
				}
				connectionsCounter.Inc(1)
				connections <- connection
			}
		}
	}()

	logger.Println("waiting for syslog connection")
	for {
		select {
		case conn := <-connections:
			s.connections = append(s.connections, conn)
			go s.handleConnection(conn)
		case err := <-errors:
			logger.Println("failed to accept connection:", err)
		case <-s.shutdown:
			goto exit
		}
	}

exit:
	logger.Println("exiting listen loop")
	return nil
}

func handleInterrupt(s *server) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGUSR1)

	go func() {
		for sig := range signals {
			switch sig {
			case syscall.SIGINT:
				logger.Println("got kill signal, closing", len(s.connections), "connections")
				s.stop()
				os.Exit(0)
				break
			case syscall.SIGUSR1:
				logger.Println("refreshing kafka metadata")
				err := s.client.RefreshAllMetadata()
				if err != nil {
					logger.Println("error refreshing metadata.", err)
				}
				break
			}
		}
	}()
}

func registerMetrics() {
	metrics.Register("openConnections count", connectionsCounter)
	metrics.Register("sentMessages", sendMeter)
	metrics.Register("sendMessageErrors count", sendErrorsMeter)
	metrics.Register("sendMessage", sendTimer)
}

func main() {
	flag.Parse()

	cfg = &config{
		port:        *port,
		verbose:     *verbose,
		publish:     *publish,
		topic:       *topic,
		zkstring:    *zkstring,
		riemannHost: *riemannHost,
	}

	logger.Println("starting with config:", cfg)

	registerMetrics()

	go riemann.Report(metrics.DefaultRegistry, time.Second, cfg.riemannHost)

	server := &server{
		connections: []*openConnection{},
		shutdown:    make(chan bool),
	}

	handleInterrupt(server)

	if err := server.start(); err != nil {
		logger.Println("failed to listen:", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>/*
Package firetest provides utilities for Firebase testing

*/
package firetest

import (
	"encoding/base32"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"
)

const (
	missingJSONExtension = "append .json to your request URI to use the REST API"
	missingBody          = `{"error":"Error: No data supplied."}`
	invalidJSON          = `{"error":"Invalid data; couldn't parse JSON object, array, or value. Perhaps you're using invalid characters in your key names."}`
)

// Firetest is a Firebase server implementation
type Firetest struct {
	// URL of form http://ipaddr:port with no trailing slash
	URL string

	listener net.Listener
	db       *treeDB
}

// New creates a new Firetest server
func New() *Firetest {
	return &Firetest{
		db: newTree(),
	}
}

func sanitizePath(p string) string {
	s := strings.Trim(p, "/")
	return strings.TrimSuffix(s, ".json")
}

// Start starts the server
func (ft *Firetest) Start() {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		if l, err = net.Listen("tcp6", "[::1]:0"); err != nil {
			panic(fmt.Errorf("failed to listen on a port: %v", err))
		}
	}
	ft.listener = l

	s := http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		ft.serveHTTP(w, req)
	})}
	go func() {
		if err := s.Serve(l); err != nil {
			log.Printf("error serving: %s", err)
		}

		ft.Close()
	}()
	ft.URL = "http://" + ft.listener.Addr().String()
}

// Close closes the server
func (ft *Firetest) Close() {
	if ft.listener != nil {
		ft.listener.Close()
	}
}

func (ft *Firetest) serveHTTP(w http.ResponseWriter, req *http.Request) {
	if !strings.HasSuffix(req.URL.String(), ".json") {
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte(missingJSONExtension))
		return
	}

	switch req.Method {
	case "PUT":
		ft.set(w, req)
	case "PATCH":
		ft.update(w, req)
	case "POST":
		ft.create(w, req)
	case "GET":
		ft.get(w, req)
	case "DELETE":
		ft.del(w, req)
	default:
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.Println("not implemented yet")
	}
}

func (ft *Firetest) set(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil || len(body) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(missingBody))
		return
	}

	var n node
	if err := json.Unmarshal(body, &n); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(invalidJSON))
		return
	}

	ft.db.add(sanitizePath(req.URL.Path), &n)
	w.Write(body)
}

func (ft *Firetest) update(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil || len(body) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(missingBody))
		return
	}

	var n node
	if err := json.Unmarshal(body, &n); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(invalidJSON))
		return
	}

	ft.db.update(sanitizePath(req.URL.Path), &n)
	w.Write(body)
}

func (ft *Firetest) create(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil || len(body) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(missingBody))
		return
	}

	var n node
	if err := json.Unmarshal(body, &n); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(invalidJSON))
		return
	}

	src := []byte(fmt.Sprint(time.Now().UnixNano()))
	name := "~" + base32.StdEncoding.EncodeToString(src)
	path := fmt.Sprintf("%s/%s", sanitizePath(req.URL.Path), name)
	ft.db.add(path, &n)

	rtn := map[string]string{"name": name}
	if err := json.NewEncoder(w).Encode(rtn); err != nil {
		log.Printf("Error encoding json: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
	}
}

func (ft *Firetest) del(w http.ResponseWriter, req *http.Request) {
	ft.db.del(sanitizePath(req.URL.Path))
}

func (ft *Firetest) get(w http.ResponseWriter, req *http.Request) {
	n := ft.db.get(sanitizePath(req.URL.Path))
	w.Header().Add("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(n); err != nil {
		log.Printf("Error encoding json: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
	}
}
<commit_msg>adding methods to directly interact with the firetest database<commit_after>/*
Package firetest provides utilities for Firebase testing

*/
package firetest

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"
)

const (
	missingJSONExtension = "append .json to your request URI to use the REST API"
	missingBody          = `{"error":"Error: No data supplied."}`
	invalidJSON          = `{"error":"Invalid data; couldn't parse JSON object, array, or value. Perhaps you're using invalid characters in your key names."}`
)

// Firetest is a Firebase server implementation
type Firetest struct {
	// URL of form http://ipaddr:port with no trailing slash
	URL string

	listener net.Listener
	db       *treeDB
}

// New creates a new Firetest server
func New() *Firetest {
	return &Firetest{
		db: newTree(),
	}
}

func sanitizePath(p string) string {
	s := strings.Trim(p, "/")
	return strings.TrimSuffix(s, ".json")
}

// Start starts the server
func (ft *Firetest) Start() {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		if l, err = net.Listen("tcp6", "[::1]:0"); err != nil {
			panic(fmt.Errorf("failed to listen on a port: %v", err))
		}
	}
	ft.listener = l

	s := http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		ft.serveHTTP(w, req)
	})}
	go func() {
		if err := s.Serve(l); err != nil {
			log.Printf("error serving: %s", err)
		}

		ft.Close()
	}()
	ft.URL = "http://" + ft.listener.Addr().String()
}

// Close closes the server
func (ft *Firetest) Close() {
	if ft.listener != nil {
		ft.listener.Close()
	}
}

func (ft *Firetest) serveHTTP(w http.ResponseWriter, req *http.Request) {
	if !strings.HasSuffix(req.URL.String(), ".json") {
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte(missingJSONExtension))
		return
	}

	switch req.Method {
	case "PUT":
		ft.set(w, req)
	case "PATCH":
		ft.update(w, req)
	case "POST":
		ft.create(w, req)
	case "GET":
		ft.get(w, req)
	case "DELETE":
		ft.del(w, req)
	default:
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.Println("not implemented yet")
	}
}

// Set writes data to at the given location.
// This will overwrite any data at this location and all child locations.
//
// Reference https://www.firebase.com/docs/rest/api/#section-put
func (ft *Firetest) Set(path string, v interface{}) {
	ft.db.add(sanitizePath(path), newNode(v))
}

func (ft *Firetest) set(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil || len(body) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(missingBody))
		return
	}

	var v interface{}
	if err := json.Unmarshal(body, &v); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(invalidJSON))
		return
	}

	ft.Set(req.URL.Path, v)
	w.Write(body)
}

// Update writes the enumerated children to this the given location.
// This will overwrite only children enumerated in the "value" parameter
// and will leave others untouched. Note that the update function is equivalent
// to calling Set() on the named children; it does not recursively update children
// if they are objects. Passing null as a value for a child is equivalent to
// calling remove() on that child.
//
// Reference https://www.firebase.com/docs/rest/api/#section-patch
func (ft *Firetest) Update(path string, v interface{}) {
	path = sanitizePath(path)
	if v == nil {
		ft.db.del(path)
	} else {
		ft.db.update(path, newNode(v))
	}
}

func (ft *Firetest) update(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil || len(body) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(missingBody))
		return
	}

	var v interface{}
	if err := json.Unmarshal(body, &v); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(invalidJSON))
		return
	}
	ft.Update(req.URL.Path, v)
	w.Write(body)
}

// Create generates a new child under the given location
// using a unique name and returns the name
//
// Reference https://www.firebase.com/docs/rest/api/#section-post
func (ft *Firetest) Create(path string, v interface{}) string {
	src := []byte(fmt.Sprint(time.Now().UnixNano()))
	name := "~" + base64.StdEncoding.EncodeToString(src)
	path = fmt.Sprintf("%s/%s", sanitizePath(path), name)

	ft.db.add(path, newNode(v))
	return name
}

func (ft *Firetest) create(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	if err != nil || len(body) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(missingBody))
		return
	}

	var v interface{}
	if err := json.Unmarshal(body, &v); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(invalidJSON))
		return
	}

	name := ft.Create(req.URL.Path, v)
	rtn := map[string]string{"name": name}
	if err := json.NewEncoder(w).Encode(rtn); err != nil {
		log.Printf("Error encoding json: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
	}
}

// Delete removes the data at the requested location.
// Any data at child locations will also be deleted.
//
// Reference https://www.firebase.com/docs/rest/api/#section-delete
func (ft *Firetest) Delete(path string) {
	ft.db.del(sanitizePath(path))
}

func (ft *Firetest) del(w http.ResponseWriter, req *http.Request) {
	ft.Delete(req.URL.Path)
}

// Get retrieves the data and all its children at the
// requested location
//
// Reference https://www.firebase.com/docs/rest/api/#section-get
func (ft *Firetest) Get(path string) (v interface{}) {
	n := ft.db.get(sanitizePath(path))
	if n != nil {
		v = n.objectify()
	}
	return v
}

func (ft *Firetest) get(w http.ResponseWriter, req *http.Request) {
	w.Header().Add("Content-Type", "application/json")

	v := ft.Get(req.URL.Path)
	if err := json.NewEncoder(w).Encode(v); err != nil {
		log.Printf("Error encoding json: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2015  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package test

import (
	"bytes"
	"errors"
	"io"
	"os"
	"os/exec"
	"path"

	"github.com/tsuna/gohbase"
	"github.com/tsuna/gohbase/hrpc"
	"golang.org/x/net/context"
)

// This error is returned when the HBASE_HOME environment variable is unset
var ErrHomeUnset = errors.New("Environment variable HBASE_HOME is not set")

// getShellCmd returns a new shell subprocess (already started) along with its
// stdin
func getShellCmd() (*exec.Cmd, io.WriteCloser, error) {
	hbaseHome := os.Getenv("HBASE_HOME")
	if len(hbaseHome) == 0 {
		return nil, nil, ErrHomeUnset
	}
	hbaseShell := path.Join(hbaseHome, "bin", "hbase")
	cmd := exec.Command(hbaseShell, "shell")
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, nil, err
	}

	err = cmd.Start()
	if err != nil {
		stdin.Close()
		return nil, nil, err
	}
	return cmd, stdin, nil
}

// CreateTable finds the HBase shell via the HBASE_HOME environment variable,
// and creates the given table with the given families
func CreateTable(host, table string, cFamilies []string) error {
	// If the table exists, delete it
	DeleteTable(host, table)
	// Don't check the error, since one will be returned if the table doesn't
	// exist

	cmd, stdin, err := getShellCmd()
	if err != nil {
		return err
	}

	var buf bytes.Buffer
	buf.WriteString("create '" + table + "'")

	for _, f := range cFamilies {
		buf.WriteString(", '")
		buf.WriteString(f)
		buf.WriteString("'")
	}
	buf.WriteString("\n")

	stdin.Write(buf.Bytes())
	stdin.Write([]byte("exit\n"))

	return cmd.Wait()
}

// DeleteTable finds the HBase shell via the HBASE_HOME environment variable,
// and disables and drops the given table
func DeleteTable(host, table string) error {
	ac := gohbase.NewClient(host, gohbase.Admin())
	dit := hrpc.NewDisableTable(context.Background(), []byte(table))
	_, err := ac.SendRPC(dit)
	if err != nil {
		return err
	}

	det := hrpc.NewDeleteTable(context.Background(), []byte(table))
	_, err = ac.SendRPC(det)
	if err != nil {
		return err
	}
	return nil
}

// LaunchRegionServers uses the script local-regionservers.sh to create new
// RegionServers. Fails silently if server already exists.
// Ex. LaunchRegions([]string{"2", "3"}) launches two servers with id=2,3
func LaunchRegionServers(servers []string) {
	hh := os.Getenv("HBASE_HOME")
	servers = append([]string{"start"}, servers...)
	exec.Command(hh+"/bin/local-regionservers.sh", servers...).Run()
}

// StopRegionServers uses the script local-regionservers.sh to stop existing
// RegionServers. Fails silently if server isn't running.
func StopRegionServers(servers []string) {
	hh := os.Getenv("HBASE_HOME")
	servers = append([]string{"stop"}, servers...)
	exec.Command(hh+"/bin/local-regionservers.sh", servers...).Run()
}
<commit_msg>test: Don't fail if we try to disable a table already disabled.<commit_after>// Copyright (C) 2015  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package test

import (
	"bytes"
	"errors"
	"io"
	"os"
	"os/exec"
	"path"
	"strings"

	"github.com/tsuna/gohbase"
	"github.com/tsuna/gohbase/hrpc"
	"golang.org/x/net/context"
)

// This error is returned when the HBASE_HOME environment variable is unset
var ErrHomeUnset = errors.New("Environment variable HBASE_HOME is not set")

// getShellCmd returns a new shell subprocess (already started) along with its
// stdin
func getShellCmd() (*exec.Cmd, io.WriteCloser, error) {
	hbaseHome := os.Getenv("HBASE_HOME")
	if len(hbaseHome) == 0 {
		return nil, nil, ErrHomeUnset
	}
	hbaseShell := path.Join(hbaseHome, "bin", "hbase")
	cmd := exec.Command(hbaseShell, "shell")
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, nil, err
	}

	err = cmd.Start()
	if err != nil {
		stdin.Close()
		return nil, nil, err
	}
	return cmd, stdin, nil
}

// CreateTable finds the HBase shell via the HBASE_HOME environment variable,
// and creates the given table with the given families
func CreateTable(host, table string, cFamilies []string) error {
	// If the table exists, delete it
	DeleteTable(host, table)
	// Don't check the error, since one will be returned if the table doesn't
	// exist

	cmd, stdin, err := getShellCmd()
	if err != nil {
		return err
	}

	var buf bytes.Buffer
	buf.WriteString("create '" + table + "'")

	for _, f := range cFamilies {
		buf.WriteString(", '")
		buf.WriteString(f)
		buf.WriteString("'")
	}
	buf.WriteString("\n")

	stdin.Write(buf.Bytes())
	stdin.Write([]byte("exit\n"))

	return cmd.Wait()
}

// DeleteTable finds the HBase shell via the HBASE_HOME environment variable,
// and disables and drops the given table
func DeleteTable(host, table string) error {
	// TODO: We leak this client.
	ac := gohbase.NewClient(host, gohbase.Admin())
	dit := hrpc.NewDisableTable(context.Background(), []byte(table))
	_, err := ac.SendRPC(dit)
	if err != nil {
		if !strings.Contains(err.Error(), "TableNotEnabledException") {
			return err
		}
	}

	det := hrpc.NewDeleteTable(context.Background(), []byte(table))
	_, err = ac.SendRPC(det)
	if err != nil {
		return err
	}
	return nil
}

// LaunchRegionServers uses the script local-regionservers.sh to create new
// RegionServers. Fails silently if server already exists.
// Ex. LaunchRegions([]string{"2", "3"}) launches two servers with id=2,3
func LaunchRegionServers(servers []string) {
	hh := os.Getenv("HBASE_HOME")
	servers = append([]string{"start"}, servers...)
	exec.Command(hh+"/bin/local-regionservers.sh", servers...).Run()
}

// StopRegionServers uses the script local-regionservers.sh to stop existing
// RegionServers. Fails silently if server isn't running.
func StopRegionServers(servers []string) {
	hh := os.Getenv("HBASE_HOME")
	servers = append([]string{"stop"}, servers...)
	exec.Command(hh+"/bin/local-regionservers.sh", servers...).Run()
}
<|endoftext|>


<commit_before>package main

import (
        "code.google.com/p/go.net/websocket"
        "net"
        "net/http"
        "log"
	"encoding/json"
	"fmt"
	"strings"
	"text/template"
	"./uuid"
)

const (
	HOST_NAME = "localhost"
	PORT_NUMBER = "8080"
	APPSERVER_API_PREFIX = "/notify/"
)

type Client struct {
	Websocket *websocket.Conn
	UAID string
	Ip   string
	Port float64
}

type Channel struct {
	uaid string

	ChannelID string `json:"channelID"`
	Version string   `json:"version"`
}

// Mapping from a UAID to the Client object
var gConnectedClients map[string]*Client;

// Mapping from a UAID to all channels owned by that UAID
var gUAIDToChannel map[string][]*Channel;

// Mapping from a ChannelID to the cooresponding Channel
var gChannelIDToChannel map[string] *Channel;

func handleRegister(client *Client, f map[string]interface{}) {
	log.Println(" -> handleRegister");

	if (f["channelID"] == nil) {
		log.Println("channelID is missing!");
		return;
	}

	var channelID = f["channelID"].(string);

	// TODO https!
	var pushEndpoint = "http://" + HOST_NAME + ":" + PORT_NUMBER + APPSERVER_API_PREFIX + channelID;

	channel := &Channel{client.UAID, channelID, ""};

	gUAIDToChannel[client.UAID] = append(gUAIDToChannel[client.UAID], channel);
	gChannelIDToChannel[channelID] = channel;
	log.Println("gChannelIDToChannel ", channelID, gChannelIDToChannel[channelID]);

	type RegisterResponse struct {
		Name string          `json:"messageType"`
		Status int           `json:"status"`
		PushEndpoint string  `json:"pushEndpoint"`
		ChannelID string     `json:"channelID"`
	}

	register := RegisterResponse{"register", 200, pushEndpoint, channelID}

	j, err := json.Marshal(register);
	if err != nil {
		log.Println("Could not convert hello response to json %s",err)
		return;
        }

	if err = websocket.Message.Send(client.Websocket, string(j)); err != nil {
		// we could not send the message to a peer
		log.Println("Could not send message to ", client.Websocket, err.Error())
	}
}

func handleHello(client *Client, f map[string]interface{}) {
	log.Println(" -> handleHello");

	status := 200;

	if (f["uaid"] == nil) {
		uaid, err := uuid.GenUUID()
		if  err != nil {
			status = 400;
			log.Println("GenUUID error %s",err)
		}
		client.UAID = uaid;
	} else {
		client.UAID = f["uaid"].(string)

        // BUG(nikhilm): Does not deal with sending
        // a new UAID if their is a channel that was sent
        // by the UA which the server does not know about.
        // Which means in the case of this memory only server
        // it should actually always send a new UAID when it was
        // restarted
        if (f["channelIDs"] != nil) {
            for _, foo := range f["channelIDs"].([]interface {}) {
                channelID := foo.(string)
                log.Println("Got CHID ", channelID);
                c := &Channel{client.UAID, channelID, ""};
                gUAIDToChannel[client.UAID] = append(gUAIDToChannel[client.UAID], c);
                gChannelIDToChannel[channelID] = c;
            }
        }
	}

	if (f["interface"] != nil) {
		m := f["interface"].(map[string]interface{})
		client.Ip = m["ip"].(string);
		client.Port = m["port"].(float64);
	}

	type HelloResponse struct {
		Name string          `json:"messageType"`
		Status int           `json:"status"`
		UAID string          `json:"uaid"`
		Channels []*Channel  `json:"channelIDs"`
	}

	hello := HelloResponse{"hello", status, client.UAID, gUAIDToChannel[client.UAID]}

	j, err := json.Marshal(hello);
	if err != nil {
		log.Println("Could not convert hello response to json %s",err)
		return;
        }

	log.Println("going to send:  \n  ", string(j));
	if err = websocket.Message.Send(client.Websocket, string(j)); err != nil {
		log.Println("Could not send message to ", client.Websocket, err.Error())
	}
}

func handleAck(client *Client, f map[string]interface{}) {
	log.Println(" -> ack");
}

func pushHandler(ws *websocket.Conn) {

    log.Println("pushHandler: new websocket connection");

	client := &Client{ws, "", "", 0}

	for {
		var f map[string]interface{}

		var err error
		if err = websocket.JSON.Receive(ws, &f); err != nil {
			log.Println("Websocket Disconnected.", err.Error())
			break;
		}

		log.Println("hi!");

		switch f["messageType"] {
		case "hello":
			handleHello(client, f);
			gConnectedClients[client.UAID] = client;
			break;

		case "register":
			handleRegister(client, f);
			break;

		case "ack":
			handleAck(client, f);
			break;
		default:
			log.Println(" -> Unknown", f);
			break;
		}
	}
	log.Println("Closing Websocket!");
	ws.Close();
	gConnectedClients[client.UAID].Websocket = nil;
}

func notifyHandler(w http.ResponseWriter, r *http.Request) {
    log.Println("Got notification from app server ", r.URL);

	if r.Method != "PUT" {
	    log.Println("NOT A PUT")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Method must be PUT."))
		return
	}

	channelID := strings.Replace(r.URL.Path, APPSERVER_API_PREFIX, "", 1)
	log.Println("channelID ", channelID);

	if (strings.Contains(channelID, "/")) {
	    log.Println("Could not find a valid channelID");
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Could not find a valid channelID."))
		return;
	}

	value := r.FormValue("version")

	if value == "" {
	    log.Println("Could not find version");
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Could not find version."))
		return;
	}

	channel, found := gChannelIDToChannel[channelID];
    if !found {
        log.Println("Could not find channel " + channelID);
        return;
    }
	channel.Version = value;

	client := gConnectedClients[channel.uaid];
	if (client == nil) {
		log.Println("no known client for the channel.");
	} else if (client.Websocket == nil) {
		wakeupClient(client)
	} else {
		sendNotificationToClient(client, channel)
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK"))
}

func wakeupClient(client *Client) {

	log.Println("wakeupClient: ", client);
	service := fmt.Sprintf("%s:%g", client.Ip, client.Port)

	udpAddr, err := net.ResolveUDPAddr("up4", service)
	if err != nil {
		log.Println("ResolveUDPAddr error ", err.Error())
		return
	}

	conn, err := net.DialUDP("udp", nil, udpAddr)
	if err != nil {
		log.Println("DialUDP error ", err.Error())
		return
	}

	_, err = conn.Write([]byte("anything"))
	if err != nil {
		log.Println("UDP Write error ", err.Error())
		return
	}

}

func sendNotificationToClient(client *Client, channel *Channel)  {

	type NotificationResponse struct {
		Name string          `json:"messageType"`
		Channels []Channel   `json:"updates"`
	}

	var channels []Channel;
	channels = append(channels, *channel);

	notification := NotificationResponse{"notification",  channels}

	j, err := json.Marshal(notification);
	if err != nil {
		log.Println("Could not convert hello response to json %s",err)
		return;
        }

	log.Println("going to send:  \n  ", string(j));
	if err = websocket.Message.Send(client.Websocket, string(j)); err != nil {
		log.Println("Could not send message to ", channel, err.Error())
	}
    log.Println("Sent notification to client");

}

func admin(w http.ResponseWriter, r *http.Request) {

	type User struct {
		UAID string
		Connected bool
		Channels []*Channel
	}

	type Arguments struct {
		PushEndpointPrefix   string
		Users []User
	}

// TODO https!
	arguments := Arguments{"http://" + HOST_NAME + ":" + PORT_NUMBER + APPSERVER_API_PREFIX, nil}

	for k := range gUAIDToChannel {
		connected := gConnectedClients[k] != nil;
		u := User{k, connected, gUAIDToChannel[k]};
		arguments.Users = append(arguments.Users, u);
	}

	t := template.New("users.template")
	s1, _ := t.ParseFiles("templates/users.template");
	s1.Execute(w, arguments)
}

func main() {

	gUAIDToChannel = make(map[string][]*Channel)
	gChannelIDToChannel = make(map[string]*Channel)

	gConnectedClients = make(map[string]*Client)

	//http.Handle("/", http.FileServer(http.Dir(".")))

	http.HandleFunc("/admin", admin)

	http.Handle("/", websocket.Handler(pushHandler))

	http.HandleFunc(APPSERVER_API_PREFIX, notifyHandler);

	log.Println("Listening on", HOST_NAME + ":" + PORT_NUMBER );
	log.Fatal(http.ListenAndServe(HOST_NAME  + ":" + PORT_NUMBER , nil))
}

<commit_msg>Sent code through go fmt<commit_after>package main

import (
	"./uuid"
	"code.google.com/p/go.net/websocket"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"text/template"
)

const (
	HOST_NAME = "localhost"
	PORT_NUMBER = "8080"
	APPSERVER_API_PREFIX = "/notify/"
)

type Client struct {
	Websocket *websocket.Conn
	UAID      string
	Ip        string
	Port      float64
}

type Channel struct {
	uaid string

	ChannelID string `json:"channelID"`
	Version   string `json:"version"`
}

// Mapping from a UAID to the Client object
var gConnectedClients map[string]*Client

// Mapping from a UAID to all channels owned by that UAID
var gUAIDToChannel map[string][]*Channel

// Mapping from a ChannelID to the cooresponding Channel
var gChannelIDToChannel map[string]*Channel

func handleRegister(client *Client, f map[string]interface{}) {
	log.Println(" -> handleRegister")

	if f["channelID"] == nil {
		log.Println("channelID is missing!")
		return
	}

	var channelID = f["channelID"].(string)

	// TODO https!
	var pushEndpoint = "http://" + HOST_NAME + ":" + PORT_NUMBER + APPSERVER_API_PREFIX + channelID

	channel := &Channel{client.UAID, channelID, ""}

	gUAIDToChannel[client.UAID] = append(gUAIDToChannel[client.UAID], channel)
	gChannelIDToChannel[channelID] = channel
	log.Println("gChannelIDToChannel ", channelID, gChannelIDToChannel[channelID])

	type RegisterResponse struct {
		Name         string `json:"messageType"`
		Status       int    `json:"status"`
		PushEndpoint string `json:"pushEndpoint"`
		ChannelID    string `json:"channelID"`
	}

	register := RegisterResponse{"register", 200, pushEndpoint, channelID}

	j, err := json.Marshal(register)
	if err != nil {
		log.Println("Could not convert hello response to json %s", err)
		return
	}

	if err = websocket.Message.Send(client.Websocket, string(j)); err != nil {
		// we could not send the message to a peer
		log.Println("Could not send message to ", client.Websocket, err.Error())
	}
}

func handleHello(client *Client, f map[string]interface{}) {
	log.Println(" -> handleHello")

	status := 200

	if f["uaid"] == nil {
		uaid, err := uuid.GenUUID()
		if err != nil {
			status = 400
			log.Println("GenUUID error %s", err)
		}
		client.UAID = uaid
	} else {
		client.UAID = f["uaid"].(string)

		// BUG(nikhilm): Does not deal with sending
		// a new UAID if their is a channel that was sent
		// by the UA which the server does not know about.
		// Which means in the case of this memory only server
		// it should actually always send a new UAID when it was
		// restarted
		if f["channelIDs"] != nil {
			for _, foo := range f["channelIDs"].([]interface{}) {
				channelID := foo.(string)
				log.Println("Got CHID ", channelID)
				c := &Channel{client.UAID, channelID, ""}
				gUAIDToChannel[client.UAID] = append(gUAIDToChannel[client.UAID], c)
				gChannelIDToChannel[channelID] = c
			}
		}
	}

	if f["interface"] != nil {
		m := f["interface"].(map[string]interface{})
		client.Ip = m["ip"].(string)
		client.Port = m["port"].(float64)
	}

	type HelloResponse struct {
		Name     string     `json:"messageType"`
		Status   int        `json:"status"`
		UAID     string     `json:"uaid"`
		Channels []*Channel `json:"channelIDs"`
	}

	hello := HelloResponse{"hello", status, client.UAID, gUAIDToChannel[client.UAID]}

	j, err := json.Marshal(hello)
	if err != nil {
		log.Println("Could not convert hello response to json %s", err)
		return
	}

	log.Println("going to send:  \n  ", string(j))
	if err = websocket.Message.Send(client.Websocket, string(j)); err != nil {
		log.Println("Could not send message to ", client.Websocket, err.Error())
	}
}

func handleAck(client *Client, f map[string]interface{}) {
	log.Println(" -> ack")
}

func pushHandler(ws *websocket.Conn) {

	log.Println("pushHandler: new websocket connection")

	client := &Client{ws, "", "", 0}

	for {
		var f map[string]interface{}

		var err error
		if err = websocket.JSON.Receive(ws, &f); err != nil {
			log.Println("Websocket Disconnected.", err.Error())
			break
		}

		log.Println("hi!")

		switch f["messageType"] {
		case "hello":
			handleHello(client, f)
			gConnectedClients[client.UAID] = client
			break

		case "register":
			handleRegister(client, f)
			break

		case "ack":
			handleAck(client, f)
			break
		default:
			log.Println(" -> Unknown", f)
			break
		}
	}
	log.Println("Closing Websocket!")
	ws.Close()
	gConnectedClients[client.UAID].Websocket = nil
}

func notifyHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("Got notification from app server ", r.URL)

	if r.Method != "PUT" {
		log.Println("NOT A PUT")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Method must be PUT."))
		return
	}

	channelID := strings.Replace(r.URL.Path, APPSERVER_API_PREFIX, "", 1)
	log.Println("channelID ", channelID)

	if strings.Contains(channelID, "/") {
		log.Println("Could not find a valid channelID")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Could not find a valid channelID."))
		return
	}

	value := r.FormValue("version")

	if value == "" {
		log.Println("Could not find version")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Could not find version."))
		return
	}

	channel, found := gChannelIDToChannel[channelID]
	if !found {
		log.Println("Could not find channel " + channelID)
		return
	}
	channel.Version = value

	client := gConnectedClients[channel.uaid]
	if client == nil {
		log.Println("no known client for the channel.")
	} else if client.Websocket == nil {
		wakeupClient(client)
	} else {
		sendNotificationToClient(client, channel)
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK"))
}

func wakeupClient(client *Client) {

	log.Println("wakeupClient: ", client)
	service := fmt.Sprintf("%s:%g", client.Ip, client.Port)

	udpAddr, err := net.ResolveUDPAddr("up4", service)
	if err != nil {
		log.Println("ResolveUDPAddr error ", err.Error())
		return
	}

	conn, err := net.DialUDP("udp", nil, udpAddr)
	if err != nil {
		log.Println("DialUDP error ", err.Error())
		return
	}

	_, err = conn.Write([]byte("anything"))
	if err != nil {
		log.Println("UDP Write error ", err.Error())
		return
	}

}

func sendNotificationToClient(client *Client, channel *Channel) {

	type NotificationResponse struct {
		Name     string    `json:"messageType"`
		Channels []Channel `json:"updates"`
	}

	var channels []Channel
	channels = append(channels, *channel)

	notification := NotificationResponse{"notification", channels}

	j, err := json.Marshal(notification)
	if err != nil {
		log.Println("Could not convert hello response to json %s", err)
		return
	}

	log.Println("going to send:  \n  ", string(j))
	if err = websocket.Message.Send(client.Websocket, string(j)); err != nil {
		log.Println("Could not send message to ", channel, err.Error())
	}
	log.Println("Sent notification to client")

}

func admin(w http.ResponseWriter, r *http.Request) {

	type User struct {
		UAID      string
		Connected bool
		Channels  []*Channel
	}

	type Arguments struct {
		PushEndpointPrefix string
		Users              []User
	}

	// TODO https!
	arguments := Arguments{"http://" + HOST_NAME + ":" + PORT_NUMBER + APPSERVER_API_PREFIX, nil}

	for k := range gUAIDToChannel {
		connected := gConnectedClients[k] != nil
		u := User{k, connected, gUAIDToChannel[k]}
		arguments.Users = append(arguments.Users, u)
	}

	t := template.New("users.template")
	s1, _ := t.ParseFiles("templates/users.template")
	s1.Execute(w, arguments)
}

func main() {

	gUAIDToChannel = make(map[string][]*Channel)
	gChannelIDToChannel = make(map[string]*Channel)

	gConnectedClients = make(map[string]*Client)

	//http.Handle("/", http.FileServer(http.Dir(".")))

	http.HandleFunc("/admin", admin)

	http.Handle("/", websocket.Handler(pushHandler))

	http.HandleFunc(APPSERVER_API_PREFIX, notifyHandler)

	log.Println("Listening on", HOST_NAME+":"+PORT_NUMBER)
	log.Fatal(http.ListenAndServe(HOST_NAME+":"+PORT_NUMBER, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
)

var port = flag.Int("port", 8080, "which port to listen on")
var authApiKey string
var dbUrl string

var rootMetrics = map[string]string{
	"agentVersion": "host.agent_version",
	"os":           "host.os",
	"python":       "host.python_version",
	"uuid":         "host.uuid",

	"system.load.1":       "system.load.1",
	"system.load.5":       "system.load.5",
	"system.load.15":      "system.load.15",
	"system.load.norm.1":  "system.load.norm.1",
	"system.load.norm.5":  "system.load.norm.5",
	"system.load.norm.15": "system.load.norm.15",

	"cpuIdle":   "system.cpu.idle",
	"cpuUser":   "system.cpu.user",
	"cpuWait":   "system.cpu.iowait",
	"cpuSystem": "system.cpu.system",
	"cpuStolen": "system.cpu.stolen",

	"memBuffers":       "system.mem.buffered",
	"memPhysPctUsable": "system.mem.pct_usable",
	"memShared":        "system.mem.shared",
	"memPhysTotal":     "system.mem.total",
	"memPhysUsable":    "system.mem.usable",
	"memCached":        "system.mem.cached",
	"memPhysUsed":      "system.mem.used",
	"memPhysFree":      "system.mem.free",

	"memSwapFree":    "system.swap.free",
	"memSwapTotal":   "system.swap.total",
	"memSwapUsed":    "system.swap.used",
	"memSwapPctFree": "system.swap.pct_free",
}

var diskMetrics = []string{
	"device",
	"total",
	"used",
	"free",
	"in_use",
	"mount",
}

type Metric struct {
	Name    string          `json:"name"`
	Columns []string        `json:"columns"`
	Points  [][]interface{} `json:"points"`
}

type StatsdSeries struct {
	Series []*StatsdMetric `json:"series"`
}

type StatsdMetric struct {
	Tags     []string        `json:"tags"`
	Metric   string          `json:"metric"`
	Interval float64         `json:"interval"`
	Host     string          `json:"host"`
	Points   [][]interface{} `json:"points"`
	Type     string          `json:"type"`
}

func NewMetric(host, name string, timestamp uint64, value interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "value", "hostname"}
	points := [][]interface{}{{timestamp, value, host}}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" {
				points[0][2] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func NewMetricGroup(host, name string, timestamp uint64, values map[string]interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "hostname"}
	points := [][]interface{}{{timestamp, host}}
	for k, v := range values {
		columns = append(columns, k)
		points[0] = append(points[0], v)
	}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" {
				points[0][1] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func PushMetrics(metrics []*Metric) {
	if len(metrics) == 0 {
		return
	}
	log.Printf("Pushing %d metrics to InfluxDB\n", len(metrics))

	body, err := json.Marshal(metrics)
	if err != nil {
		log.Println(err)
		return
	}

	resp, err := http.Post(dbUrl, "application/json", bytes.NewBuffer(body))
	if err != nil {
		log.Println(err)
	} else if resp.StatusCode != 200 {
		log.Println("Got Response: ", resp.Status)
	}
}

func mapStatsd(series []*StatsdMetric) []*Metric {
	metrics := make([]*Metric, len(series))
	host := ""
	for i, metric := range series {
		if len(metric.Host) > 0 {
			host = metric.Host
		}
		columns := []string{"time", "value", "hostname"}
		points := metric.Points
		for i, _ := range points {
			points[i][0] = uint64(points[i][0].(float64) * 1000)
			points[i] = append(points[i], host)
		}
		columns = append(columns, "metric_interval", "metric_type")
		for i, _ := range points {
			points[i] = append(points[i], metric.Interval, metric.Type)
		}
		if metric.Tags != nil {
			for _, tag := range metric.Tags {
				split := strings.SplitN(tag, ":", 2)
				if split[0] == "hostname" {
					for i, _ := range points {
						points[i][2] = split[1]
					}
				} else {
					columns = append(columns, split[0])
					for i, _ := range points {
						points[i] = append(points[i], split[1])
					}
				}
			}
		}
		metrics[i] = &Metric{metric.Metric, columns, points}
	}
	for _, metric := range metrics {
		for _, points := range metric.Points {
			if len(points[2].(string)) == 0 {
				points[2] = host
			}
		}
	}
	return metrics
}

func mapMetrics(data map[string]interface{}) []*Metric {
	metrics := []*Metric{}
	host := data["internalHostname"].(string)
	timestamp := uint64(data["collection_timestamp"].(float64) * 1000)

	delete(data, "apiKey")
	delete(data, "internalHostname")
	delete(data, "collection_timestamp")

	values := make(map[string]map[string]interface{})

	log.Printf("Parsing metrics for: %s\n", host)
	for key, value := range data {
		name, ok := rootMetrics[key]
		if ok {
			index := strings.LastIndexAny(name, ".")
			group_name := name[:index]
			group := values[group_name]
			if group == nil {
				group = make(map[string]interface{})
				values[group_name] = group
			}
			group[name[index+1:]] = value
			delete(data, key)
		}
	}
	for name, group := range values {
		metrics = append(metrics, NewMetricGroup(host, name, timestamp, group, nil))
	}

	metrics = append(metrics, mapDiskMetrics("system.disk", timestamp, host, data["diskUsage"].([]interface{})))
	delete(data, "diskUsage")
	metrics = append(metrics, mapDiskMetrics("system.fs.inodes", timestamp, host, data["inodes"].([]interface{})))
	delete(data, "inodes")

	debug, err := json.MarshalIndent(data, "", "  ")
	if err == nil {
		fmt.Println(string(debug))
	}
	return metrics
}

func mapDiskMetrics(name string, timestamp uint64, host string, data []interface{}) *Metric {
	columns := append([]string{"time", "hostname"}, diskMetrics...)
	points := make([][]interface{}, len(data))
	metric := &Metric{name, columns, points}
	for i, disk := range data {
		fields := disk.([]interface{})
		percent := fields[4].(string)
		parse, _ := strconv.ParseFloat(percent[:len(percent)-1], 64)
		fields[4] = parse / 100.0
		points[i] = append([]interface{}{timestamp, host}, fields...)
	}
	return metric
}

func handleIntake(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	data := make(map[string]interface{})
	err = json.Unmarshal(buf.Bytes(), &data)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	metrics := mapMetrics(data)
	go PushMetrics(metrics)

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApi(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	series := StatsdSeries{}
	err = json.Unmarshal(buf.Bytes(), &series)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	metrics := mapStatsd(series.Series)
	go PushMetrics(metrics)

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApiKey(w http.ResponseWriter, req *http.Request) bool {
	if req.UserAgent() == "Datadog-Status-Check" {
		io.WriteString(w, "STILL-ALIVE\n")
		return true
	}
	err := req.ParseForm()
	if err != nil {
		log.Println("Error parsing form:", err)
		http.Error(w, err.Error(), 500)
		return true
	}
	values := req.Form
	api_key := values.Get("api_key")
	delete(values, "api_key")
	if len(authApiKey) > 0 && api_key != authApiKey {
		log.Println("Got bad API key:", api_key)
		http.Error(w, "Bad API Key", 403)
		return true
	}
	return false
}

func main() {
	flag.Parse()
	authApiKey = os.Getenv("API_KEY")
	if len(authApiKey) == 0 {
		log.Println("Warning: API_KEY is blank")
	}
	dbUrl = os.Getenv("DB_URL")
	if len(dbUrl) == 0 {
		dbUrl = "http://localhost:8086/db/datadog/series?u=root&p=root"
	}

	log.Println("dd-house listening on", *port)

	http.HandleFunc("/intake", handleIntake)
	http.HandleFunc("/api/v1/series/", handleApi)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(*port), nil))
}
<commit_msg>Push integration metrics to influxdb<commit_after>package main

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
)

var port = flag.Int("port", 8080, "which port to listen on")
var authApiKey string
var dbUrl string

var rootMetrics = map[string]string{
	"agentVersion": "host.agent_version",
	"os":           "host.os",
	"python":       "host.python_version",
	"uuid":         "host.uuid",

	"system.load.1":       "system.load.1",
	"system.load.5":       "system.load.5",
	"system.load.15":      "system.load.15",
	"system.load.norm.1":  "system.load.norm.1",
	"system.load.norm.5":  "system.load.norm.5",
	"system.load.norm.15": "system.load.norm.15",

	"cpuIdle":   "system.cpu.idle",
	"cpuUser":   "system.cpu.user",
	"cpuWait":   "system.cpu.iowait",
	"cpuSystem": "system.cpu.system",
	"cpuStolen": "system.cpu.stolen",

	"memBuffers":       "system.mem.buffered",
	"memPhysPctUsable": "system.mem.pct_usable",
	"memShared":        "system.mem.shared",
	"memPhysTotal":     "system.mem.total",
	"memPhysUsable":    "system.mem.usable",
	"memCached":        "system.mem.cached",
	"memPhysUsed":      "system.mem.used",
	"memPhysFree":      "system.mem.free",

	"memSwapFree":    "system.swap.free",
	"memSwapTotal":   "system.swap.total",
	"memSwapUsed":    "system.swap.used",
	"memSwapPctFree": "system.swap.pct_free",
}

var diskMetrics = []string{
	"device",
	"total",
	"used",
	"free",
	"in_use",
	"mount",
}

type Metric struct {
	Name    string          `json:"name"`
	Columns []string        `json:"columns"`
	Points  [][]interface{} `json:"points"`
}

type StatsdSeries struct {
	Series []*StatsdMetric `json:"series"`
}

type StatsdMetric struct {
	Tags     []string        `json:"tags"`
	Metric   string          `json:"metric"`
	Interval float64         `json:"interval"`
	Host     string          `json:"host"`
	Points   [][]interface{} `json:"points"`
	Type     string          `json:"type"`
}

func NewMetric(host, name string, timestamp uint64, value interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "value", "hostname"}
	points := [][]interface{}{{timestamp, value, host}}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" {
				points[0][2] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func NewMetricGroup(host, name string, timestamp uint64, values map[string]interface{}, tags map[string]string) *Metric {
	columns := []string{"time", "hostname"}
	points := [][]interface{}{{timestamp, host}}
	for k, v := range values {
		columns = append(columns, k)
		points[0] = append(points[0], v)
	}
	if tags != nil {
		for k, v := range tags {
			if k == "hostname" && len(v) > 0 {
				points[0][1] = v
			} else {
				columns = append(columns, k)
				points[0] = append(points[0], v)
			}
		}
	}
	return &Metric{name, columns, points}
}

func PushMetrics(metrics []*Metric) {
	if len(metrics) == 0 {
		return
	}
	log.Printf("Pushing %d metrics to InfluxDB\n", len(metrics))

	body, err := json.Marshal(metrics)
	if err != nil {
		log.Println(err)
		return
	}

	resp, err := http.Post(dbUrl, "application/json", bytes.NewBuffer(body))
	if err != nil {
		log.Println(err)
	} else if resp.StatusCode != 200 {
		log.Println("Got Response: ", resp.Status)
	}
}

func mapStatsd(series []*StatsdMetric) []*Metric {
	metrics := make([]*Metric, len(series))
	host := ""
	for i, metric := range series {
		if len(metric.Host) > 0 {
			host = metric.Host
		}
		columns := []string{"time", "value", "hostname"}
		points := metric.Points
		for i, _ := range points {
			points[i][0] = uint64(points[i][0].(float64) * 1000)
			points[i] = append(points[i], host)
		}
		columns = append(columns, "metric_interval", "metric_type")
		for i, _ := range points {
			points[i] = append(points[i], metric.Interval, metric.Type)
		}
		if metric.Tags != nil {
			for _, tag := range metric.Tags {
				split := strings.SplitN(tag, ":", 2)
				if split[0] == "hostname" {
					for i, _ := range points {
						points[i][2] = split[1]
					}
				} else {
					columns = append(columns, split[0])
					for i, _ := range points {
						points[i] = append(points[i], split[1])
					}
				}
			}
		}
		metrics[i] = &Metric{metric.Metric, columns, points}
	}
	for _, metric := range metrics {
		for _, points := range metric.Points {
			if len(points[2].(string)) == 0 {
				points[2] = host
			}
		}
	}
	return metrics
}

func mapMetrics(data map[string]interface{}) []*Metric {
	metrics := []*Metric{}
	host := data["internalHostname"].(string)
	timestamp := uint64(data["collection_timestamp"].(float64) * 1000)

	delete(data, "apiKey")
	delete(data, "internalHostname")
	delete(data, "collection_timestamp")

	values := make(map[string]map[string]interface{})

	log.Printf("Parsing metrics for: %s\n", host)
	for key, value := range data {
		name, ok := rootMetrics[key]
		if ok {
			index := strings.LastIndexAny(name, ".")
			group_name := name[:index]
			group := values[group_name]
			if group == nil {
				group = make(map[string]interface{})
				values[group_name] = group
			}
			group[name[index+1:]] = value
			delete(data, key)
		}
	}
	for name, group := range values {
		metrics = append(metrics, NewMetricGroup(host, name, timestamp, group, nil))
	}

	metrics = append(metrics, mapDiskMetrics("system.disk", timestamp, host, data["diskUsage"].([]interface{})))
	delete(data, "diskUsage")
	metrics = append(metrics, mapDiskMetrics("system.fs.inodes", timestamp, host, data["inodes"].([]interface{})))
	delete(data, "inodes")
	metrics = append(metrics, mapExtraMetrics(host, data["metrics"].([]interface{}))...)
	delete(data, "metrics")

	debug, err := json.MarshalIndent(data, "", "  ")
	if err == nil {
		fmt.Println(string(debug))
	}
	return metrics
}

func mapDiskMetrics(name string, timestamp uint64, host string, data []interface{}) *Metric {
	columns := append([]string{"time", "hostname"}, diskMetrics...)
	points := make([][]interface{}, len(data))
	metric := &Metric{name, columns, points}
	for i, disk := range data {
		fields := disk.([]interface{})
		percent := fields[4].(string)
		parse, _ := strconv.ParseFloat(percent[:len(percent)-1], 64)
		fields[4] = parse / 100.0
		points[i] = append([]interface{}{timestamp, host}, fields...)
	}
	return metric
}

func mapExtraMetrics(host string, data []interface{}) []*Metric {
	metrics := make([]*Metric, len(data))
	for i, tmp := range data {
		metric := tmp.([]interface{})
		tags := make(map[string]string)
		fields := metric[3].(map[string]interface{})
		for k, v := range fields {
			if k == "tags" {
				tags2 := v.([]interface{})
				for _, tag := range tags2 {
					split := strings.SplitN(tag.(string), ":", 2)
					tags[split[0]] = split[1]
				}
			} else {
				tags[k] = v.(string)
			}
		}
		metrics[i] = NewMetric(host, metric[0].(string), uint64(metric[1].(float64)*1000), metric[2], tags)
	}
	return metrics
}

func handleIntake(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	data := make(map[string]interface{})
	err = json.Unmarshal(buf.Bytes(), &data)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	metrics := mapMetrics(data)
	go PushMetrics(metrics)

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApi(w http.ResponseWriter, req *http.Request) {
	if handleApiKey(w, req) {
		return
	}

	log.Println(req.Method, req.URL.Path)

	body := req.Body
	if req.Header.Get("Content-Encoding") == "deflate" {
		var err error
		body, err = zlib.NewReader(body)
		if err != nil {
			log.Println(err)
			io.WriteString(w, `{"status":"failed"}`)
			return
		}
	}
	buf := bytes.NewBuffer([]byte{})
	_, err := io.Copy(buf, body)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	series := StatsdSeries{}
	err = json.Unmarshal(buf.Bytes(), &series)
	if err != nil {
		log.Println(err)
		io.WriteString(w, `{"status":"failed"}`)
		return
	}

	metrics := mapStatsd(series.Series)
	go PushMetrics(metrics)

	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func handleApiKey(w http.ResponseWriter, req *http.Request) bool {
	if req.UserAgent() == "Datadog-Status-Check" {
		io.WriteString(w, "STILL-ALIVE\n")
		return true
	}
	err := req.ParseForm()
	if err != nil {
		log.Println("Error parsing form:", err)
		http.Error(w, err.Error(), 500)
		return true
	}
	values := req.Form
	api_key := values.Get("api_key")
	delete(values, "api_key")
	if len(authApiKey) > 0 && api_key != authApiKey {
		log.Println("Got bad API key:", api_key)
		http.Error(w, "Bad API Key", 403)
		return true
	}
	return false
}

func main() {
	flag.Parse()
	authApiKey = os.Getenv("API_KEY")
	if len(authApiKey) == 0 {
		log.Println("Warning: API_KEY is blank")
	}
	dbUrl = os.Getenv("DB_URL")
	if len(dbUrl) == 0 {
		dbUrl = "http://localhost:8086/db/datadog/series?u=root&p=root"
	}

	log.Println("dd-house listening on", *port)

	http.HandleFunc("/intake", handleIntake)
	http.HandleFunc("/api/v1/series/", handleApi)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(*port), nil))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"time"

	"github.com/influxdata/influxdb/client/v2"
)

// AkamaiPayload is a Golang representation of the Cloudmonitor JSON datastructure
type AkamaiPayload struct {
	// Content Provider ID
	CP string `json:"cp"`
	// Defines the format of the payload (?)
	Format  string `json:"format"`
	ID      string `json:"id"`
	Start   string `json:"start"`
	Type    string `json:"type"`
	Version string `json:"version"`

	Geo     GeoStruct     `json:"geo"`
	Message MessageStruct `json:"message"`
	NetPerf NetPerfStruct `json:"netPerf"`
	Network NetworkStruct `json:"network"`
	ReqHdr  ReqHdrStruct  `json:"reqHdr"`
	RespHdr RespHdrStruct `json:"respHdr"`
}

// GeoStruct is used for storing the JSON subfields
type GeoStruct struct {
	City    string `json:"city"`
	Country string `json:"country"`
	Lat     string `json:"lat"`
	Long    string `json:"long"`
	Region  string `json:"region"`
}

// MessageStruct is used for storing the JSON subfields
type MessageStruct struct {
	UA        string `json:"UA"`
	Bytes     string `json:"bytes"`
	CliIP     string `json:"cliIP"`
	FwdHost   string `json:"fwdHost"`
	Proto     string `json:"proto"`
	ProtoVer  string `json:"protoVer"`
	ReqHost   string `json:"reqHost"`
	ReqMethod string `json:"reqMethod"`
	ReqPath   string `json:"reqPath"`
	ReqPort   string `json:"reqPort"`
	RespCT    string `json:"respCT"`
	RespLen   string `json:"respLen"`
	Status    string `json:"status"`
}

// NetPerfStruct is used for storing the JSON subfields
type NetPerfStruct struct {
	Asnum        string `json:"asnum"`
	CacheStatus  string `json:"cacheStatus"`
	DownloadTime string `json:"downloadTime"`
	EdgeIP       string `json:"edgeIP"`
	FirstByte    string `json:"firstByte"`
	LastByte     string `json:"lastByte"`
	LastMileRTT  string `json:"lastMileRTT"`
}

// NetworkStruct is used for storing the JSON subfields
type NetworkStruct struct {
	Asnum       string `json:"asnum"`
	EdgeIP      string `json:"edgeIP"`
	Network     string `json:"network"`
	NetworkType string `json:"networkType"`
}

// ReqHdrStruct is used for storing the JSON subfields
type ReqHdrStruct struct {
	Cookie string `json:"cookie"`
}

// RespHdrStruct is used for storing the JSON subfields
type RespHdrStruct struct {
	Server  string `json:"server"`
	ContEnc string `json:"contEnc"`
}

// CreateObjects creates a list of AkamaiPayloads from a raw byte slice
func CreateObjects(jsonFile []byte) (_ []AkamaiPayload, err error) {
	var arrayObject []AkamaiPayload
	err = json.Unmarshal(jsonFile, &arrayObject)
	return arrayObject, err
}

// Handle parses the incoming request data
func Handle(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "OK! I got you Bro! -> ")
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Fatal(err)
	}
	obj, err := CreateObjects(body)
	if err != nil {
		log.Fatal(err)
	}
	c, err := client.NewHTTPClient(client.HTTPConfig{
		Addr: "http://influxDB:8086",
	})

	// Create a new point batch
	bp, err := client.NewBatchPoints(client.BatchPointsConfig{
		Database: "metrics",
		//Precision: "s",
	})
	if err != nil {
		log.Fatal(err)
	}

	for key, o := range obj {
		//payload[0].Geo["city"]
		fmt.Println("Server is:", o.RespHdr.Server, "KEY:", key)

		if err != nil {
			log.Fatal(err)
		}
		// Create a point and add to batch
		tags := map[string]string{
			"cp":          o.CP,
			"format":      o.Format,
			"city":        o.Geo.City,
			"country":     o.Geo.Country,
			"status":      o.Message.Status,
			"cacheStatus": o.NetPerf.CacheStatus,
			"proto":       o.Message.Proto,
			"protoVer":    o.Message.ProtoVer,
			"reqPort":     o.Message.ReqPort,
			"reqPath":     o.Message.ReqPath,
			"reqMethod":   o.Message.ReqMethod,
		}
		lat, err := strconv.ParseFloat(o.Geo.Lat, 64)
		long, err := strconv.ParseFloat(o.Geo.Long, 64)
		downloadTime, err := strconv.ParseInt(o.NetPerf.DownloadTime, 10, 32)
		status, err := strconv.ParseInt(o.Message.Status, 10, 32)

		if err != nil {
			log.Fatal(err)
		}
		fields := map[string]interface{}{
			"lat":          lat,
			"long":         long,
			"downloadtime": downloadTime,
			"status":       status,
		}

		pt, err := client.NewPoint("measurement", tags, fields, time.Now())
		if err != nil {
			log.Fatal(err)
		}
		bp.AddPoint(pt)
	}

	// Write the batch
	if err := c.Write(bp); err != nil {
		log.Fatal(err)
	}
}

func main() {

	http.HandleFunc("/", Handle)
	http.ListenAndServe(":9143", nil)
}
<commit_msg>add tags and fields<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"time"

	"github.com/influxdata/influxdb/client/v2"
)

// AkamaiPayload is a Golang representation of the Cloudmonitor JSON datastructure
type AkamaiPayload struct {
	// Content Provider ID
	CP string `json:"cp"`
	// Defines the format of the payload (?)
	Format  string `json:"format"`
	ID      string `json:"id"`
	Start   string `json:"start"`
	Type    string `json:"type"`
	Version string `json:"version"`

	Geo     GeoStruct     `json:"geo"`
	Message MessageStruct `json:"message"`
	NetPerf NetPerfStruct `json:"netPerf"`
	Network NetworkStruct `json:"network"`
	ReqHdr  ReqHdrStruct  `json:"reqHdr"`
	RespHdr RespHdrStruct `json:"respHdr"`
}

// GeoStruct is used for storing the JSON subfields
type GeoStruct struct {
	City    string `json:"city"`
	Country string `json:"country"`
	Lat     string `json:"lat"`
	Long    string `json:"long"`
	Region  string `json:"region"`
}

// MessageStruct is used for storing the JSON subfields
type MessageStruct struct {
	UA        string `json:"UA"`
	Bytes     string `json:"bytes"`
	CliIP     string `json:"cliIP"`
	FwdHost   string `json:"fwdHost"`
	Proto     string `json:"proto"`
	ProtoVer  string `json:"protoVer"`
	ReqHost   string `json:"reqHost"`
	ReqMethod string `json:"reqMethod"`
	ReqPath   string `json:"reqPath"`
	ReqPort   string `json:"reqPort"`
	RespCT    string `json:"respCT"`
	RespLen   string `json:"respLen"`
	Status    string `json:"status"`
}

// NetPerfStruct is used for storing the JSON subfields
type NetPerfStruct struct {
	Asnum        string `json:"asnum"`
	CacheStatus  string `json:"cacheStatus"`
	DownloadTime string `json:"downloadTime"`
	EdgeIP       string `json:"edgeIP"`
	FirstByte    string `json:"firstByte"`
	LastByte     string `json:"lastByte"`
	LastMileRTT  string `json:"lastMileRTT"`
}

// NetworkStruct is used for storing the JSON subfields
type NetworkStruct struct {
	Asnum       string `json:"asnum"`
	EdgeIP      string `json:"edgeIP"`
	Network     string `json:"network"`
	NetworkType string `json:"networkType"`
}

// ReqHdrStruct is used for storing the JSON subfields
type ReqHdrStruct struct {
	Cookie string `json:"cookie"`
}

// RespHdrStruct is used for storing the JSON subfields
type RespHdrStruct struct {
	Server  string `json:"server"`
	ContEnc string `json:"contEnc"`
}

// CreateObjects creates a list of AkamaiPayloads from a raw byte slice
func CreateObjects(jsonFile []byte) (_ []AkamaiPayload, err error) {
	var arrayObject []AkamaiPayload
	err = json.Unmarshal(jsonFile, &arrayObject)
	return arrayObject, err
}

// Handle parses the incoming request data
func Handle(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "OK! I got you Bro! -> ")
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Fatal(err)
	}
	obj, err := CreateObjects(body)
	if err != nil {
		log.Fatal(err)
	}
	c, err := client.NewHTTPClient(client.HTTPConfig{
		Addr: "http://influxDB:8086",
	})

	// Create a new point batch
	bp, err := client.NewBatchPoints(client.BatchPointsConfig{
		Database: "metrics",
		//Precision: "s",
	})
	if err != nil {
		log.Fatal(err)
	}

	for key, o := range obj {
		//payload[0].Geo["city"]
		fmt.Println("Server is:", o.RespHdr.Server, "KEY:", key)

		if err != nil {
			log.Fatal(err)
		}
		// Create a point and add to batch
		tags := map[string]string{
			"cp":          o.CP,
			"format":      o.Format,
			"city":        o.Geo.City,
			"country":     o.Geo.Country,
			"status":      o.Message.Status,
			"cacheStatus": o.NetPerf.CacheStatus,
			"proto":       o.Message.Proto,
			"protoVer":    o.Message.ProtoVer,
			"reqPort":     o.Message.ReqPort,
			"reqPath":     o.Message.ReqPath,
			"reqMethod":   o.Message.ReqMethod,
			"region":      o.Geo.Region,
			"UA": 			o.Message.UA,
			"reqHost":		o.Message.ReqHost,
			"respCT": 		o.Message.RespCT,
			"netPerfasnum": o.NetPerf.Asnum,
			"netPerfedgeIP":o.NetPerf.EdgeIP,
			"networkAsum":	o.Network.Asnum,
			"networkEdgeIP": o.Network.EdgeIP,
			"network":		o.Network.Network,
			"networkType":	o.Network.NetworkType,
			"cookie":		o.ReqHdr.Cookie,
			"server":		o.RespHdr.Server,
			"contEnc":		o.RespHdr.ContEnc,
			"type":			o.Type,
			"version":		o.Version,
		}

		lat, err := strconv.ParseFloat(o.Geo.Lat, 64)
		long, err := strconv.ParseFloat(o.Geo.Long, 64)
		downloadTime, err := strconv.ParseInt(o.NetPerf.DownloadTime, 10, 32)
		status, err := strconv.ParseInt(o.Message.Status, 10, 32)
		bytes, err := strconv.ParseInt(o.Message.Bytes, 10, 64)
		respLen, err := strconv.ParseInt(o.Message.RespLen, 10, 32)
		lastMileRTT, err := strconv.ParseFloat(o.NetPerf.LastMileRTT,64)
		start, err := strconv.ParseFloat(o.Start,64)

		if err != nil {
			log.Fatal(err)
		}
		fields := map[string]interface{}{
			"lat":          lat,
			"long":         long,
			"downloadtime": downloadTime,
			"status":       status,
			"bytes":		bytes,
			"respLen":		respLen,
			"lastMileRTT":	lastMileRTT,
			"start":		start,
		}

		pt, err := client.NewPoint("measurement", tags, fields, time.Now())
		if err != nil {
			log.Fatal(err)
		}
		bp.AddPoint(pt)
	}

	// Write the batch
	if err := c.Write(bp); err != nil {
		log.Fatal(err)
	}
}

func main() {

	http.HandleFunc("/", Handle)
	http.ListenAndServe(":9143", nil)
}
<|endoftext|>


<commit_before>package recipe

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"

	"github.com/fatih/color"
	"github.com/miclip/mybrewgo/hoputils"
	"github.com/miclip/mybrewgo/ingredients"
	"github.com/miclip/mybrewgo/ui"
	"github.com/miclip/mybrewgo/utils"

	yaml "gopkg.in/yaml.v2"
)

const (
	yieldToPoints           = 0.46
	gravityBase             = 0.00001
	preBoilGallonsFactorEst = 1.2
)

// Recipe ...
type Recipe struct {
	Name         string
	Version      int
	Batch        float64
	Style        string
	Efficiency   float64
	Method       string
	BoilTime     float64
	Hops         []ingredients.Hop
	Fermentables []ingredients.Fermentable
	Yeasts       []ingredients.Yeast
}

// Create adds a recipe via the cli
func Create(ui ui.UI) *Recipe {
	recipe := &Recipe{}
	recipe.Name = ui.AskForText("Recipe Name:")
	recipe.Version, _ = ui.AskForInt("Version Number:")
	recipe.Batch = ui.AskForFloat("Batch Size:")
	recipe.BoilTime = ui.AskForFloat("Boil Time:")
	recipe.Efficiency = ui.AskForFloat("Efficiency:")
	recipe.Method = ui.AskForText("Method:")
	recipe.Style = ui.AskForText("Style:")
	color.White("Ingredients...")
	addIngedients := true
	for addIngedients {
		i := ui.AskForText("Add Ingredient, Fermentable (f), Hop (h), Yeast (y), Save/Exit (s):")
		if i == "s" {
			break
		}
		if strings.ToLower(i) != "f" && strings.ToLower(i) != "h" && strings.ToLower(i) != "y" {
			color.Red("Invalid Ingredient, must be either h,f,y or s")
		}
	}
	return recipe
}

// OpenRecipe ...
func OpenRecipe(fileName string) (recipe *Recipe, err error) {
	filePath, _ := filepath.Abs(fileName)
	color.White("Reading recipe file %v", filePath)
	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read file with: %v", err)
	}
	if strings.Contains(strings.ToLower(fileName), ".yml") {
		return UnmarshalRecipeYML(data)
	}
	return UnmarshalRecipeJSON(data)
}

// UnmarshalRecipeYML ...
func UnmarshalRecipeYML(recipeYamlData []byte) (recipe *Recipe, err error) {
	err = yaml.Unmarshal(recipeYamlData, &recipe)
	if err != nil {
		return nil, fmt.Errorf("Failed to unmarshal yml recipe data with, %v", err)
	}
	return recipe, nil
}

// UnmarshalRecipeJSON ...
func UnmarshalRecipeJSON(recipeJSONData []byte) (recipe *Recipe, err error) {
	err = json.Unmarshal(recipeJSONData, &recipe)
	if err != nil {
		return nil, fmt.Errorf("Failed to unmarshal json recipe data with, %v", err)
	}
	return recipe, nil
}

// EstimatedPreBoilVolume estimates the preboil volume
func (r *Recipe) EstimatedPreBoilVolume() float64 {
	return r.Batch * preBoilGallonsFactorEst
}

// OriginalGravity calculates the original gravity
func (r *Recipe) OriginalGravity() float64 {
	og := 0.0
	if r.Efficiency == 0 || r.Batch == 0 || len(r.Fermentables) == 0 {
		return og
	}
	for i := range r.Fermentables {
		og = og +
			r.Fermentables[i].PointsByAmount()*(r.Efficiency*gravityBase)
	}
	return og/r.Batch + 1
}

// OriginalGravityPoints converts OG in Specific Gravity to Points
func (r *Recipe) OriginalGravityPoints() float64 {
	return (r.OriginalGravity() - 1) * 1000
}

// BoilSpecificGravity calculates the specific gravity post boil
func (r *Recipe) BoilSpecificGravity() float64 {
	return r.EstimatedPreBoilVolume()/r.Batch*(r.OriginalGravity()-1) + 1
}

// Color for recipe in SRM
func (r *Recipe) Color() float64 {
	color := 0.0
	if len(r.Fermentables) == 0 {
		return color
	}
	for i := range r.Fermentables {
		color = color + r.Fermentables[i].ColorMCU()
	}
	return ((color / r.Batch) * 0.2) + 8.4
}

// InternationalBitteringUnits calculates IBU for the recipe using a variation on Tinseth’s formula that
// incorporates a gravity/time adjustment instead of the bigness factor as documented by Randy Mosher
// in the "Brewer's Companion". No IBU formula is perfect so expect variations.
func (r *Recipe) InternationalBitteringUnits() float64 {
	hopUtils := hoputils.NewHopUtilizations()
	ibu := 0.0
	if len(r.Hops) == 0 {
		return 0.0
	}
	for i := range r.Hops {
		ibu = ibu + r.Hops[i].InternationalBitteringUnits(hopUtils, r.Batch, r.OriginalGravity())
	}
	return ibu
}

// EstimatedFinalGravity ...
func (r *Recipe) EstimatedFinalGravity() float64 {
	averageAttenuation := 0.0
	if len(r.Yeasts) == 0 {
		return 0.0
	}
	for i := range r.Yeasts {
		averageAttenuation = averageAttenuation + r.Yeasts[i].Attenuation/100
	}
	averageAttenuation = averageAttenuation / float64(len(r.Yeasts))

	return r.OriginalGravityPoints()*(1-averageAttenuation)/1000 + 1
}

// AlcoholByVolume ...
func (r *Recipe) AlcoholByVolume() float64 {
	if r.EstimatedFinalGravity() == 0 {
		return 0.0
	}
	return ((1.05 * (r.OriginalGravity() - r.EstimatedFinalGravity())) / r.EstimatedFinalGravity()) / 0.79 * 100
}

// AlcoholByWeight ...
func (r *Recipe) AlcoholByWeight() float64 {
	if r.EstimatedFinalGravity() == 0 {
		return 0.0
	}
	return (0.79 * r.AlcoholByVolume()) / r.EstimatedFinalGravity()
}

// Print writes recipe and ingredient details to stdout
func (r *Recipe) Print(ui ui.UI) {
	ui.PrintLinef("Recipe: %s Version: %d", r.Name, r.Version)
	ui.PrintLinef("Style: %s", r.Style)
	ui.PrintLinef("Batch Size: %v Boil Time: %v", r.Batch, r.BoilTime)
	ui.PrintLinef("OG: %v FG: %v IBU: %v ABV: %v SRM: %v", utils.Round(r.OriginalGravity(), .5, 3), utils.Round(r.EstimatedFinalGravity(), .5, 3),
		utils.Round(r.InternationalBitteringUnits(), .5, 1), utils.Round(r.AlcoholByVolume(), .5, 1), utils.Round(r.Color(), .5, 1))
	if len(r.Fermentables) > 0 {
		ui.PrintLinef("Fermentables: ")
	}
	for _, v := range r.Fermentables {
		v.Print(ui)
	}
	if len(r.Hops) > 0 {
		ui.PrintLinef("Hops: ")
	}
	for _, v := range r.Hops {
		v.Print(ui)
	}
	if len(r.Yeasts) > 0 {
		ui.PrintLinef("Yeasts: ")
	}
	for _, v := range r.Yeasts {
		v.Print(ui)
	}
}
<commit_msg>Remove stdout from recipe<commit_after>package recipe

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"

	"github.com/fatih/color"
	"github.com/miclip/mybrewgo/hoputils"
	"github.com/miclip/mybrewgo/ingredients"
	"github.com/miclip/mybrewgo/ui"
	"github.com/miclip/mybrewgo/utils"

	yaml "gopkg.in/yaml.v2"
)

const (
	yieldToPoints           = 0.46
	gravityBase             = 0.00001
	preBoilGallonsFactorEst = 1.2
)

// Recipe ...
type Recipe struct {
	Name         string
	Version      int
	Batch        float64
	Style        string
	Efficiency   float64
	Method       string
	BoilTime     float64
	Hops         []ingredients.Hop
	Fermentables []ingredients.Fermentable
	Yeasts       []ingredients.Yeast
}

// Create adds a recipe via the cli
func Create(ui ui.UI) *Recipe {
	recipe := &Recipe{}
	recipe.Name = ui.AskForText("Recipe Name:")
	recipe.Version, _ = ui.AskForInt("Version Number:")
	recipe.Batch = ui.AskForFloat("Batch Size:")
	recipe.BoilTime = ui.AskForFloat("Boil Time:")
	recipe.Efficiency = ui.AskForFloat("Efficiency:")
	recipe.Method = ui.AskForText("Method:")
	recipe.Style = ui.AskForText("Style:")
	color.White("Ingredients...")
	addIngedients := true
	for addIngedients {
		i := ui.AskForText("Add Ingredient, Fermentable (f), Hop (h), Yeast (y), Save/Exit (s):")
		if i == "s" {
			break
		}
		if strings.ToLower(i) != "f" && strings.ToLower(i) != "h" && strings.ToLower(i) != "y" {
			color.Red("Invalid Ingredient, must be either h,f,y or s")
		}
	}
	return recipe
}

// OpenRecipe ...
func OpenRecipe(fileName string) (recipe *Recipe, err error) {
	filePath, _ := filepath.Abs(fileName)
	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read file with: %v", err)
	}
	if strings.Contains(strings.ToLower(fileName), ".yml") {
		return UnmarshalRecipeYML(data)
	}
	return UnmarshalRecipeJSON(data)
}

// UnmarshalRecipeYML ...
func UnmarshalRecipeYML(recipeYamlData []byte) (recipe *Recipe, err error) {
	err = yaml.Unmarshal(recipeYamlData, &recipe)
	if err != nil {
		return nil, fmt.Errorf("Failed to unmarshal yml recipe data with, %v", err)
	}
	return recipe, nil
}

// UnmarshalRecipeJSON ...
func UnmarshalRecipeJSON(recipeJSONData []byte) (recipe *Recipe, err error) {
	err = json.Unmarshal(recipeJSONData, &recipe)
	if err != nil {
		return nil, fmt.Errorf("Failed to unmarshal json recipe data with, %v", err)
	}
	return recipe, nil
}

// EstimatedPreBoilVolume estimates the preboil volume
func (r *Recipe) EstimatedPreBoilVolume() float64 {
	return r.Batch * preBoilGallonsFactorEst
}

// OriginalGravity calculates the original gravity
func (r *Recipe) OriginalGravity() float64 {
	og := 0.0
	if r.Efficiency == 0 || r.Batch == 0 || len(r.Fermentables) == 0 {
		return og
	}
	for i := range r.Fermentables {
		og = og +
			r.Fermentables[i].PointsByAmount()*(r.Efficiency*gravityBase)
	}
	return og/r.Batch + 1
}

// OriginalGravityPoints converts OG in Specific Gravity to Points
func (r *Recipe) OriginalGravityPoints() float64 {
	return (r.OriginalGravity() - 1) * 1000
}

// BoilSpecificGravity calculates the specific gravity post boil
func (r *Recipe) BoilSpecificGravity() float64 {
	return r.EstimatedPreBoilVolume()/r.Batch*(r.OriginalGravity()-1) + 1
}

// Color for recipe in SRM
func (r *Recipe) Color() float64 {
	color := 0.0
	if len(r.Fermentables) == 0 {
		return color
	}
	for i := range r.Fermentables {
		color = color + r.Fermentables[i].ColorMCU()
	}
	return ((color / r.Batch) * 0.2) + 8.4
}

// InternationalBitteringUnits calculates IBU for the recipe using a variation on Tinseth’s formula that
// incorporates a gravity/time adjustment instead of the bigness factor as documented by Randy Mosher
// in the "Brewer's Companion". No IBU formula is perfect so expect variations.
func (r *Recipe) InternationalBitteringUnits() float64 {
	hopUtils := hoputils.NewHopUtilizations()
	ibu := 0.0
	if len(r.Hops) == 0 {
		return 0.0
	}
	for i := range r.Hops {
		ibu = ibu + r.Hops[i].InternationalBitteringUnits(hopUtils, r.Batch, r.OriginalGravity())
	}
	return ibu
}

// EstimatedFinalGravity ...
func (r *Recipe) EstimatedFinalGravity() float64 {
	averageAttenuation := 0.0
	if len(r.Yeasts) == 0 {
		return 0.0
	}
	for i := range r.Yeasts {
		averageAttenuation = averageAttenuation + r.Yeasts[i].Attenuation/100
	}
	averageAttenuation = averageAttenuation / float64(len(r.Yeasts))

	return r.OriginalGravityPoints()*(1-averageAttenuation)/1000 + 1
}

// AlcoholByVolume ...
func (r *Recipe) AlcoholByVolume() float64 {
	if r.EstimatedFinalGravity() == 0 {
		return 0.0
	}
	return ((1.05 * (r.OriginalGravity() - r.EstimatedFinalGravity())) / r.EstimatedFinalGravity()) / 0.79 * 100
}

// AlcoholByWeight ...
func (r *Recipe) AlcoholByWeight() float64 {
	if r.EstimatedFinalGravity() == 0 {
		return 0.0
	}
	return (0.79 * r.AlcoholByVolume()) / r.EstimatedFinalGravity()
}

// Print writes recipe and ingredient details to stdout
func (r *Recipe) Print(ui ui.UI) {
	ui.PrintLinef("Recipe: %s Version: %d", r.Name, r.Version)
	ui.PrintLinef("Style: %s", r.Style)
	ui.PrintLinef("Batch Size: %v Boil Time: %v", r.Batch, r.BoilTime)
	ui.PrintLinef("OG: %v FG: %v IBU: %v ABV: %v SRM: %v", utils.Round(r.OriginalGravity(), .5, 3), utils.Round(r.EstimatedFinalGravity(), .5, 3),
		utils.Round(r.InternationalBitteringUnits(), .5, 1), utils.Round(r.AlcoholByVolume(), .5, 1), utils.Round(r.Color(), .5, 1))
	if len(r.Fermentables) > 0 {
		ui.PrintLinef("Fermentables: ")
	}
	for _, v := range r.Fermentables {
		v.Print(ui)
	}
	if len(r.Hops) > 0 {
		ui.PrintLinef("Hops: ")
	}
	for _, v := range r.Hops {
		v.Print(ui)
	}
	if len(r.Yeasts) > 0 {
		ui.PrintLinef("Yeasts: ")
	}
	for _, v := range r.Yeasts {
		v.Print(ui)
	}
}
<|endoftext|>