Go言語でbashの標準入力と標準出力・標準エラー出力をフックする

こんな情報必要とする人いないかもしれないですが、Golangでの標準入力、標準出力のストリームデータの扱いが少しややこしかったので備忘のために記録しておきます。

やったこと

Linux上のBashに対する標準入力および標準出力・標準エラー出力の内容をGoのプログラムで取り込み、記録とか、その他活用ができるようにしてみました。
Linuxだと、標準入力と標準出力・標準エラー出力の内容を記録するコマンドとして、scriptコマンドがあるのでそれで十分なところもあるかと思いますが、それと近いことをGoのプログラムだけで実現し、自由にデータを操作できるようになるといいことあるかもとの思いから試してみました。

前提知識

GolangでのOSコマンドの実行

Goプログラムでコマンドを実行するには以下のexec.Commandで実行可能です。

qiita_sample1.go

package main

import (
	"fmt"
	"os/exec"
)

func main() {
	cmd := exec.Command("echo", "Hello")
	output, _ := cmd.Output()
	fmt.Printf(string(output))
}

実行してその内容を1回出力するだけなら上記のように簡単にかけますが、bash上での繰り返し行われる操作を継続して取り扱うにはストリーム処理が必要となります。

Stdin Stdout Stderrの扱い

詳しくはこの辺りの記事が非常に参考になります。
http://ascii.jp/elem/000/001/260/1260449/index-2.html

Commandの実行結果の値をリアルタイムに処理していくには以下のようにCommandのStdoutPipeを使って標準出力の内容をパイプで取り込み、bufioのScannerで1行ずつ読み込んで処理させます。

qiita_sample2.go

package main

import (
	"bufio"
	"fmt"
	"os/exec"
)

func main() {
	cmd := exec.Command("echo", "Hello")
	stdout, _ := cmd.StdoutPipe()
	cmd.Start()

	go func() {
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			line := scanner.Text()
			fmt.Println(line)
		}
	}()
	cmd.Wait()
}

上記は標準出力をPipeで取得してスキャンして出力させる例ですが、同様にStderrPipe()とかStdinPipe()使えば標準入力、標準エラー出力も取り扱うことが可能です。
標準入力の場合はコマンドライン上の標準入力(os.Stdin)の内容をScannerで取り込んで処理させるイメージです。

標準入力の例

qiita_sample3.go

package main

import (
	"bufio"
	"fmt"
	"os"
)

func main() {
	stdin := bufio.NewScanner(os.Stdin)
	for stdin.Scan() {
		line := stdin.Text()
		fmt.Println(line)

	}
}

TeeReaderを使ったストリーム処理の分岐

StdoutPipeとか使って取り出したものをScannerで呼び出して処理させると、
ScannerでScanすると取り出したタイミングでバッファから削除されていくので処理は1直列しかできません。

そこで、今回やりたいような、bashで実行したコマンドの内容を通常通り標準出力に出しつつ、ファイル書き出しとかその他用途でも利用できるよう取り込むにはio.TeeReaderが有用です。
これを使うことで、1つの入力を2つの出力先に分割することができます。Linuxのteeコマンドの処理と同様のことができます。

qiita_sample4.go

package main

import (
        "bufio"
        "fmt"
        "io"
        "os"
        "os/exec"
)

func main() {
        cmd := exec.Command("echo", "Hello")
        stdout, _ := cmd.StdoutPipe()

        filename := "/tmp/qiita_test.log"
        file, _ := os.Create(filename)
        defer file.Close()

        stdout2 := io.TeeReader(stdout, file)

        cmd.Start()

        go func() {
                scanner := bufio.NewScanner(stdout2)
                for scanner.Scan() {
                        line := scanner.Text()
                        fmt.Println(line)
                }
        }()
        cmd.Wait()
}

上記を実行すると、標準出力にもコマンド結果が出力され、ファイルにも結果が出力されます

$ go run qiita_sample4.go
Hello
$ cat /tmp/qiita_test.log
Hello

実装

では、このようなことを活用して目的のbashの標準入力および標準出力・エラー出力をそれぞれピックアップして処理できるようにしてみます。

ストリームのやり取りを図示すると以下のイメージです。

なお、exec.Command("bash")で実行させると、ttyが紐付かなかったり、Tabによる補完等ができないので、ptyというパッケージを使って、ターミナル起動できるようにしています。(https://github.com/kr/ptyのページのShellの例を流用しています。)

さらに、Scannerで処理させる際、デフォルトのScan処理だとバッファ内のデータを処理する際、改行コードとみなす文字が「¥n」、「¥r¥n」のみのようで、「¥r」のみの場合、改行されてないとみなされ永遠にScan()で行の情報が取り出せないのでScan処理に¥rを処理できるようコードをかませています。
この処理はこちらを参考に(そのまま流用)させてもらっています。

qiita_sample5.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/kr/pty"
	"golang.org/x/crypto/ssh/terminal"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
)

// scanner split custom (for split \r treating as new line)
func customScan(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	var i int
	if i = bytes.IndexByte(data, '\n'); i >= 0 {
		return i + 1, dropCR(data[0:i]), nil
	}
	if i = bytes.IndexByte(data, '\r'); i >= 0 {
		return i + 1, data[0:i], nil
	}
	if atEOF {
		return len(data), dropCR(data), nil
	}
	return 0, nil, nil
}

// dropCR drops a terminal \r from the data.
func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

func execShell() error {
	// Create arbitrary command.
	c := exec.Command("bash")

	//ここからptyでのbashの起動関連の処理
	// Start the command with a pty.
	ptmx, err := pty.Start(c)
	if err != nil {
		return err
	}
	// Make sure to close the pty at the end.
	defer func() { _ = ptmx.Close() }()

	// Handle pty size.
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGWINCH)
	go func() {
		for range ch {
			if err := pty.InheritSize(os.Stdin, ptmx); err != nil {
				log.Printf("error resizing pty: %s", err)
			}
		}
	}()
	ch <- syscall.SIGWINCH // Initial resize.

	// Set stdin in raw mode.
	oldState, err := terminal.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		panic(err)
	}
	defer func() { _ = terminal.Restore(int(os.Stdin.Fd()), oldState) }() // Best effort.

	//ここまで

	// 標準入力の内容をptyで起動したbashに引き渡すのと、stdinのバッファに複製して取得
	stdin := io.TeeReader(os.Stdin, ptmx)
	in_scanner := bufio.NewScanner(stdin)
	go func() { // Scan処理は基本的に処理を待つので入力この処理はgo routineでバックで回す必要あり
		// 標準入力の情報をScanしてそのまま標準出力する
		in_scanner.Split(customScan)
		for in_scanner.Scan() {
			text := in_scanner.Text()
			fmt.Printf("[in_scanner]> %s\r\n", text)
		}
	}()

	// bashの実行結果の標準出力・標準エラー出力の内容をos.Stdout(コマンドラインの標準出力)に引き渡すのと、stdoutのバッファに複製して取得
	stdout := io.TeeReader(ptmx, os.Stdout)
	out_scanner := bufio.NewScanner(stdout)

	out_scanner.Split(customScan)
	for out_scanner.Scan() {
		text := out_scanner.Text()
		fmt.Printf("[out_scanner]> %s\r\n", text)
	}

	return nil
}

func main() {
	if err := execShell(); err != nil {
		log.Fatal(err)
	}
}

実行の結果は以下のようになり、ややこしいですが、scannerで別途取得した情報には[in_scanner]もしくは[out_scanner]の接頭文字がついたカスタマイズ内容が出力されていることがわかります。
今回のサンプルでは単純にprintするだけですが、Scanで読み込んだ内容は自由に加工して処理させることができるので他の用途でも活用できると思います。

$ go run qiita_sample5.go
bash-3.2$ whoami   // 入力
[out_scanner]> bash-3.2$ whoami //out_scannerから読み取った標準出力の情報
[in_scanner]> whoami  //in_scannerから読み取った標準入力の情報
ike_dai  //実際の標準出力の内容
[out_scanner]> ike_dai  //out_scannerから読み取った標準出力の情報
bash-3.2$ exit[in_scanner]> exit  

[out_scanner]> bash-3.2$ exit
exit
[out_scanner]> exit

まとめ

かなり低レイヤでの処理の扱いとなり、結構ややこしい感じですが、この辺りコントロールできるようになると面白いですね。bashの仕組みの勉強にもなりました。
他にもわかりやすい書き方とかはあるかも？なので1参考情報としていただければ。