1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182
|
package ssh
import (
"fmt"
"github.com/viant/toolbox"
"io/ioutil"
"os"
"path"
"sort"
"strings"
)
//ReplayCommand represent a replay command
type ReplayCommand struct {
Stdin string
Index int
Stdout []string
Error string
}
//replayCommands represnets command grouped by stdin
type ReplayCommands struct {
Commands map[string]*ReplayCommand
Keys []string
BaseDir string
}
//Register register stdin and corresponding stdout conversation
func (c *ReplayCommands) Register(stdin, stdout string) {
if _, has := c.Commands[stdin]; !has {
c.Commands[stdin] = &ReplayCommand{
Stdin: stdin,
Stdout: make([]string, 0),
}
c.Keys = append(c.Keys, stdin)
}
c.Commands[stdin].Stdout = append(c.Commands[stdin].Stdout, stdout)
}
//return stdout pointed by index and increases index or empty string if exhausted
func (c *ReplayCommands) Next(stdin string) string {
var stdout = c.Commands[stdin].Stdout
var index = c.Commands[stdin].Index
if index < len(stdout) {
c.Commands[stdin].Index++
return stdout[index]
}
return ""
}
func (c *ReplayCommands) Enable(source interface{}) (err error) {
switch value := source.(type) {
case *service:
value.replayCommands = c
value.recordSession = true
case *multiCommandSession:
value.replayCommands = c
value.recordSession = true
default:
err = fmt.Errorf("unsupported type: %T", source)
}
return err
}
func (c *ReplayCommands) Disable(source interface{}) (err error) {
switch value := source.(type) {
case *service:
value.replayCommands = nil
value.recordSession = false
case *multiCommandSession:
value.replayCommands = nil
value.recordSession = false
default:
err = fmt.Errorf("unsupported type: %T", source)
}
return err
}
//Store stores replay command in the base directory
func (c *ReplayCommands) Store() error {
err := toolbox.CreateDirIfNotExist(c.BaseDir)
if err != nil {
return err
}
for i, key := range c.Keys {
var command = c.Commands[key]
var filenamePrefix = path.Join(c.BaseDir, fmt.Sprintf("%03d", i+1))
var stdinFilename = filenamePrefix + "_000.stdin"
err := ioutil.WriteFile(stdinFilename, []byte(command.Stdin), 0644)
if err != nil {
return err
}
for j, stdout := range command.Stdout {
var stdoutFilename = fmt.Sprintf("%v_%03d.stdout", filenamePrefix, j+1)
err := ioutil.WriteFile(stdoutFilename, []byte(stdout), 0644)
if err != nil {
return err
}
}
}
return nil
}
//Load loads replay command from base directory
func (c *ReplayCommands) Load() error {
parent, err := os.Open(c.BaseDir)
if err != nil {
return err
}
files, err := parent.Readdir(1000)
if err != nil {
return err
}
var stdinMap = make(map[string]string)
var stdoutMap = make(map[string]string)
for _, candidate := range files {
ext := path.Ext(candidate.Name())
var contentMap map[string]string
if ext == ".stdin" {
contentMap = stdinMap
} else if ext == ".stdout" {
contentMap = stdoutMap
} else {
continue
}
var filename = path.Join(c.BaseDir, candidate.Name())
content, err := ioutil.ReadFile(filename)
if err != nil {
return nil
}
contentMap[candidate.Name()] = string(content)
}
for key, stdin := range stdinMap {
var prefix = key[:len(key)-10]
var candidateKeys = toolbox.MapKeysToStringSlice(stdoutMap)
sort.Strings(candidateKeys)
for _, candidateKey := range candidateKeys {
if strings.HasPrefix(candidateKey, prefix) {
stdout := stdoutMap[candidateKey]
c.Register(stdin, stdout)
}
}
}
return nil
}
//Shell returns command shell
func (c *ReplayCommands) Shell() string {
for _, candidate := range c.Commands {
if strings.HasPrefix(candidate.Stdin, "PS1=") && len(candidate.Stdout) > 0 {
return candidate.Stdout[0]
}
}
return ""
}
//System returns system name
func (c *ReplayCommands) System() string {
for _, candidate := range c.Commands {
if strings.HasPrefix(candidate.Stdin, "uname -s") && len(candidate.Stdout) > 0 {
return strings.ToLower(candidate.Stdout[0])
}
}
return ""
}
//NewReplayCommands create a new replay commands or error if provided basedir does not exists and can not be created
func NewReplayCommands(basedir string) (*ReplayCommands, error) {
if !toolbox.FileExists(basedir) {
err := os.MkdirAll(basedir, 0744)
if err != nil {
return nil, err
}
}
return &ReplayCommands{
Commands: make(map[string]*ReplayCommand),
Keys: make([]string, 0),
BaseDir: basedir,
}, nil
}
|