--- checks SSH hostkeys for weak keys. May use blacklists from e.g.
-- http://packages.debian.org/source/testing/openssh-blacklist
-- http://packages.debian.org/source/testing/openssh-blacklist-extra
--
-- The blacklist filename is openssh.blacklist.$algorithm-$bits
--
--@output
--|_ SSH weak key: 2048 f0:58:ce:f4:aa:a4:59:1c:8e:dd:4d:07:44:c8:25:11 (RSA)

require("stdnse")
require("shortport")
require("datafiles")
require("ssh1")

id = "SSH weak key"
author = "Sven Klemm <sven@c3d2.de>"
description = "Show weak SSH hostkeys"
license = "Same as Nmap--See http://nmap.org/book/man-legal.html"
categories = {"default","safe","vuln"}
runlevel = 2.0

hostrule = function( host ) return true end

local weak_keys = {}
local mutex = nmap.mutex(id)

--- read list of weak ssh key fingerprints
local read_datafile = function( db )
  mutex "lock"
  if not weak_keys[db] then
    local file = "openssh.blacklist." .. db

    local status, parsed = datafiles.parse_file( file, {["^[a-f0-9]+"]= "" })
    if status then
      weak_keys[db] = parsed
    else
      weak_keys[db] = {}
      stdnse.print_debug( 3, "No weak SSH key datafile found for %s.", db )
    end
  end
  mutex "done"
end

--- get value of the first entry of a table
local table_first_key = function( t )
  for key in pairs(t) do
    return key
  end
end

--- Derive the fingerprint length of a datafile from the first entry of the
-- table of this algorithm. This allows different fingerprint lengths for 
-- each key algorithm - key length pair.
--@param t table with fingerprints
local fingerprint_length = function( t )
  if table_first_key( t ) then return #table_first_key( t ) end
end

--- check whether this is a known weak key for a specific algorithm
local check_key = function( db, fp )
  if not weak_keys[db] then read_datafile( db ) end
  datalength = fingerprint_length( weak_keys[db] )
  if datalength then
    fp = fp:sub( #fp - datalength + 1, #fp )
    if weak_keys[db][fp] then
      return true
    end
  end
end

local supported_algorithms = {RSA=true,DSA=true}

--- lookup whether key is a known weak key
local known_weak_key = function( key )
  if supported_algorithms[key.algorithm] then
    local db = ("%s-%d"):format( key.algorithm, key.bits )
    return check_key( db, stdnse.tohex( key.fingerprint ) )
  else
    stdnse.print_debug( "Unsupported key algorithm: %s", key.algorithm )
    return false
  end
end

action = function( host )
  if not nmap.registry['SSH Hostkey'] then return end
  local output = {}
  local keys = nmap.registry['SSH Hostkey'][host.ip]

  for _, key in ipairs( keys ) do
    if known_weak_key( key ) then
      table.insert(output, ssh1.fingerprint_hex(key.fingerprint,key.algorithm,key.bits))
    end
  end

  if #output > 0 then
    return table.concat( output, '\n' )
  end
end