<commit_before>package ice

import (
	"io"
	"net"
	"os"
	"strings"
	"sync"

	"github.com/pion/logging"
	"github.com/pion/stun"
)

// UDPMux allows multiple connections to go over a single UDP port
type UDPMux interface {
	io.Closer
	GetConn(ufrag string) (net.PacketConn, error)
	RemoveConnByUfrag(ufrag string)
}

// UDPMuxDefault is an implementation of the interface
type UDPMuxDefault struct {
	params UDPMuxParams

	closedChan chan struct{}
	closeOnce  sync.Once

	// conns is a map of all udpMuxedConn indexed by ufrag|network|candidateType
	conns map[string]*udpMuxedConn

	addressMapMu sync.RWMutex
	addressMap   map[string]*udpMuxedConn

	// buffer pool to recycle buffers for net.UDPAddr encodes/decodes
	pool *sync.Pool

	mu sync.Mutex
}

const maxAddrSize = 512

// UDPMuxParams are parameters for UDPMux.
type UDPMuxParams struct {
	Logger  logging.LeveledLogger
	UDPConn *net.UDPConn
}

// NewUDPMuxDefault creates an implementation of UDPMux
func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault {
	if params.Logger == nil {
		params.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
	}

	m := &UDPMuxDefault{
		addressMap: map[string]*udpMuxedConn{},
		params:     params,
		conns:      make(map[string]*udpMuxedConn),
		closedChan: make(chan struct{}, 1),
		pool: &sync.Pool{
			New: func() interface{} {
				// big enough buffer to fit both packet and address
				return newBufferHolder(receiveMTU + maxAddrSize)
			},
		},
	}

	go m.connWorker()

	return m
}

// LocalAddr returns the listening address of this UDPMuxDefault
func (m *UDPMuxDefault) LocalAddr() net.Addr {
	return m.params.UDPConn.LocalAddr()
}

// GetConn returns a PacketConn given the connection's ufrag and network
// creates the connection if an existing one can't be found
func (m *UDPMuxDefault) GetConn(ufrag string) (net.PacketConn, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.IsClosed() {
		return nil, io.ErrClosedPipe
	}

	if c, ok := m.conns[ufrag]; ok {
		return c, nil
	}

	c := m.createMuxedConn(ufrag)
	go func() {
		<-c.CloseChannel()
		m.removeConn(ufrag)
	}()
	m.conns[ufrag] = c
	return c, nil
}

// RemoveConnByUfrag stops and removes the muxed packet connection
func (m *UDPMuxDefault) RemoveConnByUfrag(ufrag string) {
	m.mu.Lock()
	removedConns := make([]*udpMuxedConn, 0)
	for key := range m.conns {
		if key != ufrag {
			continue
		}

		c := m.conns[key]
		delete(m.conns, key)
		if c != nil {
			removedConns = append(removedConns, c)
		}
	}
	// keep lock section small to avoid deadlock with conn lock
	m.mu.Unlock()

	m.addressMapMu.Lock()
	defer m.addressMapMu.Unlock()

	for _, c := range removedConns {
		addresses := c.getAddresses()
		for _, addr := range addresses {
			delete(m.addressMap, addr)
		}
	}
}

// IsClosed returns true if the mux had been closed
func (m *UDPMuxDefault) IsClosed() bool {
	select {
	case <-m.closedChan:
		return true
	default:
		return false
	}
}

// Close the mux, no further connections could be created
func (m *UDPMuxDefault) Close() error {
	var err error
	m.closeOnce.Do(func() {
		m.mu.Lock()
		defer m.mu.Unlock()

		for _, c := range m.conns {
			_ = c.Close()
		}
		m.conns = make(map[string]*udpMuxedConn)
		close(m.closedChan)
	})
	return err
}

func (m *UDPMuxDefault) removeConn(key string) {
	m.mu.Lock()
	c := m.conns[key]
	delete(m.conns, key)
	// keep lock section small to avoid deadlock with conn lock
	m.mu.Unlock()

	if c == nil {
		return
	}

	m.addressMapMu.Lock()
	defer m.addressMapMu.Unlock()

	addresses := c.getAddresses()
	for _, addr := range addresses {
		delete(m.addressMap, addr)
	}
}

func (m *UDPMuxDefault) writeTo(buf []byte, raddr net.Addr) (n int, err error) {
	return m.params.UDPConn.WriteTo(buf, raddr)
}

func (m *UDPMuxDefault) registerConnForAddress(conn *udpMuxedConn, addr string) {
	if m.IsClosed() {
		return
	}

	m.addressMapMu.Lock()
	defer m.addressMapMu.Unlock()

	existing, ok := m.addressMap[addr]
	if ok {
		existing.removeAddress(addr)
	}
	m.addressMap[addr] = conn

	m.params.Logger.Debugf("Registered %s for %s", addr, conn.params.Key)
}

func (m *UDPMuxDefault) createMuxedConn(key string) *udpMuxedConn {
	c := newUDPMuxedConn(&udpMuxedConnParams{
		Mux:       m,
		Key:       key,
		AddrPool:  m.pool,
		LocalAddr: m.LocalAddr(),
		Logger:    m.params.Logger,
	})
	return c
}

func (m *UDPMuxDefault) connWorker() {
	logger := m.params.Logger

	defer func() {
		_ = m.Close()
	}()

	buf := make([]byte, receiveMTU)
	for {
		n, addr, err := m.params.UDPConn.ReadFrom(buf)
		if m.IsClosed() {
			return
		} else if err != nil {
			if os.IsTimeout(err) {
				continue
			} else if err != io.EOF {
				logger.Errorf("could not read udp packet: %v", err)
			}

			return
		}

		udpAddr, ok := addr.(*net.UDPAddr)
		if !ok {
			logger.Errorf("underlying PacketConn did not return a UDPAddr")
			return
		}

		// If we have already seen this address dispatch to the appropriate destination
		m.addressMapMu.Lock()
		destinationConn := m.addressMap[addr.String()]
		m.addressMapMu.Unlock()

		// If we haven't seen this address before but is a STUN packet lookup by ufrag
		if destinationConn == nil && stun.IsMessage(buf[:n]) {
			msg := &stun.Message{
				Raw: append([]byte{}, buf[:n]...),
			}

			if err = msg.Decode(); err != nil {
				m.params.Logger.Warnf("Failed to handle decode ICE from %s: %v\n", addr.String(), err)
				continue
			}

			attr, stunAttrErr := msg.Get(stun.AttrUsername)
			if stunAttrErr != nil {
				m.params.Logger.Warnf("No Username attribute in STUN message from %s\n", addr.String())
				continue
			}

			ufrag := strings.Split(string(attr), ":")[0]

			m.mu.Lock()
			destinationConn = m.conns[ufrag]
			m.mu.Unlock()
		}

		if destinationConn == nil {
			m.params.Logger.Tracef("dropping packet from %s, addr: %s", udpAddr.String(), addr.String())
			continue
		}

		if err = destinationConn.writePacket(buf[:n], udpAddr); err != nil {
			m.params.Logger.Errorf("could not write packet: %v", err)
		}
	}
}

type bufferHolder struct {
	buffer []byte
}

func newBufferHolder(size int) *bufferHolder {
	return &bufferHolder{
		buffer: make([]byte, size),
	}
}
<commit_msg>UDPMuxParams use net.PacketConn instead of UDPConn<commit_after>package ice

import (
	"io"
	"net"
	"os"
	"strings"
	"sync"

	"github.com/pion/logging"
	"github.com/pion/stun"
)

// UDPMux allows multiple connections to go over a single UDP port
type UDPMux interface {
	io.Closer
	GetConn(ufrag string) (net.PacketConn, error)
	RemoveConnByUfrag(ufrag string)
}

// UDPMuxDefault is an implementation of the interface
type UDPMuxDefault struct {
	params UDPMuxParams

	closedChan chan struct{}
	closeOnce  sync.Once

	// conns is a map of all udpMuxedConn indexed by ufrag|network|candidateType
	conns map[string]*udpMuxedConn

	addressMapMu sync.RWMutex
	addressMap   map[string]*udpMuxedConn

	// buffer pool to recycle buffers for net.UDPAddr encodes/decodes
	pool *sync.Pool

	mu sync.Mutex
}

const maxAddrSize = 512

// UDPMuxParams are parameters for UDPMux.
type UDPMuxParams struct {
	Logger  logging.LeveledLogger
	UDPConn net.PacketConn
}

// NewUDPMuxDefault creates an implementation of UDPMux
func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault {
	if params.Logger == nil {
		params.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
	}

	m := &UDPMuxDefault{
		addressMap: map[string]*udpMuxedConn{},
		params:     params,
		conns:      make(map[string]*udpMuxedConn),
		closedChan: make(chan struct{}, 1),
		pool: &sync.Pool{
			New: func() interface{} {
				// big enough buffer to fit both packet and address
				return newBufferHolder(receiveMTU + maxAddrSize)
			},
		},
	}

	go m.connWorker()

	return m
}

// LocalAddr returns the listening address of this UDPMuxDefault
func (m *UDPMuxDefault) LocalAddr() net.Addr {
	return m.params.UDPConn.LocalAddr()
}

// GetConn returns a PacketConn given the connection's ufrag and network
// creates the connection if an existing one can't be found
func (m *UDPMuxDefault) GetConn(ufrag string) (net.PacketConn, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.IsClosed() {
		return nil, io.ErrClosedPipe
	}

	if c, ok := m.conns[ufrag]; ok {
		return c, nil
	}

	c := m.createMuxedConn(ufrag)
	go func() {
		<-c.CloseChannel()
		m.removeConn(ufrag)
	}()
	m.conns[ufrag] = c
	return c, nil
}

// RemoveConnByUfrag stops and removes the muxed packet connection
func (m *UDPMuxDefault) RemoveConnByUfrag(ufrag string) {
	m.mu.Lock()
	removedConns := make([]*udpMuxedConn, 0)
	for key := range m.conns {
		if key != ufrag {
			continue
		}

		c := m.conns[key]
		delete(m.conns, key)
		if c != nil {
			removedConns = append(removedConns, c)
		}
	}
	// keep lock section small to avoid deadlock with conn lock
	m.mu.Unlock()

	m.addressMapMu.Lock()
	defer m.addressMapMu.Unlock()

	for _, c := range removedConns {
		addresses := c.getAddresses()
		for _, addr := range addresses {
			delete(m.addressMap, addr)
		}
	}
}

// IsClosed returns true if the mux had been closed
func (m *UDPMuxDefault) IsClosed() bool {
	select {
	case <-m.closedChan:
		return true
	default:
		return false
	}
}

// Close the mux, no further connections could be created
func (m *UDPMuxDefault) Close() error {
	var err error
	m.closeOnce.Do(func() {
		m.mu.Lock()
		defer m.mu.Unlock()

		for _, c := range m.conns {
			_ = c.Close()
		}
		m.conns = make(map[string]*udpMuxedConn)
		close(m.closedChan)
	})
	return err
}

func (m *UDPMuxDefault) removeConn(key string) {
	m.mu.Lock()
	c := m.conns[key]
	delete(m.conns, key)
	// keep lock section small to avoid deadlock with conn lock
	m.mu.Unlock()

	if c == nil {
		return
	}

	m.addressMapMu.Lock()
	defer m.addressMapMu.Unlock()

	addresses := c.getAddresses()
	for _, addr := range addresses {
		delete(m.addressMap, addr)
	}
}

func (m *UDPMuxDefault) writeTo(buf []byte, raddr net.Addr) (n int, err error) {
	return m.params.UDPConn.WriteTo(buf, raddr)
}

func (m *UDPMuxDefault) registerConnForAddress(conn *udpMuxedConn, addr string) {
	if m.IsClosed() {
		return
	}

	m.addressMapMu.Lock()
	defer m.addressMapMu.Unlock()

	existing, ok := m.addressMap[addr]
	if ok {
		existing.removeAddress(addr)
	}
	m.addressMap[addr] = conn

	m.params.Logger.Debugf("Registered %s for %s", addr, conn.params.Key)
}

func (m *UDPMuxDefault) createMuxedConn(key string) *udpMuxedConn {
	c := newUDPMuxedConn(&udpMuxedConnParams{
		Mux:       m,
		Key:       key,
		AddrPool:  m.pool,
		LocalAddr: m.LocalAddr(),
		Logger:    m.params.Logger,
	})
	return c
}

func (m *UDPMuxDefault) connWorker() {
	logger := m.params.Logger

	defer func() {
		_ = m.Close()
	}()

	buf := make([]byte, receiveMTU)
	for {
		n, addr, err := m.params.UDPConn.ReadFrom(buf)
		if m.IsClosed() {
			return
		} else if err != nil {
			if os.IsTimeout(err) {
				continue
			} else if err != io.EOF {
				logger.Errorf("could not read udp packet: %v", err)
			}

			return
		}

		udpAddr, ok := addr.(*net.UDPAddr)
		if !ok {
			logger.Errorf("underlying PacketConn did not return a UDPAddr")
			return
		}

		// If we have already seen this address dispatch to the appropriate destination
		m.addressMapMu.Lock()
		destinationConn := m.addressMap[addr.String()]
		m.addressMapMu.Unlock()

		// If we haven't seen this address before but is a STUN packet lookup by ufrag
		if destinationConn == nil && stun.IsMessage(buf[:n]) {
			msg := &stun.Message{
				Raw: append([]byte{}, buf[:n]...),
			}

			if err = msg.Decode(); err != nil {
				m.params.Logger.Warnf("Failed to handle decode ICE from %s: %v\n", addr.String(), err)
				continue
			}

			attr, stunAttrErr := msg.Get(stun.AttrUsername)
			if stunAttrErr != nil {
				m.params.Logger.Warnf("No Username attribute in STUN message from %s\n", addr.String())
				continue
			}

			ufrag := strings.Split(string(attr), ":")[0]

			m.mu.Lock()
			destinationConn = m.conns[ufrag]
			m.mu.Unlock()
		}

		if destinationConn == nil {
			m.params.Logger.Tracef("dropping packet from %s, addr: %s", udpAddr.String(), addr.String())
			continue
		}

		if err = destinationConn.writePacket(buf[:n], udpAddr); err != nil {
			m.params.Logger.Errorf("could not write packet: %v", err)
		}
	}
}

type bufferHolder struct {
	buffer []byte
}

func newBufferHolder(size int) *bufferHolder {
	return &bufferHolder{
		buffer: make([]byte, size),
	}
}
<|endoftext|>


<commit_before>package ssh

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"os/user"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"

	"github.com/blacknon/lssh/conf"
	termbox "github.com/nsf/termbox-go"
	"github.com/shavac/gexpect"
)

func tmpFileName() string {
	var n uint64
	binary.Read(rand.Reader, binary.LittleEndian, &n)
	return strconv.FormatUint(n, 36) + ".lssh.tmp"
}

func outColorStrings(num int, inStrings string) string {
	returnStr := ""
	color := num % 5
	switch color {
	// Red
	case 1:
		returnStr = fmt.Sprintf("\x1b[31m%s\x1b[0m", inStrings)
	// Yellow
	case 2:
		returnStr = fmt.Sprintf("\x1b[33m%s\x1b[0m", inStrings)
	// Blue
	case 3:
		returnStr = fmt.Sprintf("\x1b[34m%s\x1b[0m", inStrings)
	// Magenta
	case 4:
		returnStr = fmt.Sprintf("\x1b[35m%s\x1b[0m", inStrings)
	// Cyan
	case 0:
		returnStr = fmt.Sprintf("\x1b[36m%s\x1b[0m", inStrings)

	}
	return returnStr
}

// OS ssh wrapper(terminal connect)
func ConnectSshTerminal(connectServer string, confList conf.Config) int {
	// Get ssh config value
	connectHost := confList.Server[connectServer].User + "@" + confList.Server[connectServer].Addr
	connectPort := "22"
	if confList.Server[connectServer].Port != "" {
		connectPort = confList.Server[connectServer].Port
	}
	connectPass := confList.Server[connectServer].Pass
	connectKey := confList.Server[connectServer].Key

	// ssh command Args
	// "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' connectUser@connectAddr -p connectPort"
	sshCmd := "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' " + connectHost + " -p " + connectPort
	if connectKey != "" {
		// "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' -i connectKey connectUser@connectAddr -p connectPort"
		sshCmd = "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' -i " + connectKey + " " + connectHost + " -p " + connectPort
	}

	// log Enable
	execCmd := ""
	if confList.Log.Enable == true {
		logDirPath := confList.Log.Dir
		execOS := runtime.GOOS
		execCmd = "/usr/bin/script"

		// ~ replace User current Directory
		usr, _ := user.Current()
		logDirPath = strings.Replace(logDirPath, "~", usr.HomeDir, 1)

		// mkdir logDIr
		if err := os.MkdirAll(logDirPath, 0700); err != nil {
			fmt.Println(err)
			return 1
		}

		// Golang time.format YYYYmmdd_HHMMSS = "20060102_150405".(https://golang.org/src/time/format.go)
		logFile := time.Now().Format("20060102_150405") + "_" + connectServer + ".log"
		logFilePATH := logDirPath + "/" + logFile
		awkCmd := ">(awk '{print strftime(\"%F %T \") $0}{fflush()}'>>" + logFilePATH + ")"

		// OS check
		if execOS == "linux" || execOS == "android" {
			execCmd = "/usr/bin/script -qf -c \"" + sshCmd + "\" " + awkCmd
		} else {
			execCmd = "/usr/bin/script -qF " + awkCmd + " " + sshCmd
		}

	} else {
		execCmd = sshCmd
	}

	// exec ssh command
	child, _ := gexpect.NewSubProcess("/bin/bash", "-c", execCmd)

	if err := child.Start(); err != nil {
		fmt.Println(err)
		return 1
	}
	defer child.Close()

	// Terminal Size Change Trap
	signal_chan := make(chan os.Signal, 1)
	signal.Notify(signal_chan,
		syscall.SIGWINCH)

	go func() {
		for {
			s := <-signal_chan
			switch s {
			case syscall.SIGWINCH:
				// Get terminal window size
				if err := termbox.Init(); err != nil {
					panic(err)
				}
				width, hight := termbox.Size()
				termbox.Close()
				//fmt.Println(width, hight)
			}
		}
	}()

	// Password Input
	if connectPass != "" {
		passwordPrompt := "word:"
		if idx, _ := child.ExpectTimeout(20*time.Second, regexp.MustCompile(passwordPrompt)); idx >= 0 {
			child.SendLine(connectPass)

		} else {
			fmt.Println("Not Connected")
			return 1
		}
	}

	// timeout
	child.InteractTimeout(2419200 * time.Second)
	return 0
}

// exec ssh command function
func execCommandOverSsh(connectServerNum int, connectServer string, connectServerHeadLength int, listSum int, confList conf.Config, terminalMode bool, parallelMode bool, stdinTempPath string, execRemoteCmd ...string) int {
	connectPort := "22"
	if confList.Server[connectServer].Port != "" {
		connectPort = confList.Server[connectServer].Port
	}
	connectHostPort := confList.Server[connectServer].Addr + ":" + connectPort
	connectPass := confList.Server[connectServer].Pass
	connectKey := confList.Server[connectServer].Key

	// Set ssh client config
	config := &ssh.ClientConfig{}
	if connectKey != "" {
		// Read PublicKey
		buffer, err := ioutil.ReadFile(connectKey)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Error:%s%n", err)
			return 1
		}
		key, err := ssh.ParsePrivateKey(buffer)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Error:%s%n", err)
			return 1
		}

		// Create ssh client config for KeyAuth
		config = &ssh.ClientConfig{
			User: confList.Server[connectServer].User,
			Auth: []ssh.AuthMethod{
				ssh.PublicKeys(key)},
			Timeout: 60 * time.Second,
		}
	} else {
		// Create ssh client config for PasswordAuth
		config = &ssh.ClientConfig{
			User: confList.Server[connectServer].User,
			Auth: []ssh.AuthMethod{
				ssh.Password(connectPass)},
			Timeout: 60 * time.Second,
		}
	}

	// New Connext create
	conn, err := ssh.Dial("tcp", connectHostPort, config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot connect %v: %v \n", connectHostPort, err)
		return 1
	}
	defer conn.Close()

	// New Session
	session, err := conn.NewSession()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot open new session: %v \n", err)
		return 1
	}
	defer session.Close()

	go func() {
		time.Sleep(2419200 * time.Second)
		conn.Close()
	}()

	if terminalMode == true {
		modes := ssh.TerminalModes{
			ssh.ECHO:          0,
			ssh.TTY_OP_ISPEED: 14400,
			ssh.TTY_OP_OSPEED: 14400,
		}

		// Get terminal window size
		if err := termbox.Init(); err != nil {
			panic(err)
		}
		width, hight := termbox.Size()
		termbox.Close()

		if err := session.RequestPty("xterm", hight, width, modes); err != nil {
			session.Close()
			fmt.Errorf("request for pseudo terminal failed: %s \n", err)
		}
	}

	// stdin tmp file Open.
	stdinTempRead, _ := os.OpenFile(stdinTempPath, os.O_RDONLY, 0600)
	session.Stdin = stdinTempRead

	// exec command join
	execRemoteCmdString := strings.Join(execRemoteCmd, " ")

	if listSum == 1 {
		session.Stdout = os.Stdout
		session.Stderr = os.Stderr

		err = session.Run(execRemoteCmdString)

		if err != nil {
			fmt.Fprint(os.Stderr, err)
			if ee, ok := err.(*ssh.ExitError); ok {
				return ee.ExitStatus()
			}
		}
		return 0
	} else if parallelMode == true {
		// stdout and stderr to stdoutBuf
		var stdoutBuf bytes.Buffer

		session.Stdout = &stdoutBuf
		session.Stderr = &stdoutBuf

		// commandStatus: Chan can not continuously read buffer in for.
		//                For this reason, the processing end is detected using a variable.
		commandStatus := true

		// Exec Command(parallel)
		go func() {
			err = session.Run(execRemoteCmdString)
			commandStatus = false
		}()

		// var "x" is Readed Byte position from Buffer
		x := 0
		for {
			time.Sleep(100 * time.Millisecond)

			stdoutBufToStr := stdoutBuf.String()

			if len(stdoutBufToStr) == 0 {
				continue
			}

			stdoutBufToByte := []byte(stdoutBufToStr)
			stdoutBufArray := regexp.MustCompile("\r\n|\n\r|\n|\r").Split(string(stdoutBufToByte[x:]), -1)
			x = len(stdoutBufToStr)

			for i, v := range stdoutBufArray {
				if i == len(stdoutBufArray)-1 {
					break
				}

				lineHeader := fmt.Sprintf("%-*s", connectServerHeadLength, connectServer)
				fmt.Println(outColorStrings(connectServerNum, lineHeader)+":: ", v)

			}
			if commandStatus == true {
				continue
			} else {
				break
			}
		}

		if err != nil {
			fmt.Fprint(os.Stderr, err)
			if ee, ok := err.(*ssh.ExitError); ok {
				return ee.ExitStatus()
			}
		}
	} else {
		// stdout and stderr to stdoutBuf
		var stdoutBuf bytes.Buffer

		session.Stdout = &stdoutBuf
		session.Stderr = &stdoutBuf

		// Exec Command(for loop)
		err = session.Run(execRemoteCmdString)
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			if ee, ok := err.(*ssh.ExitError); ok {
				return ee.ExitStatus()
			}
		}

		stdoutBufArray := regexp.MustCompile("\r\n|\n\r|\n|\r").Split(stdoutBuf.String(), -1)

		for i, v := range stdoutBufArray {
			if i == len(stdoutBufArray)-1 {
				break
			}

			lineHeader := fmt.Sprintf("%-*s", connectServerHeadLength, connectServer)
			fmt.Println(outColorStrings(connectServerNum, lineHeader)+":: ", v)

		}

	}

	return 0
}

// remote ssh server exec command only
func ConnectSshCommand(connectServerList []string, confList conf.Config, terminalMode bool, parallelMode bool, execRemoteCmd ...string) int {
	// Create tmp file
	stdinTemp, err := ioutil.TempFile("", tmpFileName())
	if err != nil {
		panic(err)
	}
	defer os.Remove(stdinTemp.Name())

	// Stdin only pipes
	if terminal.IsTerminal(syscall.Stdin) == false {
		io.Copy(stdinTemp, os.Stdin)
	}

	// get connect server name max length
	connectServerMaxLength := 0
	for _, connectServerName := range connectServerList {
		if connectServerMaxLength < len(connectServerName) {
			connectServerMaxLength = len(connectServerName)
		}
	}

	connectServerCount := len(connectServerList)
	if connectServerCount > 1 {
		finished := make(chan bool)

		// for command exec
		x := 1
		for _, connectServer := range connectServerList {
			y := x
			targetServer := connectServer
			go func() {
				execCommandOverSsh(y, targetServer, connectServerMaxLength, connectServerCount, confList, terminalMode, parallelMode, stdinTemp.Name(), execRemoteCmd...)
				finished <- true
			}()
			x++
		}

		for i := 1; i <= connectServerCount; i++ {
			<-finished
		}
	} else {
		for i, connectServer := range connectServerList {
			execCommandOverSsh(i, connectServer, connectServerMaxLength, connectServerCount, confList, terminalMode, parallelMode, stdinTemp.Name(), execRemoteCmd...)
		}
	}

	return 0
}
<commit_msg>signal winch after Windows size change at connect ssh<commit_after>package ssh

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"os/user"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"

	"github.com/blacknon/lssh/conf"
	termbox "github.com/nsf/termbox-go"
	"github.com/shavac/gexpect"
)

func tmpFileName() string {
	var n uint64
	binary.Read(rand.Reader, binary.LittleEndian, &n)
	return strconv.FormatUint(n, 36) + ".lssh.tmp"
}

func outColorStrings(num int, inStrings string) string {
	returnStr := ""
	color := num % 5
	switch color {
	// Red
	case 1:
		returnStr = fmt.Sprintf("\x1b[31m%s\x1b[0m", inStrings)
	// Yellow
	case 2:
		returnStr = fmt.Sprintf("\x1b[33m%s\x1b[0m", inStrings)
	// Blue
	case 3:
		returnStr = fmt.Sprintf("\x1b[34m%s\x1b[0m", inStrings)
	// Magenta
	case 4:
		returnStr = fmt.Sprintf("\x1b[35m%s\x1b[0m", inStrings)
	// Cyan
	case 0:
		returnStr = fmt.Sprintf("\x1b[36m%s\x1b[0m", inStrings)

	}
	return returnStr
}

// OS ssh wrapper(terminal connect)
func ConnectSshTerminal(connectServer string, confList conf.Config) int {
	// Get ssh config value
	connectHost := confList.Server[connectServer].User + "@" + confList.Server[connectServer].Addr
	connectPort := "22"
	if confList.Server[connectServer].Port != "" {
		connectPort = confList.Server[connectServer].Port
	}
	connectPass := confList.Server[connectServer].Pass
	connectKey := confList.Server[connectServer].Key

	// ssh command Args
	// "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' connectUser@connectAddr -p connectPort"
	sshCmd := "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' " + connectHost + " -p " + connectPort
	if connectKey != "" {
		// "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' -i connectKey connectUser@connectAddr -p connectPort"
		sshCmd = "/usr/bin/ssh -o 'StrictHostKeyChecking no' -o 'NumberOfPasswordPrompts 1' -i " + connectKey + " " + connectHost + " -p " + connectPort
	}

	// log Enable
	execCmd := ""
	if confList.Log.Enable == true {
		logDirPath := confList.Log.Dir
		execOS := runtime.GOOS
		execCmd = "/usr/bin/script"

		// ~ replace User current Directory
		usr, _ := user.Current()
		logDirPath = strings.Replace(logDirPath, "~", usr.HomeDir, 1)

		// mkdir logDIr
		if err := os.MkdirAll(logDirPath, 0700); err != nil {
			fmt.Println(err)
			return 1
		}

		// Golang time.format YYYYmmdd_HHMMSS = "20060102_150405".(https://golang.org/src/time/format.go)
		logFile := time.Now().Format("20060102_150405") + "_" + connectServer + ".log"
		logFilePATH := logDirPath + "/" + logFile
		awkCmd := ">(awk '{print strftime(\"%F %T \") $0}{fflush()}'>>" + logFilePATH + ")"

		// OS check
		if execOS == "linux" || execOS == "android" {
			execCmd = "/usr/bin/script -qf -c \"" + sshCmd + "\" " + awkCmd
		} else {
			execCmd = "/usr/bin/script -qF " + awkCmd + " " + sshCmd
		}

	} else {
		execCmd = sshCmd
	}

	// exec ssh command
	child, _ := gexpect.NewSubProcess("/bin/bash", "-c", execCmd)

	if err := child.Start(); err != nil {
		fmt.Println(err)
		return 1
	}
	defer child.Close()

	// Terminal Size Change Trap
	signal_chan := make(chan os.Signal, 1)
	signal.Notify(signal_chan,
		syscall.SIGWINCH)

	//go func() {
		//for {
			//s := <-signal_chan
			//switch s {
			//case syscall.SIGWINCH:
			//	// Get terminal window size
			//	if err := termbox.Init(); err != nil {
			//		panic(err)
			//	}
			//	width, hight := termbox.Size()
			//	termbox.Close()
			//	//fmt.Println(width, hight)
			//}
		//}
	//}()

	// Password Input
	if connectPass != "" {
		passwordPrompt := "word:"
		if idx, _ := child.ExpectTimeout(20*time.Second, regexp.MustCompile(passwordPrompt)); idx >= 0 {
			child.SendLine(connectPass)

		} else {
			fmt.Println("Not Connected")
			return 1
		}
	}

	// timeout
	child.InteractTimeout(2419200 * time.Second)
	return 0
}

// exec ssh command function
func execCommandOverSsh(connectServerNum int, connectServer string, connectServerHeadLength int, listSum int, confList conf.Config, terminalMode bool, parallelMode bool, stdinTempPath string, execRemoteCmd ...string) int {
	connectPort := "22"
	if confList.Server[connectServer].Port != "" {
		connectPort = confList.Server[connectServer].Port
	}
	connectHostPort := confList.Server[connectServer].Addr + ":" + connectPort
	connectPass := confList.Server[connectServer].Pass
	connectKey := confList.Server[connectServer].Key

	// Set ssh client config
	config := &ssh.ClientConfig{}
	if connectKey != "" {
		// Read PublicKey
		buffer, err := ioutil.ReadFile(connectKey)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Error:%s%n", err)
			return 1
		}
		key, err := ssh.ParsePrivateKey(buffer)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Error:%s%n", err)
			return 1
		}

		// Create ssh client config for KeyAuth
		config = &ssh.ClientConfig{
			User: confList.Server[connectServer].User,
			Auth: []ssh.AuthMethod{
				ssh.PublicKeys(key)},
			Timeout: 60 * time.Second,
		}
	} else {
		// Create ssh client config for PasswordAuth
		config = &ssh.ClientConfig{
			User: confList.Server[connectServer].User,
			Auth: []ssh.AuthMethod{
				ssh.Password(connectPass)},
			Timeout: 60 * time.Second,
		}
	}

	// New Connext create
	conn, err := ssh.Dial("tcp", connectHostPort, config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot connect %v: %v \n", connectHostPort, err)
		return 1
	}
	defer conn.Close()

	// New Session
	session, err := conn.NewSession()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot open new session: %v \n", err)
		return 1
	}
	defer session.Close()

	go func() {
		time.Sleep(2419200 * time.Second)
		conn.Close()
	}()

	if terminalMode == true {
		modes := ssh.TerminalModes{
			ssh.ECHO:          0,
			ssh.TTY_OP_ISPEED: 14400,
			ssh.TTY_OP_OSPEED: 14400,
		}

		// Get terminal window size
		if err := termbox.Init(); err != nil {
			panic(err)
		}
		width, hight := termbox.Size()
		termbox.Close()

		if err := session.RequestPty("xterm", hight, width, modes); err != nil {
			session.Close()
			fmt.Errorf("request for pseudo terminal failed: %s \n", err)
		}
	}

	// stdin tmp file Open.
	stdinTempRead, _ := os.OpenFile(stdinTempPath, os.O_RDONLY, 0600)
	session.Stdin = stdinTempRead

	// exec command join
	execRemoteCmdString := strings.Join(execRemoteCmd, " ")

	if listSum == 1 {
		session.Stdout = os.Stdout
		session.Stderr = os.Stderr

		err = session.Run(execRemoteCmdString)

		if err != nil {
			fmt.Fprint(os.Stderr, err)
			if ee, ok := err.(*ssh.ExitError); ok {
				return ee.ExitStatus()
			}
		}
		return 0
	} else if parallelMode == true {
		// stdout and stderr to stdoutBuf
		var stdoutBuf bytes.Buffer

		session.Stdout = &stdoutBuf
		session.Stderr = &stdoutBuf

		// commandStatus: Chan can not continuously read buffer in for.
		//                For this reason, the processing end is detected using a variable.
		commandStatus := true

		// Exec Command(parallel)
		go func() {
			err = session.Run(execRemoteCmdString)
			commandStatus = false
		}()

		// var "x" is Readed Byte position from Buffer
		x := 0
		for {
			time.Sleep(100 * time.Millisecond)

			stdoutBufToStr := stdoutBuf.String()

			if len(stdoutBufToStr) == 0 {
				continue
			}

			stdoutBufToByte := []byte(stdoutBufToStr)
			stdoutBufArray := regexp.MustCompile("\r\n|\n\r|\n|\r").Split(string(stdoutBufToByte[x:]), -1)
			x = len(stdoutBufToStr)

			for i, v := range stdoutBufArray {
				if i == len(stdoutBufArray)-1 {
					break
				}

				lineHeader := fmt.Sprintf("%-*s", connectServerHeadLength, connectServer)
				fmt.Println(outColorStrings(connectServerNum, lineHeader)+":: ", v)

			}
			if commandStatus == true {
				continue
			} else {
				break
			}
		}

		if err != nil {
			fmt.Fprint(os.Stderr, err)
			if ee, ok := err.(*ssh.ExitError); ok {
				return ee.ExitStatus()
			}
		}
	} else {
		// stdout and stderr to stdoutBuf
		var stdoutBuf bytes.Buffer

		session.Stdout = &stdoutBuf
		session.Stderr = &stdoutBuf

		// Exec Command(for loop)
		err = session.Run(execRemoteCmdString)
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			if ee, ok := err.(*ssh.ExitError); ok {
				return ee.ExitStatus()
			}
		}

		stdoutBufArray := regexp.MustCompile("\r\n|\n\r|\n|\r").Split(stdoutBuf.String(), -1)

		for i, v := range stdoutBufArray {
			if i == len(stdoutBufArray)-1 {
				break
			}

			lineHeader := fmt.Sprintf("%-*s", connectServerHeadLength, connectServer)
			fmt.Println(outColorStrings(connectServerNum, lineHeader)+":: ", v)

		}

	}

	return 0
}

// remote ssh server exec command only
func ConnectSshCommand(connectServerList []string, confList conf.Config, terminalMode bool, parallelMode bool, execRemoteCmd ...string) int {
	// Create tmp file
	stdinTemp, err := ioutil.TempFile("", tmpFileName())
	if err != nil {
		panic(err)
	}
	defer os.Remove(stdinTemp.Name())

	// Stdin only pipes
	if terminal.IsTerminal(syscall.Stdin) == false {
		io.Copy(stdinTemp, os.Stdin)
	}

	// get connect server name max length
	connectServerMaxLength := 0
	for _, connectServerName := range connectServerList {
		if connectServerMaxLength < len(connectServerName) {
			connectServerMaxLength = len(connectServerName)
		}
	}

	connectServerCount := len(connectServerList)
	if connectServerCount > 1 {
		finished := make(chan bool)

		// for command exec
		x := 1
		for _, connectServer := range connectServerList {
			y := x
			targetServer := connectServer
			go func() {
				execCommandOverSsh(y, targetServer, connectServerMaxLength, connectServerCount, confList, terminalMode, parallelMode, stdinTemp.Name(), execRemoteCmd...)
				finished <- true
			}()
			x++
		}

		for i := 1; i <= connectServerCount; i++ {
			<-finished
		}
	} else {
		for i, connectServer := range connectServerList {
			execCommandOverSsh(i, connectServer, connectServerMaxLength, connectServerCount, confList, terminalMode, parallelMode, stdinTemp.Name(), execRemoteCmd...)
		}
	}

	return 0
}
<|endoftext|>


<commit_before>package golib

import (
	"bytes"

	"github.com/lunixbochs/vtclean"
	"golang.org/x/text/width"
)

// StringLength returns the number of normalized utf8-runes within the cleaned string.
// Clean means the string is stripped of terminal escape characters and color codes.
func StringLength(str string) (strlen int) {
	str = vtclean.Clean(str, false)
	for str != "" {
		_, rest, runeWidth := ReadRune(str)
		str = rest
		strlen += runeWidth
	}
	// Alternative:
	// strlen = utf8.RuneCountInString(str)
	// Alternative:
	// var ia norm.Iter
	// ia.InitString(norm.NFKD, str)
	return
}

// ReadRune reads one utf8 rune from the input string and provides information
// about the character width of the read rune.
func ReadRune(input string) (theRune string, rest string, runeWidth int) {
	prop, size := width.LookupString(input)
	rest = input[size:]
	theRune = input[:size]
	switch prop.Kind() {
	case width.EastAsianFullwidth, width.EastAsianHalfwidth, width.EastAsianWide:
		runeWidth = 2
	default:
		runeWidth = 1
	}
	return
}

// Substring returns a substring of the given input string, but the indices
// iFrom and iTo point to normalized utf8-runes within the cleaned string.
// Clean means the string is stripped of terminal escape characters and color codes.
// The total number of normalized utf8 runes in the clean string can be obtained from
// the StringLength() function.
// All runes and special codes will be preserved in the output string.
func Substring(str string, iFrom int, iTo int) string {

	// Find the start in the input string
	buf := bytes.NewBuffer(make([]byte, 0, len(str)))
	textWidth := 0
	cleanedLen := 0
	for str != "" && textWidth < iFrom {
		runeStr, rest, runeWidth := ReadRune(str)
		buf.Write([]byte(runeStr))

		cleaned := vtclean.Clean(buf.String(), false)
		if len(cleaned) > cleanedLen {
			// A visible rune was added
			cleanedLen = len(cleaned)
			textWidth += runeWidth
		}
		if textWidth >= iFrom {
			break
		}
		str = rest
	}

	iLen := iTo - iFrom
	possibleColor := false

	// Find the end in the input string
	to := 0
	suffix := []byte(str)
	buf.Reset()
	textWidth = 0
	cleanedLen = 0
	for str != "" && textWidth < iLen {
		runeStr, rest, runeWidth := ReadRune(str)
		buf.Write([]byte(runeStr))

		cleaned := vtclean.Clean(buf.String(), false)
		if len(cleaned) > cleanedLen {
			// A visible rune was added
			cleanedLen = len(cleaned)
			textWidth += runeWidth
		}
		if len(cleaned) != buf.Len() {
			// Contains invisible escape characters, possibly color codes
			possibleColor = true
		}
		if textWidth > iLen {
			if runeWidth == 2 {
				// Splitting wide character in the end, pad with a space
				suffix[to] = ' '
				to++
			}
			break
		}
		str = rest
		to += len(runeStr)
	}

	result := string(suffix[:to])
	if possibleColor {
		// Might contain color codes, make sure to disable colors at the end
		result += "\033[0m"
	}
	return result
}
<commit_msg>Added convenience method to compare two string slices<commit_after>package golib

import (
	"bytes"

	"github.com/lunixbochs/vtclean"
	"golang.org/x/text/width"
)

// StringLength returns the number of normalized utf8-runes within the cleaned string.
// Clean means the string is stripped of terminal escape characters and color codes.
func StringLength(str string) (strlen int) {
	str = vtclean.Clean(str, false)
	for str != "" {
		_, rest, runeWidth := ReadRune(str)
		str = rest
		strlen += runeWidth
	}
	// Alternative:
	// strlen = utf8.RuneCountInString(str)
	// Alternative:
	// var ia norm.Iter
	// ia.InitString(norm.NFKD, str)
	return
}

// ReadRune reads one utf8 rune from the input string and provides information
// about the character width of the read rune.
func ReadRune(input string) (theRune string, rest string, runeWidth int) {
	prop, size := width.LookupString(input)
	rest = input[size:]
	theRune = input[:size]
	switch prop.Kind() {
	case width.EastAsianFullwidth, width.EastAsianHalfwidth, width.EastAsianWide:
		runeWidth = 2
	default:
		runeWidth = 1
	}
	return
}

// Substring returns a substring of the given input string, but the indices
// iFrom and iTo point to normalized utf8-runes within the cleaned string.
// Clean means the string is stripped of terminal escape characters and color codes.
// The total number of normalized utf8 runes in the clean string can be obtained from
// the StringLength() function.
// All runes and special codes will be preserved in the output string.
func Substring(str string, iFrom int, iTo int) string {

	// Find the start in the input string
	buf := bytes.NewBuffer(make([]byte, 0, len(str)))
	textWidth := 0
	cleanedLen := 0
	for str != "" && textWidth < iFrom {
		runeStr, rest, runeWidth := ReadRune(str)
		buf.Write([]byte(runeStr))

		cleaned := vtclean.Clean(buf.String(), false)
		if len(cleaned) > cleanedLen {
			// A visible rune was added
			cleanedLen = len(cleaned)
			textWidth += runeWidth
		}
		if textWidth >= iFrom {
			break
		}
		str = rest
	}

	iLen := iTo - iFrom
	possibleColor := false

	// Find the end in the input string
	to := 0
	suffix := []byte(str)
	buf.Reset()
	textWidth = 0
	cleanedLen = 0
	for str != "" && textWidth < iLen {
		runeStr, rest, runeWidth := ReadRune(str)
		buf.Write([]byte(runeStr))

		cleaned := vtclean.Clean(buf.String(), false)
		if len(cleaned) > cleanedLen {
			// A visible rune was added
			cleanedLen = len(cleaned)
			textWidth += runeWidth
		}
		if len(cleaned) != buf.Len() {
			// Contains invisible escape characters, possibly color codes
			possibleColor = true
		}
		if textWidth > iLen {
			if runeWidth == 2 {
				// Splitting wide character in the end, pad with a space
				suffix[to] = ' '
				to++
			}
			break
		}
		str = rest
		to += len(runeStr)
	}

	result := string(suffix[:to])
	if possibleColor {
		// Might contain color codes, make sure to disable colors at the end
		result += "\033[0m"
	}
	return result
}

func EqualStrings(a, b []string) bool {
	switch {
	case len(a) != len(b):
		return false
	case len(a) == 0:
		return true
	case &a[0] == &b[0]:
		// Compare the address of the array backing the slices
		return true
	}
	// Last resort: compare every string pair
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package runner

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/carlescere/goback"
	"github.com/harrowio/harrow/config"
	"github.com/harrowio/harrow/logger"
	"github.com/pkg/errors"
)

const imageName = "harrow-baseimage"

type LXD struct {
	config  *config.Config
	connURL *url.URL
	log     logger.Logger

	containerUuid string
}

func (lxd *LXD) MakeContainer() error {

	lxd.log.Info().Msgf("making new container: %s", lxd.containerName())

	session, err := lxd.sshSession()
	if err != nil {
		return err
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc launch %s -p default -p docker %s", imageName, lxd.containerName())
	lxd.log.Info().Msgf("running %s", cmdStr)

	output, err := session.CombinedOutput(cmdStr)
	if err != nil {
		return errors.Wrapf(err, "error starting lxd container: %s", string(output))
	}

	if !strings.Contains(string(output), fmt.Sprintf("Starting %s", lxd.containerName())) {
		return errors.Wrap(err, "expected note about 'Starting <name>', wasn't recieved")
	}

	return nil
}

func (lxd *LXD) DestroyContainer() error {

	lxd.log.Info().Msgf("destroying container: %s", lxd.containerName())

	b := &goback.SimpleBackoff{Min: 1 * time.Minute, Max: 10 * time.Minute, Factor: 2}

	session, err := lxd.sshSession()
	if err != nil {
		return err
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc delete --force %s", lxd.containerName())

	for {
		lxd.log.Info().Msgf("running %s", cmdStr)
		output, err := session.CombinedOutput(cmdStr)

		if err == nil {
			return nil // container deleted successfully
		}

		if err != nil {
			if strings.Contains(string(output), "error: not found") {
				return nil // container does not exist and error was "error: not found" (exit status: 1)
			}
			if err := goback.Wait(b); err != nil { // goback.ErrMaxAttemptsExceeded incase we're over-time
				return errors.Wrap(err, "max attempts exceeded waiting to destroy container")
			}
			return nil
		}

	}
}

func (lxd *LXD) WaitForContainerNetworking(d time.Duration) error {
	lxd.log.Info().Msgf("waiting for container %s networking (max %s)", lxd.containerName(), d)

	b := &goback.SimpleBackoff{Min: 5 * time.Second, Max: d, Factor: 2}
	for {
		err := lxd.CheckContainerNetworking()
		if err != nil {
			if err := goback.Wait(b); err != nil { // goback.ErrMaxAttemptsExceeded incase we're over-time
				return errors.Wrap(err, "max attempts exceeded waiting for networking to come up")
			}
		}
		break
	}
	lxd.log.Info().Msgf("container networking for %s is up!", lxd.containerName())
	return nil
}

func (lxd *LXD) ContainerExists() (bool, error) {

	lxd.log.Debug().Msgf("checking if container %s exists", lxd.containerName())

	session, err := lxd.sshSession()
	if err != nil {
		return false, errors.Wrap(err, "error getting session to check container existence")
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc info %s", lxd.containerName())

	output, err := session.CombinedOutput(cmdStr)
	if err == nil {
		return true, nil // container exists, we could get it's info
	}
	if err != nil {
		if strings.Contains(string(output), "error: not found") {
			return false, nil // container does not exist and error was "error: not found" (exit status: 1)
		}
		return false, errors.Wrapf(err, "error calling lxd info %s: %s", lxd.containerName(), output)
	}
	return false, nil
}

func (lxd *LXD) MaintainConnection(lost chan<- error) {
	lxd.log.Info().Msg("opening long lived connection to the container")

	session, err := lxd.sshSession()
	if err != nil {
		lost <- errors.Wrap(err, "unable to open ssh session")
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc exec %s -- cat", lxd.containerName())
	lxd.log.Info().Msgf("running %s", cmdStr)

	// Set up terminal modes
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	// Request pseudo terminal
	if err := session.RequestPty("xterm", 40, 80, modes); err != nil {
		lost <- errors.Wrapf(err, "upstream server rejected requets for a pty: %s", err)
	}

	err = session.Run(cmdStr)
	lost <- errors.Wrapf(err, "error maintaining long running command on container: %s", err)

}

func (lxd *LXD) CheckContainerNetworking() error {

	lxd.log.Info().Msg("checking for container networking")

	session, err := lxd.sshSession()
	if err != nil {
		return err
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc list %s --format json", lxd.containerName())
	lxd.log.Info().Msgf("running %s", cmdStr)

	output, err := session.CombinedOutput(cmdStr)
	if err != nil {
		return errors.Wrapf(err, "error getting container info from lxc: ", string(output))
	}

	// This messy struct maps to the LXC struct. We could import LXC's go package
	// to get this, but then we'd have a dependency on something the size of the
	// universe. #notworthit
	containerInfo := []struct {
		State struct {
			Network struct {
				Eth0 struct {
					Addresses []struct {
						Family  string
						Address string
					}
				}
			}
		}
	}{}

	if err := json.Unmarshal(output, &containerInfo); err != nil {
		return errors.Wrap(err, "error unmarshalling the container information from lxc")
	}

	for _, info := range containerInfo {
		for _, addr := range info.State.Network.Eth0.Addresses {
			if addr.Family == "inet" {
				return nil
			}
		}
	}

	return fmt.Errorf("no inet device found on container %s", lxd.containerName())
}

func (lxd *LXD) sshClient() (*ssh.Client, error) {
	sshConfig, err := lxd.config.GetSshConfig()
	if err != nil {
		return nil, errors.Wrap(err, "could not get ssh configuration")
	}
	client, err := ssh.Dial("tcp", lxd.connURL.Host, sshConfig)
	if err != nil {
		lxd.log.Fatal().Msgf("Unable to open ssh connection: %s", err)
	}
	return client, nil
}

func (lxd *LXD) sshSession() (*ssh.Session, error) {
	client, err := lxd.sshClient()
	if err != nil {
		return nil, errors.Wrap(err, "could not get ssh client")
	}
	session, err := client.NewSession()
	if err != nil {
		return nil, errors.Wrap(err, "could not start ssh session")
	}
	return session, nil
}

func (lxd *LXD) containerName() string {
	return fmt.Sprintf("%s-%s", lxd.prefix(), lxd.containerUuid)
}

func (lxd *LXD) prefix() string {
	return "container"
}
<commit_msg>Ensure that the runner does not fail on flakey SSH connections.<commit_after>package runner

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/carlescere/goback"
	"github.com/harrowio/harrow/config"
	"github.com/harrowio/harrow/logger"
	"github.com/pkg/errors"
)

const imageName = "harrow-baseimage"

type LXD struct {
	config  *config.Config
	connURL *url.URL
	log     logger.Logger

	containerUuid string
}

func (lxd *LXD) MakeContainer() error {

	lxd.log.Info().Msgf("making new container: %s", lxd.containerName())

	session, err := lxd.sshSession()
	if err != nil {
		return err
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc launch %s -p default -p docker %s", imageName, lxd.containerName())
	lxd.log.Info().Msgf("running %s", cmdStr)

	output, err := session.CombinedOutput(cmdStr)
	if err != nil {
		return errors.Wrapf(err, "error starting lxd container: %s", string(output))
	}

	if !strings.Contains(string(output), fmt.Sprintf("Starting %s", lxd.containerName())) {
		return errors.Wrap(err, "expected note about 'Starting <name>', wasn't recieved")
	}

	return nil
}

func (lxd *LXD) DestroyContainer() error {

	lxd.log.Info().Msgf("destroying container: %s", lxd.containerName())

	b := &goback.SimpleBackoff{Min: 1 * time.Minute, Max: 10 * time.Minute, Factor: 2}

	session, err := lxd.sshSession()
	if err != nil {
		return err
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc delete --force %s", lxd.containerName())

	for {
		lxd.log.Info().Msgf("running %s", cmdStr)
		output, err := session.CombinedOutput(cmdStr)

		if err == nil {
			return nil // container deleted successfully
		}

		if err != nil {
			if strings.Contains(string(output), "error: not found") {
				return nil // container does not exist and error was "error: not found" (exit status: 1)
			}
			if strings.Contains(string(output), "error: Could not remove LV named") {
				// container exists in LXD DB but underlaying storage has been removed, newer versions
				// of LXD handle this more gracefully
				return nil
			}
			if errBack := goback.Wait(b); errBack != nil { // goback.ErrMaxAttemptsExceeded incase we're over-time
				return errors.Wrap(err, "max attempts exceeded waiting to destroy container")
			} else {
				lxd.log.Info().Msgf("error destroying container, will retry: %s", err)
			}
			return nil

		}

	}
}

func (lxd *LXD) WaitForContainerNetworking(d time.Duration) error {
	lxd.log.Info().Msgf("waiting for container %s networking (max %s)", lxd.containerName(), d)

	b := &goback.SimpleBackoff{Min: 5 * time.Second, Max: d, Factor: 2}
	for {
		err := lxd.CheckContainerNetworking()
		if err != nil {
			if errBack := goback.Wait(b); errBack != nil { // goback.ErrMaxAttemptsExceeded incase we're over-time
				return errors.Wrap(err, "max attempts exceeded waiting for networking to come up")
			} else {
				lxd.log.Info().Msgf("error waiting for container networking, will retry: %s", err)
			}
		}
		break
	}
	lxd.log.Info().Msgf("container networking for %s is up!", lxd.containerName())
	return nil
}

func (lxd *LXD) ContainerExists() (bool, error) {

	lxd.log.Debug().Msgf("checking if container %s exists", lxd.containerName())

	session, err := lxd.sshSession()
	if err != nil {
		return false, errors.Wrap(err, "error getting session to check container existence")
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc info %s", lxd.containerName())

	output, err := session.CombinedOutput(cmdStr)
	if err == nil {
		return true, nil // container exists, we could get it's info
	}
	if err != nil {
		if strings.Contains(string(output), "error: not found") {
			return false, nil // container does not exist and error was "error: not found" (exit status: 1)
		}
		return false, errors.Wrapf(err, "error calling lxd info %s: %s", lxd.containerName(), output)
	}
	return false, nil
}

func (lxd *LXD) MaintainConnection(lost chan<- error) {
	lxd.log.Info().Msg("opening long lived connection to the container")

	session, err := lxd.sshSession()
	if err != nil {
		lost <- errors.Wrap(err, "unable to open ssh session")
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc exec %s -- cat", lxd.containerName())
	lxd.log.Info().Msgf("running %s", cmdStr)

	// Set up terminal modes
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	// Request pseudo terminal
	if err := session.RequestPty("xterm", 40, 80, modes); err != nil {
		lost <- errors.Wrapf(err, "upstream server rejected requets for a pty: %s", err)
	}

	err = session.Run(cmdStr)
	lost <- errors.Wrapf(err, "error maintaining long running command on container: %s", err)

}

func (lxd *LXD) CheckContainerNetworking() error {

	lxd.log.Info().Msg("checking for container networking")

	session, err := lxd.sshSession()
	if err != nil {
		return err
	}
	defer session.Close()

	cmdStr := fmt.Sprintf("lxc list %s --format json", lxd.containerName())
	lxd.log.Info().Msgf("running %s", cmdStr)

	output, err := session.CombinedOutput(cmdStr)
	if err != nil {
		return errors.Wrapf(err, "error getting container info from lxc: ", string(output))
	}

	// This messy struct maps to the LXC struct. We could import LXC's go package
	// to get this, but then we'd have a dependency on something the size of the
	// universe. #notworthit
	containerInfo := []struct {
		State struct {
			Network struct {
				Eth0 struct {
					Addresses []struct {
						Family  string
						Address string
					}
				}
			}
		}
	}{}

	if err := json.Unmarshal(output, &containerInfo); err != nil {
		return errors.Wrap(err, "error unmarshalling the container information from lxc")
	}

	for _, info := range containerInfo {
		for _, addr := range info.State.Network.Eth0.Addresses {
			if addr.Family == "inet" {
				return nil
			}
		}
	}

	return fmt.Errorf("no inet device found on container %s", lxd.containerName())
}

func (lxd *LXD) sshClient() (*ssh.Client, error) {

	b := &goback.SimpleBackoff{Min: 1 * time.Second, Max: 60 * time.Minute, Factor: 2}

	sshConfig, err := lxd.config.GetSshConfig()
	if err != nil {
		return nil, errors.Wrap(err, "could not get ssh configuration")
	}
	for {
		client, err := ssh.Dial("tcp", lxd.connURL.Host, sshConfig)
		if err != nil {
			if errBack := goback.Wait(b); errBack != nil { // goback.ErrMaxAttemptsExceeded incase we're over-time
				return nil, errors.Wrap(err, "max attempts exceeded waiting to dial ssh to host")
			} else {
				lxd.log.Info().Msgf("error dialing ssh connection, will retry: %s", err)
			}
		} else {
			return client, nil
		}
	}
}

func (lxd *LXD) sshSession() (*ssh.Session, error) {

	b := &goback.SimpleBackoff{Min: 1 * time.Second, Max: 60 * time.Minute, Factor: 2}

	for {
		client, err := lxd.sshClient()
		if err != nil {
			return nil, errors.Wrap(err, "could not get ssh client")
		}
		session, err := client.NewSession()
		if err != nil {
			if errBack := goback.Wait(b); errBack != nil { // goback.ErrMaxAttemptsExceeded incase we're over-time
				return nil, errors.Wrap(err, "max attempts exceeded waiting to start ssh session")
			} else {
				lxd.log.Info().Msgf("error starting ssh session, will retry: %s", err)
			}
		} else {
			return session, nil
		}
	}
}

func (lxd *LXD) containerName() string {
	return fmt.Sprintf("%s-%s", lxd.prefix(), lxd.containerUuid)
}

func (lxd *LXD) prefix() string {
	return "container"
}
<|endoftext|>


<commit_before>package ui

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/lhz/considerate/hvsc"
	"github.com/lhz/considerate/player"
	"github.com/nsf/termbox-go"
)

const (
	MODE_BROWSE = iota
	MODE_SEARCH
)

var w, h int
var list *List
var mode int
var debugInfo string
var searchTerm []rune
var searchCursorPos int
var quit = false

func Setup() {
	err := termbox.Init()
	if err != nil {
		log.Panicln(err)
	}

	w, h = termbox.Size()

	list = NewList(h - 2)

	termbox.SetInputMode(termbox.InputEsc | termbox.InputMouse)
	draw()
}

func Run() {
	defer termbox.Close()

	mode = MODE_BROWSE

	tickChan := time.NewTicker(1 * time.Second).C

	eventChan := make(chan termbox.Event)
	go func() {
		for {
			eventChan <- termbox.PollEvent()
		}
	}()

	for !quit {
		debugInfo = ""
		select {
		case ev := <-eventChan:
			switch ev.Type {
			case termbox.EventKey:
				keyEvent(ev)
				draw()
			case termbox.EventResize:
				resizeEvent(ev)
				draw()
			case termbox.EventError:
				panic(ev.Err)
			}
		case <-tickChan:
			if player.Playing {
				draw()
			}
		}
	}
}

func resizeEvent(ev termbox.Event) {
	w, h = ev.Width, ev.Height
	draw()
}

func keyEvent(ev termbox.Event) {
	if mode == MODE_SEARCH {
		keyEventSearch(ev)
		return
	}
	switch ev.Key {
	case termbox.KeyCtrlC, termbox.KeyEsc:
		player.Stop()
		quit = true
	case termbox.KeyPgup:
		list.PrevPage()
	case termbox.KeyPgdn:
		list.NextPage()
	case termbox.KeyArrowUp:
		list.PrevTune()
	case termbox.KeyArrowDown:
		list.NextTune()
	case termbox.KeyArrowLeft:
		player.PrevSong()
	case termbox.KeyArrowRight:
		player.NextSong()
	case termbox.KeyEnter:
		player.Play(list.CurrentItem().TuneIndex, 1)
	case termbox.KeySpace:
		list.NextTune()
		player.Play(list.CurrentItem().TuneIndex, 1)
	case termbox.KeyDelete:
		player.Stop()
	case 0:
		if n := strings.IndexRune("0123456789", ev.Ch); n > 0 {
			if n <= player.CurrentTune.Header.Songs {
				player.PlaySub(n)
			}
			return
		}
		if ev.Ch == '/' {
			mode = MODE_SEARCH
		}
	default:
		debugInfo = string(ev.Key)
	}
}

func keyEventSearch(ev termbox.Event) {
	switch ev.Key {
	case termbox.KeyCtrlC, termbox.KeyEsc:
		mode = MODE_BROWSE
	case termbox.KeyArrowLeft:
		if searchCursorPos > 0 {
			searchCursorPos--
		}
	case termbox.KeyArrowRight:
		if searchCursorPos < len(searchTerm) {
			searchCursorPos++
		}
	case termbox.KeyEnter:
		if hvsc.Filter(string(searchTerm)) {
			list = NewList(h - 2)
			mode = MODE_BROWSE
		}
	case termbox.KeySpace:
		searchInsert(rune(' '))
	case 0:
		searchInsert(ev.Ch)
	case 0x7F:
		if searchCursorPos > 0 {
			searchTerm = append(searchTerm[0:searchCursorPos-1], searchTerm[searchCursorPos:len(searchTerm)]...)
			searchCursorPos--
		}
	default:
		debugInfo = string(ev.Key)
	}
}

func searchInsert(ch rune) {
	searchTerm = append(searchTerm, rune(' '))
	copy(searchTerm[searchCursorPos+1:], searchTerm[searchCursorPos:])
	searchTerm[searchCursorPos] = ch
	searchCursorPos++
}

func draw() {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	drawHeader()
	drawList()
	drawTuneInfo()
	drawFooter()
	termbox.Flush()
}

func drawHeader() {
	var header string
	bg := termbox.ColorBlack
	switch mode {
	case MODE_BROWSE:
		header = fmt.Sprintf("Browse: %s", string(searchTerm))
		termbox.HideCursor()
	case MODE_SEARCH:
		header = fmt.Sprintf("Search: %s", string(searchTerm))
		bg = termbox.ColorMagenta
		termbox.SetCursor(8+searchCursorPos, 0)
	}
	header = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len([]rune(header))), header, "")
	writeAt(0, 0, header, termbox.ColorWhite|termbox.AttrBold, bg)
}

func drawFooter() {
	footer := fmt.Sprintf("%d/%d  %s", list.CurrentItem().TuneIndex+1, hvsc.NumFilteredTunes, debugInfo)
	footer = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len([]rune(footer))), footer, "")
	writeAt(0, h-1, footer, termbox.ColorWhite, termbox.ColorBlack)
}

func drawTuneInfo() {
	if !player.Playing {
		return
	}
	tune := player.CurrentTune
	ox := 42
	oy := 2
	fg := termbox.ColorDefault
	bg := termbox.ColorDefault

	writeAt(ox, oy+0, fmt.Sprintf("Title:    %s", tune.Title()), fg, bg)
	writeAt(ox, oy+1, fmt.Sprintf("Author:   %s", tune.Header.Author), fg, bg)
	writeAt(ox, oy+2, fmt.Sprintf("Released: %s", tune.Header.Released), fg, bg)

	writeAt(ox, oy+4, fmt.Sprintf("Tune: %d/%d  Length: %s  Time: %s",
		player.CurrentSong, tune.Header.Songs, player.SongLength(), player.Elapsed()), fg, bg)

	for i, line := range tune.Info {
		if oy+6+i > h-2 {
			break
		}
		writeAt(ox, oy+6+i, line, fg, bg)
	}
}

func drawList() {
	for i, item := range list.CurrentPage() {
		fg := termbox.ColorDefault
		bg := termbox.ColorDefault
		if i == list.PagePos {
			fg = termbox.ColorWhite
			bg = termbox.ColorBlue
		}
		if item.Type == ITEM_TUNE {
			tune := hvsc.FilteredTunes[item.TuneIndex]
			if player.CurrentTune != nil && player.CurrentTune.Index == tune.Index {
				fg |= termbox.AttrBold
			}
			writeAt(0, 1+i, fmt.Sprintf(" %-32s %4s", tune.ListName(), tune.Year()), fg, bg)
		} else {
			folder := []rune(item.Name[1:])
			if len(folder) > 40 {
				folder = folder[len(folder)-40 : len(folder)]
				folder[0] = '…'
			}
			writeAt(0, 1+i, string(folder), fg|termbox.AttrBold, bg)
		}
	}
}

func writeAt(x, y int, value string, fg, bg termbox.Attribute) {
	i := 0
	for _, c := range value {
		termbox.SetCell(x+i, y, c, fg, bg)
		i++
	}
}
<commit_msg>Esc when playing just stops the tune.<commit_after>package ui

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/lhz/considerate/hvsc"
	"github.com/lhz/considerate/player"
	"github.com/nsf/termbox-go"
)

const (
	MODE_BROWSE = iota
	MODE_SEARCH
)

var w, h int
var list *List
var mode int
var debugInfo string
var searchTerm []rune
var searchCursorPos int
var quit = false

func Setup() {
	err := termbox.Init()
	if err != nil {
		log.Panicln(err)
	}

	w, h = termbox.Size()

	list = NewList(h - 2)

	termbox.SetInputMode(termbox.InputEsc | termbox.InputMouse)
	draw()
}

func Run() {
	defer termbox.Close()

	mode = MODE_BROWSE

	tickChan := time.NewTicker(1 * time.Second).C

	eventChan := make(chan termbox.Event)
	go func() {
		for {
			eventChan <- termbox.PollEvent()
		}
	}()

	for !quit {
		debugInfo = ""
		select {
		case ev := <-eventChan:
			switch ev.Type {
			case termbox.EventKey:
				keyEvent(ev)
				draw()
			case termbox.EventResize:
				resizeEvent(ev)
				draw()
			case termbox.EventError:
				panic(ev.Err)
			}
		case <-tickChan:
			if player.Playing {
				draw()
			}
		}
	}
}

func resizeEvent(ev termbox.Event) {
	w, h = ev.Width, ev.Height
	draw()
}

func keyEvent(ev termbox.Event) {
	if mode == MODE_SEARCH {
		keyEventSearch(ev)
		return
	}
	switch ev.Key {
	case termbox.KeyCtrlC, termbox.KeyEsc:
		if player.Playing {
			player.Stop()
		} else {
			quit = true
		}
	case termbox.KeyPgup:
		list.PrevPage()
	case termbox.KeyPgdn:
		list.NextPage()
	case termbox.KeyArrowUp:
		list.PrevTune()
	case termbox.KeyArrowDown:
		list.NextTune()
	case termbox.KeyArrowLeft:
		player.PrevSong()
	case termbox.KeyArrowRight:
		player.NextSong()
	case termbox.KeyEnter:
		player.Play(list.CurrentItem().TuneIndex, 1)
	case termbox.KeySpace:
		list.NextTune()
		player.Play(list.CurrentItem().TuneIndex, 1)
	case termbox.KeyDelete:
		player.Stop()
	case 0:
		if n := strings.IndexRune("0123456789", ev.Ch); n > 0 {
			if player.Playing && n <= player.CurrentTune.Header.Songs {
				player.PlaySub(n)
			}
			return
		}
		if ev.Ch == '/' {
			mode = MODE_SEARCH
		}
	default:
		debugInfo = string(ev.Key)
	}
}

func keyEventSearch(ev termbox.Event) {
	switch ev.Key {
	case termbox.KeyCtrlC, termbox.KeyEsc:
		mode = MODE_BROWSE
	case termbox.KeyArrowLeft:
		if searchCursorPos > 0 {
			searchCursorPos--
		}
	case termbox.KeyArrowRight:
		if searchCursorPos < len(searchTerm) {
			searchCursorPos++
		}
	case termbox.KeyEnter:
		if hvsc.Filter(string(searchTerm)) {
			list = NewList(h - 2)
			mode = MODE_BROWSE
		}
	case termbox.KeySpace:
		searchInsert(rune(' '))
	case 0:
		searchInsert(ev.Ch)
	case 0x7F:
		if searchCursorPos > 0 {
			searchTerm = append(searchTerm[0:searchCursorPos-1], searchTerm[searchCursorPos:len(searchTerm)]...)
			searchCursorPos--
		}
	default:
		debugInfo = string(ev.Key)
	}
}

func searchInsert(ch rune) {
	searchTerm = append(searchTerm, rune(' '))
	copy(searchTerm[searchCursorPos+1:], searchTerm[searchCursorPos:])
	searchTerm[searchCursorPos] = ch
	searchCursorPos++
}

func draw() {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	drawHeader()
	drawList()
	drawTuneInfo()
	drawFooter()
	termbox.Flush()
}

func drawHeader() {
	var header string
	bg := termbox.ColorBlack
	switch mode {
	case MODE_BROWSE:
		header = fmt.Sprintf("Browse: %s", string(searchTerm))
		termbox.HideCursor()
	case MODE_SEARCH:
		header = fmt.Sprintf("Search: %s", string(searchTerm))
		bg = termbox.ColorMagenta
		termbox.SetCursor(8+searchCursorPos, 0)
	}
	header = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len([]rune(header))), header, "")
	writeAt(0, 0, header, termbox.ColorWhite|termbox.AttrBold, bg)
}

func drawFooter() {
	footer := fmt.Sprintf("%d/%d  %s", list.CurrentItem().TuneIndex+1, hvsc.NumFilteredTunes, debugInfo)
	footer = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len([]rune(footer))), footer, "")
	writeAt(0, h-1, footer, termbox.ColorWhite, termbox.ColorBlack)
}

func drawTuneInfo() {
	if !player.Playing {
		return
	}
	tune := player.CurrentTune
	ox := 42
	oy := 2
	fg := termbox.ColorDefault
	bg := termbox.ColorDefault

	writeAt(ox, oy+0, fmt.Sprintf("Title:    %s", tune.Title()), fg, bg)
	writeAt(ox, oy+1, fmt.Sprintf("Author:   %s", tune.Header.Author), fg, bg)
	writeAt(ox, oy+2, fmt.Sprintf("Released: %s", tune.Header.Released), fg, bg)

	writeAt(ox, oy+4, fmt.Sprintf("Tune: %d/%d  Length: %s  Time: %s",
		player.CurrentSong, tune.Header.Songs, player.SongLength(), player.Elapsed()), fg, bg)

	for i, line := range tune.Info {
		if oy+6+i > h-2 {
			break
		}
		writeAt(ox, oy+6+i, line, fg, bg)
	}
}

func drawList() {
	for i, item := range list.CurrentPage() {
		fg := termbox.ColorDefault
		bg := termbox.ColorDefault
		if i == list.PagePos {
			fg = termbox.ColorWhite
			bg = termbox.ColorBlue
		}
		if item.Type == ITEM_TUNE {
			tune := hvsc.FilteredTunes[item.TuneIndex]
			if player.CurrentTune != nil && player.CurrentTune.Index == tune.Index {
				fg |= termbox.AttrBold
			}
			writeAt(0, 1+i, fmt.Sprintf(" %-32s %4s", tune.ListName(), tune.Year()), fg, bg)
		} else {
			folder := []rune(item.Name[1:])
			if len(folder) > 40 {
				folder = folder[len(folder)-40 : len(folder)]
				folder[0] = '…'
			}
			writeAt(0, 1+i, string(folder), fg|termbox.AttrBold, bg)
		}
	}
}

func writeAt(x, y int, value string, fg, bg termbox.Attribute) {
	i := 0
	for _, c := range value {
		termbox.SetCell(x+i, y, c, fg, bg)
		i++
	}
}
<|endoftext|>


<commit_before>package ui

import (
	"fmt"
	"log"

	"github.com/lhz/considerate/hvsc"
	"github.com/lhz/considerate/player"
	"github.com/nsf/termbox-go"
)

const (
	MODE_BROWSE = iota
	MODE_SEARCH
)

var w, h int
var listOffset, listPos, lh, ly int
var mode int
var searchTerm, debugInfo string
var quit = false

func Setup() {
	err := termbox.Init()
	if err != nil {
		log.Panicln(err)
	}

	w, h = termbox.Size()
	listOffset = 0
	listPos = 0
	ly = 1
	lh = h - 2

	termbox.SetInputMode(termbox.InputEsc | termbox.InputMouse)
	draw()
}

func Run() {
	defer termbox.Close()

	mode = MODE_BROWSE

	eventQueue := make(chan termbox.Event)
	go func() {
		for {
			eventQueue <- termbox.PollEvent()
		}
	}()

	for !quit {
		debugInfo = ""
		select {
		case ev := <-eventQueue:
			switch ev.Type {
			case termbox.EventKey:
				keyEvent(ev)
			case termbox.EventResize:
				resizeEvent(ev)
			case termbox.EventError:
				panic(ev.Err)
			}
		}
		draw()
	}
}

func resizeEvent(ev termbox.Event) {
	w, h = ev.Width, ev.Height
	draw()
}

func keyEvent(ev termbox.Event) {
	switch ev.Key {
	case termbox.KeyCtrlC, termbox.KeyEsc:
		stopTune()
		quit = true
	case termbox.KeyPgup:
		pageUp()
	case termbox.KeyPgdn:
		pageDown()
	case termbox.KeyArrowUp:
		moveUp()
	case termbox.KeyArrowDown:
		moveDown()
	case termbox.KeyEnter:
		switch mode {
		case MODE_BROWSE:
			selectTune()
		case MODE_SEARCH:
			hvsc.Filter(searchTerm)
			listOffset = 0
			listPos = 0
			mode = MODE_BROWSE
		}
	case termbox.KeyDelete:
		if mode == MODE_BROWSE {
			stopTune()
		}
	case 0:
		switch ev.Ch {
		case '/':
			switch mode {
			case MODE_BROWSE:
				mode = MODE_SEARCH
				searchTerm = ""
			}
		default:
			if mode == MODE_SEARCH {
				searchTerm = searchTerm + string(ev.Ch)
			}
		}
	case 0x7F:
		if mode == MODE_SEARCH {
			if len(searchTerm) > 0 {
				searchTerm = searchTerm[0 : len(searchTerm)-1]
			}
		}
	default:
		debugInfo = string(ev.Key)
	}
}

func moveUp() {
	listPos--
	if listPos < 0 {
		if listOffset > 0 {
			listPos = lh - 1
			pageUp()
		} else {
			listPos = 0
		}
	}
}

func moveDown() {
	if listOffset+listPos == hvsc.NumFilteredTunes-1 {
		return
	}
	listPos++
	if listPos >= lh {
		listPos = 0
		pageDown()
	}
}

func pageUp() {
	listOffset -= lh
	if listOffset < 0 {
		listOffset = 0
	}
}

func pageDown() {
	listOffset += lh
	if listOffset > hvsc.NumFilteredTunes-1 {
		listOffset -= lh
	}
}

func selectTune() {
	player.Play(listOffset+listPos, 1)
}

func stopTune() {
	player.Stop()
}

func draw() {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	drawHeader()
	drawList()
	drawTuneInfo()
	drawFooter()
	termbox.Flush()
}

func drawHeader() {
	var header string
	bg := termbox.ColorBlack
	switch mode {
	case MODE_BROWSE:
		header = "Browse"
		termbox.HideCursor()
	case MODE_SEARCH:
		header = fmt.Sprintf("Search: %s", searchTerm)
		//bg = termbox.ColorGreen
		termbox.SetCursor(8+len(searchTerm), 0)
	}
	header = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len(header)), header, "")
	writeAt(0, 0, header, termbox.ColorWhite, bg)
}

func drawFooter() {
	footer := fmt.Sprintf("%d/%d  %s", listOffset+listPos+1, hvsc.NumFilteredTunes, debugInfo)
	footer = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len(footer)), footer, "")
	writeAt(0, h-1, footer, termbox.ColorWhite, termbox.ColorBlack)
}

func drawTuneInfo() {
	if !player.Playing {
		return
	}
	tune := player.CurrentTune
	ox := w - 80
	oy := 2
	fg := termbox.ColorDefault
	bg := termbox.ColorDefault

	writeAt(ox, oy+0, fmt.Sprintf("Title:    %s", tune.Header.Name), fg, bg)
	writeAt(ox, oy+1, fmt.Sprintf("Author:   %s", tune.Header.Author), fg, bg)
	writeAt(ox, oy+2, fmt.Sprintf("Released: %s", tune.Header.Released), fg, bg)

	writeAt(ox, oy+4, fmt.Sprintf("Tune: %d/%d  Time: %s", tune.Header.StartSong, tune.Header.Songs, player.Elapsed()), fg, bg)
}

func drawList() {
	for y := 0; y < lh; y++ {
		if y+listOffset >= hvsc.NumFilteredTunes {
			break
		}
		tune := hvsc.FilteredTunes[y+listOffset]
		fg := termbox.ColorDefault
		bg := termbox.ColorDefault
		if player.CurrentTune != nil && player.CurrentTune.Index == tune.Index {
			fg |= termbox.AttrBold
		}
		if y == listPos {
			bg = termbox.ColorBlue
		}
		writeAt(0, ly+y, listName(tune), fg, bg)
	}
}

func listName(tune hvsc.SidTune) string {
	return fmt.Sprintf("%4s %s", tune.Year(), tune.Path[1:len(tune.Path)-4])
}

func writeAt(x, y int, value string, fg, bg termbox.Attribute) {
	i := 0
	for _, c := range value {
		termbox.SetCell(x+i, y, c, fg, bg)
		i++
	}
}
<commit_msg>Added ticker to update elapsed time every second.<commit_after>package ui

import (
	"fmt"
	"log"
	"time"

	"github.com/lhz/considerate/hvsc"
	"github.com/lhz/considerate/player"
	"github.com/nsf/termbox-go"
)

const (
	MODE_BROWSE = iota
	MODE_SEARCH
)

var w, h int
var listOffset, listPos, lh, ly int
var mode int
var searchTerm, debugInfo string
var quit = false

func Setup() {
	err := termbox.Init()
	if err != nil {
		log.Panicln(err)
	}

	w, h = termbox.Size()
	listOffset = 0
	listPos = 0
	ly = 1
	lh = h - 2

	termbox.SetInputMode(termbox.InputEsc | termbox.InputMouse)
	draw()
}

func Run() {
	defer termbox.Close()

	mode = MODE_BROWSE

	tickChan := time.NewTicker(1 * time.Second).C

	eventChan := make(chan termbox.Event)
	go func() {
		for {
			eventChan <- termbox.PollEvent()
		}
	}()

	for !quit {
		debugInfo = ""
		select {
		case ev := <-eventChan:
			switch ev.Type {
			case termbox.EventKey:
				keyEvent(ev)
			case termbox.EventResize:
				resizeEvent(ev)
			case termbox.EventError:
				panic(ev.Err)
			}
		case <-tickChan:
		}
		draw()
	}
}

func resizeEvent(ev termbox.Event) {
	w, h = ev.Width, ev.Height
	draw()
}

func keyEvent(ev termbox.Event) {
	switch ev.Key {
	case termbox.KeyCtrlC, termbox.KeyEsc:
		stopTune()
		quit = true
	case termbox.KeyPgup:
		pageUp()
	case termbox.KeyPgdn:
		pageDown()
	case termbox.KeyArrowUp:
		moveUp()
	case termbox.KeyArrowDown:
		moveDown()
	case termbox.KeyEnter:
		switch mode {
		case MODE_BROWSE:
			selectTune()
		case MODE_SEARCH:
			hvsc.Filter(searchTerm)
			listOffset = 0
			listPos = 0
			mode = MODE_BROWSE
		}
	case termbox.KeyDelete:
		if mode == MODE_BROWSE {
			stopTune()
		}
	case 0:
		switch ev.Ch {
		case '/':
			switch mode {
			case MODE_BROWSE:
				mode = MODE_SEARCH
				searchTerm = ""
			}
		default:
			if mode == MODE_SEARCH {
				searchTerm = searchTerm + string(ev.Ch)
			}
		}
	case 0x7F:
		if mode == MODE_SEARCH {
			if len(searchTerm) > 0 {
				searchTerm = searchTerm[0 : len(searchTerm)-1]
			}
		}
	default:
		debugInfo = string(ev.Key)
	}
}

func moveUp() {
	listPos--
	if listPos < 0 {
		if listOffset > 0 {
			listPos = lh - 1
			pageUp()
		} else {
			listPos = 0
		}
	}
}

func moveDown() {
	if listOffset+listPos == hvsc.NumFilteredTunes-1 {
		return
	}
	listPos++
	if listPos >= lh {
		listPos = 0
		pageDown()
	}
}

func pageUp() {
	listOffset -= lh
	if listOffset < 0 {
		listOffset = 0
	}
}

func pageDown() {
	listOffset += lh
	if listOffset > hvsc.NumFilteredTunes-1 {
		listOffset -= lh
	}
}

func selectTune() {
	player.Play(listOffset+listPos, 1)
}

func stopTune() {
	player.Stop()
}

func draw() {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	drawHeader()
	drawList()
	drawTuneInfo()
	drawFooter()
	termbox.Flush()
}

func drawHeader() {
	var header string
	bg := termbox.ColorBlack
	switch mode {
	case MODE_BROWSE:
		header = "Browse"
		termbox.HideCursor()
	case MODE_SEARCH:
		header = fmt.Sprintf("Search: %s", searchTerm)
		//bg = termbox.ColorGreen
		termbox.SetCursor(8+len(searchTerm), 0)
	}
	header = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len(header)), header, "")
	writeAt(0, 0, header, termbox.ColorWhite, bg)
}

func drawFooter() {
	footer := fmt.Sprintf("%d/%d  %s", listOffset+listPos+1, hvsc.NumFilteredTunes, debugInfo)
	footer = fmt.Sprintf(fmt.Sprintf("%%s%%%ds", w-len(footer)), footer, "")
	writeAt(0, h-1, footer, termbox.ColorWhite, termbox.ColorBlack)
}

func drawTuneInfo() {
	if !player.Playing {
		return
	}
	tune := player.CurrentTune
	ox := w - 80
	oy := 2
	fg := termbox.ColorDefault
	bg := termbox.ColorDefault

	writeAt(ox, oy+0, fmt.Sprintf("Title:    %s", tune.Header.Name), fg, bg)
	writeAt(ox, oy+1, fmt.Sprintf("Author:   %s", tune.Header.Author), fg, bg)
	writeAt(ox, oy+2, fmt.Sprintf("Released: %s", tune.Header.Released), fg, bg)

	writeAt(ox, oy+4, fmt.Sprintf("Tune: %d/%d  Time: %s", tune.Header.StartSong, tune.Header.Songs, player.Elapsed()), fg, bg)
}

func drawList() {
	for y := 0; y < lh; y++ {
		if y+listOffset >= hvsc.NumFilteredTunes {
			break
		}
		tune := hvsc.FilteredTunes[y+listOffset]
		fg := termbox.ColorDefault
		bg := termbox.ColorDefault
		if player.CurrentTune != nil && player.CurrentTune.Index == tune.Index {
			fg |= termbox.AttrBold
		}
		if y == listPos {
			bg = termbox.ColorBlue
		}
		writeAt(0, ly+y, listName(tune), fg, bg)
	}
}

func listName(tune hvsc.SidTune) string {
	return fmt.Sprintf("%4s %s", tune.Year(), tune.Path[1:len(tune.Path)-4])
}

func writeAt(x, y int, value string, fg, bg termbox.Attribute) {
	i := 0
	for _, c := range value {
		termbox.SetCell(x+i, y, c, fg, bg)
		i++
	}
}
<|endoftext|>


<commit_before><commit_msg>Log S3 destination<commit_after><|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Package unzipit allows you to easily unpack *.tar.gz, *.tar.bzip2, *.zip and *.tar files.
// There are not CGO involved nor hard dependencies of any type.
package unzipit

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
)

var (
	magicZIP  = []byte{0x50, 0x4b, 0x03, 0x04}
	magicGZ   = []byte{0x1f, 0x8b}
	magicBZIP = []byte{0x42, 0x5a}
	magicTAR  = []byte{0x75, 0x73, 0x74, 0x61, 0x72} // at offset 257
)

// Check whether a file has the magic number for tar, gzip, bzip2 or zip files
//
// Note that this function does not advance the Reader.
//
// 50 4b 03 04 for pkzip format
// 1f 8b for .gz
// 42 5a for bzip
// 75 73 74 61 72 at offset 257 for tar files
func magicNumber(reader *bufio.Reader, offset int) (string, error) {
	headerBytes, err := reader.Peek(offset + 5)
	if err != nil {
		return "", err
	}

	magic := headerBytes[offset : offset+5]

	if bytes.Equal(magicTAR, magic) {
		return "tar", nil
	}

	if bytes.Equal(magicZIP, magic[0:4]) {
		return "zip", nil
	}

	if bytes.Equal(magicGZ, magic[0:2]) {
		return "gzip", nil
	} else if bytes.Equal(magicBZIP, magic[0:2]) {
		return "bzip", nil
	}

	return "", nil
}

// Unpack unpacks a compressed and archived file and places result output in destination
// path.
//
// File formats supported are:
//   - .tar.gz
//   - .tar.bzip2
//   - .zip
//   - .tar
//
// If it cannot recognize the file format, it will save the file, as is, to the
// destination path.
func Unpack(file *os.File, destPath string) (string, error) {
	if file == nil {
		return "", errors.New("You must provide a valid file to unpack")
	}

	var err error
	if destPath == "" {
		destPath, err = ioutil.TempDir(os.TempDir(), "unpackit-")
		if err != nil {
			return "", err
		}
	}

	// Makes sure despPath exists
	os.MkdirAll(destPath, 0740)

	r := bufio.NewReader(file)
	return UnpackStream(r, destPath)
}

// UnpackStream unpacks a compressed stream. Note that if the stream is a using ZIP
// compression (but only ZIP compression), it's going to get buffered in its entirety
// to memory prior to decompression.
func UnpackStream(reader io.Reader, destPath string) (string, error) {
	r := bufio.NewReader(reader)

	// Reads magic number from the stream so we can better determine how to proceed
	ftype, err := magicNumber(r, 0)
	if err != nil {
		return "", err
	}

	var decompressingReader *bufio.Reader
	switch ftype {
	case "gzip":
		decompressingReader, err = GunzipStream(r)
		if err != nil {
			return "", err
		}
	case "bzip":
		decompressingReader, err = Bunzip2Stream(r)
		if err != nil {
			return "", err
		}
	case "zip":
		// Like TAR, ZIP is also an archiving format, therefore we can just return
		// after it finishes
		return UnzipStream(r, destPath)
	default:
		decompressingReader = r
	}

	// Check magic number in offset 257 too see if this is also a TAR file
	ftype, err = magicNumber(decompressingReader, 257)
	if ftype == "tar" {
		return Untar(decompressingReader, destPath)
	}

	// If it's not a TAR archive then save it to disk as is.
	destRawFile := filepath.Join(destPath, sanitize(path.Base("tarstream")))

	// Creates destination file
	destFile, err := os.Create(destRawFile)
	if err != nil {
		return "", err
	}
	defer destFile.Close()

	// Copies data to destination file
	if _, err := io.Copy(destFile, decompressingReader); err != nil {
		return "", err
	}

	return destPath, nil
}

// Bunzip2 decompresses a bzip2 file and returns the decompressed stream
func Bunzip2(file *os.File) (*bufio.Reader, error) {
	freader := bufio.NewReader(file)
	bzip2Reader, err := Bunzip2Stream(freader)
	if err != nil {
		return nil, err
	}

	return bufio.NewReader(bzip2Reader), nil
}

// Bunzip2Stream unpacks a bzip2 stream
func Bunzip2Stream(reader io.Reader) (*bufio.Reader, error) {
	return bufio.NewReader(bzip2.NewReader(reader)), nil
}

// Gunzip decompresses a gzip file and returns the decompressed stream
func Gunzip(file *os.File) (*bufio.Reader, error) {
	freader := bufio.NewReader(file)
	gunzipReader, err := GunzipStream(freader)
	if err != nil {
		return nil, err
	}

	return bufio.NewReader(gunzipReader), nil
}

// GunzipStream unpacks a gzipped stream
func GunzipStream(reader io.Reader) (*bufio.Reader, error) {
	var decompressingReader *gzip.Reader
	var err error
	if decompressingReader, err = gzip.NewReader(reader); err != nil {
		return nil, err
	}

	return bufio.NewReader(decompressingReader), nil
}

// Unzip decompresses and unarchives a ZIP archive, returning the final path or an error
func Unzip(file *os.File, destPath string) (string, error) {
	fstat, err := file.Stat()
	if err != nil {
		return "", err
	}

	zr, err := zip.NewReader(file, fstat.Size())
	if err != nil {
		return "", err
	}

	return unpackZip(zr, destPath)
}

// UnzipStream unpacks a ZIP stream. Because of the nature of the ZIP format,
// the stream is copied to memory before decompression.
func UnzipStream(r io.Reader, destPath string) (string, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}

	memReader := bytes.NewReader(data)
	zr, err := zip.NewReader(memReader, int64(len(data)))
	if err != nil {
		return "", err
	}

	return unpackZip(zr, destPath)
}

func unpackZip(zr *zip.Reader, destPath string) (string, error) {
	// Iterate through the files in the archive,
	// printing some of their contents.
	pathSeparator := string(os.PathSeparator)
	for _, f := range zr.File {
		rc, err := f.Open()
		if err != nil {
			return "", err
		}
		defer rc.Close()

		filePath := sanitize(f.Name)
		sepInd := strings.LastIndex(filePath, pathSeparator)

		// If the file is a subdirectory, it creates it before attempting to
		// create the actual file
		if sepInd > -1 {
			directory := filePath[0:sepInd]
			os.MkdirAll(filepath.Join(destPath, directory), 0740)
		}

		file, err := os.Create(filepath.Join(destPath, filePath))
		if err != nil {
			return "", err
		}
		defer file.Close()

		if _, err := io.CopyN(file, rc, int64(f.UncompressedSize64)); err != nil {
			return "", err
		}
		defer rc.Close()
	}
	return destPath, nil
}

// Untar unarchives a TAR archive and returns the final destination path or an error
func Untar(data io.Reader, destPath string) (string, error) {
	// Makes sure destPath exists
	os.MkdirAll(destPath, 0740)

	tr := tar.NewReader(data)

	// Iterate through the files in the archive.
	rootdir := destPath
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}

		if err != nil {
			return rootdir, err
		}

		if hdr.FileInfo().IsDir() {
			d := filepath.Join(destPath, sanitize(hdr.Name))
			if rootdir == destPath {
				rootdir = d
			}
			os.MkdirAll(d, os.FileMode(hdr.Mode))
			continue
		}

		path := filepath.Join(destPath, sanitize(hdr.Name))
		parentDir, _ := filepath.Split(path)

		err = os.MkdirAll(parentDir, 0740)
		if err != nil {
			return rootdir, err
		}

		file, err := os.Create(path)
		if err != nil {
			return rootdir, err
		}

		defer file.Close()

		if _, err := io.Copy(file, tr); err != nil {
			return rootdir, err
		}
	}

	return rootdir, nil
}

// Sanitizes name to avoid overwriting sensitive system files when unarchiving
func sanitize(name string) string {
	// Gets rid of volume drive label in Windows
	if len(name) > 1 && name[1] == ':' && runtime.GOOS == "windows" {
		name = name[2:]
	}

	name = filepath.Clean(name)
	name = filepath.ToSlash(name)
	for strings.HasPrefix(name, "../") {
		name = name[3:]
	}
	return name
}
<commit_msg>Cleans up code a little bit<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Package unzipit allows you to easily unpack *.tar.gz, *.tar.bzip2, *.zip and *.tar files.
// There are not CGO involved nor hard dependencies of any type.
package unzipit

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
)

var (
	magicZIP  = []byte{0x50, 0x4b, 0x03, 0x04}
	magicGZ   = []byte{0x1f, 0x8b}
	magicBZIP = []byte{0x42, 0x5a}
	magicTAR  = []byte{0x75, 0x73, 0x74, 0x61, 0x72} // at offset 257
)

// Check whether a file has the magic number for tar, gzip, bzip2 or zip files
//
// Note that this function does not advance the Reader.
//
// 50 4b 03 04 for pkzip format
// 1f 8b for .gz
// 42 5a for bzip
// 75 73 74 61 72 at offset 257 for tar files
func magicNumber(reader *bufio.Reader, offset int) (string, error) {
	headerBytes, err := reader.Peek(offset + 5)
	if err != nil {
		return "", err
	}

	magic := headerBytes[offset : offset+5]

	if bytes.Equal(magicTAR, magic) {
		return "tar", nil
	}

	if bytes.Equal(magicZIP, magic[0:4]) {
		return "zip", nil
	}

	if bytes.Equal(magicGZ, magic[0:2]) {
		return "gzip", nil
	} else if bytes.Equal(magicBZIP, magic[0:2]) {
		return "bzip", nil
	}

	return "", nil
}

// Unpack unpacks a compressed and archived file and places result output in destination
// path.
//
// File formats supported are:
//   - .tar.gz
//   - .tar.bzip2
//   - .zip
//   - .tar
//
// If it cannot recognize the file format, it will save the file, as is, to the
// destination path.
func Unpack(file *os.File, destPath string) (string, error) {
	if file == nil {
		return "", errors.New("You must provide a valid file to unpack")
	}

	var err error
	if destPath == "" {
		destPath, err = ioutil.TempDir(os.TempDir(), "unpackit-")
		if err != nil {
			return "", err
		}
	}

	// Makes sure despPath exists
	os.MkdirAll(destPath, 0740)

	r := bufio.NewReader(file)
	return UnpackStream(r, destPath)
}

// UnpackStream unpacks a compressed stream. Note that if the stream is a using ZIP
// compression (but only ZIP compression), it's going to get buffered in its entirety
// to memory prior to decompression.
func UnpackStream(reader io.Reader, destPath string) (string, error) {
	r := bufio.NewReader(reader)

	// Reads magic number from the stream so we can better determine how to proceed
	ftype, err := magicNumber(r, 0)
	if err != nil {
		return "", err
	}

	var decompressingReader *bufio.Reader
	switch ftype {
	case "gzip":
		decompressingReader, err = GunzipStream(r)
		if err != nil {
			return "", err
		}
	case "bzip":
		decompressingReader, err = Bunzip2Stream(r)
		if err != nil {
			return "", err
		}
	case "zip":
		// Like TAR, ZIP is also an archiving format, therefore we can just return
		// after it finishes
		return UnzipStream(r, destPath)
	default:
		decompressingReader = r
	}

	// Check magic number in offset 257 too see if this is also a TAR file
	ftype, err = magicNumber(decompressingReader, 257)
	if ftype == "tar" {
		return Untar(decompressingReader, destPath)
	}

	// If it's not a TAR archive then save it to disk as is.
	destRawFile := filepath.Join(destPath, sanitize(path.Base("tarstream")))

	// Creates destination file
	destFile, err := os.Create(destRawFile)
	if err != nil {
		return "", err
	}
	defer destFile.Close()

	// Copies data to destination file
	if _, err := io.Copy(destFile, decompressingReader); err != nil {
		return "", err
	}

	return destPath, nil
}

// Bunzip2 decompresses a bzip2 file and returns the decompressed stream
func Bunzip2(file *os.File) (*bufio.Reader, error) {
	freader := bufio.NewReader(file)
	bzip2Reader, err := Bunzip2Stream(freader)
	if err != nil {
		return nil, err
	}

	return bufio.NewReader(bzip2Reader), nil
}

// Bunzip2Stream unpacks a bzip2 stream
func Bunzip2Stream(reader io.Reader) (*bufio.Reader, error) {
	return bufio.NewReader(bzip2.NewReader(reader)), nil
}

// Gunzip decompresses a gzip file and returns the decompressed stream
func Gunzip(file *os.File) (*bufio.Reader, error) {
	freader := bufio.NewReader(file)
	gunzipReader, err := GunzipStream(freader)
	if err != nil {
		return nil, err
	}

	return bufio.NewReader(gunzipReader), nil
}

// GunzipStream unpacks a gzipped stream
func GunzipStream(reader io.Reader) (*bufio.Reader, error) {
	var decompressingReader *gzip.Reader
	var err error
	if decompressingReader, err = gzip.NewReader(reader); err != nil {
		return nil, err
	}

	return bufio.NewReader(decompressingReader), nil
}

// Unzip decompresses and unarchives a ZIP archive, returning the final path or an error
func Unzip(file *os.File, destPath string) (string, error) {
	fstat, err := file.Stat()
	if err != nil {
		return "", err
	}

	zr, err := zip.NewReader(file, fstat.Size())
	if err != nil {
		return "", err
	}

	return unpackZip(zr, destPath)
}

// UnzipStream unpacks a ZIP stream. Because of the nature of the ZIP format,
// the stream is copied to memory before decompression.
func UnzipStream(r io.Reader, destPath string) (string, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}

	memReader := bytes.NewReader(data)
	zr, err := zip.NewReader(memReader, int64(len(data)))
	if err != nil {
		return "", err
	}

	return unpackZip(zr, destPath)
}

func unpackZip(zr *zip.Reader, destPath string) (string, error) {
	// Iterate through the files in the archive,
	// printing some of their contents.
	pathSeparator := string(os.PathSeparator)
	for _, f := range zr.File {
		rc, err := f.Open()
		if err != nil {
			return "", err
		}
		defer rc.Close()

		filePath := sanitize(f.Name)
		sepInd := strings.LastIndex(filePath, pathSeparator)

		// If the file is a subdirectory, it creates it before attempting to
		// create the actual file
		if sepInd > -1 {
			directory := filePath[0:sepInd]
			os.MkdirAll(filepath.Join(destPath, directory), 0740)
		}

		file, err := os.Create(filepath.Join(destPath, filePath))
		if err != nil {
			return "", err
		}
		defer file.Close()

		if _, err := io.CopyN(file, rc, int64(f.UncompressedSize64)); err != nil {
			return "", err
		}
		defer rc.Close()
	}
	return destPath, nil
}

// Untar unarchives a TAR archive and returns the final destination path or an error
func Untar(data io.Reader, destPath string) (string, error) {
	// Makes sure destPath exists
	os.MkdirAll(destPath, 0740)

	tr := tar.NewReader(data)

	// Iterate through the files in the archive.
	rootdir := destPath
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}

		if err != nil {
			return rootdir, err
		}

		fp := filepath.Join(destPath, sanitize(hdr.Name))
		if hdr.FileInfo().IsDir() {
			if rootdir == destPath {
				rootdir = fp
			}
			os.MkdirAll(fp, os.FileMode(hdr.Mode))
			continue
		}

		parentDir, _ := filepath.Split(fp)
		err = os.MkdirAll(parentDir, 0740)
		if err != nil {
			return rootdir, err
		}

		file, err := os.Create(fp)
		if err != nil {
			return rootdir, err
		}

		defer file.Close()

		if _, err := io.Copy(file, tr); err != nil {
			return rootdir, err
		}
	}

	return rootdir, nil
}

// Sanitizes name to avoid overwriting sensitive system files when unarchiving
func sanitize(name string) string {
	// Gets rid of volume drive label in Windows
	if len(name) > 1 && name[1] == ':' && runtime.GOOS == "windows" {
		name = name[2:]
	}

	name = filepath.Clean(name)
	name = filepath.ToSlash(name)
	for strings.HasPrefix(name, "../") {
		name = name[3:]
	}
	return name
}
<|endoftext|>