<commit_before>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bytes"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"sort"
	"strings"
	"text/template"
	"time"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cactus/go-statsd-client/statsd"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"

	"github.com/letsencrypt/boulder/cmd"
	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/mail"
	"github.com/letsencrypt/boulder/sa"
)

type emailContent struct {
	ExpirationDate   time.Time
	DaysToExpiration int
	DNSNames         string
}

type mailer struct {
	stats         statsd.Statter
	log           *blog.AuditLogger
	dbMap         *gorp.DbMap
	mailer        mail.Mailer
	emailTemplate *template.Template
	nagTimes      []time.Duration
	limit         int
}

func (m *mailer) sendNags(parsedCert *x509.Certificate, contacts []core.AcmeURL) error {
	expiresIn := int(parsedCert.NotAfter.Sub(time.Now()).Hours()/24) + 1
	emails := []string{}
	for _, contact := range contacts {
		if contact.Scheme == "mailto" {
			emails = append(emails, contact.Opaque)
		}
	}
	if len(emails) > 0 {
		email := emailContent{
			ExpirationDate:   parsedCert.NotAfter,
			DaysToExpiration: expiresIn,
			DNSNames:         strings.Join(parsedCert.DNSNames, ", "),
		}
		msgBuf := new(bytes.Buffer)
		err := m.emailTemplate.Execute(msgBuf, email)
		if err != nil {
			m.stats.Inc("Mailer.Errors.SendingNag.TemplateFailure", 1, 1.0)
			return err
		}
		err = m.mailer.SendMail(emails, msgBuf.String())
		if err != nil {
			m.stats.Inc("Mailer.Errors.SendingNag.SendFailure", 1, 1.0)
			return err
		}
		m.stats.Inc("Mailer.Expiration.Sent", int64(len(emails)), 1.0)
	}
	return nil
}

func (m *mailer) updateCertStatus(serial string) error {
	// Update CertificateStatus object
	tx, err := m.dbMap.Begin()
	if err != nil {
		m.log.Err(fmt.Sprintf("Error opening transaction for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}

	csObj, err := tx.Get(&core.CertificateStatus{}, serial)
	if err != nil {
		m.log.Err(fmt.Sprintf("Error fetching status for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}
	certStatus := csObj.(*core.CertificateStatus)
	certStatus.LastExpirationNagSent = time.Now()

	_, err = tx.Update(certStatus)
	if err != nil {
		m.log.Err(fmt.Sprintf("Error updating status for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}

	err = tx.Commit()
	if err != nil {
		m.log.Err(fmt.Sprintf("Error commiting transaction for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}

	return nil
}

func (m *mailer) processCerts(certs []core.Certificate) {
	m.log.Info(fmt.Sprintf("expiration-mailer: Found %d certificates, starting sending messages", len(certs)))
	for _, cert := range certs {
		regObj, err := m.dbMap.Get(&core.Registration{}, cert.RegistrationID)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error fetching registration %d: %s", cert.RegistrationID, err))
			m.stats.Inc("Mailer.Errors.GetRegistration", 1, 1.0)
			continue
		}
		reg := regObj.(*core.Registration)
		parsedCert, err := x509.ParseCertificate(cert.DER)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error parsing certificate %s: %s", cert.Serial, err))
			m.stats.Inc("Mailer.Errors.ParseCertificate", 1, 1.0)
			continue
		}
		err = m.sendNags(parsedCert, reg.Contact)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error sending nag emails: %s", err))
			m.stats.Inc("Mailer.Errors.SendingNags", 1, 1.0)
			continue
		}
		err = m.updateCertStatus(cert.Serial)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error updating certificate status for %s: %s", cert.Serial, err))
			m.stats.Inc("Mailer.Errors.UpdateCertificateStatus", 1, 1.0)
			continue
		}
	}
	m.log.Info("expiration-mailer: Finished sending messages")
	return
}

func (m *mailer) findExpiringCertificates() error {
	now := time.Now()
	// E.g. m.NagTimes = [1, 3, 7, 14] days from expiration
	for i, expiresIn := range m.nagTimes {
		left := now
		if i > 0 {
			left = left.Add(m.nagTimes[i-1])
		}
		right := now.Add(expiresIn)

		m.log.Info(fmt.Sprintf("expiration-mailer: Searching for certificates that expire between %s and %s", left, right))
		var certs []core.Certificate
		_, err := m.dbMap.Select(
			&certs,
			`SELECT cert.* FROM certificates AS cert
			 JOIN certificateStatus AS cs
			 ON cs.serial = cert.serial
			 WHERE cert.expires > :cutoffA
			 AND cert.expires < :cutoffB
			 AND cert.status != "revoked"
			 AND cs.lastExpirationNagSent <= :nagCutoff
       ORDER BY cert.expires ASC
			 LIMIT :limit`,
			map[string]interface{}{
				"cutoffA":   left,
				"cutoffB":   right,
				"nagCutoff": time.Now().Add(-expiresIn),
				"limit":     m.limit,
			},
		)
		if err != nil {
			m.log.Err(fmt.Sprintf("expiration-mailer: Error loading certificates: %s", err))
			return err // fatal
		}
		if len(certs) > 0 {
			m.processCerts(certs)
		}
	}

	return nil
}

type durationSlice []time.Duration

func (ds durationSlice) Len() int {
	return len(ds)
}

func (ds durationSlice) Less(a, b int) bool {
	return ds[a] < ds[b]
}

func (ds durationSlice) Swap(a, b int) {
	ds[a], ds[b] = ds[b], ds[a]
}

func main() {
	app := cmd.NewAppShell("expiration-mailer")

	app.App.Flags = append(app.App.Flags, cli.IntFlag{
		Name:   "cert_limit",
		Value:  100,
		EnvVar: "CERT_LIMIT",
		Usage:  "Count of certificates to process per expiration period",
	})

	app.Config = func(c *cli.Context, config cmd.Config) cmd.Config {
		if c.GlobalInt("cert_limit") > 0 {
			config.Mailer.CertLimit = c.GlobalInt("cert_limit")
		}
		return config
	}

	app.Action = func(c cmd.Config) {
		// Set up logging
		stats, err := statsd.NewClient(c.Statsd.Server, c.Statsd.Prefix)
		cmd.FailOnError(err, "Couldn't connect to statsd")

		auditlogger, err := blog.Dial(c.Syslog.Network, c.Syslog.Server, c.Syslog.Tag, stats)
		cmd.FailOnError(err, "Could not connect to Syslog")

		// AUDIT[ Error Conditions ] 9cc4d537-8534-4970-8665-4b382abe82f3
		defer auditlogger.AuditPanic()

		blog.SetAuditLogger(auditlogger)

		auditlogger.Info(app.VersionString())

		go cmd.DebugServer(c.Mailer.DebugAddr)

		// Configure DB
		dbMap, err := sa.NewDbMap(c.Mailer.DBDriver, c.Mailer.DBConnect)
		cmd.FailOnError(err, "Could not connect to database")

		// Load email template
		emailTmpl, err := ioutil.ReadFile(c.Mailer.EmailTemplate)
		cmd.FailOnError(err, fmt.Sprintf("Could not read email template file [%s]", c.Mailer.EmailTemplate))
		tmpl, err := template.New("expiry-email").Parse(string(emailTmpl))
		cmd.FailOnError(err, "Could not parse email template")

		mailClient := mail.New(c.Mailer.Server, c.Mailer.Port, c.Mailer.Username, c.Mailer.Password)

		var nags durationSlice
		for _, nagDuration := range c.Mailer.NagTimes {
			dur, err := time.ParseDuration(nagDuration)
			if err != nil {
				auditlogger.Err(fmt.Sprintf("Failed to parse nag duration string [%s]: %s", nagDuration, err))
				return
			}
			nags = append(nags, dur)
		}
		// Make sure durations are sorted in increasing order
		sort.Sort(nags)

		m := mailer{
			stats:         stats,
			log:           auditlogger,
			dbMap:         dbMap,
			mailer:        &mailClient,
			emailTemplate: tmpl,
			nagTimes:      nags,
			limit:         c.Mailer.CertLimit,
		}

		auditlogger.Info("expiration-mailer: Starting")
		err = m.findExpiringCertificates()
		cmd.FailOnError(err, "expiration-mailer has failed")
	}

	app.Run()
}
<commit_msg>Fix formatting<commit_after>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bytes"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"sort"
	"strings"
	"text/template"
	"time"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cactus/go-statsd-client/statsd"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"

	"github.com/letsencrypt/boulder/cmd"
	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/mail"
	"github.com/letsencrypt/boulder/sa"
)

type emailContent struct {
	ExpirationDate   time.Time
	DaysToExpiration int
	DNSNames         string
}

type mailer struct {
	stats         statsd.Statter
	log           *blog.AuditLogger
	dbMap         *gorp.DbMap
	mailer        mail.Mailer
	emailTemplate *template.Template
	nagTimes      []time.Duration
	limit         int
}

func (m *mailer) sendNags(parsedCert *x509.Certificate, contacts []core.AcmeURL) error {
	expiresIn := int(parsedCert.NotAfter.Sub(time.Now()).Hours()/24) + 1
	emails := []string{}
	for _, contact := range contacts {
		if contact.Scheme == "mailto" {
			emails = append(emails, contact.Opaque)
		}
	}
	if len(emails) > 0 {
		email := emailContent{
			ExpirationDate:   parsedCert.NotAfter,
			DaysToExpiration: expiresIn,
			DNSNames:         strings.Join(parsedCert.DNSNames, ", "),
		}
		msgBuf := new(bytes.Buffer)
		err := m.emailTemplate.Execute(msgBuf, email)
		if err != nil {
			m.stats.Inc("Mailer.Errors.SendingNag.TemplateFailure", 1, 1.0)
			return err
		}
		err = m.mailer.SendMail(emails, msgBuf.String())
		if err != nil {
			m.stats.Inc("Mailer.Errors.SendingNag.SendFailure", 1, 1.0)
			return err
		}
		m.stats.Inc("Mailer.Expiration.Sent", int64(len(emails)), 1.0)
	}
	return nil
}

func (m *mailer) updateCertStatus(serial string) error {
	// Update CertificateStatus object
	tx, err := m.dbMap.Begin()
	if err != nil {
		m.log.Err(fmt.Sprintf("Error opening transaction for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}

	csObj, err := tx.Get(&core.CertificateStatus{}, serial)
	if err != nil {
		m.log.Err(fmt.Sprintf("Error fetching status for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}
	certStatus := csObj.(*core.CertificateStatus)
	certStatus.LastExpirationNagSent = time.Now()

	_, err = tx.Update(certStatus)
	if err != nil {
		m.log.Err(fmt.Sprintf("Error updating status for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}

	err = tx.Commit()
	if err != nil {
		m.log.Err(fmt.Sprintf("Error commiting transaction for certificate %s: %s", serial, err))
		tx.Rollback()
		return err
	}

	return nil
}

func (m *mailer) processCerts(certs []core.Certificate) {
	m.log.Info(fmt.Sprintf("expiration-mailer: Found %d certificates, starting sending messages", len(certs)))
	for _, cert := range certs {
		regObj, err := m.dbMap.Get(&core.Registration{}, cert.RegistrationID)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error fetching registration %d: %s", cert.RegistrationID, err))
			m.stats.Inc("Mailer.Errors.GetRegistration", 1, 1.0)
			continue
		}
		reg := regObj.(*core.Registration)
		parsedCert, err := x509.ParseCertificate(cert.DER)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error parsing certificate %s: %s", cert.Serial, err))
			m.stats.Inc("Mailer.Errors.ParseCertificate", 1, 1.0)
			continue
		}
		err = m.sendNags(parsedCert, reg.Contact)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error sending nag emails: %s", err))
			m.stats.Inc("Mailer.Errors.SendingNags", 1, 1.0)
			continue
		}
		err = m.updateCertStatus(cert.Serial)
		if err != nil {
			m.log.Err(fmt.Sprintf("Error updating certificate status for %s: %s", cert.Serial, err))
			m.stats.Inc("Mailer.Errors.UpdateCertificateStatus", 1, 1.0)
			continue
		}
	}
	m.log.Info("expiration-mailer: Finished sending messages")
	return
}

func (m *mailer) findExpiringCertificates() error {
	now := time.Now()
	// E.g. m.NagTimes = [1, 3, 7, 14] days from expiration
	for i, expiresIn := range m.nagTimes {
		left := now
		if i > 0 {
			left = left.Add(m.nagTimes[i-1])
		}
		right := now.Add(expiresIn)

		m.log.Info(fmt.Sprintf("expiration-mailer: Searching for certificates that expire between %s and %s", left, right))
		var certs []core.Certificate
		_, err := m.dbMap.Select(
			&certs,
			`SELECT cert.* FROM certificates AS cert
			 JOIN certificateStatus AS cs
			 ON cs.serial = cert.serial
			 WHERE cert.expires > :cutoffA
			 AND cert.expires < :cutoffB
			 AND cert.status != "revoked"
			 AND cs.lastExpirationNagSent <= :nagCutoff
			 ORDER BY cert.expires ASC
			 LIMIT :limit`,
			map[string]interface{}{
				"cutoffA":   left,
				"cutoffB":   right,
				"nagCutoff": time.Now().Add(-expiresIn),
				"limit":     m.limit,
			},
		)
		if err != nil {
			m.log.Err(fmt.Sprintf("expiration-mailer: Error loading certificates: %s", err))
			return err // fatal
		}
		if len(certs) > 0 {
			m.processCerts(certs)
		}
	}

	return nil
}

type durationSlice []time.Duration

func (ds durationSlice) Len() int {
	return len(ds)
}

func (ds durationSlice) Less(a, b int) bool {
	return ds[a] < ds[b]
}

func (ds durationSlice) Swap(a, b int) {
	ds[a], ds[b] = ds[b], ds[a]
}

func main() {
	app := cmd.NewAppShell("expiration-mailer")

	app.App.Flags = append(app.App.Flags, cli.IntFlag{
		Name:   "cert_limit",
		Value:  100,
		EnvVar: "CERT_LIMIT",
		Usage:  "Count of certificates to process per expiration period",
	})

	app.Config = func(c *cli.Context, config cmd.Config) cmd.Config {
		if c.GlobalInt("cert_limit") > 0 {
			config.Mailer.CertLimit = c.GlobalInt("cert_limit")
		}
		return config
	}

	app.Action = func(c cmd.Config) {
		// Set up logging
		stats, err := statsd.NewClient(c.Statsd.Server, c.Statsd.Prefix)
		cmd.FailOnError(err, "Couldn't connect to statsd")

		auditlogger, err := blog.Dial(c.Syslog.Network, c.Syslog.Server, c.Syslog.Tag, stats)
		cmd.FailOnError(err, "Could not connect to Syslog")

		// AUDIT[ Error Conditions ] 9cc4d537-8534-4970-8665-4b382abe82f3
		defer auditlogger.AuditPanic()

		blog.SetAuditLogger(auditlogger)

		auditlogger.Info(app.VersionString())

		go cmd.DebugServer(c.Mailer.DebugAddr)

		// Configure DB
		dbMap, err := sa.NewDbMap(c.Mailer.DBDriver, c.Mailer.DBConnect)
		cmd.FailOnError(err, "Could not connect to database")

		// Load email template
		emailTmpl, err := ioutil.ReadFile(c.Mailer.EmailTemplate)
		cmd.FailOnError(err, fmt.Sprintf("Could not read email template file [%s]", c.Mailer.EmailTemplate))
		tmpl, err := template.New("expiry-email").Parse(string(emailTmpl))
		cmd.FailOnError(err, "Could not parse email template")

		mailClient := mail.New(c.Mailer.Server, c.Mailer.Port, c.Mailer.Username, c.Mailer.Password)

		var nags durationSlice
		for _, nagDuration := range c.Mailer.NagTimes {
			dur, err := time.ParseDuration(nagDuration)
			if err != nil {
				auditlogger.Err(fmt.Sprintf("Failed to parse nag duration string [%s]: %s", nagDuration, err))
				return
			}
			nags = append(nags, dur)
		}
		// Make sure durations are sorted in increasing order
		sort.Sort(nags)

		m := mailer{
			stats:         stats,
			log:           auditlogger,
			dbMap:         dbMap,
			mailer:        &mailClient,
			emailTemplate: tmpl,
			nagTimes:      nags,
			limit:         c.Mailer.CertLimit,
		}

		auditlogger.Info("expiration-mailer: Starting")
		err = m.findExpiringCertificates()
		cmd.FailOnError(err, "expiration-mailer has failed")
	}

	app.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"runtime"
	"strings"

	"github.com/ready-steady/linear"

	"../internal"
)

var (
	outputFile     = flag.String("o", "", "an output file (required)")
	parameterIndex = flag.String("s", "[]", "the parameters to sweep")
	numberOfNodes  = flag.Uint("n", 10, "the number of nodes per parameter")
	defaultNode    = flag.Float64("d", 0.5, "the default value of parameters")
)

func main() {
	internal.Run(command)
}

func command(config *internal.Config) error {
	output, err := internal.Create(*outputFile)
	if err != nil {
		return err
	}
	defer output.Close()

	problem, err := internal.NewProblem(config)
	if err != nil {
		return err
	}

	target, err := internal.NewTarget(problem)
	if err != nil {
		return err
	}

	points, err := generate(target, config.Interpolation.Rule)
	if err != nil {
		return err
	}

	ni, no := target.Dimensions()
	np := uint(len(points)) / ni

	if config.Verbose {
		fmt.Printf("Evaluating the model with reduction %.2f at %v points...\n",
			config.Probability.VarThreshold, np)
	}

	values := internal.Invoke(target, points, uint(runtime.GOMAXPROCS(0)))

	if config.Verbose {
		fmt.Println("Done.")
	}

	if err := output.Put("values", values, no, np); err != nil {
		return err
	}
	if err := output.Put("points", points, ni, np); err != nil {
		return err
	}

	return nil
}

func generate(target internal.Target, rule string) ([]float64, error) {
	ni, _ := target.Dimensions()
	nn := *numberOfNodes

	index, err := detect(target)
	if err != nil {
		return nil, err
	}

	steady := []float64{*defaultNode}

	sweep := make([]float64, nn)
	switch rule {
	case "closed":
		for i := uint(0); i < nn; i++ {
			sweep[i] = float64(i) / float64(nn-1)
		}
	case "open":
		for i := uint(0); i < nn; i++ {
			sweep[i] = float64(i+1) / float64(nn+1)
		}
	default:
		return nil, errors.New("the sweep rule is unknown")
	}

	parameters := make([][]float64, ni)
	for i := uint(0); i < ni; i++ {
		parameters[i] = steady
	}
	for _, i := range index {
		parameters[i] = sweep
	}

	return linear.Tensor(parameters...), nil
}

func detect(target internal.Target) ([]uint, error) {
	ni, _ := target.Dimensions()

	index := []uint{}

	decoder := json.NewDecoder(strings.NewReader(*parameterIndex))
	if err := decoder.Decode(&index); err != nil {
		return nil, err
	}

	if len(index) == 0 {
		index = make([]uint, ni)
		for i := uint(0); i < ni; i++ {
			index[i] = i
		}
	}

	for _, i := range index {
		if i >= ni {
			return nil, errors.New(fmt.Sprintf("the indices should be less that %v", ni))
		}
	}

	return index, nil
}
<commit_msg>Fix a typo<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"runtime"
	"strings"

	"github.com/ready-steady/linear"

	"../internal"
)

var (
	outputFile     = flag.String("o", "", "an output file (required)")
	parameterIndex = flag.String("s", "[]", "the parameters to sweep")
	numberOfNodes  = flag.Uint("n", 10, "the number of nodes per parameter")
	defaultNode    = flag.Float64("d", 0.5, "the default value of parameters")
)

func main() {
	internal.Run(command)
}

func command(config *internal.Config) error {
	output, err := internal.Create(*outputFile)
	if err != nil {
		return err
	}
	defer output.Close()

	problem, err := internal.NewProblem(config)
	if err != nil {
		return err
	}

	target, err := internal.NewTarget(problem)
	if err != nil {
		return err
	}

	points, err := generate(target, config.Interpolation.Rule)
	if err != nil {
		return err
	}

	ni, no := target.Dimensions()
	np := uint(len(points)) / ni

	if config.Verbose {
		fmt.Printf("Evaluating the model with reduction %.2f at %v points...\n",
			config.Probability.VarThreshold, np)
	}

	values := internal.Invoke(target, points, uint(runtime.GOMAXPROCS(0)))

	if config.Verbose {
		fmt.Println("Done.")
	}

	if err := output.Put("values", values, no, np); err != nil {
		return err
	}
	if err := output.Put("points", points, ni, np); err != nil {
		return err
	}

	return nil
}

func generate(target internal.Target, rule string) ([]float64, error) {
	ni, _ := target.Dimensions()
	nn := *numberOfNodes

	index, err := detect(target)
	if err != nil {
		return nil, err
	}

	steady := []float64{*defaultNode}

	sweep := make([]float64, nn)
	switch rule {
	case "closed":
		for i := uint(0); i < nn; i++ {
			sweep[i] = float64(i) / float64(nn-1)
		}
	case "open":
		for i := uint(0); i < nn; i++ {
			sweep[i] = float64(i+1) / float64(nn+1)
		}
	default:
		return nil, errors.New("the sweep rule is unknown")
	}

	parameters := make([][]float64, ni)
	for i := uint(0); i < ni; i++ {
		parameters[i] = steady
	}
	for _, i := range index {
		parameters[i] = sweep
	}

	return linear.Tensor(parameters...), nil
}

func detect(target internal.Target) ([]uint, error) {
	ni, _ := target.Dimensions()

	index := []uint{}

	decoder := json.NewDecoder(strings.NewReader(*parameterIndex))
	if err := decoder.Decode(&index); err != nil {
		return nil, err
	}

	if len(index) == 0 {
		index = make([]uint, ni)
		for i := uint(0); i < ni; i++ {
			index[i] = i
		}
	}

	for _, i := range index {
		if i >= ni {
			return nil, errors.New(fmt.Sprintf("the indices should be less than %v", ni))
		}
	}

	return index, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/hdm/golang-mtbl"
	"github.com/hdm/inetdata-parsers"
	"os"
	"runtime"
	"strings"
)

func usage() {
	fmt.Println("Usage: " + os.Args[0] + " [options]")
	fmt.Println("")
	fmt.Println("Creates a MTBL database from a CSV input.")
	fmt.Println("")
	fmt.Println("Options:")
	flag.PrintDefaults()
}

func mergeFunc(key []byte, val0 []byte, val1 []byte) (mergedVal []byte) {
	return []byte(string(val0) + " " + string(val1))
}

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())
	os.Setenv("LC_ALL", "C")

	flag.Usage = func() { usage() }

	index_key := flag.Int("k", 1, "The field index to use as the key")
	index_val := flag.Int("v", 2, "The field index to use as the value")
	reverse_key := flag.Bool("r", false, "Store the key in reverse order")
	max_fields := flag.Int("M", -1, "The maximum number of fields to parse with the delimiter")
	compression := flag.String("c", "snappy", "The compression type to use (none, snappy, zlib, lz4, lz4hc)")
	delimiter := flag.String("d", ",", "The delimiter to use as a field separator")
	sort_skip := flag.Bool("S", false, "Skip the sorting phase and assume keys are in pre-sorted order")
	sort_tmp := flag.String("t", "", "The temporary directory to use for the sorting phase")
	sort_mem := flag.Uint64("m", 1, "The maximum amount of memory to use, in gigabytes, for the sorting phase")
	version := flag.Bool("version", false, "Show the version and build timestamp")

	flag.Parse()

	if *version {
		inetdata.PrintVersion("inetdata-csv2mtbl")
		os.Exit(0)
	}

	if len(flag.Args()) != 1 {
		usage()
		os.Exit(1)
	}

	fname := flag.Args()[0]

	sort_opt := mtbl.SorterOptions{Merge: mergeFunc, MaxMemory: 1000000000}
	sort_opt.MaxMemory *= *sort_mem
	if len(*sort_tmp) > 0 {
		sort_opt.TempDir = *sort_tmp
	}

	compression_alg, ok := inetdata.MTBLCompressionTypes[*compression]
	if !ok {
		fmt.Fprintf(os.Stderr, "Invalid compression algorithm: %s\n", *compression)
		os.Exit(1)
	}

	s := mtbl.SorterInit(&sort_opt)
	defer s.Destroy()

	w, we := mtbl.WriterInit(fname, &mtbl.WriterOptions{Compression: compression_alg})
	defer w.Destroy()

	if we != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", we)
		os.Exit(1)
	}

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		raw := strings.TrimSpace(scanner.Text())
		if len(raw) == 0 {
			continue
		}

		bits := strings.SplitN(raw, *delimiter, *max_fields)

		if len(bits) < *index_key {
			fmt.Fprintf(os.Stderr, "No key: %s\n", raw)
			continue
		}

		if len(bits) < *index_val {
			fmt.Fprintf(os.Stderr, "No value: %s\n", raw)
			continue
		}

		kstr := bits[*index_key-1]
		if len(kstr) == 0 {
			continue
		}

		vstr := bits[*index_val-1]
		if len(vstr) == 0 {
			continue
		}

		if *reverse_key {
			kstr = inetdata.ReverseKey(kstr)
		}

		if *sort_skip {
			if e := w.Add([]byte(kstr), []byte(vstr)); e != nil {
				fmt.Printf("Failed to add %v -> %v: %v\n", kstr, vstr, e)
			}
		} else {
			if e := s.Add([]byte(kstr), []byte(vstr)); e != nil {
				fmt.Printf("Failed to add %v -> %v: %v\n", kstr, vstr, e)
			}
		}
	}

	if !*sort_skip {
		if e := s.Write(w); e != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", e)
			os.Exit(1)
		}
	}
}
<commit_msg>csv2mtbl: Handle long line<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"runtime"
	"strings"

	mtbl "github.com/hdm/golang-mtbl"
	"github.com/hdm/inetdata-parsers"
)

func usage() {
	fmt.Println("Usage: " + os.Args[0] + " [options]")
	fmt.Println("")
	fmt.Println("Creates a MTBL database from a CSV input.")
	fmt.Println("")
	fmt.Println("Options:")
	flag.PrintDefaults()
}

func mergeFunc(key []byte, val0 []byte, val1 []byte) (mergedVal []byte) {
	return []byte(string(val0) + " " + string(val1))
}

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())
	os.Setenv("LC_ALL", "C")

	flag.Usage = func() { usage() }

	index_key := flag.Int("k", 1, "The field index to use as the key")
	index_val := flag.Int("v", 2, "The field index to use as the value")
	reverse_key := flag.Bool("r", false, "Store the key in reverse order")
	max_fields := flag.Int("M", -1, "The maximum number of fields to parse with the delimiter")
	compression := flag.String("c", "snappy", "The compression type to use (none, snappy, zlib, lz4, lz4hc)")
	delimiter := flag.String("d", ",", "The delimiter to use as a field separator")
	sort_skip := flag.Bool("S", false, "Skip the sorting phase and assume keys are in pre-sorted order")
	sort_tmp := flag.String("t", "", "The temporary directory to use for the sorting phase")
	sort_mem := flag.Uint64("m", 1, "The maximum amount of memory to use, in gigabytes, for the sorting phase")
	version := flag.Bool("version", false, "Show the version and build timestamp")

	flag.Parse()

	if *version {
		inetdata.PrintVersion("inetdata-csv2mtbl")
		os.Exit(0)
	}

	if len(flag.Args()) != 1 {
		usage()
		os.Exit(1)
	}

	fname := flag.Args()[0]

	sort_opt := mtbl.SorterOptions{Merge: mergeFunc, MaxMemory: 1000000000}
	sort_opt.MaxMemory *= *sort_mem
	if len(*sort_tmp) > 0 {
		sort_opt.TempDir = *sort_tmp
	}

	compression_alg, ok := inetdata.MTBLCompressionTypes[*compression]
	if !ok {
		fmt.Fprintf(os.Stderr, "Invalid compression algorithm: %s\n", *compression)
		os.Exit(1)
	}

	s := mtbl.SorterInit(&sort_opt)
	defer s.Destroy()

	w, we := mtbl.WriterInit(fname, &mtbl.WriterOptions{Compression: compression_alg})
	defer w.Destroy()

	if we != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", we)
		os.Exit(1)
	}

	scanner := bufio.NewScanner(os.Stdin)
	// Tune Scanner's value for MaxScanTokenSize which defaults to 65,536
	// Lines longer than MaxScanTokenSize will cause the Scanner to fail
	// Set the intial buffsize to twice default (4096) since we're here
	buf := make([]byte, 0, 8*1024)
	scanner.Buffer(buf, 96*1024)

	var current_line uint = 1
	for scanner.Scan() {
		current_line++
		raw := strings.TrimSpace(scanner.Text())
		if len(raw) == 0 {
			continue
		}

		bits := strings.SplitN(raw, *delimiter, *max_fields)

		if len(bits) < *index_key {
			fmt.Fprintf(os.Stderr, "No key: %s\n", raw)
			continue
		}

		if len(bits) < *index_val {
			fmt.Fprintf(os.Stderr, "No value: %s\n", raw)
			continue
		}

		kstr := bits[*index_key-1]
		if len(kstr) == 0 {
			continue
		}

		vstr := bits[*index_val-1]
		if len(vstr) == 0 {
			continue
		}

		if *reverse_key {
			kstr = inetdata.ReverseKey(kstr)
		}

		if *sort_skip {
			if e := w.Add([]byte(kstr), []byte(vstr)); e != nil {
				fmt.Printf("Failed to add %v -> %v: %v\n", kstr, vstr, e)
			}
		} else {
			if e := s.Add([]byte(kstr), []byte(vstr)); e != nil {
				fmt.Printf("Failed to add %v -> %v: %v\n", kstr, vstr, e)
			}
		}
	}

	if err := scanner.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "Error while processing line %d : %s\n", current_line, err)
		os.Exit(1)
	}

	if !*sort_skip {
		if e := s.Write(w); e != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", e)
			os.Exit(1)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/hdm/inetdata-parsers/utils"
	"os"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const ZONE_MODE_UNKNOWN = 0
const ZONE_MODE_COM = 1 // Also NET, ORG, INFO, MOBI
const ZONE_MODE_BIZ = 2 // Also XXX
const ZONE_MODE_SK = 3
const ZONE_MODE_US = 4
const ZONE_MODE_CZDS = 5

var zone_mode = 0
var zone_name = ""
var zone_matched = false

var match_ipv6 = regexp.MustCompile(`^((([0-9A-Fa-f]{1,4}:){7}([0-9A-Fa-f]{1,4}|:))|(([0-9A-Fa-f]{1,4}:){6}(:[0-9A-Fa-f]{1,4}|((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3})|:))|(([0-9A-Fa-f]{1,4}:){5}(((:[0-9A-Fa-f]{1,4}){1,2})|:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3})|:))|(([0-9A-Fa-f]{1,4}:){4}(((:[0-9A-Fa-f]{1,4}){1,3})|((:[0-9A-Fa-f]{1,4})?:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(([0-9A-Fa-f]{1,4}:){3}(((:[0-9A-Fa-f]{1,4}){1,4})|((:[0-9A-Fa-f]{1,4}){0,2}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(([0-9A-Fa-f]{1,4}:){2}(((:[0-9A-Fa-f]{1,4}){1,5})|((:[0-9A-Fa-f]{1,4}){0,3}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(([0-9A-Fa-f]{1,4}:){1}(((:[0-9A-Fa-f]{1,4}){1,6})|((:[0-9A-Fa-f]{1,4}){0,4}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(:(((:[0-9A-Fa-f]{1,4}){1,7})|((:[0-9A-Fa-f]{1,4}){0,5}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:)))(%.+)?$`)

var match_ipv4 = regexp.MustCompile(`^(?:(?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})[.](?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})[.](?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})[.](?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2}))$`)

var split_ws = regexp.MustCompile(`\s+`)

var output_count int64 = 0
var input_count int64 = 0
var stdout_lock sync.Mutex
var wg sync.WaitGroup

type OutputKey struct {
	Key  string
	Vals []string
}

type OutputChannels struct {
	IP4          chan string
	IP6          chan string
	InverseNames chan string
	Names        chan string
}

func usage() {
	fmt.Println("Usage: " + os.Args[0] + " [options]")
	fmt.Println("")
	fmt.Println("Reads a zone file from stdin, generates CSV files keyed off domain names, including ")
	fmt.Println("forward, inverse, and glue addresses for IPv4 and IPv6.")
	fmt.Println("")
	fmt.Println("Options:")
	flag.PrintDefaults()
}

func showProgress(quit chan int) {
	start := time.Now()
	for {
		select {
		case <-quit:
			fmt.Fprintf(os.Stderr, "[*] Complete\n")
			return
		case <-time.After(time.Second * 1):
			icount := atomic.LoadInt64(&input_count)
			ocount := atomic.LoadInt64(&output_count)

			if icount == 0 && ocount == 0 {
				// Reset start, so that we show stats only from our first input
				start = time.Now()
				continue
			}
			elapsed := time.Since(start)
			if elapsed.Seconds() > 1.0 {
				fmt.Fprintf(os.Stderr, "[*] [inetdata-zone2csv] Read %d and wrote %d records in %d seconds (%d/s in, %d/s out)\n",
					icount,
					ocount,
					int(elapsed.Seconds()),
					int(float64(icount)/elapsed.Seconds()),
					int(float64(ocount)/elapsed.Seconds()))
			}
		}
	}
}

func outputWriter(fd *os.File, c chan string) {
	for r := range c {
		fd.Write([]byte(r))
		atomic.AddInt64(&output_count, 1)
	}
	wg.Done()
}

func writeRecord(c_names chan string, name string, rtype string, value string) {
	switch rtype {
	case "ns":
		c_names <- fmt.Sprintf("%s,%s,%s\n", name, rtype, value)

	case "a":
		if match_ipv4.Match([]byte(value)) {
			c_names <- fmt.Sprintf("%s,%s,%s\n", name, rtype, value)
		}

	case "aaaa":
		if match_ipv6.Match([]byte(value)) {
			c_names <- fmt.Sprintf("%s,%s,%s\n", name, rtype, value)
		}
	}
}

func normalizeName(name string) string {
	// Leave empty names alone
	if len(name) == 0 {
		return name
	}

	// Leave IP addresses alone
	if match_ipv4.Match([]byte(name)) || match_ipv6.Match([]byte(name)) {
		return name
	}

	if name[len(name)-1:] == "." {
		// Remove the trailing dot
		name = name[:len(name)-1]
	} else {
		// Add the domain to complete the name
		name = fmt.Sprintf("%s.%s", name, zone_name)
	}
	return name
}

func parseZoneCOM(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 3 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[1], normalizeName(bits[2])
	writeRecord(c_names, name, rtype, value)
}

func parseZoneBIZ(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 5 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[3], normalizeName(bits[4])
	writeRecord(c_names, name, rtype, value)
}

func parseZoneUS(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 4 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[2], normalizeName(bits[3])
	writeRecord(c_names, name, rtype, value)
}

func parseZoneSK(raw string, c_names chan string) {
	bits := strings.SplitN(strings.ToLower(raw), ";", -1)
	if len(bits) < 5 {
		return
	}

	name := normalizeName(bits[0])
	if len(name) == 0 {
		return
	}

	ns1, ns2, ns3, ns4 := normalizeName(bits[5]), normalizeName(bits[6]), normalizeName(bits[7]), normalizeName(bits[8])

	if len(ns1) > 0 {
		writeRecord(c_names, name, "ns", ns1)
	}

	if len(ns2) > 0 {
		writeRecord(c_names, name, "ns", ns2)
	}

	if len(ns3) > 0 {
		writeRecord(c_names, name, "ns", ns3)
	}

	if len(ns4) > 0 {
		writeRecord(c_names, name, "ns", ns4)
	}
}

func parseZoneCZDS(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 5 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[3], normalizeName(bits[4])
	writeRecord(c_names, name, rtype, value)
}

func inputParser(c chan string, c_names chan string) {

	lines_read := 0
	for r := range c {

		raw := strings.TrimSpace(r)

		if len(raw) == 0 {
			continue
		}

		atomic.AddInt64(&input_count, 1)

		if zone_mode != ZONE_MODE_UNKNOWN && zone_matched == false {
			zone_matched = true

			// Spawn more parsers
			for i := 0; i < runtime.NumCPU()-1; i++ {
				fmt.Fprintf(os.Stderr, "Launching goroutine!\n")
				go inputParser(c, c_names)
				wg.Add(1)
			}
		}

		switch zone_mode {
		case ZONE_MODE_UNKNOWN:

			// Versign Zone Format
			if strings.Contains(raw, "$ORIGIN COM.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "com"
				continue
			}

			if strings.Contains(raw, "$ORIGIN INFO.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "info"
				continue
			}

			if strings.Contains(raw, "$ORIGIN MOBI.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "mobi"
				continue
			}

			if strings.Contains(raw, "$ORIGIN NET.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "net"
				continue
			}

			if strings.Contains(raw, "$ORIGIN org.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "org"
				continue
			}

			// US zone format
			if strings.Contains(raw, "US. IN SOA A.CCTLD.US HOSTMASTER.NEUSTAR.US") {
				zone_mode = ZONE_MODE_US
				zone_name = "us"
				continue
			}

			// BIZ/XXX zone format
			if strings.Contains(raw, "BIZ.			900	IN	SOA	A.GTLD.BIZ.") {
				zone_mode = ZONE_MODE_BIZ
				zone_name = "biz"
				continue
			}

			if strings.Contains(raw, "xxx.	86400	in	soa	a0.xxx.afilias-nst.info.") {
				zone_mode = ZONE_MODE_BIZ
				zone_name = "xxx"
				continue
			}

			// SK static zone
			if strings.Contains(raw, "domena;ID reg;ID drzitela;NEW") {
				zone_mode = ZONE_MODE_SK
				zone_name = "sk"
				continue
			}

			// CZDS
			if matched, _ := regexp.Match(`^[a-zA-Z0-9\-]+\.\s+\d+\s+in\s+soa\s+`, []byte(raw)); matched {
				zone_mode = ZONE_MODE_CZDS
				zone_name = ""
				continue
			}

			lines_read++

			if lines_read > 100 {
				fmt.Fprintf(os.Stderr, "[-] Could not determine zone format, giving up: %s\n", raw)
				os.Exit(1)
			}

		case ZONE_MODE_COM:
			parseZoneCOM(raw, c_names)

		case ZONE_MODE_BIZ:
			parseZoneBIZ(raw, c_names)

		case ZONE_MODE_SK:
			parseZoneSK(raw, c_names)

		case ZONE_MODE_US:
			parseZoneUS(raw, c_names)

		case ZONE_MODE_CZDS:
			parseZoneCZDS(raw, c_names)

		default:
			panic("Unknown zone mode")
		}
	}

	wg.Done()
}

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())
	os.Setenv("LC_ALL", "C")

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

	// Progress tracker
	quit := make(chan int)
	go showProgress(quit)

	// Write output
	c_names := make(chan string, 1000)
	go outputWriter(os.Stdout, c_names)

	// Read input
	c_inp := make(chan string, 1000)
	go inputParser(c_inp, c_names)
	wg.Add(1)

	// Reader closers c_inp on completion
	e := utils.ReadLines(os.Stdin, c_inp)
	if e != nil {
		fmt.Fprintf(os.Stderr, "Error reading input: %s\n", e)
	}

	// Wait for the input parser to finish
	wg.Wait()

	// Close the output channel
	close(c_names)

	// Wait for the channel writers to finish
	wg.Add(1)
	wg.Wait()

	// Stop the main process monitoring
	quit <- 0
}
<commit_msg>Remove debug statement<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/hdm/inetdata-parsers/utils"
	"os"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const ZONE_MODE_UNKNOWN = 0
const ZONE_MODE_COM = 1 // Also NET, ORG, INFO, MOBI
const ZONE_MODE_BIZ = 2 // Also XXX
const ZONE_MODE_SK = 3
const ZONE_MODE_US = 4
const ZONE_MODE_CZDS = 5

var zone_mode = 0
var zone_name = ""
var zone_matched = false

var match_ipv6 = regexp.MustCompile(`^((([0-9A-Fa-f]{1,4}:){7}([0-9A-Fa-f]{1,4}|:))|(([0-9A-Fa-f]{1,4}:){6}(:[0-9A-Fa-f]{1,4}|((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3})|:))|(([0-9A-Fa-f]{1,4}:){5}(((:[0-9A-Fa-f]{1,4}){1,2})|:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3})|:))|(([0-9A-Fa-f]{1,4}:){4}(((:[0-9A-Fa-f]{1,4}){1,3})|((:[0-9A-Fa-f]{1,4})?:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(([0-9A-Fa-f]{1,4}:){3}(((:[0-9A-Fa-f]{1,4}){1,4})|((:[0-9A-Fa-f]{1,4}){0,2}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(([0-9A-Fa-f]{1,4}:){2}(((:[0-9A-Fa-f]{1,4}){1,5})|((:[0-9A-Fa-f]{1,4}){0,3}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(([0-9A-Fa-f]{1,4}:){1}(((:[0-9A-Fa-f]{1,4}){1,6})|((:[0-9A-Fa-f]{1,4}){0,4}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:))|(:(((:[0-9A-Fa-f]{1,4}){1,7})|((:[0-9A-Fa-f]{1,4}){0,5}:((25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(\.(25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}))|:)))(%.+)?$`)

var match_ipv4 = regexp.MustCompile(`^(?:(?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})[.](?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})[.](?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})[.](?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2}))$`)

var split_ws = regexp.MustCompile(`\s+`)

var output_count int64 = 0
var input_count int64 = 0
var stdout_lock sync.Mutex
var wg sync.WaitGroup

type OutputKey struct {
	Key  string
	Vals []string
}

type OutputChannels struct {
	IP4          chan string
	IP6          chan string
	InverseNames chan string
	Names        chan string
}

func usage() {
	fmt.Println("Usage: " + os.Args[0] + " [options]")
	fmt.Println("")
	fmt.Println("Reads a zone file from stdin, generates CSV files keyed off domain names, including ")
	fmt.Println("forward, inverse, and glue addresses for IPv4 and IPv6.")
	fmt.Println("")
	fmt.Println("Options:")
	flag.PrintDefaults()
}

func showProgress(quit chan int) {
	start := time.Now()
	for {
		select {
		case <-quit:
			fmt.Fprintf(os.Stderr, "[*] Complete\n")
			return
		case <-time.After(time.Second * 1):
			icount := atomic.LoadInt64(&input_count)
			ocount := atomic.LoadInt64(&output_count)

			if icount == 0 && ocount == 0 {
				// Reset start, so that we show stats only from our first input
				start = time.Now()
				continue
			}
			elapsed := time.Since(start)
			if elapsed.Seconds() > 1.0 {
				fmt.Fprintf(os.Stderr, "[*] [inetdata-zone2csv] Read %d and wrote %d records in %d seconds (%d/s in, %d/s out)\n",
					icount,
					ocount,
					int(elapsed.Seconds()),
					int(float64(icount)/elapsed.Seconds()),
					int(float64(ocount)/elapsed.Seconds()))
			}
		}
	}
}

func outputWriter(fd *os.File, c chan string) {
	for r := range c {
		fd.Write([]byte(r))
		atomic.AddInt64(&output_count, 1)
	}
	wg.Done()
}

func writeRecord(c_names chan string, name string, rtype string, value string) {
	switch rtype {
	case "ns":
		c_names <- fmt.Sprintf("%s,%s,%s\n", name, rtype, value)

	case "a":
		if match_ipv4.Match([]byte(value)) {
			c_names <- fmt.Sprintf("%s,%s,%s\n", name, rtype, value)
		}

	case "aaaa":
		if match_ipv6.Match([]byte(value)) {
			c_names <- fmt.Sprintf("%s,%s,%s\n", name, rtype, value)
		}
	}
}

func normalizeName(name string) string {
	// Leave empty names alone
	if len(name) == 0 {
		return name
	}

	// Leave IP addresses alone
	if match_ipv4.Match([]byte(name)) || match_ipv6.Match([]byte(name)) {
		return name
	}

	if name[len(name)-1:] == "." {
		// Remove the trailing dot
		name = name[:len(name)-1]
	} else {
		// Add the domain to complete the name
		name = fmt.Sprintf("%s.%s", name, zone_name)
	}
	return name
}

func parseZoneCOM(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 3 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[1], normalizeName(bits[2])
	writeRecord(c_names, name, rtype, value)
}

func parseZoneBIZ(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 5 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[3], normalizeName(bits[4])
	writeRecord(c_names, name, rtype, value)
}

func parseZoneUS(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 4 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[2], normalizeName(bits[3])
	writeRecord(c_names, name, rtype, value)
}

func parseZoneSK(raw string, c_names chan string) {
	bits := strings.SplitN(strings.ToLower(raw), ";", -1)
	if len(bits) < 5 {
		return
	}

	name := normalizeName(bits[0])
	if len(name) == 0 {
		return
	}

	ns1, ns2, ns3, ns4 := normalizeName(bits[5]), normalizeName(bits[6]), normalizeName(bits[7]), normalizeName(bits[8])

	if len(ns1) > 0 {
		writeRecord(c_names, name, "ns", ns1)
	}

	if len(ns2) > 0 {
		writeRecord(c_names, name, "ns", ns2)
	}

	if len(ns3) > 0 {
		writeRecord(c_names, name, "ns", ns3)
	}

	if len(ns4) > 0 {
		writeRecord(c_names, name, "ns", ns4)
	}
}

func parseZoneCZDS(raw string, c_names chan string) {
	bits := split_ws.Split(strings.ToLower(raw), -1)
	if len(bits) != 5 {
		return
	}

	name, rtype, value := normalizeName(bits[0]), bits[3], normalizeName(bits[4])
	writeRecord(c_names, name, rtype, value)
}

func inputParser(c chan string, c_names chan string) {

	lines_read := 0
	for r := range c {

		raw := strings.TrimSpace(r)

		if len(raw) == 0 {
			continue
		}

		atomic.AddInt64(&input_count, 1)

		if zone_mode != ZONE_MODE_UNKNOWN && zone_matched == false {
			zone_matched = true

			// Spawn more parsers
			for i := 0; i < runtime.NumCPU()-1; i++ {
				go inputParser(c, c_names)
				wg.Add(1)
			}
		}

		switch zone_mode {
		case ZONE_MODE_UNKNOWN:

			// Versign Zone Format
			if strings.Contains(raw, "$ORIGIN COM.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "com"
				continue
			}

			if strings.Contains(raw, "$ORIGIN INFO.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "info"
				continue
			}

			if strings.Contains(raw, "$ORIGIN MOBI.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "mobi"
				continue
			}

			if strings.Contains(raw, "$ORIGIN NET.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "net"
				continue
			}

			if strings.Contains(raw, "$ORIGIN org.") {
				zone_mode = ZONE_MODE_COM
				zone_name = "org"
				continue
			}

			// US zone format
			if strings.Contains(raw, "US. IN SOA A.CCTLD.US HOSTMASTER.NEUSTAR.US") {
				zone_mode = ZONE_MODE_US
				zone_name = "us"
				continue
			}

			// BIZ/XXX zone format
			if strings.Contains(raw, "BIZ.			900	IN	SOA	A.GTLD.BIZ.") {
				zone_mode = ZONE_MODE_BIZ
				zone_name = "biz"
				continue
			}

			if strings.Contains(raw, "xxx.	86400	in	soa	a0.xxx.afilias-nst.info.") {
				zone_mode = ZONE_MODE_BIZ
				zone_name = "xxx"
				continue
			}

			// SK static zone
			if strings.Contains(raw, "domena;ID reg;ID drzitela;NEW") {
				zone_mode = ZONE_MODE_SK
				zone_name = "sk"
				continue
			}

			// CZDS
			if matched, _ := regexp.Match(`^[a-zA-Z0-9\-]+\.\s+\d+\s+in\s+soa\s+`, []byte(raw)); matched {
				zone_mode = ZONE_MODE_CZDS
				zone_name = ""
				continue
			}

			lines_read++

			if lines_read > 100 {
				fmt.Fprintf(os.Stderr, "[-] Could not determine zone format, giving up: %s\n", raw)
				os.Exit(1)
			}

		case ZONE_MODE_COM:
			parseZoneCOM(raw, c_names)

		case ZONE_MODE_BIZ:
			parseZoneBIZ(raw, c_names)

		case ZONE_MODE_SK:
			parseZoneSK(raw, c_names)

		case ZONE_MODE_US:
			parseZoneUS(raw, c_names)

		case ZONE_MODE_CZDS:
			parseZoneCZDS(raw, c_names)

		default:
			panic("Unknown zone mode")
		}
	}

	wg.Done()
}

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())
	os.Setenv("LC_ALL", "C")

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

	// Progress tracker
	quit := make(chan int)
	go showProgress(quit)

	// Write output
	c_names := make(chan string, 1000)
	go outputWriter(os.Stdout, c_names)

	// Read input
	c_inp := make(chan string, 1000)
	go inputParser(c_inp, c_names)
	wg.Add(1)

	// Reader closers c_inp on completion
	e := utils.ReadLines(os.Stdin, c_inp)
	if e != nil {
		fmt.Fprintf(os.Stderr, "Error reading input: %s\n", e)
	}

	// Wait for the input parser to finish
	wg.Wait()

	// Close the output channel
	close(c_names)

	// Wait for the channel writers to finish
	wg.Add(1)
	wg.Wait()

	// Stop the main process monitoring
	quit <- 0
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"text/template"

	"github.com/pkg/errors"

	"github.com/micromdm/go4/env"
)

func runServer(args []string) error {
	flagset := flag.NewFlagSet("go4up", flag.ExitOnError)
	var (
		flAppName   = flagset.String("name", "example", "name of app")
		flOutputDir = flagset.String(
			"output",
			filepath.Join(gopath(), "src", "github.com", "micromdm", *flAppName),
			"path to output",
		)
	)

	flagset.Usage = usageFor(flagset, "go4up server [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	dir := filepath.Join(*flOutputDir, "cmd", *flAppName)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return errors.Wrapf(err, "creating output directory %s", dir)
	}

	tmplArgs := struct {
		Name string
	}{
		Name: *flAppName,
	}
	makefileBuf := new(bytes.Buffer)
	var makefileTmpl = template.Must(template.New("test").Parse(serverMakefileTemplate))
	if err := makefileTmpl.Execute(makefileBuf, tmplArgs); err != nil {
		return errors.Wrap(err, "execute makefile template")
	}

	serverBuf := new(bytes.Buffer)
	var serveTmpl = template.Must(template.New("serve").Parse(serverServeTemplate))
	if err := serveTmpl.Execute(serverBuf, tmplArgs); err != nil {
		return errors.Wrap(err, "execute serve.go template")
	}

	dockerfileBuf := new(bytes.Buffer)
	var dockerfileTmpl = template.Must(template.New("dockerfile").Parse(dockerfileTemplate))
	if err := dockerfileTmpl.Execute(dockerfileBuf, tmplArgs); err != nil {
		return errors.Wrap(err, "execute Dockerfile template")
	}

	gitignorePath := filepath.Join(*flOutputDir, ".gitignore")
	if err := ioutil.WriteFile(gitignorePath, []byte(gitignoreTemplate), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", gitignorePath)
	}

	makefilePath := filepath.Join(*flOutputDir, "Makefile")
	if err := ioutil.WriteFile(makefilePath, makefileBuf.Bytes(), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", makefilePath)
	}

	dockerfilePath := filepath.Join(*flOutputDir, "Dockerfile")
	if err := ioutil.WriteFile(dockerfilePath, dockerfileBuf.Bytes(), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", dockerfilePath)
	}

	mainPath := filepath.Join(dir, fmt.Sprintf("%s.go", *flAppName))
	if err := ioutil.WriteFile(mainPath, []byte(serverMainTemplate), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", mainPath)
	}

	servePath := filepath.Join(dir, "serve.go")
	if err := ioutil.WriteFile(servePath, serverBuf.Bytes(), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", servePath)
	}

	return nil

}

func gopath() string {
	home := env.String("HOME", "~/")
	return env.String("GOPATH", filepath.Join(home, "go"))
}

const serverMainTemplate = `package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"text/tabwriter"

	"github.com/micromdm/go4/version"
)

func runVersion(args []string) error {
	version.PrintFull()
	return nil
}

func usageFor(fs *flag.FlagSet, short string) func() {
	return func() {
		fmt.Fprintf(os.Stderr, "USAGE\n")
		fmt.Fprintf(os.Stderr, "  %s\n", short)
		fmt.Fprintf(os.Stderr, "\n")
		fmt.Fprintf(os.Stderr, "FLAGS\n")
		w := tabwriter.NewWriter(os.Stderr, 0, 2, 2, ' ', 0)
		fs.VisitAll(func(f *flag.Flag) {
			fmt.Fprintf(w, "\t-%s %s\t%s\n", f.Name, f.DefValue, f.Usage)
		})
		w.Flush()
		fmt.Fprintf(os.Stderr, "\n")
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "USAGE\n")
	fmt.Fprintf(os.Stderr, "  %s <mode> --help\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\n")
	fmt.Fprintf(os.Stderr, "MODES\n")
	fmt.Fprintf(os.Stderr, "  serve        Run the server\n")
	fmt.Fprintf(os.Stderr, "  version      Print full version information\n")
	fmt.Fprintf(os.Stderr, "\n")
	fmt.Fprintf(os.Stderr, "VERSION\n")
	fmt.Fprintf(os.Stderr, "  %s\n", version.Version().Version)
	fmt.Fprintf(os.Stderr, "\n")
}

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(1)
	}

	var run func([]string) error
	switch strings.ToLower(os.Args[1]) {
	case "serve":
		run = runServe
	case "version":
		run = runVersion
	case "help", "-h", "--help":
		usage()
		return
	default:
		usage()
		os.Exit(1)
	}

	if err := run(os.Args[2:]); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}
`

const serverMakefileTemplate = `all: build

.PHONY: build

ifndef ($(GOPATH))
	GOPATH = $(HOME)/go
endif

PATH := $(GOPATH)/bin:$(PATH)
VERSION = $(shell git describe --tags --always --dirty)
BRANCH = $(shell git rev-parse --abbrev-ref HEAD)
REVISION = $(shell git rev-parse HEAD)
REVSHORT = $(shell git rev-parse --short HEAD)
USER = $(shell whoami)
GOVERSION = $(shell go version | awk '{print $$3}')
NOW	= $(shell date -u +"%Y-%m-%dT%H:%M:%SZ")
SHELL = /bin/bash

ifneq ($(OS), Windows_NT)
	CURRENT_PLATFORM = linux
	ifeq ($(shell uname), Darwin)
		SHELL := /bin/bash
		CURRENT_PLATFORM = darwin
	endif
else
	CURRENT_PLATFORM = windows
endif

BUILD_VERSION = "\
	-X github.com/micromdm/{{.Name}}/vendor/github.com/micromdm/go4/version.appName=${APP_NAME} \
	-X github.com/micromdm/{{.Name}}/vendor/github.com/micromdm/go4/version.version=${VERSION} \
	-X github.com/micromdm/{{.Name}}/vendor/github.com/micromdm/go4/version.branch=${BRANCH} \
	-X github.com/micromdm/{{.Name}}/vendor/github.com/micromdm/go4/version.buildUser=${USER} \
	-X github.com/micromdm/{{.Name}}/vendor/github.com/micromdm/go4/version.buildDate=${NOW} \
	-X github.com/micromdm/{{.Name}}/vendor/github.com/micromdm/go4/version.revision=${REVISION} \
	-X github.com/micromdm/{{.Name}}/vendor/github.com/micromdm/go4/version.goVersion=${GOVERSION}"

WORKSPACE = ${GOPATH}/src/github.com/micromdm/{{.Name}}
check-deps:
ifneq ($(shell test -e ${WORKSPACE}/Gopkg.lock && echo -n yes), yes)
	@echo "folder is clonded in the wrong place, copying to a Go Workspace"
	@echo "See: https://golang.org/doc/code.html#Workspaces"
	@git clone git@github.com:micromdm/{{.Name}} ${WORKSPACE}
	@echo "cd to ${WORKSPACE} and run make deps again."
	@exit 1
endif
ifneq ($(shell pwd), $(WORKSPACE))
	@echo "cd to ${WORKSPACE} and run make deps again."
	@exit 1
endif

deps: check-deps
	go get -u github.com/golang/dep/...
	dep ensure -vendor-only

test:
	go test -cover -race -v $(shell go list ./... | grep -v /vendor/)

build: {{.Name}}

clean:
	rm -rf build/
	rm -f *.zip

.pre-build:
	mkdir -p build/darwin
	mkdir -p build/linux

INSTALL_STEPS := \
	install-{{.Name}}

install-local: $(INSTALL_STEPS)

.pre-{{.Name}}:
	$(eval APP_NAME = {{.Name}})

{{.Name}}: .pre-build .pre-{{.Name}}
	go build -i -o build/$(CURRENT_PLATFORM)/{{.Name}} -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}

install-{{.Name}}: .pre-{{.Name}}
	go install -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}

xp-{{.Name}}: .pre-build .pre-{{.Name}}
	GOOS=darwin go build -i -o build/darwin/{{.Name}} -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}
	GOOS=linux CGO_ENABLED=0 go build -i -o build/linux/{{.Name}}  -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}

release-zip: xp-{{.Name}}
	zip -r {{.Name}}_${VERSION}.zip build/

# TODO remove after bootstrap is done.
git-init:
	git init
	git add -A
	git commit -m "first commit."

dep-init:
	dep init
	git add Gopkg.*
	git commit -m "Initialize Go dependencies."

init: git-init dep-init
`

const serverServeTemplate = `
package main

import (
	"flag"
	stdlog "log"
	"net/http"
	"os"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/groob/finalizer/logutil"
	"github.com/micromdm/go4/env"
	"github.com/micromdm/go4/httputil"
	"github.com/pkg/errors"
)

func runServe(args []string) error {
	flagset := flag.NewFlagSet("{{.Name}}", flag.ExitOnError)
	var (
		flConfigPath = flagset.String("config-dir", env.String("CONFIG_DIR", "/var/micromdm/{{.Name}}"), "Path to server config directory.")
		flLogFormat  = flagset.String("log-format", env.String("LOG_FORMAT", "logfmt"), "Enable structured logging. Supported formats: logfmt, json.")
		flLogLevel   = flagset.String("log-level", env.String("LOG_LEVEL", "info"), "Log level. Either info or debug.")
		flHTTPDebug  = flagset.Bool("http-debug", false, "Enable debug for http(dumps full request).")
		flHTTPAddr   = flagset.String("http-addr", env.String("HTTP_ADDR", ":https"), "HTTP(s) listen address of http server. Defaults to :443 or :8080 if tls=false")
		flTLS        = flagset.Bool("tls", env.Bool("USE_TLS", true), "Serve HTTPS.")
		flTLSCert    = flagset.String("tls-cert", env.String("TLS_CERT", ""), "Path to TLS certificate.")
		flTLSKey     = flagset.String("tls-key", env.String("TLS_KEY", ""), "Path to TLS private key.")
		flTLSDomain  = flagset.String("tls-domain", env.String("TLS_DOMAIN", ""), "Automatically fetch certs from Let's Encrypt for this domain. Format must be server.acme.co")
	)

	flagset.Usage = usageFor(flagset, "{{.Name}} serve [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	var (
		logger     log.Logger
		httpLogger log.Logger
	)
	{
		w := log.NewSyncWriter(os.Stderr)
		switch *flLogFormat {
		case "json":
			logger = log.NewJSONLogger(w)
		default:
			logger = log.NewLogfmtLogger(w)
		}
		stdlog.SetOutput(log.NewStdlibAdapter(logger))
		logger = log.With(logger, "ts", log.DefaultTimestampUTC)
		if *flLogLevel == "debug" {
			logger = level.NewFilter(logger, level.AllowDebug())
		} else {
			logger = level.NewFilter(logger, level.AllowInfo())
		}
		httpLogger = log.With(logger, "component", "http_logger")
		logger = log.With(logger, "caller", log.Caller(4))
	}

	mux := http.NewServeMux()
	var handler http.Handler
	if *flHTTPDebug {
		handler = httputil.HTTPDebugMiddleware(os.Stdout, true, httpLogger.Log)(mux)
	} else {
		handler = mux
	}
	handler = logutil.NewHTTPLogger(httpLogger).Middleware(handler)

	serveOpts := httputil.Simple(
		*flConfigPath,
		handler,
		*flHTTPAddr,
		*flTLSCert,
		*flTLSKey,
		*flTLS,
		logger,
		*flTLSDomain,
	)

	err := httputil.ListenAndServe(serveOpts...)
	return errors.Wrap(err, "calling ListenAndServe")
}
`

const gitignoreTemplate = `.DS_Store
build/
vendor/
*.zip
*.tar.gz
`

const dockerfileTemplate = `FROM alpine

RUN apk --update add \
    ca-certificates

COPY ./build/linux/{{.Name}} /usr/bin/{{.Name}}

CMD ["{{.Name}}"]
`
<commit_msg>derive org<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/pkg/errors"

	"github.com/micromdm/go4/env"
)

func runServer(args []string) error {
	flagset := flag.NewFlagSet("go4up", flag.ExitOnError)
	var (
		flOutputDir = flagset.String(
			"output",
			filepath.Join(gopath(), "src", "github.com", "micromdm", "example"),
			"path to output",
		)
	)

	flagset.Usage = usageFor(flagset, "go4up server [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}
	dirs := strings.Split(*flOutputDir, "/")
	appName := dirs[len(dirs)-1]
	orgName := dirs[len(dirs)-2]

	dir := filepath.Join(*flOutputDir, "cmd", appName)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return errors.Wrapf(err, "creating output directory %s", dir)
	}

	tmplArgs := struct {
		Name string
		Org  string
	}{
		Name: appName,
		Org:  orgName,
	}
	makefileBuf := new(bytes.Buffer)
	var makefileTmpl = template.Must(template.New("test").Parse(serverMakefileTemplate))
	if err := makefileTmpl.Execute(makefileBuf, tmplArgs); err != nil {
		return errors.Wrap(err, "execute makefile template")
	}

	serverBuf := new(bytes.Buffer)
	var serveTmpl = template.Must(template.New("serve").Parse(serverServeTemplate))
	if err := serveTmpl.Execute(serverBuf, tmplArgs); err != nil {
		return errors.Wrap(err, "execute serve.go template")
	}

	dockerfileBuf := new(bytes.Buffer)
	var dockerfileTmpl = template.Must(template.New("dockerfile").Parse(dockerfileTemplate))
	if err := dockerfileTmpl.Execute(dockerfileBuf, tmplArgs); err != nil {
		return errors.Wrap(err, "execute Dockerfile template")
	}

	gitignorePath := filepath.Join(*flOutputDir, ".gitignore")
	if err := ioutil.WriteFile(gitignorePath, []byte(gitignoreTemplate), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", gitignorePath)
	}

	makefilePath := filepath.Join(*flOutputDir, "Makefile")
	if err := ioutil.WriteFile(makefilePath, makefileBuf.Bytes(), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", makefilePath)
	}

	dockerfilePath := filepath.Join(*flOutputDir, "Dockerfile")
	if err := ioutil.WriteFile(dockerfilePath, dockerfileBuf.Bytes(), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", dockerfilePath)
	}

	mainPath := filepath.Join(dir, fmt.Sprintf("%s.go", appName))
	if err := ioutil.WriteFile(mainPath, []byte(serverMainTemplate), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", mainPath)
	}

	servePath := filepath.Join(dir, "serve.go")
	if err := ioutil.WriteFile(servePath, serverBuf.Bytes(), 0644); err != nil {
		return errors.Wrapf(err, "writing file %s", servePath)
	}

	return nil

}

func gopath() string {
	home := env.String("HOME", "~/")
	return env.String("GOPATH", filepath.Join(home, "go"))
}

const serverMainTemplate = `package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"text/tabwriter"

	"github.com/micromdm/go4/version"
)

func runVersion(args []string) error {
	version.PrintFull()
	return nil
}

func usageFor(fs *flag.FlagSet, short string) func() {
	return func() {
		fmt.Fprintf(os.Stderr, "USAGE\n")
		fmt.Fprintf(os.Stderr, "  %s\n", short)
		fmt.Fprintf(os.Stderr, "\n")
		fmt.Fprintf(os.Stderr, "FLAGS\n")
		w := tabwriter.NewWriter(os.Stderr, 0, 2, 2, ' ', 0)
		fs.VisitAll(func(f *flag.Flag) {
			fmt.Fprintf(w, "\t-%s %s\t%s\n", f.Name, f.DefValue, f.Usage)
		})
		w.Flush()
		fmt.Fprintf(os.Stderr, "\n")
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "USAGE\n")
	fmt.Fprintf(os.Stderr, "  %s <mode> --help\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\n")
	fmt.Fprintf(os.Stderr, "MODES\n")
	fmt.Fprintf(os.Stderr, "  serve        Run the server\n")
	fmt.Fprintf(os.Stderr, "  version      Print full version information\n")
	fmt.Fprintf(os.Stderr, "\n")
	fmt.Fprintf(os.Stderr, "VERSION\n")
	fmt.Fprintf(os.Stderr, "  %s\n", version.Version().Version)
	fmt.Fprintf(os.Stderr, "\n")
}

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(1)
	}

	var run func([]string) error
	switch strings.ToLower(os.Args[1]) {
	case "serve":
		run = runServe
	case "version":
		run = runVersion
	case "help", "-h", "--help":
		usage()
		return
	default:
		usage()
		os.Exit(1)
	}

	if err := run(os.Args[2:]); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}
`

const serverMakefileTemplate = `all: build

.PHONY: build

ifndef ($(GOPATH))
	GOPATH = $(HOME)/go
endif

PATH := $(GOPATH)/bin:$(PATH)
VERSION = $(shell git describe --tags --always --dirty)
BRANCH = $(shell git rev-parse --abbrev-ref HEAD)
REVISION = $(shell git rev-parse HEAD)
REVSHORT = $(shell git rev-parse --short HEAD)
USER = $(shell whoami)
GOVERSION = $(shell go version | awk '{print $$3}')
NOW	= $(shell date -u +"%Y-%m-%dT%H:%M:%SZ")
SHELL = /bin/bash

ifneq ($(OS), Windows_NT)
	CURRENT_PLATFORM = linux
	ifeq ($(shell uname), Darwin)
		SHELL := /bin/bash
		CURRENT_PLATFORM = darwin
	endif
else
	CURRENT_PLATFORM = windows
endif

BUILD_VERSION = "\
	-X github.com/{{.Org}}/{{.Name}}/vendor/github.com/micromdm/go4/version.appName=${APP_NAME} \
	-X github.com/{{.Org}}/{{.Name}}/vendor/github.com/micromdm/go4/version.version=${VERSION} \
	-X github.com/{{.Org}}/{{.Name}}/vendor/github.com/micromdm/go4/version.branch=${BRANCH} \
	-X github.com/{{.Org}}/{{.Name}}/vendor/github.com/micromdm/go4/version.buildUser=${USER} \
	-X github.com/{{.Org}}/{{.Name}}/vendor/github.com/micromdm/go4/version.buildDate=${NOW} \
	-X github.com/{{.Org}}/{{.Name}}/vendor/github.com/micromdm/go4/version.revision=${REVISION} \
	-X github.com/{{.Org}}/{{.Name}}/vendor/github.com/micromdm/go4/version.goVersion=${GOVERSION}"

WORKSPACE = ${GOPATH}/src/github.com/{{.Org}}/{{.Name}}
check-deps:
ifneq ($(shell test -e ${WORKSPACE}/Gopkg.lock && echo -n yes), yes)
	@echo "folder is clonded in the wrong place, copying to a Go Workspace"
	@echo "See: https://golang.org/doc/code.html#Workspaces"
	@git clone git@github.com:{{.Org}}/{{.Name}} ${WORKSPACE}
	@echo "cd to ${WORKSPACE} and run make deps again."
	@exit 1
endif
ifneq ($(shell pwd), $(WORKSPACE))
	@echo "cd to ${WORKSPACE} and run make deps again."
	@exit 1
endif

deps: check-deps
	go get -u github.com/golang/dep/...
	dep ensure -vendor-only

test:
	go test -cover -race -v $(shell go list ./... | grep -v /vendor/)

build: {{.Name}}

clean:
	rm -rf build/
	rm -f *.zip

.pre-build:
	mkdir -p build/darwin
	mkdir -p build/linux

INSTALL_STEPS := \
	install-{{.Name}}

install-local: $(INSTALL_STEPS)

.pre-{{.Name}}:
	$(eval APP_NAME = {{.Name}})

{{.Name}}: .pre-build .pre-{{.Name}}
	go build -i -o build/$(CURRENT_PLATFORM)/{{.Name}} -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}

install-{{.Name}}: .pre-{{.Name}}
	go install -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}

xp-{{.Name}}: .pre-build .pre-{{.Name}}
	GOOS=darwin go build -i -o build/darwin/{{.Name}} -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}
	GOOS=linux CGO_ENABLED=0 go build -i -o build/linux/{{.Name}}  -ldflags ${BUILD_VERSION} ./cmd/{{.Name}}

release-zip: xp-{{.Name}}
	zip -r {{.Name}}_${VERSION}.zip build/

# TODO remove after bootstrap is done.
git-init:
	git init
	git add -A
	git commit -m "first commit."

dep-init:
	dep init
	git add Gopkg.*
	git commit -m "Initialize Go dependencies."

init: git-init dep-init
`

const serverServeTemplate = `
package main

import (
	"flag"
	stdlog "log"
	"net/http"
	"os"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/groob/finalizer/logutil"
	"github.com/micromdm/go4/env"
	"github.com/micromdm/go4/httputil"
	"github.com/pkg/errors"
)

func runServe(args []string) error {
	flagset := flag.NewFlagSet("{{.Name}}", flag.ExitOnError)
	var (
		flConfigPath = flagset.String("config-dir", env.String("CONFIG_DIR", "/var/{{.Org}}/{{.Name}}"), "Path to server config directory.")
		flLogFormat  = flagset.String("log-format", env.String("LOG_FORMAT", "logfmt"), "Enable structured logging. Supported formats: logfmt, json.")
		flLogLevel   = flagset.String("log-level", env.String("LOG_LEVEL", "info"), "Log level. Either info or debug.")
		flHTTPDebug  = flagset.Bool("http-debug", false, "Enable debug for http(dumps full request).")
		flHTTPAddr   = flagset.String("http-addr", env.String("HTTP_ADDR", ":https"), "HTTP(s) listen address of http server. Defaults to :443 or :8080 if tls=false")
		flTLS        = flagset.Bool("tls", env.Bool("USE_TLS", true), "Serve HTTPS.")
		flTLSCert    = flagset.String("tls-cert", env.String("TLS_CERT", ""), "Path to TLS certificate.")
		flTLSKey     = flagset.String("tls-key", env.String("TLS_KEY", ""), "Path to TLS private key.")
		flTLSDomain  = flagset.String("tls-domain", env.String("TLS_DOMAIN", ""), "Automatically fetch certs from Let's Encrypt for this domain. Format must be server.acme.co")
	)

	flagset.Usage = usageFor(flagset, "{{.Name}} serve [flags]")
	if err := flagset.Parse(args); err != nil {
		return err
	}

	var (
		logger     log.Logger
		httpLogger log.Logger
	)
	{
		w := log.NewSyncWriter(os.Stderr)
		switch *flLogFormat {
		case "json":
			logger = log.NewJSONLogger(w)
		default:
			logger = log.NewLogfmtLogger(w)
		}
		stdlog.SetOutput(log.NewStdlibAdapter(logger))
		logger = log.With(logger, "ts", log.DefaultTimestampUTC)
		if *flLogLevel == "debug" {
			logger = level.NewFilter(logger, level.AllowDebug())
		} else {
			logger = level.NewFilter(logger, level.AllowInfo())
		}
		httpLogger = log.With(logger, "component", "http_logger")
		logger = log.With(logger, "caller", log.Caller(4))
	}

	mux := http.NewServeMux()
	var handler http.Handler
	if *flHTTPDebug {
		handler = httputil.HTTPDebugMiddleware(os.Stdout, true, httpLogger.Log)(mux)
	} else {
		handler = mux
	}
	handler = logutil.NewHTTPLogger(httpLogger).Middleware(handler)

	serveOpts := httputil.Simple(
		*flConfigPath,
		handler,
		*flHTTPAddr,
		*flTLSCert,
		*flTLSKey,
		*flTLS,
		logger,
		*flTLSDomain,
	)

	err := httputil.ListenAndServe(serveOpts...)
	return errors.Wrap(err, "calling ListenAndServe")
}
`

const gitignoreTemplate = `.DS_Store
build/
vendor/
*.zip
*.tar.gz
`

const dockerfileTemplate = `FROM alpine

RUN apk --update add \
    ca-certificates

COPY ./build/linux/{{.Name}} /usr/bin/{{.Name}}

CMD ["{{.Name}}"]
`
<|endoftext|>


<commit_before>// Based on ssh/terminal:
// Copyright 2011 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.

// +build windows

package terminal

import (
	"syscall"
	"unsafe"
)

var kernel32 = syscall.NewLazyDLL("kernel32.dll")

var (
	procGetConsoleMode = kernel32.NewProc("GetConsoleMode")
)

// IsTerminal returns true if the given file descriptor is a terminal.
func IsTerminal(fd int) bool {
	var st uint32
	r, _, e := syscall.Syscall(procGetConsoleMode.Addr(), 2, uintptr(fd), uintptr(unsafe.Pointer(&st)), 0)
	return r != 0 && e == 0
}
<commit_msg>terminal_windows: Use syscall.Handle for IsTerminal argument.<commit_after>// Based on ssh/terminal:
// Copyright 2011 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.

// +build windows

package terminal

import (
	"syscall"
	"unsafe"
)

var kernel32 = syscall.NewLazyDLL("kernel32.dll")

var (
	procGetConsoleMode = kernel32.NewProc("GetConsoleMode")
)

// IsTerminal returns true if the given file descriptor is a terminal.
func IsTerminal(fd syscall.Handle) bool {
	var st uint32
	r, _, e := syscall.Syscall(procGetConsoleMode.Addr(), 2, uintptr(fd), uintptr(unsafe.Pointer(&st)), 0)
	return r != 0 && e == 0
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"

	"strings"

	"github.com/joho/godotenv"
)

func main() {
	var showHelp bool
	flag.BoolVar(&showHelp, "h", false, "show help")
	var rawEnvFilenames string
	flag.StringVar(&rawEnvFilenames, "f", "", "comma separated paths to .env files")

	flag.Parse()

	usage := `
Run a process with a env setup from a .env file

godotenv [-f ENV_FILE_PATHS] COMMAND_ARGS

ENV_FILE_PATHS: comma separated paths to .env files
COMMAND_ARGS: command and args you want to run

example
  godotenv -f /path/to/something/.env,/another/path/.env fortune
`
	// if no args or -h flag
	// print usage and return
	args := flag.Args()
	if showHelp || len(args) == 0 {
		fmt.Println(usage)
		return
	}

	// load env
	var envFilenames []string
	if rawEnvFilenames != "" {
		envFilenames = strings.Split(rawEnvFilenames, ",")
	}

	// take rest of args and "exec" them
	cmd := args[0]
	cmdArgs := args[1:len(args)]

	err := godotenv.Exec(envFilenames, cmd, cmdArgs)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>run `gofmt -w -s ./..`<commit_after>package main

import (
	"flag"
	"fmt"
	"log"

	"strings"

	"github.com/joho/godotenv"
)

func main() {
	var showHelp bool
	flag.BoolVar(&showHelp, "h", false, "show help")
	var rawEnvFilenames string
	flag.StringVar(&rawEnvFilenames, "f", "", "comma separated paths to .env files")

	flag.Parse()

	usage := `
Run a process with a env setup from a .env file

godotenv [-f ENV_FILE_PATHS] COMMAND_ARGS

ENV_FILE_PATHS: comma separated paths to .env files
COMMAND_ARGS: command and args you want to run

example
  godotenv -f /path/to/something/.env,/another/path/.env fortune
`
	// if no args or -h flag
	// print usage and return
	args := flag.Args()
	if showHelp || len(args) == 0 {
		fmt.Println(usage)
		return
	}

	// load env
	var envFilenames []string
	if rawEnvFilenames != "" {
		envFilenames = strings.Split(rawEnvFilenames, ",")
	}

	// take rest of args and "exec" them
	cmd := args[0]
	cmdArgs := args[1:]

	err := godotenv.Exec(envFilenames, cmd, cmdArgs)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/kataras/iris"
)

func main() {
	// first way:
	// simple way for simple things
	// PreListen before a station is listening ( iris.Listen/TLS...)
	iris.Plugins.PreListen(func(s *iris.Framework) {
		for _, route := range s.Lookups() {
			fmt.Printf("Func: Route Method: %s | Subdomain %s | Path: %s is going to be registed with %d handler(s). \n", route.Method(), route.Subdomain(), route.Path(), len(route.Middleware()))
		}
	})

	// second way:
	// structured way for more things
	plugin := myPlugin{}
	iris.Plugins.Add(plugin)

	iris.Get("/first_route", aHandler)

	iris.Post("/second_route", aHandler)

	iris.Put("/third_route", aHandler)

	iris.Get("/fourth_route", aHandler)

	iris.Listen(":8080")
}

func aHandler(ctx *iris.Context) {
	ctx.Write("Hello from: %s", ctx.PathString())
}

type myPlugin struct{}

// PostListen after a station is listening ( iris.Listen/TLS...)
func (pl myPlugin) PostListen(s *iris.Framework) {
	fmt.Printf("myPlugin: server is listening on host: %s", s.Servers.Main().Host())
}

//list:
/*
	Activate(iris.PluginContainer)
	GetName() string
	GetDescription() string
	PreListen(*iris.Framework)
	PostListen(*iris.Framework)
	PreClose(*iris.Framework)
	PreDownload(thePlugin iris.Plugin, downloadUrl string)
	// for custom events
	On(string,...func())
	Call(string)
*/
<commit_msg>fix plugin example (unused custom callbacks removed)<commit_after>package main

import (
	"fmt"

	"github.com/kataras/iris"
)

func main() {
	// first way:
	// simple way for simple things
	// PreListen before a station is listening ( iris.Listen/TLS...)
	iris.Plugins.PreListen(func(s *iris.Framework) {
		for _, route := range s.Lookups() {
			fmt.Printf("Func: Route Method: %s | Subdomain %s | Path: %s is going to be registed with %d handler(s). \n", route.Method(), route.Subdomain(), route.Path(), len(route.Middleware()))
		}
	})

	// second way:
	// structured way for more things
	plugin := myPlugin{}
	iris.Plugins.Add(plugin)

	iris.Get("/first_route", aHandler)

	iris.Post("/second_route", aHandler)

	iris.Put("/third_route", aHandler)

	iris.Get("/fourth_route", aHandler)

	iris.Listen(":8080")
}

func aHandler(ctx *iris.Context) {
	ctx.Write("Hello from: %s", ctx.PathString())
}

type myPlugin struct{}

// PostListen after a station is listening ( iris.Listen/TLS...)
func (pl myPlugin) PostListen(s *iris.Framework) {
	fmt.Printf("myPlugin: server is listening on host: %s", s.Servers.Main().Host())
}

//list:
/*
	Activate(iris.PluginContainer)
	GetName() string
	GetDescription() string
	PreListen(*iris.Framework)
	PostListen(*iris.Framework)
	PreClose(*iris.Framework)
	PreDownload(thePlugin iris.Plugin, downloadUrl string)
	// for custom events use go-events, https://github.com/kataras/go-events
*/
<|endoftext|>


<commit_before>package extra

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"time"

	"github.com/pkg/errors"

	"github.com/jmoiron/sqlx"

	"github.com/elpinal/coco3/extra/ast"
	"github.com/elpinal/coco3/extra/parser" // Only for ParseError.
	"github.com/elpinal/coco3/extra/typed"
	"github.com/elpinal/coco3/extra/types"
)

type Env struct {
	cmds map[string]typed.Command
	Option
}

type Option struct {
	DB *sqlx.DB
}

func New(opt Option) Env {
	return Env{
		Option: opt,
		cmds: map[string]typed.Command{
			"exec":    execCommand,
			"execenv": execenvCommand, // exec with env
			"cd":      cdCommand,
			"exit":    exitCommand,
			"free":    freeCommand,
			"history": historyCommand,

			"remove": removeCommand,

			"ls":   lsCommand,
			"man":  manCommand,
			"make": makeCommand,

			"git":   gitCommand,
			"cargo": cargoCommand,
			"go":    goCommand,
			"stack": stackCommand,
			"lein":  leinCommand,
			"ocaml": ocamlCommand,

			"vim":    vimCommand,
			"emacs":  emacsCommand,
			"screen": screenCommand,

			"cnp":  cnpCommand,
			"gvmn": gvmnCommand,
			"vvmn": vvmnCommand,
		},
	}
}

func WithoutDefault() Env {
	return Env{cmds: make(map[string]typed.Command)}
}

func (e *Env) Bind(name string, c typed.Command) {
	e.cmds[name] = c
}

func (e *Env) Eval(command *ast.Command) (err error) {
	if command == nil {
		return nil
	}
	tc, found := e.cmds[command.Name.Lit]
	if !found {
		return &parser.ParseError{
			Msg:    fmt.Sprintf("no such typed command: %q", command.Name.Lit),
			Line:   command.Name.Line,
			Column: command.Name.Column,
		}
	}
	if len(command.Args) != len(tc.Params) {
		return &parser.ParseError{
			Msg:    fmt.Sprintf("the length of args (%d) != the one of params (%d)", len(command.Args), len(tc.Params)),
			Line:   command.Name.Line,
			Column: command.Name.Column,
		}
	}
	for i, arg := range command.Args {
		if arg.Type() != tc.Params[i] {
			return &parser.ParseError{
				Msg:    fmt.Sprintf("type mismatch: (%v) (type of %v) does not match with (%v) (expected type)", arg.Type(), arg, tc.Params[i]),
				Line:   command.Name.Line,
				Column: command.Name.Column,
			}
		}
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	defer close(c)
	defer signal.Stop(c)

	defer func() {
		r := recover()
		if r == nil {
			return
		}
		var ok bool
		// may overwrite error of tc.Fn.
		err, ok = r.(error)
		if !ok {
			panic(r)
		}
	}()

	return tc.Fn(command.Args, e.DB)
}

func toSlice(list ast.List) ([]string, error) {
	ret := make([]string, 0, list.Length())
	for {
		switch x := list.(type) {
		case *ast.Cons:
			ret = append(ret, x.Head)
			list = x.Tail
		case *ast.Empty:
			return ret, nil
		default:
			return nil, fmt.Errorf("unexpected list type: %T", x)
		}
	}
}

var execCommand = typed.Command{
	Params: []types.Type{types.String, types.StringList},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, "exec")
		}
		cmd := stdCmd(args[0].(*ast.String).Lit, cmdArgs...)
		return cmd.Run()
	},
}

var execenvCommand = typed.Command{
	Params: []types.Type{types.StringList, types.String, types.StringList},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[2].(ast.List))
		if err != nil {
			return errors.Wrap(err, "execenv")
		}
		cmd := stdCmd(args[1].(*ast.String).Lit, cmdArgs...)
		env, err := toSlice(args[0].(ast.List))
		if err != nil {
			return errors.Wrap(err, "execenv")
		}
		for _, e := range env {
			if !strings.Contains(e, "=") {
				return errors.New(`execenv: each item of the first argument must be the form "key=value"`)
			}
		}
		cmd.Env = append(os.Environ(), env...)
		return cmd.Run()
	},
}

var cdCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		return os.Chdir(args[0].(*ast.String).Lit)
	},
}

var exitCommand = typed.Command{
	Params: []types.Type{types.Int},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		n, err := strconv.Atoi(args[0].(*ast.Int).Lit)
		if err != nil {
			return err
		}
		os.Exit(n)
		return nil
	},
}

var freeCommand = typed.Command{
	Params: []types.Type{types.String, types.StringList},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, "free")
		}
		name := args[0].(*ast.String).Lit
		cmd := exec.Cmd{Path: name, Args: append([]string{name}, cmdArgs...)}
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Stdin = os.Stdin
		return cmd.Run()
	},
}

func commandArgs(name string) func([]ast.Expr, *sqlx.DB) error {
	return func(args []ast.Expr, _ *sqlx.DB) error {
		list, err := toSlice(args[0].(ast.List))
		if err != nil {
			return errors.Wrap(err, name)
		}
		return stdCmd(name, list...).Run()
	}
}

func commandsInCommand(name string) func([]ast.Expr, *sqlx.DB) error {
	return func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, name)
		}
		var cmd *exec.Cmd
		switch lit := args[0].(*ast.Ident).Lit; lit {
		case "command":
			cmd = stdCmd(name, cmdArgs...)
		default:
			cmd = stdCmd(name, append([]string{lit}, cmdArgs...)...)
		}
		return cmd.Run()
	}
}

func goCommand1() func([]ast.Expr, *sqlx.DB) error {
	return func(args []ast.Expr, _ *sqlx.DB) error {
		name := "go"
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, name)
		}
		var cmd *exec.Cmd
		switch lit := args[0].(*ast.Ident).Lit; lit {
		case "command":
			cmd = stdCmd(name, cmdArgs...)
		case "testall":
			// I can't be confident in using such
			// a subcommand-specific way.  Another suggestion might
			// be like `go test all`, where 'all' is a postfix
			// operator of './...'.
			cmd = stdCmd(name, append([]string{"test"}, append(cmdArgs, "./...")...)...)
		default:
			cmd = stdCmd(name, append([]string{lit}, cmdArgs...)...)
		}
		return cmd.Run()
	}
}

var gitCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("git"),
}

var cargoCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("cargo"),
}

var goCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     goCommand1(),
}

var stackCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("stack"),
}

var leinCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("lein"),
}

func stdCmd(name string, args ...string) *exec.Cmd {
	cmd := exec.Command(name, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	return cmd
}

func stdExec(name string, args ...string) func([]ast.Expr, *sqlx.DB) error {
	return func(_ []ast.Expr, _ *sqlx.DB) error {
		return stdCmd(name, args...).Run()
	}
}

var vimCommand = typed.Command{
	Params: []types.Type{},
	Fn:     stdExec("vim"),
}

var emacsCommand = typed.Command{
	Params: []types.Type{},
	Fn:     stdExec("emacs"),
}

func withEnv(s string, cmd *exec.Cmd) *exec.Cmd {
	if cmd.Env == nil {
		cmd.Env = os.Environ()
	}
	cmd.Env = append(cmd.Env, s)
	return cmd
}

var screenCommand = typed.Command{
	Params: []types.Type{},
	Fn: func(_ []ast.Expr, _ *sqlx.DB) error {
		return withEnv("LANG=en_US.UTF-8", stdCmd("screen")).Run()
	},
}

type execution struct {
	Time time.Time
	Line string
}

var historyCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, db *sqlx.DB) error {
		var jsonFormat bool
		var enc *json.Encoder
		switch format := e[0].(*ast.String).Lit; format {
		case "json":
			jsonFormat = true
		case "lines":
		default:
			return fmt.Errorf("history: format %q is not supported", format)
		}

		buf := bufio.NewWriter(os.Stdout)
		if jsonFormat {
			enc = json.NewEncoder(buf)
		}
		rows, err := db.Queryx("select * from command_info")
		if err != nil {
			return err
		}
		data := execution{}
		for rows.Next() {
			err := rows.StructScan(&data)
			if err != nil {
				return err
			}
			if jsonFormat {
				err := enc.Encode(data)
				if err != nil {
					return err
				}
			} else {
				buf.WriteString(data.Time.Format("Mon, 02 Jan 2006 15:04:05"))
				buf.Write([]byte("  "))
				buf.WriteString(data.Line)
				buf.WriteByte('\n')
			}
		}
		return buf.Flush()
	},
}

var lsCommand = typed.Command{
	Params: []types.Type{},
	Fn:     stdExec("ls", "--show-control-chars", "--color=auto"),
}

var manCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, _ *sqlx.DB) error {
		lit := e[0].(*ast.String).Lit
		return stdCmd("man", lit).Run()
	},
}

var removeCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, _ *sqlx.DB) error {
		lit := e[0].(*ast.String).Lit
		return stdCmd("rm", "-i", lit).Run()
	},
}

var cnpCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, _ *sqlx.DB) error {
		lit := e[0].(*ast.String).Lit
		return stdCmd("create-new-project", lit).Run()
	},
}

var gvmnCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("gvmn"),
}

var vvmnCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("vvmn"),
}

var makeCommand = typed.Command{
	Params: []types.Type{types.StringList},
	Fn:     commandArgs("make"),
}

var ocamlCommand = typed.Command{
	Params: []types.Type{types.StringList},
	Fn:     commandArgs("ocaml"),
}
<commit_msg>Add remove function<commit_after>package extra

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"time"

	"github.com/pkg/errors"

	"github.com/jmoiron/sqlx"

	"github.com/elpinal/coco3/extra/ast"
	"github.com/elpinal/coco3/extra/parser" // Only for ParseError.
	"github.com/elpinal/coco3/extra/typed"
	"github.com/elpinal/coco3/extra/types"
)

type Env struct {
	cmds map[string]typed.Command
	Option
}

type Option struct {
	DB *sqlx.DB
}

func New(opt Option) Env {
	return Env{
		Option: opt,
		cmds: map[string]typed.Command{
			"exec":    execCommand,
			"execenv": execenvCommand, // exec with env
			"cd":      cdCommand,
			"exit":    exitCommand,
			"free":    freeCommand,
			"history": historyCommand,

			"remove": removeCommand,

			"ls":   lsCommand,
			"man":  manCommand,
			"make": makeCommand,

			"git":   gitCommand,
			"cargo": cargoCommand,
			"go":    goCommand,
			"stack": stackCommand,
			"lein":  leinCommand,
			"ocaml": ocamlCommand,

			"vim":    vimCommand,
			"emacs":  emacsCommand,
			"screen": screenCommand,

			"cnp":  cnpCommand,
			"gvmn": gvmnCommand,
			"vvmn": vvmnCommand,
		},
	}
}

func WithoutDefault() Env {
	return Env{cmds: make(map[string]typed.Command)}
}

func (e *Env) Bind(name string, c typed.Command) {
	e.cmds[name] = c
}

func (e *Env) Eval(command *ast.Command) (err error) {
	if command == nil {
		return nil
	}
	tc, found := e.cmds[command.Name.Lit]
	if !found {
		return &parser.ParseError{
			Msg:    fmt.Sprintf("no such typed command: %q", command.Name.Lit),
			Line:   command.Name.Line,
			Column: command.Name.Column,
		}
	}
	if len(command.Args) != len(tc.Params) {
		return &parser.ParseError{
			Msg:    fmt.Sprintf("the length of args (%d) != the one of params (%d)", len(command.Args), len(tc.Params)),
			Line:   command.Name.Line,
			Column: command.Name.Column,
		}
	}
	for i, arg := range command.Args {
		if arg.Type() != tc.Params[i] {
			return &parser.ParseError{
				Msg:    fmt.Sprintf("type mismatch: (%v) (type of %v) does not match with (%v) (expected type)", arg.Type(), arg, tc.Params[i]),
				Line:   command.Name.Line,
				Column: command.Name.Column,
			}
		}
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	defer close(c)
	defer signal.Stop(c)

	defer func() {
		r := recover()
		if r == nil {
			return
		}
		var ok bool
		// may overwrite error of tc.Fn.
		err, ok = r.(error)
		if !ok {
			panic(r)
		}
	}()

	return tc.Fn(command.Args, e.DB)
}

func toSlice(list ast.List) ([]string, error) {
	ret := make([]string, 0, list.Length())
	for {
		switch x := list.(type) {
		case *ast.Cons:
			ret = append(ret, x.Head)
			list = x.Tail
		case *ast.Empty:
			return ret, nil
		default:
			return nil, fmt.Errorf("unexpected list type: %T", x)
		}
	}
}

var execCommand = typed.Command{
	Params: []types.Type{types.String, types.StringList},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, "exec")
		}
		cmd := stdCmd(args[0].(*ast.String).Lit, cmdArgs...)
		return cmd.Run()
	},
}

var execenvCommand = typed.Command{
	Params: []types.Type{types.StringList, types.String, types.StringList},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[2].(ast.List))
		if err != nil {
			return errors.Wrap(err, "execenv")
		}
		cmd := stdCmd(args[1].(*ast.String).Lit, cmdArgs...)
		env, err := toSlice(args[0].(ast.List))
		if err != nil {
			return errors.Wrap(err, "execenv")
		}
		for _, e := range env {
			if !strings.Contains(e, "=") {
				return errors.New(`execenv: each item of the first argument must be the form "key=value"`)
			}
		}
		cmd.Env = append(os.Environ(), env...)
		return cmd.Run()
	},
}

var cdCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		return os.Chdir(args[0].(*ast.String).Lit)
	},
}

var exitCommand = typed.Command{
	Params: []types.Type{types.Int},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		n, err := strconv.Atoi(args[0].(*ast.Int).Lit)
		if err != nil {
			return err
		}
		os.Exit(n)
		return nil
	},
}

var freeCommand = typed.Command{
	Params: []types.Type{types.String, types.StringList},
	Fn: func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, "free")
		}
		name := args[0].(*ast.String).Lit
		cmd := exec.Cmd{Path: name, Args: append([]string{name}, cmdArgs...)}
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Stdin = os.Stdin
		return cmd.Run()
	},
}

func commandArgs(name string) func([]ast.Expr, *sqlx.DB) error {
	return func(args []ast.Expr, _ *sqlx.DB) error {
		list, err := toSlice(args[0].(ast.List))
		if err != nil {
			return errors.Wrap(err, name)
		}
		return stdCmd(name, list...).Run()
	}
}

func commandsInCommand(name string) func([]ast.Expr, *sqlx.DB) error {
	return func(args []ast.Expr, _ *sqlx.DB) error {
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, name)
		}
		var cmd *exec.Cmd
		switch lit := args[0].(*ast.Ident).Lit; lit {
		case "command":
			cmd = stdCmd(name, cmdArgs...)
		default:
			cmd = stdCmd(name, append([]string{lit}, cmdArgs...)...)
		}
		return cmd.Run()
	}
}

func goCommand1() func([]ast.Expr, *sqlx.DB) error {
	return func(args []ast.Expr, _ *sqlx.DB) error {
		name := "go"
		cmdArgs, err := toSlice(args[1].(ast.List))
		if err != nil {
			return errors.Wrap(err, name)
		}
		var cmd *exec.Cmd
		switch lit := args[0].(*ast.Ident).Lit; lit {
		case "command":
			cmd = stdCmd(name, cmdArgs...)
		case "testall":
			// I can't be confident in using such
			// a subcommand-specific way.  Another suggestion might
			// be like `go test all`, where 'all' is a postfix
			// operator of './...'.
			cmd = stdCmd(name, append([]string{"test"}, append(cmdArgs, "./...")...)...)
		default:
			cmd = stdCmd(name, append([]string{lit}, cmdArgs...)...)
		}
		return cmd.Run()
	}
}

var gitCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("git"),
}

var cargoCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("cargo"),
}

var goCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     goCommand1(),
}

var stackCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("stack"),
}

var leinCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("lein"),
}

func stdCmd(name string, args ...string) *exec.Cmd {
	cmd := exec.Command(name, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	return cmd
}

func stdExec(name string, args ...string) func([]ast.Expr, *sqlx.DB) error {
	return func(_ []ast.Expr, _ *sqlx.DB) error {
		return stdCmd(name, args...).Run()
	}
}

var vimCommand = typed.Command{
	Params: []types.Type{},
	Fn:     stdExec("vim"),
}

var emacsCommand = typed.Command{
	Params: []types.Type{},
	Fn:     stdExec("emacs"),
}

func withEnv(s string, cmd *exec.Cmd) *exec.Cmd {
	if cmd.Env == nil {
		cmd.Env = os.Environ()
	}
	cmd.Env = append(cmd.Env, s)
	return cmd
}

var screenCommand = typed.Command{
	Params: []types.Type{},
	Fn: func(_ []ast.Expr, _ *sqlx.DB) error {
		return withEnv("LANG=en_US.UTF-8", stdCmd("screen")).Run()
	},
}

type execution struct {
	Time time.Time
	Line string
}

var historyCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, db *sqlx.DB) error {
		var jsonFormat bool
		var enc *json.Encoder
		switch format := e[0].(*ast.String).Lit; format {
		case "json":
			jsonFormat = true
		case "lines":
		default:
			return fmt.Errorf("history: format %q is not supported", format)
		}

		buf := bufio.NewWriter(os.Stdout)
		if jsonFormat {
			enc = json.NewEncoder(buf)
		}
		rows, err := db.Queryx("select * from command_info")
		if err != nil {
			return err
		}
		data := execution{}
		for rows.Next() {
			err := rows.StructScan(&data)
			if err != nil {
				return err
			}
			if jsonFormat {
				err := enc.Encode(data)
				if err != nil {
					return err
				}
			} else {
				buf.WriteString(data.Time.Format("Mon, 02 Jan 2006 15:04:05"))
				buf.Write([]byte("  "))
				buf.WriteString(data.Line)
				buf.WriteByte('\n')
			}
		}
		return buf.Flush()
	},
}

var lsCommand = typed.Command{
	Params: []types.Type{},
	Fn:     stdExec("ls", "--show-control-chars", "--color=auto"),
}

var manCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, _ *sqlx.DB) error {
		lit := e[0].(*ast.String).Lit
		return stdCmd("man", lit).Run()
	},
}

var removeCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, _ *sqlx.DB) error {
		lit := e[0].(*ast.String).Lit
		return stdCmd("rm", "-i", lit).Run()
	},
}

func remove(exprs []ast.Expr, _ *sqlx.DB) error {
}

var cnpCommand = typed.Command{
	Params: []types.Type{types.String},
	Fn: func(e []ast.Expr, _ *sqlx.DB) error {
		lit := e[0].(*ast.String).Lit
		return stdCmd("create-new-project", lit).Run()
	},
}

var gvmnCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("gvmn"),
}

var vvmnCommand = typed.Command{
	Params: []types.Type{types.Ident, types.StringList},
	Fn:     commandsInCommand("vvmn"),
}

var makeCommand = typed.Command{
	Params: []types.Type{types.StringList},
	Fn:     commandArgs("make"),
}

var ocamlCommand = typed.Command{
	Params: []types.Type{types.StringList},
	Fn:     commandArgs("ocaml"),
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/influxdb/influxdb/cmd/influxd/backup"
	"github.com/influxdb/influxdb/cmd/influxd/help"
	"github.com/influxdb/influxdb/cmd/influxd/restore"
	"github.com/influxdb/influxdb/cmd/influxd/run"
)

// These variables are populated via the Go linker.
var (
	version string = "0.9"
	commit  string
)

func init() {
	// If commit not set, make that clear.
	if commit == "" {
		commit = "unknown"
	}
}

func main() {
	rand.Seed(time.Now().UnixNano())

	m := NewMain()
	if err := m.Run(os.Args[1:]...); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

// Main represents the program execution.
type Main struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

// NewMain return a new instance of Main.
func NewMain() *Main {
	return &Main{
		Stdin:  os.Stdin,
		Stdout: os.Stdout,
		Stderr: os.Stderr,
	}
}

// Run determines and runs the command specified by the CLI args.
func (m *Main) Run(args ...string) error {
	name, args := ParseCommandName(args)

	// Extract name from args.
	switch name {
	case "", "run":
		cmd := run.NewCommand()

		// Tell the server the build details.
		cmd.Version = version
		cmd.Commit = commit

		if err := cmd.Run(args...); err != nil {
			return fmt.Errorf("run: %s", err)
		}

		// Wait indefinitely.
		<-(chan struct{})(nil)

	case "backup":
		name := backup.NewCommand()
		if err := name.Run(args...); err != nil {
			return fmt.Errorf("backup: %s", err)
		}
	case "restore":
		name := restore.NewCommand()
		if err := name.Run(args...); err != nil {
			return fmt.Errorf("restore: %s", err)
		}
	case "config":
		if err := run.NewPrintConfigCommand().Run(args...); err != nil {
			return fmt.Errorf("config: %s", err)
		}
	case "version":
		if err := NewVersionCommand().Run(args...); err != nil {
			return fmt.Errorf("version: %s", err)
		}
	case "help":
		if err := help.NewCommand().Run(args...); err != nil {
			return fmt.Errorf("help: %s", err)
		}
	default:
		return fmt.Errorf(`unknown command "%s"`+"\n"+`Run 'influxd help' for usage`+"\n\n", name)
	}

	return nil
}

// ParseCommandName extracts the command name and args from the args list.
func ParseCommandName(args []string) (string, []string) {
	// Retrieve command name as first argument.
	var name string
	if len(args) > 0 && !strings.HasPrefix(args[0], "-") {
		name = args[0]
	}

	// Special case -h immediately following binary name
	if len(args) > 0 && args[0] == "-h" {
		name = "help"
	}

	// If command is "help" and has an argument then rewrite args to use "-h".
	if name == "help" && len(args) > 1 {
		args[0], args[1] = args[1], "-h"
		name = args[0]
	}

	// If a named command is specified then return it with its arguments.
	if name != "" {
		return name, args[1:]
	}
	return "", args
}

// Command represents the command executed by "influxd version".
type VersionCommand struct {
	Stdout io.Writer
	Stderr io.Writer
}

// NewVersionCommand return a new instance of VersionCommand.
func NewVersionCommand() *VersionCommand {
	return &VersionCommand{
		Stdout: os.Stdout,
		Stderr: os.Stderr,
	}
}

// Run prints the current version and commit info.
func (cmd *VersionCommand) Run(args ...string) error {
	// Parse flags in case -h is specified.
	fs := flag.NewFlagSet("", flag.ContinueOnError)
	fs.Usage = func() { fmt.Fprintln(cmd.Stderr, strings.TrimSpace(versionUsage)) }
	if err := fs.Parse(args); err != nil {
		return err
	}

	// Print version info.
	fmt.Fprintf(cmd.Stdout, "InfluxDB v%s (git: %s)\n", version, commit)

	return nil
}

var versionUsage = `
usage: version

	version displays the InfluxDB version and build git commit hash
`
<commit_msg>Dump run errors to stderr<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/influxdb/influxdb/cmd/influxd/backup"
	"github.com/influxdb/influxdb/cmd/influxd/help"
	"github.com/influxdb/influxdb/cmd/influxd/restore"
	"github.com/influxdb/influxdb/cmd/influxd/run"
)

// These variables are populated via the Go linker.
var (
	version string = "0.9"
	commit  string
)

func init() {
	// If commit not set, make that clear.
	if commit == "" {
		commit = "unknown"
	}
}

func main() {
	rand.Seed(time.Now().UnixNano())

	m := NewMain()
	if err := m.Run(os.Args[1:]...); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

// Main represents the program execution.
type Main struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

// NewMain return a new instance of Main.
func NewMain() *Main {
	return &Main{
		Stdin:  os.Stdin,
		Stdout: os.Stdout,
		Stderr: os.Stderr,
	}
}

// Run determines and runs the command specified by the CLI args.
func (m *Main) Run(args ...string) error {
	name, args := ParseCommandName(args)

	// Extract name from args.
	switch name {
	case "", "run":
		cmd := run.NewCommand()

		// Tell the server the build details.
		cmd.Version = version
		cmd.Commit = commit

		if err := cmd.Run(args...); err != nil {
			return fmt.Errorf("run: %s", err)
		}

		// Wait indefinitely.
		<-(chan struct{})(nil)

	case "backup":
		name := backup.NewCommand()
		if err := name.Run(args...); err != nil {
			return fmt.Errorf("backup: %s", err)
		}
	case "restore":
		name := restore.NewCommand()
		if err := name.Run(args...); err != nil {
			return fmt.Errorf("restore: %s", err)
		}
	case "config":
		if err := run.NewPrintConfigCommand().Run(args...); err != nil {
			return fmt.Errorf("config: %s", err)
		}
	case "version":
		if err := NewVersionCommand().Run(args...); err != nil {
			return fmt.Errorf("version: %s", err)
		}
	case "help":
		if err := help.NewCommand().Run(args...); err != nil {
			return fmt.Errorf("help: %s", err)
		}
	default:
		return fmt.Errorf(`unknown command "%s"`+"\n"+`Run 'influxd help' for usage`+"\n\n", name)
	}

	return nil
}

// ParseCommandName extracts the command name and args from the args list.
func ParseCommandName(args []string) (string, []string) {
	// Retrieve command name as first argument.
	var name string
	if len(args) > 0 && !strings.HasPrefix(args[0], "-") {
		name = args[0]
	}

	// Special case -h immediately following binary name
	if len(args) > 0 && args[0] == "-h" {
		name = "help"
	}

	// If command is "help" and has an argument then rewrite args to use "-h".
	if name == "help" && len(args) > 1 {
		args[0], args[1] = args[1], "-h"
		name = args[0]
	}

	// If a named command is specified then return it with its arguments.
	if name != "" {
		return name, args[1:]
	}
	return "", args
}

// Command represents the command executed by "influxd version".
type VersionCommand struct {
	Stdout io.Writer
	Stderr io.Writer
}

// NewVersionCommand return a new instance of VersionCommand.
func NewVersionCommand() *VersionCommand {
	return &VersionCommand{
		Stdout: os.Stdout,
		Stderr: os.Stderr,
	}
}

// Run prints the current version and commit info.
func (cmd *VersionCommand) Run(args ...string) error {
	// Parse flags in case -h is specified.
	fs := flag.NewFlagSet("", flag.ContinueOnError)
	fs.Usage = func() { fmt.Fprintln(cmd.Stderr, strings.TrimSpace(versionUsage)) }
	if err := fs.Parse(args); err != nil {
		return err
	}

	// Print version info.
	fmt.Fprintf(cmd.Stdout, "InfluxDB v%s (git: %s)\n", version, commit)

	return nil
}

var versionUsage = `
usage: version

	version displays the InfluxDB version and build git commit hash
`
<|endoftext|>