<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package loader

import (
	"testing"

	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
)

func TestDir(t *testing.T) {
	testdata := map[string]string{
		"/path/to/file.txt": "/path/to",
		"-":                 "/",
	}
	for name, dir := range testdata {
		t.Run("path="+name, func(t *testing.T) {
			assert.Equal(t, dir, Dir(name))
		})
	}
}

func TestLoad(t *testing.T) {
	t.Run("Blank", func(t *testing.T) {
		_, err := Load(nil, "/", "")
		assert.EqualError(t, err, "local or remote path required")
	})

	t.Run("Protocol", func(t *testing.T) {
		_, err := Load(nil, "/", "https://httpbin.org/html")
		assert.EqualError(t, err, "imports should not contain a protocol")
	})

	t.Run("Local", func(t *testing.T) {
		fs := afero.NewMemMapFs()
		assert.NoError(t, fs.MkdirAll("/path/to", 0755))
		assert.NoError(t, afero.WriteFile(fs, "/path/to/file.txt", []byte("hi"), 0644))

		testdata := map[string]struct{ pwd, path string }{
			"Absolute": {"/path", "/path/to/file.txt"},
			"Relative": {"/path", "./to/file.txt"},
			"Adjacent": {"/path/to", "./file.txt"},
		}
		for name, data := range testdata {
			t.Run(name, func(t *testing.T) {
				src, err := Load(fs, data.pwd, data.path)
				if assert.NoError(t, err) {
					assert.Equal(t, "/path/to/file.txt", src.Filename)
					assert.Equal(t, "hi", string(src.Data))
				}
			})
		}

		t.Run("Nonexistent", func(t *testing.T) {
			_, err := Load(fs, "/", "/nonexistent")
			assert.EqualError(t, err, "open /nonexistent: file does not exist")
		})

		t.Run("Remote Lifting Denied", func(t *testing.T) {
			_, err := Load(fs, "example.com", "/etc/shadow")
			assert.EqualError(t, err, "origin (example.com) not allowed to load local file: /etc/shadow")
		})
	})

	t.Run("Remote", func(t *testing.T) {
		src, err := Load(nil, "/", "httpbin.org/html")
		if assert.NoError(t, err) {
			assert.Equal(t, src.Filename, "httpbin.org/html")
			assert.Contains(t, string(src.Data), "Herman Melville - Moby-Dick")
		}

		t.Run("Absolute", func(t *testing.T) {
			src, err := Load(nil, "httpbin.org", "httpbin.org/robots.txt")
			if assert.NoError(t, err) {
				assert.Equal(t, src.Filename, "httpbin.org/robots.txt")
				assert.Equal(t, string(src.Data), "User-agent: *\nDisallow: /deny\n")
			}
		})

		t.Run("Relative", func(t *testing.T) {
			src, err := Load(nil, "httpbin.org", "./robots.txt")
			if assert.NoError(t, err) {
				assert.Equal(t, src.Filename, "httpbin.org/robots.txt")
				assert.Equal(t, string(src.Data), "User-agent: *\nDisallow: /deny\n")
			}
		})
	})

	t.Run("No _k6=1 Fallback", func(t *testing.T) {
		src, err := Load(nil, "/", "pastebin.com/raw/zngdRRDT")
		if assert.NoError(t, err) {
			assert.Equal(t, src.Filename, "pastebin.com/raw/zngdRRDT")
			assert.Equal(t, "export function fn() {\r\n    return 1234;\r\n}", string(src.Data))
		}
	})
}
<commit_msg>Remove httpbin.org dependency for loading remote scripts<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package loader

import (
	"net/http"
	"testing"

	"github.com/loadimpact/k6/lib/testutils"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
)

func TestDir(t *testing.T) {
	testdata := map[string]string{
		"/path/to/file.txt": "/path/to",
		"-":                 "/",
	}
	for name, dir := range testdata {
		t.Run("path="+name, func(t *testing.T) {
			assert.Equal(t, dir, Dir(name))
		})
	}
}

func TestLoad(t *testing.T) {
	tb := testutils.NewHTTPMultiBin(t)
	sr := tb.Replacer.Replace

	oldHTTPTransport := http.DefaultTransport
	http.DefaultTransport = tb.HTTPTransport

	defer func() {
		tb.Cleanup()
		http.DefaultTransport = oldHTTPTransport
	}()

	t.Run("Blank", func(t *testing.T) {
		_, err := Load(nil, "/", "")
		assert.EqualError(t, err, "local or remote path required")
	})

	t.Run("Protocol", func(t *testing.T) {
		_, err := Load(nil, "/", sr("HTTPSBIN_URL/html"))
		assert.EqualError(t, err, "imports should not contain a protocol")
	})

	t.Run("Local", func(t *testing.T) {
		fs := afero.NewMemMapFs()
		assert.NoError(t, fs.MkdirAll("/path/to", 0755))
		assert.NoError(t, afero.WriteFile(fs, "/path/to/file.txt", []byte("hi"), 0644))

		testdata := map[string]struct{ pwd, path string }{
			"Absolute": {"/path", "/path/to/file.txt"},
			"Relative": {"/path", "./to/file.txt"},
			"Adjacent": {"/path/to", "./file.txt"},
		}
		for name, data := range testdata {
			t.Run(name, func(t *testing.T) {
				src, err := Load(fs, data.pwd, data.path)
				if assert.NoError(t, err) {
					assert.Equal(t, "/path/to/file.txt", src.Filename)
					assert.Equal(t, "hi", string(src.Data))
				}
			})
		}

		t.Run("Nonexistent", func(t *testing.T) {
			_, err := Load(fs, "/", "/nonexistent")
			assert.EqualError(t, err, "open /nonexistent: file does not exist")
		})

		t.Run("Remote Lifting Denied", func(t *testing.T) {
			_, err := Load(fs, "example.com", "/etc/shadow")
			assert.EqualError(t, err, "origin (example.com) not allowed to load local file: /etc/shadow")
		})
	})

	t.Run("Remote", func(t *testing.T) {
		src, err := Load(nil, "/", sr("HTTPSBIN_DOMAIN:HTTPSBIN_PORT/html"))
		if assert.NoError(t, err) {
			assert.Equal(t, src.Filename, sr("HTTPSBIN_DOMAIN:HTTPSBIN_PORT/html"))
			assert.Contains(t, string(src.Data), "Herman Melville - Moby-Dick")
		}

		t.Run("Absolute", func(t *testing.T) {
			src, err := Load(nil, sr("HTTPSBIN_DOMAIN:HTTPSBIN_PORT"), sr("HTTPSBIN_DOMAIN:HTTPSBIN_PORT/robots.txt"))
			if assert.NoError(t, err) {
				assert.Equal(t, src.Filename, sr("HTTPSBIN_DOMAIN:HTTPSBIN_PORT/robots.txt"))
				assert.Equal(t, string(src.Data), "User-agent: *\nDisallow: /deny\n")
			}
		})

		t.Run("Relative", func(t *testing.T) {
			src, err := Load(nil, sr("HTTPSBIN_DOMAIN:HTTPSBIN_PORT"), "./robots.txt")
			if assert.NoError(t, err) {
				assert.Equal(t, src.Filename, sr("HTTPSBIN_DOMAIN:HTTPSBIN_PORT/robots.txt"))
				assert.Equal(t, string(src.Data), "User-agent: *\nDisallow: /deny\n")
			}
		})
	})

	//TODO: remove pastebin.com dependency
	t.Run("No _k6=1 Fallback", func(t *testing.T) {
		src, err := Load(nil, "/", "pastebin.com/raw/zngdRRDT")
		if assert.NoError(t, err) {
			assert.Equal(t, src.Filename, "pastebin.com/raw/zngdRRDT")
			assert.Equal(t, "export function fn() {\r\n    return 1234;\r\n}", string(src.Data))
		}
	})
}
<|endoftext|>


<commit_before>// Title：模板中用的函数集
//
// Description:模板中用的函数集
//
// Author:black
//
// Createtime:2013-08-07 00:47
//
// Version:1.0
//
// 修改历史:版本号 修改日期 修改人 修改说明
//
// 1.0 2013-08-07 00:47 black 创建文档
package lessgo

import (
	"fmt"
	"html/template"
	"math/rand"
	"strings"
	"time"
)

//获取通用model的指定属性的值
func GetPropValue(model *Model, propName string) string {

	if model != nil {
		for _, prop := range model.Props {
			if prop.Name == propName {
				return prop.Value
			}
		}
		Log.Debug("找不到实体", model.Entity.Id, "的属性", propName, "对应的值")
	}

	return ""
}

//获取随机的组件id
func getComponentId(componentType string) string {

	r := rand.New(rand.NewSource(time.Now().UnixNano()))

	str := ""

	for i := 0; i < 4; i++ {
		str += fmt.Sprint(r.Intn(10))
	}

	return componentType + str
}

//模板中的int类型比较比较
func CompareInt(a, b int, compareType string) (flag bool) {

	switch compareType {

	case "eq":
		if a == b {
			flag = true
		} else {
			flag = false
		}
	case "gt":
		if a > b {
			flag = true
		} else {
			flag = false
		}
	case "ge":
		if a >= b {
			flag = true
		} else {
			flag = false
		}
	case "lt":
		if a < b {
			flag = true
		} else {
			flag = false
		}
	case "le":
		if a <= b {
			flag = true
		} else {
			flag = false
		}
	default:
		if a == b {
			flag = true
		} else {
			flag = false
		}
	}

	return flag
}

//模板中字符串比较
func CompareString(a, b string) bool {
	if a == b {
		return true
	} else {
		return false
	}
}

//替换json字符中的换行等特殊符号
func DealJsonString(str string) string {
	str = strings.Replace(str, "\n", " ", -1)
	str = strings.Replace(str, "\n\r", " ", -1)
	str = strings.Replace(str, "\r\n", " ", -1)
	str = strings.Replace(str, "\r", " ", -1)
	str = strings.Replace(str, "\"", "\\\"", -1)
	str = strings.Replace(str, "'", "\\'", -1)

	return str
}

func DealHTMLEscaper(str string) string {
	return template.HTMLEscaper(str)
}
<commit_msg>表格数据bug<commit_after>// Title：模板中用的函数集
//
// Description:模板中用的函数集
//
// Author:black
//
// Createtime:2013-08-07 00:47
//
// Version:1.0
//
// 修改历史:版本号 修改日期 修改人 修改说明
//
// 1.0 2013-08-07 00:47 black 创建文档
package lessgo

import (
	"fmt"
	"html/template"
	"math/rand"
	"strings"
	"time"
)

//获取通用model的指定属性的值
func GetPropValue(model *Model, propName string) string {

	if model != nil {
		for _, prop := range model.Props {
			if prop.Name == propName {
				return prop.Value
			}
		}
		Log.Debug("找不到实体", model.Entity.Id, "的属性", propName, "对应的值")
	}

	return ""
}

//获取随机的组件id
func getComponentId(componentType string) string {

	r := rand.New(rand.NewSource(time.Now().UnixNano()))

	str := ""

	for i := 0; i < 4; i++ {
		str += fmt.Sprint(r.Intn(10))
	}

	return componentType + str
}

//模板中的int类型比较比较
func CompareInt(a, b int, compareType string) (flag bool) {

	switch compareType {

	case "eq":
		if a == b {
			flag = true
		} else {
			flag = false
		}
	case "gt":
		if a > b {
			flag = true
		} else {
			flag = false
		}
	case "ge":
		if a >= b {
			flag = true
		} else {
			flag = false
		}
	case "lt":
		if a < b {
			flag = true
		} else {
			flag = false
		}
	case "le":
		if a <= b {
			flag = true
		} else {
			flag = false
		}
	default:
		if a == b {
			flag = true
		} else {
			flag = false
		}
	}

	return flag
}

//模板中字符串比较
func CompareString(a, b string) bool {
	if a == b {
		return true
	} else {
		return false
	}
}

//替换json字符中的换行等特殊符号
func DealJsonString(str string) string {
	str = strings.Replace(str, "\n", " ", -1)
	str = strings.Replace(str, "\n\r", " ", -1)
	str = strings.Replace(str, "\r\n", " ", -1)
	str = strings.Replace(str, "\r", " ", -1)
	str = strings.Replace(str, "\"", "\\\"", -1)
	//to fixed
	str = strings.Replace(str, "'", "\\\"", -1)

	return str
}

func DealHTMLEscaper(str string) string {
	return template.HTMLEscaper(str)
}
<|endoftext|>


<commit_before>package waveguide

import "html/template"

var tmpl = template.Must(template.New("").Parse(`
{{define "header"}}

                <title>Waveguide</title>

{{end}}

{{define "footer"}}

{{end}}

{{define "root"}}
{{template "header"}}

                	{{if .Spots}}

					Location
					Coordinates
					Conditions
					Wave Height
					Last Updated

					{{range .Spots}}

							{{.HTMLName}}

								{{if .HasCoordinates}}
									❌
									{{.FormattedCoordinates}}
								{{else}}
									<form action="/coords" method="post">
										<input type="hidden" name="path" value="{{.MswPath}}" />
										<input name="coordinates" />
										<button type="submit">Submit</button>
									</form>
								{{end}}

							{{.Cond.Stars}}
							{{.Cond.WaveHeight}}
							{{.Cond.HowLong}} ago

					{{end}}

			{{else}}
				There's no data yet. You can get some by visiting /update_all.
			{{end}}

{{template "footer"}}
{{end}}

{{define "action_response"}}
{{template "header"}}

← home

{{.Message}}

{{template "footer"}}
{{end}}

{{define "map"}}
<!DOCTYPE html>

		<title>Waveguide</title>
		<meta name="viewport" content="initial-scale=1.0">
		<meta charset="utf-8">

{{end}}
`))
<commit_msg>Factor out header and footer from map template.<commit_after>package waveguide

import "html/template"

var tmpl = template.Must(template.New("").Parse(`
{{define "header"}}
<!DOCTYPE html>

                <title>Waveguide</title>
		<meta name="viewport" content="initial-scale=1.0">
		<meta charset="utf-8">

{{end}}

{{define "footer"}}

{{end}}

{{define "root"}}
{{template "header"}}

                	{{if .Spots}}

					Location
					Coordinates
					Conditions
					Wave Height
					Last Updated

					{{range .Spots}}

							{{.HTMLName}}

								{{if .HasCoordinates}}
									❌
									{{.FormattedCoordinates}}
								{{else}}
									<form action="/coords" method="post">
										<input type="hidden" name="path" value="{{.MswPath}}" />
										<input name="coordinates" />
										<button type="submit">Submit</button>
									</form>
								{{end}}

							{{.Cond.Stars}}
							{{.Cond.WaveHeight}}
							{{.Cond.HowLong}} ago

					{{end}}

			{{else}}
				There's no data yet. You can get some by visiting /update_all.
			{{end}}

{{template "footer"}}
{{end}}

{{define "action_response"}}
{{template "header"}}

← home

{{.Message}}

{{template "footer"}}
{{end}}

{{define "map"}}
{{template "header"}}

{{template "footer"}}
{{end}}
`))
<|endoftext|>


<commit_before>package drain

import (
	"fmt"
	"log"
	"net"
	"net/url"
	"os"
)

func SendToDrain(m string, drain string) error {
	u, err := url.Parse(drain)
	if err != nil {
		log.Fatal(err)
	}
	uri := u.Host + u.Path
	switch u.Scheme {
	case "syslog":
		sendToSyslogDrain(m, uri)
	default:
		log.Println(u.Scheme + " drain type is not implemented.")
	}
	return nil
}

func sendToSyslogDrain(m string, drain string) error {
	conn, err := net.Dial("udp", drain)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()
	fmt.Fprintf(conn, m)
	return nil
}

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}
<commit_msg>fix(logger): reduce severity of log failure<commit_after>package drain

import (
	"fmt"
	"log"
	"net"
	"net/url"
	"os"
)

func SendToDrain(m string, drain string) error {
	u, err := url.Parse(drain)
	if err != nil {
		log.Fatal(err)
	}
	uri := u.Host + u.Path
	switch u.Scheme {
	case "syslog":
		sendToSyslogDrain(m, uri)
	default:
		log.Println(u.Scheme + " drain type is not implemented.")
	}
	return nil
}

func sendToSyslogDrain(m string, drain string) error {
	conn, err := net.Dial("udp", drain)
	if err != nil {
		log.Print(err)
	}
	defer conn.Close()
	fmt.Fprintf(conn, m)
	return nil
}

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"text/template"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-check"
	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/fsouza/go-dockerclient"
	"github.com/flynn/flynn/Godeps/_workspace/src/golang.org/x/crypto/ssh"
	"github.com/flynn/flynn/pkg/iotool"
	"github.com/flynn/flynn/pkg/shutdown"
	"github.com/flynn/flynn/test/arg"
	"github.com/flynn/flynn/test/cluster"
	"github.com/flynn/flynn/test/cluster/client"
)

var sshWrapper = template.Must(template.New("ssh").Parse(`
#!/bin/bash

ssh -o LogLevel=FATAL -o IdentitiesOnly=yes -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i {{.SSHKey}} "$@"
`[1:]))

var args *arg.Args
var flynnrc string
var routerIP string
var testCluster *testcluster.Client

func init() {
	args = arg.Parse()
	if args.Stream {
		args.Debug = true
	}
	log.SetFlags(log.Lshortfile)
}

func main() {
	defer shutdown.Exit()

	var err error
	if err = lookupImageURIs(); err != nil {
		log.Fatalf("could not determine image ID: %s", err)
	}

	var res *check.Result
	// defer exiting here so it runs after all other defers
	defer func() {
		if err != nil || res != nil && !res.Passed() {
			if args.DumpLogs {
				if args.Gist {
					exec.Command("flynn-host", "upload-debug-info").Run()
				} else if testCluster != nil {
					testCluster.DumpLogs(&iotool.SafeWriter{W: os.Stdout})
				}
			}
			os.Exit(1)
		}
	}()

	flynnrc = args.Flynnrc
	routerIP = args.RouterIP
	if flynnrc == "" {
		c := cluster.New(args.BootConfig, os.Stdout)
		var rootFS string
		rootFS, err = c.BuildFlynn(args.RootFS, "origin/master", false, false)
		if err != nil {
			if args.Kill {
				c.Shutdown()
			}
			log.Println("could not build flynn: ", err)
			if rootFS != "" {
				os.RemoveAll(rootFS)
			}
			return
		}
		if args.BuildRootFS {
			c.Shutdown()
			fmt.Println("Built Flynn in rootfs:", rootFS)
			return
		} else {
			defer os.RemoveAll(rootFS)
		}
		if err = c.Boot(rootFS, 3, nil, args.Kill); err != nil {
			log.Println("could not boot cluster: ", err)
			return
		}
		if args.Kill {
			defer c.Shutdown()
		}

		if err = createFlynnrc(c); err != nil {
			log.Println(err)
			return
		}
		defer os.RemoveAll(flynnrc)

		routerIP = c.RouterIP
	}

	if args.ClusterAPI != "" {
		testCluster, err = testcluster.NewClient(args.ClusterAPI)
		if err != nil {
			log.Println(err)
			return
		}
	}

	res = check.RunAll(&check.RunConf{
		Filter:           args.Run,
		Stream:           args.Stream,
		Verbose:          args.Debug,
		KeepWorkDir:      args.Debug,
		ConcurrencyLevel: 5,
	})
	fmt.Println(res)
}

var imageURIs = map[string]string{
	"test-apps":           "",
	"postgresql":          "",
	"controller-examples": "",
}

func lookupImageURIs() error {
	d, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		return err
	}
	for name := range imageURIs {
		fullName := "flynn/" + name
		image, err := d.InspectImage(fullName)
		if err != nil {
			return err
		}
		imageURIs[name] = fmt.Sprintf("https://example.com/%s?id=%s", fullName, image.ID)
	}
	return nil
}

type sshData struct {
	Key     string
	Pub     string
	Env     []string
	Cleanup func()
}

func genSSHKey() (*sshData, error) {
	keyFile, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err

	}
	defer keyFile.Close()
	rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	pem.Encode(keyFile, &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(rsaKey),
	})

	pubFile, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err
	}
	defer pubFile.Close()
	rsaPubKey, err := ssh.NewPublicKey(&rsaKey.PublicKey)
	if err != nil {
		return nil, err
	}
	if _, err := pubFile.Write(ssh.MarshalAuthorizedKey(rsaPubKey)); err != nil {
		return nil, err
	}

	wrapperFile, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err
	}
	defer wrapperFile.Close()
	if err := sshWrapper.Execute(wrapperFile, map[string]string{"SSHKey": keyFile.Name()}); err != nil {
		return nil, err
	}
	if err := wrapperFile.Chmod(0700); err != nil {
		return nil, err
	}

	return &sshData{
		Key: keyFile.Name(),
		Pub: pubFile.Name(),
		Env: []string{"GIT_SSH=" + wrapperFile.Name()},
		Cleanup: func() {
			os.RemoveAll(keyFile.Name())
			os.RemoveAll(pubFile.Name())
			os.RemoveAll(wrapperFile.Name())
		},
	}, nil
}

func createFlynnrc(c *cluster.Cluster) error {
	tmpfile, err := ioutil.TempFile("", "flynnrc-")
	if err != nil {
		return err
	}
	path := tmpfile.Name()

	config, err := c.CLIConfig()
	if err != nil {
		os.RemoveAll(path)
		return err
	}

	if err := config.SaveTo(path); err != nil {
		os.RemoveAll(path)
		return err
	}

	flynnrc = path
	return nil
}

type CmdResult struct {
	Cmd    []string
	Output string
	Err    error
}

func flynnEnv(path string) []string {
	env := os.Environ()
	res := make([]string, 0, len(env)+1)
	for _, v := range env {
		if !strings.HasPrefix(v, "FLYNNRC=") {
			res = append(res, v)
		}
	}
	res = append(res, "FLYNNRC="+path)
	return res
}

func flynnCmd(dir string, cmdArgs ...string) *exec.Cmd {
	cmd := exec.Command(args.CLI, cmdArgs...)
	cmd.Env = flynnEnv(flynnrc)
	cmd.Dir = dir
	return cmd
}

func flynn(t *check.C, dir string, args ...string) *CmdResult {
	return run(t, flynnCmd(dir, args...))
}

func debug(t *check.C, v ...interface{}) {
	t.Log(append([]interface{}{"++ ", time.Now().Format("15:04:05.000"), " "}, v...)...)
}

func debugf(t *check.C, format string, v ...interface{}) {
	t.Logf(strings.Join([]string{"++", time.Now().Format("15:04:05.000"), format}, " "), v...)
}

func run(t *check.C, cmd *exec.Cmd) *CmdResult {
	var out bytes.Buffer
	debug(t, strings.Join(append([]string{cmd.Path}, cmd.Args[1:]...), " "))
	if args.Stream {
		cmd.Stdout = io.MultiWriter(os.Stdout, &out)
		cmd.Stderr = io.MultiWriter(os.Stderr, &out)
	} else {
		cmd.Stdout = &out
		cmd.Stderr = &out
	}
	err := cmd.Run()
	res := &CmdResult{
		Cmd:    cmd.Args,
		Err:    err,
		Output: out.String(),
	}
	if !args.Stream {
		t.Log(res.Output)
	}
	return res
}

var Outputs check.Checker = outputChecker{
	&check.CheckerInfo{
		Name:   "Outputs",
		Params: []string{"result", "output"},
	},
}

type outputChecker struct {
	*check.CheckerInfo
}

func (outputChecker) Check(params []interface{}, names []string) (bool, string) {
	ok, msg, s, res := checkCmdResult(params, names)
	if !ok {
		return ok, msg
	}
	return s == res.Output, ""
}

func checkCmdResult(params []interface{}, names []string) (ok bool, msg, s string, res *CmdResult) {
	res, ok = params[0].(*CmdResult)
	if !ok {
		msg = "result must be a *CmdResult"
		return
	}
	switch v := params[1].(type) {
	case []byte:
		s = string(v)
	case string:
		s = v
	default:
		msg = "output must be a []byte or string"
		return
	}
	if res.Err != nil {
		return false, "", "", nil
	}
	ok = true
	return
}

var OutputContains check.Checker = outputContainsChecker{
	&check.CheckerInfo{
		Name:   "OutputContains",
		Params: []string{"result", "contains"},
	},
}

type outputContainsChecker struct {
	*check.CheckerInfo
}

func (outputContainsChecker) Check(params []interface{}, names []string) (bool, string) {
	ok, msg, s, res := checkCmdResult(params, names)
	if !ok {
		return ok, msg
	}
	return strings.Contains(res.Output, s), ""
}

var Succeeds check.Checker = succeedsChecker{
	&check.CheckerInfo{
		Name:   "Succeeds",
		Params: []string{"result"},
	},
}

type succeedsChecker struct {
	*check.CheckerInfo
}

func (succeedsChecker) Check(params []interface{}, names []string) (bool, string) {
	res, ok := params[0].(*CmdResult)
	if !ok {
		return false, "result must be a *CmdResult"
	}
	return res.Err == nil, ""
}

type matchesChecker struct {
	*check.CheckerInfo
}

var Matches check.Checker = &matchesChecker{
	&check.CheckerInfo{Name: "Matches", Params: []string{"value", "regex"}},
}

func (checker *matchesChecker) Check(params []interface{}, names []string) (result bool, error string) {
	return matches(params[0], params[1])
}

func matches(value, regex interface{}) (result bool, error string) {
	reStr, ok := regex.(string)
	if !ok {
		return false, "Regex must be a string"
	}
	valueStr, valueIsStr := value.(string)
	if !valueIsStr {
		if valueWithStr, valueHasStr := value.(fmt.Stringer); valueHasStr {
			valueStr, valueIsStr = valueWithStr.String(), true
		}
	}
	if valueIsStr {
		matches, err := regexp.MatchString(reStr, valueStr)
		if err != nil {
			return false, "Can't compile regex: " + err.Error()
		}
		return matches, ""
	}
	return false, "Obtained value is not a string and has no .String()"
}
<commit_msg>test: Fix image URI construction<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"text/template"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-check"
	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/fsouza/go-dockerclient"
	"github.com/flynn/flynn/Godeps/_workspace/src/golang.org/x/crypto/ssh"
	"github.com/flynn/flynn/pkg/iotool"
	"github.com/flynn/flynn/pkg/shutdown"
	"github.com/flynn/flynn/test/arg"
	"github.com/flynn/flynn/test/cluster"
	"github.com/flynn/flynn/test/cluster/client"
)

var sshWrapper = template.Must(template.New("ssh").Parse(`
#!/bin/bash

ssh -o LogLevel=FATAL -o IdentitiesOnly=yes -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i {{.SSHKey}} "$@"
`[1:]))

var args *arg.Args
var flynnrc string
var routerIP string
var testCluster *testcluster.Client

func init() {
	args = arg.Parse()
	if args.Stream {
		args.Debug = true
	}
	log.SetFlags(log.Lshortfile)
}

func main() {
	defer shutdown.Exit()

	var err error
	if err = lookupImageURIs(); err != nil {
		log.Fatalf("could not determine image ID: %s", err)
	}

	var res *check.Result
	// defer exiting here so it runs after all other defers
	defer func() {
		if err != nil || res != nil && !res.Passed() {
			if args.DumpLogs {
				if args.Gist {
					exec.Command("flynn-host", "upload-debug-info").Run()
				} else if testCluster != nil {
					testCluster.DumpLogs(&iotool.SafeWriter{W: os.Stdout})
				}
			}
			os.Exit(1)
		}
	}()

	flynnrc = args.Flynnrc
	routerIP = args.RouterIP
	if flynnrc == "" {
		c := cluster.New(args.BootConfig, os.Stdout)
		var rootFS string
		rootFS, err = c.BuildFlynn(args.RootFS, "origin/master", false, false)
		if err != nil {
			if args.Kill {
				c.Shutdown()
			}
			log.Println("could not build flynn: ", err)
			if rootFS != "" {
				os.RemoveAll(rootFS)
			}
			return
		}
		if args.BuildRootFS {
			c.Shutdown()
			fmt.Println("Built Flynn in rootfs:", rootFS)
			return
		} else {
			defer os.RemoveAll(rootFS)
		}
		if err = c.Boot(rootFS, 3, nil, args.Kill); err != nil {
			log.Println("could not boot cluster: ", err)
			return
		}
		if args.Kill {
			defer c.Shutdown()
		}

		if err = createFlynnrc(c); err != nil {
			log.Println(err)
			return
		}
		defer os.RemoveAll(flynnrc)

		routerIP = c.RouterIP
	}

	if args.ClusterAPI != "" {
		testCluster, err = testcluster.NewClient(args.ClusterAPI)
		if err != nil {
			log.Println(err)
			return
		}
	}

	res = check.RunAll(&check.RunConf{
		Filter:           args.Run,
		Stream:           args.Stream,
		Verbose:          args.Debug,
		KeepWorkDir:      args.Debug,
		ConcurrencyLevel: 5,
	})
	fmt.Println(res)
}

var imageURIs = map[string]string{
	"test-apps":           "",
	"postgresql":          "",
	"controller-examples": "",
}

func lookupImageURIs() error {
	d, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		return err
	}
	for name := range imageURIs {
		fullName := "flynn/" + name
		image, err := d.InspectImage(fullName)
		if err != nil {
			return err
		}
		imageURIs[name] = fmt.Sprintf("https://example.com?name=%s&id=%s", fullName, image.ID)
	}
	return nil
}

type sshData struct {
	Key     string
	Pub     string
	Env     []string
	Cleanup func()
}

func genSSHKey() (*sshData, error) {
	keyFile, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err

	}
	defer keyFile.Close()
	rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	pem.Encode(keyFile, &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(rsaKey),
	})

	pubFile, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err
	}
	defer pubFile.Close()
	rsaPubKey, err := ssh.NewPublicKey(&rsaKey.PublicKey)
	if err != nil {
		return nil, err
	}
	if _, err := pubFile.Write(ssh.MarshalAuthorizedKey(rsaPubKey)); err != nil {
		return nil, err
	}

	wrapperFile, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err
	}
	defer wrapperFile.Close()
	if err := sshWrapper.Execute(wrapperFile, map[string]string{"SSHKey": keyFile.Name()}); err != nil {
		return nil, err
	}
	if err := wrapperFile.Chmod(0700); err != nil {
		return nil, err
	}

	return &sshData{
		Key: keyFile.Name(),
		Pub: pubFile.Name(),
		Env: []string{"GIT_SSH=" + wrapperFile.Name()},
		Cleanup: func() {
			os.RemoveAll(keyFile.Name())
			os.RemoveAll(pubFile.Name())
			os.RemoveAll(wrapperFile.Name())
		},
	}, nil
}

func createFlynnrc(c *cluster.Cluster) error {
	tmpfile, err := ioutil.TempFile("", "flynnrc-")
	if err != nil {
		return err
	}
	path := tmpfile.Name()

	config, err := c.CLIConfig()
	if err != nil {
		os.RemoveAll(path)
		return err
	}

	if err := config.SaveTo(path); err != nil {
		os.RemoveAll(path)
		return err
	}

	flynnrc = path
	return nil
}

type CmdResult struct {
	Cmd    []string
	Output string
	Err    error
}

func flynnEnv(path string) []string {
	env := os.Environ()
	res := make([]string, 0, len(env)+1)
	for _, v := range env {
		if !strings.HasPrefix(v, "FLYNNRC=") {
			res = append(res, v)
		}
	}
	res = append(res, "FLYNNRC="+path)
	return res
}

func flynnCmd(dir string, cmdArgs ...string) *exec.Cmd {
	cmd := exec.Command(args.CLI, cmdArgs...)
	cmd.Env = flynnEnv(flynnrc)
	cmd.Dir = dir
	return cmd
}

func flynn(t *check.C, dir string, args ...string) *CmdResult {
	return run(t, flynnCmd(dir, args...))
}

func debug(t *check.C, v ...interface{}) {
	t.Log(append([]interface{}{"++ ", time.Now().Format("15:04:05.000"), " "}, v...)...)
}

func debugf(t *check.C, format string, v ...interface{}) {
	t.Logf(strings.Join([]string{"++", time.Now().Format("15:04:05.000"), format}, " "), v...)
}

func run(t *check.C, cmd *exec.Cmd) *CmdResult {
	var out bytes.Buffer
	debug(t, strings.Join(append([]string{cmd.Path}, cmd.Args[1:]...), " "))
	if args.Stream {
		cmd.Stdout = io.MultiWriter(os.Stdout, &out)
		cmd.Stderr = io.MultiWriter(os.Stderr, &out)
	} else {
		cmd.Stdout = &out
		cmd.Stderr = &out
	}
	err := cmd.Run()
	res := &CmdResult{
		Cmd:    cmd.Args,
		Err:    err,
		Output: out.String(),
	}
	if !args.Stream {
		t.Log(res.Output)
	}
	return res
}

var Outputs check.Checker = outputChecker{
	&check.CheckerInfo{
		Name:   "Outputs",
		Params: []string{"result", "output"},
	},
}

type outputChecker struct {
	*check.CheckerInfo
}

func (outputChecker) Check(params []interface{}, names []string) (bool, string) {
	ok, msg, s, res := checkCmdResult(params, names)
	if !ok {
		return ok, msg
	}
	return s == res.Output, ""
}

func checkCmdResult(params []interface{}, names []string) (ok bool, msg, s string, res *CmdResult) {
	res, ok = params[0].(*CmdResult)
	if !ok {
		msg = "result must be a *CmdResult"
		return
	}
	switch v := params[1].(type) {
	case []byte:
		s = string(v)
	case string:
		s = v
	default:
		msg = "output must be a []byte or string"
		return
	}
	if res.Err != nil {
		return false, "", "", nil
	}
	ok = true
	return
}

var OutputContains check.Checker = outputContainsChecker{
	&check.CheckerInfo{
		Name:   "OutputContains",
		Params: []string{"result", "contains"},
	},
}

type outputContainsChecker struct {
	*check.CheckerInfo
}

func (outputContainsChecker) Check(params []interface{}, names []string) (bool, string) {
	ok, msg, s, res := checkCmdResult(params, names)
	if !ok {
		return ok, msg
	}
	return strings.Contains(res.Output, s), ""
}

var Succeeds check.Checker = succeedsChecker{
	&check.CheckerInfo{
		Name:   "Succeeds",
		Params: []string{"result"},
	},
}

type succeedsChecker struct {
	*check.CheckerInfo
}

func (succeedsChecker) Check(params []interface{}, names []string) (bool, string) {
	res, ok := params[0].(*CmdResult)
	if !ok {
		return false, "result must be a *CmdResult"
	}
	return res.Err == nil, ""
}

type matchesChecker struct {
	*check.CheckerInfo
}

var Matches check.Checker = &matchesChecker{
	&check.CheckerInfo{Name: "Matches", Params: []string{"value", "regex"}},
}

func (checker *matchesChecker) Check(params []interface{}, names []string) (result bool, error string) {
	return matches(params[0], params[1])
}

func matches(value, regex interface{}) (result bool, error string) {
	reStr, ok := regex.(string)
	if !ok {
		return false, "Regex must be a string"
	}
	valueStr, valueIsStr := value.(string)
	if !valueIsStr {
		if valueWithStr, valueHasStr := value.(fmt.Stringer); valueHasStr {
			valueStr, valueIsStr = valueWithStr.String(), true
		}
	}
	if valueIsStr {
		matches, err := regexp.MatchString(reStr, valueStr)
		if err != nil {
			return false, "Can't compile regex: " + err.Error()
		}
		return matches, ""
	}
	return false, "Obtained value is not a string and has no .String()"
}
<|endoftext|>


<commit_before>// Package strparse is a simple helper package for parsing strings
package strparse

import (
	"strings"
	"unicode/utf8"
)

type Parser struct {
	Str        string
	pos, width int
}

func (p *Parser) next() rune {
	if p.pos == len(p.Str) {
		p.width = 0
		return -1
	}
	r, s := utf8.DecodeRuneInString(p.Str)
	p.pos += s
	p.width = s
	return r
}

func (p *Parser) backup() {
	if p.width > 0 {
		p.pos -= p.width
		p.width = 0
	}
}

func (p *Parser) Peek() rune {
	r := p.next()
	p.backup()
	return r
}

func (p *Parser) Get() string {
	s := p.Str[:p.pos]
	p.Clear()
	return s
}

func (p *Parser) Len() int {
	return p.pos
}

func (p *Parser) Cap() int {
	return len(p.Str)
}

func (p *Parser) Clear() {
	p.Str = p.Str[p.pos:]
	p.pos = 0
	p.width = 0
}

func (p *Parser) Accept(chars string) bool {
	if strings.IndexRune(chars, p.next()) < 0 {
		p.backup()
		return false
	}
	return true
}

func (p *Parser) AcceptRun(chars string) {
	for {
		if strings.IndexRune(chars, p.next()) < 0 {
			p.backup()
			break
		}
	}
}

func (p *Parser) Except(chars string) bool {
	if r := p.next(); r == -1 || strings.IndexRune(chars, r) >= 0 {
		p.backup()
		return false
	}
	return true
}

func (p *Parser) ExceptRun(chars string) {
	for {
		if r := p.next(); r == -1 || strings.IndexRune(chars, r) >= 0 {
			p.backup()
			break
		}
	}
}
<commit_msg>added New func<commit_after>// Package strparse is a simple helper package for parsing strings
package strparse

import (
	"strings"
	"unicode/utf8"
)

type Parser struct {
	Str        string
	pos, width int
}

func New(s string) *Parser {
	return &Parser{Str: s}
}

func (p *Parser) next() rune {
	if p.pos == len(p.Str) {
		p.width = 0
		return -1
	}
	r, s := utf8.DecodeRuneInString(p.Str)
	p.pos += s
	p.width = s
	return r
}

func (p *Parser) backup() {
	if p.width > 0 {
		p.pos -= p.width
		p.width = 0
	}
}

func (p *Parser) Peek() rune {
	r := p.next()
	p.backup()
	return r
}

func (p *Parser) Get() string {
	s := p.Str[:p.pos]
	p.Clear()
	return s
}

func (p *Parser) Len() int {
	return p.pos
}

func (p *Parser) Cap() int {
	return len(p.Str)
}

func (p *Parser) Clear() {
	p.Str = p.Str[p.pos:]
	p.pos = 0
	p.width = 0
}

func (p *Parser) Accept(chars string) bool {
	if strings.IndexRune(chars, p.next()) < 0 {
		p.backup()
		return false
	}
	return true
}

func (p *Parser) AcceptRun(chars string) {
	for {
		if strings.IndexRune(chars, p.next()) < 0 {
			p.backup()
			break
		}
	}
}

func (p *Parser) Except(chars string) bool {
	if r := p.next(); r == -1 || strings.IndexRune(chars, r) >= 0 {
		p.backup()
		return false
	}
	return true
}

func (p *Parser) ExceptRun(chars string) {
	for {
		if r := p.next(); r == -1 || strings.IndexRune(chars, r) >= 0 {
			p.backup()
			break
		}
	}
}
<|endoftext|>


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

package html

import (
	"bytes"
	"io"
	"log"
	"os"
	"strconv"
)

func init() {
	// TODO(nigeltao): Remove this when ready.
	log.Stderr("The html package is incomplete; do not use for production software.")
}

// A TokenType is the type of a Token.
type TokenType int

const (
	// Error means that an error occurred during tokenization.
	Error TokenType = iota
	// Text means a text node.
	Text
	// A StartTag looks like .
	StartTag
	// An EndTag looks like .
	EndTag
	// A SelfClosingTag tag looks like
.
	SelfClosingTag
)

// String returns a string representation of the TokenType.
func (t TokenType) String() string {
	switch t {
	case Error:
		return "Error"
	case Text:
		return "Text"
	case StartTag:
		return "StartTag"
	case EndTag:
		return "EndTag"
	case SelfClosingTag:
		return "SelfClosingTag"
	}
	return "Invalid(" + strconv.Itoa(int(t)) + ")"
}

// An Attribute is an attribute key-value pair. Key is alphabetic (and hence
// does not contain escapable characters like '&', '<' or '>'), and Val is
// unescaped (it looks like "a<b" rather than "a<b").
type Attribute struct {
	Key, Val string
}

// A Token consists of a TokenType and some Data (tag name for start and end
// tags, content for text). A tag Token may also contain a slice of Attributes.
// Data is unescaped for both tag and text Tokens (it looks like "a<b" rather
// than "a<b").
type Token struct {
	Type TokenType
	Data string
	Attr []Attribute
}

// tagString returns a string representation of a tag Token's Data and Attr.
func (t Token) tagString() string {
	if len(t.Attr) == 0 {
		return t.Data
	}
	buf := bytes.NewBuffer(nil)
	buf.WriteString(t.Data)
	for _, a := range t.Attr {
		buf.WriteByte(' ')
		buf.WriteString(a.Key)
		buf.WriteString(`="`)
		escape(buf, a.Val)
		buf.WriteByte('"')
	}
	return buf.String()
}

// String returns a string representation of the Token.
func (t Token) String() string {
	switch t.Type {
	case Error:
		return ""
	case Text:
		return EscapeString(t.Data)
	case StartTag:
		return "<" + t.tagString() + ">"
	case EndTag:
		return "</" + t.tagString() + ">"
	case SelfClosingTag:
		return "<" + t.tagString() + "/>"
	}
	return "Invalid(" + strconv.Itoa(int(t.Type)) + ")"
}

// A Tokenizer returns a stream of HTML Tokens.
type Tokenizer struct {
	// r is the source of the HTML text.
	r io.Reader
	// tt is the TokenType of the most recently read token. If tt == Error
	// then err is the error associated with trying to read that token.
	tt  TokenType
	err os.Error
	// buf[p0:p1] holds the raw data of the most recent token.
	// buf[p1:] is buffered input that will yield future tokens.
	p0, p1 int
	buf    []byte
}

// Error returns the error associated with the most recent Error token. This is
// typically os.EOF, meaning the end of tokenization.
func (z *Tokenizer) Error() os.Error {
	if z.tt != Error {
		return nil
	}
	return z.err
}

// Raw returns the unmodified text of the current token. Calling Next, Token,
// Text, TagName or TagAttr may change the contents of the returned slice.
func (z *Tokenizer) Raw() []byte {
	return z.buf[z.p0:z.p1]
}

// readByte returns the next byte from the input stream, doing a buffered read
// from z.r into z.buf if necessary. z.buf[z.p0:z.p1] remains a contiguous byte
// slice that holds all the bytes read so far for the current token.
func (z *Tokenizer) readByte() (byte, os.Error) {
	if z.p1 >= len(z.buf) {
		// Our buffer is exhausted and we have to read from z.r.
		// We copy z.buf[z.p0:z.p1] to the beginning of z.buf. If the length
		// z.p1 - z.p0 is more than half the capacity of z.buf, then we
		// allocate a new buffer before the copy.
		c := cap(z.buf)
		d := z.p1 - z.p0
		var buf1 []byte
		if 2*d > c {
			buf1 = make([]byte, d, 2*c)
		} else {
			buf1 = z.buf[0:d]
		}
		copy(buf1, z.buf[z.p0:z.p1])
		z.p0, z.p1, z.buf = 0, d, buf1[0:d]
		// Now that we have copied the live bytes to the start of the buffer,
		// we read from z.r into the remainder.
		n, err := z.r.Read(buf1[d:cap(buf1)])
		if err != nil {
			return 0, err
		}
		z.buf = buf1[0 : d+n]
	}
	x := z.buf[z.p1]
	z.p1++
	return x, nil
}

// readTo keeps reading bytes until x is found.
func (z *Tokenizer) readTo(x uint8) os.Error {
	for {
		c, err := z.readByte()
		if err != nil {
			return err
		}
		switch c {
		case x:
			return nil
		case '\\':
			_, err = z.readByte()
			if err != nil {
				return err
			}
		}
	}
	panic("unreachable")
}

// nextTag returns the next TokenType starting from the tag open state.
func (z *Tokenizer) nextTag() (tt TokenType, err os.Error) {
	c, err := z.readByte()
	if err != nil {
		return Error, err
	}
	switch {
	case c == '/':
		tt = EndTag
	// Lower-cased characters are more common in tag names, so we check for them first.
	case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z':
		tt = StartTag
	case c == '!':
		return Error, os.NewError("html: TODO(nigeltao): implement comments")
	case c == '?':
		return Error, os.NewError("html: TODO(nigeltao): implement XML processing instructions")
	default:
		return Error, os.NewError("html: TODO(nigeltao): handle malformed tags")
	}
	for {
		c, err := z.readByte()
		if err != nil {
			return Text, err
		}
		switch c {
		case '"':
			err = z.readTo('"')
			if err != nil {
				return Text, err
			}
		case '\'':
			err = z.readTo('\'')
			if err != nil {
				return Text, err
			}
		case '>':
			if z.buf[z.p1-2] == '/' && tt == StartTag {
				return SelfClosingTag, nil
			}
			return tt, nil
		}
	}
	panic("unreachable")
}

// Next scans the next token and returns its type.
func (z *Tokenizer) Next() TokenType {
	if z.err != nil {
		z.tt = Error
		return z.tt
	}
	z.p0 = z.p1
	c, err := z.readByte()
	if err != nil {
		z.tt, z.err = Error, err
		return z.tt
	}
	if c == '<' {
		z.tt, z.err = z.nextTag()
		return z.tt
	}
	for {
		c, err := z.readByte()
		if err != nil {
			z.tt, z.err = Error, err
			if err == os.EOF {
				z.tt = Text
			}
			return z.tt
		}
		if c == '<' {
			z.p1--
			z.tt = Text
			return z.tt
		}
	}
	panic("unreachable")
}

// trim returns the largest j such that z.buf[i:j] contains only white space,
// or only white space plus the final ">" or "/>" of the raw data.
func (z *Tokenizer) trim(i int) int {
	k := z.p1
	for ; i < k; i++ {
		switch z.buf[i] {
		case ' ', '\n', '\t', '\f':
			continue
		case '>':
			if i == k-1 {
				return k
			}
		case '/':
			if i == k-2 {
				return k
			}
		}
		return i
	}
	return k
}

// lower finds the largest alphabetic [a-zA-Z]* word at the start of z.buf[i:]
// and returns that word lower-cased, as well as the trimmed cursor location
// after that word.
func (z *Tokenizer) lower(i int) ([]byte, int) {
	i0 := i
loop:
	for ; i < z.p1; i++ {
		c := z.buf[i]
		// TODO(nigeltao): Check what '0' <= c && c <= '9' should do.
		switch {
		case 'A' <= c && c <= 'Z':
			z.buf[i] = c + 'a' - 'A'
		case 'a' <= c && c <= 'z':
			// No-op.
		default:
			break loop
		}
	}
	return z.buf[i0:i], z.trim(i)
}

// Text returns the raw data after unescaping.
// The contents of the returned slice may change on the next call to Next.
func (z *Tokenizer) Text() []byte {
	s := unescape(z.Raw())
	z.p0 = z.p1
	return s
}

// TagName returns the lower-cased name of a tag token (the `img` out of
// `<IMG SRC="foo">`), and whether the tag has attributes.
// The contents of the returned slice may change on the next call to Next.
func (z *Tokenizer) TagName() (name []byte, remaining bool) {
	i := z.p0 + 1
	if i >= z.p1 {
		z.p0 = z.p1
		return nil, false
	}
	if z.buf[i] == '/' {
		i++
	}
	name, z.p0 = z.lower(i)
	remaining = z.p0 != z.p1
	return
}

// TagAttr returns the lower-cased key and unescaped value of the next unparsed
// attribute for the current tag token, and whether there are more attributes.
// The contents of the returned slices may change on the next call to Next.
func (z *Tokenizer) TagAttr() (key, val []byte, remaining bool) {
	key, i := z.lower(z.p0)
	// Get past the "=\"".
	if i == z.p1 || z.buf[i] != '=' {
		return
	}
	i = z.trim(i + 1)
	if i == z.p1 || z.buf[i] != '"' {
		return
	}
	i = z.trim(i + 1)
	// Copy and unescape everything up to the closing '"'.
	dst, src := i, i
loop:
	for src < z.p1 {
		c := z.buf[src]
		switch c {
		case '"':
			src++
			break loop
		case '&':
			dst, src = unescapeEntity(z.buf, dst, src)
		case '\\':
			if src == z.p1 {
				z.buf[dst] = '\\'
				dst++
			} else {
				z.buf[dst] = z.buf[src+1]
				dst, src = dst+1, src+2
			}
		default:
			z.buf[dst] = c
			dst, src = dst+1, src+1
		}
	}
	val, z.p0 = z.buf[i:dst], z.trim(src)
	remaining = z.p0 != z.p1
	return
}

// Token returns the next Token. The result's Data and Attr values remain valid
// after subsequent Next calls.
func (z *Tokenizer) Token() Token {
	t := Token{Type: z.tt}
	switch z.tt {
	case Text:
		t.Data = string(z.Text())
	case StartTag, EndTag, SelfClosingTag:
		var (
			attr []Attribute
			a    int
		)
		name, remaining := z.TagName()
		for remaining {
			var key, val []byte
			key, val, remaining = z.TagAttr()
			if a == len(attr) {
				// Grow the attr slice.
				n := 4 + 2*a
				attr1 := make([]Attribute, n, n)
				copy(attr1, attr)
				attr = attr1
			}
			attr[a] = Attribute{string(key), string(val)}
			a++
		}
		t.Data = string(name)
		t.Attr = attr[0:a]
	}
	return t
}

// NewTokenizer returns a new HTML Tokenizer for the given Reader.
// The input is assumed to be UTF-8 encoded.
func NewTokenizer(r io.Reader) *Tokenizer {
	return &Tokenizer{
		r:   r,
		buf: make([]byte, 0, 4096),
	}
}
<commit_msg>html: disable print<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package html

import (
	"bytes"
	"io"
	"os"
	"strconv"
)

// A TokenType is the type of a Token.
type TokenType int

const (
	// Error means that an error occurred during tokenization.
	Error TokenType = iota
	// Text means a text node.
	Text
	// A StartTag looks like .
	StartTag
	// An EndTag looks like .
	EndTag
	// A SelfClosingTag tag looks like
.
	SelfClosingTag
)

// String returns a string representation of the TokenType.
func (t TokenType) String() string {
	switch t {
	case Error:
		return "Error"
	case Text:
		return "Text"
	case StartTag:
		return "StartTag"
	case EndTag:
		return "EndTag"
	case SelfClosingTag:
		return "SelfClosingTag"
	}
	return "Invalid(" + strconv.Itoa(int(t)) + ")"
}

// An Attribute is an attribute key-value pair. Key is alphabetic (and hence
// does not contain escapable characters like '&', '<' or '>'), and Val is
// unescaped (it looks like "a<b" rather than "a<b").
type Attribute struct {
	Key, Val string
}

// A Token consists of a TokenType and some Data (tag name for start and end
// tags, content for text). A tag Token may also contain a slice of Attributes.
// Data is unescaped for both tag and text Tokens (it looks like "a<b" rather
// than "a<b").
type Token struct {
	Type TokenType
	Data string
	Attr []Attribute
}

// tagString returns a string representation of a tag Token's Data and Attr.
func (t Token) tagString() string {
	if len(t.Attr) == 0 {
		return t.Data
	}
	buf := bytes.NewBuffer(nil)
	buf.WriteString(t.Data)
	for _, a := range t.Attr {
		buf.WriteByte(' ')
		buf.WriteString(a.Key)
		buf.WriteString(`="`)
		escape(buf, a.Val)
		buf.WriteByte('"')
	}
	return buf.String()
}

// String returns a string representation of the Token.
func (t Token) String() string {
	switch t.Type {
	case Error:
		return ""
	case Text:
		return EscapeString(t.Data)
	case StartTag:
		return "<" + t.tagString() + ">"
	case EndTag:
		return "</" + t.tagString() + ">"
	case SelfClosingTag:
		return "<" + t.tagString() + "/>"
	}
	return "Invalid(" + strconv.Itoa(int(t.Type)) + ")"
}

// A Tokenizer returns a stream of HTML Tokens.
type Tokenizer struct {
	// r is the source of the HTML text.
	r io.Reader
	// tt is the TokenType of the most recently read token. If tt == Error
	// then err is the error associated with trying to read that token.
	tt  TokenType
	err os.Error
	// buf[p0:p1] holds the raw data of the most recent token.
	// buf[p1:] is buffered input that will yield future tokens.
	p0, p1 int
	buf    []byte
}

// Error returns the error associated with the most recent Error token. This is
// typically os.EOF, meaning the end of tokenization.
func (z *Tokenizer) Error() os.Error {
	if z.tt != Error {
		return nil
	}
	return z.err
}

// Raw returns the unmodified text of the current token. Calling Next, Token,
// Text, TagName or TagAttr may change the contents of the returned slice.
func (z *Tokenizer) Raw() []byte {
	return z.buf[z.p0:z.p1]
}

// readByte returns the next byte from the input stream, doing a buffered read
// from z.r into z.buf if necessary. z.buf[z.p0:z.p1] remains a contiguous byte
// slice that holds all the bytes read so far for the current token.
func (z *Tokenizer) readByte() (byte, os.Error) {
	if z.p1 >= len(z.buf) {
		// Our buffer is exhausted and we have to read from z.r.
		// We copy z.buf[z.p0:z.p1] to the beginning of z.buf. If the length
		// z.p1 - z.p0 is more than half the capacity of z.buf, then we
		// allocate a new buffer before the copy.
		c := cap(z.buf)
		d := z.p1 - z.p0
		var buf1 []byte
		if 2*d > c {
			buf1 = make([]byte, d, 2*c)
		} else {
			buf1 = z.buf[0:d]
		}
		copy(buf1, z.buf[z.p0:z.p1])
		z.p0, z.p1, z.buf = 0, d, buf1[0:d]
		// Now that we have copied the live bytes to the start of the buffer,
		// we read from z.r into the remainder.
		n, err := z.r.Read(buf1[d:cap(buf1)])
		if err != nil {
			return 0, err
		}
		z.buf = buf1[0 : d+n]
	}
	x := z.buf[z.p1]
	z.p1++
	return x, nil
}

// readTo keeps reading bytes until x is found.
func (z *Tokenizer) readTo(x uint8) os.Error {
	for {
		c, err := z.readByte()
		if err != nil {
			return err
		}
		switch c {
		case x:
			return nil
		case '\\':
			_, err = z.readByte()
			if err != nil {
				return err
			}
		}
	}
	panic("unreachable")
}

// nextTag returns the next TokenType starting from the tag open state.
func (z *Tokenizer) nextTag() (tt TokenType, err os.Error) {
	c, err := z.readByte()
	if err != nil {
		return Error, err
	}
	switch {
	case c == '/':
		tt = EndTag
	// Lower-cased characters are more common in tag names, so we check for them first.
	case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z':
		tt = StartTag
	case c == '!':
		return Error, os.NewError("html: TODO(nigeltao): implement comments")
	case c == '?':
		return Error, os.NewError("html: TODO(nigeltao): implement XML processing instructions")
	default:
		return Error, os.NewError("html: TODO(nigeltao): handle malformed tags")
	}
	for {
		c, err := z.readByte()
		if err != nil {
			return Text, err
		}
		switch c {
		case '"':
			err = z.readTo('"')
			if err != nil {
				return Text, err
			}
		case '\'':
			err = z.readTo('\'')
			if err != nil {
				return Text, err
			}
		case '>':
			if z.buf[z.p1-2] == '/' && tt == StartTag {
				return SelfClosingTag, nil
			}
			return tt, nil
		}
	}
	panic("unreachable")
}

// Next scans the next token and returns its type.
func (z *Tokenizer) Next() TokenType {
	if z.err != nil {
		z.tt = Error
		return z.tt
	}
	z.p0 = z.p1
	c, err := z.readByte()
	if err != nil {
		z.tt, z.err = Error, err
		return z.tt
	}
	if c == '<' {
		z.tt, z.err = z.nextTag()
		return z.tt
	}
	for {
		c, err := z.readByte()
		if err != nil {
			z.tt, z.err = Error, err
			if err == os.EOF {
				z.tt = Text
			}
			return z.tt
		}
		if c == '<' {
			z.p1--
			z.tt = Text
			return z.tt
		}
	}
	panic("unreachable")
}

// trim returns the largest j such that z.buf[i:j] contains only white space,
// or only white space plus the final ">" or "/>" of the raw data.
func (z *Tokenizer) trim(i int) int {
	k := z.p1
	for ; i < k; i++ {
		switch z.buf[i] {
		case ' ', '\n', '\t', '\f':
			continue
		case '>':
			if i == k-1 {
				return k
			}
		case '/':
			if i == k-2 {
				return k
			}
		}
		return i
	}
	return k
}

// lower finds the largest alphabetic [a-zA-Z]* word at the start of z.buf[i:]
// and returns that word lower-cased, as well as the trimmed cursor location
// after that word.
func (z *Tokenizer) lower(i int) ([]byte, int) {
	i0 := i
loop:
	for ; i < z.p1; i++ {
		c := z.buf[i]
		// TODO(nigeltao): Check what '0' <= c && c <= '9' should do.
		switch {
		case 'A' <= c && c <= 'Z':
			z.buf[i] = c + 'a' - 'A'
		case 'a' <= c && c <= 'z':
			// No-op.
		default:
			break loop
		}
	}
	return z.buf[i0:i], z.trim(i)
}

// Text returns the raw data after unescaping.
// The contents of the returned slice may change on the next call to Next.
func (z *Tokenizer) Text() []byte {
	s := unescape(z.Raw())
	z.p0 = z.p1
	return s
}

// TagName returns the lower-cased name of a tag token (the `img` out of
// `<IMG SRC="foo">`), and whether the tag has attributes.
// The contents of the returned slice may change on the next call to Next.
func (z *Tokenizer) TagName() (name []byte, remaining bool) {
	i := z.p0 + 1
	if i >= z.p1 {
		z.p0 = z.p1
		return nil, false
	}
	if z.buf[i] == '/' {
		i++
	}
	name, z.p0 = z.lower(i)
	remaining = z.p0 != z.p1
	return
}

// TagAttr returns the lower-cased key and unescaped value of the next unparsed
// attribute for the current tag token, and whether there are more attributes.
// The contents of the returned slices may change on the next call to Next.
func (z *Tokenizer) TagAttr() (key, val []byte, remaining bool) {
	key, i := z.lower(z.p0)
	// Get past the "=\"".
	if i == z.p1 || z.buf[i] != '=' {
		return
	}
	i = z.trim(i + 1)
	if i == z.p1 || z.buf[i] != '"' {
		return
	}
	i = z.trim(i + 1)
	// Copy and unescape everything up to the closing '"'.
	dst, src := i, i
loop:
	for src < z.p1 {
		c := z.buf[src]
		switch c {
		case '"':
			src++
			break loop
		case '&':
			dst, src = unescapeEntity(z.buf, dst, src)
		case '\\':
			if src == z.p1 {
				z.buf[dst] = '\\'
				dst++
			} else {
				z.buf[dst] = z.buf[src+1]
				dst, src = dst+1, src+2
			}
		default:
			z.buf[dst] = c
			dst, src = dst+1, src+1
		}
	}
	val, z.p0 = z.buf[i:dst], z.trim(src)
	remaining = z.p0 != z.p1
	return
}

// Token returns the next Token. The result's Data and Attr values remain valid
// after subsequent Next calls.
func (z *Tokenizer) Token() Token {
	t := Token{Type: z.tt}
	switch z.tt {
	case Text:
		t.Data = string(z.Text())
	case StartTag, EndTag, SelfClosingTag:
		var (
			attr []Attribute
			a    int
		)
		name, remaining := z.TagName()
		for remaining {
			var key, val []byte
			key, val, remaining = z.TagAttr()
			if a == len(attr) {
				// Grow the attr slice.
				n := 4 + 2*a
				attr1 := make([]Attribute, n, n)
				copy(attr1, attr)
				attr = attr1
			}
			attr[a] = Attribute{string(key), string(val)}
			a++
		}
		t.Data = string(name)
		t.Attr = attr[0:a]
	}
	return t
}

// NewTokenizer returns a new HTML Tokenizer for the given Reader.
// The input is assumed to be UTF-8 encoded.
func NewTokenizer(r io.Reader) *Tokenizer {
	return &Tokenizer{
		r:   r,
		buf: make([]byte, 0, 4096),
	}
}
<|endoftext|>


<commit_before>package migrator

import (
	"log"

	"github.com/awslabs/aws-sdk-go/aws"
	"github.com/awslabs/aws-sdk-go/gen/kinesis"
	"github.com/awslabs/aws-sdk-go/gen/swf"
	. "github.com/sclasen/swfsm/sugar"
	//"github.com/awslabs/aws-sdk-go/gen/dynamodb"
	"time"
)

// TypesMigrator is composed of a DomainMigrator, a WorkflowTypeMigrator and an ActivityTypeMigrator.
type TypesMigrator struct {
	DomainMigrator       *DomainMigrator
	WorkflowTypeMigrator *WorkflowTypeMigrator
	ActivityTypeMigrator *ActivityTypeMigrator
	StreamMigrator       *StreamMigrator
}

// Migrate runs Migrate on the underlying DomainMigrator, a WorkflowTypeMigrator and ActivityTypeMigrator.
func (t *TypesMigrator) Migrate() {
	if t.ActivityTypeMigrator == nil {
		t.ActivityTypeMigrator = new(ActivityTypeMigrator)
	}
	if t.DomainMigrator == nil {
		t.DomainMigrator = new(DomainMigrator)
	}
	if t.WorkflowTypeMigrator == nil {
		t.WorkflowTypeMigrator = new(WorkflowTypeMigrator)
	}
	if t.StreamMigrator == nil {
		t.StreamMigrator = new(StreamMigrator)
	}
	t.DomainMigrator.Migrate()
	t.WorkflowTypeMigrator.Migrate()
	t.ActivityTypeMigrator.Migrate()
	t.StreamMigrator.Migrate()
}

// DomainMigrator will register or deprecate the configured domains as required.
type DomainMigrator struct {
	RegisteredDomains []swf.RegisterDomainInput
	DeprecatedDomains []swf.DeprecateDomainInput
	Client            *swf.SWF
}

// Migrate asserts that DeprecatedDomains are deprecated or deprecates them, then asserts that RegisteredDomains are registered or registers them.
func (d *DomainMigrator) Migrate() {
	for _, dd := range d.DeprecatedDomains {
		if d.isDeprecated(dd.Name) {
			log.Printf("action=migrate at=deprecate-domain domain=%s status=previously-deprecated", LS(dd.Name))
		} else {
			d.deprecate(dd)
			log.Printf("action=migrate at=deprecate-domain domain=%s status=deprecated", LS(dd.Name))
		}
	}
	for _, r := range d.RegisteredDomains {
		if d.isRegisteredNotDeprecated(r) {
			log.Printf("action=migrate at=register-domain domain=%s status=previously-registered", LS(r.Name))
		} else {
			d.register(r)
			log.Printf("action=migrate at=register-domain domain=%s status=registered", LS(r.Name))
		}
	}
}

func (d *DomainMigrator) isRegisteredNotDeprecated(rd swf.RegisterDomainInput) bool {
	desc, err := d.describe(rd.Name)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return *desc.DomainInfo.Status == swf.RegistrationStatusRegistered
}

func (d *DomainMigrator) register(rd swf.RegisterDomainInput) {
	err := d.Client.RegisterDomain(&rd)
	if err != nil {
		panic(err)
	}
}

func (d *DomainMigrator) isDeprecated(domain aws.StringValue) bool {
	desc, err := d.describe(domain)
	if err != nil {
		log.Printf("action=migrate at=is-dep domain=%s error=%s", LS(domain), err.Error())
		return false
	}

	return *desc.DomainInfo.Status == swf.RegistrationStatusDeprecated
}

func (d *DomainMigrator) deprecate(dd swf.DeprecateDomainInput) {
	err := d.Client.DeprecateDomain(&dd)
	if err != nil {
		panic(err)
	}
}

func (d *DomainMigrator) describe(domain aws.StringValue) (*swf.DomainDetail, error) {
	resp, err := d.Client.DescribeDomain(&swf.DescribeDomainInput{Name: domain})
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// WorkflowTypeMigrator will register or deprecate the configured workflow types as required.
type WorkflowTypeMigrator struct {
	RegisteredWorkflowTypes []swf.RegisterWorkflowTypeInput
	DeprecatedWorkflowTypes []swf.DeprecateWorkflowTypeInput
	Client                  *swf.SWF
}

// Migrate asserts that DeprecatedWorkflowTypes are deprecated or deprecates them, then asserts that RegisteredWorkflowTypes are registered or registers them.
func (w *WorkflowTypeMigrator) Migrate() {
	for _, dd := range w.DeprecatedWorkflowTypes {
		if w.isDeprecated(dd.Domain, dd.WorkflowType.Name, dd.WorkflowType.Version) {
			log.Printf("action=migrate at=deprecate-workflow domain=%s workflow=%s version=%s status=previously-deprecated", LS(dd.Domain), LS(dd.WorkflowType.Name), LS(dd.WorkflowType.Version))
		} else {
			w.deprecate(dd)
			log.Printf("action=migrate at=deprecate-workflow domain=%s  workflow=%s version=%s status=deprecate", LS(dd.Domain), LS(dd.WorkflowType.Name), LS(dd.WorkflowType.Version))
		}
	}
	for _, r := range w.RegisteredWorkflowTypes {
		if w.isRegisteredNotDeprecated(r) {
			log.Printf("action=migrate at=register-workflow domain=%s workflow=%s version=%s status=previously-registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		} else {
			w.register(r)
			log.Printf("action=migrate at=register-workflow domain=%s  workflow=%s version=%s status=registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		}
	}
}

func (w *WorkflowTypeMigrator) isRegisteredNotDeprecated(rd swf.RegisterWorkflowTypeInput) bool {
	desc, err := w.describe(rd.Domain, rd.Name, rd.Version)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusRegistered
}

func (w *WorkflowTypeMigrator) register(rd swf.RegisterWorkflowTypeInput) {
	err := w.Client.RegisterWorkflowType(&rd)
	if err != nil {
		panic(err)
	}
}

func (w *WorkflowTypeMigrator) isDeprecated(domain aws.StringValue, name aws.StringValue, version aws.StringValue) bool {
	desc, err := w.describe(domain, name, version)
	if err != nil {
		log.Printf("action=migrate at=is-dep domain=%s workflow=%s version=%s error=%s", LS(domain), LS(name), LS(version), err.Error())
		return false
	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusDeprecated
}

func (w *WorkflowTypeMigrator) deprecate(dd swf.DeprecateWorkflowTypeInput) {
	err := w.Client.DeprecateWorkflowType(&dd)
	if err != nil {
		panic(err)
	}
}

func (w *WorkflowTypeMigrator) describe(domain aws.StringValue, name aws.StringValue, version aws.StringValue) (*swf.WorkflowTypeDetail, error) {
	resp, err := w.Client.DescribeWorkflowType(&swf.DescribeWorkflowTypeInput{Domain: domain, WorkflowType: &swf.WorkflowType{Name: name, Version: version}})
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// ActivityTypeMigrator will register or deprecate the configured activity types as required.
type ActivityTypeMigrator struct {
	RegisteredActivityTypes []swf.RegisterActivityTypeInput
	DeprecatedActivityTypes []swf.DeprecateActivityTypeInput
	Client                  *swf.SWF
}

// Migrate asserts that DeprecatedActivityTypes are deprecated or deprecates them, then asserts that RegisteredActivityTypes are registered or registers them.
func (a *ActivityTypeMigrator) Migrate() {
	for _, d := range a.DeprecatedActivityTypes {
		if a.isDeprecated(d.Domain, d.ActivityType.Name, d.ActivityType.Version) {
			log.Printf("action=migrate at=deprecate-activity domain=%s activity=%s version=%s status=previously-deprecated", LS(d.Domain), LS(d.ActivityType.Name), LS(d.ActivityType.Version))
		} else {
			a.deprecate(d)
			log.Printf("action=migrate at=depreacate-activity domain=%s activity=%s version=%s status=deprecated", LS(d.Domain), LS(d.ActivityType.Name), LS(d.ActivityType.Version))
		}
	}
	for _, r := range a.RegisteredActivityTypes {
		if a.isRegisteredNotDeprecated(r) {
			log.Printf("action=migrate at=register-activity domain=%s activity=%s version=%s status=previously-registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		} else {
			a.register(r)
			log.Printf("action=migrate at=register-activity domain=%s activity=%s version=%s status=registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		}
	}
}

func (a *ActivityTypeMigrator) isRegisteredNotDeprecated(rd swf.RegisterActivityTypeInput) bool {
	desc, err := a.describe(rd.Domain, rd.Name, rd.Version)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusRegistered
}

func (a *ActivityTypeMigrator) register(rd swf.RegisterActivityTypeInput) {
	err := a.Client.RegisterActivityType(&rd)
	if err != nil {
		panic(err)
	}
}

func (a *ActivityTypeMigrator) isDeprecated(domain aws.StringValue, name aws.StringValue, version aws.StringValue) bool {
	desc, err := a.describe(domain, name, version)
	if err != nil {
		log.Printf("action=migrate at=is-dep domain=%s activity=%s version=%s error=%s", LS(domain), LS(name), LS(version), err.Error())
		return false
	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusDeprecated
}

func (a *ActivityTypeMigrator) deprecate(dd swf.DeprecateActivityTypeInput) {
	err := a.Client.DeprecateActivityType(&dd)
	if err != nil {
		panic(err)
	}
}

func (a *ActivityTypeMigrator) describe(domain aws.StringValue, name aws.StringValue, version aws.StringValue) (*swf.ActivityTypeDetail, error) {
	resp, err := a.Client.DescribeActivityType(&swf.DescribeActivityTypeInput{Domain: domain, ActivityType: &swf.ActivityType{Name: name, Version: version}})
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// StreamMigrator will create any Kinesis Streams required.
type StreamMigrator struct {
	Streams []kinesis.CreateStreamInput
	Client  *kinesis.Kinesis
}

// Migrate checks that the desired streams have been created and if they have not, creates them.s
func (s *StreamMigrator) Migrate() {
	for _, st := range s.Streams {
		if s.isCreated(st) {
			log.Printf("action=migrate at=create-stream stream=%s status=previously-created", LS(st.StreamName))
		} else {
			s.create(st)
			log.Printf("action=migrate at=create-stream stream=%s status=created", LS(st.StreamName))
		}
		s.awaitActive(st.StreamName, 30)
	}
}

func (s *StreamMigrator) isCreated(st kinesis.CreateStreamInput) bool {
	_, err := s.describe(st)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return true
}

func (s *StreamMigrator) create(st kinesis.CreateStreamInput) {
	err := s.Client.CreateStream(&st)
	if err != nil {
		panic(err)
	}
}

func (s *StreamMigrator) describe(st kinesis.CreateStreamInput) (*kinesis.DescribeStreamOutput, error) {
	req := kinesis.DescribeStreamInput{
		StreamName: st.StreamName,
	}
	resp, err := s.Client.DescribeStream(&req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (s *StreamMigrator) awaitActive(stream aws.StringValue, atMostSeconds int) {

	waited := 0
	status := kinesis.StreamStatusCreating
	for status != kinesis.StreamStatusActive {
		desc, err := s.Client.DescribeStream(&kinesis.DescribeStreamInput{
			StreamName: stream,
		})
		if err != nil {
			log.Printf("component=kinesis-migrator fn=awaitActive at=describe-error error=%s", err)
			panic(err)
		}
		log.Printf("component=kinesis-migrator fn=awaitActive stream=%s at=describe status=%s", *stream, *desc.StreamDescription.StreamStatus)
		status = *desc.StreamDescription.StreamStatus
		time.Sleep(1 * time.Second)
		waited++
		if waited >= atMostSeconds {
			log.Printf("component=kinesis-migrator fn=awaitActive streal=%s  at=error error=exeeeded-max-wait", *stream)
			panic("waited too long")
		}
	}
}
<commit_msg>introduce ops interfaces for swf and kinesis in migrator<commit_after>package migrator

import (
	"log"

	"github.com/awslabs/aws-sdk-go/aws"
	"github.com/awslabs/aws-sdk-go/gen/kinesis"
	"github.com/awslabs/aws-sdk-go/gen/swf"
	. "github.com/sclasen/swfsm/sugar"
	//"github.com/awslabs/aws-sdk-go/gen/dynamodb"
	"time"
)

// TypesMigrator is composed of a DomainMigrator, a WorkflowTypeMigrator and an ActivityTypeMigrator.
type TypesMigrator struct {
	DomainMigrator       *DomainMigrator
	WorkflowTypeMigrator *WorkflowTypeMigrator
	ActivityTypeMigrator *ActivityTypeMigrator
	StreamMigrator       *StreamMigrator
}

type SWFOps interface {
    DeprecateActivityType(req *swf.DeprecateActivityTypeInput) (err error)
    DeprecateDomain(req *swf.DeprecateDomainInput) (err error)
    DeprecateWorkflowType(req *swf.DeprecateWorkflowTypeInput) (err error)
    DescribeActivityType(req *swf.DescribeActivityTypeInput) (resp *swf.ActivityTypeDetail, err error)
    DescribeDomain(req *swf.DescribeDomainInput) (resp *swf.DomainDetail, err error)
    DescribeWorkflowExecution(req *swf.DescribeWorkflowExecutionInput) (resp *swf.WorkflowExecutionDetail, err error)
    DescribeWorkflowType(req *swf.DescribeWorkflowTypeInput) (resp *swf.WorkflowTypeDetail, err error)
    RegisterActivityType(req *swf.RegisterActivityTypeInput) (err error)
    RegisterDomain(req *swf.RegisterDomainInput) (err error)
    RegisterWorkflowType(req *swf.RegisterWorkflowTypeInput) (err error)
}

type KinesisOps interface {
    CreateStream(req *kinesis.CreateStreamInput) (err error)
    DescribeStream(req *kinesis.DescribeStreamInput) (resp *kinesis.DescribeStreamOutput, err error)
}

// Migrate runs Migrate on the underlying DomainMigrator, a WorkflowTypeMigrator and ActivityTypeMigrator.
func (t *TypesMigrator) Migrate() {
	if t.ActivityTypeMigrator == nil {
		t.ActivityTypeMigrator = new(ActivityTypeMigrator)
	}
	if t.DomainMigrator == nil {
		t.DomainMigrator = new(DomainMigrator)
	}
	if t.WorkflowTypeMigrator == nil {
		t.WorkflowTypeMigrator = new(WorkflowTypeMigrator)
	}
	if t.StreamMigrator == nil {
		t.StreamMigrator = new(StreamMigrator)
	}
	t.DomainMigrator.Migrate()
	t.WorkflowTypeMigrator.Migrate()
	t.ActivityTypeMigrator.Migrate()
	t.StreamMigrator.Migrate()
}

// DomainMigrator will register or deprecate the configured domains as required.
type DomainMigrator struct {
	RegisteredDomains []swf.RegisterDomainInput
	DeprecatedDomains []swf.DeprecateDomainInput
	Client            SWFOps
}

// Migrate asserts that DeprecatedDomains are deprecated or deprecates them, then asserts that RegisteredDomains are registered or registers them.
func (d *DomainMigrator) Migrate() {
	for _, dd := range d.DeprecatedDomains {
		if d.isDeprecated(dd.Name) {
			log.Printf("action=migrate at=deprecate-domain domain=%s status=previously-deprecated", LS(dd.Name))
		} else {
			d.deprecate(dd)
			log.Printf("action=migrate at=deprecate-domain domain=%s status=deprecated", LS(dd.Name))
		}
	}
	for _, r := range d.RegisteredDomains {
		if d.isRegisteredNotDeprecated(r) {
			log.Printf("action=migrate at=register-domain domain=%s status=previously-registered", LS(r.Name))
		} else {
			d.register(r)
			log.Printf("action=migrate at=register-domain domain=%s status=registered", LS(r.Name))
		}
	}
}

func (d *DomainMigrator) isRegisteredNotDeprecated(rd swf.RegisterDomainInput) bool {
	desc, err := d.describe(rd.Name)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return *desc.DomainInfo.Status == swf.RegistrationStatusRegistered
}

func (d *DomainMigrator) register(rd swf.RegisterDomainInput) {
	err := d.Client.RegisterDomain(&rd)
	if err != nil {
		panic(err)
	}
}

func (d *DomainMigrator) isDeprecated(domain aws.StringValue) bool {
	desc, err := d.describe(domain)
	if err != nil {
		log.Printf("action=migrate at=is-dep domain=%s error=%s", LS(domain), err.Error())
		return false
	}

	return *desc.DomainInfo.Status == swf.RegistrationStatusDeprecated
}

func (d *DomainMigrator) deprecate(dd swf.DeprecateDomainInput) {
	err := d.Client.DeprecateDomain(&dd)
	if err != nil {
		panic(err)
	}
}

func (d *DomainMigrator) describe(domain aws.StringValue) (*swf.DomainDetail, error) {
	resp, err := d.Client.DescribeDomain(&swf.DescribeDomainInput{Name: domain})
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// WorkflowTypeMigrator will register or deprecate the configured workflow types as required.
type WorkflowTypeMigrator struct {
	RegisteredWorkflowTypes []swf.RegisterWorkflowTypeInput
	DeprecatedWorkflowTypes []swf.DeprecateWorkflowTypeInput
	Client                  SWFOps
}

// Migrate asserts that DeprecatedWorkflowTypes are deprecated or deprecates them, then asserts that RegisteredWorkflowTypes are registered or registers them.
func (w *WorkflowTypeMigrator) Migrate() {
	for _, dd := range w.DeprecatedWorkflowTypes {
		if w.isDeprecated(dd.Domain, dd.WorkflowType.Name, dd.WorkflowType.Version) {
			log.Printf("action=migrate at=deprecate-workflow domain=%s workflow=%s version=%s status=previously-deprecated", LS(dd.Domain), LS(dd.WorkflowType.Name), LS(dd.WorkflowType.Version))
		} else {
			w.deprecate(dd)
			log.Printf("action=migrate at=deprecate-workflow domain=%s  workflow=%s version=%s status=deprecate", LS(dd.Domain), LS(dd.WorkflowType.Name), LS(dd.WorkflowType.Version))
		}
	}
	for _, r := range w.RegisteredWorkflowTypes {
		if w.isRegisteredNotDeprecated(r) {
			log.Printf("action=migrate at=register-workflow domain=%s workflow=%s version=%s status=previously-registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		} else {
			w.register(r)
			log.Printf("action=migrate at=register-workflow domain=%s  workflow=%s version=%s status=registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		}
	}
}

func (w *WorkflowTypeMigrator) isRegisteredNotDeprecated(rd swf.RegisterWorkflowTypeInput) bool {
	desc, err := w.describe(rd.Domain, rd.Name, rd.Version)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusRegistered
}

func (w *WorkflowTypeMigrator) register(rd swf.RegisterWorkflowTypeInput) {
	err := w.Client.RegisterWorkflowType(&rd)
	if err != nil {
		panic(err)
	}
}

func (w *WorkflowTypeMigrator) isDeprecated(domain aws.StringValue, name aws.StringValue, version aws.StringValue) bool {
	desc, err := w.describe(domain, name, version)
	if err != nil {
		log.Printf("action=migrate at=is-dep domain=%s workflow=%s version=%s error=%s", LS(domain), LS(name), LS(version), err.Error())
		return false
	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusDeprecated
}

func (w *WorkflowTypeMigrator) deprecate(dd swf.DeprecateWorkflowTypeInput) {
	err := w.Client.DeprecateWorkflowType(&dd)
	if err != nil {
		panic(err)
	}
}

func (w *WorkflowTypeMigrator) describe(domain aws.StringValue, name aws.StringValue, version aws.StringValue) (*swf.WorkflowTypeDetail, error) {
	resp, err := w.Client.DescribeWorkflowType(&swf.DescribeWorkflowTypeInput{Domain: domain, WorkflowType: &swf.WorkflowType{Name: name, Version: version}})
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// ActivityTypeMigrator will register or deprecate the configured activity types as required.
type ActivityTypeMigrator struct {
	RegisteredActivityTypes []swf.RegisterActivityTypeInput
	DeprecatedActivityTypes []swf.DeprecateActivityTypeInput
	Client                  SWFOps
}

// Migrate asserts that DeprecatedActivityTypes are deprecated or deprecates them, then asserts that RegisteredActivityTypes are registered or registers them.
func (a *ActivityTypeMigrator) Migrate() {
	for _, d := range a.DeprecatedActivityTypes {
		if a.isDeprecated(d.Domain, d.ActivityType.Name, d.ActivityType.Version) {
			log.Printf("action=migrate at=deprecate-activity domain=%s activity=%s version=%s status=previously-deprecated", LS(d.Domain), LS(d.ActivityType.Name), LS(d.ActivityType.Version))
		} else {
			a.deprecate(d)
			log.Printf("action=migrate at=depreacate-activity domain=%s activity=%s version=%s status=deprecated", LS(d.Domain), LS(d.ActivityType.Name), LS(d.ActivityType.Version))
		}
	}
	for _, r := range a.RegisteredActivityTypes {
		if a.isRegisteredNotDeprecated(r) {
			log.Printf("action=migrate at=register-activity domain=%s activity=%s version=%s status=previously-registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		} else {
			a.register(r)
			log.Printf("action=migrate at=register-activity domain=%s activity=%s version=%s status=registered", LS(r.Domain), LS(r.Name), LS(r.Version))
		}
	}
}

func (a *ActivityTypeMigrator) isRegisteredNotDeprecated(rd swf.RegisterActivityTypeInput) bool {
	desc, err := a.describe(rd.Domain, rd.Name, rd.Version)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusRegistered
}

func (a *ActivityTypeMigrator) register(rd swf.RegisterActivityTypeInput) {
	err := a.Client.RegisterActivityType(&rd)
	if err != nil {
		panic(err)
	}
}

func (a *ActivityTypeMigrator) isDeprecated(domain aws.StringValue, name aws.StringValue, version aws.StringValue) bool {
	desc, err := a.describe(domain, name, version)
	if err != nil {
		log.Printf("action=migrate at=is-dep domain=%s activity=%s version=%s error=%s", LS(domain), LS(name), LS(version), err.Error())
		return false
	}

	return *desc.TypeInfo.Status == swf.RegistrationStatusDeprecated
}

func (a *ActivityTypeMigrator) deprecate(dd swf.DeprecateActivityTypeInput) {
	err := a.Client.DeprecateActivityType(&dd)
	if err != nil {
		panic(err)
	}
}

func (a *ActivityTypeMigrator) describe(domain aws.StringValue, name aws.StringValue, version aws.StringValue) (*swf.ActivityTypeDetail, error) {
	resp, err := a.Client.DescribeActivityType(&swf.DescribeActivityTypeInput{Domain: domain, ActivityType: &swf.ActivityType{Name: name, Version: version}})
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// StreamMigrator will create any Kinesis Streams required.
type StreamMigrator struct {
	Streams []kinesis.CreateStreamInput
	Client  KinesisOps
}

// Migrate checks that the desired streams have been created and if they have not, creates them.s
func (s *StreamMigrator) Migrate() {
	for _, st := range s.Streams {
		if s.isCreated(st) {
			log.Printf("action=migrate at=create-stream stream=%s status=previously-created", LS(st.StreamName))
		} else {
			s.create(st)
			log.Printf("action=migrate at=create-stream stream=%s status=created", LS(st.StreamName))
		}
		s.awaitActive(st.StreamName, 30)
	}
}

func (s *StreamMigrator) isCreated(st kinesis.CreateStreamInput) bool {
	_, err := s.describe(st)
	if err != nil {
		if ae, ok := err.(aws.APIError); ok && ae.Type == ErrorTypeUnknownResourceFault {
			return false
		}

		panic(err)

	}

	return true
}

func (s *StreamMigrator) create(st kinesis.CreateStreamInput) {
	err := s.Client.CreateStream(&st)
	if err != nil {
		panic(err)
	}
}

func (s *StreamMigrator) describe(st kinesis.CreateStreamInput) (*kinesis.DescribeStreamOutput, error) {
	req := kinesis.DescribeStreamInput{
		StreamName: st.StreamName,
	}
	resp, err := s.Client.DescribeStream(&req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (s *StreamMigrator) awaitActive(stream aws.StringValue, atMostSeconds int) {

	waited := 0
	status := kinesis.StreamStatusCreating
	for status != kinesis.StreamStatusActive {
		desc, err := s.Client.DescribeStream(&kinesis.DescribeStreamInput{
			StreamName: stream,
		})
		if err != nil {
			log.Printf("component=kinesis-migrator fn=awaitActive at=describe-error error=%s", err)
			panic(err)
		}
		log.Printf("component=kinesis-migrator fn=awaitActive stream=%s at=describe status=%s", *stream, *desc.StreamDescription.StreamStatus)
		status = *desc.StreamDescription.StreamStatus
		time.Sleep(1 * time.Second)
		waited++
		if waited >= atMostSeconds {
			log.Printf("component=kinesis-migrator fn=awaitActive streal=%s  at=error error=exeeeded-max-wait", *stream)
			panic("waited too long")
		}
	}
}
<|endoftext|>


<commit_before>package osdb

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"

	"golang.org/x/text/encoding"
)

// A Subtitle with its many OSDB attributes...
type Subtitle struct {
	IDMovie            string `xmlrpc:"IDMovie"`
	IDMovieImdb        string `xmlrpc:"IDMovieImdb"`
	IDSubMovieFile     string `xmlrpc:"IDSubMovieFile"`
	IDSubtitle         string `xmlrpc:"IDSubtitle"`
	IDSubtitleFile     string `xmlrpc:"IDSubtitleFile"`
	ISO639             string `xmlrpc:"ISO639"`
	LanguageName       string `xmlrpc:"LanguageName"`
	MatchedBy          string `xmlrpc:"MatchedBy"`
	MovieByteSize      string `xmlrpc:"MovieByteSize"`
	MovieFPS           string `xmlrpc:"MovieFPS"`
	MovieHash          string `xmlrpc:"MovieHash"`
	MovieImdbRating    string `xmlrpc:"MovieImdbRating"`
	MovieKind          string `xmlrpc:"MovieKind"`
	MovieName          string `xmlrpc:"MovieName"`
	MovieNameEng       string `xmlrpc:"MovieNameEng"`
	MovieReleaseName   string `xmlrpc:"MovieReleaseName"`
	MovieTimeMS        string `xmlrpc:"MovieTimeMS"`
	MovieYear          string `xmlrpc:"MovieYear"`
	MovieFileName      string `xmlrpc:"MovieName"`
	QueryNumber        string `xmlrpc:"QueryNumber"`
	SeriesEpisode      string `xmlrpc:"SeriesEpisode"`
	SeriesIMDBParent   string `xmlrpc:"SeriesIMDBParent"`
	SeriesSeason       string `xmlrpc:"SeriesSeason"`
	SubActualCD        string `xmlrpc:"SubActualCD"`
	SubAddDate         string `xmlrpc:"SubAddDate"`
	SubAuthorComment   string `xmlrpc:"SubAuthorComment"`
	SubBad             string `xmlrpc:"SubBad"`
	SubComments        string `xmlrpc:"SubComments"`
	SubDownloadLink    string `xmlrpc:"SubDownloadLink"`
	SubDownloadsCnt    string `xmlrpc:"SubDownloadsCnt"`
	SubFeatured        string `xmlrpc:"SubFeatured"`
	SubFileName        string `xmlrpc:"SubFileName"`
	SubFormat          string `xmlrpc:"SubFormat"`
	SubHash            string `xmlrpc:"SubHash"`
	SubHD              string `xmlrpc:"SubHD"`
	SubHearingImpaired string `xmlrpc:"SubHearingImpaired"`
	SubLanguageID      string `xmlrpc:"SubLanguageID"`
	SubRating          string `xmlrpc:"SubRating"`
	SubSize            string `xmlrpc:"SubSize"`
	SubSumCD           string `xmlrpc:"SubSumCD"`
	SubEncoding        string `xmlrpc:"SubEncoding"`
	SubtitlesLink      string `xmlrpc:"SubtitlesLink"`
	UserID             string `xmlrpc:"UserID"`
	UserNickName       string `xmlrpc:"UserNickName"`
	UserRank           string `xmlrpc:"UserRank"`
	ZipDownloadLink    string `xmlrpc:"ZipDownloadLink"`
	subFilePath        string
}

func (s *Subtitle) toUploadParams() map[string]string {
	return map[string]string{
		"subhash":       s.SubHash,
		"subfilename":   s.SubFileName,
		"moviehash":     s.MovieHash,
		"moviebytesize": s.MovieByteSize,
		"moviefilename": s.MovieFileName,
	}
}

func (s *Subtitle) encodeFile() (string, error) {
	fh, err := os.Open(s.subFilePath)
	if err != nil {
		return "", err
	}
	defer fh.Close()
	dest := bytes.NewBuffer([]byte{})
	gzWriter := gzip.NewWriter(dest)
	enc := base64.NewEncoder(base64.StdEncoding, gzWriter)
	_, err = io.Copy(enc, fh)
	if err != nil {
		return "", err
	}
	// XXX DEBUG
	fmt.Println("upload content size:", dest.Len())
	return dest.String(), nil
}

// Subtitles is a collection of subtitles.
type Subtitles []Subtitle

// ByDownloads implements sort interface for Subtitles, by download count.
type ByDownloads Subtitles

func (s ByDownloads) Len() int      { return len(s) }
func (s ByDownloads) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s ByDownloads) Less(i, j int) bool {
	iCnt, err := strconv.Atoi(s[i].SubDownloadsCnt)
	if err != nil {
		return false
	}
	jCnt, err := strconv.Atoi(s[j].SubDownloadsCnt)
	if err != nil {
		return true
	}
	return iCnt > jCnt
}

// Best finds the best subsitle in a Subtitles collection. Of course
// "best" is hardly an absolute concept: here, we just take the most
// downloaded file.
func (subs Subtitles) Best() *Subtitle {
	if len(subs) > 0 {
		sort.Sort(ByDownloads(subs))
		return &subs[0]
	}
	return nil
}

// SubtitleFile contains file data as returned by OSDB's API, that is to
// say: gzip-ped and base64-encoded text.
type SubtitleFile struct {
	ID       string `xmlrpc:"idsubtitlefile"`
	Data     string `xmlrpc:"data"`
	Encoding encoding.Encoding
	reader   io.ReadCloser
}

// Reader interface for SubtitleFile. Subtitle's contents are
// decompressed, and usually encoded to UTF-8: if encoding info is
// missing, no re-encoding is done.
func (sf *SubtitleFile) Reader() (r io.ReadCloser, err error) {
	if sf.reader != nil {
		return sf.reader, err
	}

	dec := base64.NewDecoder(base64.StdEncoding, strings.NewReader(sf.Data))
	gzReader, err := gzip.NewReader(dec)
	if err != nil {
		return nil, err
	}

	if sf.Encoding == nil {
		sf.reader = gzReader
	} else {
		sf.reader = newCloseableReader(
			sf.Encoding.NewDecoder().Reader(gzReader),
			gzReader.Close,
		)
	}

	return sf.reader, nil
}

// NewSubtitles builds a Subtitles from a movie path and a slice of
// subtitles paths. Intended to be used with for osdb.HasSubtitles() and
// osdb.UploadSubtitles().
func NewSubtitles(moviePath string, subPaths []string, langID string) (Subtitles, error) {
	subs := Subtitles{}
	for _, subPath := range subPaths {
		sub, err := NewSubtitle(moviePath, subPath, langID)
		if err != nil {
			return nil, err
		}
		subs = append(subs, sub)
	}
	return subs, nil
}

// NewSubtitle builds a Subtitle struct.
func NewSubtitle(moviePath string, subPath string, langID string) (s Subtitle, err error) {
	s.subFilePath = subPath
	s.SubLanguageID = langID
	s.SubFileName = path.Base(subPath)
	// Subs are identified using md5 hashes... ¬¬
	subIO, err := os.Open(subPath)
	if err != nil {
		return
	}
	defer subIO.Close()
	h := md5.New()
	_, err = io.Copy(h, subIO)
	if err != nil {
		return
	}
	s.SubHash = fmt.Sprintf("%x", h.Sum(nil))

	// Movie filename, byte-size, & hash.
	s.MovieFileName = path.Base(moviePath)
	movieIO, err := os.Open(moviePath)
	if err != nil {
		return
	}
	defer movieIO.Close()
	stat, err := movieIO.Stat()
	if err != nil {
		return
	}
	s.MovieByteSize = strconv.FormatInt(stat.Size(), 10)
	movieHash, err := HashFile(movieIO)
	if err != nil {
		return
	}
	s.MovieHash = fmt.Sprintf("%x", movieHash)
	return
}

// Serialize Subtitle to OSDB's XMLRPC params when trying to upload.
func (subs *Subtitles) toTryUploadParams() (map[string]interface{}, error) {
	subMap := map[string]interface{}{}
	for i, s := range *subs {
		key := "cd" + strconv.Itoa(i+1) // keys are cd1, cd2, ...
		subMap[key] = s.toUploadParams()
	}

	return subMap, nil
}

// Serialize Subtitle to OSDB's XMLRPC params when uploading.
func (subs *Subtitles) toUploadParams() (map[string]interface{}, error) {
	langID := (*subs)[0].SubLanguageID
	params := map[string]interface{}{}

	params["baseinfo"] = map[string]string{
		"sublanguageid": langID,
		// FIXME add "idmovieimdb"
	}

	for i, s := range *subs {
		key := "cd" + strconv.Itoa(i+1) // keys are cd1, cd2, ...
		subParam := s.toUploadParams()
		encoded, err := s.encodeFile()
		if err != nil {
			return nil, err
		}
		subParam["subcontent"] = encoded
		params[key] = subParam
	}

	return params, nil
}

// Implement io.ReadCloser by wrapping io.Reader
type closeableReader struct {
	io.Reader
	close func() error
}

// Close the reader by calling a preset close function
func (c *closeableReader) Close() error {
	return c.close()
}

// Create a ReadCloser which will read from r and call close() upon closing
func newCloseableReader(r io.Reader, close func() error) io.ReadCloser {
	return &closeableReader{r, close}
}
<commit_msg>gzip+b64 encode subtitle contents when uploading<commit_after>package osdb

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"

	"golang.org/x/text/encoding"
)

// A Subtitle with its many OSDB attributes...
type Subtitle struct {
	IDMovie            string `xmlrpc:"IDMovie"`
	IDMovieImdb        string `xmlrpc:"IDMovieImdb"`
	IDSubMovieFile     string `xmlrpc:"IDSubMovieFile"`
	IDSubtitle         string `xmlrpc:"IDSubtitle"`
	IDSubtitleFile     string `xmlrpc:"IDSubtitleFile"`
	ISO639             string `xmlrpc:"ISO639"`
	LanguageName       string `xmlrpc:"LanguageName"`
	MatchedBy          string `xmlrpc:"MatchedBy"`
	MovieByteSize      string `xmlrpc:"MovieByteSize"`
	MovieFPS           string `xmlrpc:"MovieFPS"`
	MovieHash          string `xmlrpc:"MovieHash"`
	MovieImdbRating    string `xmlrpc:"MovieImdbRating"`
	MovieKind          string `xmlrpc:"MovieKind"`
	MovieName          string `xmlrpc:"MovieName"`
	MovieNameEng       string `xmlrpc:"MovieNameEng"`
	MovieReleaseName   string `xmlrpc:"MovieReleaseName"`
	MovieTimeMS        string `xmlrpc:"MovieTimeMS"`
	MovieYear          string `xmlrpc:"MovieYear"`
	MovieFileName      string `xmlrpc:"MovieName"`
	QueryNumber        string `xmlrpc:"QueryNumber"`
	SeriesEpisode      string `xmlrpc:"SeriesEpisode"`
	SeriesIMDBParent   string `xmlrpc:"SeriesIMDBParent"`
	SeriesSeason       string `xmlrpc:"SeriesSeason"`
	SubActualCD        string `xmlrpc:"SubActualCD"`
	SubAddDate         string `xmlrpc:"SubAddDate"`
	SubAuthorComment   string `xmlrpc:"SubAuthorComment"`
	SubBad             string `xmlrpc:"SubBad"`
	SubComments        string `xmlrpc:"SubComments"`
	SubDownloadLink    string `xmlrpc:"SubDownloadLink"`
	SubDownloadsCnt    string `xmlrpc:"SubDownloadsCnt"`
	SubFeatured        string `xmlrpc:"SubFeatured"`
	SubFileName        string `xmlrpc:"SubFileName"`
	SubFormat          string `xmlrpc:"SubFormat"`
	SubHash            string `xmlrpc:"SubHash"`
	SubHD              string `xmlrpc:"SubHD"`
	SubHearingImpaired string `xmlrpc:"SubHearingImpaired"`
	SubLanguageID      string `xmlrpc:"SubLanguageID"`
	SubRating          string `xmlrpc:"SubRating"`
	SubSize            string `xmlrpc:"SubSize"`
	SubSumCD           string `xmlrpc:"SubSumCD"`
	SubEncoding        string `xmlrpc:"SubEncoding"`
	SubtitlesLink      string `xmlrpc:"SubtitlesLink"`
	UserID             string `xmlrpc:"UserID"`
	UserNickName       string `xmlrpc:"UserNickName"`
	UserRank           string `xmlrpc:"UserRank"`
	ZipDownloadLink    string `xmlrpc:"ZipDownloadLink"`
	subFilePath        string
}

func (s *Subtitle) toUploadParams() map[string]string {
	return map[string]string{
		"subhash":       s.SubHash,
		"subfilename":   s.SubFileName,
		"moviehash":     s.MovieHash,
		"moviebytesize": s.MovieByteSize,
		"moviefilename": s.MovieFileName,
	}
}

func (s *Subtitle) encodeFile() (string, error) {
	fh, err := os.Open(s.subFilePath)
	if err != nil {
		return "", err
	}
	defer fh.Close()
	dest := new(bytes.Buffer)
	enc := base64.NewEncoder(base64.StdEncoding, dest)
	gzWriter := gzip.NewWriter(enc)
	_, err = io.Copy(gzWriter, fh)
	if err != nil {
		return "", err
	}
	gzWriter.Flush()
	gzWriter.Close()
	return dest.String(), nil
}

// Subtitles is a collection of subtitles.
type Subtitles []Subtitle

// ByDownloads implements sort interface for Subtitles, by download count.
type ByDownloads Subtitles

func (s ByDownloads) Len() int      { return len(s) }
func (s ByDownloads) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s ByDownloads) Less(i, j int) bool {
	iCnt, err := strconv.Atoi(s[i].SubDownloadsCnt)
	if err != nil {
		return false
	}
	jCnt, err := strconv.Atoi(s[j].SubDownloadsCnt)
	if err != nil {
		return true
	}
	return iCnt > jCnt
}

// Best finds the best subsitle in a Subtitles collection. Of course
// "best" is hardly an absolute concept: here, we just take the most
// downloaded file.
func (subs Subtitles) Best() *Subtitle {
	if len(subs) > 0 {
		sort.Sort(ByDownloads(subs))
		return &subs[0]
	}
	return nil
}

// SubtitleFile contains file data as returned by OSDB's API, that is to
// say: gzip-ped and base64-encoded text.
type SubtitleFile struct {
	ID       string `xmlrpc:"idsubtitlefile"`
	Data     string `xmlrpc:"data"`
	Encoding encoding.Encoding
	reader   io.ReadCloser
}

// Reader interface for SubtitleFile. Subtitle's contents are
// decompressed, and usually encoded to UTF-8: if encoding info is
// missing, no re-encoding is done.
func (sf *SubtitleFile) Reader() (r io.ReadCloser, err error) {
	if sf.reader != nil {
		return sf.reader, err
	}

	dec := base64.NewDecoder(base64.StdEncoding, strings.NewReader(sf.Data))
	gzReader, err := gzip.NewReader(dec)
	if err != nil {
		return nil, err
	}

	if sf.Encoding == nil {
		sf.reader = gzReader
	} else {
		sf.reader = newCloseableReader(
			sf.Encoding.NewDecoder().Reader(gzReader),
			gzReader.Close,
		)
	}

	return sf.reader, nil
}

// NewSubtitles builds a Subtitles from a movie path and a slice of
// subtitles paths. Intended to be used with for osdb.HasSubtitles() and
// osdb.UploadSubtitles().
func NewSubtitles(moviePath string, subPaths []string, langID string) (Subtitles, error) {
	subs := Subtitles{}
	for _, subPath := range subPaths {
		sub, err := NewSubtitle(moviePath, subPath, langID)
		if err != nil {
			return nil, err
		}
		subs = append(subs, sub)
	}
	return subs, nil
}

// NewSubtitle builds a Subtitle struct.
func NewSubtitle(moviePath string, subPath string, langID string) (s Subtitle, err error) {
	s.subFilePath = subPath
	s.SubLanguageID = langID
	s.SubFileName = path.Base(subPath)
	// Subs are identified using md5 hashes... ¬¬
	subIO, err := os.Open(subPath)
	if err != nil {
		return
	}
	defer subIO.Close()
	h := md5.New()
	_, err = io.Copy(h, subIO)
	if err != nil {
		return
	}
	s.SubHash = fmt.Sprintf("%x", h.Sum(nil))

	// Movie filename, byte-size, & hash.
	s.MovieFileName = path.Base(moviePath)
	movieIO, err := os.Open(moviePath)
	if err != nil {
		return
	}
	defer movieIO.Close()
	stat, err := movieIO.Stat()
	if err != nil {
		return
	}
	s.MovieByteSize = strconv.FormatInt(stat.Size(), 10)
	movieHash, err := HashFile(movieIO)
	if err != nil {
		return
	}
	s.MovieHash = fmt.Sprintf("%x", movieHash)
	return
}

// Serialize Subtitle to OSDB's XMLRPC params when trying to upload.
func (subs *Subtitles) toTryUploadParams() (map[string]interface{}, error) {
	subMap := map[string]interface{}{}
	for i, s := range *subs {
		key := "cd" + strconv.Itoa(i+1) // keys are cd1, cd2, ...
		subMap[key] = s.toUploadParams()
	}

	return subMap, nil
}

// Serialize Subtitle to OSDB's XMLRPC params when uploading.
func (subs *Subtitles) toUploadParams() (map[string]interface{}, error) {
	langID := (*subs)[0].SubLanguageID
	params := map[string]interface{}{}

	params["baseinfo"] = map[string]string{
		"sublanguageid": langID,
		// FIXME add "idmovieimdb"
	}

	for i, s := range *subs {
		key := "cd" + strconv.Itoa(i+1) // keys are cd1, cd2, ...
		subParam := s.toUploadParams()
		encoded, err := s.encodeFile()
		if err != nil {
			return nil, err
		}
		subParam["subcontent"] = encoded
		params[key] = subParam
	}

	return params, nil
}

// Implement io.ReadCloser by wrapping io.Reader
type closeableReader struct {
	io.Reader
	close func() error
}

// Close the reader by calling a preset close function
func (c *closeableReader) Close() error {
	return c.close()
}

// Create a ReadCloser which will read from r and call close() upon closing
func newCloseableReader(r io.Reader, close func() error) io.ReadCloser {
	return &closeableReader{r, close}
}
<|endoftext|>