<commit_before>package main

import (
    "flag"
    "fmt"
    "log/syslog"
    "strconv"
    "strings"
    "syscall"
    "time"

    "github.com/BurntSushi/toml"
    "github.com/akrennmair/gopcap"
    "github.com/nranchev/go-libGeoIP"
)

type Packet struct {
    ts      time.Time
    tuple   IpPortTuple
    payload []byte
}

type protocolType uint16

const (
    UnknownProtocol protocolType = iota
    HttpProtocol
    MysqlProtocol
    RedisProtocol
)

var protocolNames = []string{"unknown", "http", "mysql", "redis"}

type tomlConfig struct {
    Interfaces    tomlInterfaces
    RunOptions    tomlRunOptions
    Protocols     map[string]tomlProtocol
    Procs         tomlProcs
    Elasticsearch tomlMothership
    Agent         tomlAgent
    Logging       tomlLogging
}

type tomlInterfaces struct {
    Device string
}

type tomlRunOptions struct {
    Uid int
    Gid int
}

type tomlLogging struct {
    Selectors []string
}

var _Config tomlConfig
var _ConfigMeta toml.MetaData
var _GeoLite *libgeo.GeoIP

func Bytes_Ipv4_Ntoa(bytes []byte) string {
    var strarr []string = make([]string, 4)
    for i, b := range bytes {
        strarr[i] = strconv.Itoa(int(b))
    }
    return strings.Join(strarr, ".")
}

func decodePktEth(datalink int, pkt *pcap.Packet) {

    packet := new(Packet)
    var l2hlen int
    var eth_type uint16

    if datalink != pcap.LINKTYPE_LINUX_SLL {
        // we are just assuming ETH
        l2hlen = 14

        // bytes 12 and 13 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[12:14])

    } else {
        l2hlen = 16

        // bytes 14 and 15 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[14:16])
    }

    if eth_type != 0x800 {
        DEBUG("ip", "Not ipv4 packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+20 {
        DEBUG("ip", "Packet to short to be ethernet")
        return
    }

    // IP header
    iphl := int((uint16(pkt.Data[l2hlen]) & 0x0f) * 4)
    if len(pkt.Data) < l2hlen+iphl {
        DEBUG("ip", "Packet to short to be IP")
        return
    }
    iphdr := pkt.Data[l2hlen : l2hlen+iphl]

    //DEBUG("Packet timestamp: %s", pkt.Time)
    packet.ts = pkt.Time

    packet.tuple.Src_ip = Bytes_Ntohl(iphdr[12:16])
    packet.tuple.Dst_ip = Bytes_Ntohl(iphdr[16:20])

    mf := (uint8(iphdr[6]&0x20) != 0)
    frag_offset := (uint16(iphdr[7]) << 3) | (uint16(iphdr[6]) & 0x1F)

    if mf || frag_offset > 0 {
        DEBUG("ip", "Fragmented packets not yet supported")
        return
    }

    ip_length := int(Bytes_Ntohs(iphdr[2:4]))

    protocol := uint8(iphdr[9])
    if protocol != 6 {
        DEBUG("ip", "Not TCP packet received. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+iphl+20 {
        DEBUG("ip", "Packet too short to be TCP")
        return
    }

    tcphl := int((uint16(pkt.Data[l2hlen+iphl+12]) >> 4) * 4)
    if tcphl > 20 && len(pkt.Data) < l2hlen+iphl+tcphl {
        DEBUG("ip", "Packet too short for the TCP header")
        return
    }

    tcphdr := pkt.Data[l2hlen+iphl : l2hlen+iphl+tcphl]
    packet.tuple.Src_port = Bytes_Ntohs(tcphdr[0:2])
    packet.tuple.Dst_port = Bytes_Ntohs(tcphdr[2:4])

    //DEBUG(" %s:%d -> %s:%d", packet.tuple.src_ip, packet.tuple.src_port,
    //      packet.tuple.dst_ip, packet.tuple.dst_port)

    packet.payload = pkt.Data[l2hlen+iphl+tcphl : l2hlen+ip_length]

    FollowTcp(tcphdr, packet)
}

func DropPrivileges() error {
    var err error

    if !_ConfigMeta.IsDefined("runoptions", "uid") {
        // not found, no dropping privileges but no err
        return nil
    }

    if !_ConfigMeta.IsDefined("runoptions", "gid") {
        return MsgError("GID must be specified for dropping priviledges")
    }

    INFO("Switching to user: %d.%d", _Config.RunOptions.Uid, _Config.RunOptions.Gid)

    if err = syscall.Setgid(_Config.RunOptions.Gid); err != nil {
        return MsgError("setgid: %s", err.Error())
    }

    if err = syscall.Setuid(_Config.RunOptions.Uid); err != nil {
        return MsgError("setuid: %s", err.Error())
    }

    return nil
}

func main() {

    configfile := flag.String("c", "agent.conf", "Configuration file")
    file := flag.String("I", "", "file")
    loop := flag.Int("l", 0, "Loop file. 0 - loop forever")
    debugSelectorsStr := flag.String("d", "", "Enable certain debug selectors")
    oneAtAtime := flag.Bool("O", false, "Read packets one at a time (press Enter)")
    toStdout := flag.Bool("e", false, "Output to stdout instead of syslog")
    panicAfter := flag.Int("P", 0, "Panic after a given number of seconds, only for testing")

    flag.Parse()

    debugSelectors := []string{}
    logLevel := syslog.LOG_DEBUG
    if len(*debugSelectorsStr) > 0 {
        debugSelectors = strings.Split(*debugSelectorsStr, ",")
    }

    var h *pcap.Pcap
    var err error

    if _ConfigMeta, err = toml.DecodeFile(*configfile, &_Config); err != nil {
        fmt.Printf("TOML config parsing failed on %s: %s. We will use defaults.\n", *configfile, err)
    }

    if len(debugSelectors) == 0 {
        debugSelectors = _Config.Logging.Selectors
    }
    LogInit(logLevel, "", !*toStdout, debugSelectors)

    if *panicAfter > 0 {
        go func() {
            time.Sleep(time.Duration(*panicAfter) * time.Second)
            panic("Panicking as requested")
        }()
    }

    if *file != "" {
        h, err = pcap.Openoffline(*file)
        if h == nil {
            ERR("Openoffline(%s) failed: %s", *file, err)
            return
        }
    } else {
        device := _Config.Interfaces.Device

        h, err = pcap.Openlive(device, 65535, true, 0)
        if h == nil {
            ERR("pcap.Openlive failed(%s): %s", device, err)
            return
        }
    }
    defer h.Close()

    if err = DropPrivileges(); err != nil {
        CRIT(err.Error())
        return
    }

    filter := configToFilter(&_Config)
    if filter != "" {
        DEBUG("pcapfilter", "Installing filter '%s'", filter)
        err := h.Setfilter(filter)
        if err != nil {
            ERR("pcap.Setfilter failed: %s", err)
            return
        }
    }

    if err = Publisher.Init(); err != nil {
        CRIT(err.Error())
        return
    }

    if err = procWatcher.Init(&_Config.Procs); err != nil {
        CRIT(err.Error())
        return
    }

    if err = TcpInit(); err != nil {
        CRIT(err.Error())
        return
    }

    datalink := h.Datalink()
    if datalink != pcap.LINKTYPE_ETHERNET && datalink != pcap.LINKTYPE_LINUX_SLL {
        WARN("Unsuported link type: %d", datalink)
    }

    _GeoLite, err = libgeo.Load("/usr/share/GeoIP/GeoIP.dat")
    if err != nil {
        WARN("Could not load GeoIP data: %s", err.Error())
    }

    counter := 0
    live := true
    loopCount := 1
    var lastPktTime *time.Time = nil
    for live {
        var pkt *pcap.Packet
        var res int32

        if *oneAtAtime {
            fmt.Println("Press enter to read packet")
            fmt.Scanln()
        }

        pkt, res = h.NextEx()
        switch res {
        case -1:
            ERR("pcap.NextEx() error: %s", h.Geterror())
            live = false
            continue
        case -2:
            DEBUG("pcapread", "End of file")
            loopCount += 1
            if *loop > 0 && loopCount > *loop {
                // give a bit of time to the publish goroutine
                // to flush
                time.Sleep(300 * time.Millisecond)
                live = false
            }
            // reopen the file
            h, err = pcap.Openoffline(*file)
            if h == nil {
                ERR("Openoffline(%s) failed: %s", *file, err)
                return
            }
            lastPktTime = nil
            continue
        case 0:
            // timeout
            continue
        }
        if res != 1 {
            panic(fmt.Sprintf("Unexpected return code from pcap.NextEx: %d", res))
        }

        if pkt == nil {
            panic("Nil packet despite res=1")
        }

        if *file != "" {
            if lastPktTime != nil {
                sleep := pkt.Time.Sub(*lastPktTime)
                if sleep > 0 {
                    time.Sleep(sleep)
                } else {
                    WARN("Time in pcap went back in time: %d", sleep)
                }
            }
            _lastPktTime := pkt.Time
            lastPktTime = &_lastPktTime
            pkt.Time = time.Now() // overwrite what we get from the pcap
        }
        counter++

        //DEBUG("pkt received")

        decodePktEth(datalink, pkt)
    }
    INFO("Input finish. Processed %d packets. Have a nice day!", counter)
}
<commit_msg>Exit in case parsing config file fails.<commit_after>package main

import (
    "flag"
    "fmt"
    "log/syslog"
    "strconv"
    "strings"
    "syscall"
    "time"

    "github.com/BurntSushi/toml"
    "github.com/akrennmair/gopcap"
    "github.com/nranchev/go-libGeoIP"
)

type Packet struct {
    ts      time.Time
    tuple   IpPortTuple
    payload []byte
}

type protocolType uint16

const (
    UnknownProtocol protocolType = iota
    HttpProtocol
    MysqlProtocol
    RedisProtocol
)

var protocolNames = []string{"unknown", "http", "mysql", "redis"}

type tomlConfig struct {
    Interfaces    tomlInterfaces
    RunOptions    tomlRunOptions
    Protocols     map[string]tomlProtocol
    Procs         tomlProcs
    Elasticsearch tomlMothership
    Agent         tomlAgent
    Logging       tomlLogging
}

type tomlInterfaces struct {
    Device string
}

type tomlRunOptions struct {
    Uid int
    Gid int
}

type tomlLogging struct {
    Selectors []string
}

var _Config tomlConfig
var _ConfigMeta toml.MetaData
var _GeoLite *libgeo.GeoIP

func Bytes_Ipv4_Ntoa(bytes []byte) string {
    var strarr []string = make([]string, 4)
    for i, b := range bytes {
        strarr[i] = strconv.Itoa(int(b))
    }
    return strings.Join(strarr, ".")
}

func decodePktEth(datalink int, pkt *pcap.Packet) {

    packet := new(Packet)
    var l2hlen int
    var eth_type uint16

    if datalink != pcap.LINKTYPE_LINUX_SLL {
        // we are just assuming ETH
        l2hlen = 14

        // bytes 12 and 13 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[12:14])

    } else {
        l2hlen = 16

        // bytes 14 and 15 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[14:16])
    }

    if eth_type != 0x800 {
        DEBUG("ip", "Not ipv4 packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+20 {
        DEBUG("ip", "Packet to short to be ethernet")
        return
    }

    // IP header
    iphl := int((uint16(pkt.Data[l2hlen]) & 0x0f) * 4)
    if len(pkt.Data) < l2hlen+iphl {
        DEBUG("ip", "Packet to short to be IP")
        return
    }
    iphdr := pkt.Data[l2hlen : l2hlen+iphl]

    //DEBUG("Packet timestamp: %s", pkt.Time)
    packet.ts = pkt.Time

    packet.tuple.Src_ip = Bytes_Ntohl(iphdr[12:16])
    packet.tuple.Dst_ip = Bytes_Ntohl(iphdr[16:20])

    mf := (uint8(iphdr[6]&0x20) != 0)
    frag_offset := (uint16(iphdr[7]) << 3) | (uint16(iphdr[6]) & 0x1F)

    if mf || frag_offset > 0 {
        DEBUG("ip", "Fragmented packets not yet supported")
        return
    }

    ip_length := int(Bytes_Ntohs(iphdr[2:4]))

    protocol := uint8(iphdr[9])
    if protocol != 6 {
        DEBUG("ip", "Not TCP packet received. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+iphl+20 {
        DEBUG("ip", "Packet too short to be TCP")
        return
    }

    tcphl := int((uint16(pkt.Data[l2hlen+iphl+12]) >> 4) * 4)
    if tcphl > 20 && len(pkt.Data) < l2hlen+iphl+tcphl {
        DEBUG("ip", "Packet too short for the TCP header")
        return
    }

    tcphdr := pkt.Data[l2hlen+iphl : l2hlen+iphl+tcphl]
    packet.tuple.Src_port = Bytes_Ntohs(tcphdr[0:2])
    packet.tuple.Dst_port = Bytes_Ntohs(tcphdr[2:4])

    //DEBUG(" %s:%d -> %s:%d", packet.tuple.src_ip, packet.tuple.src_port,
    //      packet.tuple.dst_ip, packet.tuple.dst_port)

    packet.payload = pkt.Data[l2hlen+iphl+tcphl : l2hlen+ip_length]

    FollowTcp(tcphdr, packet)
}

func DropPrivileges() error {
    var err error

    if !_ConfigMeta.IsDefined("runoptions", "uid") {
        // not found, no dropping privileges but no err
        return nil
    }

    if !_ConfigMeta.IsDefined("runoptions", "gid") {
        return MsgError("GID must be specified for dropping priviledges")
    }

    INFO("Switching to user: %d.%d", _Config.RunOptions.Uid, _Config.RunOptions.Gid)

    if err = syscall.Setgid(_Config.RunOptions.Gid); err != nil {
        return MsgError("setgid: %s", err.Error())
    }

    if err = syscall.Setuid(_Config.RunOptions.Uid); err != nil {
        return MsgError("setuid: %s", err.Error())
    }

    return nil
}

func main() {

    configfile := flag.String("c", "agent.conf", "Configuration file")
    file := flag.String("I", "", "file")
    loop := flag.Int("l", 0, "Loop file. 0 - loop forever")
    debugSelectorsStr := flag.String("d", "", "Enable certain debug selectors")
    oneAtAtime := flag.Bool("O", false, "Read packets one at a time (press Enter)")
    toStdout := flag.Bool("e", false, "Output to stdout instead of syslog")
    panicAfter := flag.Int("P", 0, "Panic after a given number of seconds, only for testing")

    flag.Parse()

    debugSelectors := []string{}
    logLevel := syslog.LOG_DEBUG
    if len(*debugSelectorsStr) > 0 {
        debugSelectors = strings.Split(*debugSelectorsStr, ",")
    }

    var h *pcap.Pcap
    var err error

    if _ConfigMeta, err = toml.DecodeFile(*configfile, &_Config); err != nil {
        fmt.Printf("TOML config parsing failed on %s: %s. Exiting.\n", *configfile, err)
        return
    }

    if len(debugSelectors) == 0 {
        debugSelectors = _Config.Logging.Selectors
    }
    LogInit(logLevel, "", !*toStdout, debugSelectors)

    if *panicAfter > 0 {
        go func() {
            time.Sleep(time.Duration(*panicAfter) * time.Second)
            panic("Panicking as requested")
        }()
    }

    if *file != "" {
        h, err = pcap.Openoffline(*file)
        if h == nil {
            ERR("Openoffline(%s) failed: %s", *file, err)
            return
        }
    } else {
        device := _Config.Interfaces.Device

        h, err = pcap.Openlive(device, 65535, true, 0)
        if h == nil {
            ERR("pcap.Openlive failed(%s): %s", device, err)
            return
        }
    }
    defer h.Close()

    if err = DropPrivileges(); err != nil {
        CRIT(err.Error())
        return
    }

    filter := configToFilter(&_Config)
    if filter != "" {
        DEBUG("pcapfilter", "Installing filter '%s'", filter)
        err := h.Setfilter(filter)
        if err != nil {
            ERR("pcap.Setfilter failed: %s", err)
            return
        }
    }

    if err = Publisher.Init(); err != nil {
        CRIT(err.Error())
        return
    }

    if err = procWatcher.Init(&_Config.Procs); err != nil {
        CRIT(err.Error())
        return
    }

    if err = TcpInit(); err != nil {
        CRIT(err.Error())
        return
    }

    datalink := h.Datalink()
    if datalink != pcap.LINKTYPE_ETHERNET && datalink != pcap.LINKTYPE_LINUX_SLL {
        WARN("Unsuported link type: %d", datalink)
    }

    _GeoLite, err = libgeo.Load("/usr/share/GeoIP/GeoIP.dat")
    if err != nil {
        WARN("Could not load GeoIP data: %s", err.Error())
    }

    counter := 0
    live := true
    loopCount := 1
    var lastPktTime *time.Time = nil
    for live {
        var pkt *pcap.Packet
        var res int32

        if *oneAtAtime {
            fmt.Println("Press enter to read packet")
            fmt.Scanln()
        }

        pkt, res = h.NextEx()
        switch res {
        case -1:
            ERR("pcap.NextEx() error: %s", h.Geterror())
            live = false
            continue
        case -2:
            DEBUG("pcapread", "End of file")
            loopCount += 1
            if *loop > 0 && loopCount > *loop {
                // give a bit of time to the publish goroutine
                // to flush
                time.Sleep(300 * time.Millisecond)
                live = false
            }
            // reopen the file
            h, err = pcap.Openoffline(*file)
            if h == nil {
                ERR("Openoffline(%s) failed: %s", *file, err)
                return
            }
            lastPktTime = nil
            continue
        case 0:
            // timeout
            continue
        }
        if res != 1 {
            panic(fmt.Sprintf("Unexpected return code from pcap.NextEx: %d", res))
        }

        if pkt == nil {
            panic("Nil packet despite res=1")
        }

        if *file != "" {
            if lastPktTime != nil {
                sleep := pkt.Time.Sub(*lastPktTime)
                if sleep > 0 {
                    time.Sleep(sleep)
                } else {
                    WARN("Time in pcap went back in time: %d", sleep)
                }
            }
            _lastPktTime := pkt.Time
            lastPktTime = &_lastPktTime
            pkt.Time = time.Now() // overwrite what we get from the pcap
        }
        counter++

        //DEBUG("pkt received")

        decodePktEth(datalink, pkt)
    }
    INFO("Input finish. Processed %d packets. Have a nice day!", counter)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/uttamgandhi24/whisper-go/whisper"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type MigrationData struct {
	wspFile      string
	relativePath string
	exportFile   string
	measurement  string
	tags         string
	field        string
	matched      bool
}

type TagKeyValue struct {
	Tagkey   string `json:"tagkey"`
	Tagvalue string `json:"tagvalue"`
}

type MigrationConfig struct {
	Pattern     string        `json:"pattern"`
	Measurement string        `json:"measurement"`
	Tags        []TagKeyValue `json:"tags"`
	Field       string        `json:"field"`
}

// Global vars
var migrationConfig []MigrationConfig
var exportedFileNumber = 0
var (
	wspPath           = flag.String("wsp-path", "", "Whisper files folder path.")
	exportPath        = flag.String("export-path", "", "Directory to export line protocol files.")
	configFile        = flag.String("config-file", "", "Configuration file for measurement and tags.")
	keepDirStructure  = flag.Bool("keep-dir-structure", false, "Keep whisper dir structure and filenames. Otherwise export files in a numbered manner.")
	databaseName      = flag.String("database", "", "Add the influxdb database context to exported files.")
	createDatabase    = flag.Bool("create-database", false, "Add a CREATE DATABASE query to exported files.")
)

func main() {
	flag.Parse()

	// List files
	fileList := []string{}
	listWspFiles(&fileList, *wspPath)

	// Open migration config file
	loadConfigFile(&migrationConfig, *configFile)

	// Go through wsp files and figure out the tags and measurements
	var migrationData []MigrationData
	for _, wspFile := range fileList {
		data := MigrationData{}
		data.wspFile = wspFile
		data.relativePath = strings.TrimPrefix(wspFile, *wspPath)

		if *keepDirStructure {
			data.exportFile = *exportPath + strings.Replace(data.relativePath, ".wsp", ".txt", -1)
		} else {
			exportedFileNumber += 1
			data.exportFile = *exportPath + fmt.Sprintf("/%08d.txt", exportedFileNumber)
		}

		// Assign the right measurment, field and tags
		data.assignConfig()

		if data.matched {
			migrationData = append(migrationData, data)
		} else {
			fmt.Println("File didn't match any config patterns: ", data.wspFile)
		}
	}

	// Set time range to export data
	from, _ := time.Parse("2006-01-02", "2000-01-01")
	until, _ := time.Parse("2006-01-02", "2100-01-01")

	// Go through wsp files
	for _, migration := range migrationData {
		// Open whisper file with driver
		w, err := whisper.Open(migration.wspFile)
		// fmt.Println("Preparing", migration.wspFile, ">>>", migration.exportFile, "(Size", w.Header.Archives[0].Size(), ")")
		check(err)
		// Get all points
		_, wspPoints, err := w.FetchUntilTime(from, until)
		w.Close()

		// Makes sure the directory exists
		os.MkdirAll(migration.exportFile[0:strings.LastIndex(migration.exportFile, "/")] ,0755);

		// Open file and prepare writer
		f, err := os.Create(migration.exportFile)
		check(err)
		writer := bufio.NewWriter(f)

		// Print the commands and context sections
		if *createDatabase {
			_, err = writer.WriteString("# DDL\nCREATE DATABASE " + *databaseName + "\n\n")
			check(err)
		}
		_, err = writer.WriteString("# DML\n# CONTEXT-DATABASE: " + *databaseName + "\n\n")
		check(err)

		// Print all points
		for _, point := range wspPoints {
			line := migration.lineprotocol(point) + "\n"

			_, err := writer.WriteString(line)
			check(err)
		}

		// Flush writer and close file
		writer.Flush()
		f.Close()

		fmt.Println("Exported:", migration.wspFile, " >>>", migration.exportFile)
	}
}

// Check errors
func check(e error) {
    if e != nil {
        panic(e)
    }
}

// Create the list of wsp files
func listWspFiles(fileList *[]string, searchDir string) {
	err := filepath.Walk(searchDir, func(path string, f os.FileInfo, err error) error {
		if os.IsNotExist(err) { //search dir does not exist
			return nil
		}
		// Only add wsp files to the list
		if strings.HasSuffix(f.Name(), "wsp") {
			*fileList = append(*fileList, path)
		}
		return nil
	})
	if err != nil {
		fmt.Println("Error listing files:")
		fmt.Println(err)
	}
}

// Read the config file and populate migrartionData.tagConfigs
func loadConfigFile(migrationConfig *[]MigrationConfig, filename string) error {
	raw, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
	return json.Unmarshal(raw, &migrationConfig)
}

// Generate the influxdb line protocol string for a given point
func (migrationData *MigrationData) lineprotocol(point whisper.Point) string {
	var line string
	line += migrationData.measurement
	line += migrationData.tags
	line += " "
	line += migrationData.field + "=" + strconv.FormatFloat(point.Value, 'f', -1, 64)
	line += " "
	line += strconv.FormatInt(int64(point.Timestamp), 10)
	return line
}

// Get measurement, tags and field by matching the whisper filename with a
// pattern in the config file
func (migrationData *MigrationData) assignConfig() {

	wspMeasurement := strings.TrimPrefix(migrationData.wspFile, *wspPath)
	wspMeasurement = strings.TrimSuffix(wspMeasurement, ".wsp")
	wspMeasurement = strings.Replace(wspMeasurement, "/", ".", -1)
	wspMeasurement = strings.Replace(wspMeasurement, ",", "_", -1)
	wspMeasurement = strings.Replace(wspMeasurement, " ", "_", -1)

	// Filename matching
	// TODO catch strings that don't match until the end
	var tagConfig MigrationConfig
	var matched []string
	var matchedArr [][]string
	var wildcards [][]string
	filenameMatched := false
	for _, tagConfig = range migrationConfig {
		reWild := regexp.MustCompile("{{\\s*([a-zA-Z0-9]+)\\s*}}")

		// Prepare regex pattern
		pattern := strings.Replace(tagConfig.Pattern, ".", "\\.", -1)
		pattern = reWild.ReplaceAllLiteralString(pattern, "([^.]+)")

		// List the matching values (Base and groups)
		re := regexp.MustCompile(pattern)

		matchedArr = re.FindAllStringSubmatch(wspMeasurement, -1)

		if matchedArr != nil {
			filenameMatched = true
			matched = matchedArr[0]
			// List of replacement wildcards like "{{ host }}"
			wildcards = reWild.FindAllStringSubmatch(tagConfig.Pattern, -1)
			break
		}
	}

	// Exit if there was no match
	if filenameMatched == false {
		return
	} else {
		migrationData.matched = true
	}

	// Fill the migrationData object
	migrationData.measurement = tagConfig.Measurement
	migrationData.field = tagConfig.Field
	for j := 0; j < len(tagConfig.Tags); j++ {
		migrationData.tags += ","
		migrationData.tags += tagConfig.Tags[j].Tagkey
		migrationData.tags += "="
		migrationData.tags += tagConfig.Tags[j].Tagvalue
	}

	// In case measurement and field aren't set in the config file
	// if tagConfig.Measurement != "" {
	// 	mtf.Measurement = tagConfig.Measurement
	// } else {
	// 	parts := strings.Split(wspMeasurement, ".")
	// 	mtf.Measurement = parts[len(parts)-1]
	// }
	// if tagConfig.Field != "" {
	// 	mtf.Field = tagConfig.Field
	// } else {
	// 	mtf.Field = "value"
	// }
	// mtf.Tags = make([]TagKeyValue, len(tagConfig.Tags))
	// copy(mtf.Tags, tagConfig.Tags)

	// Replace "{{ wildcard }}" with matched values in order
	// (reversed to avoid overlapping of bigger numbers)
	// TODO issues if matched contains "$n"
	for i := len(matched) - 1; i > 0; i-- {
		re := regexp.MustCompile("{{\\s" + wildcards[i - 1][1] + "\\s}}")

		migrationData.measurement = re.ReplaceAllLiteralString(migrationData.measurement, matched[i])
		migrationData.field = re.ReplaceAllLiteralString(migrationData.field, matched[i])
		migrationData.tags = re.ReplaceAllLiteralString(migrationData.tags, matched[i])
	}
}
<commit_msg>Add gzip export feature<commit_after>package main

import (
	"bufio"
	"compress/gzip"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/uttamgandhi24/whisper-go/whisper"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type MigrationData struct {
	wspFile      string
	relativePath string
	exportFile   string
	measurement  string
	tags         string
	field        string
	matched      bool
}

type TagKeyValue struct {
	Tagkey   string `json:"tagkey"`
	Tagvalue string `json:"tagvalue"`
}

type MigrationConfig struct {
	Pattern     string        `json:"pattern"`
	Measurement string        `json:"measurement"`
	Tags        []TagKeyValue `json:"tags"`
	Field       string        `json:"field"`
}

// Global vars
var migrationConfig []MigrationConfig
var exportedFileNumber = 0
var (
	wspPath           = flag.String("wsp-path", "", "Whisper files folder path.")
	exportPath        = flag.String("export-path", "", "Directory to export line protocol files.")
	configFile        = flag.String("config-file", "", "Configuration file for measurement and tags.")
	keepDirStructure  = flag.Bool("keep-dir-structure", false, "Keep whisper dir structure and filenames. Otherwise export files in a numbered manner.")
	databaseName      = flag.String("database", "", "Add the influxdb database context to exported files.")
	createDatabase    = flag.Bool("create-database", false, "Add a CREATE DATABASE query to exported files.")
	gzipped           = flag.Bool("gz", false, "Export data in a gzipped file")
)

func main() {
	flag.Parse()

	// List files
	fileList := []string{}
	listWspFiles(&fileList, *wspPath)

	// Open migration config file
	loadConfigFile(&migrationConfig, *configFile)

	// Go through wsp files and figure out the tags and measurements
	var migrationData []MigrationData
	for _, wspFile := range fileList {
		data := MigrationData{}
		data.wspFile = wspFile
		data.relativePath = strings.TrimPrefix(wspFile, *wspPath)

		// Figure exported filename
		if *keepDirStructure {
			data.exportFile = *exportPath + strings.Replace(data.relativePath, ".wsp", ".txt", -1)
		} else {
			exportedFileNumber += 1
			data.exportFile = *exportPath + fmt.Sprintf("/%08d.txt", exportedFileNumber)
		}

		// Add the .gz if the export needs to be compressed
		if *gzipped {
			data.exportFile += ".gz"
		}

		// Assign the right measurment, field and tags
		data.assignConfig()

		if data.matched {
			migrationData = append(migrationData, data)
		} else {
			fmt.Println("File didn't match any config patterns: ", data.wspFile)
		}
	}

	// Warning starting exporting
	fmt.Println("----------------")
	fmt.Println("Starting exporting", len(migrationData), "metrics to ", *exportPath)
	fmt.Println("----------------")

	// Set time range to export data
	from, _ := time.Parse("2006-01-02", "2000-01-01")
	until, _ := time.Parse("2006-01-02", "2100-01-01")

	// Go through wsp files and export data
	for _, migration := range migrationData {
		// Open whisper file with driver
		w, err := whisper.Open(migration.wspFile)
		// fmt.Println("Preparing", migration.wspFile, ">>>", migration.exportFile, "(Size", w.Header.Archives[0].Size(), ")")
		check(err)
		// Get all points
		_, wspPoints, err := w.FetchUntilTime(from, until)
		w.Close()

		// Makes sure the directory exists
		os.MkdirAll(migration.exportFile[0:strings.LastIndex(migration.exportFile, "/")] ,0755);

		// Open file and prepare writer
		f, err := os.Create(migration.exportFile)
		check(err)

		var gf *gzip.Writer
		var fw *bufio.Writer
		if *gzipped {
			gf = gzip.NewWriter(f)
			fw = bufio.NewWriter(gf)
		} else {
			fw = bufio.NewWriter(f)
		}

		// Print the commands and context sections
		if *createDatabase {
			_, err = fw.WriteString("# DDL\nCREATE DATABASE " + *databaseName + "\n\n")
			check(err)
		}
		_, err = fw.WriteString("# DML\n# CONTEXT-DATABASE: " + *databaseName + "\n\n")
		check(err)

		// Print all points
		for _, point := range wspPoints {
			line := migration.lineprotocol(point) + "\n"

			_, err := fw.WriteString(line)
			check(err)
		}

		// Flush writer and close file
		fw.Flush()
		if *gzipped {
			gf.Close()
		}
		f.Close()

		fmt.Println("Exported:", migration.exportFile, "from", migration.wspFile)
	}
}

// Check errors
func check(e error) {
    if e != nil {
        panic(e)
    }
}

// Create the list of wsp files
func listWspFiles(fileList *[]string, searchDir string) {
	err := filepath.Walk(searchDir, func(path string, f os.FileInfo, err error) error {
		if os.IsNotExist(err) { //search dir does not exist
			return nil
		}
		// Only add wsp files to the list
		if strings.HasSuffix(f.Name(), "wsp") {
			*fileList = append(*fileList, path)
		}
		return nil
	})
	if err != nil {
		fmt.Println("Error listing files:")
		fmt.Println(err)
	}
}

// Read the config file and populate migrartionData.tagConfigs
func loadConfigFile(migrationConfig *[]MigrationConfig, filename string) error {
	raw, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
	return json.Unmarshal(raw, &migrationConfig)
}

// Generate the influxdb line protocol string for a given point
func (migrationData *MigrationData) lineprotocol(point whisper.Point) string {
	var line string
	line += migrationData.measurement
	line += migrationData.tags
	line += " "
	line += migrationData.field + "=" + strconv.FormatFloat(point.Value, 'f', -1, 64)
	line += " "
	line += strconv.FormatInt(int64(point.Timestamp), 10)
	return line
}

// Get measurement, tags and field by matching the whisper filename with a
// pattern in the config file
func (migrationData *MigrationData) assignConfig() {

	wspMeasurement := strings.TrimPrefix(migrationData.wspFile, *wspPath)
	wspMeasurement = strings.TrimSuffix(wspMeasurement, ".wsp")
	wspMeasurement = strings.Replace(wspMeasurement, "/", ".", -1)
	wspMeasurement = strings.Replace(wspMeasurement, ",", "_", -1)
	wspMeasurement = strings.Replace(wspMeasurement, " ", "_", -1)

	// Filename matching
	// TODO catch strings that don't match until the end
	var tagConfig MigrationConfig
	var matched []string
	var matchedArr [][]string
	var wildcards [][]string
	filenameMatched := false
	for _, tagConfig = range migrationConfig {
		reWild := regexp.MustCompile("{{\\s*([a-zA-Z0-9]+)\\s*}}")

		// Prepare regex pattern
		pattern := strings.Replace(tagConfig.Pattern, ".", "\\.", -1)
		pattern = reWild.ReplaceAllLiteralString(pattern, "([^.]+)")

		// List the matching values (Base and groups)
		re := regexp.MustCompile(pattern)

		matchedArr = re.FindAllStringSubmatch(wspMeasurement, -1)

		if matchedArr != nil {
			filenameMatched = true
			matched = matchedArr[0]
			// List of replacement wildcards like "{{ host }}"
			wildcards = reWild.FindAllStringSubmatch(tagConfig.Pattern, -1)
			break
		}
	}

	// Exit if there was no match
	if filenameMatched == false {
		return
	} else {
		migrationData.matched = true
	}

	// Fill the migrationData object
	migrationData.measurement = tagConfig.Measurement
	migrationData.field = tagConfig.Field
	for j := 0; j < len(tagConfig.Tags); j++ {
		migrationData.tags += ","
		migrationData.tags += tagConfig.Tags[j].Tagkey
		migrationData.tags += "="
		migrationData.tags += tagConfig.Tags[j].Tagvalue
	}

	// In case measurement and field aren't set in the config file
	// if tagConfig.Measurement != "" {
	// 	mtf.Measurement = tagConfig.Measurement
	// } else {
	// 	parts := strings.Split(wspMeasurement, ".")
	// 	mtf.Measurement = parts[len(parts)-1]
	// }
	// if tagConfig.Field != "" {
	// 	mtf.Field = tagConfig.Field
	// } else {
	// 	mtf.Field = "value"
	// }
	// mtf.Tags = make([]TagKeyValue, len(tagConfig.Tags))
	// copy(mtf.Tags, tagConfig.Tags)

	// Replace "{{ wildcard }}" with matched values in order
	// (reversed to avoid overlapping of bigger numbers)
	// TODO issues if matched contains "$n"
	for i := len(matched) - 1; i > 0; i-- {
		re := regexp.MustCompile("{{\\s" + wildcards[i - 1][1] + "\\s}}")

		migrationData.measurement = re.ReplaceAllLiteralString(migrationData.measurement, matched[i])
		migrationData.field = re.ReplaceAllLiteralString(migrationData.field, matched[i])
		migrationData.tags = re.ReplaceAllLiteralString(migrationData.tags, matched[i])
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/voxelbrain/goptions"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"os"
	"regexp"
	"strings"
)

// VERSION holds the Current version of spruce
var VERSION = "0.0.0" // SED MARKER FOR AUTO VERSION BUMPING
// BUILD holds CURRENT BUILD OF SPRUCE
var BUILD = "master" // updated by build.sh
// DIRTY holds Whether any uncommitted changes were found in the working copy
var DIRTY = "" // updated by build.sh

var debug bool
var trace bool

// SkipSSLValidation ...
var SkipSSLValidation bool

// Index ...
type Index struct {
	PrimaryShards []Shard
	ReplicaShards []Shard
	Name          string
}

// Node ...
type Node struct {
	Name string
	ID   string
}

// Shard ...
type Shard struct {
	Index      int
	Node       string
	Relocating string
	Status     string
}

// DEBUG ...
func DEBUG(format string, args ...interface{}) {
	if debug {
		content := fmt.Sprintf(format, args...)
		lines := strings.Split(content, "\n")
		for i, line := range lines {
			lines[i] = "DEBUG> " + line
		}
		content = strings.Join(lines, "\n")
		fmt.Fprintf(os.Stderr, "%s\n", content)
	}
}

// TRACE ...
func TRACE(format string, args ...interface{}) {
	if trace {
		DEBUG(format, args...)
	}
}

// ERROR ...
func ERROR(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "%s\n", fmt.Sprintf(format, args...))
}

// FATAL ...
func FATAL(format string, args ...interface{}) {
	ERROR(format, args...)
	os.Exit(1)
}

func trueString(str string) bool {
	re := regexp.MustCompile("^(?i:no|false|0|off)$")
	if str == "" || re.Match([]byte(str)) {
		return false
	}
	return true
}

func main() {
	var options struct {
		Debug         bool   `goptions:"-D, --debug, description='Enable debugging'"`
		Trace         bool   `goptions:"--trace, description='Enable API call tracing'"`
		Fix           bool   `goptions:"-f, --fix, description='Fix any issues detected'"`
		Host          string `goptions:"-H, --elasticsearch_url, description='ElasticSearch URL. Defaults to http://localhost:9200'"`
		SkipSSLVerify bool   `goptions:"-k, --skip-ssl-validation, description='Disable SSL certificate checking'"`
		Help          bool   `goptions:"-h, --help"`
		Version       bool   `goptions:"-v, --version"`
	}

	goptions.ParseAndFail(&options)
	if options.Help {
		goptions.PrintHelp()
		os.Exit(0)
	}

	if options.Debug {
		debug = true
	}
	if trueString(os.Getenv("DEBUG")) {
		debug = true
	}

	if options.Trace {
		trace = true
		debug = true
	}
	if trueString(os.Getenv("TRACE")) {
		trace = true
		debug = true
	}

	if options.Version {
		plus := ""
		if BUILD != "release" {
			plus = "+"
		}
		fmt.Fprintf(os.Stderr, "%s - Version %s%s (%s%s)\n", os.Args[0], VERSION, plus, BUILD, DIRTY)
		os.Exit(0)
	}

	if options.SkipSSLVerify {
		SkipSSLValidation = true
	}

	if options.Host == "" {
		options.Host = "http://localhost:9200/"
	}

	DEBUG("Checking ElasticSearch on %s", options.Host)

	nodes, err := getDataNodes(options.Host)
	if err != nil {
		FATAL(err.Error())
	}

	indices, err := getIndices(options.Host)
	if err != nil {
		FATAL(err.Error())
	}

	TRACE("Indices: %v\nNodes: %v", indices, nodes)
	for _, index := range indices {
		fmt.Printf("Checking index %s\n", index.Name)
		availNodeMaps := make([]map[string]Node, len(index.PrimaryShards))
		for i := range availNodeMaps {
			availNodeMaps[i] = map[string]Node{}
			for _, node := range nodes {
				availNodeMaps[i][node.ID] = node
			}
		}

		for _, shard := range index.PrimaryShards {
			delete(availNodeMaps[shard.Index], shard.Node)
			delete(availNodeMaps[shard.Index], shard.Relocating)
		}

		for _, shard := range index.ReplicaShards {
			delete(availNodeMaps[shard.Index], shard.Node)
			delete(availNodeMaps[shard.Index], shard.Relocating)
		}

		availNodesList := make([][]Node, len(index.PrimaryShards))
		for i := range availNodesList {
			var availNodes []Node
			for _, n := range availNodeMaps[i] {
				availNodes = append(availNodes, n)
			}
			availNodesList[i] = availNodes
		}

		for _, shard := range index.PrimaryShards {
			if shard.Status == "UNASSIGNED" {
				availNodes := availNodesList[shard.Index]
				if len(availNodes) > 0 {
					var node Node
					node, availNodes = availNodes[0], availNodes[1:]
					fmt.Printf("    %s shard %d primary is UNASSIGNED.\n", index.Name, shard.Index)
					DEBUG("Would assign %s as new node", node.Name)
					if options.Fix {
						fmt.Printf("    - bringing up %s shard %d primary on %s\n", index.Name, shard.Index, node.Name)
						err := rerouteShard(options.Host, index.Name, shard.Index, node.Name, true)
						if err != nil {
							ERROR(err.Error())
						}
					}
				} else {
					fmt.Printf("    %s shard %d primary is UNASSIGNED, and no nodes are available to assign it to\n", index.Name, shard.Index)
				}
			}
		}

		for _, shard := range index.ReplicaShards {
			if shard.Status == "UNASSIGNED" {
				availNodes := availNodesList[shard.Index]
				if len(availNodes) > 0 {
					var node Node
					node, availNodes = availNodes[0], availNodes[1:]
					fmt.Printf("    %s shard %d replica is UNASSIGNED.\n", index.Name, shard.Index)
					DEBUG("Would assign %s as new node", node.Name)
					if options.Fix {
						fmt.Printf("    - bringing up %s shard %d replica on %s\n", index.Name, shard.Index, node.Name)
						err := rerouteShard(options.Host, index.Name, shard.Index, node.Name, false)
						if err != nil {
							ERROR(err.Error())
						}
					}
				} else {
					fmt.Printf("    %s shard %d replica is UNASSIGNED, and no nodes are available to assign it to\n", index.Name, shard.Index)
				}
			}
		}
	}
}

func httpRequest(method string, url string, data io.Reader) ([]byte, error) {

	req, err := http.NewRequest(method, url, data)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/json")
	dumpReq, err := httputil.DumpRequest(req, true)
	if err != nil {
		return nil, err
	}
	TRACE("HTTP Request:\n%s", dumpReq)

	client := &http.Client{}
	client.Transport = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: SkipSSLValidation}}

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	dumpResp, err := httputil.DumpResponse(resp, true)
	if err != nil {
		return nil, err
	}
	TRACE("HTTP Response:\n%s", dumpResp)

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("%s returned %d: %s", url, resp.StatusCode, body)
	}

	return body, nil
}

func getDataNodes(host string) ([]Node, error) {
	data, err := httpRequest("GET", host+"/_nodes?pretty", nil)
	if err != nil {
		return []Node{}, err
	}

	var o interface{}
	err = json.Unmarshal(data, &o)
	if err != nil {
		return []Node{}, err
	}

	var nodelist []Node
	if obj, ok := o.(map[string]interface{}); ok {
		if nodes, ok := obj["nodes"].(map[string]interface{}); ok {
			for id, n := range nodes {
				if node, ok := n.(map[string]interface{}); ok {
					if settings, ok := node["settings"].(map[string]interface{}); ok {
						if nodeSettings, ok := settings["node"].(map[string]interface{}); ok {
							var name string
							if name, ok = nodeSettings["name"].(string); !ok {
								return []Node{}, fmt.Errorf("Could not detect node name for node %s", id)
							}
							if dataNode, ok := nodeSettings["data"].(string); ok {
								if dataNode == "true" {
									nodelist = append(nodelist, Node{Name: name, ID: id})
								}
							} else {
								return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s.settings.node.data` key", id)
							}
						} else {
							return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s.settings.node` key", id)
						}
					} else {
						return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s.settings` key", id)
					}
				} else {
					return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s` key", id)
				}
			}
		} else {
			return []Node{}, fmt.Errorf("Unexpected data type for `nodes` key")
		}
	} else {
		return []Node{}, fmt.Errorf("Unexpected data type returned from `GET /_nodes`")
	}

	return nodelist, nil
}

func rerouteShard(host string, index string, shard int, node string, primary bool) error {
	data := fmt.Sprintf(`{"commands":[{"allocate":{"index":"%s","shard":%d,"node":"%s","allow_primary":%t}}]}`, index, shard, node, primary)

	dataBuf := bytes.NewBuffer([]byte(data))
	_, err := httpRequest("POST", host+"/_cluster/reroute", dataBuf)
	return err
}

func getIndices(host string) ([]Index, error) {
	data, err := httpRequest("GET", host+"/_cluster/state?pretty", nil)
	if err != nil {
		return []Index{}, err
	}

	var o interface{}
	err = json.Unmarshal(data, &o)
	if err != nil {
		return []Index{}, err
	}

	var indexlist []Index
	if obj, ok := o.(map[string]interface{}); ok {
		if rt, ok := obj["routing_table"].(map[string]interface{}); ok {
			// indices
			if indices, ok := rt["indices"].(map[string]interface{}); ok {
				for indexName, i := range indices {
					if indexData, ok := i.(map[string]interface{}); ok {
						index := Index{Name: indexName}
						if shards, ok := indexData["shards"].(map[string]interface{}); ok {
							index.PrimaryShards = make([]Shard, len(shards))
							index.ReplicaShards = make([]Shard, len(shards))
							for shardKey, s := range shards {
								if shardList, ok := s.([]interface{}); ok {
									for i, s := range shardList {
										if shard, ok := s.(map[string]interface{}); ok {

											var primary bool
											var state string
											var node string
											var relocating string
											var shardFloat float64
											shardFloat = -1.0

											if primary, ok = shard["primary"].(bool); !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].primary", indexName, shardKey, i)
											}

											if state, ok = shard["state"].(string); !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].state", indexName, shardKey, i)
											}
											if node, ok = shard["node"].(string); shard["node"] != nil && !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].node", indexName, shardKey, i)
											}
											if relocating, ok = shard["relocating_node"].(string); shard["relocating"] != nil && !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].relocating_node", indexName, shardKey, i)
											}
											if shardFloat, ok = shard["shard"].(float64); !ok {
												return []Index{}, fmt.Errorf("Could not parse 'routing_table.indices.%s.shards.%s.[%d].shard", indexName, shardKey, i)
											}
											shardIndex := int(shardFloat)

											parsedShard := Shard{
												Index:      shardIndex,
												Status:     state,
												Node:       node,
												Relocating: relocating,
											}

											if primary {
												index.PrimaryShards[shardIndex] = parsedShard
											} else {
												index.ReplicaShards[shardIndex] = parsedShard
											}
										} else {
											return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s.shards.%s.[%d]", indexName, shardKey, i)
										}
									}
								} else {
									return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s.shards.%s", indexName, shardKey)
								}
							}
						} else {
							return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s.shards", indexName)
						}
						indexlist = append(indexlist, index)
					} else {
						return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s`", indexName)
					}
				}
			} else {
				return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices`")
			}
		} else {
			return []Index{}, fmt.Errorf("Unexpected data type for `routing_table`")
		}
	} else {
		return []Index{}, fmt.Errorf("Unexpected data type returned from `GET /_status`")
	}

	return indexlist, nil
}
<commit_msg>update release version to v0.1.0<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/voxelbrain/goptions"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"os"
	"regexp"
	"strings"
)

// VERSION holds the Current version of spruce
var VERSION = "0.1.0" // SED MARKER FOR AUTO VERSION BUMPING
// BUILD holds CURRENT BUILD OF SPRUCE
var BUILD = "master" // updated by build.sh
// DIRTY holds Whether any uncommitted changes were found in the working copy
var DIRTY = "" // updated by build.sh

var debug bool
var trace bool

// SkipSSLValidation ...
var SkipSSLValidation bool

// Index ...
type Index struct {
	PrimaryShards []Shard
	ReplicaShards []Shard
	Name          string
}

// Node ...
type Node struct {
	Name string
	ID   string
}

// Shard ...
type Shard struct {
	Index      int
	Node       string
	Relocating string
	Status     string
}

// DEBUG ...
func DEBUG(format string, args ...interface{}) {
	if debug {
		content := fmt.Sprintf(format, args...)
		lines := strings.Split(content, "\n")
		for i, line := range lines {
			lines[i] = "DEBUG> " + line
		}
		content = strings.Join(lines, "\n")
		fmt.Fprintf(os.Stderr, "%s\n", content)
	}
}

// TRACE ...
func TRACE(format string, args ...interface{}) {
	if trace {
		DEBUG(format, args...)
	}
}

// ERROR ...
func ERROR(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "%s\n", fmt.Sprintf(format, args...))
}

// FATAL ...
func FATAL(format string, args ...interface{}) {
	ERROR(format, args...)
	os.Exit(1)
}

func trueString(str string) bool {
	re := regexp.MustCompile("^(?i:no|false|0|off)$")
	if str == "" || re.Match([]byte(str)) {
		return false
	}
	return true
}

func main() {
	var options struct {
		Debug         bool   `goptions:"-D, --debug, description='Enable debugging'"`
		Trace         bool   `goptions:"--trace, description='Enable API call tracing'"`
		Fix           bool   `goptions:"-f, --fix, description='Fix any issues detected'"`
		Host          string `goptions:"-H, --elasticsearch_url, description='ElasticSearch URL. Defaults to http://localhost:9200'"`
		SkipSSLVerify bool   `goptions:"-k, --skip-ssl-validation, description='Disable SSL certificate checking'"`
		Help          bool   `goptions:"-h, --help"`
		Version       bool   `goptions:"-v, --version"`
	}

	goptions.ParseAndFail(&options)
	if options.Help {
		goptions.PrintHelp()
		os.Exit(0)
	}

	if options.Debug {
		debug = true
	}
	if trueString(os.Getenv("DEBUG")) {
		debug = true
	}

	if options.Trace {
		trace = true
		debug = true
	}
	if trueString(os.Getenv("TRACE")) {
		trace = true
		debug = true
	}

	if options.Version {
		plus := ""
		if BUILD != "release" {
			plus = "+"
		}
		fmt.Fprintf(os.Stderr, "%s - Version %s%s (%s%s)\n", os.Args[0], VERSION, plus, BUILD, DIRTY)
		os.Exit(0)
	}

	if options.SkipSSLVerify {
		SkipSSLValidation = true
	}

	if options.Host == "" {
		options.Host = "http://localhost:9200/"
	}

	DEBUG("Checking ElasticSearch on %s", options.Host)

	nodes, err := getDataNodes(options.Host)
	if err != nil {
		FATAL(err.Error())
	}

	indices, err := getIndices(options.Host)
	if err != nil {
		FATAL(err.Error())
	}

	TRACE("Indices: %v\nNodes: %v", indices, nodes)
	for _, index := range indices {
		fmt.Printf("Checking index %s\n", index.Name)
		availNodeMaps := make([]map[string]Node, len(index.PrimaryShards))
		for i := range availNodeMaps {
			availNodeMaps[i] = map[string]Node{}
			for _, node := range nodes {
				availNodeMaps[i][node.ID] = node
			}
		}

		for _, shard := range index.PrimaryShards {
			delete(availNodeMaps[shard.Index], shard.Node)
			delete(availNodeMaps[shard.Index], shard.Relocating)
		}

		for _, shard := range index.ReplicaShards {
			delete(availNodeMaps[shard.Index], shard.Node)
			delete(availNodeMaps[shard.Index], shard.Relocating)
		}

		availNodesList := make([][]Node, len(index.PrimaryShards))
		for i := range availNodesList {
			var availNodes []Node
			for _, n := range availNodeMaps[i] {
				availNodes = append(availNodes, n)
			}
			availNodesList[i] = availNodes
		}

		for _, shard := range index.PrimaryShards {
			if shard.Status == "UNASSIGNED" {
				availNodes := availNodesList[shard.Index]
				if len(availNodes) > 0 {
					var node Node
					node, availNodes = availNodes[0], availNodes[1:]
					fmt.Printf("    %s shard %d primary is UNASSIGNED.\n", index.Name, shard.Index)
					DEBUG("Would assign %s as new node", node.Name)
					if options.Fix {
						fmt.Printf("    - bringing up %s shard %d primary on %s\n", index.Name, shard.Index, node.Name)
						err := rerouteShard(options.Host, index.Name, shard.Index, node.Name, true)
						if err != nil {
							ERROR(err.Error())
						}
					}
				} else {
					fmt.Printf("    %s shard %d primary is UNASSIGNED, and no nodes are available to assign it to\n", index.Name, shard.Index)
				}
			}
		}

		for _, shard := range index.ReplicaShards {
			if shard.Status == "UNASSIGNED" {
				availNodes := availNodesList[shard.Index]
				if len(availNodes) > 0 {
					var node Node
					node, availNodes = availNodes[0], availNodes[1:]
					fmt.Printf("    %s shard %d replica is UNASSIGNED.\n", index.Name, shard.Index)
					DEBUG("Would assign %s as new node", node.Name)
					if options.Fix {
						fmt.Printf("    - bringing up %s shard %d replica on %s\n", index.Name, shard.Index, node.Name)
						err := rerouteShard(options.Host, index.Name, shard.Index, node.Name, false)
						if err != nil {
							ERROR(err.Error())
						}
					}
				} else {
					fmt.Printf("    %s shard %d replica is UNASSIGNED, and no nodes are available to assign it to\n", index.Name, shard.Index)
				}
			}
		}
	}
}

func httpRequest(method string, url string, data io.Reader) ([]byte, error) {

	req, err := http.NewRequest(method, url, data)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/json")
	dumpReq, err := httputil.DumpRequest(req, true)
	if err != nil {
		return nil, err
	}
	TRACE("HTTP Request:\n%s", dumpReq)

	client := &http.Client{}
	client.Transport = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: SkipSSLValidation}}

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	dumpResp, err := httputil.DumpResponse(resp, true)
	if err != nil {
		return nil, err
	}
	TRACE("HTTP Response:\n%s", dumpResp)

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("%s returned %d: %s", url, resp.StatusCode, body)
	}

	return body, nil
}

func getDataNodes(host string) ([]Node, error) {
	data, err := httpRequest("GET", host+"/_nodes?pretty", nil)
	if err != nil {
		return []Node{}, err
	}

	var o interface{}
	err = json.Unmarshal(data, &o)
	if err != nil {
		return []Node{}, err
	}

	var nodelist []Node
	if obj, ok := o.(map[string]interface{}); ok {
		if nodes, ok := obj["nodes"].(map[string]interface{}); ok {
			for id, n := range nodes {
				if node, ok := n.(map[string]interface{}); ok {
					if settings, ok := node["settings"].(map[string]interface{}); ok {
						if nodeSettings, ok := settings["node"].(map[string]interface{}); ok {
							var name string
							if name, ok = nodeSettings["name"].(string); !ok {
								return []Node{}, fmt.Errorf("Could not detect node name for node %s", id)
							}
							if dataNode, ok := nodeSettings["data"].(string); ok {
								if dataNode == "true" {
									nodelist = append(nodelist, Node{Name: name, ID: id})
								}
							} else {
								return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s.settings.node.data` key", id)
							}
						} else {
							return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s.settings.node` key", id)
						}
					} else {
						return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s.settings` key", id)
					}
				} else {
					return []Node{}, fmt.Errorf("Unexpected data type for `nodes.%s` key", id)
				}
			}
		} else {
			return []Node{}, fmt.Errorf("Unexpected data type for `nodes` key")
		}
	} else {
		return []Node{}, fmt.Errorf("Unexpected data type returned from `GET /_nodes`")
	}

	return nodelist, nil
}

func rerouteShard(host string, index string, shard int, node string, primary bool) error {
	data := fmt.Sprintf(`{"commands":[{"allocate":{"index":"%s","shard":%d,"node":"%s","allow_primary":%t}}]}`, index, shard, node, primary)

	dataBuf := bytes.NewBuffer([]byte(data))
	_, err := httpRequest("POST", host+"/_cluster/reroute", dataBuf)
	return err
}

func getIndices(host string) ([]Index, error) {
	data, err := httpRequest("GET", host+"/_cluster/state?pretty", nil)
	if err != nil {
		return []Index{}, err
	}

	var o interface{}
	err = json.Unmarshal(data, &o)
	if err != nil {
		return []Index{}, err
	}

	var indexlist []Index
	if obj, ok := o.(map[string]interface{}); ok {
		if rt, ok := obj["routing_table"].(map[string]interface{}); ok {
			// indices
			if indices, ok := rt["indices"].(map[string]interface{}); ok {
				for indexName, i := range indices {
					if indexData, ok := i.(map[string]interface{}); ok {
						index := Index{Name: indexName}
						if shards, ok := indexData["shards"].(map[string]interface{}); ok {
							index.PrimaryShards = make([]Shard, len(shards))
							index.ReplicaShards = make([]Shard, len(shards))
							for shardKey, s := range shards {
								if shardList, ok := s.([]interface{}); ok {
									for i, s := range shardList {
										if shard, ok := s.(map[string]interface{}); ok {

											var primary bool
											var state string
											var node string
											var relocating string
											var shardFloat float64
											shardFloat = -1.0

											if primary, ok = shard["primary"].(bool); !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].primary", indexName, shardKey, i)
											}

											if state, ok = shard["state"].(string); !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].state", indexName, shardKey, i)
											}
											if node, ok = shard["node"].(string); shard["node"] != nil && !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].node", indexName, shardKey, i)
											}
											if relocating, ok = shard["relocating_node"].(string); shard["relocating"] != nil && !ok {
												return []Index{}, fmt.Errorf("Could not parse `routing_table.indices.%s.shards.%s.[%d].relocating_node", indexName, shardKey, i)
											}
											if shardFloat, ok = shard["shard"].(float64); !ok {
												return []Index{}, fmt.Errorf("Could not parse 'routing_table.indices.%s.shards.%s.[%d].shard", indexName, shardKey, i)
											}
											shardIndex := int(shardFloat)

											parsedShard := Shard{
												Index:      shardIndex,
												Status:     state,
												Node:       node,
												Relocating: relocating,
											}

											if primary {
												index.PrimaryShards[shardIndex] = parsedShard
											} else {
												index.ReplicaShards[shardIndex] = parsedShard
											}
										} else {
											return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s.shards.%s.[%d]", indexName, shardKey, i)
										}
									}
								} else {
									return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s.shards.%s", indexName, shardKey)
								}
							}
						} else {
							return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s.shards", indexName)
						}
						indexlist = append(indexlist, index)
					} else {
						return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices.%s`", indexName)
					}
				}
			} else {
				return []Index{}, fmt.Errorf("Unexpected data type for `routing_table.indices`")
			}
		} else {
			return []Index{}, fmt.Errorf("Unexpected data type for `routing_table`")
		}
	} else {
		return []Index{}, fmt.Errorf("Unexpected data type returned from `GET /_status`")
	}

	return indexlist, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"runtime"
)

func main() {
	opts, err := ParseOpts()
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}

	addr := fmt.Sprintf("%s:%d", opts.host, opts.port)
	url := fmt.Sprintf("http://%s/", addr)

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

	// See: http://stackoverflow.com/questions/32738188/go-how-can-i-start-the-browser-after-the-server-started-listening
	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	log.Println("Kifu Notebook is running at:", url)
	if !opts.noBrowser {
		open(url)
	}

	if err := http.Serve(l, nil); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
}

// See: http://stackoverflow.com/questions/39320371/how-start-web-server-to-open-page-in-browser-in-golang
func open(url string) error {
	var cmd string
	var args []string

	switch runtime.GOOS {
	case "windows":
		cmd = "cmd"
		args = []string{"/c", "start"}
	case "darwin":
		cmd = "open"
	default: // "linux", "freebsd", "openbsd", "netbsd"
		cmd = "xdg-open"
	}
	args = append(args, url)
	return exec.Command(cmd, args...).Start()
}
<commit_msg>Use go-assets-builder<commit_after>package main

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"runtime"
)

//go:generate go-assets-builder --strip-prefix="/build" -o bindata.go build

func main() {
	opts, err := ParseOpts()
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}

	addr := fmt.Sprintf("%s:%d", opts.host, opts.port)
	url := fmt.Sprintf("http://%s/", addr)

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

	// See: http://stackoverflow.com/questions/32738188/go-how-can-i-start-the-browser-after-the-server-started-listening
	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	log.Println("Kifu Notebook is running at:", url)
	if !opts.noBrowser {
		open(url)
	}

	if err := http.Serve(l, nil); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
}

// See: http://stackoverflow.com/questions/39320371/how-start-web-server-to-open-page-in-browser-in-golang
func open(url string) error {
	var cmd string
	var args []string

	switch runtime.GOOS {
	case "windows":
		cmd = "cmd"
		args = []string{"/c", "start"}
	case "darwin":
		cmd = "open"
	default: // "linux", "freebsd", "openbsd", "netbsd"
		cmd = "xdg-open"
	}
	args = append(args, url)
	return exec.Command(cmd, args...).Start()
}
<|endoftext|>


<commit_before>package governor

import (
	"fmt"
	"github.com/fatih/color"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"net/url"
	"strings"
)

const (
	banner = `
  __ _  _____   _____ _ __ _ __   ___  _ __
 / _' |/ _ \ \ / / _ \ '__| '_ \ / _ \| '__|
| (_| | (_) \ V /  __/ |  | | | | (_) | |
 \__. |\___/ \_/ \___|_|  |_| |_|\___/|_|
  __/ |
 |___/  %s

 %s
 %s
`
)

type (
	// Server is an http gateway
	Server struct {
		i          *echo.Echo
		logger     Logger
		h          *health
		s          *setup
		config     Config
		showBanner bool
	}
)

const (
	moduleIDServer = "govserver"
)

// New creates a new Server
func New(config Config, l Logger) (*Server, error) {
	i := echo.New()
	i.HideBanner = true
	i.HTTPErrorHandler = errorHandler(i, l)
	l.Info("initialized error handling", moduleIDServer, "initialize error handler", 0, nil)
	i.Binder = requestBinder()
	l.Info("added custom request binder", moduleIDServer, "initialize request binder", 0, nil)
	i.Pre(middleware.RemoveTrailingSlash())
	if len(config.RouteRewrite) > 0 {
		l.Info("adding route rewriting rules", moduleIDServer, "add route rewrite rules", 0, config.RouteRewrite)
		i.Pre(middleware.Rewrite(config.RouteRewrite))
	}

	if config.IsDebug() {
		i.Use(middleware.LoggerWithConfig(middleware.LoggerConfig{
			Format: "time=${time_rfc3339}, method=${method}, uri=${uri}, status=${status}, latency=${latency_human}\n",
		}))
	}

	i.Use(middleware.Gzip())

	if len(config.Origins) > 0 {
		i.Use(middleware.CORSWithConfig(middleware.CORSConfig{
			AllowOrigins:     config.Origins,
			AllowCredentials: true,
		}))
	}

	i.Use(middleware.BodyLimit(config.MaxReqSize))
	i.Use(middleware.Recover())

	apiMiddlewareSkipper := func(c echo.Context) bool {
		path := c.Request().URL.EscapedPath()
		return strings.HasPrefix(path, config.BaseURL+"/") || config.BaseURL == path
	}
	if len(config.FrontendProxy) > 0 {
		targets := make([]*middleware.ProxyTarget, 0, len(config.FrontendProxy))
		for _, i := range config.FrontendProxy {
			if u, err := url.Parse(i); err == nil {
				targets = append(targets, &middleware.ProxyTarget{
					URL: u,
				})
			} else {
				l.Warn("could not add frontend proxy", moduleIDServer, "fail add proxy", 0, map[string]string{
					"proxy": i,
					"error": err.Error(),
				})
			}
		}
		if len(targets) > 0 {
			i.Use(middleware.ProxyWithConfig(middleware.ProxyConfig{
				Balancer: middleware.NewRoundRobinBalancer(targets),
				Skipper:  apiMiddlewareSkipper,
			}))
		}
	} else {
		i.Use(middleware.StaticWithConfig(middleware.StaticConfig{
			Root:    config.PublicDir,
			Index:   "index.html",
			Browse:  false,
			HTML5:   true,
			Skipper: apiMiddlewareSkipper,
		}))
	}

	i.Use(middleware.RequestID())
	l.Info("initialized middleware", moduleIDServer, "initialize middleware", 0, nil)

	healthService := newHealth()
	if err := healthService.Mount(config, l, i.Group(config.BaseURL+"/healthz")); err != nil {
		return nil, err
	}
	setupService := newSetup()
	if err := setupService.Mount(config, l, i.Group(config.BaseURL+"/setupz")); err != nil {
		return nil, err
	}

	s := &Server{
		i:          i,
		logger:     l,
		config:     config,
		h:          healthService,
		s:          setupService,
		showBanner: true,
	}

	l.Info("server instance created", moduleIDServer, "initialize server instance", 0, nil)
	return s, nil
}

// Start starts the server at the specified port
func (s *Server) Start() error {
	if s.showBanner {
		fmt.Printf(color.BlueString(banner+"\n"), color.GreenString(s.config.Version), "build version:"+color.GreenString(s.config.VersionHash), "http server on "+color.RedString(":"+s.config.Port))
	}
	s.i.Logger.Fatal(s.i.Start(":" + s.config.Port))
	return nil
}

// Must ensures that the operation must succeed
func Must(err error) {
	if err != nil {
		fmt.Println(err)
		panic(err)
	}
}
<commit_msg>Fix rule rewrite regex<commit_after>package governor

import (
	"fmt"
	"github.com/fatih/color"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"net/url"
	"strings"
)

const (
	banner = `
  __ _  _____   _____ _ __ _ __   ___  _ __
 / _' |/ _ \ \ / / _ \ '__| '_ \ / _ \| '__|
| (_| | (_) \ V /  __/ |  | | | | (_) | |
 \__. |\___/ \_/ \___|_|  |_| |_|\___/|_|
  __/ |
 |___/  %s

 %s
 %s
`
)

type (
	// Server is an http gateway
	Server struct {
		i          *echo.Echo
		logger     Logger
		h          *health
		s          *setup
		config     Config
		showBanner bool
	}
)

const (
	moduleIDServer = "govserver"
)

// New creates a new Server
func New(config Config, l Logger) (*Server, error) {
	i := echo.New()
	i.HideBanner = true
	i.HTTPErrorHandler = errorHandler(i, l)
	l.Info("initialized error handling", moduleIDServer, "initialize error handler", 0, nil)
	i.Binder = requestBinder()
	l.Info("added custom request binder", moduleIDServer, "initialize request binder", 0, nil)
	i.Pre(middleware.RemoveTrailingSlash())
	if len(config.RouteRewrite) > 0 {
		l.Info("adding route rewriting rules", moduleIDServer, "add route rewrite rules", 0, config.RouteRewrite)
		rewriteRules := make(map[string]string, len(config.RouteRewrite))
		for k, v := range config.RouteRewrite {
			rewriteRules["^"+k] = v
		}
		i.Pre(middleware.Rewrite(rewriteRules))
	}

	if config.IsDebug() {
		i.Use(middleware.LoggerWithConfig(middleware.LoggerConfig{
			Format: "time=${time_rfc3339}, method=${method}, uri=${uri}, status=${status}, latency=${latency_human}\n",
		}))
	}

	i.Use(middleware.Gzip())

	if len(config.Origins) > 0 {
		i.Use(middleware.CORSWithConfig(middleware.CORSConfig{
			AllowOrigins:     config.Origins,
			AllowCredentials: true,
		}))
	}

	i.Use(middleware.BodyLimit(config.MaxReqSize))
	i.Use(middleware.Recover())

	apiMiddlewareSkipper := func(c echo.Context) bool {
		path := c.Request().URL.EscapedPath()
		return strings.HasPrefix(path, config.BaseURL+"/") || config.BaseURL == path
	}
	if len(config.FrontendProxy) > 0 {
		targets := make([]*middleware.ProxyTarget, 0, len(config.FrontendProxy))
		for _, i := range config.FrontendProxy {
			if u, err := url.Parse(i); err == nil {
				targets = append(targets, &middleware.ProxyTarget{
					URL: u,
				})
			} else {
				l.Warn("could not add frontend proxy", moduleIDServer, "fail add proxy", 0, map[string]string{
					"proxy": i,
					"error": err.Error(),
				})
			}
		}
		if len(targets) > 0 {
			i.Use(middleware.ProxyWithConfig(middleware.ProxyConfig{
				Balancer: middleware.NewRoundRobinBalancer(targets),
				Skipper:  apiMiddlewareSkipper,
			}))
		}
	} else {
		i.Use(middleware.StaticWithConfig(middleware.StaticConfig{
			Root:    config.PublicDir,
			Index:   "index.html",
			Browse:  false,
			HTML5:   true,
			Skipper: apiMiddlewareSkipper,
		}))
	}

	i.Use(middleware.RequestID())
	l.Info("initialized middleware", moduleIDServer, "initialize middleware", 0, nil)

	healthService := newHealth()
	if err := healthService.Mount(config, l, i.Group(config.BaseURL+"/healthz")); err != nil {
		return nil, err
	}
	setupService := newSetup()
	if err := setupService.Mount(config, l, i.Group(config.BaseURL+"/setupz")); err != nil {
		return nil, err
	}

	s := &Server{
		i:          i,
		logger:     l,
		config:     config,
		h:          healthService,
		s:          setupService,
		showBanner: true,
	}

	l.Info("server instance created", moduleIDServer, "initialize server instance", 0, nil)
	return s, nil
}

// Start starts the server at the specified port
func (s *Server) Start() error {
	if s.showBanner {
		fmt.Printf(color.BlueString(banner+"\n"), color.GreenString(s.config.Version), "build version:"+color.GreenString(s.config.VersionHash), "http server on "+color.RedString(":"+s.config.Port))
	}
	s.i.Logger.Fatal(s.i.Start(":" + s.config.Port))
	return nil
}

// Must ensures that the operation must succeed
func Must(err error) {
	if err != nil {
		fmt.Println(err)
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Hajime Hoshi
//
// 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.

// +build example jsgo

// This is an example to implement an audio player.
// See examples/wav for a simpler example to play a sound file.

package main

import (
	"fmt"
	"image/color"
	"io/ioutil"
	"log"
	"time"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/audio/vorbis"
	"github.com/hajimehoshi/ebiten/audio/wav"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	raudio "github.com/hajimehoshi/ebiten/examples/resources/audio"
	"github.com/hajimehoshi/ebiten/inpututil"
)

const (
	screenWidth  = 320
	screenHeight = 240

	sampleRate = 22050
)

var (
	playerBarColor     = color.RGBA{0x80, 0x80, 0x80, 0xff}
	playerCurrentColor = color.RGBA{0xff, 0xff, 0xff, 0xff}
)

// Player represents the current audio state.
type Player struct {
	audioContext *audio.Context
	audioPlayer  *audio.Player
	current      time.Duration
	total        time.Duration
	seBytes      []byte
	seCh         chan []byte
	volume128    int
}

func playerBarRect() (x, y, w, h int) {
	w, h = 300, 4
	x = (screenWidth - w) / 2
	y = screenHeight - h - 16
	return
}

func NewPlayer(audioContext *audio.Context) (*Player, error) {
	const bytesPerSample = 4 // TODO: This should be defined in audio package
	s, err := vorbis.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Ragtime_ogg))
	if err != nil {
		return nil, err
	}
	p, err := audio.NewPlayer(audioContext, s)
	if err != nil {
		return nil, err
	}
	player := &Player{
		audioContext: audioContext,
		audioPlayer:  p,
		total:        time.Second * time.Duration(s.Length()) / bytesPerSample / sampleRate,
		volume128:    128,
		seCh:         make(chan []byte),
	}
	if player.total == 0 {
		player.total = 1
	}
	player.audioPlayer.Play()
	go func() {
		s, err := wav.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Jab_wav))
		if err != nil {
			log.Fatal(err)
			return
		}
		b, err := ioutil.ReadAll(s)
		if err != nil {
			log.Fatal(err)
			return
		}
		player.seCh <- b
	}()
	return player, nil
}

func (p *Player) update() error {
	select {
	case p.seBytes = <-p.seCh:
		close(p.seCh)
		p.seCh = nil
	default:
	}

	if p.audioPlayer.IsPlaying() {
		p.current = p.audioPlayer.Current()
	}
	p.seekBarIfNeeded()
	p.switchPlayStateIfNeeded()
	p.playSEIfNeeded()
	p.updateVolumeIfNeeded()

	if inpututil.IsKeyJustPressed(ebiten.KeyB) {
		b := ebiten.IsRunnableInBackground()
		ebiten.SetRunnableInBackground(!b)
	}
	return nil
}

func (p *Player) playSEIfNeeded() {
	if p.seBytes == nil {
		// Bytes for the SE is not loaded yet.
		return
	}

	if !inpututil.IsKeyJustPressed(ebiten.KeyP) {
		return
	}
	sePlayer, _ := audio.NewPlayerFromBytes(p.audioContext, p.seBytes)
	sePlayer.Play()
}

func (p *Player) updateVolumeIfNeeded() {
	if ebiten.IsKeyPressed(ebiten.KeyZ) {
		p.volume128--
	}
	if ebiten.IsKeyPressed(ebiten.KeyX) {
		p.volume128++
	}
	if p.volume128 < 0 {
		p.volume128 = 0
	}
	if 128 < p.volume128 {
		p.volume128 = 128
	}
	p.audioPlayer.SetVolume(float64(p.volume128) / 128)
}

func (p *Player) switchPlayStateIfNeeded() {
	if !inpututil.IsKeyJustPressed(ebiten.KeyS) {
		return
	}
	if p.audioPlayer.IsPlaying() {
		p.audioPlayer.Pause()
		return
	}
	p.audioPlayer.Play()
}

func (p *Player) seekBarIfNeeded() {
	if !inpututil.IsMouseButtonJustPressed(ebiten.MouseButtonLeft) {
		return
	}

	// Calculate the next seeking position from the current cursor position.
	x, y := ebiten.CursorPosition()
	bx, by, bw, bh := playerBarRect()
	const padding = 4
	if y < by-padding || by+bh+padding <= y {
		return
	}
	if x < bx || bx+bw <= x {
		return
	}
	pos := time.Duration(x-bx) * p.total / time.Duration(bw)
	p.current = pos
	p.audioPlayer.Seek(pos)
}

func (p *Player) close() error {
	return p.audioPlayer.Close()
}

func (p *Player) draw(screen *ebiten.Image) {
	// Draw the bar.
	x, y, w, h := playerBarRect()
	ebitenutil.DrawRect(screen, float64(x), float64(y), float64(w), float64(h), playerBarColor)

	// Draw the cursor on the bar.
	c := p.current
	cw, ch := 4, 10
	cx := int(time.Duration(w)*c/p.total) + x - cw/2
	cy := y - (ch-h)/2
	ebitenutil.DrawRect(screen, float64(cx), float64(cy), float64(cw), float64(ch), playerCurrentColor)

	// Compose the curren time text.
	m := (c / time.Minute) % 100
	s := (c / time.Second) % 60
	currentTimeStr := fmt.Sprintf("%02d:%02d", m, s)

	// Draw the debug message.
	msg := fmt.Sprintf(`TPS: %0.2f
Press S to toggle Play/Pause
Press P to play SE
Press Z or X to change volume of the music
Press B to switch the run-in-background state
Current Time: %s
Current Volume: %d/128`, ebiten.CurrentTPS(), currentTimeStr, int(p.audioPlayer.Volume()*128))
	ebitenutil.DebugPrint(screen, msg)
}

var (
	musicPlayer *Player
)

func update(screen *ebiten.Image) error {
	if err := musicPlayer.update(); err != nil {
		return err
	}

	if ebiten.IsDrawingSkipped() {
		return nil
	}

	musicPlayer.draw(screen)
	return nil
}

func main() {
	audioContext, err := audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}

	musicPlayer, err = NewPlayer(audioContext)
	if err != nil {
		log.Fatal(err)
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>examples/audio: Remove unused functions<commit_after>// Copyright 2016 Hajime Hoshi
//
// 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.

// +build example jsgo

// This is an example to implement an audio player.
// See examples/wav for a simpler example to play a sound file.

package main

import (
	"fmt"
	"image/color"
	"io/ioutil"
	"log"
	"time"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/audio/vorbis"
	"github.com/hajimehoshi/ebiten/audio/wav"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	raudio "github.com/hajimehoshi/ebiten/examples/resources/audio"
	"github.com/hajimehoshi/ebiten/inpututil"
)

const (
	screenWidth  = 320
	screenHeight = 240

	sampleRate = 22050
)

var (
	playerBarColor     = color.RGBA{0x80, 0x80, 0x80, 0xff}
	playerCurrentColor = color.RGBA{0xff, 0xff, 0xff, 0xff}
)

// Player represents the current audio state.
type Player struct {
	audioContext *audio.Context
	audioPlayer  *audio.Player
	current      time.Duration
	total        time.Duration
	seBytes      []byte
	seCh         chan []byte
	volume128    int
}

func playerBarRect() (x, y, w, h int) {
	w, h = 300, 4
	x = (screenWidth - w) / 2
	y = screenHeight - h - 16
	return
}

func NewPlayer(audioContext *audio.Context) (*Player, error) {
	const bytesPerSample = 4 // TODO: This should be defined in audio package
	s, err := vorbis.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Ragtime_ogg))
	if err != nil {
		return nil, err
	}
	p, err := audio.NewPlayer(audioContext, s)
	if err != nil {
		return nil, err
	}
	player := &Player{
		audioContext: audioContext,
		audioPlayer:  p,
		total:        time.Second * time.Duration(s.Length()) / bytesPerSample / sampleRate,
		volume128:    128,
		seCh:         make(chan []byte),
	}
	if player.total == 0 {
		player.total = 1
	}
	player.audioPlayer.Play()
	go func() {
		s, err := wav.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Jab_wav))
		if err != nil {
			log.Fatal(err)
			return
		}
		b, err := ioutil.ReadAll(s)
		if err != nil {
			log.Fatal(err)
			return
		}
		player.seCh <- b
	}()
	return player, nil
}

func (p *Player) update() error {
	select {
	case p.seBytes = <-p.seCh:
		close(p.seCh)
		p.seCh = nil
	default:
	}

	if p.audioPlayer.IsPlaying() {
		p.current = p.audioPlayer.Current()
	}
	p.seekBarIfNeeded()
	p.switchPlayStateIfNeeded()
	p.playSEIfNeeded()
	p.updateVolumeIfNeeded()

	if inpututil.IsKeyJustPressed(ebiten.KeyB) {
		b := ebiten.IsRunnableInBackground()
		ebiten.SetRunnableInBackground(!b)
	}
	return nil
}

func (p *Player) playSEIfNeeded() {
	if p.seBytes == nil {
		// Bytes for the SE is not loaded yet.
		return
	}

	if !inpututil.IsKeyJustPressed(ebiten.KeyP) {
		return
	}
	sePlayer, _ := audio.NewPlayerFromBytes(p.audioContext, p.seBytes)
	sePlayer.Play()
}

func (p *Player) updateVolumeIfNeeded() {
	if ebiten.IsKeyPressed(ebiten.KeyZ) {
		p.volume128--
	}
	if ebiten.IsKeyPressed(ebiten.KeyX) {
		p.volume128++
	}
	if p.volume128 < 0 {
		p.volume128 = 0
	}
	if 128 < p.volume128 {
		p.volume128 = 128
	}
	p.audioPlayer.SetVolume(float64(p.volume128) / 128)
}

func (p *Player) switchPlayStateIfNeeded() {
	if !inpututil.IsKeyJustPressed(ebiten.KeyS) {
		return
	}
	if p.audioPlayer.IsPlaying() {
		p.audioPlayer.Pause()
		return
	}
	p.audioPlayer.Play()
}

func (p *Player) seekBarIfNeeded() {
	if !inpututil.IsMouseButtonJustPressed(ebiten.MouseButtonLeft) {
		return
	}

	// Calculate the next seeking position from the current cursor position.
	x, y := ebiten.CursorPosition()
	bx, by, bw, bh := playerBarRect()
	const padding = 4
	if y < by-padding || by+bh+padding <= y {
		return
	}
	if x < bx || bx+bw <= x {
		return
	}
	pos := time.Duration(x-bx) * p.total / time.Duration(bw)
	p.current = pos
	p.audioPlayer.Seek(pos)
}

func (p *Player) draw(screen *ebiten.Image) {
	// Draw the bar.
	x, y, w, h := playerBarRect()
	ebitenutil.DrawRect(screen, float64(x), float64(y), float64(w), float64(h), playerBarColor)

	// Draw the cursor on the bar.
	c := p.current
	cw, ch := 4, 10
	cx := int(time.Duration(w)*c/p.total) + x - cw/2
	cy := y - (ch-h)/2
	ebitenutil.DrawRect(screen, float64(cx), float64(cy), float64(cw), float64(ch), playerCurrentColor)

	// Compose the curren time text.
	m := (c / time.Minute) % 100
	s := (c / time.Second) % 60
	currentTimeStr := fmt.Sprintf("%02d:%02d", m, s)

	// Draw the debug message.
	msg := fmt.Sprintf(`TPS: %0.2f
Press S to toggle Play/Pause
Press P to play SE
Press Z or X to change volume of the music
Press B to switch the run-in-background state
Current Time: %s
Current Volume: %d/128`, ebiten.CurrentTPS(), currentTimeStr, int(p.audioPlayer.Volume()*128))
	ebitenutil.DebugPrint(screen, msg)
}

var (
	musicPlayer *Player
)

func update(screen *ebiten.Image) error {
	if err := musicPlayer.update(); err != nil {
		return err
	}

	if ebiten.IsDrawingSkipped() {
		return nil
	}

	musicPlayer.draw(screen)
	return nil
}

func main() {
	audioContext, err := audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}

	musicPlayer, err = NewPlayer(audioContext)
	if err != nil {
		log.Fatal(err)
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Hajime Hoshi
//
// 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.

// +build example jsgo

// This is an example to implement an audio player.
// See examples/wav for a simpler example to play a sound file.

package main

import (
	"fmt"
	"image/color"
	"io/ioutil"
	"log"
	"time"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/audio/vorbis"
	"github.com/hajimehoshi/ebiten/audio/wav"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	raudio "github.com/hajimehoshi/ebiten/examples/resources/audio"
	"github.com/hajimehoshi/ebiten/inpututil"
)

const (
	screenWidth  = 320
	screenHeight = 240

	sampleRate = 22050
)

var (
	playerBarColor     = color.RGBA{0x80, 0x80, 0x80, 0xff}
	playerCurrentColor = color.RGBA{0xff, 0xff, 0xff, 0xff}
)

// Player represents the current audio state.
type Player struct {
	audioContext *audio.Context
	audioPlayer  *audio.Player
	current      time.Duration
	total        time.Duration
	seBytes      []byte
	seCh         chan []byte
	volume128    int
}

func playerBarRect() (x, y, w, h int) {
	w, h = 300, 4
	x = (screenWidth - w) / 2
	y = screenHeight - h - 16
	return
}

func NewPlayer(audioContext *audio.Context) (*Player, error) {
	const bytesPerSample = 4 // TODO: This should be defined in audio package
	s, err := vorbis.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Ragtime_ogg))
	if err != nil {
		return nil, err
	}
	p, err := audio.NewPlayer(audioContext, s)
	if err != nil {
		return nil, err
	}
	player := &Player{
		audioContext: audioContext,
		audioPlayer:  p,
		total:        time.Second * time.Duration(s.Length()) / bytesPerSample / sampleRate,
		volume128:    128,
		seCh:         make(chan []byte),
	}
	if player.total == 0 {
		player.total = 1
	}
	player.audioPlayer.Play()
	go func() {
		s, err := wav.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Jab_wav))
		if err != nil {
			log.Fatal(err)
			return
		}
		b, err := ioutil.ReadAll(s)
		if err != nil {
			log.Fatal(err)
			return
		}
		player.seCh <- b
	}()
	return player, nil
}

func (p *Player) update() error {
	select {
	case p.seBytes = <-p.seCh:
		close(p.seCh)
		p.seCh = nil
	default:
	}

	if p.audioPlayer.IsPlaying() {
		p.current = p.audioPlayer.Current()
	}
	p.seekBarIfNeeded()
	p.switchPlayStateIfNeeded()
	p.playSEIfNeeded()
	p.updateVolumeIfNeeded()

	if inpututil.IsKeyJustPressed(ebiten.KeyB) {
		b := ebiten.IsRunnableInBackground()
		ebiten.SetRunnableInBackground(!b)
	}
	return nil
}

func (p *Player) playSEIfNeeded() {
	if p.seBytes == nil {
		// Bytes for the SE is not loaded yet.
		return
	}

	if !inpututil.IsKeyJustPressed(ebiten.KeyP) {
		return
	}
	sePlayer, _ := audio.NewPlayerFromBytes(p.audioContext, p.seBytes)
	sePlayer.Play()
}

func (p *Player) updateVolumeIfNeeded() {
	if ebiten.IsKeyPressed(ebiten.KeyZ) {
		p.volume128--
	}
	if ebiten.IsKeyPressed(ebiten.KeyX) {
		p.volume128++
	}
	if p.volume128 < 0 {
		p.volume128 = 0
	}
	if 128 < p.volume128 {
		p.volume128 = 128
	}
	p.audioPlayer.SetVolume(float64(p.volume128) / 128)
}

func (p *Player) switchPlayStateIfNeeded() {
	if !inpututil.IsKeyJustPressed(ebiten.KeyS) {
		return
	}
	if p.audioPlayer.IsPlaying() {
		p.audioPlayer.Pause()
		return
	}
	p.audioPlayer.Play()
}

func (p *Player) seekBarIfNeeded() {
	if !inpututil.IsMouseButtonJustPressed(ebiten.MouseButtonLeft) {
		return
	}

	// Calculate the next seeking position from the current cursor position.
	x, y := ebiten.CursorPosition()
	bx, by, bw, bh := playerBarRect()
	const padding = 4
	if y < by-padding || by+bh+padding <= y {
		return
	}
	if x < bx || bx+bw <= x {
		return
	}
	pos := time.Duration(x-bx) * p.total / time.Duration(bw)
	p.current = pos
	p.audioPlayer.Seek(pos)
}

func (p *Player) close() error {
	return p.audioPlayer.Close()
}

func (p *Player) draw(screen *ebiten.Image) {
	// Draw the bar.
	x, y, w, h := playerBarRect()
	ebitenutil.DrawRect(screen, float64(x), float64(y), float64(w), float64(h), playerBarColor)

	// Draw the cursor on the bar.
	c := p.current
	cw, ch := 4, 10
	cx := int(time.Duration(w)*c/p.total) + x - cw/2
	cy := y - (ch-h)/2
	ebitenutil.DrawRect(screen, float64(cx), float64(cy), float64(cw), float64(ch), playerCurrentColor)

	// Compose the curren time text.
	m := (c / time.Minute) % 100
	s := (c / time.Second) % 60
	currentTimeStr := fmt.Sprintf("%02d:%02d", m, s)

	// Draw the debug message.
	msg := fmt.Sprintf(`TPS: %0.2f
Press S to toggle Play/Pause
Press P to play SE
Press Z or X to change volume of the music
Press B to switch the run-in-background state
Current Time: %s
Current Volume: %d/128`, ebiten.CurrentTPS(), currentTimeStr, int(p.audioPlayer.Volume()*128))
	ebitenutil.DebugPrint(screen, msg)
}

var (
	musicPlayer *Player
)

func update(screen *ebiten.Image) error {
	if err := musicPlayer.update(); err != nil {
		return err
	}

	if ebiten.IsDrawingSkipped() {
		return nil
	}

	musicPlayer.draw(screen)
	return nil
}

func main() {
	audioContext, err := audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}

	musicPlayer, err = NewPlayer(audioContext)
	if err != nil {
		log.Fatal(err)
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
	if err := musicPlayer.close(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>examples/audio: Bug fix: Do not close after Run<commit_after>// Copyright 2016 Hajime Hoshi
//
// 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.

// +build example jsgo

// This is an example to implement an audio player.
// See examples/wav for a simpler example to play a sound file.

package main

import (
	"fmt"
	"image/color"
	"io/ioutil"
	"log"
	"time"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/audio/vorbis"
	"github.com/hajimehoshi/ebiten/audio/wav"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	raudio "github.com/hajimehoshi/ebiten/examples/resources/audio"
	"github.com/hajimehoshi/ebiten/inpututil"
)

const (
	screenWidth  = 320
	screenHeight = 240

	sampleRate = 22050
)

var (
	playerBarColor     = color.RGBA{0x80, 0x80, 0x80, 0xff}
	playerCurrentColor = color.RGBA{0xff, 0xff, 0xff, 0xff}
)

// Player represents the current audio state.
type Player struct {
	audioContext *audio.Context
	audioPlayer  *audio.Player
	current      time.Duration
	total        time.Duration
	seBytes      []byte
	seCh         chan []byte
	volume128    int
}

func playerBarRect() (x, y, w, h int) {
	w, h = 300, 4
	x = (screenWidth - w) / 2
	y = screenHeight - h - 16
	return
}

func NewPlayer(audioContext *audio.Context) (*Player, error) {
	const bytesPerSample = 4 // TODO: This should be defined in audio package
	s, err := vorbis.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Ragtime_ogg))
	if err != nil {
		return nil, err
	}
	p, err := audio.NewPlayer(audioContext, s)
	if err != nil {
		return nil, err
	}
	player := &Player{
		audioContext: audioContext,
		audioPlayer:  p,
		total:        time.Second * time.Duration(s.Length()) / bytesPerSample / sampleRate,
		volume128:    128,
		seCh:         make(chan []byte),
	}
	if player.total == 0 {
		player.total = 1
	}
	player.audioPlayer.Play()
	go func() {
		s, err := wav.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Jab_wav))
		if err != nil {
			log.Fatal(err)
			return
		}
		b, err := ioutil.ReadAll(s)
		if err != nil {
			log.Fatal(err)
			return
		}
		player.seCh <- b
	}()
	return player, nil
}

func (p *Player) update() error {
	select {
	case p.seBytes = <-p.seCh:
		close(p.seCh)
		p.seCh = nil
	default:
	}

	if p.audioPlayer.IsPlaying() {
		p.current = p.audioPlayer.Current()
	}
	p.seekBarIfNeeded()
	p.switchPlayStateIfNeeded()
	p.playSEIfNeeded()
	p.updateVolumeIfNeeded()

	if inpututil.IsKeyJustPressed(ebiten.KeyB) {
		b := ebiten.IsRunnableInBackground()
		ebiten.SetRunnableInBackground(!b)
	}
	return nil
}

func (p *Player) playSEIfNeeded() {
	if p.seBytes == nil {
		// Bytes for the SE is not loaded yet.
		return
	}

	if !inpututil.IsKeyJustPressed(ebiten.KeyP) {
		return
	}
	sePlayer, _ := audio.NewPlayerFromBytes(p.audioContext, p.seBytes)
	sePlayer.Play()
}

func (p *Player) updateVolumeIfNeeded() {
	if ebiten.IsKeyPressed(ebiten.KeyZ) {
		p.volume128--
	}
	if ebiten.IsKeyPressed(ebiten.KeyX) {
		p.volume128++
	}
	if p.volume128 < 0 {
		p.volume128 = 0
	}
	if 128 < p.volume128 {
		p.volume128 = 128
	}
	p.audioPlayer.SetVolume(float64(p.volume128) / 128)
}

func (p *Player) switchPlayStateIfNeeded() {
	if !inpututil.IsKeyJustPressed(ebiten.KeyS) {
		return
	}
	if p.audioPlayer.IsPlaying() {
		p.audioPlayer.Pause()
		return
	}
	p.audioPlayer.Play()
}

func (p *Player) seekBarIfNeeded() {
	if !inpututil.IsMouseButtonJustPressed(ebiten.MouseButtonLeft) {
		return
	}

	// Calculate the next seeking position from the current cursor position.
	x, y := ebiten.CursorPosition()
	bx, by, bw, bh := playerBarRect()
	const padding = 4
	if y < by-padding || by+bh+padding <= y {
		return
	}
	if x < bx || bx+bw <= x {
		return
	}
	pos := time.Duration(x-bx) * p.total / time.Duration(bw)
	p.current = pos
	p.audioPlayer.Seek(pos)
}

func (p *Player) close() error {
	return p.audioPlayer.Close()
}

func (p *Player) draw(screen *ebiten.Image) {
	// Draw the bar.
	x, y, w, h := playerBarRect()
	ebitenutil.DrawRect(screen, float64(x), float64(y), float64(w), float64(h), playerBarColor)

	// Draw the cursor on the bar.
	c := p.current
	cw, ch := 4, 10
	cx := int(time.Duration(w)*c/p.total) + x - cw/2
	cy := y - (ch-h)/2
	ebitenutil.DrawRect(screen, float64(cx), float64(cy), float64(cw), float64(ch), playerCurrentColor)

	// Compose the curren time text.
	m := (c / time.Minute) % 100
	s := (c / time.Second) % 60
	currentTimeStr := fmt.Sprintf("%02d:%02d", m, s)

	// Draw the debug message.
	msg := fmt.Sprintf(`TPS: %0.2f
Press S to toggle Play/Pause
Press P to play SE
Press Z or X to change volume of the music
Press B to switch the run-in-background state
Current Time: %s
Current Volume: %d/128`, ebiten.CurrentTPS(), currentTimeStr, int(p.audioPlayer.Volume()*128))
	ebitenutil.DebugPrint(screen, msg)
}

var (
	musicPlayer *Player
)

func update(screen *ebiten.Image) error {
	if err := musicPlayer.update(); err != nil {
		return err
	}

	if ebiten.IsDrawingSkipped() {
		return nil
	}

	musicPlayer.draw(screen)
	return nil
}

func main() {
	audioContext, err := audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}

	musicPlayer, err = NewPlayer(audioContext)
	if err != nil {
		log.Fatal(err)
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>