-- hokey.hs: hOpenPGP key tool
-- Copyright © 2013-2014  Clint Adams
--
-- vim: softtabstop=4:shiftwidth=4:expandtab
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as
-- published by the Free Software Foundation, either version 3 of the
-- License, or (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.

{-# LANGUAGE DeriveGeneric #-}

import HOpenPGP.Tools.Common (banner, versioner, warranty)
import Codec.Encryption.OpenPGP.Fingerprint (fingerprint, eightOctetKeyID)
import Codec.Encryption.OpenPGP.KeyInfo (keySize)
import Codec.Encryption.OpenPGP.Signatures (verifyTKWith, verifySigWith, verifyAgainstKeys)
import Codec.Encryption.OpenPGP.Types
import Control.Applicative ((<$>),(<*>))
import Control.Arrow ((***), second)
import Control.Lens ((^.), (^?!), ix)
import Control.Monad.Trans.Writer.Lazy (execWriter, tell)
import qualified Crypto.Hash.SHA3 as SHA3
import qualified Data.Aeson as A
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BC8
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Base16 as Base16
import Data.Conduit (($=),($$), runResourceT)
import qualified Data.Conduit.Binary as CB
import Data.Conduit.Cereal (conduitGet)
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Keyring (conduitToTKsDropping)
import Data.IxSet (empty, insert)
import Data.List (unfoldr, elemIndex, findIndex, sortBy, intercalate)
import Data.Maybe (fromMaybe)
import Data.Monoid ((<>), mconcat)
import Data.Ord (comparing, Down(..))
import Data.Serialize (get)
import Data.Time.Clock (UTCTime)
import Data.Time.Clock.POSIX (getPOSIXTime, posixSecondsToUTCTime, POSIXTime)
import Data.Time.Format (formatTime)
import GHC.Generics

import Options.Applicative.Builder (command, footer, header, help, info, long, metavar, prefs, progDesc, showDefault, showHelpOnError, option, subparser, value)
import Options.Applicative.Extra (customExecParser, helper)
import Options.Applicative.Types (Parser)

import System.IO (Handle, hFlush, hPutStrLn, stderr, stdin, hSetBuffering, BufferMode(..))
import System.Locale (defaultTimeLocale)

import Text.PrettyPrint.ANSI.Leijen (colon, green, indent, linebreak, list, putDoc, red, text, yellow, (<+>), Doc) -- need 0.6.7 for hardline

data KAS = KAS {
    pubkeyalgo :: PubKeyAlgorithm
  , pubkeysize :: Int
} deriving Generic

data KeyReport = KeyReport {
    keyStatus :: String
  , keyFingerprint :: TwentyOctetFingerprint
  , keyVer :: KeyVersion
  , keyCreationTime :: TimeStamp
  , keyAlgorithmAndSize :: KAS
  , keyUIDsAndUAts :: [(String, UIDReport)]
} deriving Generic

data UIDReport = UIDReport {
    uidSelfSigHashAlgorithms :: [HashAlgorithm]
  , uidPreferredHashAlgorithms :: [[HashAlgorithm]]
  , uidKeyExpirationTimes :: [[TimeStamp]]
} deriving Generic

instance A.ToJSON KAS
instance A.ToJSON KeyReport
instance A.ToJSON UIDReport
instance A.ToJSON HashAlgorithm where
    toJSON = A.toJSON . show
instance A.ToJSON PubKeyAlgorithm where
    toJSON = A.toJSON . show
instance A.ToJSON KeyVersion where
    toJSON DeprecatedV3 = A.toJSON (3 :: Int)
    toJSON V4 = A.toJSON (4 :: Int)
instance A.ToJSON TwentyOctetFingerprint where
    toJSON = A.toJSON . show

checkKey :: Maybe UTCTime -> TK -> KeyReport
checkKey mct key = KeyReport (either (const "not-good") (const "good") processedTK) (fingerprint $ key^.tkPKP) (key^.tkPKP^.keyVersion) (key^.tkPKP^.timestamp) (KAS (key^.tkPKP^.pkalgo) (keySize (key^.tkPKP^.pubkey))) (map (second uidr) (processedOrOrig^.tkUIDs) ++ map (uatspsToString *** uidr) (processedOrOrig^.tkUAts))
    where
        processedOrOrig = either (const key) id processedTK
        processedTK = verifyTKWith (verifySigWith (verifyAgainstKeys [key])) mct . stripOlderSigs . stripOtherSigs $ key
        sigissuer (SigVOther 2 _) = OtherSigSub 666 B.empty -- this is dumb
        sigissuer (SigV3 {}) = OtherSigSub 666 B.empty -- this is dumb
        sigissuer (SigV4 _ _ _ _ xs _ _) = xs^?!ix 0^.sspPayload -- this is a horrible stack of stupid assumptions
        sigissuer (SigVOther _ _) = error "We're in the future."
        sigissuer _ = error "WTF"
        eoki = eightOctetKeyID . _tkPKP
        hashAlgo (SigV4 _ _ x _ _ _ _) = x
        phas (SigV4 _ _ _ xs _ _ _) = concatMap (\(SigSubPacket _ (PreferredHashAlgorithms x)) -> x) $ filter isPHA xs
        isPHA (SigSubPacket _ (PreferredHashAlgorithms _)) = True
        isPHA _ = False
        kets (SigV4 _ _ _ xs _ _ _) = map (\(SigSubPacket _ (KeyExpirationTime x)) -> x) $ filter isKET xs
        isKET (SigSubPacket _ (KeyExpirationTime _)) = True
        isKET _ = False
        has = map hashAlgo . alleged
        isCT (SigSubPacket _ (SigCreationTime _)) = True
        isCT _ = False
        sigcts (SigV4 _ _ _ xs _ _ _) = map (\(SigSubPacket _ (SigCreationTime x)) -> x) $ filter isCT xs
        alleged = filter (\x -> sigissuer x == Issuer (eoki key))
        newest = take 1 . sortBy (comparing (Down . take 1 . sigcts)) -- FIXME: this is terrible
        stripOtherSigs tk = tk {
            _tkUIDs = map (second alleged) (_tkUIDs tk)
          , _tkUAts = map (second alleged) (_tkUAts tk)
        }
        stripOlderSigs tk = tk {
            _tkUIDs = map (second newest) (_tkUIDs tk)
          , _tkUAts = map (second newest) (_tkUAts tk)
        }
        uatspsToString us = "<uat:[" ++ intercalate "," (map uaspToString us) ++ "]>"
        uaspToString (ImageAttribute hdr d) = hdrToString hdr ++ ':':show (B.length d) ++ ':':BC8.unpack (Base16.encode (SHA3.hash 48 d))
        uaspToString (OtherUASub t d) = "other-" ++ show t ++ ':':show (B.length d) ++ ':':BC8.unpack (Base16.encode (SHA3.hash 48 d))
        hdrToString (ImageHV1 JPEG) = "jpeg"
        hdrToString (ImageHV1 fmt) = "image-" ++ show (fromFVal fmt)
        uidr sps = UIDReport (has sps) (map phas sps) (map kets sps)

prettyKeyReport :: POSIXTime -> TK -> Doc
prettyKeyReport cpt key = do
 let keyReport = checkKey (Just (posixSecondsToUTCTime cpt)) key
 execWriter $ tell
   ( linebreak <> text "Key has potential validity:" <+> text (keyStatus keyReport)
  <> linebreak <> text "Key has fingerprint:" <+> text (show (keyFingerprint keyReport))
  <> linebreak <> text "Checking to see if key is OpenPGPv4:" <+> colorIf (green,red) (==V4) (keyVer keyReport)
  <> linebreak <> (\kas -> text "Checking to see if key is RSA or DSA (>= 2048-bit):" <+> colorIf (green,yellow) (==RSA) (pubkeyalgo kas) <+> colorIf3 (green,yellow,red) (>= 3072) (>=2048) (pubkeysize kas)) (keyAlgorithmAndSize keyReport)
  <> linebreak <> text "Checking user-ID- and user-attribute-related items:"
  <> mconcat (map (uidtrip (keyCreationTime keyReport)) (keyUIDsAndUAts keyReport))
  <> linebreak
   )
    where
        colorIf (y,n) p x = ((if p x then y else n) . text . show) x
        colorIf3 (g,y,r) p1 p2 x = ((if p1 x then g else (if p2 x then y else r)) . text . show) x
        uidtrip ts (u, ur) = linebreak <> indent 2 (text u) <> colon
          <> linebreak <> indent 4 (text "Self-sig hash algorithms" <> colon <+> (listHAs . uidSelfSigHashAlgorithms) ur)
          <> linebreak <> indent 4 (text "Preferred hash algorithms" <> colon <+> mconcat (map ((linebreak <>) . indent 2 . listPHAs) (uidPreferredHashAlgorithms ur)))
          <> linebreak <> indent 4 (text "Key expiration times" <> colon <+> mconcat (map ((linebreak <>) . indent 2 . listKETs cpt ts) (uidKeyExpirationTimes ur)))
        listHAs = list . map (\x -> ((if x `elem` [DeprecatedMD5, SHA1] then red else id) . text . show) x)
        listPHAs x = (if fSHA2Family x < ei DeprecatedMD5 x && fSHA2Family x < ei SHA1 x then green else red) . list . map (text . show) $ x
        listKETs ct ts x = colorExpiration ct ts x . list . map (text . keyExp ts) $ x
        fSHA2Family = fi (`elem` [SHA512,SHA384,SHA256,SHA224])
        ei x y = fromMaybe maxBound (elemIndex x y)
        fi x y = fromMaybe maxBound (findIndex x y)
        colorExpiration ct ts kes
            | null kes  = red
            | any (\ke -> realToFrac ts + realToFrac ke < ct) kes = red
            | any (\ke -> realToFrac ts + realToFrac ke > ct + (5*31557600)) kes = yellow
            | otherwise = green
        keyExp ts ke = durationPrettyPrinter ke ++ " = " ++ formatTime defaultTimeLocale "%c" (posixSecondsToUTCTime (realToFrac ts + realToFrac ke))

jsonReport :: POSIXTime -> TK -> BL.ByteString
jsonReport _ key = A.encode (checkKey Nothing key) -- FIXME: pass time when it matters in the JSON

-- this does not have the same sense of calendar anyone else might
durationPrettyPrinter :: TimeStamp -> String
durationPrettyPrinter = concat . unfoldr durU
    where
        durU x
         | x >= 31557600 = Just ((++"y") . show $ x `div` 31557600, x `mod` 31557600)
         | x >= 2629800 = Just ((++"m") . show $ x `div` 2629800, x `mod` 2629800)
         | x >= 86400 = Just ((++"d") . show $ x `div` 86400, x `mod` 86400)
         | x > 0 = Just ((++"s") . show $ x, 0)
         | otherwise = Nothing

data OutputFormat = Pretty | JSON
    deriving (Eq, Read, Show)

data Options = Options {
    outputFormat :: OutputFormat
}

data Command = Lint Options

lintO :: Parser Options
lintO = Options
    <$> option
        ( long "output-format"
       <> metavar "FORMAT"
       <> value Pretty
       <> showDefault
       <> help "output format" )

dispatch :: Command -> IO ()
dispatch (Lint o) = banner' stderr >> hFlush stderr >> doLint o

main :: IO ()
main = do
    hSetBuffering stderr LineBuffering
    customExecParser (prefs showHelpOnError) (info (helper <*> versioner <*> cmd) (header (banner "hokey") <> progDesc "hOpenPGP Key utility" <> footer (warranty "hokey"))) >>= dispatch

cmd :: Parser Command
cmd = subparser
    ( command "lint" (info ( Lint <$> lintO) ( progDesc "check key(s) for 'best practices'" )))

doLint :: Options -> IO ()
doLint o = do
    cpt <- getPOSIXTime
    keys <- runResourceT $ CB.sourceHandle stdin $= conduitGet get $= conduitToTKsDropping $$ CL.consume
    output (outputFormat o) cpt keys
    where
        output Pretty cpt = mapM_ (putDoc . prettyKeyReport cpt)
        output JSON cpt = BL.putStr . BL.concat . map (jsonReport cpt)

banner' :: Handle -> IO ()
banner' h = hPutStrLn h (banner "hokey" ++ "\n" ++ warranty "hokey")