<commit_before>package ftps

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
	"time"
)

type FTPS struct {
	host string

	conn net.Conn
	text *textproto.Conn

	Debug     bool
	TLSConfig tls.Config
}

func (ftps *FTPS) Connect(host string, port int) (err error) {

	ftps.host = host

	ftps.conn, err = net.Dial("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return err
	}

	ftps.text = textproto.NewConn(ftps.conn)

	_, err = ftps.response(220)
	if err != nil {
		return err
	}

	_, err = ftps.request("AUTH TLS", 234)
	if err != nil {
		return err
	}

	ftps.conn = ftps.upgradeConnToTLS(ftps.conn)
	ftps.text = textproto.NewConn(ftps.conn) // TODO use sync or something similar?

	return
}

func (ftps *FTPS) isConnEstablished() {

	if ftps.conn == nil {
		panic("Connection is not established yet")
	}
}

func (ftps *FTPS) Login(username, password string) (err error) {

	ftps.isConnEstablished()

	_, err = ftps.request(fmt.Sprintf("USER %s", username), 331)
	if err != nil {
		return err
	}

	_, err = ftps.request(fmt.Sprintf("PASS %s", password), 230)
	if err != nil {
		return err
	}

	_, err = ftps.request("TYPE I", 200) // use binary mode
	if err != nil {
		return err
	}

	_, err = ftps.request("PBSZ 0", 200)
	if err != nil {
		return err
	}

	_, err = ftps.request("PROT P", 200) // encrypt data connection
	if err != nil {
		return err
	}

	return
}

func (ftps *FTPS) request(cmd string, expected int) (message string, err error) {

	ftps.isConnEstablished()

	ftps.debugInfo("<*cmd*> " + cmd)

	_, err = ftps.text.Cmd(cmd)
	if err != nil {
		return
	}

	message, err = ftps.response(expected)

	return
}

func (ftps *FTPS) requestDataConn(cmd string, expected int) (dataConn net.Conn, err error) {

	port, err := ftps.pasv()
	if err != nil {
		return
	}

	dataConn, err = ftps.openDataConn(port)
	if err != nil {
		return nil, err
	}

	_, err = ftps.request(cmd, expected)
	if err != nil {
		dataConn.Close()
		return nil, err
	}

	dataConn = ftps.upgradeConnToTLS(dataConn)

	return
}

func (ftps *FTPS) response(expected int) (message string, err error) {

	ftps.isConnEstablished()

	code, message, err := ftps.text.ReadResponse(expected)

	ftps.debugInfo(fmt.Sprintf("<*code*> %d", code))
	ftps.debugInfo("<*message*> " + message)

	return
}

func (ftps *FTPS) upgradeConnToTLS(conn net.Conn) (upgradedConn net.Conn) {

	var tlsConn *tls.Conn
	tlsConn = tls.Client(conn, &ftps.TLSConfig)

	tlsConn.Handshake()
	upgradedConn = net.Conn(tlsConn)

	// TODO verify that TLS connection is established

	return
}

func (ftps *FTPS) pasv() (port int, err error) {

	message, err := ftps.request("PASV", 227)
	if err != nil {
		return 0, err
	}

	start := strings.Index(message, "(")
	end := strings.LastIndex(message, ")")

	if start == -1 || end == -1 {
		err = errors.New("Invalid PASV response format")
		return 0, err
	}

	pasvData := strings.Split(message[start+1:end], ",")

	portPart1, err := strconv.Atoi(pasvData[4])
	if err != nil {
		return 0, err
	}

	portPart2, err := strconv.Atoi(pasvData[5])
	if err != nil {
		return 0, err
	}

	// Recompose port
	port = int(portPart1)*256 + int(portPart2)

	return
}

func (ftps *FTPS) PrintWorkingDirectory() (directory string, err error) {

	directory, err = ftps.request("PWD", 257)
	return
}

func (ftps *FTPS) ChangeWorkingDirectory(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("CWD %s", path), 250)
	return
}

func (ftps *FTPS) MakeDirectory(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("MKD %s", path), 257)
	return
}

func (ftps *FTPS) DeleteFile(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("DELE %s", path), 250)
	return
}

func (ftps *FTPS) RemoveDirectory(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("RMD %s", path), 250)
	return
}

func (ftps *FTPS) List() (entries []Entry, err error) {

	// TODO add support for MLSD

	dataConn, err := ftps.requestDataConn("LIST -a", 150) // TODO use also -L to resolve links?
	if err != nil {
		return
	}
	defer dataConn.Close()

	reader := bufio.NewReader(dataConn)
	for {
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}

		entry, err := ftps.parseEntryLine(line)
		entries = append(entries, *entry)
	}
	dataConn.Close()

	_, err = ftps.response(226)
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) parseEntryLine(line string) (entry *Entry, err error) {

	// TODO Function mostly copied from https://github.com/jlaffaye/ftp

	fields := strings.Fields(line)
	if len(fields) < 9 {
		return nil, errors.New("Unsupported line format.")
	}

	entry = &Entry{}

	// parse type
	switch fields[0][0] {
	case '-':
		entry.Type = EntryTypeFile
	case 'd':
		entry.Type = EntryTypeFolder
	case 'l':
		entry.Type = EntryTypeLink
	default:
		return nil, errors.New("Unknown entry type.")
	}

	// parse size
	size, err := strconv.ParseUint(fields[4], 10, 0)
	if err != nil {
		return nil, err
	}
	entry.Size = size

	// parse time
	var timeStr string
	if strings.Contains(fields[7], ":") { // this year
		thisYear, _, _ := time.Now().Date()
		timeStr = fmt.Sprintf("%s %s %s %s GMT", fields[6], fields[5], strconv.Itoa(thisYear)[2:4], fields[7])
	} else { // not this year
		timeStr = fmt.Sprintf("%s %s %s 00:00 GMT", fields[6], fields[5], fields[7][2:4])
	}
	t, err := time.Parse("_2 Jan 06 15:04 MST", timeStr)
	if err != nil {
		return nil, err
	}
	entry.Time = t // TODO set timezone

	// parse name
	entry.Name = strings.Join(fields[8:], " ")

	return
}

func (ftps *FTPS) StoreFile(remoteFilepath string, data []byte) (err error) {

	dataConn, err := ftps.requestDataConn(fmt.Sprintf("STOR %s", remoteFilepath), 150)
	if err != nil {
		return
	}
	defer dataConn.Close()

	count, err := dataConn.Write(data)
	if err != nil {
		return
	}
	dataConn.Close()

	if len(data) != count {
		return errors.New("file transfer not complete")
	}

	_, err = ftps.response(226)
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) RetrieveFile(remoteFilepath, localFilepath string) (err error) {

	dataConn, err := ftps.requestDataConn(fmt.Sprintf("RETR %s", remoteFilepath), 150)
	if err != nil {
		return
	}
	defer dataConn.Close()

	file, err := os.Create(localFilepath)
	if err != nil {
		return
	}
	defer file.Close()

	_, err = io.Copy(file, dataConn)
	if err != nil {
		return
	}
	dataConn.Close()

	_, err = ftps.response(226)
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) Quit() (err error) {

	_, err = ftps.request("QUIT", 221)
	if err != nil {
		return
	}
	ftps.conn.Close()

	return
}

func (ftps *FTPS) openDataConn(port int) (dataConn net.Conn, err error) {

	dataConn, err = net.Dial("tcp", fmt.Sprintf("%s:%d", ftps.host, port))
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) debugInfo(message string) {

	if ftps.Debug {
		log.Println(message)
	}
}
<commit_msg>checked error<commit_after>package ftps

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
	"time"
)

type FTPS struct {
	host string

	conn net.Conn
	text *textproto.Conn

	Debug     bool
	TLSConfig tls.Config
}

func (ftps *FTPS) Connect(host string, port int) (err error) {

	ftps.host = host

	ftps.conn, err = net.Dial("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return err
	}

	ftps.text = textproto.NewConn(ftps.conn)

	_, err = ftps.response(220)
	if err != nil {
		return err
	}

	_, err = ftps.request("AUTH TLS", 234)
	if err != nil {
		return err
	}

	ftps.conn = ftps.upgradeConnToTLS(ftps.conn)
	ftps.text = textproto.NewConn(ftps.conn) // TODO use sync or something similar?

	return
}

func (ftps *FTPS) isConnEstablished() {

	if ftps.conn == nil {
		panic("Connection is not established yet")
	}
}

func (ftps *FTPS) Login(username, password string) (err error) {

	ftps.isConnEstablished()

	_, err = ftps.request(fmt.Sprintf("USER %s", username), 331)
	if err != nil {
		return err
	}

	_, err = ftps.request(fmt.Sprintf("PASS %s", password), 230)
	if err != nil {
		return err
	}

	_, err = ftps.request("TYPE I", 200) // use binary mode
	if err != nil {
		return err
	}

	_, err = ftps.request("PBSZ 0", 200)
	if err != nil {
		return err
	}

	_, err = ftps.request("PROT P", 200) // encrypt data connection
	if err != nil {
		return err
	}

	return
}

func (ftps *FTPS) request(cmd string, expected int) (message string, err error) {

	ftps.isConnEstablished()

	ftps.debugInfo("<*cmd*> " + cmd)

	_, err = ftps.text.Cmd(cmd)
	if err != nil {
		return
	}

	message, err = ftps.response(expected)

	return
}

func (ftps *FTPS) requestDataConn(cmd string, expected int) (dataConn net.Conn, err error) {

	port, err := ftps.pasv()
	if err != nil {
		return
	}

	dataConn, err = ftps.openDataConn(port)
	if err != nil {
		return nil, err
	}

	_, err = ftps.request(cmd, expected)
	if err != nil {
		dataConn.Close()
		return nil, err
	}

	dataConn = ftps.upgradeConnToTLS(dataConn)

	return
}

func (ftps *FTPS) response(expected int) (message string, err error) {

	ftps.isConnEstablished()

	code, message, err := ftps.text.ReadResponse(expected)

	ftps.debugInfo(fmt.Sprintf("<*code*> %d", code))
	ftps.debugInfo("<*message*> " + message)

	return
}

func (ftps *FTPS) upgradeConnToTLS(conn net.Conn) (upgradedConn net.Conn) {

	var tlsConn *tls.Conn
	tlsConn = tls.Client(conn, &ftps.TLSConfig)

	tlsConn.Handshake()
	upgradedConn = net.Conn(tlsConn)

	// TODO verify that TLS connection is established

	return
}

func (ftps *FTPS) pasv() (port int, err error) {

	message, err := ftps.request("PASV", 227)
	if err != nil {
		return 0, err
	}

	start := strings.Index(message, "(")
	end := strings.LastIndex(message, ")")

	if start == -1 || end == -1 {
		err = errors.New("Invalid PASV response format")
		return 0, err
	}

	pasvData := strings.Split(message[start+1:end], ",")

	portPart1, err := strconv.Atoi(pasvData[4])
	if err != nil {
		return 0, err
	}

	portPart2, err := strconv.Atoi(pasvData[5])
	if err != nil {
		return 0, err
	}

	// Recompose port
	port = int(portPart1)*256 + int(portPart2)

	return
}

func (ftps *FTPS) PrintWorkingDirectory() (directory string, err error) {

	directory, err = ftps.request("PWD", 257)
	return
}

func (ftps *FTPS) ChangeWorkingDirectory(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("CWD %s", path), 250)
	return
}

func (ftps *FTPS) MakeDirectory(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("MKD %s", path), 257)
	return
}

func (ftps *FTPS) DeleteFile(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("DELE %s", path), 250)
	return
}

func (ftps *FTPS) RemoveDirectory(path string) (err error) {

	_, err = ftps.request(fmt.Sprintf("RMD %s", path), 250)
	return
}

func (ftps *FTPS) List() (entries []Entry, err error) {

	// TODO add support for MLSD

	dataConn, err := ftps.requestDataConn("LIST -a", 150) // TODO use also -L to resolve links?
	if err != nil {
		return
	}
	defer dataConn.Close()

	reader := bufio.NewReader(dataConn)
	for {
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}

		entry, err := ftps.parseEntryLine(line)
		if err != nil {
			return nil, err
		}
		entries = append(entries, *entry)
	}
	dataConn.Close()

	_, err = ftps.response(226)
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) parseEntryLine(line string) (entry *Entry, err error) {

	// TODO Function mostly copied from https://github.com/jlaffaye/ftp

	fields := strings.Fields(line)
	if len(fields) < 9 {
		return nil, errors.New("Unsupported line format.")
	}

	entry = &Entry{}

	// parse type
	switch fields[0][0] {
	case '-':
		entry.Type = EntryTypeFile
	case 'd':
		entry.Type = EntryTypeFolder
	case 'l':
		entry.Type = EntryTypeLink
	default:
		return nil, errors.New("Unknown entry type.")
	}

	// parse size
	size, err := strconv.ParseUint(fields[4], 10, 0)
	if err != nil {
		return nil, err
	}
	entry.Size = size

	// parse time
	var timeStr string
	if strings.Contains(fields[7], ":") { // this year
		thisYear, _, _ := time.Now().Date()
		timeStr = fmt.Sprintf("%s %s %s %s GMT", fields[6], fields[5], strconv.Itoa(thisYear)[2:4], fields[7])
	} else { // not this year
		timeStr = fmt.Sprintf("%s %s %s 00:00 GMT", fields[6], fields[5], fields[7][2:4])
	}
	t, err := time.Parse("_2 Jan 06 15:04 MST", timeStr)
	if err != nil {
		return nil, err
	}
	entry.Time = t // TODO set timezone

	// parse name
	entry.Name = strings.Join(fields[8:], " ")

	return
}

func (ftps *FTPS) StoreFile(remoteFilepath string, data []byte) (err error) {

	dataConn, err := ftps.requestDataConn(fmt.Sprintf("STOR %s", remoteFilepath), 150)
	if err != nil {
		return
	}
	defer dataConn.Close()

	count, err := dataConn.Write(data)
	if err != nil {
		return
	}
	dataConn.Close()

	if len(data) != count {
		return errors.New("file transfer not complete")
	}

	_, err = ftps.response(226)
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) RetrieveFile(remoteFilepath, localFilepath string) (err error) {

	dataConn, err := ftps.requestDataConn(fmt.Sprintf("RETR %s", remoteFilepath), 150)
	if err != nil {
		return
	}
	defer dataConn.Close()

	file, err := os.Create(localFilepath)
	if err != nil {
		return
	}
	defer file.Close()

	_, err = io.Copy(file, dataConn)
	if err != nil {
		return
	}
	dataConn.Close()

	_, err = ftps.response(226)
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) Quit() (err error) {

	_, err = ftps.request("QUIT", 221)
	if err != nil {
		return
	}
	ftps.conn.Close()

	return
}

func (ftps *FTPS) openDataConn(port int) (dataConn net.Conn, err error) {

	dataConn, err = net.Dial("tcp", fmt.Sprintf("%s:%d", ftps.host, port))
	if err != nil {
		return
	}

	return
}

func (ftps *FTPS) debugInfo(message string) {

	if ftps.Debug {
		log.Println(message)
	}
}
<|endoftext|>


<commit_before>package ploop

// A test suite, also serving as an example of how to use the package

import (
	"github.com/dustin/go-humanize"
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

var (
	old_pwd  string
	test_dir string
	d        Ploop
	snap     string
)

const baseDelta = "root.hdd"

func TestPrepare(t *testing.T) {
	var err error

	old_pwd, err = os.Getwd()
	if err != nil {
		t.Fatalf("Getwd: %s", err)
	}

	test_dir, err = ioutil.TempDir(old_pwd, "tmp-test")
	if err != nil {
		t.Fatalf("TempDir %q: %s", old_pwd, err)
	}

	err = os.Chdir(test_dir)
	if err != nil {
		t.Fatalf("Chdir %q: %s", test_dir, err)
	}

	SetVerboseLevel(255)
	SetLogLevel(1)
	SetLogFile("ploop.log")

	t.Logf("Running tests in %s", test_dir)
}

func TestUUID(t *testing.T) {
	uuid, e := UUID()
	if e != nil {
		t.Errorf("UUID: %s", e)
	}

	t.Logf("Got uuid %s", uuid)
}

func TestCreate(t *testing.T) {
	size := "384M"
	var p CreateParam

	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}
	p.Size = s
	p.File = baseDelta

	e = Create(&p)
	if e != nil {
		t.Fatalf("Create: %s", e)
	}
}

func TestOpen(t *testing.T) {
	var e error

	d, e = Open("DiskDescriptor.xml")
	if e != nil {
		t.Errorf("Open: %s", e)
	}
}

func TestMount(t *testing.T) {
	mnt := "mnt"

	e := os.Mkdir(mnt, 0755)
	if e != nil {
		t.Fatalf("os.Mkdir: %s", e)
	}

	p := MountParam{Target: mnt}
	dev, e := d.Mount(&p)
	if e != nil {
		t.Fatalf("Mount: %s", e)
	} else {
		t.Logf("Mounted; ploop device %s", dev)
	}
}

func resize(t *testing.T, size string, offline bool) {
	if offline && testing.Short() {
		t.Skip("skipping offline resize test in short mode.")
	}
	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}

	e = d.Resize(s, offline)
	if e != nil {
		t.Fatalf("Resize to %s (%d bytes) failed: %s", size, s, e)
	}
}

func TestResizeOnlineShrink(t *testing.T) {
	resize(t, "256MB", false)
}

func TestResizeOnlineGrow(t *testing.T) {
	resize(t, "512MB", false)
}

func TestSnapshot(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created online snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func copyFile(src, dst string) error {
	return exec.Command("cp", "-a", src, dst).Run
}

func testReplace(t *testing.T) {
	var p ReplaceParam
	newDelta := baseDelta + ".new"
	e := copyFile(baseDelta, newDelta)
	if e != nil {
		t.Fatalf("copyFile: %s", e)
	}

	p.File = newDelta
	p.CurFile = baseDelta
	p.Flags = KeepName
	e = d.Replace(&p)
	if e != nil {
		t.Fatalf("Replace: %s", e)
	}
}

func TestReplaceOnline(t *testing.T) {
	testReplace(t)
}

func TestDeleteSnapshot(t *testing.T) {
	e := d.DeleteSnapshot(snap)
	if e != nil {
		t.Fatalf("DeleteSnapshot: %s", e)
	} else {
		t.Logf("Deleted snapshot %s", snap)
	}
}

func TestUmount(t *testing.T) {
	e := d.Umount()
	if e != nil {
		t.Fatalf("Umount: %s", e)
	}
}

func TestResizeOfflineShrink(t *testing.T) {
	resize(t, "256MB", true)
}

func TestResizeOfflineGrow(t *testing.T) {
	resize(t, "512MB", true)
}

func TestResizeOfflineShrinkAgain(t *testing.T) {
	resize(t, "256MB", true)
}

func TestSnapshotOffline(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created offline snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func TestReplaceOffline(t *testing.T) {
	testReplace(t)
}

func TestSwitchSnapshot(t *testing.T) {
	e := d.SwitchSnapshot(snap)
	if e != nil {
		t.Fatalf("SwitchSnapshot: %s", e)
	} else {
		t.Logf("Switched to snapshot %s", snap)
	}
}

func TestFSInfo(t *testing.T) {
	i, e := FSInfo("DiskDescriptor.xml")

	if e != nil {
		t.Errorf("FSInfo: %v", e)
	} else {
		bTotal := i.Blocks * i.Blocksize
		bAvail := i.Blocks_free * i.Blocksize
		bUsed := bTotal - bAvail

		iTotal := i.Inodes
		iAvail := i.Inodes_free
		iUsed := iTotal - iAvail

		t.Logf("\n             Size       Used      Avail Use%%\n%7s %9s %10s %10s %3d%%\n%7s %9d %10d %10d %3d%%",
			"Blocks",
			humanize.Bytes(bTotal),
			humanize.Bytes(bUsed),
			humanize.Bytes(bAvail),
			100*bUsed/bTotal,
			"Inodes",
			iTotal,
			iUsed,
			iAvail,
			100*iUsed/iTotal)
		t.Logf("\nInode ratio: 1 inode per %s of disk space",
			humanize.Bytes(bTotal/iTotal))
	}
}

func TestImageInfo(t *testing.T) {
	i, e := d.ImageInfo()
	if e != nil {
		t.Errorf("ImageInfo: %v", e)
	} else {
		t.Logf("\n              Blocks  Blocksize       Size  Ver\n%20d %10d %10s %4d",
			i.Blocks, i.Blocksize,
			humanize.Bytes(512*i.Blocks),
			i.Version)
	}

}

// TestCleanup is the last test, removing files created by previous tests
func TestCleanup(t *testing.T) {
	d.Close()
	os.Chdir(old_pwd)
	os.RemoveAll(test_dir)
}
<commit_msg>test: fix copyFile()<commit_after>package ploop

// A test suite, also serving as an example of how to use the package

import (
	"github.com/dustin/go-humanize"
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

var (
	old_pwd  string
	test_dir string
	d        Ploop
	snap     string
)

const baseDelta = "root.hdd"

func TestPrepare(t *testing.T) {
	var err error

	old_pwd, err = os.Getwd()
	if err != nil {
		t.Fatalf("Getwd: %s", err)
	}

	test_dir, err = ioutil.TempDir(old_pwd, "tmp-test")
	if err != nil {
		t.Fatalf("TempDir %q: %s", old_pwd, err)
	}

	err = os.Chdir(test_dir)
	if err != nil {
		t.Fatalf("Chdir %q: %s", test_dir, err)
	}

	SetVerboseLevel(255)
	SetLogLevel(1)
	SetLogFile("ploop.log")

	t.Logf("Running tests in %s", test_dir)
}

func TestUUID(t *testing.T) {
	uuid, e := UUID()
	if e != nil {
		t.Errorf("UUID: %s", e)
	}

	t.Logf("Got uuid %s", uuid)
}

func TestCreate(t *testing.T) {
	size := "384M"
	var p CreateParam

	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}
	p.Size = s
	p.File = baseDelta

	e = Create(&p)
	if e != nil {
		t.Fatalf("Create: %s", e)
	}
}

func TestOpen(t *testing.T) {
	var e error

	d, e = Open("DiskDescriptor.xml")
	if e != nil {
		t.Errorf("Open: %s", e)
	}
}

func TestMount(t *testing.T) {
	mnt := "mnt"

	e := os.Mkdir(mnt, 0755)
	if e != nil {
		t.Fatalf("os.Mkdir: %s", e)
	}

	p := MountParam{Target: mnt}
	dev, e := d.Mount(&p)
	if e != nil {
		t.Fatalf("Mount: %s", e)
	} else {
		t.Logf("Mounted; ploop device %s", dev)
	}
}

func resize(t *testing.T, size string, offline bool) {
	if offline && testing.Short() {
		t.Skip("skipping offline resize test in short mode.")
	}
	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}

	e = d.Resize(s, offline)
	if e != nil {
		t.Fatalf("Resize to %s (%d bytes) failed: %s", size, s, e)
	}
}

func TestResizeOnlineShrink(t *testing.T) {
	resize(t, "256MB", false)
}

func TestResizeOnlineGrow(t *testing.T) {
	resize(t, "512MB", false)
}

func TestSnapshot(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created online snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func copyFile(src, dst string) error {
	return exec.Command("cp", "-a", src, dst).Run()
}

func testReplace(t *testing.T) {
	var p ReplaceParam
	newDelta := baseDelta + ".new"
	e := copyFile(baseDelta, newDelta)
	if e != nil {
		t.Fatalf("copyFile: %s", e)
	}

	p.File = newDelta
	p.CurFile = baseDelta
	p.Flags = KeepName
	e = d.Replace(&p)
	if e != nil {
		t.Fatalf("Replace: %s", e)
	}
}

func TestReplaceOnline(t *testing.T) {
	testReplace(t)
}

func TestDeleteSnapshot(t *testing.T) {
	e := d.DeleteSnapshot(snap)
	if e != nil {
		t.Fatalf("DeleteSnapshot: %s", e)
	} else {
		t.Logf("Deleted snapshot %s", snap)
	}
}

func TestUmount(t *testing.T) {
	e := d.Umount()
	if e != nil {
		t.Fatalf("Umount: %s", e)
	}
}

func TestResizeOfflineShrink(t *testing.T) {
	resize(t, "256MB", true)
}

func TestResizeOfflineGrow(t *testing.T) {
	resize(t, "512MB", true)
}

func TestResizeOfflineShrinkAgain(t *testing.T) {
	resize(t, "256MB", true)
}

func TestSnapshotOffline(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created offline snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func TestReplaceOffline(t *testing.T) {
	testReplace(t)
}

func TestSwitchSnapshot(t *testing.T) {
	e := d.SwitchSnapshot(snap)
	if e != nil {
		t.Fatalf("SwitchSnapshot: %s", e)
	} else {
		t.Logf("Switched to snapshot %s", snap)
	}
}

func TestFSInfo(t *testing.T) {
	i, e := FSInfo("DiskDescriptor.xml")

	if e != nil {
		t.Errorf("FSInfo: %v", e)
	} else {
		bTotal := i.Blocks * i.Blocksize
		bAvail := i.Blocks_free * i.Blocksize
		bUsed := bTotal - bAvail

		iTotal := i.Inodes
		iAvail := i.Inodes_free
		iUsed := iTotal - iAvail

		t.Logf("\n             Size       Used      Avail Use%%\n%7s %9s %10s %10s %3d%%\n%7s %9d %10d %10d %3d%%",
			"Blocks",
			humanize.Bytes(bTotal),
			humanize.Bytes(bUsed),
			humanize.Bytes(bAvail),
			100*bUsed/bTotal,
			"Inodes",
			iTotal,
			iUsed,
			iAvail,
			100*iUsed/iTotal)
		t.Logf("\nInode ratio: 1 inode per %s of disk space",
			humanize.Bytes(bTotal/iTotal))
	}
}

func TestImageInfo(t *testing.T) {
	i, e := d.ImageInfo()
	if e != nil {
		t.Errorf("ImageInfo: %v", e)
	} else {
		t.Logf("\n              Blocks  Blocksize       Size  Ver\n%20d %10d %10s %4d",
			i.Blocks, i.Blocksize,
			humanize.Bytes(512*i.Blocks),
			i.Version)
	}

}

// TestCleanup is the last test, removing files created by previous tests
func TestCleanup(t *testing.T) {
	d.Close()
	os.Chdir(old_pwd)
	os.RemoveAll(test_dir)
}
<|endoftext|>


<commit_before>// +build gofuzz

package merkletree

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"io"
)

// Fuzz is called by go-fuzz to look for inputs to BuildReaderProof that will
// not verify correctly.
func Fuzz(data []byte) int {
	// Use the first two bytes to determine the proof index.
	if len(data) < 2 {
		return -1
	}
	index := 256*uint64(data[0]) + uint64(data[1])
	data = data[2:]

	// Build a reader proof for index 'index' using the remaining data as input
	// to the reader. '64' is chosen as the only input size because that is the
	// size relevant to the Sia project.
	merkleRoot, proofSet, numLeaves, err := BuildReaderProof(bytes.NewReader(data), sha256.New(), 64, index)
	if err != nil {
		return 0
	}
	if !VerifyProof(sha256.New(), merkleRoot, proofSet, index, numLeaves) {
		panic("verification failed!")
	}

	// Output is more interesting when there is enough data to contain the
	// index.
	if uint64(len(data)) > 64*index {
		return 1
	}
	return 0
}

// FuzzReadSubTreesWithProof can be used by go-fuzz to test creating a merkle
// tree from cached subTrees and creating/proving a merkle proof on this tree.
func FuzzReadSubTreesWithProof(data []byte) int {
	// Use the first two bytes to determine the proof index.
	if len(data) < 2 {
		return -1
	}
	// Use the first two bytes to determine the proof index.
	index := 256*uint64(data[0]) + uint64(data[1])
	tree := New(sha256.New())
	tree.SetIndex(index)
	data = data[2:]

	// 32 is the length of a sha256 hash and 4 bytes are used for the height of
	// the subTree.
	subTreeSize := 4 + 36
	err := tree.readSubTrees(bytes.NewReader(data), subTreeSize)
	if err != io.ErrUnexpectedEOF {
		return 0
	} else if err == io.ErrUnexpectedEOF && len(data) < subTreeSize {
		return -1
	}

	// Create and verify the proof.
	merkleRoot, proofSet, proofIndex, numLeaves := tree.Prove()
	if !VerifyProof(sha256.New(), merkleRoot, proofSet, proofIndex, numLeaves) {
		panic("verification failed!")
	}
	// Output is more interesting when there is enough data to contain the
	// index.
	if uint64(len(data)) > 64*index {
		return 1
	}
	return 0
}

// FuzzReadSubTreesNoProof can be used by go-fuzz to test creating a merkle
// tree from cached subTrees.
func FuzzReadSubTreesNoProof(data []byte) int {
	tree := New(sha256.New())

	// 32 is the length of a sha256 hash and 4 bytes are used for the height of
	// the subTree.
	subTreeSize := 4 + 36
	err := tree.readSubTrees(bytes.NewReader(data), subTreeSize)
	if err != io.ErrUnexpectedEOF {
		return 0
	} else if err == io.ErrUnexpectedEOF && len(data) < subTreeSize {
		return -1
	}
	if tree.head != nil && tree.Root() == nil {
		panic("root shouldn't be nil for a non-empty tree")
	}
	// The data is better if it contains many subTrees.
	if len(data) > 100*subTreeSize {
		return 1
	}
	return 0
}

// readSubTrees is a helper function that maps the data from a io.Reader to
// subTrees and adds them to a tree.
func (t *Tree) readSubTrees(r io.Reader, subTreeSize int) error {
	if subTreeSize <= 4 {
		return errors.New(`the subTree must at least contain 4 bytes for the height
		plus some data for the sum`)
	}
	for {
		subTree := make([]byte, subTreeSize)
		_, readErr := io.ReadFull(r, subTree)
		if readErr == io.EOF {
			// All data has been read.
			break
		} else if readErr != nil {
			return readErr
		}
		// The first 4 bytes of the subTree are mapped to a height in range
		// [0,50].
		height := int(binary.LittleEndian.Uint32(subTree[:4])) % 51
		sum := subTree[4:]
		if height > 0 {
			if err := t.PushSubTree(height, sum); err != nil {
				return err
			}
		} else {
			t.Push(sum)
		}
	}
	return nil
}
<commit_msg>fix fuzzy crashes<commit_after>// +build gofuzz

package merkletree

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"io"
)

// Fuzz is called by go-fuzz to look for inputs to BuildReaderProof that will
// not verify correctly.
func Fuzz(data []byte) int {
	// Use the first two bytes to determine the proof index.
	if len(data) < 2 {
		return -1
	}
	index := 256*uint64(data[0]) + uint64(data[1])
	data = data[2:]

	// Build a reader proof for index 'index' using the remaining data as input
	// to the reader. '64' is chosen as the only input size because that is the
	// size relevant to the Sia project.
	merkleRoot, proofSet, numLeaves, err := BuildReaderProof(bytes.NewReader(data), sha256.New(), 64, index)
	if err != nil {
		return 0
	}
	if !VerifyProof(sha256.New(), merkleRoot, proofSet, index, numLeaves) {
		panic("verification failed!")
	}

	// Output is more interesting when there is enough data to contain the
	// index.
	if uint64(len(data)) > 64*index {
		return 1
	}
	return 0
}

// FuzzReadSubTreesWithProof can be used by go-fuzz to test creating a merkle
// tree from cached subTrees and creating/proving a merkle proof on this tree.
func FuzzReadSubTreesWithProof(data []byte) int {
	// Use the first two bytes to determine the proof index.
	if len(data) < 2 {
		return -1
	}
	// Use the first two bytes to determine the proof index.
	index := 256*uint64(data[0]) + uint64(data[1])
	tree := New(sha256.New())
	tree.SetIndex(index)
	data = data[2:]

	// 32 is the length of a sha256 hash and 4 bytes are used for the height of
	// the subTree.
	subTreeSize := 4 + 36
	err := tree.readSubTrees(bytes.NewReader(data), subTreeSize)
	if err != io.ErrUnexpectedEOF {
		return 0
	} else if err == io.ErrUnexpectedEOF && len(data) < subTreeSize {
		return -1
	}

	// Create and verify the proof.
	merkleRoot, proofSet, proofIndex, numLeaves := tree.Prove()
	if len(proofSet) == 0 {
		// proofIndex wasn't reached while creating proof.
		return 0
	}
	if !VerifyProof(sha256.New(), merkleRoot, proofSet, proofIndex, numLeaves) {
		panic("verification failed!")
	}
	// Output is more interesting when there is enough data to contain the
	// index.
	if uint64(len(data)) > uint64(subTreeSize)*index {
		return 1
	}
	return 0
}

// FuzzReadSubTreesNoProof can be used by go-fuzz to test creating a merkle
// tree from cached subTrees.
func FuzzReadSubTreesNoProof(data []byte) int {
	tree := New(sha256.New())

	// 32 is the length of a sha256 hash and 4 bytes are used for the height of
	// the subTree.
	subTreeSize := 4 + 36
	err := tree.readSubTrees(bytes.NewReader(data), subTreeSize)
	if err != io.ErrUnexpectedEOF {
		return 0
	} else if err == io.ErrUnexpectedEOF && len(data) < subTreeSize {
		return -1
	}
	if tree.head != nil && tree.Root() == nil {
		panic("root shouldn't be nil for a non-empty tree")
	}
	// The data is better if it contains many subTrees.
	if len(data) > 100*subTreeSize {
		return 1
	}
	return 0
}

// readSubTrees is a helper function that maps the data from a io.Reader to
// subTrees and adds them to a tree.
func (t *Tree) readSubTrees(r io.Reader, subTreeSize int) error {
	if subTreeSize <= 4 {
		return errors.New(`the subTree must at least contain 4 bytes for the height
		plus some data for the sum`)
	}
	for {
		subTree := make([]byte, subTreeSize)
		_, readErr := io.ReadFull(r, subTree)
		if readErr == io.EOF {
			// All data has been read.
			break
		} else if readErr != nil {
			return readErr
		}
		// The first 4 bytes of the subTree are mapped to a height in range
		// [0,50].
		height := int(binary.LittleEndian.Uint32(subTree[:4])) % 51
		sum := subTree[4:]
		if height > 0 {
			if err := t.PushSubTree(height, sum); err != nil {
				return err
			}
		} else {
			t.Push(sum)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"sync"
	"time"
)

var (
	debug                  bool
	verbose                bool
	info                   bool
	quiet                  bool
	force                  bool
	usemove                bool
	usecacheFallback       bool
	retryGitCommands       bool
	pfMode                 bool
	pfLocation             string
	dryRun                 bool
	check4update           bool
	checkSum               bool
	moduleDirParam         string
	cacheDirParam          string
	branchParam            string
	moduleParam            string
	configFile             string
	config                 ConfigSettings
	mutex                  sync.Mutex
	empty                  struct{}
	syncGitCount           int
	syncForgeCount         int
	needSyncGitCount       int
	needSyncForgeCount     int
	syncGitTime            float64
	syncForgeTime          float64
	ioGitTime              float64
	ioForgeTime            float64
	forgeJsonParseTime     float64
	metadataJsonParseTime  float64
	gmetadataJsonParseTime float64
	buildtime              string
	uniqueForgeModules     map[string]ForgeModule
	latestForgeModules     LatestForgeModules
	maxworker              int
	maxExtractworker       int
)

type LatestForgeModules struct {
	sync.RWMutex
	m map[string]string
}

// ConfigSettings contains the key value pairs from the g10k config file
type ConfigSettings struct {
	CacheDir                 string `yaml:"cachedir"`
	ForgeCacheDir            string
	ModulesCacheDir          string
	EnvCacheDir              string
	Git                      Git
	Forge                    Forge
	Sources                  map[string]Source
	Timeout                  int  `yaml:"timeout"`
	IgnoreUnreachableModules bool `yaml:"ignore_unreachable_modules"`
	Maxworker                int  `yaml:"maxworker"`
	MaxExtractworker         int  `yaml:"maxextractworker"`
	UseCacheFallback         bool `yaml:"use_cache_fallback"`
	RetryGitCommands         bool `yaml:"retry_git_commands"`
}

type Forge struct {
	Baseurl string `yaml:"baseurl"`
}

type Git struct {
	privateKey string `yaml:"private_key"`
	username   string
}

// Source contains basic information about a Puppet environment repository
type Source struct {
	Remote             string
	Basedir            string
	Prefix             string
	PrivateKey         string `yaml:"private_key"`
	ForceForgeVersions bool   `yaml:"force_forge_versions"`
	WarnMissingBranch  bool   `yaml:"warn_if_branch_is_missing"`
	ExitIfUnreachable  bool   `yaml:"exit_if_unreachable"`
}

// Puppetfile contains the key value pairs from the Puppetfile
type Puppetfile struct {
	moduleDir     string
	forgeBaseURL  string
	forgeCacheTtl time.Duration
	forgeModules  map[string]ForgeModule
	gitModules    map[string]GitModule
	privateKey    string
	source        string
	workDir       string
}

// ForgeModule contains information (Version, Name, Author, md5 checksum, file size of the tar.gz archive, Forge BaseURL if custom) about a Puppetlabs Forge module
type ForgeModule struct {
	version   string
	name      string
	author    string
	md5sum    string
	fileSize  int64
	baseUrl   string
	cacheTtl  time.Duration
	sha256sum string
}

// GitModule contains information about a Git Puppet module
type GitModule struct {
	privateKey        string
	git               string
	branch            string
	tag               string
	commit            string
	ref               string
	link              bool
	ignoreUnreachable bool
	fallback          []string
	installPath       string
}

// ForgeResult is returned by queryForgeAPI and contains if and which version of the Puppetlabs Forge module needs to be downloaded
type ForgeResult struct {
	needToGet     bool
	versionNumber string
	md5sum        string
	fileSize      int64
}

// ExecResult contains the exit code and output of an external command (e.g. git)
type ExecResult struct {
	returnCode int
	output     string
}

func main() {

	var (
		configFileFlag = flag.String("config", "", "which config file to use")
		versionFlag    = flag.Bool("version", false, "show build time and version number")
	)
	flag.StringVar(&branchParam, "branch", "", "which git branch of the Puppet environment to update, e.g. core_foobar")
	flag.StringVar(&moduleParam, "module", "", "which module of the Puppet environment to update, e.g. stdlib")
	flag.StringVar(&moduleDirParam, "moduledir", "", "allows overriding of Puppetfile specific moduledir setting, the folder in which Puppet modules will be extracted")
	flag.StringVar(&cacheDirParam, "cachedir", "", "allows overriding of the g10k config file cachedir setting, the folder in which g10k will download git repositories and Forge modules")
	flag.IntVar(&maxworker, "maxworker", 50, "how many Goroutines are allowed to run in parallel for Git and Forge module resolving")
	flag.IntVar(&maxExtractworker, "maxextractworker", 20, "how many Goroutines are allowed to run in parallel for local Git and Forge module extracting processes (git clone, untar and gunzip)")
	flag.BoolVar(&pfMode, "puppetfile", false, "install all modules from Puppetfile in cwd")
	flag.StringVar(&pfLocation, "puppetfilelocation", "./Puppetfile", "which Puppetfile to use in -puppetfile mode")
	flag.BoolVar(&force, "force", false, "purge the Puppet environment directory and do a full sync")
	flag.BoolVar(&dryRun, "dryrun", false, "do not modify anything, just print what would be changed")
	flag.BoolVar(&usemove, "usemove", false, "do not use hardlinks to populate your Puppet environments with Puppetlabs Forge modules. Instead uses simple move commands and purges the Forge cache directory after each run! (Useful for g10k runs inside a Docker container)")
	flag.BoolVar(&check4update, "check4update", false, "only check if the is newer version of the Puppet module avaialable. Does implicitly set dryrun to true")
	flag.BoolVar(&checkSum, "checksum", false, "get the md5 check sum for each Puppetlabs Forge module and verify the integrity of the downloaded archive. Increases g10k run time!")
	flag.BoolVar(&debug, "debug", false, "log debug output, defaults to false")
	flag.BoolVar(&verbose, "verbose", false, "log verbose output, defaults to false")
	flag.BoolVar(&info, "info", false, "log info output, defaults to false")
	flag.BoolVar(&quiet, "quiet", false, "no output, defaults to false")
	flag.BoolVar(&usecacheFallback, "usecachefallback", false, "if g10k should try to use its cache for sources and modules instead of failing")
	flag.BoolVar(&retryGitCommands, "retrygitcommands", false, "if g10k should purge the local repository and retry a failed git command (clone or remote update) instead of failing")
	flag.Parse()

	configFile = *configFileFlag
	version := *versionFlag

	if version {
		fmt.Println("g10k Version 0.4 Build time:", buildtime, "UTC")
		os.Exit(0)
	}

	if check4update {
		dryRun = true
	}

	// check for git executable dependency
	if _, err := exec.LookPath("git"); err != nil {
		Fatalf("Error: could not find 'git' executable in PATH")
	}

	target := ""
	before := time.Now()
	if len(configFile) > 0 {
		if usemove {
			Fatalf("Error: -usemove parameter is only allowed in -puppetfile mode!")
		}
		if pfMode {
			Fatalf("Error: -puppetfile parameter is not allowed with -config parameter!")
		}
		if usecacheFallback {
			config.UseCacheFallback = true
		}
		Debugf("Using as config file: " + configFile)
		config = readConfigfile(configFile)
		target = configFile
		if len(branchParam) > 0 {
			resolvePuppetEnvironment(branchParam)
			target += " with branch " + branchParam
		} else {
			resolvePuppetEnvironment("")
		}
	} else {
		if pfMode {
			Debugf("Trying to use as Puppetfile: " + pfLocation)
			sm := make(map[string]Source)
			sm["cmdlineparam"] = Source{Basedir: "."}
			cachedir := "/tmp/g10k"
			if len(os.Getenv("g10k_cachedir")) > 0 {
				cachedir = os.Getenv("g10k_cachedir")
				cachedir = checkDirAndCreate(cachedir, "cachedir environment variable g10k_cachedir")
				Debugf("Found environment variable g10k_cachedir set to: " + cachedir)
			} else if len(cacheDirParam) > 0 {
				Debugf("Using -cachedir parameter set to : " + cacheDirParam)
				cachedir = checkDirAndCreate(cacheDirParam, "cachedir CLI param")
			} else {
				cachedir = checkDirAndCreate(cachedir, "cachedir default value")
			}
			//config = ConfigSettings{CacheDir: cachedir, ForgeCacheDir: cachedir, ModulesCacheDir: cachedir, EnvCacheDir: cachedir, Forge:{Baseurl: "https://forgeapi.puppetlabs.com"}, Sources: sm}
			forgeDefaultSettings := Forge{Baseurl: "https://forgeapi.puppetlabs.com"}
			config = ConfigSettings{CacheDir: cachedir, ForgeCacheDir: cachedir, ModulesCacheDir: cachedir, EnvCacheDir: cachedir, Sources: sm, Forge: forgeDefaultSettings, Maxworker: maxworker, UseCacheFallback: usecacheFallback, MaxExtractworker: maxExtractworker}
			target = pfLocation
			puppetfile := readPuppetfile(target, "", "cmdlineparam", false)
			puppetfile.workDir = "."
			pfm := make(map[string]Puppetfile)
			pfm["cmdlineparam"] = puppetfile
			resolvePuppetfile(pfm)
		} else {
			Fatalf("Error: you need to specify at least a config file or use the Puppetfile mode\nExample call: " + os.Args[0] + " -config test.yaml or " + os.Args[0] + " -puppetfile\n")
		}
	}

	if usemove {
		// we can not reuse the Forge cache at all when -usemove gets used, because we can not delete the -latest link for some reason
		defer purgeDir(config.ForgeCacheDir, "main() -puppetfile mode with -usemove parameter")
	}

	// DEBUG
	//pf := make(map[string]Puppetfile)
	//pf["core_fullmanaged"] = readPuppetfile("/tmp/core/core_fullmanaged/", "/home/andpaul/dev/go/src/github.com/xorpaul/g10k/portal_envs")
	//pf["itodsi_corosync"] = readPuppetfile("/tmp/itodsi/itodsi_corosync/", "/home/andpaul/dev/go/src/github.com/xorpaul/g10k/portal_envs")
	//resolvePuppetfile(pf)
	//resolveGitRepositories(config)
	//resolveForgeModules(configSettings.forge)
	//doModuleInstallOrNothing("camptocamp-postfix-1.2.2", "/tmp/g10k/camptocamp-postfix-1.2.2")
	//doModuleInstallOrNothing("saz-resolv_conf-latest")
	//readModuleMetadata("/tmp/g10k/forge/camptocamp-postfix-1.2.2/metadata.json")

	Debugf("Forge response JSON parsing took " + strconv.FormatFloat(forgeJsonParseTime, 'f', 4, 64) + " seconds")
	Debugf("Forge modules metadata.json parsing took " + strconv.FormatFloat(metadataJsonParseTime, 'f', 4, 64) + " seconds")

	if !check4update && !quiet {
		fmt.Println("Synced", target, "with", syncGitCount, "git repositories and", syncForgeCount, "Forge modules in "+strconv.FormatFloat(time.Since(before).Seconds(), 'f', 1, 64)+"s with git ("+strconv.FormatFloat(syncGitTime, 'f', 1, 64)+"s sync, I/O", strconv.FormatFloat(ioGitTime, 'f', 1, 64)+"s) and Forge ("+strconv.FormatFloat(syncForgeTime, 'f', 1, 64)+"s query+download, I/O", strconv.FormatFloat(ioForgeTime, 'f', 1, 64)+"s) using", strconv.Itoa(config.Maxworker), "resolv and", strconv.Itoa(config.MaxExtractworker), "extract workers")
	}
	if dryRun && (needSyncForgeCount > 0 || needSyncGitCount > 0) {
		os.Exit(1)
	}
}
<commit_msg>bump version to 0.4.1<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"sync"
	"time"
)

var (
	debug                  bool
	verbose                bool
	info                   bool
	quiet                  bool
	force                  bool
	usemove                bool
	usecacheFallback       bool
	retryGitCommands       bool
	pfMode                 bool
	pfLocation             string
	dryRun                 bool
	check4update           bool
	checkSum               bool
	moduleDirParam         string
	cacheDirParam          string
	branchParam            string
	moduleParam            string
	configFile             string
	config                 ConfigSettings
	mutex                  sync.Mutex
	empty                  struct{}
	syncGitCount           int
	syncForgeCount         int
	needSyncGitCount       int
	needSyncForgeCount     int
	syncGitTime            float64
	syncForgeTime          float64
	ioGitTime              float64
	ioForgeTime            float64
	forgeJsonParseTime     float64
	metadataJsonParseTime  float64
	gmetadataJsonParseTime float64
	buildtime              string
	uniqueForgeModules     map[string]ForgeModule
	latestForgeModules     LatestForgeModules
	maxworker              int
	maxExtractworker       int
)

type LatestForgeModules struct {
	sync.RWMutex
	m map[string]string
}

// ConfigSettings contains the key value pairs from the g10k config file
type ConfigSettings struct {
	CacheDir                 string `yaml:"cachedir"`
	ForgeCacheDir            string
	ModulesCacheDir          string
	EnvCacheDir              string
	Git                      Git
	Forge                    Forge
	Sources                  map[string]Source
	Timeout                  int  `yaml:"timeout"`
	IgnoreUnreachableModules bool `yaml:"ignore_unreachable_modules"`
	Maxworker                int  `yaml:"maxworker"`
	MaxExtractworker         int  `yaml:"maxextractworker"`
	UseCacheFallback         bool `yaml:"use_cache_fallback"`
	RetryGitCommands         bool `yaml:"retry_git_commands"`
}

type Forge struct {
	Baseurl string `yaml:"baseurl"`
}

type Git struct {
	privateKey string `yaml:"private_key"`
	username   string
}

// Source contains basic information about a Puppet environment repository
type Source struct {
	Remote             string
	Basedir            string
	Prefix             string
	PrivateKey         string `yaml:"private_key"`
	ForceForgeVersions bool   `yaml:"force_forge_versions"`
	WarnMissingBranch  bool   `yaml:"warn_if_branch_is_missing"`
	ExitIfUnreachable  bool   `yaml:"exit_if_unreachable"`
}

// Puppetfile contains the key value pairs from the Puppetfile
type Puppetfile struct {
	moduleDir     string
	forgeBaseURL  string
	forgeCacheTtl time.Duration
	forgeModules  map[string]ForgeModule
	gitModules    map[string]GitModule
	privateKey    string
	source        string
	workDir       string
}

// ForgeModule contains information (Version, Name, Author, md5 checksum, file size of the tar.gz archive, Forge BaseURL if custom) about a Puppetlabs Forge module
type ForgeModule struct {
	version   string
	name      string
	author    string
	md5sum    string
	fileSize  int64
	baseUrl   string
	cacheTtl  time.Duration
	sha256sum string
}

// GitModule contains information about a Git Puppet module
type GitModule struct {
	privateKey        string
	git               string
	branch            string
	tag               string
	commit            string
	ref               string
	link              bool
	ignoreUnreachable bool
	fallback          []string
	installPath       string
}

// ForgeResult is returned by queryForgeAPI and contains if and which version of the Puppetlabs Forge module needs to be downloaded
type ForgeResult struct {
	needToGet     bool
	versionNumber string
	md5sum        string
	fileSize      int64
}

// ExecResult contains the exit code and output of an external command (e.g. git)
type ExecResult struct {
	returnCode int
	output     string
}

func main() {

	var (
		configFileFlag = flag.String("config", "", "which config file to use")
		versionFlag    = flag.Bool("version", false, "show build time and version number")
	)
	flag.StringVar(&branchParam, "branch", "", "which git branch of the Puppet environment to update, e.g. core_foobar")
	flag.StringVar(&moduleParam, "module", "", "which module of the Puppet environment to update, e.g. stdlib")
	flag.StringVar(&moduleDirParam, "moduledir", "", "allows overriding of Puppetfile specific moduledir setting, the folder in which Puppet modules will be extracted")
	flag.StringVar(&cacheDirParam, "cachedir", "", "allows overriding of the g10k config file cachedir setting, the folder in which g10k will download git repositories and Forge modules")
	flag.IntVar(&maxworker, "maxworker", 50, "how many Goroutines are allowed to run in parallel for Git and Forge module resolving")
	flag.IntVar(&maxExtractworker, "maxextractworker", 20, "how many Goroutines are allowed to run in parallel for local Git and Forge module extracting processes (git clone, untar and gunzip)")
	flag.BoolVar(&pfMode, "puppetfile", false, "install all modules from Puppetfile in cwd")
	flag.StringVar(&pfLocation, "puppetfilelocation", "./Puppetfile", "which Puppetfile to use in -puppetfile mode")
	flag.BoolVar(&force, "force", false, "purge the Puppet environment directory and do a full sync")
	flag.BoolVar(&dryRun, "dryrun", false, "do not modify anything, just print what would be changed")
	flag.BoolVar(&usemove, "usemove", false, "do not use hardlinks to populate your Puppet environments with Puppetlabs Forge modules. Instead uses simple move commands and purges the Forge cache directory after each run! (Useful for g10k runs inside a Docker container)")
	flag.BoolVar(&check4update, "check4update", false, "only check if the is newer version of the Puppet module avaialable. Does implicitly set dryrun to true")
	flag.BoolVar(&checkSum, "checksum", false, "get the md5 check sum for each Puppetlabs Forge module and verify the integrity of the downloaded archive. Increases g10k run time!")
	flag.BoolVar(&debug, "debug", false, "log debug output, defaults to false")
	flag.BoolVar(&verbose, "verbose", false, "log verbose output, defaults to false")
	flag.BoolVar(&info, "info", false, "log info output, defaults to false")
	flag.BoolVar(&quiet, "quiet", false, "no output, defaults to false")
	flag.BoolVar(&usecacheFallback, "usecachefallback", false, "if g10k should try to use its cache for sources and modules instead of failing")
	flag.BoolVar(&retryGitCommands, "retrygitcommands", false, "if g10k should purge the local repository and retry a failed git command (clone or remote update) instead of failing")
	flag.Parse()

	configFile = *configFileFlag
	version := *versionFlag

	if version {
		fmt.Println("g10k Version 0.4.1 Build time:", buildtime, "UTC")
		os.Exit(0)
	}

	if check4update {
		dryRun = true
	}

	// check for git executable dependency
	if _, err := exec.LookPath("git"); err != nil {
		Fatalf("Error: could not find 'git' executable in PATH")
	}

	target := ""
	before := time.Now()
	if len(configFile) > 0 {
		if usemove {
			Fatalf("Error: -usemove parameter is only allowed in -puppetfile mode!")
		}
		if pfMode {
			Fatalf("Error: -puppetfile parameter is not allowed with -config parameter!")
		}
		if usecacheFallback {
			config.UseCacheFallback = true
		}
		Debugf("Using as config file: " + configFile)
		config = readConfigfile(configFile)
		target = configFile
		if len(branchParam) > 0 {
			resolvePuppetEnvironment(branchParam)
			target += " with branch " + branchParam
		} else {
			resolvePuppetEnvironment("")
		}
	} else {
		if pfMode {
			Debugf("Trying to use as Puppetfile: " + pfLocation)
			sm := make(map[string]Source)
			sm["cmdlineparam"] = Source{Basedir: "."}
			cachedir := "/tmp/g10k"
			if len(os.Getenv("g10k_cachedir")) > 0 {
				cachedir = os.Getenv("g10k_cachedir")
				cachedir = checkDirAndCreate(cachedir, "cachedir environment variable g10k_cachedir")
				Debugf("Found environment variable g10k_cachedir set to: " + cachedir)
			} else if len(cacheDirParam) > 0 {
				Debugf("Using -cachedir parameter set to : " + cacheDirParam)
				cachedir = checkDirAndCreate(cacheDirParam, "cachedir CLI param")
			} else {
				cachedir = checkDirAndCreate(cachedir, "cachedir default value")
			}
			//config = ConfigSettings{CacheDir: cachedir, ForgeCacheDir: cachedir, ModulesCacheDir: cachedir, EnvCacheDir: cachedir, Forge:{Baseurl: "https://forgeapi.puppetlabs.com"}, Sources: sm}
			forgeDefaultSettings := Forge{Baseurl: "https://forgeapi.puppetlabs.com"}
			config = ConfigSettings{CacheDir: cachedir, ForgeCacheDir: cachedir, ModulesCacheDir: cachedir, EnvCacheDir: cachedir, Sources: sm, Forge: forgeDefaultSettings, Maxworker: maxworker, UseCacheFallback: usecacheFallback, MaxExtractworker: maxExtractworker}
			target = pfLocation
			puppetfile := readPuppetfile(target, "", "cmdlineparam", false)
			puppetfile.workDir = "."
			pfm := make(map[string]Puppetfile)
			pfm["cmdlineparam"] = puppetfile
			resolvePuppetfile(pfm)
		} else {
			Fatalf("Error: you need to specify at least a config file or use the Puppetfile mode\nExample call: " + os.Args[0] + " -config test.yaml or " + os.Args[0] + " -puppetfile\n")
		}
	}

	if usemove {
		// we can not reuse the Forge cache at all when -usemove gets used, because we can not delete the -latest link for some reason
		defer purgeDir(config.ForgeCacheDir, "main() -puppetfile mode with -usemove parameter")
	}

	// DEBUG
	//pf := make(map[string]Puppetfile)
	//pf["core_fullmanaged"] = readPuppetfile("/tmp/core/core_fullmanaged/", "/home/andpaul/dev/go/src/github.com/xorpaul/g10k/portal_envs")
	//pf["itodsi_corosync"] = readPuppetfile("/tmp/itodsi/itodsi_corosync/", "/home/andpaul/dev/go/src/github.com/xorpaul/g10k/portal_envs")
	//resolvePuppetfile(pf)
	//resolveGitRepositories(config)
	//resolveForgeModules(configSettings.forge)
	//doModuleInstallOrNothing("camptocamp-postfix-1.2.2", "/tmp/g10k/camptocamp-postfix-1.2.2")
	//doModuleInstallOrNothing("saz-resolv_conf-latest")
	//readModuleMetadata("/tmp/g10k/forge/camptocamp-postfix-1.2.2/metadata.json")

	Debugf("Forge response JSON parsing took " + strconv.FormatFloat(forgeJsonParseTime, 'f', 4, 64) + " seconds")
	Debugf("Forge modules metadata.json parsing took " + strconv.FormatFloat(metadataJsonParseTime, 'f', 4, 64) + " seconds")

	if !check4update && !quiet {
		fmt.Println("Synced", target, "with", syncGitCount, "git repositories and", syncForgeCount, "Forge modules in "+strconv.FormatFloat(time.Since(before).Seconds(), 'f', 1, 64)+"s with git ("+strconv.FormatFloat(syncGitTime, 'f', 1, 64)+"s sync, I/O", strconv.FormatFloat(ioGitTime, 'f', 1, 64)+"s) and Forge ("+strconv.FormatFloat(syncForgeTime, 'f', 1, 64)+"s query+download, I/O", strconv.FormatFloat(ioForgeTime, 'f', 1, 64)+"s) using", strconv.Itoa(config.Maxworker), "resolv and", strconv.Itoa(config.MaxExtractworker), "extract workers")
	}
	if dryRun && (needSyncForgeCount > 0 || needSyncGitCount > 0) {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package gamedayapi

import (
	"log"
	"os"
	"net/http"
	"io/ioutil"
	"encoding/xml"
	s "strings"
)

type Game struct {
	AwayAPMP			string		`xml:"away_ampm,attr"`
	AwayCode			string		`xml:"away_code,attr"`
	AwayLoss			string		`xml:"away_loss,attr"`
	AwayTeamCity		string		`xml:"away_team_city,attr"`
	AwayTeamId			string		`xml:"away_team_id,attr"`
	AwayTeamName		string		`xml:"away_team_name,attr"`
	AwayTime			string		`xml:"away_time,attr"`
	AwayTimezone		string		`xml:"away_time_zone,attr"`
	AwayWin				string		`xml:"away_win,attr"`
	HomeAMPM			string		`xml:"home_ampm,attr"`
	HomeCode			string		`xml:"home_code,attr"`
	HomeLoss			string		`xml:"home_loss,attr"`
	HomeTeamCity		string		`xml:"home_team_city,attr"`
	HomeTeamId			string		`xml:"home_team_id,attr"`
	HomeTeamName		string		`xml:"home_team_name,attr"`
	HomeTime			string		`xml:"home_time,attr"`
	HomeTimezone		string		`xml:"home_time_zone,attr"`
	HomeWin				string		`xml:"home_win,attr"`
	Id					string		`xml:"id,attr"`
	GamePk				string		`xml:"game_pk,attr"`
	Timezone			string		`xml:"time_zone,attr"`
	Venue				string		`xml:"venue,attr"`
	GameDataDirectory	string		`xml:"game_data_directory,attr"`

	boxScore BoxScore
}

func GameFor(teamCode string, date string) *Game {
	epg := EpgFor(date)
	game := epg.GameForTeam(teamCode)
	return game
}

func (game *Game) BoxScore() *BoxScore {
	if len(game.boxScore.GameId) == 0 {
		filePath := BaseCachePath() + game.GameDataDirectory + "/boxscore.xml"
		if _, err := os.Stat(filePath); os.IsNotExist(err) {
			fetchAndCache(filePath, &game.boxScore)
		} else {
			body, _ := ioutil.ReadFile(filePath)
			xml.Unmarshal(body, &game.boxScore)
		}
	}
	return &game.boxScore
}

func fetchAndCache(filePath string, val interface{}) {
	resp, err := http.Get(GamedayHostname + filePath)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}

	xml.Unmarshal(body, val)
	cacheFile(filePath, body)
}

func cacheFile(filePath string, body []byte) {
	localCachePath := BaseCachePath() + filePath[0:s.LastIndex(filePath, "/")]
	log.Println(localCachePath)
	os.MkdirAll(localCachePath, (os.FileMode)(0775))
	log.Println(filePath[s.LastIndex(filePath, "/"):])
	f, err := os.Create(localCachePath + filePath[s.LastIndex(filePath, "/"):])
	f.Write(body)
	check(err)
	defer f.Close()
}

<commit_msg>Fixing caching to get it to the correct local directory.<commit_after>package gamedayapi

import (
	"log"
	"os"
	"net/http"
	"io/ioutil"
	"encoding/xml"
	s "strings"
)

type Game struct {
	AwayAPMP			string		`xml:"away_ampm,attr"`
	AwayCode			string		`xml:"away_code,attr"`
	AwayLoss			string		`xml:"away_loss,attr"`
	AwayTeamCity		string		`xml:"away_team_city,attr"`
	AwayTeamId			string		`xml:"away_team_id,attr"`
	AwayTeamName		string		`xml:"away_team_name,attr"`
	AwayTime			string		`xml:"away_time,attr"`
	AwayTimezone		string		`xml:"away_time_zone,attr"`
	AwayWin				string		`xml:"away_win,attr"`
	HomeAMPM			string		`xml:"home_ampm,attr"`
	HomeCode			string		`xml:"home_code,attr"`
	HomeLoss			string		`xml:"home_loss,attr"`
	HomeTeamCity		string		`xml:"home_team_city,attr"`
	HomeTeamId			string		`xml:"home_team_id,attr"`
	HomeTeamName		string		`xml:"home_team_name,attr"`
	HomeTime			string		`xml:"home_time,attr"`
	HomeTimezone		string		`xml:"home_time_zone,attr"`
	HomeWin				string		`xml:"home_win,attr"`
	Id					string		`xml:"id,attr"`
	GamePk				string		`xml:"game_pk,attr"`
	Timezone			string		`xml:"time_zone,attr"`
	Venue				string		`xml:"venue,attr"`
	GameDataDirectory	string		`xml:"game_data_directory,attr"`

	boxScore BoxScore
}

func GameFor(teamCode string, date string) *Game {
	epg := EpgFor(date)
	game := epg.GameForTeam(teamCode)
	return game
}

func (game *Game) BoxScore() *BoxScore {
	if len(game.boxScore.GameId) == 0 {
		filePath := game.GameDataDirectory + "/boxscore.xml"
		if _, err := os.Stat(filePath); os.IsNotExist(err) {
			fetchAndCache(filePath, &game.boxScore)
		} else {
			body, _ := ioutil.ReadFile(filePath)
			xml.Unmarshal(body, &game.boxScore)
		}
	}
	return &game.boxScore
}

func fetchAndCache(filePath string, val interface{}) {
	resp, err := http.Get(GamedayHostname + filePath)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}

	xml.Unmarshal(body, val)
	cacheFile(filePath, body)
}

func cacheFile(filePath string, body []byte) {
	localCachePath := BaseCachePath() + filePath[0:s.LastIndex(filePath, "/")]
	os.MkdirAll(localCachePath, (os.FileMode)(0775))
	f, err := os.Create(localCachePath + filePath[s.LastIndex(filePath, "/"):])
	f.Write(body)
	check(err)
	defer f.Close()
}
<|endoftext|>


<commit_before>package tftest

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	tfjson "github.com/hashicorp/terraform-json"
)

// WorkingDir represents a distinct working directory that can be used for
// running tests. Each test should construct its own WorkingDir by calling
// NewWorkingDir or RequireNewWorkingDir on its package's singleton
// tftest.Helper.
type WorkingDir struct {
	h         *Helper
	baseDir   string
	configDir string
}

// Close deletes the directories and files created to represent the receiving
// working directory. After this method is called, the working directory object
// is invalid and may no longer be used.
func (wd *WorkingDir) Close() error {
	return os.RemoveAll(wd.baseDir)
}

// SetConfig sets a new configuration for the working directory.
//
// This must be called at least once before any call to Init, Plan, Apply, or
// Destroy to establish the configuration. Any previously-set configuration is
// discarded and any saved plan is cleared.
func (wd *WorkingDir) SetConfig(cfg string) error {
	// Each call to SetConfig creates a new directory under our baseDir.
	// We create them within so that our final cleanup step will delete them
	// automatically without any additional tracking.
	configDir, err := ioutil.TempDir(wd.baseDir, "config")
	if err != nil {
		return err
	}
	configFilename := filepath.Join(configDir, "test.tf")
	err = ioutil.WriteFile(configFilename, []byte(cfg), 0700)
	if err != nil {
		return err
	}
	wd.configDir = configDir

	// Changing configuration invalidates any saved plan.
	err = wd.ClearPlan()
	if err != nil {
		return err
	}
	return nil
}

// RequireSetConfig is a variant of SetConfig that will fail the test via the
// given TestControl if the configuration cannot be set.
func (wd *WorkingDir) RequireSetConfig(t TestControl, cfg string) {
	t.Helper()
	if err := wd.SetConfig(cfg); err != nil {
		t := testingT{t}
		t.Fatalf("failed to set config: %s", err)
	}
}

// ClearState deletes any Terraform state present in the working directory.
//
// Any remote objects tracked by the state are not destroyed first, so this
// will leave them dangling in the remote system.
func (wd *WorkingDir) ClearState() error {
	err := os.Remove(filepath.Join(wd.baseDir, "terraform.tfstate"))
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearState is a variant of ClearState that will fail the test via the
// given TestControl if the state cannot be cleared.
func (wd *WorkingDir) RequireClearState(t TestControl) {
	t.Helper()
	if err := wd.ClearState(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear state: %s", err)
	}
}

// ClearPlan deletes any saved plan present in the working directory.
func (wd *WorkingDir) ClearPlan() error {
	err := os.Remove(wd.planFilename())
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearPlan is a variant of ClearPlan that will fail the test via the
// given TestControl if the plan cannot be cleared.
func (wd *WorkingDir) RequireClearPlan(t TestControl) {
	t.Helper()
	if err := wd.ClearPlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear plan: %s", err)
	}
}

func (wd *WorkingDir) init(pluginDir string) error {
	return wd.runTerraform("init", "-plugin-dir="+pluginDir, wd.configDir)
}

// Init runs "terraform init" for the given working directory, forcing Terraform
// to use the current version of the plugin under test.
func (wd *WorkingDir) Init() error {
	if wd.configDir == "" {
		return fmt.Errorf("must call SetConfig before Init")
	}
	return wd.init(wd.h.PluginDir())
}

// RequireInit is a variant of Init that will fail the test via the given
// TestControl if init fails.
func (wd *WorkingDir) RequireInit(t TestControl) {
	t.Helper()
	if err := wd.Init(); err != nil {
		t := testingT{t}
		t.Fatalf("init failed: %s", err)
	}
}

// InitPrevious runs "terraform init" for the given working directory, forcing
// Terraform to use the previous version of the plugin under test.
//
// This method will panic if no previous plugin version is available. Use
// HasPreviousVersion or RequirePreviousVersion on the test helper singleton
// to check this first.
func (wd *WorkingDir) InitPrevious() error {
	if wd.configDir == "" {
		return fmt.Errorf("must call SetConfig before InitPrevious")
	}
	return wd.init(wd.h.PreviousPluginDir())
}

// RequireInitPrevious is a variant of InitPrevious that will fail the test
// via the given TestControl if init fails.
func (wd *WorkingDir) RequireInitPrevious(t TestControl) {
	t.Helper()
	if err := wd.InitPrevious(); err != nil {
		t := testingT{t}
		t.Fatalf("init failed: %s", err)
	}
}

func (wd *WorkingDir) planFilename() string {
	return filepath.Join(wd.baseDir, "tfplan")
}

// CreatePlan runs "terraform plan" to create a saved plan file, which if successful
// will then be used for the next call to Apply.
func (wd *WorkingDir) CreatePlan() error {
	return wd.runTerraform("plan", "-out=tfplan", wd.configDir)
}

// RequireCreatePlan is a variant of CreatePlan that will fail the test via
// the given TestControl if plan creation fails.
func (wd *WorkingDir) RequireCreatePlan(t TestControl) {
	t.Helper()
	if err := wd.CreatePlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to create plan: %s", err)
	}
}

// Apply runs "terraform apply". If CreatePlan has previously completed
// successfully and the saved plan has not been cleared in the meantime then
// this will apply the saved plan. Otherwise, it will implicitly create a new
// plan and apply it.
func (wd *WorkingDir) Apply() error {
	args := []string{"apply"}
	if wd.HasSavedPlan() {
		args = append(args, "tfplan")
	} else {
		args = append(args, "-auto-approve")
		args = append(args, wd.configDir)
	}
	return wd.runTerraform(args...)
}

// RequireApply is a variant of Apply that will fail the test via
// the given TestControl if the apply operation fails.
func (wd *WorkingDir) RequireApply(t TestControl) {
	t.Helper()
	if err := wd.Apply(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to apply: %s", err)
	}
}

// Destroy runs "terraform destroy". It does not consider or modify any saved
// plan, and is primarily for cleaning up at the end of a test run.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) Destroy() error {
	args := []string{"destroy", "-auto-approve", wd.configDir}
	return wd.runTerraform(args...)
}

// RequireDestroy is a variant of Destroy that will fail the test via
// the given TestControl if the destroy operation fails.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) RequireDestroy(t TestControl) {
	t.Helper()
	if err := wd.Destroy(); err != nil {
		t := testingT{t}
		t.Logf("WARNING: destroy failed, so remote objects may still exist and be subject to billing")
		t.Fatalf("failed to destroy: %s", err)
	}
}

// HasSavedPlan returns true if there is a saved plan in the working directory. If
// so, a subsequent call to Apply will apply that saved plan.
func (wd *WorkingDir) HasSavedPlan() bool {
	_, err := os.Stat(wd.planFilename())
	return err == nil
}

// SavedPlan returns an object describing the current saved plan file, if any.
//
// If no plan is saved or if the plan file cannot be read, SavedPlan returns
// an error.
func (wd *WorkingDir) SavedPlan() (*tfjson.Plan, error) {
	if !wd.HasSavedPlan() {
		return nil, fmt.Errorf("there is no current saved plan")
	}

	var ret tfjson.Plan
	err := wd.runTerraformJSON(&ret, "show", "-json", wd.planFilename())
	if err != nil {
		return nil, err
	}
	if err := ret.Validate(); err != nil {
		return nil, err
	}
	return &ret, nil
}

// RequireSavedPlan is a variant of SavedPlan that will fail the test via
// the given TestControl if the plan cannot be read.
func (wd *WorkingDir) RequireSavedPlan(t TestControl) *tfjson.Plan {
	t.Helper()
	ret, err := wd.SavedPlan()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read saved plan: %s", err)
	}
	return ret
}

// State returns an object describing the current state.
//
// If the state cannot be read, State returns an error.
func (wd *WorkingDir) State() (*tfjson.State, error) {
	var ret tfjson.State
	err := wd.runTerraformJSON(&ret, "show", "-json")
	if err != nil {
		return nil, err
	}
	if err := ret.Validate(); err != nil {
		return nil, err
	}
	return &ret, nil
}

// RequireState is a variant of State that will fail the test via
// the given TestControl if the state cannot be read.
func (wd *WorkingDir) RequireState(t TestControl) *tfjson.State {
	t.Helper()
	ret, err := wd.State()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read state plan: %s", err)
	}
	return ret
}
<commit_msg>add import functionality<commit_after>package tftest

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	tfjson "github.com/hashicorp/terraform-json"
)

// WorkingDir represents a distinct working directory that can be used for
// running tests. Each test should construct its own WorkingDir by calling
// NewWorkingDir or RequireNewWorkingDir on its package's singleton
// tftest.Helper.
type WorkingDir struct {
	h         *Helper
	baseDir   string
	configDir string
}

// Close deletes the directories and files created to represent the receiving
// working directory. After this method is called, the working directory object
// is invalid and may no longer be used.
func (wd *WorkingDir) Close() error {
	return os.RemoveAll(wd.baseDir)
}

// SetConfig sets a new configuration for the working directory.
//
// This must be called at least once before any call to Init, Plan, Apply, or
// Destroy to establish the configuration. Any previously-set configuration is
// discarded and any saved plan is cleared.
func (wd *WorkingDir) SetConfig(cfg string) error {
	// Each call to SetConfig creates a new directory under our baseDir.
	// We create them within so that our final cleanup step will delete them
	// automatically without any additional tracking.
	configDir, err := ioutil.TempDir(wd.baseDir, "config")
	if err != nil {
		return err
	}
	configFilename := filepath.Join(configDir, "test.tf")
	err = ioutil.WriteFile(configFilename, []byte(cfg), 0700)
	if err != nil {
		return err
	}
	wd.configDir = configDir

	// Changing configuration invalidates any saved plan.
	err = wd.ClearPlan()
	if err != nil {
		return err
	}
	return nil
}

// RequireSetConfig is a variant of SetConfig that will fail the test via the
// given TestControl if the configuration cannot be set.
func (wd *WorkingDir) RequireSetConfig(t TestControl, cfg string) {
	t.Helper()
	if err := wd.SetConfig(cfg); err != nil {
		t := testingT{t}
		t.Fatalf("failed to set config: %s", err)
	}
}

// ClearState deletes any Terraform state present in the working directory.
//
// Any remote objects tracked by the state are not destroyed first, so this
// will leave them dangling in the remote system.
func (wd *WorkingDir) ClearState() error {
	err := os.Remove(filepath.Join(wd.baseDir, "terraform.tfstate"))
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearState is a variant of ClearState that will fail the test via the
// given TestControl if the state cannot be cleared.
func (wd *WorkingDir) RequireClearState(t TestControl) {
	t.Helper()
	if err := wd.ClearState(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear state: %s", err)
	}
}

// ClearPlan deletes any saved plan present in the working directory.
func (wd *WorkingDir) ClearPlan() error {
	err := os.Remove(wd.planFilename())
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearPlan is a variant of ClearPlan that will fail the test via the
// given TestControl if the plan cannot be cleared.
func (wd *WorkingDir) RequireClearPlan(t TestControl) {
	t.Helper()
	if err := wd.ClearPlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear plan: %s", err)
	}
}

func (wd *WorkingDir) init(pluginDir string) error {
	return wd.runTerraform("init", "-plugin-dir="+pluginDir, wd.configDir)
}

// Init runs "terraform init" for the given working directory, forcing Terraform
// to use the current version of the plugin under test.
func (wd *WorkingDir) Init() error {
	if wd.configDir == "" {
		return fmt.Errorf("must call SetConfig before Init")
	}
	return wd.init(wd.h.PluginDir())
}

// RequireInit is a variant of Init that will fail the test via the given
// TestControl if init fails.
func (wd *WorkingDir) RequireInit(t TestControl) {
	t.Helper()
	if err := wd.Init(); err != nil {
		t := testingT{t}
		t.Fatalf("init failed: %s", err)
	}
}

// InitPrevious runs "terraform init" for the given working directory, forcing
// Terraform to use the previous version of the plugin under test.
//
// This method will panic if no previous plugin version is available. Use
// HasPreviousVersion or RequirePreviousVersion on the test helper singleton
// to check this first.
func (wd *WorkingDir) InitPrevious() error {
	if wd.configDir == "" {
		return fmt.Errorf("must call SetConfig before InitPrevious")
	}
	return wd.init(wd.h.PreviousPluginDir())
}

// RequireInitPrevious is a variant of InitPrevious that will fail the test
// via the given TestControl if init fails.
func (wd *WorkingDir) RequireInitPrevious(t TestControl) {
	t.Helper()
	if err := wd.InitPrevious(); err != nil {
		t := testingT{t}
		t.Fatalf("init failed: %s", err)
	}
}

func (wd *WorkingDir) planFilename() string {
	return filepath.Join(wd.baseDir, "tfplan")
}

// CreatePlan runs "terraform plan" to create a saved plan file, which if successful
// will then be used for the next call to Apply.
func (wd *WorkingDir) CreatePlan() error {
	return wd.runTerraform("plan", "-out=tfplan", wd.configDir)
}

// RequireCreatePlan is a variant of CreatePlan that will fail the test via
// the given TestControl if plan creation fails.
func (wd *WorkingDir) RequireCreatePlan(t TestControl) {
	t.Helper()
	if err := wd.CreatePlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to create plan: %s", err)
	}
}

// Apply runs "terraform apply". If CreatePlan has previously completed
// successfully and the saved plan has not been cleared in the meantime then
// this will apply the saved plan. Otherwise, it will implicitly create a new
// plan and apply it.
func (wd *WorkingDir) Apply() error {
	args := []string{"apply"}
	if wd.HasSavedPlan() {
		args = append(args, "tfplan")
	} else {
		args = append(args, "-auto-approve")
		args = append(args, wd.configDir)
	}
	return wd.runTerraform(args...)
}

// RequireApply is a variant of Apply that will fail the test via
// the given TestControl if the apply operation fails.
func (wd *WorkingDir) RequireApply(t TestControl) {
	t.Helper()
	if err := wd.Apply(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to apply: %s", err)
	}
}

// Destroy runs "terraform destroy". It does not consider or modify any saved
// plan, and is primarily for cleaning up at the end of a test run.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) Destroy() error {
	args := []string{"destroy", "-auto-approve", wd.configDir}
	return wd.runTerraform(args...)
}

// RequireDestroy is a variant of Destroy that will fail the test via
// the given TestControl if the destroy operation fails.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) RequireDestroy(t TestControl) {
	t.Helper()
	if err := wd.Destroy(); err != nil {
		t := testingT{t}
		t.Logf("WARNING: destroy failed, so remote objects may still exist and be subject to billing")
		t.Fatalf("failed to destroy: %s", err)
	}
}

// HasSavedPlan returns true if there is a saved plan in the working directory. If
// so, a subsequent call to Apply will apply that saved plan.
func (wd *WorkingDir) HasSavedPlan() bool {
	_, err := os.Stat(wd.planFilename())
	return err == nil
}

// SavedPlan returns an object describing the current saved plan file, if any.
//
// If no plan is saved or if the plan file cannot be read, SavedPlan returns
// an error.
func (wd *WorkingDir) SavedPlan() (*tfjson.Plan, error) {
	if !wd.HasSavedPlan() {
		return nil, fmt.Errorf("there is no current saved plan")
	}

	var ret tfjson.Plan
	err := wd.runTerraformJSON(&ret, "show", "-json", wd.planFilename())
	if err != nil {
		return nil, err
	}
	if err := ret.Validate(); err != nil {
		return nil, err
	}
	return &ret, nil
}

// RequireSavedPlan is a variant of SavedPlan that will fail the test via
// the given TestControl if the plan cannot be read.
func (wd *WorkingDir) RequireSavedPlan(t TestControl) *tfjson.Plan {
	t.Helper()
	ret, err := wd.SavedPlan()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read saved plan: %s", err)
	}
	return ret
}

// State returns an object describing the current state.
//
// If the state cannot be read, State returns an error.
func (wd *WorkingDir) State() (*tfjson.State, error) {
	var ret tfjson.State
	err := wd.runTerraformJSON(&ret, "show", "-json")
	if err != nil {
		return nil, err
	}
	if err := ret.Validate(); err != nil {
		return nil, err
	}
	return &ret, nil
}

// RequireState is a variant of State that will fail the test via
// the given TestControl if the state cannot be read.
func (wd *WorkingDir) RequireState(t TestControl) *tfjson.State {
	t.Helper()
	ret, err := wd.State()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read state plan: %s", err)
	}
	return ret
}

// Import runs terraform import
func (wd *WorkingDir) Import(resource, id string) error {
	args := []string{"import", "-config=" + wd.configDir, resource, id}
	return wd.runTerraform(args...)
}

// RequireImport is a variant of Import that will fail the test via
// the given TestControl if the import is non successful.
func (wd *WorkingDir) RequireImport(t TestControl, resource, id string) {
	t.Helper()
	if err := wd.Import(resource, id); err != nil {
		t := testingT{t}
		t.Fatalf("failed to import: %s", err)
	}
}
<|endoftext|>


<commit_before>// Package crowdauth provides middleware for Crowd SSO logins
//
// Goals:
//  1) drop-in authentication against Crowd SSO
//  2) make it easy to use Crowd SSO as part of your own auth flow
package crowdauth // import "go.jona.me/crowd/crowdauth"

import (
	"go.jona.me/crowd"
	"html/template"
	"log"
	"net/http"
	"time"
)

type SSO struct {
	CrowdApp            *crowd.Crowd
	LoginPage           AuthLoginPage
	LoginTemplate       *template.Template
	ClientAddressFinder ClientAddressFinder
	CookieConfig        crowd.CookieConfig
}

// The AuthLoginPage type extends the normal http.HandlerFunc type
// with a boolean return to indicate login success or failure.
type AuthLoginPage func(http.ResponseWriter, *http.Request, *SSO) bool

// ClientAddressFinder type represents a function that returns the
// end-user's IP address, allowing you to handle cases where the address
// is masked by proxy servers etc by checking headers or whatever to find
// the user's address
type ClientAddressFinder func(*http.Request) (string, error)

var authErr string = "unauthorized, login required"

func DefaultClientAddressFinder(r *http.Request) (addr string, err error) {
	return r.RemoteAddr, nil
}

// New returns creates a new instance of SSO
func New(user string, password string, crowdURL string) (s *SSO, err error) {
	s = &SSO{}
	s.LoginPage = loginPage
	s.ClientAddressFinder = DefaultClientAddressFinder
	s.LoginTemplate = template.Must(template.New("authPage").Parse(defLoginPage))

	cwd, err := crowd.New(user, password, crowdURL)
	if err != nil {
		return s, err
	}
	s.CrowdApp = &cwd
	s.CookieConfig, err = s.CrowdApp.GetCookieConfig()
	if err != nil {
		return s, err
	}

	return s, nil
}

// Handler provides HTTP middleware using http.Handler chaining
// that requires user authentication via Atlassian Crowd SSO.
func (s *SSO) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.loginHandler(w, r) == false {
			return
		}
		h.ServeHTTP(w, r)
	})

}

func (s *SSO) loginHandler(w http.ResponseWriter, r *http.Request) bool {
	ck, err := r.Cookie(s.CookieConfig.Name)

	if err == http.ErrNoCookie {
		// no cookie so show login page if GET
		// if POST check if login and handle
		// if fail, show login page with message
		if r.Method == "GET" {
			s.LoginPage(w, r, s)
		} else if r.Method == "POST" {
			authOK := s.LoginPage(w, r, s)
			if authOK == true {
				// Redirect for fresh pass through auth etc on success
				http.Redirect(w, r, r.RequestURI, http.StatusTemporaryRedirect)
				return false
			} else {
				log.Printf("crowdauth: authentication failed\n")
			}
		} else {
			http.Error(w, authErr, http.StatusUnauthorized)
		}
		return false
	} else {
		// validate cookie or show login page
		host, err := s.ClientAddressFinder(r)
		if err != nil {
			log.Printf("crowdauth: could not get remote addr: %s\n", err)
			return false
		}

		_, err = s.CrowdApp.ValidateSession(ck.Value, host)
		if err != nil {
			log.Printf("crowdauth: could not validate cookie, deleting because: %s\n", err)
			s.EndSession(w, r)
			s.LoginPage(w, r, s)
			return false
		}

		// valid cookie so fallthrough
	}
	return true
}

func (s *SSO) Login(user string, pass string, addr string) (cs crowd.Session, err error) {
	cs, err = s.CrowdApp.NewSession(user, pass, addr)
	return cs, err
}

func (s *SSO) Logout(w http.ResponseWriter, r *http.Request, newURL string) {
	s.EndSession(w, r)
	http.Redirect(w, r, newURL, http.StatusTemporaryRedirect)
}

func (s *SSO) StartSession(w http.ResponseWriter, cs crowd.Session) (err error) {
	ck := http.Cookie{
		Name:    s.CookieConfig.Name,
		Domain:  s.CookieConfig.Domain,
		Secure:  s.CookieConfig.Secure,
		Value:   cs.Token,
		Expires: cs.Expires,
	}
	http.SetCookie(w, &ck)
	return nil
}

func (s *SSO) EndSession(w http.ResponseWriter, r *http.Request) {
	currentCookie, _ := r.Cookie(s.CookieConfig.Name)
	newCookie := &http.Cookie{
		Name:    s.CookieConfig.Name,
		Domain:  s.CookieConfig.Domain,
		Secure:  s.CookieConfig.Secure,
		MaxAge:  -1,
		Expires: time.Unix(1, 0),
		Value:   "LOGGED-OUT",
	}
	s.CrowdApp.InvalidateSession(currentCookie.Value)

	log.Printf("Got cookie to remove: %+v\n", currentCookie)
	log.Printf("Removal cookie: %+v\n", newCookie)
	http.SetCookie(w, newCookie)
}

func loginPage(w http.ResponseWriter, r *http.Request, s *SSO) bool {
	if r.Method == "GET" { // show login page and bail
		showLoginPage(w, s)
		return false
	} else if r.Method == "POST" {
		user := r.FormValue("username")
		pass := r.FormValue("password")
		host, err := s.ClientAddressFinder(r)
		if err != nil {
			log.Printf("crowdauth: could not get remote addr: %s\n", err)
			showLoginPage(w, s)
			return false
		}

		sess, err := s.Login(user, pass, host)
		if err != nil {
			log.Printf("crowdauth: login/new session failed: %s\n", err)
			showLoginPage(w, s)
			return false
		}

		err = s.StartSession(w, sess)
		if err != nil {
			log.Printf("crowdauth: could not save session: %s\n", err)
			showLoginPage(w, s)
			return false
		}
	} else {
		return false
	}

	return true
}

func showLoginPage(w http.ResponseWriter, s *SSO) {
	err := s.LoginTemplate.ExecuteTemplate(w, "authPage", nil)
	if err != nil {
		log.Printf("crowdauth: could not exec template: %s\n", err)
	}
}
<commit_msg>fix doc typo<commit_after>// Package crowdauth provides middleware for Crowd SSO logins
//
// Goals:
//  1) drop-in authentication against Crowd SSO
//  2) make it easy to use Crowd SSO as part of your own auth flow
package crowdauth // import "go.jona.me/crowd/crowdauth"

import (
	"go.jona.me/crowd"
	"html/template"
	"log"
	"net/http"
	"time"
)

type SSO struct {
	CrowdApp            *crowd.Crowd
	LoginPage           AuthLoginPage
	LoginTemplate       *template.Template
	ClientAddressFinder ClientAddressFinder
	CookieConfig        crowd.CookieConfig
}

// The AuthLoginPage type extends the normal http.HandlerFunc type
// with a boolean return to indicate login success or failure.
type AuthLoginPage func(http.ResponseWriter, *http.Request, *SSO) bool

// ClientAddressFinder type represents a function that returns the
// end-user's IP address, allowing you to handle cases where the address
// is masked by proxy servers etc by checking headers or whatever to find
// the user's address
type ClientAddressFinder func(*http.Request) (string, error)

var authErr string = "unauthorized, login required"

func DefaultClientAddressFinder(r *http.Request) (addr string, err error) {
	return r.RemoteAddr, nil
}

// New creates a new instance of SSO
func New(user string, password string, crowdURL string) (s *SSO, err error) {
	s = &SSO{}
	s.LoginPage = loginPage
	s.ClientAddressFinder = DefaultClientAddressFinder
	s.LoginTemplate = template.Must(template.New("authPage").Parse(defLoginPage))

	cwd, err := crowd.New(user, password, crowdURL)
	if err != nil {
		return s, err
	}
	s.CrowdApp = &cwd
	s.CookieConfig, err = s.CrowdApp.GetCookieConfig()
	if err != nil {
		return s, err
	}

	return s, nil
}

// Handler provides HTTP middleware using http.Handler chaining
// that requires user authentication via Atlassian Crowd SSO.
func (s *SSO) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.loginHandler(w, r) == false {
			return
		}
		h.ServeHTTP(w, r)
	})

}

func (s *SSO) loginHandler(w http.ResponseWriter, r *http.Request) bool {
	ck, err := r.Cookie(s.CookieConfig.Name)

	if err == http.ErrNoCookie {
		// no cookie so show login page if GET
		// if POST check if login and handle
		// if fail, show login page with message
		if r.Method == "GET" {
			s.LoginPage(w, r, s)
		} else if r.Method == "POST" {
			authOK := s.LoginPage(w, r, s)
			if authOK == true {
				// Redirect for fresh pass through auth etc on success
				http.Redirect(w, r, r.RequestURI, http.StatusTemporaryRedirect)
				return false
			} else {
				log.Printf("crowdauth: authentication failed\n")
			}
		} else {
			http.Error(w, authErr, http.StatusUnauthorized)
		}
		return false
	} else {
		// validate cookie or show login page
		host, err := s.ClientAddressFinder(r)
		if err != nil {
			log.Printf("crowdauth: could not get remote addr: %s\n", err)
			return false
		}

		_, err = s.CrowdApp.ValidateSession(ck.Value, host)
		if err != nil {
			log.Printf("crowdauth: could not validate cookie, deleting because: %s\n", err)
			s.EndSession(w, r)
			s.LoginPage(w, r, s)
			return false
		}

		// valid cookie so fallthrough
	}
	return true
}

func (s *SSO) Login(user string, pass string, addr string) (cs crowd.Session, err error) {
	cs, err = s.CrowdApp.NewSession(user, pass, addr)
	return cs, err
}

func (s *SSO) Logout(w http.ResponseWriter, r *http.Request, newURL string) {
	s.EndSession(w, r)
	http.Redirect(w, r, newURL, http.StatusTemporaryRedirect)
}

func (s *SSO) StartSession(w http.ResponseWriter, cs crowd.Session) (err error) {
	ck := http.Cookie{
		Name:    s.CookieConfig.Name,
		Domain:  s.CookieConfig.Domain,
		Secure:  s.CookieConfig.Secure,
		Value:   cs.Token,
		Expires: cs.Expires,
	}
	http.SetCookie(w, &ck)
	return nil
}

func (s *SSO) EndSession(w http.ResponseWriter, r *http.Request) {
	currentCookie, _ := r.Cookie(s.CookieConfig.Name)
	newCookie := &http.Cookie{
		Name:    s.CookieConfig.Name,
		Domain:  s.CookieConfig.Domain,
		Secure:  s.CookieConfig.Secure,
		MaxAge:  -1,
		Expires: time.Unix(1, 0),
		Value:   "LOGGED-OUT",
	}
	s.CrowdApp.InvalidateSession(currentCookie.Value)

	log.Printf("Got cookie to remove: %+v\n", currentCookie)
	log.Printf("Removal cookie: %+v\n", newCookie)
	http.SetCookie(w, newCookie)
}

func loginPage(w http.ResponseWriter, r *http.Request, s *SSO) bool {
	if r.Method == "GET" { // show login page and bail
		showLoginPage(w, s)
		return false
	} else if r.Method == "POST" {
		user := r.FormValue("username")
		pass := r.FormValue("password")
		host, err := s.ClientAddressFinder(r)
		if err != nil {
			log.Printf("crowdauth: could not get remote addr: %s\n", err)
			showLoginPage(w, s)
			return false
		}

		sess, err := s.Login(user, pass, host)
		if err != nil {
			log.Printf("crowdauth: login/new session failed: %s\n", err)
			showLoginPage(w, s)
			return false
		}

		err = s.StartSession(w, sess)
		if err != nil {
			log.Printf("crowdauth: could not save session: %s\n", err)
			showLoginPage(w, s)
			return false
		}
	} else {
		return false
	}

	return true
}

func showLoginPage(w http.ResponseWriter, s *SSO) {
	err := s.LoginTemplate.ExecuteTemplate(w, "authPage", nil)
	if err != nil {
		log.Printf("crowdauth: could not exec template: %s\n", err)
	}
}
<|endoftext|>