<commit_before>// bulk_load_cassandra loads a Cassandra daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

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

	"github.com/gocql/gocql"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/common"
	"github.com/influxdata/influxdb-comparisons/util/report"
	"strconv"
	"strings"
)

// Program option vars:
var (
	daemonUrl      string
	workers        int
	batchSize      int
	doLoad         bool
	writeTimeout   time.Duration
	reportDatabase string
	reportHost     string
	reportUser     string
	reportPassword string
	reportTagsCSV  string
	compressor     string
)

// Global vars
var (
	batchChan      chan *gocql.Batch
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	reportTags     [][2]string
	reportHostname string
)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:9042", "Cassandra URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.DurationVar(&writeTimeout, "write-timeout", 60*time.Second, "Write timeout.")
	flag.StringVar(&compressor, "compressor", "DeflateCompressor", "Table compressor: DeflateCompressor, LZ4Compressor or SnappyCompressor ")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.StringVar(&reportDatabase, "report-database", "database_benchmarks", "Database name where to store result metrics")
	flag.StringVar(&reportHost, "report-host", "", "Host to send result metrics")
	flag.StringVar(&reportUser, "report-user", "", "User for host to send result metrics")
	flag.StringVar(&reportPassword, "report-password", "", "User password for Host to send result metrics")
	flag.StringVar(&reportTagsCSV, "report-tags", "", "Comma separated k:v tags to send  alongside result metrics")

	flag.Parse()

	if reportHost != "" {
		fmt.Printf("results report destination: %v\n", reportHost)
		fmt.Printf("results report database: %v\n", reportDatabase)

		var err error
		reportHostname, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("hostname for results report: %v\n", reportHostname)

		if reportTagsCSV != "" {
			pairs := strings.Split(reportTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				reportTags = append(reportTags, tagpair)
			}
		}
		fmt.Printf("results report tags: %v\n", reportTags)
	}
}

func main() {
	if doLoad {
		log.Println("Creating keyspace")
		createKeyspace(daemonUrl)
	}

	var session *gocql.Session

	if doLoad {
		cluster := gocql.NewCluster(daemonUrl)
		cluster.Keyspace = "measurements"
		cluster.Timeout = writeTimeout
		cluster.Consistency = gocql.One
		cluster.ProtoVersion = 4
		var err error
		session, err = cluster.CreateSession()
		if err != nil {
			log.Fatal(err)
		}
		defer session.Close()
	}

	batchChan = make(chan *gocql.Batch, workers)
	inputDone = make(chan struct{})
	log.Println("Starting workers")
	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processBatches(session)
	}

	start := time.Now()
	itemsRead, bytesRead, valuesRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)

	itemsRate := float64(itemsRead) / float64(took.Seconds())
	bytesRate := float64(bytesRead) / float64(took.Seconds())
	valuesRate := float64(valuesRead) / float64(took.Seconds())

	fmt.Printf("loaded %d items in %fsec with %d workers (mean point rate %.2f/s, mean value rate %.2f/s, %.2fMB/sec from stdin)\n", itemsRead, took.Seconds(), workers, itemsRate, valuesRate, bytesRate/(1<<20))

	if reportHost != "" {
		//append db specific tags to custom tags
		reportTags = append(reportTags, [2]string{"write_timeout", strconv.Itoa(int(writeTimeout))})

		reportParams := &report.LoadReportParams{
			ReportParams: report.ReportParams{
				DBType:             "Cassandra",
				ReportDatabaseName: reportDatabase,
				ReportHost:         reportHost,
				ReportUser:         reportUser,
				ReportPassword:     reportPassword,
				ReportTags:         reportTags,
				Hostname:           reportHostname,
				DestinationUrl:     daemonUrl,
				Workers:            workers,
				ItemLimit:          -1,
			},
			IsGzip:    false,
			BatchSize: batchSize,
		}
		err := report.ReportLoadResult(reportParams, itemsRead, valuesRate, bytesRate, took)

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

// scan reads lines from stdin. It expects input in the Cassandra CQL format.
func scan(session *gocql.Session, itemsPerBatch int) (int64, int64, int64) {
	var batch *gocql.Batch
	if doLoad {
		batch = session.NewBatch(gocql.LoggedBatch)
	}

	var n int
	var err error
	var itemsRead, bytesRead int64
	var totalPoints, totalValues int64

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line := scanner.Text()
		totalPoints, totalValues, err = common.CheckTotalValues(line)
		if totalPoints > 0 || totalValues > 0 {
			continue
		}
		if err != nil {
			log.Fatal(err)
		}
		itemsRead++
		bytesRead += int64(len(scanner.Bytes()))

		if !doLoad {
			continue
		}

		batch.Query(string(scanner.Bytes()))

		n++
		if n >= itemsPerBatch {
			batchChan <- batch
			batch = session.NewBatch(gocql.LoggedBatch)
			n = 0
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading input: %s", err.Error())
	}

	// Finished reading input, make sure last batch goes out.
	if n > 0 {
		batchChan <- batch
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)
	//cassandra's schema stores each value separately, point is represented in series_id
	if itemsRead != totalPoints {
		log.Fatalf("Incorrent number of read items: %d, expected: %d:", itemsRead, totalPoints)
	}

	return itemsRead, bytesRead, totalValues
}

// processBatches reads byte buffers from batchChan and writes them to the target server, while tracking stats on the write.
func processBatches(session *gocql.Session) {
	for batch := range batchChan {
		if !doLoad {
			continue
		}

		// Write the batch.
		err := session.ExecuteBatch(batch)
		if err != nil {
			log.Fatalf("Error writing: %s\n", err.Error())
		}
	}
	workersGroup.Done()
}

var tableOptionsFmt = "with compaction = {'class': 'TimeWindowCompactionStrategy', 'compaction_window_size': 1, 'compaction_window_unit': 'DAYS'} and compression = {'class':'%s'}"

var createTablesCQLDevops = []string{
	"CREATE table measurements.cpu(time bigint,hostname TEXT,region TEXT,datacenter TEXT,rack TEXT,os TEXT,arch TEXT,team TEXT,service TEXT,service_version TEXT,service_environment TEXT,usage_user double,usage_system double,usage_idle double,usage_nice double,usage_iowait double,usage_irq double,usage_softirq double,usage_steal double,usage_guest double,usage_guest_nice double, primary key(hostname, time)) %s;",
	"CREATE table measurements.diskio(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, serial TEXT, reads bigint, writes bigint, read_bytes bigint, write_bytes bigint, read_time bigint, write_time bigint, io_time bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.disk(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, path TEXT, fstype TEXT, total bigint, free bigint, used bigint, used_percent bigint, inodes_total bigint, inodes_free bigint, inodes_used bigint, primary key(hostname, time)) %s;",
	"CREATE table measurements.kernel(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, boot_time bigint, interrupts bigint, context_switches bigint, processes_forked bigint, disk_pages_in bigint, disk_pages_out bigint, primary key(hostname, time)) %s;",
	"CREATE table measurements.mem(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, total bigint, available bigint, used bigint, free bigint, cached bigint, buffered bigint, used_percent double, available_percent double, buffered_percent double, primary key(hostname, time)) %s;",
	"CREATE table measurements.Net(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, interface TEXT, total_connections_received bigint, expired_keys bigint, evicted_keys bigint, keyspace_hits bigint, keyspace_misses bigint, instantaneous_ops_per_sec bigint, instantaneous_input_kbps bigint, instantaneous_output_kbps bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.nginx(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, port TEXT, server TEXT, accepts bigint, active bigint, handled bigint, reading bigint, requests bigint, waiting bigint, writing bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.postgresl(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, numbackends bigint, xact_commit bigint, xact_rollback bigint, blks_read bigint, blks_hit bigint, tup_returned bigint, tup_fetched bigint, tup_inserted bigint, tup_updated bigint, tup_deleted bigint, conflicts bigint, temp_files bigint, temp_bytes bigint, deadlocks bigint, blk_read_time bigint, blk_write_time bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.redis(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, port TEXT, server TEXT, uptime_in_seconds bigint, total_connections_received bigint, expired_keys bigint, evicted_keys bigint, keyspace_hits bigint, keyspace_misses bigint, instantaneous_ops_per_sec bigint, instantaneous_input_kbps bigint, instantaneous_output_kbps bigint, connected_clients bigint, used_memory bigint, used_memory_rss bigint, used_memory_peak bigint, used_memory_lua bigint, rdb_changes_since_last_save bigint, sync_full bigint, sync_partial_ok bigint, sync_partial_err bigint, pubsub_channels bigint, pubsub_patterns bigint, latest_fork_usec bigint, connected_slaves bigint, master_repl_offset bigint, repl_backlog_active bigint, repl_backlog_size bigint, repl_backlog_histlen bigint, mem_fragmentation_ratio bigint, used_cpu_sys bigint, used_cpu_user bigint, used_cpu_sys_children bigint, used_cpu_user_children bigint , primary key(hostname, time)) %s;",
}

func createKeyspace(daemon_url string) {
	cluster := gocql.NewCluster(daemon_url)
	cluster.Consistency = gocql.Quorum
	cluster.ProtoVersion = 4
	cluster.Timeout = writeTimeout
	session, err := cluster.CreateSession()
	if err != nil {
		log.Fatal(err)
	}
	defer session.Close()

	if err := session.Query(`create keyspace measurements with replication = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 };`).Exec(); err != nil {
		log.Print("if you know what you are doing, drop the keyspace with a command line:")
		log.Print("echo 'drop keyspace measurements;' | cqlsh <host>")
		log.Fatal(err)
	}
	tableOptions := fmt.Sprintf(tableOptionsFmt, compressor)

	for _, tableCQLFormat := range createTablesCQLDevops {
		tableCQL := fmt.Sprintf(tableCQLFormat, tableOptions)
		if err := session.Query(tableCQL).Exec(); err != nil {
			log.Fatal(err)
		}
	}
}
<commit_msg>LZ4 as default compressor<commit_after>// bulk_load_cassandra loads a Cassandra daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

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

	"github.com/gocql/gocql"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/common"
	"github.com/influxdata/influxdb-comparisons/util/report"
	"strconv"
	"strings"
)

// Program option vars:
var (
	daemonUrl      string
	workers        int
	batchSize      int
	doLoad         bool
	writeTimeout   time.Duration
	reportDatabase string
	reportHost     string
	reportUser     string
	reportPassword string
	reportTagsCSV  string
	compressor     string
)

// Global vars
var (
	batchChan      chan *gocql.Batch
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	reportTags     [][2]string
	reportHostname string
)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:9042", "Cassandra URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.DurationVar(&writeTimeout, "write-timeout", 60*time.Second, "Write timeout.")
	flag.StringVar(&compressor, "compressor", "LZ4Compressor", "Table compressor: DeflateCompressor, LZ4Compressor or SnappyCompressor ")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.StringVar(&reportDatabase, "report-database", "database_benchmarks", "Database name where to store result metrics")
	flag.StringVar(&reportHost, "report-host", "", "Host to send result metrics")
	flag.StringVar(&reportUser, "report-user", "", "User for host to send result metrics")
	flag.StringVar(&reportPassword, "report-password", "", "User password for Host to send result metrics")
	flag.StringVar(&reportTagsCSV, "report-tags", "", "Comma separated k:v tags to send  alongside result metrics")

	flag.Parse()

	if reportHost != "" {
		fmt.Printf("results report destination: %v\n", reportHost)
		fmt.Printf("results report database: %v\n", reportDatabase)

		var err error
		reportHostname, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("hostname for results report: %v\n", reportHostname)

		if reportTagsCSV != "" {
			pairs := strings.Split(reportTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				reportTags = append(reportTags, tagpair)
			}
		}
		fmt.Printf("results report tags: %v\n", reportTags)
	}
}

func main() {
	if doLoad {
		log.Println("Creating keyspace")
		createKeyspace(daemonUrl)
	}

	var session *gocql.Session

	if doLoad {
		cluster := gocql.NewCluster(daemonUrl)
		cluster.Keyspace = "measurements"
		cluster.Timeout = writeTimeout
		cluster.Consistency = gocql.One
		cluster.ProtoVersion = 4
		var err error
		session, err = cluster.CreateSession()
		if err != nil {
			log.Fatal(err)
		}
		defer session.Close()
	}

	batchChan = make(chan *gocql.Batch, workers)
	inputDone = make(chan struct{})
	log.Println("Starting workers")
	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processBatches(session)
	}

	start := time.Now()
	itemsRead, bytesRead, valuesRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)

	itemsRate := float64(itemsRead) / float64(took.Seconds())
	bytesRate := float64(bytesRead) / float64(took.Seconds())
	valuesRate := float64(valuesRead) / float64(took.Seconds())

	fmt.Printf("loaded %d items in %fsec with %d workers (mean point rate %.2f/s, mean value rate %.2f/s, %.2fMB/sec from stdin)\n", itemsRead, took.Seconds(), workers, itemsRate, valuesRate, bytesRate/(1<<20))

	if reportHost != "" {
		//append db specific tags to custom tags
		reportTags = append(reportTags, [2]string{"write_timeout", strconv.Itoa(int(writeTimeout))})

		reportParams := &report.LoadReportParams{
			ReportParams: report.ReportParams{
				DBType:             "Cassandra",
				ReportDatabaseName: reportDatabase,
				ReportHost:         reportHost,
				ReportUser:         reportUser,
				ReportPassword:     reportPassword,
				ReportTags:         reportTags,
				Hostname:           reportHostname,
				DestinationUrl:     daemonUrl,
				Workers:            workers,
				ItemLimit:          -1,
			},
			IsGzip:    false,
			BatchSize: batchSize,
		}
		err := report.ReportLoadResult(reportParams, itemsRead, valuesRate, bytesRate, took)

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

// scan reads lines from stdin. It expects input in the Cassandra CQL format.
func scan(session *gocql.Session, itemsPerBatch int) (int64, int64, int64) {
	var batch *gocql.Batch
	if doLoad {
		batch = session.NewBatch(gocql.LoggedBatch)
	}

	var n int
	var err error
	var itemsRead, bytesRead int64
	var totalPoints, totalValues int64

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line := scanner.Text()
		totalPoints, totalValues, err = common.CheckTotalValues(line)
		if totalPoints > 0 || totalValues > 0 {
			continue
		}
		if err != nil {
			log.Fatal(err)
		}
		itemsRead++
		bytesRead += int64(len(scanner.Bytes()))

		if !doLoad {
			continue
		}

		batch.Query(string(scanner.Bytes()))

		n++
		if n >= itemsPerBatch {
			batchChan <- batch
			batch = session.NewBatch(gocql.LoggedBatch)
			n = 0
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading input: %s", err.Error())
	}

	// Finished reading input, make sure last batch goes out.
	if n > 0 {
		batchChan <- batch
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)
	//cassandra's schema stores each value separately, point is represented in series_id
	if itemsRead != totalPoints {
		log.Fatalf("Incorrent number of read items: %d, expected: %d:", itemsRead, totalPoints)
	}

	return itemsRead, bytesRead, totalValues
}

// processBatches reads byte buffers from batchChan and writes them to the target server, while tracking stats on the write.
func processBatches(session *gocql.Session) {
	for batch := range batchChan {
		if !doLoad {
			continue
		}

		// Write the batch.
		err := session.ExecuteBatch(batch)
		if err != nil {
			log.Fatalf("Error writing: %s\n", err.Error())
		}
	}
	workersGroup.Done()
}

var tableOptionsFmt = "with compaction = {'class': 'TimeWindowCompactionStrategy', 'compaction_window_size': 1, 'compaction_window_unit': 'DAYS'} and compression = {'class':'%s'}"

var createTablesCQLDevops = []string{
	"CREATE table measurements.cpu(time bigint,hostname TEXT,region TEXT,datacenter TEXT,rack TEXT,os TEXT,arch TEXT,team TEXT,service TEXT,service_version TEXT,service_environment TEXT,usage_user double,usage_system double,usage_idle double,usage_nice double,usage_iowait double,usage_irq double,usage_softirq double,usage_steal double,usage_guest double,usage_guest_nice double, primary key(hostname, time)) %s;",
	"CREATE table measurements.diskio(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, serial TEXT, reads bigint, writes bigint, read_bytes bigint, write_bytes bigint, read_time bigint, write_time bigint, io_time bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.disk(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, path TEXT, fstype TEXT, total bigint, free bigint, used bigint, used_percent bigint, inodes_total bigint, inodes_free bigint, inodes_used bigint, primary key(hostname, time)) %s;",
	"CREATE table measurements.kernel(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, boot_time bigint, interrupts bigint, context_switches bigint, processes_forked bigint, disk_pages_in bigint, disk_pages_out bigint, primary key(hostname, time)) %s;",
	"CREATE table measurements.mem(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, total bigint, available bigint, used bigint, free bigint, cached bigint, buffered bigint, used_percent double, available_percent double, buffered_percent double, primary key(hostname, time)) %s;",
	"CREATE table measurements.Net(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, interface TEXT, total_connections_received bigint, expired_keys bigint, evicted_keys bigint, keyspace_hits bigint, keyspace_misses bigint, instantaneous_ops_per_sec bigint, instantaneous_input_kbps bigint, instantaneous_output_kbps bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.nginx(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, port TEXT, server TEXT, accepts bigint, active bigint, handled bigint, reading bigint, requests bigint, waiting bigint, writing bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.postgresl(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, numbackends bigint, xact_commit bigint, xact_rollback bigint, blks_read bigint, blks_hit bigint, tup_returned bigint, tup_fetched bigint, tup_inserted bigint, tup_updated bigint, tup_deleted bigint, conflicts bigint, temp_files bigint, temp_bytes bigint, deadlocks bigint, blk_read_time bigint, blk_write_time bigint , primary key(hostname, time)) %s;",
	"CREATE table measurements.redis(time bigint, hostname TEXT, region TEXT, datacenter TEXT, rack TEXT, os TEXT, arch TEXT, team TEXT, service TEXT, service_version TEXT, service_environment TEXT, port TEXT, server TEXT, uptime_in_seconds bigint, total_connections_received bigint, expired_keys bigint, evicted_keys bigint, keyspace_hits bigint, keyspace_misses bigint, instantaneous_ops_per_sec bigint, instantaneous_input_kbps bigint, instantaneous_output_kbps bigint, connected_clients bigint, used_memory bigint, used_memory_rss bigint, used_memory_peak bigint, used_memory_lua bigint, rdb_changes_since_last_save bigint, sync_full bigint, sync_partial_ok bigint, sync_partial_err bigint, pubsub_channels bigint, pubsub_patterns bigint, latest_fork_usec bigint, connected_slaves bigint, master_repl_offset bigint, repl_backlog_active bigint, repl_backlog_size bigint, repl_backlog_histlen bigint, mem_fragmentation_ratio bigint, used_cpu_sys bigint, used_cpu_user bigint, used_cpu_sys_children bigint, used_cpu_user_children bigint , primary key(hostname, time)) %s;",
}

func createKeyspace(daemon_url string) {
	cluster := gocql.NewCluster(daemon_url)
	cluster.Consistency = gocql.Quorum
	cluster.ProtoVersion = 4
	cluster.Timeout = writeTimeout
	session, err := cluster.CreateSession()
	if err != nil {
		log.Fatal(err)
	}
	defer session.Close()

	if err := session.Query(`create keyspace measurements with replication = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 };`).Exec(); err != nil {
		log.Print("if you know what you are doing, drop the keyspace with a command line:")
		log.Print("echo 'drop keyspace measurements;' | cqlsh <host>")
		log.Fatal(err)
	}
	tableOptions := fmt.Sprintf(tableOptionsFmt, compressor)

	for _, tableCQLFormat := range createTablesCQLDevops {
		tableCQL := fmt.Sprintf(tableCQLFormat, tableOptions)
		if err := session.Query(tableCQL).Exec(); err != nil {
			log.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/auth"
	"github.com/mesos/mesos-go/auth/sasl"
	"github.com/mesos/mesos-go/auth/sasl/mech"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/mesos/mesos-go/scheduler"
	"github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"

	"github.com/mesosphere/etcd-mesos/rpc"
	etcdscheduler "github.com/mesosphere/etcd-mesos/scheduler"
)

func parseIP(address string) net.IP {
	addr, err := net.LookupIP(address)
	if err != nil {
		log.Fatal(err)
	}
	if len(addr) < 1 {
		log.Fatalf("failed to parse IP from address '%v'", address)
	}
	return addr[0]
}

func main() {
	clusterName :=
		flag.String("cluster-name", "default", "Unique name of this etcd cluster")
	master :=
		flag.String("master", "127.0.0.1:5050", "Master address <ip:port>")
	zkFrameworkPersist :=
		flag.String("zk-framework-persist", "", "Zookeeper URI of the form zk://host1:port1,host2:port2/chroot/path")
	taskCount :=
		flag.Int("cluster-size", 5, "Total task count to run")
	adminPort :=
		flag.Int("admin-port", 23400, "Binding port for admin interface")
	reseedTimeout :=
		flag.Int("reseed-timeout", 240, "Seconds of etcd livelock to wait for before attempting a cluster re-seed")
	autoReseed :=
		flag.Bool("auto-reseed", true, "Perform automatic cluster reseed when the "+
			"cluster has been livelocked for -reseed-timeout seconds")
	artifactPort :=
		flag.Int("artifact-port", 12300, "Binding port for artifact server")
	sandboxDisk :=
		flag.Float64("sandbox-disk-limit", 4096, "Max disk usage for the etcd mesos sandbox in MB")
	sandboxCpu :=
		flag.Float64("sandbox-cpu-limit", 4, "Max cpu usage for the etcd mesos sandbox in MB")
	sandboxMem :=
		flag.Float64("sandbox-mem-limit", 2048, "Max memory usage for the etcd mesos sandbox in MB")
	executorPath :=
		flag.String("executor-bin", "./bin/etcd-mesos-executor", "Path to executor binary")
	etcdPath :=
		flag.String("etcd-bin", "./bin/etcd", "Path to etcd binary")
	etcdctlPath :=
		flag.String("etcdctl-bin", "./bin/etcdctl", "Path to etcdctl binary")
	address :=
		flag.String("address", "", "Binding address for scheduler and artifact server")
	driverPort :=
		flag.Int("driver-port", 0, "Binding port for scheduler driver")
	mesosAuthPrincipal :=
		flag.String("mesos-authentication-principal", "", "Mesos authentication principal")
	mesosAuthSecretFile :=
		flag.String("mesos-authentication-secret-file", "", "Mesos authentication secret file")
	authProvider :=
		flag.String("mesos-authentication-provider", sasl.ProviderName,
			fmt.Sprintf("Authentication provider to use, default is SASL that supports mechanisms: %+v", mech.ListSupported()))
	singleInstancePerSlave :=
		flag.Bool("single-instance-per-slave", true, "Only allow one etcd instance to be started per slave")
	testMode :=
		flag.Bool("test-mode", false, "disable forced values for -single-instance-per-slave and "+
			"-zk-framework-persist")
	failoverTimeoutSeconds :=
		flag.Float64("failover-timeout-seconds", 60*60*24*7, "Mesos framework failover timeout in seconds")
	weburi := flag.String("framework-weburi", "", "A URI that points to a web-based interface for interacting with the framework.")

	flag.Parse()

	if *zkFrameworkPersist == "" && !*testMode {
		log.Fatal("No value provided for -zk-framework-persist ! This can be " +
			"overridden by experts using the -test-mode=true argument, but several " +
			"runtime guarantees no longer hold, and all tasks will be orphaned when " +
			"this process exits.")
	}

	if !*singleInstancePerSlave && !*testMode {
		log.Fatal("-single-instance-per-slave=false is dangerous because it may lead to " +
			"multiple etcd instances in the same cluster on a single node, amplifying " +
			"the cost of a single node being lost, livelock, and data loss.  This can " +
			"be overridden by passing the -test-mode=true argument, at your peril.")
	}

	if *address == "" {
		hostname, err := os.Hostname()
		if err == nil {
			*address = hostname
		} else {
			log.Errorf("Could not set default binding to hostname.  Defaulting to 127.0.0.1")
			*address = "127.0.0.1"
		}
	}

	if *weburi == "" {
		*weburi = fmt.Sprintf("http://%s:%d/stats", *address, *adminPort)
	}

	executorUris := []*mesos.CommandInfo_URI{}
	execUri, err := etcdscheduler.ServeExecutorArtifact(*executorPath, *address, *artifactPort)
	if err != nil {
		log.Errorf("Could not stat executor binary: %v", err)
		return
	}
	executorUris = append(executorUris, &mesos.CommandInfo_URI{
		Value:      execUri,
		Executable: proto.Bool(true),
	})
	etcdUri, err := etcdscheduler.ServeExecutorArtifact(*etcdPath, *address, *artifactPort)
	if err != nil {
		log.Errorf("Could not stat etcd binary: %v", err)
		return
	}
	executorUris = append(executorUris, &mesos.CommandInfo_URI{
		Value:      etcdUri,
		Executable: proto.Bool(true),
	})
	etcdctlUri, err := etcdscheduler.ServeExecutorArtifact(*etcdctlPath, *address, *artifactPort)
	if err != nil {
		log.Errorf("Could not stat etcd binary: %v", err)
		return
	}
	executorUris = append(executorUris, &mesos.CommandInfo_URI{
		Value:      etcdctlUri,
		Executable: proto.Bool(true),
	})

	go http.ListenAndServe(fmt.Sprintf("%s:%d", *address, *artifactPort), nil)
	log.V(2).Info("Serving executor artifacts...")

	bindingAddress := parseIP(*address)

	// chillFactor is the number of seconds that are slept for to allow for
	// convergence across the cluster during mutations.
	chillFactor := 10
	etcdScheduler := etcdscheduler.NewEtcdScheduler(
		*taskCount,
		chillFactor,
		*reseedTimeout,
		*autoReseed,
		executorUris,
		*singleInstancePerSlave,
		*sandboxDisk,
		*sandboxCpu,
		*sandboxMem,
	)
	etcdScheduler.ExecutorPath = *executorPath
	etcdScheduler.Master = *master
	etcdScheduler.ClusterName = *clusterName
	etcdScheduler.ZkConnect = *zkFrameworkPersist

	fwinfo := &mesos.FrameworkInfo{
		User:            proto.String(""), // Mesos-go will fill in user.
		Name:            proto.String("etcd-" + etcdScheduler.ClusterName),
		Checkpoint:      proto.Bool(true),
		FailoverTimeout: proto.Float64(*failoverTimeoutSeconds),
		WebuiUrl:        weburi,
	}

	cred := (*mesos.Credential)(nil)
	if *mesosAuthPrincipal != "" {
		fwinfo.Principal = proto.String(*mesosAuthPrincipal)
		secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
		if err != nil {
			log.Fatal(err)
		}
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
			Secret:    secret,
		}
	}

	zkServers, zkChroot, err := rpc.ParseZKURI(*zkFrameworkPersist)
	etcdScheduler.ZkServers = zkServers
	etcdScheduler.ZkChroot = zkChroot
	if err != nil && *zkFrameworkPersist != "" {
		log.Fatalf("Error parsing zookeeper URI of %s: %s", *zkFrameworkPersist, err)
	} else if *zkFrameworkPersist != "" {
		previous, err := rpc.GetPreviousFrameworkID(
			zkServers,
			zkChroot,
			etcdScheduler.ClusterName,
		)
		if err != nil && err != zk.ErrNoNode {
			log.Fatalf("Could not retrieve previous framework ID: %s", err)
		} else if err == zk.ErrNoNode {
			log.Info("No previous persisted framework ID exists in zookeeper.")
		} else {
			log.Infof("Found stored framework ID in Zookeeper, "+
				"attempting to re-use: %s", previous)
			fwinfo.Id = &mesos.FrameworkID{
				Value: proto.String(previous),
			}
		}
	}

	config := scheduler.DriverConfig{
		Scheduler:      etcdScheduler,
		Framework:      fwinfo,
		Master:         etcdScheduler.Master,
		Credential:     cred,
		BindingAddress: bindingAddress,
		BindingPort:    uint16(*driverPort),
		WithAuthContext: func(ctx context.Context) context.Context {
			ctx = auth.WithLoginProvider(ctx, *authProvider)
			ctx = sasl.WithBindingAddress(ctx, bindingAddress)
			return ctx
		},
	}

	driver, err := scheduler.NewMesosSchedulerDriver(config)

	if err != nil {
		log.Errorln("Unable to create a SchedulerDriver ", err.Error())
	}

	go etcdScheduler.SerialLauncher(driver)
	go etcdScheduler.PeriodicLaunchRequestor()
	go etcdScheduler.AdminHTTP(*adminPort, driver)

	if stat, err := driver.Run(); err != nil {
		log.Infof("Framework stopped with status %s and error: %s",
			stat.String(),
			err.Error())
	}
}
<commit_msg>Fix protobuf weburi<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/auth"
	"github.com/mesos/mesos-go/auth/sasl"
	"github.com/mesos/mesos-go/auth/sasl/mech"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/mesos/mesos-go/scheduler"
	"github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"

	"github.com/mesosphere/etcd-mesos/rpc"
	etcdscheduler "github.com/mesosphere/etcd-mesos/scheduler"
)

func parseIP(address string) net.IP {
	addr, err := net.LookupIP(address)
	if err != nil {
		log.Fatal(err)
	}
	if len(addr) < 1 {
		log.Fatalf("failed to parse IP from address '%v'", address)
	}
	return addr[0]
}

func main() {
	clusterName :=
		flag.String("cluster-name", "default", "Unique name of this etcd cluster")
	master :=
		flag.String("master", "127.0.0.1:5050", "Master address <ip:port>")
	zkFrameworkPersist :=
		flag.String("zk-framework-persist", "", "Zookeeper URI of the form zk://host1:port1,host2:port2/chroot/path")
	taskCount :=
		flag.Int("cluster-size", 5, "Total task count to run")
	adminPort :=
		flag.Int("admin-port", 23400, "Binding port for admin interface")
	reseedTimeout :=
		flag.Int("reseed-timeout", 240, "Seconds of etcd livelock to wait for before attempting a cluster re-seed")
	autoReseed :=
		flag.Bool("auto-reseed", true, "Perform automatic cluster reseed when the "+
			"cluster has been livelocked for -reseed-timeout seconds")
	artifactPort :=
		flag.Int("artifact-port", 12300, "Binding port for artifact server")
	sandboxDisk :=
		flag.Float64("sandbox-disk-limit", 4096, "Max disk usage for the etcd mesos sandbox in MB")
	sandboxCpu :=
		flag.Float64("sandbox-cpu-limit", 4, "Max cpu usage for the etcd mesos sandbox in MB")
	sandboxMem :=
		flag.Float64("sandbox-mem-limit", 2048, "Max memory usage for the etcd mesos sandbox in MB")
	executorPath :=
		flag.String("executor-bin", "./bin/etcd-mesos-executor", "Path to executor binary")
	etcdPath :=
		flag.String("etcd-bin", "./bin/etcd", "Path to etcd binary")
	etcdctlPath :=
		flag.String("etcdctl-bin", "./bin/etcdctl", "Path to etcdctl binary")
	address :=
		flag.String("address", "", "Binding address for scheduler and artifact server")
	driverPort :=
		flag.Int("driver-port", 0, "Binding port for scheduler driver")
	mesosAuthPrincipal :=
		flag.String("mesos-authentication-principal", "", "Mesos authentication principal")
	mesosAuthSecretFile :=
		flag.String("mesos-authentication-secret-file", "", "Mesos authentication secret file")
	authProvider :=
		flag.String("mesos-authentication-provider", sasl.ProviderName,
			fmt.Sprintf("Authentication provider to use, default is SASL that supports mechanisms: %+v", mech.ListSupported()))
	singleInstancePerSlave :=
		flag.Bool("single-instance-per-slave", true, "Only allow one etcd instance to be started per slave")
	testMode :=
		flag.Bool("test-mode", false, "disable forced values for -single-instance-per-slave and "+
			"-zk-framework-persist")
	failoverTimeoutSeconds :=
		flag.Float64("failover-timeout-seconds", 60*60*24*7, "Mesos framework failover timeout in seconds")
	weburi := flag.String("framework-weburi", "", "A URI that points to a web-based interface for interacting with the framework.")

	flag.Parse()

	if *zkFrameworkPersist == "" && !*testMode {
		log.Fatal("No value provided for -zk-framework-persist ! This can be " +
			"overridden by experts using the -test-mode=true argument, but several " +
			"runtime guarantees no longer hold, and all tasks will be orphaned when " +
			"this process exits.")
	}

	if !*singleInstancePerSlave && !*testMode {
		log.Fatal("-single-instance-per-slave=false is dangerous because it may lead to " +
			"multiple etcd instances in the same cluster on a single node, amplifying " +
			"the cost of a single node being lost, livelock, and data loss.  This can " +
			"be overridden by passing the -test-mode=true argument, at your peril.")
	}

	if *address == "" {
		hostname, err := os.Hostname()
		if err == nil {
			*address = hostname
		} else {
			log.Errorf("Could not set default binding to hostname.  Defaulting to 127.0.0.1")
			*address = "127.0.0.1"
		}
	}

	if *weburi == "" {
		*weburi = fmt.Sprintf("http://%s:%d/stats", *address, *adminPort)
	}

	executorUris := []*mesos.CommandInfo_URI{}
	execUri, err := etcdscheduler.ServeExecutorArtifact(*executorPath, *address, *artifactPort)
	if err != nil {
		log.Errorf("Could not stat executor binary: %v", err)
		return
	}
	executorUris = append(executorUris, &mesos.CommandInfo_URI{
		Value:      execUri,
		Executable: proto.Bool(true),
	})
	etcdUri, err := etcdscheduler.ServeExecutorArtifact(*etcdPath, *address, *artifactPort)
	if err != nil {
		log.Errorf("Could not stat etcd binary: %v", err)
		return
	}
	executorUris = append(executorUris, &mesos.CommandInfo_URI{
		Value:      etcdUri,
		Executable: proto.Bool(true),
	})
	etcdctlUri, err := etcdscheduler.ServeExecutorArtifact(*etcdctlPath, *address, *artifactPort)
	if err != nil {
		log.Errorf("Could not stat etcd binary: %v", err)
		return
	}
	executorUris = append(executorUris, &mesos.CommandInfo_URI{
		Value:      etcdctlUri,
		Executable: proto.Bool(true),
	})

	go http.ListenAndServe(fmt.Sprintf("%s:%d", *address, *artifactPort), nil)
	log.V(2).Info("Serving executor artifacts...")

	bindingAddress := parseIP(*address)

	// chillFactor is the number of seconds that are slept for to allow for
	// convergence across the cluster during mutations.
	chillFactor := 10
	etcdScheduler := etcdscheduler.NewEtcdScheduler(
		*taskCount,
		chillFactor,
		*reseedTimeout,
		*autoReseed,
		executorUris,
		*singleInstancePerSlave,
		*sandboxDisk,
		*sandboxCpu,
		*sandboxMem,
	)
	etcdScheduler.ExecutorPath = *executorPath
	etcdScheduler.Master = *master
	etcdScheduler.ClusterName = *clusterName
	etcdScheduler.ZkConnect = *zkFrameworkPersist

	fwinfo := &mesos.FrameworkInfo{
		User:            proto.String(""), // Mesos-go will fill in user.
		Name:            proto.String("etcd-" + etcdScheduler.ClusterName),
		Checkpoint:      proto.Bool(true),
		FailoverTimeout: proto.Float64(*failoverTimeoutSeconds),
		WebuiUrl:        proto.String(*weburi),
	}

	cred := (*mesos.Credential)(nil)
	if *mesosAuthPrincipal != "" {
		fwinfo.Principal = proto.String(*mesosAuthPrincipal)
		secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
		if err != nil {
			log.Fatal(err)
		}
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
			Secret:    secret,
		}
	}

	zkServers, zkChroot, err := rpc.ParseZKURI(*zkFrameworkPersist)
	etcdScheduler.ZkServers = zkServers
	etcdScheduler.ZkChroot = zkChroot
	if err != nil && *zkFrameworkPersist != "" {
		log.Fatalf("Error parsing zookeeper URI of %s: %s", *zkFrameworkPersist, err)
	} else if *zkFrameworkPersist != "" {
		previous, err := rpc.GetPreviousFrameworkID(
			zkServers,
			zkChroot,
			etcdScheduler.ClusterName,
		)
		if err != nil && err != zk.ErrNoNode {
			log.Fatalf("Could not retrieve previous framework ID: %s", err)
		} else if err == zk.ErrNoNode {
			log.Info("No previous persisted framework ID exists in zookeeper.")
		} else {
			log.Infof("Found stored framework ID in Zookeeper, "+
				"attempting to re-use: %s", previous)
			fwinfo.Id = &mesos.FrameworkID{
				Value: proto.String(previous),
			}
		}
	}

	config := scheduler.DriverConfig{
		Scheduler:      etcdScheduler,
		Framework:      fwinfo,
		Master:         etcdScheduler.Master,
		Credential:     cred,
		BindingAddress: bindingAddress,
		BindingPort:    uint16(*driverPort),
		WithAuthContext: func(ctx context.Context) context.Context {
			ctx = auth.WithLoginProvider(ctx, *authProvider)
			ctx = sasl.WithBindingAddress(ctx, bindingAddress)
			return ctx
		},
	}

	driver, err := scheduler.NewMesosSchedulerDriver(config)

	if err != nil {
		log.Errorln("Unable to create a SchedulerDriver ", err.Error())
	}

	go etcdScheduler.SerialLauncher(driver)
	go etcdScheduler.PeriodicLaunchRequestor()
	go etcdScheduler.AdminHTTP(*adminPort, driver)

	if stat, err := driver.Run(); err != nil {
		log.Infof("Framework stopped with status %s and error: %s",
			stat.String(),
			err.Error())
	}
}
<|endoftext|>


<commit_before>package mccli

import (
	"fmt"
	"os"

	"github.com/codegangsta/cli"
	"github.com/materials-commons/mcstore/cmd/pkg/mc"
	"path/filepath"
)

var (
	createProjectCommand = cli.Command{
		Name:    "project",
		Aliases: []string{"proj", "p"},
		Usage:   "Create a new project",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "directory, dir, d",
				Usage: "The base directory for the project",
			},
			cli.BoolFlag{
				Name:  "upload, up, u",
				Usage: "Upload project after creating it",
			},
			cli.IntFlag{
				Name:  "parallel, n",
				Value: 3,
				Usage: "Number of simultaneous uploads to perform, defaults to 3",
			},
		},
		Action: createProjectCLI,
	}
)

// createProjectCLI implements the create project command.
func createProjectCLI(c *cli.Context) {
	if len(c.Args()) != 1 {
		fmt.Println("You must specify a project name")
		os.Exit(1)
	}
	projectName := c.Args()[0]

	dirPath := c.String("directory")
	dirPath, err := filepath.Abs(dirPath)
	if err != nil {
		fmt.Println("Unable to create absolute directory path: ", err)
		os.Exit(1)
	}

	dirPath = filepath.Clean(dirPath)

	if !validateDirectoryPath(dirPath) {
		os.Exit(1)
	}

	client := mc.NewClientAPI()
	if err := client.CreateProject(projectName, dirPath); err != nil {
		fmt.Println("Unable to create project:", err)
		os.Exit(1)
	}

	fmt.Println("Project successfully created.")

	if c.Bool("upload") {
		numThreads := getNumThreads(c)
		if err := client.UploadProject(projectName, numThreads); err != nil {
			fmt.Println("Project upload failed:", err)
			os.Exit(1)
		}
		fmt.Println("Project successfully uploaded.")
	}
}

// validateDirectoryPath checks that the given directory path exists.
func validateDirectoryPath(path string) bool {
	if path == "" {
		fmt.Println("You must specify a local directory path where the project files are located.")
		return false
	}

	if _, err := os.Stat(path); err != nil {
		fmt.Println("Directory doesn't exist or you don't have access")
		return false
	}

	return true
}
<commit_msg>Format code.<commit_after>package mccli

import (
	"fmt"
	"os"

	"path/filepath"

	"github.com/codegangsta/cli"
	"github.com/materials-commons/mcstore/cmd/pkg/mc"
)

var (
	createProjectCommand = cli.Command{
		Name:    "project",
		Aliases: []string{"proj", "p"},
		Usage:   "Create a new project",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "directory, dir, d",
				Usage: "The base directory for the project",
			},
			cli.BoolFlag{
				Name:  "upload, up, u",
				Usage: "Upload project after creating it",
			},
			cli.IntFlag{
				Name:  "parallel, n",
				Value: 3,
				Usage: "Number of simultaneous uploads to perform, defaults to 3",
			},
		},
		Action: createProjectCLI,
	}
)

// createProjectCLI implements the create project command.
func createProjectCLI(c *cli.Context) {
	if len(c.Args()) != 1 {
		fmt.Println("You must specify a project name")
		os.Exit(1)
	}
	projectName := c.Args()[0]

	dirPath := c.String("directory")
	dirPath, err := filepath.Abs(dirPath)
	if err != nil {
		fmt.Println("Unable to create absolute directory path: ", err)
		os.Exit(1)
	}

	dirPath = filepath.Clean(dirPath)

	if !validateDirectoryPath(dirPath) {
		os.Exit(1)
	}

	client := mc.NewClientAPI()
	if err := client.CreateProject(projectName, dirPath); err != nil {
		fmt.Println("Unable to create project:", err)
		os.Exit(1)
	}

	fmt.Println("Project successfully created.")

	if c.Bool("upload") {
		numThreads := getNumThreads(c)
		if err := client.UploadProject(projectName, numThreads); err != nil {
			fmt.Println("Project upload failed:", err)
			os.Exit(1)
		}
		fmt.Println("Project successfully uploaded.")
	}
}

// validateDirectoryPath checks that the given directory path exists.
func validateDirectoryPath(path string) bool {
	if path == "" {
		fmt.Println("You must specify a local directory path where the project files are located.")
		return false
	}

	if _, err := os.Stat(path); err != nil {
		fmt.Println("Directory doesn't exist or you don't have access")
		return false
	}

	return true
}
<|endoftext|>


<commit_before>// Copyright 2013-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// modprobe - Add and remove modules from the Linux Kernel
//
// Synopsis:
//     modprobe [-n] modulename [parameters...]
//     modprobe [-n] -a modulename...
//
// Author:
//     Roland Kammerer <dev.rck@gmail.com>
package main

import (
	"flag"
	"log"
	"os"
	"strings"

	"github.com/u-root/u-root/pkg/kmodule"
)

const cmd = "modprobe [-an] modulename[s] [parameters...]"

var (
	dryRun     = flag.Bool("n", false, "Dry run")
	all        = flag.Bool("a", false, "Insert all module names on the command line.")
	verboseAll = flag.Bool("va", false, "Insert all module names on the command line.")
	rootDir    = flag.String("d", "/", "Root directory for modules")
)

func init() {
	defUsage := flag.Usage
	flag.Usage = func() {
		os.Args[0] = cmd
		defUsage()
	}
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		log.Println("Usage: ERROR: one module and optional module options.")
		flag.Usage()
		os.Exit(1)
	}

	opts := kmodule.ProbeOpts{
		RootDir: *rootDir,
	}
	if *dryRun {
		log.Println("Unique dependencies in load order, already loaded ones get skipped:")
		opts.DryRunCB = func(modPath string) {
			log.Println(modPath)
		}
	}

	// -va is just an alias for -a
	*all = *all || *verboseAll
	if *all {
		modNames := flag.Args()
		for _, modName := range modNames {
			if err := kmodule.ProbeOptions(modName, "", opts); err != nil {
				log.Printf("modprobe: Could not load module %q: %v", modName, err)
			}
		}
		os.Exit(0)
	}

	modName := flag.Args()[0]
	modOptions := strings.Join(flag.Args()[1:], " ")

	if err := kmodule.ProbeOptions(modName, modOptions, opts); err != nil {
		log.Fatalf("modprobe: Could not load module %q: %v", modName, err)
	}
}
<commit_msg>cmds/modprobe: add -S <kernel_version> parameter<commit_after>// Copyright 2013-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// modprobe - Add and remove modules from the Linux Kernel
//
// Synopsis:
//     modprobe [-n] modulename [parameters...]
//     modprobe [-n] -a modulename...
//
// Author:
//     Roland Kammerer <dev.rck@gmail.com>
package main

import (
	"flag"
	"log"
	"os"
	"strings"

	"github.com/u-root/u-root/pkg/kmodule"
)

const cmd = "modprobe [-an] modulename[s] [parameters...]"

var (
	dryRun     = flag.Bool("n", false, "Dry run")
	all        = flag.Bool("a", false, "Insert all module names on the command line.")
	verboseAll = flag.Bool("va", false, "Insert all module names on the command line.")
	rootDir    = flag.String("d", "/", "Root directory for modules")
	kernelVer  = flag.String("S", "", "Set kernel version instead of using uname")
)

func init() {
	defUsage := flag.Usage
	flag.Usage = func() {
		os.Args[0] = cmd
		defUsage()
	}
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		log.Println("Usage: ERROR: one module and optional module options.")
		flag.Usage()
		os.Exit(1)
	}

	opts := kmodule.ProbeOpts{
		RootDir: *rootDir,
		KVer:    *kernelVer,
	}
	if *dryRun {
		log.Println("Unique dependencies in load order, already loaded ones get skipped:")
		opts.DryRunCB = func(modPath string) {
			log.Println(modPath)
		}
	}

	// -va is just an alias for -a
	*all = *all || *verboseAll
	if *all {
		modNames := flag.Args()
		for _, modName := range modNames {
			if err := kmodule.ProbeOptions(modName, "", opts); err != nil {
				log.Printf("modprobe: Could not load module %q: %v", modName, err)
			}
		}
		os.Exit(0)
	}

	modName := flag.Args()[0]
	modOptions := strings.Join(flag.Args()[1:], " ")

	if err := kmodule.ProbeOptions(modName, modOptions, opts); err != nil {
		log.Fatalf("modprobe: Could not load module %q: %v", modName, err)
	}
}
<|endoftext|>


<commit_before>package websockets

import (
	"log"

	"golang.org/x/net/websocket"
	"time"
)

// Server represents Websocket server
type Server struct {
	messages  []*Message
	clients   map[int]*Client
	addCh     chan *Client
	delCh     chan *Client
	sendAllCh chan *Message
	doneCh    chan bool
	errCh     chan error
}

// NewServer creates new chat server
func NewServer() *Server {
	messages := []*Message{}
	clients := make(map[int]*Client)
	addCh := make(chan *Client)
	delCh := make(chan *Client)
	sendAllCh := make(chan *Message)
	doneCh := make(chan bool)
	errCh := make(chan error)

	return &Server{
		messages,
		clients,
		addCh,
		delCh,
		sendAllCh,
		doneCh,
		errCh,
	}
}

// Add adds a new client
func (s *Server) Add(c *Client) {
	s.addCh <- c
}

// Del deletes a client
func (s *Server) Del(c *Client) {
	s.delCh <- c
}

// SendAll sends a broadcast message
func (s *Server) SendAll(msg *Message) {
	s.sendAllCh <- msg
}

// Done closes the server
func (s *Server) Done() {
	s.doneCh <- true
}

// Err sends an error
func (s *Server) Err(err error) {
	s.errCh <- err
}

func (s *Server) sendPastMessages(c *Client) {
	for _, msg := range s.messages {
		c.Write(msg)
	}
}

func (s *Server) sendAll(msg *Message) {
	for _, c := range s.clients {
		go c.Write(msg)
	}
}

// Pulse sends a pulse message to keep the connections alive
func (s *Server) Pulse() {
	for _, c := range s.clients {
		go c.Ping()
	}
}

// OnConnected is the function to be passed to http.Handle(), wrapped in a
// websocket.Handler().
func (s *Server) OnConnected(ws *websocket.Conn) {
	defer func() {
		log.Print("Deferring")
		err := ws.Close()
		if err != nil {
			s.errCh <- err
		}
	}()

	client := NewClient(ws, s)
	s.Add(client)
	client.Listen()
}

// Listen and serve.
// It serves client connection and broadcast request.
func (s *Server) Listen() {
	log.Println("Websocket handler initialized")

	// Setup the worst Ping implentation of all time
	go func() {
		for {
			s.Pulse()
			time.Sleep(time.Second * 55)
		}
	}()

	for {
		select {

		// Add new a client
		case c := <-s.addCh:
			log.Println("Added new client")
			s.clients[c.id] = c
			log.Println("Now", len(s.clients), "clients connected.")
			s.sendPastMessages(c)

		// del a client
		case c := <-s.delCh:
			log.Println("Delete client")
			delete(s.clients, c.id)

		// broadcast message for all clients
		case msg := <-s.sendAllCh:
			log.Println("Send all:", msg)
			s.messages = append(s.messages, msg)
			s.sendAll(msg)

		case err := <-s.errCh:
			log.Println("Error:", err.Error())

		case <-s.doneCh:
			return
		}
	}
}
<commit_msg>Change websocket replay to just send one message<commit_after>package websockets

import (
	"log"

	"golang.org/x/net/websocket"
	"time"
)

// Server represents Websocket server
type Server struct {
	messages  []*Message
	clients   map[int]*Client
	addCh     chan *Client
	delCh     chan *Client
	sendAllCh chan *Message
	doneCh    chan bool
	errCh     chan error
}

// NewServer creates new chat server
func NewServer() *Server {
	messages := []*Message{}
	clients := make(map[int]*Client)
	addCh := make(chan *Client)
	delCh := make(chan *Client)
	sendAllCh := make(chan *Message)
	doneCh := make(chan bool)
	errCh := make(chan error)

	return &Server{
		messages,
		clients,
		addCh,
		delCh,
		sendAllCh,
		doneCh,
		errCh,
	}
}

// Add adds a new client
func (s *Server) Add(c *Client) {
	s.addCh <- c
}

// Del deletes a client
func (s *Server) Del(c *Client) {
	s.delCh <- c
}

// SendAll sends a broadcast message
func (s *Server) SendAll(msg *Message) {
	s.sendAllCh <- msg
}

// Done closes the server
func (s *Server) Done() {
	s.doneCh <- true
}

// Err sends an error
func (s *Server) Err(err error) {
	s.errCh <- err
}

func (s *Server) sendPastMessages(c *Client) {
	// We only need to send the last message, and only if something actually
	// exists.
	if len(s.messages) != 0 {
		c.Write(s.messages[len(s.messages)-1])
	}
}

func (s *Server) sendAll(msg *Message) {
	for _, c := range s.clients {
		go c.Write(msg)
	}
}

// Pulse sends a pulse message to keep the connections alive
func (s *Server) Pulse() {
	for _, c := range s.clients {
		go c.Ping()
	}
}

// OnConnected is the function to be passed to http.Handle(), wrapped in a
// websocket.Handler().
func (s *Server) OnConnected(ws *websocket.Conn) {
	defer func() {
		log.Print("Deferring")
		err := ws.Close()
		if err != nil {
			s.errCh <- err
		}
	}()

	client := NewClient(ws, s)
	s.Add(client)
	client.Listen()
}

// Listen and serve.
// It serves client connection and broadcast request.
func (s *Server) Listen() {
	log.Println("Websocket handler initialized")

	// Setup the worst Ping implentation of all time
	go func() {
		for {
			s.Pulse()
			time.Sleep(time.Second * 55)
		}
	}()

	for {
		select {

		// Add new a client
		case c := <-s.addCh:
			log.Println("Added new client")
			s.clients[c.id] = c
			log.Println("Now", len(s.clients), "clients connected.")
			s.sendPastMessages(c)

		// del a client
		case c := <-s.delCh:
			log.Println("Delete client")
			delete(s.clients, c.id)

		// broadcast message for all clients
		case msg := <-s.sendAllCh:
			log.Println("Send all:", msg)
			s.messages = append(s.messages, msg)
			s.sendAll(msg)

		case err := <-s.errCh:
			log.Println("Error:", err.Error())

		case <-s.doneCh:
			return
		}
	}
}
<|endoftext|>


<commit_before>package wsflate

import (
	"bytes"

	"github.com/gobwas/httphead"
	"github.com/gobwas/ws"
)

// Extension contains logic of compression extension parameters negotiation
// made during HTTP WebSocket handshake.
// It might be reused between different upgrades (but not concurrently) with
// Reset() being called after each.
type Extension struct {
	// Parameters is specification of extension parameters server is going to
	// accept.
	Parameters Parameters

	accepted bool
	params   Parameters
}

// Negotiate parses given HTTP header option and returns (if any) header option
// which describes accepted parameters.
//
// It may return zero option (i.e. one which Size() returns 0) alongside with
// nil error.
func (n *Extension) Negotiate(opt httphead.Option) (accept httphead.Option, err error) {
	if !bytes.Equal(opt.Name, ExtensionNameBytes) {
		return
	}
	if n.accepted {
		// Negotiate might be called multiple times during upgrade.
		// We stick to first one accepted extension since they must be passed
		// in ordered by preference.
		return
	}

	want := n.Parameters

	// NOTE: Parse() resets params inside, so no worries.
	if err = n.params.Parse(opt); err != nil {
		return
	}
	{
		offer := n.params.ServerMaxWindowBits
		want := want.ServerMaxWindowBits
		if offer > want {
			// A server declines an extension negotiation offer
			// with this parameter if the server doesn't support
			// it.
			return
		}
	}
	{
		// If a received extension negotiation offer has the
		// "client_max_window_bits" extension parameter, the server MAY
		// include the "client_max_window_bits" extension parameter in the
		// corresponding extension negotiation response to the offer.
		offer := n.params.ClientMaxWindowBits
		want := want.ClientMaxWindowBits
		if want > offer {
			return
		}
	}
	{
		offer := n.params.ServerNoContextTakeover
		want := want.ServerNoContextTakeover
		if offer && !want {
			return
		}
	}

	n.accepted = true

	return want.Option(), nil
}

// Accepted returns parameters parsed during last negotiation and a flag that
// reports whether they were accepted.
func (n *Extension) Accepted() (_ Parameters, accepted bool) {
	return n.params, n.accepted
}

// Reset resets extension for further reuse.
func (n *Extension) Reset() {
	n.accepted = false
	n.params = Parameters{}
}

var ErrUnexpectedCompressionBit = ws.ProtocolError(
	"control frame or non-first fragment of data contains compression bit set",
)

// UnsetBit clears the Per-Message Compression bit in header h and returns its
// modified copy. It reports whether compression bit was set in header h.
// It returns non-nil error if compression bit has unexpected value.
func UnsetBit(h ws.Header) (_ ws.Header, wasSet bool, err error) {
	var s MessageState
	h, err := s.UnsetBits(h)
	return h, s.IsCompressed(), err
}

// SetBit sets the Per-Message Compression bit in header h and returns its
// modified copy.
// It returns non-nil error if compression bit has unexpected value.
func SetBit(h ws.Header) (_ ws.Header, err error) {
	var s MessageState
	s.SetCompressed(true)
	return s.SetBits(h)
}

// MessageState holds message compression state.
//
// It is consulted during SetBits(h) call to make a decision whether we must
// set the Per-Message Compression bit for given header h argument.
// It is updated during UnsetBits(h) to reflect compression state of a message
// represented by header h argument.
// It can also be consulted/updated directly by calling
// IsCompressed()/SetCompressed().
//
// In general MessageState should be used when there is no direct access to
// connection to read frame from, but it is still needed to know if message
// being read is compressed. For other cases SetBit() and UnsetBit() should be
// used instead.
//
// NOTE: the compression state is updated during UnsetBits(h) only when header
// h argunent represents data (text or binary) frame.
type MessageState struct {
	compressed bool
}

// SetCompressed marks message as "compressed" or "uncompressed".
// See https://tools.ietf.org/html/rfc7692#section-6
func (s *MessageState) SetCompressed(v bool) {
	s.compressed = v
}

// IsCompressed reports whether message is "compressed".
// See https://tools.ietf.org/html/rfc7692#section-6
func (s *MessageState) IsCompressed() bool {
	return s.compressed
}

// UnsetBits changes RSV bits of the given frame header h as if compression
// extension was negotiated. It returns modified copy of h and error if header
// is malformed from the RFC perspective.
func (s *MessageState) UnsetBits(h ws.Header) (ws.Header, error) {
	r1, r2, r3 := ws.RsvBits(h.Rsv)
	switch {
	case h.OpCode.IsData() && h.OpCode != ws.OpContinuation:
		h.Rsv = ws.Rsv(false, r2, r3)
		s.SetCompressed(r1)
		return h, nil

	case r1:
		// An endpoint MUST NOT set the "Per-Message Compressed"
		// bit of control frames and non-first fragments of a data
		// message. An endpoint receiving such a frame MUST _Fail
		// the WebSocket Connection_.
		return h, ErrUnexpectedCompressionBit

	default:
		// NOTE: do not change the state of s.compressed since UnsetBits()
		// might also be called for (intermediate) control frames.
		return h, nil
	}
}

// SetBits changes RSV bits of the frame header h which is being send as if
// compression extension was negotiated. It returns modified copy of h and
// error if header is malformed from the RFC perspective.
func (s *MessageState) SetBits(h ws.Header) (ws.Header, error) {
	r1, r2, r3 := ws.RsvBits(h.Rsv)
	if r1 {
		return h, ErrUnexpectedCompressionBit
	}
	if !h.OpCode.IsData() || h.OpCode == ws.OpContinuation {
		// An endpoint MUST NOT set the "Per-Message Compressed"
		// bit of control frames and non-first fragments of a data
		// message. An endpoint receiving such a frame MUST _Fail
		// the WebSocket Connection_.
		return h, nil
	}
	if s.IsCompressed() {
		h.Rsv = ws.Rsv(true, r2, r3)
	}
	return h, nil
}
<commit_msg>wsflate: fix no new variables typo<commit_after>package wsflate

import (
	"bytes"

	"github.com/gobwas/httphead"
	"github.com/gobwas/ws"
)

// Extension contains logic of compression extension parameters negotiation
// made during HTTP WebSocket handshake.
// It might be reused between different upgrades (but not concurrently) with
// Reset() being called after each.
type Extension struct {
	// Parameters is specification of extension parameters server is going to
	// accept.
	Parameters Parameters

	accepted bool
	params   Parameters
}

// Negotiate parses given HTTP header option and returns (if any) header option
// which describes accepted parameters.
//
// It may return zero option (i.e. one which Size() returns 0) alongside with
// nil error.
func (n *Extension) Negotiate(opt httphead.Option) (accept httphead.Option, err error) {
	if !bytes.Equal(opt.Name, ExtensionNameBytes) {
		return
	}
	if n.accepted {
		// Negotiate might be called multiple times during upgrade.
		// We stick to first one accepted extension since they must be passed
		// in ordered by preference.
		return
	}

	want := n.Parameters

	// NOTE: Parse() resets params inside, so no worries.
	if err = n.params.Parse(opt); err != nil {
		return
	}
	{
		offer := n.params.ServerMaxWindowBits
		want := want.ServerMaxWindowBits
		if offer > want {
			// A server declines an extension negotiation offer
			// with this parameter if the server doesn't support
			// it.
			return
		}
	}
	{
		// If a received extension negotiation offer has the
		// "client_max_window_bits" extension parameter, the server MAY
		// include the "client_max_window_bits" extension parameter in the
		// corresponding extension negotiation response to the offer.
		offer := n.params.ClientMaxWindowBits
		want := want.ClientMaxWindowBits
		if want > offer {
			return
		}
	}
	{
		offer := n.params.ServerNoContextTakeover
		want := want.ServerNoContextTakeover
		if offer && !want {
			return
		}
	}

	n.accepted = true

	return want.Option(), nil
}

// Accepted returns parameters parsed during last negotiation and a flag that
// reports whether they were accepted.
func (n *Extension) Accepted() (_ Parameters, accepted bool) {
	return n.params, n.accepted
}

// Reset resets extension for further reuse.
func (n *Extension) Reset() {
	n.accepted = false
	n.params = Parameters{}
}

var ErrUnexpectedCompressionBit = ws.ProtocolError(
	"control frame or non-first fragment of data contains compression bit set",
)

// UnsetBit clears the Per-Message Compression bit in header h and returns its
// modified copy. It reports whether compression bit was set in header h.
// It returns non-nil error if compression bit has unexpected value.
func UnsetBit(h ws.Header) (_ ws.Header, wasSet bool, err error) {
	var s MessageState
	h, err = s.UnsetBits(h)
	return h, s.IsCompressed(), err
}

// SetBit sets the Per-Message Compression bit in header h and returns its
// modified copy.
// It returns non-nil error if compression bit has unexpected value.
func SetBit(h ws.Header) (_ ws.Header, err error) {
	var s MessageState
	s.SetCompressed(true)
	return s.SetBits(h)
}

// MessageState holds message compression state.
//
// It is consulted during SetBits(h) call to make a decision whether we must
// set the Per-Message Compression bit for given header h argument.
// It is updated during UnsetBits(h) to reflect compression state of a message
// represented by header h argument.
// It can also be consulted/updated directly by calling
// IsCompressed()/SetCompressed().
//
// In general MessageState should be used when there is no direct access to
// connection to read frame from, but it is still needed to know if message
// being read is compressed. For other cases SetBit() and UnsetBit() should be
// used instead.
//
// NOTE: the compression state is updated during UnsetBits(h) only when header
// h argunent represents data (text or binary) frame.
type MessageState struct {
	compressed bool
}

// SetCompressed marks message as "compressed" or "uncompressed".
// See https://tools.ietf.org/html/rfc7692#section-6
func (s *MessageState) SetCompressed(v bool) {
	s.compressed = v
}

// IsCompressed reports whether message is "compressed".
// See https://tools.ietf.org/html/rfc7692#section-6
func (s *MessageState) IsCompressed() bool {
	return s.compressed
}

// UnsetBits changes RSV bits of the given frame header h as if compression
// extension was negotiated. It returns modified copy of h and error if header
// is malformed from the RFC perspective.
func (s *MessageState) UnsetBits(h ws.Header) (ws.Header, error) {
	r1, r2, r3 := ws.RsvBits(h.Rsv)
	switch {
	case h.OpCode.IsData() && h.OpCode != ws.OpContinuation:
		h.Rsv = ws.Rsv(false, r2, r3)
		s.SetCompressed(r1)
		return h, nil

	case r1:
		// An endpoint MUST NOT set the "Per-Message Compressed"
		// bit of control frames and non-first fragments of a data
		// message. An endpoint receiving such a frame MUST _Fail
		// the WebSocket Connection_.
		return h, ErrUnexpectedCompressionBit

	default:
		// NOTE: do not change the state of s.compressed since UnsetBits()
		// might also be called for (intermediate) control frames.
		return h, nil
	}
}

// SetBits changes RSV bits of the frame header h which is being send as if
// compression extension was negotiated. It returns modified copy of h and
// error if header is malformed from the RFC perspective.
func (s *MessageState) SetBits(h ws.Header) (ws.Header, error) {
	r1, r2, r3 := ws.RsvBits(h.Rsv)
	if r1 {
		return h, ErrUnexpectedCompressionBit
	}
	if !h.OpCode.IsData() || h.OpCode == ws.OpContinuation {
		// An endpoint MUST NOT set the "Per-Message Compressed"
		// bit of control frames and non-first fragments of a data
		// message. An endpoint receiving such a frame MUST _Fail
		// the WebSocket Connection_.
		return h, nil
	}
	if s.IsCompressed() {
		h.Rsv = ws.Rsv(true, r2, r3)
	}
	return h, nil
}
<|endoftext|>


<commit_before>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/aquasecurity/kube-bench/check"
	"github.com/golang/glog"
	"github.com/spf13/viper"
)

var (
	errmsgs string
)

func runChecks(nodetype check.NodeType) {
	var summary check.Summary
	var file string
	var err error
	var typeConf *viper.Viper

	switch nodetype {
	case check.MASTER:
		file = masterFile
	case check.NODE:
		file = nodeFile
	case check.FEDERATED:
		file = federatedFile
	}

	runningVersion, err := getKubeVersion()
	if err != nil && kubeVersion == "" {
		exitWithError(fmt.Errorf("Version check failed: %s\nAlternatively, you can specify the version with --version", err))
	}
	path, err := getConfigFilePath(kubeVersion, runningVersion, file)
	if err != nil {
		exitWithError(fmt.Errorf("can't find %s controls file in %s: %v", nodetype, cfgDir, err))
	}

	def := filepath.Join(path, file)
	in, err := ioutil.ReadFile(def)
	if err != nil {
		exitWithError(fmt.Errorf("error opening %s controls file: %v", nodetype, err))
	}

	glog.V(1).Info(fmt.Sprintf("Using benchmark file: %s\n", def))

	// Merge kubernetes version specific config if any.
	viper.SetConfigFile(path + "/config.yaml")
	err = viper.MergeInConfig()
	if err != nil {
		if os.IsNotExist(err) {
			glog.V(2).Info(fmt.Sprintf("No version-specific config.yaml file in %s", path))
		} else {
			exitWithError(fmt.Errorf("couldn't read config file %s: %v", path+"/config.yaml", err))
		}
	} else {
		glog.V(1).Info(fmt.Sprintf("Using config file: %s\n", viper.ConfigFileUsed()))
	}

	// Get the set of exectuables and config files we care about on this type of node. This also
	// checks that the executables we need for the node type are running.
	typeConf = viper.Sub(string(nodetype))
	binmap := getBinaries(typeConf)
	confmap := getConfigFiles(typeConf)
	svcmap := getServiceFiles(typeConf)

	// Variable substitutions. Replace all occurrences of variables in controls files.
	s := string(in)
	s = makeSubstitutions(s, "bin", binmap)
	s = makeSubstitutions(s, "conf", confmap)
	s = makeSubstitutions(s, "svc", svcmap)

	controls, err := check.NewControls(nodetype, []byte(s))
	if err != nil {
		exitWithError(fmt.Errorf("error setting up %s controls: %v", nodetype, err))
	}

	if groupList != "" && checkList == "" {
		ids := cleanIDs(groupList)
		summary = controls.RunGroup(ids...)
	} else if checkList != "" && groupList == "" {
		ids := cleanIDs(checkList)
		summary = controls.RunChecks(ids...)
	} else if checkList != "" && groupList != "" {
		exitWithError(fmt.Errorf("group option and check option can't be used together"))
	} else {
		summary = controls.RunGroup()
	}

	// if we successfully ran some tests and it's json format, ignore the warnings
	if (summary.Fail > 0 || summary.Warn > 0 || summary.Pass > 0) && jsonFmt {
		out, err := controls.JSON()
		if err != nil {
			exitWithError(fmt.Errorf("failed to output in JSON format: %v", err))
		}

		fmt.Println(string(out))
	} else {
		// if we want to store in PostgreSQL, convert to JSON and save it
		if (summary.Fail > 0 || summary.Warn > 0 || summary.Pass > 0) && pgSQL {
			out, err := controls.JSON()
			if err != nil {
				exitWithError(fmt.Errorf("failed to output in JSON format: %v", err))
			}

			savePgsql(string(out))
		} else {
			prettyPrint(controls, summary)
		}
	}
}

// colorPrint outputs the state in a specific colour, along with a message string
func colorPrint(state check.State, s string) {
	colors[state].Printf("[%s] ", state)
	fmt.Printf("%s", s)
}

// prettyPrint outputs the results to stdout in human-readable format
func prettyPrint(r *check.Controls, summary check.Summary) {
	// Print check results.
	if !noResults {
		colorPrint(check.INFO, fmt.Sprintf("%s %s\n", r.ID, r.Text))
		for _, g := range r.Groups {
			colorPrint(check.INFO, fmt.Sprintf("%s %s\n", g.ID, g.Text))
			for _, c := range g.Checks {
				colorPrint(c.State, fmt.Sprintf("%s %s\n", c.ID, c.Text))
			}
		}

		fmt.Println()
	}

	// Print remediations.
	if !noRemediations {
		if summary.Fail > 0 || summary.Warn > 0 {
			colors[check.WARN].Printf("== Remediations ==\n")
			for _, g := range r.Groups {
				for _, c := range g.Checks {
					if c.State != check.PASS {
						fmt.Printf("%s %s\n", c.ID, c.Remediation)
					}
				}
			}
			fmt.Println()
		}
	}

	// Print summary setting output color to highest severity.
	if !noSummary {
		var res check.State
		if summary.Fail > 0 {
			res = check.FAIL
		} else if summary.Warn > 0 {
			res = check.WARN
		} else {
			res = check.PASS
		}

		colors[res].Printf("== Summary ==\n")
		fmt.Printf("%d checks PASS\n%d checks FAIL\n%d checks WARN\n",
			summary.Pass, summary.Fail, summary.Warn,
		)
	}
}
<commit_msg>Only get runningVersion if --version has not been provided<commit_after>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/aquasecurity/kube-bench/check"
	"github.com/golang/glog"
	"github.com/spf13/viper"
)

var (
	errmsgs string
)

func runChecks(nodetype check.NodeType) {
	var summary check.Summary
	var file string
	var err error
	var typeConf *viper.Viper

	switch nodetype {
	case check.MASTER:
		file = masterFile
	case check.NODE:
		file = nodeFile
	case check.FEDERATED:
		file = federatedFile
	}

	runningVersion := ""
	if kubeVersion == "" {
		runningVersion, err = getKubeVersion()
		if err != nil {
			exitWithError(fmt.Errorf("Version check failed: %s\nAlternatively, you can specify the version with --version", err))
		}
	}
	path, err := getConfigFilePath(kubeVersion, runningVersion, file)
	if err != nil {
		exitWithError(fmt.Errorf("can't find %s controls file in %s: %v", nodetype, cfgDir, err))
	}

	def := filepath.Join(path, file)
	in, err := ioutil.ReadFile(def)
	if err != nil {
		exitWithError(fmt.Errorf("error opening %s controls file: %v", nodetype, err))
	}

	glog.V(1).Info(fmt.Sprintf("Using benchmark file: %s\n", def))

	// Merge kubernetes version specific config if any.
	viper.SetConfigFile(path + "/config.yaml")
	err = viper.MergeInConfig()
	if err != nil {
		if os.IsNotExist(err) {
			glog.V(2).Info(fmt.Sprintf("No version-specific config.yaml file in %s", path))
		} else {
			exitWithError(fmt.Errorf("couldn't read config file %s: %v", path+"/config.yaml", err))
		}
	} else {
		glog.V(1).Info(fmt.Sprintf("Using config file: %s\n", viper.ConfigFileUsed()))
	}

	// Get the set of exectuables and config files we care about on this type of node. This also
	// checks that the executables we need for the node type are running.
	typeConf = viper.Sub(string(nodetype))
	binmap := getBinaries(typeConf)
	confmap := getConfigFiles(typeConf)
	svcmap := getServiceFiles(typeConf)

	// Variable substitutions. Replace all occurrences of variables in controls files.
	s := string(in)
	s = makeSubstitutions(s, "bin", binmap)
	s = makeSubstitutions(s, "conf", confmap)
	s = makeSubstitutions(s, "svc", svcmap)

	controls, err := check.NewControls(nodetype, []byte(s))
	if err != nil {
		exitWithError(fmt.Errorf("error setting up %s controls: %v", nodetype, err))
	}

	if groupList != "" && checkList == "" {
		ids := cleanIDs(groupList)
		summary = controls.RunGroup(ids...)
	} else if checkList != "" && groupList == "" {
		ids := cleanIDs(checkList)
		summary = controls.RunChecks(ids...)
	} else if checkList != "" && groupList != "" {
		exitWithError(fmt.Errorf("group option and check option can't be used together"))
	} else {
		summary = controls.RunGroup()
	}

	// if we successfully ran some tests and it's json format, ignore the warnings
	if (summary.Fail > 0 || summary.Warn > 0 || summary.Pass > 0) && jsonFmt {
		out, err := controls.JSON()
		if err != nil {
			exitWithError(fmt.Errorf("failed to output in JSON format: %v", err))
		}

		fmt.Println(string(out))
	} else {
		// if we want to store in PostgreSQL, convert to JSON and save it
		if (summary.Fail > 0 || summary.Warn > 0 || summary.Pass > 0) && pgSQL {
			out, err := controls.JSON()
			if err != nil {
				exitWithError(fmt.Errorf("failed to output in JSON format: %v", err))
			}

			savePgsql(string(out))
		} else {
			prettyPrint(controls, summary)
		}
	}
}

// colorPrint outputs the state in a specific colour, along with a message string
func colorPrint(state check.State, s string) {
	colors[state].Printf("[%s] ", state)
	fmt.Printf("%s", s)
}

// prettyPrint outputs the results to stdout in human-readable format
func prettyPrint(r *check.Controls, summary check.Summary) {
	// Print check results.
	if !noResults {
		colorPrint(check.INFO, fmt.Sprintf("%s %s\n", r.ID, r.Text))
		for _, g := range r.Groups {
			colorPrint(check.INFO, fmt.Sprintf("%s %s\n", g.ID, g.Text))
			for _, c := range g.Checks {
				colorPrint(c.State, fmt.Sprintf("%s %s\n", c.ID, c.Text))
			}
		}

		fmt.Println()
	}

	// Print remediations.
	if !noRemediations {
		if summary.Fail > 0 || summary.Warn > 0 {
			colors[check.WARN].Printf("== Remediations ==\n")
			for _, g := range r.Groups {
				for _, c := range g.Checks {
					if c.State != check.PASS {
						fmt.Printf("%s %s\n", c.ID, c.Remediation)
					}
				}
			}
			fmt.Println()
		}
	}

	// Print summary setting output color to highest severity.
	if !noSummary {
		var res check.State
		if summary.Fail > 0 {
			res = check.FAIL
		} else if summary.Warn > 0 {
			res = check.WARN
		} else {
			res = check.PASS
		}

		colors[res].Printf("== Summary ==\n")
		fmt.Printf("%d checks PASS\n%d checks FAIL\n%d checks WARN\n",
			summary.Pass, summary.Fail, summary.Warn,
		)
	}
}
<|endoftext|>