<commit_before>package main

import (
	"flag"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/cloudfoundry-incubator/garden/server"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/container_pool"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/network_pool"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/port_pool"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/quota_manager"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/uid_pool"
	"github.com/cloudfoundry/gunk/command_runner"
)

var listenNetwork = flag.String(
	"listenNetwork",
	"unix",
	"how to listen on the address (unix, tcp, etc.)",
)

var listenAddr = flag.String(
	"listenAddr",
	"/tmp/warden.sock",
	"address to listen on",
)

var snapshotsPath = flag.String(
	"snapshots",
	"",
	"directory in which to store container state to persist through restarts",
)

var backendName = flag.String(
	"backend",
	"linux",
	"which backend to use (linux or fake)",
)

var binPath = flag.String(
	"bin",
	"",
	"directory containing backend-specific scripts (i.e. ./create.sh)",
)

var depotPath = flag.String(
	"depot",
	"",
	"directory in which to store containers",
)

var rootFSPath = flag.String(
	"rootfs",
	"",
	"directory of the rootfs for the containers",
)

var disableQuotas = flag.Bool(
	"disableQuotas",
	false,
	"disable disk quotas",
)

var containerGraceTime = flag.Duration(
	"containerGraceTime",
	0,
	"time after which to destroy idle containers",
)

var debug = flag.Bool(
	"debug",
	false,
	"show low-level command output",
)

func main() {
	flag.Parse()

	maxProcs := runtime.NumCPU()
	prevMaxProcs := runtime.GOMAXPROCS(maxProcs)

	log.Println("set GOMAXPROCS to", maxProcs, "was", prevMaxProcs)

	if *binPath == "" {
		log.Fatalln("must specify -bin with linux backend")
	}

	if *depotPath == "" {
		log.Fatalln("must specify -depot with linux backend")
	}

	if *rootFSPath == "" {
		log.Fatalln("must specify -rootfs with linux backend")
	}

	uidPool := uid_pool.New(10000, 256)

	_, ipNet, err := net.ParseCIDR("10.254.0.0/22")
	if err != nil {
		log.Fatalln("error parsing CIDR:", err)
	}

	networkPool := network_pool.New(ipNet)

	// TODO: base on ephemeral port range
	portPool := port_pool.New(61000, 6501)

	var runner command_runner.CommandRunner

	runner = command_runner.New(*debug)

	quotaManager, err := quota_manager.New(*depotPath, *binPath, runner)
	if err != nil {
		log.Fatalln("error creating quota manager:", err)
	}

	if *disableQuotas {
		quotaManager.Disable()
	}

	pool := container_pool.New(
		*binPath,
		*depotPath,
		*rootFSPath,
		uidPool,
		networkPool,
		portPool,
		runner,
		quotaManager,
	)

	backend := linux_backend.New(pool, *snapshotsPath)

	log.Println("setting up backend")

	err = backend.Setup()
	if err != nil {
		log.Fatalln("failed to set up backend:", err)
	}

	log.Println("starting server; listening with", *listenNetwork, "on", *listenAddr)

	graceTime := *containerGraceTime

	wardenServer := server.New(*listenNetwork, *listenAddr, graceTime, backend)

	err = wardenServer.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
	}

	signals := make(chan os.Signal, 1)

	go func() {
		<-signals
		log.Println("stopping...")
		wardenServer.Stop()
		os.Exit(0)
	}()

	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	select {}
}
<commit_msg>use linux_command_runner<commit_after>package main

import (
	"flag"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/cloudfoundry-incubator/garden/server"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/container_pool"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/network_pool"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/port_pool"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/quota_manager"
	"github.com/cloudfoundry-incubator/warden-linux/linux_backend/uid_pool"
	"github.com/cloudfoundry/gunk/command_runner/linux_command_runner"
)

var listenNetwork = flag.String(
	"listenNetwork",
	"unix",
	"how to listen on the address (unix, tcp, etc.)",
)

var listenAddr = flag.String(
	"listenAddr",
	"/tmp/warden.sock",
	"address to listen on",
)

var snapshotsPath = flag.String(
	"snapshots",
	"",
	"directory in which to store container state to persist through restarts",
)

var binPath = flag.String(
	"bin",
	"",
	"directory containing backend-specific scripts (i.e. ./create.sh)",
)

var depotPath = flag.String(
	"depot",
	"",
	"directory in which to store containers",
)

var rootFSPath = flag.String(
	"rootfs",
	"",
	"directory of the rootfs for the containers",
)

var disableQuotas = flag.Bool(
	"disableQuotas",
	false,
	"disable disk quotas",
)

var containerGraceTime = flag.Duration(
	"containerGraceTime",
	0,
	"time after which to destroy idle containers",
)

var debug = flag.Bool(
	"debug",
	false,
	"show low-level command output",
)

func main() {
	flag.Parse()

	maxProcs := runtime.NumCPU()
	prevMaxProcs := runtime.GOMAXPROCS(maxProcs)

	log.Println("set GOMAXPROCS to", maxProcs, "was", prevMaxProcs)

	if *binPath == "" {
		log.Fatalln("must specify -bin with linux backend")
	}

	if *depotPath == "" {
		log.Fatalln("must specify -depot with linux backend")
	}

	if *rootFSPath == "" {
		log.Fatalln("must specify -rootfs with linux backend")
	}

	uidPool := uid_pool.New(10000, 256)

	_, ipNet, err := net.ParseCIDR("10.254.0.0/22")
	if err != nil {
		log.Fatalln("error parsing CIDR:", err)
	}

	networkPool := network_pool.New(ipNet)

	// TODO: base on ephemeral port range
	portPool := port_pool.New(61000, 6501)

	runner := linux_command_runner.New(*debug)

	quotaManager, err := quota_manager.New(*depotPath, *binPath, runner)
	if err != nil {
		log.Fatalln("error creating quota manager:", err)
	}

	if *disableQuotas {
		quotaManager.Disable()
	}

	pool := container_pool.New(
		*binPath,
		*depotPath,
		*rootFSPath,
		uidPool,
		networkPool,
		portPool,
		runner,
		quotaManager,
	)

	backend := linux_backend.New(pool, *snapshotsPath)

	log.Println("setting up backend")

	err = backend.Setup()
	if err != nil {
		log.Fatalln("failed to set up backend:", err)
	}

	log.Println("starting server; listening with", *listenNetwork, "on", *listenAddr)

	graceTime := *containerGraceTime

	wardenServer := server.New(*listenNetwork, *listenAddr, graceTime, backend)

	err = wardenServer.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
	}

	signals := make(chan os.Signal, 1)

	go func() {
		<-signals
		log.Println("stopping...")
		wardenServer.Stop()
		os.Exit(0)
	}()

	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	select {}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"github.com/codegangsta/cli"
	dupes "github.com/danmarg/undupes/libdupes"
	"github.com/dustin/go-humanize"
	"github.com/siddontang/go-log/log"
	"os"
	"regexp"
	"strconv"
	"strings"
)

func main() {
	app := cli.NewApp()
	app.Name = "undupes"
	app.Usage = "manage duplicate files"
	app.Author = "Daniel Margolis"
	app.Email = "dan@af0.net"
	app.Version = "0.1"
	app.Commands = []cli.Command{
		{
			Name:      "interactive",
			ShortName: "i",
			Usage:     "interactive mode",
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "v",
					Usage: "log level",
				},
			},
			Action: func(c *cli.Context) {
				log.SetLevel(c.Int("v"))
				if err := runInteractive(c.Bool("dry_run")); err != nil {
					fmt.Println(err)
				}
			},
		},
		{
			Name:      "print",
			ShortName: "p",
			Usage:     "print duplicates",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "directory, d",
					Usage: "directory in which to find duplicates (required)",
				},
				cli.StringFlag{
					Name:  "output, o",
					Usage: "output file",
				},
			},
			Action: func(c *cli.Context) {
				if err := runPrint(c.String("directory"), c.String("output")); err != nil {
					fmt.Println(err)
				}
			},
		},
		{
			Name:      "auto",
			ShortName: "a",
			Usage:     "automatically resolve duplicates",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "directory, d",
					Usage: "directory in which to find duplicates (required)",
				},
				cli.StringFlag{
					Name:  "prefer, p",
					Usage: "prefer to keep files matching this pattern",
				},
				cli.StringFlag{
					Name:  "over, o",
					Usage: "used with --prefer; will restrict preferred files to those where the duplicate matches --over",
				},
				cli.BoolFlag{
					Name:  "invert, i",
					Usage: "invert matching logic; preferred files with be prioritized for deletion rather than retention",
				},
				cli.BoolFlag{
					Name:  "dry_run",
					Usage: "simulate (log but don't delete files)",
				},
				cli.IntFlag{
					Name:  "v",
					Usage: "log level",
				},
			},
			Action: func(c *cli.Context) {
				// Check for required arguments.
				if c.String("directory") == "" {
					fmt.Println("--directory is required")
					return
				}
				if c.String("prefer") == "" {
					fmt.Println("--prefer is required")
					return
				}
				log.SetLevel(c.Int("v"))
				// Compile regexps.
				var prefer, over *regexp.Regexp
				var err error
				if prefer, err = regexp.Compile(c.String("prefer")); err != nil {
					fmt.Println("invalid regexp specified for --prefer")
					return
				}
				if c.String("over") != "" {
					if over, err = regexp.Compile(c.String("over")); err != nil {
						fmt.Println("invalid regexp specified for --over")
						return
					}
				}
				// Do deduplication.
				if err = runAutomatic(c.Bool("dry_run"), c.String("directory"), prefer, over, c.Bool("invert")); err != nil {
					fmt.Println(err)
				}
			},
		},
	}
	app.RunAndExitOnError()
}

func remove(dryRun bool, file string) {
	if dryRun {
		log.Info("DRY RUN: remove %s", file)
	} else {
		if err := os.Remove(file); err != nil {
			log.Warn("Error deleting %n: %v", file, err)
		}
	}

}

func getInput(prompt string, validator func(string) bool) (string, error) {
	// Reader to read from user input.
	reader := bufio.NewReader(os.Stdin)
	var (
		val string
		err error
	)
	for {
		fmt.Printf(prompt)
		val, err = reader.ReadString('\n')
		if err != nil {
			return "", err
		}
		val = strings.TrimSpace(val)
		if validator(val) {
			break
		}
	}
	return val, nil
}

func getDupesAndPrintSummary(root string) ([]dupes.Info, error) {
	if !os.IsPathSeparator(root[len(root)-1]) {
		root = fmt.Sprintf("%s%c", root, os.PathSeparator)
	}
	fmt.Printf("Indexing...")
	dupes, err := dupes.Dupes(root)
	if err != nil {
		return dupes, err
	}
	fcount := 0
	tsize := uint64(0)
	for _, i := range dupes {
		fcount += len(i.Names) - 1
		tsize += i.Size * uint64(len(i.Names)-1)
	}

	fmt.Printf("\rFound %d sets of duplicate files\n", len(dupes))
	fmt.Printf("Total file count: %d\n", fcount)
	fmt.Printf("Total size used: %s\n", humanize.Bytes(tsize))
	return dupes, err
}

func runInteractive(dryRun bool) error {
	// Get parent dir.
	root, err := getInput("Enter parent directory to scan for duplicates in: ", func(f string) bool {
		i, err := os.Stat(f)
		if err != nil {
			return false
		}
		return i.IsDir()
	})
	if err != nil {
		return err
	}
	dupes, err := getDupesAndPrintSummary(root)
	if err != nil {
		return err
	}

	fmt.Printf("\nReviewing results:\nFor each duplicate fileset, select 'f' to delete all but the first file, 'a' to keep all files, or 'n' (e.g. 2) to delete all except the second file.\n")

	for i, dupe := range dupes {
		keep := -1
		names := ""
		for j, n := range dupe.Names {
			names += fmt.Sprintf("%d: %s\n", j+1, n)
		}
		_, err := getInput(fmt.Sprintf("\n%d of %d  %s:\n%s\nKeep [F]irst/[a]ll/[n]th? ", i+1, len(dupes), humanize.Bytes(dupe.Size*uint64(len(dupe.Names)-1)), names), func(v string) bool {
			switch strings.ToLower(v) {
			case "f":
				keep = 0
				return true
			case "a":
				return true
			default:
				k, err := strconv.ParseInt(v, 10, 32)
				if err != nil {
					return false
				}
				if k < 1 || int(k) > len(dupe.Names) {
					return false
				}
				keep = int(k) - 1
				return true
			}
		})
		if err != nil {
			log.Fatal("%v", err)
		}
		if keep >= 0 {
			for i, n := range dupe.Names {
				if i != keep {
					remove(dryRun, n)
				}
			}
		}
	}
	return nil
}

func runPrint(root string, output string) error {
	dupes, err := getDupesAndPrintSummary(root)
	if err != nil {
		return err
	}

	var f *os.File
	if output != "" {
		if f, err = os.Create(output); err != nil {
			return err
		}
	}
	defer func() {
		if f != nil {
			if err := f.Close(); err != nil {
				log.Warn("%v", err)
			}
		}
	}()
	for _, dupe := range dupes {
		l := fmt.Sprintf("%s * %d => %s\n", humanize.Bytes(dupe.Size), len(dupe.Names), strings.Join(dupe.Names, ", "))
		if f != nil {
			f.Write([]byte(l))
		} else {
			fmt.Printf(l)
		}
	}
	return nil
}

func runAutomatic(dryRun bool, root string, prefer *regexp.Regexp, over *regexp.Regexp, invert bool) error {
	dupes, err := getDupesAndPrintSummary(root)
	if err != nil {
		return err
	}

	for _, dupe := range dupes {
		p := make(map[string]struct{})
		o := make(map[string]struct{})
		for _, n := range dupe.Names {
			nb := []byte(n)
			pm := prefer.Match(nb)
			om := false
			if over != nil {
				om = over.Match(nb)
			}
			if pm && om {
				log.Warn("both --prefer and --over matched %s", n)
			}
			if pm {
				p[n] = struct{}{}
			} else if om {
				o[n] = struct{}{}
			}
		}
		if len(p) > 0 { // If we found a preferred match.
			// Generate debug line.
			dbg := fmt.Sprintf("processing %s\n\tprefer: ", strings.Join(dupe.Names, ", "))
			for k := range p {
				dbg += k + ", "
			}
			if len(o) > 0 {
				dbg += "\n\tover: "
				for k := range o {
					dbg += k + ", "
				}
			}
			log.Debug("%s", dbg)
			// Logic is here.
			if over != nil {
				if len(o) > 0 {
					// If prefer and over are both specified, and both match, remove the non-preferred matches.
					if invert {
						for n := range p {
							remove(dryRun, n)
						}
					} else {
						for n := range o {
							remove(dryRun, n)
						}
					}
				}
			} else {
				// If over is not specified, keep only the preferred, but (for the case of --invert) only when preferred is not everything.
				if len(p) < len(dupe.Names) {
					for _, n := range dupe.Names {
						if _, ok := p[n]; ok && invert || !ok && !invert {
							remove(dryRun, n)
						}
					}
				}

			}
		}
	}

	return nil
}
<commit_msg>Logging improvements.<commit_after>package main

import (
	"bufio"
	"fmt"
	"github.com/cenkalti/log"
	"github.com/codegangsta/cli"
	dupes "github.com/danmarg/undupes/libdupes"
	"github.com/dustin/go-humanize"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// Horrible hack to convert -v to log Levels.
var orderedLevels []log.Level = []log.Level{log.DEBUG, log.INFO, log.NOTICE, log.WARNING, log.ERROR, log.CRITICAL}

func setLogLevel(l int) error {
	if l >= 0 && l < len(orderedLevels) {
		log.SetLevel(orderedLevels[l])
		return nil
	}
	return fmt.Errorf("invalid log level specified")
}

func main() {
	app := cli.NewApp()
	app.Name = "undupes"
	app.Usage = "manage duplicate files"
	app.Author = "Daniel Margolis"
	app.Email = "dan@af0.net"
	app.Version = "0.1"
	app.Commands = []cli.Command{
		{
			Name:      "interactive",
			ShortName: "i",
			Usage:     "interactive mode",
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "v",
					Usage: "log level",
					Value: 3, // Don't print as much in interactive mode.
				},
			},
			Action: func(c *cli.Context) {
				if err := setLogLevel(c.Int("v")); err != nil {
					fmt.Println(err)
					return
				}
				if err := runInteractive(c.Bool("dry_run")); err != nil {
					fmt.Println(err)
				}
			},
		},
		{
			Name:      "print",
			ShortName: "p",
			Usage:     "print duplicates",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "directory, d",
					Usage: "directory in which to find duplicates (required)",
				},
				cli.StringFlag{
					Name:  "output, o",
					Usage: "output file",
				},
			},
			Action: func(c *cli.Context) {
				if err := runPrint(c.String("directory"), c.String("output")); err != nil {
					fmt.Println(err)
				}
			},
		},
		{
			Name:      "auto",
			ShortName: "a",
			Usage:     "automatically resolve duplicates",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "directory, d",
					Usage: "directory in which to find duplicates (required)",
				},
				cli.StringFlag{
					Name:  "prefer, p",
					Usage: "prefer to keep files matching this pattern",
				},
				cli.StringFlag{
					Name:  "over, o",
					Usage: "used with --prefer; will restrict preferred files to those where the duplicate matches --over",
				},
				cli.BoolFlag{
					Name:  "invert, i",
					Usage: "invert matching logic; preferred files with be prioritized for deletion rather than retention",
				},
				cli.BoolFlag{
					Name:  "dry_run",
					Usage: "simulate (log but don't delete files)",
				},
				cli.IntFlag{
					Name:  "v",
					Usage: "log level",
					Value: 1,
				},
			},
			Action: func(c *cli.Context) {
				// Check for required arguments.
				if c.String("directory") == "" {
					fmt.Println("--directory is required")
					return
				}
				if c.String("prefer") == "" {
					fmt.Println("--prefer is required")
					return
				}
				if err := setLogLevel(c.Int("v")); err != nil {
					fmt.Println(err)
					return
				}
				// Compile regexps.
				var prefer, over *regexp.Regexp
				var err error
				if prefer, err = regexp.Compile(c.String("prefer")); err != nil {
					fmt.Println("invalid regexp specified for --prefer")
					return
				}
				if c.String("over") != "" {
					if over, err = regexp.Compile(c.String("over")); err != nil {
						fmt.Println("invalid regexp specified for --over")
						return
					}
				}
				// Do deduplication.
				if err = runAutomatic(c.Bool("dry_run"), c.String("directory"), prefer, over, c.Bool("invert")); err != nil {
					fmt.Println(err)
				}
			},
		},
	}
	app.RunAndExitOnError()
}

func remove(dryRun bool, file string) {
	if dryRun {
		log.Noticef("DRY RUN: delete %s", file)
	} else {
		if err := os.Remove(file); err != nil {
			log.Warningf("error deleting %n: %v", file, err)
		} else {
			log.Noticef("deleted %n", file)
		}
	}

}

func getInput(prompt string, validator func(string) bool) (string, error) {
	// Reader to read from user input.
	reader := bufio.NewReader(os.Stdin)
	var (
		val string
		err error
	)
	for {
		fmt.Printf(prompt)
		val, err = reader.ReadString('\n')
		if err != nil {
			return "", err
		}
		val = strings.TrimSpace(val)
		if validator(val) {
			break
		}
	}
	return val, nil
}

func getDupesAndPrintSummary(root string) ([]dupes.Info, error) {
	if !os.IsPathSeparator(root[len(root)-1]) {
		root = fmt.Sprintf("%s%c", root, os.PathSeparator)
	}
	fmt.Printf("Indexing...")
	dupes, err := dupes.Dupes(root)
	if err != nil {
		return dupes, err
	}
	fcount := 0
	tsize := uint64(0)
	for _, i := range dupes {
		fcount += len(i.Names) - 1
		tsize += i.Size * uint64(len(i.Names)-1)
	}

	fmt.Printf("\rFound %d sets of duplicate files\n", len(dupes))
	fmt.Printf("Total file count: %d\n", fcount)
	fmt.Printf("Total size used: %s\n", humanize.Bytes(tsize))
	return dupes, err
}

func runInteractive(dryRun bool) error {
	// Get parent dir.
	root, err := getInput("Enter parent directory to scan for duplicates in: ", func(f string) bool {
		i, err := os.Stat(f)
		if err != nil {
			return false
		}
		return i.IsDir()
	})
	if err != nil {
		return err
	}
	dupes, err := getDupesAndPrintSummary(root)
	if err != nil {
		return err
	}

	fmt.Printf("\nReviewing results:\nFor each duplicate fileset, select 'f' to delete all but the first file, 'a' to keep all files, or 'n' (e.g. 2) to delete all except the second file.\n")

	for i, dupe := range dupes {
		keep := -1
		names := ""
		for j, n := range dupe.Names {
			names += fmt.Sprintf("%d: %s\n", j+1, n)
		}
		_, err := getInput(fmt.Sprintf("\n%d of %d  %s:\n%s\nKeep [F]irst/[a]ll/[n]th? ", i+1, len(dupes), humanize.Bytes(dupe.Size*uint64(len(dupe.Names)-1)), names), func(v string) bool {
			switch strings.ToLower(v) {
			case "f":
				keep = 0
				return true
			case "a":
				return true
			default:
				k, err := strconv.ParseInt(v, 10, 32)
				if err != nil {
					return false
				}
				if k < 1 || int(k) > len(dupe.Names) {
					return false
				}
				keep = int(k) - 1
				return true
			}
		})
		if err != nil {
			log.Fatalln(err)
		}
		if keep >= 0 {
			for i, n := range dupe.Names {
				if i != keep {
					remove(dryRun, n)
				}
			}
		}
	}
	return nil
}

func runPrint(root string, output string) error {
	dupes, err := getDupesAndPrintSummary(root)
	if err != nil {
		return err
	}

	var f *os.File
	if output != "" {
		if f, err = os.Create(output); err != nil {
			return err
		}
	}
	defer func() {
		if f != nil {
			if err := f.Close(); err != nil {
				log.Warningln(err)
			}
		}
	}()
	for _, dupe := range dupes {
		l := fmt.Sprintf("%s * %d => %s\n", humanize.Bytes(dupe.Size), len(dupe.Names), strings.Join(dupe.Names, ", "))
		if f != nil {
			f.Write([]byte(l))
		} else {
			fmt.Printf(l)
		}
	}
	return nil
}

func runAutomatic(dryRun bool, root string, prefer *regexp.Regexp, over *regexp.Regexp, invert bool) error {
	dupes, err := getDupesAndPrintSummary(root)
	if err != nil {
		return err
	}

	for _, dupe := range dupes {
		p := make(map[string]struct{})
		o := make(map[string]struct{})
		for _, n := range dupe.Names {
			nb := []byte(n)
			pm := prefer.Match(nb)
			om := false
			if over != nil {
				om = over.Match(nb)
			}
			if pm && om {
				log.Warningf("both --prefer and --over matched %s", n)
			}
			if pm {
				p[n] = struct{}{}
			} else if om {
				o[n] = struct{}{}
			}
		}
		if len(p) > 0 { // If we found a preferred match.
			// Generate debug line.
			dbg := fmt.Sprintf("processing %s\n\tprefer: ", strings.Join(dupe.Names, ", "))
			for k := range p {
				dbg += k + ", "
			}
			if len(o) > 0 {
				dbg += "\n\tover: "
				for k := range o {
					dbg += k + ", "
				}
			}
			log.Debugf("%s", dbg)
			// Logic is here.
			if over != nil {
				if len(o) > 0 {
					// If prefer and over are both specified, and both match, remove the non-preferred matches.
					if invert {
						for n := range p {
							remove(dryRun, n)
						}
					} else {
						for n := range o {
							remove(dryRun, n)
						}
					}
				}
			} else {
				// If over is not specified, keep only the preferred, but (for the case of --invert) only when preferred is not everything.
				if len(p) < len(dupe.Names) {
					for _, n := range dupe.Names {
						if _, ok := p[n]; ok && invert || !ok && !invert {
							remove(dryRun, n)
						}
					}
				}

			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/docopt/docopt-go"
	"github.com/mitchellh/go-homedir"
)

var (
	HOME, _           = homedir.Dir()
	ENVIRONMENTS_FILE = filepath.Join(HOME, ".gs3pload", "envs.json")
	VERSION           = "0.0.1"
)

func setBucket(name string, environmentType string, environmentName string) string {
	var path, root string

	pathSplit := strings.SplitN(name, "/", 2)
	name = pathSplit[0]
	switch name {
	case "packages":
		root = fmt.Sprintf("%s.%s", name, environmentName)
	case "certs", "images", "stacks":
		root = fmt.Sprintf("%s.%s", environmentName, name)
	default:
		root = name
	}

	if len(pathSplit) > 1 {
		path = fmt.Sprintf("%s/%s", root, pathSplit[1])
	} else {
		path = root
	}

	return fmt.Sprintf("%s://%s/", environmentType, path)
}

// Execute executes a regular command splitted in strings
func Execute(command []string) error {
	cmd := exec.Command(command[0], command[1:]...)

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// Copy files to a given bucket and environment
func Copy(config string, bucket string, files []string, environment Environment) error {
	command := []string{"gsutil", "cp"}
	command = append(command, files...)
	command = append(command, bucket)

	os.Setenv("BOTO_CONFIG", config)
	return Execute(command)
}

// Public set public-read permissions for the given files
func Public(config string, bucket string, files []string, environment Environment) error {
	command := []string{"gsutil", "acl", "set", "public-read"}
	for _, file := range files {
		filePath := fmt.Sprintf("%s%s", bucket, file)
		command = append(command, filePath)
	}
	os.Setenv("BOTO_CONFIG", config)
	return Execute(command)
}

// Push given files to multiple environments
func Push(environments Environments, bucketName string, files []string, public bool) error {
	for _, environment := range environments {
		fmt.Printf("---> Pushing to %s environment...\n", environment.Name)
		bucket := setBucket(bucketName, environment.Type, environment.Name)
		config := fmt.Sprintf("%s.boto", filepath.Join(HOME, ".gs3pload", environment.Name))

		err := Copy(config, bucket, files, environment)
		if err != nil {
			fmt.Printf("Push failed on %s. %s\n", environment.Name, err)
			continue
		}

		if public {
			err = Public(config, bucket, files, environment)
			if err != nil {
				fmt.Printf("Set as public failed on %s. %s\n", environment.Name, err)
				continue
			}
		}
	}
	return nil
}

func main() {
	usage := `gs3pload. Upload files to multiple S3 or Google Storage
buckets at once.

Bucket names "packages", "stacks", "certs" and "images" are reserved and resolved
by environment domain.

Usage:
  gs3pload push <bucket> <name>... [-p | --public]
  gs3pload -h | --help
  gs3pload -v | --version

Options:
  -h --help        Show help.
  -p --public      Set files as public.
  -v --version     Show version.
`

	arguments, _ := docopt.Parse(usage, nil, true, fmt.Sprintf("gs3pload %s", VERSION), false)

	push := arguments["push"].(bool)
	bucketName := arguments["<bucket>"].(string)
	fileNames := arguments["<name>"].([]string)
	public := arguments["--public"].(bool)

	environments := Environments{}
	err := environments.Fetch()
	if err != nil {
		fmt.Println(err)
		return
	}

	if push {
		err = Push(environments, bucketName, fileNames, public)
		if err != nil {
			fmt.Println(err)
			return
		}
	}
}
<commit_msg>Introduce Recursive and Backup options<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/docopt/docopt-go"
	"github.com/mitchellh/go-homedir"
)

var (
	HOME, _           = homedir.Dir()
	ENVIRONMENTS_FILE = filepath.Join(HOME, ".gs3pload", "envs.json")
	VERSION           = "0.0.1"
)

func setBucket(name string, environmentType string, environmentName string) string {
	var path, root string

	pathSplit := strings.SplitN(name, "/", 2)
	name = pathSplit[0]
	switch name {
	case "packages":
		root = fmt.Sprintf("%s.%s", name, environmentName)
	case "certs", "images", "stacks":
		root = fmt.Sprintf("%s.%s", environmentName, name)
	default:
		root = name
	}

	if len(pathSplit) > 1 {
		path = fmt.Sprintf("%s/%s", root, pathSplit[1])
	} else {
		path = root
	}

	return fmt.Sprintf("%s://%s/", environmentType, path)
}

// Execute executes a regular command splitted in strings
func Execute(command []string) error {
	cmd := exec.Command(command[0], command[1:]...)

	cmd.Stdout = os.Stdout
	return cmd.Run()
}

// Copy files to a given bucket and environment
func Copy(config string, bucket string, files []string, recursive bool, environment Environment) error {
	command := []string{"gsutil", "cp"}
	if recursive {
		command = append(command, "-R")
	}
	command = append(command, files...)
	command = append(command, bucket)

	os.Setenv("BOTO_CONFIG", config)
	return Execute(command)
}

// Public set public-read permissions for the given files
func Public(config string, bucket string, files []string) error {
	command := []string{"gsutil", "acl", "set", "public-read"}
	for _, file := range files {
		filePath := fmt.Sprintf("%s%s", bucket, file)
		command = append(command, filePath)
	}
	os.Setenv("BOTO_CONFIG", config)
	return Execute(command)
}

// DaisyChain copy an object from one path to another, where these
// can belong to different buckets or environments
func DaisyChain(config string, originPath, destPath string, recursive bool) error {
	command := []string{"gsutil", "cp", "-D", "-p"}
	if recursive {
		command = append(command, "-R")
	}
	command = append(command, originPath)
	command = append(command, destPath)
	os.Setenv("BOTO_CONFIG", config)
	return Execute(command)
}

// Backup given files within the same bucket with a .bak extension
func Backup(config string, files []string, bucket string, recursive bool) {
	fmt.Printf("---> Creating backups on %s...\n", bucket)
	for _, file := range files {
		filePath := fmt.Sprintf("%s%s", bucket, file)
		backupPath := fmt.Sprintf("%s%s%s", bucket, file, ".bak")
		if err := DaisyChain(config, filePath, backupPath, recursive); err != nil {
			fmt.Printf("Skipping backup of %s (Did it exist before?). \n", filePath)
			continue
		}
	}
}

// Push given files to multiple environments
func Push(environments Environments, bucketName string, files []string, recursive, public, backup bool) error {
	for _, environment := range environments {
		fmt.Printf("---> Pushing to %s environment...\n", environment.Name)
		bucket := setBucket(bucketName, environment.Type, environment.Name)
		config := fmt.Sprintf("%s.boto", filepath.Join(HOME, ".gs3pload", environment.Name))

		if backup {
			Backup(config, files, bucket, recursive)
		}

		if err := Copy(config, bucket, files, recursive, environment); err != nil {
			fmt.Printf("Push failed on %s. %s\n", environment.Name, err)
			continue
		}

		if public {
			if err := Public(config, bucket, files); err != nil {
				fmt.Printf("Set as public failed on %s. %s\n", environment.Name, err)
				continue
			}
		}
	}
	return nil
}

func main() {
	usage := `gs3pload. Upload files to multiple S3 or Google Storage
buckets at once.

Bucket names "packages", "stacks", "certs" and "images" are reserved and resolved
by environment domain.

Usage:
  gs3pload push <bucket> <name>... [-r | --recursive] [-p | --public] [-b | --backup]
  gs3pload -h | --help
  gs3pload -v | --version

Options:
  -h --help        Show help.
  -p --public      Set files as public.
  -r --recursive   Do a recursive copy.
  -b --backup      Create backup of pushed files if they existed.
  -v --version     Show version.
`

	arguments, _ := docopt.Parse(usage, nil, true, fmt.Sprintf("gs3pload %s", VERSION), false)

	push := arguments["push"].(bool)
	bucketName := arguments["<bucket>"].(string)
	fileNames := arguments["<name>"].([]string)
	recursive := arguments["--recursive"].(bool)
	public := arguments["--public"].(bool)
	backup := arguments["--backup"].(bool)

	environments := Environments{}
	err := environments.Fetch()
	if err != nil {
		fmt.Println(err)
		return
	}

	if push {
		err = Push(environments, bucketName, fileNames, recursive, public, backup)
		if err != nil {
			fmt.Println(err)
			return
		}
	}
}
<|endoftext|>


<commit_before>/* 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/. */

// Based on github.com/jmhodges/certificatetransparency/tools/lecsv

package main

import (
  "database/sql"
  "flag"
  "fmt"
  "log"
  "net/url"
  "os"
  "runtime"
  "strings"
  "sync"
  "time"

  _ "github.com/go-sql-driver/mysql"

  "github.com/go-gorp/gorp"
  "github.com/google/certificate-transparency/go"
  "github.com/google/certificate-transparency/go/client"
  "github.com/jcjones/ct-sql/censysdata"
  "github.com/jcjones/ct-sql/sqldb"
)

var (
  logUrl     = flag.String("log", "", "URL of the CT Log")
  censysPath = flag.String("censysJson", "", "Path to a Censys.io certificate json dump")
  dbConnect  = flag.String("dbConnect", "", "DB Connection String")
  verbose    = flag.Bool("v", false, "verbose output")
  offset     = flag.Uint64("offset", 0, "offset from the beginning")
  offsetByte = flag.Uint64("offsetByte", 0, "byte offset from the beginning, only for censysJson and not compatible with offset")
  limit      = flag.Uint64("limit", 0, "limit processing to this many entries")

)

// OperationStatus contains the current state of a large operation (i.e.
// download or tree hash).
type OperationStatus struct {
  // Start contains the requested starting index of the operation.
  Start int64
  // Current contains the greatest index that has been processed.
  Current int64
  // Length contains the total number of entries.
  Length int64
}

func (status OperationStatus) Percentage() float32 {
  total := float32(status.Length - status.Start)
  done := float32(status.Current - status.Start)

  if total == 0 {
    return 100
  }
  return done * 100 / total
}

// Taken from Boulder
func recombineURLForDB(dbConnect string) (string, error) {
  dbConnect = strings.TrimSpace(dbConnect)
  dbURL, err := url.Parse(dbConnect)
  if err != nil {
    return "", err
  }

  if dbURL.Scheme != "mysql+tcp" {
    format := "given database connection string was not a mysql+tcp:// URL, was %#v"
    return "", fmt.Errorf(format, dbURL.Scheme)
  }

  dsnVals, err := url.ParseQuery(dbURL.RawQuery)
  if err != nil {
    return "", err
  }

  dsnVals.Set("parseTime", "true")

  // Required to make UPDATE return the number of rows matched,
  // instead of the number of rows changed by the UPDATE.
  dsnVals.Set("clientFoundRows", "true")

  // Ensures that MySQL/MariaDB warnings are treated as errors. This
  // avoids a number of nasty edge conditions we could wander
  // into. Common things this discovers includes places where data
  // being sent had a different type than what is in the schema,
  // strings being truncated, writing null to a NOT NULL column, and
  // so on. See
  // <https://dev.mysql.com/doc/refman/5.0/en/sql-mode.html#sql-mode-strict>.
  dsnVals.Set("strict", "true")

  user := dbURL.User.Username()
  passwd, hasPass := dbURL.User.Password()
  dbConn := ""
  if user != "" {
    dbConn = url.QueryEscape(user)
  }
  if hasPass {
    dbConn += ":" + passwd
  }
  dbConn += "@tcp(" + dbURL.Host + ")"
  return dbConn + dbURL.EscapedPath() + "?" + dsnVals.Encode(), nil
}

func clearLine() {
  fmt.Printf("\x1b[80D\x1b[2K")
}

func displayProgress(statusChan chan OperationStatus, wg *sync.WaitGroup) {
  wg.Add(1)

  go func() {
    defer wg.Done()
    symbols := []string{"|", "/", "-", "\\"}
    symbolIndex := 0

    status, ok := <-statusChan
    if !ok {
      return
    }

    ticker := time.NewTicker(200 * time.Millisecond)
    defer ticker.Stop()

    isInteractive := strings.Contains(os.Getenv("TERM"), "xterm")

    if !isInteractive {
      ticker.Stop()
    }

    for {
      select {
      case status, ok = <-statusChan:
        if !ok {
          return
        }
      case <-ticker.C:
        symbolIndex = (symbolIndex + 1) % len(symbols)
      }

      if isInteractive {
        clearLine()
        fmt.Printf("%s %.1f%% (%d of %d)", symbols[symbolIndex], status.Percentage(), status.Current, status.Length)
      } else {
        fmt.Println(fmt.Printf("%.1f%% (%d of %d)", status.Percentage(), status.Current, status.Length))
      }
    }
  }()
}

func insertCTWorker(entries <-chan ct.LogEntry, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) {
  wg.Add(1)
  defer wg.Done()
  for ep := range entries {
    err := db.InsertCTEntry(&ep)
    if err != nil {
      log.Printf("Problem inserting certificate: index: %d log: %s error: %s", ep.Index, *logUrl, err)
    }
  }
}

func insertCensysWorker(entries <-chan censysdata.CensysEntry, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) {
  wg.Add(1)
  defer wg.Done()
  for ep := range entries {
    if ep.Valid_nss {
      err := db.InsertCensysEntry(&ep)
      if err != nil {
        log.Printf("Problem inserting certificate: index: %d error: %s", ep.Offset, err)
      }
    }
  }
}

// DownloadRange downloads log entries from the given starting index till one
// less than upTo. If status is not nil then status updates will be written to
// it until the function is complete, when it will be closed. The log entries
// are provided to an output channel.
func downloadCTRangeToChannel(ctLog *client.LogClient, outEntries chan<- ct.LogEntry,
                            status chan<- OperationStatus, start, upTo int64) (int64, error) {
  if outEntries == nil {
    return 0, fmt.Errorf("No output channel provided")
  }
  defer close(outEntries)
  if status != nil {
    defer close(status)
  }

  index := start
  for index < upTo {
    if status != nil {
      status <- OperationStatus{start, index, upTo}
    }

    max := index + 2000
    if max >= upTo {
      max = upTo - 1
    }
    rawEnts, err := ctLog.GetEntries(index, max)
    if err != nil {
      return index, err
    }

    for _, ent := range rawEnts {
      outEntries <- ent
      if (ent.Index) != index {
        return index, fmt.Errorf("Index mismatch, local: %v, remote: %v", index, ent.Index)
      }

      index++
    }
  }

  return index, nil
}

func downloadLog(ctLogUrl *url.URL, ctLog *client.LogClient, db *sqldb.EntriesDatabase) (error) {
  if *offsetByte > 0 {
    return fmt.Errorf("Cannot set offsetByte for CT log downloads")
  }

  fmt.Printf("Fetching signed tree head... ")
  sth, err := ctLog.GetSTH()
  if err != nil {
    return err
  }

  // Set pointer in DB, now that we've verified the log works
  err = db.SetLog(fmt.Sprintf("%s%s", ctLogUrl.Host, ctLogUrl.Path))
  if err != nil {
    log.Fatalf("unable to set Certificate Log: %s", err)
  }

  var origCount uint64
  // Now we're OK to use the DB
  if *offset > 0 {
    log.Printf("Starting from offset %d", *offset)
    origCount = *offset
  } else {
    log.Printf("Counting existing entries... ")
    origCount, err = db.Count()
    if err != nil {
      err = fmt.Errorf("Failed to read entries file: %s", err)
      return err
    }
  }

  fmt.Printf("%d total entries at %s\n", sth.TreeSize, sqldb.Uint64ToTimestamp(sth.Timestamp).Format(time.ANSIC))
  if origCount == sth.TreeSize {
    fmt.Printf("Nothing to do\n")
    return nil
  }

  endPos := sth.TreeSize
  if *limit > 0 && endPos > origCount + *limit {
    endPos = origCount + *limit
  }

  fmt.Printf("Going from %d to %d\n", origCount, endPos)

  entryChan := make(chan ct.LogEntry, 100)
  statusChan := make(chan OperationStatus, 1)
  wg := new(sync.WaitGroup)

  displayProgress(statusChan, wg)
  for i := 0; i < runtime.NumCPU(); i++ {
    go insertCTWorker(entryChan, db, wg)
  }
  _, err = downloadCTRangeToChannel(ctLog, entryChan, statusChan, int64(origCount), int64(endPos))
  wg.Wait()

  clearLine()
  if err != nil {
    err = fmt.Errorf("Error while downloading: %s", err)
    return err
  }

  return nil
}

func processImporter(importer *censysdata.Importer, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) error {
  entryChan := make(chan censysdata.CensysEntry, 100)
  defer close(entryChan)
  statusChan := make(chan OperationStatus, 1)
  defer close(statusChan)
  wg.Add(1)
  defer wg.Done()

  displayProgress(statusChan, wg)
  for i := 0; i < runtime.NumCPU(); i++ {
    go insertCensysWorker(entryChan, db, wg)
  }

  startOffset := *offsetByte
  maxOffset, err := importer.Size()
  if err != nil {
    return err
  }

  //
  // Fast forward
  //
  if *offset >0 && *offsetByte>0 {
    return fmt.Errorf("You may not set both offset and offsetByte")
  }

  if *offset > 0 {
    err = importer.SeekLine(*offset)
    if err != nil {
      return err
    }
  }

  if *offsetByte > 0 {
    err = importer.SeekByte(*offsetByte)
    if err != nil {
      return err
    }
  }

  log.Printf("Starting import from offset=%d, line_limit=%d / size=%d.", importer.ByteOffset(), *limit, maxOffset)

  // We've already fast-forwarded, so start at 0.
  for count := uint64(0);; count++ {
    if *limit > uint64(0) && count >= *limit {
      return nil
    }
    ent, err := importer.NextEntry()
    if err != nil || ent == nil {
      return err
    }
    statusChan <- OperationStatus{int64(startOffset), int64(ent.Offset), int64(maxOffset)}
    entryChan <- *ent
  }

  return nil
}

func main() {
  flag.Parse()
  log.SetFlags(0)
  log.SetPrefix("")
  dbConnectStr, err := recombineURLForDB(*dbConnect)
  if err != nil {
    log.Printf("unable to parse %s: %s", *dbConnect, err)
  }

  if len(dbConnectStr) == 0 || (censysPath == nil && logUrl == nil) {
    flag.Usage()
    os.Exit(2)
  }

  db, err := sql.Open("mysql", dbConnectStr)
  if err != nil {
    log.Fatalf("unable to open SQL: %s: %s", dbConnectStr, err)
  }
  if err = db.Ping(); err != nil {
    log.Fatalf("unable to ping SQL: %s: %s", dbConnectStr, err)
  }

  dialect := gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF8"}
  dbMap := &gorp.DbMap{Db: db, Dialect: dialect}
  entriesDb := &sqldb.EntriesDatabase{DbMap: dbMap, Verbose: *verbose}
  err = entriesDb.InitTables()
  if err != nil {
    log.Fatalf("unable to prepare SQL: %s: %s", dbConnectStr, err)
  }

  if logUrl != nil && len(*logUrl) > 5 {
    ctLogUrl, err := url.Parse(*logUrl)
    if err != nil {
      log.Fatalf("unable to set Certificate Log: %s", err)
    }

    ctLog := client.New(*logUrl)

    err = downloadLog(ctLogUrl, ctLog, entriesDb)
    if err != nil {
      log.Fatalf("error while updating CT entries: %s", err)
    }

    os.Exit(0)
  }

  if censysPath != nil && len(*censysPath) > 5 {
    importer := &censysdata.Importer{}
    err = importer.OpenFile(*censysPath)
    if err != nil {
      log.Fatalf("unable to open Censys file: %s", err)
    }
    defer importer.Close()

    wg := new(sync.WaitGroup)
    err := processImporter(importer, entriesDb, wg)

    if err != nil {
      log.Fatalf("error while running importer: %s", err)
    }

    wg.Wait()
    os.Exit(0)
  }

  // Didn't include a mandatory action, so print usage and exit.
  flag.Usage()
  os.Exit(2)
}
<commit_msg>Don't spam status<commit_after>/* 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/. */

// Based on github.com/jmhodges/certificatetransparency/tools/lecsv

package main

import (
  "database/sql"
  "flag"
  "fmt"
  "log"
  "net/url"
  "os"
  "runtime"
  "strings"
  "sync"
  "time"

  _ "github.com/go-sql-driver/mysql"

  "github.com/go-gorp/gorp"
  "github.com/google/certificate-transparency/go"
  "github.com/google/certificate-transparency/go/client"
  "github.com/jcjones/ct-sql/censysdata"
  "github.com/jcjones/ct-sql/sqldb"
)

var (
  logUrl     = flag.String("log", "", "URL of the CT Log")
  censysPath = flag.String("censysJson", "", "Path to a Censys.io certificate json dump")
  dbConnect  = flag.String("dbConnect", "", "DB Connection String")
  verbose    = flag.Bool("v", false, "verbose output")
  offset     = flag.Uint64("offset", 0, "offset from the beginning")
  offsetByte = flag.Uint64("offsetByte", 0, "byte offset from the beginning, only for censysJson and not compatible with offset")
  limit      = flag.Uint64("limit", 0, "limit processing to this many entries")

)

// OperationStatus contains the current state of a large operation (i.e.
// download or tree hash).
type OperationStatus struct {
  // Start contains the requested starting index of the operation.
  Start int64
  // Current contains the greatest index that has been processed.
  Current int64
  // Length contains the total number of entries.
  Length int64
}

func (status OperationStatus) Percentage() float32 {
  total := float32(status.Length - status.Start)
  done := float32(status.Current - status.Start)

  if total == 0 {
    return 100
  }
  return done * 100 / total
}

// Taken from Boulder
func recombineURLForDB(dbConnect string) (string, error) {
  dbConnect = strings.TrimSpace(dbConnect)
  dbURL, err := url.Parse(dbConnect)
  if err != nil {
    return "", err
  }

  if dbURL.Scheme != "mysql+tcp" {
    format := "given database connection string was not a mysql+tcp:// URL, was %#v"
    return "", fmt.Errorf(format, dbURL.Scheme)
  }

  dsnVals, err := url.ParseQuery(dbURL.RawQuery)
  if err != nil {
    return "", err
  }

  dsnVals.Set("parseTime", "true")

  // Required to make UPDATE return the number of rows matched,
  // instead of the number of rows changed by the UPDATE.
  dsnVals.Set("clientFoundRows", "true")

  // Ensures that MySQL/MariaDB warnings are treated as errors. This
  // avoids a number of nasty edge conditions we could wander
  // into. Common things this discovers includes places where data
  // being sent had a different type than what is in the schema,
  // strings being truncated, writing null to a NOT NULL column, and
  // so on. See
  // <https://dev.mysql.com/doc/refman/5.0/en/sql-mode.html#sql-mode-strict>.
  dsnVals.Set("strict", "true")

  user := dbURL.User.Username()
  passwd, hasPass := dbURL.User.Password()
  dbConn := ""
  if user != "" {
    dbConn = url.QueryEscape(user)
  }
  if hasPass {
    dbConn += ":" + passwd
  }
  dbConn += "@tcp(" + dbURL.Host + ")"
  return dbConn + dbURL.EscapedPath() + "?" + dsnVals.Encode(), nil
}

func clearLine() {
  fmt.Printf("\x1b[80D\x1b[2K")
}

func displayProgress(statusChan chan OperationStatus, wg *sync.WaitGroup) {
  wg.Add(1)

  go func() {
    defer wg.Done()
    symbols := []string{"|", "/", "-", "\\"}
    symbolIndex := 0

    status, ok := <-statusChan
    if !ok {
      return
    }

    ticker := time.NewTicker(200 * time.Millisecond)
    defer ticker.Stop()

    isInteractive := strings.Contains(os.Getenv("TERM"), "xterm")

    if !isInteractive {
      ticker.Stop()
    }

    for {
      select {
      case status, ok = <-statusChan:
        if !ok {
          return
        }
      case <-ticker.C:
        symbolIndex = (symbolIndex + 1) % len(symbols)
      }

      if isInteractive {
        clearLine()
        fmt.Printf("%s %.1f%% (%d of %d)", symbols[symbolIndex], status.Percentage(), status.Current, status.Length)
      } else {
        fmt.Println(fmt.Printf("%.1f%% (%d of %d)", status.Percentage(), status.Current, status.Length))
      }
    }
  }()
}

func insertCTWorker(entries <-chan ct.LogEntry, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) {
  wg.Add(1)
  defer wg.Done()
  for ep := range entries {
    err := db.InsertCTEntry(&ep)
    if err != nil {
      log.Printf("Problem inserting certificate: index: %d log: %s error: %s", ep.Index, *logUrl, err)
    }
  }
}

func insertCensysWorker(entries <-chan censysdata.CensysEntry, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) {
  wg.Add(1)
  defer wg.Done()
  for ep := range entries {
    if ep.Valid_nss {
      err := db.InsertCensysEntry(&ep)
      if err != nil {
        log.Printf("Problem inserting certificate: index: %d error: %s", ep.Offset, err)
      }
    }
  }
}

// DownloadRange downloads log entries from the given starting index till one
// less than upTo. If status is not nil then status updates will be written to
// it until the function is complete, when it will be closed. The log entries
// are provided to an output channel.
func downloadCTRangeToChannel(ctLog *client.LogClient, outEntries chan<- ct.LogEntry,
                            status chan<- OperationStatus, start, upTo int64) (int64, error) {
  if outEntries == nil {
    return 0, fmt.Errorf("No output channel provided")
  }
  defer close(outEntries)
  if status != nil {
    defer close(status)
  }

  index := start
  for index < upTo {
    if status != nil {
      status <- OperationStatus{start, index, upTo}
    }

    max := index + 2000
    if max >= upTo {
      max = upTo - 1
    }
    rawEnts, err := ctLog.GetEntries(index, max)
    if err != nil {
      return index, err
    }

    for _, ent := range rawEnts {
      outEntries <- ent
      if (ent.Index) != index {
        return index, fmt.Errorf("Index mismatch, local: %v, remote: %v", index, ent.Index)
      }

      index++
    }
  }

  return index, nil
}

func downloadLog(ctLogUrl *url.URL, ctLog *client.LogClient, db *sqldb.EntriesDatabase) (error) {
  if *offsetByte > 0 {
    return fmt.Errorf("Cannot set offsetByte for CT log downloads")
  }

  fmt.Printf("Fetching signed tree head... ")
  sth, err := ctLog.GetSTH()
  if err != nil {
    return err
  }

  // Set pointer in DB, now that we've verified the log works
  err = db.SetLog(fmt.Sprintf("%s%s", ctLogUrl.Host, ctLogUrl.Path))
  if err != nil {
    log.Fatalf("unable to set Certificate Log: %s", err)
  }

  var origCount uint64
  // Now we're OK to use the DB
  if *offset > 0 {
    log.Printf("Starting from offset %d", *offset)
    origCount = *offset
  } else {
    log.Printf("Counting existing entries... ")
    origCount, err = db.Count()
    if err != nil {
      err = fmt.Errorf("Failed to read entries file: %s", err)
      return err
    }
  }

  fmt.Printf("%d total entries at %s\n", sth.TreeSize, sqldb.Uint64ToTimestamp(sth.Timestamp).Format(time.ANSIC))
  if origCount == sth.TreeSize {
    fmt.Printf("Nothing to do\n")
    return nil
  }

  endPos := sth.TreeSize
  if *limit > 0 && endPos > origCount + *limit {
    endPos = origCount + *limit
  }

  fmt.Printf("Going from %d to %d\n", origCount, endPos)

  entryChan := make(chan ct.LogEntry, 100)
  statusChan := make(chan OperationStatus, 1)
  wg := new(sync.WaitGroup)

  displayProgress(statusChan, wg)
  for i := 0; i < runtime.NumCPU(); i++ {
    go insertCTWorker(entryChan, db, wg)
  }
  _, err = downloadCTRangeToChannel(ctLog, entryChan, statusChan, int64(origCount), int64(endPos))
  wg.Wait()

  clearLine()
  if err != nil {
    err = fmt.Errorf("Error while downloading: %s", err)
    return err
  }

  return nil
}

func processImporter(importer *censysdata.Importer, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) error {
  entryChan := make(chan censysdata.CensysEntry, 100)
  defer close(entryChan)
  statusChan := make(chan OperationStatus, 1)
  defer close(statusChan)
  wg.Add(1)
  defer wg.Done()

  displayProgress(statusChan, wg)
  for i := 0; i < runtime.NumCPU(); i++ {
    go insertCensysWorker(entryChan, db, wg)
  }

  startOffset := *offsetByte
  maxOffset, err := importer.Size()
  if err != nil {
    return err
  }

  //
  // Fast forward
  //
  if *offset >0 && *offsetByte>0 {
    return fmt.Errorf("You may not set both offset and offsetByte")
  }

  if *offset > 0 {
    err = importer.SeekLine(*offset)
    if err != nil {
      return err
    }
  }

  if *offsetByte > 0 {
    err = importer.SeekByte(*offsetByte)
    if err != nil {
      return err
    }
  }

  log.Printf("Starting import from offset=%d, line_limit=%d / size=%d.", importer.ByteOffset(), *limit, maxOffset)

  // We've already fast-forwarded, so start at 0.
  for count := uint64(0);; count++ {
    if *limit > uint64(0) && count >= *limit {
      return nil
    }
    ent, err := importer.NextEntry()
    if err != nil || ent == nil {
      return err
    }
    if count % 128 == 0 {
      statusChan <- OperationStatus{int64(startOffset), int64(ent.Offset), int64(maxOffset)}
    }
    entryChan <- *ent
  }

  return nil
}

func main() {
  flag.Parse()
  log.SetFlags(0)
  log.SetPrefix("")
  dbConnectStr, err := recombineURLForDB(*dbConnect)
  if err != nil {
    log.Printf("unable to parse %s: %s", *dbConnect, err)
  }

  if len(dbConnectStr) == 0 || (censysPath == nil && logUrl == nil) {
    flag.Usage()
    os.Exit(2)
  }

  db, err := sql.Open("mysql", dbConnectStr)
  if err != nil {
    log.Fatalf("unable to open SQL: %s: %s", dbConnectStr, err)
  }
  if err = db.Ping(); err != nil {
    log.Fatalf("unable to ping SQL: %s: %s", dbConnectStr, err)
  }

  dialect := gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF8"}
  dbMap := &gorp.DbMap{Db: db, Dialect: dialect}
  entriesDb := &sqldb.EntriesDatabase{DbMap: dbMap, Verbose: *verbose}
  err = entriesDb.InitTables()
  if err != nil {
    log.Fatalf("unable to prepare SQL: %s: %s", dbConnectStr, err)
  }

  if logUrl != nil && len(*logUrl) > 5 {
    ctLogUrl, err := url.Parse(*logUrl)
    if err != nil {
      log.Fatalf("unable to set Certificate Log: %s", err)
    }

    ctLog := client.New(*logUrl)

    err = downloadLog(ctLogUrl, ctLog, entriesDb)
    if err != nil {
      log.Fatalf("error while updating CT entries: %s", err)
    }

    os.Exit(0)
  }

  if censysPath != nil && len(*censysPath) > 5 {
    importer := &censysdata.Importer{}
    err = importer.OpenFile(*censysPath)
    if err != nil {
      log.Fatalf("unable to open Censys file: %s", err)
    }
    defer importer.Close()

    wg := new(sync.WaitGroup)
    err := processImporter(importer, entriesDb, wg)

    if err != nil {
      log.Fatalf("error while running importer: %s", err)
    }

    wg.Wait()
    os.Exit(0)
  }

  // Didn't include a mandatory action, so print usage and exit.
  flag.Usage()
  os.Exit(2)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
)

var (
	client      *http.Client
	registryURL string
)

type repos struct {
	Query      string `json:"query"`
	NumResults int    `json:"num_results"`
	Results    []repo `json:"results"`
}

type repo struct {
	Name        string `json:"name"`
	Description string `json:"description"`
}

type tag string

type image string

type tags map[tag]image

type imageNode struct {
	tags     []tag
	children []*imageNode
}

func init() {
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client = &http.Client{Transport: transport}
}

func doGet(path string) []byte {
	res, err := client.Get(registryURL + "/v1/" + path)
	if err != nil {
		panic(err.Error())
	}
	if res.StatusCode != 200 {
		return []byte{}
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err.Error())
	}
	res.Body.Close()
	return body
}

func getRepos() repos {
	log.Print("Fetching repos...")
	var repos repos
	json.Unmarshal(doGet("search"), &repos)
	log.Printf("%v repo(s) fetched", repos.NumResults)
	return repos
}

func getTags(name string) map[tag]image {
	log.Printf("Fetching tags for %s ...", name)
	var tags tags
	json.Unmarshal(doGet("repositories/"+name+"/tags"), &tags)
	log.Printf("%v tags fetched for repo %s", len(tags), name)
	return tags
}

func getAncestry(id image) []image {
	log.Printf("Fetching ancestry for %s ...", id)
	var ancestry []image
	json.Unmarshal(doGet("images/"+string(id)+"/ancestry"), &ancestry)
	log.Printf("%v ancestors fetched for repo %s", len(ancestry), id)
	return ancestry
}

func fqTag(name string, t tag) tag {
	canonicalName := strings.TrimPrefix(name, "library/")
	return tag(canonicalName + ":" + string(t))
}

func printTree(root *imageNode, level int) {
	if len(root.tags) > 0 || len(root.children) > 1 {
		fmt.Printf("%s%v\n", strings.Repeat(" ", level), root.tags)
		level = level +1
	}
	for _, child := range root.children {
		printTree(child, level)
	}
}

func main() {
	var (
		tagsByImage = make(map[image][]tag)
		images      = make(map[image]*imageNode)
		roots       []*imageNode
	)
	if len(registryURL) == 0 {
		registryURL = os.Getenv("REGISTRY_URL")
	}
	if len(registryURL) == 0 {
		log.Fatal("No registry URL provided, use the environment variable REGISTRY_URL to set it")
	}
	if len(os.Getenv("REGISTREE_DEBUG")) > 0 {
		log.SetOutput(os.Stderr)
	} else {
		log.SetOutput(ioutil.Discard)
	}
	for _, repo := range getRepos().Results {
		name := repo.Name
		for t, id := range getTags(name) {
			tags, _ := tagsByImage[id]
			tagsByImage[id] = append(tags, fqTag(name, t))
		}
	}
	for imageId := range tagsByImage {
		var previousNode *imageNode
		for _, ancestryId := range getAncestry(imageId) {
			if node, ok := images[ancestryId]; ok {
				if previousNode != nil {
					node.children = append(node.children, previousNode)
				}
				previousNode = nil
				break
			}
			node := &imageNode{}
			if tags, ok := tagsByImage[ancestryId]; ok {
				node.tags = tags
			}
			if previousNode != nil {
				node.children = []*imageNode{previousNode}
			}
			images[ancestryId] = node
			previousNode = node
		}
		if previousNode != nil {
			roots = append(roots, previousNode)
		}
	}
	for _, root := range roots {
		printTree(root, 0)
	}

}
<commit_msg>parallelize<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
)

var (
	client      *http.Client
	registryURL string
)

type repos struct {
	Query      string `json:"query"`
	NumResults int    `json:"num_results"`
	Results    []repo `json:"results"`
}

type repo struct {
	Name        string `json:"name"`
	Description string `json:"description"`
}

type tag string

type image string

type tags map[tag]image

type imageNode struct {
	tags     []tag
	children []*imageNode
}

func init() {
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client = &http.Client{Transport: transport}
}

func doGet(path string) []byte {
	res, err := client.Get(registryURL + "/v1/" + path)
	if err != nil {
		panic(err.Error())
	}
	if res.StatusCode != 200 {
		return []byte{}
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err.Error())
	}
	res.Body.Close()
	return body
}

func getRepos() repos {
	log.Print("Fetching repos...")
	var repos repos
	json.Unmarshal(doGet("search"), &repos)
	log.Printf("%v repo(s) fetched", repos.NumResults)
	return repos
}

func getTags(name string) map[tag]image {
	log.Printf("Fetching tags for %s ...", name)
	var (
		rawTags tags
		fqTags  tags
	)
	json.Unmarshal(doGet("repositories/"+name+"/tags"), &rawTags)
	log.Printf("%v tags fetched for repo %s", len(rawTags), name)
	fqTags = make(map[tag]image)
	for tag, id := range rawTags {
		fqTags[fqTag(name, tag)] = id
	}
	return fqTags
}

func getAncestry(id image) []image {
	log.Printf("Fetching ancestry for %s ...", id)
	var ancestry []image
	json.Unmarshal(doGet("images/"+string(id)+"/ancestry"), &ancestry)
	log.Printf("%v ancestors fetched for repo %s", len(ancestry), id)
	return ancestry
}

func fqTag(name string, t tag) tag {
	canonicalName := strings.TrimPrefix(name, "library/")
	return tag(canonicalName + ":" + string(t))
}

func printTree(root *imageNode, level int) {
	if len(root.tags) > 0 || len(root.children) > 1 {
		fmt.Printf("%s%v\n", strings.Repeat(" ", level), root.tags)
		level = level + 1
	}
	for _, child := range root.children {
		printTree(child, level)
	}
}

func main() {
	var (
		remaining   int
		tagsCh      = make(chan tags)
		tagsByImage = make(map[image][]tag)
		ancestryCh  = make(chan []image)
		images      = make(map[image]*imageNode)
		roots       []*imageNode
	)
	if len(registryURL) == 0 {
		registryURL = os.Getenv("REGISTRY_URL")
	}
	if len(registryURL) == 0 {
		log.Fatal("No registry URL provided, use the environment variable REGISTRY_URL to set it")
	}
	if len(os.Getenv("REGISTREE_DEBUG")) > 0 {
		log.SetOutput(os.Stderr)
	} else {
		log.SetOutput(ioutil.Discard)
	}
	// get tags in parallel
	for _, repo := range getRepos().Results {
		remaining = remaining + 1
		go func(name string) { tagsCh <- getTags(name) }(repo.Name)
	}
	// group them as they are fetched
	for remaining != 0 {
		for tag, id := range <-tagsCh {
			tags, _ := tagsByImage[id]
			tagsByImage[id] = append(tags, tag)
		}
		remaining = remaining - 1
	}
	// get ancestries in parallel
	for imageId := range tagsByImage {
		remaining = remaining + 1
		go func(id image) { ancestryCh <- getAncestry(id) }(imageId)
	}
	// process them as they arrive
	for remaining != 0 {
		var (
			ancestry     = <-ancestryCh
			previousNode *imageNode
		)
		for _, id := range ancestry {
			if node, ok := images[id]; ok {
				if previousNode != nil {
					node.children = append(node.children, previousNode)
				}
				previousNode = nil
				break
			}
			node := &imageNode{}
			if tags, ok := tagsByImage[id]; ok {
				node.tags = tags
			}
			if previousNode != nil {
				node.children = []*imageNode{previousNode}
			}
			images[id] = node
			previousNode = node
		}
		if previousNode != nil {
			roots = append(roots, previousNode)
		}
		remaining = remaining - 1
	}
	// dump all the trees
	for _, root := range roots {
		printTree(root, 0)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strconv"

	"github.com/jessevdk/go-flags"
	"github.com/yuya-takeyama/argf"
)

func usage() {
	os.Stderr.WriteString(`
Usage: csvp [OPTION]... [FILE]...
Print selected parts of CSV from each FILE to standard output.

Options:
  -i, --indexes=LIST
                 select only these indexes
  -h, --headers=LIST
                 select only these headers
  -D, --output-delimiter=STRING
                 use STRING as the output delimiter (default: \t)
  --help
                 display this help text and exit
  --version
                 output version information and exit
`[1:])
}

func version() {
	os.Stderr.WriteString(`
0.6.0
`[1:])
}

type Option struct {
	IndexesList     string `short:"i" long:"indexes"`
	HeadersList     string `short:"h" long:"headers"`
	Delimiter       string `short:"d" long:"delimiter" default:","`
	OutputDelimiter string `short:"D" long:"output-delimiter" default:"\t"`
	IsHelp          bool   `          long:"help"`
	IsVersion       bool   `          long:"version"`
	Files           []string
}

func parseOption(args []string) (opt *Option, err error) {
	opt = &Option{}
	flag := flags.NewParser(opt, flags.PassDoubleDash)

	opt.Files, err = flag.ParseArgs(args)
	if err != nil {
		return nil, err
	}
	return opt, nil
}

func toDelimiter(s string) (r rune, err error) {
	s, err = strconv.Unquote(`"` + s + `"`)
	if err != nil {
		return 0, err
	}

	runes := []rune(s)
	if len(runes) != 1 {
		return 0, fmt.Errorf("the delimiter must be a single character")
	}
	return runes[0], nil
}

func newCSVScannerFromOption(opt *Option) (c *CSVScanner, err error) {
	var selector Selector
	switch {
	case opt.IndexesList != "" && opt.HeadersList != "":
		return nil, fmt.Errorf("only one type of list may be specified")
	case opt.IndexesList != "":
		selector = NewIndexes(opt.IndexesList)
	case opt.HeadersList != "":
		selector = NewHeaders(opt.HeadersList)
	default:
		selector = NewAll()
	}

	reader, err := argf.From(opt.Files)
	if err != nil {
		return nil, err
	}

	delimiter, err := toDelimiter(opt.Delimiter)
	if err != nil {
		return nil, err
	}

	c = NewCSVScanner(selector, reader)
	c.SetDelimiter(delimiter)
	c.SetOutputDelimiter(opt.OutputDelimiter)
	return c, nil
}

func do(c *CSVScanner) error {
	for c.Scan() {
		fmt.Println(c.Text())
	}
	return c.Err()
}

func printErr(err error) {
	fmt.Fprintln(os.Stderr, "cspv:", err)
}

func guideToHelp() {
	os.Stderr.WriteString(`
Try 'cspv --help' for more information.
`[1:])
}

func _main() int {
	opt, err := parseOption(os.Args[1:])
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	switch {
	case opt.IsHelp:
		usage()
		return 0
	case opt.IsVersion:
		version()
		return 0
	}

	c, err := newCSVScannerFromOption(opt)
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	if err = do(c); err != nil {
		printErr(err)
		return 1
	}
	return 0
}

func main() {
	e := _main()
	os.Exit(e)
}
<commit_msg>Implement -t, --tsv<commit_after>package main

import (
	"fmt"
	"os"
	"strconv"

	"github.com/jessevdk/go-flags"
	"github.com/yuya-takeyama/argf"
)

func usage() {
	os.Stderr.WriteString(`
Usage: csvp [OPTION]... [FILE]...
Print selected parts of CSV from each FILE to standard output.

Options:
  -i, --indexes=LIST
                 select only these indexes
  -h, --headers=LIST
                 select only these headers
  -D, --output-delimiter=STRING
                 use STRING as the output delimiter (default: \t)
  --help
                 display this help text and exit
  --version
                 output version information and exit
`[1:])
}

func version() {
	os.Stderr.WriteString(`
0.6.0
`[1:])
}

type Option struct {
	IndexesList     string `short:"i" long:"indexes"`
	HeadersList     string `short:"h" long:"headers"`
	Delimiter       string `short:"d" long:"delimiter" default:","`
	IsTSV           bool   `short:"t" long:"tsv"`
	OutputDelimiter string `short:"D" long:"output-delimiter" default:"\t"`
	IsHelp          bool   `          long:"help"`
	IsVersion       bool   `          long:"version"`
	Files           []string
}

func parseOption(args []string) (opt *Option, err error) {
	opt = &Option{}
	flag := flags.NewParser(opt, flags.PassDoubleDash)

	opt.Files, err = flag.ParseArgs(args)
	if err != nil {
		return nil, err
	}
	return opt, nil
}

func toDelimiter(s string) (r rune, err error) {
	s, err = strconv.Unquote(`"` + s + `"`)
	if err != nil {
		return 0, err
	}

	runes := []rune(s)
	if len(runes) != 1 {
		return 0, fmt.Errorf("the delimiter must be a single character")
	}
	return runes[0], nil
}

func newCSVScannerFromOption(opt *Option) (c *CSVScanner, err error) {
	var selector Selector
	switch {
	case opt.IndexesList != "" && opt.HeadersList != "":
		return nil, fmt.Errorf("only one type of list may be specified")
	case opt.IndexesList != "":
		selector = NewIndexes(opt.IndexesList)
	case opt.HeadersList != "":
		selector = NewHeaders(opt.HeadersList)
	default:
		selector = NewAll()
	}

	reader, err := argf.From(opt.Files)
	if err != nil {
		return nil, err
	}

	c = NewCSVScanner(selector, reader)
	c.SetOutputDelimiter(opt.OutputDelimiter)
	switch {
	case opt.IsTSV:
		c.SetDelimiter('\t')
	default:
		delimiter, err := toDelimiter(opt.Delimiter)
		if err != nil {
			return nil, err
		}
		c.SetDelimiter(delimiter)
	}
	return c, nil
}

func do(c *CSVScanner) error {
	for c.Scan() {
		fmt.Println(c.Text())
	}
	return c.Err()
}

func printErr(err error) {
	fmt.Fprintln(os.Stderr, "cspv:", err)
}

func guideToHelp() {
	os.Stderr.WriteString(`
Try 'cspv --help' for more information.
`[1:])
}

func _main() int {
	opt, err := parseOption(os.Args[1:])
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	switch {
	case opt.IsHelp:
		usage()
		return 0
	case opt.IsVersion:
		version()
		return 0
	}

	c, err := newCSVScannerFromOption(opt)
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	if err = do(c); err != nil {
		printErr(err)
		return 1
	}
	return 0
}

func main() {
	e := _main()
	os.Exit(e)
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"

	"github.com/boltdb/bolt"
	"github.com/cenkalti/boltbrowser/boltbrowser"
	clog "github.com/cenkalti/log"
	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/console"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/rainrpc"
	"github.com/hokaccha/go-prettyjson"
	"github.com/mitchellh/go-homedir"
	"github.com/urfave/cli"
	"gopkg.in/yaml.v2"
)

var (
	app = cli.NewApp()
	clt *rainrpc.Client
	log = logger.New("rain")
)

func main() {
	app.Version = clientversion.Version
	app.Usage = "BitTorrent client"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "config, c",
			Usage: "read config from `FILE`",
		},
		cli.StringFlag{
			Name:  "cpuprofile",
			Usage: "write cpu profile to `FILE`",
		},
		cli.StringFlag{
			Name:  "memprofile",
			Usage: "write memory profile to `FILE`",
		},
		cli.StringFlag{
			Name:  "pprof",
			Usage: "run pprof server on `ADDR`",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "enable debug log",
		},
		cli.StringFlag{
			Name:  "logfile",
			Usage: "write log to `FILE`",
		},
	}
	app.Before = handleBeforeCommand
	app.After = handleAfterCommand
	app.Commands = []cli.Command{
		{
			Name:  "server",
			Usage: "run rpc server and torrent client",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "config",
					Usage: "path to the config file",
					Value: "~/.rain/config.yaml",
				},
			},
			Action: handleServer,
		},
		{
			Name:  "client",
			Usage: "send rpc request to server",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "url",
					Usage: "URL of RPC server",
					Value: "http://127.0.0.1:" + strconv.Itoa(rainrpc.DefaultServerConfig.Port),
				},
			},
			Before: handleBeforeClient,
			Subcommands: []cli.Command{
				{
					Name:   "list",
					Usage:  "list torrents",
					Action: handleList,
				},
				{
					Name:   "add",
					Usage:  "add torrent or magnet",
					Action: handleAdd,
				},
				{
					Name:   "remove",
					Usage:  "remove torrent",
					Action: handleRemove,
				},
				{
					Name:   "stats",
					Usage:  "get stats of torrent",
					Action: handleStats,
				},
				{
					Name:   "trackers",
					Usage:  "get trackers of torrent",
					Action: handleTrackers,
				},
				{
					Name:   "peers",
					Usage:  "get peers of torrent",
					Action: handlePeers,
				},
				{
					Name:   "start",
					Usage:  "start torrent",
					Action: handleStart,
				},
				{
					Name:   "stop",
					Usage:  "stop",
					Action: handleStop,
				},
				{
					Name:   "console",
					Usage:  "show client console",
					Action: handleConsole,
				},
			},
		},
		{
			Name:   "boltbrowser",
			Hidden: true,
			Action: handleBoltBrowser,
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func handleBoltBrowser(c *cli.Context) error {
	db, err := bolt.Open(c.Args().Get(0), 0600, nil)
	if err != nil {
		return err
	}
	boltbrowser.Browse(db, false)
	return db.Close()
}

func handleBeforeCommand(c *cli.Context) error {
	cpuprofile := c.GlobalString("cpuprofile")
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
	}
	pprofFlag := c.GlobalString("pprof")
	if pprofFlag != "" {
		go func() {
			log.Notice(http.ListenAndServe(pprofFlag, nil))
		}()
	}
	logFile := c.GlobalString("logfile")
	if logFile != "" {
		f, err := os.Create(logFile)
		if err != nil {
			log.Fatal("could not create log file: ", err)
		}
		logger.SetHandler(clog.NewFileHandler(f))
	}
	if c.GlobalBool("debug") {
		logger.SetLevel(clog.DEBUG)
	}
	return nil
}

func handleAfterCommand(c *cli.Context) error {
	if c.GlobalString("cpuprofile") != "" {
		pprof.StopCPUProfile()
	}
	return nil
}

func handleServer(c *cli.Context) error {
	cfg := rainrpc.DefaultServerConfig

	configPath := c.String("config")
	if configPath != "" {
		cp, err := homedir.Expand(configPath)
		if err != nil {
			return err
		}
		b, err := ioutil.ReadFile(cp) // nolint: gosec
		if os.IsNotExist(err) {
			log.Noticef("config file not found at %q, using default config", cp)
		} else if err != nil {
			return err
		} else {
			err = yaml.Unmarshal(b, &cfg)
			if err != nil {
				return err
			}
			log.Infoln("config loaded from:", cp)
			b, err = yaml.Marshal(&cfg)
			if err != nil {
				return err
			}
			log.Debug("\n" + string(b))
		}
	}

	srv, err := rainrpc.NewServer(cfg)
	if err != nil {
		return err
	}
	err = srv.Start()
	if err != nil {
		return err
	}
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	s := <-ch
	log.Noticef("received %s, stopping server", s)

	memprofile := c.GlobalString("memprofile")
	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatal(err)
		}
		err = pprof.WriteHeapProfile(f)
		if err != nil {
			log.Fatal(err)
		}
		err = f.Close()
		if err != nil {
			log.Fatal(err)
		}
	}

	return srv.Stop()
}

func handleBeforeClient(c *cli.Context) error {
	clt = rainrpc.NewClient(c.String("url"))
	return nil
}

func handleList(c *cli.Context) error {
	resp, err := clt.ListTorrents()
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleAdd(c *cli.Context) error {
	var b []byte
	var marshalErr error
	arg := c.Args().Get(0)
	if strings.HasPrefix(arg, "magnet:") {
		resp, err := clt.AddMagnet(arg)
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	} else {
		f, err := os.Open(arg) // nolint: gosec
		if err != nil {
			return err
		}
		resp, err := clt.AddTorrent(f)
		_ = f.Close()
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	}
	if marshalErr != nil {
		return marshalErr
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleRemove(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.RemoveTorrent(id)
	return err
}

func handleStats(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentStats(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleTrackers(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentTrackers(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handlePeers(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentPeers(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleStart(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StartTorrent(id)
	return err
}

func handleStop(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StopTorrent(id)
	return err
}

func handleConsole(c *cli.Context) error {
	con := console.New(clt)
	return con.Run()
}
<commit_msg>enable blocking profiler with flag<commit_after>package main

import (
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"

	"github.com/boltdb/bolt"
	"github.com/cenkalti/boltbrowser/boltbrowser"
	clog "github.com/cenkalti/log"
	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/console"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/rainrpc"
	"github.com/hokaccha/go-prettyjson"
	"github.com/mitchellh/go-homedir"
	"github.com/urfave/cli"
	"gopkg.in/yaml.v2"
)

var (
	app = cli.NewApp()
	clt *rainrpc.Client
	log = logger.New("rain")
)

func main() {
	app.Version = clientversion.Version
	app.Usage = "BitTorrent client"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "config, c",
			Usage: "read config from `FILE`",
		},
		cli.StringFlag{
			Name:  "cpuprofile",
			Usage: "write cpu profile to `FILE`",
		},
		cli.StringFlag{
			Name:  "memprofile",
			Usage: "write memory profile to `FILE`",
		},
		cli.IntFlag{
			Name:  "blockprofile",
			Usage: "enable blocking profiler",
		},
		cli.StringFlag{
			Name:  "pprof",
			Usage: "run pprof server on `ADDR`",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "enable debug log",
		},
		cli.StringFlag{
			Name:  "logfile",
			Usage: "write log to `FILE`",
		},
	}
	app.Before = handleBeforeCommand
	app.After = handleAfterCommand
	app.Commands = []cli.Command{
		{
			Name:  "server",
			Usage: "run rpc server and torrent client",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "config",
					Usage: "path to the config file",
					Value: "~/.rain/config.yaml",
				},
			},
			Action: handleServer,
		},
		{
			Name:  "client",
			Usage: "send rpc request to server",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "url",
					Usage: "URL of RPC server",
					Value: "http://127.0.0.1:" + strconv.Itoa(rainrpc.DefaultServerConfig.Port),
				},
			},
			Before: handleBeforeClient,
			Subcommands: []cli.Command{
				{
					Name:   "list",
					Usage:  "list torrents",
					Action: handleList,
				},
				{
					Name:   "add",
					Usage:  "add torrent or magnet",
					Action: handleAdd,
				},
				{
					Name:   "remove",
					Usage:  "remove torrent",
					Action: handleRemove,
				},
				{
					Name:   "stats",
					Usage:  "get stats of torrent",
					Action: handleStats,
				},
				{
					Name:   "trackers",
					Usage:  "get trackers of torrent",
					Action: handleTrackers,
				},
				{
					Name:   "peers",
					Usage:  "get peers of torrent",
					Action: handlePeers,
				},
				{
					Name:   "start",
					Usage:  "start torrent",
					Action: handleStart,
				},
				{
					Name:   "stop",
					Usage:  "stop",
					Action: handleStop,
				},
				{
					Name:   "console",
					Usage:  "show client console",
					Action: handleConsole,
				},
			},
		},
		{
			Name:   "boltbrowser",
			Hidden: true,
			Action: handleBoltBrowser,
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func handleBoltBrowser(c *cli.Context) error {
	db, err := bolt.Open(c.Args().Get(0), 0600, nil)
	if err != nil {
		return err
	}
	boltbrowser.Browse(db, false)
	return db.Close()
}

func handleBeforeCommand(c *cli.Context) error {
	cpuprofile := c.GlobalString("cpuprofile")
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
	}
	pprofFlag := c.GlobalString("pprof")
	if pprofFlag != "" {
		go func() {
			log.Notice(http.ListenAndServe(pprofFlag, nil))
		}()
	}
	logFile := c.GlobalString("logfile")
	if logFile != "" {
		f, err := os.Create(logFile)
		if err != nil {
			log.Fatal("could not create log file: ", err)
		}
		logger.SetHandler(clog.NewFileHandler(f))
	}
	blockProfile := c.GlobalInt("blockprofile")
	if blockProfile != 0 {
		runtime.SetBlockProfileRate(blockProfile)
	}
	if c.GlobalBool("debug") {
		logger.SetLevel(clog.DEBUG)
	}
	return nil
}

func handleAfterCommand(c *cli.Context) error {
	if c.GlobalString("cpuprofile") != "" {
		pprof.StopCPUProfile()
	}
	return nil
}

func handleServer(c *cli.Context) error {
	cfg := rainrpc.DefaultServerConfig

	configPath := c.String("config")
	if configPath != "" {
		cp, err := homedir.Expand(configPath)
		if err != nil {
			return err
		}
		b, err := ioutil.ReadFile(cp) // nolint: gosec
		if os.IsNotExist(err) {
			log.Noticef("config file not found at %q, using default config", cp)
		} else if err != nil {
			return err
		} else {
			err = yaml.Unmarshal(b, &cfg)
			if err != nil {
				return err
			}
			log.Infoln("config loaded from:", cp)
			b, err = yaml.Marshal(&cfg)
			if err != nil {
				return err
			}
			log.Debug("\n" + string(b))
		}
	}

	srv, err := rainrpc.NewServer(cfg)
	if err != nil {
		return err
	}
	err = srv.Start()
	if err != nil {
		return err
	}
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	s := <-ch
	log.Noticef("received %s, stopping server", s)

	memprofile := c.GlobalString("memprofile")
	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatal(err)
		}
		err = pprof.WriteHeapProfile(f)
		if err != nil {
			log.Fatal(err)
		}
		err = f.Close()
		if err != nil {
			log.Fatal(err)
		}
	}

	return srv.Stop()
}

func handleBeforeClient(c *cli.Context) error {
	clt = rainrpc.NewClient(c.String("url"))
	return nil
}

func handleList(c *cli.Context) error {
	resp, err := clt.ListTorrents()
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleAdd(c *cli.Context) error {
	var b []byte
	var marshalErr error
	arg := c.Args().Get(0)
	if strings.HasPrefix(arg, "magnet:") {
		resp, err := clt.AddMagnet(arg)
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	} else {
		f, err := os.Open(arg) // nolint: gosec
		if err != nil {
			return err
		}
		resp, err := clt.AddTorrent(f)
		_ = f.Close()
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	}
	if marshalErr != nil {
		return marshalErr
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleRemove(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.RemoveTorrent(id)
	return err
}

func handleStats(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentStats(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleTrackers(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentTrackers(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handlePeers(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentPeers(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleStart(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StartTorrent(id)
	return err
}

func handleStop(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StopTorrent(id)
	return err
}

func handleConsole(c *cli.Context) error {
	con := console.New(clt)
	return con.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"syscall"

	"github.com/codegangsta/cli"
)

const (
	VERSION = `0.0.1`
)

func run(c *cli.Context) {
	args := c.Args()
	if len(c.Args()) == 0 {
		cli.ShowAppHelp(c)
		os.Exit(1)
	}

	var cmd *exec.Cmd
	if len(args) > 1 {
		cmd = exec.Command(args[0], args[1:len(args)]...)
	} else {
		cmd = exec.Command(args[0])
	}

	var (
		out bytes.Buffer
		err error
	)

	cmd.Stderr = &out
	if !c.Bool(`suppress-stdout`) {
		cmd.Stdout = &out
	}

	if err = cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "Could not start process: %v", err)
		os.Exit(1)
	}

	err = cmd.Wait()
	// Exited non-zero
	if exitError, ok := err.(*exec.ExitError); ok {
		if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
			exitCode := status.ExitStatus()
			if !intInSlice(exitCode, c.IntSlice(`suppress-exit-code`)) {
				fmt.Fprintf(os.Stderr, "%s", out.String())
				os.Exit(exitCode)
			}
		} else {
			fmt.Fprintf(os.Stderr, "Unhandled error, sending output: %s\n%v", err, out)
			os.Exit(1)
		}
	}
}

func intInSlice(i int, list []int) bool {
	for _, r := range list {
		if r == i {
			return true
		}
	}

	return false
}

func main() {
	app := cli.NewApp()
	app.Name = `crononag`
	app.Usage = `wrap commands running from cron to limit nagging emails`
	app.Version = VERSION
	app.Flags = []cli.Flag{
		cli.IntSliceFlag{
			Name:  `E, suppress-exit-code`,
			Value: &cli.IntSlice{0},
			Usage: `suppress output for specified exit code(s), may specify multiple times, defaults to 0`,
		},
		cli.BoolFlag{
			Name:  `O, suppress-stdout`,
			Usage: `suppress stdout, even on error`,
		},
	}
	app.Action = run
	app.Run(os.Args)
}
<commit_msg>Hide help command, in case someone wants to wrap a command named `help`<commit_after>package main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"syscall"

	"github.com/codegangsta/cli"
)

const (
	VERSION = `0.0.1`
)

func run(c *cli.Context) {
	args := c.Args()
	if len(c.Args()) == 0 {
		cli.ShowAppHelp(c)
		os.Exit(1)
	}

	var cmd *exec.Cmd
	if len(args) > 1 {
		cmd = exec.Command(args[0], args[1:len(args)]...)
	} else {
		cmd = exec.Command(args[0])
	}

	var (
		out bytes.Buffer
		err error
	)

	cmd.Stderr = &out
	if !c.Bool(`suppress-stdout`) {
		cmd.Stdout = &out
	}

	if err = cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "Could not start process: %v", err)
		os.Exit(1)
	}

	err = cmd.Wait()
	// Exited non-zero
	if exitError, ok := err.(*exec.ExitError); ok {
		if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
			exitCode := status.ExitStatus()
			if !intInSlice(exitCode, c.IntSlice(`suppress-exit-code`)) {
				fmt.Fprintf(os.Stderr, "%s", out.String())
				os.Exit(exitCode)
			}
		} else {
			fmt.Fprintf(os.Stderr, "Unhandled error, sending output: %s\n%v", err, out)
			os.Exit(1)
		}
	}
}

func intInSlice(i int, list []int) bool {
	for _, r := range list {
		if r == i {
			return true
		}
	}

	return false
}

func main() {
	app := cli.NewApp()
	app.Name = `crononag`
	app.Usage = `wrap commands running from cron to limit nagging emails`
	app.Version = VERSION
	app.HideHelp = true
	app.Flags = []cli.Flag{
		cli.IntSliceFlag{
			Name:  `E, suppress-exit-code`,
			Value: &cli.IntSlice{0},
			Usage: `suppress output for specified exit code(s), may specify multiple times, defaults to 0`,
		},
		cli.BoolFlag{
			Name:  `O, suppress-stdout`,
			Usage: `suppress stdout, even on error`,
		},
	}
	app.Action = run
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"sort"
	"sync"
)

func main() {
	const (
		oldRev = "HEAD~1"
		newRev = "HEAD"
	)

	var (
		vcs      git
		err      error
		wg       sync.WaitGroup
		newFset  *token.FileSet
		oldFset  *token.FileSet
		newDecls map[string]decls
		oldDecls map[string]decls
	)

	wg.Add(1)
	go func() {
		oldFset, oldDecls, err = parse(vcs, oldRev)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error parsing %s: %s\n", oldRev, err.Error())
			os.Exit(1)
		}
		wg.Done()
	}()

	wg.Add(1)
	go func() {
		newFset, newDecls, err = parse(vcs, newRev)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error parsing %s: %s\n", newRev, err.Error())
			os.Exit(1)
		}
		wg.Done()
	}()

	wg.Wait()

	for pkgName, decls := range oldDecls {
		if _, ok := newDecls[pkgName]; ok {
			err, changes := diff(decls, newDecls[pkgName])
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				if derr, ok := err.(*diffError); ok {
					ast.Fprint(os.Stderr, oldFset, derr.bdecl, ast.NotNilFilter)
					ast.Fprint(os.Stderr, newFset, derr.adecl, ast.NotNilFilter)
				}
				return
			}

			sort.Sort(byID(changes))
			for _, change := range changes {
				fmt.Println(change)
			}
		}
	}
}

func parse(vcs vcs, rev string) (*token.FileSet, map[string]decls, error) {
	files, err := vcs.ReadDir(rev, "")
	if err != nil {
		return nil, nil, err
	}

	fset := token.NewFileSet()
	decls := make(map[string]decls) // package to id to decls
	for _, file := range files {
		contents, err := vcs.ReadFile(rev, file)
		if err != nil {
			return nil, nil, fmt.Errorf("could not read file %s at revision %s: %s", file, rev, err)
		}

		filename := rev + ":" + file
		src, err := parser.ParseFile(fset, filename, contents, 0)
		if err != nil {
			return nil, nil, fmt.Errorf("could not parse file %s at revision %s: %s", file, rev, err)
		}

		pkgName := src.Name.Name

		if decls[pkgName] == nil {
			decls[pkgName] = make(map[string]ast.Decl)
		}
		for id, decl := range getDecls(src.Decls) {
			decls[pkgName][id] = decl
		}
	}

	return fset, decls, nil
}

func getDecls(astDecls []ast.Decl) decls {
	decls := make(map[string]ast.Decl)
	for _, astDecl := range astDecls {
		switch d := astDecl.(type) {
		case *ast.GenDecl:
			for i := range d.Specs {
				var (
					id string
					// gdecl splits declaration blocks into individual declarations to view
					// only changed declarations, instead of all, I don't imagine it's needed
					// for TypeSpec (just ValueSpec
					gdecl *ast.GenDecl
				)
				switch s := d.Specs[i].(type) {
				case *ast.ValueSpec:
					// var / const
					id = s.Names[0].Name
					gdecl = &ast.GenDecl{Tok: d.Tok, Specs: []ast.Spec{s}}
				case *ast.TypeSpec:
					// type struct/interface/etc
					id = s.Name.Name
					gdecl = &ast.GenDecl{Tok: d.Tok, Specs: []ast.Spec{s}}
				default:
					// import or possibly other
					continue
				}
				if ast.IsExported(id) {
					decls[id] = gdecl
				}
			}
		case *ast.FuncDecl:
			// function or method
			var (
				id   string = d.Name.Name
				recv string
			)
			// check if we have a receiver (and not just `func () Method() {}`)
			if d.Recv != nil && len(d.Recv.List) > 0 {
				expr := d.Recv.List[0].Type
				switch e := expr.(type) {
				case *ast.Ident:
					recv = e.Name
				case *ast.StarExpr:
					recv = e.X.(*ast.Ident).Name
				}
				id = recv + "." + id
			}
			// If it's exported and it's either not a receiver OR the receiver is also exported
			if ast.IsExported(id) && recv == "" || ast.IsExported(recv) {
				// We're not interested in the body, nil it, alternatively we could set an
				// Body.List, but that included parenthesis on different lines when printed
				astDecl.(*ast.FuncDecl).Body = nil
				decls[id] = astDecl
			}
		default:
			panic(fmt.Errorf("Unknown decl type: %#v", astDecl))
		}
	}
	return decls
}
<commit_msg>Add profiling to phases<commit_after>package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"sort"
	"sync"
	"time"
)

func main() {
	const (
		oldRev = "HEAD~1"
		newRev = "HEAD"
	)

	var (
		vcs      git
		err      error
		wg       sync.WaitGroup
		newFset  *token.FileSet
		oldFset  *token.FileSet
		newDecls map[string]decls
		oldDecls map[string]decls
	)

	start := time.Now()
	wg.Add(1)
	go func() {
		oldFset, oldDecls, err = parse(vcs, oldRev)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error parsing %s: %s\n", oldRev, err.Error())
			os.Exit(1)
		}
		wg.Done()
	}()

	wg.Add(1)
	go func() {
		newFset, newDecls, err = parse(vcs, newRev)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error parsing %s: %s\n", newRev, err.Error())
			os.Exit(1)
		}
		wg.Done()
	}()

	wg.Wait()
	parseTime := time.Since(start)

	var (
		diffTime time.Duration
		sortTime time.Duration
	)
	for pkgName, decls := range oldDecls {
		if _, ok := newDecls[pkgName]; ok {
			start = time.Now()
			err, changes := diff(decls, newDecls[pkgName])
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				if derr, ok := err.(*diffError); ok {
					ast.Fprint(os.Stderr, oldFset, derr.bdecl, ast.NotNilFilter)
					ast.Fprint(os.Stderr, newFset, derr.adecl, ast.NotNilFilter)
				}
				return
			}
			diffTime += time.Since(start)

			start = time.Now()
			sort.Sort(byID(changes))
			for _, change := range changes {
				fmt.Println(change)
			}
			sortTime += time.Since(start)
		}
	}

	fmt.Printf("Parse time: %v, Diff time: %v, Sort time: %v", parseTime, diffTime, sortTime)
}

func parse(vcs vcs, rev string) (*token.FileSet, map[string]decls, error) {
	files, err := vcs.ReadDir(rev, "")
	if err != nil {
		return nil, nil, err
	}

	fset := token.NewFileSet()
	decls := make(map[string]decls) // package to id to decls
	for _, file := range files {
		contents, err := vcs.ReadFile(rev, file)
		if err != nil {
			return nil, nil, fmt.Errorf("could not read file %s at revision %s: %s", file, rev, err)
		}

		filename := rev + ":" + file
		src, err := parser.ParseFile(fset, filename, contents, 0)
		if err != nil {
			return nil, nil, fmt.Errorf("could not parse file %s at revision %s: %s", file, rev, err)
		}

		pkgName := src.Name.Name

		if decls[pkgName] == nil {
			decls[pkgName] = make(map[string]ast.Decl)
		}
		for id, decl := range getDecls(src.Decls) {
			decls[pkgName][id] = decl
		}
	}

	return fset, decls, nil
}

func getDecls(astDecls []ast.Decl) decls {
	decls := make(map[string]ast.Decl)
	for _, astDecl := range astDecls {
		switch d := astDecl.(type) {
		case *ast.GenDecl:
			for i := range d.Specs {
				var (
					id string
					// gdecl splits declaration blocks into individual declarations to view
					// only changed declarations, instead of all, I don't imagine it's needed
					// for TypeSpec (just ValueSpec
					gdecl *ast.GenDecl
				)
				switch s := d.Specs[i].(type) {
				case *ast.ValueSpec:
					// var / const
					id = s.Names[0].Name
					gdecl = &ast.GenDecl{Tok: d.Tok, Specs: []ast.Spec{s}}
				case *ast.TypeSpec:
					// type struct/interface/etc
					id = s.Name.Name
					gdecl = &ast.GenDecl{Tok: d.Tok, Specs: []ast.Spec{s}}
				default:
					// import or possibly other
					continue
				}
				if ast.IsExported(id) {
					decls[id] = gdecl
				}
			}
		case *ast.FuncDecl:
			// function or method
			var (
				id   string = d.Name.Name
				recv string
			)
			// check if we have a receiver (and not just `func () Method() {}`)
			if d.Recv != nil && len(d.Recv.List) > 0 {
				expr := d.Recv.List[0].Type
				switch e := expr.(type) {
				case *ast.Ident:
					recv = e.Name
				case *ast.StarExpr:
					recv = e.X.(*ast.Ident).Name
				}
				id = recv + "." + id
			}
			// If it's exported and it's either not a receiver OR the receiver is also exported
			if ast.IsExported(id) && recv == "" || ast.IsExported(recv) {
				// We're not interested in the body, nil it, alternatively we could set an
				// Body.List, but that included parenthesis on different lines when printed
				astDecl.(*ast.FuncDecl).Body = nil
				decls[id] = astDecl
			}
		default:
			panic(fmt.Errorf("Unknown decl type: %#v", astDecl))
		}
	}
	return decls
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
)

var (
	listenAddress  = flag.String("listen", ":80", "Server Listen Address")
	sshAddress     = flag.String("ssh", "127.0.0.1:22", "SSH Server Address")
	defaultAddress = flag.String("default", "127.0.0.1:8080", "Default Server Address")
)

func usage() {
	fmt.Println("Switcher 1.0.0")
	fmt.Println("usage: switcher [options]\n")

	fmt.Println("Options:")
	fmt.Println("  --listen   <:80>            Server Listen Address")
	fmt.Println("  --ssh      <127.0.0.1:22>   SSH Server Address")
	fmt.Println("  --default  <127.0.0.1:8080>  Default Server Address\n")

	fmt.Println("Examples:")
	fmt.Println("  To serve SSH(127.0.0.1:22) and HTTP(127.0.0.1:8080) on port 80")
	fmt.Println("  $ switcher\n")

	fmt.Println("  To serve SSH(127.0.0.1:2222) and HTTPS(127.0.0.1:443) on port 443")
	fmt.Println("  $ switcher --listen :443 --ssh 127.0.0.1:2222 --default 127.0.0.1:443")
}

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

	mux := NewMux()

	mux.Handle(SSH(*sshAddress))
	mux.Handle(TCP(*defaultAddress))

	log.Printf("[INFO] listen: %s\n", *listenAddress)
	err := mux.ListenAndServe(*listenAddress)
	if err != nil {
		log.Fatalf("[FATAL] listen: %s\n", err)
	}
}
<commit_msg>Bumped Verson to 1.0.1<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
)

var (
	listenAddress  = flag.String("listen", ":80", "Server Listen Address")
	sshAddress     = flag.String("ssh", "127.0.0.1:22", "SSH Server Address")
	defaultAddress = flag.String("default", "127.0.0.1:8080", "Default Server Address")
)

func usage() {
	fmt.Println("Switcher 1.0.1")
	fmt.Println("usage: switcher [options]\n")

	fmt.Println("Options:")
	fmt.Println("  --listen   <:80>            Server Listen Address")
	fmt.Println("  --ssh      <127.0.0.1:22>   SSH Server Address")
	fmt.Println("  --default  <127.0.0.1:8080>  Default Server Address\n")

	fmt.Println("Examples:")
	fmt.Println("  To serve SSH(127.0.0.1:22) and HTTP(127.0.0.1:8080) on port 80")
	fmt.Println("  $ switcher\n")

	fmt.Println("  To serve SSH(127.0.0.1:2222) and HTTPS(127.0.0.1:443) on port 443")
	fmt.Println("  $ switcher --listen :443 --ssh 127.0.0.1:2222 --default 127.0.0.1:443")
}

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

	mux := NewMux()

	mux.Handle(SSH(*sshAddress))
	mux.Handle(TCP(*defaultAddress))

	log.Printf("[INFO] listen: %s\n", *listenAddress)
	err := mux.ListenAndServe(*listenAddress)
	if err != nil {
		log.Fatalf("[FATAL] listen: %s\n", err)
	}
}
<|endoftext|>