HostKeyEntries classes

5 files changed, 65 insertions, 21 deletions

diff --git a/lib/net/ssh.rb b/lib/net/ssh.rb
index 8b8b7b9..bbae909 100644
--- a/lib/net/ssh.rb
+++ b/lib/net/ssh.rb
@@ -121,7 +121,7 @@ module Net
     # * :forward_agent => set to true if you want the SSH agent connection to
     #   be forwarded
     # * :known_hosts => a custom object holding known hosts records.
-    #   It must implement #search_for and add in a similiar manner as KnownHosts.
+    #   It must implement #search_for and `add` in a similiar manner as KnownHosts.
     # * :global_known_hosts_file => the location of the global known hosts
     #   file. Set to an array if you want to specify multiple global known
     #   hosts files. Defaults to %w(/etc/ssh/ssh_known_hosts /etc/ssh/ssh_known_hosts2).
diff --git a/lib/net/ssh/known_hosts.rb b/lib/net/ssh/known_hosts.rb
index 7fbda9b..7c8b3d5 100644
--- a/lib/net/ssh/known_hosts.rb
+++ b/lib/net/ssh/known_hosts.rb
@@ -3,21 +3,58 @@ require 'openssl'
 require 'base64'
 require 'net/ssh/buffer'
 require 'net/ssh/authentication/ed25519_loader'
+require 'byebug'
 
 module Net
   module SSH
-    class HostCertPub
-      def ssh_type
-        "ecdsa-sha2-nistp256-cert-v01@openssh.com"
-      end
+    module HostKeyEntries
+      class PubKey < Delegator
+        def initialize(key)
+          @key = key
+        end
+
+        def ssh_type
+          @key.ssh_type
+        end
+
+        def ssh_types
+          [ssh_type]
+        end
+
+        def to_blob
+          @key.to_blob
+        end
+
+        def __getobj__
+          Kernel.warn("Calling Net::SSH::Buffer methods on HostKeyEntries PubKey is deprecated")
+          @key
+        end
 
-      def initialize(content)
-        @content = content
+        def matches_key?(server_key)
+          @key.ssh_type == server_key.ssh_type && @key.to_blob == server_key.to_blob
+        end
       end
 
-      def matches?(server_key)
-        certblob = Buffer.new(server_key).read_key
-        certblob.signature_valid? && (certblob.signature_key.to_blob == @content.to_blob)
+      class CertAuthority
+        def ssh_types
+          %w[
+            ecdsa-sha2-nistp256-cert-v01@openssh.com
+            ecdsa-sha2-nistp384-cert-v01@openssh.com
+            ecdsa-sha2-nistp521-cert-v01@openssh.com
+          ]
+        end
+  
+        def initialize(key)
+          @key = key
+        end
+
+        def matches_key?(server_key)
+          if ssh_types.include?(server_key.ssh_type)
+            server_key.signature_valid? && (server_key.signature_key.to_blob == @key.to_blob)
+          else
+            false
+          end
+        end
       end
     end
 
@@ -142,10 +179,16 @@ module Net
 
         File.open(source) do |file|
           file.each_line do |line|
-            if line.start_with?("@cert-authority ")
-              cert_auth, hosts, type, key_content = line.split(' ')
+            marker = nil
+            hosts, type, key_content, comment, last = line.split(' ')
+
+            if ['@cert-authority'].include?(hosts)
+              marker, hosts, type, key_content, comment = hosts, type, key_content, comment, last
+            end
+
+            if marker == "@cert-authority"
               blob = key_content.unpack("m*").first
-              keys << HostCertPub.new(Net::SSH::Buffer.new(blob).read_key )
+              keys << HostKeyEntries::CertAuthority.new(Net::SSH::Buffer.new(blob).read_key)
             else
               hosts, type, key_content = line.split(' ')
               # Skip empty line or one that is commented
@@ -162,7 +205,7 @@ module Net
               next unless found
 
               blob = key_content.unpack("m*").first
-              keys << Net::SSH::Buffer.new(blob).read_key
+              keys << HostKeyEntries::PubKey.new(Net::SSH::Buffer.new(blob).read_key)
             end
           end
         end
diff --git a/lib/net/ssh/transport/algorithms.rb b/lib/net/ssh/transport/algorithms.rb
index 7408d40..f0d8d15 100644
--- a/lib/net/ssh/transport/algorithms.rb
+++ b/lib/net/ssh/transport/algorithms.rb
@@ -278,7 +278,7 @@ module Net
             # existing known key for the host has preference.
 
             existing_keys = session.host_keys
-            host_keys = existing_keys.map { |key| key.ssh_type }.uniq
+            host_keys = existing_keys.flat_map { |key| key.respond_to?(:ssh_types) ? key.ssh_types : [key.ssh_type] }.uniq
             algorithms[:host_key].each do |name|
               host_keys << name unless host_keys.include?(name)
             end
diff --git a/lib/net/ssh/verifiers/always.rb b/lib/net/ssh/verifiers/always.rb
index 48c70bb..0c86589 100644
--- a/lib/net/ssh/verifiers/always.rb
+++ b/lib/net/ssh/verifiers/always.rb
@@ -21,12 +21,13 @@ module Net
 
           # If we found any matches, check to see that the key type and
           # blob also match.
+
           found = host_keys.any? do |key|
-            key.ssh_type == arguments[:key].ssh_type &&
-            (
-              (key.respond_to?(:to_blob) && key.to_blob  == arguments[:key].to_blob) ||
-              (key.respond_to?(:matches?) && key.matches?(arguments[:key].to_blob))
-            )
+            if key.respond_to?(:matches_key?)
+              key.matches_key?(arguments[:key])
+            else
+              key.ssh_type == arguments[:key].ssh_type && key.to_blob == arguments[:key].to_blob
+            end
           end
 
           # If a match was found, return true. Otherwise, raise an exception
diff --git a/test/transport/test_algorithms.rb b/test/transport/test_algorithms.rb
index d998d56..fbe96a5 100644
--- a/test/transport/test_algorithms.rb
+++ b/test/transport/test_algorithms.rb
@@ -442,7 +442,7 @@ module Transport
     def transport(transport_options={})
       @transport ||= MockTransport.new(
         {user_known_hosts_file: '/dev/null',
-        global_known_hosts_file: '/dev/null'}.merge(transport_options)
+        global_known_hosts_file: '/dev/null'}.merge(transport_options)  
       )
     end
   end