<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"strings"

	"github.com/lair-framework/api-server/client"
	"github.com/lair-framework/go-lair"
)

const (
	version = "2.0.1"
	tool    = "raw"
	usage   = `
Imports a lair JSON file into a new lair project.

Usage:
  drone-raw [options] <id> <filename>
  export LAIR_ID=<id>; drone-raw [options] <filename>
Options:
  -v              show version and exit
  -h              show usage and exit
  -k              allow insecure SSL connections
  -force-ports    disable data protection in the API server for excessive ports
  -tags           a comma separated list of tags to add to every host that is imported
`
)

func main() {
	showVersion := flag.Bool("v", false, "")
	insecureSSL := flag.Bool("k", false, "")
	forcePorts := flag.Bool("force-ports", false, "")
	tags := flag.String("tags", "", "")
	flag.Usage = func() {
		fmt.Println(usage)
	}
	flag.Parse()
	if *showVersion {
		log.Println(version)
		os.Exit(0)
	}
	lairURL := os.Getenv("LAIR_API_SERVER")
	if lairURL == "" {
		log.Fatal("Fatal: Missing LAIR_API_SERVER environment variable")
	}
	lairPID := os.Getenv("LAIR_ID")
	var filename string
	switch len(flag.Args()) {
	case 2:
		lairPID = flag.Arg(0)
		filename = flag.Arg(1)
	case 1:
		filename = flag.Arg(0)
	default:
		log.Fatal("Fatal: Missing required argument")
	}
	u, err := url.Parse(lairURL)
	if err != nil {
		log.Fatalf("Fatal: Error parsing LAIR_API_SERVER URL. Error %s", err.Error())
	}
	if u.User == nil {
		log.Fatal("Fatal: Missing username and/or password")
	}
	user := u.User.Username()
	pass, _ := u.User.Password()
	if user == "" || pass == "" {
		log.Fatal("Fatal: Missing username and/or password")
	}
	c, err := client.New(&client.COptions{
		User:               user,
		Password:           pass,
		Host:               u.Host,
		Scheme:             u.Scheme,
		InsecureSkipVerify: *insecureSSL,
	})
	if err != nil {
		log.Fatalf("Fatal: Error setting up client: Error %s", err.Error())
	}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalf("Fatal: Could not open file. Error %s", err.Error())
	}
	hostTags := []string{}
	if *tags != "" {
		hostTags = strings.Split(*tags, ",")
	}
	l := lair.Project{}
	data = bytes.TrimPrefix(data, []byte("\xef\xbb\xbf")) // Or []byte{239, 187, 191}
	if err := json.Unmarshal(data, &l); err != nil {
		log.Fatalf("Fatal: Could not parse JSON. Error %s", err.Error())
	}
	l.ID = lairPID
	l.Tool = tool

	for i := range l.Hosts {
		l.Hosts[i].Tags = append(l.Hosts[i].Tags, hostTags...)
	}

	res, err := c.ImportProject(&client.DOptions{ForcePorts: *forcePorts}, &l)
	if err != nil {
		log.Fatalf("Fatal: Unable to import project. Error %s", err)
	}
	defer res.Body.Close()
	droneRes := &client.Response{}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.Fatalf("Fatal: Error %s", err.Error())
	}
	if err := json.Unmarshal(body, droneRes); err != nil {
		log.Fatalf("Fatal: Could not unmarshal JSON. Error %s", err.Error())
	}
	if droneRes.Status == "Error" {
		log.Fatalf("Fatal: Import failed. Error %s", droneRes.Message)
	}
	log.Println("Success: Operation completed successfully")
}
<commit_msg>Update main.go<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"strings"

	"github.com/lair-framework/api-server/client"
	"github.com/lair-framework/go-lair"
)

const (
	version = "2.0.1"
	tool    = "raw"
	usage   = `
Imports a lair JSON file into a new lair project.

Usage:
  drone-raw [options] <id> <filename>
  export LAIR_ID=<id>; drone-raw [options] <filename>
Options:
  -v              show version and exit
  -h              show usage and exit
  -k              allow insecure SSL connections
  -force-ports    disable data protection in the API server for excessive ports
  -tags           a comma separated list of tags to add to every host that is imported
`
)

func main() {
	showVersion := flag.Bool("v", false, "")
	insecureSSL := flag.Bool("k", false, "")
	forcePorts := flag.Bool("force-ports", false, "")
	tags := flag.String("tags", "", "")
	flag.Usage = func() {
		fmt.Println(usage)
	}
	flag.Parse()
	if *showVersion {
		log.Println(version)
		os.Exit(0)
	}
	lairURL := os.Getenv("LAIR_API_SERVER")
	if lairURL == "" {
		log.Fatal("Fatal: Missing LAIR_API_SERVER environment variable")
	}
	lairPID := os.Getenv("LAIR_ID")
        if lairPID == "" {
                log.Fatal("Fatal: Missing LAIR_ID")
        }
	var filename string
	switch len(flag.Args()) {
	case 2:
		lairPID = flag.Arg(0)
		filename = flag.Arg(1)
	case 1:
		filename = flag.Arg(0)
	default:
		log.Fatal("Fatal: Missing required argument")
	}
	u, err := url.Parse(lairURL)
	if err != nil {
		log.Fatalf("Fatal: Error parsing LAIR_API_SERVER URL. Error %s", err.Error())
	}
	if u.User == nil {
		log.Fatal("Fatal: Missing username and/or password")
	}
	user := u.User.Username()
	pass, _ := u.User.Password()
	if user == "" || pass == "" {
		log.Fatal("Fatal: Missing username and/or password")
	}
	c, err := client.New(&client.COptions{
		User:               user,
		Password:           pass,
		Host:               u.Host,
		Scheme:             u.Scheme,
		InsecureSkipVerify: *insecureSSL,
	})
	if err != nil {
		log.Fatalf("Fatal: Error setting up client: Error %s", err.Error())
	}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalf("Fatal: Could not open file. Error %s", err.Error())
	}
	hostTags := []string{}
	if *tags != "" {
		hostTags = strings.Split(*tags, ",")
	}
	l := lair.Project{}
	data = bytes.TrimPrefix(data, []byte("\xef\xbb\xbf")) // Or []byte{239, 187, 191}
	if err := json.Unmarshal(data, &l); err != nil {
		log.Fatalf("Fatal: Could not parse JSON. Error %s", err.Error())
	}
	l.ID = lairPID
	l.Tool = tool

	for i := range l.Hosts {
		l.Hosts[i].Tags = append(l.Hosts[i].Tags, hostTags...)
	}

	res, err := c.ImportProject(&client.DOptions{ForcePorts: *forcePorts}, &l)
	if err != nil {
		log.Fatalf("Fatal: Unable to import project. Error %s", err)
	}
	defer res.Body.Close()
	droneRes := &client.Response{}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.Fatalf("Fatal: Error %s", err.Error())
	}
	if err := json.Unmarshal(body, droneRes); err != nil {
		log.Fatalf("Fatal: Could not unmarshal JSON. Error %s", err.Error())
	}
	if droneRes.Status == "Error" {
		log.Fatalf("Fatal: Import failed. Error %s", droneRes.Message)
	}
	log.Println("Success: Operation completed successfully")
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

const (
	numPollers     = 1
	pollInterval   = 60 * time.Second
	statusInterval = 1 * time.Minute
	errTimeout     = 10 * time.Second
)

var (
	threshold = flag.Uint64("threshold", uint64(1.074*10E11), "Threshold to start serving 500's over HTTP")
	addr      = flag.String("addr", ":8080", "Listen address for HTTP")
	path      = flag.String("path", "/mnt/storage", "The path to query for disk usage")
)

// State represents the last-known state of a path
type State struct {
	path  string
	bytes uint64
}

// StateMonitor maintains a map that stores the disk usage for paths
// being
// polled, and prints the current state every updateInterval
// nanoseconds.
// It returns a chan State to which resource state should be sent.
// It also serves a HTTP Handler Func for threshold checking. It's
// probably doing too many things! :D
func StateMonitor(updateInterval time.Duration) chan<- State {
	updates := make(chan State)
	diskStatus := make(map[string]uint64)
	ticker := time.NewTicker(updateInterval)
	go func() {
		for {
			select {
			case <-ticker.C:
				logState(diskStatus)
			case s := <-updates:
				diskStatus[s.path] = s.bytes
			}
		}
	}()
	go func() {
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			bytes := diskStatus[*path]
			switch {
			case bytes == 0:
				http.Error(w, "Disk status not cached yet", http.StatusServiceUnavailable)
			case bytes > *threshold:
				err := fmt.Sprintf("ERROR: Bytes exceed threshold (%v/%v)", bytes, *threshold)
				http.Error(w, err, http.StatusInternalServerError)
			default:
				fmt.Fprintf(w, "OK: %v\n", bytes)
			}
		})
		if err := http.ListenAndServe(*addr, nil); err != nil {
			log.Fatal("ListenAndServe failed: ", err)
		}
	}()
	return updates
}

// logState prints a state map.
func logState(s map[string]uint64) {
	log.Println("Current state:")
	for k, v := range s {
		log.Printf(" %s %v", k, v)
	}
}

// Path represents a filesystem directory to be polled with du
type Path struct {
	path     string
	errCount int
}

// Poll executes du for a path
// and returns the disk usage in bytes or an error string
func (r *Path) Poll() (bytes uint64) {
	out, err := exec.Command("du", "-sbx", r.path).Output()
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	// Tidy up the line
	s := string(out)
	s = strings.TrimSpace(s)

	// Parse tabulation
	bytes_str := strings.Split(s, "\t")[0]

	// Parse uint64 from string
	bytes, err = strconv.ParseUint(bytes_str, 0, 64)
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	r.errCount = 0
	return bytes
}

// Sleep sleeps for an appropriate interval (dependent on error state)
// before sending the Path to done.
func (r *Path) Sleep(done chan<- *Path) {
	time.Sleep(pollInterval + errTimeout*time.Duration(r.errCount))
	done <- r
}

func Poller(in <-chan *Path, out chan<- *Path, status chan<- State) {
	for r := range in {
		bytes := r.Poll()
		status <- State{r.path, bytes}
		out <- r
	}
}

func main() {
	// parse CLI flags
	flag.Parse()

	// Create our input and output channels.
	pending, complete := make(chan *Path), make(chan *Path)

	// Launch the StateMonitor.
	status := StateMonitor(statusInterval)
	log.Println("State Monitor started")

	// Launch some Poller goroutines.
	for i := 0; i < numPollers; i++ {
		go Poller(pending, complete, status)
	}

	// Send the path flag to the pending queue
	go func() {
		pending <- &Path{path: *path}
	}()

	for r := range complete {
		go r.Sleep(pending)
	}

}
<commit_msg>use sync.RWMutex to read/write to the diskStatus map from 2 goroutines safely<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	numPollers     = 1
	pollInterval   = 60 * time.Second
	statusInterval = 1 * time.Minute
	errTimeout     = 10 * time.Second
)

var (
	threshold = flag.Uint64("threshold", uint64(1.074*10E11), "Threshold to start serving 500's over HTTP")
	addr      = flag.String("addr", ":8080", "Listen address for HTTP")
	path      = flag.String("path", "/mnt/storage", "The path to query for disk usage")
)

// State represents the last-known state of a path
type State struct {
	path  string
	bytes uint64
}

// StateMonitor maintains a map that stores the disk usage for paths
// being
// polled, and prints the current state every updateInterval
// nanoseconds.
// It returns a chan State to which resource state should be sent.
// It also serves a HTTP Handler Func for threshold checking. It's
// probably doing too many things! :D
func StateMonitor(updateInterval time.Duration) chan<- State {
	updates := make(chan State)
	diskStatus := make(map[string]uint64)
	diskStatusMutex := new(sync.RWMutex)
	ticker := time.NewTicker(updateInterval)
	go func() {
		for {
			select {
			case <-ticker.C:
				diskStatusMutex.Lock()
				logState(diskStatus)
				diskStatusMutex.Unlock()
			case s := <-updates:
				diskStatusMutex.Lock()
				diskStatus[s.path] = s.bytes
				diskStatusMutex.Unlock()
			}
		}
	}()
	go func() {
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			diskStatusMutex.Lock()
			bytes := diskStatus[*path]
			diskStatusMutex.Unlock()
			switch {
			case bytes == 0:
				http.Error(w, "Disk status not cached yet", http.StatusServiceUnavailable)
			case bytes > *threshold:
				err := fmt.Sprintf("ERROR: Bytes exceed threshold (%v/%v)", bytes, *threshold)
				http.Error(w, err, http.StatusInternalServerError)
			default:
				fmt.Fprintf(w, "OK: %v\n", bytes)
			}
		})
		if err := http.ListenAndServe(*addr, nil); err != nil {
			log.Fatal("ListenAndServe failed: ", err)
		}
	}()
	return updates
}

// logState prints a state map.
func logState(s map[string]uint64) {
	log.Println("Current state:")
	for k, v := range s {
		log.Printf(" %s %v", k, v)
	}
}

// Path represents a filesystem directory to be polled with du
type Path struct {
	path     string
	errCount int
}

// Poll executes du for a path
// and returns the disk usage in bytes or an error string
func (r *Path) Poll() (bytes uint64) {
	out, err := exec.Command("du", "-sbx", r.path).Output()
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	// Tidy up the line
	s := string(out)
	s = strings.TrimSpace(s)

	// Parse tabulation
	bytes_str := strings.Split(s, "\t")[0]

	// Parse uint64 from string
	bytes, err = strconv.ParseUint(bytes_str, 0, 64)
	if err != nil {
		log.Fatal(err)
		r.errCount++
	}
	r.errCount = 0
	return bytes
}

// Sleep sleeps for an appropriate interval (dependent on error state)
// before sending the Path to done.
func (r *Path) Sleep(done chan<- *Path) {
	time.Sleep(pollInterval + errTimeout*time.Duration(r.errCount))
	done <- r
}

func Poller(in <-chan *Path, out chan<- *Path, status chan<- State) {
	for r := range in {
		bytes := r.Poll()
		status <- State{r.path, bytes}
		out <- r
	}
}

func main() {
	// parse CLI flags
	flag.Parse()

	// Create our input and output channels.
	pending := make(chan *Path)
	complete := make(chan *Path)

	// Launch the StateMonitor.
	status := StateMonitor(statusInterval)
	log.Println("State Monitor started")

	// Launch some Poller goroutines.
	for i := 0; i < numPollers; i++ {
		go Poller(pending, complete, status)
	}

	// Send the path flag to the pending queue
	go func() {
		pending <- &Path{path: *path}
	}()

	for r := range complete {
		go r.Sleep(pending)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"flag"
)

var sr *string
var lim *int

const user = "dan"

func init() {
	sr = flag.String("subreddit", "ALL", "Name of the subreddit to browse. Defaults to ALL.")
	lim = flag.Int("limit", 10, "Total articles to fetch. Default is 10. 5 is the minimum number to fetch.")

	flag.Parse()
}

func main() {
	// If less than 5, set limit to 5.
	if *lim < 5 {
		lim = 5
	}

	// Start the process.
	fmt.Println("Starting process ...")
	Start(sr, *lim)
}
<commit_msg>bugs in main.go<commit_after>package main

import (
	"flag"
	"fmt"
)

var sr *string
var lim *int

const user = "dan"

func init() {
	sr = flag.String("subreddit", "ALL", "Name of the subreddit to browse. Defaults to ALL.")
	lim = flag.Int("limit", 10, "Total articles to fetch. Default is 10. 5 is the minimum number to fetch.")

	flag.Parse()
}

func main() {
	// If less than 5, set limit to 5.
	if *lim < 5 {
		*lim = 5
	}

	// Start the process.
	fmt.Println("Starting process ...")
	Start()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"database/sql"
	"encoding/csv"

	_ "github.com/mattn/go-sqlite3"
	"github.com/xwb1989/sqlparser"
)

func rowimport(db *sql.DB, table string, columns string, place []string, list []interface{}) {
	sqlstr := "INSERT INTO " + table + " (" + columns + ") VALUES (" + strings.Join(place, ",") + ");"
	_, err := db.Exec(sqlstr, list...)
	if err != nil {
		log.Println(sqlstr, err)
	}
}

func csvImport(db *sql.DB, reader *csv.Reader, table string, header []string) {
	columns := make([]string, len(header))
	place := make([]string, len(header))
	list := make([]interface{}, len(header))

	for i := range header {
		columns[i] = "c" + strconv.Itoa(i+1)
		place[i] = "?"
		list[i] = header[i]
	}
	columName := strings.Join(columns, ",")
	rowimport(db, table, columName, place, list)

	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		} else {
			if err != nil {
				log.Fatal("ERROR: ", err)
			}
		}
		for i := range header {
			list[i] = record[i]
		}
		rowimport(db, table, columName, place, list)
	}
}

func csvOpen(filename string) *csv.Reader {
	var file *os.File
	var err error
	if filename == "-" {
		file = os.Stdin
	} else {
		if filename[0] == '`' {
			filename = strings.Replace(filename, "`", "", 2)
		}
		file, err = os.Open(filename)
		if err != nil {
			log.Fatal("ERROR: ", err)
		}
	}
	reader := csv.NewReader(file)
	return reader
}

func csvRead(reader *csv.Reader) (header []string) {
	var err error
	header, err = reader.Read()
	if err != nil {
		log.Fatal("ERROR: ", err)
	}
	return header
}

func escapetable(oldname string) (newname string) {
	if oldname[0] != '`' {
		newname = "`" + oldname + "`"
	} else {
		newname = oldname
	}
	return newname
}

func rewrite(sqlstr string, oldname string, newname string) (rewrite string) {
	rewrite = strings.Replace(sqlstr, oldname, newname, -1)
	return rewrite
}

func rewriteTable(tree sqlparser.SQLNode, tablename string) string {
	rewriter := func(origin []byte) []byte {
		s := string(origin)
		if s == tablename {
			s = "_"
		}
		return []byte(s)
	}
	sqlparser.Rewrite(tree, rewriter)
	return sqlparser.String(tree)
}

func dbConnect() *sql.DB {
	db, err := sql.Open("sqlite3", ":memory:")
	if err != nil {
		log.Fatal(err)
	}
	return db
}

func dbDisconnect(db *sql.DB) {
	err := db.Close()
	if err != nil {
		log.Fatal(err)
	}
}

func dbCreate(db *sql.DB, table string, header []string) {
	columns := make([]string, len(header))
	for i := 0; i < len(header); i++ {
		columns[i] = "c" + strconv.Itoa(i+1) + " text"
	}
	c := strings.Join(columns, ",")
	sqlstr := "CREATE TABLE " + table + " ( " + c + " )"
	log.Println(sqlstr)
	_, err := db.Exec(sqlstr)
	if err != nil {
		log.Fatal(err)
	}
}

func dbSelect(db *sql.DB, writer *csv.Writer, sqlstr string) {
	sqlstr = strings.TrimSpace(sqlstr)
	if sqlstr == "" {
		log.Fatal("ERROR: no SQL statement")
	}
	log.Println(sqlstr)
	rows, err := db.Query(sqlstr)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()
	columns, err := rows.Columns()
	if err != nil {
		log.Fatal(err)
	}
	values := make([]sql.RawBytes, len(columns))
	results := make([]string, len(columns))
	scanArgs := make([]interface{}, len(values))
	for i := range values {
		scanArgs[i] = &values[i]
	}
	for rows.Next() {
		err = rows.Scan(scanArgs...)
		if err != nil {
			log.Fatal(err)
		}
		for i, col := range values {
			results[i] = string(col)
		}
		writer.Write(results)
		writer.Flush()
	}
}

func sqlparse(sqlstr string) []string {
	word := strings.Fields(sqlstr)
	tablenames := make([]string, 0, 1)
	for i := 0; i < len(word); i++ {
		if element := strings.ToUpper(word[i]); element == "FROM" || element == "JOIN" {
			tablenames = append(tablenames, word[i+1])
		}
	}
	return tablenames
}

func getSeparator(sepString string) (sepRune rune) {
	sepString = `'` + sepString + `'`
	sepRunes, err := strconv.Unquote(sepString)
	if err != nil {
		log.Fatal(sepString, ": ", err)
	}
	sepRune = ([]rune(sepRunes))[0]

	return sepRune
}

func main() {
	var (
		inSep  string
		outSep string
	)
	flag.StringVar(&inSep, "input-delimiter", ",", "Field delimiter for input.")
	flag.StringVar(&inSep, "d", ",", "Field delimiter for input.")
	flag.StringVar(&outSep, "output-delimiter", ",", "Field delimiter for output.")
	flag.StringVar(&outSep, "D", ",", "Field delimiter for output.")
	flag.Parse()
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(2)
	}
	sqlstr := flag.Args()[0]
	writer := csv.NewWriter(os.Stdout)
	writer.Comma = getSeparator(outSep)
	readerComma := getSeparator(inSep)

	db := dbConnect()
	defer dbDisconnect(db)

	tablenames := sqlparse(sqlstr)
	for _, tablename := range tablenames {
		rtable := escapetable(tablename)
		sqlstr = rewrite(sqlstr, tablename, rtable)
		reader := csvOpen(tablename)
		reader.Comma = readerComma
		reader.FieldsPerRecord = -1
		header := csvRead(reader)
		dbCreate(db, rtable, header)
		csvImport(db, reader, rtable, header)
	}
	dbSelect(db, writer, sqlstr)
}
<commit_msg>Added support postgresql database driver.<commit_after>package main

import (
	"flag"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"database/sql"
	"encoding/csv"

	_ "github.com/lib/pq"
	_ "github.com/mattn/go-sqlite3"
	"github.com/xwb1989/sqlparser"
)

var (
	dbdriver string
	dbdsn    string
)

func rowimport(db *sql.DB, table string, columns string, place []string, list []interface{}) {
	sqlstr := "INSERT INTO " + table + " (" + columns + ") VALUES (" + strings.Join(place, ",") + ");"
	_, err := db.Exec(sqlstr, list...)
	if err != nil {
		log.Println(sqlstr, err)
	}
}

func csvImport(db *sql.DB, reader *csv.Reader, table string, header []string) {
	columns := make([]string, len(header))
	place := make([]string, len(header))
	list := make([]interface{}, len(header))

	for i := range header {
		columns[i] = "c" + strconv.Itoa(i+1)
		if dbdriver == "postgres" {
			place[i] = "$" + strconv.Itoa(i+1)
		} else {
			place[i] = "?"
		}
		list[i] = header[i]
	}
	columName := strings.Join(columns, ",")
	rowimport(db, table, columName, place, list)

	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		} else {
			if err != nil {
				log.Fatal("ERROR: ", err)
			}
		}
		for i := range header {
			list[i] = record[i]
		}
		rowimport(db, table, columName, place, list)
	}
}

func csvOpen(filename string) (*csv.Reader, error) {
	var file *os.File
	var err error
	if filename == "-" {
		file = os.Stdin
	} else {
		if filename[0] == '`' {
			filename = strings.Replace(filename, "`", "", 2)
		}
		file, err = os.Open(filename)
		if err != nil {
			// log.Fatal("ERROR: ", err)
			return nil, err
		}
	}
	reader := csv.NewReader(file)
	return reader, err
}

func csvRead(reader *csv.Reader) (header []string) {
	var err error
	header, err = reader.Read()
	if err != nil {
		log.Fatal("ERROR: ", err)
	}
	return header
}

func escapetable(oldname string) (newname string) {
	if dbdriver == "postgres" {
		if oldname[0] != '"' {
			newname = "\"" + oldname + "\""
		} else {
			newname = oldname
		}
	} else {
		if oldname[0] != '`' {
			newname = "`" + oldname + "`"
		} else {
			newname = oldname
		}
	}
	return newname
}

func rewrite(sqlstr string, oldname string, newname string) (rewrite string) {
	rewrite = strings.Replace(sqlstr, oldname, newname, -1)
	return rewrite
}

func rewriteTable(tree sqlparser.SQLNode, tablename string) string {
	rewriter := func(origin []byte) []byte {
		s := string(origin)
		if s == tablename {
			s = "_"
		}
		return []byte(s)
	}
	sqlparser.Rewrite(tree, rewriter)
	return sqlparser.String(tree)
}

func dbConnect(driver, dsn string) *sql.DB {
	db, err := sql.Open(driver, dsn)
	if err != nil {
		log.Fatal(err)
	}
	return db
}

func dbDisconnect(db *sql.DB) {
	err := db.Close()
	if err != nil {
		log.Fatal(err)
	}
}

func dbCreate(db *sql.DB, table string, header []string) {
	columns := make([]string, len(header))
	for i := 0; i < len(header); i++ {
		columns[i] = "c" + strconv.Itoa(i+1) + " text"
	}
	c := strings.Join(columns, ",")
	sqlstr := "CREATE TEMP TABLE " + table + " ( " + c + " )"
	log.Println(sqlstr)
	_, err := db.Exec(sqlstr)
	if err != nil {
		log.Fatal(err)
	}
}

func dbSelect(db *sql.DB, writer *csv.Writer, sqlstr string) {
	sqlstr = strings.TrimSpace(sqlstr)
	if sqlstr == "" {
		log.Fatal("ERROR: no SQL statement")
	}
	log.Println(sqlstr)
	rows, err := db.Query(sqlstr)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()
	columns, err := rows.Columns()
	if err != nil {
		log.Fatal(err)
	}
	values := make([]sql.RawBytes, len(columns))
	results := make([]string, len(columns))
	scanArgs := make([]interface{}, len(values))
	for i := range values {
		scanArgs[i] = &values[i]
	}
	for rows.Next() {
		err = rows.Scan(scanArgs...)
		if err != nil {
			log.Fatal(err)
		}
		for i, col := range values {
			results[i] = string(col)
		}
		writer.Write(results)
		writer.Flush()
	}
}

func sqlparse(sqlstr string) []string {
	word := strings.Fields(sqlstr)
	tablenames := make([]string, 0, 1)
	for i := 0; i < len(word); i++ {
		if element := strings.ToUpper(word[i]); element == "FROM" || element == "JOIN" {
			tablenames = append(tablenames, word[i+1])
		}
	}
	return tablenames
}

func getSeparator(sepString string) (sepRune rune) {
	sepString = `'` + sepString + `'`
	sepRunes, err := strconv.Unquote(sepString)
	if err != nil {
		log.Fatal(sepString, ": ", err)
	}
	sepRune = ([]rune(sepRunes))[0]

	return sepRune
}

func main() {
	var (
		inSep  string
		outSep string
	)
	flag.StringVar(&dbdriver, "dbdriver", "sqlite3", "database driver.")
	flag.StringVar(&dbdsn, "dbdsn", ":memory:", "database connection option.")
	flag.StringVar(&inSep, "input-delimiter", ",", "Field delimiter for input.")
	flag.StringVar(&inSep, "d", ",", "Field delimiter for input.")
	flag.StringVar(&outSep, "output-delimiter", ",", "Field delimiter for output.")
	flag.StringVar(&outSep, "D", ",", "Field delimiter for output.")
	flag.Parse()
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(2)
	}
	sqlstr := flag.Args()[0]
	writer := csv.NewWriter(os.Stdout)
	writer.Comma = getSeparator(outSep)
	readerComma := getSeparator(inSep)

	db := dbConnect(dbdriver, dbdsn)
	defer dbDisconnect(db)

	tablenames := sqlparse(sqlstr)
	for _, tablename := range tablenames {
		rtable := escapetable(tablename)
		sqlstr = rewrite(sqlstr, tablename, rtable)
		reader, err := csvOpen(tablename)
		if err != nil {
			continue
		}
		reader.Comma = readerComma
		reader.FieldsPerRecord = -1
		header := csvRead(reader)
		dbCreate(db, rtable, header)
		csvImport(db, reader, rtable, header)
	}
	dbSelect(db, writer, sqlstr)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"

	"github.com/maxmind/geoip2-csv-converter/convert"
)

func main() {
	input := flag.String("block-file", "", "The path to the block CSV file to use as input (REQUIRED)")
	output := flag.String("output-file", "", "The path to the output CSV (REQUIRED)")
	ipRange := flag.Bool("include-range", false, "Include the IP range of the network in string format")
	intRange := flag.Bool("include-integer-range", false, "Include the IP range of the network in integer format")
	cidr := flag.Bool("include-cidr", false, "Include the network in CIDR format")

	flag.Parse()

	var errors []string

	if *input == "" {
		errors = append(errors, "-block-file is required")
	}

	if *output == "" {
		errors = append(errors, "-output-file is required")
	}

	if !*ipRange && !*intRange && !*cidr {
		errors = append(errors, "At least one of -include-* param is required")
	}

	if len(errors) != 0 {
		printHelp(errors)
		os.Exit(1)
	}

	err := convert.ConvertFile(*input, *output, *cidr, *ipRange, *intRange)
	if err != nil {
		fmt.Fprintf(flag.CommandLine.Output(), "Error: %v\n", err)
	}
}

func printHelp(errors []string) {
	for _, message := range errors {
		fmt.Fprintln(flag.CommandLine.Output(), message)
	}
	flag.PrintDefaults()
}
<commit_msg>Show unknown arguments in error message<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/maxmind/geoip2-csv-converter/convert"
)

func main() {
	input := flag.String("block-file", "", "The path to the block CSV file to use as input (REQUIRED)")
	output := flag.String("output-file", "", "The path to the output CSV (REQUIRED)")
	ipRange := flag.Bool("include-range", false, "Include the IP range of the network in string format")
	intRange := flag.Bool("include-integer-range", false, "Include the IP range of the network in integer format")
	cidr := flag.Bool("include-cidr", false, "Include the network in CIDR format")

	flag.Parse()

	var errors []string

	if *input == "" {
		errors = append(errors, "-block-file is required")
	}

	if *output == "" {
		errors = append(errors, "-output-file is required")
	}

	if !*ipRange && !*intRange && !*cidr {
		errors = append(errors, "At least one of -include-* param is required")
	}

	args := flag.Args()
	if len(args) > 0 {
		errors = append(errors, "unknown argument(s): "+strings.Join(args, ", "))
	}

	if len(errors) != 0 {
		printHelp(errors)
		os.Exit(1)
	}

	err := convert.ConvertFile(*input, *output, *cidr, *ipRange, *intRange)
	if err != nil {
		fmt.Fprintf(flag.CommandLine.Output(), "Error: %v\n", err)
	}
}

func printHelp(errors []string) {
	for _, message := range errors {
		fmt.Fprintln(flag.CommandLine.Output(), message)
	}
	flag.PrintDefaults()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	//"os"
)

func main() {
	// Change path to location of ImgurGo config
	configFile := "config/default.conf.json"
    //if _, err := os.Stat(configFile); err != nil {
	//	fmt.Printf("Configuration file %s does not exist", configFile)
	//	os.Exit(-1)
	//}
    fmt.Printf("got here")
	config := NewConfiguration(configFile)
    fmt.Printf("got here2")
	server := CreateServer(config)
    fmt.Printf("got here3")
	server.initServer()
    fmt.Printf("got here4")
}
<commit_msg>use println<commit_after>package main

import (
	"fmt"
	//"os"
)

func main() {
	// Change path to location of ImgurGo config
	configFile := "config/default.conf.json"
    //if _, err := os.Stat(configFile); err != nil {
	//	fmt.Printf("Configuration file %s does not exist", configFile)
	//	os.Exit(-1)
	//}
    fmt.Println("got here")
	config := NewConfiguration(configFile)
    fmt.Println("got here2")
	server := CreateServer(config)
    fmt.Println("got here3")
	server.initServer()
    fmt.Println("got here4")
}
<|endoftext|>


<commit_before>package main

import (
  "log"
  "fmt"
  "io/ioutil"
  "encoding/json"
)

type ServerConfig struct {
  Port int `json:"port"`
  UrlPrefix string `json:"urlPrefix"`
}

type RecordConfig struct {
  SaveTo string `json:"saveTo"`
  NotifyConsole bool `json:"notifyConsole"`
}

type Request struct {
  Method string `json:"method"`
  Url string `json:"url"`
}

type Response struct {
  Status string `json:"status"`
  Body string `json:"body"`
}

type Route struct {
  Request Request
  Response Response
}

type Routes []Route

type Config struct {
  Server ServerConfig
  Record RecordConfig
  Routes Routes
}

func main() {
    configFileContent, err := ioutil.ReadFile("gomock.json")

    if err != nil {
      log.Fatal(err)
    }

    var config Config
    json.Unmarshal(configFileContent, config)

    fmt.Print(config)
}
<commit_msg>apply gofmt<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
)

type ServerConfig struct {
	Port      int    `json:"port"`
	UrlPrefix string `json:"urlPrefix"`
}

type RecordConfig struct {
	SaveTo        string `json:"saveTo"`
	NotifyConsole bool   `json:"notifyConsole"`
}

type Request struct {
	Method string `json:"method"`
	Url    string `json:"url"`
}

type Response struct {
	Status string `json:"status"`
	Body   string `json:"body"`
}

type Route struct {
	Request  Request
	Response Response
}

type Routes []Route

type Config struct {
	Server ServerConfig
	Record RecordConfig
	Routes Routes
}

func main() {
	configFileContent, err := ioutil.ReadFile("gomock.json")

	if err != nil {
		log.Fatal(err)
	}

	var config Config
	json.Unmarshal(configFileContent, config)

	fmt.Print(config)
}
<|endoftext|>


<commit_before>// gcvis is a tool to assist you visualising the operation of
// the go runtime garbage collector.
//
// usage:
//
//     gcvis program [arguments]...
package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strings"

	"github.com/pkg/browser"

	"golang.org/x/crypto/ssh/terminal"
)

var iface = flag.String("i", "127.0.0.1", "specify interface to use. defaults to 127.0.0.1.")
var port = flag.String("p", "0", "specify port to use.")

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s: command <args>...\n", os.Args[0])
		flag.PrintDefaults()
	}

	var pipeRead io.ReadCloser
	var subcommand *SubCommand

	flag.Parse()
	if len(flag.Args()) < 1 {
		if terminal.IsTerminal(int(os.Stdin.Fd())) {
			flag.Usage()
			return
		} else {
			pipeRead = os.Stdin
		}
	} else {
		subcommand := NewSubCommand(flag.Args())
		pipeRead = subcommand.PipeRead
		go subcommand.Run()
	}

	parser := NewParser(pipeRead)
	gcvisGraph := NewGraph(strings.Join(flag.Args(), " "), GCVIS_TMPL)
	server := NewHttpServer(*iface, *port, &gcvisGraph)

	go parser.Run()
	go server.Start()

	url := server.Url()
	log.Printf("opening browser window, if this fails, navigate to %s", url)
	browser.OpenURL(url)

	for {
		select {
		case gcTrace := <-parser.GcChan:
			gcvisGraph.AddGCTraceGraphPoint(gcTrace)
		case scvgTrace := <-parser.ScvgChan:
			gcvisGraph.AddScavengerGraphPoint(scvgTrace)
		case output := <-parser.NoMatchChan:
			fmt.Fprintln(os.Stderr, output)
		case <-parser.done:
			if parser.Err != nil {
				fmt.Fprintf(os.Stderr, parser.Err.Error())
				os.Exit(1)
			}

			break
		}
	}

	if subcommand != nil && subcommand.Err() != nil {
		fmt.Fprintf(os.Stderr, subcommand.Err().Error())
		os.Exit(1)
	}

	os.Exit(0)
}
<commit_msg>fixup! Use stdin if there is no argument given to gcvis<commit_after>// gcvis is a tool to assist you visualising the operation of
// the go runtime garbage collector.
//
// usage:
//
//     gcvis program [arguments]...
package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strings"

	"github.com/pkg/browser"

	"golang.org/x/crypto/ssh/terminal"
)

var iface = flag.String("i", "127.0.0.1", "specify interface to use. defaults to 127.0.0.1.")
var port = flag.String("p", "0", "specify port to use.")

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s: command <args>...\n", os.Args[0])
		flag.PrintDefaults()
	}

	var pipeRead io.ReadCloser
	var subcommand *SubCommand

	flag.Parse()
	if len(flag.Args()) < 1 {
		if terminal.IsTerminal(int(os.Stdin.Fd())) {
			flag.Usage()
			return
		} else {
			pipeRead = os.Stdin
		}
	} else {
		subcommand = NewSubCommand(flag.Args())
		pipeRead = subcommand.PipeRead
		go subcommand.Run()
	}

	parser := NewParser(pipeRead)
	gcvisGraph := NewGraph(strings.Join(flag.Args(), " "), GCVIS_TMPL)
	server := NewHttpServer(*iface, *port, &gcvisGraph)

	go parser.Run()
	go server.Start()

	url := server.Url()
	log.Printf("opening browser window, if this fails, navigate to %s", url)
	browser.OpenURL(url)

	for {
		select {
		case gcTrace := <-parser.GcChan:
			gcvisGraph.AddGCTraceGraphPoint(gcTrace)
		case scvgTrace := <-parser.ScvgChan:
			gcvisGraph.AddScavengerGraphPoint(scvgTrace)
		case output := <-parser.NoMatchChan:
			fmt.Fprintln(os.Stderr, output)
		case <-parser.done:
			if parser.Err != nil {
				fmt.Fprintf(os.Stderr, parser.Err.Error())
				os.Exit(1)
			}

			break
		}
	}

	if subcommand != nil && subcommand.Err() != nil {
		fmt.Fprintf(os.Stderr, subcommand.Err().Error())
		os.Exit(1)
	}

	os.Exit(0)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"

	"github.com/Smile-SA/argoos/apiutils"
)

func sig() {
	c := make(chan os.Signal, 0)
	signal.Notify(c, os.Interrupt)

	// Block until a signal is received.
	s := <-c
	apiutils.StopRollout()
	log.Println("Got signal", s)
	os.Exit(0)
}

// Action is sent each time the registry sends an event.
func Action(w http.ResponseWriter, r *http.Request) {
	log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
	c, _ := ioutil.ReadAll(r.Body)
	events := apiutils.GetEvents(c)
	for _, e := range events.Events {
		apiutils.ImpactedDeployments(e)
	}
}

// Health return always "ok" with 200 OK. Usefull to check liveness.
func Health(w http.ResponseWriter, r *http.Request) {
	log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
	w.Write([]byte("ok"))
}

func main() {
	host := ":3000"

	if v := os.Getenv("KUBE_MASTER_URL"); len(v) > 0 {
		apiutils.KubeMasterURL = v
	}

	if v := os.Getenv("SKIP_SSL_VERIFICATION"); strings.ToUpper(v) == "TRUE" {
		apiutils.SkipSSLVerification = true
		// Certificates
		if v := os.Getenv("CA_FILE"); len(v) > 0 {
			apiutils.CAFile = v
		}
		if v := os.Getenv("CERT_FILE"); len(v) > 0 {
			apiutils.CertFile = v
		}
		if v := os.Getenv("KEY_FILE"); len(v) > 0 {
			apiutils.KeyFile = v
		}
	}

	if v := os.Getenv("LISTEN"); len(v) > 0 {
		host = v
	}

	flag.StringVar(&apiutils.KubeMasterURL,
		"master",
		apiutils.KubeMasterURL,
		"Kube master scheme://host:port")
	flag.BoolVar(&apiutils.SkipSSLVerification,
		"skip-ssl-verification",
		apiutils.SkipSSLVerification,
		"Skip SSL verification for kubernetes api")
	flag.StringVar(&host,
		"listen",
		host,
		"Listen interface, could be host:port, or :port")

	// certs
	flag.StringVar(&apiutils.CAFile,
		"ca-file",
		apiutils.CAFile,
		"Certificate Authority certificate file path")
	flag.StringVar(&apiutils.CertFile,
		"cert-file",
		apiutils.CertFile,
		"Client certificate file path")
	flag.StringVar(&apiutils.KeyFile,
		"key-file",
		apiutils.KeyFile,
		"Client private certificate file path")

	flag.Parse()

	go sig()
	apiutils.StartRollout()

	log.Println("Starting")
	http.HandleFunc("/healthz", Health)
	http.HandleFunc("/event", Action)
	log.Fatal(http.ListenAndServe(host, nil))

}
<commit_msg>Add SSL capabilities<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"

	"github.com/Smile-SA/argoos/apiutils"
)

// signal handling, if server should stop, cleanup goroutines.
func sig() {
	c := make(chan os.Signal, 0)
	signal.Notify(c, os.Interrupt)

	// Block until a signal is received.
	s := <-c
	apiutils.StopRollout()
	log.Println("Got signal", s)
	os.Exit(0)
}

// Action is sent each time the registry sends an event.
func Action(w http.ResponseWriter, r *http.Request) {
	log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
	c, _ := ioutil.ReadAll(r.Body)
	events := apiutils.GetEvents(c)
	for _, e := range events.Events {
		apiutils.ImpactedDeployments(e)
	}
}

// Health return always "ok" with 200 OK. Usefull to check liveness.
func Health(w http.ResponseWriter, r *http.Request) {
	log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
	w.Write([]byte("ok"))
}

func main() {
	host := ":3000"
	var servercert, serverkey string

	if v := os.Getenv("KUBE_MASTER_URL"); len(v) > 0 {
		apiutils.KubeMasterURL = v
	}

	if v := os.Getenv("SKIP_SSL_VERIFICATION"); strings.ToUpper(v) == "TRUE" {
		apiutils.SkipSSLVerification = true
		// Certificates
		if v := os.Getenv("CA_FILE"); len(v) > 0 {
			apiutils.CAFile = v
		}
		if v := os.Getenv("CERT_FILE"); len(v) > 0 {
			apiutils.CertFile = v
		}
		if v := os.Getenv("KEY_FILE"); len(v) > 0 {
			apiutils.KeyFile = v
		}
	}

	if v := os.Getenv("SERVER_CERT"); len(v) > 0 {
		servercert = v
	}

	if v := os.Getenv("SERVER_KEY"); len(v) > 0 {
		serverkey = v
	}

	if v := os.Getenv("LISTEN"); len(v) > 0 {
		host = v
	}

	flag.StringVar(&apiutils.KubeMasterURL,
		"master",
		apiutils.KubeMasterURL,
		"Kube master scheme://host:port")
	flag.BoolVar(&apiutils.SkipSSLVerification,
		"skip-ssl-verification",
		apiutils.SkipSSLVerification,
		"Skip SSL verification for kubernetes api")
	flag.StringVar(&host,
		"listen",
		host,
		"Listen interface, could be host:port, or :port")

	// certs
	flag.StringVar(&apiutils.CAFile,
		"ca-file",
		apiutils.CAFile,
		"Certificate Authority certificate file path (only if using https to contact kubernetes api)")
	flag.StringVar(&apiutils.CertFile,
		"cert-file",
		apiutils.CertFile,
		"Client certificate file path (client authentication only)")
	flag.StringVar(&apiutils.KeyFile,
		"key-file",
		apiutils.KeyFile,
		"Client private key file path (client authentication only)")

	// argoos can serve https
	flag.StringVar(&servercert,
		"server-cert",
		servercert,
		"Server certificate to serve SSL")

	flag.StringVar(&serverkey,
		"server-key",
		serverkey,
		"Server key to server SSL")

	flag.Parse()

	go sig()
	apiutils.StartRollout()

	log.Println("Starting")

	http.HandleFunc("/healthz", Health)
	http.HandleFunc("/event", Action)

	if len(serverkey) > 0 && len(servercert) > 0 {
		log.Fatal(http.ListenAndServeTLS(host, servercert, serverkey, nil))
	} else {
		log.Fatal(http.ListenAndServe(host, nil))
	}

}
<|endoftext|>


<commit_before>package main

// FIXME write some stats every now and again - packets per second etc

// FIXME detect overflows?

// FIXME write dt in the file?

// FIXME don't dump stats file if empty?

// FIXME have a max entries parameter
// - dump the stats if get beyond that many entries?
// - or just refuse to add any more?
// - or add to 0.0.0.0 entry?

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"log/syslog"
	"net"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"sync"
	"syscall"
	"time"
)

const (
	// Time format to output in CSV files
	// http://stackoverflow.com/questions/804118/best-timestamp-format-for-csv-excel
	CsvTimeFormat = "2006-01-02 15:04:05"
	// Format to use when making filenames
	FileTimeFormat = "2006-01-02-150405"
	// Number of packets to keep in the AddPackets Queue
	AddPacketsQueueSize = 8
)

// Globals
var (
	// Flags
	IPv4SourceGroup  = flag.Int("ip4-src-group", 0, "NFLog Group to read IPv4 packets and account the source address")
	IPv4DestGroup    = flag.Int("ip4-dst-group", 0, "NFLog Group to read IPv4 packets and account the destination address")
	IPv4PrefixLength = flag.Int("ip4-prefix-length", 32, "Size of the IPv4 prefix to account to, default is /32")
	IPv6SourceGroup  = flag.Int("ip6-src-group", 0, "NFLog Group to read IPv6 packets and account the source address")
	IPv6DestGroup    = flag.Int("ip6-dst-group", 0, "NFLog Group to read IPv6 packets and account the destination address")
	IPv6PrefixLength = flag.Int("ip6-prefix-length", 64, "Size of the IPv6 prefix to account to, default is /64")
	Cpus             = flag.Int("cpus", 0, "Number of CPUs to use - default 0 is all of them")
	UseSyslog        = flag.Bool("syslog", false, "Use Syslog for logging")
	Debug            = flag.Bool("debug", false, "Print every single packet that arrives")
	Interval         = flag.Duration("interval", time.Minute*5, "Interval to log stats")
	LogDirectory     = flag.String("log-directory", "/var/log/accounting", "Directory to write accounting files to.")
	CpuProfile       = flag.String("cpuprofile", "", "Write cpu profile to file")

	// Globals
	BaseName       = path.Base(os.Args[0])
	Version        = "0.1"
	DefaultMapSize = 1024
)

// Information about one direction for a single IP or range
type HalfIpStats struct {
	Bytes   int64
	Packets int64
}

// Collect info about a single IP (or range)
type IpStats struct {
	Source HalfIpStats
	Dest   HalfIpStats
}

// Accounting for IP addresses
//
// The key is the net.IP which is a []byte converted to a string so it
// can be used as a hash key
type IpMap map[string]*IpStats

// Dump the IpMap to the output as a CSV
func (Ips IpMap) Dump(w io.Writer, when time.Time) error {
	wb := bufio.NewWriter(w)
	whenString := when.Format(CsvTimeFormat)
	_, err := fmt.Fprintf(wb, "Time,IP,SourceBytes,SourcePackets,DestBytes,DestPackets\n")
	if err != nil {
		return err
	}
	for key, stat := range Ips {
		ip := net.IP(key)
		_, err := fmt.Fprintf(wb, "%s,%s,%d,%d,%d,%d\n", whenString, ip, stat.Source.Bytes, stat.Source.Packets, stat.Dest.Bytes, stat.Dest.Packets)
		if err != nil {
			return err
		}
	}
	return wb.Flush()
}

// Returns the time rounded down to *Interval
func FlooredTime(t time.Time) time.Time {
	ns := t.UnixNano()
	ns /= int64(*Interval)
	ns *= int64(*Interval)
	return time.Unix(0, ns)
}

// AddPacket holds the info about one packet
//
// Addr is a net.IP which is a []byte converted into a string This
// won't be a nice UTF-8 string but will preserve the bytes and can be
// used as a hash key
type AddPacket struct {
	Direction IpDirection
	Addr      string
	Length    int
}

// Accounting
type Accounting struct {
	StartPeriod       time.Time
	EndPeriod         time.Time
	Ips               IpMap
	engineWg          sync.WaitGroup
	dumpStatsWg       sync.WaitGroup
	engineStop        chan struct{}
	dumpStatsStop     chan struct{}
	processAddPackets chan []AddPacket
	returnAddPackets  chan []AddPacket
	flip              chan bool
	oldIps            chan IpMap
}

func NewAccounting() *Accounting {
	a := &Accounting{
		Ips:               make(IpMap, DefaultMapSize),
		engineStop:        make(chan struct{}),
		dumpStatsStop:     make(chan struct{}),
		processAddPackets: make(chan []AddPacket, AddPacketsQueueSize),
		returnAddPackets:  make(chan []AddPacket, AddPacketsQueueSize),
		flip:              make(chan bool, 1),
		oldIps:            make(chan IpMap, 1),
	}

	// Make some buffers
	for i := 0; i < AddPacketsQueueSize; i++ {
		a.returnAddPackets <- make([]AddPacket, 0, 128)
	}

	return a
}

// Make a new map returning the old one
func (a *Accounting) newMaps() IpMap {
	a.flip <- true
	return <-a.oldIps
}

// Account a single packet in a thread safe way
func (a *Accounting) Packet(Direction IpDirection, Addr string, Length int) {
	stat := a.Ips[Addr]
	if stat == nil {
		stat = &IpStats{}
		a.Ips[Addr] = stat
	}
	if Direction == IpSource {
		stat.Source.Bytes += int64(Length)
		stat.Source.Packets += 1
	} else {
		stat.Dest.Bytes += int64(Length)
		stat.Dest.Packets += 1
	}
	if *Debug {
		log.Printf("IP message %s Addr %s Size %d", Direction, net.IP(Addr), Length)
	}
}

// Accounting engine to account packets
//
// It has sole control over the Ips map which means that no locking is
// required
func (a *Accounting) Engine() {
	defer a.engineWg.Done()
	for {
		select {
		// Process a bunch of packets
		case ps := <-a.processAddPackets:
			for _, p := range ps {
				a.Packet(p.Direction, p.Addr, p.Length)
			}
			a.returnAddPackets <- ps

		// Flip the map
		case <-a.flip:
			a.oldIps <- a.Ips
			a.Ips = make(IpMap, DefaultMapSize)

		// Stop the engine
		case <-a.engineStop:
			if *Debug {
				log.Printf("Engine stop")
			}
			return
		}
	}
}

// Dump the current stats to a file
//
// We do this carefully writing to a .tmp file and renaming
func (a *Accounting) DumpFile() {
	fileLeaf := a.StartPeriod.Format(FileTimeFormat) + "_" + a.EndPeriod.Format(FileTimeFormat)
	filePath := path.Join(*LogDirectory, fileLeaf)
	fileTmp := filePath + ".tmp"
	fileCsv := filePath + ".csv"
	_, err := os.Lstat(fileCsv)
	if err == nil {
		log.Printf("Output file %q already exists", fileCsv)
		return
	}
	log.Printf("Dumping stats to %s", fileCsv)
	fd, err := os.Create(fileTmp)
	if err != nil {
		log.Printf("Failed to open stats file: %s", err)
		return
	}
	Ips := a.newMaps()
	err = Ips.Dump(fd, a.StartPeriod)
	err2 := fd.Close()
	if err != nil {
		log.Printf("Failed to write to stats file: %s", err)
		return
	}
	if err2 != nil {
		log.Printf("Failed to close stats file: %s", err)
		return
	}
	err = os.Rename(fileTmp, fileCsv)
	if err != nil {
		log.Printf("Failed to rename %q to %q: %s", fileTmp, fileCsv, err)
		return
	}
	log.Printf("Stats file %q written", fileCsv)
}

// Schedules file dumping dump for the end of the interval
func (a *Accounting) DumpStats() {
	defer a.dumpStatsWg.Done()
	a.StartPeriod = time.Now()
	a.EndPeriod = FlooredTime(a.StartPeriod).Add(*Interval)
	for {
		if *Debug {
			log.Printf("Next stats dump at %s", a.EndPeriod)
		}
		select {
		case <-time.After(a.EndPeriod.Sub(time.Now())):
		case <-a.dumpStatsStop:
			if *Debug {
				log.Printf("DumpStats stop")
			}
			a.EndPeriod = time.Now()
			a.DumpFile()
			return
		}
		a.DumpFile()
		a.StartPeriod = a.EndPeriod
		a.EndPeriod = a.StartPeriod.Add(*Interval)
	}
}

// Starts the accounting
func (a *Accounting) Start() {
	log.Printf("Starting accounting")
	a.engineWg.Add(1)
	go a.Engine()
	a.dumpStatsWg.Add(1)
	go a.DumpStats()
	log.Printf("Started accounting")
}

// Stops the accounting saving the stats so far
func (a *Accounting) Stop() {
	log.Printf("Stopping accounting")
	close(a.dumpStatsStop)
	if *Debug {
		log.Printf("Wait for dump stats stop")
	}
	a.dumpStatsWg.Wait()
	close(a.engineStop)
	if *Debug {
		log.Printf("Wait for engine stop")
	}
	a.engineWg.Wait()
	log.Printf("Stopped accounting")
}

// usage prints the syntax
func usage() {
	fmt.Fprintf(os.Stderr,
		"%s ver %s\n\n"+
			"Usage: %s [flags]\n\n"+
			"flags:\n\n",
		BaseName, Version, BaseName)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n")
}

// main
func main() {
	flag.Usage = usage
	flag.Parse()

	const Day = 24 * time.Hour

	// Check interval
	if *Interval < time.Second {
		log.Fatalf("Interval must be >= 1 Second")
	}
	if *Interval >= Day {
		if (*Interval % Day) != 0 {
			log.Fatalf("Interval %s isn't a whole number of days", *Interval)
		}
	} else {
		if (Day % *Interval) != 0 {
			log.Fatalf("Interval %s doesn't divide a day exactly", *Interval)
		}
	}

	// Make output directory
	err := os.MkdirAll(*LogDirectory, 0750)
	if err != nil {
		log.Fatalf("Failed to make log directory %q: %s", *LogDirectory, err)
	}

	if *Cpus <= 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		runtime.GOMAXPROCS(*Cpus)
	}

	if *UseSyslog {
		w, err := syslog.New(syslog.LOG_INFO, BaseName)
		if err != nil {
			log.Fatalf("Failed to start syslog: %s", err)
		}
		log.SetFlags(0)
		log.SetOutput(w)
	}

	// Setup profiling if desired
	if *CpuProfile != "" {
		log.Printf("Starting cpu profiler on %q", *CpuProfile)
		f, err := os.Create(*CpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	a := NewAccounting()
	config := func(Group int, IpVersion byte, Direction IpDirection, MaskBits int) {
		if Group > 0 {
			log.Printf("Monitoring NFLog multicast group %d for IPv%d %s mask /%d", Group, IpVersion, Direction, MaskBits)
			NewNfLog(Group, IpVersion, Direction, MaskBits, a)
		}
	}

	config(*IPv4DestGroup, 4, IpDest, *IPv4PrefixLength)
	config(*IPv4SourceGroup, 4, IpSource, *IPv4PrefixLength)
	config(*IPv6DestGroup, 6, IpDest, *IPv6PrefixLength)
	config(*IPv6SourceGroup, 6, IpSource, *IPv6PrefixLength)

	if nflogs.Count() == 0 {
		log.Fatal("Not monitoring any groups - exiting")
	}

	// Loop forever accounting stuff
	a.Start()

	// Exit neatly on interrupt
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT)
	signal.Notify(ch, syscall.SIGTERM)
	signal.Notify(ch, syscall.SIGQUIT)
	s := <-ch
	log.Printf("%s received - shutting down", s)
	a.Stop()
	nflogs.Stop()
	log.Printf("Exit")
}
<commit_msg>Add -memprofile flag<commit_after>package main

// FIXME write some stats every now and again - packets per second etc

// FIXME detect overflows?

// FIXME write dt in the file?

// FIXME don't dump stats file if empty?

// FIXME have a max entries parameter
// - dump the stats if get beyond that many entries?
// - or just refuse to add any more?
// - or add to 0.0.0.0 entry?

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"log/syslog"
	"net"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"sync"
	"syscall"
	"time"
)

const (
	// Time format to output in CSV files
	// http://stackoverflow.com/questions/804118/best-timestamp-format-for-csv-excel
	CsvTimeFormat = "2006-01-02 15:04:05"
	// Format to use when making filenames
	FileTimeFormat = "2006-01-02-150405"
	// Number of packets to keep in the AddPackets Queue
	AddPacketsQueueSize = 8
)

// Globals
var (
	// Flags
	IPv4SourceGroup  = flag.Int("ip4-src-group", 0, "NFLog Group to read IPv4 packets and account the source address")
	IPv4DestGroup    = flag.Int("ip4-dst-group", 0, "NFLog Group to read IPv4 packets and account the destination address")
	IPv4PrefixLength = flag.Int("ip4-prefix-length", 32, "Size of the IPv4 prefix to account to, default is /32")
	IPv6SourceGroup  = flag.Int("ip6-src-group", 0, "NFLog Group to read IPv6 packets and account the source address")
	IPv6DestGroup    = flag.Int("ip6-dst-group", 0, "NFLog Group to read IPv6 packets and account the destination address")
	IPv6PrefixLength = flag.Int("ip6-prefix-length", 64, "Size of the IPv6 prefix to account to, default is /64")
	Cpus             = flag.Int("cpus", 0, "Number of CPUs to use - default 0 is all of them")
	UseSyslog        = flag.Bool("syslog", false, "Use Syslog for logging")
	Debug            = flag.Bool("debug", false, "Print every single packet that arrives")
	Interval         = flag.Duration("interval", time.Minute*5, "Interval to log stats")
	LogDirectory     = flag.String("log-directory", "/var/log/accounting", "Directory to write accounting files to.")
	CpuProfile       = flag.String("cpuprofile", "", "Write cpu profile to file")
	MemProfile       = flag.String("memprofile", "", "write memory profile to this file")

	// Globals
	BaseName       = path.Base(os.Args[0])
	Version        = "0.1"
	DefaultMapSize = 1024
)

// Information about one direction for a single IP or range
type HalfIpStats struct {
	Bytes   int64
	Packets int64
}

// Collect info about a single IP (or range)
type IpStats struct {
	Source HalfIpStats
	Dest   HalfIpStats
}

// Accounting for IP addresses
//
// The key is the net.IP which is a []byte converted to a string so it
// can be used as a hash key
type IpMap map[string]*IpStats

// Dump the IpMap to the output as a CSV
func (Ips IpMap) Dump(w io.Writer, when time.Time) error {
	wb := bufio.NewWriter(w)
	whenString := when.Format(CsvTimeFormat)
	_, err := fmt.Fprintf(wb, "Time,IP,SourceBytes,SourcePackets,DestBytes,DestPackets\n")
	if err != nil {
		return err
	}
	for key, stat := range Ips {
		ip := net.IP(key)
		_, err := fmt.Fprintf(wb, "%s,%s,%d,%d,%d,%d\n", whenString, ip, stat.Source.Bytes, stat.Source.Packets, stat.Dest.Bytes, stat.Dest.Packets)
		if err != nil {
			return err
		}
	}
	return wb.Flush()
}

// Returns the time rounded down to *Interval
func FlooredTime(t time.Time) time.Time {
	ns := t.UnixNano()
	ns /= int64(*Interval)
	ns *= int64(*Interval)
	return time.Unix(0, ns)
}

// AddPacket holds the info about one packet
//
// Addr is a net.IP which is a []byte converted into a string This
// won't be a nice UTF-8 string but will preserve the bytes and can be
// used as a hash key
type AddPacket struct {
	Direction IpDirection
	Addr      string
	Length    int
}

// Accounting
type Accounting struct {
	StartPeriod       time.Time
	EndPeriod         time.Time
	Ips               IpMap
	engineWg          sync.WaitGroup
	dumpStatsWg       sync.WaitGroup
	engineStop        chan struct{}
	dumpStatsStop     chan struct{}
	processAddPackets chan []AddPacket
	returnAddPackets  chan []AddPacket
	flip              chan bool
	oldIps            chan IpMap
}

func NewAccounting() *Accounting {
	a := &Accounting{
		Ips:               make(IpMap, DefaultMapSize),
		engineStop:        make(chan struct{}),
		dumpStatsStop:     make(chan struct{}),
		processAddPackets: make(chan []AddPacket, AddPacketsQueueSize),
		returnAddPackets:  make(chan []AddPacket, AddPacketsQueueSize),
		flip:              make(chan bool, 1),
		oldIps:            make(chan IpMap, 1),
	}

	// Make some buffers
	for i := 0; i < AddPacketsQueueSize; i++ {
		a.returnAddPackets <- make([]AddPacket, 0, 128)
	}

	return a
}

// Make a new map returning the old one
func (a *Accounting) newMaps() IpMap {
	a.flip <- true
	return <-a.oldIps
}

// Account a single packet in a thread safe way
func (a *Accounting) Packet(Direction IpDirection, Addr string, Length int) {
	stat := a.Ips[Addr]
	if stat == nil {
		stat = &IpStats{}
		a.Ips[Addr] = stat
	}
	if Direction == IpSource {
		stat.Source.Bytes += int64(Length)
		stat.Source.Packets += 1
	} else {
		stat.Dest.Bytes += int64(Length)
		stat.Dest.Packets += 1
	}
	if *Debug {
		log.Printf("IP message %s Addr %s Size %d", Direction, net.IP(Addr), Length)
	}
}

// Accounting engine to account packets
//
// It has sole control over the Ips map which means that no locking is
// required
func (a *Accounting) Engine() {
	defer a.engineWg.Done()
	for {
		select {
		// Process a bunch of packets
		case ps := <-a.processAddPackets:
			for _, p := range ps {
				a.Packet(p.Direction, p.Addr, p.Length)
			}
			a.returnAddPackets <- ps

		// Flip the map
		case <-a.flip:
			a.oldIps <- a.Ips
			a.Ips = make(IpMap, DefaultMapSize)

		// Stop the engine
		case <-a.engineStop:
			if *Debug {
				log.Printf("Engine stop")
			}
			return
		}
	}
}

// Dump the current stats to a file
//
// We do this carefully writing to a .tmp file and renaming
func (a *Accounting) DumpFile() {
	fileLeaf := a.StartPeriod.Format(FileTimeFormat) + "_" + a.EndPeriod.Format(FileTimeFormat)
	filePath := path.Join(*LogDirectory, fileLeaf)
	fileTmp := filePath + ".tmp"
	fileCsv := filePath + ".csv"
	_, err := os.Lstat(fileCsv)
	if err == nil {
		log.Printf("Output file %q already exists", fileCsv)
		return
	}
	log.Printf("Dumping stats to %s", fileCsv)
	fd, err := os.Create(fileTmp)
	if err != nil {
		log.Printf("Failed to open stats file: %s", err)
		return
	}
	Ips := a.newMaps()
	err = Ips.Dump(fd, a.StartPeriod)
	err2 := fd.Close()
	if err != nil {
		log.Printf("Failed to write to stats file: %s", err)
		return
	}
	if err2 != nil {
		log.Printf("Failed to close stats file: %s", err)
		return
	}
	err = os.Rename(fileTmp, fileCsv)
	if err != nil {
		log.Printf("Failed to rename %q to %q: %s", fileTmp, fileCsv, err)
		return
	}
	log.Printf("Stats file %q written", fileCsv)
}

// Schedules file dumping dump for the end of the interval
func (a *Accounting) DumpStats() {
	defer a.dumpStatsWg.Done()
	a.StartPeriod = time.Now()
	a.EndPeriod = FlooredTime(a.StartPeriod).Add(*Interval)
	for {
		if *Debug {
			log.Printf("Next stats dump at %s", a.EndPeriod)
		}
		select {
		case <-time.After(a.EndPeriod.Sub(time.Now())):
		case <-a.dumpStatsStop:
			if *Debug {
				log.Printf("DumpStats stop")
			}
			a.EndPeriod = time.Now()
			a.DumpFile()
			return
		}
		a.DumpFile()
		a.StartPeriod = a.EndPeriod
		a.EndPeriod = a.StartPeriod.Add(*Interval)
	}
}

// Starts the accounting
func (a *Accounting) Start() {
	log.Printf("Starting accounting")
	a.engineWg.Add(1)
	go a.Engine()
	a.dumpStatsWg.Add(1)
	go a.DumpStats()
	log.Printf("Started accounting")
}

// Stops the accounting saving the stats so far
func (a *Accounting) Stop() {
	log.Printf("Stopping accounting")
	close(a.dumpStatsStop)
	if *Debug {
		log.Printf("Wait for dump stats stop")
	}
	a.dumpStatsWg.Wait()
	close(a.engineStop)
	if *Debug {
		log.Printf("Wait for engine stop")
	}
	a.engineWg.Wait()
	log.Printf("Stopped accounting")
}

// usage prints the syntax
func usage() {
	fmt.Fprintf(os.Stderr,
		"%s ver %s\n\n"+
			"Usage: %s [flags]\n\n"+
			"flags:\n\n",
		BaseName, Version, BaseName)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n")
}

// main
func main() {
	flag.Usage = usage
	flag.Parse()

	const Day = 24 * time.Hour

	// Check interval
	if *Interval < time.Second {
		log.Fatalf("Interval must be >= 1 Second")
	}
	if *Interval >= Day {
		if (*Interval % Day) != 0 {
			log.Fatalf("Interval %s isn't a whole number of days", *Interval)
		}
	} else {
		if (Day % *Interval) != 0 {
			log.Fatalf("Interval %s doesn't divide a day exactly", *Interval)
		}
	}

	// Make output directory
	err := os.MkdirAll(*LogDirectory, 0750)
	if err != nil {
		log.Fatalf("Failed to make log directory %q: %s", *LogDirectory, err)
	}

	if *Cpus <= 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		runtime.GOMAXPROCS(*Cpus)
	}

	if *UseSyslog {
		w, err := syslog.New(syslog.LOG_INFO, BaseName)
		if err != nil {
			log.Fatalf("Failed to start syslog: %s", err)
		}
		log.SetFlags(0)
		log.SetOutput(w)
	}

	// Setup profiling if desired
	if *CpuProfile != "" {
		log.Printf("Starting cpu profiler on %q", *CpuProfile)
		f, err := os.Create(*CpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	a := NewAccounting()
	config := func(Group int, IpVersion byte, Direction IpDirection, MaskBits int) {
		if Group > 0 {
			log.Printf("Monitoring NFLog multicast group %d for IPv%d %s mask /%d", Group, IpVersion, Direction, MaskBits)
			NewNfLog(Group, IpVersion, Direction, MaskBits, a)
		}
	}

	config(*IPv4DestGroup, 4, IpDest, *IPv4PrefixLength)
	config(*IPv4SourceGroup, 4, IpSource, *IPv4PrefixLength)
	config(*IPv6DestGroup, 6, IpDest, *IPv6PrefixLength)
	config(*IPv6SourceGroup, 6, IpSource, *IPv6PrefixLength)

	if nflogs.Count() == 0 {
		log.Fatal("Not monitoring any groups - exiting")
	}

	// Loop forever accounting stuff
	a.Start()

	// Exit neatly on interrupt
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT)
	signal.Notify(ch, syscall.SIGTERM)
	signal.Notify(ch, syscall.SIGQUIT)
	s := <-ch
	log.Printf("%s received - shutting down", s)
	if *MemProfile != "" {
		log.Printf("Writing memory profile %q\n", *MemProfile)
		f, err := os.Create(*MemProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
		return
	}
	a.Stop()
	nflogs.Stop()
	log.Printf("Exit")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/mlafeldt/chef-runner/cookbook"
	"github.com/mlafeldt/chef-runner/driver"
	"github.com/mlafeldt/chef-runner/driver/kitchen"
	"github.com/mlafeldt/chef-runner/driver/ssh"
	"github.com/mlafeldt/chef-runner/driver/vagrant"
	"github.com/mlafeldt/chef-runner/log"
	"github.com/mlafeldt/chef-runner/provisioner"
	"github.com/mlafeldt/chef-runner/provisioner/chefsolo"
	"github.com/mlafeldt/chef-runner/util"
)

func logLevel() log.Level {
	l := log.LevelInfo
	e := os.Getenv("CHEF_RUNNER_LOG")
	if e == "" {
		return l
	}
	m := map[string]log.Level{
		"debug": log.LevelDebug,
		"info":  log.LevelInfo,
		"warn":  log.LevelWarn,
		"error": log.LevelError,
	}
	if v, ok := m[strings.ToLower(e)]; ok {
		l = v
	}
	return l
}

func abort(v ...interface{}) {
	log.Error(v...)
	os.Exit(1)
}

func buildRunList(cookbookName string, recipes []string) ([]string, error) {
	runList := []string{}
	for _, r := range recipes {
		var recipeName string
		if strings.Contains(r, "::") {
			recipeName = r
		} else {
			if cookbookName == "" {
				log.Errorf("cannot add recipe `%s` to run list\n", r)
				return nil, errors.New("cookbook name required")
			}
			if path.Dir(r) == "recipes" && path.Ext(r) == ".rb" {
				recipeName = cookbookName + "::" + util.BaseName(r, ".rb")
			} else {
				recipeName = cookbookName + "::" + r
			}
		}
		runList = append(runList, recipeName)
	}
	return runList, nil
}

func uploadFiles(drv driver.Driver) error {
	log.Info("Uploading local files to machine. This may take a while...")
	log.Debugf("Uploading files from %s to %s on machine\n",
		provisioner.SandboxPath, provisioner.RootPath)
	return drv.Upload(provisioner.RootPath, provisioner.SandboxPath+"/")
}

func installChef(drv driver.Driver, p provisioner.Provisioner) error {
	installCmd := p.InstallCommand()
	if len(installCmd) == 0 {
		log.Info("Skipping installation of Chef")
		return nil
	}
	log.Info("Installing Chef")
	return drv.RunCommand(installCmd)
}

func runChef(drv driver.Driver, p provisioner.Provisioner) error {
	log.Infof("Running Chef using %s\n", drv)
	return drv.RunCommand(p.ProvisionCommand())
}

func main() {
	startTime := time.Now()

	log.SetLevel(logLevel())

	flags, err := ParseFlags(os.Args[1:])
	if err != nil {
		abort(err)
	}

	if flags.ShowVersion {
		fmt.Printf("chef-runner %s %s\n", VersionString(), TargetString())
		os.Exit(0)
	}

	log.Infof("Starting chef-runner (%s %s)\n", VersionString(), TargetString())

	var attributes string
	if flags.JSONFile != "" {
		data, err := ioutil.ReadFile(flags.JSONFile)
		if err != nil {
			abort(err)
		}
		attributes = string(data)
	}

	// 1) Run default recipe if no recipes are passed
	// 2) Use run list from JSON file if present, overriding 1)
	// 3) Use run list from command line if present, overriding 1) and 2)
	recipes := flags.Recipes
	if len(recipes) == 0 {
		// TODO: parse actual JSON data
		if strings.Contains(attributes, "run_list") {
			log.Infof("Using run list from %s\n", flags.JSONFile)
		} else {
			recipes = []string{"default"}
		}
	}

	var runList []string
	if len(recipes) > 0 {
		cb, err := cookbook.NewCookbook(".")
		if err != nil {
			abort(err)
		}
		log.Debugf("Cookbook = %s\n", cb)

		if runList, err = buildRunList(cb.Name, recipes); err != nil {
			abort(err)
		}
		log.Infof("Run list is %s\n", runList)
	}

	var p provisioner.Provisioner
	p = chefsolo.Provisioner{
		RunList:     runList,
		Attributes:  attributes,
		Format:      flags.Format,
		LogLevel:    flags.LogLevel,
		UseSudo:     true,
		ChefVersion: flags.ChefVersion,
	}

	log.Debugf("Provisioner = %+v\n", p)

	if err := p.CreateSandbox(); err != nil {
		abort(err)
	}

	var drv driver.Driver
	if flags.Host != "" {
		drv, err = ssh.NewDriver(flags.Host)
	} else if flags.Kitchen != "" {
		drv, err = kitchen.NewDriver(flags.Kitchen)
	} else {
		drv, err = vagrant.NewDriver(flags.Machine)
	}
	if err != nil {
		abort(err)
	}

	if err := uploadFiles(drv); err != nil {
		abort(err)
	}

	if err := installChef(drv, p); err != nil {
		abort(err)
	}

	if err := runChef(drv, p); err != nil {
		abort(err)
	}

	log.Info("chef-runner finished in", time.Now().Sub(startTime))
}
<commit_msg>Extract findDriver function<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/mlafeldt/chef-runner/cookbook"
	"github.com/mlafeldt/chef-runner/driver"
	"github.com/mlafeldt/chef-runner/driver/kitchen"
	"github.com/mlafeldt/chef-runner/driver/ssh"
	"github.com/mlafeldt/chef-runner/driver/vagrant"
	"github.com/mlafeldt/chef-runner/log"
	"github.com/mlafeldt/chef-runner/provisioner"
	"github.com/mlafeldt/chef-runner/provisioner/chefsolo"
	"github.com/mlafeldt/chef-runner/util"
)

func logLevel() log.Level {
	l := log.LevelInfo
	e := os.Getenv("CHEF_RUNNER_LOG")
	if e == "" {
		return l
	}
	m := map[string]log.Level{
		"debug": log.LevelDebug,
		"info":  log.LevelInfo,
		"warn":  log.LevelWarn,
		"error": log.LevelError,
	}
	if v, ok := m[strings.ToLower(e)]; ok {
		l = v
	}
	return l
}

func abort(v ...interface{}) {
	log.Error(v...)
	os.Exit(1)
}

func findDriver(flags *Flags) (driver.Driver, error) {
	if flags.Host != "" {
		return ssh.NewDriver(flags.Host)
	}
	if flags.Kitchen != "" {
		return kitchen.NewDriver(flags.Kitchen)
	}
	return vagrant.NewDriver(flags.Machine)
}

func buildRunList(cookbookName string, recipes []string) ([]string, error) {
	runList := []string{}
	for _, r := range recipes {
		var recipeName string
		if strings.Contains(r, "::") {
			recipeName = r
		} else {
			if cookbookName == "" {
				log.Errorf("cannot add recipe `%s` to run list\n", r)
				return nil, errors.New("cookbook name required")
			}
			if path.Dir(r) == "recipes" && path.Ext(r) == ".rb" {
				recipeName = cookbookName + "::" + util.BaseName(r, ".rb")
			} else {
				recipeName = cookbookName + "::" + r
			}
		}
		runList = append(runList, recipeName)
	}
	return runList, nil
}

func uploadFiles(drv driver.Driver) error {
	log.Info("Uploading local files to machine. This may take a while...")
	log.Debugf("Uploading files from %s to %s on machine\n",
		provisioner.SandboxPath, provisioner.RootPath)
	return drv.Upload(provisioner.RootPath, provisioner.SandboxPath+"/")
}

func installChef(drv driver.Driver, p provisioner.Provisioner) error {
	installCmd := p.InstallCommand()
	if len(installCmd) == 0 {
		log.Info("Skipping installation of Chef")
		return nil
	}
	log.Info("Installing Chef")
	return drv.RunCommand(installCmd)
}

func runChef(drv driver.Driver, p provisioner.Provisioner) error {
	log.Infof("Running Chef using %s\n", drv)
	return drv.RunCommand(p.ProvisionCommand())
}

func main() {
	startTime := time.Now()

	log.SetLevel(logLevel())

	flags, err := ParseFlags(os.Args[1:])
	if err != nil {
		abort(err)
	}

	if flags.ShowVersion {
		fmt.Printf("chef-runner %s %s\n", VersionString(), TargetString())
		os.Exit(0)
	}

	log.Infof("Starting chef-runner (%s %s)\n", VersionString(), TargetString())

	var attributes string
	if flags.JSONFile != "" {
		data, err := ioutil.ReadFile(flags.JSONFile)
		if err != nil {
			abort(err)
		}
		attributes = string(data)
	}

	// 1) Run default recipe if no recipes are passed
	// 2) Use run list from JSON file if present, overriding 1)
	// 3) Use run list from command line if present, overriding 1) and 2)
	recipes := flags.Recipes
	if len(recipes) == 0 {
		// TODO: parse actual JSON data
		if strings.Contains(attributes, "run_list") {
			log.Infof("Using run list from %s\n", flags.JSONFile)
		} else {
			recipes = []string{"default"}
		}
	}

	var runList []string
	if len(recipes) > 0 {
		cb, err := cookbook.NewCookbook(".")
		if err != nil {
			abort(err)
		}
		log.Debugf("Cookbook = %s\n", cb)

		if runList, err = buildRunList(cb.Name, recipes); err != nil {
			abort(err)
		}
		log.Infof("Run list is %s\n", runList)
	}

	var p provisioner.Provisioner
	p = chefsolo.Provisioner{
		RunList:     runList,
		Attributes:  attributes,
		Format:      flags.Format,
		LogLevel:    flags.LogLevel,
		UseSudo:     true,
		ChefVersion: flags.ChefVersion,
	}

	log.Debugf("Provisioner = %+v\n", p)

	if err := p.CreateSandbox(); err != nil {
		abort(err)
	}

	drv, err := findDriver(flags)
	if err != nil {
		abort(err)
	}

	if err := uploadFiles(drv); err != nil {
		abort(err)
	}

	if err := installChef(drv, p); err != nil {
		abort(err)
	}

	if err := runChef(drv, p); err != nil {
		abort(err)
	}

	log.Info("chef-runner finished in", time.Now().Sub(startTime))
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"

	bundle "github.com/TykTechnologies/tyk-cli/bundle"
)

// tyk-cli <module> <submodule> <command> [--options] args...

var module, submodule, command string

var bundleOutput, privKey string
var forceInsecure, showVersion *bool

func init() {
	if len(os.Args) == 1 {
		fmt.Println("No module specified!")
		os.Exit(1)
	}
	if len(os.Args) == 2 {
		fmt.Println("No command specified!")
		os.Exit(1)
	}

	module = os.Args[1]
	command = os.Args[2]

	os.Args = os.Args[2:]

	flag.StringVar(&bundleOutput, "output", "", "Bundle output")
	flag.StringVar(&privKey, "key", "", "Key for bundle signature")
	forceInsecure = flag.Bool("y", false, "Skip bundle signing")
	showVersion = flag.Bool("v", false, "Show version")

	flag.Parse()
}

// main is the entrypoint.
func main() {
	fmt.Println("tyk-cli:", flag.CommandLine, os.Args)

	fmt.Println("module =", module)
	fmt.Println("command =", command)

	if *showVersion {
		fmt.Println("tyk-cli", VERSION)
		os.Exit(0)
	}

	var err error

	switch module {
	case "bundle":
		fmt.Println("Using bundle module.")
		err = bundle.Bundle(command, bundleOutput, privKey, forceInsecure)
	default:
		err = errors.New("Invalid module")
	}

	if err != nil {
		fmt.Println("Error:", err)
		os.Exit(1)
	}
}
<commit_msg>Cleanup.<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"

	bundle "github.com/TykTechnologies/tyk-cli/bundle"
)

// tyk-cli <module> <submodule> <command> [--options] args...

var module, submodule, command string

var bundleOutput, privKey string
var forceInsecure, showVersion *bool

func init() {
	if len(os.Args) == 1 {
		fmt.Println("No module specified!")
		os.Exit(1)
	}
	if len(os.Args) == 2 {
		fmt.Println("No command specified!")
		os.Exit(1)
	}

	module = os.Args[1]
	command = os.Args[2]

	os.Args = os.Args[2:]

	flag.StringVar(&bundleOutput, "output", "", "Bundle output")
	flag.StringVar(&privKey, "key", "", "Key for bundle signature")
	forceInsecure = flag.Bool("y", false, "Skip bundle signing")
	showVersion = flag.Bool("v", false, "Show version")

	flag.Parse()
}

// main is the entrypoint.
func main() {
	if *showVersion {
		fmt.Println("tyk-cli", VERSION)
		os.Exit(0)
	}

	var err error

	switch module {
	case "bundle":
		fmt.Println("Using bundle module.")
		err = bundle.Bundle(command, bundleOutput, privKey, forceInsecure)
	default:
		err = errors.New("Invalid module")
	}

	if err != nil {
		fmt.Println("Error:", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/codegangsta/cli"
	plumber "github.com/qadium/plumber/cli"
	"os"
)

func createRequiredArgCheck(check func(args cli.Args) bool, message string) func(c *cli.Context) error {
	return func(c *cli.Context) error {
		if !check(c.Args()) {
			fmt.Println(message)
			return errors.New(message)
		}
		return nil
	}
}

func exactly(num int) func(args cli.Args) bool {
	return func(args cli.Args) bool {
		return len(args) == num
	}
}

func atLeast(num int) func(args cli.Args) bool {
	return func(args cli.Args) bool {
		return len(args) >= num
	}
}

func main() {
	versionString := versionString()
	app := cli.NewApp()
	app.Name = "plumber"
	app.Usage = "a command line tool for managing distributed data pipelines"
	app.Version = versionString
	// app.Flags = []cli.Flag{
	// 	cli.StringFlag{
	// 		Name:   "server, s",
	// 		Value:  "/var/run/plumber.sock",
	// 		Usage:  "location of plumber server socket",
	// 		EnvVar: "LINK_SERVER",
	// 	},
	// }
	app.Commands = []cli.Command{
		{
			Name:   "add",
			Usage:  "add a plumber-enabled bundle to a pipeline",
			Before: createRequiredArgCheck(atLeast(2), "Please provide both a pipeline name and a bundle path."),
			Action: func(c *cli.Context) {
				pipeline := c.Args()[0]
				bundles := c.Args()[1:]
				if err := plumber.Add(pipeline, bundles...); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:   "create",
			Usage:  "create a pipeline managed by plumber",
			Before: createRequiredArgCheck(exactly(1), "Please provide a pipeline name."),
			Action: func(c *cli.Context) {
				path := c.Args().First()
				if err := plumber.Create(path); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:  "bootstrap",
			Usage: "bootstrap local setup for use with plumber",
			Description: `The bootstrap command builds the latest manager for use with plumber.
This packages the manager into a minimal container for use on localhost.

When running the pipeline on Google Cloud, the manager container is
pushed to your project's private repository.`,
			Action: func(c *cli.Context) {
				if err := plumber.Bootstrap(); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:   "start",
			Usage:  "start a pipeline managed by plumber",
			Before: createRequiredArgCheck(exactly(1), "Please provide a pipeline name."),
			Action: func(c *cli.Context) {
				pipeline := c.Args().First()
				if err := plumber.Start(pipeline); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:   "bundle",
			Usage:  "bundle a node for use in a pipeline managed by plumber",
			Before: createRequiredArgCheck(exactly(1), "Please provide a bundle path."),
			Action: func(c *cli.Context) {
				path := c.Args().First()
				if err := plumber.Bundle(path); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:  "version",
			Usage: "more detailed version information for plumber",
			Action: func(c *cli.Context) {
				fmt.Println("plumber version:", versionString)
				fmt.Println("git commit:", GitCommit)
			},
		},
	}
	app.Run(os.Args)
}

func versionString() string {
	versionString := Version
	if VersionPrerelease != "" {
		versionString += "-" + VersionPrerelease
	}
	return versionString
}
<commit_msg>main cli now takes gce project id as argument<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/codegangsta/cli"
	plumber "github.com/qadium/plumber/cli"
	"os"
)

func createRequiredArgCheck(check func(args cli.Args) bool, message string) func(c *cli.Context) error {
	return func(c *cli.Context) error {
		if !check(c.Args()) {
			fmt.Println(message)
			return errors.New(message)
		}
		return nil
	}
}

func exactly(num int) func(args cli.Args) bool {
	return func(args cli.Args) bool {
		return len(args) == num
	}
}

func atLeast(num int) func(args cli.Args) bool {
	return func(args cli.Args) bool {
		return len(args) >= num
	}
}

func main() {
	versionString := versionString()
	app := cli.NewApp()
	app.Name = "plumber"
	app.Usage = "a command line tool for managing distributed data pipelines"
	app.Version = versionString
	// app.Flags = []cli.Flag{
	// 	cli.StringFlag{
	// 		Name:   "server, s",
	// 		Value:  "/var/run/plumber.sock",
	// 		Usage:  "location of plumber server socket",
	// 		EnvVar: "LINK_SERVER",
	// 	},
	// }
	app.Commands = []cli.Command{
		{
			Name:   "add",
			Usage:  "add a plumber-enabled bundle to a pipeline",
			Before: createRequiredArgCheck(atLeast(2), "Please provide both a pipeline name and a bundle path."),
			Action: func(c *cli.Context) {
				pipeline := c.Args()[0]
				bundles := c.Args()[1:]
				if err := plumber.Add(pipeline, bundles...); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:   "create",
			Usage:  "create a pipeline managed by plumber",
			Before: createRequiredArgCheck(exactly(1), "Please provide a pipeline name."),
			Action: func(c *cli.Context) {
				path := c.Args().First()
				if err := plumber.Create(path); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:  "bootstrap",
			Usage: "bootstrap local setup for use with plumber",
			Description: `The bootstrap command builds the latest manager for use with plumber.
This packages the manager into a minimal container for use on localhost.

When running the pipeline on Google Cloud, the manager container is
pushed to your project's private repository.`,
			Action: func(c *cli.Context) {
				if err := plumber.Bootstrap(); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:   "start",
			Usage:  "start a pipeline managed by plumber",
			Flags:  []cli.Flag {
				cli.StringFlag{
					Name: "gce",
					Value: "",
					Usage: "Google Cloud project ID",
				},
			},
			Before: createRequiredArgCheck(exactly(1), "Please provide a pipeline name."),
			Action: func(c *cli.Context) {
				pipeline := c.Args().First()
				gce := c.String("gce")
				if err := plumber.Start(pipeline, gce); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:   "bundle",
			Usage:  "bundle a node for use in a pipeline managed by plumber",
			Before: createRequiredArgCheck(exactly(1), "Please provide a bundle path."),
			Action: func(c *cli.Context) {
				path := c.Args().First()
				if err := plumber.Bundle(path); err != nil {
					panic(err)
				}
			},
		},
		{
			Name:  "version",
			Usage: "more detailed version information for plumber",
			Action: func(c *cli.Context) {
				fmt.Println("plumber version:", versionString)
				fmt.Println("git commit:", GitCommit)
			},
		},
	}
	app.Run(os.Args)
}

func versionString() string {
	versionString := Version
	if VersionPrerelease != "" {
		versionString += "-" + VersionPrerelease
	}
	return versionString
}
<|endoftext|>


<commit_before>package main

import (
	"time"
	"encoding/json"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"log"
	"os"
	//"io"
	"fmt"
	"database/sql"
	_ "github.com/lib/pq"
	"gopkg.in/antage/eventsource.v1"
)

var conn *sql.DB
var es eventsource.EventSource

func SetHeaders(w *http.ResponseWriter) {
	(*w).Header().Set("Access-Control-Allow-Origin", "*")
	(*w).Header().Set("Content-Type", "application/json")
}

func GetBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM blog")
	data := []BlogPost{}
	for rows.Next() {
		post := BlogPost{}
		rows.Scan(&post.Id, &post.Titel, &post.Text, &post.Auteur, &post.Img_url, &post.Ctime, &post.Image)
		data = append(data, post)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetPost(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT * FROM blog WHERE id = $1 LIMIT 1", ps.ByName("id"))
	data := BlogPost{}
	row.Scan(&data.Id, &data.Titel, &data.Text, &data.Auteur, &data.Img_url, &data.Ctime, &data.Image)
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata ORDER BY tijd DESC LIMIT 1")
	spin := SpinData{}
	rows.Next()
	rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
	buf,_ := json.Marshal(spin)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinBatterij(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT batterij FROM spindata ORDER BY tijd DESC LIMIT 1")
	var data int
	row.Scan(&data)
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinMode(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT mode FROM spindata ORDER BY tijd DESC LIMIT 1")
	var data string
	row.Scan(&data)
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata")
	data := []SpinData{}
	for rows.Next() {
		spin := SpinData{}
		rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
		data = append(data, spin)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinBatterij(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT batterij FROM spindata")
	data := make([]int, 0)
	var scanInt int
	for rows.Next() {
		rows.Scan(&scanInt)
		data = append(data, scanInt)
	}
	buf,_ := json.Marshal(data)
	fmt.Printf(string(buf))
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinMode(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT mode FROM spindata")
	data := make([]string, 0)
	var scanStr string
	for rows.Next() {
		rows.Scan(&scanStr)
		data = append(data, scanStr)
	}
	buf,_ := json.Marshal(data)
	fmt.Printf(string(buf))
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT * FROM servodata ORDER BY tijd DESC LIMIT 1")
	servo := ServoData{}
	row.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
	buf,_ := json.Marshal(servo)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM servodata")
	data := []ServoData{}
	for rows.Next() {
		servo := ServoData{}
		rows.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
		data = append(data, servo)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func Test(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	buf,_ := json.Marshal("test")
	SetHeaders(&w)
	w.Write(buf)
}

func PostBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	//r.ParseMultipartForm(32 << 20)
	/*file, handler, err := r.FormFile("uploadfile")
	defer file.Close()
	if err == nil {
		fmt.Fprintf(w, "%v", handler.Header)
		f, err := os.OpenFile("./img/"+handler.Filename, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			w.WriteHeader(500)
			return
		}
		defer f.Close()
		io.Copy(f, file)
	}

	err = nil*/

	//_,err := conn.Query("INSERT INTO blog (titel, text, auteur, ctime, image) VALUES ($1, $2, $3, $4, $5)", r.FormValue("titel"), r.FormValue("text"), r.FormValue("auteur"), time.Now(), "http://idp-api.herokuapp.com/img/"+handler.Filename)
	_,err := conn.Query("INSERT INTO blog (titel, text, auteur, ctime) VALUES ($1, $2, $3, $4)", r.FormValue("onderwerp"), r.FormValue("bericht"), r.FormValue("naam"), time.Now())
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.Header().Set("Content-Type", "text/html")
	w.WriteHeader(201)
	w.Write([]byte("<meta http-equiv=\"refresh\" content=\"1; url=http://knightspider.herokuapp.com/#/blog\">successful"))
}

func PostSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO spindata (tijd, mode, hellingsgraad, snelheid, batterij, balloncount) VALUES ($1, $2, $3, $4, $5, $6)", time.Now(),
		r.FormValue("mode"), r.FormValue("hellingsgraad"), r.FormValue("snelheid"), r.FormValue("batterij"), r.FormValue("ballonCount"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO servodata (servo_id, tijd, voltage, positie, load, temperatuur) VALUES ($1, $2, $3, $4, $5, $6)",
		r.FormValue("servo_id"), time.Now(), r.FormValue("voltage"), r.FormValue("positie"), r.FormValue("load"), r.FormValue("Temperatuur"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostLog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	es.SendEventMessage(r.FormValue("log"), "", "")
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func main() {
	conn,_ = sql.Open("postgres", os.Getenv("DATABASE_URL"))
	defer conn.Close()

	es = eventsource.New(
		&eventsource.Settings{
			Timeout: 5 * time.Second,
			CloseOnTimeout: false,
			IdleTimeout: 30 * time.Minute,
		},
		func(req *http.Request) [][]byte {
			return [][]byte{
				[]byte("Content-Type: text/event-stream"),
				[]byte("X-Accel-Buffering: no"),
				[]byte("Access-Control-Allow-Origin: *"),
			}
		},
	)
	defer es.Close()

	router := httprouter.New()
	router.GET("/test", Test)
	router.GET("/blog", GetBlog)
	router.GET("/blog/:id", GetPost)
	router.GET("/spin/latest", GetLatestSpinData)
	router.GET("/spin/latest/batterij", GetLatestSpinBatterij)
	router.GET("/spin/latest/mode", GetLatestSpinMode)
	router.GET("/spin/archive", GetArchivedSpinData)
	router.GET("/spin/archive/batterij", GetArchivedSpinBatterij)
	router.GET("/spin/archive/mode", GetArchivedSpinMode)
	router.GET("/servo/latest", GetLatestServoData)
	router.GET("/servo/archive", GetArchivedServoData)
	router.POST("/blog", PostBlog)
	router.POST("/spin", PostSpinData)
	router.POST("/servo", PostServoData)
	router.POST("/log", PostLog)

	http.Handle("/subscribe", es)
	http.Handle("/", router)

	port := os.Getenv("PORT")
	if port == "" {
		port = "3000"
	}
	fmt.Printf("Starting server at localhost:%s...", port)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}<commit_msg>stripped unnecessary SSE data<commit_after>package main

import (
	"time"
	"encoding/json"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"log"
	"os"
	//"io"
	"fmt"
	"database/sql"
	_ "github.com/lib/pq"
	"gopkg.in/antage/eventsource.v1"
)

var conn *sql.DB
var es eventsource.EventSource

func SetHeaders(w *http.ResponseWriter) {
	(*w).Header().Set("Access-Control-Allow-Origin", "*")
	(*w).Header().Set("Content-Type", "application/json")
}

func GetBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM blog")
	data := []BlogPost{}
	for rows.Next() {
		post := BlogPost{}
		rows.Scan(&post.Id, &post.Titel, &post.Text, &post.Auteur, &post.Img_url, &post.Ctime, &post.Image)
		data = append(data, post)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetPost(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT * FROM blog WHERE id = $1 LIMIT 1", ps.ByName("id"))
	data := BlogPost{}
	row.Scan(&data.Id, &data.Titel, &data.Text, &data.Auteur, &data.Img_url, &data.Ctime, &data.Image)
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata ORDER BY tijd DESC LIMIT 1")
	spin := SpinData{}
	rows.Next()
	rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
	buf,_ := json.Marshal(spin)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinBatterij(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT batterij FROM spindata ORDER BY tijd DESC LIMIT 1")
	var data int
	row.Scan(&data)
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinMode(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT mode FROM spindata ORDER BY tijd DESC LIMIT 1")
	var data string
	row.Scan(&data)
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata")
	data := []SpinData{}
	for rows.Next() {
		spin := SpinData{}
		rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
		data = append(data, spin)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinBatterij(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT batterij FROM spindata")
	data := make([]int, 0)
	var scanInt int
	for rows.Next() {
		rows.Scan(&scanInt)
		data = append(data, scanInt)
	}
	buf,_ := json.Marshal(data)
	fmt.Printf(string(buf))
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinMode(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT mode FROM spindata")
	data := make([]string, 0)
	var scanStr string
	for rows.Next() {
		rows.Scan(&scanStr)
		data = append(data, scanStr)
	}
	buf,_ := json.Marshal(data)
	fmt.Printf(string(buf))
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	row := conn.QueryRow("SELECT * FROM servodata ORDER BY tijd DESC LIMIT 1")
	servo := ServoData{}
	row.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
	buf,_ := json.Marshal(servo)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM servodata")
	data := []ServoData{}
	for rows.Next() {
		servo := ServoData{}
		rows.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
		data = append(data, servo)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func Test(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	buf,_ := json.Marshal("test")
	SetHeaders(&w)
	w.Write(buf)
}

func PostBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	//r.ParseMultipartForm(32 << 20)
	/*file, handler, err := r.FormFile("uploadfile")
	defer file.Close()
	if err == nil {
		fmt.Fprintf(w, "%v", handler.Header)
		f, err := os.OpenFile("./img/"+handler.Filename, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			w.WriteHeader(500)
			return
		}
		defer f.Close()
		io.Copy(f, file)
	}

	err = nil*/

	//_,err := conn.Query("INSERT INTO blog (titel, text, auteur, ctime, image) VALUES ($1, $2, $3, $4, $5)", r.FormValue("titel"), r.FormValue("text"), r.FormValue("auteur"), time.Now(), "http://idp-api.herokuapp.com/img/"+handler.Filename)
	_,err := conn.Query("INSERT INTO blog (titel, text, auteur, ctime) VALUES ($1, $2, $3, $4)", r.FormValue("onderwerp"), r.FormValue("bericht"), r.FormValue("naam"), time.Now())
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.Header().Set("Content-Type", "text/html")
	w.WriteHeader(201)
	w.Write([]byte("<meta http-equiv=\"refresh\" content=\"1; url=http://knightspider.herokuapp.com/#/blog\">successful"))
}

func PostSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO spindata (tijd, mode, hellingsgraad, snelheid, batterij, balloncount) VALUES ($1, $2, $3, $4, $5, $6)", time.Now(),
		r.FormValue("mode"), r.FormValue("hellingsgraad"), r.FormValue("snelheid"), r.FormValue("batterij"), r.FormValue("ballonCount"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO servodata (servo_id, tijd, voltage, positie, load, temperatuur) VALUES ($1, $2, $3, $4, $5, $6)",
		r.FormValue("servo_id"), time.Now(), r.FormValue("voltage"), r.FormValue("positie"), r.FormValue("load"), r.FormValue("Temperatuur"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostLog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	es.SendEventMessage(r.FormValue("log"), "", "")
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func main() {
	conn,_ = sql.Open("postgres", os.Getenv("DATABASE_URL"))
	defer conn.Close()

	es = eventsource.New(
		&eventsource.Settings{
			Timeout: 5 * time.Second,
			CloseOnTimeout: false,
			IdleTimeout: 30 * time.Minute,
		},
		func(req *http.Request) [][]byte {
			return [][]byte{
				[]byte("X-Accel-Buffering: no"),
				[]byte("Access-Control-Allow-Origin: *"),
			}
		},
	)
	defer es.Close()

	router := httprouter.New()
	router.GET("/test", Test)
	router.GET("/blog", GetBlog)
	router.GET("/blog/:id", GetPost)
	router.GET("/spin/latest", GetLatestSpinData)
	router.GET("/spin/latest/batterij", GetLatestSpinBatterij)
	router.GET("/spin/latest/mode", GetLatestSpinMode)
	router.GET("/spin/archive", GetArchivedSpinData)
	router.GET("/spin/archive/batterij", GetArchivedSpinBatterij)
	router.GET("/spin/archive/mode", GetArchivedSpinMode)
	router.GET("/servo/latest", GetLatestServoData)
	router.GET("/servo/archive", GetArchivedServoData)
	router.POST("/blog", PostBlog)
	router.POST("/spin", PostSpinData)
	router.POST("/servo", PostServoData)
	router.POST("/log", PostLog)

	http.Handle("/subscribe", es)
	http.Handle("/", router)

	port := os.Getenv("PORT")
	if port == "" {
		port = "3000"
	}
	fmt.Printf("Starting server at localhost:%s...", port)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}<|endoftext|>