class Net::SSH::Transport::Kex::DiffieHellmanGroup1SHA1

def build_signature_buffer(result)
  response = Net::SSH::Buffer.new
  response.write_string data[:client_version_string],
                        data[:server_version_string],
                        data[:client_algorithm_packet],
                        data[:server_algorithm_packet],
                        result[:key_blob]
  response.write_bignum dh.pub_key,
                        result[:server_dh_pubkey],
                        result[:shared_secret]
  response
end

def confirm_newkeys #:nodoc:
  # send own NEWKEYS message first (the wodSSHServer won't send first)
  response = Net::SSH::Buffer.new
  response.write_byte(NEWKEYS)
  connection.send_message(response)
  # wait for the server's NEWKEYS message
  buffer = connection.next_message
  raise Net::SSH::Exception, "expected NEWKEYS" unless buffer.type == NEWKEYS
end

def exchange_keys
  result = send_kexinit
  verify_server_key(result[:server_key])
  session_id = verify_signature(result)
  confirm_newkeys
  return { :session_id        => session_id, 
           :server_key        => result[:server_key],
           :shared_secret     => result[:shared_secret],
           :hashing_algorithm => digester }
end

def generate_key #:nodoc:
  dh = OpenSSL::PKey::DH.new
  dh.p, dh.g = get_parameters
  dh.priv_key = OpenSSL::BN.rand(data[:need_bytes] * 8)
  dh.generate_key! until dh.valid?
  dh
end

def generate_key_fingerprint(key)
  blob = Net::SSH::Buffer.from(:key, key).to_s
  fingerprint = OpenSSL::Digest::MD5.hexdigest(blob).scan(/../).join(":")
  [blob, fingerprint]
rescue ::Exception => e
  [nil, "(could not generate fingerprint: #{e.message})"]
end

def get_g
  G
end

def get_message_types
  [KEXDH_INIT, KEXDH_REPLY]
end

def get_p
  OpenSSL::BN.new(P_s, P_r)
end

def get_parameters
  [p, g]
end

def initialize(algorithms, connection, data)
  @p = get_p
  @g = get_g
  @digester = OpenSSL::Digest::SHA1
  @algorithms = algorithms
  @connection = connection
  @data = data.dup
  @dh = generate_key
  @logger = @data.delete(:logger)
end

def send_kexinit #:nodoc:
  init, reply = get_message_types
  # send the KEXDH_INIT message
  buffer = Net::SSH::Buffer.from(:byte, init, :bignum, dh.pub_key)
  connection.send_message(buffer)
  # expect the KEXDH_REPLY message
  buffer = connection.next_message
  raise Net::SSH::Exception, "expected REPLY" unless buffer.type == reply
  result = Hash.new
  result[:key_blob] = buffer.read_string
  result[:server_key] = Net::SSH::Buffer.new(result[:key_blob]).read_key
  result[:server_dh_pubkey] = buffer.read_bignum
  result[:shared_secret] = OpenSSL::BN.new(dh.compute_key(result[:server_dh_pubkey]), 2)
  sig_buffer = Net::SSH::Buffer.new(buffer.read_string)
  sig_type = sig_buffer.read_string
  if sig_type != algorithms.host_key
    raise Net::SSH::Exception,
      "host key algorithm mismatch for signature " +
      "'#{sig_type}' != '#{algorithms.host_key}'"
  end
  result[:server_sig] = sig_buffer.read_string
  return result
end

def verify_server_key(key) #:nodoc:
  if key.ssh_type != algorithms.host_key
    raise Net::SSH::Exception,
      "host key algorithm mismatch " +
      "'#{key.ssh_type}' != '#{algorithms.host_key}'"
  end
  blob, fingerprint = generate_key_fingerprint(key)
  unless connection.host_key_verifier.verify(:key => key, :key_blob => blob, :fingerprint => fingerprint, :session => connection)
    raise Net::SSH::Exception, "host key verification failed"
  end
end

def verify_signature(result) #:nodoc:
  response = build_signature_buffer(result)
  hash = @digester.digest(response.to_s)
  unless result[:server_key].ssh_do_verify(result[:server_sig], hash)
    raise Net::SSH::Exception, "could not verify server signature"
  end
  return hash
end