oneterm/backend/session/parser.go

package session

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/veops/go-ansiterm"
	mysql "github.com/veops/oneterm/db"
	"github.com/veops/oneterm/logger"
	"github.com/veops/oneterm/model"
	"go.uber.org/zap"
)

var (
	enterMarks = [][]byte{
		[]byte("\x1b[?1049h"),
		[]byte("\x1b[?1048h"),
		[]byte("\x1b[?1047h"),
		[]byte("\x1b[?47h"),
		[]byte("\x1b[?25l"),
	}

	exitMarks = [][]byte{
		[]byte("\x1b[?1049l"),
		[]byte("\x1b[?1048l"),
		[]byte("\x1b[?1047l"),
		[]byte("\x1b[?47l"),
		[]byte("\x1b[?25h"),
	}

	screenMarks = [][]byte{
		{0x1b, 0x5b, 0x4b, 0x0d, 0x0a},
		{0x1b, 0x5b, 0x34, 0x6c},
	}
)

func NewParser(sessionId string, w, h int) *Parser {
	screen := ansiterm.NewScreen(w, h)
	stream := ansiterm.InitByteStream(screen, false)
	stream.Attach(screen)
	p := &Parser{
		OutputStream: stream,
		isEdit:       false,
		first:        true,
	}
	return p
}

type Parser struct {
	OutputStream *ansiterm.ByteStream
	Input        []byte
	Output       []byte
	SessionId    string
	Cmds         []*model.Command
	first        bool
	prompt       string
	isEdit       bool
	lastCmd      string
	lastRes      string
}

func (p *Parser) AddInput(bs []byte) (cmd string, forbidden bool) {
	if p.first {
		p.GetOutput()
		p.first = false
	}
	p.Input = append(p.Input, bs...)
	if !bytes.HasSuffix(p.Input, []byte("\r")) {
		return
	}
	cmd = p.GetCmd()
	fmt.Println("----------------------cmd", cmd)
	p.Reset()
	filter := ""
	if filter, forbidden = p.IsForbidden(cmd); forbidden {
		cmd = filter
		return
	}
	p.lastCmd = cmd
	p.WriteDb()
	return
}

func (p *Parser) IsForbidden(cmd string) (string, bool) {
	for _, c := range p.Cmds {
		if c.IsRe {
			if c.Re.MatchString(cmd) {
				return fmt.Sprintf("Regex: %s", c.Cmd), true
			}
		} else {
			if strings.Contains(cmd, c.Cmd) {
				return c.Cmd, true
			}
		}
	}
	return "", false
}

func (p *Parser) WriteDb() {
	if p.lastCmd == "" {
		return
	}
	m := &model.SessionCmd{
		SessionId: p.SessionId,
		Cmd:       p.lastCmd,
		Result:    p.lastRes,
	}
	err := mysql.DB.Model(m).Create(m).Error
	if err != nil {
		logger.L().Error("write session cmd failed", zap.Error(err))
	}
}

func (p *Parser) AddOutput(bs []byte) {
	fmt.Println("-----------out", string(bs))
	if !p.isEdit {
		p.Output = append(p.Output, bs...)
		end := bytes.LastIndex(p.Output, []byte("\r"))
		if end < 0 {
			return
		}
		begin := end - 1
		for ; begin > 0; begin-- {
			if p.Output[begin] == '\r' {
				break
			}
		}
		if begin+1 > end-1 {
			return
		}
		p.prompt = string(p.Output[begin+1 : end-1])
	}
}

func (p *Parser) GetCmd() string {
	s := p.GetOutput()
	// TODO: some promot may change with its dir
	fmt.Println("============", s)
	fmt.Println("============", p.prompt)
	return strings.TrimPrefix(s, p.prompt)
}

func (p *Parser) Resize(w, h int) {
	p.OutputStream.Listener.Resize(w, h)
}

func (p *Parser) Reset() {
	p.OutputStream.Listener.Reset()
	p.Output = nil
	p.Input = nil
}

func (p *Parser) GetOutput() string {
	p.OutputStream.Feed(p.Output)

	res := parseOutput(p.OutputStream.Listener.Display())
	if len(res) == 0 {
		return ""
	}
	p.lastRes = res[len(res)-1]
	return p.lastRes
}

func parseOutput(data []string) (output []string) {
	for _, line := range data {
		if strings.TrimSpace(line) != "" {
			output = append(output, line)
		}
	}
	return output
}

func (p *Parser) State(b []byte) bool {
	if !p.isEdit && IsEditEnterMode(b) {
		if !isNewScreen(b) {
			p.isEdit = true
		}
	}
	if p.isEdit && IsEditExitMode(b) {
		p.Reset()
		p.isEdit = false
	}
	return p.isEdit
}

func isNewScreen(p []byte) bool {
	return matchMark(p, screenMarks)
}

func IsEditEnterMode(p []byte) bool {
	return matchMark(p, enterMarks)
}

func IsEditExitMode(p []byte) bool {
	return matchMark(p, exitMarks)
}

func matchMark(p []byte, marks [][]byte) bool {
	for _, item := range marks {
		if bytes.Contains(p, item) {
			return true
		}
	}
	return false
}