<commit_before>package gearcmd

import (
	"bufio"
	"bytes"
	"container/ring"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/Clever/gearcmd/argsparser"
	"github.com/Clever/gearcmd/baseworker"
	"gopkg.in/Clever/kayvee-go.v3/logger"
)

// TaskConfig defines the configuration for the task.
type TaskConfig struct {
	FunctionName string
	FunctionCmd  string
	WarningLines int
	ParseArgs    bool
	CmdTimeout   time.Duration
	RetryCount   int
	Halt         chan struct{}
}

var (
	lg = logger.New("gearcmd")
	// legacy logger used to maintain existing alerts.
	// Once all workers are migrated to using gearcmd >= v0.5.0 and the alarms are switched over,
	// then we can remove this logger
	legacyLg = logger.New("gearman")
)

// Process runs the Gearman job by running the configured task.
// We need to implement the Task interface so we return (byte[], error)
// though the byte[] is always nil.
func (conf TaskConfig) Process(job baseworker.Job) (b []byte, returnErr error) {
	jobID := getJobID(job)
	if jobID == "" {
		jobID = strconv.Itoa(rand.Int())
		lg.InfoD("rand-job-id", logger.M{"msg": "no job id parsed, random assigned."})
	}

	jobData := string(job.Data())
	data := logger.M{
		"function": conf.FunctionName,
		"job_id":   jobID,
		"job_data": jobData,
	}

	// This wraps the actual processing to do some logging
	lg.InfoD("START", data)
	start := time.Now()

	for try := 0; try < conf.RetryCount+1; try++ {
		// We create a temporary directory to be used as the work directory of the process.
		// A new work directory is created for every retry of the process.
		// We try to use MEOS_SANDBOX, the default will be the system temp directory.
		tempDirPath, err := ioutil.TempDir(os.Getenv("MESOS_SANDBOX"),
			fmt.Sprintf("%s-%s-%d-", conf.FunctionName, jobID, try))
		if err != nil {
			lg.CriticalD("tempdir-failure", logger.M{"error": err.Error()})
			return nil, err
		}
		defer os.RemoveAll(tempDirPath)

		// insert the job id and the work directory path into the environment
		extraEnvVars := []string{
			fmt.Sprintf("JOB_ID=%s", jobID),
			fmt.Sprintf("WORK_DIR=%s", tempDirPath)}

		err = conf.doProcess(job, extraEnvVars, try)
		end := time.Now()
		data["type"] = "gauge"

		// Return if the job was successful.
		if err == nil {
			lg.InfoD("SUCCESS", logger.M{
				"type":     "counter",
				"function": conf.FunctionName})
			legacyLg.InfoD("success", logger.M{
				"type":     "counter",
				"function": conf.FunctionName})
			data["value"] = 1
			data["success"] = true
			lg.InfoD("END", data)
			// Hopefully none of our jobs last long enough for a uint64...
			// Note that we cannot use lg.GaugeIntD because duration is uint64
			lg.InfoD("duration", logger.M{
				"value":    uint64(end.Sub(start).Seconds() * 1000),
				"type":     "gauge",
				"function": conf.FunctionName,
				"job_id":   jobID,
				"job_data": jobData})
			return nil, nil
		}

		data["value"] = 0
		data["success"] = false
		data["error_message"] = err.Error()
		returnErr = err

		if try != conf.RetryCount {
			lg.ErrorD("RETRY", data)
		}
	}

	lg.InfoD("FAILURE", logger.M{"type": "counter", "function": conf.FunctionName})
	legacyLg.InfoD("failure", logger.M{"type": "counter", "function": conf.FunctionName})
	lg.ErrorD("END", data)
	return nil, returnErr
}

// getJobID returns the jobId from the job handle
func getJobID(job baseworker.Job) string {
	splits := strings.Split(job.Handle(), ":")
	return splits[len(splits)-1]
}

func (conf TaskConfig) doProcess(job baseworker.Job, envVars []string, tryCount int) error {
	defer func() {
		// If we panicked then set the panic message as a warning. Gearman-go will
		// handle marking this job as failed.
		if r := recover(); r != nil {
			err := r.(error)
			job.SendWarning([]byte(err.Error()))
		}
	}()

	// shutdownTicker will effectively control the executution of the ticker.
	shutdownTicker := make(chan interface{})
	defer func() {
		shutdownTicker <- 1
	}()

	// every minute we will output a heartbeat kayvee log for the job.
	tickUnit := time.Minute
	ticker := time.NewTicker(tickUnit)
	go func() {
		defer ticker.Stop()
		units := 0
		for {
			select {
			case <-shutdownTicker:
				close(shutdownTicker)
				return
			case <-ticker.C:
				units++
				lg.GaugeIntD("heartbeat", units, logger.M{
					"try_number": tryCount,
					"function":   job.Fn(),
					"job_id":     getJobID(job),
					"unit":       tickUnit.String(),
				})
			}
		}
	}()

	var args []string
	var err error
	if conf.ParseArgs {
		args, err = argsparser.ParseArgs(string(job.Data()))
		if err != nil {
			return fmt.Errorf("Failed to parse args: %s", err.Error())
		}
	} else {
		args = []string{string(job.Data())}
	}
	cmd := exec.Command(conf.FunctionCmd, args...)

	// insert provided env vars into the job
	cmd.Env = append(os.Environ(), envVars...)

	// create new pgid for this process so we can later kill all subprocess launched by it
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	// Write the stdout and stderr of the process to both this process' stdout and stderr
	// and also write it to a byte buffer so that we can return it with the Gearman job
	// data as necessary.
	var stderrbuf bytes.Buffer
	cmd.Stderr = io.MultiWriter(os.Stderr, &stderrbuf)
	defer sendStderrWarnings(&stderrbuf, job, conf.WarningLines)

	stdoutReader, stdoutWriter := io.Pipe()
	cmd.Stdout = io.MultiWriter(os.Stdout, stdoutWriter)

	done := make(chan error)
	// Track when the job has started so that we don't try and sigterm a nil process
	started := make(chan struct{})
	go func() {
		defer close(done)

		finishedProcessingStdout := make(chan error)
		go func() {
			finishedProcessingStdout <- streamToGearman(stdoutReader, job)
		}()

		if err := cmd.Start(); err != nil {
			done <- err
			return
		}
		close(started)
		// Save the cmdErr. We want to process stdout and stderr before we return it
		cmdErr := cmd.Wait()
		stdoutWriter.Close()

		stdoutErr := <-finishedProcessingStdout
		if cmdErr != nil {
			done <- cmdErr
		} else if stdoutErr != nil {
			done <- stdoutErr
		}
	}()
	<-started

	// No timeout
	if conf.CmdTimeout == 0 {
		select {
		case err := <-done:
			// Will be nil if the channel was closed without any errors
			return err
		case <-conf.Halt:
			if err := stopProcess(cmd.Process, conf.CmdTimeout); err != nil {
				return fmt.Errorf("error stopping process: %s", err)
			}
			return fmt.Errorf("killed process due to sigterm")
		}
	}
	select {
	case err := <-done:
		// Will be nil if the channel was closed without any errors
		return err
	case <-conf.Halt:
		if err := stopProcess(cmd.Process, conf.CmdTimeout); err != nil {
			return fmt.Errorf("error stopping process: %s", err)
		}
		return nil
	case <-time.After(conf.CmdTimeout):
		if err := stopProcess(cmd.Process, 0); err != nil {
			return fmt.Errorf("error timing out process after %s: %s", conf.CmdTimeout.String(), err)
		}
		return fmt.Errorf("process timed out after %s", conf.CmdTimeout.String())
	}
}

// stopProcess kills a given process. It's second argument is a grace period, if > 0, SIGTERM will
// be sent first folloed by SIGKILL after the grace period is over. If the grace period is 0, this
// is simply a hard SIGKILL.
func stopProcess(p *os.Process, gracePeriod time.Duration) error {
	lg.InfoD("killing-pid", logger.M{"pid": p.Pid})
	if gracePeriod > 0 {
		// we use SIGTERM so that the subprocess can gracefully exit
		if err := syscall.Kill(p.Pid, syscall.SIGTERM); err != nil {
			return fmt.Errorf("unable to send SIGTERM, error: %s", err)
		}
		time.Sleep(gracePeriod)
	}
	if err := syscall.Kill(p.Pid, syscall.SIGKILL); err != nil {
		// The graceful shutdown may have completed, so we don't error in that case.
		if gracePeriod == 0 || !strings.Contains(err.Error(), "no such process") {
			return fmt.Errorf("unable to send SIGKILL, error: %s", err)
		}
	}
	return nil
}

// This function streams the reader to the Gearman job (through job.SendData())
func streamToGearman(reader io.Reader, job baseworker.Job) error {
	buffer := make([]byte, 1024)
	for {
		n, err := reader.Read(buffer)
		// Process the data before processing the error (as per the io.Reader docs)
		if n > 0 {
			job.SendData(buffer[:n])
		}
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// sendStderrWarnings sends the last X lines in the stderr output and to the job's warnings
// field
func sendStderrWarnings(buffer io.Reader, job baseworker.Job, warningLines int) error {
	scanner := bufio.NewScanner(buffer)
	// Create a circular buffer for the last X lines
	lastStderrLines := ring.New(warningLines)
	for scanner.Scan() {
		lastStderrLines = lastStderrLines.Next()
		lastStderrLines.Value = scanner.Bytes()
	}
	// Walk forward through the buffer to get all the last X entries. Note that we call next first
	// so that we start at the oldest entry.
	for i := 0; i < lastStderrLines.Len(); i++ {
		if lastStderrLines = lastStderrLines.Next(); lastStderrLines.Value != nil {
			job.SendWarning(append(lastStderrLines.Value.([]byte), byte('\n')))
		}
	}
	return scanner.Err()
}
<commit_msg>write all stderr lines at once<commit_after>package gearcmd

import (
	"bufio"
	"bytes"
	"container/ring"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/Clever/gearcmd/argsparser"
	"github.com/Clever/gearcmd/baseworker"
	"gopkg.in/Clever/kayvee-go.v3/logger"
)

// TaskConfig defines the configuration for the task.
type TaskConfig struct {
	FunctionName string
	FunctionCmd  string
	WarningLines int
	ParseArgs    bool
	CmdTimeout   time.Duration
	RetryCount   int
	Halt         chan struct{}
}

var (
	lg = logger.New("gearcmd")
	// legacy logger used to maintain existing alerts.
	// Once all workers are migrated to using gearcmd >= v0.5.0 and the alarms are switched over,
	// then we can remove this logger
	legacyLg = logger.New("gearman")
)

// Process runs the Gearman job by running the configured task.
// We need to implement the Task interface so we return (byte[], error)
// though the byte[] is always nil.
func (conf TaskConfig) Process(job baseworker.Job) (b []byte, returnErr error) {
	jobID := getJobID(job)
	if jobID == "" {
		jobID = strconv.Itoa(rand.Int())
		lg.InfoD("rand-job-id", logger.M{"msg": "no job id parsed, random assigned."})
	}

	jobData := string(job.Data())
	data := logger.M{
		"function": conf.FunctionName,
		"job_id":   jobID,
		"job_data": jobData,
	}

	// This wraps the actual processing to do some logging
	lg.InfoD("START", data)
	start := time.Now()

	for try := 0; try < conf.RetryCount+1; try++ {
		// We create a temporary directory to be used as the work directory of the process.
		// A new work directory is created for every retry of the process.
		// We try to use MEOS_SANDBOX, the default will be the system temp directory.
		tempDirPath, err := ioutil.TempDir(os.Getenv("MESOS_SANDBOX"),
			fmt.Sprintf("%s-%s-%d-", conf.FunctionName, jobID, try))
		if err != nil {
			lg.CriticalD("tempdir-failure", logger.M{"error": err.Error()})
			return nil, err
		}
		defer os.RemoveAll(tempDirPath)

		// insert the job id and the work directory path into the environment
		extraEnvVars := []string{
			fmt.Sprintf("JOB_ID=%s", jobID),
			fmt.Sprintf("WORK_DIR=%s", tempDirPath)}

		err = conf.doProcess(job, extraEnvVars, try)
		end := time.Now()
		data["type"] = "gauge"

		// Return if the job was successful.
		if err == nil {
			lg.InfoD("SUCCESS", logger.M{
				"type":     "counter",
				"function": conf.FunctionName})
			legacyLg.InfoD("success", logger.M{
				"type":     "counter",
				"function": conf.FunctionName})
			data["value"] = 1
			data["success"] = true
			lg.InfoD("END", data)
			// Hopefully none of our jobs last long enough for a uint64...
			// Note that we cannot use lg.GaugeIntD because duration is uint64
			lg.InfoD("duration", logger.M{
				"value":    uint64(end.Sub(start).Seconds() * 1000),
				"type":     "gauge",
				"function": conf.FunctionName,
				"job_id":   jobID,
				"job_data": jobData})
			return nil, nil
		}

		data["value"] = 0
		data["success"] = false
		data["error_message"] = err.Error()
		returnErr = err

		if try != conf.RetryCount {
			lg.ErrorD("RETRY", data)
		}
	}

	lg.InfoD("FAILURE", logger.M{"type": "counter", "function": conf.FunctionName})
	legacyLg.InfoD("failure", logger.M{"type": "counter", "function": conf.FunctionName})
	lg.ErrorD("END", data)
	return nil, returnErr
}

// getJobID returns the jobId from the job handle
func getJobID(job baseworker.Job) string {
	splits := strings.Split(job.Handle(), ":")
	return splits[len(splits)-1]
}

func (conf TaskConfig) doProcess(job baseworker.Job, envVars []string, tryCount int) error {
	defer func() {
		// If we panicked then set the panic message as a warning. Gearman-go will
		// handle marking this job as failed.
		if r := recover(); r != nil {
			err := r.(error)
			job.SendWarning([]byte(err.Error()))
		}
	}()

	// shutdownTicker will effectively control the executution of the ticker.
	shutdownTicker := make(chan interface{})
	defer func() {
		shutdownTicker <- 1
	}()

	// every minute we will output a heartbeat kayvee log for the job.
	tickUnit := time.Minute
	ticker := time.NewTicker(tickUnit)
	go func() {
		defer ticker.Stop()
		units := 0
		for {
			select {
			case <-shutdownTicker:
				close(shutdownTicker)
				return
			case <-ticker.C:
				units++
				lg.GaugeIntD("heartbeat", units, logger.M{
					"try_number": tryCount,
					"function":   job.Fn(),
					"job_id":     getJobID(job),
					"unit":       tickUnit.String(),
				})
			}
		}
	}()

	var args []string
	var err error
	if conf.ParseArgs {
		args, err = argsparser.ParseArgs(string(job.Data()))
		if err != nil {
			return fmt.Errorf("Failed to parse args: %s", err.Error())
		}
	} else {
		args = []string{string(job.Data())}
	}
	cmd := exec.Command(conf.FunctionCmd, args...)

	// insert provided env vars into the job
	cmd.Env = append(os.Environ(), envVars...)

	// create new pgid for this process so we can later kill all subprocess launched by it
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	// Write the stdout and stderr of the process to both this process' stdout and stderr
	// and also write it to a byte buffer so that we can return it with the Gearman job
	// data as necessary.
	var stderrbuf bytes.Buffer
	cmd.Stderr = io.MultiWriter(os.Stderr, &stderrbuf)
	defer sendStderrWarnings(&stderrbuf, job, conf.WarningLines)

	stdoutReader, stdoutWriter := io.Pipe()
	cmd.Stdout = io.MultiWriter(os.Stdout, stdoutWriter)

	done := make(chan error)
	// Track when the job has started so that we don't try and sigterm a nil process
	started := make(chan struct{})
	go func() {
		defer close(done)

		finishedProcessingStdout := make(chan error)
		go func() {
			finishedProcessingStdout <- streamToGearman(stdoutReader, job)
		}()

		if err := cmd.Start(); err != nil {
			done <- err
			return
		}
		close(started)
		// Save the cmdErr. We want to process stdout and stderr before we return it
		cmdErr := cmd.Wait()
		stdoutWriter.Close()

		stdoutErr := <-finishedProcessingStdout
		if cmdErr != nil {
			done <- cmdErr
		} else if stdoutErr != nil {
			done <- stdoutErr
		}
	}()
	<-started

	// No timeout
	if conf.CmdTimeout == 0 {
		select {
		case err := <-done:
			// Will be nil if the channel was closed without any errors
			return err
		case <-conf.Halt:
			if err := stopProcess(cmd.Process, conf.CmdTimeout); err != nil {
				return fmt.Errorf("error stopping process: %s", err)
			}
			return fmt.Errorf("killed process due to sigterm")
		}
	}
	select {
	case err := <-done:
		// Will be nil if the channel was closed without any errors
		return err
	case <-conf.Halt:
		if err := stopProcess(cmd.Process, conf.CmdTimeout); err != nil {
			return fmt.Errorf("error stopping process: %s", err)
		}
		return nil
	case <-time.After(conf.CmdTimeout):
		if err := stopProcess(cmd.Process, 0); err != nil {
			return fmt.Errorf("error timing out process after %s: %s", conf.CmdTimeout.String(), err)
		}
		return fmt.Errorf("process timed out after %s", conf.CmdTimeout.String())
	}
}

// stopProcess kills a given process. It's second argument is a grace period, if > 0, SIGTERM will
// be sent first folloed by SIGKILL after the grace period is over. If the grace period is 0, this
// is simply a hard SIGKILL.
func stopProcess(p *os.Process, gracePeriod time.Duration) error {
	lg.InfoD("killing-pid", logger.M{"pid": p.Pid})
	if gracePeriod > 0 {
		// we use SIGTERM so that the subprocess can gracefully exit
		if err := syscall.Kill(p.Pid, syscall.SIGTERM); err != nil {
			return fmt.Errorf("unable to send SIGTERM, error: %s", err)
		}
		time.Sleep(gracePeriod)
	}
	if err := syscall.Kill(p.Pid, syscall.SIGKILL); err != nil {
		// The graceful shutdown may have completed, so we don't error in that case.
		if gracePeriod == 0 || !strings.Contains(err.Error(), "no such process") {
			return fmt.Errorf("unable to send SIGKILL, error: %s", err)
		}
	}
	return nil
}

// This function streams the reader to the Gearman job (through job.SendData())
func streamToGearman(reader io.Reader, job baseworker.Job) error {
	buffer := make([]byte, 1024)
	for {
		n, err := reader.Read(buffer)
		// Process the data before processing the error (as per the io.Reader docs)
		if n > 0 {
			job.SendData(buffer[:n])
		}
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// sendStderrWarnings sends the last X lines in the stderr output and to the job's warnings
// field
func sendStderrWarnings(buffer io.Reader, job baseworker.Job, warningLines int) error {
	scanner := bufio.NewScanner(buffer)
	// Create a circular buffer for the last X lines
	lastStderrLines := ring.New(warningLines)
	for scanner.Scan() {
		lastStderrLines = lastStderrLines.Next()
		lastStderrLines.Value = scanner.Bytes()
	}
	// Walk forward through the buffer to get all the last X entries. Note that we call next first
	// so that we start at the oldest entry.
	stderrbuf := []byte{}
	for i := 0; i < lastStderrLines.Len(); i++ {
		if lastStderrLines = lastStderrLines.Next(); lastStderrLines.Value != nil {
			stderrbuf = append(stderrbuf, lastStderrLines.Value.([]byte)...)
			stderrbuf = append(stderrbuf, '\n')
		}
	}
	job.SendWarning(stderrbuf)
	return scanner.Err()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"go/importer"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

type SliceFlag []string

func (sf *SliceFlag) Set(value string) error {
	*sf = append(*sf, value)
	return nil
}

func (sf SliceFlag) String() string {
	return strings.Join(sf, ", ")
}

func (sf SliceFlag) Get() interface{} {
	return sf
}

func main() {
	var packageFlag SliceFlag
	flag.Var(&packageFlag, "p", "package to list types of. can be used multiple times")
	outputFile := flag.String("o", "-", "File to output to. Blank or - for stdin")
	templateFile := flag.String("t", "", "File to use as template for sprintf. if blank, just list the types")
	fmtStr := flag.String("f", "%s", "Format string to use on each type before sending to the template")

	flag.Parse()
	var err error

	tmpl := "%s"
	if *templateFile != "" {
		var bytes []byte
		bytes, err = ioutil.ReadFile(*templateFile)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		tmpl = string(bytes)
	}

	var wr io.WriteCloser = os.Stdout
	if *outputFile != "" && *outputFile != "-" {
		wr, err = os.Create(*outputFile)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}
	importer := importer.Default()

	// aaaallllrighty that's all the flag stuff outta the way
	// now we read all the packages and fmt.Fprintf(wr, tmpl, types)
	var types []string

	for _, p := range packageFlag {
		pkg, err := importer.Import(p)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		pkgName := pkg.Name()
		pkgPath := pkg.Path()
		scope := pkg.Scope()
		names := scope.Names()
		for _, name := range names {
			obj := scope.Lookup(name)
			inScopeRef := fmt.Sprintf("%s.%s", pkgName, name)
			fullNameWithPath := fmt.Sprintf("%s.%s", pkgPath, name)
			if obj.Type().String() == fullNameWithPath {
				types = append(types, fmt.Sprintf(*fmtStr, inScopeRef))
			}
		}

	}

	fmt.Fprintf(wr, tmpl, strings.Join(types, "\n"))
	err = wr.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

}
<commit_msg>Try specifying the gc importer?<commit_after>package main

import (
	"flag"
	"fmt"
	"go/importer"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

type SliceFlag []string

func (sf *SliceFlag) Set(value string) error {
	*sf = append(*sf, value)
	return nil
}

func (sf SliceFlag) String() string {
	return strings.Join(sf, ", ")
}

func (sf SliceFlag) Get() interface{} {
	return sf
}

func main() {
	var packageFlag SliceFlag
	flag.Var(&packageFlag, "p", "package to list types of. can be used multiple times")
	outputFile := flag.String("o", "-", "File to output to. Blank or - for stdin")
	templateFile := flag.String("t", "", "File to use as template for sprintf. if blank, just list the types")
	fmtStr := flag.String("f", "%s", "Format string to use on each type before sending to the template")

	flag.Parse()
	var err error

	tmpl := "%s"
	if *templateFile != "" {
		var bytes []byte
		bytes, err = ioutil.ReadFile(*templateFile)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		tmpl = string(bytes)
	}

	var wr io.WriteCloser = os.Stdout
	if *outputFile != "" && *outputFile != "-" {
		wr, err = os.Create(*outputFile)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}
	importer := importer.Default()

	// aaaallllrighty that's all the flag stuff outta the way
	// now we read all the packages and fmt.Fprintf(wr, tmpl, types)
	var types []string

	for _, p := range packageFlag {
		pkg, err := importer.For("gc", nil).Import(p)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		pkgName := pkg.Name()
		pkgPath := pkg.Path()
		scope := pkg.Scope()
		names := scope.Names()
		for _, name := range names {
			obj := scope.Lookup(name)
			inScopeRef := fmt.Sprintf("%s.%s", pkgName, name)
			fullNameWithPath := fmt.Sprintf("%s.%s", pkgPath, name)
			if obj.Type().String() == fullNameWithPath {
				types = append(types, fmt.Sprintf(*fmtStr, inScopeRef))
			}
		}

	}

	fmt.Fprintf(wr, tmpl, strings.Join(types, "\n"))
	err = wr.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

}
<|endoftext|>


<commit_before>// Copyright 2018 The Prometheus Authors
// 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 main

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

	"github.com/prometheus/common/log"
	"gopkg.in/alecthomas/kingpin.v2"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/snmp_exporter/config"
)

// Generate a snmp_exporter config and write it out.
func generateConfig(nodes *Node, nameToNode map[string]*Node) {
	outputPath, err := filepath.Abs(*outputPath)
	if err != nil {
		log.Fatal("Unable to determine absolute path for output")
	}

	content, err := ioutil.ReadFile("generator.yml")
	if err != nil {
		log.Fatalf("Error reading yml config: %s", err)
	}
	cfg := &Config{}
	err = yaml.UnmarshalStrict(content, cfg)
	if err != nil {
		log.Fatalf("Error parsing yml config: %s", err)
	}

	outputConfig := config.Config{}
	for name, m := range cfg.Modules {
		log.Infof("Generating config for module %s", name)
		// Give each module a copy of the tree so that it can be modified.
		mNodes := nodes.Copy()
		// Build the map with new pointers.
		mNameToNode := map[string]*Node{}
		walkNode(mNodes, func(n *Node) {
			mNameToNode[n.Oid] = n
			mNameToNode[n.Label] = n
		})
		outputConfig[name] = generateConfigModule(m, mNodes, mNameToNode)
		outputConfig[name].WalkParams = m.WalkParams
		log.Infof("Generated %d metrics for module %s", len(outputConfig[name].Metrics), name)
	}

	config.DoNotHideSecrets = true
	out, err := yaml.Marshal(outputConfig)
	config.DoNotHideSecrets = false
	if err != nil {
		log.Fatalf("Error marshaling yml: %s", err)
	}

	// Check the generated config to catch auth/version issues.
	err = yaml.UnmarshalStrict(out, &config.Config{})
	if err != nil {
		log.Fatalf("Error parsing generated config: %s", err)
	}

	f, err := os.Create(outputPath)
	if err != nil {
		log.Fatalf("Error opening output file: %s", err)
	}
	out = append([]byte("# WARNING: This file was auto-generated using snmp_exporter generator, manual changes will be lost.\n"), out...)
	_, err = f.Write(out)
	if err != nil {
		log.Fatalf("Error writing to output file: %s", err)
	}
	log.Infof("Config written to %s", outputPath)
}

var (
	failOnParseErrors  = kingpin.Flag("fail-on-parse-errors", "Exit with a non-zero status if there are MIB parsing errors").Default("false").Bool()
	generateCommand    = kingpin.Command("generate", "Generate snmp.yml from generator.yml")
	outputPath         = generateCommand.Flag("output-path", "Path to to write resulting config file").Default("snmp.yml").Short('o').String()
	parseErrorsCommand = kingpin.Command("parse_errors", "Debug: Print the parse errors output by NetSNMP")
	dumpCommand        = kingpin.Command("dump", "Debug: Dump the parsed and prepared MIBs")
)

func main() {
	log.AddFlags(kingpin.CommandLine)
	kingpin.HelpFlag.Short('h')
	command := kingpin.Parse()

	parseErrors := initSNMP()
	parseErrorCount := len(parseErrors)

	if parseErrorCount != 0 {
		log.Warnf("NetSNMP reported %d parse errors", len(strings.Split(parseErrors, "\n")))
	}

	nodes := getMIBTree()
	nameToNode := prepareTree(nodes)

	switch command {
	case generateCommand.FullCommand():
		generateConfig(nodes, nameToNode)
	case parseErrorsCommand.FullCommand():
		fmt.Println(parseErrors)
	case dumpCommand.FullCommand():
		walkNode(nodes, func(n *Node) {
			t := n.Type
			if n.FixedSize != 0 {
				t = fmt.Sprintf("%s(%d)", n.Type, n.FixedSize)
			}
			implied := ""
			if n.ImpliedIndex {
				implied = "(implied)"
			}
			fmt.Printf("%s %s %s %q %q %s%s %v %s\n",
				n.Oid, n.Label, t, n.TextualConvention, n.Hint, n.Indexes, implied, n.EnumValues, n.Description)
		})
	}
	if *failOnParseErrors && parseErrorCount != 0 {
		os.Exit(1)
	}
}
<commit_msg>Handle trailing linefeed in NetSNMP output adding 1 to the error count (#398)<commit_after>// Copyright 2018 The Prometheus Authors
// 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 main

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

	"github.com/prometheus/common/log"
	"gopkg.in/alecthomas/kingpin.v2"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/snmp_exporter/config"
)

// Generate a snmp_exporter config and write it out.
func generateConfig(nodes *Node, nameToNode map[string]*Node) {
	outputPath, err := filepath.Abs(*outputPath)
	if err != nil {
		log.Fatal("Unable to determine absolute path for output")
	}

	content, err := ioutil.ReadFile("generator.yml")
	if err != nil {
		log.Fatalf("Error reading yml config: %s", err)
	}
	cfg := &Config{}
	err = yaml.UnmarshalStrict(content, cfg)
	if err != nil {
		log.Fatalf("Error parsing yml config: %s", err)
	}

	outputConfig := config.Config{}
	for name, m := range cfg.Modules {
		log.Infof("Generating config for module %s", name)
		// Give each module a copy of the tree so that it can be modified.
		mNodes := nodes.Copy()
		// Build the map with new pointers.
		mNameToNode := map[string]*Node{}
		walkNode(mNodes, func(n *Node) {
			mNameToNode[n.Oid] = n
			mNameToNode[n.Label] = n
		})
		outputConfig[name] = generateConfigModule(m, mNodes, mNameToNode)
		outputConfig[name].WalkParams = m.WalkParams
		log.Infof("Generated %d metrics for module %s", len(outputConfig[name].Metrics), name)
	}

	config.DoNotHideSecrets = true
	out, err := yaml.Marshal(outputConfig)
	config.DoNotHideSecrets = false
	if err != nil {
		log.Fatalf("Error marshaling yml: %s", err)
	}

	// Check the generated config to catch auth/version issues.
	err = yaml.UnmarshalStrict(out, &config.Config{})
	if err != nil {
		log.Fatalf("Error parsing generated config: %s", err)
	}

	f, err := os.Create(outputPath)
	if err != nil {
		log.Fatalf("Error opening output file: %s", err)
	}
	out = append([]byte("# WARNING: This file was auto-generated using snmp_exporter generator, manual changes will be lost.\n"), out...)
	_, err = f.Write(out)
	if err != nil {
		log.Fatalf("Error writing to output file: %s", err)
	}
	log.Infof("Config written to %s", outputPath)
}

var (
	failOnParseErrors  = kingpin.Flag("fail-on-parse-errors", "Exit with a non-zero status if there are MIB parsing errors").Default("false").Bool()
	generateCommand    = kingpin.Command("generate", "Generate snmp.yml from generator.yml")
	outputPath         = generateCommand.Flag("output-path", "Path to to write resulting config file").Default("snmp.yml").Short('o').String()
	parseErrorsCommand = kingpin.Command("parse_errors", "Debug: Print the parse errors output by NetSNMP")
	dumpCommand        = kingpin.Command("dump", "Debug: Dump the parsed and prepared MIBs")
)

func main() {
	log.AddFlags(kingpin.CommandLine)
	kingpin.HelpFlag.Short('h')
	command := kingpin.Parse()

	parseOutput := strings.TrimSpace(initSNMP())
	parseErrors := len(parseOutput) != 0
	if parseErrors {
		log.Warnf("NetSNMP reported %d parse error(s)", len(strings.Split(parseOutput, "\n")))
	}

	nodes := getMIBTree()
	nameToNode := prepareTree(nodes)

	switch command {
	case generateCommand.FullCommand():
		generateConfig(nodes, nameToNode)
	case parseErrorsCommand.FullCommand():
		fmt.Println(parseOutput)
	case dumpCommand.FullCommand():
		walkNode(nodes, func(n *Node) {
			t := n.Type
			if n.FixedSize != 0 {
				t = fmt.Sprintf("%s(%d)", n.Type, n.FixedSize)
			}
			implied := ""
			if n.ImpliedIndex {
				implied = "(implied)"
			}
			fmt.Printf("%s %s %s %q %q %s%s %v %s\n",
				n.Oid, n.Label, t, n.TextualConvention, n.Hint, n.Indexes, implied, n.EnumValues, n.Description)
		})
	}
	if *failOnParseErrors && parseErrors {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/valyala/fasthttp"
	"net"
	"net/url"
	"os"
	"time"
)

var bytesSlash = []byte("/") // heap optimization
var byteZero = []byte{0}     // heap optimization

// HTTPClient is a reusable HTTP Client.
type HTTPClient struct {
	client     fasthttp.Client
	Host       []byte
	HostString string
	debug      int
}

// HTTPClientDoOptions wraps options uses when calling `Do`.
type HTTPClientDoOptions struct {
	Debug                int
	PrettyPrintResponses bool
}

// NewHTTPClient creates a new HTTPClient.
func NewHTTPClient(host string, debug int, timeout time.Duration) *HTTPClient {
	return &HTTPClient{
		client: fasthttp.Client{
			Name: "query_benchmarker",
			Dial: func(addr string) (net.Conn, error) {
				return fasthttp.DialTimeout(addr, timeout)
			},
			MaxIdleConnDuration: 1*time.Hour,
			ReadTimeout: 60*time.Second,
			WriteTimeout: 60*time.Second,
		},
		Host:       []byte(host),
		HostString: host,
		debug:      debug,
	}
}

// Do performs the action specified by the given Query. It uses fasthttp, and
// tries to minimize heap allocations.
func (w *HTTPClient) Do(q *Query, opts *HTTPClientDoOptions) (lag float64, err error) {
	// populate uri from the reusable byte slice:
	uri := make([]byte, 0, 100)
	uri = append(uri, w.Host...)
	uri = append(uri, bytesSlash...)
	uri = append(uri, q.Path...)

	// populate a request with data from the Query:
	req := fasthttp.AcquireRequest()
	defer fasthttp.ReleaseRequest(req)

	req.Header.SetMethodBytes(q.Method)
	req.Header.SetRequestURIBytes(uri)
	req.SetBody(q.Body)
	// Perform the request while tracking latency:
	resp := fasthttp.AcquireResponse()
	defer fasthttp.ReleaseResponse(resp)
	start := time.Now()
	err = w.client.Do(req, resp)
	lag = float64(time.Since(start).Nanoseconds()) / 1e6 // milliseconds

	if (err != nil || resp.StatusCode() != fasthttp.StatusOK) && opts.Debug == 5 {
		values, _ := url.ParseQuery(string(uri))
		fmt.Printf("debug: url: %s, path %s, parsed url - %s\n", string(uri), q.Path, values)
	}

	// Check that the status code was 200 OK:
	if err == nil {
		sc := resp.StatusCode()
		if sc != fasthttp.StatusOK {
			err = fmt.Errorf("Invalid write response (status %d): %s", sc, resp.Body())
			return
		}
	}

	if opts != nil {
		// Print debug messages, if applicable:
		switch opts.Debug {
		case 1:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms\n", q.HumanLabel, lag)
		case 2:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms -- %s\n", q.HumanLabel, lag, q.HumanDescription)
		case 3:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms -- %s\n", q.HumanLabel, lag, q.HumanDescription)
			fmt.Fprintf(os.Stderr, "debug:   request: %s\n", string(q.String()))
		case 4:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms -- %s\n", q.HumanLabel, lag, q.HumanDescription)
			fmt.Fprintf(os.Stderr, "debug:   request: %s\n", string(q.String()))
			fmt.Fprintf(os.Stderr, "debug:   response: %s\n", string(resp.Body()))
		default:
		}

		// Pretty print JSON responses, if applicable:
		if opts.PrettyPrintResponses {
			// InfluxQL responses are in JSON and can be pretty-printed here.
			// Flux responses are just simple CSV.

			prefix := fmt.Sprintf("ID %d: ", q.ID)
			if json.Valid(resp.Body()) {
				var pretty bytes.Buffer
				err = json.Indent(&pretty, resp.Body(), prefix, "  ")
				if err != nil {
					return
				}

				_, err = fmt.Fprintf(os.Stderr, "%s%s\n", prefix, pretty)
				if err != nil {
					return
				}
			} else {
				_, err = fmt.Fprintf(os.Stderr, "%s%s\n", prefix, resp.Body())
				if err != nil {
					return
				}
			}
		}
	}

	return lag, err
}
<commit_msg>Added read and write timeouts<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/valyala/fasthttp"
	"net"
	"net/url"
	"os"
	"time"
)

var bytesSlash = []byte("/") // heap optimization
var byteZero = []byte{0}     // heap optimization

// HTTPClient is a reusable HTTP Client.
type HTTPClient struct {
	client     fasthttp.Client
	Host       []byte
	HostString string
	debug      int
}

// HTTPClientDoOptions wraps options uses when calling `Do`.
type HTTPClientDoOptions struct {
	Debug                int
	PrettyPrintResponses bool
}

// NewHTTPClient creates a new HTTPClient.
func NewHTTPClient(host string, debug int, timeout time.Duration) *HTTPClient {
	return &HTTPClient{
		client: fasthttp.Client{
			Name: "query_benchmarker",
			Dial: func(addr string) (net.Conn, error) {
				return fasthttp.DialTimeout(addr, timeout)
			},
			MaxIdleConnDuration: 1 * time.Hour,
			ReadTimeout:         60 * time.Second,
			WriteTimeout:        60 * time.Second,
		},
		Host:       []byte(host),
		HostString: host,
		debug:      debug,
	}
}

// Do performs the action specified by the given Query. It uses fasthttp, and
// tries to minimize heap allocations.
func (w *HTTPClient) Do(q *Query, opts *HTTPClientDoOptions) (lag float64, err error) {
	// populate uri from the reusable byte slice:
	uri := make([]byte, 0, 100)
	uri = append(uri, w.Host...)
	uri = append(uri, bytesSlash...)
	uri = append(uri, q.Path...)

	// populate a request with data from the Query:
	req := fasthttp.AcquireRequest()
	defer fasthttp.ReleaseRequest(req)

	req.Header.SetMethodBytes(q.Method)
	req.Header.SetRequestURIBytes(uri)
	req.SetBody(q.Body)
	// Perform the request while tracking latency:
	resp := fasthttp.AcquireResponse()
	defer fasthttp.ReleaseResponse(resp)
	start := time.Now()
	err = w.client.Do(req, resp)
	lag = float64(time.Since(start).Nanoseconds()) / 1e6 // milliseconds

	if (err != nil || resp.StatusCode() != fasthttp.StatusOK) && opts.Debug == 5 {
		values, _ := url.ParseQuery(string(uri))
		fmt.Printf("debug: url: %s, path %s, parsed url - %s\n", string(uri), q.Path, values)
	}

	// Check that the status code was 200 OK:
	if err == nil {
		sc := resp.StatusCode()
		if sc != fasthttp.StatusOK {
			err = fmt.Errorf("Invalid write response (status %d): %s", sc, resp.Body())
			return
		}
	}

	if opts != nil {
		// Print debug messages, if applicable:
		switch opts.Debug {
		case 1:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms\n", q.HumanLabel, lag)
		case 2:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms -- %s\n", q.HumanLabel, lag, q.HumanDescription)
		case 3:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms -- %s\n", q.HumanLabel, lag, q.HumanDescription)
			fmt.Fprintf(os.Stderr, "debug:   request: %s\n", string(q.String()))
		case 4:
			fmt.Fprintf(os.Stderr, "debug: %s in %7.2fms -- %s\n", q.HumanLabel, lag, q.HumanDescription)
			fmt.Fprintf(os.Stderr, "debug:   request: %s\n", string(q.String()))
			fmt.Fprintf(os.Stderr, "debug:   response: %s\n", string(resp.Body()))
		default:
		}

		// Pretty print JSON responses, if applicable:
		if opts.PrettyPrintResponses {
			// InfluxQL responses are in JSON and can be pretty-printed here.
			// Flux responses are just simple CSV.

			prefix := fmt.Sprintf("ID %d: ", q.ID)
			if json.Valid(resp.Body()) {
				var pretty bytes.Buffer
				err = json.Indent(&pretty, resp.Body(), prefix, "  ")
				if err != nil {
					return
				}

				_, err = fmt.Fprintf(os.Stderr, "%s%s\n", prefix, pretty)
				if err != nil {
					return
				}
			} else {
				_, err = fmt.Fprintf(os.Stderr, "%s%s\n", prefix, resp.Body())
				if err != nil {
					return
				}
			}
		}
	}

	return lag, err
}
<|endoftext|>


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

package main

import (
	"bytes"
	"container/vector"
	"exec"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"patch"
	"path/filepath"
	"sort"
	"strings"
)

var checkSync = flag.Bool("checksync", true, "check whether repository is out of sync")

func usage() {
	fmt.Fprintf(os.Stderr, "usage: hgpatch [options] [patchfile]\n")
	flag.PrintDefaults()
	os.Exit(2)
}

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

	args := flag.Args()
	var data []byte
	var err os.Error
	switch len(args) {
	case 0:
		data, err = ioutil.ReadAll(os.Stdin)
	case 1:
		data, err = ioutil.ReadFile(args[0])
	default:
		usage()
	}
	chk(err)

	pset, err := patch.Parse(data)
	chk(err)

	// Change to hg root directory, because
	// patch paths are relative to root.
	root, err := hgRoot()
	chk(err)
	chk(os.Chdir(root))

	// Make sure there are no pending changes on the server.
	if *checkSync && hgIncoming() {
		fmt.Fprintf(os.Stderr, "incoming changes waiting; run hg sync first\n")
		os.Exit(2)
	}

	// Make sure we won't be editing files with local pending changes.
	dirtylist, err := hgModified()
	chk(err)
	dirty := make(map[string]bool)
	for _, f := range dirtylist {
		dirty[f] = true
	}
	conflict := make(map[string]bool)
	for _, f := range pset.File {
		if f.Verb == patch.Delete || f.Verb == patch.Rename {
			if dirty[f.Src] {
				conflict[f.Src] = true
			}
		}
		if f.Verb != patch.Delete {
			if dirty[f.Dst] {
				conflict[f.Dst] = true
			}
		}
	}
	if len(conflict) > 0 {
		fmt.Fprintf(os.Stderr, "cannot apply patch to locally modified files:\n")
		for name := range conflict {
			fmt.Fprintf(os.Stderr, "\t%s\n", name)
		}
		os.Exit(2)
	}

	// Apply changes in memory.
	op, err := pset.Apply(ioutil.ReadFile)
	chk(err)

	// Write changes to disk copy: order of commands matters.
	// Accumulate undo log as we go, in case there is an error.
	// Also accumulate list of modified files to print at end.
	changed := make(map[string]int)

	// Copy, Rename create the destination file, so they
	// must happen before we write the data out.
	// A single patch may have a Copy and a Rename
	// with the same source, so we have to run all the
	// Copy in one pass, then all the Rename.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Copy {
			makeParent(o.Dst)
			chk(hgCopy(o.Dst, o.Src))
			undoRevert(o.Dst)
			changed[o.Dst] = 1
		}
	}
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Rename {
			makeParent(o.Dst)
			chk(hgRename(o.Dst, o.Src))
			undoRevert(o.Dst)
			undoRevert(o.Src)
			changed[o.Src] = 1
			changed[o.Dst] = 1
		}
	}

	// Run Delete before writing to files in case one of the
	// deleted paths is becoming a directory.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Delete {
			chk(hgRemove(o.Src))
			undoRevert(o.Src)
			changed[o.Src] = 1
		}
	}

	// Write files.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Delete {
			continue
		}
		if o.Verb == patch.Add {
			makeParent(o.Dst)
			changed[o.Dst] = 1
		}
		if o.Data != nil {
			chk(ioutil.WriteFile(o.Dst, o.Data, 0644))
			if o.Verb == patch.Add {
				undoRm(o.Dst)
			} else {
				undoRevert(o.Dst)
			}
			changed[o.Dst] = 1
		}
		if o.Mode != 0 {
			chk(os.Chmod(o.Dst, uint32(o.Mode&0755)))
			undoRevert(o.Dst)
			changed[o.Dst] = 1
		}
	}

	// hg add looks at the destination file, so it must happen
	// after we write the data out.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Add {
			chk(hgAdd(o.Dst))
			undoRevert(o.Dst)
			changed[o.Dst] = 1
		}
	}

	// Finished editing files.  Write the list of changed files to stdout.
	list := make([]string, len(changed))
	i := 0
	for f := range changed {
		list[i] = f
		i++
	}
	sort.Strings(list)
	for _, f := range list {
		fmt.Printf("%s\n", f)
	}
}

// make parent directory for name, if necessary
func makeParent(name string) {
	parent, _ := filepath.Split(name)
	chk(mkdirAll(parent, 0755))
}

// Copy of os.MkdirAll but adds to undo log after
// creating a directory.
func mkdirAll(path string, perm uint32) os.Error {
	dir, err := os.Lstat(path)
	if err == nil {
		if dir.IsDirectory() {
			return nil
		}
		return &os.PathError{"mkdir", path, os.ENOTDIR}
	}

	i := len(path)
	for i > 0 && path[i-1] == '/' { // Skip trailing slashes.
		i--
	}

	j := i
	for j > 0 && path[j-1] != '/' { // Scan backward over element.
		j--
	}

	if j > 0 {
		err = mkdirAll(path[0:j-1], perm)
		if err != nil {
			return err
		}
	}

	err = os.Mkdir(path, perm)
	if err != nil {
		// Handle arguments like "foo/." by
		// double-checking that directory doesn't exist.
		dir, err1 := os.Lstat(path)
		if err1 == nil && dir.IsDirectory() {
			return nil
		}
		return err
	}
	undoRm(path)
	return nil
}

// If err != nil, process the undo log and exit.
func chk(err os.Error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		runUndo()
		os.Exit(2)
	}
}

// Undo log
type undo func() os.Error

var undoLog vector.Vector // vector of undo

func undoRevert(name string) { undoLog.Push(undo(func() os.Error { return hgRevert(name) })) }

func undoRm(name string) { undoLog.Push(undo(func() os.Error { return os.Remove(name) })) }

func runUndo() {
	for i := undoLog.Len() - 1; i >= 0; i-- {
		if err := undoLog.At(i).(undo)(); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err)
		}
	}
}

// hgRoot returns the root directory of the repository.
func hgRoot() (string, os.Error) {
	out, err := run([]string{"hg", "root"}, nil)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(out), nil
}

// hgIncoming returns true if hg sync will pull in changes.
func hgIncoming() bool {
	// hg -q incoming exits 0 when there is nothing incoming, 1 otherwise.
	_, err := run([]string{"hg", "-q", "incoming"}, nil)
	return err == nil
}

// hgModified returns a list of the modified files in the
// repository.
func hgModified() ([]string, os.Error) {
	out, err := run([]string{"hg", "status", "-n"}, nil)
	if err != nil {
		return nil, err
	}
	return strings.Split(strings.TrimSpace(out), "\n"), nil
}

// hgAdd adds name to the repository.
func hgAdd(name string) os.Error {
	_, err := run([]string{"hg", "add", name}, nil)
	return err
}

// hgRemove removes name from the repository.
func hgRemove(name string) os.Error {
	_, err := run([]string{"hg", "rm", name}, nil)
	return err
}

// hgRevert reverts name.
func hgRevert(name string) os.Error {
	_, err := run([]string{"hg", "revert", name}, nil)
	return err
}

// hgCopy copies src to dst in the repository.
// Note that the argument order matches io.Copy, not "hg cp".
func hgCopy(dst, src string) os.Error {
	_, err := run([]string{"hg", "cp", src, dst}, nil)
	return err
}

// hgRename renames src to dst in the repository.
// Note that the argument order matches io.Copy, not "hg mv".
func hgRename(dst, src string) os.Error {
	_, err := run([]string{"hg", "mv", src, dst}, nil)
	return err
}

func dup(a []string) []string {
	b := make([]string, len(a))
	copy(b, a)
	return b
}

var lookPathCache = make(map[string]string)

// run runs the command argv, resolving argv[0] if necessary by searching $PATH.
// It provides input on standard input to the command.
func run(argv []string, input []byte) (out string, err os.Error) {
	if len(argv) < 1 {
		return "", &runError{dup(argv), os.EINVAL}
	}

	prog, ok := lookPathCache[argv[0]]
	if !ok {
		prog, err = exec.LookPath(argv[0])
		if err != nil {
			return "", &runError{dup(argv), err}
		}
		lookPathCache[argv[0]] = prog
	}

	cmd := exec.Command(prog, argv[1:]...)
	if len(input) > 0 {
		cmd.Stdin = bytes.NewBuffer(input)
	}
	bs, err := cmd.CombinedOutput()
	if err != nil {
		return "", &runError{dup(argv), err}
	}
	return string(bs), nil
}

// A runError represents an error that occurred while running a command.
type runError struct {
	cmd []string
	err os.Error
}

func (e *runError) String() string { return strings.Join(e.cmd, " ") + ": " + e.err.String() }
<commit_msg>cmd/hgpatch: remove container/vector<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"exec"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"patch"
	"path/filepath"
	"sort"
	"strings"
)

var checkSync = flag.Bool("checksync", true, "check whether repository is out of sync")

func usage() {
	fmt.Fprintf(os.Stderr, "usage: hgpatch [options] [patchfile]\n")
	flag.PrintDefaults()
	os.Exit(2)
}

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

	args := flag.Args()
	var data []byte
	var err os.Error
	switch len(args) {
	case 0:
		data, err = ioutil.ReadAll(os.Stdin)
	case 1:
		data, err = ioutil.ReadFile(args[0])
	default:
		usage()
	}
	chk(err)

	pset, err := patch.Parse(data)
	chk(err)

	// Change to hg root directory, because
	// patch paths are relative to root.
	root, err := hgRoot()
	chk(err)
	chk(os.Chdir(root))

	// Make sure there are no pending changes on the server.
	if *checkSync && hgIncoming() {
		fmt.Fprintf(os.Stderr, "incoming changes waiting; run hg sync first\n")
		os.Exit(2)
	}

	// Make sure we won't be editing files with local pending changes.
	dirtylist, err := hgModified()
	chk(err)
	dirty := make(map[string]bool)
	for _, f := range dirtylist {
		dirty[f] = true
	}
	conflict := make(map[string]bool)
	for _, f := range pset.File {
		if f.Verb == patch.Delete || f.Verb == patch.Rename {
			if dirty[f.Src] {
				conflict[f.Src] = true
			}
		}
		if f.Verb != patch.Delete {
			if dirty[f.Dst] {
				conflict[f.Dst] = true
			}
		}
	}
	if len(conflict) > 0 {
		fmt.Fprintf(os.Stderr, "cannot apply patch to locally modified files:\n")
		for name := range conflict {
			fmt.Fprintf(os.Stderr, "\t%s\n", name)
		}
		os.Exit(2)
	}

	// Apply changes in memory.
	op, err := pset.Apply(ioutil.ReadFile)
	chk(err)

	// Write changes to disk copy: order of commands matters.
	// Accumulate undo log as we go, in case there is an error.
	// Also accumulate list of modified files to print at end.
	changed := make(map[string]int)

	// Copy, Rename create the destination file, so they
	// must happen before we write the data out.
	// A single patch may have a Copy and a Rename
	// with the same source, so we have to run all the
	// Copy in one pass, then all the Rename.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Copy {
			makeParent(o.Dst)
			chk(hgCopy(o.Dst, o.Src))
			undoRevert(o.Dst)
			changed[o.Dst] = 1
		}
	}
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Rename {
			makeParent(o.Dst)
			chk(hgRename(o.Dst, o.Src))
			undoRevert(o.Dst)
			undoRevert(o.Src)
			changed[o.Src] = 1
			changed[o.Dst] = 1
		}
	}

	// Run Delete before writing to files in case one of the
	// deleted paths is becoming a directory.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Delete {
			chk(hgRemove(o.Src))
			undoRevert(o.Src)
			changed[o.Src] = 1
		}
	}

	// Write files.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Delete {
			continue
		}
		if o.Verb == patch.Add {
			makeParent(o.Dst)
			changed[o.Dst] = 1
		}
		if o.Data != nil {
			chk(ioutil.WriteFile(o.Dst, o.Data, 0644))
			if o.Verb == patch.Add {
				undoRm(o.Dst)
			} else {
				undoRevert(o.Dst)
			}
			changed[o.Dst] = 1
		}
		if o.Mode != 0 {
			chk(os.Chmod(o.Dst, uint32(o.Mode&0755)))
			undoRevert(o.Dst)
			changed[o.Dst] = 1
		}
	}

	// hg add looks at the destination file, so it must happen
	// after we write the data out.
	for i := range op {
		o := &op[i]
		if o.Verb == patch.Add {
			chk(hgAdd(o.Dst))
			undoRevert(o.Dst)
			changed[o.Dst] = 1
		}
	}

	// Finished editing files.  Write the list of changed files to stdout.
	list := make([]string, len(changed))
	i := 0
	for f := range changed {
		list[i] = f
		i++
	}
	sort.Strings(list)
	for _, f := range list {
		fmt.Printf("%s\n", f)
	}
}

// make parent directory for name, if necessary
func makeParent(name string) {
	parent, _ := filepath.Split(name)
	chk(mkdirAll(parent, 0755))
}

// Copy of os.MkdirAll but adds to undo log after
// creating a directory.
func mkdirAll(path string, perm uint32) os.Error {
	dir, err := os.Lstat(path)
	if err == nil {
		if dir.IsDirectory() {
			return nil
		}
		return &os.PathError{"mkdir", path, os.ENOTDIR}
	}

	i := len(path)
	for i > 0 && path[i-1] == '/' { // Skip trailing slashes.
		i--
	}

	j := i
	for j > 0 && path[j-1] != '/' { // Scan backward over element.
		j--
	}

	if j > 0 {
		err = mkdirAll(path[0:j-1], perm)
		if err != nil {
			return err
		}
	}

	err = os.Mkdir(path, perm)
	if err != nil {
		// Handle arguments like "foo/." by
		// double-checking that directory doesn't exist.
		dir, err1 := os.Lstat(path)
		if err1 == nil && dir.IsDirectory() {
			return nil
		}
		return err
	}
	undoRm(path)
	return nil
}

// If err != nil, process the undo log and exit.
func chk(err os.Error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		runUndo()
		os.Exit(2)
	}
}

// Undo log
type undo func() os.Error

var undoLog []undo

func undoRevert(name string) {
	undoLog = append(undoLog, undo(func() os.Error { return hgRevert(name) }))
}

func undoRm(name string) { undoLog = append(undoLog, undo(func() os.Error { return os.Remove(name) })) }

func runUndo() {
	for i := len(undoLog) - 1; i >= 0; i-- {
		if err := undoLog[i](); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err)
		}
	}
}

// hgRoot returns the root directory of the repository.
func hgRoot() (string, os.Error) {
	out, err := run([]string{"hg", "root"}, nil)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(out), nil
}

// hgIncoming returns true if hg sync will pull in changes.
func hgIncoming() bool {
	// hg -q incoming exits 0 when there is nothing incoming, 1 otherwise.
	_, err := run([]string{"hg", "-q", "incoming"}, nil)
	return err == nil
}

// hgModified returns a list of the modified files in the
// repository.
func hgModified() ([]string, os.Error) {
	out, err := run([]string{"hg", "status", "-n"}, nil)
	if err != nil {
		return nil, err
	}
	return strings.Split(strings.TrimSpace(out), "\n"), nil
}

// hgAdd adds name to the repository.
func hgAdd(name string) os.Error {
	_, err := run([]string{"hg", "add", name}, nil)
	return err
}

// hgRemove removes name from the repository.
func hgRemove(name string) os.Error {
	_, err := run([]string{"hg", "rm", name}, nil)
	return err
}

// hgRevert reverts name.
func hgRevert(name string) os.Error {
	_, err := run([]string{"hg", "revert", name}, nil)
	return err
}

// hgCopy copies src to dst in the repository.
// Note that the argument order matches io.Copy, not "hg cp".
func hgCopy(dst, src string) os.Error {
	_, err := run([]string{"hg", "cp", src, dst}, nil)
	return err
}

// hgRename renames src to dst in the repository.
// Note that the argument order matches io.Copy, not "hg mv".
func hgRename(dst, src string) os.Error {
	_, err := run([]string{"hg", "mv", src, dst}, nil)
	return err
}

func dup(a []string) []string {
	b := make([]string, len(a))
	copy(b, a)
	return b
}

var lookPathCache = make(map[string]string)

// run runs the command argv, resolving argv[0] if necessary by searching $PATH.
// It provides input on standard input to the command.
func run(argv []string, input []byte) (out string, err os.Error) {
	if len(argv) < 1 {
		return "", &runError{dup(argv), os.EINVAL}
	}

	prog, ok := lookPathCache[argv[0]]
	if !ok {
		prog, err = exec.LookPath(argv[0])
		if err != nil {
			return "", &runError{dup(argv), err}
		}
		lookPathCache[argv[0]] = prog
	}

	cmd := exec.Command(prog, argv[1:]...)
	if len(input) > 0 {
		cmd.Stdin = bytes.NewBuffer(input)
	}
	bs, err := cmd.CombinedOutput()
	if err != nil {
		return "", &runError{dup(argv), err}
	}
	return string(bs), nil
}

// A runError represents an error that occurred while running a command.
type runError struct {
	cmd []string
	err os.Error
}

func (e *runError) String() string { return strings.Join(e.cmd, " ") + ": " + e.err.String() }
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

var integrationTest bool

func init() {
	flag.BoolVar(&integrationTest, "integration-test", false, "Run integration test with actual Slack API")
}

func TestIntegration(t *testing.T) {
	if !integrationTest {
		t.Skip()
	}

	slackURL := os.Getenv("SLACK_URL")
	assert.NotEmpty(t, slackURL)

	notifications := NewNotifications(slackURL)

	userMap, err := ReadUsersFile(filepath.Join(testDataDir, "real-user.yml"))
	assert.Nil(t, err)
	users := NewUsers(userMap)

	parsedTemplates, err := ParseTemplates(testTemplatesDir)
	assert.Nil(t, err)
	templates := NewTemplates(parsedTemplates)

	shouter := NewShouter(users, templates, notifications)
	shoutHandler := NewShoutHandler(shouter)

	ts := httptest.NewServer(shoutHandler)
	defer ts.Close()

	requestBody := `
{
  "recipients":["baker"],
  "template":"years-old",
  "data":{
    "name":"Rufus",
    "years":"eleventeen"
  }
}`
	res, err := http.Post(fmt.Sprintf("%s/%s", ts.URL, "shout"),
		"application/json", strings.NewReader(requestBody))

	assert.Nil(t, err)
	assert.Equal(t, http.StatusAccepted, res.StatusCode)
}
<commit_msg>Update integration test to actual env var<commit_after>package main

import (
	"flag"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

var integrationTest bool

func init() {
	flag.BoolVar(&integrationTest, "integration-test", false, "Run integration test with actual Slack API")
}

func TestIntegration(t *testing.T) {
	if !integrationTest {
		t.Skip()
	}

	slackURL := os.Getenv("SHOUT_SLACK_URL")
	assert.NotEmpty(t, slackURL)

	notifications := NewNotifications(slackURL)

	userMap, err := ReadUsersFile(filepath.Join(testDataDir, "real-user.yml"))
	assert.Nil(t, err)
	users := NewUsers(userMap)

	parsedTemplates, err := ParseTemplates(testTemplatesDir)
	assert.Nil(t, err)
	templates := NewTemplates(parsedTemplates)

	shouter := NewShouter(users, templates, notifications)
	shoutHandler := NewShoutHandler(shouter)

	ts := httptest.NewServer(shoutHandler)
	defer ts.Close()

	requestBody := `
{
  "recipients":["baker"],
  "template":"years-old",
  "data":{
    "name":"Rufus",
    "years":"eleventeen"
  }
}`
	res, err := http.Post(fmt.Sprintf("%s/%s", ts.URL, "shout"),
		"application/json", strings.NewReader(requestBody))

	assert.Nil(t, err)
	assert.Equal(t, http.StatusAccepted, res.StatusCode)
}
<|endoftext|>


<commit_before>/*
	Copyright 2017 by Christian Hüning (christianhuening@googlemail.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 glk8smain

import (
	"fmt"
	"gitlab.informatik.haw-hamburg.de/icc/gl-k8s-integrator/usecases"
	"gitlab.informatik.haw-hamburg.de/icc/gl-k8s-integrator/webhooklistener"
	"log"
	"os"
	"os/signal"
	"syscall"
)

func Main() {
	log.Println("Gitlab K8s Integrator starting up!")
	if os.Getenv("GITLAB_HOSTNAME") == "" {
		log.Fatalln("Please provide GITLAB_HOSTNAME env!")
	}
	if os.Getenv("GITLAB_API_VERSION") == "" {
		log.Fatalln("Please provide GITLAB_API_VERSION env!")
	}
	if os.Getenv("GITLAB_PRIVATE_TOKEN") == "" {
		log.Fatalln("Please provide GITLAB_PRIVATE_TOKEN env!")
	}

	quit := make(chan int)

	// Handle System signals
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		s := <-sigc
		switch s {
		case os.Interrupt:
			quit <- 0
		case syscall.SIGTERM:
			quit <- 0
		}
	}()

	// listen in sep. routine
	go webhooklistener.Listen(quit)
	go usecases.StartRecurringSyncTimer()
	log.Println("Gitlab K8s Integrator listening!")
	// Wait until server signals quit
	select {
	case <-quit:
		fmt.Println("Goodbye!")
	}

}
<commit_msg>finished 1st version of graylog integration<commit_after>/*
	Copyright 2017 by Christian Hüning (christianhuening@googlemail.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 glk8smain

import (
	"fmt"
	"gitlab.informatik.haw-hamburg.de/icc/gl-k8s-integrator/usecases"
	"gitlab.informatik.haw-hamburg.de/icc/gl-k8s-integrator/webhooklistener"
	"log"
	"os"
	"os/signal"
	"syscall"
)

func Main() {
	log.Println("Gitlab K8s Integrator starting up!")
	if os.Getenv("GITLAB_HOSTNAME") == "" {
		log.Fatalln("Please provide GITLAB_HOSTNAME env!")
	}
	if os.Getenv("GITLAB_API_VERSION") == "" {
		log.Fatalln("Please provide GITLAB_API_VERSION env!")
	}
	if os.Getenv("GITLAB_PRIVATE_TOKEN") == "" {
		log.Fatalln("Please provide GITLAB_PRIVATE_TOKEN env!")
	}

	quit := make(chan int)

	// Handle System signals
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		s := <-sigc
		switch s {
		case os.Interrupt:
			quit <- 0
		case syscall.SIGTERM:
			quit <- 0
		}
	}()

	// listen in sep. routine
	go webhooklistener.Listen(quit)
	go usecases.StartRecurringSyncTimer()
	log.Println("Gitlab K8s Integrator listening!")

	// Perform 1st sync now:
	go usecases.PerformGlK8sSync()
	// Wait until server signals quit
	select {
	case <-quit:
		fmt.Println("Goodbye!")
	}

}
<|endoftext|>


<commit_before>package qln

import (
	"fmt"
	"log"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcutil"
	"github.com/mit-dci/lit/lndc"
	"github.com/mit-dci/lit/lnutil"
	"github.com/mit-dci/lit/portxo"
)

// Gets the list of ports where LitNode is listening for incoming connections, & the connection key
func (nd *LitNode) GetLisAddressAndPorts() (*btcutil.AddressWitnessPubKeyHash, []string, error) {
    idPriv := nd.IdKey()
    myId := btcutil.Hash160(idPriv.PubKey().SerializeCompressed())
    lisAdr, err := btcutil.NewAddressWitnessPubKeyHash(myId, nd.Param)
	if err != nil {
		return nil, nil, err
	}
	nd.RemoteMtx.Lock()
	ports := nd.LisIpPorts
	nd.RemoteMtx.Unlock()
	return lisAdr, ports, nil
}

// TCPListener starts a litNode listening for incoming LNDC connections
func (nd *LitNode) TCPListener(
	lisIpPort string) (*btcutil.AddressWitnessPubKeyHash, error) {
	idPriv := nd.IdKey()
	listener, err := lndc.NewListener(nd.IdKey(), lisIpPort)
	if err != nil {
		return nil, err
	}
	myId := btcutil.Hash160(idPriv.PubKey().SerializeCompressed())

	// listening pubkey hash should be in a format independent of subwallets.
	lisAdr, err := btcutil.NewAddressWitnessPubKeyHash(
		myId, &chaincfg.TestNet3Params)
	// Really we need to make LN-specific encodings here ^^

	if err != nil {
		return nil, err
	}
	fmt.Printf("Listening on %s\n", listener.Addr().String())
	fmt.Printf("Listening with base58 address: %s lnid: %x\n",
		lisAdr.String(), myId[:16])

	go func() {
		for {
			netConn, err := listener.Accept() // this blocks
			if err != nil {
				log.Printf("Listener error: %s\n", err.Error())
				continue
			}
			newConn, ok := netConn.(*lndc.LNDConn)
			if !ok {
				fmt.Printf("Got something that wasn't a LNDC")
				continue
			}
			fmt.Printf("Incomming connection from %x on %s\n",
				newConn.RemotePub.SerializeCompressed(), newConn.RemoteAddr().String())

			// don't save host/port for incomming connections
			peerIdx, err := nd.GetPeerIdx(newConn.RemotePub, "")
			if err != nil {
				log.Printf("Listener error: %s\n", err.Error())
				continue
			}

			nd.RemoteMtx.Lock()
			var peer RemotePeer
			peer.Idx = peerIdx
			peer.Con = newConn
			nd.RemoteCons[peerIdx] = &peer
			nd.RemoteMtx.Unlock()

			// each connection to a peer gets its own LNDCReader
			go nd.LNDCReader(&peer)
		}
	}()
	nd.RemoteMtx.Lock()
	nd.LisIpPorts = append(nd.LisIpPorts, lisIpPort)
	nd.RemoteMtx.Unlock()
	return lisAdr, nil
}

// DialPeer makes an outgoing connection to another node.
func (nd *LitNode) DialPeer(lnAdr *lndc.LNAdr) error {
	// get my private ID key
	idPriv := nd.IdKey()

	// Assign remote connection
	newConn := new(lndc.LNDConn)

	var id []byte
	if lnAdr.PubKey != nil {
		id = lnAdr.PubKey.SerializeCompressed()
	} else {
		id = lnAdr.Base58Adr.ScriptAddress()
	}

	err := newConn.Dial(idPriv, lnAdr.NetAddr.String(), id)
	if err != nil {
		return err
	}

	// if connect is successful, either query for already existing peer index, or
	// if the peer is new, make an new index, and save the hostname&port

	// figure out peer index, or assign new one for new peer.  Since
	// we're connecting out, also specify the hostname&port
	peerIdx, err := nd.GetPeerIdx(newConn.RemotePub, newConn.RemoteAddr().String())
	if err != nil {
		return err
	}

	nd.RemoteMtx.Lock()
	var p RemotePeer
	p.Con = newConn
	p.Idx = peerIdx
	nd.RemoteCons[peerIdx] = &p
	nd.RemoteMtx.Unlock()

	// each connection to a peer gets its own LNDCReader
	go nd.LNDCReader(&p)

	return nil
}

// OutMessager takes messages from the outbox and sends them to the ether. net.
func (nd *LitNode) OutMessager() {
	for {
		msg := <-nd.OmniOut
		if !nd.ConnectedToPeer(msg.PeerIdx) {
			fmt.Printf("message type %x to peer %d but not connected\n",
				msg.MsgType, msg.PeerIdx)
			continue
		}

		rawmsg := append([]byte{msg.MsgType}, msg.Data...)
		nd.RemoteMtx.Lock() // not sure this is needed...
		n, err := nd.RemoteCons[msg.PeerIdx].Con.Write(rawmsg)
		if err != nil {
			fmt.Printf("error writing to peer %d: %s\n", msg.PeerIdx, err.Error())
		} else {
			fmt.Printf("type %x %d bytes to peer %d\n", msg.MsgType, n, msg.PeerIdx)
		}
		nd.RemoteMtx.Unlock()
	}
}

type PeerInfo struct {
	PeerNumber uint32
	RemoteHost string
}

func (nd *LitNode) GetConnectedPeerList() []PeerInfo {
	nd.RemoteMtx.Lock()
	nd.RemoteMtx.Unlock() //TODO: This unlock is in the wrong place...?
	var peers []PeerInfo
	for k, v := range nd.RemoteCons {
		var newPeer PeerInfo
		newPeer.PeerNumber = k
		newPeer.RemoteHost = v.Con.RemoteAddr().String()
		peers = append(peers, newPeer)
	}
	return peers
}

// ConnectedToPeer checks whether you're connected to a specific peer
func (nd *LitNode) ConnectedToPeer(peer uint32) bool {
	nd.RemoteMtx.Lock()
	_, ok := nd.RemoteCons[peer]
	nd.RemoteMtx.Unlock()
	return ok
}

// IdKey returns the identity private key
func (nd *LitNode) IdKey() *btcec.PrivateKey {
	var kg portxo.KeyGen
	kg.Depth = 5
	kg.Step[0] = 44 | 1<<31
	kg.Step[1] = 0 | 1<<31
	kg.Step[2] = 9 | 1<<31
	kg.Step[3] = 0 | 1<<31
	kg.Step[4] = 0 | 1<<31
	return nd.SubWallet.GetPriv(kg)
}

// SendChat sends a text string to a peer
func (nd *LitNode) SendChat(peer uint32, chat string) error {
	if !nd.ConnectedToPeer(peer) {
		return fmt.Errorf("Not connected to peer %d", peer)
	}

	outMsg := new(lnutil.LitMsg)
	outMsg.MsgType = lnutil.MSGID_TEXTCHAT
	outMsg.PeerIdx = peer
	outMsg.Data = []byte(chat)
	nd.OmniOut <- outMsg

	return nil
}
<commit_msg>fix formatting, set litNode adr to testnet3 params<commit_after>package qln

import (
	"fmt"
	"log"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcutil"
	"github.com/mit-dci/lit/lndc"
	"github.com/mit-dci/lit/lnutil"
	"github.com/mit-dci/lit/portxo"
)

// Gets the list of ports where LitNode is listening for incoming connections,
// & the connection key
func (nd *LitNode) GetLisAddressAndPorts() (
	*btcutil.AddressWitnessPubKeyHash, []string, error) {

	idPriv := nd.IdKey()
	myId := btcutil.Hash160(idPriv.PubKey().SerializeCompressed())
	// litNode addresses transcend network parameters
	lisAdr, err := btcutil.NewAddressWitnessPubKeyHash(myId,
		&chaincfg.TestNet3Params)
	if err != nil {
		return nil, nil, err
	}
	nd.RemoteMtx.Lock()
	ports := nd.LisIpPorts
	nd.RemoteMtx.Unlock()
	return lisAdr, ports, nil
}

// TCPListener starts a litNode listening for incoming LNDC connections
func (nd *LitNode) TCPListener(
	lisIpPort string) (*btcutil.AddressWitnessPubKeyHash, error) {
	idPriv := nd.IdKey()
	listener, err := lndc.NewListener(nd.IdKey(), lisIpPort)
	if err != nil {
		return nil, err
	}
	myId := btcutil.Hash160(idPriv.PubKey().SerializeCompressed())

	// listening pubkey hash should be in a format independent of subwallets.
	lisAdr, err := btcutil.NewAddressWitnessPubKeyHash(
		myId, &chaincfg.TestNet3Params)
	// Really we need to make LN-specific encodings here ^^

	if err != nil {
		return nil, err
	}
	fmt.Printf("Listening on %s\n", listener.Addr().String())
	fmt.Printf("Listening with base58 address: %s lnid: %x\n",
		lisAdr.String(), myId[:16])

	go func() {
		for {
			netConn, err := listener.Accept() // this blocks
			if err != nil {
				log.Printf("Listener error: %s\n", err.Error())
				continue
			}
			newConn, ok := netConn.(*lndc.LNDConn)
			if !ok {
				fmt.Printf("Got something that wasn't a LNDC")
				continue
			}
			fmt.Printf("Incomming connection from %x on %s\n",
				newConn.RemotePub.SerializeCompressed(), newConn.RemoteAddr().String())

			// don't save host/port for incomming connections
			peerIdx, err := nd.GetPeerIdx(newConn.RemotePub, "")
			if err != nil {
				log.Printf("Listener error: %s\n", err.Error())
				continue
			}

			nd.RemoteMtx.Lock()
			var peer RemotePeer
			peer.Idx = peerIdx
			peer.Con = newConn
			nd.RemoteCons[peerIdx] = &peer
			nd.RemoteMtx.Unlock()

			// each connection to a peer gets its own LNDCReader
			go nd.LNDCReader(&peer)
		}
	}()
	nd.RemoteMtx.Lock()
	nd.LisIpPorts = append(nd.LisIpPorts, lisIpPort)
	nd.RemoteMtx.Unlock()
	return lisAdr, nil
}

// DialPeer makes an outgoing connection to another node.
func (nd *LitNode) DialPeer(lnAdr *lndc.LNAdr) error {
	// get my private ID key
	idPriv := nd.IdKey()

	// Assign remote connection
	newConn := new(lndc.LNDConn)

	var id []byte
	if lnAdr.PubKey != nil {
		id = lnAdr.PubKey.SerializeCompressed()
	} else {
		id = lnAdr.Base58Adr.ScriptAddress()
	}

	err := newConn.Dial(idPriv, lnAdr.NetAddr.String(), id)
	if err != nil {
		return err
	}

	// if connect is successful, either query for already existing peer index, or
	// if the peer is new, make an new index, and save the hostname&port

	// figure out peer index, or assign new one for new peer.  Since
	// we're connecting out, also specify the hostname&port
	peerIdx, err := nd.GetPeerIdx(newConn.RemotePub, newConn.RemoteAddr().String())
	if err != nil {
		return err
	}

	nd.RemoteMtx.Lock()
	var p RemotePeer
	p.Con = newConn
	p.Idx = peerIdx
	nd.RemoteCons[peerIdx] = &p
	nd.RemoteMtx.Unlock()

	// each connection to a peer gets its own LNDCReader
	go nd.LNDCReader(&p)

	return nil
}

// OutMessager takes messages from the outbox and sends them to the ether. net.
func (nd *LitNode) OutMessager() {
	for {
		msg := <-nd.OmniOut
		if !nd.ConnectedToPeer(msg.PeerIdx) {
			fmt.Printf("message type %x to peer %d but not connected\n",
				msg.MsgType, msg.PeerIdx)
			continue
		}

		rawmsg := append([]byte{msg.MsgType}, msg.Data...)
		nd.RemoteMtx.Lock() // not sure this is needed...
		n, err := nd.RemoteCons[msg.PeerIdx].Con.Write(rawmsg)
		if err != nil {
			fmt.Printf("error writing to peer %d: %s\n", msg.PeerIdx, err.Error())
		} else {
			fmt.Printf("type %x %d bytes to peer %d\n", msg.MsgType, n, msg.PeerIdx)
		}
		nd.RemoteMtx.Unlock()
	}
}

type PeerInfo struct {
	PeerNumber uint32
	RemoteHost string
}

func (nd *LitNode) GetConnectedPeerList() []PeerInfo {
	nd.RemoteMtx.Lock()
	nd.RemoteMtx.Unlock() //TODO: This unlock is in the wrong place...?
	var peers []PeerInfo
	for k, v := range nd.RemoteCons {
		var newPeer PeerInfo
		newPeer.PeerNumber = k
		newPeer.RemoteHost = v.Con.RemoteAddr().String()
		peers = append(peers, newPeer)
	}
	return peers
}

// ConnectedToPeer checks whether you're connected to a specific peer
func (nd *LitNode) ConnectedToPeer(peer uint32) bool {
	nd.RemoteMtx.Lock()
	_, ok := nd.RemoteCons[peer]
	nd.RemoteMtx.Unlock()
	return ok
}

// IdKey returns the identity private key
func (nd *LitNode) IdKey() *btcec.PrivateKey {
	var kg portxo.KeyGen
	kg.Depth = 5
	kg.Step[0] = 44 | 1<<31
	kg.Step[1] = 0 | 1<<31
	kg.Step[2] = 9 | 1<<31
	kg.Step[3] = 0 | 1<<31
	kg.Step[4] = 0 | 1<<31
	return nd.SubWallet.GetPriv(kg)
}

// SendChat sends a text string to a peer
func (nd *LitNode) SendChat(peer uint32, chat string) error {
	if !nd.ConnectedToPeer(peer) {
		return fmt.Errorf("Not connected to peer %d", peer)
	}

	outMsg := new(lnutil.LitMsg)
	outMsg.MsgType = lnutil.MSGID_TEXTCHAT
	outMsg.PeerIdx = peer
	outMsg.Data = []byte(chat)
	nd.OmniOut <- outMsg

	return nil
}
<|endoftext|>


<commit_before>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2014 Douglas Hall
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package r900

import (
	"bytes"
	"fmt"
	"math"
	"strconv"

	"github.com/bemasher/gf"
	"github.com/bemasher/rtlamr/decode"
	"github.com/bemasher/rtlamr/parse"
)

const (
	PayloadSymbols = 42
)

func NewPacketConfig(symbolLength int) (cfg decode.PacketConfig) {
	cfg.DataRate = 32768

	cfg.CenterFreq = 912380000

	cfg.SymbolLength = symbolLength
	cfg.SymbolLength2 = cfg.SymbolLength << 1

	cfg.SampleRate = cfg.DataRate * cfg.SymbolLength

	cfg.PreambleSymbols = 32
	cfg.PacketSymbols = 116

	cfg.PreambleLength = cfg.PreambleSymbols * cfg.SymbolLength2
	cfg.PacketLength = cfg.PacketSymbols * cfg.SymbolLength2

	cfg.BlockSize = decode.NextPowerOf2(cfg.PreambleLength)
	cfg.BlockSize2 = cfg.BlockSize << 1

	cfg.BufferLength = cfg.PacketLength + cfg.BlockSize

	cfg.Preamble = "00000000000000001110010101100100"

	return
}

type Parser struct {
	decode.Decoder
	field *gf.Field
	rsBuf [31]byte

	csum      []float64
	filtered  [][3]float64
	quantized []byte
}

func NewParser(symbolLength int, fastMag bool) (p Parser) {
	p.Decoder = decode.NewDecoder(NewPacketConfig(symbolLength), fastMag)

	// GF of order 32, polynomial 37, generator 2.
	p.field = gf.NewField(32, 37, 2)

	p.csum = make([]float64, p.Decoder.Cfg.BufferLength+1)
	p.filtered = make([][3]float64, p.Decoder.Cfg.BufferLength)
	p.quantized = make([]byte, p.Decoder.Cfg.BufferLength)

	return
}

func (p Parser) Dec() decode.Decoder {
	return p.Decoder
}

func (p Parser) Cfg() decode.PacketConfig {
	return p.Decoder.Cfg
}

// Perform matched filtering.
func (p Parser) Filter() {
	// This function computes the convolution of each symbol kernel with the
	// signal. The naive approach requires for each symbol to calculate the
	// summation of samples between a pair of indices.

	// 0 |--------|
	// 1   |--------|
	// 2     |--------|
	// 3       |--------|

	// To avoid redundant calculations we compute the cumulative sum of the
	// signal. This reduces each summation to the difference between the two
	// indices of the cumulative sum.

	var sum float64
	for idx, v := range p.Decoder.Signal {
		sum += v
		p.csum[idx+1] = sum
	}

	// There are six symbols, composed of three base symbols and their bitwise
	// inversions. Compute the convolution of each base symbol with the
	// signal.

	// 1100 -> 0011
	// 1010 -> 0101
	// 1001 -> 0110

	// This is basically unreadable because of a lot of algebraic
	// simplification but is necessary for efficiency.
	for idx := 0; idx < p.Decoder.Cfg.BufferLength-p.Decoder.Cfg.SymbolLength*4; idx++ {
		c0 := p.csum[idx]
		c1 := p.csum[idx+p.Decoder.Cfg.SymbolLength] * 2
		c2 := p.csum[idx+p.Decoder.Cfg.SymbolLength*2] * 2
		c3 := p.csum[idx+p.Decoder.Cfg.SymbolLength*3] * 2
		c4 := p.csum[idx+p.Decoder.Cfg.SymbolLength*4]

		p.filtered[idx][0] = c2 - c4 - c0           // 1100
		p.filtered[idx][1] = c1 - c2 + c3 - c4 - c0 // 1010
		p.filtered[idx][2] = c1 - c3 + c4 - c0      // 1001
	}
}

// Determine the symbol that exists at each sample of the signal.
func (p Parser) Quantize() {
	// 0 0011, 3 1100
	// 1 0101, 4 1010
	// 2 0110, 5 1001

	for idx, vec := range p.filtered {
		argmax := byte(0)
		max := math.Abs(vec[0])

		// If v1 is larger than v0, update max and argmax.
		if v1 := math.Abs(vec[1]); v1 > max {
			max = v1
			argmax = 1
		}

		// If v2 is larger than the greater of v1 or v0, update max and argmax.
		if v2 := math.Abs(vec[2]); v2 > max {
			max = v2
			argmax = 2
		}

		// Set the output symbol index.
		p.quantized[idx] = argmax

		// If the sign is negative, jump to the index of the inverted symbol.
		if vec[argmax] > 0 {
			p.quantized[idx] += 3
		}
	}
}

// Given a list of indices the preamble exists at, decode and parse a message.
func (p Parser) Parse(indices []int) (msgs []parse.Message) {
	p.Filter()
	p.Quantize()

	preambleLength := p.Decoder.Cfg.PreambleLength
	symbolLength := p.Decoder.Cfg.SymbolLength

	symbols := make([]byte, 21)
	zeros := make([]byte, 5)

	seen := make(map[string]bool)

	for _, preambleIdx := range indices {
		if preambleIdx > p.Decoder.Cfg.BlockSize {
			break
		}

		payloadIdx := preambleIdx + preambleLength
		var digits string
		for idx := 0; idx < PayloadSymbols*4*p.Decoder.Cfg.SymbolLength; idx += symbolLength * 4 {
			qIdx := payloadIdx + idx

			digits += strconv.Itoa(int(p.quantized[qIdx]))
		}

		var (
			bits      string
			badSymbol bool
		)
		for idx := 0; idx < len(digits); idx += 2 {
			symbol, _ := strconv.ParseInt(digits[idx:idx+2], 6, 32)
			if symbol > 31 {
				badSymbol = true
				break
			}
			symbols[idx>>1] = byte(symbol)
			bits += fmt.Sprintf("%05b", symbol)
		}

		if badSymbol || seen[bits] {
			continue
		}

		seen[bits] = true

		copy(p.rsBuf[:], symbols[:16])
		copy(p.rsBuf[26:], symbols[16:])
		syndromes := p.field.Syndrome(p.rsBuf[:], 5, 29)

		if !bytes.Equal(zeros, syndromes) {
			continue
		}

		id, _ := strconv.ParseUint(bits[:32], 2, 32)
		unkn1, _ := strconv.ParseUint(bits[32:40], 2, 8)
		unkn2, _ := strconv.ParseUint(bits[40:48], 2, 8)
		consumption, _ := strconv.ParseUint(bits[48:72], 2, 24)
		unkn3, _ := strconv.ParseUint(bits[72:74], 2, 2)
		unkn4, _ := strconv.ParseUint(bits[74:80], 2, 6)

		var r900 R900

		r900.ID = uint32(id)
		r900.Unkn1 = uint8(unkn1)
		r900.Unkn2 = uint8(unkn2)
		r900.Consumption = uint32(consumption)
		r900.Unkn3 = uint8(unkn3)
		r900.Unkn4 = uint8(unkn4)

		msgs = append(msgs, r900)
	}

	return
}

type R900 struct {
	ID          uint32 `xml:",attr"` // 32 bits
	Unkn1       uint8  `xml:",attr"` // 8 bits
	Unkn2       uint8  `xml:",attr"` // 8 bits
	Consumption uint32 `xml:",attr"` // 24 bits
	Unkn3       uint8  `xml:",attr"` // 2 bits
	Unkn4       uint8  `xml:",attr"` // 6 bits
}

func (r900 R900) MsgType() string {
	return "R900"
}

func (r900 R900) MeterID() uint32 {
	return r900.ID
}

func (r900 R900) MeterType() uint8 {
	return r900.Unkn1
}

func (r900 R900) String() string {
	return fmt.Sprintf("{ID:%10d Unkn1:0x%02X Unkn2:0x%02X Consumption:%8d Unkn3:0x%02X Unkn4:0x%02X}",
		r900.ID,
		r900.Unkn1,
		r900.Unkn2,
		r900.Consumption,
		r900.Unkn3,
		r900.Unkn4,
	)
}

func (r900 R900) Record() (r []string) {
	r = append(r, strconv.FormatUint(uint64(r900.ID), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn1), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn2), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Consumption), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn3), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn4), 10))

	return
}
<commit_msg>Fix galois field package path.<commit_after>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2014 Douglas Hall
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package r900

import (
	"bytes"
	"fmt"
	"math"
	"strconv"

	"github.com/bemasher/rtlamr/decode"
	"github.com/bemasher/rtlamr/parse"
	"github.com/bemasher/rtlamr/r900/gf"
)

const (
	PayloadSymbols = 42
)

func NewPacketConfig(symbolLength int) (cfg decode.PacketConfig) {
	cfg.DataRate = 32768

	cfg.CenterFreq = 912380000

	cfg.SymbolLength = symbolLength
	cfg.SymbolLength2 = cfg.SymbolLength << 1

	cfg.SampleRate = cfg.DataRate * cfg.SymbolLength

	cfg.PreambleSymbols = 32
	cfg.PacketSymbols = 116

	cfg.PreambleLength = cfg.PreambleSymbols * cfg.SymbolLength2
	cfg.PacketLength = cfg.PacketSymbols * cfg.SymbolLength2

	cfg.BlockSize = decode.NextPowerOf2(cfg.PreambleLength)
	cfg.BlockSize2 = cfg.BlockSize << 1

	cfg.BufferLength = cfg.PacketLength + cfg.BlockSize

	cfg.Preamble = "00000000000000001110010101100100"

	return
}

type Parser struct {
	decode.Decoder
	field *gf.Field
	rsBuf [31]byte

	csum      []float64
	filtered  [][3]float64
	quantized []byte
}

func NewParser(symbolLength int, fastMag bool) (p Parser) {
	p.Decoder = decode.NewDecoder(NewPacketConfig(symbolLength), fastMag)

	// GF of order 32, polynomial 37, generator 2.
	p.field = gf.NewField(32, 37, 2)

	p.csum = make([]float64, p.Decoder.Cfg.BufferLength+1)
	p.filtered = make([][3]float64, p.Decoder.Cfg.BufferLength)
	p.quantized = make([]byte, p.Decoder.Cfg.BufferLength)

	return
}

func (p Parser) Dec() decode.Decoder {
	return p.Decoder
}

func (p Parser) Cfg() decode.PacketConfig {
	return p.Decoder.Cfg
}

// Perform matched filtering.
func (p Parser) Filter() {
	// This function computes the convolution of each symbol kernel with the
	// signal. The naive approach requires for each symbol to calculate the
	// summation of samples between a pair of indices.

	// 0 |--------|
	// 1   |--------|
	// 2     |--------|
	// 3       |--------|

	// To avoid redundant calculations we compute the cumulative sum of the
	// signal. This reduces each summation to the difference between the two
	// indices of the cumulative sum.

	var sum float64
	for idx, v := range p.Decoder.Signal {
		sum += v
		p.csum[idx+1] = sum
	}

	// There are six symbols, composed of three base symbols and their bitwise
	// inversions. Compute the convolution of each base symbol with the
	// signal.

	// 1100 -> 0011
	// 1010 -> 0101
	// 1001 -> 0110

	// This is basically unreadable because of a lot of algebraic
	// simplification but is necessary for efficiency.
	for idx := 0; idx < p.Decoder.Cfg.BufferLength-p.Decoder.Cfg.SymbolLength*4; idx++ {
		c0 := p.csum[idx]
		c1 := p.csum[idx+p.Decoder.Cfg.SymbolLength] * 2
		c2 := p.csum[idx+p.Decoder.Cfg.SymbolLength*2] * 2
		c3 := p.csum[idx+p.Decoder.Cfg.SymbolLength*3] * 2
		c4 := p.csum[idx+p.Decoder.Cfg.SymbolLength*4]

		p.filtered[idx][0] = c2 - c4 - c0           // 1100
		p.filtered[idx][1] = c1 - c2 + c3 - c4 - c0 // 1010
		p.filtered[idx][2] = c1 - c3 + c4 - c0      // 1001
	}
}

// Determine the symbol that exists at each sample of the signal.
func (p Parser) Quantize() {
	// 0 0011, 3 1100
	// 1 0101, 4 1010
	// 2 0110, 5 1001

	for idx, vec := range p.filtered {
		argmax := byte(0)
		max := math.Abs(vec[0])

		// If v1 is larger than v0, update max and argmax.
		if v1 := math.Abs(vec[1]); v1 > max {
			max = v1
			argmax = 1
		}

		// If v2 is larger than the greater of v1 or v0, update max and argmax.
		if v2 := math.Abs(vec[2]); v2 > max {
			max = v2
			argmax = 2
		}

		// Set the output symbol index.
		p.quantized[idx] = argmax

		// If the sign is negative, jump to the index of the inverted symbol.
		if vec[argmax] > 0 {
			p.quantized[idx] += 3
		}
	}
}

// Given a list of indices the preamble exists at, decode and parse a message.
func (p Parser) Parse(indices []int) (msgs []parse.Message) {
	p.Filter()
	p.Quantize()

	preambleLength := p.Decoder.Cfg.PreambleLength
	symbolLength := p.Decoder.Cfg.SymbolLength

	symbols := make([]byte, 21)
	zeros := make([]byte, 5)

	seen := make(map[string]bool)

	for _, preambleIdx := range indices {
		if preambleIdx > p.Decoder.Cfg.BlockSize {
			break
		}

		payloadIdx := preambleIdx + preambleLength
		var digits string
		for idx := 0; idx < PayloadSymbols*4*p.Decoder.Cfg.SymbolLength; idx += symbolLength * 4 {
			qIdx := payloadIdx + idx

			digits += strconv.Itoa(int(p.quantized[qIdx]))
		}

		var (
			bits      string
			badSymbol bool
		)
		for idx := 0; idx < len(digits); idx += 2 {
			symbol, _ := strconv.ParseInt(digits[idx:idx+2], 6, 32)
			if symbol > 31 {
				badSymbol = true
				break
			}
			symbols[idx>>1] = byte(symbol)
			bits += fmt.Sprintf("%05b", symbol)
		}

		if badSymbol || seen[bits] {
			continue
		}

		seen[bits] = true

		copy(p.rsBuf[:], symbols[:16])
		copy(p.rsBuf[26:], symbols[16:])
		syndromes := p.field.Syndrome(p.rsBuf[:], 5, 29)

		if !bytes.Equal(zeros, syndromes) {
			continue
		}

		id, _ := strconv.ParseUint(bits[:32], 2, 32)
		unkn1, _ := strconv.ParseUint(bits[32:40], 2, 8)
		unkn2, _ := strconv.ParseUint(bits[40:48], 2, 8)
		consumption, _ := strconv.ParseUint(bits[48:72], 2, 24)
		unkn3, _ := strconv.ParseUint(bits[72:74], 2, 2)
		unkn4, _ := strconv.ParseUint(bits[74:80], 2, 6)

		var r900 R900

		r900.ID = uint32(id)
		r900.Unkn1 = uint8(unkn1)
		r900.Unkn2 = uint8(unkn2)
		r900.Consumption = uint32(consumption)
		r900.Unkn3 = uint8(unkn3)
		r900.Unkn4 = uint8(unkn4)

		msgs = append(msgs, r900)
	}

	return
}

type R900 struct {
	ID          uint32 `xml:",attr"` // 32 bits
	Unkn1       uint8  `xml:",attr"` // 8 bits
	Unkn2       uint8  `xml:",attr"` // 8 bits
	Consumption uint32 `xml:",attr"` // 24 bits
	Unkn3       uint8  `xml:",attr"` // 2 bits
	Unkn4       uint8  `xml:",attr"` // 6 bits
}

func (r900 R900) MsgType() string {
	return "R900"
}

func (r900 R900) MeterID() uint32 {
	return r900.ID
}

func (r900 R900) MeterType() uint8 {
	return r900.Unkn1
}

func (r900 R900) String() string {
	return fmt.Sprintf("{ID:%10d Unkn1:0x%02X Unkn2:0x%02X Consumption:%8d Unkn3:0x%02X Unkn4:0x%02X}",
		r900.ID,
		r900.Unkn1,
		r900.Unkn2,
		r900.Consumption,
		r900.Unkn3,
		r900.Unkn4,
	)
}

func (r900 R900) Record() (r []string) {
	r = append(r, strconv.FormatUint(uint64(r900.ID), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn1), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn2), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Consumption), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn3), 10))
	r = append(r, strconv.FormatUint(uint64(r900.Unkn4), 10))

	return
}
<|endoftext|>


<commit_before>package main
// Package main runs the petulant-lana server.

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"os"
	"strconv"
	"strings"
)

// Define the types

type configuration struct {
	Name           string `json: "name"`
	Url            string `json: "url"`
	CallbackSecret string `json: "callbacksecret"`
	BasePrice      int    `json: "baseprice"`
	MinimumPrice   int    `json: "minprice"`
	ApiKey         string `json: "coinbasekey"`
}

type transactionResult struct {
	Success bool `json:"success"`
	Button  struct {
		Code string `json:"code"`
	} `json:"button"`
}

type callbackResult struct {
	Order struct {
		Filename string `json:"custom"`
	} `json:"order"`
}

// Create the configuration
var config = configuration{}

// Do stuff

// Get an appropriate name for the file.
func newFileName(fname string) string {
	result := strings.Replace(strings.Replace(fname, "/", "-", -1), " ", "-", -1)
	if _, err := os.Stat("f/" + result); os.IsNotExist(err) {
		if _, err := os.Stat("tmp/" + result); os.IsNotExist(err) {
			// Don't do anything.
		} else {
			result = newFileName("p" + result)
		}
	} else {
		result = newFileName("p" + result)
	}
	return result
}

// Create a coinbase button.
func createButton(n string, p int) string {
	coinbaserequest := "{ \"button\": {" +
		"\"name\": \"One-Time Hosting Purchase\"," +
		"\"type\": \"buy_now\"," +
		"\"price_string\": \"" + strconv.FormatFloat(float64(p)/float64(100000000), 'f', 8, 64) + "\"," +
		"\"price_currency_iso\": \"BTC\"," +
		"\"custom\": \"" + n + "\"," +
		"\"callback_url\": \"whatever\"," +
		"\"description\": \"Indefinite storage of the provided file. Your file will be available at: http://btcdl.bearbin.net/f/" + n + " when the transaction processes.\"," +
		"\"type\": \"buy_now\"," +
		"\"style\": \"custom_large\"" +
		"} }"
	fmt.Println(coinbaserequest)
	request_body := bytes.NewBuffer([]byte(coinbaserequest))

	client := &http.Client{}
	req, err := http.NewRequest("POST", "https://coinbase.com/api/v1/buttons?api_key="+config.ApiKey, request_body)
	if err != nil {
		fmt.Println(err)
	}

	req.Header.Add("content-type", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		fmt.Println(err)
	}

	response_body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
	}
	defer resp.Body.Close()
	res := transactionResult{}
	fmt.Println(string(response_body))
	err = json.Unmarshal(response_body, &res)
	return res.Button.Code

}

// hello world, the web server
func upload(w http.ResponseWriter, req *http.Request) {

	// Get the form file.
	file, header, err := req.FormFile("file")
	if err != nil {
		fmt.Println(err)
		return
	}

	// Get the name for the file.
	fileName := newFileName(header.Filename)
	log.Print("Uploaded new file: ", fileName)

	data, err := ioutil.ReadAll(file)
	if err != nil {
		log.Print(err)
		return
	}
	err = ioutil.WriteFile("tmp/"+fileName, data, 0777)
	if err != nil {
		log.Print(err)
		return
	}

	// Get file size.
	supfil, _ := os.Stat("tmp/" + fileName)
	fileSize := math.Floor(float64(supfil.Size()) / 1024)
	price := int(math.Floor(float64(config.BasePrice) * (fileSize/1024)))
	if price < config.MinimumPrice {
		price = config.MinimumPrice
	}
	// Redirect the user.
	http.Redirect(w, req, "https://coinbase.com/checkouts/"+createButton(fileName, price), 302)

}

func coinbaseCallback(w http.ResponseWriter, req *http.Request) {
	fmt.Println("LELELELE")
	body, _ := ioutil.ReadAll(req.Body)
	res := callbackResult{}
	fmt.Println(body)
	json.Unmarshal([]byte(body), &res)
	fmt.Println(res.Order.Filename)
	os.Rename("tmp/"+res.Order.Filename, "f/"+res.Order.Filename)
}

func MainPage(w http.ResponseWriter, req *http.Request) {
	t, _ := template.ParseFiles("index.html")
	t.Execute(w, "")
}

func main() {
	// Inititalize the config.
	configFile, err := os.Open("config.json")
	if err != nil {
		log.Fatal("Failed to open config: ", err)
	}
	decoder := json.NewDecoder(configFile)

	err = decoder.Decode(&config)
	if err != nil {
		log.Fatal("Failed to open config: ", err)
	}

	// Main page
	http.HandleFunc("/", MainPage)
	// Upload page
	http.HandleFunc("/upload", upload)
	// Coinbase callback
	http.HandleFunc("/wheatver", coinbaseCallback)
	// Static files
	http.Handle("/f/", http.FileServer(http.Dir("")))

	// Try and serve port 80.
	err = http.ListenAndServe(":80", nil)
	if err != nil {
		// Failed for some reason, try port 8080
		log.Print("Failed to bind to port 80, trying 8080.")
		err := http.ListenAndServe(":8080", nil)
		if err != nil {
			// Failed.
			log.Fatal("ListenAndServe: ", err)
		}
	}
}
<commit_msg>Add a better filename collision prevention mechanism.<commit_after>package main
// Package main runs the petulant-lana server.

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"os"
	"rand"
	"strconv"
	"strings"
)

// Define the types

type configuration struct {
	Name           string `json: "name"`
	Url            string `json: "url"`
	CallbackSecret string `json: "callbacksecret"`
	BasePrice      int    `json: "baseprice"`
	MinimumPrice   int    `json: "minprice"`
	ApiKey         string `json: "coinbasekey"`
}

type transactionResult struct {
	Success bool `json:"success"`
	Button  struct {
		Code string `json:"code"`
	} `json:"button"`
}

type callbackResult struct {
	Order struct {
		Filename string `json:"custom"`
	} `json:"order"`
}

// Create the configuration
var config = configuration{}

// Do stuff

// Get an appropriate name for the file.
func newFileName(fname string) string {
	// First, remove slashes and spaces, replace with dashes.
	newName := strings.Replace(strings.Replace(fname, "/", "-", -1), " ", "-", -1)

	// Does the current file already exist, in storage.
	if _, err := os.Stat("f/" + newName); os.IsNotExist(err) {
		// Does the current file already exist, in temporary storage.
		if _, err := os.Stat("tmp/" + newName); os.IsNotExist(err) {
			// Don't do anything.
		} else {
			// Add a random number onto the front of the filename.
			// This is not the best method, but it does for now.
			randomLetter := fmt.Sprint(rand.Int())
			newName = newFileName(randomLetter + newName)
		}
	} else {
		// Add a random number onto the front of the filename.
		// This is not the best method, but it does for now.
		randomLetter := fmt.Sprint(rand.Int())
		newName = newFileName(randomLetter + newName)
	}
	return newName
}

// Create a coinbase button.
func createButton(n string, p int) string {
	coinbaserequest := "{ \"button\": {" +
		"\"name\": \"One-Time Hosting Purchase\"," +
		"\"type\": \"buy_now\"," +
		"\"price_string\": \"" + strconv.FormatFloat(float64(p)/float64(100000000), 'f', 8, 64) + "\"," +
		"\"price_currency_iso\": \"BTC\"," +
		"\"custom\": \"" + n + "\"," +
		"\"callback_url\": \"whatever\"," +
		"\"description\": \"Indefinite storage of the provided file. Your file will be available at: http://btcdl.bearbin.net/f/" + n + " when the transaction processes.\"," +
		"\"type\": \"buy_now\"," +
		"\"style\": \"custom_large\"" +
		"} }"
	fmt.Println(coinbaserequest)
	request_body := bytes.NewBuffer([]byte(coinbaserequest))

	client := &http.Client{}
	req, err := http.NewRequest("POST", "https://coinbase.com/api/v1/buttons?api_key="+config.ApiKey, request_body)
	if err != nil {
		fmt.Println(err)
	}

	req.Header.Add("content-type", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		fmt.Println(err)
	}

	response_body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
	}
	defer resp.Body.Close()
	res := transactionResult{}
	fmt.Println(string(response_body))
	err = json.Unmarshal(response_body, &res)
	return res.Button.Code

}

// hello world, the web server
func upload(w http.ResponseWriter, req *http.Request) {

	// Get the form file.
	file, header, err := req.FormFile("file")
	if err != nil {
		fmt.Println(err)
		return
	}

	// Get the name for the file.
	fileName := newFileName(header.Filename)
	log.Print("Uploaded new file: ", fileName)

	data, err := ioutil.ReadAll(file)
	if err != nil {
		log.Print(err)
		return
	}
	err = ioutil.WriteFile("tmp/"+fileName, data, 0777)
	if err != nil {
		log.Print(err)
		return
	}

	// Get file size.
	supfil, _ := os.Stat("tmp/" + fileName)
	fileSize := math.Floor(float64(supfil.Size()) / 1024)
	price := int(math.Floor(float64(config.BasePrice) * (fileSize/1024)))
	if price < config.MinimumPrice {
		price = config.MinimumPrice
	}
	// Redirect the user.
	http.Redirect(w, req, "https://coinbase.com/checkouts/"+createButton(fileName, price), 302)

}

func coinbaseCallback(w http.ResponseWriter, req *http.Request) {
	fmt.Println("LELELELE")
	body, _ := ioutil.ReadAll(req.Body)
	res := callbackResult{}
	fmt.Println(body)
	json.Unmarshal([]byte(body), &res)
	fmt.Println(res.Order.Filename)
	os.Rename("tmp/"+res.Order.Filename, "f/"+res.Order.Filename)
}

func MainPage(w http.ResponseWriter, req *http.Request) {
	t, _ := template.ParseFiles("index.html")
	t.Execute(w, "")
}

func main() {
	// Inititalize the config.
	configFile, err := os.Open("config.json")
	if err != nil {
		log.Fatal("Failed to open config: ", err)
	}
	decoder := json.NewDecoder(configFile)

	err = decoder.Decode(&config)
	if err != nil {
		log.Fatal("Failed to open config: ", err)
	}

	// Main page
	http.HandleFunc("/", MainPage)
	// Upload page
	http.HandleFunc("/upload", upload)
	// Coinbase callback
	http.HandleFunc("/wheatver", coinbaseCallback)
	// Static files
	http.Handle("/f/", http.FileServer(http.Dir("")))

	// Try and serve port 80.
	err = http.ListenAndServe(":80", nil)
	if err != nil {
		// Failed for some reason, try port 8080
		log.Print("Failed to bind to port 80, trying 8080.")
		err := http.ListenAndServe(":8080", nil)
		if err != nil {
			// Failed.
			log.Fatal("ListenAndServe: ", err)
		}
	}
}
<|endoftext|>