<commit_before>package tdigest

import (
	"fmt"
	"log"
	"math"
	"math/rand"
)

// A TDigest is an efficient data structure for computing streaming
// approximate quantiles of a dataset. It supports two methods:
//
// Add will add a value to the TDigest, updating all quantiles. A
// weight can be specified; use weight of 1 if you don't care about
// weighting your dataset.
//
// Quantile(q) will estimate the qth quantile value of the
// dataset. The input value of q should be in the range [0.0, 1.0]; if
// it is outside that range, it will be clipped into it automatically.
//
// Calling Quantile on a TDigest with no data should return NaN.
type TDigest interface {
	Add(val float64, weight int)
	Quantile(q float64) (val float64)
}

// New produces a new TDigest. The input compression value, which
// should be >= 1.0, will control how aggressively the TDigest
// compresses data together.
//
// The original TDigest paper suggests using a value of 100 for a good
// balance between precision and efficiency. It will land at very
// small (think like 1e-6 percentile points) errors at extreme points
// in the distribution, and compression ratios of around 500 for large
// data sets (1 millionish datapoints).
func New(compression float64) TDigest {
	return newCentroidSet(compression)
}

// centroid is a simple container for a mean,count pair.
type centroid struct {
	mean  float64
	count int
}

func (c *centroid) String() string {
	return fmt.Sprintf("c{%f x%d}", c.mean, c.count)
}

type centroidSet struct {
	centroids   []*centroid
	compression float64
	countTotal  int

	reclusterAt int
}

func newCentroidSet(compression float64) *centroidSet {
	return &centroidSet{
		centroids:   make([]*centroid, 0),
		compression: compression,
		countTotal:  0,
	}
}

// Find the indexes of centroids which have the minimum distance to the
// input value.
//
// TODO: Use a better data structure to avoid this loop.
func (cs *centroidSet) nearest(val float64) []int {
	var (
		nearestDist float64
		thisDist    float64
		result      []int
	)
	for i, c := range cs.centroids {
		thisDist = math.Abs(val - c.mean)
		// on the first loop, the nearest hasn't yet been found, so
		// mark the first centroid as a 'winner.'
		if len(result) == 0 {
			nearestDist = thisDist
			result = []int{i}
			continue
		}
		if thisDist < nearestDist {
			// we have a new winner!
			nearestDist = thisDist
			result = []int{i}
			continue
		}
		if thisDist == nearestDist {
			// we have a tie
			result = append(result, i)
			continue
		}
		if thisDist > nearestDist {
			// Since cs.centroids is sorted by mean, this means we
			// have passed the best spot, so we may as well break
			break
		}
	}
	return result
}

// returns the maximum weight that can be placed at specified index
func (cs *centroidSet) weightLimit(idx int) int {
	ptile := cs.quantileOf(idx)
	limit := int(4 * cs.compression * ptile * (1 - ptile) * float64(len(cs.centroids)))
	return limit
}

// checks whether the centroid has room for more weight
func (cs *centroidSet) centroidHasRoom(idx int) bool {
	return cs.centroids[idx].count < cs.weightLimit(idx)
}

// find which centroid to add the value to (by index)
func (cs *centroidSet) findAddTarget(val float64) int {
	var (
		nearest  []int = cs.nearest(val)
		eligible []int
	)

	for _, c := range nearest {
		// if there is room for more weight at this centroid...
		if cs.centroidHasRoom(c) {
			eligible = append(eligible, c)
		}
	}

	if len(eligible) == 0 {
		return -1
	}
	if len(eligible) == 1 {
		return eligible[0]
	}

	// Multiple eligible centroids to add to. They must be equidistant
	// from this value. Four cases are possible:
	//
	//   1. All eligible centroids' means are less than val
	//   2. Some eligible centroids' means are less than val, some are greater
	//   3. All eligible centroids' means are exactly equal to val
	//   4. All eligible centroids' means are greater than val
	//
	// If 1, then we should take the highest indexed centroid to
	// preserve ordering.  If 4, we should take the lowest for the
	// same reason. If 2 or 3, we can pick randomly.

	var anyLesser, anyGreater bool
	for _, c := range eligible {
		m := cs.centroids[c].mean
		if m < val {
			anyLesser = true
		} else if m > val {
			anyGreater = true
		}
	}

	// case 1: all are less, none are greater. Take highest one.
	if anyLesser && !anyGreater {
		greatest := eligible[0]
		for _, c := range eligible[1:] {
			if c > greatest {
				greatest = c
			}
		}
		return greatest
	}

	// case 4: all are greater, none are less. Take lowest one.
	if !anyLesser && anyGreater {
		least := eligible[0]
		for _, c := range eligible[1:] {
			if c < least {
				least = c
			}
		}
		return least
	}

	// case 2 or 3: Anything else. Pick randomly.
	return eligible[rand.Intn(len(eligible))]
}

func (cs *centroidSet) addNewCentroid(mean float64, weight int) {
	var idx int = len(cs.centroids)

	for i, c := range cs.centroids {
		// add in sorted order
		if mean < c.mean {
			idx = i
			break
		}
	}

	cs.centroids = append(cs.centroids, nil)
	copy(cs.centroids[idx+1:], cs.centroids[idx:])
	cs.centroids[idx] = &centroid{mean, weight}
}

// Add a value to the centroidSet.
func (cs *centroidSet) Add(val float64, weight int) {
	cs.countTotal += weight
	var idx = cs.findAddTarget(val)

	if idx == -1 {
		cs.addNewCentroid(val, weight)
		return
	}

	c := cs.centroids[idx]

	limit := cs.weightLimit(idx)
	// how much weight will we be adding?
	// if adding this node to this centroid would put it over the
	// weight limit, just add the most we can and recur with the remainder
	if c.count+weight > limit {
		log.Printf("splitting value %f weight=%d", val, weight)
		add := limit - c.count
		if add < 0 {
			// this node was already overweight
			add = 0
		}
		remainder := weight - add

		c.count += add
		c.mean = c.mean + float64(add)*(val-c.mean)/float64(c.count)

		cs.Add(val, remainder)
	} else {
		c.count += weight
		c.mean = c.mean + float64(weight)*(val-c.mean)/float64(c.count)
	}

}

// returns the approximate quantile that a particular centroid
// represents
func (cs *centroidSet) quantileOf(idx int) float64 {
	total := 0
	for _, c := range cs.centroids[:idx] {
		total += c.count
	}
	return (float64(cs.centroids[idx].count/2) + float64(total)) / float64(cs.countTotal)
}

// Quantile returns the approximate value at a quantile (eg the 99th
// percentile value would be centroidSet.quantileValue(0.99))
func (cs *centroidSet) Quantile(q float64) float64 {
	var n = len(cs.centroids)
	if n == 0 {
		return math.NaN()
	}
	if n == 1 {
		return cs.centroids[0].mean
	}

	if q < 0 {
		q = 0
	} else if q > 1 {
		q = 1
	}

	// rescale into count units instead of 0 to 1 units
	q = float64(cs.countTotal) * q
	// find the first centroid which straddles q
	var (
		qTotal float64 = 0
		i      int
	)
	for i = 0; i < n && float64(cs.centroids[i].count)/2+qTotal < q; i++ {
		qTotal += float64(cs.centroids[i].count)
	}

	if i == 0 {
		// special case 1: the targeted quantile is before the
		// left-most centroid. extrapolate from the slope from
		// centroid0 to centroid1.
		c0 := cs.centroids[0]
		c1 := cs.centroids[1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - float64(c0.count)/2 // this is negative
		return c0.mean + slope*deltaQ
	}
	if i == n {
		// special case 2: the targeted quantile is from the
		// right-most centroid. extrapolate from the slope at the
		// right edge.
		c0 := cs.centroids[n-2]
		c1 := cs.centroids[n-1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - (qTotal - float64(c1.count)/2)
		return c1.mean + slope*deltaQ
	}
	// common case: targeted quantile is between 2 centroids
	c0 := cs.centroids[i-1]
	c1 := cs.centroids[i]
	slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
	deltaQ := q - (float64(c1.count)/2 + qTotal)
	return c1.mean + slope*deltaQ
}
<commit_msg>Optimize centroidSet.nearest<commit_after>package tdigest

import (
	"fmt"
	"log"
	"math"
	"math/rand"
)

// A TDigest is an efficient data structure for computing streaming
// approximate quantiles of a dataset. It supports two methods:
//
// Add will add a value to the TDigest, updating all quantiles. A
// weight can be specified; use weight of 1 if you don't care about
// weighting your dataset.
//
// Quantile(q) will estimate the qth quantile value of the
// dataset. The input value of q should be in the range [0.0, 1.0]; if
// it is outside that range, it will be clipped into it automatically.
//
// Calling Quantile on a TDigest with no data should return NaN.
type TDigest interface {
	Add(val float64, weight int)
	Quantile(q float64) (val float64)
}

// New produces a new TDigest. The input compression value, which
// should be >= 1.0, will control how aggressively the TDigest
// compresses data together.
//
// The original TDigest paper suggests using a value of 100 for a good
// balance between precision and efficiency. It will land at very
// small (think like 1e-6 percentile points) errors at extreme points
// in the distribution, and compression ratios of around 500 for large
// data sets (1 millionish datapoints).
func New(compression float64) TDigest {
	return newCentroidSet(compression)
}

// centroid is a simple container for a mean,count pair.
type centroid struct {
	mean  float64
	count int
}

func (c *centroid) String() string {
	return fmt.Sprintf("c{%f x%d}", c.mean, c.count)
}

type centroidSet struct {
	centroids   []*centroid
	compression float64
	countTotal  int

	reclusterAt int
}

func newCentroidSet(compression float64) *centroidSet {
	return &centroidSet{
		centroids:   make([]*centroid, 0),
		compression: compression,
		countTotal:  0,
	}
}

// Find the indexes of centroids which have the minimum distance to the
// input value.
//
// TODO: Use a better data structure to avoid this loop.
func (cs *centroidSet) nearest(val float64) []int {
	var (
		nearestDist float64 = math.Inf(+1)
		thisDist    float64
		delta       float64
		result      []int = make([]int, 0)
	)
	for i, c := range cs.centroids {
		thisDist = val - c.mean
		if thisDist < 0 {
			thisDist *= -1
		}

		delta = thisDist - nearestDist
		switch {
		case delta < 0:
			// we have a new winner!
			nearestDist = thisDist
			result = result[0:0] // wipe result
			result = append(result, i)
		case delta == 0:
			// we have a tie
			result = append(result, i)
		default:
			// Since cs.centroids is sorted by mean, this means we
			// have passed the best spot, so we may as well break
			break
		}
	}
	return result
}

// returns the maximum weight that can be placed at specified index
func (cs *centroidSet) weightLimit(idx int) int {
	ptile := cs.quantileOf(idx)
	limit := int(4 * cs.compression * ptile * (1 - ptile) * float64(len(cs.centroids)))
	return limit
}

// checks whether the centroid has room for more weight
func (cs *centroidSet) centroidHasRoom(idx int) bool {
	return cs.centroids[idx].count < cs.weightLimit(idx)
}

// find which centroid to add the value to (by index)
func (cs *centroidSet) findAddTarget(val float64) int {
	var (
		nearest  []int = cs.nearest(val)
		eligible []int
	)

	for _, c := range nearest {
		// if there is room for more weight at this centroid...
		if cs.centroidHasRoom(c) {
			eligible = append(eligible, c)
		}
	}

	if len(eligible) == 0 {
		return -1
	}
	if len(eligible) == 1 {
		return eligible[0]
	}

	// Multiple eligible centroids to add to. They must be equidistant
	// from this value. Four cases are possible:
	//
	//   1. All eligible centroids' means are less than val
	//   2. Some eligible centroids' means are less than val, some are greater
	//   3. All eligible centroids' means are exactly equal to val
	//   4. All eligible centroids' means are greater than val
	//
	// If 1, then we should take the highest indexed centroid to
	// preserve ordering.  If 4, we should take the lowest for the
	// same reason. If 2 or 3, we can pick randomly.

	var anyLesser, anyGreater bool
	for _, c := range eligible {
		m := cs.centroids[c].mean
		if m < val {
			anyLesser = true
		} else if m > val {
			anyGreater = true
		}
	}

	// case 1: all are less, none are greater. Take highest one.
	if anyLesser && !anyGreater {
		greatest := eligible[0]
		for _, c := range eligible[1:] {
			if c > greatest {
				greatest = c
			}
		}
		return greatest
	}

	// case 4: all are greater, none are less. Take lowest one.
	if !anyLesser && anyGreater {
		least := eligible[0]
		for _, c := range eligible[1:] {
			if c < least {
				least = c
			}
		}
		return least
	}

	// case 2 or 3: Anything else. Pick randomly.
	return eligible[rand.Intn(len(eligible))]
}

func (cs *centroidSet) addNewCentroid(mean float64, weight int) {
	var idx int = len(cs.centroids)

	for i, c := range cs.centroids {
		// add in sorted order
		if mean < c.mean {
			idx = i
			break
		}
	}

	cs.centroids = append(cs.centroids, nil)
	copy(cs.centroids[idx+1:], cs.centroids[idx:])
	cs.centroids[idx] = &centroid{mean, weight}
}

// Add a value to the centroidSet.
func (cs *centroidSet) Add(val float64, weight int) {
	cs.countTotal += weight
	var idx = cs.findAddTarget(val)

	if idx == -1 {
		cs.addNewCentroid(val, weight)
		return
	}

	c := cs.centroids[idx]

	limit := cs.weightLimit(idx)
	// how much weight will we be adding?
	// if adding this node to this centroid would put it over the
	// weight limit, just add the most we can and recur with the remainder
	if c.count+weight > limit {
		log.Printf("splitting value %f weight=%d", val, weight)
		add := limit - c.count
		if add < 0 {
			// this node was already overweight
			add = 0
		}
		remainder := weight - add

		c.count += add
		c.mean = c.mean + float64(add)*(val-c.mean)/float64(c.count)

		cs.Add(val, remainder)
	} else {
		c.count += weight
		c.mean = c.mean + float64(weight)*(val-c.mean)/float64(c.count)
	}

}

// returns the approximate quantile that a particular centroid
// represents
func (cs *centroidSet) quantileOf(idx int) float64 {
	total := 0
	for _, c := range cs.centroids[:idx] {
		total += c.count
	}
	return (float64(cs.centroids[idx].count/2) + float64(total)) / float64(cs.countTotal)
}

// Quantile returns the approximate value at a quantile (eg the 99th
// percentile value would be centroidSet.quantileValue(0.99))
func (cs *centroidSet) Quantile(q float64) float64 {
	var n = len(cs.centroids)
	if n == 0 {
		return math.NaN()
	}
	if n == 1 {
		return cs.centroids[0].mean
	}

	if q < 0 {
		q = 0
	} else if q > 1 {
		q = 1
	}

	// rescale into count units instead of 0 to 1 units
	q = float64(cs.countTotal) * q
	// find the first centroid which straddles q
	var (
		qTotal float64 = 0
		i      int
	)
	for i = 0; i < n && float64(cs.centroids[i].count)/2+qTotal < q; i++ {
		qTotal += float64(cs.centroids[i].count)
	}

	if i == 0 {
		// special case 1: the targeted quantile is before the
		// left-most centroid. extrapolate from the slope from
		// centroid0 to centroid1.
		c0 := cs.centroids[0]
		c1 := cs.centroids[1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - float64(c0.count)/2 // this is negative
		return c0.mean + slope*deltaQ
	}
	if i == n {
		// special case 2: the targeted quantile is from the
		// right-most centroid. extrapolate from the slope at the
		// right edge.
		c0 := cs.centroids[n-2]
		c1 := cs.centroids[n-1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - (qTotal - float64(c1.count)/2)
		return c1.mean + slope*deltaQ
	}
	// common case: targeted quantile is between 2 centroids
	c0 := cs.centroids[i-1]
	c1 := cs.centroids[i]
	slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
	deltaQ := q - (float64(c1.count)/2 + qTotal)
	return c1.mean + slope*deltaQ
}
<|endoftext|>


<commit_before>package telegram

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	elasticsearch7 "github.com/elastic/go-elasticsearch/v7"
	"github.com/go-redis/redis"
	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
	logging "github.com/op/go-logging"

	"yubari/pixiv"
)

const (
	tgDeleteTube = "tg_delete"
	tgPixivTube  = "tg_pixiv"
)

type Config struct {
	Token          string  `json:"token"`
	SelfID         int64   `json:"selfID"`
	WhitelistChats []int64 `json:"whitelistChats"`
	ComicPath      string  `json:"comicPath"`
	DeleteDelay    string  `json:"deleteDelay"`
}

type DownloadPixiv struct {
	ChatID    int64
	MessageID int
	PixivID   uint64
}

type Bot struct {
	Name           string
	SelfID         int64
	WhitelistChats []int64
	ComicPath      string
	PixivPath      string
	TwitterImgPath string
	DeleteDelay    time.Duration
	Client         *tgbotapi.BotAPI
	Queue          *bt.Pool
	logger         *logging.Logger
	redis          *redis.Client
	es             *elasticsearch7.Client
}

func NewBot(cfg *Config) (b *Bot, err error) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		return nil, fmt.Errorf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		return nil, fmt.Errorf("delete delay error: %+v", err)
	}

	b = &Bot{
		Name:           bot.Self.UserName,
		SelfID:         cfg.SelfID,
		WhitelistChats: cfg.WhitelistChats,
		ComicPath:      cfg.ComicPath,
		DeleteDelay:    delay,
		Client:         bot,
	}
	return
}

func (b *Bot) WithLogger(logger *logging.Logger) *Bot {
	b.logger = logger
	return b
}

func (b *Bot) WithRedis(rds *redis.Client) *Bot {
	b.redis = rds
	return b
}

func (b *Bot) WithPixivImg(imgPath string) *Bot {
	b.PixivPath = imgPath
	return b
}

func (b *Bot) WithTwitterImg(imgPath string) *Bot {
	b.TwitterImgPath = imgPath
	return b
}

func (b *Bot) WithQueue(queue *bt.Pool) *Bot {
	b.Queue = queue
	return b
}

func (b *Bot) WithES(es *elasticsearch7.Client) *Bot {
	b.es = es
	return b
}

func (b *Bot) putQueue(msg []byte, tube string) {
	conn, err := b.Queue.Get()
	if err != nil {
		b.logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(tube)
	_, err = conn.Put(msg, 1, b.DeleteDelay, time.Minute)
	if err != nil {
		b.logger.Errorf("%+v", err)
		return
	}
}

func (b *Bot) isAuthedChat(c *tgbotapi.Chat) bool {
	for _, w := range b.WhitelistChats {
		if c.ID == w {
			return true
		}
	}
	return false
}

func (b *Bot) Send(chat int64, msg string) (tgbotapi.Message, error) {
	b.logger.Debugf("[%d]%s", chat, msg)
	message := tgbotapi.NewMessage(chat, msg)
	message.DisableNotification = true
	return b.Client.Send(message)
}

func (b *Bot) SendPixivIllust(target int64, id uint64) {
	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData("⭕", buildReactionData("pixivIllust", strconv.FormatUint(id, 10), "like")),
		tgbotapi.NewInlineKeyboardButtonData("❌", buildReactionData("pixivIllust", strconv.FormatUint(id, 10), "diss")),
	)
	msg := tgbotapi.NewMessage(target, pixiv.URLWithID(id))
	msg.ReplyMarkup = tgbotapi.NewInlineKeyboardMarkup(row)
	msg.DisableNotification = true
	_, err := b.Client.Send(msg)
	if err != nil {
		b.logger.Errorf("%+v", err)
	}
}

func (b *Bot) startDownloadPixiv() {
	for {
		conn, err := b.Queue.Get()
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(tgPixivTube)
		job, err := conn.Reserve()
		if err != nil {
			b.logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &DownloadPixiv{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			b.logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				b.logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}

		sizes, errs := pixiv.Download(msg.PixivID, b.PixivPath)
		for i := range sizes {
			if errs[i] != nil {
				err = errs[i]
				break
			}
			if sizes[i] == 0 {
				continue
			}
			b.logger.Debugf("download pixiv %d_p%d: %d bytes", msg.PixivID, i, sizes[i])
		}
		if err != nil {
			b.logger.Errorf("%+v", err)
			continue
		}

		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.ChatID,
			MessageID: msg.MessageID,
		}
		_, err = b.Client.DeleteMessage(delMsg)
		if err != nil {
			b.logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				b.logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		b.Queue.Release(conn, false)
	}
}

func (b *Bot) startDeleteMessage() {
	for {
		conn, err := b.Queue.Get()
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(tgDeleteTube)
		job, err := conn.Reserve()
		if err != nil {
			b.logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			b.logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				b.logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		b.logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = b.Client.DeleteMessage(delMsg)
		if err != nil {
			b.logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				b.logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		b.Queue.Release(conn, false)
	}
}

func (b *Bot) Start() {
	go b.startDeleteMessage()
	go b.startDownloadPixiv()

	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := b.Client.GetUpdatesChan(u)
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				b.logger.Infof(
					"recv:(%d)[%s]reaction:{%s}",
					update.CallbackQuery.Message.Chat.ID,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic", "pixiv":
					go onReaction(b, update.CallbackQuery)
				case "pixivIllust":
					if !b.isAuthedChat(update.CallbackQuery.Message.Chat) {
						b.logger.Warning("reaction from illegal chat, ignore")
						break
					}
					go onReactionSelf(b, update.CallbackQuery)
				case "search":
					go onReactionSearch(b, update.CallbackQuery)
				default:
				}
				continue
			} else {
				continue
			}

			if message.Chat.IsGroup() {
				b.logger.Infof(
					"recv:(%d)[%s:%s]{%s}",
					message.Chat.ID,
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				b.logger.Infof(
					"recv:(%d)[%s]{%s}",
					message.Chat.ID,
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(b, message)
				case "roll":
					go onRoll(b, message)
				case "comic":
					go onComic(b, message)
				case "pic":
					go onPic(b, message)
				case "pixiv":
					go onPixiv(b, message)
				case "search":
					go onSearch(b, message)
				default:
					b.logger.Infof("ignore unknown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				go checkRepeat(b, message)
				go checkPixiv(b, message)
				go checkSave(b, message)
			}
		}
		b.logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func (b *Bot) probate(_type, _id string) error {
	b.logger.Noticef("%s: %s", _type, _id)
	switch _type {
	case "comic":
		fileName := "nhentai.net@" + _id + ".epub"
		return os.Rename(
			filepath.Join(b.ComicPath, fileName),
			filepath.Join(b.ComicPath, "probation", fileName),
		)
	case "pic":
		return os.Rename(
			filepath.Join(b.TwitterImgPath, _id),
			filepath.Join(b.TwitterImgPath, "probation", _id),
		)
	case "pixiv":
		return os.Rename(
			filepath.Join(b.PixivPath, _id),
			filepath.Join(b.PixivPath, "probation", _id),
		)
	default:
		return fmt.Errorf("prohibit unkown type")
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}
<commit_msg>fix: defer to bury/del job for deleting msg<commit_after>package telegram

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	elasticsearch7 "github.com/elastic/go-elasticsearch/v7"
	"github.com/go-redis/redis"
	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
	logging "github.com/op/go-logging"

	"yubari/pixiv"
)

const (
	tgDeleteTube = "tg_delete"
	tgPixivTube  = "tg_pixiv"
)

type Config struct {
	Token          string  `json:"token"`
	SelfID         int64   `json:"selfID"`
	WhitelistChats []int64 `json:"whitelistChats"`
	ComicPath      string  `json:"comicPath"`
	DeleteDelay    string  `json:"deleteDelay"`
}

type DownloadPixiv struct {
	ChatID    int64
	MessageID int
	PixivID   uint64
}

type Bot struct {
	Name           string
	SelfID         int64
	WhitelistChats []int64
	ComicPath      string
	PixivPath      string
	TwitterImgPath string
	DeleteDelay    time.Duration
	Client         *tgbotapi.BotAPI
	Queue          *bt.Pool
	logger         *logging.Logger
	redis          *redis.Client
	es             *elasticsearch7.Client
}

func NewBot(cfg *Config) (b *Bot, err error) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		return nil, fmt.Errorf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		return nil, fmt.Errorf("delete delay error: %+v", err)
	}

	b = &Bot{
		Name:           bot.Self.UserName,
		SelfID:         cfg.SelfID,
		WhitelistChats: cfg.WhitelistChats,
		ComicPath:      cfg.ComicPath,
		DeleteDelay:    delay,
		Client:         bot,
	}
	return
}

func (b *Bot) WithLogger(logger *logging.Logger) *Bot {
	b.logger = logger
	return b
}

func (b *Bot) WithRedis(rds *redis.Client) *Bot {
	b.redis = rds
	return b
}

func (b *Bot) WithPixivImg(imgPath string) *Bot {
	b.PixivPath = imgPath
	return b
}

func (b *Bot) WithTwitterImg(imgPath string) *Bot {
	b.TwitterImgPath = imgPath
	return b
}

func (b *Bot) WithQueue(queue *bt.Pool) *Bot {
	b.Queue = queue
	return b
}

func (b *Bot) WithES(es *elasticsearch7.Client) *Bot {
	b.es = es
	return b
}

func (b *Bot) putQueue(msg []byte, tube string) {
	conn, err := b.Queue.Get()
	if err != nil {
		b.logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(tube)
	_, err = conn.Put(msg, 1, b.DeleteDelay, time.Minute)
	if err != nil {
		b.logger.Errorf("%+v", err)
		return
	}
}

func (b *Bot) isAuthedChat(c *tgbotapi.Chat) bool {
	for _, w := range b.WhitelistChats {
		if c.ID == w {
			return true
		}
	}
	return false
}

func (b *Bot) Send(chat int64, msg string) (tgbotapi.Message, error) {
	b.logger.Debugf("[%d]%s", chat, msg)
	message := tgbotapi.NewMessage(chat, msg)
	message.DisableNotification = true
	return b.Client.Send(message)
}

func (b *Bot) SendPixivIllust(target int64, id uint64) {
	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData("⭕", buildReactionData("pixivIllust", strconv.FormatUint(id, 10), "like")),
		tgbotapi.NewInlineKeyboardButtonData("❌", buildReactionData("pixivIllust", strconv.FormatUint(id, 10), "diss")),
	)
	msg := tgbotapi.NewMessage(target, pixiv.URLWithID(id))
	msg.ReplyMarkup = tgbotapi.NewInlineKeyboardMarkup(row)
	msg.DisableNotification = true
	_, err := b.Client.Send(msg)
	if err != nil {
		b.logger.Errorf("%+v", err)
	}
}

func (b *Bot) startDownloadPixiv() {
	for {
		conn, err := b.Queue.Get()
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(tgPixivTube)
		job, err := conn.Reserve()
		if err != nil {
			b.logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &DownloadPixiv{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			b.logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				b.logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}

		sizes, errs := pixiv.Download(msg.PixivID, b.PixivPath)
		for i := range sizes {
			if errs[i] != nil {
				err = errs[i]
				break
			}
			if sizes[i] == 0 {
				continue
			}
			b.logger.Debugf("download pixiv %d_p%d: %d bytes", msg.PixivID, i, sizes[i])
		}
		if err != nil {
			b.logger.Errorf("%+v", err)
			continue
		}

		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.ChatID,
			MessageID: msg.MessageID,
		}
		_, err = b.Client.DeleteMessage(delMsg)
		if err != nil {
			b.logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				b.logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		b.Queue.Release(conn, false)
	}
}

func (b *Bot) startDeleteMessage() {
	for {
		conn, err := b.Queue.Get()
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(tgDeleteTube)
		job, err := conn.Reserve()
		if err != nil {
			b.logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}

		func() {
			var err error
			defer func() {
				if err != nil {
					b.logger.Errorf("%+v", err)
					if e := conn.Bury(job.ID, 0); e != nil {
						b.logger.Errorf("%+v", err)
					}
					time.Sleep(3 * time.Second)
				} else {
					if e := conn.Delete(job.ID); e != nil {
						b.logger.Errorf("%+v", err)
						time.Sleep(3 * time.Second)
					}
				}
			}()

			msg := &tgbotapi.Message{}
			err = json.Unmarshal(job.Body, msg)
			if err != nil {
				return
			}

			if msg.Chat == nil {
				err = fmt.Errorf("err msg with no chat: %+v", msg)
				return
			}
			delMsg := tgbotapi.DeleteMessageConfig{
				ChatID:    msg.Chat.ID,
				MessageID: msg.MessageID,
			}
			b.logger.Infof("del:[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))
			_, err = b.Client.DeleteMessage(delMsg)

		}()
		b.Queue.Release(conn, false)
	}
}

func (b *Bot) Start() {
	go b.startDeleteMessage()
	go b.startDownloadPixiv()

	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := b.Client.GetUpdatesChan(u)
		if err != nil {
			b.logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				b.logger.Infof(
					"recv:(%d)[%s]reaction:{%s}",
					update.CallbackQuery.Message.Chat.ID,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic", "pixiv":
					go onReaction(b, update.CallbackQuery)
				case "pixivIllust":
					if !b.isAuthedChat(update.CallbackQuery.Message.Chat) {
						b.logger.Warning("reaction from illegal chat, ignore")
						break
					}
					go onReactionSelf(b, update.CallbackQuery)
				case "search":
					go onReactionSearch(b, update.CallbackQuery)
				default:
				}
				continue
			} else {
				continue
			}

			if message.Chat.IsGroup() {
				b.logger.Infof(
					"recv:(%d)[%s:%s]{%s}",
					message.Chat.ID,
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				b.logger.Infof(
					"recv:(%d)[%s]{%s}",
					message.Chat.ID,
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(b, message)
				case "roll":
					go onRoll(b, message)
				case "comic":
					go onComic(b, message)
				case "pic":
					go onPic(b, message)
				case "pixiv":
					go onPixiv(b, message)
				case "search":
					go onSearch(b, message)
				default:
					b.logger.Infof("ignore unknown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				go checkRepeat(b, message)
				go checkPixiv(b, message)
				go checkSave(b, message)
			}
		}
		b.logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func (b *Bot) probate(_type, _id string) error {
	b.logger.Noticef("%s: %s", _type, _id)
	switch _type {
	case "comic":
		fileName := "nhentai.net@" + _id + ".epub"
		return os.Rename(
			filepath.Join(b.ComicPath, fileName),
			filepath.Join(b.ComicPath, "probation", fileName),
		)
	case "pic":
		return os.Rename(
			filepath.Join(b.TwitterImgPath, _id),
			filepath.Join(b.TwitterImgPath, "probation", _id),
		)
	case "pixiv":
		return os.Rename(
			filepath.Join(b.PixivPath, _id),
			filepath.Join(b.PixivPath, "probation", _id),
		)
	default:
		return fmt.Errorf("prohibit unkown type")
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}
<|endoftext|>


<commit_before>// +build linux,ppc64le

package qemu

import (
	"fmt"
	"strconv"

	"github.com/hyperhq/runv/hypervisor"
)

const (
	QEMU_SYSTEM_EXE    = "qemu-system-ppc64le"
	VM_MIN_MEMORY_SIZE = 256 // On ppc64le the minimum memory size of a VM is 256 MiB
)

func (qc *QemuContext) arguments(ctx *hypervisor.VmContext) []string {
	if ctx.Boot == nil {
		ctx.Boot = &hypervisor.BootConfig{
			CPU:    1,
			Memory: VM_MIN_MEMORY_SIZE,
			Kernel: hypervisor.DefaultKernel,
			Initrd: hypervisor.DefaultInitrd,
		}
	}
	boot := ctx.Boot
	qc.cpus = boot.CPU

	// Currently the default memory size is fixed to 128 MiB.
	// TODO: Check with PPC team for a better solution
	if boot.Memory < VM_MIN_MEMORY_SIZE {
		boot.Memory = VM_MIN_MEMORY_SIZE
	}

	var memParams, cpuParams string
	if boot.HotAddCpuMem {
		memParams = fmt.Sprintf("size=%d,slots=1,maxmem=%dM", boot.Memory, hypervisor.DefaultMaxMem) // TODO set maxmem to the total memory of the system
		cpuParams = fmt.Sprintf("cpus=%d,maxcpus=%d", boot.CPU, hypervisor.DefaultMaxCpus)           // TODO set it to the cpus of the system
	} else {
		memParams = strconv.Itoa(boot.Memory)
		cpuParams = strconv.Itoa(boot.CPU)
	}

	return []string{
		"-machine", "pseries,accel=kvm,usb=off", "-global", "kvm-pit.lost_tick_policy=discard", "-cpu", "host",
		"-kernel", boot.Kernel, "-initrd", boot.Initrd,
		"-machine", "pseries,accel=kvm,usb=off", "-cpu", "host",
		"-realtime", "mlock=off", "-no-user-config", "-nodefaults",
		"-rtc", "base=utc,driftfix=slew", "-no-reboot", "-display", "none", "-boot", "strict=on",
		"-m", memParams, "-smp", cpuParams,
		"-qmp", fmt.Sprintf("unix:%s,server,nowait", qc.qmpSockName), "-serial", fmt.Sprintf("unix:%s,server,nowait", ctx.ConsoleSockName),
		"-device", "virtio-serial-pci,id=virtio-serial0,bus=pci.0,addr=0x2", "-device", "virtio-scsi-pci,id=scsi0,bus=pci.0,addr=0x3",
		"-chardev", fmt.Sprintf("socket,id=charch0,path=%s,server,nowait", ctx.HyperSockName),
		"-device", "virtserialport,bus=virtio-serial0.0,nr=1,chardev=charch0,id=channel0,name=sh.hyper.channel.0",
		"-chardev", fmt.Sprintf("socket,id=charch1,path=%s,server,nowait", ctx.TtySockName),
		"-device", "virtserialport,bus=virtio-serial0.0,nr=2,chardev=charch1,id=channel1,name=sh.hyper.channel.1",
		"-fsdev", fmt.Sprintf("local,id=virtio9p,path=%s,security_model=none", ctx.ShareDir),
		"-device", fmt.Sprintf("virtio-9p-pci,fsdev=virtio9p,mount_tag=%s", hypervisor.ShareDirTag),
	}
}
<commit_msg>ppc64le: remove duplicate qemu argument<commit_after>// +build linux,ppc64le

package qemu

import (
	"fmt"
	"strconv"

	"github.com/hyperhq/runv/hypervisor"
)

const (
	QEMU_SYSTEM_EXE    = "qemu-system-ppc64le"
	VM_MIN_MEMORY_SIZE = 256 // On ppc64le the minimum memory size of a VM is 256 MiB
)

func (qc *QemuContext) arguments(ctx *hypervisor.VmContext) []string {
	if ctx.Boot == nil {
		ctx.Boot = &hypervisor.BootConfig{
			CPU:    1,
			Memory: VM_MIN_MEMORY_SIZE,
			Kernel: hypervisor.DefaultKernel,
			Initrd: hypervisor.DefaultInitrd,
		}
	}
	boot := ctx.Boot
	qc.cpus = boot.CPU

	// Currently the default memory size is fixed to 128 MiB.
	// TODO: Check with PPC team for a better solution
	if boot.Memory < VM_MIN_MEMORY_SIZE {
		boot.Memory = VM_MIN_MEMORY_SIZE
	}

	var memParams, cpuParams string
	if boot.HotAddCpuMem {
		memParams = fmt.Sprintf("size=%d,slots=1,maxmem=%dM", boot.Memory, hypervisor.DefaultMaxMem) // TODO set maxmem to the total memory of the system
		cpuParams = fmt.Sprintf("cpus=%d,maxcpus=%d", boot.CPU, hypervisor.DefaultMaxCpus)           // TODO set it to the cpus of the system
	} else {
		memParams = strconv.Itoa(boot.Memory)
		cpuParams = strconv.Itoa(boot.CPU)
	}

	return []string{
		"-machine", "pseries,accel=kvm,usb=off", "-global", "kvm-pit.lost_tick_policy=discard", "-cpu", "host",
		"-kernel", boot.Kernel, "-initrd", boot.Initrd,
		"-realtime", "mlock=off", "-no-user-config", "-nodefaults",
		"-rtc", "base=utc,driftfix=slew", "-no-reboot", "-display", "none", "-boot", "strict=on",
		"-m", memParams, "-smp", cpuParams,
		"-qmp", fmt.Sprintf("unix:%s,server,nowait", qc.qmpSockName), "-serial", fmt.Sprintf("unix:%s,server,nowait", ctx.ConsoleSockName),
		"-device", "virtio-serial-pci,id=virtio-serial0,bus=pci.0,addr=0x2", "-device", "virtio-scsi-pci,id=scsi0,bus=pci.0,addr=0x3",
		"-chardev", fmt.Sprintf("socket,id=charch0,path=%s,server,nowait", ctx.HyperSockName),
		"-device", "virtserialport,bus=virtio-serial0.0,nr=1,chardev=charch0,id=channel0,name=sh.hyper.channel.0",
		"-chardev", fmt.Sprintf("socket,id=charch1,path=%s,server,nowait", ctx.TtySockName),
		"-device", "virtserialport,bus=virtio-serial0.0,nr=2,chardev=charch1,id=channel1,name=sh.hyper.channel.1",
		"-fsdev", fmt.Sprintf("local,id=virtio9p,path=%s,security_model=none", ctx.ShareDir),
		"-device", fmt.Sprintf("virtio-9p-pci,fsdev=virtio9p,mount_tag=%s", hypervisor.ShareDirTag),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os/exec"
	"path/filepath"

	"github.com/pkg/errors"
)

// Plugin defines the PyPi plugin parameters
type Plugin struct {
	Repository    string
	Username      string
	Password      string
	SetupFile     string
	Distributions []string
	SkipBuild     bool
	DistDir       string
}

func (p Plugin) buildCommand() *exec.Cmd {
	// Set the default of distributions in here
	// as CLI package still has issues with string slice defaults
	distributions := []string{"sdist"}
	if len(p.Distributions) > 0 {
		distributions = p.Distributions
	}
	args := []string{p.SetupFile}
	for i := range distributions {
		args = append(args, distributions[i])
	}
	return exec.Command("python3", args...)
}

func (p Plugin) uploadCommand() *exec.Cmd {
	args := []string{}
	args = append(args, "upload")
	args = append(args, "--repository-url")
	args = append(args, p.Repository)
	args = append(args, "--username")
	args = append(args, p.Username)
	args = append(args, "--password")
	args = append(args, p.Password)
	args = append(args, filepath.Join(filepath.Dir(p.DistDir), "/*"))

	return exec.Command("twine", args...)
}

// Exec runs the plugin - doing the necessary setup.py modifications
func (p Plugin) Exec() error {
	if !p.SkipBuild {
		out, err := p.buildCommand().CombinedOutput()
		if err != nil {
			return errors.Wrap(err, string(out))
		}
		log.Printf("Output: %s", out)
	}

	out, err := p.uploadCommand().CombinedOutput()
	if err != nil {
		return errors.Wrap(err, string(out))
	}
	log.Printf("Output: %s", out)

	return nil
}
<commit_msg>simplify distdir filepath<commit_after>package main

import (
	"log"
	"os/exec"
	"path/filepath"

	"github.com/pkg/errors"
)

// Plugin defines the PyPi plugin parameters
type Plugin struct {
	Repository    string
	Username      string
	Password      string
	SetupFile     string
	Distributions []string
	SkipBuild     bool
	DistDir       string
}

func (p Plugin) buildCommand() *exec.Cmd {
	// Set the default of distributions in here
	// as CLI package still has issues with string slice defaults
	distributions := []string{"sdist"}
	if len(p.Distributions) > 0 {
		distributions = p.Distributions
	}
	args := []string{p.SetupFile}
	for i := range distributions {
		args = append(args, distributions[i])
	}
	return exec.Command("python3", args...)
}

func (p Plugin) uploadCommand() *exec.Cmd {
	args := []string{}
	args = append(args, "upload")
	args = append(args, "--repository-url")
	args = append(args, p.Repository)
	args = append(args, "--username")
	args = append(args, p.Username)
	args = append(args, "--password")
	args = append(args, p.Password)
	args = append(args, filepath.Join(p.DistDir, "/*"))

	return exec.Command("twine", args...)
}

// Exec runs the plugin - doing the necessary setup.py modifications
func (p Plugin) Exec() error {
	if !p.SkipBuild {
		out, err := p.buildCommand().CombinedOutput()
		if err != nil {
			return errors.Wrap(err, string(out))
		}
		log.Printf("Output: %s", out)
	}

	out, err := p.uploadCommand().CombinedOutput()
	if err != nil {
		return errors.Wrap(err, string(out))
	}
	log.Printf("Output: %s", out)

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// Plugin defines the Helm plugin parameters.
type Plugin struct {
	Debug        bool     `envconfig:"DEBUG"`
	ShowEnv      bool     `envconfig:"SHOW_ENV"`
	Wait         bool     `envconfig:"WAIT"`
	Recreate     bool     `envconfig:"RECREATE_PODS" default:"false"`
	WaitTimeout  uint32   `envconfig:"WAIT_TIMEOUT" default:"300"`
	Actions      []string `envconfig:"ACTIONS" required:"true"`
	AuthKey      string   `envconfig:"AUTH_KEY"`
	KeyPath      string   `envconfig:"KEY_PATH"`
	Zone         string   `envconfig:"ZONE"`
	Region       string   `envconfig:"REGION"`
	Cluster      string   `envconfig:"CLUSTER"`
	Project      string   `envconfig:"PROJECT"`
	Namespace    string   `envconfig:"NAMESPACE"`
	ChartRepo    string   `envconfig:"CHART_REPO"`
	Bucket       string   `envconfig:"BUCKET"`
	ChartPath    string   `envconfig:"CHART_PATH" required:"true"`
	ChartVersion string   `envconfig:"CHART_VERSION"`
	Release      string   `envconfig:"RELEASE"`
	Package      string   `envconfig:"PACKAGE"`
	Values       []string `envconfig:"VALUES"`
	ValueFiles   []string `envconfig:"VALUE_FILES"`
}

const (
	gcloudBin  = "gcloud"
	gsutilBin  = "gsutil"
	kubectlBin = "kubectl"
	helmBin    = "helm"

	lintPkg   = "lint"
	createPkg = "create"
	pushPkg   = "push"
	pullPkg   = "pull"
	deployPkg = "deploy"
	testPkg   = "test"

	updateWaitTime = 10 * time.Second
	updateRetries  = 10
)

// Exec executes the plugin step.
func (p Plugin) Exec() error {
	// only setup project when needed args are provided
	if p.Project != "" && p.Cluster != "" && (p.Zone != "" || p.Region != "") {
		if err := setupProject(p.Project, p.Cluster, p.Zone, p.Region, p.Debug); err != nil {
			return err
		}
	}

	for _, a := range p.Actions {
		switch a {
		case lintPkg:
			if err := p.lintPackage(); err != nil {
				return err
			}
		case createPkg:
			if err := p.createPackage(); err != nil {
				return err
			}
		case pushPkg:
			if err := p.pushPackage(); err != nil {
				return err
			}
		case pullPkg:
			if err := p.pullPackage(); err != nil {
				return err
			}
		case deployPkg:
			if err := p.deployPackage(); err != nil {
				return err
			}
		case testPkg:
			if err := p.testPackage(); err != nil {
				return err
			}
		default:
			return errors.New("unknown action")
		}
	}

	return nil
}

// setupProject setups gcloud project.
func setupProject(project, cluster, zone, region string, debug bool) error {
	// project configuration
	cmd := exec.Command(gcloudBin, "config", "set", "project", project)
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not the configure the project with glcoud: %v", err)
	}

	// cluster configuration
	cmd = exec.Command(gcloudBin, "container", "clusters", "get-credentials", cluster, "--zone", zone)
	if region != "" {
		// override zone when region is set
		cmd = exec.Command(gcloudBin, "container", "clusters", "get-credentials", cluster, "--region", region)
	}

	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not configure the cluster with glcoud: %v", err)
	}

	return nil
}

// setupAuth configures gcloud to use the given authFile
func setupAuth(authFile string, debug bool) error {
	if err := os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", authFile); err != nil {
		return fmt.Errorf("could not set GOOGLE_APPLICATION_CREDENTIALS env variable: %v", err)
	}

	// authorization
	cmd := exec.Command(gcloudBin, "auth", "activate-service-account", fmt.Sprintf("--key-file=%s", authFile))
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not authorize with glcoud: %v", err)
	}
	return nil
}

// createPackage creates Helm package for Kubernetes.
// helm package --version $PLUGIN_CHART_VERSION $PLUGIN_CHART_PATH
func (p Plugin) createPackage() error {
	return run(exec.Command(helmBin, "package", "--version", p.ChartVersion, p.ChartPath), p.Debug)
}

// cpPackage copies a file from SOURCE to DEST
// gsutil cp SOURCE DEST
func (p Plugin) cpPackage(source string, dest string) error {
	return run(exec.Command(gsutilBin, "cp", source, dest), p.Debug)
}

// cpPackage pulls helm chart from Google Storage to local
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pullPackage() error {
	return p.cpPackage(
		fmt.Sprintf("gs://%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	)
}

// pushPackage pushes Helm package to the Google Storage.
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pushPackage() error {
	return p.cpPackage(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("gs://%s", p.Bucket),
	)
}

// helm lint $CHARTPATH -i
func (p Plugin) lintPackage() error {
	return run(exec.Command(helmBin, "lint", p.ChartPath), p.Debug)
}

// helm upgrade $PACKAGE $PACKAGE-$PLUGIN_CHART_VERSION.tgz -i
func (p Plugin) deployPackage() error {
	args := []string{
		helmBin,
		"upgrade",
		p.Release,
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	}
	if len(p.ValueFiles) > 0 {
		for _, f := range p.ValueFiles {
			args = append(args, "-f", f)
		}
	}
	if len(p.Values) > 0 {
		args = append(args, "--set", strings.Join(p.Values, ","))
	}
	if p.Recreate {
		args = append(args, "--recreate-pods")
	}
	args = append(args, "--install")
	args = append(args, "--namespace", p.Namespace)

	if p.Wait {
		args = append(args, "--wait", "--timeout", strconv.Itoa(int(p.WaitTimeout)))
	}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

// helm test $PACKAGE
func (p Plugin) testPackage() error {
	args := []string{helmBin, "test", p.Release, "--cleanup", "--timeout", strconv.Itoa(int(p.WaitTimeout))}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

type semVer struct {
	Version string `json:"version"`
}
type helmVersions struct {
	Client semVer `json:"client"`
	Server semVer `json:"server"`
}

func (p Plugin) movePkg() error {
	if err := os.Mkdir(p.Bucket, os.ModeDir); err != nil {
		return err
	}
	return cp(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
	)
}

// cp copies file
func cp(src, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}
	// no need to check errors on read only file, we already got everything
	// we need from the filesystem, so nothing can go wrong now.
	defer s.Close()
	d, err := os.Create(dst)
	if err != nil {
		return err
	}
	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}
	return d.Close()
}

// scanNamed maps named regex groups to a golang map
func scanNamed(str string, rg *regexp.Regexp) (map[string]string, error) {
	result := make(map[string]string)
	for _, match := range rg.FindAllStringSubmatch(str, -1) {
		for i, name := range rg.SubexpNames() {
			if i != 0 && match[i] != "" && name != "" {
				result[name] = match[i]
			}
		}
	}

	if len(result) == 0 {
		return nil, errors.New("emtpy resultset")
	}

	return result, nil
}

func run(cmd *exec.Cmd, debug bool) error {
	if debug {
		log.Printf("running: %s", strings.Join(cmd.Args, " "))
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmd.Run()
}
<commit_msg>fix missing unit<commit_after>package main

import (
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// Plugin defines the Helm plugin parameters.
type Plugin struct {
	Debug        bool     `envconfig:"DEBUG"`
	ShowEnv      bool     `envconfig:"SHOW_ENV"`
	Wait         bool     `envconfig:"WAIT"`
	Recreate     bool     `envconfig:"RECREATE_PODS" default:"false"`
	WaitTimeout  uint32   `envconfig:"WAIT_TIMEOUT" default:"300s"`
	Actions      []string `envconfig:"ACTIONS" required:"true"`
	AuthKey      string   `envconfig:"AUTH_KEY"`
	KeyPath      string   `envconfig:"KEY_PATH"`
	Zone         string   `envconfig:"ZONE"`
	Region       string   `envconfig:"REGION"`
	Cluster      string   `envconfig:"CLUSTER"`
	Project      string   `envconfig:"PROJECT"`
	Namespace    string   `envconfig:"NAMESPACE"`
	ChartRepo    string   `envconfig:"CHART_REPO"`
	Bucket       string   `envconfig:"BUCKET"`
	ChartPath    string   `envconfig:"CHART_PATH" required:"true"`
	ChartVersion string   `envconfig:"CHART_VERSION"`
	Release      string   `envconfig:"RELEASE"`
	Package      string   `envconfig:"PACKAGE"`
	Values       []string `envconfig:"VALUES"`
	ValueFiles   []string `envconfig:"VALUE_FILES"`
}

const (
	gcloudBin  = "gcloud"
	gsutilBin  = "gsutil"
	kubectlBin = "kubectl"
	helmBin    = "helm"

	lintPkg   = "lint"
	createPkg = "create"
	pushPkg   = "push"
	pullPkg   = "pull"
	deployPkg = "deploy"
	testPkg   = "test"

	updateWaitTime = 10 * time.Second
	updateRetries  = 10
)

// Exec executes the plugin step.
func (p Plugin) Exec() error {
	// only setup project when needed args are provided
	if p.Project != "" && p.Cluster != "" && (p.Zone != "" || p.Region != "") {
		if err := setupProject(p.Project, p.Cluster, p.Zone, p.Region, p.Debug); err != nil {
			return err
		}
	}

	for _, a := range p.Actions {
		switch a {
		case lintPkg:
			if err := p.lintPackage(); err != nil {
				return err
			}
		case createPkg:
			if err := p.createPackage(); err != nil {
				return err
			}
		case pushPkg:
			if err := p.pushPackage(); err != nil {
				return err
			}
		case pullPkg:
			if err := p.pullPackage(); err != nil {
				return err
			}
		case deployPkg:
			if err := p.deployPackage(); err != nil {
				return err
			}
		case testPkg:
			if err := p.testPackage(); err != nil {
				return err
			}
		default:
			return errors.New("unknown action")
		}
	}

	return nil
}

// setupProject setups gcloud project.
func setupProject(project, cluster, zone, region string, debug bool) error {
	// project configuration
	cmd := exec.Command(gcloudBin, "config", "set", "project", project)
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not the configure the project with glcoud: %v", err)
	}

	// cluster configuration
	cmd = exec.Command(gcloudBin, "container", "clusters", "get-credentials", cluster, "--zone", zone)
	if region != "" {
		// override zone when region is set
		cmd = exec.Command(gcloudBin, "container", "clusters", "get-credentials", cluster, "--region", region)
	}

	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not configure the cluster with glcoud: %v", err)
	}

	return nil
}

// setupAuth configures gcloud to use the given authFile
func setupAuth(authFile string, debug bool) error {
	if err := os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", authFile); err != nil {
		return fmt.Errorf("could not set GOOGLE_APPLICATION_CREDENTIALS env variable: %v", err)
	}

	// authorization
	cmd := exec.Command(gcloudBin, "auth", "activate-service-account", fmt.Sprintf("--key-file=%s", authFile))
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not authorize with glcoud: %v", err)
	}
	return nil
}

// createPackage creates Helm package for Kubernetes.
// helm package --version $PLUGIN_CHART_VERSION $PLUGIN_CHART_PATH
func (p Plugin) createPackage() error {
	return run(exec.Command(helmBin, "package", "--version", p.ChartVersion, p.ChartPath), p.Debug)
}

// cpPackage copies a file from SOURCE to DEST
// gsutil cp SOURCE DEST
func (p Plugin) cpPackage(source string, dest string) error {
	return run(exec.Command(gsutilBin, "cp", source, dest), p.Debug)
}

// cpPackage pulls helm chart from Google Storage to local
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pullPackage() error {
	return p.cpPackage(
		fmt.Sprintf("gs://%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	)
}

// pushPackage pushes Helm package to the Google Storage.
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pushPackage() error {
	return p.cpPackage(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("gs://%s", p.Bucket),
	)
}

// helm lint $CHARTPATH -i
func (p Plugin) lintPackage() error {
	return run(exec.Command(helmBin, "lint", p.ChartPath), p.Debug)
}

// helm upgrade $PACKAGE $PACKAGE-$PLUGIN_CHART_VERSION.tgz -i
func (p Plugin) deployPackage() error {
	args := []string{
		helmBin,
		"upgrade",
		p.Release,
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	}
	if len(p.ValueFiles) > 0 {
		for _, f := range p.ValueFiles {
			args = append(args, "-f", f)
		}
	}
	if len(p.Values) > 0 {
		args = append(args, "--set", strings.Join(p.Values, ","))
	}
	if p.Recreate {
		args = append(args, "--recreate-pods")
	}
	args = append(args, "--install")
	args = append(args, "--namespace", p.Namespace)

	if p.Wait {
		args = append(args, "--wait", "--timeout", strconv.Itoa(int(p.WaitTimeout)))
	}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

// helm test $PACKAGE
func (p Plugin) testPackage() error {
	args := []string{helmBin, "test", p.Release, "--cleanup", "--timeout", strconv.Itoa(int(p.WaitTimeout))}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

type semVer struct {
	Version string `json:"version"`
}
type helmVersions struct {
	Client semVer `json:"client"`
	Server semVer `json:"server"`
}

func (p Plugin) movePkg() error {
	if err := os.Mkdir(p.Bucket, os.ModeDir); err != nil {
		return err
	}
	return cp(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
	)
}

// cp copies file
func cp(src, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}
	// no need to check errors on read only file, we already got everything
	// we need from the filesystem, so nothing can go wrong now.
	defer s.Close()
	d, err := os.Create(dst)
	if err != nil {
		return err
	}
	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}
	return d.Close()
}

// scanNamed maps named regex groups to a golang map
func scanNamed(str string, rg *regexp.Regexp) (map[string]string, error) {
	result := make(map[string]string)
	for _, match := range rg.FindAllStringSubmatch(str, -1) {
		for i, name := range rg.SubexpNames() {
			if i != 0 && match[i] != "" && name != "" {
				result[name] = match[i]
			}
		}
	}

	if len(result) == 0 {
		return nil, errors.New("emtpy resultset")
	}

	return result, nil
}

func run(cmd *exec.Cmd, debug bool) error {
	if debug {
		log.Printf("running: %s", strings.Join(cmd.Args, " "))
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmd.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/alecthomas/template"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

var HELM_BIN = "/bin/helm"
var KUBECONFIG = "/root/.kube/kubeconfig"
var CONFIG = "/root/.kube/config"

type (
	// Config maps the params we need to run Helm
	Config struct {
		APIServer     string   `json:"api_server"`
		Token         string   `json:"token"`
		HelmCommand   []string `json:"helm_command"`
		SkipTLSVerify bool     `json:"tls_skip_verify"`
		Namespace     string   `json:"namespace"`
		Release       string   `json:"release"`
		Chart         string   `json:"chart"`
		Values        string   `json:"values"`
		Debug         bool     `json:"debug"`
		DryRun        bool     `json:"dry_run"`
		Secrets       []string `json:"secrets"`
		Prefix        string   `json:"prefix"`
	}
	// Plugin default
	Plugin struct {
		Config Config
	}
)

func setHelmHelp(p *Plugin) {
	p.Config.HelmCommand = []string{""}
}
func setDeleteEventCommand(p *Plugin) {
	upgrade := make([]string, 2)
	upgrade[0] = "delete"
	upgrade[1] = p.Config.Release

	p.Config.HelmCommand = upgrade
}

func setPushEventCommand(p *Plugin) {
	upgrade := make([]string, 2)
	upgrade[0] = "upgrade"
	upgrade[1] = "--install"
	if p.Config.Release != "" {
		upgrade = append(upgrade, p.Config.Release)
	}
	upgrade = append(upgrade, p.Config.Chart)
	if p.Config.Values != "" {
		upgrade = append(upgrade, "--set")
		upgrade = append(upgrade, p.Config.Values)
	}
	if p.Config.DryRun {
		upgrade = append(upgrade, "--dry-run")
	}
	if p.Config.Debug {
		upgrade = append(upgrade, "--debug")
	}
	p.Config.HelmCommand = upgrade

}

func setHelmCommand(p *Plugin) {
	buildEvent := os.Getenv("DRONE_BUILD_EVENT")
	switch buildEvent {
	case "push":
		setPushEventCommand(p)
	case "delete":
		setDeleteEventCommand(p)
	default:
		setHelmHelp(p)
	}

}

// Exec default method
func (p *Plugin) Exec() error {
	resolveSecrets(p)
	if p.Config.APIServer == "" {
		return fmt.Errorf("Error: API Server is needed to deploy.")
	}
	if p.Config.Token == "" {
		return fmt.Errorf("Error: Token is needed to deploy.")
	}
	initialiseKubeconfig(&p.Config, KUBECONFIG, CONFIG)

	if p.Config.Debug {
		p.debug()
	}

	init := make([]string, 1)
	init[0] = "init"
	err := runCommand(init)
	if err != nil {
		return fmt.Errorf("Error running helm comand: " + strings.Join(init[:], " "))
	}

	setHelmCommand(p)
	if p.Config.Debug {
		log.Println("helm comand: " + strings.Join(p.Config.HelmCommand[:], " "))
	}
	err = runCommand(p.Config.HelmCommand)
	if err != nil {
		return fmt.Errorf("Error running helm comand: " + strings.Join(p.Config.HelmCommand[:], " "))
	}
	return nil
}

func initialiseKubeconfig(params *Config, source string, target string) error {
	t, _ := template.ParseFiles(source)
	f, err := os.Create(target)
	err = t.Execute(f, params)
	f.Close()
	return err
}

func runCommand(params []string) error {
	cmd := new(exec.Cmd)
	cmd = exec.Command(HELM_BIN, params...)

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err := cmd.Run()
	return err
}

func resolveSecrets(p *Plugin) {
	p.Config.Values = resolveEnvVar(p.Config.Values, p.Config.Prefix)
	p.Config.APIServer = resolveEnvVar("${API_SERVER}", p.Config.Prefix)
	p.Config.Token = resolveEnvVar("${TOKEN}", p.Config.Prefix)
}

func getEnvVars(envvars string) [][]string {
	re := regexp.MustCompile(`\$(\{?(\w+)\}?)\.?`)
	extracted := re.FindAllStringSubmatch(envvars, -1)
	return extracted
}

func resolveEnvVar(key string, prefix string) string {
	envvars := getEnvVars(key)
	return replaceEnvvars(envvars, prefix, key)
}

func replaceEnvvars(envvars [][]string, prefix string, s string) string {
	fmt.Println(envvars)
	fmt.Println(prefix)
	fmt.Println(s)
	for _, envvar := range envvars {
		envvarName := envvar[0]
		envvarKey := envvar[2]
		envval := os.Getenv(envvarKey)
		if prefix != "" {
			envval = os.Getenv(prefix + "_" + envvarKey)
		}

		if strings.Contains(s, envvarKey) {
			s = strings.Replace(s, envvarName, envval, -1)
		}
	}
	fmt.Println(s)
	return s
}

func (p *Plugin) debug() {
	fmt.Println(p)
	// debug env vars
	for _, e := range os.Environ() {
		fmt.Println(e)
	}
	// debug plugin obj
	fmt.Printf("Api server: %s \n", p.Config.APIServer)
	fmt.Printf("Values: %s \n", p.Config.Values)
	fmt.Printf("Values: %s \n", p.Config.Secrets)

	kubeconfig, err := ioutil.ReadFile(KUBECONFIG)
	if err == nil {
		fmt.Println(string(kubeconfig))
	}
	config, err := ioutil.ReadFile(CONFIG)
	if err == nil {
		fmt.Println(string(config))
	}

}
<commit_msg>envval envvarkey<commit_after>package main

import (
	"fmt"
	"github.com/alecthomas/template"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

var HELM_BIN = "/bin/helm"
var KUBECONFIG = "/root/.kube/kubeconfig"
var CONFIG = "/root/.kube/config"

type (
	// Config maps the params we need to run Helm
	Config struct {
		APIServer     string   `json:"api_server"`
		Token         string   `json:"token"`
		HelmCommand   []string `json:"helm_command"`
		SkipTLSVerify bool     `json:"tls_skip_verify"`
		Namespace     string   `json:"namespace"`
		Release       string   `json:"release"`
		Chart         string   `json:"chart"`
		Values        string   `json:"values"`
		Debug         bool     `json:"debug"`
		DryRun        bool     `json:"dry_run"`
		Secrets       []string `json:"secrets"`
		Prefix        string   `json:"prefix"`
	}
	// Plugin default
	Plugin struct {
		Config Config
	}
)

func setHelmHelp(p *Plugin) {
	p.Config.HelmCommand = []string{""}
}
func setDeleteEventCommand(p *Plugin) {
	upgrade := make([]string, 2)
	upgrade[0] = "delete"
	upgrade[1] = p.Config.Release

	p.Config.HelmCommand = upgrade
}

func setPushEventCommand(p *Plugin) {
	upgrade := make([]string, 2)
	upgrade[0] = "upgrade"
	upgrade[1] = "--install"
	if p.Config.Release != "" {
		upgrade = append(upgrade, p.Config.Release)
	}
	upgrade = append(upgrade, p.Config.Chart)
	if p.Config.Values != "" {
		upgrade = append(upgrade, "--set")
		upgrade = append(upgrade, p.Config.Values)
	}
	if p.Config.DryRun {
		upgrade = append(upgrade, "--dry-run")
	}
	if p.Config.Debug {
		upgrade = append(upgrade, "--debug")
	}
	p.Config.HelmCommand = upgrade

}

func setHelmCommand(p *Plugin) {
	buildEvent := os.Getenv("DRONE_BUILD_EVENT")
	switch buildEvent {
	case "push":
		setPushEventCommand(p)
	case "delete":
		setDeleteEventCommand(p)
	default:
		setHelmHelp(p)
	}

}

// Exec default method
func (p *Plugin) Exec() error {
	resolveSecrets(p)
	if p.Config.APIServer == "" {
		return fmt.Errorf("Error: API Server is needed to deploy.")
	}
	if p.Config.Token == "" {
		return fmt.Errorf("Error: Token is needed to deploy.")
	}
	initialiseKubeconfig(&p.Config, KUBECONFIG, CONFIG)

	if p.Config.Debug {
		p.debug()
	}

	init := make([]string, 1)
	init[0] = "init"
	err := runCommand(init)
	if err != nil {
		return fmt.Errorf("Error running helm comand: " + strings.Join(init[:], " "))
	}

	setHelmCommand(p)
	if p.Config.Debug {
		log.Println("helm comand: " + strings.Join(p.Config.HelmCommand[:], " "))
	}
	err = runCommand(p.Config.HelmCommand)
	if err != nil {
		return fmt.Errorf("Error running helm comand: " + strings.Join(p.Config.HelmCommand[:], " "))
	}
	return nil
}

func initialiseKubeconfig(params *Config, source string, target string) error {
	t, _ := template.ParseFiles(source)
	f, err := os.Create(target)
	err = t.Execute(f, params)
	f.Close()
	return err
}

func runCommand(params []string) error {
	cmd := new(exec.Cmd)
	cmd = exec.Command(HELM_BIN, params...)

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err := cmd.Run()
	return err
}

func resolveSecrets(p *Plugin) {
	p.Config.Values = resolveEnvVar(p.Config.Values, p.Config.Prefix)
	p.Config.APIServer = resolveEnvVar("${API_SERVER}", p.Config.Prefix)
	p.Config.Token = resolveEnvVar("${TOKEN}", p.Config.Prefix)
}

func getEnvVars(envvars string) [][]string {
	re := regexp.MustCompile(`\$(\{?(\w+)\}?)\.?`)
	extracted := re.FindAllStringSubmatch(envvars, -1)
	return extracted
}

func resolveEnvVar(key string, prefix string) string {
	envvars := getEnvVars(key)
	return replaceEnvvars(envvars, prefix, key)
}

func replaceEnvvars(envvars [][]string, prefix string, s string) string {
	fmt.Println(envvars)
	fmt.Println(prefix)
	fmt.Println(s)
	for _, envvar := range envvars {
		envvarName := envvar[0]
		envvarKey := envvar[2]
		envval := os.Getenv(envvarKey)
		if prefix != "" {
			envval = os.Getenv(prefix + "_" + envvarKey)
		}
		fmt.Println(envval)
		fmt.Println(envvarKey)
		if strings.Contains(s, envvarKey) {
			s = strings.Replace(s, envvarName, envval, -1)
		}
	}
	fmt.Println(s)
	return s
}

func (p *Plugin) debug() {
	fmt.Println(p)
	// debug env vars
	for _, e := range os.Environ() {
		fmt.Println(e)
	}
	// debug plugin obj
	fmt.Printf("Api server: %s \n", p.Config.APIServer)
	fmt.Printf("Values: %s \n", p.Config.Values)
	fmt.Printf("Values: %s \n", p.Config.Secrets)

	kubeconfig, err := ioutil.ReadFile(KUBECONFIG)
	if err == nil {
		fmt.Println(string(kubeconfig))
	}
	config, err := ioutil.ReadFile(CONFIG)
	if err == nil {
		fmt.Println(string(config))
	}

}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/NoahShen/gotunnelme/src/gotunnelme"
	"github.com/appleboy/com/random"
	"github.com/appleboy/drone-facebook/template"
	"github.com/fatih/color"
	"github.com/line/line-bot-sdk-go/linebot"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"golang.org/x/crypto/acme/autocert"
)

const defaultPreviewImageURL = "https://cdn4.iconfinder.com/data/icons/miu/24/device-camera-recorder-video-glyph-256.png"

type (
	// Repo information.
	Repo struct {
		Owner string
		Name  string
	}

	// Build information.
	Build struct {
		Tag      string
		Event    string
		Number   int
		Commit   string
		Branch   string
		Author   string
		Email    string
		Message  string
		Status   string
		Link     string
		Started  float64
		Finished float64
	}

	// Config for the plugin.
	Config struct {
		ChannelToken  string
		ChannelSecret string
		To            []string
		Delimiter     string
		Message       []string
		Image         []string
		Video         []string
		Audio         []string
		Sticker       []string
		Location      []string
		MatchEmail    bool
		Port          int
		Tunnel        bool
		Debug         bool
		Domain        string
		AutoTLS       bool
		Host          []string
	}

	// Plugin values.
	Plugin struct {
		Repo   Repo
		Build  Build
		Config Config
	}

	// Audio format
	Audio struct {
		URL      string
		Duration int
	}

	// Location format
	Location struct {
		Title     string
		Address   string
		Latitude  float64
		Longitude float64
	}
)

var (
	// ReceiveCount is receive notification count
	ReceiveCount int64
	// SendCount is send notification count
	SendCount int64
)

func init() {
	// Support metrics
	m := NewMetrics()
	prometheus.MustRegister(m)
}

func trimElement(keys []string) []string {
	var newKeys []string

	for _, value := range keys {
		value = strings.Trim(value, " ")
		if len(value) == 0 {
			continue
		}
		newKeys = append(newKeys, value)
	}

	return newKeys
}

func convertImage(value, delimiter string) []string {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		values = append(values, values[0])
	}

	return values
}

func convertVideo(value, delimiter string) []string {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		values = append(values, defaultPreviewImageURL)
	}

	return values
}

func convertAudio(value, delimiter string) (Audio, bool) {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		return Audio{}, true
	}

	duration, err := strconv.Atoi(values[1])

	if err != nil {
		log.Println(err.Error())
		return Audio{}, true
	}

	return Audio{
		URL:      values[0],
		Duration: duration,
	}, false
}

func convertSticker(value, delimiter string) ([]string, bool) {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		return []string{}, true
	}

	return values, false
}

func convertLocation(value, delimiter string) (Location, bool) {
	var latitude, longitude float64
	var err error
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 4 {
		return Location{}, true
	}

	latitude, err = strconv.ParseFloat(values[2], 64)

	if err != nil {
		log.Println(err.Error())
		return Location{}, true
	}

	longitude, err = strconv.ParseFloat(values[3], 64)

	if err != nil {
		log.Println(err.Error())
		return Location{}, true
	}

	return Location{
		Title:     values[0],
		Address:   values[1],
		Latitude:  latitude,
		Longitude: longitude,
	}, false
}

func parseTo(to []string, authorEmail string, matchEmail bool, delimiter string) []string {
	var emails []string
	var ids []string
	attachEmail := true

	for _, value := range trimElement(to) {
		idArray := trimElement(strings.Split(value, delimiter))

		// check match author email
		if len(idArray) > 1 {
			if email := idArray[1]; email != authorEmail {
				continue
			}

			emails = append(emails, idArray[0])
			attachEmail = false
			continue
		}

		ids = append(ids, idArray[0])
	}

	if matchEmail == true && attachEmail == false {
		return emails
	}

	for _, value := range emails {
		ids = append(ids, value)
	}

	return ids
}

// Bot is new Line Bot clien.
func (p Plugin) Bot() (*linebot.Client, error) {
	if len(p.Config.ChannelToken) == 0 || len(p.Config.ChannelSecret) == 0 {
		log.Println("missing line bot config")

		return nil, errors.New("missing line bot config")
	}

	return linebot.New(p.Config.ChannelSecret, p.Config.ChannelToken)
}

func (p Plugin) getTunnelDomain() (string, error) {
	var domain string
	if p.Config.Domain != "" {
		if len(p.Config.Domain) < 4 || len(p.Config.Domain) > 63 {
			return "", errors.New("tunnel host name must be lowercase and between 4 and 63 alphanumeric characters")
		}
		domain = p.Config.Domain
	} else {
		domain = strings.ToLower(random.String(10))
	}

	return domain, nil
}

// Handler is http handler.
func (p Plugin) Handler(bot *linebot.Client) *http.ServeMux {
	mux := http.NewServeMux()

	// Setup HTTP Server for receiving requests from LINE platform
	mux.HandleFunc("/callback", func(w http.ResponseWriter, req *http.Request) {
		events, err := bot.ParseRequest(req)
		if err != nil {
			if err == linebot.ErrInvalidSignature {
				w.WriteHeader(400)
			} else {
				w.WriteHeader(500)
			}
			return
		}
		for _, event := range events {
			if event.Type == linebot.EventTypeMessage {
				switch message := event.Message.(type) {
				case *linebot.TextMessage:
					log.Printf("User ID is %v\n", event.Source.UserID)
					ReceiveCount++
					if message.Text == "test" {
						SendCount++
						if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage("count + 1")).Do(); err != nil {
							log.Print(err)
						}
					}
					if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(message.Text)).Do(); err != nil {
						log.Print(err)
					}
				}
			}
		}
	})

	mux.HandleFunc("/metrics", func(w http.ResponseWriter, req *http.Request) {
		promhttp.Handler().ServeHTTP(w, req)
	})

	// Setup HTTP Server for receiving requests from LINE platform
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		fmt.Fprintln(w, "Welcome to Line webhook page.")
	})

	return mux
}

// Webhook support line callback service.
func (p Plugin) Webhook() error {
	readyToListen := false
	bot, err := p.Bot()

	if err != nil {
		return err
	}
	mux := p.Handler(bot)

	if p.Config.Tunnel {
		if p.Config.Debug {
			gotunnelme.Debug = true
		}

		domain, err := p.getTunnelDomain()
		if err != nil {
			panic(err)
		}

		tunnel := gotunnelme.NewTunnel()
		url, err := tunnel.GetUrl(domain)
		if err != nil {
			panic("Could not get localtunnel.me URL. " + err.Error())
		}
		go func() {
			for !readyToListen {
				time.Sleep(1 * time.Second)
			}
			c := color.New(color.FgYellow)
			c.Println("Tunnel URL:", url)
			err := tunnel.CreateTunnel(p.Config.Port)
			if err != nil {
				panic("Could not create tunnel. " + err.Error())
			}
		}()
	}

	readyToListen = true
	if p.Config.Port != 443 && !p.Config.AutoTLS {
		log.Println("Line Webhook Server Listin on " + strconv.Itoa(p.Config.Port) + " port")
		if err := http.ListenAndServe(":"+strconv.Itoa(p.Config.Port), mux); err != nil {
			log.Fatal(err)
		}
	}

	if !p.Config.AutoTLS && len(p.Config.Host) != 0 {
		log.Println("Line Webhook Server Listin on 443 port, hostname: " + strings.Join(p.Config.Host, ", "))
		return http.Serve(autocert.NewListener(p.Config.Host...), mux)
	}

	return nil
}

// Exec executes the plugin.
func (p Plugin) Exec() error {

	bot, err := p.Bot()

	if err != nil {
		return err
	}

	if len(p.Config.To) == 0 {
		log.Println("missing line user config")

		return errors.New("missing line user config")
	}

	var message []string
	if len(p.Config.Message) > 0 {
		message = p.Config.Message
	} else {
		message = p.Message(p.Repo, p.Build)
	}

	// Initial messages array.
	var messages []linebot.Message

	for _, value := range trimElement(message) {
		txt, err := template.RenderTrim(value, p)
		if err != nil {
			return err
		}

		messages = append(messages, linebot.NewTextMessage(txt))
	}

	// Add image message
	for _, value := range trimElement(p.Config.Image) {
		values := convertImage(value, p.Config.Delimiter)

		messages = append(messages, linebot.NewImageMessage(values[0], values[1]))
	}

	// Add image message.
	for _, value := range trimElement(p.Config.Video) {
		values := convertVideo(value, p.Config.Delimiter)

		messages = append(messages, linebot.NewVideoMessage(values[0], values[1]))
	}

	// Add Audio message.
	for _, value := range trimElement(p.Config.Audio) {
		audio, empty := convertAudio(value, p.Config.Delimiter)

		if empty == true {
			continue
		}

		messages = append(messages, linebot.NewAudioMessage(audio.URL, audio.Duration))
	}

	// Add Sticker message.
	for _, value := range trimElement(p.Config.Sticker) {
		sticker, empty := convertSticker(value, p.Config.Delimiter)

		if empty == true {
			continue
		}

		messages = append(messages, linebot.NewStickerMessage(sticker[0], sticker[1]))
	}

	// check Location array.
	for _, value := range trimElement(p.Config.Location) {
		location, empty := convertLocation(value, p.Config.Delimiter)

		if empty == true {
			continue
		}

		messages = append(messages, linebot.NewLocationMessage(location.Title, location.Address, location.Latitude, location.Longitude))
	}

	ids := parseTo(p.Config.To, p.Build.Email, p.Config.MatchEmail, p.Config.Delimiter)

	// Send messages to multiple users at any time.
	if _, err := bot.Multicast(ids, messages...).Do(); err != nil {
		log.Println(err.Error())
	}

	return nil
}

// Message is line default message.
func (p Plugin) Message(repo Repo, build Build) []string {
	return []string{fmt.Sprintf("[%s] <%s> (%s)『%s』by %s",
		build.Status,
		build.Link,
		build.Branch,
		build.Message,
		build.Author,
	)}
}
<commit_msg>fix: boolean flag.<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/NoahShen/gotunnelme/src/gotunnelme"
	"github.com/appleboy/com/random"
	"github.com/appleboy/drone-facebook/template"
	"github.com/fatih/color"
	"github.com/line/line-bot-sdk-go/linebot"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"golang.org/x/crypto/acme/autocert"
)

const defaultPreviewImageURL = "https://cdn4.iconfinder.com/data/icons/miu/24/device-camera-recorder-video-glyph-256.png"

type (
	// Repo information.
	Repo struct {
		Owner string
		Name  string
	}

	// Build information.
	Build struct {
		Tag      string
		Event    string
		Number   int
		Commit   string
		Branch   string
		Author   string
		Email    string
		Message  string
		Status   string
		Link     string
		Started  float64
		Finished float64
	}

	// Config for the plugin.
	Config struct {
		ChannelToken  string
		ChannelSecret string
		To            []string
		Delimiter     string
		Message       []string
		Image         []string
		Video         []string
		Audio         []string
		Sticker       []string
		Location      []string
		MatchEmail    bool
		Port          int
		Tunnel        bool
		Debug         bool
		Domain        string
		AutoTLS       bool
		Host          []string
	}

	// Plugin values.
	Plugin struct {
		Repo   Repo
		Build  Build
		Config Config
	}

	// Audio format
	Audio struct {
		URL      string
		Duration int
	}

	// Location format
	Location struct {
		Title     string
		Address   string
		Latitude  float64
		Longitude float64
	}
)

var (
	// ReceiveCount is receive notification count
	ReceiveCount int64
	// SendCount is send notification count
	SendCount int64
)

func init() {
	// Support metrics
	m := NewMetrics()
	prometheus.MustRegister(m)
}

func trimElement(keys []string) []string {
	var newKeys []string

	for _, value := range keys {
		value = strings.Trim(value, " ")
		if len(value) == 0 {
			continue
		}
		newKeys = append(newKeys, value)
	}

	return newKeys
}

func convertImage(value, delimiter string) []string {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		values = append(values, values[0])
	}

	return values
}

func convertVideo(value, delimiter string) []string {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		values = append(values, defaultPreviewImageURL)
	}

	return values
}

func convertAudio(value, delimiter string) (Audio, bool) {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		return Audio{}, true
	}

	duration, err := strconv.Atoi(values[1])

	if err != nil {
		log.Println(err.Error())
		return Audio{}, true
	}

	return Audio{
		URL:      values[0],
		Duration: duration,
	}, false
}

func convertSticker(value, delimiter string) ([]string, bool) {
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 2 {
		return []string{}, true
	}

	return values, false
}

func convertLocation(value, delimiter string) (Location, bool) {
	var latitude, longitude float64
	var err error
	values := trimElement(strings.Split(value, delimiter))

	if len(values) < 4 {
		return Location{}, true
	}

	latitude, err = strconv.ParseFloat(values[2], 64)

	if err != nil {
		log.Println(err.Error())
		return Location{}, true
	}

	longitude, err = strconv.ParseFloat(values[3], 64)

	if err != nil {
		log.Println(err.Error())
		return Location{}, true
	}

	return Location{
		Title:     values[0],
		Address:   values[1],
		Latitude:  latitude,
		Longitude: longitude,
	}, false
}

func parseTo(to []string, authorEmail string, matchEmail bool, delimiter string) []string {
	var emails []string
	var ids []string
	attachEmail := true

	for _, value := range trimElement(to) {
		idArray := trimElement(strings.Split(value, delimiter))

		// check match author email
		if len(idArray) > 1 {
			if email := idArray[1]; email != authorEmail {
				continue
			}

			emails = append(emails, idArray[0])
			attachEmail = false
			continue
		}

		ids = append(ids, idArray[0])
	}

	if matchEmail == true && attachEmail == false {
		return emails
	}

	for _, value := range emails {
		ids = append(ids, value)
	}

	return ids
}

// Bot is new Line Bot clien.
func (p Plugin) Bot() (*linebot.Client, error) {
	if len(p.Config.ChannelToken) == 0 || len(p.Config.ChannelSecret) == 0 {
		log.Println("missing line bot config")

		return nil, errors.New("missing line bot config")
	}

	return linebot.New(p.Config.ChannelSecret, p.Config.ChannelToken)
}

func (p Plugin) getTunnelDomain() (string, error) {
	var domain string
	if p.Config.Domain != "" {
		if len(p.Config.Domain) < 4 || len(p.Config.Domain) > 63 {
			return "", errors.New("tunnel host name must be lowercase and between 4 and 63 alphanumeric characters")
		}
		domain = p.Config.Domain
	} else {
		domain = strings.ToLower(random.String(10))
	}

	return domain, nil
}

// Handler is http handler.
func (p Plugin) Handler(bot *linebot.Client) *http.ServeMux {
	mux := http.NewServeMux()

	// Setup HTTP Server for receiving requests from LINE platform
	mux.HandleFunc("/callback", func(w http.ResponseWriter, req *http.Request) {
		events, err := bot.ParseRequest(req)
		if err != nil {
			if err == linebot.ErrInvalidSignature {
				w.WriteHeader(400)
			} else {
				w.WriteHeader(500)
			}
			return
		}
		for _, event := range events {
			if event.Type == linebot.EventTypeMessage {
				switch message := event.Message.(type) {
				case *linebot.TextMessage:
					log.Printf("User ID is %v\n", event.Source.UserID)
					ReceiveCount++
					if message.Text == "test" {
						SendCount++
						if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage("count + 1")).Do(); err != nil {
							log.Print(err)
						}
					}
					if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(message.Text)).Do(); err != nil {
						log.Print(err)
					}
				}
			}
		}
	})

	mux.HandleFunc("/metrics", func(w http.ResponseWriter, req *http.Request) {
		promhttp.Handler().ServeHTTP(w, req)
	})

	// Setup HTTP Server for receiving requests from LINE platform
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		fmt.Fprintln(w, "Welcome to Line webhook page.")
	})

	return mux
}

// Webhook support line callback service.
func (p Plugin) Webhook() error {
	readyToListen := false
	bot, err := p.Bot()

	if err != nil {
		return err
	}
	mux := p.Handler(bot)

	if p.Config.Tunnel {
		if p.Config.Debug {
			gotunnelme.Debug = true
		}

		domain, err := p.getTunnelDomain()
		if err != nil {
			panic(err)
		}

		tunnel := gotunnelme.NewTunnel()
		url, err := tunnel.GetUrl(domain)
		if err != nil {
			panic("Could not get localtunnel.me URL. " + err.Error())
		}
		go func() {
			for !readyToListen {
				time.Sleep(1 * time.Second)
			}
			c := color.New(color.FgYellow)
			c.Println("Tunnel URL:", url)
			err := tunnel.CreateTunnel(p.Config.Port)
			if err != nil {
				panic("Could not create tunnel. " + err.Error())
			}
		}()
	}

	readyToListen = true
	if p.Config.Port != 443 && !p.Config.AutoTLS {
		log.Println("Line Webhook Server Listin on " + strconv.Itoa(p.Config.Port) + " port")
		if err := http.ListenAndServe(":"+strconv.Itoa(p.Config.Port), mux); err != nil {
			log.Fatal(err)
		}
	}

	if p.Config.AutoTLS && len(p.Config.Host) != 0 {
		log.Println("Line Webhook Server Listin on 443 port, hostname: " + strings.Join(p.Config.Host, ", "))
		return http.Serve(autocert.NewListener(p.Config.Host...), mux)
	}

	return nil
}

// Exec executes the plugin.
func (p Plugin) Exec() error {

	bot, err := p.Bot()

	if err != nil {
		return err
	}

	if len(p.Config.To) == 0 {
		log.Println("missing line user config")

		return errors.New("missing line user config")
	}

	var message []string
	if len(p.Config.Message) > 0 {
		message = p.Config.Message
	} else {
		message = p.Message(p.Repo, p.Build)
	}

	// Initial messages array.
	var messages []linebot.Message

	for _, value := range trimElement(message) {
		txt, err := template.RenderTrim(value, p)
		if err != nil {
			return err
		}

		messages = append(messages, linebot.NewTextMessage(txt))
	}

	// Add image message
	for _, value := range trimElement(p.Config.Image) {
		values := convertImage(value, p.Config.Delimiter)

		messages = append(messages, linebot.NewImageMessage(values[0], values[1]))
	}

	// Add image message.
	for _, value := range trimElement(p.Config.Video) {
		values := convertVideo(value, p.Config.Delimiter)

		messages = append(messages, linebot.NewVideoMessage(values[0], values[1]))
	}

	// Add Audio message.
	for _, value := range trimElement(p.Config.Audio) {
		audio, empty := convertAudio(value, p.Config.Delimiter)

		if empty == true {
			continue
		}

		messages = append(messages, linebot.NewAudioMessage(audio.URL, audio.Duration))
	}

	// Add Sticker message.
	for _, value := range trimElement(p.Config.Sticker) {
		sticker, empty := convertSticker(value, p.Config.Delimiter)

		if empty == true {
			continue
		}

		messages = append(messages, linebot.NewStickerMessage(sticker[0], sticker[1]))
	}

	// check Location array.
	for _, value := range trimElement(p.Config.Location) {
		location, empty := convertLocation(value, p.Config.Delimiter)

		if empty == true {
			continue
		}

		messages = append(messages, linebot.NewLocationMessage(location.Title, location.Address, location.Latitude, location.Longitude))
	}

	ids := parseTo(p.Config.To, p.Build.Email, p.Config.MatchEmail, p.Config.Delimiter)

	// Send messages to multiple users at any time.
	if _, err := bot.Multicast(ids, messages...).Do(); err != nil {
		log.Println(err.Error())
	}

	return nil
}

// Message is line default message.
func (p Plugin) Message(repo Repo, build Build) []string {
	return []string{fmt.Sprintf("[%s] <%s> (%s)『%s』by %s",
		build.Status,
		build.Link,
		build.Branch,
		build.Message,
		build.Author,
	)}
}
<|endoftext|>


<commit_before>package discovery

import (
	"context"
	"time"

	cid "github.com/ipfs/go-cid"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	mh "github.com/multiformats/go-multihash"
)

// RoutingDiscovery is an implementation of discovery using ContentRouting
type RoutingDiscovery struct {
	routing.ContentRouting
}

func NewRoutingDiscovery(router routing.ContentRouting) *RoutingDiscovery {
	return &RoutingDiscovery{router}
}

func (d *RoutingDiscovery) Advertise(ctx context.Context, ns string, opts ...Option) (time.Duration, error) {
	cid, err := nsToCid(ns)
	if err != nil {
		return 0, err
	}

	err = d.Provide(ctx, cid, true)
	if err != nil {
		return 0, err
	}

	// this is the dht provide validity
	return 24 * time.Hour, nil
}

func (d *RoutingDiscovery) FindPeers(ctx context.Context, ns string, opts ...Option) (<-chan pstore.PeerInfo, error) {
	var options Options
	err := options.Apply(opts...)
	if err != nil {
		return nil, err
	}

	limit := options.Limit
	if limit == 0 {
		limit = 100 // that's just arbitrary, but FindProvidersAsync needs a count
	}

	cid, err := nsToCid(ns)
	if err != nil {
		return nil, err
	}

	return d.FindProvidersAsync(ctx, cid, limit), nil
}

func nsToCid(ns string) (cid.Cid, error) {
	h, err := mh.Encode([]byte(ns), mh.SHA2_256)
	if err != nil {
		return cid.Undef, err
	}

	return cid.NewCidV1(cid.Raw, h), nil
}
<commit_msg>document the use of SHA256 as the cid hash function<commit_after>package discovery

import (
	"context"
	"time"

	cid "github.com/ipfs/go-cid"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	mh "github.com/multiformats/go-multihash"
)

// RoutingDiscovery is an implementation of discovery using ContentRouting
// Namespaces are translated to Cids using the SHA256 hash.
type RoutingDiscovery struct {
	routing.ContentRouting
}

func NewRoutingDiscovery(router routing.ContentRouting) *RoutingDiscovery {
	return &RoutingDiscovery{router}
}

func (d *RoutingDiscovery) Advertise(ctx context.Context, ns string, opts ...Option) (time.Duration, error) {
	cid, err := nsToCid(ns)
	if err != nil {
		return 0, err
	}

	err = d.Provide(ctx, cid, true)
	if err != nil {
		return 0, err
	}

	// this is the dht provide validity
	return 24 * time.Hour, nil
}

func (d *RoutingDiscovery) FindPeers(ctx context.Context, ns string, opts ...Option) (<-chan pstore.PeerInfo, error) {
	var options Options
	err := options.Apply(opts...)
	if err != nil {
		return nil, err
	}

	limit := options.Limit
	if limit == 0 {
		limit = 100 // that's just arbitrary, but FindProvidersAsync needs a count
	}

	cid, err := nsToCid(ns)
	if err != nil {
		return nil, err
	}

	return d.FindProvidersAsync(ctx, cid, limit), nil
}

func nsToCid(ns string) (cid.Cid, error) {
	h, err := mh.Encode([]byte(ns), mh.SHA2_256)
	if err != nil {
		return cid.Undef, err
	}

	return cid.NewCidV1(cid.Raw, h), nil
}
<|endoftext|>