-- hkt.hs: hOpenPGP key tool
-- Copyright © 2013-2021  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 Codec.Encryption.OpenPGP.Fingerprint (eightOctetKeyID, fingerprint)
import Codec.Encryption.OpenPGP.KeyInfo (pkalgoAbbrev, pubkeySize)
import Codec.Encryption.OpenPGP.KeySelection
  ( parseEightOctetKeyId
  , parseFingerprint
  )
import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Signatures
  ( verifyAgainstKeyring
  , verifySigWith
  , verifyTKWith
  )
import Codec.Encryption.OpenPGP.Types
import Control.Applicative ((<|>), optional)
import Control.Arrow ((&&&))
import Control.Lens ((^.), (^..), _1, _2)
import Control.Monad.Trans.Resource (MonadResource, MonadThrow)
import qualified Data.Aeson as A
import Data.Binary (get, put)
import Data.Binary.Put (runPut)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Conduit (ConduitM, (.|), runConduitRes)
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Filter
  ( FilterPredicates(RTKFilterPredicate)
  , conduitTKFilter
  )
import Data.Conduit.OpenPGP.Keyring (conduitToTKsDropping, sinkKeyringMap)
import Data.Conduit.Serialization.Binary (conduitGet)
import Data.Data.Lens (biplate)
import Data.Either (rights)
import Data.Graph.Inductive.Graph (Graph(mkGraph), Path, emap, prettyPrint)
import Data.Graph.Inductive.PatriciaTree (Gr)
import Data.Graph.Inductive.Query.SP (sp)
import Data.GraphViz (GraphvizParams(..), graphToDot, nonClusteredParams)
import Data.GraphViz.Attributes (toLabel)
import Data.GraphViz.Types (printDotGraph)
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HashMap
import qualified Data.IxSet.Typed as IxSet
import Data.List (nub, sort)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, listToMaybe, mapMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Lazy.IO as TLIO
import Data.Time.Clock.POSIX (getPOSIXTime, posixSecondsToUTCTime)
import Data.Tuple (swap)
import qualified Data.Yaml as Y
import GHC.Generics
import HOpenPGP.Tools.Common
  ( banner
  , keyMatchesEightOctetKeyId
  , keyMatchesFingerprint
  , keyMatchesUIDSubString
  , versioner
  , warranty
  )
import HOpenPGP.Tools.Parser (parseTKExp)
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
  , switch
  , value
  )
import Options.Applicative.Extra (customExecParser, helper, hsubparser)
import Options.Applicative.Types (Parser)

import Data.Text.Prettyprint.Convert.AnsiWlPprint (toAnsiWlPprint)
import Data.Text.Prettyprint.Doc ((<+>), fillSep, hardline, list, pretty)
import Data.Text.Prettyprint.Doc.Render.Text (hPutDoc, putDoc)
import System.IO (BufferMode(..), Handle, hFlush, hSetBuffering, stderr)

grabMatchingKeysConduit ::
     (MonadResource m, MonadThrow m)
  => FilePath
  -> Bool
  -> Text
  -> ConduitM () TK m ()
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 srch
    eeok = parseEightOctetKeyId srch
    ufp = RTKFilterPredicate (parseE srch)
    parseE =
      either (error . ("filter parse error: " ++)) id . parseTKExp . T.unpack -- this should be more specialized

grabMatchingKeys :: FilePath -> Bool -> Text -> IO [TK]
grabMatchingKeys fp filt srch =
  runConduitRes $ grabMatchingKeysConduit fp filt srch .| CL.consume

grabMatchingKeysKeyring :: FilePath -> Bool -> Text -> IO Keyring
grabMatchingKeysKeyring fp filt srch =
  runConduitRes $ grabMatchingKeysConduit fp filt srch .| sinkKeyringMap

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

data TKey =
  TKey
    { publickey :: Key
    , uids :: [Text]
    , 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 <*>
      show .
      pretty .
      fingerprint

showTKey :: TKey -> IO ()
showTKey key =
  putDoc $ pretty "pub  " <+> sizeabbrevkeyid (publickey key) <> hardline <>
  mconcat
    (map
       (\x ->
          pretty "uid                           " <+> pretty (T.unpack x) <>
          hardline)
       (uids key)) <>
  mconcat
    (map (\x -> pretty "sub  " <+> sizeabbrevkeyid x <> hardline) (subkeys key)) <>
  hardline
  where
    sizeabbrevkeyid k =
      pretty (maybe "unknown" show (keysize k)) <>
      pretty (keyalgoabbreviation k) <>
      pretty "/" <>
      pretty (fpr 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 <$>
  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") <*>
  (fromMaybe "" <$> optional (argument str (metavar "TARGET" <> targetHelp))) <*>
  pure "" <*>
  pure ""
  where
    targetHelp =
      helpDoc . Just . toAnsiWlPprint $ pretty "target (which keys to output)*"

graphO :: String -> Parser Options
graphO homedir =
  Options <$>
  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") <*>
  (fromMaybe "" <$> optional (argument str (metavar "TARGET" <> targetHelp))) <*>
  pure "" <*>
  pure ""
  where
    ofhelp =
      helpDoc . Just . toAnsiWlPprint $ pretty "output format" <> hardline <>
      list (map (pretty . show) ofchoices)
    ofchoices = [minBound .. maxBound] :: [GraphOutputFormat]
    targetHelp =
      helpDoc . Just . toAnsiWlPprint $ pretty "target (which keys to graph)*"

findPathsO :: String -> Parser Options
findPathsO homedir =
  Options <$>
  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 . toAnsiWlPprint $
      pretty "target (which keys to use in pathfinding)*"
    fromHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "from (which keys to use for the source of paths)*"
    toHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "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 "hkt" <*> cmd homedir)
       (headerDoc (Just (toAnsiWlPprint (banner "hkt"))) <>
        progDesc "hOpenPGP Keyring Tool" <>
        footerDoc (Just (toAnsiWlPprint (warranty "hkt"))))) >>=
    dispatch

cmd :: String -> Parser Command
cmd homedir =
  hsubparser
    (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 =
      toAnsiWlPprint $ hardline <>
      fillSep
        [ pretty "*if --filter is not specified, this must be"
        , pretty "a fingerprint,"
        , pretty "an eight-octet key ID,"
        , pretty "or a substring of a UID (including an empty string)"
        ] <>
      hardline <>
      fillSep
        [ pretty "if --filter is specified, it must be"
        , pretty "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
  let ttarget1 = T.pack . target1
  keys' <- grabMatchingKeys (keyring o) (targetIsFilter o) (ttarget1 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
  let ttarget1 = T.pack . target1
  keys <- grabMatchingKeys (keyring o) (targetIsFilter o) (ttarget1 o)
  case pathsOutputFormat o of
    Unstructured -> mapM_ (BL.putStr . putTK') keys
    JSON -> BL.putStr . A.encode $ keys
    YAML -> B.putStr . Y.encode $ 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
  let ttarget1 = T.pack . target1
  cpt <- getPOSIXTime
  kr <- grabMatchingKeysKeyring (keyring o) (targetIsFilter o) (ttarget1 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 (pretty l)]}

buildMaps :: [TK] -> (KeyMaps, Int)
buildMaps =
  foldr mapsInsertions (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 -> (KeyMaps, Int) -> (KeyMaps, Int)
mapsInsertions tk (KeyMaps k2f f2i i2f, i) =
  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
   in (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

doFindPaths :: Options -> IO ()
doFindPaths o = do
  let ttarget1 = T.pack . target1
      ttarget2 = T.pack . target2
      ttarget3 = T.pack . target3
  cpt <- getPOSIXTime
  kr <- grabMatchingKeysKeyring (keyring o) (targetIsFilter o) (ttarget1 o)
    -- FIXME: seriously clean this up
  keys1 <-
    runConduitRes $ CL.sourceList (IxSet.toList kr) .|
    (if filt
       then conduitTKFilter (ufpt (ttarget2 o))
       else CL.filter (matchAny (ttarget2 o))) .|
    CL.consume
  keys2 <-
    runConduitRes $ CL.sourceList (IxSet.toList kr) .|
    (if filt
       then conduitTKFilter (ufpt (ttarget3 o))
       else CL.filter (matchAny (ttarget3 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) ->
             fromMaybe [] (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 -- FIXME: do something about this
     -> do
      putStrLn . unlines $ map (show . ((,) =<< length)) paths
      putStrLn . unlines $
        map
          (\x ->
             maybe (show x) show $ HashMap.lookup x i2f >>= \y ->
               return (x, pretty y))
          (nub (sort (concat paths)))
    JSON -> BL.putStr . A.encode $ paf
    YAML -> B.putStr . Y.encode $ paf
  putStrLn ""
  where
    filt = targetIsFilter o
    matchAny srch tk =
      either (const False) id $
      fmap (keyMatchesFingerprint True tk) (parseFingerprint srch) <|>
      fmap
        (keyMatchesEightOctetKeyId True tk . Right)
        (parseEightOctetKeyId srch) <|>
      return (keyMatchesUIDSubString srch tk)
    ufpt srch = RTKFilterPredicate (parseE srch)
    parseE e =
      either (error . ("filter parse error: " ++)) id (parseTKExp (T.unpack 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"
