-- 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 DeriveFunctor, DeriveGeneric, FlexibleInstances #-}

import HOpenPGP.Tools.Common (banner, versioner, warranty)
import Codec.Encryption.OpenPGP.Expirations (getKeyExpirationTimesFromSignature)
import Codec.Encryption.OpenPGP.Fingerprint (fingerprint, eightOctetKeyID)
import Codec.Encryption.OpenPGP.KeyInfo (pubkeySize, pkalgoAbbrev)
import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Signatures (verifyTKWith, verifySigWith, verifyAgainstKeys)
import Codec.Encryption.OpenPGP.Types
import Control.Applicative ((<$>),(<*>), pure)
import Control.Arrow ((***), second)
import Control.Error.Util (hush)
import Control.Lens ((^.), (&), _1, _2, mapped, over)
import Control.Monad.Trans.Resource (runResourceT)
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 (($=),($$))
import qualified Data.Conduit.Binary as CB
import Data.Conduit.Cereal (conduitGet, conduitPut)
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Keyring (conduitToTKsDropping)
import Data.List (unfoldr, elemIndex, findIndex, nub, sort, sortBy, intercalate)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, mapMaybe, listToMaybe)
import Data.Monoid ((<>), mconcat, Monoid(..))
import Data.Ord (comparing, Down(..))
import Data.Serialize (get, put)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Text.Encoding (decodeUtf8With) -- FIXME: this should be made unnecessary by an API break in hOpenPGP
import Data.Text.Encoding.Error (lenientDecode)
import Data.Time.Clock.POSIX (getPOSIXTime, posixSecondsToUTCTime, POSIXTime)
import Data.Time.Format (formatTime)
import qualified Data.Yaml as Y
import GHC.Generics

import Options.Applicative.Builder (auto, command, footerDoc, headerDoc, helpDoc, 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, stderr, stdin, stdout, hSetBuffering, BufferMode(..))
import System.Locale (defaultTimeLocale)

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

data KAS = KAS {
    pubkeyalgo :: Colored PubKeyAlgorithm
  , pubkeysize :: Colored (Maybe Int)
  , stringrep :: String
} deriving Generic

data Color = Green | Yellow | Red
  deriving (Eq, Generic, Ord)

data Colored a = Colored {
    color :: Maybe Color
  , explanation :: Maybe String
  , val :: a
} deriving (Functor, Generic)

newtype FakeMap a b = FakeMap { unFakeMap :: [(a, b)] }

data KeyReport = KeyReport {
    keyStatus :: String
  , keyFingerprint :: TwentyOctetFingerprint
  , keyVer :: Colored KeyVersion
  , keyCreationTime :: TimeStamp
  , keyAlgorithmAndSize :: KAS
  , keyUIDsAndUAts :: FakeMap String (Colored UIDReport)
  , keyBestOf :: Maybe UIDReport
} deriving Generic

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

data RevocationStatus = RevocationStatus {
    isRevoked :: Bool
  , revocationCode :: String
  , revocationReason :: Text
} deriving Generic

instance A.ToJSON KAS
instance A.ToJSON Color
instance (A.ToJSON a) => A.ToJSON (Colored a)
instance A.ToJSON KeyReport
instance A.ToJSON UIDReport
instance A.ToJSON b => A.ToJSON (FakeMap String b) where
    toJSON = A.toJSON . Map.fromList . unFakeMap
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
instance A.ToJSON RevocationStatus

instance Monoid UIDReport where
    mempty = UIDReport [] [] [] []
    mappend (UIDReport a b c d) (UIDReport a' b' c' d') = UIDReport (a++a') (b++b') (c++c') (d <> d')

checkKey :: Maybe POSIXTime -> TK -> KeyReport
checkKey mpt key = (\x -> x { keyBestOf = populateBestOf x }) KeyReport {
    keyStatus = either id (const "good") processedTK
  , keyFingerprint = key^.tkKey._1 & fingerprint
  , keyVer = key^.tkKey._1.keyVersion & colorizeKV
  , keyCreationTime = key^.tkKey._1.timestamp
  , keyAlgorithmAndSize = KAS {
      pubkeyalgo = key^.tkKey._1.pkalgo & colorizePKA
    , pubkeysize = key^.tkKey._1.pubkey & colorizePKS . pubkeySize
    , stringrep = (key^.tkKey._1.pubkey & either (const "unknown") show . pubkeySize) ++ (key^.tkKey._1.pkalgo & pkalgoAbbrev)
  }
  , keyUIDsAndUAts = FakeMap (map (\(x,y) -> (x, uidr (Just x) y)) (processedOrOrig^.tkUIDs) ++ map (uatspsToString *** uidr Nothing) (processedOrOrig^.tkUAts))
  , keyBestOf = Nothing
}
    where
        processedOrOrig = either (const key) id processedTK
        processedTK = verifyTKWith (verifySigWith (verifyAgainstKeys [key])) (fmap posixSecondsToUTCTime mpt) . stripOlderSigs . stripOtherSigs $ key
        colorizeKV kv = uncurry Colored (if kv == V4 then (Just Green, Nothing) else (Just Red, Just "not a V4 key")) kv
        colorizePKA pka = uncurry Colored (if pka == RSA then (Just Green, Nothing) else (Just Yellow, Just "not an RSA key")) pka
        colorizePKS pks = uncurry Colored (colorizePKS' pks) (hush pks)
        colorizePKS' (Right pks)
          | pks >= 3072 = (Just Green, Nothing)
          | pks >= 2048 = (Just Yellow, Just "Public key size between 2048 and 3072 bits")
          | otherwise = (Just Red, Just "Public key size under 2048 bits")
        colorizePKS' (Left _) = (Just Red, Just "public key algorithm not understood")
        colorizePHAs x = uncurry Colored (if fSHA2Family x < ei DeprecatedMD5 x && fSHA2Family x < ei SHA1 x then (Just Green, Nothing) else (Just Red, Just "weak hash with higher preference")) x
        fSHA2Family = fi (`elem` [SHA512,SHA384,SHA256,SHA224])
        ei x y = fromMaybe maxBound (elemIndex x y)
        fi x y = fromMaybe maxBound (findIndex x y)
        colorizeKETs ct ts kes
            | null kes  = Colored (Just Red) (Just "no expiration set") kes
            | any (\ke -> realToFrac ts + realToFrac ke < ct) kes = Colored (Just Red) (Just "expiration passed") kes
            | any (\ke -> realToFrac ts + realToFrac ke > ct + (5*31557600)) kes = Colored (Just Yellow) (Just "expiration too far in future") kes
            | otherwise = Colored (Just Green) Nothing kes
        eoki pkp
            | pkp^.keyVersion == V4 = hush . eightOctetKeyID $ pkp
            | pkp^.keyVersion == DeprecatedV3 && elem (pkp^.pkalgo) [RSA,DeprecatedRSASignOnly] = hush . eightOctetKeyID $ pkp
            | otherwise = Nothing
        phas (SigV4 _ _ _ xs _ _ _) = colorizePHAs . concatMap (\(SigSubPacket _ (PreferredHashAlgorithms x)) -> x) $ filter isPHA xs
        phas _ = Colored Nothing Nothing []
        isPHA (SigSubPacket _ (PreferredHashAlgorithms _)) = True
        isPHA _ = False
        has = map (colorizeHA . hashAlgo) . alleged
        colorizeHA ha = uncurry Colored (if ha `elem` [DeprecatedMD5, SHA1] then (Just Red, Just "weak hash algorithm") else (Nothing, Nothing)) ha
        isCT (SigSubPacket _ (SigCreationTime _)) = True
        isCT _ = False
        sigcts (SigV4 _ _ _ xs _ _ _) = map (\(SigSubPacket _ (SigCreationTime x)) -> x) $ filter isCT xs
        alleged = filter (\x -> ((==) <$> sigissuer x <*> eoki (key^.tkKey._1)) == Just True)
        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 Nothing sps = Colored Nothing Nothing (UIDReport (has sps) (map phas sps) (map (colorizeKETs (fromMaybe 0 mpt) (key^.tkKey._1.timestamp) . getKeyExpirationTimesFromSignature) sps) (findRevocationReason sps)) -- should that be 0?
        uidr (Just u) sps = colorizeUID u (UIDReport (has sps) (map phas sps) (map (colorizeKETs (fromMaybe 0 mpt) (key^.tkKey._1.timestamp) . getKeyExpirationTimesFromSignature) sps) (findRevocationReason sps)) -- should that be 0?
        populateBestOf krep = Just (UIDReport <$> best . uidSelfSigHashAlgorithms <*> best . uidPreferredHashAlgorithms <*> best . uidKeyExpirationTimes <*> pure [] $ mconcat (justTheUIDRs krep))
        justTheUIDRs = map (decolorize . snd) . unFakeMap . keyUIDsAndUAts
        best = take 1 . sortBy (comparing color)
        decolorize (Colored _ _ x) = x
        colorizeUID u
            | '(' `elem` u = Colored (Just Yellow) (Just "parenthesis in uid")  -- FIXME: be more discerning
            | not ('<' `elem` u) = Colored (Just Yellow) (Just "no left angle bracket in uid")  -- FIXME: be more discerning
            | otherwise = Colored Nothing Nothing
        findRevocationReason = concatMap grabReasons . filter isCertRevocationSig
        isCertRevocationSig (SigV4 CertRevocationSig _ _ _ _ _ _) = True
        isCertRevocationSig _ = False
        grabReasons (SigV4 CertRevocationSig _ _ has _ _ _) = mapMaybe (grabReasons' . _sspPayload) has
        grabReasons _ = []
        grabReasons' (ReasonForRevocation a b) = Just (RevocationStatus True (show a) (decodeUtf8With lenientDecode b))
        grabReasons' _ = Nothing

prettyKeyReport :: POSIXTime -> TK -> Doc
prettyKeyReport cpt key = do
 let keyReport = checkKey (Just cpt) key
 execWriter $ tell
   ( linebreak <> text "Key has potential validity:" <+> text (keyStatus keyReport)
  <> linebreak <> text "Key has fingerprint:" <+> text (show (SpacedFingerprint (keyFingerprint keyReport)))
  <> linebreak <> text "Checking to see if key is OpenPGPv4:" <+> coloredToColor (text . show) (keyVer keyReport)
  <> linebreak <> (\kas -> text "Checking to see if key is RSA or DSA (>= 2048-bit):" <+> coloredToColor (text . show) (pubkeyalgo kas) <+> coloredToColor (text . maybe "unknown" show) (pubkeysize kas)) (keyAlgorithmAndSize keyReport)
  <> linebreak <> text "Checking user-ID- and user-attribute-related items:"
  <> mconcat (map (uidtrip (keyCreationTime keyReport) . gottabeabetterway) (unFakeMap (keyUIDsAndUAts keyReport)))
  <> linebreak
   )
    where
        coloredToColor f (Colored (Just Green) _ x) = green (f x)
        coloredToColor f (Colored (Just Yellow) _ x) = yellow (f x)
        coloredToColor f (Colored (Just Red) _ x) = red (f x)
        coloredToColor f (Colored Nothing _ x) = f x
        uidtrip ts (u, ur)
           | null (uidRevocationStatus ur) = linebreak <> indent 2 (coloredToColor text u) <> colon
          <> linebreak <> indent 4 (text "Self-sig hash algorithms" <> colon <+> (list . map (coloredToColor (text . show)) . uidSelfSigHashAlgorithms) ur)
          <> linebreak <> indent 4 (text "Preferred hash algorithms" <> colon <+> mconcat (map ((linebreak <>) . indent 2 . coloredToColor (text . show)) (uidPreferredHashAlgorithms ur)))
          <> linebreak <> indent 4 (text "Key expiration times" <> colon <+> mconcat (map ((linebreak <>) . indent 2 . coloredToColor list . fmap (map (text . keyExp ts))) (uidKeyExpirationTimes ur)))
           | otherwise = linebreak <> indent 2 (coloredToColor text u) <> colon <+> text "[revoked]"
          <> linebreak <> indent 4 (text "Revocation code" <> colon <+> list (map (text . revocationCode) (uidRevocationStatus ur)))
          <> linebreak <> indent 4 (text "Revocation reason" <> colon <+> list (map (text . T.unpack . revocationReason) (uidRevocationStatus ur)))
        keyExp ts ke = durationPrettyPrinter ke ++ " = " ++ formatTime defaultTimeLocale "%c" (posixSecondsToUTCTime (realToFrac ts + realToFrac ke))
        gottabeabetterway (a, Colored x y z) = (Colored x y a, z)

jsonReport :: POSIXTime -> TK -> BL.ByteString
jsonReport ps = A.encode . checkKey (Just ps)

yamlReport :: POSIXTime -> TK -> B.ByteString
yamlReport ps = Y.encode . (:[]) . checkKey (Just ps)

-- 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 LintOutputFormat = Pretty | JSON | YAML
    deriving (Bounded, Enum, Eq, Read, Show)

data LintOptions = LintOptions {
    lintOutputFormat :: LintOutputFormat
}

data Command = CmdLint LintOptions | CmdCanonicalize

lintO :: Parser LintOptions
lintO = LintOptions
    <$> option auto
        ( long "output-format"
       <> metavar "FORMAT"
       <> value Pretty
       <> showDefault
       <> ofHelp
        )
    where
        ofHelp = helpDoc . Just $ text "output format" <> hardline <> list (map (text . show) ofchoices)
        ofchoices = [minBound..maxBound] :: [LintOutputFormat]


dispatch :: Command -> IO ()
dispatch (CmdLint o) = banner' stderr >> hFlush stderr >> doLint o
dispatch (CmdCanonicalize) = banner' stderr >> hFlush stderr >> doCanonicalize

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

cmd :: Parser Command
cmd = subparser
    ( command "canonicalize" (info ( pure CmdCanonicalize) ( progDesc "arrange key components in a canonical ordering" ))
   <> command "lint" (info ( CmdLint <$> lintO) ( progDesc "check key(s) for 'best practices'" ))
    )

doLint :: LintOptions -> IO ()
doLint o = do
    cpt <- getPOSIXTime
    keys <- runResourceT $ CB.sourceHandle stdin $= conduitGet get $= conduitToTKsDropping $$ CL.consume
    output (lintOutputFormat o) cpt keys
    where
        output Pretty cpt = mapM_ (putDoc . prettyKeyReport cpt)
        output JSON cpt = mapM_ (BL.putStr . flip BL.append (BL.singleton 0x0a) . jsonReport cpt)
        output YAML cpt = mapM_ (B.putStr . yamlReport cpt)

doCanonicalize :: IO ()
doCanonicalize = do
    runResourceT $ CB.sourceHandle stdin $= conduitGet get $=
                   conduitToTKsDropping $= CL.map canonicalize $=
                   conduitPut put $$ CB.sinkHandle stdout
    where
        canonicalize (TK k r ui ua s) = TK k
                                           (sort r)
                                           (indepthsort ui)
                                           (indepthsort ua)
                                           (indepthsort s)
        indepthsort :: (Ord a, Ord b) => [(a,[b])] -> [(a,[b])]
        indepthsort = nub . sort . over (mapped._2) sort

banner' :: Handle -> IO ()
banner' h = hPutDoc h (banner "hokey" <> hardline <> warranty "hokey" <> hardline)

sigissuer :: SignaturePayload -> Maybe EightOctetKeyId
getIssuer :: SigSubPacketPayload -> Maybe EightOctetKeyId
hashAlgo :: SignaturePayload -> HashAlgorithm

sigissuer (SigVOther 2 _) = Nothing
sigissuer (SigV3 {}) = Nothing
sigissuer (SigV4 _ _ _ ys xs _ _) = listToMaybe . mapMaybe (getIssuer . _sspPayload) $ (ys++xs) -- FIXME: what should this be if there are multiple matches?
sigissuer (SigVOther _ _) = error "We're in the future." -- FIXME

getIssuer (Issuer i) = Just i
getIssuer _ = Nothing

hashAlgo (SigV4 _ _ x _ _ _ _) = x
hashAlgo _ = error "V3 sig not supported here"