class Nucleon::Util::SSH

# File lib/core/util/ssh.rb, line 209
def self.close(hostname = nil, user = nil)
  if hostname && user.nil? # Assume we entered a session id
    if @@sessions.has_key?(hostname)
      @@sessions[hostname].close
      @@sessions.delete(hostname)  
    end
    
  elsif hostname && user # Generate session id from args
    session_id = session_id(hostname, user)
    
    if @@sessions.has_key?(session_id)
      @@sessions[session_id].close
      @@sessions.delete(session_id)  
    end
    
  else # Close all connections
    @@sessions.keys.each do |id|
      @@sessions[id].close
      @@sessions.delete(id)      
    end
  end
end

# File lib/core/util/ssh.rb, line 195
def self.close_session(hostname, user)
  session_id = session_id(hostname, user)
  
  if @@sessions.has_key?(session_id)
    begin # Don't care about errors here
      @@sessions[session_id].close
    rescue
    end
    @@sessions.delete(session_id)
  end
end

# File lib/core/util/ssh.rb, line 290
def self.download(hostname, user, remote_path, local_path, options = {})
  config = Config.ensure(options)
  
  require 'net/scp'
  
  # Accepted options:
  # * :recursive - the +remote+ parameter refers to a remote directory, which
  # should be downloaded to a new directory named +local+ on the local
  # machine.
  # * :preserve - the atime and mtime of the file should be preserved.
  # * :verbose - the process should result in verbose output on the server
  # end (useful for debugging).
  #
  config.init(:recursive, true)
  config.init(:preserve, true)
  config.init(:verbose, true)
  
  blocking = config.delete(:blocking, true)
  
  session(hostname, user) do |ssh|
    if blocking
      ssh.scp.download!(remote_path, local_path, config.export) do |ch, name, received, total|
        yield(name, received, total) if block_given?
      end
    else
      ssh.scp.download(remote_path, local_path, config.export)
    end
  end
end

# File lib/core/util/ssh.rb, line 234
def self.exec(hostname, user, commands)
  results = []
      
  begin
    session(hostname, user) do |ssh|
      Data.array(commands).each do |command|
        command = command.flatten.join(' ') if command.is_a?(Array)
        command = command.to_s
        result  = Shell::Result.new(command)
        
        logger.info(">> running SSH: #{command}")
            
        ssh.open_channel do |ssh_channel|
          ssh_channel.request_pty
          ssh_channel.exec(command) do |channel, success|
            unless success
              raise "Could not execute command: #{command.inspect}"
            end

            channel.on_data do |ch, data|
              result.append_output(data)
              yield(:output, command, data) if block_given?
            end

            channel.on_extended_data do |ch, type, data|
              next unless type == 1
              result.append_errors(data)
              yield(:error, command, data) if block_given?
            end

            channel.on_request('exit-status') do |ch, data|
              result.status = data.read_long
            end

            channel.on_request('exit-signal') do |ch, data|
              result.status = 255
            end
          end
        end
        logger.warn("`#{command}` messages: #{result.errors}") if result.errors.length > 0
        logger.warn("`#{command}` status: #{result.status}") unless result.status == 0
     
        results << result
        ssh.loop          
      end
    end
  rescue Net::SSH::HostKeyMismatch => error
    error.remember_host!
    sleep 0.2
    retry
  end
  results  
end

# File lib/core/util/ssh.rb, line 26
def self.generate(options = {})
  config = Config.ensure(options)
  
  private_key  = config.get(:private_key, nil)
  original_key = nil
  key_comment  = config.get(:comment, '')
  passphrase   = config.get(:passphrase, nil)
  force        = config.get(:force, false)    
  
  if private_key.nil?
    key_type    = config.get(:type, "RSA")
    key_bits    = config.get(:bits, 2048)
    
    key_data = SSHKey.generate({
      :type       => key_type, 
      :bits       => key_bits, 
      :comment    => key_comment, 
      :passphrase => passphrase
    })
    is_new = true
    
  else
    if private_key.include?('PRIVATE KEY')
      original_key = private_key   
    else
      original_key = Disk.read(private_key)
    end
    
    if original_key
      encrypted = original_key.include?('ENCRYPTED')
      key_data  = SSHKey.new(original_key, { 
        :comment    => key_comment, 
        :passphrase => passphrase 
      }) if force || ! encrypted || passphrase 
    end      
    is_new = false
  end
  
  return nil unless key_data && ! key_data.ssh_public_key.empty?
  Keypair.new(key_data, is_new, original_key, passphrase)
end

# File lib/core/util/ssh.rb, line 189
def self.init_session(hostname, user, port = 22, private_key = nil, options = {})
  session(hostname, user, port, private_key, true, options)  
end

# File lib/core/util/ssh.rb, line 13
def self.key_path
  unless @@key_path
    home_path  = ( ENV['USER'] == 'root' ? '/root' : ENV['HOME'] ) # In case we are using sudo
    @@key_path = File.join(home_path, '.ssh')
  
    FileUtils.mkdir(@@key_path) unless File.directory?(@@key_path)
  end
  @@key_path
end

# File lib/core/util/ssh.rb, line 151
def self.session(hostname, user, port = 22, private_key = nil, reset = false, options = {})
  require 'net/ssh'
  
  session_id  = session_id(hostname, user)
  config      = Config.ensure(options)
  
  ssh_options = Config.new({
    :user_known_hosts_file => [ File.join(key_path, 'known_hosts'), File.join(key_path, 'known_hosts2') ],
    :auth_methods          => [ 'publickey' ],
    :paranoid              => :very
  }).import(Util::Data.subset(config, config.keys - [ :keypair, :key_dir, :key_name ]))
  
  if private_key
    auth_id = [ session_id, private_key ].join('_')      
    
    if ! @@auth[auth_id] && keypair = unlock_private_key(private_key, config)
      @@auth[auth_id] = keypair
    end
    config[:keypair] = @@auth[auth_id] # Reset so caller can access updated keypair      
    
    if @@auth[auth_id].is_a?(String)
      ssh_options[:keys_only] = false
      ssh_options[:keys]      = [ @@auth[auth_id] ]  
    else
      ssh_options[:keys_only] = true
      ssh_options[:key_data]  = [ @@auth[auth_id].private_key ]
    end
  end
  
  ssh_options[:port] = port    
  
  if reset || ! @@sessions.has_key?(session_id)
    @@sessions[session_id] = Net::SSH.start(hostname, user, ssh_options.export)
  end
  yield(@@sessions[session_id]) if block_given? && @@sessions[session_id]    
  @@sessions[session_id] 
end

# File lib/core/util/ssh.rb, line 145
def self.session_id(hostname, user)
  "#{hostname}-#{user}"  
end

# File lib/core/util/ssh.rb, line 364
def self.terminal(hostname, user, options = {})
  config   = Config.ensure(options)
  ssh_path = nucleon_locate("ssh")
  
  raise Errors::SSHUnavailable unless ssh_path
  
  port         = config.get(:port, 22)
  private_keys = config.get(:private_keys, File.join(ENV['HOME'], '.ssh', 'id_rsa'))
  
  command_options = [
    "#{user}@#{hostname}",
    "-p", port.to_s,
    "-o", "Compression=yes",
    "-o", "DSAAuthentication=yes",
    "-o", "LogLevel=FATAL",
    "-o", "StrictHostKeyChecking=no",
    "-o", "UserKnownHostsFile=/dev/null",
    "-o", "IdentitiesOnly=yes"
  ]

  Util::Data.array(private_keys).each do |path|
    command_options += [ "-i", File.expand_path(path) ]
  end
  
  if config.get(:forward_x11, false)
    command_options += [
      "-o", "ForwardX11=yes",
      "-o", "ForwardX11Trusted=yes"
    ]
  end

  command_options += [ "-o", "ProxyCommand=#{config[:proxy_command]}" ] if config.get(:proxy_command, false)
  command_options += [ "-o", "ForwardAgent=yes" ] if config.get(:forward_agent, false)
  
  command_options.concat(Util::Data.array(config[:extra_args])) if config.get(:extra_args, false)

  #---

  logger.info("Executing SSH in subprocess: #{command_options.inspect}")
  
  process = ChildProcess.build('ssh', *command_options)
  process.io.inherit!
  
  process.start
  process.wait
  process.exit_code
end

# File lib/core/util/ssh.rb, line 415
def self.unlock_private_key(private_key, options = {})
  require 'net/ssh'
  
  config   = Config.ensure(options)
  keypair  = config.get(:keypair, nil)
  key_dir  = config.get(:key_dir, nil)
  key_name = config.get(:key_name, 'default')
  no_file  = ENV['NUCLEON_NO_SSH_KEY_SAVE']
  
  password    = nil
  tmp_key_dir = nil
  
  if private_key
    keypair = nil if keypair && ! keypair.private_key
    
    unless no_file
      if key_dir && key_name && File.exists?(private_key) && loaded_private_key = Util::Disk.read(private_key)
        FileUtils.mkdir_p(key_dir)
    
        loaded_private_key =~ /BEGIN\s+([A-Z]+)\s+/
      
        local_key_type = $1
        local_key_name = Keypair.render(local_key_type, key_name)
        local_key_path = File.join(key_dir, local_key_name)
      
        keypair = generate({ :private_key => local_key_path }) if File.exists?(local_key_path)
      end
    end
    
    unless keypair
      key_manager_logger       = ::Logger.new(STDERR)
      key_manager_logger.level = ::Logger::FATAL
      key_manager              = Net::SSH::Authentication::KeyManager.new(key_manager_logger)        
      
      key_manager.each_identity do |identity|
        if identity.comment == private_key
          # Feed the key to the system password manager if it exists
          keypair = private_key
        end
      end
      key_manager.finish
    
      until keypair
        keypair = generate({
          :private_key => private_key,
          :passphrase  => password
        })        
        password = ui.ask("Enter passphrase for #{private_key}: ", { :echo => false }) unless keypair
      end
      
      unless no_file
        if key_dir && key_name && ! keypair.is_a?(String)
          key_files = keypair.store(key_dir, key_name, false)
        
          if key_files && File.exists?(key_files[:private_key])
            keypair = generate({ :private_key => key_files[:private_key] })
          end
        end
      end
    end
  end
  keypair      
end

# File lib/core/util/ssh.rb, line 322
def self.upload(hostname, user, local_path, remote_path, options = {})
  config = Config.ensure(options)
  
  require 'net/scp'
  
  # Accepted options:
  # * :recursive - the +local+ parameter refers to a local directory, which
  # should be uploaded to a new directory named +remote+ on the remote
  # server.
  # * :preserve - the atime and mtime of the file should be preserved.
  # * :verbose - the process should result in verbose output on the server
  # end (useful for debugging).
  # * :chunk_size - the size of each "chunk" that should be sent. Defaults
  # to 2048. Changing this value may improve throughput at the expense
  # of decreasing interactivity.
  #
  config.init(:recursive, true)
  config.init(:preserve, true)
  config.init(:verbose, true)
  config.init(:chunk_size, 2048)
  
  blocking = config.delete(:blocking, true)
  
  session(hostname, user) do |ssh|
    if blocking
      ssh.scp.upload!(local_path, remote_path, config.export) do |ch, name, sent, total|
        yield(name, sent, total) if block_given?
      end
    else
      ssh.scp.upload(local_path, remote_path, config.export)
    end
  end
end

# File lib/core/util/ssh.rb, line 71
def self.valid?(public_ssh_key)
  SSHKey.valid_ssh_public_key?(public_ssh_key)
end