{-# LANGUAGE OverloadedStrings #-}

-- Module      : KnownHosts
-- Description : Parses SSH known_hosts file into local data objects
-- Copyright   : (c) Kevin Quick, 2016
-- License     : BSD-3
-- Maintainer  : quick@sparq.org
-- Stability   : stable
-- Portability : POSIX
--
-- Remote hosts specified in the .ssh/known_hosts file can take one of
-- the following forms:
--
--   * 192.168.3.1
--   * ::1
--   * shortname
--   * sys.domain.name
--
-- Any of the above can be enclosed in square brackets and followed by
-- :PORTNUM
--
-- Multiples of any of the above (with or without port specifications)
-- can occur with comma separators, representing equivalent addresses
-- for the line.
--
-- The known_hosts key file can be hashed (ssh-keygen -H), in which
-- case all of the above (including port specifications) are replaced
-- by hash values of the form "|1|partone|part2=".  Note that comma
-- separators are not used for multiple hash values: multiple lines
-- are used instead.
--
-- This library parses all of the above forms, returning an array of
-- 'SSHRemote' specifications.  In addition, this library provides a
-- 'targetAddr' function to get the actual target address from an
-- entry.


module Network.SSH.KnownHosts
    ( readKnownHosts
    , readKnownHostsFile
    , parseRemotes
    , SSHRemoteAddr(..), SSHRemoteAddrs(..)
    , KeyAlgorithm, SSHKey, DNSName, SSHRemote(..)
    , IPv4, IPv6
    , targetAddr
    ) where

import           Data.IP.Internal
import           Data.List (intercalate)
import           Data.Maybe
import           Data.Monoid ((<>))
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import           System.Environment
import           Text.Read


type DNSName = [T.Text]

data SSHRemoteAddr = RemoteV4 IPv4
                   | RemoteV4Port IPv4 Int
                   | RemoteV6 IPv6
                   | RemoteV6Port IPv6 Int
                   | RemoteDNS [T.Text]
                   | RemoteDNSPort [T.Text] Int
                   | RemoteHashed String
                 deriving (Show, Eq)

data SSHRemoteAddrs = SSH1Addr SSHRemoteAddr
                    | SSHAddrs [SSHRemoteAddr]
                    deriving (Show, Eq)

type KeyAlgorithm = String
type SSHKey = String

data SSHRemote = SSHRemote SSHRemoteAddrs KeyAlgorithm SSHKey
               deriving (Show, Eq)

-- | targetAddr returns the target address and port, using the best
-- (most explicit) form of the address in the case where there are
-- multiple addresses for the same remote system.
targetAddr :: SSHRemote -> (T.Text, Int)
targetAddr (SSHRemote (SSH1Addr a) _ _) =
  let addr (RemoteV4 a)        = (T.pack $ show a, 22)
      addr (RemoteV4Port a p)  = (T.pack $ show a, p)
      addr (RemoteV6 a)        = (T.pack $ show a, 22)
      addr (RemoteV6Port a p)  = (T.pack $ show a, p)
      addr (RemoteDNS d)       = (dotPack $ reverse d, 22)
      addr (RemoteDNSPort d p) = (dotPack $ reverse d, p)
      addr (RemoteHashed h)    = ("[{hashed}]",22)  -- unuseable
  in addr a
targetAddr (SSHRemote (SSHAddrs m) a b) =
  targetAddr $ SSHRemote (SSH1Addr $ foldr1 findBest m) a b
  where findBest   best@(RemoteV4Port _ _)  _ = best
        findBest _ best@(RemoteV4Port _ _)    = best
        findBest   best@(RemoteV4 _)        _ = best
        findBest _ best@(RemoteV4 _)          = best
        findBest   best@(RemoteV6Port _ _)  _ = best
        findBest _ best@(RemoteV6Port _ _)    = best
        findBest   best@(RemoteV6 _)        _ = best
        findBest _ best@(RemoteV6 _)          = best
        findBest   best@(RemoteDNSPort _ _) _ = best
        findBest _ best@(RemoteDNSPort _ _)   = best
        findBest   best@(RemoteDNS _)       _ = best
        findBest _ best@(RemoteDNS _)         = best
        findBest best _                       = best

dotPack = T.intercalate (T.pack ".")

splitDNS :: T.Text -> [T.Text]
splitDNS n = let split = T.splitOn "." n
                 dottedQuad = and [ length split == 4
                                  , T.null $ snd $ T.span isNumOrDot n]
                 isNumOrDot = flip elem ("0123456789." :: String)
             in if dottedQuad then [] else reverse split

parseAddr :: T.Text -> SSHRemoteAddrs
parseAddr addrSpec =
  let addrs = T.splitOn "," addrSpec
  in if null (tail addrs) then SSH1Addr (parse1Addr addrSpec)
     else SSHAddrs (map parse1Addr addrs)

parse1Addr :: T.Text -> SSHRemoteAddr
parse1Addr addrSpec =
  let hasPort = T.head addrSpec == '[' && T.head (snd parts) == ']'
      hashedAddr = T.head addrSpec == '|'
      (a',p) = if hasPort
               then (T.tail $ fst parts,
                     read (T.unpack $ T.tail $ T.tail $ snd parts))
               else (addrSpec, 22)
      parts = T.break (== ']') addrSpec
      dns = splitDNS a'
      addr = case readMaybe (T.unpack a') of
               Just a -> if hasPort
                         then RemoteV4Port a p
                         else RemoteV4 a
               Nothing -> case readMaybe (T.unpack a') of
                           Just a -> if hasPort
                                     then RemoteV6Port a p
                                     else RemoteV6 a
                           Nothing -> if hashedAddr
                                      then RemoteHashed $ T.unpack addrSpec
                                      else if hasPort
                                           then RemoteDNSPort dns p
                                           else RemoteDNS dns
  in addr


parseRemote :: T.Text -> SSHRemote
parseRemote line =
  let w = T.words line
      expfmt = "address algorithm key"
      addr = parseAddr $ head w
      alg = T.unpack $ head $ tail w
      hash = T.unpack $ head $ tail $ tail w
  in if length w == 3
     then SSHRemote addr alg hash
     else error $ "Expected \"" <> expfmt <> "\" but got: " <> show line

-- | Passed the specification of remote ssh known hosts and returns
-- the parsed SSHRemote structures representing those hosts.  The
-- input text block is expected to be the raw contents of a
-- known_hosts file (e.g. comparable to a @Data.Text.IO.readFile
-- "$HOME/.ssh/known_hosts@) and the output is the parsed data
-- structures.
parseRemotes :: T.Text -> [SSHRemote]
parseRemotes = map parseRemote . T.lines

-- | Reads the current user's known_hosts file and returns the
-- SSHRemote structures representing the contents of that file (by
-- calling 'parseRemotes').
readKnownHosts :: IO [SSHRemote]
readKnownHosts =
  do home <- getEnv "HOME"
     readKnownHostsFile $ home <> "/.ssh/known_hosts"

-- | Reads the specified file and returns the SSHRemote structures
-- representing the contents of that file (by calling 'parseRemotes').
readKnownHostsFile :: String -> IO [SSHRemote]
readKnownHostsFile f = parseRemotes <$> TIO.readFile f