<commit_before>package youtube

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

func SetLogOutput(w io.Writer) {
	log.SetOutput(w)
}

func NewYoutube(debug bool) *Youtube {
	return &Youtube{DebugMode: debug, DownloadPercent: make(chan int64, 100)}
}

type stream map[string]string

type Youtube struct {
	DebugMode         bool
	StreamList        []stream
	VideoID           string
	videoInfo         string
	DownloadPercent   chan int64
	contentLength     float64
	totalWrittenBytes float64
	downloadLevel     float64
}

func (y *Youtube) DecodeURL(url string) error {
	err := y.findVideoID(url)
	if err != nil {
		return fmt.Errorf("findVideoID error=%s", err)
	}

	err = y.getVideoInfo()
	if err != nil {
		return fmt.Errorf("getVideoInfo error=%s", err)
	}

	err = y.parseVideoInfo()
	if err != nil {
		return fmt.Errorf("parse video info failed, err=%s", err)
	}

	return nil
}

func (y *Youtube) StartDownload(destFile string) error {
	//download highest resolution on [0]
	targetStream := y.StreamList[0]
	url := targetStream["url"] + "&signature=" + targetStream["sig"]
	y.log(fmt.Sprintln("Download url=", url))

	y.log(fmt.Sprintln("Download to file=", destFile))
	err := y.videoDLWorker(destFile, url)
	return err
}

func (y *Youtube) parseVideoInfo() error {
	answer, err := url.ParseQuery(y.videoInfo)
	if err != nil {
		return err
	}

	status, ok := answer["status"]
	if !ok {
		err = fmt.Errorf("no response status found in the server's answer")
		return err
	}
	if status[0] == "fail" {
		reason, ok := answer["reason"]
		if ok {
			err = fmt.Errorf("'fail' response status found in the server's answer, reason: '%s'", reason[0])
		} else {
			err = errors.New(fmt.Sprint("'fail' response status found in the server's answer, no reason given"))
		}
		return err
	}
	if status[0] != "ok" {
		err = fmt.Errorf("non-success response status found in the server's answer (status: '%s')", status)
		return err
	}

	// read the streams map
	streamMap, ok := answer["url_encoded_fmt_streamMap"]
	if !ok {
		err = errors.New(fmt.Sprint("no stream map found in the server's answer"))
		return err
	}

	// read each stream
	streamsList := strings.Split(streamMap[0], ",")

	var streams []stream
	for streamPos, streamRaw := range streamsList {
		streamQry, err := url.ParseQuery(streamRaw)
		if err != nil {
			log.Printf("An error occured while decoding one of the video's stream's information: stream %d: %s\n", streamPos, err)
			continue
		}
		var sig string
		if _, exist := streamQry["sig"]; exist {
			sig = streamQry["sig"][0]
		}

		streams = append(streams, stream{
			"quality": streamQry["quality"][0],
			"type":    streamQry["type"][0],
			"url":     streamQry["url"][0],
			"sig":     sig,
			"title":   answer["title"][0],
			"author":  answer["author"][0],
		})
		y.log(fmt.Sprintf("Stream found: quality '%s', format '%s'", streamQry["quality"][0], streamQry["type"][0]))
	}

	y.StreamList = streams
	return nil
}

func (y *Youtube) getVideoInfo() error {
	url := "http://youtube.com/get_video_info?video_id=" + y.VideoID
	y.log(fmt.Sprintf("url: %s", url))
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	y.videoInfo = string(body)
	return nil
}

func (y *Youtube) findVideoID(url string) error {
	videoID := url
	if strings.Contains(videoID, "youtu") || strings.ContainsAny(videoID, "\"?&/<%=") {
		reList := []*regexp.Regexp{
			regexp.MustCompile(`(?:v|embed|watch\?v)(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`([^"&?/=%]{11})`),
		}
		for _, re := range reList {
			if isMatch := re.MatchString(videoID); isMatch {
				subs := re.FindStringSubmatch(videoID)
				videoID = subs[1]
			}
		}
	}
	log.Printf("Found video id: '%s'", videoID)
	y.VideoID = videoID
	if strings.ContainsAny(videoID, "?&/<%=") {
		return errors.New("invalid characters in video id")
	}
	if len(videoID) < 10 {
		return errors.New("the video id must be at least 10 characters long")
	}
	return nil
}

func (y *Youtube) Write(p []byte) (n int, err error) {
	n = len(p)
	y.totalWrittenBytes = y.totalWrittenBytes + float64(n)
	currentPercent := ((y.totalWrittenBytes / y.contentLength) * 100)
	if (y.downloadLevel <= currentPercent) && (y.downloadLevel < 100) {
		y.downloadLevel++
		y.DownloadPercent <- int64(y.downloadLevel)
	}
	return
}
func (y *Youtube) videoDLWorker(destFile string, target string) error {
	resp, err := http.Get(target)
	if err != nil {
		log.Printf("Http.Get\nerror: %s\ntarget: %s\n", err, target)
		return err
	}
	defer resp.Body.Close()
	y.contentLength = float64(resp.ContentLength)

	if resp.StatusCode != 200 {
		log.Printf("reading answer: non 200[code=%v] status code received: '%v'", resp.StatusCode, err)
		return errors.New("non 200 status code received")
	}
	err = os.MkdirAll(filepath.Dir(destFile), 666)
	if err != nil {
		return err
	}
	out, err := os.Create(destFile)
	if err != nil {
		return err
	}
	mw := io.MultiWriter(out, y)
	_, err = io.Copy(mw, resp.Body)
	if err != nil {
		log.Println("download video err=", err)
		return err
	}
	return nil
}

func (y *Youtube) log(logText string) {
	if y.DebugMode {
		log.Println(logText)
	}
}
<commit_msg>youtube change the url var name to 'url_encoded_fmt_stream_map'<commit_after>package youtube

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

func SetLogOutput(w io.Writer) {
	log.SetOutput(w)
}

func NewYoutube(debug bool) *Youtube {
	return &Youtube{DebugMode: debug, DownloadPercent: make(chan int64, 100)}
}

type stream map[string]string

type Youtube struct {
	DebugMode         bool
	StreamList        []stream
	VideoID           string
	videoInfo         string
	DownloadPercent   chan int64
	contentLength     float64
	totalWrittenBytes float64
	downloadLevel     float64
}

func (y *Youtube) DecodeURL(url string) error {
	err := y.findVideoID(url)
	if err != nil {
		return fmt.Errorf("findVideoID error=%s", err)
	}

	err = y.getVideoInfo()
	if err != nil {
		return fmt.Errorf("getVideoInfo error=%s", err)
	}

	err = y.parseVideoInfo()
	if err != nil {
		return fmt.Errorf("parse video info failed, err=%s", err)
	}

	return nil
}

func (y *Youtube) StartDownload(destFile string) error {
	//download highest resolution on [0]
	targetStream := y.StreamList[0]
	url := targetStream["url"] + "&signature=" + targetStream["sig"]
	y.log(fmt.Sprintln("Download url=", url))

	y.log(fmt.Sprintln("Download to file=", destFile))
	err := y.videoDLWorker(destFile, url)
	return err
}

func (y *Youtube) parseVideoInfo() error {
	answer, err := url.ParseQuery(y.videoInfo)
	if err != nil {
		return err
	}

	status, ok := answer["status"]
	if !ok {
		err = fmt.Errorf("no response status found in the server's answer")
		return err
	}
	if status[0] == "fail" {
		reason, ok := answer["reason"]
		if ok {
			err = fmt.Errorf("'fail' response status found in the server's answer, reason: '%s'", reason[0])
		} else {
			err = errors.New(fmt.Sprint("'fail' response status found in the server's answer, no reason given"))
		}
		return err
	}
	if status[0] != "ok" {
		err = fmt.Errorf("non-success response status found in the server's answer (status: '%s')", status)
		return err
	}

	// read the streams map
	streamMap, ok := answer["url_encoded_fmt_stream_map"]
	if !ok {
		err = errors.New(fmt.Sprint("no stream map found in the server's answer"))
		return err
	}

	// read each stream
	streamsList := strings.Split(streamMap[0], ",")

	var streams []stream
	for streamPos, streamRaw := range streamsList {
		streamQry, err := url.ParseQuery(streamRaw)
		if err != nil {
			log.Printf("An error occured while decoding one of the video's stream's information: stream %d: %s\n", streamPos, err)
			continue
		}
		var sig string
		if _, exist := streamQry["sig"]; exist {
			sig = streamQry["sig"][0]
		}

		streams = append(streams, stream{
			"quality": streamQry["quality"][0],
			"type":    streamQry["type"][0],
			"url":     streamQry["url"][0],
			"sig":     sig,
			"title":   answer["title"][0],
			"author":  answer["author"][0],
		})
		y.log(fmt.Sprintf("Stream found: quality '%s', format '%s'", streamQry["quality"][0], streamQry["type"][0]))
	}

	y.StreamList = streams
	return nil
}

func (y *Youtube) getVideoInfo() error {
	url := "http://youtube.com/get_video_info?video_id=" + y.VideoID
	y.log(fmt.Sprintf("url: %s", url))
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	y.videoInfo = string(body)
	return nil
}

func (y *Youtube) findVideoID(url string) error {
	videoID := url
	if strings.Contains(videoID, "youtu") || strings.ContainsAny(videoID, "\"?&/<%=") {
		reList := []*regexp.Regexp{
			regexp.MustCompile(`(?:v|embed|watch\?v)(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`([^"&?/=%]{11})`),
		}
		for _, re := range reList {
			if isMatch := re.MatchString(videoID); isMatch {
				subs := re.FindStringSubmatch(videoID)
				videoID = subs[1]
			}
		}
	}
	log.Printf("Found video id: '%s'", videoID)
	y.VideoID = videoID
	if strings.ContainsAny(videoID, "?&/<%=") {
		return errors.New("invalid characters in video id")
	}
	if len(videoID) < 10 {
		return errors.New("the video id must be at least 10 characters long")
	}
	return nil
}

func (y *Youtube) Write(p []byte) (n int, err error) {
	n = len(p)
	y.totalWrittenBytes = y.totalWrittenBytes + float64(n)
	currentPercent := ((y.totalWrittenBytes / y.contentLength) * 100)
	if (y.downloadLevel <= currentPercent) && (y.downloadLevel < 100) {
		y.downloadLevel++
		y.DownloadPercent <- int64(y.downloadLevel)
	}
	return
}
func (y *Youtube) videoDLWorker(destFile string, target string) error {
	resp, err := http.Get(target)
	if err != nil {
		log.Printf("Http.Get\nerror: %s\ntarget: %s\n", err, target)
		return err
	}
	defer resp.Body.Close()
	y.contentLength = float64(resp.ContentLength)

	if resp.StatusCode != 200 {
		log.Printf("reading answer: non 200[code=%v] status code received: '%v'", resp.StatusCode, err)
		return errors.New("non 200 status code received")
	}
	err = os.MkdirAll(filepath.Dir(destFile), 666)
	if err != nil {
		return err
	}
	out, err := os.Create(destFile)
	if err != nil {
		return err
	}
	mw := io.MultiWriter(out, y)
	_, err = io.Copy(mw, resp.Body)
	if err != nil {
		log.Println("download video err=", err)
		return err
	}
	return nil
}

func (y *Youtube) log(logText string) {
	if y.DebugMode {
		log.Println(logText)
	}
}
<|endoftext|>


<commit_before>package youtube

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"golang.org/x/net/proxy"
)

//SetLogOutput :Set logger writer
func SetLogOutput(w io.Writer) {
	log.SetOutput(w)
}

//NewYoutube :Initialize youtube package object
func NewYoutube(debug bool) *Youtube {
	return &Youtube{DebugMode: debug, DownloadPercent: make(chan int64, 100)}
}

func NewYoutubeWithSocks5Proxy(debug bool, socks5Proxy string) *Youtube {
	return &Youtube{DebugMode: debug, DownloadPercent: make(chan int64, 100), Socks5Proxy: socks5Proxy}
}

type stream map[string]string

type Youtube struct {
	DebugMode         bool
	StreamList        []stream
	VideoID           string
	videoInfo         string
	DownloadPercent   chan int64
	Socks5Proxy       string
	contentLength     float64
	totalWrittenBytes float64
	downloadLevel     float64
}

//DecodeURL : Decode youtube URL to retrieval video information.
func (y *Youtube) DecodeURL(url string) error {
	err := y.findVideoID(url)
	if err != nil {
		return fmt.Errorf("findVideoID error=%s", err)
	}

	err = y.getVideoInfo()
	if err != nil {
		return fmt.Errorf("getVideoInfo error=%s", err)
	}

	err = y.parseVideoInfo()
	if err != nil {
		return fmt.Errorf("parse video info failed, err=%s", err)
	}

	return nil
}

//StartDownload : Starting download video to specific address.
func (y *Youtube) StartDownload(destFile string) error {
	//download highest resolution on [0]
	err := errors.New("Empty stream list")
	for _, v := range y.StreamList {
		url := v["url"]
		y.log(fmt.Sprintln("Download url=", url))

		y.log(fmt.Sprintln("Download to file=", destFile))
		err = y.videoDLWorker(destFile, url)
		if err != nil {
			y.log(fmt.Sprintln("Err: ", err.Error()))
			return err
		}
	}
	return err
}

//StartDownloadWithQuality : Starting download video with specific quality.
func (y *Youtube) StartDownloadWithQuality(destFile string, quality string) error {
	//download highest resolution on [0]
	err := errors.New("Empty stream list")
	for _, v := range y.StreamList {
		if strings.Compare(v["quality"], quality) == 0 {
			url := v["url"]
			y.log(fmt.Sprintln("Download url=", url))
			y.log(fmt.Sprintln("Download to file=", destFile))
			err = y.videoDLWorker(destFile, url)
			if err == nil {
				break
			}
		}
	}

	if err != nil {
		return y.StartDownload(destFile)
	}
	return err
}

//StartDownloadFile : Starting download video on my download.
func (y *Youtube) StartDownloadFile() error {
	//download highest resolution on [0]
	err := errors.New("Empty stream list")
	for _, stream := range y.StreamList {
		streamUrl := stream["url"]
		y.log(fmt.Sprintln("Download url=", streamUrl))

		// Find out what the file name should be.
		fileName := sanitizeFilename(stream["title"])
		fileName += pickIdealFileExtension(stream["type"])

		usr, _ := user.Current()
		destFile := filepath.Join(filepath.Join(usr.HomeDir, "Movies", "youtubedr"), fileName)
		y.log(fmt.Sprintln("Download to file=", destFile))

		err = y.videoDLWorker(destFile, streamUrl)
		if err == nil {
			return nil
		}
	}
	return err
}

func pickIdealFileExtension(mediaType string) string {
	defaultExtension := ".mov"

	mediaType, _, err := mime.ParseMediaType(mediaType)
	if err != nil {
		return defaultExtension
	}

	// Rely on hardcoded canonical mime types, as the ones provided by Go aren't exhaustive [1].
	// This seems to be a recurring problem for youtube downloaders, see [2].
	// The implementation is based on mozilla's list [3], IANA [4] and Youtube's support [5].
	// [1] https://github.com/golang/go/blob/ed7888aea6021e25b0ea58bcad3f26da2b139432/src/mime/type.go#L60
	// [2] https://github.com/ZiTAL/youtube-dl/blob/master/mime.types
	// [3] https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/MIME_types/Common_types
	// [4] https://www.iana.org/assignments/media-types/media-types.xhtml#video
	// [5] https://support.google.com/youtube/troubleshooter/2888402?hl=en
	canonicals := map[string]string{
		"video/quicktime":  ".mov",
		"video/x-msvideo":  ".avi",
		"video/x-matroska": ".mkv",
		"video/mpeg":       ".mpeg",
		"video/webm":       ".webm",
		"video/3gpp2":      ".3g2",
		"video/x-flv":      ".flv",
		"video/3gpp":       ".3gp",
		"video/mp4":        ".mp4",
		"video/ogg":        ".ogv",
		"video/mp2t":       ".ts",
	}

	if extension, ok := canonicals[mediaType]; ok {
		return extension
	}

	// Our last resort is to ask the operating system, but these give multiple results and are rarely canonical.
	extensions, err := mime.ExtensionsByType(mediaType)
	if err != nil || extensions == nil {
		return defaultExtension
	}

	return extensions[0]
}

func sanitizeFilename(fileName string) string {
	// Characters not allowed on mac
	//	:/
	// Characters not allowed on linux
	//	/
	// Characters not allowed on windows
	//	<>:"/\|?*

	// Ref https://docs.microsoft.com/en-us/windows/win32/fileio/naming-a-file#naming-conventions

	fileName = regexp.MustCompile(`[:/<>\:"\\|?*]`).ReplaceAllString(fileName, "")
	fileName = regexp.MustCompile(`\s+`).ReplaceAllString(fileName, " ")

	return fileName
}

func (y *Youtube) parseVideoInfo() error {
	answer, err := url.ParseQuery(y.videoInfo)
	if err != nil {
		return err
	}

	status, ok := answer["status"]
	if !ok {
		err = fmt.Errorf("no response status found in the server's answer")
		return err
	}
	if status[0] == "fail" {
		reason, ok := answer["reason"]
		if ok {
			err = fmt.Errorf("'fail' response status found in the server's answer, reason: '%s'", reason[0])
		} else {
			err = errors.New(fmt.Sprint("'fail' response status found in the server's answer, no reason given"))
		}
		return err
	}
	if status[0] != "ok" {
		err = fmt.Errorf("non-success response status found in the server's answer (status: '%s')", status)
		return err
	}

	// read the streams map
	streamMap, ok := answer["player_response"]
	if !ok {
		err = errors.New(fmt.Sprint("no stream map found in the server's answer."))
		return err
	}

	// Get video title and author.
	title, author := getVideoTitleAuthor(answer)

	var prData PlayerResponseData
	if err := json.Unmarshal([]byte(streamMap[0]), &prData); err != nil {
		fmt.Println(err)
		panic("Player response json data has changed.")
	}

	// Get video download link
	if prData.PlayabilityStatus.Status == "UNPLAYABLE" {
		//Cannot playback on embedded video screen, could not download.
		return errors.New(fmt.Sprint("Cannot playback and download, reason:", prData.PlayabilityStatus.Reason))
	}

	var streams []stream
	for streamPos, streamRaw := range prData.StreamingData.Formats {
		if streamRaw.MimeType == "" {
			y.log(fmt.Sprintf("An error occurred while decoding one of the video's stream's information: stream %d.\n", streamPos))
			continue
		}
		streamUrl := streamRaw.URL
		if streamUrl == "" {
			cipher := streamRaw.Cipher
			decipheredUrl, err := y.decipher(cipher)
			if err != nil {
				return err
			}
			streamUrl = decipheredUrl
		}

		streams = append(streams, stream{
			"quality": streamRaw.Quality,
			"type":    streamRaw.MimeType,
			"url":     streamUrl,

			"title":  title,
			"author": author,
		})
		y.log(fmt.Sprintf("Title: %s Author: %s Stream found: quality '%s', format '%s'", title, author, streamRaw.Quality, streamRaw.MimeType))
	}

	y.StreamList = streams
	if len(y.StreamList) == 0 {
		return errors.New(fmt.Sprint("no stream list found in the server's answer"))
	}
	return nil
}

func (y *Youtube) getHTTPClient() (*http.Client, error) {
	// setup a http client
	httpTransport := &http.Transport{
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		IdleConnTimeout:       60 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	httpClient := &http.Client{Transport: httpTransport}

	if len(y.Socks5Proxy) == 0 {
		y.log("Using http without proxy.")
		return httpClient, nil
	}

	dialer, err := proxy.SOCKS5("tcp", y.Socks5Proxy, nil, proxy.Direct)
	if err != nil {
		fmt.Fprintln(os.Stderr, "can't connect to the proxy:", err)
		return nil, err
	}
	// set our socks5 as the dialer
	dc := dialer.(interface {
		DialContext(ctx context.Context, network, addr string) (net.Conn, error)
	})
	httpTransport.DialContext = dc.DialContext

	y.log(fmt.Sprintf("Using http with proxy %s.", y.Socks5Proxy))

	return httpClient, nil
}

func (y *Youtube) getVideoInfo() error {
	eurl := "https://youtube.googleapis.com/v/" + y.VideoID
	url := "https://youtube.com/get_video_info?video_id=" + y.VideoID + "&eurl=" + eurl
	y.log(fmt.Sprintf("url: %s", url))

	httpClient, err := y.getHTTPClient()
	if err != nil {
		return err
	}

	resp, err := httpClient.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	y.videoInfo = string(body)
	return nil
}

func (y *Youtube) findVideoID(url string) error {
	videoID := url
	if strings.Contains(videoID, "youtu") || strings.ContainsAny(videoID, "\"?&/<%=") {
		reList := []*regexp.Regexp{
			regexp.MustCompile(`(?:v|embed|watch\?v)(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`([^"&?/=%]{11})`),
		}
		for _, re := range reList {
			if isMatch := re.MatchString(videoID); isMatch {
				subs := re.FindStringSubmatch(videoID)
				videoID = subs[1]
			}
		}
	}
	y.log(fmt.Sprintf("Found video id: '%s'", videoID))
	y.VideoID = videoID
	if strings.ContainsAny(videoID, "?&/<%=") {
		return errors.New("invalid characters in video id")
	}
	if len(videoID) < 10 {
		return errors.New("the video id must be at least 10 characters long")
	}
	return nil
}

func (y *Youtube) Write(p []byte) (n int, err error) {
	n = len(p)
	y.totalWrittenBytes = y.totalWrittenBytes + float64(n)
	currentPercent := ((y.totalWrittenBytes / y.contentLength) * 100)
	if (y.downloadLevel <= currentPercent) && (y.downloadLevel < 100) {
		y.downloadLevel++
		y.DownloadPercent <- int64(y.downloadLevel)
	}
	return
}
func (y *Youtube) videoDLWorker(destFile string, target string) error {

	httpClient, err := y.getHTTPClient()
	if err != nil {
		return err
	}

	resp, err := httpClient.Get(target)
	if err != nil {
		y.log(fmt.Sprintf("Http.Get\nerror: %s\ntarget: %s\n", err, target))
		return err
	}
	defer resp.Body.Close()
	y.contentLength = float64(resp.ContentLength)

	if resp.StatusCode != 200 {
		y.log(fmt.Sprintf("reading answer: non 200[code=%v] status code received: '%v'", resp.StatusCode, err))
		return errors.New("non 200 status code received")
	}
	err = os.MkdirAll(filepath.Dir(destFile), 0755)
	if err != nil {
		return err
	}
	out, err := os.Create(destFile)
	if err != nil {
		return err
	}
	mw := io.MultiWriter(out, y)
	_, err = io.Copy(mw, resp.Body)
	if err != nil {
		y.log(fmt.Sprintln("download video err=", err))
		return err
	}
	return nil
}

func (y *Youtube) log(logText string) {
	if y.DebugMode {
		log.Println(logText)
	}
}

func getVideoTitleAuthor(in url.Values) (string, string) {
	playResponse, ok := in["player_response"]
	if !ok {
		return "", ""
	}
	personMap := make(map[string]interface{})

	if err := json.Unmarshal([]byte(playResponse[0]), &personMap); err != nil {
		panic(err)
	}

	s := personMap["videoDetails"]
	myMap := s.(map[string]interface{})
	// fmt.Println("-->", myMap["title"], "oooo:", myMap["author"])
	if title, ok := myMap["title"]; ok {
		if author, ok := myMap["author"]; ok {
			return title.(string), author.(string)
		}
	}

	return "", ""
}
<commit_msg>fix lint errors<commit_after>package youtube

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"golang.org/x/net/proxy"
)

//SetLogOutput :Set logger writer
func SetLogOutput(w io.Writer) {
	log.SetOutput(w)
}

//NewYoutube :Initialize youtube package object
func NewYoutube(debug bool) *Youtube {
	return &Youtube{DebugMode: debug, DownloadPercent: make(chan int64, 100)}
}

func NewYoutubeWithSocks5Proxy(debug bool, socks5Proxy string) *Youtube {
	return &Youtube{DebugMode: debug, DownloadPercent: make(chan int64, 100), Socks5Proxy: socks5Proxy}
}

type stream map[string]string

type Youtube struct {
	DebugMode         bool
	StreamList        []stream
	VideoID           string
	videoInfo         string
	DownloadPercent   chan int64
	Socks5Proxy       string
	contentLength     float64
	totalWrittenBytes float64
	downloadLevel     float64
}

//DecodeURL : Decode youtube URL to retrieval video information.
func (y *Youtube) DecodeURL(url string) error {
	err := y.findVideoID(url)
	if err != nil {
		return fmt.Errorf("findVideoID error=%s", err)
	}

	err = y.getVideoInfo()
	if err != nil {
		return fmt.Errorf("getVideoInfo error=%s", err)
	}

	err = y.parseVideoInfo()
	if err != nil {
		return fmt.Errorf("parse video info failed, err=%s", err)
	}

	return nil
}

//StartDownload : Starting download video to specific address.
func (y *Youtube) StartDownload(destFile string) error {
	//download highest resolution on [0]
	err := errors.New("Empty stream list")
	for _, v := range y.StreamList {
		url := v["url"]
		y.log(fmt.Sprintln("Download url=", url))

		y.log(fmt.Sprintln("Download to file=", destFile))
		err = y.videoDLWorker(destFile, url)
		if err != nil {
			y.log(fmt.Sprintln("Err: ", err.Error()))
			return err
		}
	}
	return err
}

//StartDownloadWithQuality : Starting download video with specific quality.
func (y *Youtube) StartDownloadWithQuality(destFile string, quality string) error {
	//download highest resolution on [0]
	err := errors.New("Empty stream list")
	for _, v := range y.StreamList {
		if strings.Compare(v["quality"], quality) == 0 {
			url := v["url"]
			y.log(fmt.Sprintln("Download url=", url))
			y.log(fmt.Sprintln("Download to file=", destFile))
			err = y.videoDLWorker(destFile, url)
			if err == nil {
				break
			}
		}
	}

	if err != nil {
		return y.StartDownload(destFile)
	}
	return err
}

//StartDownloadFile : Starting download video on my download.
func (y *Youtube) StartDownloadFile() error {
	//download highest resolution on [0]
	err := errors.New("Empty stream list")
	for _, stream := range y.StreamList {
		streamUrl := stream["url"]
		y.log(fmt.Sprintln("Download url=", streamUrl))

		// Find out what the file name should be.
		fileName := sanitizeFilename(stream["title"])
		fileName += pickIdealFileExtension(stream["type"])

		usr, _ := user.Current()
		destFile := filepath.Join(filepath.Join(usr.HomeDir, "Movies", "youtubedr"), fileName)
		y.log(fmt.Sprintln("Download to file=", destFile))

		err = y.videoDLWorker(destFile, streamUrl)
		if err == nil {
			return nil
		}
	}
	return err
}

func pickIdealFileExtension(mediaType string) string {
	defaultExtension := ".mov"

	mediaType, _, err := mime.ParseMediaType(mediaType)
	if err != nil {
		return defaultExtension
	}

	// Rely on hardcoded canonical mime types, as the ones provided by Go aren't exhaustive [1].
	// This seems to be a recurring problem for youtube downloaders, see [2].
	// The implementation is based on mozilla's list [3], IANA [4] and Youtube's support [5].
	// [1] https://github.com/golang/go/blob/ed7888aea6021e25b0ea58bcad3f26da2b139432/src/mime/type.go#L60
	// [2] https://github.com/ZiTAL/youtube-dl/blob/master/mime.types
	// [3] https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/MIME_types/Common_types
	// [4] https://www.iana.org/assignments/media-types/media-types.xhtml#video
	// [5] https://support.google.com/youtube/troubleshooter/2888402?hl=en
	canonicals := map[string]string{
		"video/quicktime":  ".mov",
		"video/x-msvideo":  ".avi",
		"video/x-matroska": ".mkv",
		"video/mpeg":       ".mpeg",
		"video/webm":       ".webm",
		"video/3gpp2":      ".3g2",
		"video/x-flv":      ".flv",
		"video/3gpp":       ".3gp",
		"video/mp4":        ".mp4",
		"video/ogg":        ".ogv",
		"video/mp2t":       ".ts",
	}

	if extension, ok := canonicals[mediaType]; ok {
		return extension
	}

	// Our last resort is to ask the operating system, but these give multiple results and are rarely canonical.
	extensions, err := mime.ExtensionsByType(mediaType)
	if err != nil || extensions == nil {
		return defaultExtension
	}

	return extensions[0]
}

func sanitizeFilename(fileName string) string {
	// Characters not allowed on mac
	//	:/
	// Characters not allowed on linux
	//	/
	// Characters not allowed on windows
	//	<>:"/\|?*

	// Ref https://docs.microsoft.com/en-us/windows/win32/fileio/naming-a-file#naming-conventions

	fileName = regexp.MustCompile(`[:/<>\:"\\|?*]`).ReplaceAllString(fileName, "")
	fileName = regexp.MustCompile(`\s+`).ReplaceAllString(fileName, " ")

	return fileName
}

func (y *Youtube) parseVideoInfo() error {
	answer, err := url.ParseQuery(y.videoInfo)
	if err != nil {
		return err
	}

	status, ok := answer["status"]
	if !ok {
		err = fmt.Errorf("no response status found in the server's answer")
		return err
	}
	if status[0] == "fail" {
		reason, ok := answer["reason"]
		if ok {
			err = fmt.Errorf("'fail' response status found in the server's answer, reason: '%s'", reason[0])
		} else {
			err = errors.New("'fail' response status found in the server's answer, no reason given")
		}
		return err
	}
	if status[0] != "ok" {
		err = fmt.Errorf("non-success response status found in the server's answer (status: '%s')", status)
		return err
	}

	// read the streams map
	streamMap, ok := answer["player_response"]
	if !ok {
		err = errors.New("no stream map found in the server's answer")
		return err
	}

	// Get video title and author.
	title, author := getVideoTitleAuthor(answer)

	var prData PlayerResponseData
	if err := json.Unmarshal([]byte(streamMap[0]), &prData); err != nil {
		fmt.Println(err)
		panic("Player response json data has changed.")
	}

	// Get video download link
	if prData.PlayabilityStatus.Status == "UNPLAYABLE" {
		//Cannot playback on embedded video screen, could not download.
		return errors.New(fmt.Sprint("Cannot playback and download, reason:", prData.PlayabilityStatus.Reason))
	}

	var streams []stream
	for streamPos, streamRaw := range prData.StreamingData.Formats {
		if streamRaw.MimeType == "" {
			y.log(fmt.Sprintf("An error occurred while decoding one of the video's stream's information: stream %d.\n", streamPos))
			continue
		}
		streamUrl := streamRaw.URL
		if streamUrl == "" {
			cipher := streamRaw.Cipher
			decipheredUrl, err := y.decipher(cipher)
			if err != nil {
				return err
			}
			streamUrl = decipheredUrl
		}

		streams = append(streams, stream{
			"quality": streamRaw.Quality,
			"type":    streamRaw.MimeType,
			"url":     streamUrl,

			"title":  title,
			"author": author,
		})
		y.log(fmt.Sprintf("Title: %s Author: %s Stream found: quality '%s', format '%s'", title, author, streamRaw.Quality, streamRaw.MimeType))
	}

	y.StreamList = streams
	if len(y.StreamList) == 0 {
		return errors.New("no stream list found in the server's answer")
	}
	return nil
}

func (y *Youtube) getHTTPClient() (*http.Client, error) {
	// setup a http client
	httpTransport := &http.Transport{
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		IdleConnTimeout:       60 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	httpClient := &http.Client{Transport: httpTransport}

	if len(y.Socks5Proxy) == 0 {
		y.log("Using http without proxy.")
		return httpClient, nil
	}

	dialer, err := proxy.SOCKS5("tcp", y.Socks5Proxy, nil, proxy.Direct)
	if err != nil {
		fmt.Fprintln(os.Stderr, "can't connect to the proxy:", err)
		return nil, err
	}
	// set our socks5 as the dialer
	dc := dialer.(interface {
		DialContext(ctx context.Context, network, addr string) (net.Conn, error)
	})
	httpTransport.DialContext = dc.DialContext

	y.log(fmt.Sprintf("Using http with proxy %s.", y.Socks5Proxy))

	return httpClient, nil
}

func (y *Youtube) getVideoInfo() error {
	eurl := "https://youtube.googleapis.com/v/" + y.VideoID
	url := "https://youtube.com/get_video_info?video_id=" + y.VideoID + "&eurl=" + eurl
	y.log(fmt.Sprintf("url: %s", url))

	httpClient, err := y.getHTTPClient()
	if err != nil {
		return err
	}

	resp, err := httpClient.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	y.videoInfo = string(body)
	return nil
}

func (y *Youtube) findVideoID(url string) error {
	videoID := url
	if strings.Contains(videoID, "youtu") || strings.ContainsAny(videoID, "\"?&/<%=") {
		reList := []*regexp.Regexp{
			regexp.MustCompile(`(?:v|embed|watch\?v)(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`(?:=|/)([^"&?/=%]{11})`),
			regexp.MustCompile(`([^"&?/=%]{11})`),
		}
		for _, re := range reList {
			if isMatch := re.MatchString(videoID); isMatch {
				subs := re.FindStringSubmatch(videoID)
				videoID = subs[1]
			}
		}
	}
	y.log(fmt.Sprintf("Found video id: '%s'", videoID))
	y.VideoID = videoID
	if strings.ContainsAny(videoID, "?&/<%=") {
		return errors.New("invalid characters in video id")
	}
	if len(videoID) < 10 {
		return errors.New("the video id must be at least 10 characters long")
	}
	return nil
}

func (y *Youtube) Write(p []byte) (n int, err error) {
	n = len(p)
	y.totalWrittenBytes = y.totalWrittenBytes + float64(n)
	currentPercent := ((y.totalWrittenBytes / y.contentLength) * 100)
	if (y.downloadLevel <= currentPercent) && (y.downloadLevel < 100) {
		y.downloadLevel++
		y.DownloadPercent <- int64(y.downloadLevel)
	}
	return
}
func (y *Youtube) videoDLWorker(destFile string, target string) error {

	httpClient, err := y.getHTTPClient()
	if err != nil {
		return err
	}

	resp, err := httpClient.Get(target)
	if err != nil {
		y.log(fmt.Sprintf("Http.Get\nerror: %s\ntarget: %s\n", err, target))
		return err
	}
	defer resp.Body.Close()
	y.contentLength = float64(resp.ContentLength)

	if resp.StatusCode != 200 {
		y.log(fmt.Sprintf("reading answer: non 200[code=%v] status code received: '%v'", resp.StatusCode, err))
		return errors.New("non 200 status code received")
	}
	err = os.MkdirAll(filepath.Dir(destFile), 0755)
	if err != nil {
		return err
	}
	out, err := os.Create(destFile)
	if err != nil {
		return err
	}
	mw := io.MultiWriter(out, y)
	_, err = io.Copy(mw, resp.Body)
	if err != nil {
		y.log(fmt.Sprintln("download video err=", err))
		return err
	}
	return nil
}

func (y *Youtube) log(logText string) {
	if y.DebugMode {
		log.Println(logText)
	}
}

func getVideoTitleAuthor(in url.Values) (string, string) {
	playResponse, ok := in["player_response"]
	if !ok {
		return "", ""
	}
	personMap := make(map[string]interface{})

	if err := json.Unmarshal([]byte(playResponse[0]), &personMap); err != nil {
		panic(err)
	}

	s := personMap["videoDetails"]
	myMap := s.(map[string]interface{})
	// fmt.Println("-->", myMap["title"], "oooo:", myMap["author"])
	if title, ok := myMap["title"]; ok {
		if author, ok := myMap["author"]; ok {
			return title.(string), author.(string)
		}
	}

	return "", ""
}
<|endoftext|>


<commit_before>package cmd

import (
	"errors"
	"fmt"
	"strconv"

	"github.com/PuerkitoBio/gred/resp"
	"github.com/PuerkitoBio/gred/srv"
)

const (
	// WrongNumberOfArgsFmt is a string that holds the normalized error message
	// for when the number of arguments of a command is invalid.
	WrongNumberOfArgsFmt = "ERR wrong number of arguments for '%s' command"
)

var (
	// PongVal is a sentinel value used to return the PONG response for the PING command.
	PongVal = resp.Pong{}

	// OKVal is a sentinel value used to indicate that the standard OK simple string
	// response should be returned.
	OKVal = resp.OK{}

	ErrNotInteger        = errors.New("ERR value is not an integer or out of range")
	ErrNotFloat          = errors.New("ERR value is not a valid float")
	ErrInvalidValType    = errors.New("WRONGTYPE Operation against a key holding the wrong kind of value")
	ErrSyntax            = errors.New("ERR syntax error")
	ErrNoSuchKey         = errors.New("ERR no such key")
	ErrOutOfRange        = errors.New("ERR index out of range")
	ErrHashFieldNotInt   = errors.New("ERR hash value is not an integer")
	ErrHashFieldNotFloat = errors.New("ERR hash value is not a valid float")

	ErrQuit = errors.New("quit")
)

// Commands holds the list of registered commands.
var Commands = make(map[string]Cmd)

// Register registers a command name with an implementation.
func Register(name string, c Cmd) {
	if name == "" {
		panic("cmds: call Register with empty command name")
	}
	if _, ok := Commands[name]; ok {
		panic(fmt.Sprintf("cmds: command %s already registered", name))
	}
	Commands[name] = c
}

// Cmd defines the common method required to implement a basic command.
// It is insufficient to implement this sole interface. A command must also
// implement one of the more specific {Srv,DB}Cmd interfaces.
type Cmd interface {
	Parse(string, []string) ([]string, []int64, []float64, error)
}

type SrvFn func([]string, []int64, []float64) (interface{}, error)

type SrvCmd interface {
	Cmd
	Exec([]string, []int64, []float64) (interface{}, error)
}

func NewSrvCmd(arg *ArgDef, fn SrvFn) SrvCmd {
	return &srvCmd{
		arg,
		fn,
	}
}

type srvCmd struct {
	*ArgDef
	fn SrvFn
}

func (s *srvCmd) Exec(args []string, ints []int64, floats []float64) (interface{}, error) {
	return s.fn(args, ints, floats)
}

type DBCmd interface {
	Cmd
	ExecWithDB(srv.DB, []string, []int64, []float64) (interface{}, error)
}

type ArgFn func([]string, []int64, []float64) error

type ArgDef struct {
	FloatIndices     []int
	IntIndices       []int
	MinArgs, MaxArgs int
	ValidateFn       ArgFn
}

func (a *ArgDef) Parse(name string, args []string) ([]string, []int64, []float64, error) {
	l := len(args)
	if l < a.MinArgs || (l > a.MaxArgs && a.MaxArgs >= 0) {
		return nil, nil, nil, fmt.Errorf(WrongNumberOfArgsFmt, name)
	}

	// Parse integers
	intix := a.IntIndices
	ints := make([]int64, len(intix))
	for i, ix := range intix {
		if ix < 0 {
			ix = l + ix
		}
		val, err := strconv.ParseInt(args[ix], 10, 64)
		if err != nil {
			return nil, nil, nil, ErrNotInteger
		}
		ints[i] = val
	}

	// Parse floats
	fix := a.FloatIndices
	floats := make([]float64, len(fix))
	for i, ix := range fix {
		if ix < 0 {
			ix = l + ix
		}
		val, err := strconv.ParseFloat(args[ix], 64)
		if err != nil {
			return nil, nil, nil, ErrNotFloat
		}
		floats[i] = val
	}

	if a.ValidateFn != nil {
		err := a.ValidateFn(args, ints, floats)
		if err != nil {
			return nil, nil, nil, err
		}
	}
	return args, ints, floats, nil
}

var _ DBCmd = (*singleKeyCmd)(nil)

type KeyFn func(srv.Key, []string, []int64, []float64) (interface{}, error)

func NewSingleKeyCmd(arg *ArgDef, noKeyFlag srv.NoKeyFlag, fn KeyFn) DBCmd {
	return &singleKeyCmd{
		ArgDef: arg,
		noKey:  noKeyFlag,
		fn:     fn,
	}
}

type singleKeyCmd struct {
	*ArgDef
	noKey srv.NoKeyFlag
	fn    KeyFn
}

func (c *singleKeyCmd) ExecWithDB(db srv.DB, args []string, ints []int64, floats []float64) (interface{}, error) {
	k, def := db.LockGetKey(args[0], c.noKey)
	defer def()

	return c.fn(k, args, ints, floats)
}

var _ DBCmd = (*dbCmd)(nil)

type DBFn func(srv.DB, []string, []int64, []float64) (interface{}, error)

func NewDBCmd(arg *ArgDef, fn DBFn) DBCmd {
	return &dbCmd{
		ArgDef: arg,
		fn:     fn,
	}
}

type dbCmd struct {
	*ArgDef
	fn DBFn
}

func (d *dbCmd) ExecWithDB(db srv.DB, args []string, ints []int64, floats []float64) (interface{}, error) {
	return d.fn(db, args, ints, floats)
}
<commit_msg>comments<commit_after>package cmd

import (
	"errors"
	"fmt"
	"strconv"

	"github.com/PuerkitoBio/gred/resp"
	"github.com/PuerkitoBio/gred/srv"
)

const (
	// WrongNumberOfArgsFmt is a string that holds the normalized error message
	// for when the number of arguments of a command is invalid.
	WrongNumberOfArgsFmt = "ERR wrong number of arguments for '%s' command"
)

var (
	// PongVal is a sentinel value used to return the PONG response for the PING command.
	PongVal = resp.Pong{}

	// OKVal is a sentinel value used to indicate that the standard OK simple string
	// response should be returned.
	OKVal = resp.OK{}

	// ErrNotInteger is returned when an argument that is expected to be an integer
	// cannot be parsed as an integer.
	ErrNotInteger = errors.New("ERR value is not an integer or out of range")

	// ErrNotFloat is returned when an argument that is expected to be a float
	// cannot be parsed as a float.
	ErrNotFloat = errors.New("ERR value is not a valid float")

	// ErrInvalidValType is returned when the key's value is not of the expected type.
	ErrInvalidValType = errors.New("WRONGTYPE Operation against a key holding the wrong kind of value")

	// ErrSyntax is returned when an argument doesn't have the expected allowed syntax.
	ErrSyntax = errors.New("ERR syntax error")

	// ErrNoSuchKey is returned when a command is attempted against a non-existing key.
	ErrNoSuchKey = errors.New("ERR no such key")

	// ErrOutOfRange is returned when an argument is out of range.
	ErrOutOfRange = errors.New("ERR index out of range")

	// ErrHashFieldNotInt is returned when an increment operation is attempted on
	// a hash field that does not contain an integer value.
	ErrHashFieldNotInt = errors.New("ERR hash value is not an integer")

	// ErrHashFieldNotFloat is returned when an increment operation is attempted on
	// a hash field that does not contain a float value.
	ErrHashFieldNotFloat = errors.New("ERR hash value is not a valid float")

	// ErrQuit is a sentinel error value to indicate that the network connection
	// should be closed, as requested by the client.
	ErrQuit = errors.New("quit")
)

// Commands holds the list of registered commands.
var Commands = make(map[string]Cmd)

// Register registers a command name with an implementation.
func Register(name string, c Cmd) {
	if name == "" {
		panic("cmds: call Register with empty command name")
	}
	if _, ok := Commands[name]; ok {
		panic(fmt.Sprintf("cmds: command %s already registered", name))
	}
	Commands[name] = c
}

// Cmd defines the common method required to implement a basic command.
// It is insufficient to implement this sole interface. A command must also
// implement one of the more specific {Srv,DB}Cmd interfaces.
type Cmd interface {
	Parse(string, []string) ([]string, []int64, []float64, error)
}

// SrvFn defines the function signature required for the SrvCmd implementation.
type SrvFn func([]string, []int64, []float64) (interface{}, error)

// SrvCmd defines the methods required to implement a server command.
type SrvCmd interface {
	Cmd
	Exec([]string, []int64, []float64) (interface{}, error)
}

// NewSrvCmd creates a new SrvCmd value with the specified argument definition
// and execution function.
func NewSrvCmd(arg *ArgDef, fn SrvFn) SrvCmd {
	return &srvCmd{
		arg,
		fn,
	}
}

type srvCmd struct {
	*ArgDef
	fn SrvFn
}

func (s *srvCmd) Exec(args []string, ints []int64, floats []float64) (interface{}, error) {
	return s.fn(args, ints, floats)
}

// DBCmd defines the methods required to implement a database command.
type DBCmd interface {
	Cmd
	ExecWithDB(srv.DB, []string, []int64, []float64) (interface{}, error)
}

type ArgFn func([]string, []int64, []float64) error

type ArgDef struct {
	FloatIndices     []int
	IntIndices       []int
	MinArgs, MaxArgs int
	ValidateFn       ArgFn
}

func (a *ArgDef) Parse(name string, args []string) ([]string, []int64, []float64, error) {
	l := len(args)
	if l < a.MinArgs || (l > a.MaxArgs && a.MaxArgs >= 0) {
		return nil, nil, nil, fmt.Errorf(WrongNumberOfArgsFmt, name)
	}

	// Parse integers
	intix := a.IntIndices
	ints := make([]int64, len(intix))
	for i, ix := range intix {
		if ix < 0 {
			ix = l + ix
		}
		val, err := strconv.ParseInt(args[ix], 10, 64)
		if err != nil {
			return nil, nil, nil, ErrNotInteger
		}
		ints[i] = val
	}

	// Parse floats
	fix := a.FloatIndices
	floats := make([]float64, len(fix))
	for i, ix := range fix {
		if ix < 0 {
			ix = l + ix
		}
		val, err := strconv.ParseFloat(args[ix], 64)
		if err != nil {
			return nil, nil, nil, ErrNotFloat
		}
		floats[i] = val
	}

	if a.ValidateFn != nil {
		err := a.ValidateFn(args, ints, floats)
		if err != nil {
			return nil, nil, nil, err
		}
	}
	return args, ints, floats, nil
}

var _ DBCmd = (*singleKeyCmd)(nil)

type KeyFn func(srv.Key, []string, []int64, []float64) (interface{}, error)

func NewSingleKeyCmd(arg *ArgDef, noKeyFlag srv.NoKeyFlag, fn KeyFn) DBCmd {
	return &singleKeyCmd{
		ArgDef: arg,
		noKey:  noKeyFlag,
		fn:     fn,
	}
}

type singleKeyCmd struct {
	*ArgDef
	noKey srv.NoKeyFlag
	fn    KeyFn
}

func (c *singleKeyCmd) ExecWithDB(db srv.DB, args []string, ints []int64, floats []float64) (interface{}, error) {
	k, def := db.LockGetKey(args[0], c.noKey)
	defer def()

	return c.fn(k, args, ints, floats)
}

var _ DBCmd = (*dbCmd)(nil)

type DBFn func(srv.DB, []string, []int64, []float64) (interface{}, error)

func NewDBCmd(arg *ArgDef, fn DBFn) DBCmd {
	return &dbCmd{
		ArgDef: arg,
		fn:     fn,
	}
}

type dbCmd struct {
	*ArgDef
	fn DBFn
}

func (d *dbCmd) ExecWithDB(db srv.DB, args []string, ints []int64, floats []float64) (interface{}, error) {
	return d.fn(db, args, ints, floats)
}
<|endoftext|>


<commit_before>package filer

type ReaderPattern struct {
	isStreaming    bool
	lastReadOffset int64
}

// For streaming read: only cache the first chunk
// For random read: only fetch the requested range, instead of the whole chunk

func NewReaderPattern() *ReaderPattern {
	return &ReaderPattern{
		isStreaming:    true,
		lastReadOffset: 0,
	}
}

func (rp *ReaderPattern) MonitorReadAt(offset int64, size int) {
	if rp.lastReadOffset > offset {
		rp.isStreaming = false
	}
	rp.lastReadOffset = offset
}

func (rp *ReaderPattern) IsStreamingMode() bool {
	return rp.isStreaming
}

func (rp *ReaderPattern) IsRandomMode() bool {
	return !rp.isStreaming
}
<commit_msg>detect non streaming mode on the first read<commit_after>package filer

type ReaderPattern struct {
	isStreaming    bool
	lastReadOffset int64
}

// For streaming read: only cache the first chunk
// For random read: only fetch the requested range, instead of the whole chunk

func NewReaderPattern() *ReaderPattern {
	return &ReaderPattern{
		isStreaming:    true,
		lastReadOffset: -1,
	}
}

func (rp *ReaderPattern) MonitorReadAt(offset int64, size int) {
	if rp.lastReadOffset > offset {
		rp.isStreaming = false
	}
	if rp.lastReadOffset == -1 {
		if offset != 0 {
			rp.isStreaming = false
		}
	}
	rp.lastReadOffset = offset
}

func (rp *ReaderPattern) IsStreamingMode() bool {
	return rp.isStreaming
}

func (rp *ReaderPattern) IsRandomMode() bool {
	return !rp.isStreaming
}
<|endoftext|>


<commit_before>package spy

import (
	"reflect"
	"github.com/Sirupsen/logrus"
	"github.com/Jeffail/tunny"
)

type Crawler struct {
	*Config
	*logrus.Logger
	*Stats

	Concurrency int

	Spider ISpider
	Scheduler IScheduler
	Fetcher IFetcher
	*SpiderMiddlewareManager
	*ItemPipelineManager

	crawling bool
	Paused bool

	*tunny.WorkPool
}

func NewCrawler(spider ISpider, scheduler IScheduler) *Crawler {
	concurrency := 100

	return &Crawler{
		Spider: spider,
		Scheduler: scheduler,
		Concurrency: concurrency,
		WorkPool: tunny.CreatePoolGeneric(concurrency)
	}
}

func (c *Crawler) Craw() {
	// TODO: signal handling for SIGTERM, SIGINT, SIGBREAK

	if c.crawling {
		panic("Crawling already taking place")
	}
	c.crawling = true

	c.WorkPool.Open()
}

func (c *Crawler) Stop() {
	c.WorkPool.Close()
}

func (c *Crawler) OpenSpider() {
	startRequests, err := c.SpiderMiddlewareManager.ProcessStartRequests(c.Spider.StartRequests(), c.Spider)

	c.Scheduler.Open(c.Spider)

	SpiderOpened.Pub(c.Spider)

	c.nextRequest()
}

func (c *Crawler) CloseSpider() {

}

func (c *Crawler) nextRequest() {
	if c.Paused {
		return
	}

	for !c.needsBackout() {
		request := c.Scheduler.NextRequest()
		if request == nil {
			break
		}
		c.Fetch(request)
	}

}

func (c *Crawler) needsBackout() {
	return !c.crawling // TODO
}

func (c *Crawler) Fetch(request *Request) {
	result, err := c.Fetcher.Fetch(request, c.Spider)
	if err != nil {
		c.EnqueueScrape(nil, err, request)
		return
	}
	if result.Response != nil {
		result.Response.Request = request // tie request to response received
		c.Logger.WithFields(logrus.Fields{
			"action": "crawl",
			"status": result.Response.StatusCode,
			"request": request,
		}).Debugf("Crawled request %s, status %d", request, result.Response.StatusCode)
		ResponseReceived.Pub(c.Spider, request, result.Response)

		c.EnqueueScrape(result.Response, nil, request)
	} else { // fetcher can return request, i.e., redirect
		c.Fetch(result.Request)
	}
}

func (c *Crawler) EnqueueScrape(response *Response, err error, request *Request) {

}

func (c *Crawler) Scrape(response *Response, err error, request *Request) {
	result, err := c.SpiderMiddlewareManager.ScrapeResponse(request, response, c.Spider)
	if err == nil {
		for _, req := range result.Requests {
			c.processSpiderRequest(req)
		}
		for _, item := range result.Items {
			c.processSpiderItem(item)
		}

	} else {
		if err == ErrSpiderClosed {
			return
		}
		c.Logger.WithError(err).Errorf("Processing request %s (referer: %s)", request, request.Header.Get("Referer"))
		SpiderError.Pub(c.Spider, response, err)
		c.Stats.Inc("SpiderError/"+reflect.TypeOf(err).Name())
	}
}

func (c *Crawler) processSpiderRequest(request *Request) {
	c.WorkPool.SendWorkAsync(func() {

	}, nil)
}

func (c *Crawler) processSpiderItem(item *Item, response *Response) {
	c.WorkPool.SendWorkAsync(func() {
		resultItem, err := c.ItemPipelineManager.ProcessItem(item, c.Spider)

		if err == ErrItemDropped {
			c.Logger.WithFields(logrus.Fields{
				"action": "drop",
				"item": item,
			}).Warnf("Dropped item %s", item)
			ItemDropped.Pub(c.Spider, response, item)
		} else if err != nil {
			c.Logger.WithError(err).Errorf("Processing item %s", item)
		} else {
			c.Logger.WithFields(logrus.Fields{
				"action": "scrap",
				"item": resultItem,
				"src": response,
			}).Debugf("Scraped item %s from %s", resultItem, response)
		}
	}, nil)
}
<commit_msg>Update<commit_after>package spy

import (
	"reflect"
	"github.com/Sirupsen/logrus"
	"github.com/Jeffail/tunny"
)

type Crawler struct {
	*Config
	*logrus.Logger
	*Stats

	Concurrency int

	Spider ISpider
	Scheduler IScheduler
	Fetcher IFetcher
	*SpiderMiddlewareManager
	*ItemPipelineManager

	crawling bool
	Paused bool

	*tunny.WorkPool
}

func NewCrawler(spider ISpider, scheduler IScheduler) *Crawler {
	concurrency := 100

	return &Crawler{
		Spider: spider,
		Scheduler: scheduler,
		Concurrency: concurrency,
		WorkPool: tunny.CreatePoolGeneric(concurrency),
	}
}

func (c *Crawler) Craw() {
	// TODO: signal handling for SIGTERM, SIGINT, SIGBREAK

	if c.crawling {
		panic("Crawling already taking place")
	}
	c.crawling = true

	c.WorkPool.Open()
}

func (c *Crawler) Stop() {
	c.WorkPool.Close()
}

func (c *Crawler) OpenSpider() {
	startRequests, err := c.SpiderMiddlewareManager.ProcessStartRequests(c.Spider.StartRequests(), c.Spider)

	c.Scheduler.Open(c.Spider)

	SpiderOpened.Pub(c.Spider)

	for _, req := range startRequests {
		c.Scheduler.EnqueueRequest(req)
	}

	c.scheduleRequests()
}

func (c *Crawler) CloseSpider() {

}

func (c *Crawler) scheduleRequests() {
	if c.Paused {
		return
	}

	for !c.needsBackout() {
		request := c.Scheduler.NextRequest()
		if request == nil {
			break
		}
		c.fetch(request)
	}
}

func (c *Crawler) needsBackout() bool {
	return !c.crawling // TODO
}

func (c *Crawler) fetch(request *Request) {
	result, err := c.Fetcher.Fetch(request, c.Spider)
	if err != nil {
		c.enqueueScrape(nil, err, request)
		return
	}
	if result.Response != nil {
		result.Response.Request = request // tie request to response received
		c.Logger.WithFields(logrus.Fields{
			"action": "crawl",
			"status": result.Response.StatusCode,
			"request": request,
		}).Debugf("Crawled request %s, status %d", request, result.Response.StatusCode)
		ResponseReceived.Pub(c.Spider, request, result.Response)

		c.enqueueScrape(result.Response, nil, request)
	} else { // fetcher can return request, i.e., redirect
		c.fetch(result.Request)
	}
}

func (c *Crawler) enqueueScrape(response *Response, err error, request *Request) {
	c.WorkPool.SendWorkAsync(func() {
		c.scrape(response, err, request)
	}, nil)
}

func (c *Crawler) scrape(response *Response, err error, request *Request) {
	result, err := c.SpiderMiddlewareManager.ScrapeResponse(request, response, c.Spider)

	if err == nil {
		for _, req := range result.Requests {
			c.processSpiderRequest(req)
		}
		for _, item := range result.Items {
			c.processSpiderItem(item, response)
		}

	} else {
		if err == ErrSpiderClosed {
			return
		}
		c.Logger.WithError(err).Errorf("Processing request %s (referer: %s)", request, request.Header.Get("Referer"))
		SpiderError.Pub(c.Spider, response, err)
		c.Stats.Inc("SpiderError/"+reflect.TypeOf(err).Name())
	}
}

func (c *Crawler) processSpiderRequest(request *Request) {
	c.WorkPool.SendWorkAsync(func() {

	}, nil)
}

func (c *Crawler) processSpiderItem(item *Item, response *Response) {
	c.WorkPool.SendWorkAsync(func() {
		resultItem, err := c.ItemPipelineManager.ProcessItem(item, c.Spider)

		if err == ErrItemDropped {
			c.Logger.WithFields(logrus.Fields{
				"action": "drop",
				"item": item,
			}).Warnf("Dropped item %s", item)
			ItemDropped.Pub(c.Spider, response, item)
		} else if err != nil {
			c.Logger.WithError(err).Errorf("Processing item %s", item)
		} else {
			c.Logger.WithFields(logrus.Fields{
				"action": "scrap",
				"item": resultItem,
				"src": response,
			}).Debugf("Scraped item %s from %s", resultItem, response)
		}
	}, nil)
}
<|endoftext|>


<commit_before>/*
 Copyright 2020 Qiniu Cloud (qiniu.com)

 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 bytecode

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/qiniu/x/log"
)

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

func init() {
	log.SetOutputLevel(log.Ldebug)
	if opCallGoFunc != SymbolFunc || opCallGoFuncv != SymbolFuncv {
		panic("opCallGoFunc != SymbolFunc || opCallGoFuncv != SymbolFuncv")
	}
}

func Strcat(a, b string) string {
	return a + b
}

func execStrcat(arity int, p *Context) {
	args := p.GetArgs(2)
	ret := Strcat(args[0].(string), args[1].(string))
	p.Ret(2, ret)
}

func execSprint(arity int, p *Context) {
	args := p.GetArgs(arity)
	s := fmt.Sprint(args...)
	p.Ret(arity, s)
}

func execSprintf(arity int, p *Context) {
	args := p.GetArgs(arity)
	s := fmt.Sprintf(args[0].(string), args[1:]...)
	p.Ret(arity, s)
}

func execErrorf(arity int, p *Context) {
	args := p.GetArgs(arity)
	s := fmt.Errorf(args[0].(string), args[1:]...)
	p.Ret(arity, s)
}

// I is a Go package instance.
var I = NewGoPackage("")

func init() {
	I.RegisterFuncvs(
		I.Funcv("Sprint", fmt.Sprint, execSprint),
		I.Funcv("Sprintf", fmt.Sprintf, execSprintf),
		I.Funcv("Errorf", fmt.Errorf, execErrorf),
	)
	I.RegisterFuncs(
		I.Func("strcat", Strcat, execStrcat),
	)
	I.RegisterVars(
		I.Var("x", new(int)),
	)
	I.RegisterConsts(
		I.Const("true", reflect.Bool, true),
		I.Const("false", reflect.Bool, false),
		I.Const("nil", ConstUnboundPtr, nil),
	)
	I.RegisterTypes(
		I.Rtype(reflect.TypeOf((*Context)(nil))),
		I.Rtype(reflect.TypeOf((*Code)(nil))),
		I.Rtype(reflect.TypeOf((*Stack)(nil))),
		I.Type("rune", TyRune),
	)
	_ = I.PkgPath()
}

func TestVarAndConst(t *testing.T) {
	if ci, ok := I.FindConst("true"); !ok || ci.Value != true {
		t.Fatal("FindConst failed:", ci.Value)
	}
	if ci, ok := I.FindConst("nil"); !ok || ci.Kind != ConstUnboundPtr {
		t.Fatal("FindConst failed:", ci.Kind)
	}
	if addr, ok := I.FindVar("x"); !ok || addr != 0 {
		t.Fatal("FindVar failed:", addr)
	}
}

func TestSprint(t *testing.T) {
	sprint, ok := I.FindFuncv("Sprint")
	if !ok {
		t.Fatal("FindFuncv failed: Sprint")
	}

	code := newBuilder().
		Push(5).
		Push("32").
		CallGoFuncv(sprint, 2, 2).
		Resolve()

	ctx := NewContext(code)
	ctx.Exec(0, code.Len())
	if v := checkPop(ctx); v != "532" {
		t.Fatal("5 `32` sprint != 532, ret =", v)
	}
}

func TestSprintf(t *testing.T) {
	sprintf, ok := I.FindFuncv("Sprintf")
	strcat, ok2 := I.FindFunc("strcat")
	if !ok || !ok2 {
		t.Fatal("FindFunc failed: Sprintf/strcat")
	}
	_ = defaultImpl.GetGoFuncType(strcat)
	_ = defaultImpl.GetGoFuncvType(sprintf)

	code := newBuilder().
		Push("Hello, %v, %d, %s").
		Push(1.3).
		Push(1).
		Push("x").
		Push("sw").
		CallGoFunc(strcat, 2).
		CallGoFuncv(sprintf, 4, 4).
		Resolve()

	ctx := NewContext(code)
	ctx.Exec(0, code.Len())
	if v := checkPop(ctx); v != "Hello, 1.3, 1, xsw" {
		t.Fatal("format 1.3 1 `x` `sw` strcat sprintf != `Hello, 1.3, 1, xsw`, ret =", v)
	}
}

func TestLargeArity(t *testing.T) {
	sprint, kind, ok := defaultImpl.FindGoPackage("").Find("Sprint")
	if !ok || kind != SymbolFuncv {
		t.Fatal("Find failed: Sprint")
	}

	b := newBuilder()
	ret := ""
	for i := 0; i < bitsFuncvArityMax+1; i++ {
		b.Push("32")
		ret += "32"
	}
	code := b.
		CallGoFuncv(GoFuncvAddr(sprint), bitsFuncvArityMax+1, bitsFuncvArityMax+1).
		Resolve()

	ctx := NewContext(code)
	ctx.Exec(0, code.Len())
	if v := checkPop(ctx); v != ret {
		t.Fatal("32 times(1024) sprint != `32` times(1024), ret =", v)
	}
}

func TestType(t *testing.T) {
	typ, ok := I.FindType("Context")
	if !ok {
		t.Fatal("FindType failed: Context not found")
	}
	fmt.Println(typ)

	typ, ok = FindGoPackage("").FindType("rune")
	if !ok || typ != TyRune {
		t.Fatal("FindType failed: rune not found")
	}
	fmt.Println(typ)
}
<commit_msg>add TestGoVar: LoadGoVar, StoreGoVar, AddrGoVar<commit_after>/*
 Copyright 2020 Qiniu Cloud (qiniu.com)

 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 bytecode

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/qiniu/x/log"
)

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

func init() {
	log.SetOutputLevel(log.Ldebug)
	if opCallGoFunc != SymbolFunc || opCallGoFuncv != SymbolFuncv {
		panic("opCallGoFunc != SymbolFunc || opCallGoFuncv != SymbolFuncv")
	}
}

func Strcat(a, b string) string {
	return a + b
}

func execStrcat(arity int, p *Context) {
	args := p.GetArgs(2)
	ret := Strcat(args[0].(string), args[1].(string))
	p.Ret(2, ret)
}

func execSprint(arity int, p *Context) {
	args := p.GetArgs(arity)
	s := fmt.Sprint(args...)
	p.Ret(arity, s)
}

func execSprintf(arity int, p *Context) {
	args := p.GetArgs(arity)
	s := fmt.Sprintf(args[0].(string), args[1:]...)
	p.Ret(arity, s)
}

func execErrorf(arity int, p *Context) {
	args := p.GetArgs(arity)
	s := fmt.Errorf(args[0].(string), args[1:]...)
	p.Ret(arity, s)
}

// I is a Go package instance.
var I = NewGoPackage("")

func init() {
	I.RegisterFuncvs(
		I.Funcv("Sprint", fmt.Sprint, execSprint),
		I.Funcv("Sprintf", fmt.Sprintf, execSprintf),
		I.Funcv("Errorf", fmt.Errorf, execErrorf),
	)
	I.RegisterFuncs(
		I.Func("strcat", Strcat, execStrcat),
	)
	I.RegisterVars(
		I.Var("x", new(int)),
		I.Var("y", new(int)),
	)
	I.RegisterConsts(
		I.Const("true", reflect.Bool, true),
		I.Const("false", reflect.Bool, false),
		I.Const("nil", ConstUnboundPtr, nil),
	)
	I.RegisterTypes(
		I.Rtype(reflect.TypeOf((*Context)(nil))),
		I.Rtype(reflect.TypeOf((*Code)(nil))),
		I.Rtype(reflect.TypeOf((*Stack)(nil))),
		I.Type("rune", TyRune),
	)
	_ = I.PkgPath()
}

func TestVarAndConst(t *testing.T) {
	if ci, ok := I.FindConst("true"); !ok || ci.Value != true {
		t.Fatal("FindConst failed:", ci.Value)
	}
	if ci, ok := I.FindConst("nil"); !ok || ci.Kind != ConstUnboundPtr {
		t.Fatal("FindConst failed:", ci.Kind)
	}
	if addr, ok := I.FindVar("x"); !ok || addr != 0 {
		t.Fatal("FindVar failed:", addr)
	}
}

func TestGoVar(t *testing.T) {
	sprint, ok := I.FindFuncv("Sprint")
	strcat, ok2 := I.FindFunc("strcat")
	if !ok || !ok2 {
		t.Fatal("FindFunc failed: Sprintf/strcat")
	}

	pkg := NewGoPackage("pkg")

	pkg.RegisterVars(
		I.Var("x", new(string)),
		I.Var("y", new(string)),
		I.Var("i", new(int)),
		I.Var("j", new(int)),
	)

	x, ok := pkg.FindVar("x")
	if !ok {
		t.Fatal("FindVar failed:", x)
	}
	y, ok := pkg.FindVar("y")
	if !ok {
		t.Fatal("FindVar failed:", y)
	}
	i, ok := pkg.FindVar("i")
	if !ok {
		t.Fatal("FindVar failed:", i)
	}
	j, ok := pkg.FindVar("j")
	if !ok {
		t.Fatal("FindVar failed:", j)
	}

	b := newBuilder()
	code := b.
		Push(5).
		Push("32").
		CallGoFuncv(sprint, 2, 2).
		StoreGoVar(x). // x = sprint(5, "32")
		Push("78").
		LoadGoVar(x).
		CallGoFunc(strcat, 2).
		StoreGoVar(y). // y = strcat("78", x)
		Push(5).
		StoreGoVar(i).
		LoadGoVar(i).
		AddrGoVar(j).
		AddrOp(Int, OpAssign).
		Push(5).
		AddrGoVar(j).
		AddrOp(Int, OpAddAssign).
		Resolve()

	ctx := NewContext(code)
	ctx.Exec(0, code.Len())

	if v := reflect.ValueOf(govars[x].Addr).Elem().String(); v != "532" {
		t.Fatal("x != 532, x = ", v)
	}
	if v := reflect.ValueOf(govars[y].Addr).Elem().String(); v != "78532" {
		t.Fatal("y != 78532, x =", v)
	}
	if v := reflect.ValueOf(govars[i].Addr).Elem().Int(); v != 5 {
		t.Fatal("i !=5, x =", i)
	}
	if v := reflect.ValueOf(govars[j].Addr).Elem().Int(); v != 10 {
		t.Fatal("j !=10, x =", j)
	}
}

func TestSprint(t *testing.T) {
	sprint, ok := I.FindFuncv("Sprint")
	if !ok {
		t.Fatal("FindFuncv failed: Sprint")
	}

	code := newBuilder().
		Push(5).
		Push("32").
		CallGoFuncv(sprint, 2, 2).
		Resolve()

	ctx := NewContext(code)
	ctx.Exec(0, code.Len())
	if v := checkPop(ctx); v != "532" {
		t.Fatal("5 `32` sprint != 532, ret =", v)
	}
}

func TestSprintf(t *testing.T) {
	sprintf, ok := I.FindFuncv("Sprintf")
	strcat, ok2 := I.FindFunc("strcat")
	if !ok || !ok2 {
		t.Fatal("FindFunc failed: Sprintf/strcat")
	}
	_ = defaultImpl.GetGoFuncType(strcat)
	_ = defaultImpl.GetGoFuncvType(sprintf)

	code := newBuilder().
		Push("Hello, %v, %d, %s").
		Push(1.3).
		Push(1).
		Push("x").
		Push("sw").
		CallGoFunc(strcat, 2).
		CallGoFuncv(sprintf, 4, 4).
		Resolve()

	ctx := NewContext(code)
	ctx.Exec(0, code.Len())
	if v := checkPop(ctx); v != "Hello, 1.3, 1, xsw" {
		t.Fatal("format 1.3 1 `x` `sw` strcat sprintf != `Hello, 1.3, 1, xsw`, ret =", v)
	}
}

func TestLargeArity(t *testing.T) {
	sprint, kind, ok := defaultImpl.FindGoPackage("").Find("Sprint")
	if !ok || kind != SymbolFuncv {
		t.Fatal("Find failed: Sprint")
	}

	b := newBuilder()
	ret := ""
	for i := 0; i < bitsFuncvArityMax+1; i++ {
		b.Push("32")
		ret += "32"
	}
	code := b.
		CallGoFuncv(GoFuncvAddr(sprint), bitsFuncvArityMax+1, bitsFuncvArityMax+1).
		Resolve()

	ctx := NewContext(code)
	ctx.Exec(0, code.Len())
	if v := checkPop(ctx); v != ret {
		t.Fatal("32 times(1024) sprint != `32` times(1024), ret =", v)
	}
}

func TestType(t *testing.T) {
	typ, ok := I.FindType("Context")
	if !ok {
		t.Fatal("FindType failed: Context not found")
	}
	fmt.Println(typ)

	typ, ok = FindGoPackage("").FindType("rune")
	if !ok || typ != TyRune {
		t.Fatal("FindType failed: rune not found")
	}
	fmt.Println(typ)
}
<|endoftext|>


<commit_before>package core_test

import (
	"github.com/blackbass1988/access_logs_stats/core"
	"testing"
)

func TestYamlConfig(t *testing.T) {
	filepath := "../config.yaml.example"
	config, err := core.NewConfig(filepath)

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

	testConfig(t, config)
}

func TestJsonConfig(t *testing.T) {
	filepath := "../config.json.example"

	config, err := core.NewConfig(filepath)

	if err != nil {
		t.Fatal(err)
	}
	testConfig(t, config)

}

func testConfig(t *testing.T, config core.Config) {

	if config.InputDsn != "file:foo.txt" {
		t.Error(
			"expected file:foo.txt",
			"actual ",
			config.InputDsn,
		)
	}

	if config.Rex.String() == "" {
		t.Error(
			"expected != \"\"",
			"actual ",
			config.Rex.String(),
		)
	}

	if config.Period.String() != "10s" {
		t.Error(
			"expected 10s",
			"actual ",
			config.Period.String(),
		)
	}

	if len(config.Outputs) != 2 {
		t.Error(
			"expected 2 senders",
			"actual ",
			len(config.Outputs),
		)
	}

	for _, output := range config.Outputs {

		if output.Type == "" {
			t.Error("type must != '' but was ''")
		}

		if output.Settings == nil {
			t.Error("settings was nil")
		}

	}

	if len(config.Filters) != 1 {
		t.Error("filters count expected 1 but was ", len(config.Filters))
		t.FailNow()
	}

	f := config.Filters[0]

	if f.String() != ".+" {
		t.Errorf("filter. Expected [.+] . Actual [%s]", f.String())
	}

	if f.Prefix != "prefix2_" {
		t.Error("filter Prefix. Expected prefix2_ . Actual ", f.Prefix)
	}

	if len(f.Items) != 2 {
		t.Error("filter Items. Expected 2. Actual ", len(f.Items))
	}

	oneItem := f.Items[0]
	if oneItem.Field != "code" {
		t.Error("first filter item must be code but was ", oneItem.Field)
	}

	if len(oneItem.Metrics) != 4 {
		t.Error("len(oneItem.Metrics) must be 3 but was ", len(oneItem.Metrics))
	}

	if len(config.Counts) != 2 {
		t.Error("config.Counts. Expected 2. Actual ", len(config.Counts))
	}

	if _, ok := config.Counts["code"]; !ok {
		t.Error("config.Aggregates[code]. Expected code. Actual ", config.Counts["code"])
	}

	if _, ok := config.Counts["time"]; !ok {
		t.Error("config.Aggregates[time]. Expected time. Actual ", config.Counts["time"])
	}

	if len(config.Aggregates) != 1 {
		t.Error("config.Counts. Expected 1. Actual ", len(config.Aggregates))
	}

	if _, ok := config.Aggregates["time"]; !ok {
		t.Error("config.Aggregates[time]. Expected time. Actual ", config.Aggregates["time"])
	}
}
<commit_msg>cyclo in core/config_test.go<commit_after>package core_test

import (
	"github.com/blackbass1988/access_logs_stats/core"
	"testing"
)

func TestYamlConfig(t *testing.T) {
	filepath := "../config.yaml.example"
	config, err := core.NewConfig(filepath)

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

	testConfig(t, config)
}

func TestJsonConfig(t *testing.T) {
	filepath := "../config.json.example"

	config, err := core.NewConfig(filepath)

	if err != nil {
		t.Fatal(err)
	}
	testConfig(t, config)

}

func testConfig(t *testing.T, config core.Config) {

	if config.InputDsn != "file:foo.txt" {
		t.Error(
			"expected file:foo.txt",
			"actual ",
			config.InputDsn,
		)
	}

	if config.Rex.String() == "" {
		t.Error(
			"expected != \"\"",
			"actual ",
			config.Rex.String(),
		)
	}

	if config.Period.String() != "10s" {
		t.Error(
			"expected 10s",
			"actual ",
			config.Period.String(),
		)
	}

	if len(config.Outputs) != 2 {
		t.Error(
			"expected 2 senders",
			"actual ",
			len(config.Outputs),
		)
	}

	for _, output := range config.Outputs {

		if output.Type == "" {
			t.Error("type must != '' but was ''")
		}

		if output.Settings == nil {
			t.Error("settings was nil")
		}

	}

	if len(config.Filters) != 1 {
		t.Error("filters count expected 1 but was ", len(config.Filters))
		t.FailNow()
	}

	f := config.Filters[0]

	checkFilter(t, f)
	checkConfig(t, config)
}

func checkFilter(t *testing.T, f *core.Filter) {
	if f.String() != ".+" {
		t.Errorf("filter. Expected [.+] . Actual [%s]", f.String())
	}

	if f.Prefix != "prefix2_" {
		t.Error("filter Prefix. Expected prefix2_ . Actual ", f.Prefix)
	}

	if len(f.Items) != 2 {
		t.Error("filter Items. Expected 2. Actual ", len(f.Items))
	}

	oneItem := f.Items[0]
	if oneItem.Field != "code" {
		t.Error("first filter item must be code but was ", oneItem.Field)
	}

	if len(oneItem.Metrics) != 4 {
		t.Error("len(oneItem.Metrics) must be 3 but was ", len(oneItem.Metrics))
	}
}

func checkConfig(t *testing.T, config core.Config) {

	if len(config.Counts) != 2 {
		t.Error("config.Counts. Expected 2. Actual ", len(config.Counts))
	}

	if _, ok := config.Counts["code"]; !ok {
		t.Error("config.Aggregates[code]. Expected code. Actual ", config.Counts["code"])
	}

	if _, ok := config.Counts["time"]; !ok {
		t.Error("config.Aggregates[time]. Expected time. Actual ", config.Counts["time"])
	}

	if len(config.Aggregates) != 1 {
		t.Error("config.Counts. Expected 1. Actual ", len(config.Aggregates))
	}

	if _, ok := config.Aggregates["time"]; !ok {
		t.Error("config.Aggregates[time]. Expected time. Actual ", config.Aggregates["time"])
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"

	_ "github.com/Shopify/themekit/cmd/static" // This will import the asset bundle
	"github.com/Shopify/themekit/src/cmdutil"
	"github.com/Shopify/themekit/src/colors"
	"github.com/Shopify/themekit/src/shopify"
	"github.com/Shopify/themekit/src/static"
)

var newCmd = &cobra.Command{
	Use:   "new",
	Short: "New will create theme using Shopify Timber",
	Long: `New will create a new theme on Shopify, generate a minimal template of
	all the required files that a theme needs to be functional and then setup
  your config file for working with your new theme.

  For more documentation please see http://shopify.github.io/themekit/commands/#new
  `,
	RunE: func(cmd *cobra.Command, args []string) error {
		// This is a hack to get around theme ID validation for the list operation which doesn't need it
		flags.ThemeID = "1337"
		return cmdutil.ForDefaultClient(flags, args, func(ctx *cmdutil.Ctx) error {
			return newTheme(ctx, static.Unbundle)
		})
	},
}

func newTheme(ctx *cmdutil.Ctx, generate func(ctx *cmdutil.Ctx) error) error {
	theme, err := ctx.Client.CreateNewTheme(ctx.Flags.Name)
	if err != nil {
		if err == shopify.ErrThemeNameRequired {
			return fmt.Errorf("a theme name is required, please use the --name flag to define it")
		}
		return err
	}
	ctx.Log.Printf("[%s] theme created", colors.Yellow(ctx.Env.Domain))

	ctx.Env.ThemeID = fmt.Sprintf("%v", theme.ID)
	if err := createConfig(ctx); err != nil {
		return err
	}
	ctx.Log.Printf("[%s] config created", colors.Yellow(ctx.Env.Domain))

	if err := generate(ctx); err != nil {
		return err
	}

	ctx.Log.Printf("[%s] uploading new files to shopify", colors.Yellow(ctx.Env.Domain))
	return deploy(ctx)
}
<commit_msg>Update help description for new command<commit_after>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"

	_ "github.com/Shopify/themekit/cmd/static" // This will import the asset bundle
	"github.com/Shopify/themekit/src/cmdutil"
	"github.com/Shopify/themekit/src/colors"
	"github.com/Shopify/themekit/src/shopify"
	"github.com/Shopify/themekit/src/static"
)

var newCmd = &cobra.Command{
	Use:   "new",
	Short: "New will create theme using Shopify Timber in the same directory where it gets called from",
	Long: `New will create a new theme on Shopify, generate a minimal template of
	all the required files that a theme needs to be functional and then setup
	your config file for working with your new theme. Note: by default the new command will generate files in
	the same directory it's called from. Use the --dir flag to specify a custom directory where the generated files
	should be placed.

  For more documentation please see http://shopify.github.io/themekit/commands/#new
  `,
	RunE: func(cmd *cobra.Command, args []string) error {
		// This is a hack to get around theme ID validation for the list operation which doesn't need it
		flags.ThemeID = "1337"
		return cmdutil.ForDefaultClient(flags, args, func(ctx *cmdutil.Ctx) error {
			return newTheme(ctx, static.Unbundle)
		})
	},
}

func newTheme(ctx *cmdutil.Ctx, generate func(ctx *cmdutil.Ctx) error) error {
	theme, err := ctx.Client.CreateNewTheme(ctx.Flags.Name)
	if err != nil {
		if err == shopify.ErrThemeNameRequired {
			return fmt.Errorf("a theme name is required, please use the --name flag to define it")
		}
		return err
	}
	ctx.Log.Printf("[%s] theme created", colors.Yellow(ctx.Env.Domain))

	ctx.Env.ThemeID = fmt.Sprintf("%v", theme.ID)
	if err := createConfig(ctx); err != nil {
		return err
	}
	ctx.Log.Printf("[%s] config created", colors.Yellow(ctx.Env.Domain))

	if err := generate(ctx); err != nil {
		return err
	}

	ctx.Log.Printf("[%s] uploading new files to shopify", colors.Yellow(ctx.Env.Domain))
	return deploy(ctx)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"strings"
	"syscall"

	"github.com/bfirsh/whalebrew/packages"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh/terminal"
)

func init() {
	RootCmd.AddCommand(runCommand)
}

var runCommand = &cobra.Command{
	Use:                "run PACKAGEPATH [ARGS ...]",
	Short:              "Run a package",
	DisableFlagParsing: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) < 1 {
			return cmd.Help()
		}

		pkg, err := packages.LoadPackageFromPath(args[0])
		if err != nil {
			return err
		}
		dockerPath, err := exec.LookPath("docker")
		if err != nil {
			return err
		}

		cwd, err := os.Getwd()
		if err != nil {
			return err
		}
		dockerArgs := []string{
			dockerPath,
			"run",
			"--interactive",
			"--rm",
			"--workdir", "/workdir",
			"-v", fmt.Sprintf("%s:/workdir", cwd),
		}
		if terminal.IsTerminal(int(os.Stdin.Fd())) {
			dockerArgs = append(dockerArgs, "--tty")
		}
		for _, volume := range pkg.Volumes {
			// special case expanding home directory
			if strings.HasPrefix(volume, "~/") {
				user, err := user.Current()
				if err != nil {
					return err
				}
				volume = user.HomeDir + volume[1:]
			}
			dockerArgs = append(dockerArgs, "-v")
			dockerArgs = append(dockerArgs, volume)
		}
		for _, envvar := range pkg.Environment {
			dockerArgs = append(dockerArgs, "-e")
			dockerArgs = append(dockerArgs, envvar)
		}
		for _, portmap := range pkg.Ports {
			dockerArgs = append(dockerArgs, "-p")
			dockerArgs = append(dockerArgs, portmap)
		}
		dockerArgs = append(dockerArgs, pkg.Image)
		dockerArgs = append(dockerArgs, args[1:]...)

		return syscall.Exec(dockerPath, dockerArgs, os.Environ())
	},
}
<commit_msg>Support to run and return status code on Windows<commit_after>package cmd

import (
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"strings"
	"syscall"

	"runtime"

	"github.com/bfirsh/whalebrew/packages"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh/terminal"
)

func init() {
	RootCmd.AddCommand(runCommand)
}

var runCommand = &cobra.Command{
	Use:                "run PACKAGEPATH [ARGS ...]",
	Short:              "Run a package",
	DisableFlagParsing: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) < 1 {
			return cmd.Help()
		}

		pkg, err := packages.LoadPackageFromPath(args[0])
		if err != nil {
			return err
		}
		dockerPath, err := exec.LookPath("docker")
		if err != nil {
			return err
		}

		cwd, err := os.Getwd()
		if err != nil {
			return err
		}
		dockerArgs := []string{
			dockerPath,
			"run",
			"--interactive",
			"--rm",
			"--workdir", "/workdir",
			"-v", fmt.Sprintf("%s:/workdir", cwd),
		}
		if terminal.IsTerminal(int(os.Stdin.Fd())) {
			dockerArgs = append(dockerArgs, "--tty")
		}
		for _, volume := range pkg.Volumes {
			// special case expanding home directory
			if strings.HasPrefix(volume, "~/") {
				user, err := user.Current()
				if err != nil {
					return err
				}
				volume = user.HomeDir + volume[1:]
			}
			dockerArgs = append(dockerArgs, "-v")
			dockerArgs = append(dockerArgs, volume)
		}
		for _, envvar := range pkg.Environment {
			dockerArgs = append(dockerArgs, "-e")
			dockerArgs = append(dockerArgs, envvar)
		}
		for _, portmap := range pkg.Ports {
			dockerArgs = append(dockerArgs, "-p")
			dockerArgs = append(dockerArgs, portmap)
		}
		dockerArgs = append(dockerArgs, pkg.Image)
		dockerArgs = append(dockerArgs, args[1:]...)

		if runtime.GOOS == "windows" {
			dockerCmd := exec.Command(dockerPath, dockerArgs[1:]...)
			dockerCmd.Env = os.Environ()
			dockerCmd.Stdin = os.Stdin
			dockerCmd.Stdout = os.Stdout
			dockerCmd.Stderr = os.Stderr

			exitStatus := 1
			if err := dockerCmd.Run(); err != nil {
				if exitError, ok := err.(*exec.ExitError); ok {
					if ws, ok := exitError.Sys().(syscall.WaitStatus); ok {
						exitStatus = ws.ExitStatus()
					}
				}
			} else {
				ws := dockerCmd.ProcessState.Sys().(syscall.WaitStatus)
				exitStatus = ws.ExitStatus()
			}
			os.Exit(exitStatus)
		}

		return syscall.Exec(dockerPath, dockerArgs, os.Environ())
	},
}
<|endoftext|>


<commit_before>// Package tls provides the functionality to Ponzu systems to encrypt HTTP traffic
// through the ability to generate self-signed certificates for local development
// and fetch/update production certificates from Let's Encrypt.
package tls

import (
	"crypto/tls"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/ponzu-cms/ponzu/system/db"
	"golang.org/x/crypto/acme/autocert"
)

var m autocert.Manager

// setup attempts to locate or create the cert cache directory and the certs for TLS encryption
func setup() {
	pwd, err := os.Getwd()
	if err != nil {
		log.Fatalln("Couldn't find working directory to locate or save certificates.")
	}

	cache := autocert.DirCache(filepath.Join(pwd, "system", "tls", "certs"))
	if _, err := os.Stat(string(cache)); os.IsNotExist(err) {
		err := os.MkdirAll(string(cache), os.ModePerm|os.ModeDir)
		if err != nil {
			log.Fatalln("Couldn't create cert directory at", cache)
		}
	}

	// get host/domain and email from Config to use for TLS request to Let's encryption.
	// we will fail fatally if either are not found since Let's Encrypt will rate-limit
	// and sending incomplete requests is wasteful and guaranteed to fail its check
	host, err := db.Config("domain")
	if err != nil {
		log.Fatalln("Error identifying host/domain during TLS set-up.", err)
	}

	if host == nil {
		log.Fatalln("No 'domain' field set in Configuration. Please add a domain before attempting to make certificates.")
	}
	fmt.Println("Using", string(host), "as host/domain for certificate...")
	fmt.Println("NOTE: if the host/domain is not configured properly or is unreachable, HTTPS set-up will fail.")

	email, err := db.Config("admin_email")
	if err != nil {
		log.Fatalln("Error identifying admin email during TLS set-up.", err)
	}

	if email == nil {
		log.Fatalln("No 'admin_email' field set in Configuration. Please add an admin email before attempting to make certificates.")
	}
	fmt.Println("Using", string(email), "as contact email for certificate...")

	m = autocert.Manager{
		Prompt:      autocert.AcceptTOS,
		Cache:       cache,
		HostPolicy:  autocert.HostWhitelist(string(host)),
		RenewBefore: time.Hour * 24 * 30,
		Email:       string(email),
	}

}

// Enable runs the setup for creating or locating production certificates and
// starts the TLS server
func Enable() {
	setup()

	server := &http.Server{
		Addr:      fmt.Sprintf(":%s", db.ConfigCache("https_port").(string)),
		TLSConfig: &tls.Config{GetCertificate: m.GetCertificate},
	}

	log.Fatalln(server.ListenAndServeTLS("", ""))
}
<commit_msg>[core] system/tls: implement handler for "http-01" challenge<commit_after>// Package tls provides the functionality to Ponzu systems to encrypt HTTP traffic
// through the ability to generate self-signed certificates for local development
// and fetch/update production certificates from Let's Encrypt.
package tls

import (
	"crypto/tls"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/ponzu-cms/ponzu/system/db"
	"golang.org/x/crypto/acme/autocert"
)

// newManager attempts to locate or create the cert cache directory and the
// certs for TLS encryption and returns an autocert.Manager
func newManager() autocert.Manager {
	pwd, err := os.Getwd()
	if err != nil {
		log.Fatalln("Couldn't find working directory to locate or save certificates.")
	}

	cache := autocert.DirCache(filepath.Join(pwd, "system", "tls", "certs"))
	if _, err := os.Stat(string(cache)); os.IsNotExist(err) {
		err := os.MkdirAll(string(cache), os.ModePerm|os.ModeDir)
		if err != nil {
			log.Fatalln("Couldn't create cert directory at", cache)
		}
	}

	// get host/domain and email from Config to use for TLS request to Let's encryption.
	// we will fail fatally if either are not found since Let's Encrypt will rate-limit
	// and sending incomplete requests is wasteful and guaranteed to fail its check
	host, err := db.Config("domain")
	if err != nil {
		log.Fatalln("Error identifying host/domain during TLS set-up.", err)
	}

	if host == nil {
		log.Fatalln("No 'domain' field set in Configuration. Please add a domain before attempting to make certificates.")
	}
	fmt.Println("Using", string(host), "as host/domain for certificate...")
	fmt.Println("NOTE: if the host/domain is not configured properly or is unreachable, HTTPS set-up will fail.")

	email, err := db.Config("admin_email")
	if err != nil {
		log.Fatalln("Error identifying admin email during TLS set-up.", err)
	}

	if email == nil {
		log.Fatalln("No 'admin_email' field set in Configuration. Please add an admin email before attempting to make certificates.")
	}
	fmt.Println("Using", string(email), "as contact email for certificate...")

	return autocert.Manager{
		Prompt:      autocert.AcceptTOS,
		Cache:       cache,
		HostPolicy:  autocert.HostWhitelist(string(host)),
		RenewBefore: time.Hour * 24 * 30,
		Email:       string(email),
	}
}

// Enable runs the setup for creating or locating production certificates and
// starts the TLS server
func Enable() {
	m := newManager()

	server := &http.Server{
		Addr:      fmt.Sprintf(":%s", db.ConfigCache("https_port").(string)),
		TLSConfig: &tls.Config{GetCertificate: m.GetCertificate},
	}

	// launch http listener for "http-01" ACME challenge
	go http.ListenAndServe(":http", m.HTTPHandler(nil))

	log.Fatalln(server.ListenAndServeTLS("", ""))
}
<|endoftext|>


<commit_before>package ergo

import (
	"net/http"

	"github.com/wlMalk/ergo/validation"
)

type Request struct {
	*http.Request
	Input      map[string]validation.Valuer
	pathParams map[string]string
	operation  *Operation // Operation object
}

func NewRequest(httpRequest *http.Request) *Request {
	return &Request{
		Request:    httpRequest,
		pathParams: map[string]string{},
		Input:      map[string]validation.Valuer{},
	}
}

// Req returns the request.
func (req *Request) Req() *http.Request {
	return req.Request
}

// Param returns the input parameter value by its name.
func (req *Request) Param(name string) validation.Valuer {
	return req.Input[name]
}

// ParamOk returns the input parameter value by its name.
func (req *Request) ParamOk(name string) (validation.Valuer, bool) {
	p, ok := req.Input[name]
	return p, ok
}

// Params returns a map of input parameters values by their names.
// If no names given then it returns r.Input
func (req *Request) Params(names ...string) map[string]validation.Valuer {
	if len(names) == 0 {
		return req.Input
	}
	params := map[string]validation.Valuer{}
	for _, n := range names {
		p, ok := req.Input[n]
		if !ok {
			continue
		}
		params[n] = p
	}
	return params
}

func (req *Request) GetOperation() Operationer {
	return req.operation
}

func (req *Request) SetPathParams(params map[string]string) {
	req.pathParams = params
}
<commit_msg>Expose PathParams<commit_after>package ergo

import (
	"net/http"

	"github.com/wlMalk/ergo/validation"
)

type Request struct {
	*http.Request
	Input      map[string]validation.Valuer
	PathParams map[string]string
	operation  *Operation // Operation object
}

func NewRequest(httpRequest *http.Request) *Request {
	return &Request{
		Request:    httpRequest,
		PathParams: map[string]string{},
		Input:      map[string]validation.Valuer{},
	}
}

// Req returns the request.
func (req *Request) Req() *http.Request {
	return req.Request
}

// Param returns the input parameter value by its name.
func (req *Request) Param(name string) validation.Valuer {
	return req.Input[name]
}

// ParamOk returns the input parameter value by its name.
func (req *Request) ParamOk(name string) (validation.Valuer, bool) {
	p, ok := req.Input[name]
	return p, ok
}

// Params returns a map of input parameters values by their names.
// If no names given then it returns r.Input
func (req *Request) Params(names ...string) map[string]validation.Valuer {
	if len(names) == 0 {
		return req.Input
	}
	params := map[string]validation.Valuer{}
	for _, n := range names {
		p, ok := req.Input[n]
		if !ok {
			continue
		}
		params[n] = p
	}
	return params
}

func (req *Request) GetOperation() Operationer {
	return req.operation
}
<|endoftext|>


<commit_before>package openstack

import (
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
)

func TestAccNetworkingV2Port_importBasic(t *testing.T) {
	resourceName := "openstack_networking_port_v2.port_1"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckNetworkingV2PortDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccNetworkingV2Port_basic,
			},

			resource.TestStep{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}
<commit_msg>provider/openstack: Ignore fixed_ip when importing ports (#13563)<commit_after>package openstack

import (
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
)

func TestAccNetworkingV2Port_importBasic(t *testing.T) {
	resourceName := "openstack_networking_port_v2.port_1"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckNetworkingV2PortDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccNetworkingV2Port_basic,
			},

			resource.TestStep{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"fixed_ip",
				},
			},
		},
	})
}
<|endoftext|>


<commit_before>package sms // import "github.com/mndrix/sms-over-xmpp"
import (
	"crypto/rand"
	"encoding/base32"
	"encoding/xml"
	"log"
	"strings"

	xco "github.com/mndrix/go-xco"
	"github.com/pkg/errors"
)

// gatewayProcess is the piece which sits between the XMPP and HTTP
// processes, translating between their different protocols.
type gatewayProcess struct {
	// fields shared with Component. see docs there
	config     Config
	receiptFor map[string]*xco.Message
	smsRx      <-chan rxSms
	xmppRx     <-chan *xco.Message
	xmppTx     chan<- *xco.Message
}

func (g *gatewayProcess) run() <-chan struct{} {
	healthCh := make(chan struct{})
	go g.loop(healthCh)
	return healthCh
}

func (g *gatewayProcess) loop(healthCh chan<- struct{}) {
	defer func() { close(healthCh) }()

	for {
		select {
		case rxSms := <-g.smsRx:
			var err error
			errCh := rxSms.ErrCh()

			switch x := rxSms.(type) {
			case *rxSmsMessage:
				err = g.sms2xmpp(x.sms)
			case *rxSmsStatus:
				switch x.status {
				case smsDelivered:
					err = g.smsDelivered(x.id)
				default:
					log.Panicf("unexpected SMS status: %d", x.status)
				}
			default:
				log.Panicf("unexpected rxSms type: %#v", rxSms)
			}

			go func() { errCh <- err }()
		case msg := <-g.xmppRx:
			err := g.xmpp2sms(msg)
			if err != nil {
				log.Printf("ERROR: converting XMPP to SMS: %s", err)
				return
			}
		}
	}
}

func (g *gatewayProcess) sms2xmpp(sms *Sms) error {

	// convert author's phone number into XMPP address
	from, err := g.config.PhoneToAddress(sms.From)
	switch err {
	case nil:
		// all is well. proceed
	case ErrIgnoreMessage:
		msg := "ignored based on From address"
		log.Println(msg)
		return nil
	default:
		return errors.Wrap(err, "From address "+sms.From)
	}

	// convert recipient's phone number into XMPP address
	to, err := g.config.PhoneToAddress(sms.To)
	switch err {
	case nil:
		// all is well. proceed
	case ErrIgnoreMessage:
		msg := "ignored based on To address"
		log.Println(msg)
		return nil
	default:
		return errors.Wrap(err, "To address "+sms.To)
	}

	// deliver message over XMPP
	msg := &xco.Message{
		XMLName: xml.Name{
			Local: "message",
			Space: "jabber:component:accept",
		},

		Header: xco.Header{
			From: from,
			To:   to,
			ID:   NewId(),
		},
		Type: "chat",
		Body: sms.Body,
	}
	go func() { g.xmppTx <- msg }()
	return nil
}

func (g *gatewayProcess) smsDelivered(smsId string) error {
	if receipt, ok := g.receiptFor[smsId]; ok {
		go func() { g.xmppTx <- receipt }()
		log.Printf("Sent SMS delivery receipt")
		delete(g.receiptFor, smsId)
	}
	return nil
}

func (g *gatewayProcess) xmpp2sms(m *xco.Message) error {
	// convert recipient address into a phone number
	toPhone, err := g.config.AddressToPhone(m.To)
	switch err {
	case nil:
		// all is well. we'll continue below
	case ErrIgnoreMessage:
		return nil
	default:
		return errors.Wrap(err, "converting 'to' address to phone")
	}

	// convert author's address into a phone number
	fromPhone, err := g.config.AddressToPhone(m.From)
	switch err {
	case nil:
		// all is well. we'll continue below
	case ErrIgnoreMessage:
		return nil
	default:
		return errors.Wrap(err, "converting 'from' address to phone")
	}

	// choose an SMS provider
	provider, err := g.config.SmsProvider()
	switch err {
	case nil:
		// all is well. we'll continue below
	case ErrIgnoreMessage:
		return nil
	default:
		return errors.Wrap(err, "choosing an SMS provider")
	}

	// send the message
	id, err := provider.SendSms(&Sms{
		From: fromPhone,
		To:   toPhone,
		Body: m.Body,
	})
	if err != nil {
		return errors.Wrap(err, "sending SMS")
	}
	log.Printf("Sent SMS with ID %s", id)

	// prepare to handle delivery receipts
	if m.ReceiptRequest != nil && id != "" {
		receipt := xco.Message{
			Header: xco.Header{
				From: m.Header.To,
				To:   m.Header.From,
				ID:   NewId(),
			},
			ReceiptAck: &xco.ReceiptAck{
				Id: m.Header.ID,
			},
			XMLName: m.XMLName,
		}
		if len(g.receiptFor) > 10 { // don't get too big
			log.Printf("clearing pending receipts queue")
			g.receiptFor = make(map[string]*xco.Message)
		}
		g.receiptFor[id] = &receipt
		log.Printf("Waiting to send receipt: %#v", receipt)
	}

	return nil
}

// NewId generates a random string which is suitable as an XMPP stanza
// ID.  The string contains enough entropy to be universally unique.
func NewId() string {
	// generate 128 random bits (6 more than standard UUID)
	bytes := make([]byte, 16)
	_, err := rand.Read(bytes)
	if err != nil {
		panic(err)
	}

	// convert them to base 32 encoding
	s := base32.StdEncoding.EncodeToString(bytes)
	return strings.ToLower(strings.TrimRight(s, "="))
}
<commit_msg>Simplify sms2xmpp<commit_after>package sms // import "github.com/mndrix/sms-over-xmpp"
import (
	"crypto/rand"
	"encoding/base32"
	"encoding/xml"
	"log"
	"strings"

	xco "github.com/mndrix/go-xco"
	"github.com/pkg/errors"
)

// gatewayProcess is the piece which sits between the XMPP and HTTP
// processes, translating between their different protocols.
type gatewayProcess struct {
	// fields shared with Component. see docs there
	config     Config
	receiptFor map[string]*xco.Message
	smsRx      <-chan rxSms
	xmppRx     <-chan *xco.Message
	xmppTx     chan<- *xco.Message
}

func (g *gatewayProcess) run() <-chan struct{} {
	healthCh := make(chan struct{})
	go g.loop(healthCh)
	return healthCh
}

func (g *gatewayProcess) loop(healthCh chan<- struct{}) {
	defer func() { close(healthCh) }()

	for {
		select {
		case rxSms := <-g.smsRx:
			var err error
			errCh := rxSms.ErrCh()

			switch x := rxSms.(type) {
			case *rxSmsMessage:
				err = g.sms2xmpp(x.sms)
			case *rxSmsStatus:
				switch x.status {
				case smsDelivered:
					err = g.smsDelivered(x.id)
				default:
					log.Panicf("unexpected SMS status: %d", x.status)
				}
			default:
				log.Panicf("unexpected rxSms type: %#v", rxSms)
			}

			go func() { errCh <- err }()
		case msg := <-g.xmppRx:
			err := g.xmpp2sms(msg)
			if err != nil {
				log.Printf("ERROR: converting XMPP to SMS: %s", err)
				return
			}
		}
	}
}

func (g *gatewayProcess) sms2xmpp(sms *Sms) error {
	var err error
	msg := &xco.Message{
		XMLName: xml.Name{
			Local: "message",
			Space: "jabber:component:accept",
		},
		Header: xco.Header{
			ID: NewId(),
		},
		Type: "chat",
		Body: sms.Body,
	}

	// convert author's phone number into XMPP address
	msg.Header.From, err = g.config.PhoneToAddress(sms.From)
	if err == nil {
		msg.Header.To, err = g.config.PhoneToAddress(sms.To)
	}
	switch err {
	case nil:
		go func() { g.xmppTx <- msg }()
	case ErrIgnoreMessage:
		log.Println("ignoring message based on phone number")
	default:
		return errors.Wrap(err, "sms2xmpp")
	}

	return nil
}

func (g *gatewayProcess) smsDelivered(smsId string) error {
	if receipt, ok := g.receiptFor[smsId]; ok {
		go func() { g.xmppTx <- receipt }()
		log.Printf("Sent SMS delivery receipt")
		delete(g.receiptFor, smsId)
	}
	return nil
}

func (g *gatewayProcess) xmpp2sms(m *xco.Message) error {
	// convert recipient address into a phone number
	toPhone, err := g.config.AddressToPhone(m.To)
	switch err {
	case nil:
		// all is well. we'll continue below
	case ErrIgnoreMessage:
		return nil
	default:
		return errors.Wrap(err, "converting 'to' address to phone")
	}

	// convert author's address into a phone number
	fromPhone, err := g.config.AddressToPhone(m.From)
	switch err {
	case nil:
		// all is well. we'll continue below
	case ErrIgnoreMessage:
		return nil
	default:
		return errors.Wrap(err, "converting 'from' address to phone")
	}

	// choose an SMS provider
	provider, err := g.config.SmsProvider()
	switch err {
	case nil:
		// all is well. we'll continue below
	case ErrIgnoreMessage:
		return nil
	default:
		return errors.Wrap(err, "choosing an SMS provider")
	}

	// send the message
	id, err := provider.SendSms(&Sms{
		From: fromPhone,
		To:   toPhone,
		Body: m.Body,
	})
	if err != nil {
		return errors.Wrap(err, "sending SMS")
	}
	log.Printf("Sent SMS with ID %s", id)

	// prepare to handle delivery receipts
	if m.ReceiptRequest != nil && id != "" {
		receipt := xco.Message{
			Header: xco.Header{
				From: m.Header.To,
				To:   m.Header.From,
				ID:   NewId(),
			},
			ReceiptAck: &xco.ReceiptAck{
				Id: m.Header.ID,
			},
			XMLName: m.XMLName,
		}
		if len(g.receiptFor) > 10 { // don't get too big
			log.Printf("clearing pending receipts queue")
			g.receiptFor = make(map[string]*xco.Message)
		}
		g.receiptFor[id] = &receipt
		log.Printf("Waiting to send receipt: %#v", receipt)
	}

	return nil
}

// NewId generates a random string which is suitable as an XMPP stanza
// ID.  The string contains enough entropy to be universally unique.
func NewId() string {
	// generate 128 random bits (6 more than standard UUID)
	bytes := make([]byte, 16)
	_, err := rand.Read(bytes)
	if err != nil {
		panic(err)
	}

	// convert them to base 32 encoding
	s := base32.StdEncoding.EncodeToString(bytes)
	return strings.ToLower(strings.TrimRight(s, "="))
}
<|endoftext|>