<commit_before>package cmd

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/sahib/brig/cmd/tabwriter"

	"github.com/fatih/color"
	"github.com/sahib/brig/client"
	"github.com/urfave/cli"
)

func handleReset(ctx *cli.Context, ctl *client.Client) error {
	force := ctx.Bool("force")
	rev := ctx.Args().First()
	path := ""

	if len(ctx.Args()) > 1 {
		path = ctx.Args().Get(1)
	}

	if err := ctl.Reset(path, rev, force); err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("reset: %v", err)}
	}

	return nil
}

func commitName(cmt *client.Commit) string {
	if cmt == nil {
		return ""
	}

	if len(cmt.Tags) > 0 {
		return strings.ToUpper(cmt.Tags[0])
	}

	return cmt.Hash.ShortB58()
}

func handleHistory(ctx *cli.Context, ctl *client.Client) error {
	path := ctx.Args().First()

	history, err := ctl.History(path)
	if err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("history: %v", err)}
	}

	if _, err := ctl.Stat(path); err != nil {
		fmt.Printf("%s %s",
			color.YellowString("WARNING:"),
			`This file is not part of this commit, but there's still history for it.
         Most likely this file was moved or removed in the past.

`)
	}

	tabW := tabwriter.NewWriter(
		os.Stdout, 0, 0, 2, ' ',
		tabwriter.StripEscape,
	)

	containsMoves := false
	for _, entry := range history {
		for _, detail := range entry.Mask {
			if detail == "moved" {
				containsMoves = true
				break
			}
		}

		if containsMoves {
			break
		}
	}

	if len(history) != 0 {
		if containsMoves {
			fmt.Fprintf(tabW, "CHANGE\tFROM\tTO\tHOW\tWHEN\t\n")
		} else {
			fmt.Fprintf(tabW, "CHANGE\tFROM\tTO\t\tWHEN\t\n")
		}
	}

	for idx, entry := range history {
		what := ""
		printLine := true

		for _, detail := range entry.Mask {
			// If it was moved, let's display what moved.
			if detail == "moved" && idx+1 < len(history) {
				src := history[idx+1].Path
				dst := entry.Path

				if entry.ReferTo != "" {
					dst = entry.ReferTo
				}

				what = fmt.Sprintf(
					"%s → %s", color.RedString(src), color.RedString(dst),
				)
			}

			// Only display empty changes if nothing happened.
			if detail == "none" && !ctx.Bool("empty") {
				printLine = false
			}
		}
		if !printLine {
			continue
		}

		changeDesc := color.YellowString(strings.Join(entry.Mask, ", "))
		when := color.MagentaString(entry.Head.Date.Format(time.Stamp))

		fmt.Fprintf(
			tabW,
			"%s\t%s\t%s\t%s\t%s\t\n",
			changeDesc,
			color.CyanString(commitName(entry.Next)),
			color.GreenString(commitName(entry.Head)),
			what,
			when,
		)
	}

	return tabW.Flush()
}

// makePathAbbrev tries to abbreviate the `dst` path if
// both are in the same directory.
func makePathAbbrev(src, dst string) string {
	if path.Dir(src) == path.Dir(dst) {
		return path.Base(dst)
	}

	relPath, err := filepath.Rel(path.Dir(src), dst)
	if err != nil {
		fmt.Println("Failed to get relatie path: ", err)
		return dst
	}

	// We could also possibly check here if relPath is longer than dst
	// and only display the relative version then. But being consistent
	// is more valueable here I think.
	return relPath
}

func printDiffTree(diff *client.Diff) {
	const (
		diffTypeNone = iota
		diffTypeAdded
		diffTypeRemoved
		diffTypeMissing
		diffTypeMoved
		diffTypeIgnored
		diffTypeConflict
		diffTypeMerged
	)

	type diffEntry struct {
		typ  int
		pair client.DiffPair
	}

	entries := []client.StatInfo{}
	types := make(map[string]diffEntry)

	// Singular types:
	for _, info := range diff.Added {
		types[info.Path] = diffEntry{typ: diffTypeAdded}
		entries = append(entries, info)
	}
	for _, info := range diff.Removed {
		types[info.Path] = diffEntry{typ: diffTypeRemoved}
		entries = append(entries, info)
	}
	for _, info := range diff.Missing {
		types[info.Path] = diffEntry{typ: diffTypeMissing}
		entries = append(entries, info)
	}
	for _, info := range diff.Ignored {
		types[info.Path] = diffEntry{typ: diffTypeIgnored}
		entries = append(entries, info)
	}

	// Pair types:
	for _, pair := range diff.Moved {
		types[pair.Dst.Path] = diffEntry{
			typ:  diffTypeMoved,
			pair: pair,
		}
		entries = append(entries, pair.Dst)
	}
	for _, pair := range diff.Conflict {
		types[pair.Dst.Path] = diffEntry{
			typ:  diffTypeConflict,
			pair: pair,
		}
		entries = append(entries, pair.Dst)
	}
	for _, pair := range diff.Merged {
		types[pair.Dst.Path] = diffEntry{
			typ:  diffTypeMerged,
			pair: pair,
		}
		entries = append(entries, pair.Dst)
	}

	if len(entries) == 0 {
		// Nothing to show:
		return
	}

	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Path < entries[j].Path
	})

	// Called to format each name in the resulting tree:
	formatter := func(n *treeNode) string {
		if n.name == "/" {
			return color.MagentaString("•")
		}

		if diffEntry, ok := types[n.entry.Path]; ok {
			switch diffEntry.typ {
			case diffTypeAdded:
				return color.GreenString(" + " + n.name)
			case diffTypeRemoved:
				return color.RedString(" - " + n.name)
			case diffTypeMissing:
				return color.MagentaString(" _ " + n.name)
			case diffTypeIgnored:
				return color.YellowString(" * " + n.name)
			case diffTypeMoved:
				dstPath := makePathAbbrev(
					diffEntry.pair.Src.Path,
					diffEntry.pair.Dst.Path,
				)
				name := fmt.Sprintf(
					" %s → %s",
					dstPath,
					path.Base(diffEntry.pair.Src.Path),
				)
				return color.BlueString(name)
			case diffTypeMerged:
				dstPath := makePathAbbrev(
					diffEntry.pair.Src.Path,
					diffEntry.pair.Dst.Path,
				)
				name := fmt.Sprintf(
					" %s ⇄ %s",
					path.Base(diffEntry.pair.Src.Path),
					dstPath,
				)
				return color.CyanString(name)
			case diffTypeConflict:
				dstPath := makePathAbbrev(
					diffEntry.pair.Src.Path,
					diffEntry.pair.Dst.Path,
				)
				name := fmt.Sprintf(
					" %s ⚡ %s",
					path.Base(diffEntry.pair.Src.Path),
					dstPath,
				)
				return color.MagentaString(name)
			}
		}

		return n.name
	}

	// Render the tree:
	showTree(entries, &treeCfg{
		format:  formatter,
		showPin: false,
	})
}

func printDiff(diff *client.Diff) {
	simpleSection := func(heading string, infos []client.StatInfo) {
		if len(infos) == 0 {
			return
		}

		fmt.Println(heading)
		for _, info := range infos {
			fmt.Printf("  %s\n", info.Path)
		}

		fmt.Println()
	}

	pairSection := func(heading string, infos []client.DiffPair) {
		if len(infos) == 0 {
			return
		}

		fmt.Println(heading)
		for _, pair := range infos {
			fmt.Printf("  %s <-> %s\n", pair.Src.Path, pair.Dst.Path)
		}

		fmt.Println()
	}

	simpleSection(color.GreenString("Added:"), diff.Added)
	simpleSection(color.YellowString("Ignored:"), diff.Ignored)
	simpleSection(color.RedString("Removed:"), diff.Removed)
	simpleSection(color.RedString("Missing:"), diff.Missing)

	pairSection(color.BlueString("Moved:"), diff.Moved)
	pairSection(color.CyanString("Resolveable Conflicts:"), diff.Merged)
	pairSection(color.MagentaString("Conflicts:"), diff.Conflict)
}

func handleDiff(ctx *cli.Context, ctl *client.Client) error {
	if ctx.NArg() > 4 {
		fmt.Println("More than four arguments can't be handled.")
	}

	self, err := ctl.Whoami()
	if err != nil {
		return err
	}

	localName := self.CurrentUser
	remoteName := self.CurrentUser

	localRev := "HEAD"
	remoteRev := "CURR"

	switch n := ctx.NArg(); {
	case n >= 4:
		localRev = ctx.Args().Get(3)
		fallthrough
	case n >= 3:
		remoteRev = ctx.Args().Get(2)
		fallthrough
	case n >= 2:
		localName = ctx.Args().Get(1)
		fallthrough
	case n >= 1:
		// When comparing with others,
		// we should always compare with their latest state.
		localRev = "CURR"
		remoteName = ctx.Args().Get(0)
	}

	needFetch := !ctx.Bool("offline")
	diff, err := ctl.MakeDiff(localName, remoteName, localRev, remoteRev, needFetch)
	if err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("diff: %v", err)}
	}

	if ctx.Bool("list") {
		printDiff(diff)
	} else {
		printDiffTree(diff)
	}

	return nil
}

func handleFetch(ctx *cli.Context, ctl *client.Client) error {
	who := ctx.Args().First()
	return ctl.Fetch(who)
}

func handleSync(ctx *cli.Context, ctl *client.Client) error {
	who := ctx.Args().First()

	needFetch := true
	if ctx.Bool("no-fetch") {
		needFetch = false
	}

	return ctl.Sync(who, needFetch)
}

func handleStatus(ctx *cli.Context, ctl *client.Client) error {
	self, err := ctl.Whoami()
	if err != nil {
		return err
	}

	curr := self.CurrentUser
	diff, err := ctl.MakeDiff(curr, curr, "HEAD", "CURR", false)
	if err != nil {
		return err
	}

	if ctx.Bool("tree") {
		printDiffTree(diff)
	} else {
		printDiff(diff)
	}

	return nil
}

func handleBecome(ctx *cli.Context, ctl *client.Client) error {
	becomeSelf := ctx.Bool("self")
	if !becomeSelf && ctx.NArg() < 1 {
		return fmt.Errorf("become needs at least one argument without -s")
	}

	whoami, err := ctl.Whoami()
	if err != nil {
		return err
	}

	who := ctx.Args().First()
	if becomeSelf {
		who = whoami.Owner
	}

	if whoami.CurrentUser == who {
		fmt.Printf("You are already %s.\n", color.GreenString(who))
		return nil
	}

	if err := ctl.Become(who); err != nil {
		return err
	}

	suffix := "Changes will be local only."
	if who == whoami.Owner {
		suffix = "Welcome back!"
	}

	fmt.Printf(
		"You are viewing %s's data now. %s\n",
		color.GreenString(who), suffix,
	)
	return nil
}

func handleCommit(ctx *cli.Context, ctl *client.Client) error {
	msg := ""
	if msg = ctx.String("message"); msg == "" {
		msg = fmt.Sprintf("manual commit")
	}

	if err := ctl.MakeCommit(msg); err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("commit: %v", err)}
	}

	return nil
}

func handleTag(ctx *cli.Context, ctl *client.Client) error {
	if ctx.Bool("delete") {
		name := ctx.Args().Get(0)

		if err := ctl.Untag(name); err != nil {
			return ExitCode{
				UnknownError,
				fmt.Sprintf("untag: %v", err),
			}
		}
	} else {
		if len(ctx.Args()) < 2 {
			return ExitCode{BadArgs, "tag needs at least two arguments"}
		}

		rev := ctx.Args().Get(0)
		name := ctx.Args().Get(1)

		if err := ctl.Tag(rev, name); err != nil {
			return ExitCode{
				UnknownError,
				fmt.Sprintf("tag: %v", err),
			}
		}
	}

	return nil
}

func handleLog(ctx *cli.Context, ctl *client.Client) error {
	entries, err := ctl.Log()
	if err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("commit: %v", err)}
	}

	for _, entry := range entries {
		tags := ""
		if len(entry.Tags) > 0 {
			tags = fmt.Sprintf(" (%s)", strings.Join(entry.Tags, ", "))
		}

		msg := entry.Msg
		if msg == "" {
			msg = color.RedString("•")
		}

		entry.Hash.ShortB58()

		fmt.Printf(
			"%s %s %s%s\n",
			color.GreenString(entry.Hash.ShortB58()),
			color.YellowString(entry.Date.Format(time.Stamp)),
			msg,
			color.CyanString(tags),
		)
	}

	return nil
}
<commit_msg>cmd: diff: fix faulty logic of previous bug fix<commit_after>package cmd

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/sahib/brig/cmd/tabwriter"

	"github.com/fatih/color"
	"github.com/sahib/brig/client"
	"github.com/urfave/cli"
)

func handleReset(ctx *cli.Context, ctl *client.Client) error {
	force := ctx.Bool("force")
	rev := ctx.Args().First()
	path := ""

	if len(ctx.Args()) > 1 {
		path = ctx.Args().Get(1)
	}

	if err := ctl.Reset(path, rev, force); err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("reset: %v", err)}
	}

	return nil
}

func commitName(cmt *client.Commit) string {
	if cmt == nil {
		return ""
	}

	if len(cmt.Tags) > 0 {
		return strings.ToUpper(cmt.Tags[0])
	}

	return cmt.Hash.ShortB58()
}

func handleHistory(ctx *cli.Context, ctl *client.Client) error {
	path := ctx.Args().First()

	history, err := ctl.History(path)
	if err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("history: %v", err)}
	}

	if _, err := ctl.Stat(path); err != nil {
		fmt.Printf("%s %s",
			color.YellowString("WARNING:"),
			`This file is not part of this commit, but there's still history for it.
         Most likely this file was moved or removed in the past.

`)
	}

	tabW := tabwriter.NewWriter(
		os.Stdout, 0, 0, 2, ' ',
		tabwriter.StripEscape,
	)

	containsMoves := false
	for _, entry := range history {
		for _, detail := range entry.Mask {
			if detail == "moved" {
				containsMoves = true
				break
			}
		}

		if containsMoves {
			break
		}
	}

	if len(history) != 0 {
		if containsMoves {
			fmt.Fprintf(tabW, "CHANGE\tFROM\tTO\tHOW\tWHEN\t\n")
		} else {
			fmt.Fprintf(tabW, "CHANGE\tFROM\tTO\t\tWHEN\t\n")
		}
	}

	for idx, entry := range history {
		what := ""
		printLine := true

		for _, detail := range entry.Mask {
			// If it was moved, let's display what moved.
			if detail == "moved" && idx+1 < len(history) {
				src := history[idx+1].Path
				dst := entry.Path

				if entry.ReferTo != "" {
					dst = entry.ReferTo
				}

				what = fmt.Sprintf(
					"%s → %s", color.RedString(src), color.RedString(dst),
				)
			}

			// Only display empty changes if nothing happened.
			if detail == "none" && !ctx.Bool("empty") {
				printLine = false
			}
		}
		if !printLine {
			continue
		}

		changeDesc := color.YellowString(strings.Join(entry.Mask, ", "))
		when := color.MagentaString(entry.Head.Date.Format(time.Stamp))

		fmt.Fprintf(
			tabW,
			"%s\t%s\t%s\t%s\t%s\t\n",
			changeDesc,
			color.CyanString(commitName(entry.Next)),
			color.GreenString(commitName(entry.Head)),
			what,
			when,
		)
	}

	return tabW.Flush()
}

// makePathAbbrev tries to abbreviate the `dst` path if
// both are in the same directory.
func makePathAbbrev(src, dst string) string {
	if path.Dir(src) == path.Dir(dst) {
		return path.Base(dst)
	}

	relPath, err := filepath.Rel(path.Dir(src), dst)
	if err != nil {
		fmt.Println("Failed to get relatie path: ", err)
		return dst
	}

	// We could also possibly check here if relPath is longer than dst
	// and only display the relative version then. But being consistent
	// is more valueable here I think.
	return relPath
}

func printDiffTree(diff *client.Diff) {
	const (
		diffTypeNone = iota
		diffTypeAdded
		diffTypeRemoved
		diffTypeMissing
		diffTypeMoved
		diffTypeIgnored
		diffTypeConflict
		diffTypeMerged
	)

	type diffEntry struct {
		typ  int
		pair client.DiffPair
	}

	entries := []client.StatInfo{}
	types := make(map[string]diffEntry)

	// Singular types:
	for _, info := range diff.Added {
		types[info.Path] = diffEntry{typ: diffTypeAdded}
		entries = append(entries, info)
	}
	for _, info := range diff.Removed {
		types[info.Path] = diffEntry{typ: diffTypeRemoved}
		entries = append(entries, info)
	}
	for _, info := range diff.Missing {
		types[info.Path] = diffEntry{typ: diffTypeMissing}
		entries = append(entries, info)
	}
	for _, info := range diff.Ignored {
		types[info.Path] = diffEntry{typ: diffTypeIgnored}
		entries = append(entries, info)
	}

	// Pair types:
	for _, pair := range diff.Moved {
		types[pair.Dst.Path] = diffEntry{
			typ:  diffTypeMoved,
			pair: pair,
		}
		entries = append(entries, pair.Dst)
	}
	for _, pair := range diff.Conflict {
		types[pair.Dst.Path] = diffEntry{
			typ:  diffTypeConflict,
			pair: pair,
		}
		entries = append(entries, pair.Dst)
	}
	for _, pair := range diff.Merged {
		types[pair.Dst.Path] = diffEntry{
			typ:  diffTypeMerged,
			pair: pair,
		}
		entries = append(entries, pair.Dst)
	}

	if len(entries) == 0 {
		// Nothing to show:
		return
	}

	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Path < entries[j].Path
	})

	// Called to format each name in the resulting tree:
	formatter := func(n *treeNode) string {
		if n.name == "/" {
			return color.MagentaString("•")
		}

		if diffEntry, ok := types[n.entry.Path]; ok {
			switch diffEntry.typ {
			case diffTypeAdded:
				return color.GreenString(" + " + n.name)
			case diffTypeRemoved:
				return color.RedString(" - " + n.name)
			case diffTypeMissing:
				return color.MagentaString(" _ " + n.name)
			case diffTypeIgnored:
				return color.YellowString(" * " + n.name)
			case diffTypeMoved:
				dstPath := makePathAbbrev(
					diffEntry.pair.Src.Path,
					diffEntry.pair.Dst.Path,
				)
				name := fmt.Sprintf(
					" %s → %s",
					dstPath,
					path.Base(diffEntry.pair.Src.Path),
				)
				return color.BlueString(name)
			case diffTypeMerged:
				dstPath := makePathAbbrev(
					diffEntry.pair.Src.Path,
					diffEntry.pair.Dst.Path,
				)
				name := fmt.Sprintf(
					" %s ⇄ %s",
					path.Base(diffEntry.pair.Src.Path),
					dstPath,
				)
				return color.CyanString(name)
			case diffTypeConflict:
				dstPath := makePathAbbrev(
					diffEntry.pair.Src.Path,
					diffEntry.pair.Dst.Path,
				)
				name := fmt.Sprintf(
					" %s ⚡ %s",
					path.Base(diffEntry.pair.Src.Path),
					dstPath,
				)
				return color.MagentaString(name)
			}
		}

		return n.name
	}

	// Render the tree:
	showTree(entries, &treeCfg{
		format:  formatter,
		showPin: false,
	})
}

func printDiff(diff *client.Diff) {
	simpleSection := func(heading string, infos []client.StatInfo) {
		if len(infos) == 0 {
			return
		}

		fmt.Println(heading)
		for _, info := range infos {
			fmt.Printf("  %s\n", info.Path)
		}

		fmt.Println()
	}

	pairSection := func(heading string, infos []client.DiffPair) {
		if len(infos) == 0 {
			return
		}

		fmt.Println(heading)
		for _, pair := range infos {
			fmt.Printf("  %s <-> %s\n", pair.Src.Path, pair.Dst.Path)
		}

		fmt.Println()
	}

	simpleSection(color.GreenString("Added:"), diff.Added)
	simpleSection(color.YellowString("Ignored:"), diff.Ignored)
	simpleSection(color.RedString("Removed:"), diff.Removed)
	simpleSection(color.RedString("Missing:"), diff.Missing)

	pairSection(color.BlueString("Moved:"), diff.Moved)
	pairSection(color.CyanString("Resolveable Conflicts:"), diff.Merged)
	pairSection(color.MagentaString("Conflicts:"), diff.Conflict)
}

func handleDiff(ctx *cli.Context, ctl *client.Client) error {
	if ctx.NArg() > 4 {
		fmt.Println("More than four arguments can't be handled.")
	}

	self, err := ctl.Whoami()
	if err != nil {
		return err
	}

	localName := self.CurrentUser
	remoteName := self.CurrentUser

	remoteRev := "CURR"
	localRev := "CURR"

	if ctx.NArg() == 0 {
		// Special case: When typing brig diff we want to show
		// the diff from our CURR to HEAD only.
		localRev = "HEAD"
	}

	switch n := ctx.NArg(); {
	case n >= 4:
		localRev = ctx.Args().Get(3)
		fallthrough
	case n >= 3:
		remoteRev = ctx.Args().Get(2)
		fallthrough
	case n >= 2:
		localName = ctx.Args().Get(1)
		fallthrough
	case n >= 1:
		remoteName = ctx.Args().Get(0)
	}

	needFetch := !ctx.Bool("offline")
	diff, err := ctl.MakeDiff(localName, remoteName, localRev, remoteRev, needFetch)
	if err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("diff: %v", err)}
	}

	if ctx.Bool("list") {
		printDiff(diff)
	} else {
		printDiffTree(diff)
	}

	return nil
}

func handleFetch(ctx *cli.Context, ctl *client.Client) error {
	who := ctx.Args().First()
	return ctl.Fetch(who)
}

func handleSync(ctx *cli.Context, ctl *client.Client) error {
	who := ctx.Args().First()

	needFetch := true
	if ctx.Bool("no-fetch") {
		needFetch = false
	}

	return ctl.Sync(who, needFetch)
}

func handleStatus(ctx *cli.Context, ctl *client.Client) error {
	self, err := ctl.Whoami()
	if err != nil {
		return err
	}

	curr := self.CurrentUser
	diff, err := ctl.MakeDiff(curr, curr, "HEAD", "CURR", false)
	if err != nil {
		return err
	}

	if ctx.Bool("tree") {
		printDiffTree(diff)
	} else {
		printDiff(diff)
	}

	return nil
}

func handleBecome(ctx *cli.Context, ctl *client.Client) error {
	becomeSelf := ctx.Bool("self")
	if !becomeSelf && ctx.NArg() < 1 {
		return fmt.Errorf("become needs at least one argument without -s")
	}

	whoami, err := ctl.Whoami()
	if err != nil {
		return err
	}

	who := ctx.Args().First()
	if becomeSelf {
		who = whoami.Owner
	}

	if whoami.CurrentUser == who {
		fmt.Printf("You are already %s.\n", color.GreenString(who))
		return nil
	}

	if err := ctl.Become(who); err != nil {
		return err
	}

	suffix := "Changes will be local only."
	if who == whoami.Owner {
		suffix = "Welcome back!"
	}

	fmt.Printf(
		"You are viewing %s's data now. %s\n",
		color.GreenString(who), suffix,
	)
	return nil
}

func handleCommit(ctx *cli.Context, ctl *client.Client) error {
	msg := ""
	if msg = ctx.String("message"); msg == "" {
		msg = fmt.Sprintf("manual commit")
	}

	if err := ctl.MakeCommit(msg); err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("commit: %v", err)}
	}

	return nil
}

func handleTag(ctx *cli.Context, ctl *client.Client) error {
	if ctx.Bool("delete") {
		name := ctx.Args().Get(0)

		if err := ctl.Untag(name); err != nil {
			return ExitCode{
				UnknownError,
				fmt.Sprintf("untag: %v", err),
			}
		}
	} else {
		if len(ctx.Args()) < 2 {
			return ExitCode{BadArgs, "tag needs at least two arguments"}
		}

		rev := ctx.Args().Get(0)
		name := ctx.Args().Get(1)

		if err := ctl.Tag(rev, name); err != nil {
			return ExitCode{
				UnknownError,
				fmt.Sprintf("tag: %v", err),
			}
		}
	}

	return nil
}

func handleLog(ctx *cli.Context, ctl *client.Client) error {
	entries, err := ctl.Log()
	if err != nil {
		return ExitCode{UnknownError, fmt.Sprintf("commit: %v", err)}
	}

	for _, entry := range entries {
		tags := ""
		if len(entry.Tags) > 0 {
			tags = fmt.Sprintf(" (%s)", strings.Join(entry.Tags, ", "))
		}

		msg := entry.Msg
		if msg == "" {
			msg = color.RedString("•")
		}

		entry.Hash.ShortB58()

		fmt.Printf(
			"%s %s %s%s\n",
			color.GreenString(entry.Hash.ShortB58()),
			color.YellowString(entry.Date.Format(time.Stamp)),
			msg,
			color.CyanString(tags),
		)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

func x11(n string, args ...string) error {
	cmd := exec.Command(n, args...)
	cmd.Env = append(os.Environ(), "DISPLAY=:0")
	cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("X11 start %v %v: %v", n, args, err)
	}
	return nil
}

func setup() error {
	if err := os.Symlink("/usr/local/bin/bash", "/bin/bash"); err != nil {
		return err
	}
	if err := os.Symlink("/lib/ld-linux-x86-64.so.2", "/lib64/ld-linux-x86-64.so.2"); err != nil {
		return err
	}
	go func() {
		cmd := exec.Command("Xfbdev")
		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
		if err := cmd.Run(); err != nil {
			log.Fatalf("X11 startup: %v", err)
		}
	}()
	for {
		s, err := filepath.Glob("/tmp/.X*/X?")
		if err != nil {
			return err
		}
		if len(s) > 0 {
			break
		}
		time.Sleep(time.Second)
	}
	for _, f := range []string{"wingo", "flwm", "midori", "opera-12"} {
		log.Printf("Run %v", f)
		go x11(f)
	}
	// we block on the aterm. When the aterm exits, we do too.
	if err := x11("aterm"); err != nil {
		return err
	}
	return nil
}

func main() {
	if err := setup(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>cleanup: xinit<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

func x11(n string, args ...string) error {
	cmd := exec.Command(n, args...)
	cmd.Env = append(os.Environ(), "DISPLAY=:0")
	cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("X11 start %v %v: %v", n, args, err)
	}
	return nil
}

func setup() error {
	if err := os.Symlink("/usr/local/bin/bash", "/bin/bash"); err != nil {
		return err
	}
	if err := os.Symlink("/lib/ld-linux-x86-64.so.2", "/lib64/ld-linux-x86-64.so.2"); err != nil {
		return err
	}
	go func() {
		cmd := exec.Command("Xfbdev")
		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
		if err := cmd.Run(); err != nil {
			log.Fatalf("X11 startup: %v", err)
		}
	}()
	for {
		s, err := filepath.Glob("/tmp/.X*/X?")
		if err != nil {
			return err
		}
		if len(s) > 0 {
			break
		}
		time.Sleep(time.Second)
	}
	for _, f := range []string{"wingo", "flwm", "midori", "opera-12"} {
		log.Printf("Run %v", f)
		go x11(f)
	}

	// we block on the aterm. When the aterm exits, we do too.
	return x11("aterm")
}

func main() {
	if err := setup(); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package peppy

import (
	"database/sql"
	"strconv"
	"strings"

	"git.zxq.co/ripple/rippleapi/common"

	"git.zxq.co/x/getrank"
	"github.com/gin-gonic/gin"
	"github.com/jmoiron/sqlx"
	"gopkg.in/thehowl/go-osuapi.v1"
)

// GetScores retrieve information about the top 100 scores of a specified beatmap.
func GetScores(c *gin.Context, db *sqlx.DB) {
	if c.Query("b") == "" {
		c.JSON(200, defaultResponse)
		return
	}
	var beatmapMD5 string
	err := db.Get(&beatmapMD5, "SELECT beatmap_md5 FROM beatmaps WHERE beatmap_id = ? LIMIT 1", c.Query("b"))
	switch {
	case err == sql.ErrNoRows:
		c.JSON(200, defaultResponse)
		return
	case err != nil:
		c.Error(err)
		c.JSON(200, defaultResponse)
		return
	}
	var sb = "scores.score"
	if rankable(c.Query("m")) {
		sb = "scores.pp"
	}
	var (
		extraWhere  string
		extraParams []interface{}
	)
	if c.Query("u") != "" {
		w, p := genUser(c, db)
		extraWhere = "AND " + w
		extraParams = append(extraParams, p)
	}
	rows, err := db.Query(`
SELECT
	scores.id, scores.score, users.username, scores.300_count, scores.100_count,
	scores.50_count, scores.misses_count, scores.gekis_count, scores.katus_count,
	scores.max_combo, scores.full_combo, scores.mods, users.id, scores.time, scores.pp,
	scores.accuracy
FROM scores
INNER JOIN users ON users.id = scores.userid
WHERE scores.completed = '3'
  AND users.privileges & 1 > 0
  AND scores.beatmap_md5 = ?
  AND scores.play_mode = ?
  `+extraWhere+`
ORDER BY `+sb+` DESC LIMIT `+strconv.Itoa(common.InString(1, c.Query("limit"), 100, 50)),
		append([]interface{}{beatmapMD5, genmodei(c.Query("m"))}, extraParams...)...)
	if err != nil {
		c.Error(err)
		c.JSON(200, defaultResponse)
		return
	}
	var results []osuapi.GSScore
	for rows.Next() {
		var (
			s         osuapi.GSScore
			fullcombo bool
			mods      int
			date      common.UnixTimestamp
			accuracy  float64
		)
		err := rows.Scan(
			&s.ScoreID, &s.Score.Score, &s.Username, &s.Count300, &s.Count100,
			&s.Count50, &s.CountMiss, &s.CountGeki, &s.CountKatu,
			&s.MaxCombo, &fullcombo, &mods, &s.UserID, &date, &s.PP,
			&accuracy,
		)
		if err != nil {
			if err != sql.ErrNoRows {
				c.Error(err)
			}
			continue
		}
		s.FullCombo = osuapi.OsuBool(fullcombo)
		s.Mods = osuapi.Mods(mods)
		s.Date = osuapi.MySQLDate(date)
		s.Rank = strings.ToUpper(getrank.GetRank(osuapi.Mode(genmodei(c.Query("m"))), s.Mods,
			accuracy, s.Count300, s.Count100, s.Count50, s.CountMiss))
		results = append(results, s)
	}
	c.JSON(200, results)
	return
}
<commit_msg>Add mods to get_scores<commit_after>package peppy

import (
	"database/sql"
	"strconv"
	"strings"

	"git.zxq.co/ripple/rippleapi/common"

	"git.zxq.co/x/getrank"
	"github.com/gin-gonic/gin"
	"github.com/jmoiron/sqlx"
	"gopkg.in/thehowl/go-osuapi.v1"
)

// GetScores retrieve information about the top 100 scores of a specified beatmap.
func GetScores(c *gin.Context, db *sqlx.DB) {
	if c.Query("b") == "" {
		c.JSON(200, defaultResponse)
		return
	}
	var beatmapMD5 string
	err := db.Get(&beatmapMD5, "SELECT beatmap_md5 FROM beatmaps WHERE beatmap_id = ? LIMIT 1", c.Query("b"))
	switch {
	case err == sql.ErrNoRows:
		c.JSON(200, defaultResponse)
		return
	case err != nil:
		c.Error(err)
		c.JSON(200, defaultResponse)
		return
	}
	var sb = "scores.score"
	if rankable(c.Query("m")) {
		sb = "scores.pp"
	}
	var (
		extraWhere  string
		extraParams []interface{}
	)
	if c.Query("u") != "" {
		w, p := genUser(c, db)
		extraWhere = "AND " + w
		extraParams = append(extraParams, p)
	}
	mods := common.Int(c.Query("mods"))
	rows, err := db.Query(`
SELECT
	scores.id, scores.score, users.username, scores.300_count, scores.100_count,
	scores.50_count, scores.misses_count, scores.gekis_count, scores.katus_count,
	scores.max_combo, scores.full_combo, scores.mods, users.id, scores.time, scores.pp,
	scores.accuracy
FROM scores
INNER JOIN users ON users.id = scores.userid
WHERE scores.completed = '3'
  AND users.privileges & 1 > 0
  AND scores.beatmap_md5 = ?
  AND scores.play_mode = ?
  AND scores.mods & ? = ?
  `+extraWhere+`
ORDER BY `+sb+` DESC LIMIT `+strconv.Itoa(common.InString(1, c.Query("limit"), 100, 50)),
		append([]interface{}{beatmapMD5, genmodei(c.Query("m")), mods, mods}, extraParams...)...)
	if err != nil {
		c.Error(err)
		c.JSON(200, defaultResponse)
		return
	}
	var results []osuapi.GSScore
	for rows.Next() {
		var (
			s         osuapi.GSScore
			fullcombo bool
			mods      int
			date      common.UnixTimestamp
			accuracy  float64
		)
		err := rows.Scan(
			&s.ScoreID, &s.Score.Score, &s.Username, &s.Count300, &s.Count100,
			&s.Count50, &s.CountMiss, &s.CountGeki, &s.CountKatu,
			&s.MaxCombo, &fullcombo, &mods, &s.UserID, &date, &s.PP,
			&accuracy,
		)
		if err != nil {
			if err != sql.ErrNoRows {
				c.Error(err)
			}
			continue
		}
		s.FullCombo = osuapi.OsuBool(fullcombo)
		s.Mods = osuapi.Mods(mods)
		s.Date = osuapi.MySQLDate(date)
		s.Rank = strings.ToUpper(getrank.GetRank(osuapi.Mode(genmodei(c.Query("m"))), s.Mods,
			accuracy, s.Count300, s.Count100, s.Count50, s.CountMiss))
		results = append(results, s)
	}
	c.JSON(200, results)
	return
}
<|endoftext|>


<commit_before>package handler

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/m-lab/annotation-service/metrics"
	"github.com/m-lab/annotation-service/parser"
	"github.com/m-lab/annotation-service/search"
	"github.com/m-lab/etl/annotation"
)

func init() {
	// Always prepend the filename and line number.
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

// A mutex to make sure that we are not reading from the dataset
// pointer while trying to update it
var currentDataMutex = &sync.RWMutex{}

// This is a pointer to a GeoDataset struct containing the absolute
// latest data for the annotator to search and reply with
var CurrentGeoDataset *parser.GeoDataset = nil

// This is the base in which we should encode the timestamp when we
// are creating the keys for the mapt to return for batch requests
const encodingBase = 36

// A function to set up any handlers that are needed, including url
// handlers and pubsub handlers
func SetupHandlers() {
	http.HandleFunc("/annotate", Annotate)
	http.HandleFunc("/batch_annotate", BatchAnnotate)
	go waitForDownloaderMessages()
}

// Annotate is a URL handler that looks up IP address and puts
// metadata out to the response encoded in json format.
func Annotate(w http.ResponseWriter, r *http.Request) {
	// Setup timers and counters for prometheus metrics.
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.Metrics_requestTimes.Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	metrics.Metrics_activeRequests.Inc()
	metrics.Metrics_totalRequests.Inc()
	defer metrics.Metrics_activeRequests.Dec()

	data, err := ValidateAndParse(r)
	if err != nil {
		fmt.Fprintf(w, "Invalid request")
		return
	}

	result := GetMetadataForSingleIP(data)
	encodedResult, err := json.Marshal(result)
	if err != nil {
		fmt.Fprintf(w, "Unknown JSON Encoding Error")
		return
	}
	fmt.Fprint(w, string(encodedResult))
}

// ValidateAndParse takes a request and validates the URL parameters,
// verifying that it has a valid ip address and time. Then, it uses
// that to construct a RequestData struct and returns the pointer.
func ValidateAndParse(r *http.Request) (*annotation.RequestData, error) {
	query := r.URL.Query()

	time_milli, err := strconv.ParseInt(query.Get("since_epoch"), 10, 64)
	if err != nil {
		return nil, errors.New("Invalid time")
	}

	ip := query.Get("ip_addr")

	newIP := net.ParseIP(ip)
	if newIP == nil {
		return nil, errors.New("Invalid IP address")
	}
	if newIP.To4() != nil {
		return &annotation.RequestData{ip, 4, time.Unix(time_milli, 0)}, nil
	}
	return &annotation.RequestData{ip, 6, time.Unix(time_milli, 0)}, nil
}

// BatchAnnotate is a URL handler that expects the body of the request
// to contain a JSON encoded slice of annotation.RequestDatas. It will
// look up all the ip addresses and bundle them into a map of metadata
// structs (with the keys being the ip concatenated with the base 36
// encoded timestamp) and send them back, again JSON encoded.
func BatchAnnotate(w http.ResponseWriter, r *http.Request) {
	// Setup timers and counters for prometheus metrics.
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.Metrics_requestTimes.Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	metrics.Metrics_activeRequests.Inc()
	metrics.Metrics_totalRequests.Inc()
	defer metrics.Metrics_activeRequests.Dec()

	dataSlice, err := BatchValidateAndParse(r.Body)
	r.Body.Close()

	if err != nil {
		fmt.Println(err)
		fmt.Fprintf(w, "Invalid Request!")
		return
	}

	responseMap := make(map[string]*annotation.GeoData)
	for _, data := range dataSlice {
		responseMap[data.IP+strconv.FormatInt(data.Timestamp.Unix(), encodingBase)] = GetMetadataForSingleIP(&data)
	}
	encodedResult, err := json.Marshal(responseMap)
	if err != nil {
		fmt.Fprintf(w, "Unknown JSON Encoding Error")
		return
	}
	fmt.Fprint(w, string(encodedResult))

}

// BatchValidateAndParse will take a reader (likely the body of a
// request) containing the JSON encoded array of
// annotation.RequestDatas. It will then validate that json and use it to
// construct a slice of annotation.RequestDatas, which it will return. If
// it encounters an error, then it will return nil and that error.
func BatchValidateAndParse(source io.Reader) ([]annotation.RequestData, error) {
	jsonBuffer, err := ioutil.ReadAll(source)
	validatedData := []annotation.RequestData{}
	if err != nil {
		return nil, err
	}
	uncheckedData := []annotation.RequestData{}

	err = json.Unmarshal(jsonBuffer, &uncheckedData)
	if err != nil {
		return nil, err
	}
	for _, data := range uncheckedData {
		newIP := net.ParseIP(data.IP)
		if newIP == nil {
			return nil, errors.New("Invalid IP address.")
		}
		ipType := 6
		if newIP.To4() != nil {
			ipType = 4
		}
		validatedData = append(validatedData, annotation.RequestData{data.IP, ipType, data.Timestamp})
	}
	return validatedData, nil
}

// GetMetadataForSingleIP takes a pointer to a annotation.RequestData
// struct and will use it to fetch the appropriate associated
// metadata, returning a pointer. It is gaurenteed to return a non-nil
// pointer, even if it cannot find the appropriate metadata.
func GetMetadataForSingleIP(request *annotation.RequestData) *annotation.GeoData {
	metrics.Metrics_totalLookups.Inc()
	if CurrentGeoDataset == nil {
		// TODO: Block until the value is not nil
		return nil
	}
	// TODO: Figure out which table to use based on time
	err := errors.New("Unknown IP Format!")
	currentDataMutex.RLock()
	defer currentDataMutex.RUnlock()
	var node parser.IPNode
	// TODO: Push this logic down to searchlist (after binary search is implemented)
	if request.IPFormat == 4 {
		node, err = search.SearchList(CurrentGeoDataset.IP4Nodes, request.IP)
	} else if request.IPFormat == 6 {
		node, err = search.SearchList(CurrentGeoDataset.IP6Nodes, request.IP)
	}

	if err != nil {
		log.Println(err)
		//TODO metric here
		return nil
	}

	return ConvertIPNodeToGeoData(node, CurrentGeoDataset.LocationNodes)
}

// ConvertIPNodeToGeoData takes a parser.IPNode, plus a list of
// locationNodes. It will then use that data to fill in a GeoData
// struct and return its pointer.
func ConvertIPNodeToGeoData(ipNode parser.IPNode, locationNodes []parser.LocationNode) *annotation.GeoData {
	locNode := parser.LocationNode{}
	if ipNode.LocationIndex >= 0 {
		locNode = locationNodes[ipNode.LocationIndex]
	}
	return &annotation.GeoData{
		Geo: &annotation.GeolocationIP{
			Continent_code: locNode.ContinentCode,
			Country_code:   locNode.CountryCode,
			Country_name:   locNode.CountryName,
			Postal_code:    ipNode.PostalCode,
			Metro_code:     locNode.MetroCode,
			City:           locNode.CityName,
			Latitude:       ipNode.Latitude,
			Longitude:      ipNode.Longitude,
		},
		ASN: &annotation.IPASNData{},
	}

}
<commit_msg>use binary search<commit_after>package handler

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/m-lab/annotation-service/metrics"
	"github.com/m-lab/annotation-service/parser"
	"github.com/m-lab/annotation-service/search"
	"github.com/m-lab/etl/annotation"
)

func init() {
	// Always prepend the filename and line number.
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

// A mutex to make sure that we are not reading from the dataset
// pointer while trying to update it
var currentDataMutex = &sync.RWMutex{}

// This is a pointer to a GeoDataset struct containing the absolute
// latest data for the annotator to search and reply with
var CurrentGeoDataset *parser.GeoDataset = nil

// This is the base in which we should encode the timestamp when we
// are creating the keys for the mapt to return for batch requests
const encodingBase = 36

// A function to set up any handlers that are needed, including url
// handlers and pubsub handlers
func SetupHandlers() {
	http.HandleFunc("/annotate", Annotate)
	http.HandleFunc("/batch_annotate", BatchAnnotate)
	go waitForDownloaderMessages()
}

// Annotate is a URL handler that looks up IP address and puts
// metadata out to the response encoded in json format.
func Annotate(w http.ResponseWriter, r *http.Request) {
	// Setup timers and counters for prometheus metrics.
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.Metrics_requestTimes.Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	metrics.Metrics_activeRequests.Inc()
	metrics.Metrics_totalRequests.Inc()
	defer metrics.Metrics_activeRequests.Dec()

	data, err := ValidateAndParse(r)
	if err != nil {
		fmt.Fprintf(w, "Invalid request")
		return
	}

	result := GetMetadataForSingleIP(data)
	encodedResult, err := json.Marshal(result)
	if err != nil {
		fmt.Fprintf(w, "Unknown JSON Encoding Error")
		return
	}
	fmt.Fprint(w, string(encodedResult))
}

// ValidateAndParse takes a request and validates the URL parameters,
// verifying that it has a valid ip address and time. Then, it uses
// that to construct a RequestData struct and returns the pointer.
func ValidateAndParse(r *http.Request) (*annotation.RequestData, error) {
	query := r.URL.Query()

	time_milli, err := strconv.ParseInt(query.Get("since_epoch"), 10, 64)
	if err != nil {
		return nil, errors.New("Invalid time")
	}

	ip := query.Get("ip_addr")

	newIP := net.ParseIP(ip)
	if newIP == nil {
		return nil, errors.New("Invalid IP address")
	}
	if newIP.To4() != nil {
		return &annotation.RequestData{ip, 4, time.Unix(time_milli, 0)}, nil
	}
	return &annotation.RequestData{ip, 6, time.Unix(time_milli, 0)}, nil
}

// BatchAnnotate is a URL handler that expects the body of the request
// to contain a JSON encoded slice of annotation.RequestDatas. It will
// look up all the ip addresses and bundle them into a map of metadata
// structs (with the keys being the ip concatenated with the base 36
// encoded timestamp) and send them back, again JSON encoded.
func BatchAnnotate(w http.ResponseWriter, r *http.Request) {
	// Setup timers and counters for prometheus metrics.
	timerStart := time.Now()
	defer func(tStart time.Time) {
		metrics.Metrics_requestTimes.Observe(float64(time.Since(tStart).Nanoseconds()))
	}(timerStart)
	metrics.Metrics_activeRequests.Inc()
	metrics.Metrics_totalRequests.Inc()
	defer metrics.Metrics_activeRequests.Dec()

	dataSlice, err := BatchValidateAndParse(r.Body)
	r.Body.Close()

	if err != nil {
		fmt.Println(err)
		fmt.Fprintf(w, "Invalid Request!")
		return
	}

	responseMap := make(map[string]*annotation.GeoData)
	for _, data := range dataSlice {
		responseMap[data.IP+strconv.FormatInt(data.Timestamp.Unix(), encodingBase)] = GetMetadataForSingleIP(&data)
	}
	encodedResult, err := json.Marshal(responseMap)
	if err != nil {
		fmt.Fprintf(w, "Unknown JSON Encoding Error")
		return
	}
	fmt.Fprint(w, string(encodedResult))

}

// BatchValidateAndParse will take a reader (likely the body of a
// request) containing the JSON encoded array of
// annotation.RequestDatas. It will then validate that json and use it to
// construct a slice of annotation.RequestDatas, which it will return. If
// it encounters an error, then it will return nil and that error.
func BatchValidateAndParse(source io.Reader) ([]annotation.RequestData, error) {
	jsonBuffer, err := ioutil.ReadAll(source)
	validatedData := []annotation.RequestData{}
	if err != nil {
		return nil, err
	}
	uncheckedData := []annotation.RequestData{}

	err = json.Unmarshal(jsonBuffer, &uncheckedData)
	if err != nil {
		return nil, err
	}
	for _, data := range uncheckedData {
		newIP := net.ParseIP(data.IP)
		if newIP == nil {
			return nil, errors.New("Invalid IP address.")
		}
		ipType := 6
		if newIP.To4() != nil {
			ipType = 4
		}
		validatedData = append(validatedData, annotation.RequestData{data.IP, ipType, data.Timestamp})
	}
	return validatedData, nil
}

// GetMetadataForSingleIP takes a pointer to a annotation.RequestData
// struct and will use it to fetch the appropriate associated
// metadata, returning a pointer. It is gaurenteed to return a non-nil
// pointer, even if it cannot find the appropriate metadata.
func GetMetadataForSingleIP(request *annotation.RequestData) *annotation.GeoData {
	metrics.Metrics_totalLookups.Inc()
	if CurrentGeoDataset == nil {
		// TODO: Block until the value is not nil
		return nil
	}
	// TODO: Figure out which table to use based on time
	err := errors.New("Unknown IP Format!")
	currentDataMutex.RLock()
	defer currentDataMutex.RUnlock()
	var node parser.IPNode
	// TODO: Push this logic down to searchlist (after binary search is implemented)
	if request.IPFormat == 4 {
		node, err = search.SearchBinary(
			CurrentGeoDataset.IP4Nodes, request.IP)
	} else if request.IPFormat == 6 {
		node, err = search.SearchBinary(
			CurrentGeoDataset.IP6Nodes, request.IP)
	}

	if err != nil {
		log.Println(err)
		//TODO metric here
		return nil
	}

	return ConvertIPNodeToGeoData(node, CurrentGeoDataset.LocationNodes)
}

// ConvertIPNodeToGeoData takes a parser.IPNode, plus a list of
// locationNodes. It will then use that data to fill in a GeoData
// struct and return its pointer.
func ConvertIPNodeToGeoData(ipNode parser.IPNode, locationNodes []parser.LocationNode) *annotation.GeoData {
	locNode := parser.LocationNode{}
	if ipNode.LocationIndex >= 0 {
		locNode = locationNodes[ipNode.LocationIndex]
	}
	return &annotation.GeoData{
		Geo: &annotation.GeolocationIP{
			Continent_code: locNode.ContinentCode,
			Country_code:   locNode.CountryCode,
			Country_name:   locNode.CountryName,
			Postal_code:    ipNode.PostalCode,
			Metro_code:     locNode.MetroCode,
			City:           locNode.CityName,
			Latitude:       ipNode.Latitude,
			Longitude:      ipNode.Longitude,
		},
		ASN: &annotation.IPASNData{},
	}

}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"code.google.com/p/go.crypto/scrypt"
	"crypto/aes"
	"crypto/sha256"
	"encoding/hex"
	"github.com/piotrnar/gocoin/btc"
	"log"
	"math/big"
)

func sha256Twice(b []byte) []byte {
	h := sha256.New()
	h.Write(b)
	hashedOnce := h.Sum(nil)
	h.Reset()
	h.Write(hashedOnce)
	return h.Sum(nil)
}

func verifyPassphrase(encryptedKey string, passphrase string) bool {
	dec := btc.Decodeb58(encryptedKey)[:39] // trim to length 39 (not sure why needed)
	if dec == nil {
		log.Fatal("Cannot decode base58 string " + encryptedKey)
	}

	// log.Printf("Decoded base58 string to %s (length %d)", hex.EncodeToString(dec), len(dec))

	if dec[0] == 0x01 && dec[1] == 0x42 {
		log.Print("EC multiply mode not used")
		log.Fatal("TODO: implement decryption when EC multiply mode not used")
	} else if dec[0] == 0x01 && dec[1] == 0x43 {
		// log.Print("EC multiply mode used")

		ownerSalt := dec[7:15]
		hasLotSequence := dec[2]&0x04 == 0x04

		// log.Printf("Owner salt: %s", hex.EncodeToString(ownerSalt))
		// log.Printf("Has lot/sequence: %t", hasLotSequence)

		prefactorA, err := scrypt.Key([]byte(passphrase), ownerSalt, 16384, 8, 8, 32)
		if prefactorA == nil {
			log.Fatal(err)
		}

		var passFactor []byte

		if hasLotSequence {
			prefactorB := bytes.Join([][]byte{prefactorA, ownerSalt}, nil)

			passFactor = sha256Twice(prefactorB)

			lotNumber := int(ownerSalt[4])*4096 + int(ownerSalt[5])*16 + int(ownerSalt[6])/16
			sequenceNumber := int(ownerSalt[6]&0x0f)*256 + int(ownerSalt[7])

			log.Printf("Lot number: %d", lotNumber)
			log.Printf("Sequence number: %d", sequenceNumber)
		} else {
			passFactor = prefactorA
		}

		// log.Printf("passfactor: %s (length %d)", hex.EncodeToString(passFactor), len(passFactor))

		passpoint, err := btc.PublicFromPrivate(passFactor, true)
		if passpoint == nil {
			log.Fatal(err)
		}

		// log.Printf("passpoint: %s", hex.EncodeToString(passpoint))

		encryptedpart1 := dec[15:23]
		encryptedpart2 := dec[23:39]

		addresshashplusownerentropy := bytes.Join([][]byte{dec[3:7], ownerSalt[:8]}, nil)

		derived, err := scrypt.Key(passpoint, addresshashplusownerentropy, 1024, 1, 1, 64)
		if derived == nil {
			log.Fatal(err)
		}

		derivedhalf2 := derived[32:]

		h, err := aes.NewCipher(derivedhalf2)
		if h == nil {
			log.Fatal(err)
		}

		unencryptedpart2 := make([]byte, 16)
		h.Decrypt(unencryptedpart2, encryptedpart2)
		for i := range unencryptedpart2 {
			unencryptedpart2[i] ^= derived[i+16]
		}

		encryptedpart1 = bytes.Join([][]byte{encryptedpart1, unencryptedpart2[:8]}, nil)

		unencryptedpart1 := make([]byte, 16)
		h.Decrypt(unencryptedpart1, encryptedpart1)
		for i := range unencryptedpart1 {
			unencryptedpart1[i] ^= derived[i]
		}

		seeddb := bytes.Join([][]byte{unencryptedpart1[:16], unencryptedpart2[8:]}, nil)

		factorb := sha256Twice(seeddb)

		// log.Printf("passfactor: %s", hex.EncodeToString(passFactor))
		// log.Printf("factorb: %s", hex.EncodeToString(factorb))

		bigN, success := new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141", 16)
		if !success {
			log.Fatal("Failed to create Int for N")
		}

		passFactorBig := new(big.Int).SetBytes(passFactor)
		factorbBig := new(big.Int).SetBytes(factorb)

		privKey := new(big.Int)
		privKey.Mul(passFactorBig, factorbBig)
		privKey.Mod(privKey, bigN)

		pubKey, err := btc.PublicFromPrivate(privKey.Bytes(), false)
		if pubKey == nil {
			log.Fatal(err)
		}

		addr := btc.NewAddrFromPubkey(pubKey, 0).String()

		addrHashed := sha256Twice([]byte(addr))

		if addrHashed[0] != dec[3] || addrHashed[1] != dec[4] || addrHashed[2] != dec[5] || addrHashed[3] != dec[6] {
			return false
		}

		log.Printf("Address: %s", addr)
		log.Printf("Private key: %s", hex.EncodeToString(privKey.Bytes()))
		return true
	}

	log.Fatal("Malformed byte slice")
	return false
}

func main() {
	// what about unicode.PrintRanges?

	printableAscii := " !\"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~€‚ƒ„…†‡ˆ‰Š‹ŒŽ‘’“”•–—˜™š›œžŸ ¡¢£¤¥¦§¨©ª«¬®¯°±²³´µ¶·¸¹º»¼½¾¿ÀÁÂÃÄÅÆÇÈÉÊËÌÍÎÏÐÑÒÓÔÕÖ×ØÙÚÛÜÝÞßàáâãäåæçèéêëìíîïðñòóôõö÷øùúûüýþÿ"
	numTried := 0

	// assumes 3 characters, for http://www.reddit.com/r/Bitcoin/comments/1q5wu7/this_paper_wallet_contains_01125_btc_and_is_bip/
	for _, rune1 := range printableAscii {
		for _, rune2 := range printableAscii {
			for _, rune3 := range printableAscii {
				guess := string(rune1) + string(rune2) + string(rune3)

				if verifyPassphrase("6PfYU8C5sLGsjDNWsCHRYD6G5noFmc184Q4owtnfvXrUdpsfNkeTq2HDV8", guess) {
					return
				}

				if numTried%10 == 0 {
					log.Printf("%d passphrases tried", numTried)
				}

				numTried++
			}
		}
	}
}
<commit_msg>Add backwards printable characters, to search from the other end<commit_after>package main

import (
	"bytes"
	"code.google.com/p/go.crypto/scrypt"
	"crypto/aes"
	"crypto/sha256"
	"encoding/hex"
	"github.com/piotrnar/gocoin/btc"
	"log"
	"math/big"
)

func sha256Twice(b []byte) []byte {
	h := sha256.New()
	h.Write(b)
	hashedOnce := h.Sum(nil)
	h.Reset()
	h.Write(hashedOnce)
	return h.Sum(nil)
}

func verifyPassphrase(encryptedKey string, passphrase string) bool {
	dec := btc.Decodeb58(encryptedKey)[:39] // trim to length 39 (not sure why needed)
	if dec == nil {
		log.Fatal("Cannot decode base58 string " + encryptedKey)
	}

	// log.Printf("Decoded base58 string to %s (length %d)", hex.EncodeToString(dec), len(dec))

	if dec[0] == 0x01 && dec[1] == 0x42 {
		log.Print("EC multiply mode not used")
		log.Fatal("TODO: implement decryption when EC multiply mode not used")
	} else if dec[0] == 0x01 && dec[1] == 0x43 {
		// log.Print("EC multiply mode used")

		ownerSalt := dec[7:15]
		hasLotSequence := dec[2]&0x04 == 0x04

		// log.Printf("Owner salt: %s", hex.EncodeToString(ownerSalt))
		// log.Printf("Has lot/sequence: %t", hasLotSequence)

		prefactorA, err := scrypt.Key([]byte(passphrase), ownerSalt, 16384, 8, 8, 32)
		if prefactorA == nil {
			log.Fatal(err)
		}

		var passFactor []byte

		if hasLotSequence {
			prefactorB := bytes.Join([][]byte{prefactorA, ownerSalt}, nil)

			passFactor = sha256Twice(prefactorB)

			lotNumber := int(ownerSalt[4])*4096 + int(ownerSalt[5])*16 + int(ownerSalt[6])/16
			sequenceNumber := int(ownerSalt[6]&0x0f)*256 + int(ownerSalt[7])

			log.Printf("Lot number: %d", lotNumber)
			log.Printf("Sequence number: %d", sequenceNumber)
		} else {
			passFactor = prefactorA
		}

		// log.Printf("passfactor: %s (length %d)", hex.EncodeToString(passFactor), len(passFactor))

		passpoint, err := btc.PublicFromPrivate(passFactor, true)
		if passpoint == nil {
			log.Fatal(err)
		}

		// log.Printf("passpoint: %s", hex.EncodeToString(passpoint))

		encryptedpart1 := dec[15:23]
		encryptedpart2 := dec[23:39]

		addresshashplusownerentropy := bytes.Join([][]byte{dec[3:7], ownerSalt[:8]}, nil)

		derived, err := scrypt.Key(passpoint, addresshashplusownerentropy, 1024, 1, 1, 64)
		if derived == nil {
			log.Fatal(err)
		}

		derivedhalf2 := derived[32:]

		h, err := aes.NewCipher(derivedhalf2)
		if h == nil {
			log.Fatal(err)
		}

		unencryptedpart2 := make([]byte, 16)
		h.Decrypt(unencryptedpart2, encryptedpart2)
		for i := range unencryptedpart2 {
			unencryptedpart2[i] ^= derived[i+16]
		}

		encryptedpart1 = bytes.Join([][]byte{encryptedpart1, unencryptedpart2[:8]}, nil)

		unencryptedpart1 := make([]byte, 16)
		h.Decrypt(unencryptedpart1, encryptedpart1)
		for i := range unencryptedpart1 {
			unencryptedpart1[i] ^= derived[i]
		}

		seeddb := bytes.Join([][]byte{unencryptedpart1[:16], unencryptedpart2[8:]}, nil)

		factorb := sha256Twice(seeddb)

		// log.Printf("passfactor: %s", hex.EncodeToString(passFactor))
		// log.Printf("factorb: %s", hex.EncodeToString(factorb))

		bigN, success := new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141", 16)
		if !success {
			log.Fatal("Failed to create Int for N")
		}

		passFactorBig := new(big.Int).SetBytes(passFactor)
		factorbBig := new(big.Int).SetBytes(factorb)

		privKey := new(big.Int)
		privKey.Mul(passFactorBig, factorbBig)
		privKey.Mod(privKey, bigN)

		pubKey, err := btc.PublicFromPrivate(privKey.Bytes(), false)
		if pubKey == nil {
			log.Fatal(err)
		}

		addr := btc.NewAddrFromPubkey(pubKey, 0).String()

		addrHashed := sha256Twice([]byte(addr))

		if addrHashed[0] != dec[3] || addrHashed[1] != dec[4] || addrHashed[2] != dec[5] || addrHashed[3] != dec[6] {
			return false
		}

		log.Printf("Address: %s", addr)
		log.Printf("Private key: %s", hex.EncodeToString(privKey.Bytes()))
		return true
	}

	log.Fatal("Malformed byte slice")
	return false
}

func main() {
	// what about unicode.PrintRanges?

	//printableAscii := " !\"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~€‚ƒ„…†‡ˆ‰Š‹ŒŽ‘’“”•–—˜™š›œžŸ ¡¢£¤¥¦§¨©ª«¬®¯°±²³´µ¶·¸¹º»¼½¾¿ÀÁÂÃÄÅÆÇÈÉÊËÌÍÎÏÐÑÒÓÔÕÖ×ØÙÚÛÜÝÞßàáâãäåæçèéêëìíîïðñòóôõö÷øùúûüýþÿ"
	printableAscii := "ÿþýüûúùø÷öõôóòñðïîíìëêéèçæåäãâáàßÞÝÜÛÚÙØ×ÖÕÔÓÒÑÐÏÎÍÌËÊÉÈÇÆÅÄÃÂÁÀ¿¾½¼»º¹¸·¶µ´³²±°¯®¬«ª©¨§¦¥¤£¢¡ Ÿžœ›š™˜—–•”“’‘ŽŒ‹Š‰ˆ‡†…„ƒ‚€~}|{zyxwvutsrqponmlkjihgfedcba`_^]\\[ZYXWVUTSRQPONMLKJIHGFEDCBA@?>=<;:9876543210/.-,+*)('&%$#\"! "
	numTried := 0

	// assumes 3 characters, for http://www.reddit.com/r/Bitcoin/comments/1q5wu7/this_paper_wallet_contains_01125_btc_and_is_bip/
	for _, rune1 := range printableAscii {
		for _, rune2 := range printableAscii {
			for _, rune3 := range printableAscii {
				guess := string(rune1) + string(rune2) + string(rune3)

				if verifyPassphrase("6PfYU8C5sLGsjDNWsCHRYD6G5noFmc184Q4owtnfvXrUdpsfNkeTq2HDV8", guess) {
					return
				}

				if numTried%10 == 0 {
					log.Printf("%d passphrases tried", numTried)
				}

				numTried++
			}
		}
	}
}
<|endoftext|>


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

// Package handler provides functions for responding to specific client
// requests by the ePoxy boot server.
package handler

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"path"
	"strings"
	"time"

	"github.com/gorilla/mux"
	"github.com/m-lab/epoxy/extension"
	"github.com/m-lab/epoxy/metrics"
	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/epoxy/template"
)

// Config provides access to Host records.
type Config interface {
	Save(host *storage.Host) error
	Load(name string) (*storage.Host, error)
}

// Env holds data necessary for executing handler functions.
type Env struct {
	// Config provides access to Host records.
	Config Config
	// ServerAddr is the host:port of the public service. Used to generate absolute URLs.
	ServerAddr string
	// AllowForwardedRequests changes how the ePoxy server evaluates and applies
	// the Host IP whitelist to incoming requests. Typically, the ePoxy server
	// allows an operation when the request "remote address" matches the target Host
	// IP. AllowForwardedRequests should be false unless the ePoxy server runs
	// in a trusted environment like AppEngine. When AllowForwardedRequests is true,
	// then the ePoxy server substitutes the value in the "X-Forwarded-For" request
	// header for the request "remote address".
	AllowForwardedRequests bool
}

var (
	// ErrCannotAccessHost indicates that the request should not be allowed.
	ErrCannotAccessHost = fmt.Errorf("Caller cannot access host")
)

// extractIP parses an "IP:port" string created by the Go http package and
// returns the IP address portion.
func extractIP(remoteAddr string) (string, error) {
	// url.Parse works correctly whether IP is IPv4 or IPv6 formatted.
	u, err := url.Parse("//" + remoteAddr)
	if err != nil {
		return "", err
	}
	return u.Hostname(), nil
}

// requestIsFromHost checks whether the request appears to have originated from the given host.
func (env *Env) requestIsFromHost(req *http.Request, host *storage.Host) error {
	// Typically, booting machines contact the ePoxy server directly. In this
	// case, the req.RemoteAddr contains the IP of the booting machine.
	//
	// However, when the ePoxy server runs in AppEngine, client requests are
	// forwareded by a load balancer, which adds the `X-Forwarded-For` header.
	//
	// Depending on the value of AllowForwardedRequests, we check the X-Forwarded-For
	// header (when true) or the value in RemoteAddr (when false).

	// TODO: allow requests from an administrative network.
	log.Println("Header:", req.Header.Get("X-Forwarded-For"), "vs", host.IPv4Addr)
	log.Println("Header:", req.Header.Get("X-Forwarded-For"), "vs", req.RemoteAddr)

	// Split the header into individual IPs. The first IP is the original client.
	fwdIPs := strings.Split(req.Header.Get("X-Forwarded-For"), ", ")
	// Note: Since this value can be set by the original client, we must check the other IPs.
	// There should be two IPs: one for the original client, and one for the AE load balancer.
	if env.AllowForwardedRequests && len(fwdIPs) <= 2 && fwdIPs[0] == host.IPv4Addr {
		// TODO: verify that fwdIPs[1] is an AppEngine load balancer.
		return nil
	}

	// Check RemoteAddr.
	remoteIP, err := extractIP(req.RemoteAddr)
	if err != nil {
		return ErrCannotAccessHost
	}
	// Check whether remoteIP matches the registered host IPv4Addr.
	if !env.AllowForwardedRequests && (remoteIP == host.IPv4Addr) {
		return nil
	}
	return ErrCannotAccessHost
}

// GenerateStage1IPXE creates the stage1 iPXE script for booting machines.
func (env *Env) GenerateStage1IPXE(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}

	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// TODO(soltesz):
	// * Save information sent in PostForm.

	// Generate new session IDs.
	host.GenerateSessionIDs()

	// Count all requests for the stage1 target.
	metrics.Stage1Total.WithLabelValues(host.Name).Inc()

	// Save host record to Datastore to commit session IDs.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// Generate iPXE script.
	script := template.FormatStage1IPXEScript(host, env.ServerAddr)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "text/plain; charset=us-ascii")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// TODO: consolidate logic in GenerateStage1IPXE & GenerateStage1JSON.

// GenerateStage1JSON creates the stage1 JSON epoxy_client script for booting
// machines.
func (env *Env) GenerateStage1JSON(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}

	// TODO(soltesz): re-enable once PLC updates are completed.
	/* err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}
	*/

	// TODO(soltesz):
	// * Save information sent in PostForm.

	// Generate new session IDs.
	host.GenerateSessionIDs()

	// Count all requests for the stage1 target.
	metrics.Stage1Total.WithLabelValues(host.Name).Inc()

	// Save host record to Datastore to commit session IDs.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// Generate epoxy client JSON action.
	script := template.CreateStage1Action(host, env.ServerAddr)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "application/json; charset=utf-8")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// GenerateJSONConfig creates and returns a JSON serialized nextboot.Config
// suitable for responding to stage2 or stage3 requests.
func (env *Env) GenerateJSONConfig(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]
	// TODO: Verify that the sessionID matches the host.CurrentSessionIDs.Stage2ID.
	// sessionID := mux.Vars(req)["sessionID"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// TODO(soltesz):
	// * Save information sent in PostForm, e.g. ssh host key.
	stage := path.Base(req.URL.Path)

	script := template.FormatJSONConfig(host, stage)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "application/json; charset=utf-8")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// ReceiveReport handles the last step of a boot sequence when the epoxy client reports
// success or failure. In both cases, the session ids are invalidated. In all cases,
// epoxy_client is expected to report the server's public host key.
func (env *Env) ReceiveReport(rw http.ResponseWriter, req *http.Request) {
	// TODO: Verify that the source IP maches the host IP.
	// TODO: log or save values where appropriate.
	req.ParseForm()

	// Use hostname as key to load record from Datastore.
	hostname := mux.Vars(req)["hostname"]
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// TODO(soltesz):

	// Verify sessionID matches the host record (i.e. request is authorized).
	sessionID := mux.Vars(req)["sessionID"]
	if sessionID != host.CurrentSessionIDs.ReportID {
		http.Error(rw, "Given session ID does not match host record", http.StatusForbidden)
		return
	}

	host.LastReport = time.Now()
	status := req.PostForm.Get("message")
	if status == "success" {
		// When the status is success, disable the "update" and mark the time.
		host.LastSuccess = host.LastReport
		host.UpdateEnabled = false
		// TODO: invalidate session ids.
	}

	// Save the new host state.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// TODO: log using structured JSON.
	log.Println(req.PostForm)

	// Report success with no content.
	rw.WriteHeader(http.StatusNoContent)
	return
}

// newReverseProxy creates an httputil.ReverseProxy instance with a custom Director
// that will send the given content to the target.
//
// This implementation differs from the builtin httputil.NewSingleHostReverseProxy by
// modifying how the target URL is treated and by completely overwriting the request
// body. In this way we are restricting how the request is delivered to the
// target and leveraging the response forwarding logic of the httputil.ReverseProxy.
func newReverseProxy(target *url.URL, content string) *httputil.ReverseProxy {
	director := func(req *http.Request) {
		// Overwrite request URL with target, which discards any client query parameters.
		// Overwrite the client request body with the given content.
		// Overwrite the ContentLength to match the given content.
		// Everything else (e.g. original Headers) is unchanged.
		req.URL = target
		req.Body = ioutil.NopCloser(strings.NewReader(content))
		req.ContentLength = int64(len(content))
	}
	return &httputil.ReverseProxy{Director: director}
}

// HandleExtension handles client requests to ePoxy extension URLs. The handler creates
// and sends a request to the extension service registered for the operation.
func (env *Env) HandleExtension(rw http.ResponseWriter, req *http.Request) {
	// Use hostname as key to load record from Datastore.
	hostname := mux.Vars(req)["hostname"]
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// Verify sessionID matches the host record (i.e. request is authorized).
	sessionID := mux.Vars(req)["sessionID"]
	// TODO: support multiple extensions.
	if sessionID != host.CurrentSessionIDs.ExtensionID {
		http.Error(rw, "Given session ID does not match host record", http.StatusForbidden)
		return
	}

	operation := mux.Vars(req)["operation"]
	if len(operation) == 0 {
		http.Error(rw, "Zero length operation is invalid", http.StatusBadRequest)
		return
	}
	// TODO: load extension URL from datastore.
	if _, ok := storage.Extensions[operation]; !ok {
		http.Error(rw, "Unknown Extension for operation: "+operation, http.StatusInternalServerError)
		return
	}

	webreq := extension.Request{
		V1: &extension.V1{
			Hostname:    host.Name,
			IPv4Address: host.IPv4Addr,
			LastBoot:    host.LastSessionCreation,
		},
	}

	extURL, err := url.Parse(storage.Extensions[operation])
	if err != nil {
		http.Error(rw, "Failed to parse extension URL for operation: "+operation, http.StatusInternalServerError)
		return
	}

	// TODO: track metrics about extension requests:
	//  * histogram of request latencies,
	//  * counts status codes,
	srv := newReverseProxy(extURL, webreq.Encode())
	srv.ServeHTTP(rw, req)
}
<commit_msg>Add support for storage proxy<commit_after>// Copyright 2016 ePoxy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//////////////////////////////////////////////////////////////////////////////

// Package handler provides functions for responding to specific client
// requests by the ePoxy boot server.
package handler

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"path"
	"strings"
	"time"

	"github.com/gorilla/mux"
	"github.com/kr/pretty"
	"github.com/m-lab/epoxy/extension"
	"github.com/m-lab/epoxy/metrics"
	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/epoxy/template"
	"github.com/m-lab/go/rtx"
)

// Config provides access to Host records.
type Config interface {
	Save(host *storage.Host) error
	Load(name string) (*storage.Host, error)
}

// Env holds data necessary for executing handler functions.
type Env struct {
	// Config provides access to Host records.
	Config Config
	// ServerAddr is the host:port of the public service. Used to generate absolute URLs.
	ServerAddr string
	// AllowForwardedRequests changes how the ePoxy server evaluates and applies
	// the Host IP whitelist to incoming requests. Typically, the ePoxy server
	// allows an operation when the request "remote address" matches the target Host
	// IP. AllowForwardedRequests should be false unless the ePoxy server runs
	// in a trusted environment like AppEngine. When AllowForwardedRequests is true,
	// then the ePoxy server substitutes the value in the "X-Forwarded-For" request
	// header for the request "remote address".
	AllowForwardedRequests bool
	// Project is the GCP project name in which the server is running.
	Project string
	// StoragePrefixURL is the complete prefix URL for storage proxy requests.
	StoragePrefixURL string
}

var (
	// ErrCannotAccessHost indicates that the request should not be allowed.
	ErrCannotAccessHost = fmt.Errorf("Caller cannot access host")
)

// extractIP parses an "IP:port" string created by the Go http package and
// returns the IP address portion.
func extractIP(remoteAddr string) (string, error) {
	// url.Parse works correctly whether IP is IPv4 or IPv6 formatted.
	u, err := url.Parse("//" + remoteAddr)
	if err != nil {
		return "", err
	}
	return u.Hostname(), nil
}

// requestIsFromHost checks whether the request appears to have originated from the given host.
func (env *Env) requestIsFromHost(req *http.Request, host *storage.Host) error {
	// Typically, booting machines contact the ePoxy server directly. In this
	// case, the req.RemoteAddr contains the IP of the booting machine.
	//
	// However, when the ePoxy server runs in AppEngine, client requests are
	// forwareded by a load balancer, which adds the `X-Forwarded-For` header.
	//
	// Depending on the value of AllowForwardedRequests, we check the X-Forwarded-For
	// header (when true) or the value in RemoteAddr (when false).

	// TODO: allow requests from an administrative network.
	log.Println("Header:", req.Header.Get("X-Forwarded-For"), "vs", host.IPv4Addr)
	log.Println("Header:", req.Header.Get("X-Forwarded-For"), "vs", req.RemoteAddr)

	// Split the header into individual IPs. The first IP is the original client.
	fwdIPs := strings.Split(req.Header.Get("X-Forwarded-For"), ", ")
	// Note: Since this value can be set by the original client, we must check the other IPs.
	// There should be two IPs: one for the original client, and one for the AE load balancer.
	if env.AllowForwardedRequests && len(fwdIPs) <= 2 && fwdIPs[0] == host.IPv4Addr {
		// TODO: verify that fwdIPs[1] is an AppEngine load balancer.
		return nil
	}

	// Check RemoteAddr.
	remoteIP, err := extractIP(req.RemoteAddr)
	if err != nil {
		return ErrCannotAccessHost
	}
	// Check whether remoteIP matches the registered host IPv4Addr.
	if !env.AllowForwardedRequests && (remoteIP == host.IPv4Addr) {
		return nil
	}
	return ErrCannotAccessHost
}

// GenerateStage1IPXE creates the stage1 iPXE script for booting machines.
func (env *Env) GenerateStage1IPXE(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}

	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// TODO(soltesz):
	// * Save information sent in PostForm.

	// Generate new session IDs.
	host.GenerateSessionIDs()

	// Count all requests for the stage1 target.
	metrics.Stage1Total.WithLabelValues(host.Name).Inc()

	// Save host record to Datastore to commit session IDs.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// Generate iPXE script.
	script := template.FormatStage1IPXEScript(host, env.ServerAddr)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "text/plain; charset=us-ascii")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// TODO: consolidate logic in GenerateStage1IPXE & GenerateStage1JSON.

// GenerateStage1JSON creates the stage1 JSON epoxy_client script for booting
// machines.
func (env *Env) GenerateStage1JSON(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}

	// TODO(soltesz): re-enable once PLC updates are completed.
	/* err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}
	*/

	// TODO(soltesz):
	// * Save information sent in PostForm.

	// Generate new session IDs.
	host.GenerateSessionIDs()

	// Count all requests for the stage1 target.
	metrics.Stage1Total.WithLabelValues(host.Name).Inc()

	// Save host record to Datastore to commit session IDs.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// Generate epoxy client JSON action.
	script := template.CreateStage1Action(host, env.ServerAddr)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "application/json; charset=utf-8")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// GenerateJSONConfig creates and returns a JSON serialized nextboot.Config
// suitable for responding to stage2 or stage3 requests.
func (env *Env) GenerateJSONConfig(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]
	// TODO: Verify that the sessionID matches the host.CurrentSessionIDs.Stage2ID.
	// sessionID := mux.Vars(req)["sessionID"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// TODO(soltesz):
	// * Save information sent in PostForm, e.g. ssh host key.
	stage := path.Base(req.URL.Path)

	script := template.FormatJSONConfig(host, stage)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "application/json; charset=utf-8")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// ReceiveReport handles the last step of a boot sequence when the epoxy client reports
// success or failure. In both cases, the session ids are invalidated. In all cases,
// epoxy_client is expected to report the server's public host key.
func (env *Env) ReceiveReport(rw http.ResponseWriter, req *http.Request) {
	// TODO: Verify that the source IP maches the host IP.
	// TODO: log or save values where appropriate.
	req.ParseForm()

	// Use hostname as key to load record from Datastore.
	hostname := mux.Vars(req)["hostname"]
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// TODO(soltesz):

	// Verify sessionID matches the host record (i.e. request is authorized).
	sessionID := mux.Vars(req)["sessionID"]
	if sessionID != host.CurrentSessionIDs.ReportID {
		http.Error(rw, "Given session ID does not match host record", http.StatusForbidden)
		return
	}

	host.LastReport = time.Now()
	status := req.PostForm.Get("message")
	if status == "success" {
		// When the status is success, disable the "update" and mark the time.
		host.LastSuccess = host.LastReport
		host.UpdateEnabled = false
		// TODO: invalidate session ids.
	}

	// Save the new host state.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// TODO: log using structured JSON.
	log.Println(req.PostForm)

	// Report success with no content.
	rw.WriteHeader(http.StatusNoContent)
	return
}

// newReverseProxy creates an httputil.ReverseProxy instance with a custom Director
// that will send the given content to the target.
//
// This implementation differs from the builtin httputil.NewSingleHostReverseProxy by
// modifying how the target URL is treated and by completely overwriting the request
// body. In this way we are restricting how the request is delivered to the
// target and leveraging the response forwarding logic of the httputil.ReverseProxy.
func newReverseProxy(target *url.URL, content string) *httputil.ReverseProxy {
	director := func(req *http.Request) {
		// Overwrite request URL with target, which discards any client query parameters.
		// Overwrite the client request body with the given content.
		// Overwrite the ContentLength to match the given content.
		// Everything else (e.g. original Headers) is unchanged.
		req.URL = target
		req.Body = ioutil.NopCloser(strings.NewReader(content))
		req.ContentLength = int64(len(content))
	}
	return &httputil.ReverseProxy{Director: director}
}

// HandleExtension handles client requests to ePoxy extension URLs. The handler creates
// and sends a request to the extension service registered for the operation.
func (env *Env) HandleExtension(rw http.ResponseWriter, req *http.Request) {
	// Use hostname as key to load record from Datastore.
	hostname := mux.Vars(req)["hostname"]
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	err = env.requestIsFromHost(req, host)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusForbidden)
		return
	}

	// Verify sessionID matches the host record (i.e. request is authorized).
	sessionID := mux.Vars(req)["sessionID"]
	// TODO: support multiple extensions.
	if sessionID != host.CurrentSessionIDs.ExtensionID {
		http.Error(rw, "Given session ID does not match host record", http.StatusForbidden)
		return
	}

	operation := mux.Vars(req)["operation"]
	if len(operation) == 0 {
		http.Error(rw, "Zero length operation is invalid", http.StatusBadRequest)
		return
	}
	// TODO: load extension URL from datastore.
	if _, ok := storage.Extensions[operation]; !ok {
		http.Error(rw, "Unknown Extension for operation: "+operation, http.StatusInternalServerError)
		return
	}

	webreq := extension.Request{
		V1: &extension.V1{
			Hostname:    host.Name,
			IPv4Address: host.IPv4Addr,
			LastBoot:    host.LastSessionCreation,
		},
	}

	extURL, err := url.Parse(storage.Extensions[operation])
	if err != nil {
		http.Error(rw, "Failed to parse extension URL for operation: "+operation, http.StatusInternalServerError)
		return
	}

	// TODO: track metrics about extension requests:
	//  * histogram of request latencies,
	//  * counts status codes,
	srv := newReverseProxy(extURL, webreq.Encode())
	srv.ServeHTTP(rw, req)
}

// newStorageReverseProxy allows GET requests to the epoxy GCS bucket in the
// current project.
func newStorageReverseProxy(storagePrefixURL string) *httputil.ReverseProxy {
	target, err := url.Parse(storagePrefixURL)
	rtx.Must(err, "Failed to parse static GCS URL")

	director := func(req *http.Request) {
		req.Host = target.Host
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = target.Path + "/" + req.URL.Path
		req.URL.RawQuery = "" // Reject any given query parameters.

		if _, ok := req.Header["User-Agent"]; !ok {
			// Explicitly disable User-Agent so it's not set to default value.
			req.Header.Set("User-Agent", "")
		}
		log.Println("proxy", req.Method, req.URL)
		log.Println(req.RemoteAddr, req.Method, req.Host, req.Header, req.RequestURI)
		log.Println(pretty.Sprint(req.URL))
	}
	return &httputil.ReverseProxy{Director: director}
}

// HandleProxy creates a pass-through proxy for GET requests to GCS.
func (env *Env) HandleProxy(rw http.ResponseWriter, req *http.Request) {
	if env.StoragePrefixURL == "" {
		// When no storage prefix url is given, then signal that this is unsupported.
		http.Error(rw, "StoragePrefixURL is not specified", http.StatusNotImplemented)
		return
	}
	if req.Method != http.MethodGet {
		http.Error(rw, "Wrong HTTP method", http.StatusMethodNotAllowed)
		return
	}

	log.Println("orig", req.URL)
	path := mux.Vars(req)["path"]
	log.Println("path", path)
	req.URL.Path = path

	srv := newStorageReverseProxy(env.StoragePrefixURL)
	srv.ServeHTTP(rw, req)
}
<|endoftext|>


<commit_before>package bitmap

import "reflect"

// -----------------------------------------------------------------------------

// TODO
func ValToBitMap(valArrays ...[]interface{}) map[interface{}][]int {
	var bitMap map[interface{}][]int

	for i, valArray := range valArrays {
		for _, val := range valArray {
			if _, ok := bitMap[val]; !ok {
				bitMap[val] = make([]int, len(valArray))
			}
			bitMap[val][i] = 1
		}
	}
	return bitMap
}

// TODO
func appendMatrix(a, b [][]interface{}) [][]interface{} {
	tmp := a[0]

	for i := 0; i < len(b); i++ {
		if i == len(a) {
			a = append(a, tmp)
		}
		a[i] = append(a[i], b[i]...)
	}
	return a
}

// TODO
func concatArrays(a, b [][]interface{}) [][]interface{} {
	var tmp []interface{}

	for i := 0; i < len(a); i++ {
		tmp = append(tmp, a[i]...)
	}
	for j := 0; j < len(b); j++ {
		tmp = append(tmp, b[j]...)
	}
	return append([][]interface{}{}, tmp)
}

// TODO
func Flatten(iface interface{}) [][]interface{} {
	var flatArray [][]interface{}
	concreteVal := reflect.ValueOf(iface)

	switch concreteVal.Kind() {
	case reflect.Ptr:
		return Flatten(concreteVal.Elem().Interface())
	case reflect.Struct:
		for i := 0; i < concreteVal.NumField(); i++ {
			f := concreteVal.Field(i)
			if f.Kind() == reflect.Map {
				flatArray = appendMatrix(flatArray, Flatten(f.Interface()))
			} else {
				flatArray = concatArrays(flatArray, Flatten(f.Interface()))
			}
		}
	case reflect.Map:
		for _, k := range concreteVal.MapKeys() {
			v := concreteVal.MapIndex(k)
			flatArray = append(flatArray, Flatten(v.Interface())...)
		}
	case reflect.Array, reflect.Slice:
		for i := 0; i < concreteVal.Len(); i++ {
			flatArray = append(flatArray, Flatten(concreteVal.Index(i).Interface())...)
		}
	default:
		return [][]interface{}{{concreteVal}}
	}
	return flatArray
}
<commit_msg>update ValToBitMap<commit_after>package bitmap

import "reflect"

// -----------------------------------------------------------------------------

// Bitset is a structure representing a bitmap set.
type Bitset map[string]map[string][]int

// -----------------------------------------------------------------------------

// appendMatrix is a helper function used by `Flatten`.
func appendMatrix(a, b [][]interface{}) [][]interface{} {
	tmp := a[0]

	for i := 0; i < len(b); i++ {
		if i == len(a) {
			a = append(a, tmp)
		}
		a[i] = append(a[i], b[i]...)
	}
	return a
}

// concatArrays is a helper function used by `Flatten`.
func concatArrays(a, b [][]interface{}) [][]interface{} {
	var tmp []interface{}

	for i := 0; i < len(a); i++ {
		tmp = append(tmp, a[i]...)
	}
	for j := 0; j < len(b); j++ {
		tmp = append(tmp, b[j]...)
	}
	return append([][]interface{}{}, tmp)
}

// -----------------------------------------------------------------------------

// ValToBitMap generates a bitmap set from a two-dimensional array.
// Each column name is stored in `index`.
//
// NOTE: `index` lenght MUST be equal or longer than slice's length contained
// in `valArrays`.
func ValToBitMap(valArrays [][]interface{}, index []string) (bm Bitset) {
	bm = make(Bitset)

	for i, valArray := range valArrays {
		for j, val := range valArray {
			colName := index[j]
			v := val.(reflect.Value).Interface().(string)
			if _, ok := bm[colName]; !ok {
				bm[colName] = make(map[string][]int)
			}
			if _, ok := bm[colName][v]; !ok {
				bm[colName][v] = make([]int, len(valArrays))
			}
			bm[colName][v][i] = 1
		}
	}
	return
}

// Flatten generates a two-dimensional array from several Go type, including
// composed types and nested types. The purpose of this array is to be
// converted into a bitmap set.
//
// NOTE: `struct` and `map` types MUST always be the last field of `iface`.
// NOTE: This function panics if a struct's field is not exported.
func Flatten(iface interface{}) [][]interface{} {
	var flatArray [][]interface{}
	concreteVal := reflect.ValueOf(iface)

	switch concreteVal.Kind() {
	case reflect.Ptr:
		return Flatten(concreteVal.Elem().Interface())
	case reflect.Struct:
		for i := 0; i < concreteVal.NumField(); i++ {
			f := concreteVal.Field(i)
			if f.Kind() == reflect.Map {
				flatArray = appendMatrix(flatArray, Flatten(f.Interface()))
			} else {
				flatArray = concatArrays(flatArray, Flatten(f.Interface()))
			}
		}
	case reflect.Map:
		for _, k := range concreteVal.MapKeys() {
			v := concreteVal.MapIndex(k)
			flatArray = append(flatArray, Flatten(v.Interface())...)
		}
	case reflect.Array, reflect.Slice:
		for i := 0; i < concreteVal.Len(); i++ {
			flatArray = append(flatArray, Flatten(concreteVal.Index(i).Interface())...)
		}
	default:
		return [][]interface{}{{concreteVal}}
	}
	return flatArray
}
<|endoftext|>