<commit_before>// Package pop3 provides an implementation of the Post Office
// Protocol - Version 3.
package pop3

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/mail"
	"strconv"
	"strings"
)

// Dial connects to the address on the named network.
func Dial(address string) (c *Client, err error) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		return
	}
	return NewClient(conn)
}

// Dial connects to the address on the named network using tls.
func DialTLS(address string) (c *Client, err error) {
	conn, err := tls.Dial("tcp", address, nil)
	if err != nil {
		return
	}
	return NewClient(conn)
}

// NewClient returns a new client object using an existing connection.
func NewClient(conn net.Conn) (c *Client, err error) {
	c = &Client{
		conn: conn,
		r:    bufio.NewReader(conn),
		w:    bufio.NewWriter(conn),
	}

	// Make sure we receive the server greeting
	line, err := c.ReadLine()
	if err != nil {
		return
	}
	if !IsOK(line) {
		return nil, errors.New("pop3: Server did not respond with +OK on the initial connection.")
	}
	return
}

// ReadLine reads a single line from the buffer.
func (c *Client) ReadLine() (line string, err error) {
	b, _, err := c.r.ReadLine()
	if err == io.EOF {
		return
	}
	if err != nil {
		return
	}
	line = string(b)
	return
}

// ReadLines reads from the buffer until it hits the message end dot (".").
func (c *Client) ReadLines() (lines []string, err error) {
	for {
		line, err := c.ReadLine()
		if err != nil {
			return nil, err
		}

		// Look for a dot to indicate the end of a message
		// from the server.
		if line == "." {
			break
		}
		lines = append(lines, line)
	}
	return
}

// Send writes a command to the buffer and flushes it. Does not return any
// lines from the buffer.
func (c *Client) Send(format string, args ...interface{}) (err error) {
	_, err = c.w.WriteString(fmt.Sprintf(format, args...))
	if err != nil {
		return
	}

	err = c.w.Flush()
	if err != nil {
		return
	}
	return
}

// Cmd sends a command to the server and returns a single line from the buffer.
func (c *Client) Cmd(format string, args ...interface{}) (line string, err error) {
	err = c.Send(format, args...)
	if err != nil {
		return
	}

	line, err = c.ReadLine()
	if err != nil {
		return
	}
	if !IsOK(line) {
		return "", errors.New("pop3: Server did not respond with +OK after sending a command.")
	}
	return
}

// User sends the username to the server.
func (c *Client) User(u string) (err error) {
	_, err = c.Cmd("%s %s\r\n", USER, u)
	if err != nil {
		return
	}
	return
}

// Pass sends the password to the server.
func (c *Client) Pass(p string) (err error) {
	_, err = c.Cmd("%s %s\r\n", PASS, p)
	if err != nil {
		return
	}
	return
}

// Quit sends the quit command to the server and closes the socket.
func (c *Client) Quit() (err error) {
	err = c.Send("%s\r\n", QUIT)
	if err != nil {
		return
	}
	return c.conn.Close()
}

// Auth sends the username and password to the server using the User and
// Pass methods. Noop is also called incase the server does not respond
// with invalid auth.
func (c *Client) Auth(u, p string) (err error) {
	err = c.User(u)
	if err != nil {
		return
	}

	err = c.Pass(p)
	if err != nil {
		return
	}

	// issue a dud command. Server might not respond with invalid auth
	// unless a cmd is issued
	return c.Noop()
}

// Stat retreives a listing for the current maildrop, consisting of the number
// of messages and the total size of the maildrop.
func (c *Client) Stat() (count, size int, err error) {
	line, err := c.Cmd("%s\r\n", STAT)
	if err != nil {
		return
	}

	// Number of messages in maildrop
	count, err = strconv.Atoi(strings.Fields(line)[1])
	if err != nil {
		return
	}
	if count == 0 {
		return
	}

	// Total size of messages in bytes
	size, err = strconv.Atoi(strings.Fields(line)[2])
	if err != nil {
		return
	}
	if size == 0 {
		return
	}
	return
}

// List returns the MessageList object which contains the message non unique
// id and its size.
func (c *Client) List(msg int) (list MessageList, err error) {
	line, err := c.Cmd("%s %s\r\n", LIST, msg)
	if err != nil {
		return
	}

	id, err := strconv.Atoi(strings.Fields(line)[0])
	if err != nil {
		return
	}

	size, err := strconv.Atoi(strings.Fields(line)[1])
	if err != nil {
		return
	}
	return MessageList{id, size}, nil
}

// ListAll returns a MessageList object which contains all messages in the
// maildrop.
func (c *Client) ListAll() (list []MessageList, err error) {
	_, err = c.Cmd("%s\r\n", LIST)
	if err != nil {
		return
	}

	lines, err := c.ReadLines()
	if err != nil {
		return
	}
	for _, v := range lines {
		id, err := strconv.Atoi(strings.Fields(v)[0])
		if err != nil {
			return nil, err
		}

		size, err := strconv.Atoi(strings.Fields(v)[1])
		if err != nil {
			return nil, err
		}
		list = append(list, MessageList{id, size})
	}
	return
}

// Retr downloads the given message and returns it as a mail.Message object.
func (c *Client) Retr(msg int) (m *mail.Message, err error) {
	_, err = c.Cmd("%s %s\r\n", RETR, msg)
	if err != nil {
		return
	}

	m, err = mail.ReadMessage(c.r)
	if err != nil {
		return
	}

	// mail.ReadMessage does not consume the message end dot in the buffer
	// so we must move the buffer along. Need to find a better way of
	// doing this.
	line, err := c.ReadLine()
	if err != nil {
		return
	}
	if line != "." {
		err = c.r.UnreadByte()
		if err != nil {
			return
		}
	}
	return
}

// Dele will delete the given message from the maildrop.
// Changes will only take affect after the Quit command is issued.
func (c *Client) Dele(msg int) (err error) {
	_, err = c.Cmd("%s %s\r\n", DELE, msg)
	if err != nil {
		return
	}
	return
}

// Noop will do nothing however can prolong the end of a connection.
func (c *Client) Noop() (err error) {
	_, err = c.Cmd("%s\r\n", NOOP)
	if err != nil {
		return
	}
	return
}

// Rset will unmark any messages that have being marked for deletion in
// the current session.
func (c *Client) Rset() (err error) {
	_, err = c.Cmd("%s\r\n", RSET)
	if err != nil {
		return
	}
	return
}

// Top will return a varible number of lines for a given message as a
// mail.Message object.
func (c *Client) Top(msg int, n int) (m *mail.Message, err error) {
	_, err = c.Cmd("%s %d %d\r\n", TOP, msg, n)
	if err != nil {
		return
	}

	m, err = mail.ReadMessage(c.r)
	if err != nil {
		return
	}

	// mail.ReadMessage does not consume the message end dot in the buffer
	// so we must move the buffer along. Need to find a better way of
	// doing this.
	line, err := c.ReadLine()
	if err != nil {

		return
	}
	if line != "." {
		err = c.r.UnreadByte()
		if err != nil {
			return
		}
	}
	return
}

// Uidl will return a MessageUidl object which contains the message non
// unique id and a unique id.
func (c *Client) Uidl(msg int) (list MessageUidl, err error) {
	line, err := c.Cmd("%s %s\r\n", UIDL, msg)
	if err != nil {
		return
	}

	id, err := strconv.Atoi(strings.Fields(line)[1])
	if err != nil {
		return
	}
	return MessageUidl{id, strings.Fields(line)[2]}, nil
}

// UidlAll will return a MessageUidl object which contains all messages in
// the maildrop.
func (c *Client) UidlAll() (list []MessageUidl, err error) {
	_, err = c.Cmd("%s\r\n", UIDL)
	if err != nil {
		return
	}

	lines, err := c.ReadLines()
	if err != nil {
		return
	}
	for _, v := range lines {
		id, err := strconv.Atoi(strings.Fields(v)[0])
		if err != nil {
			return nil, err
		}
		list = append(list, MessageUidl{id, strings.Fields(v)[1]})
	}
	return
}
<commit_msg>Refactor the code some more.<commit_after>// Package pop3 provides an implementation of the Post Office
// Protocol - Version 3.
package pop3

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/mail"
	"strconv"
	"strings"
)

// Dial connects to the address on the named network.
func Dial(address string) (c *Client, err error) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		return
	}
	return NewClient(conn)
}

// Dial connects to the address on the named network using tls.
func DialTLS(address string) (c *Client, err error) {
	conn, err := tls.Dial("tcp", address, nil)
	if err != nil {
		return
	}
	return NewClient(conn)
}

// NewClient returns a new client object using an existing connection.
func NewClient(conn net.Conn) (c *Client, err error) {
	c = &Client{
		conn: conn,
		r:    bufio.NewReader(conn),
		w:    bufio.NewWriter(conn),
	}

	// Make sure we receive the server greeting
	line, err := c.ReadLine()
	if err != nil {
		return
	}
	if !IsOK(line) {
		return nil, errors.New("pop3: Server did not respond with +OK on the initial connection.")
	}
	return
}

// ReadLine reads a single line from the buffer.
func (c *Client) ReadLine() (line string, err error) {
	b, _, err := c.r.ReadLine()
	if err == io.EOF {
		return
	}
	if err != nil {
		return
	}
	line = string(b)
	return
}

// ReadLines reads from the buffer until it hits the message end dot (".").
func (c *Client) ReadLines() (lines []string, err error) {
	for {
		line, err := c.ReadLine()
		if err != nil {
			return nil, err
		}

		// Look for a dot to indicate the end of a message
		// from the server.
		if line == "." {
			break
		}
		lines = append(lines, line)
	}
	return
}

// Send writes a command to the buffer and flushes it. Does not return any
// lines from the buffer.
func (c *Client) Send(format string, args ...interface{}) (err error) {
	_, err = c.w.WriteString(fmt.Sprintf(format, args...))
	if err != nil {
		return
	}
	return c.w.Flush()
}

// Cmd sends a command to the server and returns a single line from the buffer.
func (c *Client) Cmd(format string, args ...interface{}) (line string, err error) {
	if err = c.Send(format, args...); err != nil {
		return
	}

	line, err = c.ReadLine()
	if err != nil {
		return
	}
	if !IsOK(line) {
		return "", errors.New("pop3: Server did not respond with +OK after sending a command.")
	}
	return
}

// User sends the username to the server.
func (c *Client) User(u string) (err error) {
	if _, err = c.Cmd("%s %s\r\n", USER, u); err != nil {
		return
	}
	return
}

// Pass sends the password to the server.
func (c *Client) Pass(p string) (err error) {
	if _, err = c.Cmd("%s %s\r\n", PASS, p); err != nil {
		return
	}
	return
}

// Quit sends the quit command to the server and closes the socket.
func (c *Client) Quit() (err error) {
	if err = c.Send("%s\r\n", QUIT); err != nil {
		return
	}
	return c.conn.Close()
}

// Auth sends the username and password to the server using the User and
// Pass methods. Noop is also called incase the server does not respond
// with invalid auth.
func (c *Client) Auth(u, p string) (err error) {
	if err = c.User(u); err != nil {
		return
	}

	if err = c.Pass(p); err != nil {
		return
	}

	// issue a dud command. Server might not respond with invalid auth
	// unless a cmd is issued
	return c.Noop()
}

// Stat retreives a listing for the current maildrop, consisting of the number
// of messages and the total size of the maildrop.
func (c *Client) Stat() (count, size int, err error) {
	line, err := c.Cmd("%s\r\n", STAT)
	if err != nil {
		return
	}

	// Number of messages in maildrop
	count, err = strconv.Atoi(strings.Fields(line)[1])
	if err != nil {
		return
	}
	if count == 0 {
		return
	}

	// Total size of messages in bytes
	size, err = strconv.Atoi(strings.Fields(line)[2])
	if err != nil {
		return
	}
	if size == 0 {
		return
	}
	return
}

// List returns the MessageList object which contains the message non unique
// id and its size.
func (c *Client) List(msg int) (list MessageList, err error) {
	line, err := c.Cmd("%s %s\r\n", LIST, msg)
	if err != nil {
		return
	}

	id, err := strconv.Atoi(strings.Fields(line)[0])
	if err != nil {
		return
	}

	size, err := strconv.Atoi(strings.Fields(line)[1])
	if err != nil {
		return
	}
	return MessageList{id, size}, nil
}

// ListAll returns a MessageList object which contains all messages in the
// maildrop.
func (c *Client) ListAll() (list []MessageList, err error) {
	if _, err = c.Cmd("%s\r\n", LIST); err != nil {
		return
	}

	lines, err := c.ReadLines()
	if err != nil {
		return
	}
	for _, v := range lines {
		id, err := strconv.Atoi(strings.Fields(v)[0])
		if err != nil {
			return nil, err
		}

		size, err := strconv.Atoi(strings.Fields(v)[1])
		if err != nil {
			return nil, err
		}
		list = append(list, MessageList{id, size})
	}
	return
}

// Retr downloads the given message and returns it as a mail.Message object.
func (c *Client) Retr(msg int) (m *mail.Message, err error) {
	if _, err = c.Cmd("%s %s\r\n", RETR, msg); err != nil {
		return
	}

	m, err = mail.ReadMessage(c.r)
	if err != nil {
		return
	}

	// mail.ReadMessage does not consume the message end dot in the buffer
	// so we must move the buffer along. Need to find a better way of
	// doing this.
	line, err := c.ReadLine()
	if err != nil {
		return
	}
	if line != "." {
		err = c.r.UnreadByte()
		if err != nil {
			return
		}
	}
	return
}

// Dele will delete the given message from the maildrop.
// Changes will only take affect after the Quit command is issued.
func (c *Client) Dele(msg int) (err error) {
	if _, err = c.Cmd("%s %s\r\n", DELE, msg); err != nil {
		return
	}
	return
}

// Noop will do nothing however can prolong the end of a connection.
func (c *Client) Noop() (err error) {
	if _, err = c.Cmd("%s\r\n", NOOP); err != nil {
		return
	}
	return
}

// Rset will unmark any messages that have being marked for deletion in
// the current session.
func (c *Client) Rset() (err error) {
	if _, err = c.Cmd("%s\r\n", RSET); err != nil {
		return
	}
	return
}

// Top will return a varible number of lines for a given message as a
// mail.Message object.
func (c *Client) Top(msg int, n int) (m *mail.Message, err error) {
	if _, err = c.Cmd("%s %d %d\r\n", TOP, msg, n); err != nil {
		return
	}

	m, err = mail.ReadMessage(c.r)
	if err != nil {
		return
	}

	// mail.ReadMessage does not consume the message end dot in the buffer
	// so we must move the buffer along. Need to find a better way of
	// doing this.
	line, err := c.ReadLine()
	if err != nil {

		return
	}
	if line != "." {
		if err = c.r.UnreadByte(); err != nil {
			return
		}
	}
	return
}

// Uidl will return a MessageUidl object which contains the message non
// unique id and a unique id.
func (c *Client) Uidl(msg int) (list MessageUidl, err error) {
	line, err := c.Cmd("%s %s\r\n", UIDL, msg)
	if err != nil {
		return
	}

	id, err := strconv.Atoi(strings.Fields(line)[1])
	if err != nil {
		return
	}
	return MessageUidl{id, strings.Fields(line)[2]}, nil
}

// UidlAll will return a MessageUidl object which contains all messages in
// the maildrop.
func (c *Client) UidlAll() (list []MessageUidl, err error) {
	if _, err = c.Cmd("%s\r\n", UIDL); err != nil {
		return
	}

	lines, err := c.ReadLines()
	if err != nil {
		return
	}
	for _, v := range lines {
		id, err := strconv.Atoi(strings.Fields(v)[0])
		if err != nil {
			return nil, err
		}
		list = append(list, MessageUidl{id, strings.Fields(v)[1]})
	}
	return
}
<|endoftext|>


<commit_before>package modd

import (
	"bufio"
	"io"
	"os"
	"os/exec"
	"regexp"
	"sync"
	"syscall"
	"time"

	"github.com/cortesi/modd/conf"
	"github.com/cortesi/termlog"
)

// MinRestart is the minimum amount of time between daemon restarts
const MinRestart = 1 * time.Second

func getShell() string {
	sh := os.Getenv("SHELL")
	if sh == "" {
		if _, err := os.Stat("/bin/sh"); err == nil {
			sh = "/bin/sh"
		}
	}
	return sh
}

func logOutput(fp io.ReadCloser, out func(string, ...interface{})) {
	r := bufio.NewReader(fp)
	for {
		line, _, err := r.ReadLine()
		if err != nil {
			return
		}
		out(string(line))
	}
}

// RunProc runs a process to completion, sending output to log
func RunProc(cmd string, log termlog.Stream) error {
	log.Header()
	sh := getShell()
	c := exec.Command(sh, "-c", cmd)
	stdo, err := c.StdoutPipe()
	if err != nil {
		return err
	}
	stde, err := c.StderrPipe()
	if err != nil {
		return err
	}
	go logOutput(stde, log.Warn)
	go logOutput(stdo, log.Say)
	err = c.Start()
	if err != nil {
		return err
	}
	err = c.Wait()
	if err != nil {
		log.Shout("%s", c.ProcessState.String())
		return err
	}
	// FIXME: rusage stats here
	log.NoticeAs("cmdstats", "run time: %s", c.ProcessState.UserTime())
	return nil
}

// RunPreps runs all commands in sequence. Stops if any command returns an error.
func RunPreps(preps []conf.Prep, log termlog.TermLog) error {
	for _, p := range preps {
		err := RunProc(
			p.Command,
			log.Stream(p.Command),
		)
		if err != nil {
			return err
		}
	}
	return nil
}

type daemon struct {
	conf conf.Daemon
	log  termlog.Stream
	cmd  *exec.Cmd
	stop bool
}

func (d *daemon) Run() {
	d.log.Header()
	var lastStart time.Time
	for d.stop != true {
		since := time.Now().Sub(lastStart)
		if since < MinRestart {
			time.Sleep(MinRestart - since)
		}
		lastStart = time.Now()
		sh := getShell()
		c := exec.Command(sh, "-c", d.conf.Command)
		stdo, err := c.StdoutPipe()
		if err != nil {
			d.log.Shout("%s", err)
			continue
		}
		stde, err := c.StderrPipe()
		if err != nil {
			d.log.Shout("%s", err)
			continue
		}
		go logOutput(stde, d.log.Warn)
		go logOutput(stdo, d.log.Say)
		err = c.Start()
		if err != nil {
			d.log.Shout("%s", err)
			continue
		}
		d.cmd = c
		err = c.Wait()
		if err != nil {
			d.log.Shout("%s", c.ProcessState.String())
			continue
		}
	}
}

func (d *daemon) Restart() {
	if d.cmd != nil {
		d.log.Header()
		d.cmd.Process.Signal(syscall.SIGHUP)
	}
}

func (d *daemon) Shutdown(sig os.Signal) {
	d.stop = true
	if d.cmd != nil {
		d.cmd.Process.Signal(sig)
		d.cmd.Wait()
	}
}

// DaemonPen is a group of daemons, managed as a unit.
type DaemonPen struct {
	daemons *[]daemon
	sync.Mutex
}

var ws = regexp.MustCompile(`\s\s+`)

const lineLimit = 80
const postamble = "..."

// niceName tries to produce a nicer process name. We condense whitespace to
// make commands split over multiple lines with indentation more legible, and
// limit the line length to 80 characters.
func niceName(in string) string {
	in = ws.ReplaceAllString(in, " ")
	if len(in) > lineLimit-len(postamble) {
		post := termlog.DefaultPalette.Say.SprintFunc()(postamble)
		return in[:lineLimit-len(postamble)] + post
	}
	return in
}

// Start starts set of daemons, each specified by a command
func (dp *DaemonPen) Start(daemons []conf.Daemon, log termlog.TermLog) {
	dp.Lock()
	defer dp.Unlock()
	d := make([]daemon, len(daemons))
	for i, dmn := range daemons {
		d[i] = daemon{
			conf: dmn,
			log:  log.Stream(niceName(dmn.Command)),
		}
		go d[i].Run()
	}
	dp.daemons = &d
}

// Restart all daemons in the pen
func (dp *DaemonPen) Restart() {
	dp.Lock()
	defer dp.Unlock()
	if dp.daemons != nil {
		for _, d := range *dp.daemons {
			d.Restart()
		}
	}
}

// Shutdown all daemons in the pen
func (dp *DaemonPen) Shutdown(sig os.Signal) {
	dp.Lock()
	defer dp.Unlock()
	if dp.daemons != nil {
		for _, d := range *dp.daemons {
			d.Shutdown(sig)
		}
	}
}
<commit_msg>Note if a process is prep or daemon in logs<commit_after>package modd

import (
	"bufio"
	"io"
	"os"
	"os/exec"
	"regexp"
	"sync"
	"syscall"
	"time"

	"github.com/cortesi/modd/conf"
	"github.com/cortesi/termlog"
)

// MinRestart is the minimum amount of time between daemon restarts
const MinRestart = 1 * time.Second

const lineLimit = 80
const postamble = "..."

// niceHeader tries to produce a nicer process name. We condense whitespace to
// make commands split over multiple lines with indentation more legible, and
// limit the line length to 80 characters.
func niceHeader(preamble string, in string) string {
	pre := termlog.DefaultPalette.Timestamp.SprintFunc()(preamble)
	post := ""
	in = ws.ReplaceAllString(in, " ")
	if len(in) > lineLimit-len(postamble) {
		post = termlog.DefaultPalette.Say.SprintFunc()(postamble)
		in = in[:lineLimit-len(postamble)]
	}
	in = termlog.DefaultPalette.Header.SprintFunc()(in)
	return pre + in + post
}

func getShell() string {
	sh := os.Getenv("SHELL")
	if sh == "" {
		if _, err := os.Stat("/bin/sh"); err == nil {
			sh = "/bin/sh"
		}
	}
	return sh
}

func logOutput(fp io.ReadCloser, out func(string, ...interface{})) {
	r := bufio.NewReader(fp)
	for {
		line, _, err := r.ReadLine()
		if err != nil {
			return
		}
		out(string(line))
	}
}

// RunProc runs a process to completion, sending output to log
func RunProc(cmd string, log termlog.Stream) error {
	log.Header()
	sh := getShell()
	c := exec.Command(sh, "-c", cmd)
	stdo, err := c.StdoutPipe()
	if err != nil {
		return err
	}
	stde, err := c.StderrPipe()
	if err != nil {
		return err
	}
	go logOutput(stde, log.Warn)
	go logOutput(stdo, log.Say)
	err = c.Start()
	if err != nil {
		return err
	}
	err = c.Wait()
	if err != nil {
		log.Shout("%s", c.ProcessState.String())
		return err
	}
	// FIXME: rusage stats here
	log.NoticeAs("cmdstats", "run time: %s", c.ProcessState.UserTime())
	return nil
}

// RunPreps runs all commands in sequence. Stops if any command returns an error.
func RunPreps(preps []conf.Prep, log termlog.TermLog) error {
	for _, p := range preps {
		err := RunProc(
			p.Command,
			log.Stream(niceHeader("prep: ", p.Command)),
		)
		if err != nil {
			return err
		}
	}
	return nil
}

type daemon struct {
	conf conf.Daemon
	log  termlog.Stream
	cmd  *exec.Cmd
	stop bool
}

func (d *daemon) Run() {
	d.log.Header()
	var lastStart time.Time
	for d.stop != true {
		since := time.Now().Sub(lastStart)
		if since < MinRestart {
			time.Sleep(MinRestart - since)
		}
		lastStart = time.Now()
		sh := getShell()
		c := exec.Command(sh, "-c", d.conf.Command)
		stdo, err := c.StdoutPipe()
		if err != nil {
			d.log.Shout("%s", err)
			continue
		}
		stde, err := c.StderrPipe()
		if err != nil {
			d.log.Shout("%s", err)
			continue
		}
		go logOutput(stde, d.log.Warn)
		go logOutput(stdo, d.log.Say)
		err = c.Start()
		if err != nil {
			d.log.Shout("%s", err)
			continue
		}
		d.cmd = c
		err = c.Wait()
		if err != nil {
			d.log.Shout("%s", c.ProcessState.String())
			continue
		}
	}
}

func (d *daemon) Restart() {
	if d.cmd != nil {
		d.log.Header()
		d.cmd.Process.Signal(syscall.SIGHUP)
	}
}

func (d *daemon) Shutdown(sig os.Signal) {
	d.stop = true
	if d.cmd != nil {
		d.cmd.Process.Signal(sig)
		d.cmd.Wait()
	}
}

// DaemonPen is a group of daemons, managed as a unit.
type DaemonPen struct {
	daemons *[]daemon
	sync.Mutex
}

var ws = regexp.MustCompile(`\s\s+`)

// Start starts set of daemons, each specified by a command
func (dp *DaemonPen) Start(daemons []conf.Daemon, log termlog.TermLog) {
	dp.Lock()
	defer dp.Unlock()
	d := make([]daemon, len(daemons))
	for i, dmn := range daemons {
		d[i] = daemon{
			conf: dmn,
			log: log.Stream(
				niceHeader("daemon: ", dmn.Command),
			),
		}
		go d[i].Run()
	}
	dp.daemons = &d
}

// Restart all daemons in the pen
func (dp *DaemonPen) Restart() {
	dp.Lock()
	defer dp.Unlock()
	if dp.daemons != nil {
		for _, d := range *dp.daemons {
			d.Restart()
		}
	}
}

// Shutdown all daemons in the pen
func (dp *DaemonPen) Shutdown(sig os.Signal) {
	dp.Lock()
	defer dp.Unlock()
	if dp.daemons != nil {
		for _, d := range *dp.daemons {
			d.Shutdown(sig)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"sync"

	"gopkg.in/itchio/rsync-go.v0"
	"gopkg.in/kothar/brotli-go.v0/dec"

	"golang.org/x/crypto/ssh"

	"github.com/dustin/go-humanize"
	"github.com/itchio/butler/bio"
	"github.com/itchio/butler/wharf"
)

func push(src string, repoSpec string) {
	err := doPush(src, repoSpec)
	if err != nil {
		bio.Die(err.Error())
	}
}

var (
	addressRe   = regexp.MustCompile("[^:]:.*")
	defaultPort = 22
)

func doPush(src string, repoSpec string) error {
	address := *pushAddress
	if !addressRe.MatchString(address) {
		address = fmt.Sprintf("%s:%d", address, defaultPort)
	}

	conn, err := wharf.Connect(address, *pushIdentity, version)
	if err != nil {
		return err
	}
	defer conn.Close()

	closed := make(chan bool)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		for req := range conn.Reqs {
			payload, err := wharf.GetPayload(req)
			if err != nil {
				log.Printf("Error receiving payload: %s\n", err.Error())
				conn.Close()
			}

			switch v := payload.(type) {
			case bio.LogEntry:
				log.Printf("remote: %s\n", v.Message)
			case bio.EndOfSources:
				wg.Done()
			default:
				log.Printf("Server sent us unknown req %s\n", req.Type)
			}
		}
	}()

	go func() {
		for req := range conn.Chans {
			req := req

			payload, err := bio.Unmarshal(req.ExtraData())
			if err != nil {
				log.Printf("error while decoding chanreq %s payload: %s\n", req.ChannelType(), err)
				conn.Close()
				return
			}

			switch v := payload.(type) {
			case bio.SourceFile:
				wg.Add(1)
				go func() {
					defer wg.Done()
					err := handleSourceFile(src, conn, req, v)
					log.Printf("done handling sourcefile %s\n", v.Path)
					if err != nil {
						panic(err)
					}
				}()
			default:
				log.Printf("unknown channel type req'd from remote: %s\n", req.ChannelType())
				req.Reject(ssh.UnknownChannelType, "")
			}

			if err != nil {
				log.Printf("Error while handling channel :%s\n", err.Error())
				conn.Close()
				return
			}
		}
		log.Println("done handling chans, closing")
		close(closed)
	}()

	up := bio.Target{RepoSpec: repoSpec}
	ok, _, err := conn.SendRequest("butler/target", true, up)
	if err != nil {
		return fmt.Errorf("Could not send target: %s", err.Error())
	}

	if !ok {
		return fmt.Errorf("Could not find upload to replace from '%s'", repoSpec)
	}
	bio.Log("Locked onto upload target")

	done := make(chan bool)
	errs := make(chan error)

	var totalSize uint64

	fileList := make(map[string]os.FileInfo)

	filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			errs <- err
			return err
		}

		if info.IsDir() && len(path) > 1 && strings.HasPrefix(info.Name(), ".") {
			log.Printf("Skipping %s because it starts with .\n", path)
			return filepath.SkipDir
		}

		if !info.IsDir() {
			fileList[path] = info
		}
		return nil
	})

	for _, info := range fileList {
		fileSize := uint64(info.Size())
		totalSize += fileSize
	}
	log.Printf("Have %s in %d files\n", humanize.Bytes(totalSize), len(fileList))

	go func() {
		wg.Wait()
		done <- true
	}()

	select {
	case err := <-errs:
		log.Println("<-errs :(")
		return err
	case <-done:
		log.Println("<-done !")
		return nil
	case <-closed:
		log.Println("<-closed!")
		return errors.New("remote closed the connection")
	}
}

func handleSourceFile(src string, conn *wharf.Conn, req ssh.NewChannel, sf bio.SourceFile) (err error) {
	ch, reqs, err := req.Accept()
	if err != nil {
		return
	}
	go ssh.DiscardRequests(reqs)

	br := dec.NewBrotliReader(ch)
	defer br.Close()
	sig := make([]rsync.BlockHash, 0)

	gdec := gob.NewDecoder(br)
	var recipient interface{}

	for {
		err := gdec.Decode(&recipient)
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}

		switch v := recipient.(type) {
		case rsync.BlockHash:
			sig = append(sig, v)
		case bio.MD5Hash:
			_ = v.Hash
		default:
			return fmt.Errorf("wat")
		}
	}

	// create delta + send it

	path := path.Join(src, sf.Path)

	fr, err := os.Open(path)
	if err != nil {
		log.Printf("warn: could not open %s - deleted?\n", path)
		return nil
	}
	defer fr.Close()

	out, err := conn.OpenCompressedChannel("butler/patch-file", &bio.FilePatched{
		Path: path,
	})
	if err != nil {
		return err
	}

	h := md5.New()
	var newMD5 []byte

	nops := 0

	err = func() (err error) {
		defer out.CloseWrite()

		opWriter := func(op rsync.Operation) error {
			nops++
			log.Printf("sending op %d for %s\n", nops, path)
			return out.Send(op)
		}

		tr := io.TeeReader(fr, h)

		rs := &rsync.RSync{}
		err = rs.CreateDelta(tr, sig, opWriter)
		if err != nil {
			return fmt.Errorf("while creating delta for %s: %s", path, err.Error())
		}

		newMD5 = h.Sum(nil)
		out.Send(bio.MD5Hash{Hash: newMD5})
		return
	}()

	if err != nil {
		return err
	}

	log.Printf("%8s | %x | %s", humanize.Bytes(uint64(out.BytesWritten())), newMD5, path)

	// wait for read end to close
	log.Println("joining....")
	out.Join()
	log.Println("joined!")

	return
}
<commit_msg>be less chatty<commit_after>package main

import (
	"crypto/md5"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"sync"

	"gopkg.in/itchio/rsync-go.v0"
	"gopkg.in/kothar/brotli-go.v0/dec"

	"golang.org/x/crypto/ssh"

	"github.com/dustin/go-humanize"
	"github.com/itchio/butler/bio"
	"github.com/itchio/butler/wharf"
)

func push(src string, repoSpec string) {
	err := doPush(src, repoSpec)
	if err != nil {
		bio.Die(err.Error())
	}
}

var (
	addressRe   = regexp.MustCompile("[^:]:.*")
	defaultPort = 22
)

func doPush(src string, repoSpec string) error {
	address := *pushAddress
	if !addressRe.MatchString(address) {
		address = fmt.Sprintf("%s:%d", address, defaultPort)
	}

	conn, err := wharf.Connect(address, *pushIdentity, version)
	if err != nil {
		return err
	}
	defer conn.Close()

	closed := make(chan bool)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		for req := range conn.Reqs {
			payload, err := wharf.GetPayload(req)
			if err != nil {
				log.Printf("Error receiving payload: %s\n", err.Error())
				conn.Close()
			}

			switch v := payload.(type) {
			case bio.LogEntry:
				log.Printf("remote: %s\n", v.Message)
			case bio.EndOfSources:
				wg.Done()
			default:
				log.Printf("Server sent us unknown req %s\n", req.Type)
			}
		}
	}()

	go func() {
		for req := range conn.Chans {
			req := req

			payload, err := bio.Unmarshal(req.ExtraData())
			if err != nil {
				log.Printf("error while decoding chanreq %s payload: %s\n", req.ChannelType(), err)
				conn.Close()
				return
			}

			switch v := payload.(type) {
			case bio.SourceFile:
				wg.Add(1)
				go func() {
					defer wg.Done()
					err := handleSourceFile(src, conn, req, v)
					log.Printf("done handling sourcefile %s\n", v.Path)
					if err != nil {
						panic(err)
					}
				}()
			default:
				log.Printf("unknown channel type req'd from remote: %s\n", req.ChannelType())
				req.Reject(ssh.UnknownChannelType, "")
			}

			if err != nil {
				log.Printf("Error while handling channel :%s\n", err.Error())
				conn.Close()
				return
			}
		}
		log.Println("done handling chans, closing")
		close(closed)
	}()

	up := bio.Target{RepoSpec: repoSpec}
	ok, _, err := conn.SendRequest("butler/target", true, up)
	if err != nil {
		return fmt.Errorf("Could not send target: %s", err.Error())
	}

	if !ok {
		return fmt.Errorf("Could not find upload to replace from '%s'", repoSpec)
	}
	bio.Log("Locked onto upload target")

	done := make(chan bool)
	errs := make(chan error)

	var totalSize uint64

	fileList := make(map[string]os.FileInfo)

	filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			errs <- err
			return err
		}

		if info.IsDir() && len(path) > 1 && strings.HasPrefix(info.Name(), ".") {
			log.Printf("Skipping %s because it starts with .\n", path)
			return filepath.SkipDir
		}

		if !info.IsDir() {
			fileList[path] = info
		}
		return nil
	})

	for _, info := range fileList {
		fileSize := uint64(info.Size())
		totalSize += fileSize
	}
	log.Printf("Have %s in %d files\n", humanize.Bytes(totalSize), len(fileList))

	go func() {
		wg.Wait()
		done <- true
	}()

	select {
	case err := <-errs:
		log.Println("<-errs :(")
		return err
	case <-done:
		log.Println("<-done !")
		return nil
	case <-closed:
		log.Println("<-closed!")
		return errors.New("remote closed the connection")
	}
}

func handleSourceFile(src string, conn *wharf.Conn, req ssh.NewChannel, sf bio.SourceFile) (err error) {
	ch, reqs, err := req.Accept()
	if err != nil {
		return
	}
	go ssh.DiscardRequests(reqs)

	br := dec.NewBrotliReader(ch)
	defer br.Close()
	sig := make([]rsync.BlockHash, 0)

	gdec := gob.NewDecoder(br)
	var recipient interface{}

	for {
		err := gdec.Decode(&recipient)
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}

		switch v := recipient.(type) {
		case rsync.BlockHash:
			sig = append(sig, v)
		case bio.MD5Hash:
			_ = v.Hash
		default:
			return fmt.Errorf("wat")
		}
	}

	// create delta + send it

	path := path.Join(src, sf.Path)

	fr, err := os.Open(path)
	if err != nil {
		log.Printf("warn: could not open %s - deleted?\n", path)
		return nil
	}
	defer fr.Close()

	out, err := conn.OpenCompressedChannel("butler/patch-file", &bio.FilePatched{
		Path: path,
	})
	if err != nil {
		return err
	}

	h := md5.New()
	var newMD5 []byte

	nops := 0

	err = func() (err error) {
		defer out.CloseWrite()

		opWriter := func(op rsync.Operation) error {
			nops++
			log.Printf("sending op %d for %s\n", nops, path)
			return out.Send(op)
		}

		tr := io.TeeReader(fr, h)

		rs := &rsync.RSync{}
		err = rs.CreateDelta(tr, sig, opWriter)
		if err != nil {
			return fmt.Errorf("while creating delta for %s: %s", path, err.Error())
		}

		newMD5 = h.Sum(nil)
		out.Send(bio.MD5Hash{Hash: newMD5})
		return
	}()

	if err != nil {
		return err
	}

	// wait for read end to close
	out.Join()

	log.Printf("%8s | %x | %s", humanize.Bytes(uint64(out.BytesWritten())), newMD5, path)

	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"gopkg.in/yaml.v2"

	"github.com/daviddengcn/go-colortext"
	"github.com/phrase/phraseapp-go/phraseapp"
)

type PushCommand struct {
	Credentials
}

func (cmd *PushCommand) Run() error {
	if cmd.Debug {
		// suppresses content output
		cmd.Debug = false
		Debug = true
	}
	client, err := ClientFromCmdCredentials(cmd.Credentials)
	if err != nil {
		return err
	}

	sources, err := SourcesFromConfig(cmd)
	if err != nil {
		return err
	}

	for _, source := range sources {
		err := source.Push(client)
		if err != nil {
			return err
		}
	}
	return nil
}

type Sources []*Source

type Source struct {
	File          string      `yaml:"file,omitempty"`
	ProjectID     string      `yaml:"project_id,omitempty"`
	AccessToken   string      `yaml:"access_token,omitempty"`
	FileFormat    string      `yaml:"file_format,omitempty"`
	Params        *phraseapp.UploadParams `yaml:"params"`

	RemoteLocales []*phraseapp.Locale
	Extension     string
}

var separator = string(os.PathSeparator)

func (source *Source) CheckPreconditions() error {
	if err := ValidPath(source.File, source.FileFormat, ""); err != nil {
		return err
	}

	duplicatedPlaceholders := []string{}
	for _, name := range []string{"<locale_name>", "<locale_code>", "<tag>"} {
		if strings.Count(source.File, name) > 1 {
			duplicatedPlaceholders = append(duplicatedPlaceholders, name)
		}
	}

	starCount := strings.Count(source.File, "*")
	recCount := strings.Count(source.File, "**")

	if recCount == 0 && starCount > 1 || starCount-(recCount*2) > 1 {
		duplicatedPlaceholders = append(duplicatedPlaceholders, "*")
	}

	if recCount > 1 {
		duplicatedPlaceholders = append(duplicatedPlaceholders, "**")
	}

	if len(duplicatedPlaceholders) > 0 {
		dups := strings.Join(duplicatedPlaceholders, ", ")
		return fmt.Errorf(fmt.Sprintf("%s can only occur once in a file pattern!", dups))
	}

	return nil
}

func (source *Source) Push(client *phraseapp.Client) error {
	if err := source.CheckPreconditions(); err != nil {
		return err
	}

	source.Extension = filepath.Ext(source.File)

	remoteLocales, err := RemoteLocales(client, source.ProjectID)
	if err != nil {
		return err
	}
	source.RemoteLocales = remoteLocales

	localeFiles, err := source.LocaleFiles()
	if err != nil {
		return err
	}

	for _, localeFile := range localeFiles {
		fmt.Println("Uploading", localeFile.RelPath())

		if !localeFile.ExistsRemote {
			localeDetails, err := source.createLocale(client, localeFile)
			if err == nil {
				localeFile.ID = localeDetails.ID
				localeFile.RFC = localeDetails.Code
				localeFile.Name = localeDetails.Name
			} else {
				fmt.Printf("failed to create locale: %s\n", err)
				continue
			}
		}

		err = source.uploadFile(client, localeFile)
		if err != nil {
			return err
		} else {
			sharedMessage("push", localeFile)
		}

		if Debug {
			fmt.Fprintln(os.Stderr, strings.Repeat("-", 10))
		}

	}

	return nil
}

func (source *Source) createLocale(client *phraseapp.Client, localeFile *LocaleFile) (*phraseapp.LocaleDetails, error) {
	if localeFile.RFC == "" {
		return nil, fmt.Errorf("no locale code specified")
	}

	localeParams := new(phraseapp.LocaleParams)

	if localeFile.Name != "" {
		localeParams.Name = &localeFile.Name
	} else if localeFile.RFC != "" {
		localeParams.Name = &localeFile.RFC
	}

	localeName := source.replacePlaceholderInParams(localeFile)
	if localeName != "" && localeName != localeFile.RFC {
		localeParams.Name = &localeName
	}

	if localeFile.RFC != "" {
		localeParams.Code = &localeFile.RFC
	}

	localeDetails, err := client.LocaleCreate(source.ProjectID, localeParams)
	if err != nil {
		return nil, err
	}
	return localeDetails, nil
}

func (source *Source) replacePlaceholderInParams(localeFile *LocaleFile) string {
	if localeFile.RFC != "" && strings.Contains(source.GetLocaleID(), "<locale_code>") {
		return strings.Replace(source.GetLocaleID(), "<locale_code>", localeFile.RFC, 1)
	}
	return ""
}

func (source *Source) uploadFile(client *phraseapp.Client, localeFile *LocaleFile) error {
	if Debug {
		fmt.Fprintln(os.Stdout, "Source file pattern:", source.File)
		fmt.Fprintln(os.Stdout, "Actual file location:", localeFile.Path)
	}

	params := new(phraseapp.UploadParams)
	*params = *source.Params

	params.File = &localeFile.Path

	if params.LocaleID == nil {
		switch {
		case localeFile.ID != "":
			params.LocaleID = &localeFile.ID
		case localeFile.RFC != "":
			params.LocaleID = &localeFile.RFC
		}
	}

	aUpload, err := client.UploadCreate(source.ProjectID, params)
	if err != nil {
		return err
	}

	printSummary(&aUpload.Summary)

	return nil
}

func (source *Source) SystemFiles() ([]string, error) {
	filePaths := []string{}
	if strings.Contains(source.File, "**") {
		rec, err := source.recurse()
		if err != nil {
			return nil, err
		}
		filePaths = rec
	}

	globFiles, err := source.glob()
	if err != nil {
		return nil, err
	}

	for _, f := range globFiles {
		if !Contains(filePaths, f) {
			filePaths = append(filePaths, f)
		}
	}

	return filePaths, nil
}

func (source *Source) glob() ([]string, error) {
	withoutPlaceholder := placeholderRegexp.ReplaceAllString(source.File, "*")
	tokens := splitPathToTokens(withoutPlaceholder)

	fileHead := tokens[len(tokens)-1]
	if strings.HasPrefix(fileHead, ".") {
		tokens[len(tokens)-1] = "*" + fileHead
	}
	pattern := strings.Join(tokens, separator)

	files, err := filepath.Glob(pattern)
	if err != nil {
		return nil, err
	}

	if Debug {
		fmt.Fprintln(os.Stderr, "Found", len(files), "files matching the source pattern", pattern)
	}

	return files, nil
}

func (source *Source) recurse() ([]string, error) {
	files := []string{}
	err := filepath.Walk(source.root(), func(path string, f os.FileInfo, err error) error {
		if err != nil {
			errmsg := fmt.Sprintf("%s for pattern: %s", err, source.File)
			ReportError("Push Error", errmsg)
			return fmt.Errorf(errmsg)
		}
		if strings.HasSuffix(f.Name(), source.Extension) {
			files = append(files, path)
		}
		return nil
	})

	if err != nil {
		return nil, err
	}

	return files, nil
}

func (source *Source) root() string {
	parts := splitPathToTokens(source.File)
	rootParts := TakeWhile(parts, func(x string) bool {
		return x != "**"
	})
	root := strings.Join(rootParts, separator)
	if root == "" {
		root = "."
	}
	return root
}

// Return all locale files from disk that match the source pattern.
func (source *Source) LocaleFiles() (LocaleFiles, error) {
	filePaths, err := source.SystemFiles()
	if err != nil {
		return nil, err
	}

	tokens := splitPathToTokens(source.File)

	var localeFiles LocaleFiles
	for _, path := range filePaths {

		pathTokens := splitPathToTokens(path)
		if len(pathTokens) < len(tokens) {
			continue
		}
		localeFile := Reduce(tokens, pathTokens)

		absolutePath, err := filepath.Abs(path)
		if err != nil {
			return nil, err
		}
		localeFile.Path = absolutePath

		locale := source.getRemoteLocaleForLocaleFile(localeFile)
		if locale != nil {
			localeFile.ExistsRemote = true
			localeFile.RFC = locale.Code
			localeFile.Name = locale.Name
			localeFile.ID = locale.ID
		}

		if Debug {
			fmt.Println(fmt.Sprintf(
				"RFC:'%s', Name:'%s', Tag;'%s', Pattern:'%s'",
				localeFile.RFC, localeFile.Name, localeFile.Tag,
			))
		}

		localeFiles = append(localeFiles, localeFile)
	}

	if len(localeFiles) <= 0 {
		abs, err := filepath.Abs(source.File)
		if err != nil {
			return nil, err
		}
		errmsg := fmt.Sprintf("Could not find any files on your system that matches: '%s'", abs)
		ReportError("Push Error", errmsg)
		return nil, fmt.Errorf(errmsg)
	}
	return localeFiles, nil
}

func (source *Source) getRemoteLocaleForLocaleFile(localeFile *LocaleFile) *phraseapp.Locale {
	for _, remote := range source.RemoteLocales {
		if remote.Name == source.GetLocaleID() || remote.ID == source.GetLocaleID() {
			return remote
		}

		localeName := source.replacePlaceholderInParams(localeFile)
		if localeName != "" && strings.Contains(remote.Name, localeName) {
			return remote
		}

		if remote.Name == localeFile.Name {
			return remote
		}

		if remote.Name == localeFile.RFC {
			return remote
		}
	}
	return nil
}

func splitPathToTokens(s string) []string {
	tokens := []string{}
	for _, token := range strings.Split(s, separator) {
		if token == "." || token == "" {
			continue
		}
		tokens = append(tokens, token)
	}
	return tokens
}

func Reduce(tokens, pathTokens []string) *LocaleFile {
	tagged := map[string]string{}

	for idx, token := range tokens {
		pathToken := pathTokens[idx]
		if token == "*" {
			continue
		}
		if token == "**" {
			break
		}
		tagged = tag(tagged, token, pathToken)
	}

	if Contains(tokens, "**") {
		offset := 1
		for idx := len(tokens) - 1; idx >= 0; idx-- {
			token := tokens[idx]
			pathToken := pathTokens[len(pathTokens)-offset]
			offset += 1

			if token == "*" {
				continue
			}
			if token == "**" {
				break
			}

			tagged = tag(tagged, token, pathToken)
		}
	}

	return &LocaleFile{
		Name: tagged["locale_name"],
		RFC:  tagged["locale_code"],
		Tag:  tagged["tag"],
	}
}

func tag(tagged map[string]string, token, pathToken string) map[string]string {
	groups := placeholderRegexp.FindAllString(token, -1)
	if len(groups) <= 0 {
		return tagged
	}

	match := strings.Replace(token, ".", "[.]", -1)
	if strings.HasPrefix(match, "*") {
		match = strings.Replace(match, "*", ".*", -1)
	}

	for _, group := range groups {
		replacer := fmt.Sprintf("(?P%s.+)", group)
		match = strings.Replace(match, group, replacer, 1)
	}

	if match == "" {
		return tagged
	}

	tmpRegexp, err := regexp.Compile(match)
	if err != nil {
		return tagged
	}

	namedMatches := tmpRegexp.SubexpNames()
	subMatches := tmpRegexp.FindStringSubmatch(pathToken)
	for i, subMatch := range subMatches {
		if subMatch != "" {
			tagged[namedMatches[i]] = strings.Trim(subMatch, separator)
		}
	}

	return tagged
}

// Print out
func printSummary(summary *phraseapp.SummaryType) {
	newItems := []int64{
		summary.LocalesCreated,
		summary.TranslationsUpdated,
		summary.TranslationKeysCreated,
		summary.TranslationsCreated,
	}
	var changed bool
	for _, item := range newItems {
		if item > 0 {
			changed = true
		}
	}
	if changed || Debug {
		printMessage("Locales created: ", fmt.Sprintf("%d", summary.LocalesCreated))
		printMessage(" - Keys created: ", fmt.Sprintf("%d", summary.TranslationKeysCreated))
		printMessage(" - Translations created: ", fmt.Sprintf("%d", summary.TranslationsCreated))
		printMessage(" - Translations updated: ", fmt.Sprintf("%d", summary.TranslationsUpdated))
		fmt.Print("\n")
	}
}

func printMessage(msg, stat string) {
	fmt.Print(msg)
	ct.Foreground(ct.Green, true)
	fmt.Print(stat)
	ct.ResetColor()
}

// Configuration
type PushConfig struct {
	Phraseapp struct {
		AccessToken string `yaml:"access_token"`
		ProjectID   string `yaml:"project_id"`
		FileFormat  string `yaml:"file_format,omitempty"`
		Push        struct {
			Sources Sources
		}
	}
}

func SourcesFromConfig(cmd *PushCommand) (Sources, error) {
	content, err := ConfigContent()
	if err != nil {
		return nil, err
	}

	var config *PushConfig

	err = yaml.Unmarshal([]byte(content), &config)
	if err != nil {
		return nil, err
	}

	token := config.Phraseapp.AccessToken
	if cmd.Token != "" {
		token = cmd.Token
	}
	projectId := config.Phraseapp.ProjectID
	fileFormat := config.Phraseapp.FileFormat

	if &config.Phraseapp.Push == nil || config.Phraseapp.Push.Sources == nil {
		errmsg := "no sources for upload specified"
		ReportError("Push Error", errmsg)
		return nil, fmt.Errorf(errmsg)
	}

	sources := config.Phraseapp.Push.Sources

	validSources := []*Source{}
	for _, source := range sources {
		if source == nil {
			continue
		}
		if source.ProjectID == "" {
			source.ProjectID = projectId
		}
		if source.AccessToken == "" {
			source.AccessToken = token
		}
		if source.Params == nil {
			source.Params = new(phraseapp.UploadParams)
		}

		if source.Params.FileFormat == nil {
			switch {
			case source.FileFormat != "":
				source.Params.FileFormat = &source.FileFormat
			case fileFormat != "":
				source.Params.FileFormat = &fileFormat
			}
		}
		validSources = append(validSources, source)
	}

	if len(validSources) <= 0 {
		errmsg := "no sources could be identified! Refine the sources list in your config"
		ReportError("Push Error", errmsg)
		return nil, fmt.Errorf(errmsg)
	}

	return validSources, nil
}

func (source *Source) GetLocaleID() string {
	if source.Params != nil && source.Params.LocaleID != nil {
		return *source.Params.LocaleID
	}
	return ""
}
<commit_msg>fixed globbing vs recursive walk<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"gopkg.in/yaml.v2"

	"github.com/daviddengcn/go-colortext"
	"github.com/phrase/phraseapp-go/phraseapp"
)

type PushCommand struct {
	Credentials
}

func (cmd *PushCommand) Run() error {
	if cmd.Debug {
		// suppresses content output
		cmd.Debug = false
		Debug = true
	}
	client, err := ClientFromCmdCredentials(cmd.Credentials)
	if err != nil {
		return err
	}

	sources, err := SourcesFromConfig(cmd)
	if err != nil {
		return err
	}

	for _, source := range sources {
		err := source.Push(client)
		if err != nil {
			return err
		}
	}
	return nil
}

type Sources []*Source

type Source struct {
	File          string      `yaml:"file,omitempty"`
	ProjectID     string      `yaml:"project_id,omitempty"`
	AccessToken   string      `yaml:"access_token,omitempty"`
	FileFormat    string      `yaml:"file_format,omitempty"`
	Params        *phraseapp.UploadParams `yaml:"params"`

	RemoteLocales []*phraseapp.Locale
	Extension     string
}

var separator = string(os.PathSeparator)

func (source *Source) CheckPreconditions() error {
	if err := ValidPath(source.File, source.FileFormat, ""); err != nil {
		return err
	}

	duplicatedPlaceholders := []string{}
	for _, name := range []string{"<locale_name>", "<locale_code>", "<tag>"} {
		if strings.Count(source.File, name) > 1 {
			duplicatedPlaceholders = append(duplicatedPlaceholders, name)
		}
	}

	starCount := strings.Count(source.File, "*")
	recCount := strings.Count(source.File, "**")

	if recCount == 0 && starCount > 1 || starCount-(recCount*2) > 1 {
		duplicatedPlaceholders = append(duplicatedPlaceholders, "*")
	}

	if recCount > 1 {
		duplicatedPlaceholders = append(duplicatedPlaceholders, "**")
	}

	if len(duplicatedPlaceholders) > 0 {
		dups := strings.Join(duplicatedPlaceholders, ", ")
		return fmt.Errorf(fmt.Sprintf("%s can only occur once in a file pattern!", dups))
	}

	return nil
}

func (source *Source) Push(client *phraseapp.Client) error {
	if err := source.CheckPreconditions(); err != nil {
		return err
	}

	source.Extension = filepath.Ext(source.File)

	remoteLocales, err := RemoteLocales(client, source.ProjectID)
	if err != nil {
		return err
	}
	source.RemoteLocales = remoteLocales

	localeFiles, err := source.LocaleFiles()
	if err != nil {
		return err
	}

	for _, localeFile := range localeFiles {
		fmt.Println("Uploading", localeFile.RelPath())

		if !localeFile.ExistsRemote {
			localeDetails, err := source.createLocale(client, localeFile)
			if err == nil {
				localeFile.ID = localeDetails.ID
				localeFile.RFC = localeDetails.Code
				localeFile.Name = localeDetails.Name
			} else {
				fmt.Printf("failed to create locale: %s\n", err)
				continue
			}
		}

		err = source.uploadFile(client, localeFile)
		if err != nil {
			return err
		} else {
			sharedMessage("push", localeFile)
		}

		if Debug {
			fmt.Fprintln(os.Stderr, strings.Repeat("-", 10))
		}

	}

	return nil
}

func (source *Source) createLocale(client *phraseapp.Client, localeFile *LocaleFile) (*phraseapp.LocaleDetails, error) {
	if localeFile.RFC == "" {
		return nil, fmt.Errorf("no locale code specified")
	}

	localeParams := new(phraseapp.LocaleParams)

	if localeFile.Name != "" {
		localeParams.Name = &localeFile.Name
	} else if localeFile.RFC != "" {
		localeParams.Name = &localeFile.RFC
	}

	localeName := source.replacePlaceholderInParams(localeFile)
	if localeName != "" && localeName != localeFile.RFC {
		localeParams.Name = &localeName
	}

	if localeFile.RFC != "" {
		localeParams.Code = &localeFile.RFC
	}

	localeDetails, err := client.LocaleCreate(source.ProjectID, localeParams)
	if err != nil {
		return nil, err
	}
	return localeDetails, nil
}

func (source *Source) replacePlaceholderInParams(localeFile *LocaleFile) string {
	if localeFile.RFC != "" && strings.Contains(source.GetLocaleID(), "<locale_code>") {
		return strings.Replace(source.GetLocaleID(), "<locale_code>", localeFile.RFC, 1)
	}
	return ""
}

func (source *Source) uploadFile(client *phraseapp.Client, localeFile *LocaleFile) error {
	if Debug {
		fmt.Fprintln(os.Stdout, "Source file pattern:", source.File)
		fmt.Fprintln(os.Stdout, "Actual file location:", localeFile.Path)
	}

	params := new(phraseapp.UploadParams)
	*params = *source.Params

	params.File = &localeFile.Path

	if params.LocaleID == nil {
		switch {
		case localeFile.ID != "":
			params.LocaleID = &localeFile.ID
		case localeFile.RFC != "":
			params.LocaleID = &localeFile.RFC
		}
	}

	aUpload, err := client.UploadCreate(source.ProjectID, params)
	if err != nil {
		return err
	}

	printSummary(&aUpload.Summary)

	return nil
}

func (source *Source) SystemFiles() ([]string, error) {
	if strings.Contains(source.File, "**") {
		return source.recurse()
	}

	return source.glob()
}

func (source *Source) glob() ([]string, error) {
	withoutPlaceholder := placeholderRegexp.ReplaceAllString(source.File, "*")
	tokens := splitPathToTokens(withoutPlaceholder)

	fileHead := tokens[len(tokens)-1]
	if strings.HasPrefix(fileHead, ".") {
		tokens[len(tokens)-1] = "*" + fileHead
	}
	pattern := strings.Join(tokens, separator)

	files, err := filepath.Glob(pattern)
	if err != nil {
		return nil, err
	}

	if Debug {
		fmt.Fprintln(os.Stderr, "Found", len(files), "files matching the source pattern", pattern)
	}

	return files, nil
}

func (source *Source) recurse() ([]string, error) {
	files := []string{}
	err := filepath.Walk(source.root(), func(path string, f os.FileInfo, err error) error {
		if err != nil {
			errmsg := fmt.Sprintf("%s for pattern: %s", err, source.File)
			ReportError("Push Error", errmsg)
			return fmt.Errorf(errmsg)
		}
		if !f.Mode().IsDir() && strings.HasSuffix(f.Name(), source.Extension) {
			files = append(files, path)
		}
		return nil
	})

	return files, err
}

func (source *Source) root() string {
	parts := splitPathToTokens(source.File)
	rootParts := TakeWhile(parts, func(x string) bool {
		return x != "**"
	})
	root := strings.Join(rootParts, separator)
	if root == "" {
		root = "."
	}
	return root
}

// Return all locale files from disk that match the source pattern.
func (source *Source) LocaleFiles() (LocaleFiles, error) {
	filePaths, err := source.SystemFiles()
	if err != nil {
		return nil, err
	}

	tokens := splitPathToTokens(source.File)

	var localeFiles LocaleFiles
	for _, path := range filePaths {

		pathTokens := splitPathToTokens(path)
		if len(pathTokens) < len(tokens) {
			continue
		}
		localeFile := Reduce(tokens, pathTokens)

		absolutePath, err := filepath.Abs(path)
		if err != nil {
			return nil, err
		}
		localeFile.Path = absolutePath

		locale := source.getRemoteLocaleForLocaleFile(localeFile)
		if locale != nil {
			localeFile.ExistsRemote = true
			localeFile.RFC = locale.Code
			localeFile.Name = locale.Name
			localeFile.ID = locale.ID
		}

		if Debug {
			fmt.Println(fmt.Sprintf(
				"RFC:'%s', Name:'%s', Tag;'%s', Pattern:'%s'",
				localeFile.RFC, localeFile.Name, localeFile.Tag,
			))
		}

		localeFiles = append(localeFiles, localeFile)
	}

	if len(localeFiles) <= 0 {
		abs, err := filepath.Abs(source.File)
		if err != nil {
			return nil, err
		}
		errmsg := fmt.Sprintf("Could not find any files on your system that matches: '%s'", abs)
		ReportError("Push Error", errmsg)
		return nil, fmt.Errorf(errmsg)
	}
	return localeFiles, nil
}

func (source *Source) getRemoteLocaleForLocaleFile(localeFile *LocaleFile) *phraseapp.Locale {
	for _, remote := range source.RemoteLocales {
		if remote.Name == source.GetLocaleID() || remote.ID == source.GetLocaleID() {
			return remote
		}

		localeName := source.replacePlaceholderInParams(localeFile)
		if localeName != "" && strings.Contains(remote.Name, localeName) {
			return remote
		}

		if remote.Name == localeFile.Name {
			return remote
		}

		if remote.Name == localeFile.RFC {
			return remote
		}
	}
	return nil
}

func splitPathToTokens(s string) []string {
	tokens := []string{}
	for _, token := range strings.Split(s, separator) {
		if token == "." || token == "" {
			continue
		}
		tokens = append(tokens, token)
	}
	return tokens
}

func Reduce(tokens, pathTokens []string) *LocaleFile {
	tagged := map[string]string{}

	for idx, token := range tokens {
		pathToken := pathTokens[idx]
		if token == "*" {
			continue
		}
		if token == "**" {
			break
		}
		tagged = tag(tagged, token, pathToken)
	}

	if Contains(tokens, "**") {
		offset := 1
		for idx := len(tokens) - 1; idx >= 0; idx-- {
			token := tokens[idx]
			pathToken := pathTokens[len(pathTokens)-offset]
			offset += 1

			if token == "*" {
				continue
			}
			if token == "**" {
				break
			}

			tagged = tag(tagged, token, pathToken)
		}
	}

	return &LocaleFile{
		Name: tagged["locale_name"],
		RFC:  tagged["locale_code"],
		Tag:  tagged["tag"],
	}
}

func tag(tagged map[string]string, token, pathToken string) map[string]string {
	groups := placeholderRegexp.FindAllString(token, -1)
	if len(groups) <= 0 {
		return tagged
	}

	match := strings.Replace(token, ".", "[.]", -1)
	if strings.HasPrefix(match, "*") {
		match = strings.Replace(match, "*", ".*", -1)
	}

	for _, group := range groups {
		replacer := fmt.Sprintf("(?P%s.+)", group)
		match = strings.Replace(match, group, replacer, 1)
	}

	if match == "" {
		return tagged
	}

	tmpRegexp, err := regexp.Compile(match)
	if err != nil {
		return tagged
	}

	namedMatches := tmpRegexp.SubexpNames()
	subMatches := tmpRegexp.FindStringSubmatch(pathToken)
	for i, subMatch := range subMatches {
		if subMatch != "" {
			tagged[namedMatches[i]] = strings.Trim(subMatch, separator)
		}
	}

	return tagged
}

// Print out
func printSummary(summary *phraseapp.SummaryType) {
	newItems := []int64{
		summary.LocalesCreated,
		summary.TranslationsUpdated,
		summary.TranslationKeysCreated,
		summary.TranslationsCreated,
	}
	var changed bool
	for _, item := range newItems {
		if item > 0 {
			changed = true
		}
	}
	if changed || Debug {
		printMessage("Locales created: ", fmt.Sprintf("%d", summary.LocalesCreated))
		printMessage(" - Keys created: ", fmt.Sprintf("%d", summary.TranslationKeysCreated))
		printMessage(" - Translations created: ", fmt.Sprintf("%d", summary.TranslationsCreated))
		printMessage(" - Translations updated: ", fmt.Sprintf("%d", summary.TranslationsUpdated))
		fmt.Print("\n")
	}
}

func printMessage(msg, stat string) {
	fmt.Print(msg)
	ct.Foreground(ct.Green, true)
	fmt.Print(stat)
	ct.ResetColor()
}

// Configuration
type PushConfig struct {
	Phraseapp struct {
		AccessToken string `yaml:"access_token"`
		ProjectID   string `yaml:"project_id"`
		FileFormat  string `yaml:"file_format,omitempty"`
		Push        struct {
			Sources Sources
		}
	}
}

func SourcesFromConfig(cmd *PushCommand) (Sources, error) {
	content, err := ConfigContent()
	if err != nil {
		return nil, err
	}

	var config *PushConfig

	err = yaml.Unmarshal([]byte(content), &config)
	if err != nil {
		return nil, err
	}

	token := config.Phraseapp.AccessToken
	if cmd.Token != "" {
		token = cmd.Token
	}
	projectId := config.Phraseapp.ProjectID
	fileFormat := config.Phraseapp.FileFormat

	if &config.Phraseapp.Push == nil || config.Phraseapp.Push.Sources == nil {
		errmsg := "no sources for upload specified"
		ReportError("Push Error", errmsg)
		return nil, fmt.Errorf(errmsg)
	}

	sources := config.Phraseapp.Push.Sources

	validSources := []*Source{}
	for _, source := range sources {
		if source == nil {
			continue
		}
		if source.ProjectID == "" {
			source.ProjectID = projectId
		}
		if source.AccessToken == "" {
			source.AccessToken = token
		}
		if source.Params == nil {
			source.Params = new(phraseapp.UploadParams)
		}

		if source.Params.FileFormat == nil {
			switch {
			case source.FileFormat != "":
				source.Params.FileFormat = &source.FileFormat
			case fileFormat != "":
				source.Params.FileFormat = &fileFormat
			}
		}
		validSources = append(validSources, source)
	}

	if len(validSources) <= 0 {
		errmsg := "no sources could be identified! Refine the sources list in your config"
		ReportError("Push Error", errmsg)
		return nil, fmt.Errorf(errmsg)
	}

	return validSources, nil
}

func (source *Source) GetLocaleID() string {
	if source.Params != nil && source.Params.LocaleID != nil {
		return *source.Params.LocaleID
	}
	return ""
}
<|endoftext|>


<commit_before>package main

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

var cmdPush = &Command{
	Run:   runPush,
	Usage: "push path name",
	Short: "Deploy single artifact from a local directory",
	Long: `
Deploy single artifact from a local directory

Examples:

  force push connectedApps name
`,
}

var pxml = `<?xml version="1.0" encoding="UTF-8"?>
<Package xmlns="http://soap.sforce.com/2006/04/metadata">
    <types>
        <members>%s</members>
        <name>%s</name>
    </types>
    <version>29.0</version>
</Package>`

type metapath struct {
    path string
    name string
}

var metapaths = []metapath{
    metapath{"classes"      , "ApexClass"},
    metapath{"objects"      , "CustomObject"},
    metapath{"tabs"         , "CustomTab"},
    metapath{"flexipages"   , "FlexiPage"},
    metapath{"components"   , "ApexComponent"},
    metapath{"triggers"     , "ApexTrigger"},
    metapath{"pages"        , "ApexPage"},
}

func getPathForMeta(metaname string) string {
    for _, mp := range metapaths {
        if mp.name == metaname {
            return mp.path
        }
    }

    // Unknown, so use metaname
    return metaname
}

func getMetaForPath(path string) string {
    for _, mp := range metapaths {
        if mp.path == path {
            return mp.name
        }
    }

    // Unknown, so use path
    return path
}

func runPush(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.printUsage()
		return
	}

    if len(args) == 1 {
        pushByPath(args[0])
    }

	if len(args) == 2 {
        pushByName(args)
    }
}

func pushByName(args []string) {
	wd, _ := os.Getwd()
	root := filepath.Join(wd, "metadata")
	//if len(args) == 1 {
	//	root, _ = filepath.Abs(args[0])
	//}

	if _, err := os.Stat(filepath.Join(root, "package.xml")); os.IsNotExist(err) {
        root = filepath.Join(wd, "src")
        if _, err := os.Stat(filepath.Join(root, "package.xml")); os.IsNotExist(err) {
            ErrorAndExit("Must specify a directory that contains metadata files")
        }
	}

	if _, err := os.Stat(filepath.Join(root, args[0])); os.IsNotExist(err) {
		ErrorAndExit("Folder " + args[0] + " not found, must specify a metadata folder")
	}

    objType := getMetaForPath(args[0])

	found := false
	err := filepath.Walk(filepath.Join(root, args[0]), func(path string, f os.FileInfo, err error) error {
		if f.Mode().IsRegular() {
            fname := strings.ToLower(f.Name())
            fname = strings.TrimSuffix(fname, filepath.Ext(fname))
			if strings.ToLower(fname) == strings.ToLower(args[1]) {
				found = true
			}
		}
		return nil
	})
	if err != nil {
		ErrorAndExit(err.Error())
	}
	if !found {
		ErrorAndExit("Could not find " + args[1] + " in " + args[0])
	}

	files := make(ForceMetadataFiles)

	err = os.Rename(filepath.Join(root, "package.xml"), filepath.Join(root, "package.copy.xml"))

	if err := ioutil.WriteFile(filepath.Join(root, "package.xml"), []byte(fmt.Sprintf(pxml, args[1], objType)), 0644); err != nil {
		ErrorAndExit(err.Error())
	}
	if err != nil {
		ErrorAndExit(err.Error())
	}

	err = filepath.Walk(root, func(path string, f os.FileInfo, err error) error {
		if f.Mode().IsRegular() {
            fname := strings.ToLower(f.Name())
            fname = strings.TrimSuffix(fname, filepath.Ext(fname))
			if fname == strings.ToLower(args[1]) {
				data, err := ioutil.ReadFile(path)
				if err != nil {
					ErrorAndExit(err.Error())
				}
				files[strings.Replace(path, fmt.Sprintf("%s/", root), "", -1)] = data

                path += "-meta.xml"
				data, err = ioutil.ReadFile(path)
				if err != nil {
					ErrorAndExit(err.Error())
				}
				files[strings.Replace(path, fmt.Sprintf("%s/", root), "", -1)] = data
			}
            if f.Name() == "package.xml" {
				data, err := ioutil.ReadFile(path)
				if err != nil {
					ErrorAndExit(err.Error())
				}
				files[strings.Replace(path, fmt.Sprintf("%s/", root), "", -1)] = data
            }
		}
		return nil
	})
	if err != nil {
		ErrorAndExit(err.Error())
	}

    deployFiles(files)

	fmt.Printf("Pushed %s to Force.com\n", args[1])
}

// Push metadata object by path to a file
func pushByPath(fpath string) {
    fpath, err := filepath.Abs(fpath)
    if err != nil {
        ErrorAndExit("Cound not find " + fpath)
    }
    if _, err := os.Stat(fpath); err != nil {
        ErrorAndExit("Cound not open " + fpath)
    }

    hasMeta := true
    fname := filepath.Base(fpath)
    fname = strings.TrimSuffix(fname, filepath.Ext(fname))
    fdir := filepath.Dir(fpath)
    typePath := filepath.Base(fdir)
    srcDir := filepath.Dir(fdir)
    metaType := getMetaForPath(typePath)
    // Should be present since we worked back to srcDir
    frel, _ := filepath.Rel(srcDir, fpath)

    // Try to find meta file
    fmeta := fpath + "-meta.xml"
    fmetarel := ""
    if _, err := os.Stat(fmeta); err != nil {
        if os.IsNotExist(err) {
            hasMeta = false
        } else {
            ErrorAndExit("Cound not open " + fmeta)
        }
    } else {
        // Should be present since we worked back to srcDir
        fmetarel, _ = filepath.Rel(srcDir, fmeta)
    }

    // DEBUG
    // fmt.Println("Dir: " + fdir)
    // fmt.Println("fname: " + fname)
    // fmt.Println("Dir end: " + typePath)
    // fmt.Println("srcDir: " + srcDir)
    // fmt.Println("Type: " + metaType)
    // fmt.Println("relPath: " + frel)

	files := make(ForceMetadataFiles)

    fdata, err := ioutil.ReadFile(fpath)
    files[frel] = fdata
    // IF META EXISTS
    if hasMeta {
        fdata, err = ioutil.ReadFile(fmeta)
        files[fmetarel] = fdata
    }

    files["package.xml"] = []byte(fmt.Sprintf(pxml, fname, metaType))

    deployFiles(files)
}

func deployFiles(files ForceMetadataFiles) {
	force, _ := ActiveForce()
	var DeploymentOptions ForceDeployOptions
	successes, problems, err := force.Metadata.Deploy(files, DeploymentOptions)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	fmt.Printf("\nFailures - %d\n", len(problems))
	for _, problem := range problems {
		if problem.FullName == "" {
			fmt.Println(problem.Problem)
		} else {
			fmt.Printf("ERROR with %s:\n %s\n", problem.FullName, problem.Problem)
		}
	}

	fmt.Printf("\nSuccesses - %d\n", len(successes))
	for _, success := range successes {
		if success.FullName != "package.xml" {
			verb := "unchanged"
			if success.Changed {
				verb = "changed"
			} else if success.Deleted {
				verb = "deleted"
			} else if success.Created {
				verb = "created"
			}
			fmt.Printf("%s\n\tstatus: %s\n\tid=%s\n", success.FullName, verb, success.Id)
		}
	}

}
<commit_msg>Clean up pushByName<commit_after>package main

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

var cmdPush = &Command{
	Run:   runPush,
	Usage: "push path name",
	Short: "Deploy single artifact from a local directory",
	Long: `
Deploy single artifact from a local directory

Examples:

  force push connectedApps name
`,
}

var pxml = `<?xml version="1.0" encoding="UTF-8"?>
<Package xmlns="http://soap.sforce.com/2006/04/metadata">
    <types>
        <members>%s</members>
        <name>%s</name>
    </types>
    <version>29.0</version>
</Package>`

type metapath struct {
    path string
    name string
}

var metapaths = []metapath{
    metapath{"classes"      , "ApexClass"},
    metapath{"objects"      , "CustomObject"},
    metapath{"tabs"         , "CustomTab"},
    metapath{"flexipages"   , "FlexiPage"},
    metapath{"components"   , "ApexComponent"},
    metapath{"triggers"     , "ApexTrigger"},
    metapath{"pages"        , "ApexPage"},
}

func getPathForMeta(metaname string) string {
    for _, mp := range metapaths {
        if strings.ToLower(mp.name) == strings.ToLower(metaname) {
            return mp.path
        }
    }

    // Unknown, so use metaname
    return metaname
}

func getMetaForPath(path string) string {
    for _, mp := range metapaths {
        if mp.path == path {
            return mp.name
        }
    }

    // Unknown, so use path
    return path
}

func runPush(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.printUsage()
		return
	}

    if len(args) == 1 {
        fpath := args[0]
        pushByPath(fpath)

        fmt.Printf("Pushed %s to Force.com\n", fpath)
    }

	if len(args) == 2 {
        // If arg[0] is already path or meta, the method will return arg[0]
        objPath := getPathForMeta(args[0])
        objName := args[1]
        pushByName(objPath, objName)

        fmt.Printf("Pushed %s to Force.com\n", objName)
    }
}

func pushByName(objPath string, objName string) {
	wd, _ := os.Getwd()

    // First try for metadata directory
	root := filepath.Join(wd, "metadata")
	if _, err := os.Stat(filepath.Join(root, "package.xml")); os.IsNotExist(err) {
        // If not found, try for src directory
        root = filepath.Join(wd, "src")
        if _, err := os.Stat(filepath.Join(root, "package.xml")); os.IsNotExist(err) {
            ErrorAndExit("Current directory must contain a src or metadata directory")
        }
	}

	if _, err := os.Stat(filepath.Join(root, objPath)); os.IsNotExist(err) {
		ErrorAndExit("Folder " + objPath + " not found, must specify valid metadata")
	}

    // Find file by walking directory and ignoring extension
	found := false
    fpath := ""
	err := filepath.Walk(filepath.Join(root, objPath), func(path string, f os.FileInfo, err error) error {
		if f.Mode().IsRegular() {
            fname := strings.ToLower(f.Name())
            fname = strings.TrimSuffix(fname, filepath.Ext(fname))
			if strings.ToLower(fname) == strings.ToLower(objName) {
				found = true
                fpath = filepath.Join(root, objPath, f.Name())
			}
		}
		return nil
	})
	if err != nil {
		ErrorAndExit(err.Error())
	}
	if !found {
		ErrorAndExit("Could not find " + objName + " in " + objPath)
	}

    pushByPath(fpath)
}

// Push metadata object by path to a file
func pushByPath(fpath string) {
    fpath, err := filepath.Abs(fpath)
    if err != nil {
        ErrorAndExit("Cound not find " + fpath)
    }
    if _, err := os.Stat(fpath); err != nil {
        ErrorAndExit("Cound not open " + fpath)
    }

    hasMeta := true
    fname := filepath.Base(fpath)
    fname = strings.TrimSuffix(fname, filepath.Ext(fname))
    fdir := filepath.Dir(fpath)
    typePath := filepath.Base(fdir)
    srcDir := filepath.Dir(fdir)
    metaType := getMetaForPath(typePath)
    // Should be present since we worked back to srcDir
    frel, _ := filepath.Rel(srcDir, fpath)

    // Try to find meta file
    fmeta := fpath + "-meta.xml"
    fmetarel := ""
    if _, err := os.Stat(fmeta); err != nil {
        if os.IsNotExist(err) {
            hasMeta = false
        } else {
            ErrorAndExit("Cound not open " + fmeta)
        }
    } else {
        // Should be present since we worked back to srcDir
        fmetarel, _ = filepath.Rel(srcDir, fmeta)
    }

    // DEBUG
    // fmt.Println("Dir: " + fdir)
    // fmt.Println("fname: " + fname)
    // fmt.Println("Dir end: " + typePath)
    // fmt.Println("srcDir: " + srcDir)
    // fmt.Println("Type: " + metaType)
    // fmt.Println("relPath: " + frel)

	files := make(ForceMetadataFiles)

    fdata, err := ioutil.ReadFile(fpath)
    files[frel] = fdata
    if hasMeta {
        fdata, err = ioutil.ReadFile(fmeta)
        files[fmetarel] = fdata
    }

    files["package.xml"] = []byte(fmt.Sprintf(pxml, fname, metaType))

    deployFiles(files)
}

func deployFiles(files ForceMetadataFiles) {
	force, _ := ActiveForce()
	var DeploymentOptions ForceDeployOptions
	successes, problems, err := force.Metadata.Deploy(files, DeploymentOptions)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	fmt.Printf("\nFailures - %d\n", len(problems))
	for _, problem := range problems {
		if problem.FullName == "" {
			fmt.Println(problem.Problem)
		} else {
			fmt.Printf("ERROR with %s:\n %s\n", problem.FullName, problem.Problem)
		}
	}

	fmt.Printf("\nSuccesses - %d\n", len(successes))
	for _, success := range successes {
		if success.FullName != "package.xml" {
			verb := "unchanged"
			if success.Changed {
				verb = "changed"
			} else if success.Deleted {
				verb = "deleted"
			} else if success.Created {
				verb = "created"
			}
			fmt.Printf("%s\n\tstatus: %s\n\tid=%s\n", success.FullName, verb, success.Id)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/doozr/goslack"
	"github.com/doozr/qbot/command"
	"github.com/doozr/qbot/dispatch"
	"github.com/doozr/qbot/notification"
	"github.com/doozr/qbot/queue"
	"github.com/doozr/qbot/usercache"
	"github.com/doozr/qbot/util"
)

// Version is the current release version
const Version = "1.2"

func listen(name string, connection *goslack.Connection, messageChan dispatch.MessageChan, userChan dispatch.UserChan) {

	for {
		// read each incoming message
		e := <-connection.RealTime

		// see if we're mentioned
		if e.Type == "message" {
			m, err := e.RtmMessage()
			if err != nil {
				log.Println(err)
				continue
			}

			directedAtUs := strings.HasPrefix(m.Text, name) || strings.HasPrefix(m.Text, "<@"+connection.ID+">")
			if directedAtUs {
				_, m.Text = util.StringPop(m.Text)
				messageChan <- m
			} else if util.IsPrivateChannel(m.Channel) {
				messageChan <- m
			}
		}

		// see if it's a user update
		if e.Type == "user_change" {
			uc, err := e.RtmUserChange()
			if err != nil {
				log.Println(err)
				continue
			}
			userChan <- uc.User
		}
	}
}

func main() {
	if len(os.Args) < 3 {
		fmt.Println("Usage: qbot <token> <data file>")
		os.Exit(1)
	}

	log.Printf("Qbot version %s", Version)

	// Get command line parameters
	token := os.Args[1]
	filename := os.Args[2]

	// Instantiate state
	connection := connectToSlack(token)
	userCache := getUserList(connection)
	name := getBotName(userCache, connection.ID)
	q := loadQueue(filename)

	// Set up command and response processors
	notifications := notification.New(userCache)
	commands := command.New(notifications, userCache)

	// Create channels
	messageChan := make(dispatch.MessageChan, 100)
	saveChan := make(dispatch.SaveChan, 5)
	notifyChan := make(dispatch.NotifyChan, 5)
	userChan := make(dispatch.UserChan, 5)

	// Start goroutines
	go dispatch.Message(name, q, commands, messageChan, saveChan, notifyChan)
	go dispatch.Save(filename, saveChan)
	go dispatch.Notify(connection, notifyChan)
	go dispatch.User(userCache, userChan)

	// Dispatch incoming events
	log.Println("Ready to receive events")
	listen(name, connection, messageChan, userChan)
}

func connectToSlack(token string) (connection *goslack.Connection) {
	log.Print("Connecting to Slack")
	connection, err := goslack.New(token)
	if err != nil {
		log.Fatal(err)
	}
	return
}

func getUserList(connection *goslack.Connection) (userCache *usercache.UserCache) {
	log.Println("Getting user list")
	users, err := connection.GetUserList()
	if err != nil {
		log.Fatal(err)
	}
	userCache = usercache.New(users)
	return
}

func getBotName(userCache *usercache.UserCache, id string) (name string) {
	name = userCache.GetUserName(id)
	if name == "" {
		log.Fatal("Could not get username of bot")
	}
	log.Printf("Responding to requests directed to @%s", name)
	return
}

func loadQueue(filename string) (q queue.Queue) {
	log.Printf("Attempting to load queue from %s", filename)
	q, err := queue.Load(filename)
	if err != nil {
		log.Fatalf("Error loading queue: %s", err)
	}
	return
}
<commit_msg>Bump to v1.3<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/doozr/goslack"
	"github.com/doozr/qbot/command"
	"github.com/doozr/qbot/dispatch"
	"github.com/doozr/qbot/notification"
	"github.com/doozr/qbot/queue"
	"github.com/doozr/qbot/usercache"
	"github.com/doozr/qbot/util"
)

// Version is the current release version
const Version = "1.3"

func listen(name string, connection *goslack.Connection, messageChan dispatch.MessageChan, userChan dispatch.UserChan) {

	for {
		// read each incoming message
		e := <-connection.RealTime

		// see if we're mentioned
		if e.Type == "message" {
			m, err := e.RtmMessage()
			if err != nil {
				log.Println(err)
				continue
			}

			directedAtUs := strings.HasPrefix(m.Text, name) || strings.HasPrefix(m.Text, "<@"+connection.ID+">")
			if directedAtUs {
				_, m.Text = util.StringPop(m.Text)
				messageChan <- m
			} else if util.IsPrivateChannel(m.Channel) {
				messageChan <- m
			}
		}

		// see if it's a user update
		if e.Type == "user_change" {
			uc, err := e.RtmUserChange()
			if err != nil {
				log.Println(err)
				continue
			}
			userChan <- uc.User
		}
	}
}

func main() {
	if len(os.Args) < 3 {
		fmt.Println("Usage: qbot <token> <data file>")
		os.Exit(1)
	}

	log.Printf("Qbot version %s", Version)

	// Get command line parameters
	token := os.Args[1]
	filename := os.Args[2]

	// Instantiate state
	connection := connectToSlack(token)
	userCache := getUserList(connection)
	name := getBotName(userCache, connection.ID)
	q := loadQueue(filename)

	// Set up command and response processors
	notifications := notification.New(userCache)
	commands := command.New(notifications, userCache)

	// Create channels
	messageChan := make(dispatch.MessageChan, 100)
	saveChan := make(dispatch.SaveChan, 5)
	notifyChan := make(dispatch.NotifyChan, 5)
	userChan := make(dispatch.UserChan, 5)

	// Start goroutines
	go dispatch.Message(name, q, commands, messageChan, saveChan, notifyChan)
	go dispatch.Save(filename, saveChan)
	go dispatch.Notify(connection, notifyChan)
	go dispatch.User(userCache, userChan)

	// Dispatch incoming events
	log.Println("Ready to receive events")
	listen(name, connection, messageChan, userChan)
}

func connectToSlack(token string) (connection *goslack.Connection) {
	log.Print("Connecting to Slack")
	connection, err := goslack.New(token)
	if err != nil {
		log.Fatal(err)
	}
	return
}

func getUserList(connection *goslack.Connection) (userCache *usercache.UserCache) {
	log.Println("Getting user list")
	users, err := connection.GetUserList()
	if err != nil {
		log.Fatal(err)
	}
	userCache = usercache.New(users)
	return
}

func getBotName(userCache *usercache.UserCache, id string) (name string) {
	name = userCache.GetUserName(id)
	if name == "" {
		log.Fatal("Could not get username of bot")
	}
	log.Printf("Responding to requests directed to @%s", name)
	return
}

func loadQueue(filename string) (q queue.Queue) {
	log.Printf("Attempting to load queue from %s", filename)
	q, err := queue.Load(filename)
	if err != nil {
		log.Fatalf("Error loading queue: %s", err)
	}
	return
}
<|endoftext|>


<commit_before>package rbac

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

type Rbac interface {
	Assign(Role, Permission) (int64, error)
	Unassign(Role, Permission) error
	Check(permission Permission, userId int64) (bool, error)
	Reset(ensure bool)

	Permissions() PermissionManager
	Roles() RoleManager
	Users() UserManager
}

type Config struct {
	Name     string
	Host     string
	Port     int
	Username string
	Password string
}

type rbac struct {
	permissions PermissionManager
	roles       RoleManager
	users       UserManager

	db *sql.DB
}

func New(config *Config) Rbac {
	var rbac = new(rbac)

	rbac.roles = NewRoleManager(rbac)
	rbac.permissions = NewPermissionManager(rbac)
	rbac.users = NewUserManager(rbac)

	var err error
	rbac.db, err = sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?parseTime=true", config.Username, config.Password, config.Host, config.Port, config.Name))
	if err != nil {
		log.Fatal(err)
	}

	return rbac
}

func (r *rbac) Assign(role Role, permission Permission) (int64, error) {
	var err error
	var roleId int64
	var permissionId int64

	roleId, err = r.Roles().GetRoleId(role)
	if err != nil {
		return 0, err
	}

	permissionId, err = r.permissions.GetPermissionId(permission)
	if err != nil {
		return 0, err
	}

	res, err := r.db.Exec("INSERT INTO rolepermissions (RoleID, PermissionID, AssignmentDate) VALUES(?,?,?)", roleId, permissionId, time.Now().Nanosecond())
	if err != nil {
		return 0, err
	}

	insertId, _ := res.LastInsertId()

	return insertId, nil
}

func (r *rbac) Unassign(role Role, permission Permission) error {
	var err error
	var roleId int64
	var permissionId int64

	roleId, err = r.Roles().GetRoleId(role)
	if err != nil {
		return err
	}

	permissionId, err = r.permissions.GetPermissionId(permission)
	if err != nil {
		return err
	}

	_, err = r.db.Exec("DELETE FROM rolepermissions WHERE RoleId=? AND PermissionID=?", roleId, permissionId)

	if err != nil {
		return err
	}

	return nil
}

func (r rbac) Check(permission Permission, userId int64) (bool, error) {
	if userId == 0 {
		return false, fmt.Errorf("userId cannot be null")
	}

	permissionId, err := r.permissions.GetPermissionId(permission)
	if err != nil {
		return false, err
	}

	if permissionId == 0 {
		return false, fmt.Errorf("permission not found")
	}

	lastPart := `
	ON ( TR.ID = TRel.RoleID)
 							WHERE
 							TUrel.UserID=?
 							AND
 							TPdirect.ID=?
	`
	query := fmt.Sprintf(`SELECT COUNT(*) AS Result
	FROM
		userroles AS TUrel
	JOIN roles AS TRdirect ON (TRdirect.ID=TUrel.RoleID)
	JOIN roles AS TR ON ( TR.Lft BETWEEN TRdirect.Lft AND TRdirect.Rght)
	JOIN
		(permissions AS TPdirect
			JOIN permissions AS TP ON (TPdirect.Lft BETWEEN TP.Lft AND TP.Rght)
			JOIN rolepermissions AS TRel ON (TP.ID=TRel.PermissionID)
		) %s`, lastPart)

	var result int64
	err = r.db.QueryRow(query, userId, permissionId).Scan(&result)
	if err != nil {
		if err != sql.ErrNoRows {
			return false, err
		}
	}

	if result > 0 {
		return true, nil
	}

	return false, nil
}

func (r rbac) Reset(ensure bool) {
	if err := r.roles.ResetAssignments(ensure); err != nil {
		log.Fatal(err)
	}
	if err := r.roles.Reset(ensure); err != nil {
		log.Fatal(err)
	}

	if err := r.permissions.Reset(ensure); err != nil {
		log.Fatal(err)
	}

	if err := r.users.ResetAssignments(ensure); err != nil {
		log.Fatal(err)
	}
}

func (r *rbac) rootId() int64 {
	return 1
}

func (r rbac) Permissions() PermissionManager {
	return r.permissions
}

func (r rbac) Roles() RoleManager {
	return r.roles
}

func (r rbac) Users() UserManager {
	return r.users
}
<commit_msg> Changes to be committed: 	modified:   rbac.go<commit_after>package rbac

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

type Rbac interface {
	Assign(Role, Permission) (int64, error)
	Unassign(Role, Permission) error
	Check(permission Permission, userId int64) (bool, error)
	Reset(ensure bool)

	Permissions() PermissionManager
	Roles() RoleManager
	Users() UserManager
}

type Config struct {
	Name     string
	Host     string
	Port     int
	Username string
	Password string
}

type rbac struct {
	permissions PermissionManager
	roles       RoleManager
	users       UserManager

	db *sql.DB
}

func New(config *Config) Rbac {
	var rbac = new(rbac)

	rbac.roles = NewRoleManager(rbac)
	rbac.permissions = NewPermissionManager(rbac)
	rbac.users = NewUserManager(rbac)

	var err error
	rbac.db, err = sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?parseTime=true", config.Username, config.Password, config.Host, config.Port, config.Name))
	if err != nil {
		log.Fatal(err)
	}

	return rbac
}

/**
* Assigns a role to a permission (or vice-verse)
*
* @param mixed $Role
*         Id, Title and Path
* @param mixed $Permission
*         Id, Title and Path
* @return boolean inserted or existing
*
* @todo: Check for valid permissions/roles
* @todo: Implement custom error handler
 */
func (r *rbac) Assign(role Role, permission Permission) (int64, error) {
	var err error
	var roleId int64
	var permissionId int64

	roleId, err = r.Roles().GetRoleId(role)
	if err != nil {
		return 0, err
	}

	permissionId, err = r.permissions.GetPermissionId(permission)
	if err != nil {
		return 0, err
	}

	res, err := r.db.Exec("INSERT INTO rolepermissions (RoleID, PermissionID, AssignmentDate) VALUES(?,?,?)", roleId, permissionId, time.Now().Nanosecond())
	if err != nil {
		return 0, err
	}

	insertId, _ := res.LastInsertId()

	return insertId, nil
}

func (r *rbac) Unassign(role Role, permission Permission) error {
	var err error
	var roleId int64
	var permissionId int64

	roleId, err = r.Roles().GetRoleId(role)
	if err != nil {
		return err
	}

	permissionId, err = r.permissions.GetPermissionId(permission)
	if err != nil {
		return err
	}

	_, err = r.db.Exec("DELETE FROM rolepermissions WHERE RoleId=? AND PermissionID=?", roleId, permissionId)

	if err != nil {
		return err
	}

	return nil
}

func (r rbac) Check(permission Permission, userId int64) (bool, error) {
	if userId == 0 {
		return false, fmt.Errorf("userId cannot be null")
	}

	permissionId, err := r.permissions.GetPermissionId(permission)
	if err != nil {
		return false, err
	}

	if permissionId == 0 {
		return false, fmt.Errorf("permission not found")
	}

	lastPart := `
	ON ( TR.ID = TRel.RoleID)
 							WHERE
 							TUrel.UserID=?
 							AND
 							TPdirect.ID=?
	`
	query := fmt.Sprintf(`SELECT COUNT(*) AS Result
	FROM
		userroles AS TUrel
	JOIN roles AS TRdirect ON (TRdirect.ID=TUrel.RoleID)
	JOIN roles AS TR ON ( TR.Lft BETWEEN TRdirect.Lft AND TRdirect.Rght)
	JOIN
		(permissions AS TPdirect
			JOIN permissions AS TP ON (TPdirect.Lft BETWEEN TP.Lft AND TP.Rght)
			JOIN rolepermissions AS TRel ON (TP.ID=TRel.PermissionID)
		) %s`, lastPart)

	var result int64
	err = r.db.QueryRow(query, userId, permissionId).Scan(&result)
	if err != nil {
		if err != sql.ErrNoRows {
			return false, err
		}
	}

	if result > 0 {
		return true, nil
	}

	return false, nil
}

func (r rbac) Reset(ensure bool) {
	if err := r.roles.ResetAssignments(ensure); err != nil {
		log.Fatal(err)
	}
	if err := r.roles.Reset(ensure); err != nil {
		log.Fatal(err)
	}

	if err := r.permissions.Reset(ensure); err != nil {
		log.Fatal(err)
	}

	if err := r.users.ResetAssignments(ensure); err != nil {
		log.Fatal(err)
	}
}

func (r *rbac) rootId() int64 {
	return 1
}

func (r rbac) Permissions() PermissionManager {
	return r.permissions
}

func (r rbac) Roles() RoleManager {
	return r.roles
}

func (r rbac) Users() UserManager {
	return r.users
}
<|endoftext|>


<commit_before>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2015 Douglas Hall
//
// 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 main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/bemasher/rtlamr/idm"
	"github.com/bemasher/rtlamr/parse"
	"github.com/bemasher/rtlamr/r900"
	"github.com/bemasher/rtlamr/scm"
	"github.com/bemasher/rtltcp"
)

var rcvr Receiver

type Receiver struct {
	rtltcp.SDR
	p parse.Parser
}

func (rcvr *Receiver) NewReceiver() {
	switch strings.ToLower(*msgType) {
	case "scm":
		rcvr.p = scm.NewParser(*symbolLength, *decimation, *fastMag)
	case "idm":
		rcvr.p = idm.NewParser(*symbolLength, *decimation, *fastMag)
	case "r900":
		rcvr.p = r900.NewParser(*symbolLength, *decimation, *fastMag)
	default:
		log.Fatalf("Invalid message type: %q\n", *msgType)
	}

	if !*quiet {
		rcvr.p.Cfg().Log()
	}

	// Connect to rtl_tcp server.
	if err := rcvr.Connect(nil); err != nil {
		log.Fatal(err)
	}

	rcvr.HandleFlags()

	// Tell the user how many gain settings were reported by rtl_tcp.
	if !*quiet {
		log.Println("GainCount:", rcvr.SDR.Info.GainCount)
	}

	centerfreqFlagSet := false
	sampleRateFlagSet := false
	gainFlagSet := false
	flag.Visit(func(f *flag.Flag) {
		switch f.Name {
		case "centerfreq":
			centerfreqFlagSet = true
		case "samplerate":
			sampleRateFlagSet = true
		case "gainbyindex", "tunergainmode", "tunergain", "agcmode":
			gainFlagSet = true
		}
	})

	// Set some parameters for listening.
	if centerfreqFlagSet {
		rcvr.SetCenterFreq(uint32(rcvr.Flags.CenterFreq))
	} else {
		rcvr.SetCenterFreq(rcvr.p.Cfg().CenterFreq)
	}

	if !sampleRateFlagSet {
		rcvr.SetSampleRate(uint32(rcvr.p.Cfg().SampleRate))
	}
	if !gainFlagSet {
		rcvr.SetGainMode(true)
	}

	return
}

func (rcvr *Receiver) Run() {
	// Setup signal channel for interruption.
	sigint := make(chan os.Signal, 1)
	signal.Notify(sigint, os.Kill, os.Interrupt)

	// Setup time limit channel
	tLimit := make(<-chan time.Time, 1)
	if *timeLimit != 0 {
		tLimit = time.After(*timeLimit)
	}

	in, out := io.Pipe()

	go func() {
		tcpBlock := make([]byte, 16384)
		for {
			n, err := rcvr.Read(tcpBlock)
			if err != nil {
				return
			}
			out.Write(tcpBlock[:n])
		}
	}()

	block := make([]byte, rcvr.p.Cfg().BlockSize2)

	start := time.Now()
	for {
		// Exit on interrupt or time limit, otherwise receive.
		select {
		case <-sigint:
			return
		case <-tLimit:
			fmt.Println("Time Limit Reached:", time.Since(start))
			return
		default:
			// Read new sample block.
			_, err := in.Read(block)
			if err != nil {
				log.Fatal("Error reading samples: ", err)
			}

			pktFound := false
			indices := rcvr.p.Dec().Decode(block)

			for _, pkt := range rcvr.p.Parse(indices) {
				if len(meterID) > 0 && !meterID[uint(pkt.MeterID())] {
					continue
				}

				if len(meterType) > 0 && !meterType[uint(pkt.MeterType())] {
					continue
				}

				var msg parse.LogMessage
				msg.Time = time.Now()
				msg.Offset, _ = sampleFile.Seek(0, os.SEEK_CUR)
				msg.Length = rcvr.p.Cfg().BufferLength << 1
				msg.Message = pkt

				if encoder == nil {
					// A nil encoder is just plain-text output.
					if *sampleFilename == os.DevNull {
						fmt.Fprintln(logFile, msg.StringNoOffset())
					} else {
						fmt.Fprintln(logFile, msg)
					}
				} else {
					err = encoder.Encode(msg)
					if err != nil {
						log.Fatal("Error encoding message: ", err)
					}

					// The XML encoder doesn't write new lines after each
					// element, add them.
					if _, ok := encoder.(*xml.Encoder); ok {
						fmt.Fprintln(logFile)
					}
				}

				pktFound = true
				if *single {
					break
				}
			}

			if pktFound {
				if *sampleFilename != os.DevNull {
					_, err = sampleFile.Write(rcvr.p.Dec().IQ)
					if err != nil {
						log.Fatal("Error writing raw samples to file:", err)
					}
				}
				if *single {
					return
				}
			}
		}
	}
}

func init() {
	log.SetFlags(log.Lshortfile | log.Lmicroseconds)
}

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to this file")

func main() {
	rcvr.RegisterFlags()
	RegisterFlags()

	flag.Parse()
	HandleFlags()

	rcvr.NewReceiver()

	defer logFile.Close()
	defer sampleFile.Close()
	defer rcvr.Close()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	rcvr.Run()
}
<commit_msg>Fix block read, reads might span writes to the pipe.<commit_after>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2015 Douglas Hall
//
// 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 main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/bemasher/rtlamr/idm"
	"github.com/bemasher/rtlamr/parse"
	"github.com/bemasher/rtlamr/r900"
	"github.com/bemasher/rtlamr/scm"
	"github.com/bemasher/rtltcp"
)

var rcvr Receiver

type Receiver struct {
	rtltcp.SDR
	p parse.Parser
}

func (rcvr *Receiver) NewReceiver() {
	switch strings.ToLower(*msgType) {
	case "scm":
		rcvr.p = scm.NewParser(*symbolLength, *decimation, *fastMag)
	case "idm":
		rcvr.p = idm.NewParser(*symbolLength, *decimation, *fastMag)
	case "r900":
		rcvr.p = r900.NewParser(*symbolLength, *decimation, *fastMag)
	default:
		log.Fatalf("Invalid message type: %q\n", *msgType)
	}

	if !*quiet {
		rcvr.p.Cfg().Log()
	}

	// Connect to rtl_tcp server.
	if err := rcvr.Connect(nil); err != nil {
		log.Fatal(err)
	}

	rcvr.HandleFlags()

	// Tell the user how many gain settings were reported by rtl_tcp.
	if !*quiet {
		log.Println("GainCount:", rcvr.SDR.Info.GainCount)
	}

	centerfreqFlagSet := false
	sampleRateFlagSet := false
	gainFlagSet := false
	flag.Visit(func(f *flag.Flag) {
		switch f.Name {
		case "centerfreq":
			centerfreqFlagSet = true
		case "samplerate":
			sampleRateFlagSet = true
		case "gainbyindex", "tunergainmode", "tunergain", "agcmode":
			gainFlagSet = true
		}
	})

	// Set some parameters for listening.
	if centerfreqFlagSet {
		rcvr.SetCenterFreq(uint32(rcvr.Flags.CenterFreq))
	} else {
		rcvr.SetCenterFreq(rcvr.p.Cfg().CenterFreq)
	}

	if !sampleRateFlagSet {
		rcvr.SetSampleRate(uint32(rcvr.p.Cfg().SampleRate))
	}
	if !gainFlagSet {
		rcvr.SetGainMode(true)
	}

	return
}

func (rcvr *Receiver) Run() {
	// Setup signal channel for interruption.
	sigint := make(chan os.Signal, 1)
	signal.Notify(sigint, os.Kill, os.Interrupt)

	// Setup time limit channel
	tLimit := make(<-chan time.Time, 1)
	if *timeLimit != 0 {
		tLimit = time.After(*timeLimit)
	}

	in, out := io.Pipe()

	go func() {
		tcpBlock := make([]byte, 16384)
		for {
			n, err := rcvr.Read(tcpBlock)
			if err != nil {
				return
			}
			out.Write(tcpBlock[:n])
		}
	}()

	block := make([]byte, rcvr.p.Cfg().BlockSize2)

	start := time.Now()
	for {
		// Exit on interrupt or time limit, otherwise receive.
		select {
		case <-sigint:
			return
		case <-tLimit:
			fmt.Println("Time Limit Reached:", time.Since(start))
			return
		default:
			// Read new sample block.
			_, err := io.ReadFull(in, block)
			if err != nil {
				log.Fatal("Error reading samples: ", err)
			}

			pktFound := false
			indices := rcvr.p.Dec().Decode(block)

			for _, pkt := range rcvr.p.Parse(indices) {
				if len(meterID) > 0 && !meterID[uint(pkt.MeterID())] {
					continue
				}

				if len(meterType) > 0 && !meterType[uint(pkt.MeterType())] {
					continue
				}

				var msg parse.LogMessage
				msg.Time = time.Now()
				msg.Offset, _ = sampleFile.Seek(0, os.SEEK_CUR)
				msg.Length = rcvr.p.Cfg().BufferLength << 1
				msg.Message = pkt

				if encoder == nil {
					// A nil encoder is just plain-text output.
					if *sampleFilename == os.DevNull {
						fmt.Fprintln(logFile, msg.StringNoOffset())
					} else {
						fmt.Fprintln(logFile, msg)
					}
				} else {
					err = encoder.Encode(msg)
					if err != nil {
						log.Fatal("Error encoding message: ", err)
					}

					// The XML encoder doesn't write new lines after each
					// element, add them.
					if _, ok := encoder.(*xml.Encoder); ok {
						fmt.Fprintln(logFile)
					}
				}

				pktFound = true
				if *single {
					break
				}
			}

			if pktFound {
				if *sampleFilename != os.DevNull {
					_, err = sampleFile.Write(rcvr.p.Dec().IQ)
					if err != nil {
						log.Fatal("Error writing raw samples to file:", err)
					}
				}
				if *single {
					return
				}
			}
		}
	}
}

func init() {
	log.SetFlags(log.Lshortfile | log.Lmicroseconds)
}

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to this file")

func main() {
	rcvr.RegisterFlags()
	RegisterFlags()

	flag.Parse()
	HandleFlags()

	rcvr.NewReceiver()

	defer logFile.Close()
	defer sampleFile.Close()
	defer rcvr.Close()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	rcvr.Run()
}
<|endoftext|>