<commit_before>package net

import (
	"fmt"
	"os"
	"runtime"
	"testing"

	"github.com/shirou/gopsutil/internal/common"
)

func TestAddrString(t *testing.T) {
	v := Addr{IP: "192.168.0.1", Port: 8000}

	s := fmt.Sprintf("%v", v)
	if s != "{\"ip\":\"192.168.0.1\",\"port\":8000}" {
		t.Errorf("Addr string is invalid: %v", v)
	}
}

func TestNetIOCountersStatString(t *testing.T) {
	v := IOCountersStat{
		Name:      "test",
		BytesSent: 100,
	}
	e := `{"name":"test","bytesSent":100,"bytesRecv":0,"packetsSent":0,"packetsRecv":0,"errin":0,"errout":0,"dropin":0,"dropout":0,"fifoin":0,"fifoout":0}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("NetIOCountersStat string is invalid: %v", v)
	}
}

func TestNetProtoCountersStatString(t *testing.T) {
	v := ProtoCountersStat{
		Protocol: "tcp",
		Stats: map[string]int64{
			"MaxConn":      -1,
			"ActiveOpens":  4000,
			"PassiveOpens": 3000,
		},
	}
	e := `{"protocol":"tcp","stats":{"ActiveOpens":4000,"MaxConn":-1,"PassiveOpens":3000}}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("NetProtoCountersStat string is invalid: %v", v)
	}

}

func TestNetConnectionStatString(t *testing.T) {
	v := ConnectionStat{
		Fd:     10,
		Family: 10,
		Type:   10,
		Uids:   []int32{10, 10},
	}
	e := `{"fd":10,"family":10,"type":10,"localaddr":{"ip":"","port":0},"remoteaddr":{"ip":"","port":0},"status":"","uids":[10,10],"pid":0}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("NetConnectionStat string is invalid: %v", v)
	}

}

func TestNetIOCountersAll(t *testing.T) {
	v, err := IOCounters(false)
	per, err := IOCounters(true)
	if err != nil {
		t.Errorf("Could not get NetIOCounters: %v", err)
	}
	if len(v) != 1 {
		t.Errorf("Could not get NetIOCounters: %v", v)
	}
	if v[0].Name != "all" {
		t.Errorf("Invalid NetIOCounters: %v", v)
	}
	var pr uint64
	for _, p := range per {
		pr += p.PacketsRecv
	}
	if v[0].PacketsRecv != pr {
		t.Errorf("invalid sum value: %v, %v", v[0].PacketsRecv, pr)
	}
}

func TestNetIOCountersPerNic(t *testing.T) {
	v, err := IOCounters(true)
	if err != nil {
		t.Errorf("Could not get NetIOCounters: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("Could not get NetIOCounters: %v", v)
	}
	for _, vv := range v {
		if vv.Name == "" {
			t.Errorf("Invalid NetIOCounters: %v", vv)
		}
	}
}

func TestGetNetIOCountersAll(t *testing.T) {
	n := []IOCountersStat{
		{
			Name:        "a",
			BytesRecv:   10,
			PacketsRecv: 10,
		},
		{
			Name:        "b",
			BytesRecv:   10,
			PacketsRecv: 10,
			Errin:       10,
		},
	}
	ret, err := getIOCountersAll(n)
	if err != nil {
		t.Error(err)
	}
	if len(ret) != 1 {
		t.Errorf("invalid return count")
	}
	if ret[0].Name != "all" {
		t.Errorf("invalid return name")
	}
	if ret[0].BytesRecv != 20 {
		t.Errorf("invalid count bytesrecv")
	}
	if ret[0].Errin != 10 {
		t.Errorf("invalid count errin")
	}
}

func TestNetInterfaces(t *testing.T) {
	v, err := Interfaces()
	if err != nil {
		t.Errorf("Could not get NetInterfaceStat: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("Could not get NetInterfaceStat: %v", err)
	}
	for _, vv := range v {
		if vv.Name == "" {
			t.Errorf("Invalid NetInterface: %v", vv)
		}
	}
}

func TestNetProtoCountersStatsAll(t *testing.T) {
	v, err := ProtoCounters(nil)
	if err != nil {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	if len(v) == 0 {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	for _, vv := range v {
		if vv.Protocol == "" {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
		if len(vv.Stats) == 0 {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
	}
}

func TestNetProtoCountersStats(t *testing.T) {
	v, err := ProtoCounters([]string{"tcp", "ip"})
	if err != nil {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	if len(v) == 0 {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	if len(v) != 2 {
		t.Fatalf("Go incorrect number of NetProtoCounters: %v", err)
	}
	for _, vv := range v {
		if vv.Protocol != "tcp" && vv.Protocol != "ip" {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
		if len(vv.Stats) == 0 {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
	}
}

func TestNetConnections(t *testing.T) {
	if ci := os.Getenv("CI"); ci != "" { // skip if test on drone.io
		return
	}

	v, err := Connections("inet")
	if err != nil {
		t.Errorf("could not get NetConnections: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("could not get NetConnections: %v", v)
	}
	for _, vv := range v {
		if vv.Family == 0 {
			t.Errorf("invalid NetConnections: %v", vv)
		}
	}

}

func TestNetFilterCounters(t *testing.T) {
	if ci := os.Getenv("CI"); ci != "" { // skip if test on drone.io
		return
	}

	if runtime.GOOS == "linux" {
		// some test environment has not the path.
		if !common.PathExists("/proc/sys/net/netfilter/nf_conntrackCount") {
			t.SkipNow()
		}
	}

	v, err := FilterCounters()
	if err != nil {
		t.Errorf("could not get NetConnections: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("could not get NetConnections: %v", v)
	}
	for _, vv := range v {
		if vv.ConnTrackMax == 0 {
			t.Errorf("nf_conntrackMax needs to be greater than zero: %v", vv)
		}
	}

}
<commit_msg>net: Check for error object<commit_after>package net

import (
	"fmt"
	"os"
	"runtime"
	"testing"

	"github.com/shirou/gopsutil/internal/common"
)

func TestAddrString(t *testing.T) {
	v := Addr{IP: "192.168.0.1", Port: 8000}

	s := fmt.Sprintf("%v", v)
	if s != "{\"ip\":\"192.168.0.1\",\"port\":8000}" {
		t.Errorf("Addr string is invalid: %v", v)
	}
}

func TestNetIOCountersStatString(t *testing.T) {
	v := IOCountersStat{
		Name:      "test",
		BytesSent: 100,
	}
	e := `{"name":"test","bytesSent":100,"bytesRecv":0,"packetsSent":0,"packetsRecv":0,"errin":0,"errout":0,"dropin":0,"dropout":0,"fifoin":0,"fifoout":0}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("NetIOCountersStat string is invalid: %v", v)
	}
}

func TestNetProtoCountersStatString(t *testing.T) {
	v := ProtoCountersStat{
		Protocol: "tcp",
		Stats: map[string]int64{
			"MaxConn":      -1,
			"ActiveOpens":  4000,
			"PassiveOpens": 3000,
		},
	}
	e := `{"protocol":"tcp","stats":{"ActiveOpens":4000,"MaxConn":-1,"PassiveOpens":3000}}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("NetProtoCountersStat string is invalid: %v", v)
	}

}

func TestNetConnectionStatString(t *testing.T) {
	v := ConnectionStat{
		Fd:     10,
		Family: 10,
		Type:   10,
		Uids:   []int32{10, 10},
	}
	e := `{"fd":10,"family":10,"type":10,"localaddr":{"ip":"","port":0},"remoteaddr":{"ip":"","port":0},"status":"","uids":[10,10],"pid":0}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("NetConnectionStat string is invalid: %v", v)
	}

}

func TestNetIOCountersAll(t *testing.T) {
	v, err := IOCounters(false)
	if err != nil {
		t.Errorf("Could not get NetIOCounters: %v", err)
	}
	per, err := IOCounters(true)
	if err != nil {
		t.Errorf("Could not get NetIOCounters: %v", err)
	}
	if len(v) != 1 {
		t.Errorf("Could not get NetIOCounters: %v", v)
	}
	if v[0].Name != "all" {
		t.Errorf("Invalid NetIOCounters: %v", v)
	}
	var pr uint64
	for _, p := range per {
		pr += p.PacketsRecv
	}
	if v[0].PacketsRecv != pr {
		t.Errorf("invalid sum value: %v, %v", v[0].PacketsRecv, pr)
	}
}

func TestNetIOCountersPerNic(t *testing.T) {
	v, err := IOCounters(true)
	if err != nil {
		t.Errorf("Could not get NetIOCounters: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("Could not get NetIOCounters: %v", v)
	}
	for _, vv := range v {
		if vv.Name == "" {
			t.Errorf("Invalid NetIOCounters: %v", vv)
		}
	}
}

func TestGetNetIOCountersAll(t *testing.T) {
	n := []IOCountersStat{
		{
			Name:        "a",
			BytesRecv:   10,
			PacketsRecv: 10,
		},
		{
			Name:        "b",
			BytesRecv:   10,
			PacketsRecv: 10,
			Errin:       10,
		},
	}
	ret, err := getIOCountersAll(n)
	if err != nil {
		t.Error(err)
	}
	if len(ret) != 1 {
		t.Errorf("invalid return count")
	}
	if ret[0].Name != "all" {
		t.Errorf("invalid return name")
	}
	if ret[0].BytesRecv != 20 {
		t.Errorf("invalid count bytesrecv")
	}
	if ret[0].Errin != 10 {
		t.Errorf("invalid count errin")
	}
}

func TestNetInterfaces(t *testing.T) {
	v, err := Interfaces()
	if err != nil {
		t.Errorf("Could not get NetInterfaceStat: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("Could not get NetInterfaceStat: %v", err)
	}
	for _, vv := range v {
		if vv.Name == "" {
			t.Errorf("Invalid NetInterface: %v", vv)
		}
	}
}

func TestNetProtoCountersStatsAll(t *testing.T) {
	v, err := ProtoCounters(nil)
	if err != nil {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	if len(v) == 0 {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	for _, vv := range v {
		if vv.Protocol == "" {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
		if len(vv.Stats) == 0 {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
	}
}

func TestNetProtoCountersStats(t *testing.T) {
	v, err := ProtoCounters([]string{"tcp", "ip"})
	if err != nil {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	if len(v) == 0 {
		t.Fatalf("Could not get NetProtoCounters: %v", err)
	}
	if len(v) != 2 {
		t.Fatalf("Go incorrect number of NetProtoCounters: %v", err)
	}
	for _, vv := range v {
		if vv.Protocol != "tcp" && vv.Protocol != "ip" {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
		if len(vv.Stats) == 0 {
			t.Errorf("Invalid NetProtoCountersStat: %v", vv)
		}
	}
}

func TestNetConnections(t *testing.T) {
	if ci := os.Getenv("CI"); ci != "" { // skip if test on drone.io
		return
	}

	v, err := Connections("inet")
	if err != nil {
		t.Errorf("could not get NetConnections: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("could not get NetConnections: %v", v)
	}
	for _, vv := range v {
		if vv.Family == 0 {
			t.Errorf("invalid NetConnections: %v", vv)
		}
	}

}

func TestNetFilterCounters(t *testing.T) {
	if ci := os.Getenv("CI"); ci != "" { // skip if test on drone.io
		return
	}

	if runtime.GOOS == "linux" {
		// some test environment has not the path.
		if !common.PathExists("/proc/sys/net/netfilter/nf_conntrackCount") {
			t.SkipNow()
		}
	}

	v, err := FilterCounters()
	if err != nil {
		t.Errorf("could not get NetConnections: %v", err)
	}
	if len(v) == 0 {
		t.Errorf("could not get NetConnections: %v", v)
	}
	for _, vv := range v {
		if vv.ConnTrackMax == 0 {
			t.Errorf("nf_conntrackMax needs to be greater than zero: %v", vv)
		}
	}

}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/mackerel-agent/logging"
)

var logger = logging.GetLogger("metrics.plugin.inode")

// InodePlugin plugin
type InodePlugin struct{}

var dfHeaderPattern = regexp.MustCompile(
	`^Filesystem\s+`,
)

var dfColumnsPattern = regexp.MustCompile(
	`^(.+?)\s+(?:(?:\d+)\s+(?:\d+)\s+(?:\d+)\s+(?:\d+%)\s+|(?:\d+)\s+)?(\d+)\s+(\d+)\s+(\d+%)\s+(.+)$`,
)

var devicePattern = regexp.MustCompile(
	`^/dev/(.*)$`,
)

var deviceUnacceptablePattern = regexp.MustCompile(
	`[^A-Za-z0-9_-]`,
)

//  $ df -i
// Filesystem      Inodes  IUsed   IFree IUse% Mounted on
// /dev/xvda1     1310720 131197 1179523   11% /
//  $ df -i # on Mac OSX (impossible to display only inode information)
// Filesystem 512-blocks      Used Available Capacity  iused    ifree %iused  Mounted on
// /dev/disk1  974737408 176727800 797497608    19% 22154973 99687201   18%   /

// FetchMetrics interface for mackerelplugin
func (p InodePlugin) FetchMetrics() (map[string]interface{}, error) {
	cmd := exec.Command("df", "-i")
	cmd.Env = append(os.Environ(), "LANG=C")
	out, err := cmd.Output()
	if err != nil {
		logger.Warningf("'df -i' command exited with a non-zero status: '%s'", err)
		return nil, err
	}
	result := make(map[string]interface{})
	for _, line := range strings.Split(string(out), "\n") {
		if dfHeaderPattern.MatchString(line) {
			continue
		} else if matches := dfColumnsPattern.FindStringSubmatch(line); matches != nil {
			name := matches[1]
			// https://github.com/docker/docker/blob/v1.5.0/daemon/graphdriver/devmapper/deviceset.go#L981
			if regexp.MustCompile(`^/dev/mapper/docker-`).FindStringSubmatch(name) != nil {
				continue
			}
			if nameMatches := devicePattern.FindStringSubmatch(name); nameMatches != nil {
				device := deviceUnacceptablePattern.ReplaceAllString(nameMatches[1], "_")
				iused, err := strconv.ParseInt(matches[2], 0, 64)
				if err != nil {
					logger.Warningf("Failed to parse value: [%s]", matches[2])
					continue
				}
				ifree, err := strconv.ParseInt(matches[3], 0, 64)
				if err != nil {
					logger.Warningf("Failed to parse value: [%s]", matches[3])
					continue
				}
				result["inode.count."+device+".used"] = uint64(iused)
				result["inode.count."+device+".free"] = uint64(ifree)
				result["inode.count."+device+".total"] = uint64(iused + ifree)
				usedPercentage := 100.0 // 100% if both iused and ifree are 0
				if iused+ifree > 0 {
					usedPercentage = float64(iused) * 100 / float64(iused+ifree)
				}
				result["inode.percentage."+device+".used"] = usedPercentage
			}
		}
	}
	return result, nil
}

// GraphDefinition interface for mackerelplugin
func (p InodePlugin) GraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		"inode.count.#": mp.Graphs{
			Label: "Inode",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				mp.Metrics{Name: "used", Label: "used"},
				mp.Metrics{Name: "free", Label: "free"},
				mp.Metrics{Name: "total", Label: "total"},
			},
		},
		"inode.percentage.#": mp.Graphs{
			Label: "Inode Percentage",
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				mp.Metrics{Name: "used", Label: "used %"},
			},
		},
	}
}

func main() {
	flag.Parse()
	inode := InodePlugin{}
	helper := mp.NewMackerelPlugin(inode)
	helper.Run()
}
<commit_msg>improve inode plugin: this plugin takes no flags<commit_after>package main

import (
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/mackerel-agent/logging"
)

var logger = logging.GetLogger("metrics.plugin.inode")

// InodePlugin plugin
type InodePlugin struct{}

var dfHeaderPattern = regexp.MustCompile(
	`^Filesystem\s+`,
)

var dfColumnsPattern = regexp.MustCompile(
	`^(.+?)\s+(?:(?:\d+)\s+(?:\d+)\s+(?:\d+)\s+(?:\d+%)\s+|(?:\d+)\s+)?(\d+)\s+(\d+)\s+(\d+%)\s+(.+)$`,
)

var devicePattern = regexp.MustCompile(
	`^/dev/(.*)$`,
)

var deviceUnacceptablePattern = regexp.MustCompile(
	`[^A-Za-z0-9_-]`,
)

//  $ df -i
// Filesystem      Inodes  IUsed   IFree IUse% Mounted on
// /dev/xvda1     1310720 131197 1179523   11% /
//  $ df -i # on Mac OSX (impossible to display only inode information)
// Filesystem 512-blocks      Used Available Capacity  iused    ifree %iused  Mounted on
// /dev/disk1  974737408 176727800 797497608    19% 22154973 99687201   18%   /

// FetchMetrics interface for mackerelplugin
func (p InodePlugin) FetchMetrics() (map[string]interface{}, error) {
	cmd := exec.Command("df", "-i")
	cmd.Env = append(os.Environ(), "LANG=C")
	out, err := cmd.Output()
	if err != nil {
		logger.Warningf("'df -i' command exited with a non-zero status: '%s'", err)
		return nil, err
	}
	result := make(map[string]interface{})
	for _, line := range strings.Split(string(out), "\n") {
		if dfHeaderPattern.MatchString(line) {
			continue
		} else if matches := dfColumnsPattern.FindStringSubmatch(line); matches != nil {
			name := matches[1]
			// https://github.com/docker/docker/blob/v1.5.0/daemon/graphdriver/devmapper/deviceset.go#L981
			if regexp.MustCompile(`^/dev/mapper/docker-`).FindStringSubmatch(name) != nil {
				continue
			}
			if nameMatches := devicePattern.FindStringSubmatch(name); nameMatches != nil {
				device := deviceUnacceptablePattern.ReplaceAllString(nameMatches[1], "_")
				iused, err := strconv.ParseInt(matches[2], 0, 64)
				if err != nil {
					logger.Warningf("Failed to parse value: [%s]", matches[2])
					continue
				}
				ifree, err := strconv.ParseInt(matches[3], 0, 64)
				if err != nil {
					logger.Warningf("Failed to parse value: [%s]", matches[3])
					continue
				}
				result["inode.count."+device+".used"] = uint64(iused)
				result["inode.count."+device+".free"] = uint64(ifree)
				result["inode.count."+device+".total"] = uint64(iused + ifree)
				usedPercentage := 100.0 // 100% if both iused and ifree are 0
				if iused+ifree > 0 {
					usedPercentage = float64(iused) * 100 / float64(iused+ifree)
				}
				result["inode.percentage."+device+".used"] = usedPercentage
			}
		}
	}
	return result, nil
}

// GraphDefinition interface for mackerelplugin
func (p InodePlugin) GraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		"inode.count.#": mp.Graphs{
			Label: "Inode",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				mp.Metrics{Name: "used", Label: "used"},
				mp.Metrics{Name: "free", Label: "free"},
				mp.Metrics{Name: "total", Label: "total"},
			},
		},
		"inode.percentage.#": mp.Graphs{
			Label: "Inode Percentage",
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				mp.Metrics{Name: "used", Label: "used %"},
			},
		},
	}
}

func main() {
	inode := InodePlugin{}
	helper := mp.NewMackerelPlugin(inode)
	helper.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv6"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/kexec"
)

var (
	useV4              = flag.Bool("4", false, "Get a DHCPv4 lease")
	useV6              = flag.Bool("6", true, "Get a DHCPv6 lease")
	ifname             = flag.String("i", "eth0", "Interface to send packets through")
	dryRun             = flag.Bool("dryrun", false, "Do everything except assigning IP addresses, changing DNS, and kexec")
	doDebug            = flag.Bool("d", false, "Print debug output")
	skipDHCP           = flag.Bool("skip-dhcp", false, "Skip DHCP and rely on SLAAC for network configuration. This requires -netboot-url")
	overrideNetbootURL = flag.String("netboot-url", "", "Override the netboot URL normally obtained via DHCP")
	readTimeout        = flag.Int("timeout", 3, "Read timeout in seconds")
	dhcpRetries        = flag.Int("retries", 3, "Number of times a DHCP request is retried")
	userClass          = flag.String("userclass", "", "Override DHCP User Class option")
)

const (
	interfaceUpTimeout = 30 * time.Second
)

var banner = `

 _________________________________
< Net booting is so hot right now >
 ---------------------------------
        \   ^__^
         \  (oo)\_______
            (__)\       )\/\
                ||----w |
                ||     ||

`

func main() {
	flag.Parse()
	if *skipDHCP && *overrideNetbootURL == "" {
		log.Fatal("-skip-dhcp requires -netboot-url")
	}
	debug := func(string, ...interface{}) {}
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	if !*useV6 && !*useV4 {
		log.Fatal("At least one of DHCPv6 and DHCPv4 is required")
	}
	// DHCPv6
	if *useV6 {
		log.Printf("Trying to obtain a DHCPv6 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv6: IfUp failed: %v", err)
		}
		debug("DHCPv6: interface %s is up", *ifname)
		var (
			netconf  *netboot.NetConf
			bootfile string
		)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			// send a netboot request via DHCP
			modifiers := []dhcpv6.Modifier{
				dhcpv6.WithArchType(dhcpv6.EFI_X86_64),
			}
			if *userClass != "" {
				modifiers = append(modifiers, dhcpv6.WithUserClass([]byte(*userClass)))
			}
			conversation, err := netboot.RequestNetbootv6(*ifname, time.Duration(*readTimeout)*time.Second, *dhcpRetries, modifiers...)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv6: netboot request for interface %s failed: %v", *ifname, err)
			}
			// get network configuration and boot file
			netconf, bootfile, err = netboot.ConversationToNetconf(conversation)
			if err != nil {
				log.Fatalf("DHCPv6: failed to extract network configuration for %s: %v", *ifname, err)
			}
			debug("DHCPv6: network configuration: %+v", netconf)
			if !*dryRun {
				// Set up IP addresses
				log.Printf("DHCPv6: configuring network interface %s", *ifname)
				if err = netboot.ConfigureInterface(*ifname, netconf); err != nil {
					log.Fatalf("DHCPv6: cannot configure IPv6 addresses on interface %s: %v", *ifname, err)
				}
				// Set up DNS
			}
			if *overrideNetbootURL != "" {
				bootfile = *overrideNetbootURL
			}
			log.Printf("DHCPv6: boot file for interface %s is %s", *ifname, bootfile)
		}
		if *overrideNetbootURL != "" {
			bootfile = *overrideNetbootURL
		}
		debug("DHCPv6: boot file URL is %s", bootfile)
		// check for supported schemes
		if !strings.HasPrefix(bootfile, "http://") {
			log.Fatal("DHCPv6: can only handle http scheme")
		}

		log.Printf("DHCPv6: fetching boot file URL: %s", bootfile)
		resp, err := http.Get(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: http.Get of %s failed: %v", bootfile, err)
		}
		// FIXME this will not be called if something fails after this point
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Fatalf("DHCPv6: cannot read boot file from the network: %v", err)
		}
		u, err := url.Parse(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: cannot parse URL %s: %v", bootfile, err)
		}
		// extract file name component
		if strings.HasSuffix(u.Path, "/") {
			log.Fatalf("Invalid file path, cannot end with '/': %s", u.Path)
		}
		filename := filepath.Base(u.Path)
		if filename == "." || filename == "" {
			log.Fatalf("Invalid empty file name extracted from file path %s", u.Path)
		}
		if err = ioutil.WriteFile(filename, body, 0400); err != nil {
			log.Fatalf("DHCPv6: cannot write to file %s: %v", filename, err)
		}
		debug("DHCPv6: saved boot file to %s", filename)
		if !*dryRun {
			log.Printf("DHCPv6: kexec'ing into %s", filename)
			kernel, err := os.OpenFile(filename, os.O_RDONLY, 0)
			if err != nil {
				log.Fatalf("DHCPv6: cannot open file %s: %v", filename, err)
			}
			if err = kexec.FileLoad(kernel, nil /* ramfs */, "" /* cmdline */); err != nil {
				log.Fatalf("DHCPv6: kexec.FileLoad failed: %v", err)
			}
			if err = kexec.Reboot(); err != nil {
				log.Fatalf("DHCPv6: kexec.Reboot failed: %v", err)
			}
		}
	}
	// DHCPv4
	if *useV4 {
		log.Printf("Trying to obtain a DHCPv4 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv4: IfUp failed: %v", err)
		}
		debug("DHCPv4: interface %s is up", *ifname)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			log.Print("DHCPv4: sending request")
			client := dhcpv4.NewClient()
			// TODO add options to request to netboot
			conversation, err := client.Exchange(*ifname, nil)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv4: Exchange failed: %v", err)
			}
			// TODO configure the network and DNS
			// TODO extract the next server and boot file and fetch it
			// TODO kexec into the NBP
		}
	}

}
<commit_msg>Handling HTTP status code != 200<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv6"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/kexec"
)

var (
	useV4              = flag.Bool("4", false, "Get a DHCPv4 lease")
	useV6              = flag.Bool("6", true, "Get a DHCPv6 lease")
	ifname             = flag.String("i", "eth0", "Interface to send packets through")
	dryRun             = flag.Bool("dryrun", false, "Do everything except assigning IP addresses, changing DNS, and kexec")
	doDebug            = flag.Bool("d", false, "Print debug output")
	skipDHCP           = flag.Bool("skip-dhcp", false, "Skip DHCP and rely on SLAAC for network configuration. This requires -netboot-url")
	overrideNetbootURL = flag.String("netboot-url", "", "Override the netboot URL normally obtained via DHCP")
	readTimeout        = flag.Int("timeout", 3, "Read timeout in seconds")
	dhcpRetries        = flag.Int("retries", 3, "Number of times a DHCP request is retried")
	userClass          = flag.String("userclass", "", "Override DHCP User Class option")
)

const (
	interfaceUpTimeout = 30 * time.Second
)

var banner = `

 _________________________________
< Net booting is so hot right now >
 ---------------------------------
        \   ^__^
         \  (oo)\_______
            (__)\       )\/\
                ||----w |
                ||     ||

`

func main() {
	flag.Parse()
	if *skipDHCP && *overrideNetbootURL == "" {
		log.Fatal("-skip-dhcp requires -netboot-url")
	}
	debug := func(string, ...interface{}) {}
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	if !*useV6 && !*useV4 {
		log.Fatal("At least one of DHCPv6 and DHCPv4 is required")
	}
	// DHCPv6
	if *useV6 {
		log.Printf("Trying to obtain a DHCPv6 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv6: IfUp failed: %v", err)
		}
		debug("DHCPv6: interface %s is up", *ifname)
		var (
			netconf  *netboot.NetConf
			bootfile string
		)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			// send a netboot request via DHCP
			modifiers := []dhcpv6.Modifier{
				dhcpv6.WithArchType(dhcpv6.EFI_X86_64),
			}
			if *userClass != "" {
				modifiers = append(modifiers, dhcpv6.WithUserClass([]byte(*userClass)))
			}
			conversation, err := netboot.RequestNetbootv6(*ifname, time.Duration(*readTimeout)*time.Second, *dhcpRetries, modifiers...)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv6: netboot request for interface %s failed: %v", *ifname, err)
			}
			// get network configuration and boot file
			netconf, bootfile, err = netboot.ConversationToNetconf(conversation)
			if err != nil {
				log.Fatalf("DHCPv6: failed to extract network configuration for %s: %v", *ifname, err)
			}
			debug("DHCPv6: network configuration: %+v", netconf)
			if !*dryRun {
				// Set up IP addresses
				log.Printf("DHCPv6: configuring network interface %s", *ifname)
				if err = netboot.ConfigureInterface(*ifname, netconf); err != nil {
					log.Fatalf("DHCPv6: cannot configure IPv6 addresses on interface %s: %v", *ifname, err)
				}
				// Set up DNS
			}
			if *overrideNetbootURL != "" {
				bootfile = *overrideNetbootURL
			}
			log.Printf("DHCPv6: boot file for interface %s is %s", *ifname, bootfile)
		}
		if *overrideNetbootURL != "" {
			bootfile = *overrideNetbootURL
		}
		debug("DHCPv6: boot file URL is %s", bootfile)
		// check for supported schemes
		if !strings.HasPrefix(bootfile, "http://") {
			log.Fatal("DHCPv6: can only handle http scheme")
		}

		log.Printf("DHCPv6: fetching boot file URL: %s", bootfile)
		resp, err := http.Get(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: http.Get of %s failed: %v", bootfile, err)
		}
		// FIXME this will not be called if something fails after this point
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			log.Fatalf("Status code is not 200 OK: %sd", resp.StatusCode)
		}
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Fatalf("DHCPv6: cannot read boot file from the network: %v", err)
		}
		u, err := url.Parse(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: cannot parse URL %s: %v", bootfile, err)
		}
		// extract file name component
		if strings.HasSuffix(u.Path, "/") {
			log.Fatalf("Invalid file path, cannot end with '/': %s", u.Path)
		}
		filename := filepath.Base(u.Path)
		if filename == "." || filename == "" {
			log.Fatalf("Invalid empty file name extracted from file path %s", u.Path)
		}
		if err = ioutil.WriteFile(filename, body, 0400); err != nil {
			log.Fatalf("DHCPv6: cannot write to file %s: %v", filename, err)
		}
		debug("DHCPv6: saved boot file to %s", filename)
		if !*dryRun {
			log.Printf("DHCPv6: kexec'ing into %s", filename)
			kernel, err := os.OpenFile(filename, os.O_RDONLY, 0)
			if err != nil {
				log.Fatalf("DHCPv6: cannot open file %s: %v", filename, err)
			}
			if err = kexec.FileLoad(kernel, nil /* ramfs */, "" /* cmdline */); err != nil {
				log.Fatalf("DHCPv6: kexec.FileLoad failed: %v", err)
			}
			if err = kexec.Reboot(); err != nil {
				log.Fatalf("DHCPv6: kexec.Reboot failed: %v", err)
			}
		}
	}
	// DHCPv4
	if *useV4 {
		log.Printf("Trying to obtain a DHCPv4 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv4: IfUp failed: %v", err)
		}
		debug("DHCPv4: interface %s is up", *ifname)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			log.Print("DHCPv4: sending request")
			client := dhcpv4.NewClient()
			// TODO add options to request to netboot
			conversation, err := client.Exchange(*ifname, nil)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv4: Exchange failed: %v", err)
			}
			// TODO configure the network and DNS
			// TODO extract the next server and boot file and fetch it
			// TODO kexec into the NBP
		}
	}

}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcwire

import (
	"fmt"
	"io"
)

// MsgGetData implements the Message interface and represents a bitcoin
// getdata message.  It is used to request data such as blocks and transactions
// from another peer.  It should be used in response to the inv (MsgInv) message
// to request the actual data referenced by each inventory vector the receiving
// peer doesn't already have.  Each message is limited to a maximum number of
// inventory vectors, which is currently 50,000.  As a result, multiple messages
// must be used to request larger amounts of data.
//
// Use the AddInvVect function to build up the list of inventory vectors when
// sending a getdata message to another peer.
type MsgGetData struct {
	InvList []*InvVect
}

// AddInvVect adds an inventory vector to the message.
func (msg *MsgGetData) AddInvVect(iv *InvVect) error {
	if len(msg.InvList)+1 > MaxInvPerMsg {
		str := fmt.Sprintf("too many invvect in message [max %v]",
			MaxInvPerMsg)
		return messageError("MsgAddr.AddInvVect", str)
	}

	msg.InvList = append(msg.InvList, iv)
	return nil
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
func (msg *MsgGetData) BtcDecode(r io.Reader, pver uint32) error {
	count, err := readVarInt(r, pver)
	if err != nil {
		return err
	}

	// Limit to max inventory vectors per message.
	if count > MaxInvPerMsg {
		str := fmt.Sprintf("too many invvect in message [%v]", count)
		return messageError("MsgGetData.BtcDecode", str)
	}

	for i := uint64(0); i < count; i++ {
		iv := InvVect{}
		err := readInvVect(r, pver, &iv)
		if err != nil {
			return err
		}
		msg.AddInvVect(&iv)
	}

	return nil
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
func (msg *MsgGetData) BtcEncode(w io.Writer, pver uint32) error {
	// Limit to max inventory vectors per message.
	count := len(msg.InvList)
	if count > MaxInvPerMsg {
		str := fmt.Sprintf("too many invvect in message [%v]", count)
		return messageError("MsgGetData.BtcEncode", str)
	}

	err := writeVarInt(w, pver, uint64(count))
	if err != nil {
		return err
	}

	for _, iv := range msg.InvList {
		err := writeInvVect(w, pver, iv)
		if err != nil {
			return err
		}
	}

	return nil
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgGetData) Command() string {
	return cmdGetData
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgGetData) MaxPayloadLength(pver uint32) uint32 {
	// Num inventory vectors (varInt) + max allowed inventory vectors.
	return maxVarIntPayload + (MaxInvPerMsg * maxInvVectPayload)
}

// NewMsgGetData returns a new bitcoin getdata message that conforms to the
// Message interface.  See MsgGetData for details.
func NewMsgGetData() *MsgGetData {
	return &MsgGetData{}
}
<commit_msg>Correct func prefix in GetData error message.<commit_after>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcwire

import (
	"fmt"
	"io"
)

// MsgGetData implements the Message interface and represents a bitcoin
// getdata message.  It is used to request data such as blocks and transactions
// from another peer.  It should be used in response to the inv (MsgInv) message
// to request the actual data referenced by each inventory vector the receiving
// peer doesn't already have.  Each message is limited to a maximum number of
// inventory vectors, which is currently 50,000.  As a result, multiple messages
// must be used to request larger amounts of data.
//
// Use the AddInvVect function to build up the list of inventory vectors when
// sending a getdata message to another peer.
type MsgGetData struct {
	InvList []*InvVect
}

// AddInvVect adds an inventory vector to the message.
func (msg *MsgGetData) AddInvVect(iv *InvVect) error {
	if len(msg.InvList)+1 > MaxInvPerMsg {
		str := fmt.Sprintf("too many invvect in message [max %v]",
			MaxInvPerMsg)
		return messageError("MsgGetData.AddInvVect", str)
	}

	msg.InvList = append(msg.InvList, iv)
	return nil
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
func (msg *MsgGetData) BtcDecode(r io.Reader, pver uint32) error {
	count, err := readVarInt(r, pver)
	if err != nil {
		return err
	}

	// Limit to max inventory vectors per message.
	if count > MaxInvPerMsg {
		str := fmt.Sprintf("too many invvect in message [%v]", count)
		return messageError("MsgGetData.BtcDecode", str)
	}

	for i := uint64(0); i < count; i++ {
		iv := InvVect{}
		err := readInvVect(r, pver, &iv)
		if err != nil {
			return err
		}
		msg.AddInvVect(&iv)
	}

	return nil
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
func (msg *MsgGetData) BtcEncode(w io.Writer, pver uint32) error {
	// Limit to max inventory vectors per message.
	count := len(msg.InvList)
	if count > MaxInvPerMsg {
		str := fmt.Sprintf("too many invvect in message [%v]", count)
		return messageError("MsgGetData.BtcEncode", str)
	}

	err := writeVarInt(w, pver, uint64(count))
	if err != nil {
		return err
	}

	for _, iv := range msg.InvList {
		err := writeInvVect(w, pver, iv)
		if err != nil {
			return err
		}
	}

	return nil
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgGetData) Command() string {
	return cmdGetData
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgGetData) MaxPayloadLength(pver uint32) uint32 {
	// Num inventory vectors (varInt) + max allowed inventory vectors.
	return maxVarIntPayload + (MaxInvPerMsg * maxInvVectPayload)
}

// NewMsgGetData returns a new bitcoin getdata message that conforms to the
// Message interface.  See MsgGetData for details.
func NewMsgGetData() *MsgGetData {
	return &MsgGetData{}
}
<|endoftext|>


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

package log

import (
	"fmt"
	"io"
	"reflect"
	"strconv"
	"strings"
)

const escape = "\033"

// ColorAttribute defines a single SGR Code
type ColorAttribute int

// Base ColorAttributes
const (
	Reset ColorAttribute = iota
	Bold
	Faint
	Italic
	Underline
	BlinkSlow
	BlinkRapid
	ReverseVideo
	Concealed
	CrossedOut
)

// Foreground text colors
const (
	FgBlack ColorAttribute = iota + 30
	FgRed
	FgGreen
	FgYellow
	FgBlue
	FgMagenta
	FgCyan
	FgWhite
)

// Foreground Hi-Intensity text colors
const (
	FgHiBlack ColorAttribute = iota + 90
	FgHiRed
	FgHiGreen
	FgHiYellow
	FgHiBlue
	FgHiMagenta
	FgHiCyan
	FgHiWhite
)

// Background text colors
const (
	BgBlack ColorAttribute = iota + 40
	BgRed
	BgGreen
	BgYellow
	BgBlue
	BgMagenta
	BgCyan
	BgWhite
)

// Background Hi-Intensity text colors
const (
	BgHiBlack ColorAttribute = iota + 100
	BgHiRed
	BgHiGreen
	BgHiYellow
	BgHiBlue
	BgHiMagenta
	BgHiCyan
	BgHiWhite
)

var colorAttributeToString = map[ColorAttribute]string{
	Reset:        "Reset",
	Bold:         "Bold",
	Faint:        "Faint",
	Italic:       "Italic",
	Underline:    "Underline",
	BlinkSlow:    "BlinkSlow",
	BlinkRapid:   "BlinkRapid",
	ReverseVideo: "ReverseVideo",
	Concealed:    "Concealed",
	CrossedOut:   "CrossedOut",
	FgBlack:      "FgBlack",
	FgRed:        "FgRed",
	FgGreen:      "FgGreen",
	FgYellow:     "FgYellow",
	FgBlue:       "FgBlue",
	FgMagenta:    "FgMagenta",
	FgCyan:       "FgCyan",
	FgWhite:      "FgWhite",
	FgHiBlack:    "FgHiBlack",
	FgHiRed:      "FgHiRed",
	FgHiGreen:    "FgHiGreen",
	FgHiYellow:   "FgHiYellow",
	FgHiBlue:     "FgHiBlue",
	FgHiMagenta:  "FgHiMagenta",
	FgHiCyan:     "FgHiCyan",
	FgHiWhite:    "FgHiWhite",
	BgBlack:      "BgBlack",
	BgRed:        "BgRed",
	BgGreen:      "BgGreen",
	BgYellow:     "BgYellow",
	BgBlue:       "BgBlue",
	BgMagenta:    "BgMagenta",
	BgCyan:       "BgCyan",
	BgWhite:      "BgWhite",
	BgHiBlack:    "BgHiBlack",
	BgHiRed:      "BgHiRed",
	BgHiGreen:    "BgHiGreen",
	BgHiYellow:   "BgHiYellow",
	BgHiBlue:     "BgHiBlue",
	BgHiMagenta:  "BgHiMagenta",
	BgHiCyan:     "BgHiCyan",
	BgHiWhite:    "BgHiWhite",
}

func (c *ColorAttribute) String() string {
	return colorAttributeToString[*c]
}

var colorAttributeFromString = map[string]ColorAttribute{}

// ColorAttributeFromString will return a ColorAttribute given a string
func ColorAttributeFromString(from string) ColorAttribute {
	lowerFrom := strings.TrimSpace(strings.ToLower(from))
	return colorAttributeFromString[lowerFrom]
}

// ColorString converts a list of ColorAttributes to a color string
func ColorString(attrs ...ColorAttribute) string {
	return string(ColorBytes(attrs...))
}

// ColorBytes converts a list of ColorAttributes to a byte array
func ColorBytes(attrs ...ColorAttribute) []byte {
	bytes := make([]byte, 0, 20)
	bytes = append(bytes, escape[0], '[')
	if len(attrs) > 0 {
		bytes = append(bytes, strconv.Itoa(int(attrs[0]))...)
		for _, a := range attrs[1:] {
			bytes = append(bytes, ';')
			bytes = append(bytes, strconv.Itoa(int(a))...)
		}
	} else {
		bytes = append(bytes, strconv.Itoa(int(Bold))...)
	}
	bytes = append(bytes, 'm')
	return bytes
}

var levelToColor = map[Level][]byte{
	TRACE:    ColorBytes(Bold, FgCyan),
	DEBUG:    ColorBytes(Bold, FgBlue),
	INFO:     ColorBytes(Bold, FgGreen),
	WARN:     ColorBytes(Bold, FgYellow),
	ERROR:    ColorBytes(Bold, FgRed),
	CRITICAL: ColorBytes(Bold, BgMagenta),
	FATAL:    ColorBytes(Bold, BgRed),
	NONE:     ColorBytes(Reset),
}

var resetBytes = ColorBytes(Reset)
var fgCyanBytes = ColorBytes(FgCyan)
var fgGreenBytes = ColorBytes(FgGreen)
var fgBoldBytes = ColorBytes(Bold)

type protectedANSIWriterMode int

const (
	escapeAll protectedANSIWriterMode = iota
	allowColor
	removeColor
)

type protectedANSIWriter struct {
	w    io.Writer
	mode protectedANSIWriterMode
}

// Write will protect against unusual characters
func (c *protectedANSIWriter) Write(bytes []byte) (int, error) {
	end := len(bytes)
	totalWritten := 0
normalLoop:
	for i := 0; i < end; {
		lasti := i

		if c.mode == escapeAll {
			for i < end && (bytes[i] >= ' ' || bytes[i] == '\n' || bytes[i] == '\t') {
				i++
			}
		} else {
			// Allow tabs if we're not escaping everything
			for i < end && (bytes[i] >= ' ' || bytes[i] == '\t') {
				i++
			}
		}

		if i > lasti {
			written, err := c.w.Write(bytes[lasti:i])
			totalWritten += written
			if err != nil {
				return totalWritten, err
			}

		}
		if i >= end {
			break
		}

		// If we're not just escaping all we should prefix all newlines with a \t
		if c.mode != escapeAll {
			if bytes[i] == '\n' {
				written, err := c.w.Write([]byte{'\n', '\t'})
				if written > 0 {
					totalWritten++
				}
				if err != nil {
					return totalWritten, err
				}
				i++
				continue normalLoop
			}

			if bytes[i] == escape[0] && i+1 < end && bytes[i+1] == '[' {
				for j := i + 2; j < end; j++ {
					if bytes[j] >= '0' && bytes[j] <= '9' {
						continue
					}
					if bytes[j] == ';' {
						continue
					}
					if bytes[j] == 'm' {
						if c.mode == allowColor {
							written, err := c.w.Write(bytes[i : j+1])
							totalWritten += written
							if err != nil {
								return totalWritten, err
							}
						} else {
							totalWritten = j
						}
						i = j + 1
						continue normalLoop
					}
					break
				}
			}
		}

		// Process naughty character
		if _, err := fmt.Fprintf(c.w, `\%#o03d`, bytes[i]); err != nil {
			return totalWritten, err
		}
		i++
		totalWritten++
	}
	return totalWritten, nil
}

// ColorSprintf returns a colored string from a format and arguments
// arguments will be wrapped in ColoredValues to protect against color spoofing
func ColorSprintf(format string, args ...interface{}) string {
	if len(args) > 0 {
		v := make([]interface{}, len(args))
		for i := 0; i < len(v); i++ {
			v[i] = NewColoredValuePointer(&args[i])
		}
		return fmt.Sprintf(format, v...)
	}
	return format
}

// ColorFprintf will write to the provided writer similar to ColorSprintf
func ColorFprintf(w io.Writer, format string, args ...interface{}) (int, error) {
	if len(args) > 0 {
		v := make([]interface{}, len(args))
		for i := 0; i < len(v); i++ {
			v[i] = NewColoredValuePointer(&args[i])
		}
		return fmt.Fprintf(w, format, v...)
	}
	return fmt.Fprint(w, format)
}

// ColorFormatted structs provide their own colored string when formatted with ColorSprintf
type ColorFormatted interface {
	// ColorFormat provides the colored representation of the value
	ColorFormat(s fmt.State)
}

var colorFormattedType = reflect.TypeOf((*ColorFormatted)(nil)).Elem()

// ColoredValue will Color the provided value
type ColoredValue struct {
	colorBytes *[]byte
	resetBytes *[]byte
	Value      *interface{}
}

// NewColoredValue is a helper function to create a ColoredValue from a Value
// If no color is provided it defaults to Bold with standard Reset
// If a ColoredValue is provided it is not changed
func NewColoredValue(value interface{}, color ...ColorAttribute) *ColoredValue {
	return NewColoredValuePointer(&value, color...)
}

// NewColoredValuePointer is a helper function to create a ColoredValue from a Value Pointer
// If no color is provided it defaults to Bold with standard Reset
// If a ColoredValue is provided it is not changed
func NewColoredValuePointer(value *interface{}, color ...ColorAttribute) *ColoredValue {
	if val, ok := (*value).(*ColoredValue); ok {
		return val
	}
	if len(color) > 0 {
		bytes := ColorBytes(color...)
		return &ColoredValue{
			colorBytes: &bytes,
			resetBytes: &resetBytes,
			Value:      value,
		}
	}
	return &ColoredValue{
		colorBytes: &fgBoldBytes,
		resetBytes: &resetBytes,
		Value:      value,
	}

}

// NewColoredValueBytes creates a value from the provided value with color bytes
// If a ColoredValue is provided it is not changed
func NewColoredValueBytes(value interface{}, colorBytes *[]byte) *ColoredValue {
	if val, ok := value.(*ColoredValue); ok {
		return val
	}
	return &ColoredValue{
		colorBytes: colorBytes,
		resetBytes: &resetBytes,
		Value:      &value,
	}
}

// NewColoredIDValue is a helper function to create a ColoredValue from a Value
// The Value will be colored with FgCyan
// If a ColoredValue is provided it is not changed
func NewColoredIDValue(value interface{}) *ColoredValue {
	return NewColoredValueBytes(value, &fgCyanBytes)
}

// Format will format the provided value and protect against ANSI color spoofing within the value
// If the wrapped value is ColorFormatted and the format is "%-v" then its ColorString will
// be used. It is presumed that this ColorString is safe.
func (cv *ColoredValue) Format(s fmt.State, c rune) {
	if c == 'v' && s.Flag('-') {
		if val, ok := (*cv.Value).(ColorFormatted); ok {
			val.ColorFormat(s)
			return
		}
		v := reflect.ValueOf(*cv.Value)
		t := v.Type()

		if reflect.PtrTo(t).Implements(colorFormattedType) {
			vp := reflect.New(t)
			vp.Elem().Set(v)
			val := vp.Interface().(ColorFormatted)
			val.ColorFormat(s)
			return
		}
	}
	s.Write(*cv.colorBytes)
	fmt.Fprintf(&protectedANSIWriter{w: s}, fmtString(s, c), *(cv.Value))
	s.Write(*cv.resetBytes)
}

// SetColorBytes will allow a user to set the colorBytes of a colored value
func (cv *ColoredValue) SetColorBytes(colorBytes []byte) {
	cv.colorBytes = &colorBytes
}

// SetColorBytesPointer will allow a user to set the colorBytes pointer of a colored value
func (cv *ColoredValue) SetColorBytesPointer(colorBytes *[]byte) {
	cv.colorBytes = colorBytes
}

// SetResetBytes will allow a user to set the resetBytes pointer of a colored value
func (cv *ColoredValue) SetResetBytes(resetBytes []byte) {
	cv.resetBytes = &resetBytes
}

// SetResetBytesPointer will allow a user to set the resetBytes pointer of a colored value
func (cv *ColoredValue) SetResetBytesPointer(resetBytes *[]byte) {
	cv.resetBytes = resetBytes
}

func fmtString(s fmt.State, c rune) string {
	var width, precision string
	base := make([]byte, 0, 8)
	base = append(base, '%')
	for _, c := range []byte(" +-#0") {
		if s.Flag(int(c)) {
			base = append(base, c)
		}
	}
	if w, ok := s.Width(); ok {
		width = strconv.Itoa(w)
	}
	if p, ok := s.Precision(); ok {
		precision = "." + strconv.Itoa(p)
	}
	return fmt.Sprintf("%s%s%s%c", base, width, precision, c)
}

func init() {
	for attr, from := range colorAttributeToString {
		colorAttributeFromString[strings.ToLower(from)] = attr
	}
}
<commit_msg>Stop printing 03d after escaped characters in logs (#18030)<commit_after>// Copyright 2019 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package log

import (
	"fmt"
	"io"
	"reflect"
	"strconv"
	"strings"
)

const escape = "\033"

// ColorAttribute defines a single SGR Code
type ColorAttribute int

// Base ColorAttributes
const (
	Reset ColorAttribute = iota
	Bold
	Faint
	Italic
	Underline
	BlinkSlow
	BlinkRapid
	ReverseVideo
	Concealed
	CrossedOut
)

// Foreground text colors
const (
	FgBlack ColorAttribute = iota + 30
	FgRed
	FgGreen
	FgYellow
	FgBlue
	FgMagenta
	FgCyan
	FgWhite
)

// Foreground Hi-Intensity text colors
const (
	FgHiBlack ColorAttribute = iota + 90
	FgHiRed
	FgHiGreen
	FgHiYellow
	FgHiBlue
	FgHiMagenta
	FgHiCyan
	FgHiWhite
)

// Background text colors
const (
	BgBlack ColorAttribute = iota + 40
	BgRed
	BgGreen
	BgYellow
	BgBlue
	BgMagenta
	BgCyan
	BgWhite
)

// Background Hi-Intensity text colors
const (
	BgHiBlack ColorAttribute = iota + 100
	BgHiRed
	BgHiGreen
	BgHiYellow
	BgHiBlue
	BgHiMagenta
	BgHiCyan
	BgHiWhite
)

var colorAttributeToString = map[ColorAttribute]string{
	Reset:        "Reset",
	Bold:         "Bold",
	Faint:        "Faint",
	Italic:       "Italic",
	Underline:    "Underline",
	BlinkSlow:    "BlinkSlow",
	BlinkRapid:   "BlinkRapid",
	ReverseVideo: "ReverseVideo",
	Concealed:    "Concealed",
	CrossedOut:   "CrossedOut",
	FgBlack:      "FgBlack",
	FgRed:        "FgRed",
	FgGreen:      "FgGreen",
	FgYellow:     "FgYellow",
	FgBlue:       "FgBlue",
	FgMagenta:    "FgMagenta",
	FgCyan:       "FgCyan",
	FgWhite:      "FgWhite",
	FgHiBlack:    "FgHiBlack",
	FgHiRed:      "FgHiRed",
	FgHiGreen:    "FgHiGreen",
	FgHiYellow:   "FgHiYellow",
	FgHiBlue:     "FgHiBlue",
	FgHiMagenta:  "FgHiMagenta",
	FgHiCyan:     "FgHiCyan",
	FgHiWhite:    "FgHiWhite",
	BgBlack:      "BgBlack",
	BgRed:        "BgRed",
	BgGreen:      "BgGreen",
	BgYellow:     "BgYellow",
	BgBlue:       "BgBlue",
	BgMagenta:    "BgMagenta",
	BgCyan:       "BgCyan",
	BgWhite:      "BgWhite",
	BgHiBlack:    "BgHiBlack",
	BgHiRed:      "BgHiRed",
	BgHiGreen:    "BgHiGreen",
	BgHiYellow:   "BgHiYellow",
	BgHiBlue:     "BgHiBlue",
	BgHiMagenta:  "BgHiMagenta",
	BgHiCyan:     "BgHiCyan",
	BgHiWhite:    "BgHiWhite",
}

func (c *ColorAttribute) String() string {
	return colorAttributeToString[*c]
}

var colorAttributeFromString = map[string]ColorAttribute{}

// ColorAttributeFromString will return a ColorAttribute given a string
func ColorAttributeFromString(from string) ColorAttribute {
	lowerFrom := strings.TrimSpace(strings.ToLower(from))
	return colorAttributeFromString[lowerFrom]
}

// ColorString converts a list of ColorAttributes to a color string
func ColorString(attrs ...ColorAttribute) string {
	return string(ColorBytes(attrs...))
}

// ColorBytes converts a list of ColorAttributes to a byte array
func ColorBytes(attrs ...ColorAttribute) []byte {
	bytes := make([]byte, 0, 20)
	bytes = append(bytes, escape[0], '[')
	if len(attrs) > 0 {
		bytes = append(bytes, strconv.Itoa(int(attrs[0]))...)
		for _, a := range attrs[1:] {
			bytes = append(bytes, ';')
			bytes = append(bytes, strconv.Itoa(int(a))...)
		}
	} else {
		bytes = append(bytes, strconv.Itoa(int(Bold))...)
	}
	bytes = append(bytes, 'm')
	return bytes
}

var levelToColor = map[Level][]byte{
	TRACE:    ColorBytes(Bold, FgCyan),
	DEBUG:    ColorBytes(Bold, FgBlue),
	INFO:     ColorBytes(Bold, FgGreen),
	WARN:     ColorBytes(Bold, FgYellow),
	ERROR:    ColorBytes(Bold, FgRed),
	CRITICAL: ColorBytes(Bold, BgMagenta),
	FATAL:    ColorBytes(Bold, BgRed),
	NONE:     ColorBytes(Reset),
}

var resetBytes = ColorBytes(Reset)
var fgCyanBytes = ColorBytes(FgCyan)
var fgGreenBytes = ColorBytes(FgGreen)
var fgBoldBytes = ColorBytes(Bold)

type protectedANSIWriterMode int

const (
	escapeAll protectedANSIWriterMode = iota
	allowColor
	removeColor
)

type protectedANSIWriter struct {
	w    io.Writer
	mode protectedANSIWriterMode
}

// Write will protect against unusual characters
func (c *protectedANSIWriter) Write(bytes []byte) (int, error) {
	end := len(bytes)
	totalWritten := 0
normalLoop:
	for i := 0; i < end; {
		lasti := i

		if c.mode == escapeAll {
			for i < end && (bytes[i] >= ' ' || bytes[i] == '\n' || bytes[i] == '\t') {
				i++
			}
		} else {
			// Allow tabs if we're not escaping everything
			for i < end && (bytes[i] >= ' ' || bytes[i] == '\t') {
				i++
			}
		}

		if i > lasti {
			written, err := c.w.Write(bytes[lasti:i])
			totalWritten += written
			if err != nil {
				return totalWritten, err
			}

		}
		if i >= end {
			break
		}

		// If we're not just escaping all we should prefix all newlines with a \t
		if c.mode != escapeAll {
			if bytes[i] == '\n' {
				written, err := c.w.Write([]byte{'\n', '\t'})
				if written > 0 {
					totalWritten++
				}
				if err != nil {
					return totalWritten, err
				}
				i++
				continue normalLoop
			}

			if bytes[i] == escape[0] && i+1 < end && bytes[i+1] == '[' {
				for j := i + 2; j < end; j++ {
					if bytes[j] >= '0' && bytes[j] <= '9' {
						continue
					}
					if bytes[j] == ';' {
						continue
					}
					if bytes[j] == 'm' {
						if c.mode == allowColor {
							written, err := c.w.Write(bytes[i : j+1])
							totalWritten += written
							if err != nil {
								return totalWritten, err
							}
						} else {
							totalWritten = j
						}
						i = j + 1
						continue normalLoop
					}
					break
				}
			}
		}

		// Process naughty character
		if _, err := fmt.Fprintf(c.w, `\%#03o`, bytes[i]); err != nil {
			return totalWritten, err
		}
		i++
		totalWritten++
	}
	return totalWritten, nil
}

// ColorSprintf returns a colored string from a format and arguments
// arguments will be wrapped in ColoredValues to protect against color spoofing
func ColorSprintf(format string, args ...interface{}) string {
	if len(args) > 0 {
		v := make([]interface{}, len(args))
		for i := 0; i < len(v); i++ {
			v[i] = NewColoredValuePointer(&args[i])
		}
		return fmt.Sprintf(format, v...)
	}
	return format
}

// ColorFprintf will write to the provided writer similar to ColorSprintf
func ColorFprintf(w io.Writer, format string, args ...interface{}) (int, error) {
	if len(args) > 0 {
		v := make([]interface{}, len(args))
		for i := 0; i < len(v); i++ {
			v[i] = NewColoredValuePointer(&args[i])
		}
		return fmt.Fprintf(w, format, v...)
	}
	return fmt.Fprint(w, format)
}

// ColorFormatted structs provide their own colored string when formatted with ColorSprintf
type ColorFormatted interface {
	// ColorFormat provides the colored representation of the value
	ColorFormat(s fmt.State)
}

var colorFormattedType = reflect.TypeOf((*ColorFormatted)(nil)).Elem()

// ColoredValue will Color the provided value
type ColoredValue struct {
	colorBytes *[]byte
	resetBytes *[]byte
	Value      *interface{}
}

// NewColoredValue is a helper function to create a ColoredValue from a Value
// If no color is provided it defaults to Bold with standard Reset
// If a ColoredValue is provided it is not changed
func NewColoredValue(value interface{}, color ...ColorAttribute) *ColoredValue {
	return NewColoredValuePointer(&value, color...)
}

// NewColoredValuePointer is a helper function to create a ColoredValue from a Value Pointer
// If no color is provided it defaults to Bold with standard Reset
// If a ColoredValue is provided it is not changed
func NewColoredValuePointer(value *interface{}, color ...ColorAttribute) *ColoredValue {
	if val, ok := (*value).(*ColoredValue); ok {
		return val
	}
	if len(color) > 0 {
		bytes := ColorBytes(color...)
		return &ColoredValue{
			colorBytes: &bytes,
			resetBytes: &resetBytes,
			Value:      value,
		}
	}
	return &ColoredValue{
		colorBytes: &fgBoldBytes,
		resetBytes: &resetBytes,
		Value:      value,
	}

}

// NewColoredValueBytes creates a value from the provided value with color bytes
// If a ColoredValue is provided it is not changed
func NewColoredValueBytes(value interface{}, colorBytes *[]byte) *ColoredValue {
	if val, ok := value.(*ColoredValue); ok {
		return val
	}
	return &ColoredValue{
		colorBytes: colorBytes,
		resetBytes: &resetBytes,
		Value:      &value,
	}
}

// NewColoredIDValue is a helper function to create a ColoredValue from a Value
// The Value will be colored with FgCyan
// If a ColoredValue is provided it is not changed
func NewColoredIDValue(value interface{}) *ColoredValue {
	return NewColoredValueBytes(value, &fgCyanBytes)
}

// Format will format the provided value and protect against ANSI color spoofing within the value
// If the wrapped value is ColorFormatted and the format is "%-v" then its ColorString will
// be used. It is presumed that this ColorString is safe.
func (cv *ColoredValue) Format(s fmt.State, c rune) {
	if c == 'v' && s.Flag('-') {
		if val, ok := (*cv.Value).(ColorFormatted); ok {
			val.ColorFormat(s)
			return
		}
		v := reflect.ValueOf(*cv.Value)
		t := v.Type()

		if reflect.PtrTo(t).Implements(colorFormattedType) {
			vp := reflect.New(t)
			vp.Elem().Set(v)
			val := vp.Interface().(ColorFormatted)
			val.ColorFormat(s)
			return
		}
	}
	s.Write(*cv.colorBytes)
	fmt.Fprintf(&protectedANSIWriter{w: s}, fmtString(s, c), *(cv.Value))
	s.Write(*cv.resetBytes)
}

// SetColorBytes will allow a user to set the colorBytes of a colored value
func (cv *ColoredValue) SetColorBytes(colorBytes []byte) {
	cv.colorBytes = &colorBytes
}

// SetColorBytesPointer will allow a user to set the colorBytes pointer of a colored value
func (cv *ColoredValue) SetColorBytesPointer(colorBytes *[]byte) {
	cv.colorBytes = colorBytes
}

// SetResetBytes will allow a user to set the resetBytes pointer of a colored value
func (cv *ColoredValue) SetResetBytes(resetBytes []byte) {
	cv.resetBytes = &resetBytes
}

// SetResetBytesPointer will allow a user to set the resetBytes pointer of a colored value
func (cv *ColoredValue) SetResetBytesPointer(resetBytes *[]byte) {
	cv.resetBytes = resetBytes
}

func fmtString(s fmt.State, c rune) string {
	var width, precision string
	base := make([]byte, 0, 8)
	base = append(base, '%')
	for _, c := range []byte(" +-#0") {
		if s.Flag(int(c)) {
			base = append(base, c)
		}
	}
	if w, ok := s.Width(); ok {
		width = strconv.Itoa(w)
	}
	if p, ok := s.Precision(); ok {
		precision = "." + strconv.Itoa(p)
	}
	return fmt.Sprintf("%s%s%s%c", base, width, precision, c)
}

func init() {
	for attr, from := range colorAttributeToString {
		colorAttributeFromString[strings.ToLower(from)] = attr
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
)

//-------------------------------------------------------------------------
// []package_import
//-------------------------------------------------------------------------

type package_import struct {
	alias string
	path  string
}

// Parses import declarations until the first non-import declaration and fills
// `packages` array with import information.
func collect_package_imports(filename string, decls []ast.Decl, context *package_lookup_context) []package_import {
	pi := make([]package_import, 0, 16)
	for _, decl := range decls {
		if gd, ok := decl.(*ast.GenDecl); ok && gd.Tok == token.IMPORT {
			for _, spec := range gd.Specs {
				imp := spec.(*ast.ImportSpec)
				path, alias := path_and_alias(imp)
				path, ok := abs_path_for_package(filename, path, context)
				if ok && alias != "_" {
					pi = append(pi, package_import{alias, path})
				}
			}
		} else {
			break
		}
	}
	return pi
}

//-------------------------------------------------------------------------
// decl_file_cache
//
// Contains cache for top-level declarations of a file as well as its
// contents, AST and import information.
//-------------------------------------------------------------------------

type decl_file_cache struct {
	name  string // file name
	mtime int64  // last modification time

	decls     map[string]*decl // top-level declarations
	error     error            // last error
	packages  []package_import // import information
	filescope *scope

	fset    *token.FileSet
	context *package_lookup_context
}

func new_decl_file_cache(name string, context *package_lookup_context) *decl_file_cache {
	return &decl_file_cache{
		name:    name,
		context: context,
	}
}

func (f *decl_file_cache) update() {
	stat, err := os.Stat(f.name)
	if err != nil {
		f.decls = nil
		f.error = err
		f.fset = nil
		return
	}

	statmtime := stat.ModTime().UnixNano()
	if f.mtime == statmtime {
		return
	}

	f.mtime = statmtime
	f.read_file()
}

func (f *decl_file_cache) read_file() {
	var data []byte
	data, f.error = file_reader.read_file(f.name)
	if f.error != nil {
		return
	}
	data, _ = filter_out_shebang(data)

	f.process_data(data)
}

func (f *decl_file_cache) process_data(data []byte) {
	var file *ast.File
	f.fset = token.NewFileSet()
	file, f.error = parser.ParseFile(f.fset, "", data, 0)
	f.filescope = new_scope(nil)
	for _, d := range file.Decls {
		anonymify_ast(d, 0, f.filescope)
	}
	f.packages = collect_package_imports(f.name, file.Decls, f.context)
	f.decls = make(map[string]*decl, len(file.Decls))
	for _, decl := range file.Decls {
		append_to_top_decls(f.decls, decl, f.filescope)
	}
}

func append_to_top_decls(decls map[string]*decl, decl ast.Decl, scope *scope) {
	foreach_decl(decl, func(data *foreach_decl_struct) {
		class := ast_decl_class(data.decl)
		for i, name := range data.names {
			typ, v, vi := data.type_value_index(i)

			d := new_decl_full(name.Name, class, 0, typ, v, vi, scope)
			if d == nil {
				return
			}

			methodof := method_of(decl)
			if methodof != "" {
				decl, ok := decls[methodof]
				if ok {
					decl.add_child(d)
				} else {
					decl = new_decl(methodof, decl_methods_stub, scope)
					decls[methodof] = decl
					decl.add_child(d)
				}
			} else {
				decl, ok := decls[d.name]
				if ok {
					decl.expand_or_replace(d)
				} else {
					decls[d.name] = d
				}
			}
		}
	})
}

func abs_path_for_package(filename, p string, context *package_lookup_context) (string, bool) {
	dir, _ := filepath.Split(filename)
	if len(p) == 0 {
		return "", false
	}
	if p[0] == '.' {
		return fmt.Sprintf("%s.a", filepath.Join(dir, p)), true
	}
	pkg, ok := find_go_dag_package(p, dir)
	if ok {
		return pkg, true
	}
	return find_global_file(p, context)
}

func path_and_alias(imp *ast.ImportSpec) (string, string) {
	path := ""
	if imp.Path != nil {
		path = string(imp.Path.Value)
		path = path[1 : len(path)-1]
	}
	alias := ""
	if imp.Name != nil {
		alias = imp.Name.Name
	}
	return path, alias
}

func find_go_dag_package(imp, filedir string) (string, bool) {
	// Support godag directory structure
	dir, pkg := filepath.Split(imp)
	godag_pkg := filepath.Join(filedir, "..", dir, "_obj", pkg+".a")
	if file_exists(godag_pkg) {
		return godag_pkg, true
	}
	return "", false
}

// autobuild compares the mod time of the source files of the package, and if any of them is newer
// than the package object file will rebuild it.
func autobuild(p *build.Package) error {
	if p.Dir == "" {
		return fmt.Errorf("no files to build")
	}
	ps, err := os.Stat(p.PkgObj)
	if err != nil {
		// Assume package file does not exist and build for the first time.
		return build_package(p)
	}
	pt := ps.ModTime()
	fs, err := readdir(p.Dir)
	if err != nil {
		return err
	}
	for _, f := range fs {
		if f.IsDir() {
			continue
		}
		if f.ModTime().After(pt) {
			// Source file is newer than package file; rebuild.
			return build_package(p)
		}
	}
	return nil
}

// build_package builds the package by calling `go install package/import`. If everything compiles
// correctly, the newly compiled package should then be in the usual place in the `$GOPATH/pkg`
// directory, and gocode will pick it up from there.
func build_package(p *build.Package) error {
	if *g_debug {
		log.Printf("-------------------")
		log.Printf("rebuilding package %s", p.Name)
		log.Printf("package import: %s", p.ImportPath)
		log.Printf("package object: %s", p.PkgObj)
		log.Printf("package source dir: %s", p.Dir)
		log.Printf("package source files: %v", p.GoFiles)
	}
	// TODO: Should read STDERR rather than STDOUT.
	out, err := exec.Command("go", "install", p.ImportPath).Output()
	if err != nil {
		return err
	}
	if *g_debug {
		log.Printf("build out: %s\n", string(out))
	}
	return nil
}

// executes autobuild function if autobuild option is enabled, logs error and
// ignores it
func try_autobuild(p *build.Package) {
	if g_config.Autobuild {
		err := autobuild(p)
		if err != nil && *g_debug {
			log.Printf("Autobuild error: %s\n", err)
		}
	}
}

func log_found_package_maybe(imp, pkgpath string) {
	if *g_debug {
		log.Printf("Found %q at %q\n", imp, pkgpath)
	}
}

func log_build_context(context *package_lookup_context) {
	log.Printf(" GOROOT: %s\n", context.GOROOT)
	log.Printf(" GOPATH: %s\n", context.GOPATH)
	log.Printf(" GOOS: %s\n", context.GOOS)
	log.Printf(" GOARCH: %s\n", context.GOARCH)
	log.Printf(" GBProjectRoot: %q\n", context.GBProjectRoot)
	log.Printf(" lib-path: %q\n", g_config.LibPath)
}

// find_global_file returns the file path of the compiled package corresponding to the specified
// import, and a boolean stating whether such path is valid.
// TODO: Return only one value, possibly empty string if not found.
func find_global_file(imp string, context *package_lookup_context) (string, bool) {
	// gocode synthetically generates the builtin package
	// "unsafe", since the "unsafe.a" package doesn't really exist.
	// Thus, when the user request for the package "unsafe" we
	// would return synthetic global file that would be used
	// just as a key name to find this synthetic package
	if imp == "unsafe" {
		return "unsafe", true
	}

	pkgfile := fmt.Sprintf("%s.a", imp)

	// if lib-path is defined, use it
	if g_config.LibPath != "" {
		for _, p := range filepath.SplitList(g_config.LibPath) {
			pkg_path := filepath.Join(p, pkgfile)
			if file_exists(pkg_path) {
				log_found_package_maybe(imp, pkg_path)
				return pkg_path, true
			}
			// Also check the relevant pkg/OS_ARCH dir for the libpath, if provided.
			pkgdir := fmt.Sprintf("%s_%s", context.GOOS, context.GOARCH)
			pkg_path = filepath.Join(p, "pkg", pkgdir, pkgfile)
			if file_exists(pkg_path) {
				log_found_package_maybe(imp, pkg_path)
				return pkg_path, true
			}
		}
	}

	// gb-specific lookup mode, only if the root dir was found
	if g_config.PackageLookupMode == "gb" && context.GBProjectRoot != "" {
		root := context.GBProjectRoot
		pkg_path := filepath.Join(root, "pkg", context.GOOS+"-"+context.GOARCH, pkgfile)
		if file_exists(pkg_path) {
			log_found_package_maybe(imp, pkg_path)
			return pkg_path, true
		}
	}

	if context.CurrentPackagePath != "" {
		// Try vendor path first, see GO15VENDOREXPERIMENT.
		// We don't check this environment variable however, seems like there is
		// almost no harm in doing so (well.. if you experiment with vendoring,
		// gocode will fail after enabling/disabling the flag, and you'll be
		// forced to get rid of vendor binaries). But asking users to set this
		// env var is up will bring more trouble. Because we also need to pass
		// it from client to server, make sure their editors set it, etc.
		// So, whatever, let's just pretend it's always on.
		package_path := context.CurrentPackagePath
		for {
			limp := filepath.Join(package_path, "vendor", imp)
			if p, err := context.Import(limp, "", build.AllowBinary|build.FindOnly); err == nil {
				try_autobuild(p)
				if file_exists(p.PkgObj) {
					log_found_package_maybe(imp, p.PkgObj)
					return p.PkgObj, true
				}
			}
			if package_path == "" {
				break
			}
			next_path := filepath.Base(package_path)
			// let's protect ourselves from inf recursion here
			if next_path == package_path {
				break
			}
			package_path = next_path
		}
	}

	if p, err := context.Import(imp, "", build.AllowBinary|build.FindOnly); err == nil {
		try_autobuild(p)
		if file_exists(p.PkgObj) {
			log_found_package_maybe(imp, p.PkgObj)
			return p.PkgObj, true
		}
	}

	if *g_debug {
		log.Printf("Import path %q was not resolved\n", imp)
		log.Println("Gocode's build context is:")
		log_build_context(context)
	}
	return "", false
}

func package_name(file *ast.File) string {
	if file.Name != nil {
		return file.Name.Name
	}
	return ""
}

//-------------------------------------------------------------------------
// decl_cache
//
// Thread-safe collection of DeclFileCache entities.
//-------------------------------------------------------------------------

type package_lookup_context struct {
	build.Context
	GBProjectRoot      string
	CurrentPackagePath string
}

type decl_cache struct {
	cache   map[string]*decl_file_cache
	context *package_lookup_context
	sync.Mutex
}

func new_decl_cache(context *package_lookup_context) *decl_cache {
	return &decl_cache{
		cache:   make(map[string]*decl_file_cache),
		context: context,
	}
}

func (c *decl_cache) get(filename string) *decl_file_cache {
	c.Lock()
	defer c.Unlock()

	f, ok := c.cache[filename]
	if !ok {
		f = new_decl_file_cache(filename, c.context)
		c.cache[filename] = f
	}
	return f
}

func (c *decl_cache) get_and_update(filename string) *decl_file_cache {
	f := c.get(filename)
	f.update()
	return f
}
<commit_msg>Oops, use filepath.Dir instead of filepath.Base.<commit_after>package main

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
)

//-------------------------------------------------------------------------
// []package_import
//-------------------------------------------------------------------------

type package_import struct {
	alias string
	path  string
}

// Parses import declarations until the first non-import declaration and fills
// `packages` array with import information.
func collect_package_imports(filename string, decls []ast.Decl, context *package_lookup_context) []package_import {
	pi := make([]package_import, 0, 16)
	for _, decl := range decls {
		if gd, ok := decl.(*ast.GenDecl); ok && gd.Tok == token.IMPORT {
			for _, spec := range gd.Specs {
				imp := spec.(*ast.ImportSpec)
				path, alias := path_and_alias(imp)
				path, ok := abs_path_for_package(filename, path, context)
				if ok && alias != "_" {
					pi = append(pi, package_import{alias, path})
				}
			}
		} else {
			break
		}
	}
	return pi
}

//-------------------------------------------------------------------------
// decl_file_cache
//
// Contains cache for top-level declarations of a file as well as its
// contents, AST and import information.
//-------------------------------------------------------------------------

type decl_file_cache struct {
	name  string // file name
	mtime int64  // last modification time

	decls     map[string]*decl // top-level declarations
	error     error            // last error
	packages  []package_import // import information
	filescope *scope

	fset    *token.FileSet
	context *package_lookup_context
}

func new_decl_file_cache(name string, context *package_lookup_context) *decl_file_cache {
	return &decl_file_cache{
		name:    name,
		context: context,
	}
}

func (f *decl_file_cache) update() {
	stat, err := os.Stat(f.name)
	if err != nil {
		f.decls = nil
		f.error = err
		f.fset = nil
		return
	}

	statmtime := stat.ModTime().UnixNano()
	if f.mtime == statmtime {
		return
	}

	f.mtime = statmtime
	f.read_file()
}

func (f *decl_file_cache) read_file() {
	var data []byte
	data, f.error = file_reader.read_file(f.name)
	if f.error != nil {
		return
	}
	data, _ = filter_out_shebang(data)

	f.process_data(data)
}

func (f *decl_file_cache) process_data(data []byte) {
	var file *ast.File
	f.fset = token.NewFileSet()
	file, f.error = parser.ParseFile(f.fset, "", data, 0)
	f.filescope = new_scope(nil)
	for _, d := range file.Decls {
		anonymify_ast(d, 0, f.filescope)
	}
	f.packages = collect_package_imports(f.name, file.Decls, f.context)
	f.decls = make(map[string]*decl, len(file.Decls))
	for _, decl := range file.Decls {
		append_to_top_decls(f.decls, decl, f.filescope)
	}
}

func append_to_top_decls(decls map[string]*decl, decl ast.Decl, scope *scope) {
	foreach_decl(decl, func(data *foreach_decl_struct) {
		class := ast_decl_class(data.decl)
		for i, name := range data.names {
			typ, v, vi := data.type_value_index(i)

			d := new_decl_full(name.Name, class, 0, typ, v, vi, scope)
			if d == nil {
				return
			}

			methodof := method_of(decl)
			if methodof != "" {
				decl, ok := decls[methodof]
				if ok {
					decl.add_child(d)
				} else {
					decl = new_decl(methodof, decl_methods_stub, scope)
					decls[methodof] = decl
					decl.add_child(d)
				}
			} else {
				decl, ok := decls[d.name]
				if ok {
					decl.expand_or_replace(d)
				} else {
					decls[d.name] = d
				}
			}
		}
	})
}

func abs_path_for_package(filename, p string, context *package_lookup_context) (string, bool) {
	dir, _ := filepath.Split(filename)
	if len(p) == 0 {
		return "", false
	}
	if p[0] == '.' {
		return fmt.Sprintf("%s.a", filepath.Join(dir, p)), true
	}
	pkg, ok := find_go_dag_package(p, dir)
	if ok {
		return pkg, true
	}
	return find_global_file(p, context)
}

func path_and_alias(imp *ast.ImportSpec) (string, string) {
	path := ""
	if imp.Path != nil {
		path = string(imp.Path.Value)
		path = path[1 : len(path)-1]
	}
	alias := ""
	if imp.Name != nil {
		alias = imp.Name.Name
	}
	return path, alias
}

func find_go_dag_package(imp, filedir string) (string, bool) {
	// Support godag directory structure
	dir, pkg := filepath.Split(imp)
	godag_pkg := filepath.Join(filedir, "..", dir, "_obj", pkg+".a")
	if file_exists(godag_pkg) {
		return godag_pkg, true
	}
	return "", false
}

// autobuild compares the mod time of the source files of the package, and if any of them is newer
// than the package object file will rebuild it.
func autobuild(p *build.Package) error {
	if p.Dir == "" {
		return fmt.Errorf("no files to build")
	}
	ps, err := os.Stat(p.PkgObj)
	if err != nil {
		// Assume package file does not exist and build for the first time.
		return build_package(p)
	}
	pt := ps.ModTime()
	fs, err := readdir(p.Dir)
	if err != nil {
		return err
	}
	for _, f := range fs {
		if f.IsDir() {
			continue
		}
		if f.ModTime().After(pt) {
			// Source file is newer than package file; rebuild.
			return build_package(p)
		}
	}
	return nil
}

// build_package builds the package by calling `go install package/import`. If everything compiles
// correctly, the newly compiled package should then be in the usual place in the `$GOPATH/pkg`
// directory, and gocode will pick it up from there.
func build_package(p *build.Package) error {
	if *g_debug {
		log.Printf("-------------------")
		log.Printf("rebuilding package %s", p.Name)
		log.Printf("package import: %s", p.ImportPath)
		log.Printf("package object: %s", p.PkgObj)
		log.Printf("package source dir: %s", p.Dir)
		log.Printf("package source files: %v", p.GoFiles)
	}
	// TODO: Should read STDERR rather than STDOUT.
	out, err := exec.Command("go", "install", p.ImportPath).Output()
	if err != nil {
		return err
	}
	if *g_debug {
		log.Printf("build out: %s\n", string(out))
	}
	return nil
}

// executes autobuild function if autobuild option is enabled, logs error and
// ignores it
func try_autobuild(p *build.Package) {
	if g_config.Autobuild {
		err := autobuild(p)
		if err != nil && *g_debug {
			log.Printf("Autobuild error: %s\n", err)
		}
	}
}

func log_found_package_maybe(imp, pkgpath string) {
	if *g_debug {
		log.Printf("Found %q at %q\n", imp, pkgpath)
	}
}

func log_build_context(context *package_lookup_context) {
	log.Printf(" GOROOT: %s\n", context.GOROOT)
	log.Printf(" GOPATH: %s\n", context.GOPATH)
	log.Printf(" GOOS: %s\n", context.GOOS)
	log.Printf(" GOARCH: %s\n", context.GOARCH)
	log.Printf(" GBProjectRoot: %q\n", context.GBProjectRoot)
	log.Printf(" lib-path: %q\n", g_config.LibPath)
}

// find_global_file returns the file path of the compiled package corresponding to the specified
// import, and a boolean stating whether such path is valid.
// TODO: Return only one value, possibly empty string if not found.
func find_global_file(imp string, context *package_lookup_context) (string, bool) {
	// gocode synthetically generates the builtin package
	// "unsafe", since the "unsafe.a" package doesn't really exist.
	// Thus, when the user request for the package "unsafe" we
	// would return synthetic global file that would be used
	// just as a key name to find this synthetic package
	if imp == "unsafe" {
		return "unsafe", true
	}

	pkgfile := fmt.Sprintf("%s.a", imp)

	// if lib-path is defined, use it
	if g_config.LibPath != "" {
		for _, p := range filepath.SplitList(g_config.LibPath) {
			pkg_path := filepath.Join(p, pkgfile)
			if file_exists(pkg_path) {
				log_found_package_maybe(imp, pkg_path)
				return pkg_path, true
			}
			// Also check the relevant pkg/OS_ARCH dir for the libpath, if provided.
			pkgdir := fmt.Sprintf("%s_%s", context.GOOS, context.GOARCH)
			pkg_path = filepath.Join(p, "pkg", pkgdir, pkgfile)
			if file_exists(pkg_path) {
				log_found_package_maybe(imp, pkg_path)
				return pkg_path, true
			}
		}
	}

	// gb-specific lookup mode, only if the root dir was found
	if g_config.PackageLookupMode == "gb" && context.GBProjectRoot != "" {
		root := context.GBProjectRoot
		pkg_path := filepath.Join(root, "pkg", context.GOOS+"-"+context.GOARCH, pkgfile)
		if file_exists(pkg_path) {
			log_found_package_maybe(imp, pkg_path)
			return pkg_path, true
		}
	}

	if context.CurrentPackagePath != "" {
		// Try vendor path first, see GO15VENDOREXPERIMENT.
		// We don't check this environment variable however, seems like there is
		// almost no harm in doing so (well.. if you experiment with vendoring,
		// gocode will fail after enabling/disabling the flag, and you'll be
		// forced to get rid of vendor binaries). But asking users to set this
		// env var is up will bring more trouble. Because we also need to pass
		// it from client to server, make sure their editors set it, etc.
		// So, whatever, let's just pretend it's always on.
		package_path := context.CurrentPackagePath
		for {
			limp := filepath.Join(package_path, "vendor", imp)
			if p, err := context.Import(limp, "", build.AllowBinary|build.FindOnly); err == nil {
				try_autobuild(p)
				if file_exists(p.PkgObj) {
					log_found_package_maybe(imp, p.PkgObj)
					return p.PkgObj, true
				}
			}
			if package_path == "" {
				break
			}
			next_path := filepath.Dir(package_path)
			// let's protect ourselves from inf recursion here
			if next_path == package_path {
				break
			}
			package_path = next_path
		}
	}

	if p, err := context.Import(imp, "", build.AllowBinary|build.FindOnly); err == nil {
		try_autobuild(p)
		if file_exists(p.PkgObj) {
			log_found_package_maybe(imp, p.PkgObj)
			return p.PkgObj, true
		}
	}

	if *g_debug {
		log.Printf("Import path %q was not resolved\n", imp)
		log.Println("Gocode's build context is:")
		log_build_context(context)
	}
	return "", false
}

func package_name(file *ast.File) string {
	if file.Name != nil {
		return file.Name.Name
	}
	return ""
}

//-------------------------------------------------------------------------
// decl_cache
//
// Thread-safe collection of DeclFileCache entities.
//-------------------------------------------------------------------------

type package_lookup_context struct {
	build.Context
	GBProjectRoot      string
	CurrentPackagePath string
}

type decl_cache struct {
	cache   map[string]*decl_file_cache
	context *package_lookup_context
	sync.Mutex
}

func new_decl_cache(context *package_lookup_context) *decl_cache {
	return &decl_cache{
		cache:   make(map[string]*decl_file_cache),
		context: context,
	}
}

func (c *decl_cache) get(filename string) *decl_file_cache {
	c.Lock()
	defer c.Unlock()

	f, ok := c.cache[filename]
	if !ok {
		f = new_decl_file_cache(filename, c.context)
		c.cache[filename] = f
	}
	return f
}

func (c *decl_cache) get_and_update(filename string) *decl_file_cache {
	f := c.get(filename)
	f.update()
	return f
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/napsy/httpauth"
	"golang.org/x/crypto/bcrypt"
)

var (
	aaa        httpauth.Authorizer
	userDb     httpauth.GobFileAuthBackend
	bruteCount = 0
	bruteLimit = 10
	authLock   bool
	bruteLock  sync.Mutex
)

func passwordSalt(password string) ([]byte, []byte, error) {
	salt := make([]byte, 8)
	if _, err := rand.Read(salt); err != nil {
		return nil, nil, err
	}
	salted := []byte{}
	salted = append(salted, salt...)
	salted = append(salted, []byte(password)...)
	return salt, salted, nil
}

func createUser(username, password, email, role string) error {
	salt, salted, err := passwordSalt(password)
	if err != nil {
		return err
	}
	pwd, err := bcrypt.GenerateFromPassword(salted, bcrypt.DefaultCost)
	if err != nil {
		return err
	}
	if err := userDb.SaveUser(httpauth.UserData{Username: username, Hash: pwd, Salt: salt, Role: role}); err != nil {
		return err
	}
	return nil
}

func logout(rw http.ResponseWriter, req *http.Request) {
	if err := aaa.Logout(rw, req); err != nil {
		log.Printf("%v", err)
		return
	}
	http.Redirect(rw, req, "/", http.StatusSeeOther)
}

func login(rw http.ResponseWriter, req *http.Request) {
	if !canAuth() {
		http.Error(rw, "429 Too Many Requests", http.StatusTooManyRequests)
		return
	}
	brutePlus()
	if req.Method == "GET" {
		t, _ := template.ParseFiles("templates/login.html")
		t.Execute(rw, nil)
		return
	}
	username := req.PostFormValue("username")
	password := req.PostFormValue("password")

	if err := aaa.Login(rw, req, username, password, "/"); err != nil && err.Error() == "already authenticated" {
		http.Redirect(rw, req, "/", http.StatusSeeOther)
	} else if err != nil {
		fmt.Println(err)
		http.Redirect(rw, req, "/login", http.StatusSeeOther)
	}
}

func requiresAuth(role string, handleFunc func(w http.ResponseWriter, r *http.Request)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		if err := aaa.AuthorizeRole(w, r, role, true); err != nil {
			fmt.Println(err)
			http.Redirect(w, r, "/login", http.StatusSeeOther)
		}
		handleFunc(w, r)
	}
}

func brutePlus() {
	bruteLock.Lock()
	defer bruteLock.Unlock()
	bruteCount++
}

func canAuth() bool {
	bruteLock.Lock()
	defer bruteLock.Unlock()
	return !authLock
}

// brute-force password guessing protection
func bruteProtect() {
	for {
		time.Sleep(10 * time.Second)
		bruteLock.Lock()
		if bruteCount >= bruteLimit {
			authLock = true
		} else {
			authLock = false
		}
		bruteCount = 0
		bruteLock.Unlock()
	}
}
func initAuth(adminPwd string) {
	var err error
	userDb, err = httpauth.NewGobFileAuthBackend("users.db")
	if err != nil {
		panic(err)
	}

	if len(adminPwd) > 0 {
		if err := ioutil.WriteFile("users.db", []byte{}, 0666); err != nil {
			log.Printf("Couldn't create 'users.db' file: %v", err)
			os.Exit(1)
		}

		if err := createUser("admin", adminPwd, "root@localhost", "admin"); err != nil {
			log.Printf("Couldn't create admin user: %v", err)
			os.Exit(1)
		}
		log.Printf("Created admin user with password %q ...", adminPwd)
	}
	roles := make(map[string]httpauth.Role)
	roles["user"] = 30
	roles["admin"] = 80
	aaa, err = httpauth.NewAuthorizer(userDb, []byte("cookfwedfie-encryption-key"), "user", roles)
	if err != nil {
		panic(err)
	}

	go bruteProtect()
}
<commit_msg>wijiwiki: search users.db inside data/ directory<commit_after>package main

import (
	"crypto/rand"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/napsy/httpauth"
	"golang.org/x/crypto/bcrypt"
)

var (
	aaa        httpauth.Authorizer
	userDb     httpauth.GobFileAuthBackend
	bruteCount = 0
	bruteLimit = 10
	authLock   bool
	bruteLock  sync.Mutex
)

func passwordSalt(password string) ([]byte, []byte, error) {
	salt := make([]byte, 8)
	if _, err := rand.Read(salt); err != nil {
		return nil, nil, err
	}
	salted := []byte{}
	salted = append(salted, salt...)
	salted = append(salted, []byte(password)...)
	return salt, salted, nil
}

func createUser(username, password, email, role string) error {
	salt, salted, err := passwordSalt(password)
	if err != nil {
		return err
	}
	pwd, err := bcrypt.GenerateFromPassword(salted, bcrypt.DefaultCost)
	if err != nil {
		return err
	}
	if err := userDb.SaveUser(httpauth.UserData{Username: username, Hash: pwd, Salt: salt, Role: role}); err != nil {
		return err
	}
	return nil
}

func logout(rw http.ResponseWriter, req *http.Request) {
	if err := aaa.Logout(rw, req); err != nil {
		log.Printf("%v", err)
		return
	}
	http.Redirect(rw, req, "/", http.StatusSeeOther)
}

func login(rw http.ResponseWriter, req *http.Request) {
	if !canAuth() {
		http.Error(rw, "429 Too Many Requests", http.StatusTooManyRequests)
		return
	}
	brutePlus()
	if req.Method == "GET" {
		t, _ := template.ParseFiles("templates/login.html")
		t.Execute(rw, nil)
		return
	}
	username := req.PostFormValue("username")
	password := req.PostFormValue("password")

	if err := aaa.Login(rw, req, username, password, "/"); err != nil && err.Error() == "already authenticated" {
		http.Redirect(rw, req, "/", http.StatusSeeOther)
	} else if err != nil {
		fmt.Println(err)
		http.Redirect(rw, req, "/login", http.StatusSeeOther)
	}
}

func requiresAuth(role string, handleFunc func(w http.ResponseWriter, r *http.Request)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		if err := aaa.AuthorizeRole(w, r, role, true); err != nil {
			fmt.Println(err)
			http.Redirect(w, r, "/login", http.StatusSeeOther)
		}
		handleFunc(w, r)
	}
}

func brutePlus() {
	bruteLock.Lock()
	defer bruteLock.Unlock()
	bruteCount++
}

func canAuth() bool {
	bruteLock.Lock()
	defer bruteLock.Unlock()
	return !authLock
}

// brute-force password guessing protection
func bruteProtect() {
	for {
		time.Sleep(10 * time.Second)
		bruteLock.Lock()
		if bruteCount >= bruteLimit {
			authLock = true
		} else {
			authLock = false
		}
		bruteCount = 0
		bruteLock.Unlock()
	}
}
func initAuth(adminPwd string) {
	var err error
	userDb, err = httpauth.NewGobFileAuthBackend("data/users.db")
	if err != nil {
		panic(err)
	}

	if len(adminPwd) > 0 {
		if err := ioutil.WriteFile("data/users.db", []byte{}, 0666); err != nil {
			log.Printf("Couldn't create 'data/users.db' file: %v", err)
			os.Exit(1)
		}

		if err := createUser("admin", adminPwd, "root@localhost", "admin"); err != nil {
			log.Printf("Couldn't create admin user: %v", err)
			os.Exit(1)
		}
		log.Printf("Created admin user with password %q ...", adminPwd)
	}
	roles := make(map[string]httpauth.Role)
	roles["user"] = 30
	roles["admin"] = 80
	aaa, err = httpauth.NewAuthorizer(userDb, []byte("cookfwedfie-encryption-key"), "user", roles)
	if err != nil {
		panic(err)
	}

	go bruteProtect()
}
<|endoftext|>


<commit_before>package repl

import (
	"fmt"
	"io"
	"os"

	"github.com/chzyer/readline"
	"github.com/mattn/go-shellwords"
)

type (
	// REPL is a read-eval-print loop used to create a simple, minimalistic,
	// easy-to-use command line interface for masterkey.
	REPL struct {
		prompt          string
		commands        map[string]Command
		prefixCompleter *readline.PrefixCompleter
		input           io.Reader
		output          io.Writer
		rl              *readline.Instance
		stopfunc        func()
	}

	// Command is a command that can be registered with the REPL. It consists
	// of a name, an action that is run when the name is input to the REPL, and
	// a usage string.
	Command struct {
		Name   string
		Action ActionFunc
		Usage  string
	}

	// ActionFunc defines the signature of an action associated with a command.
	// Actions take on parameter, a slice of strings, representing the arguments
	// passed to the command. Actions should return a string representing the
	// result of the action, or an error if the action fails.
	ActionFunc func([]string) (string, error)
)

// New instantiates a new REPL using the provided `prompt`.
func New(prompt string) *REPL {
	return &REPL{
		commands:        make(map[string]Command),
		prefixCompleter: readline.NewPrefixCompleter(readline.PcItem("exit"), readline.PcItem("help")),
		prompt:          prompt,
		input:           os.Stdin,
		output:          os.Stdout,
	}
}

// OnStop registers a function to be called when the REPL stops.
func (r *REPL) OnStop(sf func()) {
	r.stopfunc = sf
}

// Usage returns the usage for every command in the REPL.
func (r *REPL) Usage() string {
	printstring := ""
	for _, command := range r.commands {
		printstring += "\n"
		printstring += command.Usage + "\n"
	}
	return printstring
}

// AddCommand registers the command provided in `cmd` with the REPL.
func (r *REPL) AddCommand(cmd Command) {
	r.commands[cmd.Name] = cmd

	var completers []readline.PrefixCompleterInterface
	for name := range r.commands {
		completers = append(completers, readline.PcItem(name))
	}

	r.prefixCompleter = readline.NewPrefixCompleter(completers...)
}

// eval evaluates a line that was input to the REPL.
func (r *REPL) eval(line string) (string, error) {
	args, err := shellwords.Parse(line)
	if err != nil {
		return "", err
	}
	command := args[0]

	if command == "help" {
		return r.Usage(), nil
	}

	if command == "exit" {
		if r.rl != nil {
			return "", r.rl.Close()
		}
	}

	cmd, exists := r.commands[command]
	if !exists {
		return "", fmt.Errorf("command not recognized. Type `help` for a list of commands.")
	}

	res, err := cmd.Action(args[1:])
	if err != nil {
		return "", err
	}

	return res, nil
}

// Loop starts the Read-Eval-Print loop.
func (r *REPL) Loop() error {
	rl, err := readline.NewEx(&readline.Config{
		Prompt:       r.prompt,
		AutoComplete: r.prefixCompleter,
	})

	if err != nil {
		return err
	}
	defer rl.Close()
	r.rl = rl

	for {
		line, err := r.rl.Readline()
		if err != nil {
			if err == readline.ErrInterrupt && r.stopfunc != nil {
				r.stopfunc()
			}
			break
		}
		if line != "" {
			res, err := r.eval(line)
			if err != nil {
				fmt.Fprintln(r.output, err.Error())
				continue
			}
			fmt.Fprint(r.output, res)
		}
	}
	return nil
}
<commit_msg>add `clear` command<commit_after>package repl

import (
	"fmt"
	"io"
	"os"

	"github.com/chzyer/readline"
	"github.com/mattn/go-shellwords"
)

type (
	// REPL is a read-eval-print loop used to create a simple, minimalistic,
	// easy-to-use command line interface for masterkey.
	REPL struct {
		prompt          string
		commands        map[string]Command
		prefixCompleter *readline.PrefixCompleter
		input           io.Reader
		output          io.Writer
		rl              *readline.Instance
		stopfunc        func()
	}

	// Command is a command that can be registered with the REPL. It consists
	// of a name, an action that is run when the name is input to the REPL, and
	// a usage string.
	Command struct {
		Name   string
		Action ActionFunc
		Usage  string
	}

	// ActionFunc defines the signature of an action associated with a command.
	// Actions take on parameter, a slice of strings, representing the arguments
	// passed to the command. Actions should return a string representing the
	// result of the action, or an error if the action fails.
	ActionFunc func([]string) (string, error)
)

// New instantiates a new REPL using the provided `prompt`.
func New(prompt string) *REPL {
	return &REPL{
		commands:        make(map[string]Command),
		prefixCompleter: readline.NewPrefixCompleter(readline.PcItem("exit"), readline.PcItem("help")),
		prompt:          prompt,
		input:           os.Stdin,
		output:          os.Stdout,
	}
}

// OnStop registers a function to be called when the REPL stops.
func (r *REPL) OnStop(sf func()) {
	r.stopfunc = sf
}

// Usage returns the usage for every command in the REPL.
func (r *REPL) Usage() string {
	printstring := ""
	for _, command := range r.commands {
		printstring += command.Usage + "\n"
	}
	return printstring
}

// AddCommand registers the command provided in `cmd` with the REPL.
func (r *REPL) AddCommand(cmd Command) {
	r.commands[cmd.Name] = cmd

	var completers []readline.PrefixCompleterInterface
	for name := range r.commands {
		completers = append(completers, readline.PcItem(name))
	}

	r.prefixCompleter = readline.NewPrefixCompleter(completers...)
}

// eval evaluates a line that was input to the REPL.
func (r *REPL) eval(line string) (string, error) {
	args, err := shellwords.Parse(line)
	if err != nil {
		return "", err
	}
	command := args[0]

	if command == "help" {
		return r.Usage(), nil
	}

	if command == "exit" {
		return "", r.rl.Close()
	}

	if command == "clear" {
		_, err := readline.ClearScreen(r.output)
		if err != nil {
			return "", err
		}
		return "terminal cleared", nil
	}

	cmd, exists := r.commands[command]
	if !exists {
		return "", fmt.Errorf("command not recognized. Type `help` for a list of commands.")
	}

	res, err := cmd.Action(args[1:])
	if err != nil {
		return "", err
	}

	return res, nil
}

// Loop starts the Read-Eval-Print loop.
func (r *REPL) Loop() error {
	rl, err := readline.NewEx(&readline.Config{
		Prompt:       r.prompt,
		AutoComplete: r.prefixCompleter,
	})

	if err != nil {
		return err
	}
	defer rl.Close()
	r.rl = rl

	for {
		line, err := r.rl.Readline()
		if err != nil {
			if err == readline.ErrInterrupt && r.stopfunc != nil {
				r.stopfunc()
			}
			break
		}
		if line != "" {
			res, err := r.eval(line)
			if err != nil {
				fmt.Fprintln(r.output, err.Error())
				continue
			}
			fmt.Fprint(r.output, res)
		}
	}
	return nil
}
<|endoftext|>