<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"
	"sync"
	"time"
)

const (
	errPath = "/error"
)

type session struct {
	host   string
	cancel chan struct{}
}

type stringSlice []string

func (sslice *stringSlice) String() string {
	return fmt.Sprintf("%v", *sslice)
}

func (sslice *stringSlice) Set(value string) error {
	for _, s := range strings.Split(value, ",") {
		*sslice = append(*sslice, strings.TrimSpace(s))
	}
	return nil
}

var (
	listen  string
	timeout time.Duration
	nodes   stringSlice
	hosts   chan string
	route   map[string]*session = make(map[string]*session)
	lock    sync.RWMutex
)

func errFunc(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "Session not found", http.StatusNotFound)
}

func status(w http.ResponseWriter, r *http.Request) {
	status := make(map[string]string)
	lock.RLock()
	for k, v := range route {
		status[k] = v.host
	}
	lock.RUnlock()
	json.NewEncoder(w).Encode(&status)
}

func create(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	log.Println("[NEW_REQUEST] - Waiting for a free node")
	host := <-hosts
	r.URL.Scheme = "http"
	r.URL.Host = host
	r.URL.Path = strings.Replace(r.URL.Path, "/wd/hub", "", 1)
	log.Printf("[SESSION_ATTEMPTED] [%s]\n", host)
	resp, err := http.Post(r.URL.String(), "", r.Body)
	if err != nil {
		log.Printf("[SESSION_FAILED] [%s] - [%s]\n", host, err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		hosts <- host
		return
	}
	w.WriteHeader(resp.StatusCode)
	var s struct {
		Id string `json:"sessionId"`
	}
	tee := io.TeeReader(resp.Body, w)
	json.NewDecoder(tee).Decode(&s)
	if s.Id == "" {
		hosts <- host
		return
	}
	lock.Lock()
	route[s.Id] = &session{host, onTimeout(timeout, func() { deleteSession(s.Id) })}
	lock.Unlock()
	log.Printf("[SESSION_CREATED] [%s] [%s]\n", s.Id, host)
}

func proxy(r *http.Request) {
	dump, _ := httputil.DumpRequest(r, true)
	log.Println(string(dump))
	r.URL.Scheme = "http"
	sid := strings.Split(r.URL.Path, "/")[4]
	lock.RLock()
	s, ok := route[sid]
	lock.RUnlock()
	if ok {
		close(s.cancel)
		r.URL.Host = s.host
		r.URL.Path = strings.Replace(r.URL.Path, "/wd/hub", "", 1)
		if r.Method != http.MethodDelete {
			lock.Lock()
			s.cancel = onTimeout(timeout, func() { deleteSession(sid) })
			lock.Unlock()
			return
		}
		lock.Lock()
		delete(route, sid)
		hosts <- s.host
		lock.Unlock()
		log.Printf("[SESSION_DELETED] [%s]\n", sid)
		return
	}
	r.URL.Host = listen
	r.URL.Path = errPath
}

func deleteSession(id string) {
	log.Printf("[SESSION_TIMED_OUT] [%s] - Deleting session\n", id)
	req, err := http.NewRequest(http.MethodDelete, fmt.Sprintf("http://%s/wd/hub/session/%s", listen, id), nil)
	if err != nil {
		return
	}
	http.DefaultClient.Do(req)
}

func onTimeout(t time.Duration, f func()) chan struct{} {
	cancel := make(chan struct{})
	go func() {
		select {
		case <-time.After(t):
			f()
		case <-cancel:
		}
	}()
	return cancel
}

func handlers() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc(errPath, errFunc)
	mux.HandleFunc("/status", status)
	mux.HandleFunc("/wd/hub/session", create)
	mux.Handle("/wd/hub/session/", &httputil.ReverseProxy{Director: proxy})
	return mux
}

func init() {
	flag.StringVar(&listen, "listen", ":4444", "network address to accept connections")
	flag.DurationVar(&timeout, "timeout", 60*time.Second, "session idle timeout in seconds")
	flag.Var(&nodes, "nodes", "underlying driver nodes (required)")
	flag.Parse()
}

func queue(nodes stringSlice) {
	log.Println("Initializing nodes with", nodes)
	hosts = make(chan string, len(nodes))
	for _, h := range nodes {
		hosts <- h
	}
}

func main() {
	if len(nodes) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	host, port, err := net.SplitHostPort(listen)
	if err != nil {
		log.Fatal("invalid listen argument")
	}
	if host == "" {
		host = "localhost"
	}
	listen = fmt.Sprintf("%s:%s", host, port)
	queue(nodes)
	log.Printf("Listening on %s\n", listen)
	log.Fatal(http.ListenAndServe(listen, handlers()))
}
<commit_msg>Fixed proxy on error logic<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"
	"sync"
	"time"
)

const (
	errPath = "/error"
)

type session struct {
	host   string
	cancel chan struct{}
}

type stringSlice []string

func (sslice *stringSlice) String() string {
	return fmt.Sprintf("%v", *sslice)
}

func (sslice *stringSlice) Set(value string) error {
	for _, s := range strings.Split(value, ",") {
		*sslice = append(*sslice, strings.TrimSpace(s))
	}
	return nil
}

var (
	listen  string
	timeout time.Duration
	nodes   stringSlice
	hosts   chan string
	route   map[string]*session = make(map[string]*session)
	lock    sync.RWMutex
)

func errFunc(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "Session not found", http.StatusNotFound)
}

func status(w http.ResponseWriter, r *http.Request) {
	status := make(map[string]string)
	lock.RLock()
	for k, v := range route {
		status[k] = v.host
	}
	lock.RUnlock()
	json.NewEncoder(w).Encode(&status)
}

func create(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	log.Println("[NEW_REQUEST] - Waiting for a free node")
	host := <-hosts
	r.URL.Scheme = "http"
	r.URL.Host = host
	r.URL.Path = strings.Replace(r.URL.Path, "/wd/hub", "", 1)
	log.Printf("[SESSION_ATTEMPTED] [%s]\n", host)
	resp, err := http.Post(r.URL.String(), "", r.Body)
	if err != nil {
		log.Printf("[SESSION_FAILED] [%s] - [%s]\n", host, err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		hosts <- host
		return
	}
	w.WriteHeader(resp.StatusCode)
	var s struct {
		Id string `json:"sessionId"`
	}
	tee := io.TeeReader(resp.Body, w)
	json.NewDecoder(tee).Decode(&s)
	if s.Id == "" {
		hosts <- host
		return
	}
	lock.Lock()
	route[s.Id] = &session{host, onTimeout(timeout, func() { deleteSession(s.Id) })}
	lock.Unlock()
	log.Printf("[SESSION_CREATED] [%s] [%s]\n", s.Id, host)
}

func proxy(r *http.Request) {
	dump, _ := httputil.DumpRequest(r, true)
	log.Println(string(dump))
	r.URL.Scheme = "http"
	sid := strings.Split(r.URL.Path, "/")[4]
	lock.RLock()
	s, ok := route[sid]
	lock.RUnlock()
	if ok {
		close(s.cancel)
		r.URL.Host = s.host
		r.URL.Path = strings.Replace(r.URL.Path, "/wd/hub", "", 1)
		if r.Method != http.MethodDelete {
			lock.Lock()
			s.cancel = onTimeout(timeout, func() { deleteSession(sid) })
			lock.Unlock()
			return
		}
		lock.Lock()
		delete(route, sid)
		hosts <- s.host
		lock.Unlock()
		log.Printf("[SESSION_DELETED] [%s]\n", sid)
		return
	}
	host, port, _ := net.SplitHostPort(listen)
	if host == "" {
		host = "localhost"
	}
	r.URL.Host = fmt.Sprintf("%s:%s", host, port)
	r.URL.Path = errPath
}

func deleteSession(id string) {
	log.Printf("[SESSION_TIMED_OUT] [%s] - Deleting session\n", id)
	req, err := http.NewRequest(http.MethodDelete, fmt.Sprintf("http://%s/wd/hub/session/%s", listen, id), nil)
	if err != nil {
		return
	}
	http.DefaultClient.Do(req)
}

func onTimeout(t time.Duration, f func()) chan struct{} {
	cancel := make(chan struct{})
	go func() {
		select {
		case <-time.After(t):
			f()
		case <-cancel:
		}
	}()
	return cancel
}

func handlers() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc(errPath, errFunc)
	mux.HandleFunc("/status", status)
	mux.HandleFunc("/wd/hub/session", create)
	mux.Handle("/wd/hub/session/", &httputil.ReverseProxy{Director: proxy})
	return mux
}

func init() {
	flag.StringVar(&listen, "listen", ":4444", "network address to accept connections")
	flag.DurationVar(&timeout, "timeout", 60*time.Second, "session idle timeout in seconds")
	flag.Var(&nodes, "nodes", "underlying driver nodes (required)")
	flag.Parse()
}

func queue(nodes stringSlice) {
	log.Println("Initializing nodes with", nodes)
	hosts = make(chan string, len(nodes))
	for _, h := range nodes {
		hosts <- h
	}
}

func main() {
	if len(nodes) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	queue(nodes)
	log.Printf("Listening on %s\n", listen)
	log.Fatal(http.ListenAndServe(listen, handlers()))
}
<|endoftext|>


<commit_before>package telebot

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strconv"
	"strings"
)

// Recipient is any possible endpoint you can send
// messages to: either user, group or a channel.
type Recipient interface {
	Recipient() string // must return legit Telegram chat_id or username
}

// Sendable is any object that can send itself.
//
// This is pretty cool, since it lets bots implement
// custom Sendables for complex kind of media or
// chat objects spanning across multiple messages.
//
type Sendable interface {
	Send(*Bot, Recipient, *SendOptions) (*Message, error)
}

// Send delivers media through bot b to recipient.
func (p *Photo) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
		"caption": p.Caption,
	}
	b.embedSendOptions(params, opt)

	msg, err := b.sendMedia(p, params, nil)
	if err != nil {
		return nil, err
	}

	msg.Photo.File.stealRef(&p.File)
	*p = *msg.Photo
	p.Caption = msg.Caption

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (a *Audio) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   a.Caption,
		"performer": a.Performer,
		"title":     a.Title,
		"file_name": a.FileName,
	}
	b.embedSendOptions(params, opt)

	if a.Duration != 0 {
		params["duration"] = strconv.Itoa(a.Duration)
	}

	msg, err := b.sendMedia(a, params, thumbnailToFilemap(a.Thumbnail))
	if err != nil {
		return nil, err
	}

	if msg.Audio != nil {
		msg.Audio.File.stealRef(&a.File)
		*a = *msg.Audio
		a.Caption = msg.Caption
	}

	if msg.Document != nil {
		msg.Document.File.stealRef(&a.File)
		a.File = msg.Document.File
	}

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (d *Document) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   d.Caption,
		"file_name": d.FileName,
	}
	b.embedSendOptions(params, opt)

	if d.FileSize != 0 {
		params["file_size"] = strconv.Itoa(d.FileSize)
	}
	if d.DisableTypeDetection {
		params["disable_content_type_detection"] = "true"
	}

	msg, err := b.sendMedia(d, params, thumbnailToFilemap(d.Thumbnail))
	if err != nil {
		return nil, err
	}

	msg.Document.File.stealRef(&d.File)
	*d = *msg.Document
	d.Caption = msg.Caption

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (s *Sticker) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
	}
	b.embedSendOptions(params, opt)

	msg, err := b.sendMedia(s, params, nil)
	if err != nil {
		return nil, err
	}

	msg.Sticker.File.stealRef(&s.File)
	*s = *msg.Sticker

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (v *Video) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   v.Caption,
		"file_name": v.FileName,
	}
	b.embedSendOptions(params, opt)

	if v.Duration != 0 {
		params["duration"] = strconv.Itoa(v.Duration)
	}
	if v.Width != 0 {
		params["width"] = strconv.Itoa(v.Width)
	}
	if v.Height != 0 {
		params["height"] = strconv.Itoa(v.Height)
	}
	if v.Streaming {
		params["supports_streaming"] = "true"
	}

	msg, err := b.sendMedia(v, params, thumbnailToFilemap(v.Thumbnail))
	if err != nil {
		return nil, err
	}

	if vid := msg.Video; vid != nil {
		vid.File.stealRef(&v.File)
		*v = *vid
		v.Caption = msg.Caption
	} else if doc := msg.Document; doc != nil {
		// If video has no sound, Telegram can turn it into Document (GIF)
		doc.File.stealRef(&v.File)

		v.Caption = doc.Caption
		v.MIME = doc.MIME
		v.Thumbnail = doc.Thumbnail
	}

	return msg, nil
}

// Send delivers animation through bot b to recipient.
func (a *Animation) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   a.Caption,
		"file_name": a.FileName,
	}
	b.embedSendOptions(params, opt)

	if a.Duration != 0 {
		params["duration"] = strconv.Itoa(a.Duration)
	}
	if a.Width != 0 {
		params["width"] = strconv.Itoa(a.Width)
	}
	if a.Height != 0 {
		params["height"] = strconv.Itoa(a.Height)
	}

	// file_name is required, without it animation sends as a document
	if params["file_name"] == "" && a.File.OnDisk() {
		params["file_name"] = filepath.Base(a.File.FileLocal)
	}

	msg, err := b.sendMedia(a, params, thumbnailToFilemap(a.Thumbnail))
	if err != nil {
		return nil, err
	}

	if anim := msg.Animation; anim != nil {
		anim.File.stealRef(&a.File)
		*a = *msg.Animation
	} else if doc := msg.Document; doc != nil {
		*a = Animation{
			File:      doc.File,
			Thumbnail: doc.Thumbnail,
			MIME:      doc.MIME,
			FileName:  doc.FileName,
		}
	}

	a.Caption = msg.Caption
	return msg, nil
}

// Send delivers media through bot b to recipient.
func (v *Voice) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
		"caption": v.Caption,
	}
	b.embedSendOptions(params, opt)

	if v.Duration != 0 {
		params["duration"] = strconv.Itoa(v.Duration)
	}

	msg, err := b.sendMedia(v, params, nil)
	if err != nil {
		return nil, err
	}

	msg.Voice.File.stealRef(&v.File)
	*v = *msg.Voice

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (v *VideoNote) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
	}
	b.embedSendOptions(params, opt)

	if v.Duration != 0 {
		params["duration"] = strconv.Itoa(v.Duration)
	}
	if v.Length != 0 {
		params["length"] = strconv.Itoa(v.Length)
	}

	msg, err := b.sendMedia(v, params, thumbnailToFilemap(v.Thumbnail))
	if err != nil {
		return nil, err
	}

	msg.VideoNote.File.stealRef(&v.File)
	*v = *msg.VideoNote

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (x *Location) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":     to.Recipient(),
		"latitude":    fmt.Sprintf("%f", x.Lat),
		"longitude":   fmt.Sprintf("%f", x.Lng),
		"live_period": strconv.Itoa(x.LivePeriod),
	}
	if x.HorizontalAccuracy != nil {
		params["horizontal_accuracy"] = fmt.Sprintf("%f", *x.HorizontalAccuracy)
	}
	if x.Heading != 0 {
		params["heading"] = strconv.Itoa(x.Heading)
	}
	if x.AlertRadius != 0 {
		params["proximity_alert_radius"] = strconv.Itoa(x.Heading)
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendLocation", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers media through bot b to recipient.
func (v *Venue) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":           to.Recipient(),
		"latitude":          fmt.Sprintf("%f", v.Location.Lat),
		"longitude":         fmt.Sprintf("%f", v.Location.Lng),
		"title":             v.Title,
		"address":           v.Address,
		"foursquare_id":     v.FoursquareID,
		"foursquare_type":   v.FoursquareType,
		"google_place_id":   v.GooglePlaceID,
		"google_place_type": v.GooglePlaceType,
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendVenue", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers invoice through bot b to recipient.
func (i *Invoice) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":                       to.Recipient(),
		"title":                         i.Title,
		"description":                   i.Description,
		"start_parameter":               i.Start,
		"payload":                       i.Payload,
		"provider_token":                i.Token,
		"currency":                      i.Currency,
		"max_tip_amount":                strconv.Itoa(i.MaxTipAmount),
		"need_name":                     strconv.FormatBool(i.NeedName),
		"need_phone_number":             strconv.FormatBool(i.NeedPhoneNumber),
		"need_email":                    strconv.FormatBool(i.NeedEmail),
		"need_shipping_address":         strconv.FormatBool(i.NeedShippingAddress),
		"send_phone_number_to_provider": strconv.FormatBool(i.SendPhoneNumber),
		"send_email_to_provider":        strconv.FormatBool(i.SendEmail),
		"is_flexible":                   strconv.FormatBool(i.Flexible),
	}
	if i.Photo != nil {
		if i.Photo.FileURL != "" {
			params["photo_url"] = i.Photo.FileURL
		}
		if i.PhotoSize > 0 {
			params["photo_size"] = strconv.Itoa(i.PhotoSize)
		}
		if i.Photo.Width > 0 {
			params["photo_width"] = strconv.Itoa(i.Photo.Width)
		}
		if i.Photo.Height > 0 {
			params["photo_height"] = strconv.Itoa(i.Photo.Height)
		}
	}
	if len(i.Prices) > 0 {
		data, _ := json.Marshal(i.Prices)
		params["prices"] = string(data)
	}
	if len(i.SuggestedTipAmounts) > 0 {
		params["suggested_tip_amounts"] = "[" + strings.Join(intsToStrs(i.SuggestedTipAmounts), ",") + "]"
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendInvoice", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers poll through bot b to recipient.
func (p *Poll) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":                 to.Recipient(),
		"question":                p.Question,
		"type":                    string(p.Type),
		"is_closed":               strconv.FormatBool(p.Closed),
		"is_anonymous":            strconv.FormatBool(p.Anonymous),
		"allows_multiple_answers": strconv.FormatBool(p.MultipleAnswers),
		"correct_option_id":       strconv.Itoa(p.CorrectOption),
	}
	if p.Explanation != "" {
		params["explanation"] = p.Explanation
		params["explanation_parse_mode"] = p.ParseMode
	}
	if p.OpenPeriod != 0 {
		params["open_period"] = strconv.Itoa(p.OpenPeriod)
	} else if p.CloseUnixdate != 0 {
		params["close_date"] = strconv.FormatInt(p.CloseUnixdate, 10)
	}
	b.embedSendOptions(params, opt)

	var options []string
	for _, o := range p.Options {
		options = append(options, o.Text)
	}

	opts, _ := json.Marshal(options)
	params["options"] = string(opts)

	data, err := b.Raw("sendPoll", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers dice through bot b to recipient.
func (d *Dice) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
		"emoji":   string(d.Type),
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendDice", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers game through bot b to recipient.
func (g *Game) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":         to.Recipient(),
		"game_short_name": g.Name,
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendGame", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}
<commit_msg>send video as document<commit_after>package telebot

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strconv"
	"strings"
)

// Recipient is any possible endpoint you can send
// messages to: either user, group or a channel.
type Recipient interface {
	Recipient() string // must return legit Telegram chat_id or username
}

// Sendable is any object that can send itself.
//
// This is pretty cool, since it lets bots implement
// custom Sendables for complex kind of media or
// chat objects spanning across multiple messages.
//
type Sendable interface {
	Send(*Bot, Recipient, *SendOptions) (*Message, error)
}

// Send delivers media through bot b to recipient.
func (p *Photo) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
		"caption": p.Caption,
	}
	b.embedSendOptions(params, opt)

	msg, err := b.sendMedia(p, params, nil)
	if err != nil {
		return nil, err
	}

	msg.Photo.File.stealRef(&p.File)
	*p = *msg.Photo
	p.Caption = msg.Caption

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (a *Audio) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   a.Caption,
		"performer": a.Performer,
		"title":     a.Title,
		"file_name": a.FileName,
	}
	b.embedSendOptions(params, opt)

	if a.Duration != 0 {
		params["duration"] = strconv.Itoa(a.Duration)
	}

	msg, err := b.sendMedia(a, params, thumbnailToFilemap(a.Thumbnail))
	if err != nil {
		return nil, err
	}

	if msg.Audio != nil {
		msg.Audio.File.stealRef(&a.File)
		*a = *msg.Audio
		a.Caption = msg.Caption
	}

	if msg.Document != nil {
		msg.Document.File.stealRef(&a.File)
		a.File = msg.Document.File
	}

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (d *Document) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   d.Caption,
		"file_name": d.FileName,
	}
	b.embedSendOptions(params, opt)

	if d.FileSize != 0 {
		params["file_size"] = strconv.Itoa(d.FileSize)
	}
	if d.DisableTypeDetection {
		params["disable_content_type_detection"] = "true"
	}

	msg, err := b.sendMedia(d, params, thumbnailToFilemap(d.Thumbnail))
	if err != nil {
		return nil, err
	}

	if doc := msg.Document; doc != nil {
		doc.File.stealRef(&d.File)
		*d = *doc
		d.Caption = msg.Caption
	} else if vid := msg.Video; vid != nil {
		vid.File.stealRef(&d.File)
		d.Caption = vid.Caption
		d.MIME = vid.MIME
		d.Thumbnail = vid.Thumbnail
	}

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (s *Sticker) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
	}
	b.embedSendOptions(params, opt)

	msg, err := b.sendMedia(s, params, nil)
	if err != nil {
		return nil, err
	}

	msg.Sticker.File.stealRef(&s.File)
	*s = *msg.Sticker

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (v *Video) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   v.Caption,
		"file_name": v.FileName,
	}
	b.embedSendOptions(params, opt)

	if v.Duration != 0 {
		params["duration"] = strconv.Itoa(v.Duration)
	}
	if v.Width != 0 {
		params["width"] = strconv.Itoa(v.Width)
	}
	if v.Height != 0 {
		params["height"] = strconv.Itoa(v.Height)
	}
	if v.Streaming {
		params["supports_streaming"] = "true"
	}

	msg, err := b.sendMedia(v, params, thumbnailToFilemap(v.Thumbnail))
	if err != nil {
		return nil, err
	}

	if vid := msg.Video; vid != nil {
		vid.File.stealRef(&v.File)
		*v = *vid
		v.Caption = msg.Caption
	} else if doc := msg.Document; doc != nil {
		// If video has no sound, Telegram can turn it into Document (GIF)
		doc.File.stealRef(&v.File)

		v.Caption = doc.Caption
		v.MIME = doc.MIME
		v.Thumbnail = doc.Thumbnail
	}

	return msg, nil
}

// Send delivers animation through bot b to recipient.
func (a *Animation) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":   to.Recipient(),
		"caption":   a.Caption,
		"file_name": a.FileName,
	}
	b.embedSendOptions(params, opt)

	if a.Duration != 0 {
		params["duration"] = strconv.Itoa(a.Duration)
	}
	if a.Width != 0 {
		params["width"] = strconv.Itoa(a.Width)
	}
	if a.Height != 0 {
		params["height"] = strconv.Itoa(a.Height)
	}

	// file_name is required, without it animation sends as a document
	if params["file_name"] == "" && a.File.OnDisk() {
		params["file_name"] = filepath.Base(a.File.FileLocal)
	}

	msg, err := b.sendMedia(a, params, thumbnailToFilemap(a.Thumbnail))
	if err != nil {
		return nil, err
	}

	if anim := msg.Animation; anim != nil {
		anim.File.stealRef(&a.File)
		*a = *msg.Animation
	} else if doc := msg.Document; doc != nil {
		*a = Animation{
			File:      doc.File,
			Thumbnail: doc.Thumbnail,
			MIME:      doc.MIME,
			FileName:  doc.FileName,
		}
	}

	a.Caption = msg.Caption
	return msg, nil
}

// Send delivers media through bot b to recipient.
func (v *Voice) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
		"caption": v.Caption,
	}
	b.embedSendOptions(params, opt)

	if v.Duration != 0 {
		params["duration"] = strconv.Itoa(v.Duration)
	}

	msg, err := b.sendMedia(v, params, nil)
	if err != nil {
		return nil, err
	}

	msg.Voice.File.stealRef(&v.File)
	*v = *msg.Voice

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (v *VideoNote) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
	}
	b.embedSendOptions(params, opt)

	if v.Duration != 0 {
		params["duration"] = strconv.Itoa(v.Duration)
	}
	if v.Length != 0 {
		params["length"] = strconv.Itoa(v.Length)
	}

	msg, err := b.sendMedia(v, params, thumbnailToFilemap(v.Thumbnail))
	if err != nil {
		return nil, err
	}

	msg.VideoNote.File.stealRef(&v.File)
	*v = *msg.VideoNote

	return msg, nil
}

// Send delivers media through bot b to recipient.
func (x *Location) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":     to.Recipient(),
		"latitude":    fmt.Sprintf("%f", x.Lat),
		"longitude":   fmt.Sprintf("%f", x.Lng),
		"live_period": strconv.Itoa(x.LivePeriod),
	}
	if x.HorizontalAccuracy != nil {
		params["horizontal_accuracy"] = fmt.Sprintf("%f", *x.HorizontalAccuracy)
	}
	if x.Heading != 0 {
		params["heading"] = strconv.Itoa(x.Heading)
	}
	if x.AlertRadius != 0 {
		params["proximity_alert_radius"] = strconv.Itoa(x.Heading)
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendLocation", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers media through bot b to recipient.
func (v *Venue) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":           to.Recipient(),
		"latitude":          fmt.Sprintf("%f", v.Location.Lat),
		"longitude":         fmt.Sprintf("%f", v.Location.Lng),
		"title":             v.Title,
		"address":           v.Address,
		"foursquare_id":     v.FoursquareID,
		"foursquare_type":   v.FoursquareType,
		"google_place_id":   v.GooglePlaceID,
		"google_place_type": v.GooglePlaceType,
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendVenue", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers invoice through bot b to recipient.
func (i *Invoice) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":                       to.Recipient(),
		"title":                         i.Title,
		"description":                   i.Description,
		"start_parameter":               i.Start,
		"payload":                       i.Payload,
		"provider_token":                i.Token,
		"currency":                      i.Currency,
		"max_tip_amount":                strconv.Itoa(i.MaxTipAmount),
		"need_name":                     strconv.FormatBool(i.NeedName),
		"need_phone_number":             strconv.FormatBool(i.NeedPhoneNumber),
		"need_email":                    strconv.FormatBool(i.NeedEmail),
		"need_shipping_address":         strconv.FormatBool(i.NeedShippingAddress),
		"send_phone_number_to_provider": strconv.FormatBool(i.SendPhoneNumber),
		"send_email_to_provider":        strconv.FormatBool(i.SendEmail),
		"is_flexible":                   strconv.FormatBool(i.Flexible),
	}
	if i.Photo != nil {
		if i.Photo.FileURL != "" {
			params["photo_url"] = i.Photo.FileURL
		}
		if i.PhotoSize > 0 {
			params["photo_size"] = strconv.Itoa(i.PhotoSize)
		}
		if i.Photo.Width > 0 {
			params["photo_width"] = strconv.Itoa(i.Photo.Width)
		}
		if i.Photo.Height > 0 {
			params["photo_height"] = strconv.Itoa(i.Photo.Height)
		}
	}
	if len(i.Prices) > 0 {
		data, _ := json.Marshal(i.Prices)
		params["prices"] = string(data)
	}
	if len(i.SuggestedTipAmounts) > 0 {
		params["suggested_tip_amounts"] = "[" + strings.Join(intsToStrs(i.SuggestedTipAmounts), ",") + "]"
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendInvoice", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers poll through bot b to recipient.
func (p *Poll) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":                 to.Recipient(),
		"question":                p.Question,
		"type":                    string(p.Type),
		"is_closed":               strconv.FormatBool(p.Closed),
		"is_anonymous":            strconv.FormatBool(p.Anonymous),
		"allows_multiple_answers": strconv.FormatBool(p.MultipleAnswers),
		"correct_option_id":       strconv.Itoa(p.CorrectOption),
	}
	if p.Explanation != "" {
		params["explanation"] = p.Explanation
		params["explanation_parse_mode"] = p.ParseMode
	}
	if p.OpenPeriod != 0 {
		params["open_period"] = strconv.Itoa(p.OpenPeriod)
	} else if p.CloseUnixdate != 0 {
		params["close_date"] = strconv.FormatInt(p.CloseUnixdate, 10)
	}
	b.embedSendOptions(params, opt)

	var options []string
	for _, o := range p.Options {
		options = append(options, o.Text)
	}

	opts, _ := json.Marshal(options)
	params["options"] = string(opts)

	data, err := b.Raw("sendPoll", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers dice through bot b to recipient.
func (d *Dice) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id": to.Recipient(),
		"emoji":   string(d.Type),
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendDice", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}

// Send delivers game through bot b to recipient.
func (g *Game) Send(b *Bot, to Recipient, opt *SendOptions) (*Message, error) {
	params := map[string]string{
		"chat_id":         to.Recipient(),
		"game_short_name": g.Name,
	}
	b.embedSendOptions(params, opt)

	data, err := b.Raw("sendGame", params)
	if err != nil {
		return nil, err
	}

	return extractMessage(data)
}
<|endoftext|>


<commit_before>package configuration

import (
	"flag"
	"log"

	"github.com/FogCreek/mini"
)

type hostConfiguration struct {
	TTL int64
}

type webConfiguration struct {
	Address string
	Status  bool
}

type travisConfiguration struct {
	Authenticate bool
	Token        string
}

var (
	Host   hostConfiguration
	Web    webConfiguration
	Travis travisConfiguration
)

func Process() {
	path := flag.String("config", "/etc/iago.ini", "Configuration file path")
	flag.Parse()

	config, err := mini.LoadConfiguration(*path)

	if err != nil {
		log.Fatal(err)
	}

	Host.TTL = config.IntegerFromSection("Host", "TTL", 30)

	Web.Address = config.StringFromSection("Web", "Address", "127.0.01:8000")
	Web.Status = config.BooleanFromSection("Web", "Status", true)

	Travis.Authenticate = config.BooleanFromSection("Travis", "Authenticate", false)
	Travis.Token = config.StringFromSection("Travis", "Token", "")
}
<commit_msg>Changed the default config file path to /etc/iagod.ini<commit_after>package configuration

import (
	"flag"
	"log"

	"github.com/FogCreek/mini"
)

type hostConfiguration struct {
	TTL int64
}

type webConfiguration struct {
	Address string
	Status  bool
}

type travisConfiguration struct {
	Authenticate bool
	Token        string
}

var (
	Host   hostConfiguration
	Web    webConfiguration
	Travis travisConfiguration
)

func Process() {
	path := flag.String("config", "/etc/iagod.ini", "Configuration file path")
	flag.Parse()

	config, err := mini.LoadConfiguration(*path)

	if err != nil {
		log.Fatal(err)
	}

	Host.TTL = config.IntegerFromSection("Host", "TTL", 30)

	Web.Address = config.StringFromSection("Web", "Address", "127.0.01:8000")
	Web.Status = config.BooleanFromSection("Web", "Status", true)

	Travis.Authenticate = config.BooleanFromSection("Travis", "Authenticate", false)
	Travis.Token = config.StringFromSection("Travis", "Token", "")
}
<|endoftext|>


<commit_before>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package http

import (
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"

	"sigs.k8s.io/gateway-api/conformance/utils/roundtripper"
)

// ExpectedResponse defines the response expected for a given request.
type ExpectedResponse struct {
	// Request defines the request to make.
	Request Request

	// ExpectedRequest defines the request that
	// is expected to arrive at the backend. If
	// not specified, the backend request will be
	// expected to match Request.
	ExpectedRequest *ExpectedRequest

	StatusCode int
	Backend    string
	Namespace  string
}

// Request can be used as both the request to make and a means to verify
// that echoserver received the expected request. Note that multiple header
// values can be provided, as a comma-separated value.
type Request struct {
	Host    string
	Method  string
	Path    string
	Headers map[string]string
}

// ExpectedRequest defines expected properties of a request.
type ExpectedRequest struct {
	Request

	// AbsentHeaders are names of headers that are expected
	// *not* to be present on the request.
	AbsentHeaders []string
}

// maxTimeToConsistency is the maximum time that WaitForConsistency will wait for
// requiredConsecutiveSuccesses requests to succeed in a row before failing the test.
const maxTimeToConsistency = 30 * time.Second

// requiredConsecutiveSuccesses is the number of requests that must succeed in a row
// for MakeRequestAndExpectEventuallyConsistentResponse to consider the response "consistent"
// before making additional assertions on the response body. If this number is not reached within
// maxTimeToConsistency, the test will fail.
const requiredConsecutiveSuccesses = 3

// MakeRequestAndExpectEventuallyConsistentResponse makes a request with the given parameters,
// understanding that the request may fail for some amount of time.
//
// Once the request succeeds consistently with the response having the expected status code, make
// additional assertions on the response body using the provided ExpectedResponse.
func MakeRequestAndExpectEventuallyConsistentResponse(t *testing.T, r roundtripper.RoundTripper, gwAddr string, expected ExpectedResponse) {
	t.Helper()

	if expected.Request.Method == "" {
		expected.Request.Method = "GET"
	}

	if expected.StatusCode == 0 {
		expected.StatusCode = 200
	}

	t.Logf("Making %s request to http://%s%s", expected.Request.Method, gwAddr, expected.Request.Path)

	req := roundtripper.Request{
		Method:   expected.Request.Method,
		Host:     expected.Request.Host,
		URL:      url.URL{Scheme: "http", Host: gwAddr, Path: expected.Request.Path},
		Protocol: "HTTP",
	}

	if expected.Request.Headers != nil {
		req.Headers = map[string][]string{}
		for name, value := range expected.Request.Headers {
			req.Headers[name] = []string{value}
		}
	}

	// The request expected to arrive at the backend is
	// the same as the request made, unless otherwise
	// specified.
	if expected.ExpectedRequest == nil {
		expected.ExpectedRequest = &ExpectedRequest{Request: expected.Request}
	}

	if expected.ExpectedRequest.Method == "" {
		expected.ExpectedRequest.Method = "GET"
	}

	cReq, cRes := WaitForConsistency(t, r, req, expected, requiredConsecutiveSuccesses)

	ExpectResponse(t, cReq, cRes, expected)
}

// awaitConvergence runs the given function until it returns 'true' `threshold` times in a row.
// Each failed attempt has a 1s delay; successful attempts have no delay.
func awaitConvergence(t *testing.T, threshold int, fn func() bool) {
	successes := 0
	attempts := 0
	to := time.After(maxTimeToConsistency)
	delay := time.Second
	for {
		select {
		case <-to:
			t.Fatalf("timeout while waiting after %d attempts", attempts)
		default:
		}

		completed := fn()
		attempts++
		if completed {
			successes++
			if successes >= threshold {
				return
			}
			// Skip delay if we have a success
			continue
		}

		successes = 0
		select {
		// Capture the overall timeout
		case <-to:
			t.Fatalf("timeout while waiting after %d attempts, %d/%d sucessess", attempts, successes, threshold)
			// And the per-try delay
		case <-time.After(delay):
		}
	}
}

// WaitForConsistency repeats the provided request until it completes with a response having
// the expected status code consistently. The provided threshold determines how many times in
// a row this must occur to be considered "consistent".
func WaitForConsistency(t *testing.T, r roundtripper.RoundTripper, req roundtripper.Request, expected ExpectedResponse, threshold int) (*roundtripper.CapturedRequest, *roundtripper.CapturedResponse) {
	var (
		cReq *roundtripper.CapturedRequest
		cRes *roundtripper.CapturedResponse
		err  error
	)

	awaitConvergence(t, threshold, func() bool {
		cReq, cRes, err = r.CaptureRoundTrip(req)
		if err != nil {
			t.Logf("Request failed, not ready yet: %v", err.Error())
			return false
		}

		if cRes.StatusCode != expected.StatusCode {
			t.Logf("Expected response to have status %d but got %d, not ready yet", expected.StatusCode, cRes.StatusCode)
			return false
		}

		return true
	})
	t.Logf("Request passed")

	return cReq, cRes
}

// ExpectResponse verifies that a captured request and response match the
// provided ExpectedResponse.
func ExpectResponse(t *testing.T, cReq *roundtripper.CapturedRequest, cRes *roundtripper.CapturedResponse, expected ExpectedResponse) {
	t.Helper()
	assert.Equal(t, expected.StatusCode, cRes.StatusCode, "expected status code to be %d, got %d", expected.StatusCode, cRes.StatusCode)
	if cRes.StatusCode == 200 {
		assert.Equal(t, expected.ExpectedRequest.Path, cReq.Path, "expected path to be %s, got %s", expected.ExpectedRequest.Path, cReq.Path)
		assert.Equal(t, expected.ExpectedRequest.Method, cReq.Method, "expected method to be %s, got %s", expected.ExpectedRequest.Method, cReq.Method)
		assert.Equal(t, expected.Namespace, cReq.Namespace, "expected namespace to be %s, got %s", expected.Namespace, cReq.Namespace)
		if expected.ExpectedRequest.Headers != nil {
			if cReq.Headers == nil {
				t.Error("No headers captured")
			} else {
				for name, val := range cReq.Headers {
					cReq.Headers[strings.ToLower(name)] = val
				}
				for name, expectedVal := range expected.ExpectedRequest.Headers {
					actualVal, ok := cReq.Headers[strings.ToLower(name)]
					if !ok {
						t.Errorf("Expected %s header to be set, actual headers: %v", name, cReq.Headers)
					} else if strings.Join(actualVal, ",") != expectedVal {
						t.Errorf("Expected %s header to be set to %s, got %s", name, expectedVal, strings.Join(actualVal, ","))
					}
				}
			}
		}

		// Verify that headers expected *not* to be present on the
		// request are actually not present.
		if len(expected.ExpectedRequest.AbsentHeaders) > 0 {
			for name, val := range cReq.Headers {
				cReq.Headers[strings.ToLower(name)] = val
			}

			for _, name := range expected.ExpectedRequest.AbsentHeaders {
				val, ok := cReq.Headers[strings.ToLower(name)]
				if ok {
					t.Errorf("Expected %s header to not be set, got %s", name, val)
				}
			}
		}

		if !strings.HasPrefix(cReq.Pod, expected.Backend) {
			t.Errorf("Expected pod name to start with %s, got %s", expected.Backend, cReq.Pod)
		}
	}
}
<commit_msg>move ExpectedRequest defaulting into ExpectResponse<commit_after>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package http

import (
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"

	"sigs.k8s.io/gateway-api/conformance/utils/roundtripper"
)

// ExpectedResponse defines the response expected for a given request.
type ExpectedResponse struct {
	// Request defines the request to make.
	Request Request

	// ExpectedRequest defines the request that
	// is expected to arrive at the backend. If
	// not specified, the backend request will be
	// expected to match Request.
	ExpectedRequest *ExpectedRequest

	StatusCode int
	Backend    string
	Namespace  string
}

// Request can be used as both the request to make and a means to verify
// that echoserver received the expected request. Note that multiple header
// values can be provided, as a comma-separated value.
type Request struct {
	Host    string
	Method  string
	Path    string
	Headers map[string]string
}

// ExpectedRequest defines expected properties of a request.
type ExpectedRequest struct {
	Request

	// AbsentHeaders are names of headers that are expected
	// *not* to be present on the request.
	AbsentHeaders []string
}

// maxTimeToConsistency is the maximum time that WaitForConsistency will wait for
// requiredConsecutiveSuccesses requests to succeed in a row before failing the test.
const maxTimeToConsistency = 30 * time.Second

// requiredConsecutiveSuccesses is the number of requests that must succeed in a row
// for MakeRequestAndExpectEventuallyConsistentResponse to consider the response "consistent"
// before making additional assertions on the response body. If this number is not reached within
// maxTimeToConsistency, the test will fail.
const requiredConsecutiveSuccesses = 3

// MakeRequestAndExpectEventuallyConsistentResponse makes a request with the given parameters,
// understanding that the request may fail for some amount of time.
//
// Once the request succeeds consistently with the response having the expected status code, make
// additional assertions on the response body using the provided ExpectedResponse.
func MakeRequestAndExpectEventuallyConsistentResponse(t *testing.T, r roundtripper.RoundTripper, gwAddr string, expected ExpectedResponse) {
	t.Helper()

	if expected.Request.Method == "" {
		expected.Request.Method = "GET"
	}

	if expected.StatusCode == 0 {
		expected.StatusCode = 200
	}

	t.Logf("Making %s request to http://%s%s", expected.Request.Method, gwAddr, expected.Request.Path)

	req := roundtripper.Request{
		Method:   expected.Request.Method,
		Host:     expected.Request.Host,
		URL:      url.URL{Scheme: "http", Host: gwAddr, Path: expected.Request.Path},
		Protocol: "HTTP",
	}

	if expected.Request.Headers != nil {
		req.Headers = map[string][]string{}
		for name, value := range expected.Request.Headers {
			req.Headers[name] = []string{value}
		}
	}

	cReq, cRes := WaitForConsistency(t, r, req, expected, requiredConsecutiveSuccesses)
	ExpectResponse(t, cReq, cRes, expected)
}

// awaitConvergence runs the given function until it returns 'true' `threshold` times in a row.
// Each failed attempt has a 1s delay; successful attempts have no delay.
func awaitConvergence(t *testing.T, threshold int, fn func() bool) {
	successes := 0
	attempts := 0
	to := time.After(maxTimeToConsistency)
	delay := time.Second
	for {
		select {
		case <-to:
			t.Fatalf("timeout while waiting after %d attempts", attempts)
		default:
		}

		completed := fn()
		attempts++
		if completed {
			successes++
			if successes >= threshold {
				return
			}
			// Skip delay if we have a success
			continue
		}

		successes = 0
		select {
		// Capture the overall timeout
		case <-to:
			t.Fatalf("timeout while waiting after %d attempts, %d/%d sucessess", attempts, successes, threshold)
			// And the per-try delay
		case <-time.After(delay):
		}
	}
}

// WaitForConsistency repeats the provided request until it completes with a response having
// the expected status code consistently. The provided threshold determines how many times in
// a row this must occur to be considered "consistent".
func WaitForConsistency(t *testing.T, r roundtripper.RoundTripper, req roundtripper.Request, expected ExpectedResponse, threshold int) (*roundtripper.CapturedRequest, *roundtripper.CapturedResponse) {
	var (
		cReq *roundtripper.CapturedRequest
		cRes *roundtripper.CapturedResponse
		err  error
	)

	awaitConvergence(t, threshold, func() bool {
		cReq, cRes, err = r.CaptureRoundTrip(req)
		if err != nil {
			t.Logf("Request failed, not ready yet: %v", err.Error())
			return false
		}

		if cRes.StatusCode != expected.StatusCode {
			t.Logf("Expected response to have status %d but got %d, not ready yet", expected.StatusCode, cRes.StatusCode)
			return false
		}

		return true
	})
	t.Logf("Request passed")

	return cReq, cRes
}

// ExpectResponse verifies that a captured request and response match the
// provided ExpectedResponse.
func ExpectResponse(t *testing.T, cReq *roundtripper.CapturedRequest, cRes *roundtripper.CapturedResponse, expected ExpectedResponse) {
	t.Helper()
	assert.Equal(t, expected.StatusCode, cRes.StatusCode, "expected status code to be %d, got %d", expected.StatusCode, cRes.StatusCode)
	if cRes.StatusCode == 200 {
		// The request expected to arrive at the backend is
		// the same as the request made, unless otherwise
		// specified.
		if expected.ExpectedRequest == nil {
			expected.ExpectedRequest = &ExpectedRequest{Request: expected.Request}
		}

		if expected.ExpectedRequest.Method == "" {
			expected.ExpectedRequest.Method = "GET"
		}

		assert.Equal(t, expected.ExpectedRequest.Path, cReq.Path, "expected path to be %s, got %s", expected.ExpectedRequest.Path, cReq.Path)
		assert.Equal(t, expected.ExpectedRequest.Method, cReq.Method, "expected method to be %s, got %s", expected.ExpectedRequest.Method, cReq.Method)
		assert.Equal(t, expected.Namespace, cReq.Namespace, "expected namespace to be %s, got %s", expected.Namespace, cReq.Namespace)
		if expected.ExpectedRequest.Headers != nil {
			if cReq.Headers == nil {
				t.Error("No headers captured")
			} else {
				for name, val := range cReq.Headers {
					cReq.Headers[strings.ToLower(name)] = val
				}
				for name, expectedVal := range expected.ExpectedRequest.Headers {
					actualVal, ok := cReq.Headers[strings.ToLower(name)]
					if !ok {
						t.Errorf("Expected %s header to be set, actual headers: %v", name, cReq.Headers)
					} else if strings.Join(actualVal, ",") != expectedVal {
						t.Errorf("Expected %s header to be set to %s, got %s", name, expectedVal, strings.Join(actualVal, ","))
					}
				}
			}
		}

		// Verify that headers expected *not* to be present on the
		// request are actually not present.
		if len(expected.ExpectedRequest.AbsentHeaders) > 0 {
			for name, val := range cReq.Headers {
				cReq.Headers[strings.ToLower(name)] = val
			}

			for _, name := range expected.ExpectedRequest.AbsentHeaders {
				val, ok := cReq.Headers[strings.ToLower(name)]
				if ok {
					t.Errorf("Expected %s header to not be set, got %s", name, val)
				}
			}
		}

		if !strings.HasPrefix(cReq.Pod, expected.Backend) {
			t.Errorf("Expected pod name to start with %s, got %s", expected.Backend, cReq.Pod)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017-2021 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/filter"
	amassnet "github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/PuerkitoBio/goquery"
	"github.com/caffix/stringset"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

const (
	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	httpTimeout      = 30 * time.Second
	handshakeTimeout = 5 * time.Second
)

var (
  // UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent       string
	subRE           = dns.AnySubdomainRegex()
	crawlRE         = regexp.MustCompile(`\.\w{2,6}($|\?|#)`)
	crawlFileEnds   = []string{"html", "do", "action", "cgi"}
	crawlFileStarts = []string{"js", "htm", "as", "php", "inc"}
	nameStripRE     = regexp.MustCompile(`^u[0-9a-f]{4}|20|22|25|27|2b|2f|3d|3a|40`)
)

// DefaultClient is the same HTTP client used by the package methods.
var DefaultClient *http.Client

// BasicAuth contains the data used for HTTP basic authentication.
type BasicAuth struct {
	Username string
	Password string
}

func init() {
	jar, _ := cookiejar.New(nil)
	DefaultClient = &http.Client{
		Timeout: httpTimeout,
		Transport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           amassnet.DialContext,
			MaxIdleConns:          200,
			MaxConnsPerHost:       50,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   handshakeTimeout,
			ExpectContinueTimeout: 10 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}

	switch runtime.GOOS {
	case "windows":
		UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	case "darwin":
		UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 11_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	default:
		UserAgent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	DefaultClient.Jar.SetCookies(destURL, DefaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range DefaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for the provided URL when successful.
func RequestWebPage(ctx context.Context, u string, body io.Reader, hvals map[string]string, auth *BasicAuth) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequestWithContext(ctx, method, u, body)
	if err != nil {
		return "", err
	}
	if auth != nil && auth.Username != "" && auth.Password != "" {
		req.SetBasicAuth(auth.Username, auth.Password)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)

	for k, v := range hvals {
		req.Header.Set(k, v)
	}

	resp, err := DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		err = errors.New(resp.Status)
	}
	return string(in), err
}

// Crawl will spider the web page at the URL argument looking for DNS names within the scope argument.
func Crawl(ctx context.Context, u string, scope []string, max int, f filter.Filter) ([]string, error) {
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("The context expired")
	default:
	}

	newScope := append([]string{}, scope...)

	target := subRE.FindString(u)
	if target != "" {
		var found bool
		for _, domain := range newScope {
			if target == domain {
				found = true
				break
			}
		}
		if !found {
			newScope = append(newScope, target)
		}
	}

	if f == nil {
		f = filter.NewStringFilter()
	}

	var count int
	var m sync.Mutex
	results := stringset.New()
	g := geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:        newScope,
		StartURLs:             []string{u},
		Timeout:               5 * time.Minute,
		RobotsTxtDisabled:     true,
		UserAgent:             UserAgent,
		LogDisabled:           true,
		ConcurrentRequests:    5,
		RequestDelay:          750 * time.Millisecond,
		RequestDelayRandomize: true,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			resp, err := httputil.DumpResponse(interface{}(r).(*http.Response), true)
			if err == nil {
				for _, n := range subRE.FindAllString(string(resp), -1) {
					if name := CleanName(n); whichDomain(name, scope) != "" {
						m.Lock()
						results.Insert(name)
						m.Unlock()
					}
				}
			}

			processURL := func(u *url.URL) {
				// Attempt to save the name in our results
				m.Lock()
				results.Insert(u.Hostname())
				m.Unlock()

				if s := crawlFilterURLs(u, f); s != "" {
					// Be sure the crawl has not exceeded the maximum links to be followed
					m.Lock()
					count++
					current := count
					m.Unlock()
					if max <= 0 || current < max {
						g.Get(s, g.Opt.ParseFunc)
					}
				}
			}

			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if u, err := r.JoinURL(href); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})

			r.HTMLDoc.Find("script").Each(func(i int, s *goquery.Selection) {
				if src, ok := s.Attr("src"); ok {
					if u, err := r.JoinURL(src); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})
		},
	})
	options := &client.Options{
		MaxBodySize:    100 * 1024 * 1024, // 100MB
		RetryTimes:     2,
		RetryHTTPCodes: []int{408, 500, 502, 503, 504, 522, 524},
	}
	g.Client = client.NewClient(options)
	g.Client.Client = http.DefaultClient

	done := make(chan struct{}, 2)
	go func() {
		g.Start()
		done <- struct{}{}
	}()

	var err error
	select {
	case <-ctx.Done():
		err = fmt.Errorf("The context expired during the crawl of %s", u)
	case <-done:
		if len(results.Slice()) == 0 {
			err = fmt.Errorf("No DNS names were discovered during the crawl of %s", u)
		}
	}

	return results.Slice(), err
}

func crawlFilterURLs(p *url.URL, f filter.Filter) string {
	// Check that the URL has an appropriate scheme for scraping
	if !p.IsAbs() || (p.Scheme != "http" && p.Scheme != "https") {
		return ""
	}
	// If the URL path has a file extension, check that it's of interest
	if ext := crawlRE.FindString(p.Path); ext != "" {
		ext = strings.TrimRight(ext, "?#")

		var found bool
		for _, s := range crawlFileStarts {
			if strings.HasPrefix(ext, "." + s) {
				found = true
				break
			}
		}
		for _, e := range crawlFileEnds {
			if strings.HasSuffix(ext, e) {
				found = true
				break
			}
		}
		if !found {
			return ""
		}
	}
	// Remove fragments and check if we've seen this URL before
	p.Fragment = ""
	p.RawFragment = ""
	if f.Duplicate(p.String()) {
		return ""
	}
	return p.String()
}

func whichDomain(name string, scope []string) string {
	n := strings.TrimSpace(name)

	for _, d := range scope {
		if strings.HasSuffix(n, d) {
			// fork made me do it :>
			nlen := len(n)
			dlen := len(d)
			// Check for exact match first to guard against out of bound index
			if nlen == dlen || n[nlen-dlen-1] == '.' {
				return d
			}
		}
	}
	return ""
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(ctx context.Context, addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		select {
		case <-ctx.Done():
			return names
		default:
		}

		// Set the maximum time allowed for making the connection
		tCtx, cancel := context.WithTimeout(ctx, handshakeTimeout)
		defer cancel()
		// Obtain the connection
		conn, err := amassnet.DialContext(tCtx, "tcp", net.JoinHostPort(addr, strconv.Itoa(port)))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		go func() {
			errChan <- c.Handshake()
		}()

		t := time.NewTimer(handshakeTimeout)
		select {
		case <-t.C:
			err = errors.New("Handshake timeout")
		case e := <-errChan:
			err = e
		}
		t.Stop()

		if err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}

	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode(ctx context.Context) string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage(ctx, "https://ipapi.co/json", nil, headers, nil)
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	if err := json.Unmarshal([]byte(page), &ipinfo); err != nil {
		return ""
	}
	return strings.ToLower(ipinfo.CountryCode)
}

// CleanName will clean up the names scraped from the web.
func CleanName(name string) string {
	var err error

	name, err = strconv.Unquote("\"" + strings.TrimSpace(name) + "\"")
	if err == nil {
		name = subRE.FindString(name)
	}

	name = strings.ToLower(name)
	for {
		name = strings.Trim(name, "-.")

		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	return name
}
<commit_msg>Linting<commit_after>// Copyright 2017-2021 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/filter"
	amassnet "github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/PuerkitoBio/goquery"
	"github.com/caffix/stringset"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

const (
	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	httpTimeout      = 30 * time.Second
	handshakeTimeout = 5 * time.Second
)

var (
	// UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent       string
	subRE           = dns.AnySubdomainRegex()
	crawlRE         = regexp.MustCompile(`\.\w{2,6}($|\?|#)`)
	crawlFileEnds   = []string{"html", "do", "action", "cgi"}
	crawlFileStarts = []string{"js", "htm", "as", "php", "inc"}
	nameStripRE     = regexp.MustCompile(`^u[0-9a-f]{4}|20|22|25|27|2b|2f|3d|3a|40`)
)

// DefaultClient is the same HTTP client used by the package methods.
var DefaultClient *http.Client

// BasicAuth contains the data used for HTTP basic authentication.
type BasicAuth struct {
	Username string
	Password string
}

func init() {
	jar, _ := cookiejar.New(nil)
	DefaultClient = &http.Client{
		Timeout: httpTimeout,
		Transport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           amassnet.DialContext,
			MaxIdleConns:          200,
			MaxConnsPerHost:       50,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   handshakeTimeout,
			ExpectContinueTimeout: 10 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}

	switch runtime.GOOS {
	case "windows":
		UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	case "darwin":
		UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 11_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	default:
		UserAgent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	DefaultClient.Jar.SetCookies(destURL, DefaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range DefaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for the provided URL when successful.
func RequestWebPage(ctx context.Context, u string, body io.Reader, hvals map[string]string, auth *BasicAuth) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequestWithContext(ctx, method, u, body)
	if err != nil {
		return "", err
	}
	if auth != nil && auth.Username != "" && auth.Password != "" {
		req.SetBasicAuth(auth.Username, auth.Password)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)

	for k, v := range hvals {
		req.Header.Set(k, v)
	}

	resp, err := DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		err = errors.New(resp.Status)
	}
	return string(in), err
}

// Crawl will spider the web page at the URL argument looking for DNS names within the scope argument.
func Crawl(ctx context.Context, u string, scope []string, max int, f filter.Filter) ([]string, error) {
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("The context expired")
	default:
	}

	newScope := append([]string{}, scope...)

	target := subRE.FindString(u)
	if target != "" {
		var found bool
		for _, domain := range newScope {
			if target == domain {
				found = true
				break
			}
		}
		if !found {
			newScope = append(newScope, target)
		}
	}

	if f == nil {
		f = filter.NewStringFilter()
	}

	var count int
	var m sync.Mutex
	results := stringset.New()
	g := geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:        newScope,
		StartURLs:             []string{u},
		Timeout:               5 * time.Minute,
		RobotsTxtDisabled:     true,
		UserAgent:             UserAgent,
		LogDisabled:           true,
		ConcurrentRequests:    5,
		RequestDelay:          750 * time.Millisecond,
		RequestDelayRandomize: true,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			resp, err := httputil.DumpResponse(interface{}(r).(*http.Response), true)
			if err == nil {
				for _, n := range subRE.FindAllString(string(resp), -1) {
					if name := CleanName(n); whichDomain(name, scope) != "" {
						m.Lock()
						results.Insert(name)
						m.Unlock()
					}
				}
			}

			processURL := func(u *url.URL) {
				// Attempt to save the name in our results
				m.Lock()
				results.Insert(u.Hostname())
				m.Unlock()

				if s := crawlFilterURLs(u, f); s != "" {
					// Be sure the crawl has not exceeded the maximum links to be followed
					m.Lock()
					count++
					current := count
					m.Unlock()
					if max <= 0 || current < max {
						g.Get(s, g.Opt.ParseFunc)
					}
				}
			}

			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if u, err := r.JoinURL(href); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})

			r.HTMLDoc.Find("script").Each(func(i int, s *goquery.Selection) {
				if src, ok := s.Attr("src"); ok {
					if u, err := r.JoinURL(src); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})
		},
	})
	options := &client.Options{
		MaxBodySize:    100 * 1024 * 1024, // 100MB
		RetryTimes:     2,
		RetryHTTPCodes: []int{408, 500, 502, 503, 504, 522, 524},
	}
	g.Client = client.NewClient(options)
	g.Client.Client = http.DefaultClient

	done := make(chan struct{}, 2)
	go func() {
		g.Start()
		done <- struct{}{}
	}()

	var err error
	select {
	case <-ctx.Done():
		err = fmt.Errorf("The context expired during the crawl of %s", u)
	case <-done:
		if len(results.Slice()) == 0 {
			err = fmt.Errorf("No DNS names were discovered during the crawl of %s", u)
		}
	}

	return results.Slice(), err
}

func crawlFilterURLs(p *url.URL, f filter.Filter) string {
	// Check that the URL has an appropriate scheme for scraping
	if !p.IsAbs() || (p.Scheme != "http" && p.Scheme != "https") {
		return ""
	}
	// If the URL path has a file extension, check that it's of interest
	if ext := crawlRE.FindString(p.Path); ext != "" {
		ext = strings.ToLower(ext)

		var found bool
		for _, s := range crawlFileStarts {
			if strings.HasPrefix(ext, "." + s) {
				found = true
				break
			}
		}
		for _, e := range crawlFileEnds {
			if strings.HasSuffix(ext, e) {
				found = true
				break
			}
		}
		if !found {
			return ""
		}
	}
	// Remove fragments and check if we've seen this URL before
	p.Fragment = ""
	p.RawFragment = ""
	if f.Duplicate(p.String()) {
		return ""
	}
	return p.String()
}

func whichDomain(name string, scope []string) string {
	n := strings.TrimSpace(name)

	for _, d := range scope {
		if strings.HasSuffix(n, d) {
			// fork made me do it :>
			nlen := len(n)
			dlen := len(d)
			// Check for exact match first to guard against out of bound index
			if nlen == dlen || n[nlen-dlen-1] == '.' {
				return d
			}
		}
	}
	return ""
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(ctx context.Context, addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		select {
		case <-ctx.Done():
			return names
		default:
		}

		// Set the maximum time allowed for making the connection
		tCtx, cancel := context.WithTimeout(ctx, handshakeTimeout)
		defer cancel()
		// Obtain the connection
		conn, err := amassnet.DialContext(tCtx, "tcp", net.JoinHostPort(addr, strconv.Itoa(port)))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		go func() {
			errChan <- c.Handshake()
		}()

		t := time.NewTimer(handshakeTimeout)
		select {
		case <-t.C:
			err = errors.New("Handshake timeout")
		case e := <-errChan:
			err = e
		}
		t.Stop()

		if err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}

	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode(ctx context.Context) string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage(ctx, "https://ipapi.co/json", nil, headers, nil)
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	if err := json.Unmarshal([]byte(page), &ipinfo); err != nil {
		return ""
	}
	return strings.ToLower(ipinfo.CountryCode)
}

// CleanName will clean up the names scraped from the web.
func CleanName(name string) string {
	var err error

	name, err = strconv.Unquote("\"" + strings.TrimSpace(name) + "\"")
	if err == nil {
		name = subRE.FindString(name)
	}

	name = strings.ToLower(name)
	for {
		name = strings.Trim(name, "-.")

		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	return name
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/franela/goblin"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"net/http"
	"testing"
)

type FakeResponseBodyGenerator struct {
	use string
}

func (instance *FakeResponseBodyGenerator) Generate() string {
	return instance.use
}

func (instance *FakeResponseBodyGenerator) UseString(value string) {
	instance.use = value
}

func NewFakeResponseBodyGenerator() *FakeResponseBodyGenerator {
	return &FakeResponseBodyGenerator{""}
}

func Test_Enanos_Without_Goblin(t *testing.T) {
	var fakeResponseBodyGenerator *FakeResponseBodyGenerator = NewFakeResponseBodyGenerator()
	var enanosHttpHandlerFactory *DefaultEnanosHttpHandlerFactory = NewDefaultEnanosHttpHandlerFactory(fakeResponseBodyGenerator)

	url := func(path string) (fullPath string) {
		fullPath = "http://localhost:8000" + path
		return
	}
	go func() {
		StartEnanos(fakeResponseBodyGenerator, enanosHttpHandlerFactory)
	}()
	sample := "foobar"
	fakeResponseBodyGenerator.UseString(sample)
	resp, _ := http.Get(url("/default/sneezy"))
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)
	assert.Equal(t, http.StatusOK, resp.StatusCode)
	assert.Equal(t, sample, string(body))
}

func Test_ResponseBodyGenerator(t *testing.T) {

	g := goblin.Goblin(t)
	g.Describe("Default Response Body Generator", func() {
		g.It("generates a string of the defined lenth", func() {
			maxLength := 5
			generator := NewDefaultResponseBodyGenerator(maxLength)
			value := generator.Generate()
			assert.Equal(t, maxLength, len(value))
		})
	})

	g.Describe("Random Response Body Generator", func() {
		g.It("generates a string of length between the defined min length and the defined max length", func() {
			minLength := 50
			maxLength := 500
			generator := NewRandomResponseBodyGenerator(maxLength, minLength)
			value := generator.Generate()
			assert.True(t, len(value) >= minLength && len(value) <= maxLength)
		})
	})
}

func Test_Enanos(t *testing.T) {
	g := goblin.Goblin(t)
	g.Describe("Enanos Server:", func() {

		var fakeResponseBodyGenerator *FakeResponseBodyGenerator
		var enanosHttpHandlerFactory *DefaultEnanosHttpHandlerFactory

		url := func(path string) (fullPath string) {
			fullPath = "http://localhost:8000" + path
			return
		}

		g.BeforeEach(func() {
			fakeResponseBodyGenerator = NewFakeResponseBodyGenerator()
			enanosHttpHandlerFactory = NewDefaultEnanosHttpHandlerFactory(fakeResponseBodyGenerator)
			go func() {
				StartEnanos(fakeResponseBodyGenerator, enanosHttpHandlerFactory)
			}()
		})

		g.Describe("Happy :", func() {
			g.It("GET returns 200", func() {
				resp, _ := http.Get(url("/default/happy"))
				assert.Equal(t, http.StatusOK, resp.StatusCode)
			})
		})

		g.Describe("Grumpy :", func() {
			g.It("GET returns 500", func() {
				resp, _ := http.Get(url("/default/grumpy"))
				assert.Equal(t, http.StatusInternalServerError, resp.StatusCode)
			})
		})

		g.Describe("Sneezy :", func() {
			g.It("GET returns 200", func() {
				resp, _ := http.Get(url("/default/sneezy"))
				assert.Equal(t, http.StatusOK, resp.StatusCode)
			})

			g.It("GET returns random response body", func() {
				sample := "foobar"
				fakeResponseBodyGenerator.UseString(sample)
				resp, _ := http.Get(url("/default/sneezy"))
				defer resp.Body.Close()
				body, _ := ioutil.ReadAll(resp.Body)
				assert.Equal(t, sample, string(body))
			})
		})
	})
}
<commit_msg>Some wierdness around setting up variables and go sub routines with goblin or my understanding.  For now I am just using the testing TestMain method<commit_after>package main

import (
	"github.com/franela/goblin"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"net/http"
	"os"
	"testing"
)

type FakeResponseBodyGenerator struct {
	use string
}

func (instance *FakeResponseBodyGenerator) Generate() string {
	return instance.use
}

func (instance *FakeResponseBodyGenerator) UseString(value string) {
	instance.use = value
}

func NewFakeResponseBodyGenerator() *FakeResponseBodyGenerator {
	return &FakeResponseBodyGenerator{""}
}

var fakeResponseBodyGenerator *FakeResponseBodyGenerator
var enanosHttpHandlerFactory *DefaultEnanosHttpHandlerFactory

func TestMain(m *testing.M) {
	fakeResponseBodyGenerator = NewFakeResponseBodyGenerator()
	enanosHttpHandlerFactory = NewDefaultEnanosHttpHandlerFactory(fakeResponseBodyGenerator)
	go func() {
		StartEnanos(fakeResponseBodyGenerator, enanosHttpHandlerFactory)
	}()
	os.Exit(m.Run())
}

func Test_ResponseBodyGenerator(t *testing.T) {
	g := goblin.Goblin(t)
	g.Describe("Default Response Body Generator", func() {
		g.It("generates a string of the defined lenth", func() {
			maxLength := 5
			generator := NewDefaultResponseBodyGenerator(maxLength)
			value := generator.Generate()
			assert.Equal(t, maxLength, len(value))
		})
	})

	g.Describe("Random Response Body Generator", func() {
		g.It("generates a string of length between the defined min length and the defined max length", func() {
			minLength := 50
			maxLength := 500
			generator := NewRandomResponseBodyGenerator(minLength, maxLength)
			value := generator.Generate()
			assert.True(t, len(value) >= minLength && len(value) <= maxLength)
		})
	})
}

func Test_Enanos(t *testing.T) {
	g := goblin.Goblin(t)
	g.Describe("Enanos Server:", func() {

		url := func(path string) (fullPath string) {
			fullPath = "http://localhost:8000" + path
			return
		}

		g.Describe("Happy :", func() {
			g.It("GET returns 200", func() {
				resp, _ := http.Get(url("/default/happy"))
				assert.Equal(t, http.StatusOK, resp.StatusCode)
			})
		})

		g.Describe("Grumpy :", func() {
			g.It("GET returns 500", func() {
				resp, _ := http.Get(url("/default/grumpy"))
				assert.Equal(t, http.StatusInternalServerError, resp.StatusCode)
			})
		})

		g.Describe("Sneezy :", func() {
			g.It("GET returns 200", func() {
				resp, _ := http.Get(url("/default/sneezy"))
				assert.Equal(t, http.StatusOK, resp.StatusCode)
			})

			g.It("GET returns random response body", func() {
				sample := "foobar"
				fakeResponseBodyGenerator.UseString(sample)
				resp, _ := http.Get(url("/default/sneezy"))
				defer resp.Body.Close()
				body, _ := ioutil.ReadAll(resp.Body)
				assert.Equal(t, sample, string(body))
			})
		})
	})
}
<|endoftext|>


<commit_before>package smtp

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"math/rand"
	"net/mail"
	"strings"
	"time"

	"github.com/ije/gox/valid"
)

var CRLF = []byte("\r\n")

type Mail struct {
	Subject     string
	PlainText   []byte
	Html        []byte
	Attachments []*Attachment
}

type Attachment struct {
	Name        string
	ContentType string
	io.Reader
}

func (mail *Mail) MakeBody(from *mail.Address, to AddressList) []byte {
	buf := &buffer{bytes.NewBuffer(nil)}
	buf.writeln("MIME-Version: 1.0")
	buf.writeln("Date: ", time.Now().Format(time.RFC1123Z))
	buf.writeln("Subject: ", encodeSubject(mail.Subject))
	buf.writeln("From: ", from)
	buf.writeln("To: ", to)
	var boundary string
	if len(mail.Attachments) > 0 {
		boundary = boundaryGen()
		buf.writeln("Content-Type: multipart/mixed; boundary=", boundary)
		buf.writeln()
		buf.writeln("--", boundary)
	}
	if len(mail.PlainText) > 0 && len(mail.Html) > 0 {
		cBoundary := boundaryGen()
		buf.writeln("Content-Type: multipart/alternative; boundary=", cBoundary)
		buf.writeln()
		buf.writeln("--", cBoundary)
		buf.writeTextBody(mail.PlainText)
		buf.writeln()
		buf.writeln()
		buf.writeln("--", cBoundary)
		buf.writeHtmlBody(mail.Html)
		buf.writeln()
		buf.writeln()
		buf.writeln("--", cBoundary, "--")
	} else if len(mail.PlainText) > 0 {
		buf.writeTextBody(mail.PlainText)
	} else if len(mail.Html) > 0 {
		buf.writeHtmlBody(mail.Html)
	}
	if len(mail.Attachments) > 0 {
		for _, attchment := range mail.Attachments {
			buf.writeln()
			buf.writeln()
			buf.writeln("--", boundary)
			buf.writeln("Content-Type: ", attchment.ContentType, "; name=", attchment.Name, ";")
			buf.writeln("Content-Transfer-Encoding: base64")
			buf.writeln("Content-Disposition: attachment; filename=", attchment.Name, ";")
			buf.writeln()
			encoder := base64.NewEncoder(base64.StdEncoding, buf)
			io.Copy(encoder, attchment)
			encoder.Close()
		}
		buf.writeln()
		buf.writeln()
		buf.writeln("--", boundary, "--")
	}
	return buf.Bytes()
}

func Address(a ...string) *mail.Address {
	var name string
	var address string
	if len(a) == 1 {
		address = a[0]
	} else if len(a) > 1 {
		name = a[0]
		address = a[1]
	}

	if len(address) == 0 || !valid.IsEmail(address) {
		return nil
	}

	return &mail.Address{
		Name:    name,
		Address: address,
	}
}

type AddressList []*mail.Address

func (list AddressList) String() string {
	var ss []string
	for _, addr := range list {
		ss = append(ss, addr.String())
	}
	return strings.Join(ss, ", ")
}

func (list AddressList) List() []string {
	var ss []string
	for _, addr := range list {
		ss = append(ss, addr.Address)
	}
	return ss
}

type buffer struct {
	*bytes.Buffer
}

func (buf *buffer) writeln(s ...interface{}) {
	fmt.Fprint(buf, s...)
	buf.Write(CRLF)
}

func (buf *buffer) writeTextBody(text []byte) {
	buf.writeln("Content-Type: text/plain; charset=UTF-8")

	for _, c := range text {
		if c > 127 {
			buf.writeln("Content-Transfer-Encoding: base64")
			buf.writeln()
			buf.WriteString(base64.StdEncoding.EncodeToString(text))
			return
		}
	}

	buf.writeln()
	buf.Write(text)
}

func (buf *buffer) writeHtmlBody(html []byte) {
	buf.writeln("Content-Type: text/html; charset=UTF-8")

	for _, c := range html {
		if c > 127 {
			buf.writeln("Content-Transfer-Encoding: quoted-printable")
			buf.writeln()
			var c byte
			for i, l := 0, len(html); i < l; i++ {
				if c = html[i]; c > 127 {
					fmt.Fprintf(buf, "=%X", c)
					i++
					fmt.Fprintf(buf, "=%X", html[i])
					i++
					fmt.Fprintf(buf, "=%X", html[i])
				} else if c == '=' {
					buf.WriteString("=3D")
				} else {
					buf.WriteByte(c)
				}
			}
			return
		}
	}

	buf.writeln()
	buf.Write(html)
}

func encodeSubject(subject string) string {
	for _, c := range subject {
		if c > 127 {
			return fmt.Sprintf("=?UTF-8?B?%s?=", base64.StdEncoding.EncodeToString([]byte(subject)))
		}
	}
	return subject
}

func boundaryGen() string {
	h := md5.New()
	fmt.Fprint(h, time.Now().UnixNano(), rand.Int())
	return fmt.Sprintf("--%s--", hex.EncodeToString(h.Sum(nil)))
}
<commit_msg>improve the Address function<commit_after>package smtp

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"math/rand"
	"net/mail"
	"strings"
	"time"

	"github.com/ije/gox/valid"
)

var CRLF = []byte("\r\n")

type Mail struct {
	Subject     string
	PlainText   []byte
	Html        []byte
	Attachments []*Attachment
}

type Attachment struct {
	Name        string
	ContentType string
	io.Reader
}

func (mail *Mail) MakeBody(from *mail.Address, to AddressList) []byte {
	buf := &buffer{bytes.NewBuffer(nil)}
	buf.writeln("MIME-Version: 1.0")
	buf.writeln("Date: ", time.Now().Format(time.RFC1123Z))
	buf.writeln("Subject: ", encodeSubject(mail.Subject))
	buf.writeln("From: ", from)
	buf.writeln("To: ", to)
	var boundary string
	if len(mail.Attachments) > 0 {
		boundary = boundaryGen()
		buf.writeln("Content-Type: multipart/mixed; boundary=", boundary)
		buf.writeln()
		buf.writeln("--", boundary)
	}
	if len(mail.PlainText) > 0 && len(mail.Html) > 0 {
		cBoundary := boundaryGen()
		buf.writeln("Content-Type: multipart/alternative; boundary=", cBoundary)
		buf.writeln()
		buf.writeln("--", cBoundary)
		buf.writeTextBody(mail.PlainText)
		buf.writeln()
		buf.writeln()
		buf.writeln("--", cBoundary)
		buf.writeHtmlBody(mail.Html)
		buf.writeln()
		buf.writeln()
		buf.writeln("--", cBoundary, "--")
	} else if len(mail.PlainText) > 0 {
		buf.writeTextBody(mail.PlainText)
	} else if len(mail.Html) > 0 {
		buf.writeHtmlBody(mail.Html)
	}
	if len(mail.Attachments) > 0 {
		for _, attchment := range mail.Attachments {
			buf.writeln()
			buf.writeln()
			buf.writeln("--", boundary)
			buf.writeln("Content-Type: ", attchment.ContentType, "; name=", attchment.Name, ";")
			buf.writeln("Content-Transfer-Encoding: base64")
			buf.writeln("Content-Disposition: attachment; filename=", attchment.Name, ";")
			buf.writeln()
			encoder := base64.NewEncoder(base64.StdEncoding, buf)
			io.Copy(encoder, attchment)
			encoder.Close()
		}
		buf.writeln()
		buf.writeln()
		buf.writeln("--", boundary, "--")
	}
	return buf.Bytes()
}

func Address(a ...string) *mail.Address {
	var name string
	var address string
	if len(a) == 1 {
		addr, err := mail.ParseAddress(a[0])
		if err == nil {
			return addr
		}
	} else if len(a) > 1 {
		name = a[0]
		address = a[1]
	}

	if len(address) == 0 || !valid.IsEmail(address) {
		return nil
	}

	return &mail.Address{
		Name:    name,
		Address: address,
	}
}

type AddressList []*mail.Address

func (list AddressList) String() string {
	var ss []string
	for _, addr := range list {
		ss = append(ss, addr.String())
	}
	return strings.Join(ss, ", ")
}

func (list AddressList) List() []string {
	var ss []string
	for _, addr := range list {
		ss = append(ss, addr.Address)
	}
	return ss
}

type buffer struct {
	*bytes.Buffer
}

func (buf *buffer) writeln(s ...interface{}) {
	fmt.Fprint(buf, s...)
	buf.Write(CRLF)
}

func (buf *buffer) writeTextBody(text []byte) {
	buf.writeln("Content-Type: text/plain; charset=UTF-8")

	for _, c := range text {
		if c > 127 {
			buf.writeln("Content-Transfer-Encoding: base64")
			buf.writeln()
			buf.WriteString(base64.StdEncoding.EncodeToString(text))
			return
		}
	}

	buf.writeln()
	buf.Write(text)
}

func (buf *buffer) writeHtmlBody(html []byte) {
	buf.writeln("Content-Type: text/html; charset=UTF-8")

	for _, c := range html {
		if c > 127 {
			buf.writeln("Content-Transfer-Encoding: quoted-printable")
			buf.writeln()
			var c byte
			for i, l := 0, len(html); i < l; i++ {
				if c = html[i]; c > 127 {
					fmt.Fprintf(buf, "=%X", c)
					i++
					fmt.Fprintf(buf, "=%X", html[i])
					i++
					fmt.Fprintf(buf, "=%X", html[i])
				} else if c == '=' {
					buf.WriteString("=3D")
				} else {
					buf.WriteByte(c)
				}
			}
			return
		}
	}

	buf.writeln()
	buf.Write(html)
}

func encodeSubject(subject string) string {
	for _, c := range subject {
		if c > 127 {
			return fmt.Sprintf("=?UTF-8?B?%s?=", base64.StdEncoding.EncodeToString([]byte(subject)))
		}
	}
	return subject
}

func boundaryGen() string {
	h := md5.New()
	fmt.Fprint(h, time.Now().UnixNano(), rand.Int())
	return fmt.Sprintf("--%s--", hex.EncodeToString(h.Sum(nil)))
}
<|endoftext|>