-- hop.hs: hOpenPGP-stateless OpenPGP (sop) tool
-- Copyright © 2019-2022  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 #-}
{-# LANGUAGE RecordWildCards #-}

import qualified Codec.Encryption.OpenPGP.ASCIIArmor as AA
import Codec.Encryption.OpenPGP.ASCIIArmor.Types (Armor(..), ArmorType(..))
import Codec.Encryption.OpenPGP.Fingerprint (eightOctetKeyID, fingerprint)
import Codec.Encryption.OpenPGP.Ontology (isKUF)
import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Signatures
  ( crossSignSubkeyWithRSA
  , signDataWithRSA
  , signUserIDwithRSA
  , verifyAgainstKeyring
  , verifySigWith
  , verifyTKWith
  )
import Codec.Encryption.OpenPGP.Types
import Control.Applicative ((<|>), optional, some)
import Control.Arrow ((&&&))
import Control.Error.Util (note)
import Control.Lens ((^.), (^..), _1, _2)
import Control.Monad ((>=>), forM_)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.State.Lazy (MonadState, StateT, evalStateT, get, modify)
import Control.Monad.Trans.Resource (MonadResource, MonadThrow)
import qualified Crypto.PubKey.RSA as RSA
import qualified Data.Aeson as A
import qualified Data.Binary as Bin
import Data.Binary.Get (runGet)
import Data.Binary.Put (runPut)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Lazy.Char8 as BLC8
import Data.Conduit (ConduitM, (.|), runConduitRes)
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.Combinators as CC
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Keyring
  ( conduitToTKs
  , conduitToTKsDropping
  , sinkKeyringMap
  )
import Data.Conduit.OpenPGP.Verify (conduitVerify)
import Data.Conduit.Serialization.Binary (conduitGet)
import Data.Data.Lens (biplate)
import Data.Either (fromRight, isRight, rights)
import Data.List (find)
import Data.Maybe (catMaybes, fromJust, fromMaybe, listToMaybe, mapMaybe)
import Data.Monoid ((<>))
import qualified Data.Set as S
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Text.Lazy as TL
import Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime, posixSecondsToUTCTime)
import qualified Data.Vector as V
import Data.Version (showVersion)
import qualified Data.Yaml as Y
import GHC.Generics
import HOpenPGP.Tools.Armor (doDeArmor)
import HOpenPGP.Tools.Common
  ( banner
  , keyMatchesEightOctetKeyId
  , keyMatchesFingerprint
  , keyMatchesUIDSubString
  , versioner
  , warranty
  )
import HOpenPGP.Tools.Parser (parseTKExp)
import HOpenPGP.Tools.TKUtils (processTK)
import Paths_hopenpgp_tools (version)
import System.Exit (exitFailure, exitSuccess)

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

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

data Command
  = VersionC
  | GenerateKeyC KeyGenOptions
  | ExtractCertC ExtractCertOptions
  | SignC SignOptions
  | DeArmorC
  | ArmorC ArmoringOptions

data Options =
  Options
    { keyrings :: [String]
    , outputFormat :: OutputFormat
    , sigFilter :: String
    , sigFile :: String
    , blobFile :: String
    }

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

o :: Parser Options
o =
  Options <$>
  some
    (strOption
       (long "keyring" <>
        short 'k' <> metavar "FILE" <> help "file containing keyring")) <*>
  option
    auto
    (long "output-format" <>
     metavar "FORMAT" <>
     value Unstructured <> showDefault <> help "output format") <*>
  option
    auto
    (long "signature-filter" <>
     metavar "SIGFILTER" <>
     value "sigcreationtime < now" <>
     showDefault <> help "verify only signatures which match filter spec") <*>
  argument str (metavar "SIGNATURE" <> sigHelp) <*>
  argument str (metavar "BLOB" <> blobHelp)
  where
    sigHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "file containing OpenPGP binary signatures"
    blobHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "file containing binary blob to be validated"

dispatch :: POSIXTime -> Command -> IO ()
dispatch cpt o = banner' stderr >> hFlush stderr >> dispatch' cpt o
  where
    dispatch' _ VersionC = doVersion
    dispatch' t (GenerateKeyC o) = doGenerateKey t o
    dispatch' _ (ExtractCertC o) = doExtractCert o
    dispatch' t (SignC o) = doSign t o
    dispatch' _ DeArmorC = doDeArmor
    dispatch' _ (ArmorC o) = doArmor o

main :: IO ()
main = do
  hSetBuffering stderr LineBuffering
  cpt <- getPOSIXTime
  customExecParser
    (prefs showHelpOnError)
    (info
       (helper <*> versioner "hop" <*> cmd)
       (headerDoc (Just (toAnsiWlPprint (banner "hop"))) <>
        progDesc "hOpenPGP Validator Tool" <>
        footerDoc (Just (toAnsiWlPprint (warranty "hop"))))) >>=
    dispatch cpt

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

data Vrf =
  Vrf
    { _vrfmsg :: String
    , _vrfmfpr :: Maybe TwentyOctetFingerprint
    }
  deriving (Eq, Generic, Show)

instance A.ToJSON Vrf

doV :: POSIXTime -> Options -> IO ()
doV cpt o = do
  let allkfiles = sequence_ (map CC.sourceFile (keyrings o))
  krs <-
    runConduitRes $
    allkfiles .| conduitGet Bin.get .| conduitToTKsDropping .| sinkKeyringMap
  sigs <-
    runConduitRes $
    CC.sourceFile (sigFile o) .| conduitGet Bin.get .| CC.filter v4b .|
    CC.sinkVector
  blob <- runConduitRes $ CC.sourceFile (blobFile o) .| CC.sinkLazy
  verifications <-
    runConduitRes $
    CC.yieldMany (V.cons (LiteralDataPkt BinaryData mempty 0 blob) sigs) .|
    conduitVerify krs Nothing .|
    CC.sinkList
  let verifications' = map v2v verifications
  case outputFormat o of
    Unstructured -> mapM_ print verifications'
    JSON -> BL.putStr . A.encode $ verifications'
    YAML -> B.putStr . Y.encode $ verifications'
  putStrLn ""
  case any isRight verifications of
    True -> exitSuccess
    _ -> exitFailure
  where
    v4b (SignaturePkt s@(SigV4 BinarySig _ _ _ _ _ _)) = sf s
    v4b _ = False
    v2v (Left l) = Vrf l Nothing
    v2v (Right v) =
      Vrf "verified signature" (Just (fingerprint (_verificationSigner v)))
    sf = const True

cmd :: Parser Command
cmd =
  hsubparser
    (command "armor" (info (ArmorC <$> aoP) (progDesc "Armor stdin to stdout")) <>
     command
       "dearmor"
       (info (pure DeArmorC) (progDesc "Dearmor stdin to stdout")) <>
     command
       "extract-cert"
       (info
          (ExtractCertC <$> ecoP)
          (progDesc
             "Extract a certificate from a secret key and output it to stdout")) <>
     command
       "generate-key"
       (info
          (GenerateKeyC <$> gkoP)
          (progDesc "Generate a secret key and output it to stdout")) <>
     command
       "sign"
       (info
          (SignC <$> soP)
          (progDesc "Create detached signatures and output them to stdout")) <>
     command
       "version"
       (info (pure VersionC) (progDesc "output hop version to stdout")))

armorTypes :: [(String, Maybe ArmorType)]
armorTypes =
  [ ("auto", Nothing)
  , ("sig", Just ArmorSignature)
  , ("key", Just ArmorPrivateKeyBlock)
  , ("cert", Just ArmorPublicKeyBlock)
  , ("message", Just ArmorMessage)
  ]

armorTypeReader :: String -> Either String (Maybe ArmorType)
armorTypeReader = note "unknown armor type" . flip lookup armorTypes

aoP :: Parser ArmoringOptions
aoP =
  ArmoringOptions <$>
  option
    (eitherReader armorTypeReader)
    (long "label" <> metavar "LABEL" <> armortypeHelp) <*>
  switch
    (long "allow-nested" <>
     help "do the sane thing and unconditionally armor the output")
  where
    armortypeHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "ASCII armor type" <>
      softline <> list (map (pretty . fst) armorTypes)

data ArmoringOptions =
  ArmoringOptions
    { label :: Maybe ArmorType
    , allowNested :: Bool
    }

doArmor :: ArmoringOptions -> IO ()
doArmor ArmoringOptions {..} = do
  m <- runConduitRes $ CB.sourceHandle stdin .| CL.consume
  let lbs = BL.fromChunks m
      armoredAlready = BLC8.pack "-----BEGIN PGP" == BL.take 14 lbs
      label' = guessLabel label (decodeFirstPacket lbs)
      a = Armor label' [] lbs
  BL.putStr $
    if armoredAlready && not allowNested
      then lbs
      else AA.encodeLazy [a]
  where
    decodeFirstPacket = runGet Bin.get
    guessLabel (Just l) _ = l
    guessLabel Nothing (SignaturePkt _) = ArmorSignature
    guessLabel Nothing (SecretKeyPkt _ _) = ArmorPrivateKeyBlock
    guessLabel Nothing (PublicKeyPkt _) = ArmorPublicKeyBlock
    guessLabel Nothing _ = ArmorMessage

doVersion :: IO ()
doVersion = putStrLn $ "hop " ++ showVersion version

gkoP :: Parser KeyGenOptions
gkoP =
  KeyGenOptions <$> switch (long "armor" <> help "armor the output") <*>
  switch (long "no-armor" <> help "don't armor the output") <*>
  strArgument (metavar "USERID" <> help "User ID associated with this key")

data KeyGenOptions =
  KeyGenOptions
    { armor :: Bool
    , noArmor :: Bool
    , userId :: String
    }

doGenerateKey :: POSIXTime -> KeyGenOptions -> IO ()
doGenerateKey pt KeyGenOptions {..} = do
  let ts = ThirtyTwoBitTimeStamp (floor pt)
  sk <- generateSecretKey ts RSA
  s <-
    buildKeyWith sk $ do
      addUserId ts True (T.pack userId)
      addSubkey ts [EncryptStorageKey, EncryptCommunicationsKey]
      addSubkey ts [SignDataKey]
      addSubkey ts [AuthKey]
      newkey <- get
      return newkey
  let lbs = runPut $ Bin.put s
  BL.putStr $
    if not armor && not noArmor
      then AA.encodeLazy [Armor ArmorPrivateKeyBlock [] lbs]
      else lbs

type KeyBuilder = StateT TK IO

buildKeyWith :: SecretKey -> KeyBuilder a -> IO a
buildKeyWith sk a = evalStateT a (bareTK sk)
  where
    bareTK (SecretKey pkp ska) = TK (pkp, Just ska) [] [] [] []

generateSecretKey :: ThirtyTwoBitTimeStamp -> PubKeyAlgorithm -> IO SecretKey
generateSecretKey ts RSA = do
  (pub, priv) <- liftIO $ RSA.generate 512 0x10001
  return $ SecretKey (pkp pub) (ska priv)
  where
    pkp pub = PKPayload V4 ts 0 RSA (RSAPubKey (RSA_PublicKey pub))
    ska priv = SUUnencrypted (RSAPrivateKey (RSA_PrivateKey priv)) 0 -- FIXME: calculate checksum

addUserId :: ThirtyTwoBitTimeStamp -> Bool -> Text -> KeyBuilder ()
addUserId ts primary userid = modify (newUID userid)
  where
    newUID u tk = tk {_tkUIDs = _tkUIDs tk ++ [selfsign (_tkKey tk) u]}
    selfsign (pkp, Just ska) u =
      ( u
      , [ fromRight
            undefined
            (signUserIDwithRSA
               pkp
               (UserId u)
               (hashed pkp)
               (unhashed pkp)
               (skey ska))
        ])
    hashed pkp =
      [ SigSubPacket False (SigCreationTime ts)
      , SigSubPacket False (IssuerFingerprint 4 (fingerprint pkp))
      , SigSubPacket False (KeyFlags (S.singleton CertifyKeysKey))
      , SigSubPacket False (PrimaryUserId primary)
      , SigSubPacket
          False
          (PreferredHashAlgorithms [SHA512, SHA256, SHA384, SHA224])
      , SigSubPacket
          False
          (PreferredSymmetricAlgorithms [AES256, AES192, AES128])
      ]
    unhashed pkp =
      [SigSubPacket False (Issuer (fromRight undefined (eightOctetKeyID pkp)))]
    skey (SUUnencrypted (RSAPrivateKey (RSA_PrivateKey k)) _) = k

addSubkey :: ThirtyTwoBitTimeStamp -> [KeyFlag] -> KeyBuilder ()
addSubkey ts keyflags = do
  tk <- get
  (SecretKey subpkp subska) <- liftIO $ generateSecretKey ts RSA
  let (pkp, Just ska) = _tkKey tk
      Right crossig =
        crossSignSubkeyWithRSA
          pkp
          subpkp
          (hashedwithflags pkp)
          (unhashed pkp)
          (hashed subpkp)
          (unhashed subpkp)
          (skey ska)
          (skey subska)
  modify (addIt subpkp subska crossig)
  where
    skey (SUUnencrypted (RSAPrivateKey (RSA_PrivateKey k)) _) = k
    addIt sp ss cross tk =
      tk {_tkSubs = _tkSubs tk ++ [(SecretSubkeyPkt sp ss, [cross])]}
    hashed pkp =
      [ SigSubPacket False (SigCreationTime ts)
      , SigSubPacket False (IssuerFingerprint 4 (fingerprint pkp))
      ]
    hashedwithflags pkp =
      hashed pkp ++ [SigSubPacket False (KeyFlags (S.fromList keyflags))]
    unhashed pkp =
      [SigSubPacket False (Issuer (fromRight undefined (eightOctetKeyID pkp)))]

ecoP :: Parser ExtractCertOptions
ecoP =
  ExtractCertOptions <$> switch (long "armor" <> help "armor the output") <*>
  switch (long "no-armor" <> help "don't armor the output")

data ExtractCertOptions =
  ExtractCertOptions
    { ecArmor :: Bool
    , ecNoArmor :: Bool
    }

doExtractCert :: ExtractCertOptions -> IO ()
doExtractCert ExtractCertOptions {..} = do
  kbs <- runConduitRes $ CB.sourceHandle stdin .| CL.consume
  let lbs = BL.fromChunks kbs
      isArmored =
        BLC8.pack "-----BEGIN PGP PRIVATE KEY BLOCK-----" == BL.take 37 lbs
      Right [Armor ArmorPrivateKeyBlock _ decoded] =
        AA.decodeLazy lbs :: Either String [Armor]
      lbs' =
        if isArmored
          then decoded
          else lbs
  k <-
    runConduitRes $
    CL.sourceList (BL.toChunks lbs') .| conduitGet Bin.get .| conduitToTKs .|
    CL.take 1
  let output = runPut $ Bin.put (pubToSecret (head k))
  BL.putStr $
    if not ecArmor && not ecNoArmor
      then AA.encodeLazy [Armor ArmorPublicKeyBlock [] output]
      else output
  where
    pubToSecret tk =
      tk {_tkKey = pToS (_tkKey tk), _tkSubs = map subPToS (_tkSubs tk)}
    pToS (pkp, _) = (pkp, Nothing)
    subPToS (SecretSubkeyPkt pkp _, sigs) = (PublicSubkeyPkt pkp, sigs)

soP :: Parser SignOptions
soP =
  SignOptions <$> switch (long "armor" <> help "armor the output") <*>
  switch (long "no-armor" <> help "don't armor the output") <*>
  option
    (eitherReader asTypeReader)
    (long "as" <> metavar "DATATYPE" <> astypeHelp <> value AsBinary) <*>
  some
    (strArgument
       (metavar "KEYS..." <>
        help "paths to at least one secret key, one key per filename"))
  where
    astypeHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "what to treat the input as" <>
      softline <> list (map (pretty . fst) asTypes)

data SignOptions =
  SignOptions
    { sArmor :: Bool
    , sNoArmor :: Bool
    , sAs :: AsBinaryText
    , sKeyFiles :: [String]
    }

asTypes :: [(String, AsBinaryText)]
asTypes = [("binary", AsBinary), ("text", AsText)]

data AsBinaryText
  = AsBinary
  | AsText
  deriving (Eq)

asTypeReader :: String -> Either String AsBinaryText
asTypeReader = note "unknown as type" . flip lookup asTypes

doSign :: POSIXTime -> SignOptions -> IO ()
doSign pt SignOptions {..} = do
  mbs <- runConduitRes $ CB.sourceHandle stdin .| CL.consume
  ks <- mapM grabKey sKeyFiles
  let ts = ThirtyTwoBitTimeStamp (floor pt)
      payload' = BL.fromChunks mbs
      payload =
        if sAs == AsText
          then canonicalize payload'
          else payload'
      funkeys = concatMap tkToFunKeys . rights . map (processTK (Just pt)) $ ks
      allSigningCapableKeys = filter (isSigner . fkufs) funkeys
  forM_ allSigningCapableKeys $ \k -> do
    let Right sig = signData sAs ts k payload
        sigpkt = SignaturePkt sig
        output = runPut (Bin.put sigpkt)
    BL.putStr $
      if not sArmor && not sNoArmor
        then AA.encodeLazy [Armor ArmorSignature [] output]
        else output
  where
    signData ::
         AsBinaryText
      -> ThirtyTwoBitTimeStamp
      -> FunKey
      -> BL.ByteString
      -> Either String SignaturePayload
    signData AsBinary t k =
      signDataWithRSA
        BinarySig
        (skey (fromJust (fmska k)))
        (hashed (fpkp k) t)
        (unhashed (fpkp k))
    signData AsText t k =
      signDataWithRSA
        CanonicalTextSig
        (skey (fromJust (fmska k)))
        (hashed (fpkp k) t)
        (unhashed (fpkp k))
    skey (SUUnencrypted (RSAPrivateKey (RSA_PrivateKey k)) _) =
      k {RSA.private_p = 0, RSA.private_q = 0} -- FIXME: why is this necessary?
    hashed pkp ct =
      [ SigSubPacket False (SigCreationTime ct)
      , SigSubPacket False (IssuerFingerprint 4 (fingerprint pkp))
      ]
    unhashed pkp =
      [SigSubPacket False (Issuer (fromRight undefined (eightOctetKeyID pkp)))]
    isSigner = S.member SignDataKey
    canonicalize :: BL.ByteString -> BL.ByteString
    canonicalize =
      BL.fromStrict .
      TE.encodeUtf8 .
      T.intercalate (T.pack "\r\n") . T.lines . TE.decodeUtf8 . BL.toStrict

grabKey :: String -> IO TK
grabKey fp = do
  kbs <- runConduitRes $ CB.sourceFile fp .| CL.consume
  let lbs = BL.fromChunks kbs
      isArmored =
        BLC8.pack "-----BEGIN PGP PRIVATE KEY BLOCK-----" == BL.take 37 lbs
      Right [Armor ArmorPrivateKeyBlock _ decoded] =
        AA.decodeLazy lbs :: Either String [Armor]
      lbs' =
        if isArmored
          then decoded
          else lbs
  Just k <-
    runConduitRes $
    CL.sourceList (BL.toChunks lbs') .| conduitGet Bin.get .| conduitToTKs .|
    CL.head
  return k

data FunKey =
  FunKey
    { fpkp :: PKPayload
    , fmska :: Maybe SKAddendum
    , fkufs :: S.Set KeyFlag
    }
  deriving (Show)

tkToFunKeys :: TK -> [FunKey]
tkToFunKeys (TK (pkp, mska) revs uids uats subs) =
  catMaybes (mainKey : map extract subs)
  where
    mainKey = grabASig uids >>= sig2KUFs >>= \kf -> return (FunKey pkp mska kf)
    sig2KUFs = getHasheds >=> find isKUF >=> getKUFs
    grabASig :: [(a, [b])] -> Maybe b
    grabASig = (listToMaybe >=> listToMaybe) . map snd -- FIXME: this should grab the "best" sig
    getHasheds :: SignaturePayload -> Maybe [SigSubPacket]
    getHasheds (SigV4 _ _ _ hasheds _ _ _) = Just hasheds
    getHasheds _ = Nothing
    getKUFs :: SigSubPacket -> Maybe (S.Set KeyFlag)
    getKUFs (SigSubPacket _ (KeyFlags kfs)) = Just kfs
    getKUFs _ = Nothing
    extract ((SecretSubkeyPkt spkp sska), sigs) =
      listToMaybe sigs >>= sig2KUFs >>= \kf ->
        return (FunKey spkp (Just sska) kf)
    extract ((PublicSubkeyPkt spkp), sigs) =
      listToMaybe sigs >>= sig2KUFs >>= \kf -> return (FunKey spkp Nothing kf)
    extract _ = Nothing