github.com/zmwangx/ets@v0.2.2-0.20201107170110-98f3846f5ca3/main.go (about) 1 package main 2 3 import ( 4 "bufio" 5 "bytes" 6 "fmt" 7 "io" 8 "log" 9 "os" 10 "os/exec" 11 "os/signal" 12 "regexp" 13 "syscall" 14 "time" 15 16 "github.com/creack/pty" 17 "github.com/mattn/go-runewidth" 18 "github.com/riywo/loginshell" 19 flag "github.com/spf13/pflag" 20 ) 21 22 var version = "unknown" 23 24 // Regexp to strip ANSI escape sequences from string. Credit: 25 // https://github.com/chalk/ansi-regex/blob/2b56fb0c7a07108e5b54241e8faec160d393aedb/index.js#L4-L7 26 // https://github.com/acarl005/stripansi/blob/5a71ef0e047df0427e87a79f27009029921f1f9b/stripansi.go#L7 27 var ansiEscapes = regexp.MustCompile("[\u001B\u009B][[\\]()#;?]*(?:(?:(?:[a-zA-Z\\d]*(?:;[a-zA-Z\\d]*)*)?\u0007)|(?:(?:\\d{1,4}(?:;\\d{0,4})*)?[\\dA-PRZcf-ntqry=><~]))") 28 29 func printStreamWithTimestamper(r io.Reader, timestamper *Timestamper) { 30 scanner := bufio.NewScanner(r) 31 // Split on \r\n|\r|\n, and return the line as well as the line ending (\r 32 // or \n is preserved, \r\n is collapsed to \n). Adaptation of 33 // bufio.ScanLines. 34 scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { 35 if atEOF && len(data) == 0 { 36 return 0, nil, nil 37 } 38 lfpos := bytes.IndexByte(data, '\n') 39 crpos := bytes.IndexByte(data, '\r') 40 if crpos >= 0 { 41 if lfpos < 0 || lfpos > crpos+1 { 42 // We have a CR-terminated "line". 43 return crpos + 1, data[0 : crpos+1], nil 44 } 45 if lfpos == crpos+1 { 46 // We have a CRLF-terminated line. 47 return lfpos + 1, append(data[0:crpos], '\n'), nil 48 } 49 } 50 if lfpos >= 0 { 51 // We have a LF-terminated line. 52 return lfpos + 1, data[0 : lfpos+1], nil 53 } 54 // If we're at EOF, we have a final, non-terminated line. Return it. 55 if atEOF { 56 return len(data), data, nil 57 } 58 // Request more data. 59 return 0, nil, nil 60 }) 61 for scanner.Scan() { 62 fmt.Print(timestamper.CurrentTimestampString(), " ", scanner.Text()) 63 } 64 } 65 66 func runCommandWithTimestamper(args []string, timestamper *Timestamper) error { 67 // Calculate optimal pty size, taking into account horizontal space taken up by timestamps. 68 getPtyWinsize := func() *pty.Winsize { 69 winsize, err := pty.GetsizeFull(os.Stdin) 70 if err != nil { 71 // Most likely stdin isn't a tty, in which case we don't care. 72 return winsize 73 } 74 totalCols := winsize.Cols 75 plainTimestampString := ansiEscapes.ReplaceAllString(timestamper.CurrentTimestampString(), "") 76 // Timestamp width along with one space character. 77 occupiedWidth := uint16(runewidth.StringWidth(plainTimestampString)) + 1 78 var effectiveCols uint16 = 0 79 if occupiedWidth < totalCols { 80 effectiveCols = totalCols - occupiedWidth 81 } 82 winsize.Cols = effectiveCols 83 // Best effort estimate of the effective width in pixels. 84 if totalCols > 0 { 85 winsize.X = winsize.X * effectiveCols / totalCols 86 } 87 return winsize 88 } 89 90 command := exec.Command(args[0], args[1:]...) 91 ptmx, err := pty.StartWithSize(command, getPtyWinsize()) 92 if err != nil { 93 return err 94 } 95 defer func() { _ = ptmx.Close() }() 96 97 sigs := make(chan os.Signal, 1) 98 signal.Notify(sigs, syscall.SIGWINCH, syscall.SIGINT, syscall.SIGTERM) 99 go func() { 100 for sig := range sigs { 101 switch sig { 102 case syscall.SIGWINCH: 103 if err := pty.Setsize(ptmx, getPtyWinsize()); err != nil { 104 log.Println("error resizing pty:", err) 105 } 106 107 case syscall.SIGINT: 108 _ = syscall.Kill(-command.Process.Pid, syscall.SIGINT) 109 110 case syscall.SIGTERM: 111 _ = syscall.Kill(-command.Process.Pid, syscall.SIGTERM) 112 113 default: 114 } 115 } 116 }() 117 sigs <- syscall.SIGWINCH 118 119 go func() { _, _ = io.Copy(ptmx, os.Stdin) }() 120 121 printStreamWithTimestamper(ptmx, timestamper) 122 123 return command.Wait() 124 } 125 126 func main() { 127 log.SetFlags(log.Flags() &^ (log.Ldate | log.Ltime)) 128 129 var elapsedMode = flag.BoolP("elapsed", "s", false, "show elapsed timestamps") 130 var incrementalMode = flag.BoolP("incremental", "i", false, "show incremental timestamps") 131 var format = flag.StringP("format", "f", "", "show timestamps in this format") 132 var utc = flag.BoolP("utc", "u", false, "show absolute timestamps in UTC") 133 var timezoneName = flag.StringP("timezone", "z", "", "show absolute timestamps in this timezone, e.g. America/New_York") 134 var color = flag.BoolP("color", "c", false, "show timestamps in color") 135 var printHelp = flag.BoolP("help", "h", false, "print help and exit") 136 var printVersion = flag.BoolP("version", "v", false, "print version and exit") 137 flag.CommandLine.SortFlags = false 138 flag.SetInterspersed(false) 139 flag.Usage = func() { 140 fmt.Fprintf(os.Stderr, ` 141 ets -- command output timestamper 142 143 ets prefixes each line of a command's output with a timestamp. Lines are 144 delimited by CR, LF, or CRLF. 145 146 Usage: 147 148 %s [-s | -i] [-f format] [-u | -z timezone] command [arg ...] 149 %s [options] shell_command 150 %s [options] 151 152 The three usage strings correspond to three command execution modes: 153 154 * If given a single command without whitespace(s), or a command and its 155 arguments, execute the command with exec in a pty; 156 157 * If given a single command with whitespace(s), the command is treated as 158 a shell command and executed as SHELL -c shell_command, where SHELL is 159 the current user's login shell, or sh if login shell cannot be determined; 160 161 * If given no command, output is read from stdin, and the user is 162 responsible for piping in a command's output. 163 164 There are three mutually exclusive timestamp modes: 165 166 * The default is absolute time mode, where timestamps from the wall clock 167 are shown; 168 169 * -s, --elapsed turns on elapsed time mode, where every timestamp is the 170 time elapsed from the start of the command (using a monotonic clock); 171 172 * -i, --incremental turns on incremental time mode, where every timestamp is 173 the time elapsed since the last timestamp (using a monotonic clock). 174 175 The default format of the prefixed timestamps depends on the timestamp mode 176 active. Users may supply a custom format string with the -f, --format option. 177 The format string is basically a strftime(3) format string; see the man page 178 or README for details on supported formatting directives. 179 180 The timezone for absolute timestamps can be controlled via the -u, --utc 181 and -z, --timezone options. --timezone accepts IANA time zone names, e.g., 182 America/Los_Angeles. Local time is used by default. 183 184 Options: 185 `, os.Args[0], os.Args[0], os.Args[0]) 186 flag.PrintDefaults() 187 } 188 flag.Parse() 189 190 if *printHelp { 191 flag.Usage() 192 os.Exit(0) 193 } 194 195 if *printVersion { 196 fmt.Println(version) 197 os.Exit(0) 198 } 199 200 mode := AbsoluteTimeMode 201 if *elapsedMode && *incrementalMode { 202 log.Fatal("conflicting flags --elapsed and --incremental") 203 } 204 if *elapsedMode { 205 mode = ElapsedTimeMode 206 } 207 if *incrementalMode { 208 mode = IncrementalTimeMode 209 } 210 if *format == "" { 211 if mode == AbsoluteTimeMode { 212 *format = "[%F %T]" 213 } else { 214 *format = "[%T]" 215 } 216 } 217 timezone := time.Local 218 if *utc && *timezoneName != "" { 219 log.Fatal("conflicting flags --utc and --timezone") 220 } 221 if *utc { 222 timezone = time.UTC 223 } 224 if *timezoneName != "" { 225 location, err := time.LoadLocation(*timezoneName) 226 if err != nil { 227 log.Fatal(err) 228 } 229 timezone = location 230 } 231 if *color { 232 *format = "\x1b[32m" + *format + "\x1b[0m" 233 } 234 args := flag.Args() 235 236 timestamper, err := NewTimestamper(*format, mode, timezone) 237 if err != nil { 238 log.Fatal(err) 239 } 240 241 exitCode := 0 242 if len(args) == 0 { 243 printStreamWithTimestamper(os.Stdin, timestamper) 244 } else { 245 if len(args) == 1 { 246 arg0 := args[0] 247 if matched, _ := regexp.MatchString(`\s`, arg0); matched { 248 shell, err := loginshell.Shell() 249 if err != nil { 250 shell = "sh" 251 } 252 args = []string{shell, "-c", arg0} 253 } 254 } 255 if err = runCommandWithTimestamper(args, timestamper); err != nil { 256 if exitErr, ok := err.(*exec.ExitError); ok { 257 exitCode = exitErr.ExitCode() 258 } else { 259 log.Fatal(err) 260 } 261 } 262 } 263 os.Exit(exitCode) 264 }