require 'net/ssh/transport/constants'
require 'net/ssh/transport/kex/diffie_hellman_group1_sha1'

module Net; module SSH; module Transport; module Kex

  # A key-exchange service implementing the "ecdh-sha2-nistp256"
  # key-exchange algorithm. (defined in RFC 5656)
  class EcdhSHA2NistP256 < DiffieHellmanGroup1SHA1
    include Constants, Loggable

    attr_reader :ecdh

    def digester
      OpenSSL::Digest::SHA256
    end

    def curve_name
      OpenSSL::PKey::EC::CurveNameAlias['nistp256']
    end

    def initialize(algorithms, connection, data)
      @algorithms = algorithms
      @connection = connection

      @digester = digester
      @data = data.dup
      @ecdh = generate_key
      @logger = @data.delete(:logger)
    end

    private

    def get_message_types
      [KEXECDH_INIT, KEXECDH_REPLY]
    end

    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],
                            ecdh.public_key.to_bn.to_s(2),
                            result[:server_ecdh_pubkey]
      response.write_bignum result[:shared_secret]
      response
    end

    def generate_key #:nodoc:
      OpenSSL::PKey::EC.new(curve_name).generate_key
    end

    def send_kexinit #:nodoc:
      init, reply = get_message_types

      # send the KEXECDH_INIT message
      ## byte     SSH_MSG_KEX_ECDH_INIT
      ## string   Q_C, client's ephemeral public key octet string
      buffer = Net::SSH::Buffer.from(:byte, init, :mstring, ecdh.public_key.to_bn.to_s(2))
      connection.send_message(buffer)

      # expect the following KEXECDH_REPLY message
      ## byte     SSH_MSG_KEX_ECDH_REPLY
      ## string   K_S, server's public host key
      ## string   Q_S, server's ephemeral public key octet string
      ## string   the signature on the exchange hash
      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_ecdh_pubkey] = buffer.read_string

      # compute shared secret from server's public key and client's private key
      pk = OpenSSL::PKey::EC::Point.new(OpenSSL::PKey::EC.new(curve_name).group,
                                        OpenSSL::BN.new(result[:server_ecdh_pubkey], 2))
      result[:shared_secret] = OpenSSL::BN.new(ecdh.dh_compute_key(pk), 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
  end
end; end; end; end