shell.go

// Copyright (c) 2021 Blacknon. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

package sshlib

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	"github.com/lunixbochs/vtclean"
	"golang.org/x/crypto/ssh"
	terminal "golang.org/x/term"
)

// Shell connect login shell over ssh.
func (c *Connect) Shell(session *ssh.Session) (err error) {
	// Input terminal Make raw
	var fd int
	if c.PtyRelayTty != nil {
		fd = int(c.PtyRelayTty.Fd())
	} else {
		fd = int(os.Stdin.Fd())
	}

	state, err := terminal.MakeRaw(fd)
	if err != nil {
		return
	}
	defer terminal.Restore(fd, state)

	// setup
	err = c.setupShell(session)
	if err != nil {
		return
	}

	// set tty
	if c.PtyRelayTty != nil {
		session.Stdin = c.PtyRelayTty
		session.Stdout = c.PtyRelayTty
		session.Stderr = c.PtyRelayTty
	}

	// Start shell
	err = session.Shell()
	if err != nil {
		return
	}

	// keep alive packet
	go c.SendKeepAlive(session)

	// if tty is set, get signal winch
	if c.PtyRelayTty != nil {
		go c.ChangeWinSize(session)
	}

	err = session.Wait()
	if err != nil {
		return
	}

	return
}

// Shell connect command shell over ssh.
// Used to start a shell with a specified command.
func (c *Connect) CmdShell(session *ssh.Session, command string) (err error) {
	// Input terminal Make raw
	var fd int
	if c.PtyRelayTty != nil {
		fd = int(c.PtyRelayTty.Fd())
	} else {
		fd = int(os.Stdin.Fd())
	}

	state, err := terminal.MakeRaw(fd)
	if err != nil {
		return
	}
	defer terminal.Restore(fd, state)

	// setup
	err = c.setupShell(session)
	if err != nil {
		return
	}

	// set tty
	if c.PtyRelayTty != nil {
		session.Stdin = c.PtyRelayTty
		session.Stdout = c.PtyRelayTty
		session.Stderr = c.PtyRelayTty
	}

	// Start shell
	err = session.Start(command)
	if err != nil {
		return
	}

	// keep alive packet
	go c.SendKeepAlive(session)

	err = session.Wait()
	if err != nil {
		return
	}

	return
}

func (c *Connect) ChangeWinSize(session *ssh.Session) {
	// Get terminal window size
	var fd int
	if c.PtyRelayTty != nil {
		fd = int(c.PtyRelayTty.Fd())
	} else {
		fd = int(os.Stdout.Fd())
	}
	width, height, err := terminal.GetSize(fd)
	if err != nil {
		return
	}

	// Send window size
	session.WindowChange(height, width)
}

// SetLog set up terminal log logging.
// This only happens in Connect.Shell().
func (c *Connect) SetLog(path string, timestamp bool) {
	c.logging = true
	c.logFile = path
	c.logTimestamp = timestamp
}

func (c *Connect) SetLogWithRemoveAnsiCode(path string, timestamp bool) {
	c.logging = true
	c.logFile = path
	c.logTimestamp = timestamp
	c.logRemoveAnsiCode = true
}

// logger is logging terminal log to c.logFile
func (c *Connect) logger(session *ssh.Session) (err error) {
	logfile, err := os.OpenFile(c.logFile, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0600)
	if err != nil {
		return
	}

	if !c.logTimestamp && !c.logRemoveAnsiCode {
		session.Stdout = io.MultiWriter(session.Stdout, logfile)
		session.Stderr = io.MultiWriter(session.Stderr, logfile)
	} else {
		buf := new(bytes.Buffer)
		session.Stdout = io.MultiWriter(session.Stdout, buf)
		session.Stderr = io.MultiWriter(session.Stderr, buf)

		go func() {
			preLine := []byte{}
			for {
				if buf.Len() > 0 {
					// get line
					line, err := buf.ReadBytes('\n')

					if err == io.EOF {
						preLine = append(preLine, line...)
						continue
					} else {
						printLine := string(append(preLine, line...))

						if c.logTimestamp {
							timestamp := time.Now().Format("2006/01/02 15:04:05 ") // yyyy/mm/dd HH:MM:SS
							printLine = timestamp + printLine
						}

						// remove ansi code.
						if c.logRemoveAnsiCode {
							// NOTE:
							//     In vtclean.Clean, the beginning of the line is deleted for some reason.
							//     for that reason, one character add at line head.
							printLine = "." + printLine
							printLine = vtclean.Clean(printLine, false)
						}

						fmt.Fprintf(logfile, printLine)
						preLine = []byte{}
					}
				} else {
					time.Sleep(10 * time.Millisecond)
				}
			}
		}()
	}

	return err
}

func (c *Connect) setupShell(session *ssh.Session) (err error) {
	// set FD
	stdin := GetStdin()
	session.Stdin = stdin
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr

	// Logging
	if c.logging {
		err = c.logger(session)
		if err != nil {
			log.Println(err)
		}
	}
	err = nil

	// Request tty
	err = RequestTty(session)
	if err != nil {
		return err
	}

	// x11 forwarding
	if c.ForwardX11 {
		err = c.X11Forward(session)
		if err != nil {
			log.Println(err)
		}
	}
	err = nil

	// ssh agent forwarding
	if c.ForwardAgent {
		c.ForwardSshAgent(session)
	}

	return
}