<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"syscall"
	"time"

	"github.com/ekanite/ekanite"
	"github.com/ekanite/ekanite/input"
	"github.com/ekanite/ekanite/status"
)

var (
	stats = expvar.NewMap("ekanite")
)

// Program parameters
var datadir string
var tcpIface string
var udpIface string
var caPemPath string
var caKeyPath string
var queryIface string
var batchSize int
var batchTimeout int
var indexMaxPending int
var gomaxprocs int
var numShards int
var retentionPeriod string
var cpuProfile string
var memProfile string
var inputFormat string

// Flag set
var fs *flag.FlagSet

// Types
const (
	DefaultDataDir         = "/var/opt/ekanite"
	DefaultBatchSize       = 300
	DefaultBatchTimeout    = 1000
	DefaultIndexMaxPending = 1000
	DefaultNumShards       = 4
	DefaultRetentionPeriod = "168h"
	DefaultQueryAddr       = "localhost:9950"
	DefaultHTTPQueryAddr   = "localhost:8080"
	DefaultDiagsIface      = "localhost:9951"
	DefaultTCPServer       = "localhost:5514"
	DefaultInputFormat     = "syslog"
)

func main() {
	fs = flag.NewFlagSet("", flag.ExitOnError)
	var (
		datadir         = fs.String("datadir", DefaultDataDir, "Set data directory")
		batchSize       = fs.Int("batchsize", DefaultBatchSize, "Indexing batch size")
		batchTimeout    = fs.Int("batchtime", DefaultBatchTimeout, "Indexing batch timeout, in milliseconds")
		indexMaxPending = fs.Int("maxpending", DefaultIndexMaxPending, "Maximum pending index events")
		tcpIface        = fs.String("tcp", DefaultTCPServer, "Syslog server TCP bind address in the form host:port. To disable set to empty string")
		udpIface        = fs.String("udp", "", "Syslog server UDP bind address in the form host:port. If not set, not started")
		diagIface       = fs.String("diag", DefaultDiagsIface, "expvar and pprof bind address in the form host:port. If not set, not started")
		caPemPath       = fs.String("tlspem", "", "path to CA PEM file for TLS-enabled TCP server. If not set, TLS not activated")
		caKeyPath       = fs.String("tlskey", "", "path to CA key file for TLS-enabled TCP server. If not set, TLS not activated")
		queryIface      = fs.String("query", DefaultQueryAddr, "TCP Bind address for query server in the form host:port. To disable set to empty string")
		queryIfaceHttp  = fs.String("queryhttp", DefaultHTTPQueryAddr, "TCP Bind address for http query server in the form host:port. To disable set to empty string")
		numShards       = fs.Int("numshards", DefaultNumShards, "Set number of shards per index")
		retentionPeriod = fs.String("retention", DefaultRetentionPeriod, "Data retention period. Minimum is 24 hours")
		cpuProfile      = fs.String("cpuprof", "", "Where to write CPU profiling data. Not written if not set")
		memProfile      = fs.String("memprof", "", "Where to write memory profiling data. Not written if not set")
		inputFormat     = fs.String("input", DefaultInputFormat, "Message format of input (only syslog supported)")
	)
	fs.Usage = printHelp
	fs.Parse(os.Args[1:])

	absDataDir, err := filepath.Abs(*datadir)
	if err != nil {
		log.Fatalf("failed to get absolute data path for '%s': %s", *datadir, err.Error())
	}

	// Get the retention period.
	retention, err := time.ParseDuration(*retentionPeriod)
	if err != nil {
		log.Fatalf("failed to parse retention period '%s'", *retentionPeriod)
	}

	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[ekanite] ")
	log.Printf("ekanite started using %s for index storage", absDataDir)

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("GOMAXPROCS set to", runtime.GOMAXPROCS(0))

	var diagServer *status.Service
	// Start the expvar handler if requested.
	if *diagIface != "" {
		diagServer = status.NewService(*diagIface)

		if diagServer.Start(); err != nil {
			log.Fatalf("failed to start diag server: %s", err.Error())
		}
	}

	// Create and open the Engine.
	engine := ekanite.NewEngine(absDataDir)
	if engine == nil {
		log.Fatalf("failed to create indexing engine at %s", absDataDir)
	}
	engine.NumShards = *numShards
	engine.RetentionPeriod = retention

	if err := engine.Open(); err != nil {
		log.Fatalf("failed to open engine: %s", err.Error())
	}
	log.Printf("engine opened with shard number of %d, retention period of %s",
		engine.NumShards, engine.RetentionPeriod)

	// Start the simple query server if requested.
	if *queryIface != "" {
		server := ekanite.NewServer(*queryIface, engine)
		if server == nil {
			log.Fatal("failed to create query server")
		}
		if err := server.Start(); err != nil {
			log.Fatalf("failed to start query server: %s", err.Error())
		}
		log.Printf("query server listening to %s", *queryIface)
	}

	// Start the http query server if requested.
	if *queryIfaceHttp != "" {
		server := ekanite.NewHTTPServer(*queryIfaceHttp, engine)
		if server == nil {
			log.Fatal("failed to create HTTP query server")
		}
		if err := server.Start(); err != nil {
			log.Fatalf("failed to start HTTP query server: %s", err.Error())
		}
		log.Printf("HTTP query server listening to %s", *queryIfaceHttp)
	}

	// Create and start the batcher.
	batcherTimeout := time.Duration(*batchTimeout) * time.Millisecond
	batcher := ekanite.NewBatcher(engine, *batchSize, batcherTimeout, *indexMaxPending)
	if batcher == nil {
		log.Fatal("failed to create indexing batcher")
	}

	errChan := make(chan error)
	if err := batcher.Start(errChan); err != nil {
		log.Fatalf("failed to start indexing batcher: %s", err.Error())
	}
	log.Printf("batching configured with size %d, timeout %s, max pending %d",
		*batchSize, batcherTimeout, *indexMaxPending)

	// Start draining batcher errors.
	go func() {
		for {
			select {
			case err := <-errChan:
				if err != nil {
					log.Printf("error indexing batch: %s", err.Error())
				}
			}
		}
	}()

	// Start TCP collector if requested.
	if *tcpIface != "" {
		var tlsConfig *tls.Config
		if *caPemPath != "" && *caKeyPath != "" {
			tlsConfig, err = newTLSConfig(*caPemPath, *caKeyPath)
			if err != nil {
				log.Fatalf("failed to configure TLS: %s", err.Error())
			}
			log.Printf("TLS successfully configured")
		}

		collector, err := input.NewCollector("tcp", *tcpIface, *inputFormat, tlsConfig)
		if err != nil {
			log.Fatalf("failed to create TCP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start TCP collector: %s", err.Error())
		}
		log.Printf("TCP collector listening to %s", *tcpIface)
	}

	// Start UDP collector if requested.
	if *udpIface != "" {
		collector, err := input.NewCollector("udp", *udpIface, *inputFormat, nil)
		if collector == nil {
			log.Fatalf("failed to create UDP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start UDP collector: %s", err.Error())
		}
		log.Printf("UDP collector listening to %s", *udpIface)
	}

	// Start profiling.
	startProfile(*cpuProfile, *memProfile)

	stats.Set("launch", time.Now().UTC())

	// Set up signal handling.
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Block until one of the signals above is received
	select {
	case <-signalCh:
		log.Println("signal received, shutting down...")
	}

	stopProfile()
}

func newTLSConfig(caPemPath, caKeyPath string) (*tls.Config, error) {
	var config *tls.Config

	caPem, err := ioutil.ReadFile(caPemPath)
	if err != nil {
		return nil, err
	}
	ca, err := x509.ParseCertificate(caPem)
	if err != nil {
		return nil, err
	}

	caKey, err := ioutil.ReadFile(caKeyPath)
	if err != nil {
		return nil, err
	}
	key, err := x509.ParsePKCS1PrivateKey(caKey)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	pool.AddCert(ca)

	cert := tls.Certificate{
		Certificate: [][]byte{caPem},
		PrivateKey:  key,
	}

	config = &tls.Config{
		ClientAuth:   tls.RequireAndVerifyClientCert,
		Certificates: []tls.Certificate{cert},
		ClientCAs:    pool,
	}

	config.Rand = rand.Reader

	return config, nil
}

// prof stores the file locations of active profiles.
var prof struct {
	cpu *os.File
	mem *os.File
}

// StartProfile initializes the cpu and memory profile, if specified.
func startProfile(cpuprofile, memprofile string) {
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("cpuprofile: %v", err)
		}
		log.Printf("writing CPU profile to: %s\n", cpuprofile)
		prof.cpu = f
		pprof.StartCPUProfile(prof.cpu)
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing memory profile to: %s\n", memprofile)
		prof.mem = f
		runtime.MemProfileRate = 4096
	}

}

// StopProfile closes the cpu and memory profiles if they are running.
func stopProfile() {
	if prof.cpu != nil {
		pprof.StopCPUProfile()
		prof.cpu.Close()
		log.Println("CPU profile stopped")
	}
	if prof.mem != nil {
		pprof.Lookup("heap").WriteTo(prof.mem, 0)
		prof.mem.Close()
		log.Println("memory profile stopped")
	}
}

func printHelp() {
	fmt.Println("ekanited [options]")
	fs.PrintDefaults()
}
<commit_msg>Better logging for status server<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"syscall"
	"time"

	"github.com/ekanite/ekanite"
	"github.com/ekanite/ekanite/input"
	"github.com/ekanite/ekanite/status"
)

var (
	stats = expvar.NewMap("ekanite")
)

// Program parameters
var datadir string
var tcpIface string
var udpIface string
var caPemPath string
var caKeyPath string
var queryIface string
var batchSize int
var batchTimeout int
var indexMaxPending int
var gomaxprocs int
var numShards int
var retentionPeriod string
var cpuProfile string
var memProfile string
var inputFormat string

// Flag set
var fs *flag.FlagSet

// Types
const (
	DefaultDataDir         = "/var/opt/ekanite"
	DefaultBatchSize       = 300
	DefaultBatchTimeout    = 1000
	DefaultIndexMaxPending = 1000
	DefaultNumShards       = 4
	DefaultRetentionPeriod = "168h"
	DefaultQueryAddr       = "localhost:9950"
	DefaultHTTPQueryAddr   = "localhost:8080"
	DefaultDiagsIface      = "localhost:9951"
	DefaultTCPServer       = "localhost:5514"
	DefaultInputFormat     = "syslog"
)

func main() {
	fs = flag.NewFlagSet("", flag.ExitOnError)
	var (
		datadir         = fs.String("datadir", DefaultDataDir, "Set data directory")
		batchSize       = fs.Int("batchsize", DefaultBatchSize, "Indexing batch size")
		batchTimeout    = fs.Int("batchtime", DefaultBatchTimeout, "Indexing batch timeout, in milliseconds")
		indexMaxPending = fs.Int("maxpending", DefaultIndexMaxPending, "Maximum pending index events")
		tcpIface        = fs.String("tcp", DefaultTCPServer, "Syslog server TCP bind address in the form host:port. To disable set to empty string")
		udpIface        = fs.String("udp", "", "Syslog server UDP bind address in the form host:port. If not set, not started")
		diagIface       = fs.String("diag", DefaultDiagsIface, "expvar and pprof bind address in the form host:port. If not set, not started")
		caPemPath       = fs.String("tlspem", "", "path to CA PEM file for TLS-enabled TCP server. If not set, TLS not activated")
		caKeyPath       = fs.String("tlskey", "", "path to CA key file for TLS-enabled TCP server. If not set, TLS not activated")
		queryIface      = fs.String("query", DefaultQueryAddr, "TCP Bind address for query server in the form host:port. To disable set to empty string")
		queryIfaceHttp  = fs.String("queryhttp", DefaultHTTPQueryAddr, "TCP Bind address for http query server in the form host:port. To disable set to empty string")
		numShards       = fs.Int("numshards", DefaultNumShards, "Set number of shards per index")
		retentionPeriod = fs.String("retention", DefaultRetentionPeriod, "Data retention period. Minimum is 24 hours")
		cpuProfile      = fs.String("cpuprof", "", "Where to write CPU profiling data. Not written if not set")
		memProfile      = fs.String("memprof", "", "Where to write memory profiling data. Not written if not set")
		inputFormat     = fs.String("input", DefaultInputFormat, "Message format of input (only syslog supported)")
	)
	fs.Usage = printHelp
	fs.Parse(os.Args[1:])

	absDataDir, err := filepath.Abs(*datadir)
	if err != nil {
		log.Fatalf("failed to get absolute data path for '%s': %s", *datadir, err.Error())
	}

	// Get the retention period.
	retention, err := time.ParseDuration(*retentionPeriod)
	if err != nil {
		log.Fatalf("failed to parse retention period '%s'", *retentionPeriod)
	}

	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[ekanite] ")
	log.Printf("ekanite started using %s for index storage", absDataDir)

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("GOMAXPROCS set to", runtime.GOMAXPROCS(0))

	var diagServer *status.Service
	// Start the expvar handler if requested.
	if *diagIface != "" {
		diagServer = status.NewService(*diagIface)

		if diagServer.Start(); err != nil {
			log.Fatalf("failed to start status server: %s", err.Error())
		}
	}

	// Create and open the Engine.
	engine := ekanite.NewEngine(absDataDir)
	if engine == nil {
		log.Fatalf("failed to create indexing engine at %s", absDataDir)
	}
	engine.NumShards = *numShards
	engine.RetentionPeriod = retention

	if err := engine.Open(); err != nil {
		log.Fatalf("failed to open engine: %s", err.Error())
	}
	log.Printf("engine opened with shard number of %d, retention period of %s",
		engine.NumShards, engine.RetentionPeriod)

	// Start the simple query server if requested.
	if *queryIface != "" {
		server := ekanite.NewServer(*queryIface, engine)
		if server == nil {
			log.Fatal("failed to create query server")
		}
		if err := server.Start(); err != nil {
			log.Fatalf("failed to start query server: %s", err.Error())
		}
		log.Printf("query server listening to %s", *queryIface)
	}

	// Start the http query server if requested.
	if *queryIfaceHttp != "" {
		server := ekanite.NewHTTPServer(*queryIfaceHttp, engine)
		if server == nil {
			log.Fatal("failed to create HTTP query server")
		}
		if err := server.Start(); err != nil {
			log.Fatalf("failed to start HTTP query server: %s", err.Error())
		}
		log.Printf("HTTP query server listening to %s", *queryIfaceHttp)
	}

	// Create and start the batcher.
	batcherTimeout := time.Duration(*batchTimeout) * time.Millisecond
	batcher := ekanite.NewBatcher(engine, *batchSize, batcherTimeout, *indexMaxPending)
	if batcher == nil {
		log.Fatal("failed to create indexing batcher")
	}

	errChan := make(chan error)
	if err := batcher.Start(errChan); err != nil {
		log.Fatalf("failed to start indexing batcher: %s", err.Error())
	}
	log.Printf("batching configured with size %d, timeout %s, max pending %d",
		*batchSize, batcherTimeout, *indexMaxPending)

	// Start draining batcher errors.
	go func() {
		for {
			select {
			case err := <-errChan:
				if err != nil {
					log.Printf("error indexing batch: %s", err.Error())
				}
			}
		}
	}()

	// Start TCP collector if requested.
	if *tcpIface != "" {
		var tlsConfig *tls.Config
		if *caPemPath != "" && *caKeyPath != "" {
			tlsConfig, err = newTLSConfig(*caPemPath, *caKeyPath)
			if err != nil {
				log.Fatalf("failed to configure TLS: %s", err.Error())
			}
			log.Printf("TLS successfully configured")
		}

		collector, err := input.NewCollector("tcp", *tcpIface, *inputFormat, tlsConfig)
		if err != nil {
			log.Fatalf("failed to create TCP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start TCP collector: %s", err.Error())
		}
		log.Printf("TCP collector listening to %s", *tcpIface)
	}

	// Start UDP collector if requested.
	if *udpIface != "" {
		collector, err := input.NewCollector("udp", *udpIface, *inputFormat, nil)
		if collector == nil {
			log.Fatalf("failed to create UDP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start UDP collector: %s", err.Error())
		}
		log.Printf("UDP collector listening to %s", *udpIface)
	}

	// Start profiling.
	startProfile(*cpuProfile, *memProfile)

	stats.Set("launch", time.Now().UTC())

	// Set up signal handling.
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Block until one of the signals above is received
	select {
	case <-signalCh:
		log.Println("signal received, shutting down...")
	}

	stopProfile()
}

func newTLSConfig(caPemPath, caKeyPath string) (*tls.Config, error) {
	var config *tls.Config

	caPem, err := ioutil.ReadFile(caPemPath)
	if err != nil {
		return nil, err
	}
	ca, err := x509.ParseCertificate(caPem)
	if err != nil {
		return nil, err
	}

	caKey, err := ioutil.ReadFile(caKeyPath)
	if err != nil {
		return nil, err
	}
	key, err := x509.ParsePKCS1PrivateKey(caKey)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	pool.AddCert(ca)

	cert := tls.Certificate{
		Certificate: [][]byte{caPem},
		PrivateKey:  key,
	}

	config = &tls.Config{
		ClientAuth:   tls.RequireAndVerifyClientCert,
		Certificates: []tls.Certificate{cert},
		ClientCAs:    pool,
	}

	config.Rand = rand.Reader

	return config, nil
}

// prof stores the file locations of active profiles.
var prof struct {
	cpu *os.File
	mem *os.File
}

// StartProfile initializes the cpu and memory profile, if specified.
func startProfile(cpuprofile, memprofile string) {
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("cpuprofile: %v", err)
		}
		log.Printf("writing CPU profile to: %s\n", cpuprofile)
		prof.cpu = f
		pprof.StartCPUProfile(prof.cpu)
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing memory profile to: %s\n", memprofile)
		prof.mem = f
		runtime.MemProfileRate = 4096
	}

}

// StopProfile closes the cpu and memory profiles if they are running.
func stopProfile() {
	if prof.cpu != nil {
		pprof.StopCPUProfile()
		prof.cpu.Close()
		log.Println("CPU profile stopped")
	}
	if prof.mem != nil {
		pprof.Lookup("heap").WriteTo(prof.mem, 0)
		prof.mem.Close()
		log.Println("memory profile stopped")
	}
}

func printHelp() {
	fmt.Println("ekanited [options]")
	fs.PrintDefaults()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/coyove/common/sched"
	"github.com/coyove/goflyway"
	"github.com/coyove/goflyway/v"
	"golang.org/x/crypto/acme/autocert"
)

var (
	version      = "__devel__"
	remoteAddr   string
	localAddr    string
	addr         string
	httpsProxy   string
	resetTraffic bool
	cconfig      = &goflyway.ClientConfig{}
	sconfig      = &goflyway.ServerConfig{}
)

func printHelp(a ...interface{}) {
	if len(a) > 0 {
		fmt.Printf("goflyway: ")
		fmt.Println(a...)
	}
	fmt.Println("usage: goflyway -DLhHUvkqpPtTwWy address:port")
	os.Exit(0)
}

func main() {
	sched.Verbose = false

	for i, last := 1, rune(0); i < len(os.Args); i++ {
		p := strings.TrimLeft(os.Args[i], "-")

		// HACK: ss-local compatible command flags
		if p == "fast-open" || p == "V" || p == "u" || p == "m" || p == "b" {
			if i < len(os.Args)-1 && !strings.HasPrefix(os.Args[i+1], "-") {
				i++
			}
			continue
		}

		if len(p) != len(os.Args[i]) {
			for i, c := range p {
				switch c {
				case 'h':
					printHelp()
				//case 'V':
				//	printHelp(version)
				case 'L', 'P', 'p', 'k', 't', 'T', 'W', 'H', 'U', 'D', 'c':
					last = c
				case 'v':
					v.Verbose++
				case 'q':
					v.Verbose = -1
				case 'w':
					cconfig.WebSocket = true
				case 'y':
					resetTraffic = true
				case '=':
					i++
					fallthrough
				default:
					if last == 0 {
						printHelp("illegal option --", string(c))
					}
					p = p[i:]
					goto PARSE
				}
			}
			continue
		}
	PARSE:
		if strings.HasPrefix(p, "\"") {
			if p, _ = strconv.Unquote(p); p == "" {
				printHelp("illegal option --", string(last))
			}
		}
		switch last {
		case 'D':
			cconfig.Dynamic = true
			fallthrough
		case 'L':
			switch parts := strings.Split(p, ":"); len(parts) {
			case 1:
				localAddr = ":" + parts[0]
			case 2:
				localAddr = p
			case 3:
				localAddr, remoteAddr = ":"+parts[0], parts[1]+":"+parts[2]
			case 4:
				localAddr, remoteAddr = parts[0]+":"+parts[1], parts[2]+":"+parts[3]
			default:
				printHelp("illegal option --", string(last), p)
			}
		case 'P':
			sconfig.ProxyPassAddr = p
		case 'U':
			cconfig.PathPattern = p
		case 'T':
			speed, _ := strconv.ParseInt(p, 10, 64)
			sconfig.SpeedThrot = goflyway.NewTokenBucket(speed, speed*25)
		case 'W':
			writebuffer, _ := strconv.ParseInt(p, 10, 64)
			sconfig.WriteBuffer, cconfig.WriteBuffer = writebuffer, writebuffer
		case 't':
			*(*int64)(&cconfig.Timeout), _ = strconv.ParseInt(p+"000000000", 10, 64)
			sconfig.Timeout = cconfig.Timeout
		case 'p', 'k':
			sconfig.Key, cconfig.Key = p, p
		case 'H':
			cconfig.URLHeader = p
			httpsProxy = p
		case 'c':
			buf, _ := ioutil.ReadFile(p)
			cmds := make(map[string]interface{})
			json.Unmarshal(buf, &cmds)
			cconfig.Key, cconfig.VPN = cmds["password"].(string), true
			addr = fmt.Sprintf("%v:%v", cmds["server"], cmds["server_port"])

			v.Verbose = 3
			v.Vprint(os.Args, " config: ", cmds)
		default:
			addr = p
		}
		last = 0
	}

	if addr == "" {
		if localAddr == "" {
			v.Vprint("assume you want a default server at :8100")
			addr = ":8100"
		} else {
			printHelp("missing address:port to listen/connect")
		}
	}

	if localAddr != "" && remoteAddr == "" {
		_, port, err1 := net.SplitHostPort(localAddr)
		host, _, err2 := net.SplitHostPort(addr)
		remoteAddr = host + ":" + port
		if err1 != nil || err2 != nil {
			printHelp("invalid address --", localAddr, addr)
		}
	}

	if localAddr != "" && remoteAddr != "" {
		cconfig.Bind = remoteAddr
		cconfig.Upstream = addr
		cconfig.Stat = &goflyway.Traffic{}

		if v.Verbose > 0 {
			go watchTraffic(cconfig, resetTraffic)
		}
		if cconfig.Dynamic {
			v.Vprint("dynamic: forward ", localAddr, " to * through ", addr)
		} else {
			v.Vprint("forward ", localAddr, " to ", remoteAddr, " through ", addr)
		}
		if cconfig.WebSocket {
			v.Vprint("relay: use Websocket protocol")
		}
		if a := os.Getenv("http_proxy") + os.Getenv("HTTP_PROXY"); a != "" {
			v.Vprint("note: system HTTP proxy is set to: ", a)
		}
		if a := os.Getenv("https_proxy") + os.Getenv("HTTPS_PROXY"); a != "" {
			v.Vprint("note: system HTTPS proxy is set to: ", a)
		}

		v.Eprint(goflyway.NewClient(localAddr, cconfig))
	} else if httpsProxy != "" {
		v.Vprint("server listen on ", addr, " (https://", httpsProxy, ")")
		m := &autocert.Manager{
			Cache:      autocert.DirCache("secret-dir"),
			Prompt:     autocert.AcceptTOS,
			HostPolicy: autocert.HostWhitelist(httpsProxy),
		}
		s := &http.Server{
			Addr:      addr,
			TLSConfig: m.TLSConfig(),
		}
		s.Handler = &connector{}
		v.Eprint(s.ListenAndServeTLS("", ""))
	} else {
		v.Vprint("server listen on ", addr)
		v.Eprint(goflyway.NewServer(addr, sconfig))
	}
}

func watchTraffic(cconfig *goflyway.ClientConfig, reset bool) {
	path := filepath.Join(os.TempDir(), "goflyway_traffic")

	tmpbuf, _ := ioutil.ReadFile(path)
	if len(tmpbuf) != 16 || reset {
		tmpbuf = make([]byte, 16)
	}

	cconfig.Stat.Set(int64(binary.BigEndian.Uint64(tmpbuf)), int64(binary.BigEndian.Uint64(tmpbuf[8:])))

	var lastSent, lastRecv int64
	for range time.Tick(time.Second * 5) {
		s, r := *cconfig.Stat.Sent(), *cconfig.Stat.Recv()
		sv, rv := float64(s-lastSent)/1024/1024/5, float64(r-lastRecv)/1024/1024/5
		lastSent, lastRecv = s, r

		if sv >= 0.001 || rv >= 0.001 {
			v.Vprint("client send: ", float64(s)/1024/1024, "M (", sv, "M/s), recv: ", float64(r)/1024/1024, "M (", rv, "M/s)")
		}

		binary.BigEndian.PutUint64(tmpbuf, uint64(s))
		binary.BigEndian.PutUint64(tmpbuf[8:], uint64(r))
		ioutil.WriteFile(path, tmpbuf, 0644)
	}
}

type connector struct{}

func (c *connector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "CONNECT" {
		w.WriteHeader(404)
		return
	}

	// we are inside GFW and should pass data to upstream
	host := r.URL.Host
	v.VVprint(r.URL)

	if !regexp.MustCompile(`:\d+$`).MatchString(host) {
		host += ":443"
	}

	up, err := net.Dial("tcp", host)
	if err != nil {
		v.Eprint(err)
		w.WriteHeader(500)
		return
	}

	hij, _ := w.(http.Hijacker) // No HTTP2
	proxyClient, _, err := hij.Hijack()
	if err != nil {
		v.Eprint(err)
		w.WriteHeader(500)
		return
	}

	proxyClient.Write([]byte("HTTP/1.1 101 Switching Protocols"))

	go func() {
		if _, err := io.Copy(proxyClient, up); err != nil {
			v.Eprint(err)
		}
	}()

	go func() {
		if _, err := io.Copy(up, proxyClient); err != nil {
			v.Eprint(err)
		}
	}()
}
<commit_msg>DEbug<commit_after>package main

import (
	"crypto/tls"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/coyove/common/sched"
	"github.com/coyove/goflyway"
	"github.com/coyove/goflyway/v"
	"golang.org/x/crypto/acme/autocert"
)

var (
	version      = "__devel__"
	remoteAddr   string
	localAddr    string
	addr         string
	httpsProxy   string
	resetTraffic bool
	cconfig      = &goflyway.ClientConfig{}
	sconfig      = &goflyway.ServerConfig{}
)

func printHelp(a ...interface{}) {
	if len(a) > 0 {
		fmt.Printf("goflyway: ")
		fmt.Println(a...)
	}
	fmt.Println("usage: goflyway -DLhHUvkqpPtTwWy address:port")
	os.Exit(0)
}

func main() {
	sched.Verbose = false

	for i, last := 1, rune(0); i < len(os.Args); i++ {
		p := strings.TrimLeft(os.Args[i], "-")

		// HACK: ss-local compatible command flags
		if p == "fast-open" || p == "V" || p == "u" || p == "m" || p == "b" {
			if i < len(os.Args)-1 && !strings.HasPrefix(os.Args[i+1], "-") {
				i++
			}
			continue
		}

		if len(p) != len(os.Args[i]) {
			for i, c := range p {
				switch c {
				case 'h':
					printHelp()
				//case 'V':
				//	printHelp(version)
				case 'L', 'P', 'p', 'k', 't', 'T', 'W', 'H', 'U', 'D', 'c':
					last = c
				case 'v':
					v.Verbose++
				case 'q':
					v.Verbose = -1
				case 'w':
					cconfig.WebSocket = true
				case 'y':
					resetTraffic = true
				case '=':
					i++
					fallthrough
				default:
					if last == 0 {
						printHelp("illegal option --", string(c))
					}
					p = p[i:]
					goto PARSE
				}
			}
			continue
		}
	PARSE:
		if strings.HasPrefix(p, "\"") {
			if p, _ = strconv.Unquote(p); p == "" {
				printHelp("illegal option --", string(last))
			}
		}
		switch last {
		case 'D':
			cconfig.Dynamic = true
			fallthrough
		case 'L':
			switch parts := strings.Split(p, ":"); len(parts) {
			case 1:
				localAddr = ":" + parts[0]
			case 2:
				localAddr = p
			case 3:
				localAddr, remoteAddr = ":"+parts[0], parts[1]+":"+parts[2]
			case 4:
				localAddr, remoteAddr = parts[0]+":"+parts[1], parts[2]+":"+parts[3]
			default:
				printHelp("illegal option --", string(last), p)
			}
		case 'P':
			sconfig.ProxyPassAddr = p
		case 'U':
			cconfig.PathPattern = p
		case 'T':
			speed, _ := strconv.ParseInt(p, 10, 64)
			sconfig.SpeedThrot = goflyway.NewTokenBucket(speed, speed*25)
		case 'W':
			writebuffer, _ := strconv.ParseInt(p, 10, 64)
			sconfig.WriteBuffer, cconfig.WriteBuffer = writebuffer, writebuffer
		case 't':
			*(*int64)(&cconfig.Timeout), _ = strconv.ParseInt(p+"000000000", 10, 64)
			sconfig.Timeout = cconfig.Timeout
		case 'p', 'k':
			sconfig.Key, cconfig.Key = p, p
		case 'H':
			cconfig.URLHeader = p
			httpsProxy = p
		case 'c':
			buf, _ := ioutil.ReadFile(p)
			cmds := make(map[string]interface{})
			json.Unmarshal(buf, &cmds)
			cconfig.Key, cconfig.VPN = cmds["password"].(string), true
			addr = fmt.Sprintf("%v:%v", cmds["server"], cmds["server_port"])

			v.Verbose = 3
			v.Vprint(os.Args, " config: ", cmds)
		default:
			addr = p
		}
		last = 0
	}

	if addr == "" {
		if localAddr == "" {
			v.Vprint("assume you want a default server at :8100")
			addr = ":8100"
		} else {
			printHelp("missing address:port to listen/connect")
		}
	}

	if localAddr != "" && remoteAddr == "" {
		_, port, err1 := net.SplitHostPort(localAddr)
		host, _, err2 := net.SplitHostPort(addr)
		remoteAddr = host + ":" + port
		if err1 != nil || err2 != nil {
			printHelp("invalid address --", localAddr, addr)
		}
	}

	if localAddr != "" && remoteAddr != "" {
		cconfig.Bind = remoteAddr
		cconfig.Upstream = addr
		cconfig.Stat = &goflyway.Traffic{}

		if v.Verbose > 0 {
			go watchTraffic(cconfig, resetTraffic)
		}
		if cconfig.Dynamic {
			v.Vprint("dynamic: forward ", localAddr, " to * through ", addr)
		} else {
			v.Vprint("forward ", localAddr, " to ", remoteAddr, " through ", addr)
		}
		if cconfig.WebSocket {
			v.Vprint("relay: use Websocket protocol")
		}
		if a := os.Getenv("http_proxy") + os.Getenv("HTTP_PROXY"); a != "" {
			v.Vprint("note: system HTTP proxy is set to: ", a)
		}
		if a := os.Getenv("https_proxy") + os.Getenv("HTTPS_PROXY"); a != "" {
			v.Vprint("note: system HTTPS proxy is set to: ", a)
		}

		v.Eprint(goflyway.NewClient(localAddr, cconfig))
	} else if httpsProxy != "" {
		v.Vprint("server listen on ", addr, " (https://", httpsProxy, ")")
		m := &autocert.Manager{
			Cache:      autocert.DirCache("secret-dir"),
			Prompt:     autocert.AcceptTOS,
			HostPolicy: autocert.HostWhitelist(httpsProxy),
		}
		s := &http.Server{
			Addr:         addr,
			TLSConfig:    m.TLSConfig(),
			TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)),
		}
		s.Handler = &connector{}
		v.Eprint(s.ListenAndServeTLS("", ""))
	} else {
		v.Vprint("server listen on ", addr)
		v.Eprint(goflyway.NewServer(addr, sconfig))
	}
}

func watchTraffic(cconfig *goflyway.ClientConfig, reset bool) {
	path := filepath.Join(os.TempDir(), "goflyway_traffic")

	tmpbuf, _ := ioutil.ReadFile(path)
	if len(tmpbuf) != 16 || reset {
		tmpbuf = make([]byte, 16)
	}

	cconfig.Stat.Set(int64(binary.BigEndian.Uint64(tmpbuf)), int64(binary.BigEndian.Uint64(tmpbuf[8:])))

	var lastSent, lastRecv int64
	for range time.Tick(time.Second * 5) {
		s, r := *cconfig.Stat.Sent(), *cconfig.Stat.Recv()
		sv, rv := float64(s-lastSent)/1024/1024/5, float64(r-lastRecv)/1024/1024/5
		lastSent, lastRecv = s, r

		if sv >= 0.001 || rv >= 0.001 {
			v.Vprint("client send: ", float64(s)/1024/1024, "M (", sv, "M/s), recv: ", float64(r)/1024/1024, "M (", rv, "M/s)")
		}

		binary.BigEndian.PutUint64(tmpbuf, uint64(s))
		binary.BigEndian.PutUint64(tmpbuf[8:], uint64(r))
		ioutil.WriteFile(path, tmpbuf, 0644)
	}
}

type connector struct{}

func (c *connector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "CONNECT" {
		w.WriteHeader(404)
		return
	}

	// we are inside GFW and should pass data to upstream
	host := r.URL.Host
	v.VVprint(r.URL)

	if !regexp.MustCompile(`:\d+$`).MatchString(host) {
		host += ":443"
	}

	up, err := net.Dial("tcp", host)
	if err != nil {
		v.Eprint(err)
		w.WriteHeader(500)
		return
	}

	hij, _ := w.(http.Hijacker) // No HTTP2
	proxyClient, _, err := hij.Hijack()
	if err != nil {
		v.Eprint(err)
		w.WriteHeader(500)
		return
	}

	proxyClient.Write([]byte("HTTP/1.1 101 Switching Protocols"))

	go func() {
		if _, err := io.Copy(proxyClient, up); err != nil {
			v.Eprint(err)
		}
	}()

	go func() {
		if _, err := io.Copy(up, proxyClient); err != nil {
			v.Eprint(err)
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	d "github.com/studio-b12/gowebdav"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"strings"
)

func main() {
	root := flag.String("root", os.Getenv("ROOT"), "WebDAV Endpoint [ENV.ROOT]")
	usr := flag.String("user", os.Getenv("USER"), "User [ENV.USER]")
	pw := flag.String("pw", os.Getenv("PASSWORD"), "Password [ENV.PASSWORD]")
	netrc := flag.String("netrc-file", filepath.Join(getHome(), ".netrc"), "read login from netrc file")
	method := flag.String("X", "", `Method:
	LS <PATH>
	STAT <PATH>

	MKDIR <PATH>
	MKDIRALL <PATH>

	GET <PATH> [<FILE>]
	PUT <PATH> [<FILE>]

	MV <OLD> <NEW>
	CP <OLD> <NEW>

	DEL <PATH>
	`)
	flag.Parse()

	if *root == "" {
		fail("Set WebDAV ROOT")
	}

	if l := len(flag.Args()); l == 0 || l > 2 {
		fail("Unsupported arguments")
	}

	if *pw == "" {
		if u, p := d.ReadConfig(*root, *netrc); u != "" && p != "" {
			usr = &u
			pw = &p
		}
	}

	c := d.NewClient(*root, *usr, *pw)

	cmd := getCmd(*method)

	if e := cmd(c, flag.Arg(0), flag.Arg(1)); e != nil {
		fail(e)
	}
}

func fail(err interface{}) {
	if err != nil {
		fmt.Println(err)
	}
	os.Exit(-1)
}

func getHome() string {
	if u, e := user.Current(); e != nil {
		return u.HomeDir
	}
	return os.Getenv("HOME")
}

func getCmd(method string) func(c *d.Client, p0, p1 string) error {
	switch strings.ToUpper(method) {
	case "LS", "LIST", "PROPFIND":
		return cmdLs

	case "STAT":
		return cmdStat

	case "GET", "PULL", "READ":
		return cmdGet

	case "DELETE", "RM", "DEL":
		return cmdRm

	case "MKCOL", "MKDIR":
		return cmdMkdir

	case "MKCOLALL", "MKDIRALL", "MKDIRP":
		return cmdMkdirAll

	case "RENAME", "MV", "MOVE":
		return cmdMv

	case "COPY", "CP":
		return cmdCp

	case "PUT", "PUSH", "WRITE":
		return cmdPut

	default:
		return func(c *d.Client, p0, p1 string) (err error) {
			return errors.New("Unsupported method: " + method)
		}
	}
}

func cmdLs(c *d.Client, p0, _ string) (err error) {
	files, err := c.ReadDir(p0)
	if err == nil {
		fmt.Println(fmt.Sprintf("ReadDir: '%s' entries: %d ", p0, len(files)))
		for _, f := range files {
			fmt.Println(f)
		}
	}
	return
}

func cmdStat(c *d.Client, p0, _ string) (err error) {
	file, err := c.Stat(p0)
	if err == nil {
		fmt.Println(file)
	}
	return
}

func cmdGet(c *d.Client, p0, p1 string) (err error) {
	bytes, err := c.Read(p0)
	if err == nil {
		if p1 == "" {
			p1 = filepath.Join(".", p0)
		}
		err = writeFile(p1, bytes, 0644)
		if err == nil {
			fmt.Println(fmt.Sprintf("Written %d bytes to: %s", len(bytes), p1))
		}
	}
	return
}

func cmdRm(c *d.Client, p0, _ string) (err error) {
	if err = c.Remove(p0); err == nil {
		fmt.Println("Remove: " + p0)
	}
	return
}

func cmdMkdir(c *d.Client, p0, _ string) (err error) {
	if err = c.Mkdir(p0, 0755); err == nil {
		fmt.Println("Mkdir: " + p0)
	}
	return
}

func cmdMkdirAll(c *d.Client, p0, _ string) (err error) {
	if err = c.MkdirAll(p0, 0755); err == nil {
		fmt.Println("MkdirAll: " + p0)
	}
	return
}

func cmdMv(c *d.Client, p0, p1 string) (err error) {
	if err = c.Rename(p0, p1, true); err == nil {
		fmt.Println("Rename: " + p0 + " -> " + p1)
	}
	return
}

func cmdCp(c *d.Client, p0, p1 string) (err error) {
	if err = c.Copy(p0, p1, true); err == nil {
		fmt.Println("Copy: " + p0 + " -> " + p1)
	}
	return
}

func cmdPut(c *d.Client, p0, p1 string) (err error) {
	if p1 == "" {
		p1 = filepath.Join(".", p0)
	}
	stream, err := getStream(p1)
	if err != nil {
		return
	}
	defer stream.Close()

	if err = c.WriteStream(p0, stream, 0644); err == nil {
		fmt.Println("Put: " + p1 + " -> " + p0)
	}
	return
}

func writeFile(path string, bytes []byte, mode os.FileMode) error {
	parent := filepath.Dir(path)
	if _, e := os.Stat(parent); os.IsNotExist(e) {
		if e := os.MkdirAll(parent, os.ModePerm); e != nil {
			return e
		}
	}

	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = f.Write(bytes)
	return err
}

func getStream(pathOrString string) (io.ReadCloser, error) {
	fi, err := os.Stat(pathOrString)
	if err == nil {
		if fi.IsDir() {
			return nil, &os.PathError{
				Op:   "Open",
				Path: pathOrString,
				Err:  errors.New("Path: '" + pathOrString + "' is a directory"),
			}
		}
		f, err := os.Open(pathOrString)
		if err == nil {
			return f, nil
		}
		return nil, &os.PathError{
			Op:   "Open",
			Path: pathOrString,
			Err:  err,
		}
	}
	return nopCloser{strings.NewReader(pathOrString)}, nil
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error {
	return nil
}
<commit_msg>fix uploading file with wrong content. close #30<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	d "github.com/studio-b12/gowebdav"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"strings"
)

func main() {
	root := flag.String("root", os.Getenv("ROOT"), "WebDAV Endpoint [ENV.ROOT]")
	usr := flag.String("user", os.Getenv("USER"), "User [ENV.USER]")
	pw := flag.String("pw", os.Getenv("PASSWORD"), "Password [ENV.PASSWORD]")
	netrc := flag.String("netrc-file", filepath.Join(getHome(), ".netrc"), "read login from netrc file")
	method := flag.String("X", "", `Method:
	LS <PATH>
	STAT <PATH>

	MKDIR <PATH>
	MKDIRALL <PATH>

	GET <PATH> [<FILE>]
	PUT <PATH> [<FILE>]

	MV <OLD> <NEW>
	CP <OLD> <NEW>

	DEL <PATH>
	`)
	flag.Parse()

	if *root == "" {
		fail("Set WebDAV ROOT")
	}

	if l := len(flag.Args()); l == 0 || l > 2 {
		fail("Unsupported arguments")
	}

	if *pw == "" {
		if u, p := d.ReadConfig(*root, *netrc); u != "" && p != "" {
			usr = &u
			pw = &p
		}
	}

	c := d.NewClient(*root, *usr, *pw)

	cmd := getCmd(*method)

	if e := cmd(c, flag.Arg(0), flag.Arg(1)); e != nil {
		fail(e)
	}
}

func fail(err interface{}) {
	if err != nil {
		fmt.Println(err)
	}
	os.Exit(-1)
}

func getHome() string {
	if u, e := user.Current(); e != nil {
		return u.HomeDir
	}
	return os.Getenv("HOME")
}

func getCmd(method string) func(c *d.Client, p0, p1 string) error {
	switch strings.ToUpper(method) {
	case "LS", "LIST", "PROPFIND":
		return cmdLs

	case "STAT":
		return cmdStat

	case "GET", "PULL", "READ":
		return cmdGet

	case "DELETE", "RM", "DEL":
		return cmdRm

	case "MKCOL", "MKDIR":
		return cmdMkdir

	case "MKCOLALL", "MKDIRALL", "MKDIRP":
		return cmdMkdirAll

	case "RENAME", "MV", "MOVE":
		return cmdMv

	case "COPY", "CP":
		return cmdCp

	case "PUT", "PUSH", "WRITE":
		return cmdPut

	default:
		return func(c *d.Client, p0, p1 string) (err error) {
			return errors.New("Unsupported method: " + method)
		}
	}
}

func cmdLs(c *d.Client, p0, _ string) (err error) {
	files, err := c.ReadDir(p0)
	if err == nil {
		fmt.Println(fmt.Sprintf("ReadDir: '%s' entries: %d ", p0, len(files)))
		for _, f := range files {
			fmt.Println(f)
		}
	}
	return
}

func cmdStat(c *d.Client, p0, _ string) (err error) {
	file, err := c.Stat(p0)
	if err == nil {
		fmt.Println(file)
	}
	return
}

func cmdGet(c *d.Client, p0, p1 string) (err error) {
	bytes, err := c.Read(p0)
	if err == nil {
		if p1 == "" {
			p1 = filepath.Join(".", p0)
		}
		err = writeFile(p1, bytes, 0644)
		if err == nil {
			fmt.Println(fmt.Sprintf("Written %d bytes to: %s", len(bytes), p1))
		}
	}
	return
}

func cmdRm(c *d.Client, p0, _ string) (err error) {
	if err = c.Remove(p0); err == nil {
		fmt.Println("Remove: " + p0)
	}
	return
}

func cmdMkdir(c *d.Client, p0, _ string) (err error) {
	if err = c.Mkdir(p0, 0755); err == nil {
		fmt.Println("Mkdir: " + p0)
	}
	return
}

func cmdMkdirAll(c *d.Client, p0, _ string) (err error) {
	if err = c.MkdirAll(p0, 0755); err == nil {
		fmt.Println("MkdirAll: " + p0)
	}
	return
}

func cmdMv(c *d.Client, p0, p1 string) (err error) {
	if err = c.Rename(p0, p1, true); err == nil {
		fmt.Println("Rename: " + p0 + " -> " + p1)
	}
	return
}

func cmdCp(c *d.Client, p0, p1 string) (err error) {
	if err = c.Copy(p0, p1, true); err == nil {
		fmt.Println("Copy: " + p0 + " -> " + p1)
	}
	return
}

func cmdPut(c *d.Client, p0, p1 string) (err error) {
	if p1 == "" {
		p1 = filepath.Join(".", p0)
	}
	stream, err := getStream(p1)
	if err != nil {
		return
	}
	defer stream.Close()

	if err = c.WriteStream(p0, stream, 0644); err == nil {
		fmt.Println("Put: " + p1 + " -> " + p0)
	}
	return
}

func writeFile(path string, bytes []byte, mode os.FileMode) error {
	parent := filepath.Dir(path)
	if _, e := os.Stat(parent); os.IsNotExist(e) {
		if e := os.MkdirAll(parent, os.ModePerm); e != nil {
			return e
		}
	}

	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = f.Write(bytes)
	return err
}

func getStream(pathOrString string) (io.ReadCloser, error) {

	fi, err := os.Stat(pathOrString)
	if err != nil {
		return nil, err
	}

	if fi.IsDir() {
		return nil, &os.PathError{
			Op:   "Open",
			Path: pathOrString,
			Err:  errors.New("Path: '" + pathOrString + "' is a directory"),
		}
	}

	f, err := os.Open(pathOrString)
	if err == nil {
		return f, nil
	}

	return nil, &os.PathError{
		Op:   "Open",
		Path: pathOrString,
		Err:  err,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"

	"github.com/qjcg/horeb/pkg/horeb"
	pb "github.com/qjcg/horeb/proto"
	"github.com/spf13/viper"

	"google.golang.org/grpc"
)

// Setup for global levelled logging.
var (
	info  = log.New(os.Stderr, "INFO ", log.LstdFlags)
	debug = log.New(ioutil.Discard, "DEBUG ", log.LstdFlags)

	sep = flag.String("s", " ", "rune separator")
)

func main() {
	block := flag.String("b", "geometric", "unicode block")
	debugFlag := flag.Bool("d", false, "debug logging")
	num := flag.Int("n", 5, "number of runes to print")
	version := flag.Bool("v", false, "print version")
	flag.Parse()

	conf := viper.New()
	conf.SetDefault("host", "localhost")
	conf.SetDefault("port", 9999)

	conf.SetEnvPrefix("HOREBCTL")
	conf.BindEnv("host")
	conf.BindEnv("num")
	conf.BindEnv("port")

	if *version {
		fmt.Println(horeb.Version)
		return
	}

	if *debugFlag {
		debug.SetOutput(os.Stderr)
	}

	var opts []grpc.DialOption
	opts = append(opts, grpc.WithInsecure())
	conn, err := grpc.Dial(fmt.Sprintf("%s:%d", conf.GetString("host"), conf.GetInt("port")), opts...)
	if err != nil {
		info.Fatalf("fail to dial: %v", err)
	}
	defer conn.Close()
	client := pb.NewHorebClient(conn)

	err = getRuneStream(
		client,
		&pb.RuneRequest{
			Num:   int32(*num),
			Block: *block,
		},
	)
	if err != nil {
		info.Fatal(err)
	}
}

func getRuneStream(client pb.HorebClient, rr *pb.RuneRequest) error {
	debug.Printf("SENT: %#v\n", rr)

	stream, err := client.GetStream(context.Background(), rr)
	if err != nil {
		info.Printf("%v.GetStream(_) = _, %v\n", client, err)
		return err
	}

	for {
		streamedRune, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			info.Printf("stream receive error: %v\n", err)
		}
		fmt.Printf("%s%s", streamedRune.R, *sep)
	}
	debug.Printf("RECEIVED: %d %s", rr.Num, rr.Block)

	return nil
}
<commit_msg>Tweak import formatting<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"

	"github.com/qjcg/horeb/pkg/horeb"
	pb "github.com/qjcg/horeb/proto"

	"github.com/spf13/viper"
	"google.golang.org/grpc"
)

// Setup for global levelled logging.
var (
	info  = log.New(os.Stderr, "INFO ", log.LstdFlags)
	debug = log.New(ioutil.Discard, "DEBUG ", log.LstdFlags)

	sep = flag.String("s", " ", "rune separator")
)

func main() {
	block := flag.String("b", "geometric", "unicode block")
	debugFlag := flag.Bool("d", false, "debug logging")
	num := flag.Int("n", 5, "number of runes to print")
	version := flag.Bool("v", false, "print version")
	flag.Parse()

	conf := viper.New()
	conf.SetDefault("host", "localhost")
	conf.SetDefault("port", 9999)

	conf.SetEnvPrefix("HOREBCTL")
	conf.BindEnv("host")
	conf.BindEnv("num")
	conf.BindEnv("port")

	if *version {
		fmt.Println(horeb.Version)
		return
	}

	if *debugFlag {
		debug.SetOutput(os.Stderr)
	}

	var opts []grpc.DialOption
	opts = append(opts, grpc.WithInsecure())
	conn, err := grpc.Dial(fmt.Sprintf("%s:%d", conf.GetString("host"), conf.GetInt("port")), opts...)
	if err != nil {
		info.Fatalf("fail to dial: %v", err)
	}
	defer conn.Close()
	client := pb.NewHorebClient(conn)

	err = getRuneStream(
		client,
		&pb.RuneRequest{
			Num:   int32(*num),
			Block: *block,
		},
	)
	if err != nil {
		info.Fatal(err)
	}
}

func getRuneStream(client pb.HorebClient, rr *pb.RuneRequest) error {
	debug.Printf("SENT: %#v\n", rr)

	stream, err := client.GetStream(context.Background(), rr)
	if err != nil {
		info.Printf("%v.GetStream(_) = _, %v\n", client, err)
		return err
	}

	for {
		streamedRune, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			info.Printf("stream receive error: %v\n", err)
		}
		fmt.Printf("%s%s", streamedRune.R, *sep)
	}
	debug.Printf("RECEIVED: %d %s", rr.Num, rr.Block)

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strings"
)

// RunShellCommand executes a shell command and returns the output/error
func RunShellCommand(input string) (string, error) {
	inputCmd := strings.Split(input, " ")[0]
	args := strings.Split(input, " ")[1:]

	cmd := exec.Command(inputCmd, args...)
	outputBytes := &bytes.Buffer{}
	cmd.Stdout = outputBytes
	cmd.Stderr = outputBytes
	cmd.Start()
	err := cmd.Wait() // wait for command to finish
	outstring := outputBytes.String()
	return outstring, err
}

// HandleShellCommand runs the shell command
// The openTerm argument specifies whether a terminal should be opened (for viewing output
// or interacting with stdin)
func HandleShellCommand(input string, openTerm bool) {
	inputCmd := strings.Split(input, " ")[0]
	if !openTerm {
		// Simply run the command in the background and notify the user when it's done
		messenger.Message("Running...")
		go func() {
			output, err := RunShellCommand(input)
			totalLines := strings.Split(output, "\n")

			if len(totalLines) < 3 {
				if err == nil {
					messenger.Message(inputCmd, " exited without error")
				} else {
					messenger.Message(inputCmd, " exited with error: ", err, ": ", output)
				}
			} else {
				messenger.Message(output)
			}
			// We have to make sure to redraw
			RedrawAll()
		}()
	} else {
		// Shut down the screen because we're going to interact directly with the shell
		screen.Fini()
		screen = nil

		args := strings.Split(input, " ")[1:]

		// Set up everything for the command
		cmd := exec.Command(inputCmd, args...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		// This is a trap for Ctrl-C so that it doesn't kill micro
		// Instead we trap Ctrl-C to kill the program we're running
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)
		go func() {
			for range c {
				cmd.Process.Kill()
			}
		}()

		// Start the command
		cmd.Start()
		cmd.Wait()

		// This is just so we don't return right away and let the user press enter to return
		TermMessage("")

		// Start the screen back up
		InitScreen()
	}
}

// HandleCommand handles input from the user
func HandleCommand(input string) {
	inputCmd := strings.Split(input, " ")[0]
	args := strings.Split(input, " ")[1:]

	switch inputCmd {
	case "set":
		// Set an option and we have to set it for every view
		for _, view := range views {
			SetOption(view, args)
		}
	case "run":
		// Run a shell command in the background (openTerm is false)
		HandleShellCommand(strings.Join(args, " "), false)
	case "quit":
		// This is a bit weird because micro only has one view for now so there is no way to close
		// a single view
		// Currently if multiple views were open, it would close all of them, and not check the non-mainviews
		// for unsaved changes. This, and the behavior of Ctrl-Q need to be changed when splits are implemented
		if views[mainView].CanClose("Quit anyway? (yes, no, save) ") {
			screen.Fini()
			os.Exit(0)
		}
	case "save":
		// Save the main view
		views[mainView].Save()
	case "replace":
		// This is a regex to parse the replace expression
		// We allow no quotes if there are no spaces, but if you want to search
		// for or replace an expression with spaces, you can add double quotes
		r := regexp.MustCompile(`"[^"\\]*(?:\\.[^"\\]*)*"|[^\s]*`)
		replaceCmd := r.FindAllString(strings.Join(args, " "), -1)
		if len(replaceCmd) < 2 {
			// We need to find both a search and replace expression
			messenger.Error("Invalid replace statement: " + strings.Join(args, " "))
			return
		}

		var flags string
		if len(replaceCmd) == 3 {
			// The user included some flags
			flags = replaceCmd[2]
		}

		search := string(replaceCmd[0])
		replace := string(replaceCmd[1])

		// If the search and replace expressions have quotes, we need to remove those
		if strings.HasPrefix(search, `"`) && strings.HasSuffix(search, `"`) {
			search = search[1 : len(search)-1]
		}
		if strings.HasPrefix(replace, `"`) && strings.HasSuffix(replace, `"`) {
			replace = replace[1 : len(replace)-1]
		}

		// We replace all escaped double quotes to real double quotes
		search = strings.Replace(search, `\"`, `"`, -1)
		replace = strings.Replace(replace, `\"`, `"`, -1)

		regex, err := regexp.Compile(search)
		if err != nil {
			// There was an error with the user's regex
			messenger.Error(err.Error())
			return
		}

		view := views[mainView]

		found := false
		for {
			match := regex.FindStringIndex(view.Buf.String())
			if match == nil {
				break
			}
			found = true
			if strings.Contains(flags, "c") {
				// The 'check' flag was used
				Search(search, view, true)
				view.Relocate()
				RedrawAll()
				choice, canceled := messenger.YesNoPrompt("Perform replacement? (y,n)")
				if canceled {
					if view.Cursor.HasSelection() {
						view.Cursor.SetLoc(view.Cursor.curSelection[0])
						view.Cursor.ResetSelection()
					}
					messenger.Reset()
					return
				}
				if choice {
					view.Cursor.DeleteSelection()
					view.Buf.Insert(match[0], replace)
					view.Cursor.ResetSelection()
					messenger.Reset()
				} else {
					if view.Cursor.HasSelection() {
						searchStart = view.Cursor.curSelection[1]
					} else {
						searchStart = ToCharPos(view.Cursor.x, view.Cursor.y, view.Buf)
					}
					continue
				}
			} else {
				view.Buf.Replace(match[0], match[1], replace)
			}
		}
		if !found {
			messenger.Message("Nothing matched " + search)
		}
	default:
		messenger.Error("Unknown command: " + inputCmd)
	}
}
<commit_msg>Rematch every time a replacement is made<commit_after>package main

import (
	"bytes"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strings"
)

// RunShellCommand executes a shell command and returns the output/error
func RunShellCommand(input string) (string, error) {
	inputCmd := strings.Split(input, " ")[0]
	args := strings.Split(input, " ")[1:]

	cmd := exec.Command(inputCmd, args...)
	outputBytes := &bytes.Buffer{}
	cmd.Stdout = outputBytes
	cmd.Stderr = outputBytes
	cmd.Start()
	err := cmd.Wait() // wait for command to finish
	outstring := outputBytes.String()
	return outstring, err
}

// HandleShellCommand runs the shell command
// The openTerm argument specifies whether a terminal should be opened (for viewing output
// or interacting with stdin)
func HandleShellCommand(input string, openTerm bool) {
	inputCmd := strings.Split(input, " ")[0]
	if !openTerm {
		// Simply run the command in the background and notify the user when it's done
		messenger.Message("Running...")
		go func() {
			output, err := RunShellCommand(input)
			totalLines := strings.Split(output, "\n")

			if len(totalLines) < 3 {
				if err == nil {
					messenger.Message(inputCmd, " exited without error")
				} else {
					messenger.Message(inputCmd, " exited with error: ", err, ": ", output)
				}
			} else {
				messenger.Message(output)
			}
			// We have to make sure to redraw
			RedrawAll()
		}()
	} else {
		// Shut down the screen because we're going to interact directly with the shell
		screen.Fini()
		screen = nil

		args := strings.Split(input, " ")[1:]

		// Set up everything for the command
		cmd := exec.Command(inputCmd, args...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		// This is a trap for Ctrl-C so that it doesn't kill micro
		// Instead we trap Ctrl-C to kill the program we're running
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)
		go func() {
			for range c {
				cmd.Process.Kill()
			}
		}()

		// Start the command
		cmd.Start()
		cmd.Wait()

		// This is just so we don't return right away and let the user press enter to return
		TermMessage("")

		// Start the screen back up
		InitScreen()
	}
}

// HandleCommand handles input from the user
func HandleCommand(input string) {
	inputCmd := strings.Split(input, " ")[0]
	args := strings.Split(input, " ")[1:]

	switch inputCmd {
	case "set":
		// Set an option and we have to set it for every view
		for _, view := range views {
			SetOption(view, args)
		}
	case "run":
		// Run a shell command in the background (openTerm is false)
		HandleShellCommand(strings.Join(args, " "), false)
	case "quit":
		// This is a bit weird because micro only has one view for now so there is no way to close
		// a single view
		// Currently if multiple views were open, it would close all of them, and not check the non-mainviews
		// for unsaved changes. This, and the behavior of Ctrl-Q need to be changed when splits are implemented
		if views[mainView].CanClose("Quit anyway? (yes, no, save) ") {
			screen.Fini()
			os.Exit(0)
		}
	case "save":
		// Save the main view
		views[mainView].Save()
	case "replace":
		// This is a regex to parse the replace expression
		// We allow no quotes if there are no spaces, but if you want to search
		// for or replace an expression with spaces, you can add double quotes
		r := regexp.MustCompile(`"[^"\\]*(?:\\.[^"\\]*)*"|[^\s]*`)
		replaceCmd := r.FindAllString(strings.Join(args, " "), -1)
		if len(replaceCmd) < 2 {
			// We need to find both a search and replace expression
			messenger.Error("Invalid replace statement: " + strings.Join(args, " "))
			return
		}

		var flags string
		if len(replaceCmd) == 3 {
			// The user included some flags
			flags = replaceCmd[2]
		}

		search := string(replaceCmd[0])
		replace := string(replaceCmd[1])

		// If the search and replace expressions have quotes, we need to remove those
		if strings.HasPrefix(search, `"`) && strings.HasSuffix(search, `"`) {
			search = search[1 : len(search)-1]
		}
		if strings.HasPrefix(replace, `"`) && strings.HasSuffix(replace, `"`) {
			replace = replace[1 : len(replace)-1]
		}

		// We replace all escaped double quotes to real double quotes
		search = strings.Replace(search, `\"`, `"`, -1)
		replace = strings.Replace(replace, `\"`, `"`, -1)

		regex, err := regexp.Compile(search)
		if err != nil {
			// There was an error with the user's regex
			messenger.Error(err.Error())
			return
		}

		view := views[mainView]

		found := false
		for {
			match := regex.FindStringIndex(view.Buf.String())
			if match == nil {
				break
			}
			found = true
			if strings.Contains(flags, "c") {
				// The 'check' flag was used
				Search(search, view, true)
				view.Relocate()
				if settings["syntax"].(bool) {
					view.matches = Match(view)
				}
				RedrawAll()
				choice, canceled := messenger.YesNoPrompt("Perform replacement? (y,n)")
				if canceled {
					if view.Cursor.HasSelection() {
						view.Cursor.SetLoc(view.Cursor.curSelection[0])
						view.Cursor.ResetSelection()
					}
					messenger.Reset()
					return
				}
				if choice {
					view.Cursor.DeleteSelection()
					view.Buf.Insert(match[0], replace)
					view.Cursor.ResetSelection()
					messenger.Reset()
				} else {
					if view.Cursor.HasSelection() {
						searchStart = view.Cursor.curSelection[1]
					} else {
						searchStart = ToCharPos(view.Cursor.x, view.Cursor.y, view.Buf)
					}
					continue
				}
			} else {
				view.Buf.Replace(match[0], match[1], replace)
			}
		}
		if !found {
			messenger.Message("Nothing matched " + search)
		}
	default:
		messenger.Error("Unknown command: " + inputCmd)
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/gonum/matrix/mat64"
	"github.com/sbinet/npyio"
)

func main() {
	flag.Parse()

	if len(os.Args) <= 1 {
		flag.Usage()
		os.Exit(1)
	}

	allgood := true
	for i, fname := range os.Args[1:] {
		if i > 0 {
			fmt.Printf("\n")
		}
		f, err := os.Open(fname)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()

		r, err := npyio.NewReader(f)
		if err != nil {
			log.Fatal(err)
		}

		fmt.Printf(strings.Repeat("=", 80) + "\n")
		fmt.Printf("file: %v\n", fname)
		fmt.Printf("npy-header: %v\n", r.Header)

		if len(r.Header.Descr.Shape) > 2 {
			rt := npyio.TypeFrom(r.Header.Descr.Type)
			if rt == nil {
				fmt.Printf("error: no reflect.Type for %q\n", r.Header.Descr.Type)
				allgood = false
				continue
			}
			rv := reflect.New(reflect.SliceOf(rt))
			err = r.Read(rv.Interface())
			if err != nil {
				allgood = false
				fmt.Printf("error: %v\n", err)
				continue
			}
			fmt.Printf("data = %v\n", rv.Elem().Interface())
			continue
		}

		var m mat64.Dense
		err = r.Read(&m)
		if err != nil {
			allgood = false
			fmt.Printf("error: %v\n", err)
			continue
		}

		fmt.Printf("data = %v\n", mat64.Formatted(&m, mat64.Prefix("       ")))
	}

	if !allgood {
		os.Exit(1)
	}
}
<commit_msg>npyio-ls: detect and handle .npz archives<commit_after>package main

import (
	"archive/zip"
	"bytes"
	"flag"
	"fmt"
	"io"
	"os"
	"reflect"
	"strings"

	"github.com/gonum/matrix/mat64"
	"github.com/sbinet/npyio"
)

var (
	zipMagic = [4]byte{'P', 'K', 3, 4}
)

func main() {
	flag.Parse()

	if len(os.Args) <= 1 {
		flag.Usage()
		os.Exit(1)
	}

	allgood := true
	for i, fname := range os.Args[1:] {
		if i > 0 {
			fmt.Printf("\n")
		}
		fmt.Printf(strings.Repeat("=", 80) + "\n")
		fmt.Printf("file: %v\n", fname)

		f, err := os.Open(fname)
		if err != nil {
			fmt.Printf("error: %v\n", err)
			allgood = false
		}
		defer f.Close()

		// detect .npz files (check if we find a ZIP file magic header)
		var hdr [6]byte
		_, err = f.Read(hdr[:])
		if err != nil {
			fmt.Printf("error sniffing file format: %v\n", err)
			allgood = false
			continue
		}
		_, err = f.Seek(0, 0)
		if err != nil {
			fmt.Printf("error rewinding file: %v\n", err)
			allgood = false
			continue
		}

		switch {
		case bytes.Equal(npyio.Magic[:], hdr[:]):
			allgood = display(f, fname) && allgood
		case bytes.Equal(zipMagic[:], hdr[:len(zipMagic)]):
			fi, err := f.Stat()
			if err != nil {
				fmt.Printf("error stat-ing file: %v\n", err)
				allgood = false
				continue
			}
			zr, err := zip.NewReader(f, fi.Size())
			if err != nil {
				fmt.Printf("error creating zip-reader: %v\n", err)
				allgood = false
				continue
			}

			for ii, f := range zr.File {
				r, err := f.Open()
				if err != nil {
					fmt.Printf("error opening entry %s: %v\n", f.Name, err)
					allgood = false
					continue
				}
				defer r.Close()
				if ii > 0 {
					fmt.Printf("\n")
				}
				fmt.Printf("entry: %s\n", f.Name)
				allgood = display(r, fname+"@"+f.Name) && allgood
			}
		default:
			fmt.Printf("error: unknown magic header %q\n", string(hdr[:]))
			allgood = false
			continue
		}
	}

	if !allgood {
		os.Exit(1)
	}
}

func display(f io.Reader, fname string) bool {
	r, err := npyio.NewReader(f)
	if err != nil {
		fmt.Printf("error creating reader %s: %v\n", fname, err)
		return false
	}

	fmt.Printf("npy-header: %v\n", r.Header)

	if len(r.Header.Descr.Shape) > 2 {
		rt := npyio.TypeFrom(r.Header.Descr.Type)
		if rt == nil {
			fmt.Printf("error: no reflect.Type for %q\n", r.Header.Descr.Type)
			return false
		}
		rv := reflect.New(reflect.SliceOf(rt))
		err = r.Read(rv.Interface())
		if err != nil {
			fmt.Printf("error: %v\n", err)
			return false
		}
		fmt.Printf("data = %v\n", rv.Elem().Interface())
		return true
	}

	var m mat64.Dense
	err = r.Read(&m)
	if err != nil {
		fmt.Printf("error: %v\n", err)
		return false
	}

	fmt.Printf("data = %v\n", mat64.Formatted(&m, mat64.Prefix("       ")))
	return true
}
<|endoftext|>


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

package main

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/api/compute/v1"
	gs "google.golang.org/api/storage/v1"

	"github.com/coreos/mantle/auth"
	"github.com/coreos/mantle/platform/api/aws"
	"github.com/coreos/mantle/platform/api/azure"
	"github.com/coreos/mantle/platform/api/gcloud"
	"github.com/coreos/mantle/storage"
	"github.com/coreos/mantle/storage/index"
)

var (
	releaseDryRun bool
	cmdRelease    = &cobra.Command{
		Use:   "release [options]",
		Short: "Publish a new CoreOS release.",
		Run:   runRelease,
		Long:  `Publish a new CoreOS release.`,
	}
)

func init() {
	cmdRelease.Flags().StringVar(&awsCredentialsFile, "aws-credentials", "", "AWS credentials file")
	cmdRelease.Flags().StringVar(&azureProfile, "azure-profile", "", "Azure Profile json file")
	cmdRelease.Flags().BoolVarP(&releaseDryRun, "dry-run", "n", false,
		"perform a trial run, do not make changes")
	AddSpecFlags(cmdRelease.Flags())
	root.AddCommand(cmdRelease)
}

func runRelease(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		plog.Fatal("No args accepted")
	}

	spec := ChannelSpec()
	ctx := context.Background()
	client, err := auth.GoogleClient()
	if err != nil {
		plog.Fatalf("Authentication failed: %v", err)
	}

	src, err := storage.NewBucket(client, spec.SourceURL())
	if err != nil {
		plog.Fatal(err)
	}
	src.WriteDryRun(releaseDryRun)

	if err := src.Fetch(ctx); err != nil {
		plog.Fatal(err)
	}

	// Sanity check!
	if vertxt := src.Object(src.Prefix() + "version.txt"); vertxt == nil {
		verurl := src.URL().String() + "version.txt"
		plog.Fatalf("File not found: %s", verurl)
	}

	// Register GCE image if needed.
	doGCE(ctx, client, src, &spec)

	// Make Azure images public.
	doAzure(ctx, client, src, &spec)

	// Make AWS images public.
	doAWS(ctx, client, src, &spec)

	for _, dSpec := range spec.Destinations {
		dst, err := storage.NewBucket(client, dSpec.BaseURL)
		if err != nil {
			plog.Fatal(err)
		}
		dst.WriteDryRun(releaseDryRun)

		// Fetch parent directories non-recursively to re-index it later.
		for _, prefix := range dSpec.ParentPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, false); err != nil {
				plog.Fatal(err)
			}
		}

		// Fetch and sync each destination directory.
		for _, prefix := range dSpec.FinalPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, true); err != nil {
				plog.Fatal(err)
			}

			sync := index.NewSyncIndexJob(src, dst)
			sync.DestinationPrefix(prefix)
			sync.DirectoryHTML(dSpec.DirectoryHTML)
			sync.IndexHTML(dSpec.IndexHTML)
			sync.Delete(true)
			if dSpec.Title != "" {
				sync.Name(dSpec.Title)
			}
			if err := sync.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}

		// Now refresh the parent directory indexes.
		for _, prefix := range dSpec.ParentPrefixes() {
			parent := index.NewIndexJob(dst)
			parent.Prefix(prefix)
			parent.DirectoryHTML(dSpec.DirectoryHTML)
			parent.IndexHTML(dSpec.IndexHTML)
			parent.Recursive(false)
			parent.Delete(true)
			if dSpec.Title != "" {
				parent.Name(dSpec.Title)
			}
			if err := parent.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}
	}
}

func sanitizeVersion() string {
	v := strings.Replace(specVersion, ".", "-", -1)
	return strings.Replace(v, "+", "-", -1)
}

func gceUploadImage(spec *channelSpec, api *gcloud.API, obj *gs.Object, name, desc string) string {
	plog.Noticef("Creating GCE image %s", name)
	parsedVersion, err := semver.NewVersion(specVersion)
	if err != nil {
		plog.Fatalf("couldn't parse version %s: %v", specVersion, err)
	}
	disableMultiqueue := false
	if parsedVersion.LessThan(semver.Version{Major: 1409}) {
		disableMultiqueue = true
		plog.Noticef("Not enabling multiqueue for version %v", specVersion)
	}
	op, pending, err := api.CreateImage(&gcloud.ImageSpec{
		SourceImage:           obj.MediaLink,
		Family:                spec.GCE.Family,
		Name:                  name,
		Description:           desc,
		Licenses:              spec.GCE.Licenses,
		DisableSCSIMultiqueue: disableMultiqueue,
	}, false)
	if err != nil {
		plog.Fatalf("GCE image creation failed: %v", err)
	}

	plog.Infof("Waiting for image creation to finish...")
	pending.Interval = 3 * time.Second
	pending.Progress = func(_ string, _ time.Duration, op *compute.Operation) error {
		status := strings.ToLower(op.Status)
		if op.Progress != 0 {
			plog.Infof("Image creation is %s: %s % 2d%%", status, op.StatusMessage, op.Progress)
		} else {
			plog.Infof("Image creation is %s. %s", status, op.StatusMessage)
		}
		return nil
	}
	if err := pending.Wait(); err != nil {
		plog.Fatal(err)
	}
	plog.Info("Success!")
	return op.TargetLink
}

func doGCE(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.GCE.Project == "" || spec.GCE.Image == "" {
		plog.Notice("GCE image creation disabled.")
		return
	}

	api, err := gcloud.New(&gcloud.Options{
		Project: spec.GCE.Project,
	})
	if err != nil {
		plog.Fatalf("GCE client failed: %v", err)
	}

	publishImage := func(image string) {
		if spec.GCE.Publish == "" {
			plog.Notice("GCE image name publishing disabled.")
			return
		}
		obj := gs.Object{
			Name:        src.Prefix() + spec.GCE.Publish,
			ContentType: "text/plain",
		}
		media := strings.NewReader(
			fmt.Sprintf("projects/%s/global/images/%s\n",
				spec.GCE.Project, image))
		if err := src.Upload(ctx, &obj, media); err != nil {
			plog.Fatal(err)
		}
	}

	nameVer := fmt.Sprintf("%s-%s-v", spec.GCE.Family, sanitizeVersion())
	date := time.Now().UTC()
	name := nameVer + date.Format("20060102")
	desc := fmt.Sprintf("%s, %s, %s published on %s", spec.GCE.Description,
		specVersion, specBoard, date.Format("2006-01-02"))

	images, err := api.ListImages(ctx, spec.GCE.Family+"-")
	if err != nil {
		plog.Fatal(err)
	}

	var conflicting []string
	for _, image := range images {
		if strings.HasPrefix(image.Name, nameVer) {
			conflicting = append(conflicting, image.Name)
		}
	}

	// Check for any with the same version but possibly different dates.
	if len(conflicting) > 1 {
		plog.Fatalf("Duplicate GCE images found: %v", conflicting)
	} else if len(conflicting) == 1 {
		plog.Noticef("GCE image already exists: %s", conflicting[0])
		publishImage(conflicting[0])
		return
	}

	obj := src.Object(src.Prefix() + spec.GCE.Image)
	if obj == nil {
		plog.Fatalf("GCE image not found %s%s", src.URL(), spec.GCE.Image)
	}

	if releaseDryRun {
		plog.Noticef("Would create GCE image %s", name)
		return
	}

	imageLink := gceUploadImage(spec, api, obj, name, desc)
	publishImage(name)

	var pendings []*gcloud.Pending
	for _, old := range images {
		if old.Deprecated != nil && old.Deprecated.State != "" {
			continue
		}
		plog.Noticef("Deprecating old image %s", old.Name)
		pending, err := api.DeprecateImage(old.Name, gcloud.DeprecationStateDeprecated, imageLink)
		if err != nil {
			plog.Fatal(err)
		}
		pending.Interval = 1 * time.Second
		pending.Timeout = 0
		pendings = append(pendings, pending)
	}

	if spec.GCE.Limit > 0 && len(images) > spec.GCE.Limit {
		plog.Noticef("Pruning %d GCE images.", len(images)-spec.GCE.Limit)
		plog.Notice("NOPE! JUST KIDDING, TODO")
	}

	plog.Infof("Waiting on %d operations.", len(pendings))
	for _, pending := range pendings {
		err := pending.Wait()
		if err != nil {
			plog.Fatal(err)
		}
	}
}

func doAzure(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.Azure.StorageAccount == "" {
		plog.Notice("Azure image creation disabled.")
		return
	}

	prof, err := auth.ReadAzureProfile(azureProfile)
	if err != nil {
		plog.Fatalf("failed reading Azure profile: %v", err)
	}

	// channel name should be caps for azure image
	imageName := fmt.Sprintf("%s-%s-%s", spec.Azure.Offer, strings.Title(specChannel), specVersion)

	for _, environment := range spec.Azure.Environments {
		opt := prof.SubscriptionOptions(environment.SubscriptionName)
		if opt == nil {
			plog.Fatalf("couldn't find subscription %q", environment.SubscriptionName)
		}

		api, err := azure.New(opt)
		if err != nil {
			plog.Fatalf("failed to create Azure API: %v", err)
		}

		if releaseDryRun {
			// TODO(bgilbert): check that the image exists
			plog.Printf("Would share %q on %v", imageName, environment.SubscriptionName)
			continue
		} else {
			plog.Printf("Sharing %q on %v...", imageName, environment.SubscriptionName)
		}

		if err := api.ShareImage(imageName, "public"); err != nil {
			plog.Fatalf("failed to share image %q: %v", imageName, err)
		}
	}
}

func doAWS(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.AWS.Image == "" {
		plog.Notice("AWS image creation disabled.")
		return
	}

	imageName := fmt.Sprintf("%v-%v-%v", spec.AWS.BaseName, specChannel, specVersion)
	imageName = regexp.MustCompile(`[^A-Za-z0-9()\\./_-]`).ReplaceAllLiteralString(imageName, "_")

	for _, part := range spec.AWS.Partitions {
		for _, region := range part.Regions {
			if releaseDryRun {
				plog.Printf("Checking for images in %v %v...", part.Name, region)
			} else {
				plog.Printf("Publishing images in %v %v...", part.Name, region)
			}

			api, err := aws.New(&aws.Options{
				CredentialsFile: awsCredentialsFile,
				Profile:         part.Profile,
				Region:          region,
			})
			if err != nil {
				plog.Fatalf("creating client for %v %v: %v", part.Name, region, err)
			}

			publish := func(imageName string) {
				imageID, err := api.FindImage(imageName)
				if err != nil {
					plog.Fatalf("couldn't find image %q in %v %v: %v", imageName, part.Name, region, err)
				}

				if !releaseDryRun {
					err := api.PublishImage(imageID)
					if err != nil {
						plog.Fatalf("couldn't publish image in %v %v: %v", part.Name, region, err)
					}
				}
			}
			publish(imageName)
			publish(imageName + "-hvm")
		}
	}
}
<commit_msg>cmd/plume: Drop early return if GCE image already exists<commit_after>// Copyright 2016 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/api/compute/v1"
	gs "google.golang.org/api/storage/v1"

	"github.com/coreos/mantle/auth"
	"github.com/coreos/mantle/platform/api/aws"
	"github.com/coreos/mantle/platform/api/azure"
	"github.com/coreos/mantle/platform/api/gcloud"
	"github.com/coreos/mantle/storage"
	"github.com/coreos/mantle/storage/index"
)

var (
	releaseDryRun bool
	cmdRelease    = &cobra.Command{
		Use:   "release [options]",
		Short: "Publish a new CoreOS release.",
		Run:   runRelease,
		Long:  `Publish a new CoreOS release.`,
	}
)

func init() {
	cmdRelease.Flags().StringVar(&awsCredentialsFile, "aws-credentials", "", "AWS credentials file")
	cmdRelease.Flags().StringVar(&azureProfile, "azure-profile", "", "Azure Profile json file")
	cmdRelease.Flags().BoolVarP(&releaseDryRun, "dry-run", "n", false,
		"perform a trial run, do not make changes")
	AddSpecFlags(cmdRelease.Flags())
	root.AddCommand(cmdRelease)
}

func runRelease(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		plog.Fatal("No args accepted")
	}

	spec := ChannelSpec()
	ctx := context.Background()
	client, err := auth.GoogleClient()
	if err != nil {
		plog.Fatalf("Authentication failed: %v", err)
	}

	src, err := storage.NewBucket(client, spec.SourceURL())
	if err != nil {
		plog.Fatal(err)
	}
	src.WriteDryRun(releaseDryRun)

	if err := src.Fetch(ctx); err != nil {
		plog.Fatal(err)
	}

	// Sanity check!
	if vertxt := src.Object(src.Prefix() + "version.txt"); vertxt == nil {
		verurl := src.URL().String() + "version.txt"
		plog.Fatalf("File not found: %s", verurl)
	}

	// Register GCE image if needed.
	doGCE(ctx, client, src, &spec)

	// Make Azure images public.
	doAzure(ctx, client, src, &spec)

	// Make AWS images public.
	doAWS(ctx, client, src, &spec)

	for _, dSpec := range spec.Destinations {
		dst, err := storage.NewBucket(client, dSpec.BaseURL)
		if err != nil {
			plog.Fatal(err)
		}
		dst.WriteDryRun(releaseDryRun)

		// Fetch parent directories non-recursively to re-index it later.
		for _, prefix := range dSpec.ParentPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, false); err != nil {
				plog.Fatal(err)
			}
		}

		// Fetch and sync each destination directory.
		for _, prefix := range dSpec.FinalPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, true); err != nil {
				plog.Fatal(err)
			}

			sync := index.NewSyncIndexJob(src, dst)
			sync.DestinationPrefix(prefix)
			sync.DirectoryHTML(dSpec.DirectoryHTML)
			sync.IndexHTML(dSpec.IndexHTML)
			sync.Delete(true)
			if dSpec.Title != "" {
				sync.Name(dSpec.Title)
			}
			if err := sync.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}

		// Now refresh the parent directory indexes.
		for _, prefix := range dSpec.ParentPrefixes() {
			parent := index.NewIndexJob(dst)
			parent.Prefix(prefix)
			parent.DirectoryHTML(dSpec.DirectoryHTML)
			parent.IndexHTML(dSpec.IndexHTML)
			parent.Recursive(false)
			parent.Delete(true)
			if dSpec.Title != "" {
				parent.Name(dSpec.Title)
			}
			if err := parent.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}
	}
}

func sanitizeVersion() string {
	v := strings.Replace(specVersion, ".", "-", -1)
	return strings.Replace(v, "+", "-", -1)
}

func gceUploadImage(spec *channelSpec, api *gcloud.API, obj *gs.Object, name, desc string) string {
	plog.Noticef("Creating GCE image %s", name)
	parsedVersion, err := semver.NewVersion(specVersion)
	if err != nil {
		plog.Fatalf("couldn't parse version %s: %v", specVersion, err)
	}
	disableMultiqueue := false
	if parsedVersion.LessThan(semver.Version{Major: 1409}) {
		disableMultiqueue = true
		plog.Noticef("Not enabling multiqueue for version %v", specVersion)
	}
	op, pending, err := api.CreateImage(&gcloud.ImageSpec{
		SourceImage:           obj.MediaLink,
		Family:                spec.GCE.Family,
		Name:                  name,
		Description:           desc,
		Licenses:              spec.GCE.Licenses,
		DisableSCSIMultiqueue: disableMultiqueue,
	}, false)
	if err != nil {
		plog.Fatalf("GCE image creation failed: %v", err)
	}

	plog.Infof("Waiting for image creation to finish...")
	pending.Interval = 3 * time.Second
	pending.Progress = func(_ string, _ time.Duration, op *compute.Operation) error {
		status := strings.ToLower(op.Status)
		if op.Progress != 0 {
			plog.Infof("Image creation is %s: %s % 2d%%", status, op.StatusMessage, op.Progress)
		} else {
			plog.Infof("Image creation is %s. %s", status, op.StatusMessage)
		}
		return nil
	}
	if err := pending.Wait(); err != nil {
		plog.Fatal(err)
	}
	plog.Info("Success!")
	return op.TargetLink
}

func doGCE(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.GCE.Project == "" || spec.GCE.Image == "" {
		plog.Notice("GCE image creation disabled.")
		return
	}

	api, err := gcloud.New(&gcloud.Options{
		Project: spec.GCE.Project,
	})
	if err != nil {
		plog.Fatalf("GCE client failed: %v", err)
	}

	nameVer := fmt.Sprintf("%s-%s-v", spec.GCE.Family, sanitizeVersion())
	date := time.Now().UTC()
	name := nameVer + date.Format("20060102")
	desc := fmt.Sprintf("%s, %s, %s published on %s", spec.GCE.Description,
		specVersion, specBoard, date.Format("2006-01-02"))

	images, err := api.ListImages(ctx, spec.GCE.Family+"-")
	if err != nil {
		plog.Fatal(err)
	}

	var conflicting []*compute.Image
	for _, image := range images {
		if strings.HasPrefix(image.Name, nameVer) {
			conflicting = append(conflicting, image)
		}
	}

	// Check for any with the same version but possibly different dates.
	var imageLink string
	if len(conflicting) > 1 {
		plog.Fatalf("Duplicate GCE images found: %v", conflicting)
	} else if len(conflicting) == 1 {
		image := conflicting[0]
		name = image.Name
		imageLink = image.SelfLink
		plog.Noticef("GCE image already exists: %s", name)

		if releaseDryRun {
			return
		}
	} else {
		obj := src.Object(src.Prefix() + spec.GCE.Image)
		if obj == nil {
			plog.Fatalf("GCE image not found %s%s", src.URL(), spec.GCE.Image)
		}

		if releaseDryRun {
			plog.Noticef("Would create GCE image %s", name)
			return
		}

		imageLink = gceUploadImage(spec, api, obj, name, desc)
	}

	if spec.GCE.Publish != "" {
		obj := gs.Object{
			Name:        src.Prefix() + spec.GCE.Publish,
			ContentType: "text/plain",
		}
		media := strings.NewReader(
			fmt.Sprintf("projects/%s/global/images/%s\n",
				spec.GCE.Project, name))
		if err := src.Upload(ctx, &obj, media); err != nil {
			plog.Fatal(err)
		}
	} else {
		plog.Notice("GCE image name publishing disabled.")
	}

	var pendings []*gcloud.Pending
	for _, old := range images {
		if old.Deprecated != nil && old.Deprecated.State != "" {
			continue
		}
		if old.Name == name {
			// The current image, uploaded in a previous run
			continue
		}
		plog.Noticef("Deprecating old image %s", old.Name)
		pending, err := api.DeprecateImage(old.Name, gcloud.DeprecationStateDeprecated, imageLink)
		if err != nil {
			plog.Fatal(err)
		}
		pending.Interval = 1 * time.Second
		pending.Timeout = 0
		pendings = append(pendings, pending)
	}

	if spec.GCE.Limit > 0 && len(images) > spec.GCE.Limit {
		plog.Noticef("Pruning %d GCE images.", len(images)-spec.GCE.Limit)
		plog.Notice("NOPE! JUST KIDDING, TODO")
	}

	plog.Infof("Waiting on %d operations.", len(pendings))
	for _, pending := range pendings {
		err := pending.Wait()
		if err != nil {
			plog.Fatal(err)
		}
	}
}

func doAzure(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.Azure.StorageAccount == "" {
		plog.Notice("Azure image creation disabled.")
		return
	}

	prof, err := auth.ReadAzureProfile(azureProfile)
	if err != nil {
		plog.Fatalf("failed reading Azure profile: %v", err)
	}

	// channel name should be caps for azure image
	imageName := fmt.Sprintf("%s-%s-%s", spec.Azure.Offer, strings.Title(specChannel), specVersion)

	for _, environment := range spec.Azure.Environments {
		opt := prof.SubscriptionOptions(environment.SubscriptionName)
		if opt == nil {
			plog.Fatalf("couldn't find subscription %q", environment.SubscriptionName)
		}

		api, err := azure.New(opt)
		if err != nil {
			plog.Fatalf("failed to create Azure API: %v", err)
		}

		if releaseDryRun {
			// TODO(bgilbert): check that the image exists
			plog.Printf("Would share %q on %v", imageName, environment.SubscriptionName)
			continue
		} else {
			plog.Printf("Sharing %q on %v...", imageName, environment.SubscriptionName)
		}

		if err := api.ShareImage(imageName, "public"); err != nil {
			plog.Fatalf("failed to share image %q: %v", imageName, err)
		}
	}
}

func doAWS(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.AWS.Image == "" {
		plog.Notice("AWS image creation disabled.")
		return
	}

	imageName := fmt.Sprintf("%v-%v-%v", spec.AWS.BaseName, specChannel, specVersion)
	imageName = regexp.MustCompile(`[^A-Za-z0-9()\\./_-]`).ReplaceAllLiteralString(imageName, "_")

	for _, part := range spec.AWS.Partitions {
		for _, region := range part.Regions {
			if releaseDryRun {
				plog.Printf("Checking for images in %v %v...", part.Name, region)
			} else {
				plog.Printf("Publishing images in %v %v...", part.Name, region)
			}

			api, err := aws.New(&aws.Options{
				CredentialsFile: awsCredentialsFile,
				Profile:         part.Profile,
				Region:          region,
			})
			if err != nil {
				plog.Fatalf("creating client for %v %v: %v", part.Name, region, err)
			}

			publish := func(imageName string) {
				imageID, err := api.FindImage(imageName)
				if err != nil {
					plog.Fatalf("couldn't find image %q in %v %v: %v", imageName, part.Name, region, err)
				}

				if !releaseDryRun {
					err := api.PublishImage(imageID)
					if err != nil {
						plog.Fatalf("couldn't publish image in %v %v: %v", part.Name, region, err)
					}
				}
			}
			publish(imageName)
			publish(imageName + "-hvm")
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2016 The Syncthing Authors.
//
// 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 (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"math/big"
	mr "math/rand"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/syncthing/syncthing/lib/protocol"
)

type result struct {
	id       protocol.DeviceID
	priv     *ecdsa.PrivateKey
	derBytes []byte
}

func main() {
	flag.Parse()
	prefix := strings.ToUpper(strings.Replace(flag.Arg(0), "-", "", -1))
	if len(prefix) > 7 {
		prefix = prefix[:7] + "-" + prefix[7:]
	}

	found := make(chan result)
	stop := make(chan struct{})
	var count int64

	// Print periodic progress reports.
	go printProgress(prefix, &count)

	// Run one certificate generator per CPU core.
	var wg sync.WaitGroup
	for i := 0; i < runtime.GOMAXPROCS(-1); i++ {
		wg.Add(1)
		go func() {
			generatePrefixed(prefix, &count, found, stop)
			wg.Done()
		}()
	}

	// Save the result, when one has been found.
	res := <-found
	close(stop)
	wg.Wait()

	fmt.Println("Found", res.id)
	saveCert(res.priv, res.derBytes)
	fmt.Println("Saved to cert.pem, key.pem")
}

// Try certificates until one is found that has the prefix at the start of
// the resulting device ID. Increments count atomically, sends the result to
// found, returns when stop is closed.
func generatePrefixed(prefix string, count *int64, found chan<- result, stop <-chan struct{}) {
	notBefore := time.Now()
	notAfter := time.Date(2049, 12, 31, 23, 59, 59, 0, time.UTC)

	template := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(mr.Int63()),
		Subject: pkix.Name{
			CommonName: "syncthing",
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	priv, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	for {
		select {
		case <-stop:
			return
		default:
		}

		derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, priv.Public(), priv)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		id := protocol.NewDeviceID(derBytes)
		atomic.AddInt64(count, 1)

		if strings.HasPrefix(id.String(), prefix) {
			select {
			case found <- result{id, priv, derBytes}:
			case <-stop:
			}
			return
		}
	}
}

func printProgress(prefix string, count *int64) {
	started := time.Now()
	wantBits := 5 * len(prefix)
	if wantBits > 63 {
		fmt.Printf("Want %d bits for prefix %q, refusing to boil the ocean.\n", wantBits, prefix)
		os.Exit(1)
	}
	expectedIterations := float64(int(1) << uint(wantBits))
	fmt.Printf("Want %d bits for prefix %q, about %.2g certs to test (statistically speaking)\n", wantBits, prefix, expectedIterations)

	for _ = range time.NewTicker(15 * time.Second).C {
		tried := atomic.LoadInt64(count)
		elapsed := time.Since(started)
		rate := float64(tried) / elapsed.Seconds()
		expected := timeStr(expectedIterations / rate)
		fmt.Printf("Trying %.0f certs/s, tested %d so far in %v, expect ~%s total time to complete\n", rate, tried, elapsed/time.Second*time.Second, expected)
	}
}

func saveCert(priv interface{}, derBytes []byte) {
	certOut, err := os.Create("cert.pem")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	err = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	err = certOut.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	keyOut, err := os.OpenFile("key.pem", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	block, err := pemBlockForKey(priv)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	err = pem.Encode(keyOut, block)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	err = keyOut.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func pemBlockForKey(priv interface{}) (*pem.Block, error) {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k)}, nil
	case *ecdsa.PrivateKey:
		b, err := x509.MarshalECPrivateKey(k)
		if err != nil {
			return nil, err
		}
		return &pem.Block{Type: "EC PRIVATE KEY", Bytes: b}, nil
	default:
		return nil, fmt.Errorf("unknown key type")
	}
}

func timeStr(seconds float64) string {
	if seconds < 60 {
		return fmt.Sprintf("%.0fs", seconds)
	}
	if seconds < 3600 {
		return fmt.Sprintf("%.0fm", seconds/60)
	}
	if seconds < 86400 {
		return fmt.Sprintf("%.0fh", seconds/3600)
	}
	if seconds < 86400*365 {
		return fmt.Sprintf("%.0f days", seconds/3600)
	}
	return fmt.Sprintf("%.0f years", seconds/86400/365)
}
<commit_msg>cmd/stvanity: Use Go 1.3 compatible interface<commit_after>// Copyright (C) 2016 The Syncthing Authors.
//
// 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 (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"math/big"
	mr "math/rand"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/syncthing/syncthing/lib/protocol"
)

type result struct {
	id       protocol.DeviceID
	priv     *ecdsa.PrivateKey
	derBytes []byte
}

func main() {
	flag.Parse()
	prefix := strings.ToUpper(strings.Replace(flag.Arg(0), "-", "", -1))
	if len(prefix) > 7 {
		prefix = prefix[:7] + "-" + prefix[7:]
	}

	found := make(chan result)
	stop := make(chan struct{})
	var count int64

	// Print periodic progress reports.
	go printProgress(prefix, &count)

	// Run one certificate generator per CPU core.
	var wg sync.WaitGroup
	for i := 0; i < runtime.GOMAXPROCS(-1); i++ {
		wg.Add(1)
		go func() {
			generatePrefixed(prefix, &count, found, stop)
			wg.Done()
		}()
	}

	// Save the result, when one has been found.
	res := <-found
	close(stop)
	wg.Wait()

	fmt.Println("Found", res.id)
	saveCert(res.priv, res.derBytes)
	fmt.Println("Saved to cert.pem, key.pem")
}

// Try certificates until one is found that has the prefix at the start of
// the resulting device ID. Increments count atomically, sends the result to
// found, returns when stop is closed.
func generatePrefixed(prefix string, count *int64, found chan<- result, stop <-chan struct{}) {
	notBefore := time.Now()
	notAfter := time.Date(2049, 12, 31, 23, 59, 59, 0, time.UTC)

	template := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(mr.Int63()),
		Subject: pkix.Name{
			CommonName: "syncthing",
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	priv, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	for {
		select {
		case <-stop:
			return
		default:
		}

		derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, priv.PublicKey, priv)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		id := protocol.NewDeviceID(derBytes)
		atomic.AddInt64(count, 1)

		if strings.HasPrefix(id.String(), prefix) {
			select {
			case found <- result{id, priv, derBytes}:
			case <-stop:
			}
			return
		}
	}
}

func printProgress(prefix string, count *int64) {
	started := time.Now()
	wantBits := 5 * len(prefix)
	if wantBits > 63 {
		fmt.Printf("Want %d bits for prefix %q, refusing to boil the ocean.\n", wantBits, prefix)
		os.Exit(1)
	}
	expectedIterations := float64(int(1) << uint(wantBits))
	fmt.Printf("Want %d bits for prefix %q, about %.2g certs to test (statistically speaking)\n", wantBits, prefix, expectedIterations)

	for _ = range time.NewTicker(15 * time.Second).C {
		tried := atomic.LoadInt64(count)
		elapsed := time.Since(started)
		rate := float64(tried) / elapsed.Seconds()
		expected := timeStr(expectedIterations / rate)
		fmt.Printf("Trying %.0f certs/s, tested %d so far in %v, expect ~%s total time to complete\n", rate, tried, elapsed/time.Second*time.Second, expected)
	}
}

func saveCert(priv interface{}, derBytes []byte) {
	certOut, err := os.Create("cert.pem")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	err = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	err = certOut.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	keyOut, err := os.OpenFile("key.pem", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	block, err := pemBlockForKey(priv)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	err = pem.Encode(keyOut, block)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	err = keyOut.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func pemBlockForKey(priv interface{}) (*pem.Block, error) {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k)}, nil
	case *ecdsa.PrivateKey:
		b, err := x509.MarshalECPrivateKey(k)
		if err != nil {
			return nil, err
		}
		return &pem.Block{Type: "EC PRIVATE KEY", Bytes: b}, nil
	default:
		return nil, fmt.Errorf("unknown key type")
	}
}

func timeStr(seconds float64) string {
	if seconds < 60 {
		return fmt.Sprintf("%.0fs", seconds)
	}
	if seconds < 3600 {
		return fmt.Sprintf("%.0fm", seconds/60)
	}
	if seconds < 86400 {
		return fmt.Sprintf("%.0fh", seconds/3600)
	}
	if seconds < 86400*365 {
		return fmt.Sprintf("%.0f days", seconds/3600)
	}
	return fmt.Sprintf("%.0f years", seconds/86400/365)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"io"
	"log"
	"os"

	"github.com/spf13/cobra"
)

var (
	// calllogCmd represents the calllog command
	calllogCmd = &cobra.Command{
		Use:   "logs <test-run-ref>",
		Short: "Fetch up to 10k lines of the test runs call log",
		Long: `Will fetch up to 10k lines of the test runs call log.

The call log contains:
  * time (epoch in seconds)
  * HTTP verb
  * HTTP host
  * request path
  * HTTP Status Code
  * response size (in Bytes)
  * duration (in ms)
  * request tag`,
		Run: run,
	}

	logOpts struct {
		Type    string
		Preview bool
	}
)

func init() {
	TestRunCmd.AddCommand(calllogCmd)

	calllogCmd.Flags().StringVar(&logOpts.Type, "type", "request", "type of logs")
	calllogCmd.Flags().BoolVarP(&logOpts.Preview, "preview", "p", true, "Preview of logs")
}

func run(cmd *cobra.Command, args []string) {
	if len(args) == 0 {
		log.Fatal("Missing argument: Test Run Reference")
	}

	if len(args) > 2 {
		log.Fatal("Too many arguments")
	}

	if logOpts.Type != "request" {
		log.Fatal(fmt.Sprintf("Unsupported log type %s", logOpts.Type))
	}

	client := NewClient()

	reader, err := client.TestRunCallLog(args[0], logOpts.Preview)
	if err != nil {
		log.Fatal(err)
	}

	if len(args) == 1 {
		io.Copy(os.Stdout, reader)
	} else {
		file, err := os.Create(args[1])
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()
		defer reader.Close()

		io.Copy(file, reader)
	}
}
<commit_msg>make switch between full and preview of logs easier<commit_after>package cmd

import (
	"fmt"
	"io"
	"log"
	"os"

	"github.com/spf13/cobra"
)

var (
	// calllogCmd represents the calllog command
	calllogCmd = &cobra.Command{
		Use:   "logs <test-run-ref>",
		Short: "Fetch the the test runs call log (request log)",
		Long: `Will fetch the the test runs call log (request log).

By default, you will get the first 10k lines. Using --full you
will download the entire request log.

The call log contains:
  * time (epoch in seconds)
  * HTTP verb
  * HTTP host
  * request path
  * HTTP Status Code
  * response size (in Bytes)
  * duration (in ms)
  * request tag`,
		Run: run,
	}

	logOpts struct {
		Type       string
		Full       bool
		OutputFile string
	}
)

func init() {
	TestRunCmd.AddCommand(calllogCmd)

	calllogCmd.Flags().StringVar(&logOpts.Type, "type", "request", "type of logs")
	calllogCmd.Flags().BoolVarP(&logOpts.Full, "full", "f", false, "download full logs")
	calllogCmd.Flags().StringVar(&logOpts.OutputFile, "output", "-", "save logs to file or '-' for stdout")
}

func run(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		log.Fatal("Expecting exactly one argument: Test Run Reference")
	}

	if logOpts.Type != "request" {
		log.Fatal(fmt.Sprintf("Unsupported log type %s", logOpts.Type))
	}

	client := NewClient()

	reader, err := client.TestRunCallLog(args[0], !logOpts.Full)
	if err != nil {
		log.Fatal(err)
	}

	if logOpts.OutputFile == "-" {
		io.Copy(os.Stdout, reader)
	} else {
		file, err := os.Create(logOpts.OutputFile)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()
		defer reader.Close()

		io.Copy(file, reader)
	}
}
<|endoftext|>