-- hkt.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, keyMatchesFingerprint, keyMatchesEightOctetKeyId, keyMatchesUIDSubString)
import HOpenPGP.Tools.Parser (parseTKExp)
import Codec.Encryption.OpenPGP.Fingerprint (fingerprint, eightOctetKeyID)
import Codec.Encryption.OpenPGP.KeyInfo (pubkeySize, pkalgoAbbrev)
import Codec.Encryption.OpenPGP.KeySelection (parseEightOctetKeyId, parseFingerprint)
import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Signatures (verifyTKWith, verifySigWith, verifyAgainstKeyring)
import Codec.Encryption.OpenPGP.Types
import Control.Applicative ((<$>),(<*>), optional, (<|>), pure)
import Control.Arrow ((&&&))
import Control.Lens ((^.), _1, _2, (^..))
import Control.Monad.Trans.Resource (runResourceT, MonadResource)
import qualified Control.Monad.Trans.State.Lazy as S
import qualified Data.Aeson as A
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Conduit (($=),($$), Source)
import qualified Data.Conduit.Binary as CB
import Data.Conduit.Cereal (conduitGet)
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Filter (conduitTKFilter, FilterPredicates(RTKFilterPredicate))
import Data.Conduit.OpenPGP.Keyring (conduitToTKsDropping, sinkKeyringMap)
import Data.Data.Lens (biplate)
import Data.Either (rights)
import qualified Data.IxSet as IxSet
import Data.Graph.Inductive.Graph (Graph(mkGraph), emap, Path, prettyPrint)
import Data.Graph.Inductive.PatriciaTree (Gr)
import Data.Graph.Inductive.Query.SP (sp)
import Data.GraphViz (graphToDot, nonClusteredParams, GraphvizParams(..))
import Data.GraphViz.Attributes (toLabel)
import Data.GraphViz.Types (printDotGraph)
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HashMap
import Data.List (nub, sort)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, mapMaybe, listToMaybe)
import Data.Monoid ((<>), mconcat)
import Data.Serialize (get, put, runPut)
import qualified Data.Text as T
import qualified Data.Text.Lazy.IO as TLIO
import Data.Time.Clock.POSIX (getPOSIXTime, posixSecondsToUTCTime)
import Data.Traversable (traverse)
import Data.Tuple (swap)
import qualified Data.Yaml as Y
import GHC.Generics
import System.Directory (getHomeDirectory)

import Options.Applicative.Builder (argument, auto, command, footerDoc, headerDoc, help, helpDoc, info, long, metavar, option, prefs, progDesc, showDefault, showHelpOnError, str, strOption, subparser, switch, value)
import Options.Applicative.Extra (customExecParser, helper)
import Options.Applicative.Types (Parser)

import System.IO (Handle, hFlush, stderr, hSetBuffering, BufferMode(..))
import Text.PrettyPrint.ANSI.Leijen ((<+>), (</>), hardline, hPutDoc, list, putDoc, text)

grabMatchingKeysConduit :: MonadResource m => FilePath -> Bool -> String -> Source m TK
grabMatchingKeysConduit fp filt srch = CB.sourceFile fp $= conduitGet get $= conduitToTKsDropping $= (if filt then conduitTKFilter ufp else CL.filter matchAny)
    where
        matchAny tk = either (const False) id $ fmap (keyMatchesFingerprint True tk) efp <|> fmap (keyMatchesEightOctetKeyId True tk . Right) eeok <|> return (keyMatchesUIDSubString srch tk)
        efp = parseFingerprint . T.pack $ srch
        eeok = parseEightOctetKeyId . T.pack $ srch
        ufp = RTKFilterPredicate (parseE srch)
        parseE e = either (error . ("filter parse error: "++)) id (parseTKExp e) -- this should be more specialized

grabMatchingKeys :: FilePath -> Bool -> String -> IO [TK]
grabMatchingKeys fp filt srch = runResourceT $ grabMatchingKeysConduit fp filt srch $$ CL.consume

grabMatchingKeysKeyring :: FilePath -> Bool -> String -> IO Keyring
grabMatchingKeysKeyring fp filt srch = runResourceT $ grabMatchingKeysConduit fp filt srch $$ sinkKeyringMap

data Key = Key {
    keysize :: Maybe Int
  , keyalgo :: String
  , keyalgoabbreviation :: String
  , keyid :: String
  , fpr :: String
} deriving Generic

data TKey = TKey {
    publickey :: Key
  , uids :: [String]
  , subkeys :: [Key]
} deriving Generic

instance A.ToJSON Key
instance A.ToJSON TKey

tkToTKey :: TK -> TKey
tkToTKey tk = TKey {
    publickey = mkey (tk^.tkKey._1)
  , uids = tk^.tkUIDs^..traverse._1
  , subkeys = map (mkey . \(PublicSubkeyPkt x,_) -> x) (tk^.tkSubs)
}
    where
        mkey = Key <$> either (const Nothing) Just . pubkeySize . _pubkey
                   <*> show . _pkalgo
                   <*> pkalgoAbbrev . _pkalgo
                   <*> either (const "unknown") show . eightOctetKeyID
                   <*> show . fingerprint

showTKey :: TKey -> IO ()
showTKey key = putDoc $
    text "pub  " <+> sizeabbrevkeyid (publickey key) <> hardline <>
    mconcat (map (\x -> text "uid                           " <+> text x <> hardline) (uids key)) <>
    mconcat (map (\x -> text "sub  " <+> sizeabbrevkeyid x <> hardline) (subkeys key)) <>
    hardline
    where
        sizeabbrevkeyid k = text (maybe "unknown" show (keysize k)) <> text (keyalgoabbreviation k) <> text "/0x" <> text (keyid k)

data Options = Options {
    keyring :: String
  , graphOutputFormat :: GraphOutputFormat
  , pathsOutputFormat :: PathsOutputFormat
  , targetIsFilter :: Bool
  , target1 :: String
  , target2 :: String
  , target3 :: String
}

data Command = CmdList Options | CmdExportPubkeys Options | CmdGraph Options | CmdFindPaths Options

data GraphOutputFormat = GraphViz | LossyPretty
    deriving (Bounded, Enum, Eq, Read, Show)

data PathsOutputFormat = Unstructured | JSON | YAML
    deriving (Eq, Read, Show)

listO :: String -> Parser Options
listO homedir = Options
    <$> (fromMaybe (homedir ++ "/.gnupg/pubring.gpg") <$> optional (strOption
        ( long "keyring"
       <> metavar "FILE"
       <> help "file containing keyring" )))
    <*> pure GraphViz -- unused
    <*> option auto
        ( long "output-format"
       <> metavar "FORMAT"
       <> value Unstructured
       <> showDefault
       <> help "output format" )
    <*> switch ( long "filter" <> help "treat target as filter" )
    <*> argument str ( metavar "TARGET" <> targetHelp )
    <*> pure ""
    <*> pure ""
    where
        targetHelp = helpDoc . Just $ text "target (which keys to output)*"

graphO :: String -> Parser Options
graphO homedir = Options
    <$> (fromMaybe (homedir ++ "/.gnupg/pubring.gpg") <$> optional (strOption
        ( long "keyring"
       <> metavar "FILE"
       <> help "file containing keyring" )))
    <*> option auto
        ( long "output-format"
       <> metavar "FORMAT"
       <> value GraphViz
       <> showDefault
       <> ofhelp )
    <*> pure Unstructured -- unused
    <*> switch ( long "filter" <> help "treat target as filter" )
    <*> argument str ( metavar "TARGET" <> targetHelp )
    <*> pure ""
    <*> pure ""
    where
        ofhelp = helpDoc . Just $ text "output format" <> hardline <> list (map (text . show) ofchoices)
        ofchoices = [minBound..maxBound] :: [GraphOutputFormat]
        targetHelp = helpDoc . Just $ text "target (which keys to graph)*"

findPathsO :: String -> Parser Options
findPathsO homedir = Options
    <$> (fromMaybe (homedir ++ "/.gnupg/pubring.gpg") <$> optional (strOption
        ( long "keyring"
       <> metavar "FILE"
       <> help "file containing keyring" )))
    <*> pure GraphViz -- unused
    <*> option auto
        ( long "output-format"
       <> metavar "FORMAT"
       <> value Unstructured
       <> showDefault
       <> help "output format" )
    <*> switch ( long "filter" <> help "treat targets as filter" )
    <*> argument str ( metavar "TARGET-SET" <> targetHelp )
    <*> argument str ( metavar "FROM-KEYS" <> fromHelp )
    <*> argument str ( metavar "TO-KEYS" <> toHelp )
    where
        targetHelp = helpDoc . Just $ text "target (which keys to use in pathfinding)*"
        fromHelp = helpDoc . Just $ text "from (which keys to use for the source of paths)*"
        toHelp = helpDoc . Just $ text "to (which keys to use for the destinations of paths)*"

dispatch :: Command -> IO ()
dispatch (CmdList o) = banner' stderr >> hFlush stderr >> doList o
dispatch (CmdExportPubkeys o) = banner' stderr >> hFlush stderr >> doExportPubkeys o
dispatch (CmdGraph o) = banner' stderr >> hFlush stderr >> doGraph o
dispatch (CmdFindPaths o) = banner' stderr >> hFlush stderr >> doFindPaths o

main :: IO ()
main = do
    hSetBuffering stderr LineBuffering
    homedir <- getHomeDirectory
    customExecParser (prefs showHelpOnError) (info (helper <*> versioner <*> cmd homedir) (headerDoc (Just (banner "hkt")) <> progDesc "hOpenPGP Keyring Tool" <> footerDoc (Just (warranty "hkt")))) >>= dispatch

cmd :: String -> Parser Command
cmd homedir = subparser
    ( command "export-pubkeys" (info ( CmdExportPubkeys <$> listO homedir) ( progDesc "export matching keys to stdout" <> footerDoc (Just foot) ))
   <> command "findpaths" (info ( CmdFindPaths <$> findPathsO homedir) ( progDesc "find short paths between keys" <> footerDoc (Just foot) ))
   <> command "graph" (info ( CmdGraph <$> graphO homedir) ( progDesc "graph certifications" <> footerDoc (Just foot) ))
   <> command "list" (info ( CmdList <$> listO homedir) ( progDesc "list matching keys" <> footerDoc (Just foot) ))
    )
    where
        foot = hardline <> text "*if --filter is not specified, this must be"
               </> text "a fingerprint," </> text "an eight-octet key ID,"
               </> text "or a substring of a UID (including an empty string)" <> hardline
               <> text "if --filter is specified, it must be"
               </> text "something in filter syntax (see source)."

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

doList :: Options -> IO ()
doList o = do
    keys' <- grabMatchingKeys (keyring o) (targetIsFilter o) (target1 o)
    let keys = map tkToTKey keys'
    case pathsOutputFormat o of
        Unstructured -> mapM_ showTKey keys
        JSON -> BL.putStr . A.encode $ keys
        YAML -> B.putStr . Y.encode $ keys
    putStrLn ""

doExportPubkeys :: Options -> IO ()
doExportPubkeys o = do
    keys <- grabMatchingKeys (keyring o) (targetIsFilter o) (target1 o)
    mapM_ (B.putStr . putTK') keys
    where
        putTK' key = runPut $ do
            put (PublicKey (key^.tkKey._1))
            mapM_ (put . Signature) (_tkRevs key)
            mapM_ putUid' (_tkUIDs key)
            mapM_ putUat' (_tkUAts key)
            mapM_ putSub' (_tkSubs key)
        putUid' (u, sps) = put (UserId u) >> mapM_ (put . Signature) sps
        putUat' (us, sps) = put (UserAttribute us) >> mapM_ (put . Signature) sps
        putSub' (p, sps) = put p >> mapM_ (put . Signature) sps

doGraph :: Options -> IO ()
doGraph o = do
    cpt <- getPOSIXTime
    kr <- grabMatchingKeysKeyring (keyring o) (targetIsFilter o) (target1 o)
    let g = buildKeyGraph ((buildMaps &&& id) (rights (map (verifyTKWith (verifySigWith (verifyAgainstKeyring kr)) (Just (posixSecondsToUTCTime cpt))) (IxSet.toList kr))))
    case graphOutputFormat o of
        LossyPretty -> prettyPrint g
        GraphViz -> TLIO.putStrLn . printDotGraph . graphToDot nonClusteredLabeledNodesParams $ g
    where
        nonClusteredLabeledNodesParams = nonClusteredParams { fmtNode = \(_,l) -> [toLabel $ show l] }

buildMaps :: [TK] -> (KeyMaps, Int)
buildMaps ks = S.execState (mapM_ mapsInsertions ks) (KeyMaps HashMap.empty HashMap.empty HashMap.empty, 0)

-- FIXME: this presumes no keyID collisions in the input
data KeyMaps = KeyMaps {
    _k2f :: HashMap EightOctetKeyId TwentyOctetFingerprint
  , _f2i :: HashMap TwentyOctetFingerprint Int
  , _i2f :: HashMap Int TwentyOctetFingerprint
}

mapsInsertions :: TK -> S.State (KeyMaps, Int) ()
mapsInsertions tk = do
    (KeyMaps k2f f2i i2f, i) <- S.get
    let fp = fingerprint (tk^.tkKey._1)
        keyids = rights . map eightOctetKeyID $ (tk ^.. biplate :: [PKPayload])
        i' = i + 1
        k2f' = foldr (\k m -> HashMap.insert k fp m) k2f keyids
        f2i' = HashMap.insert fp i' f2i
        i2f' = HashMap.insert i' fp i2f
    S.put (KeyMaps k2f' f2i' i2f', i')

buildKeyGraph :: ((KeyMaps, Int), [TK]) -> Gr TwentyOctetFingerprint HashAlgorithm
buildKeyGraph ((KeyMaps k2f f2i _, _), ks) = mkGraph nodes edges
    where
        nodes = map swap . HashMap.toList $ f2i
        edges = filter (not . samesies) . nub . sort . concatMap tkToEdges $ ks
        tkToEdges tk = map (\(ha, i) -> (source i, target tk, ha)) (mapMaybe (fakejoin . (hashAlgo &&& sigissuer)) (sigs tk))
        target tk = fromMaybe (error "Epic fail") (HashMap.lookup (fingerprint (tk^.tkKey._1)) f2i)
        source i = fromMaybe (-1) (HashMap.lookup i k2f >>= flip HashMap.lookup f2i)
        fakejoin (x, y) = fmap ((,) x) y
        sigs tk = concat ((tk^..tkUIDs.traverse._2) ++ (tk^..tkUAts.traverse._2))
        samesies (x,y,_) = x == y

data PaF = PaF {
    certPaths :: [Path]
  , keyFingerprints :: Map String TwentyOctetFingerprint
} deriving Generic

instance A.ToJSON PaF
instance A.ToJSON TwentyOctetFingerprint where
    toJSON = A.toJSON . show

doFindPaths :: Options -> IO ()
doFindPaths o = do
    cpt <- getPOSIXTime
    kr <- grabMatchingKeysKeyring (keyring o) (targetIsFilter o) (target1 o)
    -- FIXME: clean this up
    keys1 <- runResourceT $ CL.sourceList (IxSet.toList kr) $= (if filt then conduitTKFilter (ufpt (target2 o)) else CL.filter (matchAny (target2 o))) $$ CL.consume
    keys2 <- runResourceT $ CL.sourceList (IxSet.toList kr) $= (if filt then conduitTKFilter (ufpt (target3 o)) else CL.filter (matchAny (target3 o))) $$ CL.consume
    let ((KeyMaps k2f f2i i2f, i), ks) = (buildMaps &&& id) (rights (map (verifyTKWith (verifySigWith (verifyAgainstKeyring kr)) (Just (posixSecondsToUTCTime cpt))) (IxSet.toList kr)))
        keygraph = buildKeyGraph ((KeyMaps k2f f2i i2f, i), ks)
        keysToIs = mapMaybe (\x -> HashMap.lookup (fingerprint (x^.tkKey._1)) f2i)
        froms = keysToIs keys1
        tos = keysToIs keys2
        combos = froms >>= \f -> tos >>= \t -> return (f,t)
        paths = map (\(x,y) -> sp x y (emap (const (1.0 :: Double)) keygraph)) combos
        paf = PaF paths (Map.fromList (mapMaybe (\x -> HashMap.lookup x i2f >>= \y -> return (show x,y)) (nub (sort (concat paths)))))
    case pathsOutputFormat o of
        Unstructured -> do -- FIXME: do something about this
                           putStrLn . unlines $ map (show . ((,) =<< length)) paths
                           putStrLn . unlines $ map (\x -> maybe (show x) show $ HashMap.lookup x i2f >>= \y -> return (x, y)) (nub (sort (concat paths)))
        JSON -> BL.putStr . A.encode $ paf
        YAML -> B.putStr . Y.encode $ paf
    putStrLn ""
    where  -- FIXME: deduplicate this
        filt = targetIsFilter o
        matchAny srch tk = either (const False) id $ fmap (keyMatchesFingerprint True tk) (efp srch) <|> fmap (keyMatchesEightOctetKeyId True tk . Right) (eeok srch) <|> return (keyMatchesUIDSubString srch tk)
        efp srch = parseFingerprint . T.pack $ srch
        eeok srch = parseEightOctetKeyId . T.pack $ srch
        ufpt srch = RTKFilterPredicate (parseE srch)
        parseE e = either (error . ("filter parse error: "++)) id (parseTKExp e) -- this should be more specialized

-- FIXME: deduplicate the following code
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"