Rubyでexpectを用いて複数マシンのssh鍵を収集

開発 研究

今日、あるプロラムを複数のマシン上で走らせるために、あらかじめ各マシンのssh鍵を収集しておく必要ができた。使用するマシンのssh鍵をほとんど持っていなかったため、マニュアルでするとなると毎回sshでログインしてyesって打たなきゃならない。そんなことはやっていられなかったので、Rubyのexpectライブラリを用いてyesって入力してくれるスクリプトを書いた。
ssh鍵を取得するためのスクリプトとして、今回は2つスクリプトを作成した。1つは汎用的で使いやすいexpect on Rubyなライブラリ。もう1つはそのライブラリを用いて複数マシンから連続的にssh鍵を収集するスクリプト。
最初の expect on Rubyライブラリは次の2つのページを参考に実装した。
Young risk taker.: Rubyのexpect.rbの使い方
http://www.ksky.ne.jp/~sakae/d2/20518.html

# new_expect.rb
require 'pty'
require 'expect'

module Expect
  def spawn(cmd)
    puts "CMD: #{cmd}" if $expect_verbose
    PTY.spawn(cmd) { |r,w,pid|
      @input_stream = r
      @output_stream = w
      @child_pid = pid
      yield
    }
  end

  def expect(pat, timeout=5)
    ret = @input_stream.expect(pat, timeout) { |match|
      raise "expect %s timeout " %(pat.kind_of?(Regexp)? pat.source : pat) unless match
      @output_stream.puts(yield(match))
    }
  end
end

実装は最初のページで紹介されているサンプルとほとんど一緒。だけど、PTY.protect_signal関数は現在の実装では呼ばなくても良いとのことなのでそこは消した。
で、このライブラリを用いてssh鍵を集めるスクリプトを次のように実装した。

require "new_expect"

class SshKeyCollector
  include Expect

  def initialize(fname)
    @list_name = fname
  end

  def collect
    File::open(@list_name) { |file|
      file.each_line { |line|
        hostnameFQDN = line.strip
        collect_one hostnameFQDN

        hostname = hostnameFQDN.split(".")[0]
        collect_one hostname if hostname
      }
    }
  end

  private
  def collect_one(hostname)
    STDERR.print "ssh to #{hostname}..."

    spawn("ssh #{hostname}") {
      begin
        expect("(yes/no)?", 1) { |match| "yes" }
      rescue
        STDERR.puts "failed"
        STDERR.puts "Key from #{hostname} may have been collected"
        return
      end
      # expect("Password:") { |match| "__PASSWORD__" }
      expect(/\$$/) { |match| "exit" }
    }

    STDERR.puts "done"    
  end
end


### main
if __FILE__ == $0
  if ARGV.size != 1
    puts "Usage"
    puts "    collect_sshkey NODE_LIST_FILE_NAME"
  end

  list_name = ARGV.shift
  if !FileTest.exists? list_name
    puts "#{list_name} doesn't exist"
  end

  SshKeyCollector::new(list_name).collect
end

実行するときは、

$ ruby collect_sshkey.rb host_list

って感じで、マシン名を1行に1つずつ記述したファイルの名前を与える。
SshKeyCollector#collectでhostnameFQDNとhostnameと同じマシンに対してホスト名だけの場合とFQDNの場合で鍵を取得しているが、これはsshが2つを区別するため。SshKeyCollector#collect_oneでは「(yes/no)?」プロンプトが出てこなくても処理が続けられるようにしている。これは既にssh鍵を持っている場合にはこのプロンプトは出力されなく、その場合にも後続するマシンの鍵を集めたいため。
このスクリプトを使って、簡単に全てのマシンの鍵を集めることができた。マシン名は設定ファイルで指定するので、移植も簡単。昔、素のexpect（URL）を使って挫折したことがあったんだけど、今回はRubyだったから（？）簡単にできた。