<commit_before>/*
Copyright 2019 The Kubernetes Authors All rights reserved.

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 metricshandler

import (
	"compress/gzip"
	"context"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/pkg/errors"
	appsv1 "k8s.io/api/apps/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
	"k8s.io/klog/v2"

	ksmtypes "k8s.io/kube-state-metrics/v2/pkg/builder/types"
	metricsstore "k8s.io/kube-state-metrics/v2/pkg/metrics_store"
	"k8s.io/kube-state-metrics/v2/pkg/options"
)

// MetricsHandler is a http.Handler that exposes the main kube-state-metrics
// /metrics endpoint. It allows concurrent reconfiguration at runtime.
type MetricsHandler struct {
	opts               *options.Options
	kubeClient         kubernetes.Interface
	storeBuilder       ksmtypes.BuilderInterface
	enableGZIPEncoding bool

	cancel func()

	// mtx protects metricsWriters, curShard, and curTotalShards
	mtx            *sync.RWMutex
	metricsWriters []metricsstore.MetricsWriter
	curShard       int32
	curTotalShards int
}

// New creates and returns a new MetricsHandler with the given options.
func New(opts *options.Options, kubeClient kubernetes.Interface, storeBuilder ksmtypes.BuilderInterface, enableGZIPEncoding bool) *MetricsHandler {
	return &MetricsHandler{
		opts:               opts,
		kubeClient:         kubeClient,
		storeBuilder:       storeBuilder,
		enableGZIPEncoding: enableGZIPEncoding,
		mtx:                &sync.RWMutex{},
	}
}

// ConfigureSharding (re-)configures sharding. Re-configuration can be done
// concurrently.
func (m *MetricsHandler) ConfigureSharding(ctx context.Context, shard int32, totalShards int) {
	m.mtx.Lock()
	defer m.mtx.Unlock()

	if m.cancel != nil {
		m.cancel()
	}
	if totalShards != 1 {
		klog.InfoS("Configuring sharding of this instance to be shard index (zero-indexed) out of total shards", "shard", shard, "totalShards", totalShards)
	}
	ctx, m.cancel = context.WithCancel(ctx)
	m.storeBuilder.WithSharding(shard, totalShards)
	m.storeBuilder.WithContext(ctx)
	m.metricsWriters = m.storeBuilder.Build()
	m.curShard = shard
	m.curTotalShards = totalShards
}

// Run configures the MetricsHandler's sharding and if autosharding is enabled
// re-configures sharding on re-sharding events. Run should only be called
// once.
func (m *MetricsHandler) Run(ctx context.Context) error {
	autoSharding := len(m.opts.Pod) > 0 && len(m.opts.Namespace) > 0

	if !autoSharding {
		klog.InfoS("Autosharding disabled")
		m.ConfigureSharding(ctx, m.opts.Shard, m.opts.TotalShards)
		<-ctx.Done()
		return ctx.Err()
	}

	klog.InfoS("Autosharding enabled with pod pod_namespace", "pod", m.opts.Pod, "namespace", m.opts.Namespace)
	klog.InfoS("Auto detecting sharding settings")
	ss, err := detectStatefulSet(m.kubeClient, m.opts.Pod, m.opts.Namespace)
	if err != nil {
		return errors.Wrap(err, "detect StatefulSet")
	}
	statefulSetName := ss.Name

	labelSelectorOptions := func(o *metav1.ListOptions) {
		o.LabelSelector = fields.SelectorFromSet(ss.Labels).String()
	}

	i := cache.NewSharedIndexInformer(
		cache.NewFilteredListWatchFromClient(m.kubeClient.AppsV1().RESTClient(), "statefulsets", m.opts.Namespace, labelSelectorOptions),
		&appsv1.StatefulSet{}, 0, cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc},
	)
	i.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(o interface{}) {
			ss := o.(*appsv1.StatefulSet)
			if ss.Name != statefulSetName {
				return
			}

			shard, totalShards, err := shardingSettingsFromStatefulSet(ss, m.opts.Pod)
			if err != nil {
				klog.ErrorS(err, "Detected sharding settings from StatefulSet")
				return
			}

			m.mtx.RLock()
			shardingUnchanged := m.curShard == shard && m.curTotalShards == totalShards
			m.mtx.RUnlock()

			if shardingUnchanged {
				return
			}

			m.ConfigureSharding(ctx, shard, totalShards)
		},
		UpdateFunc: func(oldo, curo interface{}) {
			old := oldo.(*appsv1.StatefulSet)
			cur := curo.(*appsv1.StatefulSet)
			if cur.Name != statefulSetName {
				return
			}

			if old.ResourceVersion == cur.ResourceVersion {
				return
			}

			shard, totalShards, err := shardingSettingsFromStatefulSet(cur, m.opts.Pod)
			if err != nil {
				klog.ErrorS(err, "Detected sharding settings from StatefulSet")
				return
			}

			m.mtx.RLock()
			shardingUnchanged := m.curShard == shard && m.curTotalShards == totalShards
			m.mtx.RUnlock()

			if shardingUnchanged {
				return
			}

			m.ConfigureSharding(ctx, shard, totalShards)
		},
	})
	go i.Run(ctx.Done())
	if !cache.WaitForCacheSync(ctx.Done(), i.HasSynced) {
		return errors.New("waiting for informer cache to sync failed")
	}
	<-ctx.Done()
	return ctx.Err()
}

// ServeHTTP implements the http.Handler interface. It writes all generated
// metrics to the response body.
func (m *MetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	m.mtx.RLock()
	defer m.mtx.RUnlock()
	resHeader := w.Header()
	var writer io.Writer = w

	resHeader.Set("Content-Type", `text/plain; version=`+"0.0.4")

	if m.enableGZIPEncoding {
		// Gzip response if requested. Taken from
		// github.com/prometheus/client_golang/prometheus/promhttp.decorateWriter.
		reqHeader := r.Header.Get("Accept-Encoding")
		parts := strings.Split(reqHeader, ",")
		for _, part := range parts {
			part = strings.TrimSpace(part)
			if part == "gzip" || strings.HasPrefix(part, "gzip;") {
				writer = gzip.NewWriter(writer)
				resHeader.Set("Content-Encoding", "gzip")
			}
		}
	}

	for _, w := range m.metricsWriters {
		w.WriteAll(writer)
	}

	// In case we gzipped the response, we have to close the writer.
	if closer, ok := writer.(io.Closer); ok {
		closer.Close()
	}
}

func shardingSettingsFromStatefulSet(ss *appsv1.StatefulSet, podName string) (nominal int32, totalReplicas int, err error) {
	nominal, err = detectNominalFromPod(ss.Name, podName)
	if err != nil {
		return 0, 0, errors.Wrap(err, "detecting Pod nominal")
	}

	totalReplicas = 1
	replicas := ss.Spec.Replicas
	if replicas != nil {
		totalReplicas = int(*replicas)
	}

	return nominal, totalReplicas, nil
}

func detectNominalFromPod(statefulSetName, podName string) (int32, error) {
	nominalString := strings.TrimPrefix(podName, statefulSetName+"-")
	nominal, err := strconv.Atoi(nominalString)
	if err != nil {
		return 0, errors.Wrapf(err, "failed to detect shard index for Pod %s of StatefulSet %s, parsed %s", podName, statefulSetName, nominalString)
	}

	return int32(nominal), nil
}

func detectStatefulSet(kubeClient kubernetes.Interface, podName, namespaceName string) (*appsv1.StatefulSet, error) {
	p, err := kubeClient.CoreV1().Pods(namespaceName).Get(context.TODO(), podName, metav1.GetOptions{})
	if err != nil {
		return nil, errors.Wrapf(err, "retrieve pod %s for sharding", podName)
	}

	owners := p.GetOwnerReferences()
	for _, o := range owners {
		if o.APIVersion != "apps/v1" || o.Kind != "StatefulSet" || o.Controller == nil || !*o.Controller {
			continue
		}

		ss, err := kubeClient.AppsV1().StatefulSets(namespaceName).Get(context.TODO(), o.Name, metav1.GetOptions{})
		if err != nil {
			return nil, errors.Wrapf(err, "retrieve shard's StatefulSet: %s/%s", namespaceName, o.Name)
		}

		return ss, nil
	}

	return nil, errors.Errorf("no suitable statefulset found for auto detecting sharding for Pod %s/%s", namespaceName, podName)
}
<commit_msg>Apply suggestions from code review<commit_after>/*
Copyright 2019 The Kubernetes Authors All rights reserved.

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 metricshandler

import (
	"compress/gzip"
	"context"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/pkg/errors"
	appsv1 "k8s.io/api/apps/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
	"k8s.io/klog/v2"

	ksmtypes "k8s.io/kube-state-metrics/v2/pkg/builder/types"
	metricsstore "k8s.io/kube-state-metrics/v2/pkg/metrics_store"
	"k8s.io/kube-state-metrics/v2/pkg/options"
)

// MetricsHandler is a http.Handler that exposes the main kube-state-metrics
// /metrics endpoint. It allows concurrent reconfiguration at runtime.
type MetricsHandler struct {
	opts               *options.Options
	kubeClient         kubernetes.Interface
	storeBuilder       ksmtypes.BuilderInterface
	enableGZIPEncoding bool

	cancel func()

	// mtx protects metricsWriters, curShard, and curTotalShards
	mtx            *sync.RWMutex
	metricsWriters []metricsstore.MetricsWriter
	curShard       int32
	curTotalShards int
}

// New creates and returns a new MetricsHandler with the given options.
func New(opts *options.Options, kubeClient kubernetes.Interface, storeBuilder ksmtypes.BuilderInterface, enableGZIPEncoding bool) *MetricsHandler {
	return &MetricsHandler{
		opts:               opts,
		kubeClient:         kubeClient,
		storeBuilder:       storeBuilder,
		enableGZIPEncoding: enableGZIPEncoding,
		mtx:                &sync.RWMutex{},
	}
}

// ConfigureSharding (re-)configures sharding. Re-configuration can be done
// concurrently.
func (m *MetricsHandler) ConfigureSharding(ctx context.Context, shard int32, totalShards int) {
	m.mtx.Lock()
	defer m.mtx.Unlock()

	if m.cancel != nil {
		m.cancel()
	}
	if totalShards != 1 {
		klog.InfoS("Configuring sharding of this instance to be shard index (zero-indexed) out of total shards", "shard", shard, "totalShards", totalShards)
	}
	ctx, m.cancel = context.WithCancel(ctx)
	m.storeBuilder.WithSharding(shard, totalShards)
	m.storeBuilder.WithContext(ctx)
	m.metricsWriters = m.storeBuilder.Build()
	m.curShard = shard
	m.curTotalShards = totalShards
}

// Run configures the MetricsHandler's sharding and if autosharding is enabled
// re-configures sharding on re-sharding events. Run should only be called
// once.
func (m *MetricsHandler) Run(ctx context.Context) error {
	autoSharding := len(m.opts.Pod) > 0 && len(m.opts.Namespace) > 0

	if !autoSharding {
		klog.InfoS("Autosharding disabled")
		m.ConfigureSharding(ctx, m.opts.Shard, m.opts.TotalShards)
		<-ctx.Done()
		return ctx.Err()
	}

	klog.InfoS("Autosharding enabled with pod", "pod", klog.KRef(m.opts.Namespace, m.opts.Pod))
	klog.InfoS("Auto detecting sharding settings")
	ss, err := detectStatefulSet(m.kubeClient, m.opts.Pod, m.opts.Namespace)
	if err != nil {
		return errors.Wrap(err, "detect StatefulSet")
	}
	statefulSetName := ss.Name

	labelSelectorOptions := func(o *metav1.ListOptions) {
		o.LabelSelector = fields.SelectorFromSet(ss.Labels).String()
	}

	i := cache.NewSharedIndexInformer(
		cache.NewFilteredListWatchFromClient(m.kubeClient.AppsV1().RESTClient(), "statefulsets", m.opts.Namespace, labelSelectorOptions),
		&appsv1.StatefulSet{}, 0, cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc},
	)
	i.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(o interface{}) {
			ss := o.(*appsv1.StatefulSet)
			if ss.Name != statefulSetName {
				return
			}

			shard, totalShards, err := shardingSettingsFromStatefulSet(ss, m.opts.Pod)
			if err != nil {
				klog.ErrorS(err, "Detected sharding settings from StatefulSet")
				return
			}

			m.mtx.RLock()
			shardingUnchanged := m.curShard == shard && m.curTotalShards == totalShards
			m.mtx.RUnlock()

			if shardingUnchanged {
				return
			}

			m.ConfigureSharding(ctx, shard, totalShards)
		},
		UpdateFunc: func(oldo, curo interface{}) {
			old := oldo.(*appsv1.StatefulSet)
			cur := curo.(*appsv1.StatefulSet)
			if cur.Name != statefulSetName {
				return
			}

			if old.ResourceVersion == cur.ResourceVersion {
				return
			}

			shard, totalShards, err := shardingSettingsFromStatefulSet(cur, m.opts.Pod)
			if err != nil {
				klog.ErrorS(err, "Detected sharding settings from StatefulSet")
				return
			}

			m.mtx.RLock()
			shardingUnchanged := m.curShard == shard && m.curTotalShards == totalShards
			m.mtx.RUnlock()

			if shardingUnchanged {
				return
			}

			m.ConfigureSharding(ctx, shard, totalShards)
		},
	})
	go i.Run(ctx.Done())
	if !cache.WaitForCacheSync(ctx.Done(), i.HasSynced) {
		return errors.New("waiting for informer cache to sync failed")
	}
	<-ctx.Done()
	return ctx.Err()
}

// ServeHTTP implements the http.Handler interface. It writes all generated
// metrics to the response body.
func (m *MetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	m.mtx.RLock()
	defer m.mtx.RUnlock()
	resHeader := w.Header()
	var writer io.Writer = w

	resHeader.Set("Content-Type", `text/plain; version=`+"0.0.4")

	if m.enableGZIPEncoding {
		// Gzip response if requested. Taken from
		// github.com/prometheus/client_golang/prometheus/promhttp.decorateWriter.
		reqHeader := r.Header.Get("Accept-Encoding")
		parts := strings.Split(reqHeader, ",")
		for _, part := range parts {
			part = strings.TrimSpace(part)
			if part == "gzip" || strings.HasPrefix(part, "gzip;") {
				writer = gzip.NewWriter(writer)
				resHeader.Set("Content-Encoding", "gzip")
			}
		}
	}

	for _, w := range m.metricsWriters {
		w.WriteAll(writer)
	}

	// In case we gzipped the response, we have to close the writer.
	if closer, ok := writer.(io.Closer); ok {
		closer.Close()
	}
}

func shardingSettingsFromStatefulSet(ss *appsv1.StatefulSet, podName string) (nominal int32, totalReplicas int, err error) {
	nominal, err = detectNominalFromPod(ss.Name, podName)
	if err != nil {
		return 0, 0, errors.Wrap(err, "detecting Pod nominal")
	}

	totalReplicas = 1
	replicas := ss.Spec.Replicas
	if replicas != nil {
		totalReplicas = int(*replicas)
	}

	return nominal, totalReplicas, nil
}

func detectNominalFromPod(statefulSetName, podName string) (int32, error) {
	nominalString := strings.TrimPrefix(podName, statefulSetName+"-")
	nominal, err := strconv.Atoi(nominalString)
	if err != nil {
		return 0, errors.Wrapf(err, "failed to detect shard index for Pod %s of StatefulSet %s, parsed %s", podName, statefulSetName, nominalString)
	}

	return int32(nominal), nil
}

func detectStatefulSet(kubeClient kubernetes.Interface, podName, namespaceName string) (*appsv1.StatefulSet, error) {
	p, err := kubeClient.CoreV1().Pods(namespaceName).Get(context.TODO(), podName, metav1.GetOptions{})
	if err != nil {
		return nil, errors.Wrapf(err, "retrieve pod %s for sharding", podName)
	}

	owners := p.GetOwnerReferences()
	for _, o := range owners {
		if o.APIVersion != "apps/v1" || o.Kind != "StatefulSet" || o.Controller == nil || !*o.Controller {
			continue
		}

		ss, err := kubeClient.AppsV1().StatefulSets(namespaceName).Get(context.TODO(), o.Name, metav1.GetOptions{})
		if err != nil {
			return nil, errors.Wrapf(err, "retrieve shard's StatefulSet: %s/%s", namespaceName, o.Name)
		}

		return ss, nil
	}

	return nil, errors.Errorf("no suitable statefulset found for auto detecting sharding for Pod %s/%s", namespaceName, podName)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"strconv"
	"unicode/utf8"

	"../urxvtermbox"
	"github.com/mpasternacki/termbox-go"
)

type UIState struct {
	Desktops []WMDesktop
	Selected int
	Height   int
	Width    int
}

func NewUIState(desks []WMDesktop) UIState {
	st := UIState{
		Desktops: desks,
	}

	for i, desk := range desks {
		if desk.IsCurrent {
			st.Selected = i
		}
		if nw := len(desk.Windows); nw > st.Height {
			st.Height = nw
		}
		for _, win := range desk.Windows {
			if nw := utf8.RuneCountInString(win.Name); nw > st.Width {
				st.Width = nw
			}
		}
	}

	return st
}

func (ui *UIState) Desk() *WMDesktop {
	if ui.Selected < 0 {
		return nil
	}
	return &ui.Desktops[ui.Selected]
}

func (ui *UIState) Prev() {
	for cur := ui.Selected - 1; cur > 0; cur-- {
		if ui.Desktops[cur].IsVisible() {
			ui.Selected = cur
			return
		}
	}
}

func (ui *UIState) Next() {
	for cur := ui.Selected + 1; cur < len(ui.Desktops); cur++ {
		if ui.Desktops[cur].IsVisible() {
			ui.Selected = cur
			return
		}
	}
}

var indexDigits = []rune{'⁰', '¹', '²', '³', '⁴', '⁵', '⁶', '⁷', '⁸', '⁹'}

func (ui *UIState) Draw() {
	cols, rows := termbox.Size()
	fgFrame := termbox.ColorYellow
	fgTitle := termbox.ColorGreen

	if rows < ui.Height+4 {
		panic("Too little rows!")
	}
	if cols < ui.Width+2 {
		panic("Too little cols!")
	}

	// Tab bar
	termbox.SetCell(0, 2, '╭', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	col := 1
	for i, desk := range ui.Desktops {
		if !desk.IsVisible() {
			continue
		}

		fg := fgTitle // termbox.ColorDefault
		extra := termbox.Attribute(0)

		if i == ui.Selected {
			// fg = termbox.ColorWhite
		}

		if desk.IsUrgent {
			fg = termbox.ColorRed
		}

		if desk.IsCurrent {
			extra = termbox.AttrUnderline
		}

		if i == ui.Selected {
			extra = extra | termbox.AttrBold
		}

		if i < ui.Selected {
			termbox.SetCell(col, 0, '╭', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else if i == ui.Selected {
			termbox.SetCell(col, 0, '╭', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 2, '╯', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++

		index := ' '
		if i < len(indexDigits) {
			index = indexDigits[i]
		}
		if i == ui.Selected {
			termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, index, fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, index, fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++

		for _, ch := range desk.Name {
			termbox.SetCell(col, 1, ch, fg|extra, termbox.ColorDefault)
			if i == ui.Selected {
				termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
				termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
			} else {
				termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
				termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			}
			col++
		}

		if i == ui.Selected {
			termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++

		for _, ch := range strconv.Itoa(len(desk.Windows)) {
			termbox.SetCell(col, 1, ch, fgFrame, termbox.ColorDefault)
			if i == ui.Selected {
				termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
				termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
			} else {
				termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
				termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			}
			col++
		}

		if i > ui.Selected {
			termbox.SetCell(col, 0, '╮', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else if i == ui.Selected {
			termbox.SetCell(col, 0, '╮', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 2, '╰', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++
	}

	if col > ui.Width {
		ui.Width = col
	}

	for ; col < ui.Width+1; col++ {
		termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}
	termbox.SetCell(ui.Width+1, 2, '╮', fgFrame|termbox.AttrBold, termbox.ColorDefault)

	// Window List
	desk := ui.Desk()
	for i, win := range desk.Windows {
		fg := termbox.ColorDefault
		extra := termbox.Attribute(0)

		if win.IsUrgent {
			fg = termbox.ColorRed
		}

		if win.IsActive {
			extra = termbox.AttrUnderline
		}

		if i == desk.Selected {
			fg = fg | termbox.AttrReverse
		}

		termbox.SetCell(0, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		col = 1
		for _, ch := range win.Name {
			termbox.SetCell(col, i+3, ch, fg|extra, termbox.ColorDefault)
			col++
		}

		for ; col < ui.Width+1; col++ {
			termbox.SetCell(col, i+3, ' ', fg, termbox.ColorDefault)
		}
		termbox.SetCell(ui.Width+1, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}

	for i := len(desk.Windows); i < ui.Height; i++ {
		termbox.SetCell(0, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		for j := 1; j < ui.Width+1; j++ {
			termbox.SetCell(j, i+3, ' ', termbox.ColorDefault, termbox.ColorDefault)
		}
		termbox.SetCell(ui.Width+1, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}

	termbox.SetCell(0, ui.Height+3, '╰', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	for j := 1; j < ui.Width+1; j++ {
		termbox.SetCell(j, ui.Height+3, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}
	termbox.SetCell(ui.Width+1, ui.Height+3, '╯', fgFrame|termbox.AttrBold, termbox.ColorDefault)

	termbox.Flush()
}

func (ui *UIState) Main() error {
	if fini, err := urxvtermbox.TermboxUrxvt(ui.Width+2, ui.Height+4); err != nil {
		return err
	} else {
		defer fini()
	}

	if err := termbox.Init(); err != nil {
		return err
	}
	defer termbox.Close()
	termbox.SetInputMode(termbox.InputEsc)

	ui.Draw()
	for {
		switch ev := termbox.PollEvent(); ev.Type {
		case termbox.EventKey:
			switch ev.Key {
			case termbox.KeyEsc:
				return cmdCancel
			case termbox.KeyArrowLeft:
				ui.Prev()
			case termbox.KeyArrowRight:
				ui.Next()
			case termbox.KeyArrowUp:
				ui.Desk().Prev()
			case termbox.KeyArrowDown:
				ui.Desk().Next()
			case termbox.KeyTab:
				ui.Desk().NextWrap()
			case termbox.KeyEnter:
				return nil
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				return cmdCloseWindow
			default:
				switch ev.Ch {
				case 'w': // ↑
					ui.Prev()
				case 's': // ↓
					ui.Next()
				case 'a': // ←
					ui.Desk().Prev()
				case 'd': // →
					ui.Desk().Next()
				case 'q': // Esc
					return cmdCancel
				case 'e': // move to active window
					for i, desk := range ui.Desktops {
						if desk.IsCurrent {
							ui.Selected = i
							for j, win := range desk.Windows {
								if win.IsActive {
									// desk is not a pointer
									ui.Desktops[i].Selected = j
									break
								}
							}
							break
						}
					}

				case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
					if desk := int(ev.Ch - '0'); desk < len(ui.Desktops) && len(ui.Desktops[desk].Windows) > 0 {
						ui.Selected = desk
					}
				case ' ': // same as Enter
					return nil
				}
			}
		case termbox.EventError:
			return ev.Err
		}
		ui.Draw()
	}

	return errors.New("CAN'T HAPPEN")
}
<commit_msg>switcher: ! binding to cycle through urgent windows<commit_after>package main

import (
	"errors"
	"strconv"
	"unicode/utf8"

	"../urxvtermbox"
	"github.com/mpasternacki/termbox-go"
)

type UIState struct {
	Desktops []WMDesktop
	Selected int
	Height   int
	Width    int
}

func NewUIState(desks []WMDesktop) UIState {
	st := UIState{
		Desktops: desks,
	}

	for i, desk := range desks {
		if desk.IsCurrent {
			st.Selected = i
		}
		if nw := len(desk.Windows); nw > st.Height {
			st.Height = nw
		}
		for _, win := range desk.Windows {
			if nw := utf8.RuneCountInString(win.Name); nw > st.Width {
				st.Width = nw
			}
		}
	}

	return st
}

func (ui *UIState) Desk() *WMDesktop {
	if ui.Selected < 0 {
		return nil
	}
	return &ui.Desktops[ui.Selected]
}

func (ui *UIState) Prev() {
	for cur := ui.Selected - 1; cur > 0; cur-- {
		if ui.Desktops[cur].IsVisible() {
			ui.Selected = cur
			return
		}
	}
}

func (ui *UIState) Next() {
	for cur := ui.Selected + 1; cur < len(ui.Desktops); cur++ {
		if ui.Desktops[cur].IsVisible() {
			ui.Selected = cur
			return
		}
	}
}

var indexDigits = []rune{'⁰', '¹', '²', '³', '⁴', '⁵', '⁶', '⁷', '⁸', '⁹'}

func (ui *UIState) Draw() {
	cols, rows := termbox.Size()
	fgFrame := termbox.ColorYellow
	fgTitle := termbox.ColorGreen

	if rows < ui.Height+4 {
		panic("Too little rows!")
	}
	if cols < ui.Width+2 {
		panic("Too little cols!")
	}

	// Tab bar
	termbox.SetCell(0, 2, '╭', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	col := 1
	for i, desk := range ui.Desktops {
		if !desk.IsVisible() {
			continue
		}

		fg := fgTitle // termbox.ColorDefault
		extra := termbox.Attribute(0)

		if i == ui.Selected {
			// fg = termbox.ColorWhite
		}

		if desk.IsUrgent {
			fg = termbox.ColorRed
		}

		if desk.IsCurrent {
			extra = termbox.AttrUnderline
		}

		if i == ui.Selected {
			extra = extra | termbox.AttrBold
		}

		if i < ui.Selected {
			termbox.SetCell(col, 0, '╭', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else if i == ui.Selected {
			termbox.SetCell(col, 0, '╭', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 2, '╯', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++

		index := ' '
		if i < len(indexDigits) {
			index = indexDigits[i]
		}
		if i == ui.Selected {
			termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, index, fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, index, fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++

		for _, ch := range desk.Name {
			termbox.SetCell(col, 1, ch, fg|extra, termbox.ColorDefault)
			if i == ui.Selected {
				termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
				termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
			} else {
				termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
				termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			}
			col++
		}

		if i == ui.Selected {
			termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++

		for _, ch := range strconv.Itoa(len(desk.Windows)) {
			termbox.SetCell(col, 1, ch, fgFrame, termbox.ColorDefault)
			if i == ui.Selected {
				termbox.SetCell(col, 0, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
				termbox.SetCell(col, 2, ' ', fgFrame, termbox.ColorDefault)
			} else {
				termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
				termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			}
			col++
		}

		if i > ui.Selected {
			termbox.SetCell(col, 0, '╮', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else if i == ui.Selected {
			termbox.SetCell(col, 0, '╮', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 1, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
			termbox.SetCell(col, 2, '╰', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		} else {
			termbox.SetCell(col, 0, '─', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 1, ' ', fgFrame, termbox.ColorDefault)
			termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		}
		col++
	}

	if col > ui.Width {
		ui.Width = col
	}

	for ; col < ui.Width+1; col++ {
		termbox.SetCell(col, 2, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}
	termbox.SetCell(ui.Width+1, 2, '╮', fgFrame|termbox.AttrBold, termbox.ColorDefault)

	// Window List
	desk := ui.Desk()
	for i, win := range desk.Windows {
		fg := termbox.ColorDefault
		extra := termbox.Attribute(0)

		if win.IsUrgent {
			fg = termbox.ColorRed
		}

		if win.IsActive {
			extra = termbox.AttrUnderline
		}

		if i == desk.Selected {
			fg = fg | termbox.AttrReverse
		}

		termbox.SetCell(0, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		col = 1
		for _, ch := range win.Name {
			termbox.SetCell(col, i+3, ch, fg|extra, termbox.ColorDefault)
			col++
		}

		for ; col < ui.Width+1; col++ {
			termbox.SetCell(col, i+3, ' ', fg, termbox.ColorDefault)
		}
		termbox.SetCell(ui.Width+1, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}

	for i := len(desk.Windows); i < ui.Height; i++ {
		termbox.SetCell(0, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
		for j := 1; j < ui.Width+1; j++ {
			termbox.SetCell(j, i+3, ' ', termbox.ColorDefault, termbox.ColorDefault)
		}
		termbox.SetCell(ui.Width+1, i+3, '│', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}

	termbox.SetCell(0, ui.Height+3, '╰', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	for j := 1; j < ui.Width+1; j++ {
		termbox.SetCell(j, ui.Height+3, '─', fgFrame|termbox.AttrBold, termbox.ColorDefault)
	}
	termbox.SetCell(ui.Width+1, ui.Height+3, '╯', fgFrame|termbox.AttrBold, termbox.ColorDefault)

	termbox.Flush()
}

func (ui *UIState) Main() error {
	if fini, err := urxvtermbox.TermboxUrxvt(ui.Width+2, ui.Height+4); err != nil {
		return err
	} else {
		defer fini()
	}

	if err := termbox.Init(); err != nil {
		return err
	}
	defer termbox.Close()
	termbox.SetInputMode(termbox.InputEsc)

	ui.Draw()
	for {
		switch ev := termbox.PollEvent(); ev.Type {
		case termbox.EventKey:
			switch ev.Key {
			case termbox.KeyEsc:
				return cmdCancel
			case termbox.KeyArrowLeft:
				ui.Prev()
			case termbox.KeyArrowRight:
				ui.Next()
			case termbox.KeyArrowUp:
				ui.Desk().Prev()
			case termbox.KeyArrowDown:
				ui.Desk().Next()
			case termbox.KeyTab:
				ui.Desk().NextWrap()
			case termbox.KeyEnter:
				return nil
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				return cmdCloseWindow
			default:
				switch ev.Ch {
				case 'w': // ↑
					ui.Prev()
				case 's': // ↓
					ui.Next()
				case 'a': // ←
					ui.Desk().Prev()
				case 'd': // →
					ui.Desk().Next()
				case 'q': // Esc
					return cmdCancel
				case 'e': // move to active window
					for i, desk := range ui.Desktops {
						if desk.IsCurrent {
							ui.Selected = i
							for j, win := range desk.Windows {
								if win.IsActive {
									// desk is not a pointer
									ui.Desktops[i].Selected = j
									break
								}
							}
							break
						}
					}

				case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
					if desk := int(ev.Ch - '0'); desk < len(ui.Desktops) && len(ui.Desktops[desk].Windows) > 0 {
						ui.Selected = desk
					}
				case ' ': // same as Enter
					return nil
				case '!':
					// Find next urgent window
					sxw := ui.Desk().Window().XWin
					d := ui.Selected
					w := ui.Desk().Selected
					for {
						w++                                   // next window
						if w >= len(ui.Desktops[d].Windows) { // next desktop
							w = 0
							for {
								d++
								if d >= len(ui.Desktops) {
									d = 0
								}
								if len(ui.Desktops[d].Windows) > 0 {
									break
								}
							}
						}
						win := ui.Desktops[d].Windows[w]
						if win.XWin == sxw {
							// We have wrapped around, let's break
							break
						}
						if win.IsUrgent {
							ui.Selected = d
							ui.Desk().Selected = w
							break
						}
					}
				}
			}
		case termbox.EventError:
			return ev.Err
		}
		ui.Draw()
	}

	return errors.New("CAN'T HAPPEN")
}
<|endoftext|>


<commit_before>package gopisysfs

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	// pollInterval is how long we should wait between naieve polling of files
	pollInterval = 20 * time.Millisecond
	// forever is about 100 years.
	forever = 100 * 365 * 24 * time.Hour
)

var rootpath = "/"

// setRoot is designed to be called by the test cases to exercise some hard-to-change things on an actual pi.
func setRoot(rt string) {
	rootpath = rt
}

// file gets a file path inside the /sys file system,
// but it can be hooked by the test cases to use a test filesystem instead of the real /sys
func file(paths ...string) string {
	path := filepath.Join(paths...)
	if !filepath.IsAbs(path) {
		path = filepath.Join(rootpath, path)
	}
	return path
}

// readFilePanic reads a file returning the contents as a string, and panics if it cannot be read
func readFilePanic(name string) string {
	data, err := readFile(name)
	if err != nil {
		log.Panicf("Unable to read file %v: %v", name, err)
	}
	return data
}

// awaitFileCreate establishes an asynchronous poll on a file location until it exists
// at which point the returned channel will return a nil on the channel. A non-nil indicates
// an error in the polling.
func awaitFileCreate(name string, timeout time.Duration) (<-chan error, error) {

	ret := make(chan error, 1)

	if checkFile(name) {
		ret <- nil
		return ret, nil
	}

	dir := filepath.Dir(name)
	if stat, err := os.Stat(dir); err != nil || !stat.IsDir() {
		if err != nil {
			return nil, fmt.Errorf("Unable to poll for a file in a nonexistent folder %v: %v", dir, err)
		}
		return nil, fmt.Errorf("Unable to poll for a file in a non-folder %v: %v", dir, stat)
	}

	// set up notification and timeout
	tout := time.After(timeout)
	// intervals at every poll cycle
	interval := time.NewTicker(pollInterval).C
	// naieve polling system
	go func() {
		for {

			if checkFile(name) {
				// Found it!
				ret <- nil
				return
			}

			select {
			case <-tout:
				ret <- fmt.Errorf("Timed out waiting for %v after %v", name, timeout)
				return
			case <-interval:
				// ignore specific event, check actual file later
			}
		}
	}()

	return ret, nil

}

// awaitFileRemove establishes an asynchronous poll on a file location until it is removed
// at which point the returned channel will return a nil on the channel. A non-nil indicates
// an error in the polling.
func awaitFileRemove(name string, timeout time.Duration) (<-chan error, error) {

	ret := make(chan error, 1)

	// file is not there. Easy.
	if !checkFile(name) {
		ret <- nil
		return ret, nil
	}

	// set up notification and timeout
	tout := time.After(timeout)
	// intervals at every 20 milliseconds
	interval := time.NewTicker(pollInterval).C

	// naieve polling system
	go func() {
		for {

			if !checkFile(name) {
				// gone!
				ret <- nil
				return
			}

			select {
			case <-tout:
				ret <- fmt.Errorf("Timed out waiting for %v after %v", name, timeout)
				return
			case <-interval:
				// ignore specific event, check actual file later
			}
		}
	}()

	return ret, nil

}

func readStringFileAsInt(name string) (int, error) {
	data, err := readFile(name)
	if err != nil {
		return 0, err
	}
	val, err := strconv.Atoi(data)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert value %v from %v to an int: %v", data, name, err)
	}
	return val, nil
}

//readFile reads the file and returns the contents as a string (trimmed)
func readFile(name string) (string, error) {
	data, err := ioutil.ReadFile(name)
	if err != nil {
		return "", err
	}
	str := string(data)
	str = strings.TrimSpace(str)
	return str, nil
}

// readBuffer reads a file in to a byte buffer
func readBytes(name string) ([]byte, error) {
	return ioutil.ReadFile(name)
}

// writeBuffer writes a buffer in to a file
func writeBuffer(name string, data []byte) error {
	return ioutil.WriteFile(name, data, 0444)
}

// writeFile will overwrite the specified file with the given string content
func writeFile(name, text string) error {
	data := []byte(text)
	return ioutil.WriteFile(name, data, 0444)
}

func checkWritable(name string) bool {
	if stat, err := os.Stat(name); err == nil {
		// exists, but is it writable?
		mode := os.O_RDWR
		desc := "writable file"
		if stat.IsDir() {
			mode = os.O_RDONLY
			desc = "readable folder"
		}
		// Note, you can open directories as well
		file, err := os.OpenFile(name, mode, 0)
		if err != nil {
			fmt.Printf("Existing file %v but it is not a %v: %v\n", name, desc, err)
			return false
		}
		file.Close()
		// already exists
		return true
	}
	return false
}

// checkFile retuns true if the specified file exists
func checkFile(name string) bool {
	if _, err := os.Stat(name); err == nil {
		// exists, but is it writable?
		// already exists
		return true
	}
	return false
}
<commit_msg>Deep debugging<commit_after>package gopisysfs

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	// pollInterval is how long we should wait between naieve polling of files
	pollInterval = 20 * time.Millisecond
	// forever is about 100 years.
	forever = 100 * 365 * 24 * time.Hour
)

var rootpath = "/"

// setRoot is designed to be called by the test cases to exercise some hard-to-change things on an actual pi.
func setRoot(rt string) {
	rootpath = rt
}

// file gets a file path inside the /sys file system,
// but it can be hooked by the test cases to use a test filesystem instead of the real /sys
func file(paths ...string) string {
	path := filepath.Join(paths...)
	if !filepath.IsAbs(path) {
		path = filepath.Join(rootpath, path)
	}
	return path
}

// readFilePanic reads a file returning the contents as a string, and panics if it cannot be read
func readFilePanic(name string) string {
	data, err := readFile(name)
	if err != nil {
		log.Panicf("Unable to read file %v: %v", name, err)
	}
	return data
}

// awaitFileCreate establishes an asynchronous poll on a file location until it exists
// at which point the returned channel will return a nil on the channel. A non-nil indicates
// an error in the polling.
func awaitFileCreate(name string, timeout time.Duration) (<-chan error, error) {

	ret := make(chan error, 1)

	if checkFile(name) {
		ret <- nil
		return ret, nil
	}

	dir := filepath.Dir(name)
	if stat, err := os.Stat(dir); err != nil || !stat.IsDir() {
		if err != nil {
			return nil, fmt.Errorf("Unable to poll for a file in a nonexistent folder %v: %v", dir, err)
		}
		return nil, fmt.Errorf("Unable to poll for a file in a non-folder %v: %v", dir, stat)
	}

	// set up notification and timeout
	tout := time.After(timeout)
	// intervals at every poll cycle
	interval := time.NewTicker(pollInterval).C
	// naieve polling system
	go func() {
		for {

			if checkFile(name) {
				// Found it!
				ret <- nil
				return
			}

			select {
			case <-tout:
				ret <- fmt.Errorf("Timed out waiting for %v after %v", name, timeout)
				return
			case <-interval:
				// ignore specific event, check actual file later
			}
		}
	}()

	return ret, nil

}

// awaitFileRemove establishes an asynchronous poll on a file location until it is removed
// at which point the returned channel will return a nil on the channel. A non-nil indicates
// an error in the polling.
func awaitFileRemove(name string, timeout time.Duration) (<-chan error, error) {

	ret := make(chan error, 1)

	// file is not there. Easy.
	if !checkFile(name) {
		ret <- nil
		return ret, nil
	}

	// set up notification and timeout
	tout := time.After(timeout)
	// intervals at every 20 milliseconds
	interval := time.NewTicker(pollInterval).C

	// naieve polling system
	go func() {
		for {

			if !checkFile(name) {
				// gone!
				ret <- nil
				return
			}

			select {
			case <-tout:
				ret <- fmt.Errorf("Timed out waiting for %v after %v", name, timeout)
				return
			case <-interval:
				// ignore specific event, check actual file later
			}
		}
	}()

	return ret, nil

}

func readStringFileAsInt(name string) (int, error) {
	data, err := readFile(name)
	if err != nil {
		return 0, err
	}
	val, err := strconv.Atoi(data)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert value %v from %v to an int: %v", data, name, err)
	}
	return val, nil
}

//readFile reads the file and returns the contents as a string (trimmed)
func readFile(name string) (string, error) {
	data, err := ioutil.ReadFile(name)
	if err != nil {
		return "", err
	}
	str := string(data)
	str = strings.TrimSpace(str)
	return str, nil
}

// readBuffer reads a file in to a byte buffer
func readBytes(name string) ([]byte, error) {
	return ioutil.ReadFile(name)
}

// writeBuffer writes a buffer in to a file
func writeBuffer(name string, data []byte) error {
	info("Writing to %v: %v\n", name, data)
	return ioutil.WriteFile(name, data, 0444)
}

// writeFile will overwrite the specified file with the given string content
func writeFile(name, text string) error {
	info("Writing to %v: %v\n", name, text)
	data := []byte(text)
	return ioutil.WriteFile(name, data, 0444)
}

func checkWritable(name string) bool {
	if stat, err := os.Stat(name); err == nil {
		// exists, but is it writable?
		mode := os.O_RDWR
		desc := "writable file"
		if stat.IsDir() {
			mode = os.O_RDONLY
			desc = "readable folder"
		}
		// Note, you can open directories as well
		file, err := os.OpenFile(name, mode, 0)
		if err != nil {
			info("Existing file %v but it is not a %v: %v\n", name, desc, err)
			return false
		}
		file.Close()
		// already exists
		return true
	}
	return false
}

// checkFile retuns true if the specified file exists
func checkFile(name string) bool {
	if _, err := os.Stat(name); err == nil {
		// exists, but is it writable?
		// already exists
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package dhcp6

import (
	"encoding/binary"
	"errors"
	"net"
	"time"
)

const (
	// ethernet10Mb is the default IANA hardware type used in DUID generation,
	// if a hardware type cannot be parsed from a network interface.
	ethernet10Mb uint16 = 1
)

var (
	// errInvalidDUID is returned when not enough bytes are present
	// to parse a valid DUID from a byte slice.
	errInvalidDUID = errors.New("not enough bytes for valid DUID")

	// errInvalidDUIDLLT is returned when not enough bytes are present
	// to parse a valid DUIDLLT from a byte slice, or when the DUID type
	// found in the byte slice is incorrect.
	errInvalidDUIDLLT = errors.New("invalid DUID-LLT")

	// errInvalidDUIDEN is returned when not enough bytes are present
	// to parse a valid DUIDEN from a byte slice, or when the DUID type
	// found in the byte slice is incorrect.
	errInvalidDUIDEN = errors.New("invalid DUID-EN")

	// errInvalidDUIDLL is returned when not enough bytes are present
	// to parse a valid DUIDLL from a byte slice, or when the DUID type
	// found in the byte slice is incorrect.
	errInvalidDUIDLL = errors.New("invalid DUID-LL")

	// errUnknownDUID is returned when an unknown DUID type is
	// encountered, and thus, a DUID cannot be parsed.
	errUnknownDUID = errors.New("unknown DUID type")
)

var (
	// duidLLTTime is the date specified in IETF RFC 3315, Section 9.2, for use
	// with DUIT-LLT generation.  It is used to calculate a duration from an
	// input time after this date.  Dates before this time are not valid for
	// creation of DUIDLLT values.
	duidLLTTime = time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)
)

// DUIDType is a type of DHCP Unique Identifier, as defined in IETF RFC
// 3315, Section 9.  DUIDs are used to uniquely identify a client to a
// server, or vice-versa.
type DUIDType uint16

// DUIDType constants which indicate DUID-LLT, DUID-EN, or DUID-LL.
// Additional DUID types are defined in IANA's DHCPv6 parameters registry:
// http://www.iana.org/assignments/dhcpv6-parameters/dhcpv6-parameters.xhtml.
const (
	DUIDTypeLLT DUIDType = 1
	DUIDTypeEN  DUIDType = 2
	DUIDTypeLL  DUIDType = 3

	// BUG(mdlayher): add additional DUID types defined by IANA
)

// DUID represents a DHCP Unique Identifier, as defined in IETF RFC
// 3315, Section 9.  A DUID is used by a DHCP server to identify
// unique clients.  A DUID can also be used by a DHCP client to identify
// a unique server, when needed.
//
// The DUID interface represents a generic DUID, but DUIDs can be
// type-asserted to one of three specific types outlined in RFC 3315,
// Section 9.2, 9.3, and 9.4:
//   - DUIDLLT - DUID Based on Link-layer Address Plus Time
//   -  DUIDEN - DUID Assigned by Vendor Based on Enterprise Number
//   -  DUIDLL - DUID Based on Link-layer Address
//
// If further introspection of the DUID is needed, a type switch is
// recommended:
//	switch d := duid.(type) {
//	case dhcp6.DUIDLLT:
//		fmt.Println(d.Time)
//	case dhcp6.DUIDEN:
//		fmt.Println(d.EnterpriseNumber)
//	case dhcp6.DUIDLL:
//		fmt.Println(d.HardwareAddr)
//	}
type DUID Byteser

// DUIDLLT represents a DUID Based on Link-layer Address Plus Time [DUID-LLT],
// as defined in IETF RFC 3315, Section 9.2.
//
// This DUID type must only be used with clients and servers with stable,
// persistent storage.  It is the recommended DUID type for all general
// purpose computing devices.
type DUIDLLT struct {
	// Type specifies the DUID type.  For a DUIDLLT, this should always be
	// DUIDTypeLLT.
	Type DUIDType

	// HardwareType specifies an IANA-assigned hardware type, as described
	// in IETF RFC 826.
	HardwareType uint16

	// Time specifies the duration of the time this DUID was generated, minus
	// midnight (UTC), January 1, 2000.
	Time time.Duration

	// HardwareAddr specifies the hardware address for an arbitrary link-layer
	// interface on a device, used in generating the DUIDLLT.  This value
	// could represent any arbitrary interface on a system, and should not be
	// treated as a client or server's communicating hardware address.
	HardwareAddr net.HardwareAddr
}

// NewDUIDLLT generates a new DUIDLLT from an input IANA-assigned hardware
// type, time value, and a hardware address.
//
// The time value must be greater than midnight (UTC), January 1, 2000.
func NewDUIDLLT(hardwareType uint16, time time.Time, hardwareAddr net.HardwareAddr) (*DUIDLLT, error) {
	// Do not accept dates before duidLLTTime.
	if time.Before(duidLLTTime) {
		return nil, ErrInvalidDUIDLLTTime
	}

	return &DUIDLLT{
		Type:         DUIDTypeLLT,
		HardwareType: hardwareType,
		Time:         time.Sub(duidLLTTime),
		HardwareAddr: hardwareAddr,
	}, nil
}

// Bytes implements DUID, and allocates a byte slice containing the data
// from a DUIDLLT.
func (d *DUIDLLT) Bytes() []byte {
	// 2 bytes: DUID type
	// 2 bytes: hardware type
	// 4 bytes: time duration
	// N bytes: hardware address
	b := make([]byte, 8+len(d.HardwareAddr))

	binary.BigEndian.PutUint16(b[0:2], uint16(d.Type))
	binary.BigEndian.PutUint16(b[2:4], d.HardwareType)
	binary.BigEndian.PutUint32(b[4:8], uint32(d.Time/time.Second))
	copy(b[8:], d.HardwareAddr)

	return b
}

// parseDUIDLLT parses a raw byte slice into a DUIDLLT.  If the byte slice
// does not contain enough data to form a valid DUIDLLT, or another DUID type
// is indicated, errInvalidDUIDLLT is returned.
func parseDUIDLLT(b []byte) (*DUIDLLT, error) {
	// Too short to be valid DUIDLLT
	if len(b) < 8 {
		return nil, errInvalidDUIDLLT
	}

	// Verify DUID type
	dType := DUIDType(binary.BigEndian.Uint16(b[0:2]))
	if dType != DUIDTypeLLT {
		return nil, errInvalidDUIDLLT
	}

	mac := make(net.HardwareAddr, len(b[8:]))
	copy(mac, b[8:])

	return &DUIDLLT{
		Type:         dType,
		HardwareType: binary.BigEndian.Uint16(b[2:4]),
		Time:         time.Duration(binary.BigEndian.Uint32(b[4:8])) * time.Second,
		HardwareAddr: mac,
	}, nil
}

// DUIDEN represents a DUID Assigned by Vendor Based on Enterprise Number
// [DUID-EN], as defined in IETF RFC 3315, Section 9.3.  This DUID type
// uses an IANA-assigned Private Enterprise Number for a given vendor.
type DUIDEN struct {
	// Type specifies the DUID type.  For a DUIDLLT, this should always be
	// DUIDTypeLLT.
	Type DUIDType

	// EnterpriseNumber specifies an IANA-assigned vendor Private Enterprise
	// Number.
	EnterpriseNumber uint32

	// Identifier specifies a unique identifier of arbitrary length.  This
	// value is typically assigned when a device is manufactured.
	Identifier []byte
}

// NewDUIDEN generates a new DUIDEN from an input IANA-assigned Private
// Enterprise Number and a variable length unique identifier byte slice.
// type and a hardware address.
func NewDUIDEN(enterpriseNumber uint32, identifier []byte) *DUIDEN {
	return &DUIDEN{
		Type:             DUIDTypeEN,
		EnterpriseNumber: enterpriseNumber,
		Identifier:       identifier,
	}
}

// Bytes implements DUID, and allocates a byte slice containing the data
// from a DUIDEN.
func (d *DUIDEN) Bytes() []byte {
	// 2 bytes: DUID type
	// 4 bytes: enterprise number
	// N bytes: identifier
	b := make([]byte, 6+len(d.Identifier))

	binary.BigEndian.PutUint16(b[0:2], uint16(d.Type))
	binary.BigEndian.PutUint32(b[2:6], d.EnterpriseNumber)
	copy(b[6:], d.Identifier)

	return b
}

// parseDUIDEN parses a raw byte slice into a DUIDEN.  If the byte slice
// does not contain enough data to form a valid DUIDEN, or another DUID type
// is indicated, errInvalidDUIDEN is returned.
func parseDUIDEN(b []byte) (*DUIDEN, error) {
	// Too short to be valid DUIDEN
	if len(b) < 6 {
		return nil, errInvalidDUIDEN
	}

	// Verify DUID type
	dType := DUIDType(binary.BigEndian.Uint16(b[0:2]))
	if dType != DUIDTypeEN {
		return nil, errInvalidDUIDEN
	}

	id := make([]byte, len(b[6:]))
	copy(id, b[6:])

	return &DUIDEN{
		Type:             dType,
		EnterpriseNumber: binary.BigEndian.Uint32(b[2:6]),
		Identifier:       id,
	}, nil
}

// DUIDLL represents a DUID Based on Link-layer Address [DUID-LL],
// as defined in IETF RFC 3315, Section 9.4.
//
// This DUID type is recommended for devices with a
// permanently-connected network interface, but without stable,
// persistent storage.
//
// DUIDLL values are generated automatically for Servers which are not
// created with a ServerID, using the hardware type found by HardwareType
// and the hardware address of the listening network interface.
type DUIDLL struct {
	// Type specifies the DUID type.  For a DUIDLL, this should always be
	// DUIDTypeLL.
	Type DUIDType

	// HardwareType specifies an IANA-assigned hardware type, as described
	// in IETF RFC 826.
	HardwareType uint16

	// HardwareAddr specifies the hardware address for an arbitrary link-layer
	// interface on a device, used in generating the DUIDLL.  This value
	// could represent any arbitrary interface on a system, and should not be
	// treated as a client or server's communicating hardware address.
	HardwareAddr net.HardwareAddr
}

// NewDUIDLL generates a new DUIDLL from an input IANA-assigned hardware
// type and a hardware address.
func NewDUIDLL(hardwareType uint16, hardwareAddr net.HardwareAddr) *DUIDLL {
	return &DUIDLL{
		Type:         DUIDTypeLL,
		HardwareType: hardwareType,
		HardwareAddr: hardwareAddr,
	}
}

// Bytes implements DUID, and allocates a byte slice containing the data
// from a DUIDLL.
func (d *DUIDLL) Bytes() []byte {
	// 2 bytes: DUID type
	// 2 bytes: hardware type
	// N bytes: hardware address
	b := make([]byte, 4+len(d.HardwareAddr))

	binary.BigEndian.PutUint16(b[0:2], uint16(d.Type))
	binary.BigEndian.PutUint16(b[2:4], d.HardwareType)
	copy(b[4:], d.HardwareAddr)

	return b
}

// parseDUIDLL parses a raw byte slice into a DUIDLL.  If the byte slice
// does not contain enough data to form a valid DUIDLL, or another DUID type
// is indicated, errInvalidDUIDLL is returned.
func parseDUIDLL(b []byte) (*DUIDLL, error) {
	// Too short to be DUIDLL
	if len(b) < 4 {
		return nil, errInvalidDUIDLL
	}

	// Verify DUID type
	dType := DUIDType(binary.BigEndian.Uint16(b[0:2]))
	if dType != DUIDTypeLL {
		return nil, errInvalidDUIDLL
	}

	mac := make(net.HardwareAddr, len(b[4:]))
	copy(mac, b[4:])

	return &DUIDLL{
		Type:         dType,
		HardwareType: binary.BigEndian.Uint16(b[2:4]),
		HardwareAddr: mac,
	}, nil
}

// parseDUID returns the correct DUID type of the input byte slice as a
// DUID interface type.
func parseDUID(d []byte) (DUID, error) {
	// DUID must have enough bytes to determine its type
	if len(d) < 2 {
		return nil, errInvalidDUID
	}

	// BUG(mdlayher): add DUID-UUID to this in the future.
	switch DUIDType(binary.BigEndian.Uint16(d[0:2])) {
	case DUIDTypeLLT:
		return parseDUIDLLT(d)
	case DUIDTypeEN:
		return parseDUIDEN(d)
	case DUIDTypeLL:
		return parseDUIDLL(d)
	}

	return nil, errUnknownDUID
}

// interfaceDUID generates a DUIDLL for an input net.Interface, using its
// IANA-assigned hardware type and its hardware address.
func interfaceDUID(ifi *net.Interface) (DUID, error) {
	// Attempt to check for IANA hardware type, default to Ethernet (10Mb)
	// on failure (this relies on syscalls which only work on Linux)
	// Hardware types can be found here:
	// http://www.iana.org/assignments/arp-parameters/arp-parameters.xhtml.
	htype, err := HardwareType(ifi)
	if err != nil {
		// Return syscall errors
		if err != ErrParseHardwareType && err != ErrHardwareTypeNotImplemented {
			return nil, err
		}

		// Use default value if hardware type can't be found or
		// detection isn't implemented
		htype = ethernet10Mb
	}

	return NewDUIDLL(htype, ifi.HardwareAddr), nil
}
<commit_msg>duid: remove erroneous comment from NewDUIDEN<commit_after>package dhcp6

import (
	"encoding/binary"
	"errors"
	"net"
	"time"
)

const (
	// ethernet10Mb is the default IANA hardware type used in DUID generation,
	// if a hardware type cannot be parsed from a network interface.
	ethernet10Mb uint16 = 1
)

var (
	// errInvalidDUID is returned when not enough bytes are present
	// to parse a valid DUID from a byte slice.
	errInvalidDUID = errors.New("not enough bytes for valid DUID")

	// errInvalidDUIDLLT is returned when not enough bytes are present
	// to parse a valid DUIDLLT from a byte slice, or when the DUID type
	// found in the byte slice is incorrect.
	errInvalidDUIDLLT = errors.New("invalid DUID-LLT")

	// errInvalidDUIDEN is returned when not enough bytes are present
	// to parse a valid DUIDEN from a byte slice, or when the DUID type
	// found in the byte slice is incorrect.
	errInvalidDUIDEN = errors.New("invalid DUID-EN")

	// errInvalidDUIDLL is returned when not enough bytes are present
	// to parse a valid DUIDLL from a byte slice, or when the DUID type
	// found in the byte slice is incorrect.
	errInvalidDUIDLL = errors.New("invalid DUID-LL")

	// errUnknownDUID is returned when an unknown DUID type is
	// encountered, and thus, a DUID cannot be parsed.
	errUnknownDUID = errors.New("unknown DUID type")
)

var (
	// duidLLTTime is the date specified in IETF RFC 3315, Section 9.2, for use
	// with DUIT-LLT generation.  It is used to calculate a duration from an
	// input time after this date.  Dates before this time are not valid for
	// creation of DUIDLLT values.
	duidLLTTime = time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)
)

// DUIDType is a type of DHCP Unique Identifier, as defined in IETF RFC
// 3315, Section 9.  DUIDs are used to uniquely identify a client to a
// server, or vice-versa.
type DUIDType uint16

// DUIDType constants which indicate DUID-LLT, DUID-EN, or DUID-LL.
// Additional DUID types are defined in IANA's DHCPv6 parameters registry:
// http://www.iana.org/assignments/dhcpv6-parameters/dhcpv6-parameters.xhtml.
const (
	DUIDTypeLLT DUIDType = 1
	DUIDTypeEN  DUIDType = 2
	DUIDTypeLL  DUIDType = 3

	// BUG(mdlayher): add additional DUID types defined by IANA
)

// DUID represents a DHCP Unique Identifier, as defined in IETF RFC
// 3315, Section 9.  A DUID is used by a DHCP server to identify
// unique clients.  A DUID can also be used by a DHCP client to identify
// a unique server, when needed.
//
// The DUID interface represents a generic DUID, but DUIDs can be
// type-asserted to one of three specific types outlined in RFC 3315,
// Section 9.2, 9.3, and 9.4:
//   - DUIDLLT - DUID Based on Link-layer Address Plus Time
//   -  DUIDEN - DUID Assigned by Vendor Based on Enterprise Number
//   -  DUIDLL - DUID Based on Link-layer Address
//
// If further introspection of the DUID is needed, a type switch is
// recommended:
//	switch d := duid.(type) {
//	case dhcp6.DUIDLLT:
//		fmt.Println(d.Time)
//	case dhcp6.DUIDEN:
//		fmt.Println(d.EnterpriseNumber)
//	case dhcp6.DUIDLL:
//		fmt.Println(d.HardwareAddr)
//	}
type DUID Byteser

// DUIDLLT represents a DUID Based on Link-layer Address Plus Time [DUID-LLT],
// as defined in IETF RFC 3315, Section 9.2.
//
// This DUID type must only be used with clients and servers with stable,
// persistent storage.  It is the recommended DUID type for all general
// purpose computing devices.
type DUIDLLT struct {
	// Type specifies the DUID type.  For a DUIDLLT, this should always be
	// DUIDTypeLLT.
	Type DUIDType

	// HardwareType specifies an IANA-assigned hardware type, as described
	// in IETF RFC 826.
	HardwareType uint16

	// Time specifies the duration of the time this DUID was generated, minus
	// midnight (UTC), January 1, 2000.
	Time time.Duration

	// HardwareAddr specifies the hardware address for an arbitrary link-layer
	// interface on a device, used in generating the DUIDLLT.  This value
	// could represent any arbitrary interface on a system, and should not be
	// treated as a client or server's communicating hardware address.
	HardwareAddr net.HardwareAddr
}

// NewDUIDLLT generates a new DUIDLLT from an input IANA-assigned hardware
// type, time value, and a hardware address.
//
// The time value must be greater than midnight (UTC), January 1, 2000.
func NewDUIDLLT(hardwareType uint16, time time.Time, hardwareAddr net.HardwareAddr) (*DUIDLLT, error) {
	// Do not accept dates before duidLLTTime.
	if time.Before(duidLLTTime) {
		return nil, ErrInvalidDUIDLLTTime
	}

	return &DUIDLLT{
		Type:         DUIDTypeLLT,
		HardwareType: hardwareType,
		Time:         time.Sub(duidLLTTime),
		HardwareAddr: hardwareAddr,
	}, nil
}

// Bytes implements DUID, and allocates a byte slice containing the data
// from a DUIDLLT.
func (d *DUIDLLT) Bytes() []byte {
	// 2 bytes: DUID type
	// 2 bytes: hardware type
	// 4 bytes: time duration
	// N bytes: hardware address
	b := make([]byte, 8+len(d.HardwareAddr))

	binary.BigEndian.PutUint16(b[0:2], uint16(d.Type))
	binary.BigEndian.PutUint16(b[2:4], d.HardwareType)
	binary.BigEndian.PutUint32(b[4:8], uint32(d.Time/time.Second))
	copy(b[8:], d.HardwareAddr)

	return b
}

// parseDUIDLLT parses a raw byte slice into a DUIDLLT.  If the byte slice
// does not contain enough data to form a valid DUIDLLT, or another DUID type
// is indicated, errInvalidDUIDLLT is returned.
func parseDUIDLLT(b []byte) (*DUIDLLT, error) {
	// Too short to be valid DUIDLLT
	if len(b) < 8 {
		return nil, errInvalidDUIDLLT
	}

	// Verify DUID type
	dType := DUIDType(binary.BigEndian.Uint16(b[0:2]))
	if dType != DUIDTypeLLT {
		return nil, errInvalidDUIDLLT
	}

	mac := make(net.HardwareAddr, len(b[8:]))
	copy(mac, b[8:])

	return &DUIDLLT{
		Type:         dType,
		HardwareType: binary.BigEndian.Uint16(b[2:4]),
		Time:         time.Duration(binary.BigEndian.Uint32(b[4:8])) * time.Second,
		HardwareAddr: mac,
	}, nil
}

// DUIDEN represents a DUID Assigned by Vendor Based on Enterprise Number
// [DUID-EN], as defined in IETF RFC 3315, Section 9.3.  This DUID type
// uses an IANA-assigned Private Enterprise Number for a given vendor.
type DUIDEN struct {
	// Type specifies the DUID type.  For a DUIDLLT, this should always be
	// DUIDTypeLLT.
	Type DUIDType

	// EnterpriseNumber specifies an IANA-assigned vendor Private Enterprise
	// Number.
	EnterpriseNumber uint32

	// Identifier specifies a unique identifier of arbitrary length.  This
	// value is typically assigned when a device is manufactured.
	Identifier []byte
}

// NewDUIDEN generates a new DUIDEN from an input IANA-assigned Private
// Enterprise Number and a variable length unique identifier byte slice.
func NewDUIDEN(enterpriseNumber uint32, identifier []byte) *DUIDEN {
	return &DUIDEN{
		Type:             DUIDTypeEN,
		EnterpriseNumber: enterpriseNumber,
		Identifier:       identifier,
	}
}

// Bytes implements DUID, and allocates a byte slice containing the data
// from a DUIDEN.
func (d *DUIDEN) Bytes() []byte {
	// 2 bytes: DUID type
	// 4 bytes: enterprise number
	// N bytes: identifier
	b := make([]byte, 6+len(d.Identifier))

	binary.BigEndian.PutUint16(b[0:2], uint16(d.Type))
	binary.BigEndian.PutUint32(b[2:6], d.EnterpriseNumber)
	copy(b[6:], d.Identifier)

	return b
}

// parseDUIDEN parses a raw byte slice into a DUIDEN.  If the byte slice
// does not contain enough data to form a valid DUIDEN, or another DUID type
// is indicated, errInvalidDUIDEN is returned.
func parseDUIDEN(b []byte) (*DUIDEN, error) {
	// Too short to be valid DUIDEN
	if len(b) < 6 {
		return nil, errInvalidDUIDEN
	}

	// Verify DUID type
	dType := DUIDType(binary.BigEndian.Uint16(b[0:2]))
	if dType != DUIDTypeEN {
		return nil, errInvalidDUIDEN
	}

	id := make([]byte, len(b[6:]))
	copy(id, b[6:])

	return &DUIDEN{
		Type:             dType,
		EnterpriseNumber: binary.BigEndian.Uint32(b[2:6]),
		Identifier:       id,
	}, nil
}

// DUIDLL represents a DUID Based on Link-layer Address [DUID-LL],
// as defined in IETF RFC 3315, Section 9.4.
//
// This DUID type is recommended for devices with a
// permanently-connected network interface, but without stable,
// persistent storage.
//
// DUIDLL values are generated automatically for Servers which are not
// created with a ServerID, using the hardware type found by HardwareType
// and the hardware address of the listening network interface.
type DUIDLL struct {
	// Type specifies the DUID type.  For a DUIDLL, this should always be
	// DUIDTypeLL.
	Type DUIDType

	// HardwareType specifies an IANA-assigned hardware type, as described
	// in IETF RFC 826.
	HardwareType uint16

	// HardwareAddr specifies the hardware address for an arbitrary link-layer
	// interface on a device, used in generating the DUIDLL.  This value
	// could represent any arbitrary interface on a system, and should not be
	// treated as a client or server's communicating hardware address.
	HardwareAddr net.HardwareAddr
}

// NewDUIDLL generates a new DUIDLL from an input IANA-assigned hardware
// type and a hardware address.
func NewDUIDLL(hardwareType uint16, hardwareAddr net.HardwareAddr) *DUIDLL {
	return &DUIDLL{
		Type:         DUIDTypeLL,
		HardwareType: hardwareType,
		HardwareAddr: hardwareAddr,
	}
}

// Bytes implements DUID, and allocates a byte slice containing the data
// from a DUIDLL.
func (d *DUIDLL) Bytes() []byte {
	// 2 bytes: DUID type
	// 2 bytes: hardware type
	// N bytes: hardware address
	b := make([]byte, 4+len(d.HardwareAddr))

	binary.BigEndian.PutUint16(b[0:2], uint16(d.Type))
	binary.BigEndian.PutUint16(b[2:4], d.HardwareType)
	copy(b[4:], d.HardwareAddr)

	return b
}

// parseDUIDLL parses a raw byte slice into a DUIDLL.  If the byte slice
// does not contain enough data to form a valid DUIDLL, or another DUID type
// is indicated, errInvalidDUIDLL is returned.
func parseDUIDLL(b []byte) (*DUIDLL, error) {
	// Too short to be DUIDLL
	if len(b) < 4 {
		return nil, errInvalidDUIDLL
	}

	// Verify DUID type
	dType := DUIDType(binary.BigEndian.Uint16(b[0:2]))
	if dType != DUIDTypeLL {
		return nil, errInvalidDUIDLL
	}

	mac := make(net.HardwareAddr, len(b[4:]))
	copy(mac, b[4:])

	return &DUIDLL{
		Type:         dType,
		HardwareType: binary.BigEndian.Uint16(b[2:4]),
		HardwareAddr: mac,
	}, nil
}

// parseDUID returns the correct DUID type of the input byte slice as a
// DUID interface type.
func parseDUID(d []byte) (DUID, error) {
	// DUID must have enough bytes to determine its type
	if len(d) < 2 {
		return nil, errInvalidDUID
	}

	// BUG(mdlayher): add DUID-UUID to this in the future.
	switch DUIDType(binary.BigEndian.Uint16(d[0:2])) {
	case DUIDTypeLLT:
		return parseDUIDLLT(d)
	case DUIDTypeEN:
		return parseDUIDEN(d)
	case DUIDTypeLL:
		return parseDUIDLL(d)
	}

	return nil, errUnknownDUID
}

// interfaceDUID generates a DUIDLL for an input net.Interface, using its
// IANA-assigned hardware type and its hardware address.
func interfaceDUID(ifi *net.Interface) (DUID, error) {
	// Attempt to check for IANA hardware type, default to Ethernet (10Mb)
	// on failure (this relies on syscalls which only work on Linux)
	// Hardware types can be found here:
	// http://www.iana.org/assignments/arp-parameters/arp-parameters.xhtml.
	htype, err := HardwareType(ifi)
	if err != nil {
		// Return syscall errors
		if err != ErrParseHardwareType && err != ErrHardwareTypeNotImplemented {
			return nil, err
		}

		// Use default value if hardware type can't be found or
		// detection isn't implemented
		htype = ethernet10Mb
	}

	return NewDUIDLL(htype, ifi.HardwareAddr), nil
}
<|endoftext|>


<commit_before>/*
Package mote represents a connected Pimoroni Mote device, communicating over USB serial.

It allows you to configure the 4 channels and set individual pixels.

It is a port of the Pimoroni Mote Python library (https://github.com/pimoroni/mote).

The Mote device can be obtained directly from Pimoroni (https://shop.pimoroni.com/products/mote).
*/
package mote

import (
	"fmt"
	"log"

	serial "github.com/johnmccabe/go-serial-native"
)

// VID is the USB Vendor ID of the Pimoroni Mote
const VID = 5840

// PID is the USB Product ID of the Pimoroni Mote
const PID = 2244

// ProductName is the USB Product Name of the Pimoroni Mote
const ProductName = "Mote USB Dock"

// NumChannels is the number of available channel connections on the Mote device
const NumChannels = 4

// MaxPixels is the maximum addressable number of pixels across all channels
const MaxPixels = 512

// MaxPixelsPerChannel is the maximum addressable number of pixels across a single channel
const MaxPixelsPerChannel = int(MaxPixels / NumChannels)

// Mote represents a connected Pimoroni Mote device
type Mote struct {
	PortName string
	Port     *serial.Port
	Channels [NumChannels]*Channel
}

// Pixel represents a single RGB pixel
type Pixel struct {
	Red, Green, Blue int
}

// Channel represents a single channel on the Mote board
type Channel struct {
	Pixels          []Pixel
	GammaCorrection bool
}

// NewMote creates a connection to a Mote device, communicating over USB serial.
//
// It will attach to the first available Mote device unless a non-empty string `port_name` is specified at init.
//
//   - portName: override auto-detect and specify an explicit port to use. Must be a complete path ie: /dev/tty.usbmodem1234
func NewMote(portName string) *Mote {
	mote := Mote{
		PortName: portName,
	}
	if mote.PortName == "" {
		mote.PortName = *findSerialPort(VID, PID, ProductName)
	}
	if mote.PortName == "" {
		log.Fatal("unable to detect connected Mote")
	}

	options := serial.RawOptions
	options.Mode = serial.MODE_WRITE
	options.BitRate = 115200
	options.DataBits = 8
	options.Parity = serial.PARITY_NONE
	options.StopBits = 1
	options.FlowControl = 0
	p, err := options.Open(mote.PortName)
	if err != nil {
		panic(err)
	}
	mote.Port = p

	return &mote
}

// ConfigureChannel configures a channel, taking the following parameters.
//
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
//   - numPixels: Number of pixels to configure for this channel
//   - gammaCorrection: Whether to enable gamma correction
func (m *Mote) ConfigureChannel(channel, numPixels int, gammaCorrection bool) error {
	if channel > NumChannels || channel < 1 {
		return fmt.Errorf("channel index must be between 1 and 4")
	}
	if numPixels > MaxPixelsPerChannel {
		return fmt.Errorf("number of pixels can not be more than %d", MaxPixelsPerChannel)
	}

	p := []Pixel{}
	for i := 0; i < numPixels; i++ {
		p = append(p, Pixel{0, 0, 0})
	}
	c := Channel{
		Pixels:          p,
		GammaCorrection: gammaCorrection,
	}
	m.Channels[channel-1] = &c

	var b []byte
	b = append(b, []byte("mote")...)
	b = append(b, []byte("c")...)
	b = append(b, byte(channel))
	b = append(b, byte(numPixels))
	gammaCorrectionVar := 0
	if gammaCorrection {
		gammaCorrectionVar = 1
	}
	b = append(b, byte(gammaCorrectionVar))

	a, err := m.Port.Write(b)
	fmt.Printf("%v %v", a, err)
	return nil
}

// SetPixel Set the RGB colour of a single pixel, on a single channel, taking the following parameters.
//
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
//   - index: Index of pixel to set, from 0 up
//   - r: Amount of red: 0-255
//   - g: Amount of green: 0-255
//   - b: Amount of blue: 0-255
func (m *Mote) SetPixel(channel, index, r, g, b int) error {
	if channel > NumChannels || channel < 1 {
		return fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return fmt.Errorf("please set up channel %d before using it", channel)
	}
	if index >= len(m.Channels[channel-1].Pixels) {
		return fmt.Errorf("Pixel index must be < %d", m.Channels[channel-1].Pixels)
	}
	m.Channels[channel-1].Pixels[index] = Pixel{r & 0xff, g & 0xff, b & 0xff}
	return nil
}

// Show sends the pixel buffer to the hardware.
func (m *Mote) Show() {
	var b []byte
	b = append(b, []byte("mote")...)
	b = append(b, []byte("o")...)
	for channel, data := range m.Channels {
		if data == nil {
			fmt.Printf("skipping empty channel %d\n", channel+1)
			continue
		}
		for _, pixel := range data.Pixels {
			b = append(b, byte(pixel.Blue))
			b = append(b, byte(pixel.Green))
			b = append(b, byte(pixel.Red))
		}
	}
	m.Port.Write(b)
}

// Clear the buffer of a specific channels, taking the following parameter.
//
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
func (m *Mote) Clear(channel int) error {
	if channel > NumChannels || channel < 0 {
		return fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return fmt.Errorf("please set up channel %d before using it", channel)
	}
	for pixel := 0; pixel < 16; pixel++ {
		err := m.SetPixel(channel, pixel, 0, 0, 0)
		if err != nil {
			return err
		}
	}
	return nil
}

// ClearAll clears the buffers of all configured channels.
func (m *Mote) ClearAll() error {
	for channel := 1; channel < NumChannels+1; channel++ {
		if m.Channels[channel-1] != nil {
			err := m.Clear(channel)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// Close closes the serial port connected to the Mote device.
func (m *Mote) Close() {
	log.Printf("Closing port %v\n", m.PortName)
	m.Port.Close()
}

// GetPixelCount gets the number of pixels a channel is configured to use, taking the following parameters.
//   - channel: Channel, either 1, 2 3 or 4 corresponding to numbers on Mote
func (m *Mote) GetPixelCount(channel int) (int, error) {
	if channel > NumChannels || channel < 1 {
		return 0, fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return 0, fmt.Errorf("channel %d has not been set up", channel)
	}
	return len(m.Channels[channel-1].Pixels), nil
}

// GetPixel gets the RGB colour of a single pixel, on a single channel, taking the following parameters.
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
//   - index: Index of pixel to set, from 0 up
func (m *Mote) GetPixel(channel, index int) (Pixel, error) {
	if channel > NumChannels || channel < 1 {
		return Pixel{}, fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return Pixel{}, fmt.Errorf("channel %d has not been set up", channel)
	}
	if index >= len(m.Channels[channel-1].Pixels) {
		return Pixel{}, fmt.Errorf("pixel index must be < %d", len(m.Channels[channel-1].Pixels))
	}
	return m.Channels[channel-1].Pixels[index], nil
}

func findSerialPort(v, p int, n string) *string {
	ports, _ := serial.ListPorts()
	if len(ports) == 0 {
		return nil
	}
	for _, info := range ports {
		portName := info.Name()
		if vid, pid, err := info.USBVIDPID(); err == nil {
			if vid == v && pid == p {
				log.Printf("found Mote connected to port: %s\n", portName)
				return &portName
			}
		}
		product := info.USBProduct()
		if product == ProductName {
			log.Printf("found Mote connected to port: %s\n", portName)
			return &portName
		}
	}
	return nil
}
<commit_msg>use WriteCloser interface for port<commit_after>/*
Package mote represents a connected Pimoroni Mote device, communicating over USB serial.

It allows you to configure the 4 channels and set individual pixels.

It is a port of the Pimoroni Mote Python library (https://github.com/pimoroni/mote).

The Mote device can be obtained directly from Pimoroni (https://shop.pimoroni.com/products/mote).
*/
package mote

import (
	"fmt"
	"io"
	"log"

	serial "github.com/johnmccabe/go-serial-native"
)

// VID is the USB Vendor ID of the Pimoroni Mote
const VID = 5840

// PID is the USB Product ID of the Pimoroni Mote
const PID = 2244

// ProductName is the USB Product Name of the Pimoroni Mote
const ProductName = "Mote USB Dock"

// NumChannels is the number of available channel connections on the Mote device
const NumChannels = 4

// MaxPixels is the maximum addressable number of pixels across all channels
const MaxPixels = 512

// MaxPixelsPerChannel is the maximum addressable number of pixels across a single channel
const MaxPixelsPerChannel = int(MaxPixels / NumChannels)

// Mote represents a connected Pimoroni Mote device
type Mote struct {
	PortName string
	Port     io.WriteCloser
	Channels [NumChannels]*Channel
}

// Pixel represents a single RGB pixel
type Pixel struct {
	Red, Green, Blue int
}

// Channel represents a single channel on the Mote board
type Channel struct {
	Pixels          []Pixel
	GammaCorrection bool
}

// NewMote creates a connection to a Mote device, communicating over USB serial.
//
// It will attach to the first available Mote device unless a non-empty string `port_name` is specified at init.
//
//   - portName: override auto-detect and specify an explicit port to use. Must be a complete path ie: /dev/tty.usbmodem1234
func NewMote(portName string) *Mote {
	mote := Mote{
		PortName: portName,
	}
	if mote.PortName == "" {
		mote.PortName = *findSerialPort(VID, PID, ProductName)
	}
	if mote.PortName == "" {
		log.Fatal("unable to detect connected Mote")
	}

	options := serial.RawOptions
	options.Mode = serial.MODE_WRITE
	options.BitRate = 115200
	options.DataBits = 8
	options.Parity = serial.PARITY_NONE
	options.StopBits = 1
	options.FlowControl = 0
	p, err := options.Open(mote.PortName)
	if err != nil {
		panic(err)
	}
	mote.Port = p

	return &mote
}

// ConfigureChannel configures a channel, taking the following parameters.
//
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
//   - numPixels: Number of pixels to configure for this channel
//   - gammaCorrection: Whether to enable gamma correction
func (m *Mote) ConfigureChannel(channel, numPixels int, gammaCorrection bool) error {
	if channel > NumChannels || channel < 1 {
		return fmt.Errorf("channel index must be between 1 and 4")
	}
	if numPixels > MaxPixelsPerChannel {
		return fmt.Errorf("number of pixels can not be more than %d", MaxPixelsPerChannel)
	}

	p := []Pixel{}
	for i := 0; i < numPixels; i++ {
		p = append(p, Pixel{0, 0, 0})
	}
	c := Channel{
		Pixels:          p,
		GammaCorrection: gammaCorrection,
	}
	m.Channels[channel-1] = &c

	var b []byte
	b = append(b, []byte("mote")...)
	b = append(b, []byte("c")...)
	b = append(b, byte(channel))
	b = append(b, byte(numPixels))
	gammaCorrectionVar := 0
	if gammaCorrection {
		gammaCorrectionVar = 1
	}
	b = append(b, byte(gammaCorrectionVar))

	_, err := m.Port.Write(b)
	if err != nil {
		return err
	}
	return nil
}

// SetPixel Set the RGB colour of a single pixel, on a single channel, taking the following parameters.
//
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
//   - index: Index of pixel to set, from 0 up
//   - r: Amount of red: 0-255
//   - g: Amount of green: 0-255
//   - b: Amount of blue: 0-255
func (m *Mote) SetPixel(channel, index, r, g, b int) error {
	if channel > NumChannels || channel < 1 {
		return fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return fmt.Errorf("please set up channel %d before using it", channel)
	}
	if index >= len(m.Channels[channel-1].Pixels) {
		return fmt.Errorf("Pixel index must be < %d", m.Channels[channel-1].Pixels)
	}
	m.Channels[channel-1].Pixels[index] = Pixel{r & 0xff, g & 0xff, b & 0xff}
	return nil
}

// Show sends the pixel buffer to the hardware.
func (m *Mote) Show() {
	var b []byte
	b = append(b, []byte("mote")...)
	b = append(b, []byte("o")...)
	for channel, data := range m.Channels {
		if data == nil {
			fmt.Printf("skipping empty channel %d\n", channel+1)
			continue
		}
		for _, pixel := range data.Pixels {
			b = append(b, byte(pixel.Blue))
			b = append(b, byte(pixel.Green))
			b = append(b, byte(pixel.Red))
		}
	}
	m.Port.Write(b)
}

// Clear the buffer of a specific channels, taking the following parameter.
//
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
func (m *Mote) Clear(channel int) error {
	if channel > NumChannels || channel < 0 {
		return fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return fmt.Errorf("please set up channel %d before using it", channel)
	}
	for pixel := 0; pixel < 16; pixel++ {
		err := m.SetPixel(channel, pixel, 0, 0, 0)
		if err != nil {
			return err
		}
	}
	return nil
}

// ClearAll clears the buffers of all configured channels.
func (m *Mote) ClearAll() error {
	for channel := 1; channel < NumChannels+1; channel++ {
		if m.Channels[channel-1] != nil {
			err := m.Clear(channel)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// Close closes the serial port connected to the Mote device.
func (m *Mote) Close() {
	log.Printf("Closing port %v\n", m.PortName)
	m.Port.Close()
}

// GetPixelCount gets the number of pixels a channel is configured to use, taking the following parameters.
//   - channel: Channel, either 1, 2 3 or 4 corresponding to numbers on Mote
func (m *Mote) GetPixelCount(channel int) (int, error) {
	if channel > NumChannels || channel < 1 {
		return 0, fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return 0, fmt.Errorf("channel %d has not been set up", channel)
	}
	return len(m.Channels[channel-1].Pixels), nil
}

// GetPixel gets the RGB colour of a single pixel, on a single channel, taking the following parameters.
//   - channel: Channel, either 1, 2, 3 or 4 corresponding to numbers on Mote
//   - index: Index of pixel to set, from 0 up
func (m *Mote) GetPixel(channel, index int) (Pixel, error) {
	if channel > NumChannels || channel < 1 {
		return Pixel{}, fmt.Errorf("channel index must be between 1 and 4")
	}
	if m.Channels[channel-1] == nil {
		return Pixel{}, fmt.Errorf("channel %d has not been set up", channel)
	}
	if index >= len(m.Channels[channel-1].Pixels) {
		return Pixel{}, fmt.Errorf("pixel index must be < %d", len(m.Channels[channel-1].Pixels))
	}
	return m.Channels[channel-1].Pixels[index], nil
}

func findSerialPort(v, p int, n string) *string {
	ports, _ := serial.ListPorts()
	if len(ports) == 0 {
		return nil
	}
	for _, info := range ports {
		portName := info.Name()
		if vid, pid, err := info.USBVIDPID(); err == nil {
			if vid == v && pid == p {
				log.Printf("found Mote connected to port: %s\n", portName)
				return &portName
			}
		}
		product := info.USBProduct()
		if product == ProductName {
			log.Printf("found Mote connected to port: %s\n", portName)
			return &portName
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// myLG is command line looking glass that written with Go language
// it tries from its own icmp and external looking glasses tools
package main

import (
	"errors"
	"fmt"
	"github.com/briandowns/spinner"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"time"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/disc"
	"github.com/mehrdadrad/mylg/http/ping"
	"github.com/mehrdadrad/mylg/icmp"
	"github.com/mehrdadrad/mylg/lg"
	"github.com/mehrdadrad/mylg/ns"
	"github.com/mehrdadrad/mylg/packet"
	"github.com/mehrdadrad/mylg/peeringdb"
	"github.com/mehrdadrad/mylg/scan"
	"github.com/mehrdadrad/mylg/services/httpd"
	"github.com/mehrdadrad/mylg/whois"
)

// Provider represents looking glass
type Provider interface {
	Set(host, version string)
	GetDefaultNode() string
	GetNodes() []string
	ChangeNode(node string) bool
	Ping() (string, error)
	Trace() chan string
	BGP() chan string
}

const (
	version = "0.2.0"
)

var (
	// register looking glass hosts
	providers = map[string]Provider{"telia": new(lg.Telia), "level3": new(lg.Level3), "cogent": new(lg.Cogent)}
	pNames    = providerNames()
	req       = make(chan string, 1)
	nxt       = make(chan struct{}, 1)
	spin      = spinner.New(spinner.CharSets[26], 220*time.Millisecond)
	args      string
	prompt    string
	cPName    string
	nsr       *ns.Request
	c         *cli.Readline
)

// providerName
func providerNames() []string {
	pNames := []string{}
	for p := range providers {
		pNames = append(pNames, p)
	}
	return pNames
}

// validateProvider
func validateProvider(p string) (string, error) {
	pNames := []string{}
	match, _ := regexp.MatchString("("+strings.Join(pNames, "|")+")", p)
	p = strings.ToLower(p)
	if match {
		return p, nil
	}
	return "", errors.New("provider not support")

}
func init() {
	// initialize cli
	c = cli.Init("local", version)
	go c.Run(req, nxt)
	// initialize name server
	nsr = ns.NewRequest()
	go nsr.Init()
	// start web server
	go httpd.Run()
	// set default provider
	cPName = "local"
}

func main() {
	var (
		request string
		loop    = true
	)

	for loop {
		select {
		case request, loop = <-req:
			if !loop {
				break
			}
			if len(request) < 1 {
				c.Next()
				continue
			}
			subReq := cli.CMDReg.FindStringSubmatch(request)
			if len(subReq) == 0 {
				println("syntax error")
				c.Next()
				continue
			}
			prompt = c.GetPrompt()
			args = strings.TrimSpace(subReq[2])
			cmd := strings.TrimSpace(subReq[1])
			switch {
			case cmd == "hping" && cPName == "local":
				hping()
			case cmd == "ping" && cPName == "local":
				pingLocal()
			case cmd == "ping":
				pingLG()
			case cmd == "trace":
				trace()
			case cmd == "bgp":
				BGP()
			case cmd == "dig":
				nsr.Dig(args)
			case cmd == "node":
				node()
			case cmd == "local":
				nsr.Local()
				cPName = "local"
				c.SetPrompt(cPName)
			case cmd == "connect":
				connect()
			case cmd == "lg":
				c.SetPrompt("lg")
				c.UpdateCompleter("connect", pNames)
			case cmd == "ns":
				c.UpdateCompleter("connect", nsr.CountryList())
				c.UpdateCompleter("node", []string{})
				c.SetPrompt("ns")
			case cmd == "whois":
				whois.Lookup(args)
			case cmd == "peering":
				peeringdb.Search(args)
			case cmd == "scan":
				scanPorts()
			case cmd == "mode":
				mode()
			case cmd == "web":
				web()
			case cmd == "dump":
				dump()
			case cmd == "disc":
				discovery()
			case cmd == "help":
				c.Help()
			case cmd == "exit", cmd == "quit":
				c.Close(nxt)
				close(req)
			}
			// next line
			c.Next()
		}
	}
}

// node handles node cmd
func node() {
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if _, ok := providers[cPName]; ok {
			if providers[cPName].ChangeNode(args) {
				c.UpdatePromptN(args, 3)
				return
			}
		}
		println("the specified node doesn't support")
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkNode(args) {
			println("error: argument is not valid")
		} else {
			c.UpdatePromptN(args, 3)
		}
	default:
		if cPName == "local" {
			println("local doesn't support node")
		}

	}
}

// web tries to open web interface at default web browser
func web() {
	var openCmd = "open"
	println("opening default web broswer ...")
	if runtime.GOOS != "darwin" {
		openCmd = "xdg-open"
	}
	cmd := exec.Command(openCmd, "http://localhost:8080")
	err := cmd.Start()
	if err != nil {
		println("error opening default browser")
	}

}

// dump provides decoding packets
func dump() {
	p, err := packet.NewPacket(args)
	if err != nil {
		return
	}
	for l := range p.Open() {
		l.PrintPretty()
	}
}

// connect handles connect cmd
func connect() {
	var (
		pName string
		err   error
	)
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if pName, err = validateProvider(args); err != nil {
			println("provider not available")
			c.Next()
			return
		}
		cPName = pName
		if _, ok := providers[cPName]; ok {
			c.UpdatePromptN(cPName+"/"+providers[cPName].GetDefaultNode(), 2)
			go func() {
				c.UpdateCompleter("node", providers[cPName].GetNodes())
			}()
		} else {
			println("it doesn't support")
		}
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkCountry(args) {
			println("error: argument is not valid")
		} else {
			c.SetPrompt("ns/" + args)
			c.UpdateCompleter("node", nsr.NodeList())
		}
	}
}

// mode set editor mode
func mode() {
	if args == "vim" {
		c.SetVim()
	} else if args == "emacs" {
		c.SetEmacs()
	} else {
		println("the request mode doesn't support")
	}
}

// trace tries to trace from local and lg
func trace() {
	switch {
	case strings.HasPrefix(prompt, "local"):
		trace := icmp.Trace{}
		trace.Run(args)
	case strings.HasPrefix(prompt, "lg"):
		providers[cPName].Set(args, "ipv4")
		for l := range providers[cPName].Trace() {
			println(l)
		}
	}
}

// hping tries to ping a web server by http
func hping() {
	p, err := ping.NewPing(args)
	if err != nil {
		println(err.Error())
	} else {
		p.Run()
	}
}

// pingLG tries to ping through a looking glass
func pingLG() {
	spin.Prefix = "please wait "
	spin.Start()
	providers[cPName].Set(args, "ipv4")
	m, err := providers[cPName].Ping()
	spin.Stop()
	if err != nil {
		println(err.Error())
	} else {
		println(m)
	}
}

// pingLocal tries to ping from local source ip
func pingLocal() {
	pFmt := "%d bytes from %s icmp_seq=%d time=%f ms"
	p, err := icmp.NewPing(args)
	if err != nil {
		return
	}
	resp := p.Run()
	for r := range resp {
		if r.Error != nil {
			println(r.Error.Error())
			continue
		}
		msg := fmt.Sprintf(pFmt, r.Size, r.Addr, r.Sequence, r.RTT)
		println(msg)
	}
}

// scanPorts tries to scan tcp/ip ports
func scanPorts() {
	scan, err := scan.NewScan(args)
	if err != nil {
		println(err.Error())
	} else {
		scan.Run()
	}
}

// BGP tries to get BGP lookup from a LG
func BGP() {
	if cPName == "local" {
		println("no provider selected")
		return
	}
	providers[cPName].Set(args, "ipv4")
	for l := range providers[cPName].BGP() {
		println(l)
	}
}

func discovery() {
	d := disc.New()
	d.PingLan()
	if err := d.GetARPTable(); err != nil {
		println(err.Error())
		return
	}
	d.LoadOUI()
	d.PrintPretty()
}
<commit_msg>enhanced discovery command<commit_after>// myLG is command line looking glass that written with Go language
// it tries from its own icmp and external looking glasses tools
package main

import (
	"errors"
	"fmt"
	"github.com/briandowns/spinner"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/disc"
	"github.com/mehrdadrad/mylg/http/ping"
	"github.com/mehrdadrad/mylg/icmp"
	"github.com/mehrdadrad/mylg/lg"
	"github.com/mehrdadrad/mylg/ns"
	"github.com/mehrdadrad/mylg/packet"
	"github.com/mehrdadrad/mylg/peeringdb"
	"github.com/mehrdadrad/mylg/scan"
	"github.com/mehrdadrad/mylg/services/httpd"
	"github.com/mehrdadrad/mylg/whois"
)

// Provider represents looking glass
type Provider interface {
	Set(host, version string)
	GetDefaultNode() string
	GetNodes() []string
	ChangeNode(node string) bool
	Ping() (string, error)
	Trace() chan string
	BGP() chan string
}

const (
	version = "0.2.0"
)

var (
	// register looking glass hosts
	providers = map[string]Provider{"telia": new(lg.Telia), "level3": new(lg.Level3), "cogent": new(lg.Cogent)}
	pNames    = providerNames()
	req       = make(chan string, 1)
	nxt       = make(chan struct{}, 1)
	spin      = spinner.New(spinner.CharSets[26], 220*time.Millisecond)
	args      string
	prompt    string
	cPName    string
	nsr       *ns.Request
	c         *cli.Readline
)

// providerName
func providerNames() []string {
	pNames := []string{}
	for p := range providers {
		pNames = append(pNames, p)
	}
	return pNames
}

// validateProvider
func validateProvider(p string) (string, error) {
	pNames := []string{}
	match, _ := regexp.MatchString("("+strings.Join(pNames, "|")+")", p)
	p = strings.ToLower(p)
	if match {
		return p, nil
	}
	return "", errors.New("provider not support")

}
func init() {
	// initialize cli
	c = cli.Init("local", version)
	go c.Run(req, nxt)
	// initialize name server
	nsr = ns.NewRequest()
	go nsr.Init()
	// start web server
	go httpd.Run()
	// set default provider
	cPName = "local"
}

func main() {
	var (
		request string
		loop    = true
	)

	for loop {
		select {
		case request, loop = <-req:
			if !loop {
				break
			}
			if len(request) < 1 {
				c.Next()
				continue
			}
			subReq := cli.CMDReg.FindStringSubmatch(request)
			if len(subReq) == 0 {
				println("syntax error")
				c.Next()
				continue
			}
			prompt = c.GetPrompt()
			args = strings.TrimSpace(subReq[2])
			cmd := strings.TrimSpace(subReq[1])
			switch {
			case cmd == "hping" && cPName == "local":
				hping()
			case cmd == "ping" && cPName == "local":
				pingLocal()
			case cmd == "ping":
				pingLG()
			case cmd == "trace":
				trace()
			case cmd == "bgp":
				BGP()
			case cmd == "dig":
				nsr.Dig(args)
			case cmd == "node":
				node()
			case cmd == "local":
				nsr.Local()
				cPName = "local"
				c.SetPrompt(cPName)
			case cmd == "connect":
				connect()
			case cmd == "lg":
				c.SetPrompt("lg")
				c.UpdateCompleter("connect", pNames)
			case cmd == "ns":
				c.UpdateCompleter("connect", nsr.CountryList())
				c.UpdateCompleter("node", []string{})
				c.SetPrompt("ns")
			case cmd == "whois":
				whois.Lookup(args)
			case cmd == "peering":
				peeringdb.Search(args)
			case cmd == "scan":
				scanPorts()
			case cmd == "mode":
				mode()
			case cmd == "web":
				web()
			case cmd == "dump":
				dump()
			case cmd == "disc":
				discovery()
			case cmd == "help":
				c.Help()
			case cmd == "exit", cmd == "quit":
				c.Close(nxt)
				close(req)
			}
			// next line
			c.Next()
		}
	}
}

// node handles node cmd
func node() {
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if _, ok := providers[cPName]; ok {
			if providers[cPName].ChangeNode(args) {
				c.UpdatePromptN(args, 3)
				return
			}
		}
		println("the specified node doesn't support")
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkNode(args) {
			println("error: argument is not valid")
		} else {
			c.UpdatePromptN(args, 3)
		}
	default:
		if cPName == "local" {
			println("local doesn't support node")
		}

	}
}

// web tries to open web interface at default web browser
func web() {
	var openCmd = "open"
	println("opening default web broswer ...")
	if runtime.GOOS != "darwin" {
		openCmd = "xdg-open"
	}
	cmd := exec.Command(openCmd, "http://localhost:8080")
	err := cmd.Start()
	if err != nil {
		println("error opening default browser")
	}

}

// dump provides decoding packets
func dump() {
	p, err := packet.NewPacket(args)
	if p == nil || err != nil {
		return
	}
	for l := range p.Open() {
		l.PrintPretty()
	}
}

// connect handles connect cmd
func connect() {
	var (
		pName string
		err   error
	)
	switch {
	case strings.HasPrefix(prompt, "lg"):
		if pName, err = validateProvider(args); err != nil {
			println("provider not available")
			c.Next()
			return
		}
		cPName = pName
		if _, ok := providers[cPName]; ok {
			c.UpdatePromptN(cPName+"/"+providers[cPName].GetDefaultNode(), 2)
			go func() {
				c.UpdateCompleter("node", providers[cPName].GetNodes())
			}()
		} else {
			println("it doesn't support")
		}
	case strings.HasPrefix(prompt, "ns"):
		if !nsr.ChkCountry(args) {
			println("error: argument is not valid")
		} else {
			c.SetPrompt("ns/" + args)
			c.UpdateCompleter("node", nsr.NodeList())
		}
	}
}

// mode set editor mode
func mode() {
	if args == "vim" {
		c.SetVim()
	} else if args == "emacs" {
		c.SetEmacs()
	} else {
		println("the request mode doesn't support")
	}
}

// trace tries to trace from local and lg
func trace() {
	switch {
	case strings.HasPrefix(prompt, "local"):
		trace := icmp.Trace{}
		trace.Run(args)
	case strings.HasPrefix(prompt, "lg"):
		providers[cPName].Set(args, "ipv4")
		for l := range providers[cPName].Trace() {
			println(l)
		}
	}
}

// hping tries to ping a web server by http
func hping() {
	p, err := ping.NewPing(args)
	if err != nil {
		println(err.Error())
	} else {
		p.Run()
	}
}

// pingLG tries to ping through a looking glass
func pingLG() {
	spin.Prefix = "please wait "
	spin.Start()
	providers[cPName].Set(args, "ipv4")
	m, err := providers[cPName].Ping()
	spin.Stop()
	if err != nil {
		println(err.Error())
	} else {
		println(m)
	}
}

// pingLocal tries to ping from local source ip
func pingLocal() {
	pFmt := "%d bytes from %s icmp_seq=%d time=%f ms"
	p, err := icmp.NewPing(args)
	if err != nil {
		return
	}
	resp := p.Run()
	for r := range resp {
		if r.Error != nil {
			println(r.Error.Error())
			continue
		}
		msg := fmt.Sprintf(pFmt, r.Size, r.Addr, r.Sequence, r.RTT)
		println(msg)
	}
}

// scanPorts tries to scan tcp/ip ports
func scanPorts() {
	scan, err := scan.NewScan(args)
	if err != nil {
		println(err.Error())
	} else {
		scan.Run()
	}
}

// BGP tries to get BGP lookup from a LG
func BGP() {
	if cPName == "local" {
		println("no provider selected")
		return
	}
	providers[cPName].Set(args, "ipv4")
	for l := range providers[cPName].BGP() {
		println(l)
	}
}

// discovery handles disc command
func discovery() {
	var (
		wg sync.WaitGroup
		//ts = time.Now()
	)

	spin.Prefix = "please wait "
	spin.Start()

	d := disc.New(args)
	// load OUI async
	go func() {
		wg.Add(1)
		d.LoadOUI()
		wg.Done()
	}()
	d.PingLan()
	//s := 10 - int(time.Since(ts).Seconds())
	time.Sleep(5 * time.Second)
	if err := d.GetARPTable(); err != nil {
		println(err.Error())
		return
	}
	wg.Wait()
	spin.Stop()

	println("Network LAN Discovery")
	d.PrintPretty()
}
<|endoftext|>