-- Signatures.hs: OpenPGP (RFC4880) signature verification
-- Copyright © 2012-2020  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
module Codec.Encryption.OpenPGP.Signatures
  ( verifySigWith
  , verifyAgainstKeyring
  , verifyAgainstKeys
  , verifyTKWith
  , signUserIDwithRSA
  , crossSignSubkeyWithRSA
  , signDataWithRSA
  ) where

import Control.Applicative ((<|>))
import Control.Error.Util (hush)
import Control.Lens ((^.), _1)
import Control.Monad (liftM2)

import Crypto.Error (eitherCryptoError)
import Crypto.Hash (hashWith)
import qualified Crypto.Hash.Algorithms as CHA
import Crypto.Number.Serialize (i2osp, os2ip)
import qualified Crypto.PubKey.DSA as DSA
import qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import qualified Crypto.PubKey.Ed25519 as Ed25519
import qualified Crypto.PubKey.RSA.PKCS15 as P15
import qualified Crypto.PubKey.RSA.Types as RSATypes

import Data.Bifunctor (first)
import Data.Binary.Put (runPut)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as BL
import Data.Either (isRight, lefts, rights)
import Data.Function (on)
import Data.IxSet.Typed ((@=))
import qualified Data.IxSet.Typed as IxSet
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NE
import Data.Text (Text)
import Data.Time.Clock (UTCTime(..), diffUTCTime)
import Data.Time.Clock.POSIX (posixSecondsToUTCTime)

import Codec.Encryption.OpenPGP.Fingerprint (eightOctetKeyID, fingerprint)
import Codec.Encryption.OpenPGP.Internal
  ( PktStreamContext(..)
  , emptyPSC
  , issuer
  , issuerFP
  )
import Codec.Encryption.OpenPGP.Ontology
  ( isRevocationKeySSP
  , isRevokerP
  , isSubkeyBindingSig
  , isSubkeyRevocation
  )

import Codec.Encryption.OpenPGP.SerializeForSigs
  ( payloadForSig
  , putKeyforSigning
  , putPartialSigforSigning
  , putSigTrailer
  , putUforSigning
  )
import Codec.Encryption.OpenPGP.Types
import Data.Conduit.OpenPGP.Keyring.Instances ()

verifySigWith ::
     (Pkt -> Maybe UTCTime -> ByteString -> Either String Verification)
  -> Pkt
  -> PktStreamContext
  -> Maybe UTCTime
  -> Either String Verification -- FIXME: check expiration here?
verifySigWith vf sig@(SignaturePkt (SigV4 st _ _ hs _ _ _)) state mt = do
  v <- vf sig mt (payloadForSig st state)
  mapM_ (checkI (v ^. verificationSigner) . _sspPayload) hs
  return v
  where
    checkI s i@Issuer {} = checkIssuer (eightOctetKeyID s) i
    checkI s i@IssuerFingerprint {} = checkIssuerFP (fingerprint s) i
    checkI _ _ = Right True
    checkIssuer ::
         Either String EightOctetKeyId
      -> SigSubPacketPayload
      -> Either String Bool
    checkIssuer (Right signer) (Issuer i) =
      if signer == i
        then Right True
        else Left "issuer subpacket does not match"
    checkIssuer (Left err) (Issuer _) =
      Left $ "issuer subpacket cannot be checked (" ++ err ++ ")"
    checkIssuer _ _ = Right True
    checkIssuerFP ::
         TwentyOctetFingerprint -> SigSubPacketPayload -> Either String Bool
    checkIssuerFP signer (IssuerFingerprint _ i) =
      if signer == i
        then Right True
        else Left "issuer fingerprint subpacket does not match"
    checkIssuerFP _ _ = Right True
verifySigWith _ _ _ _ = Left "This should never happen (verifySigWith)."

verifyTKWith ::
     (Pkt -> PktStreamContext -> Maybe UTCTime -> Either String Verification)
  -> Maybe UTCTime
  -> TK
  -> Either String TK
verifyTKWith vsf mt key = do
  revokers <- checkRevokers key
  revs <- checkKeyRevocations revokers key
  let uids = filter (not . null . snd) . checkUidSigs $ key ^. tkUIDs -- FIXME: check revocations here?
  let uats = filter (not . null . snd) . checkUAtSigs $ key ^. tkUAts -- FIXME: check revocations here?
  let subs = concatMap checkSub $ key ^. tkSubs -- FIXME: check revocations here?
  return (TK (key ^. tkKey) revs uids uats subs)
  where
    checkRevokers =
      Right . concat . rights . map verifyRevoker . filter isRevokerP . _tkRevs
    checkKeyRevocations ::
         [(PubKeyAlgorithm, TwentyOctetFingerprint)]
      -> TK
      -> Either String [SignaturePayload]
    checkKeyRevocations rs k =
      Prelude.sequence . concatMap (filterRevs rs) . rights .
      map (liftM2 fmap (,) vSig) $
      k ^.
      tkRevs
    checkUidSigs :: [(Text, [SignaturePayload])] -> [(Text, [SignaturePayload])]
    checkUidSigs =
      map
        (\(uid, sps) ->
           (uid, (rights . map (\sp -> fmap (const sp) (vUid (uid, sp)))) sps))
    checkUAtSigs ::
         [([UserAttrSubPacket], [SignaturePayload])]
      -> [([UserAttrSubPacket], [SignaturePayload])]
    checkUAtSigs =
      map
        (\(uat, sps) ->
           (uat, (rights . map (\sp -> fmap (const sp) (vUAt (uat, sp)))) sps))
    checkSub :: (Pkt, [SignaturePayload]) -> [(Pkt, [SignaturePayload])]
    checkSub (pkt, sps) =
      if revokedSub pkt sps
        then []
        else checkSub' pkt sps
    revokedSub :: Pkt -> [SignaturePayload] -> Bool
    revokedSub _ [] = False
    revokedSub p sigs = any (vSubSig p) (filter isSubkeyRevocation sigs)
    checkSub' :: Pkt -> [SignaturePayload] -> [(Pkt, [SignaturePayload])]
    checkSub' p sps =
      let goodsigs = filter (vSubSig p) (filter isSubkeyBindingSig sps)
       in if null goodsigs
            then []
            else [(p, goodsigs)]
    getHasheds (SigV4 _ _ _ ha _ _ _) = ha
    getHasheds _ = []
    filterRevs ::
         [(PubKeyAlgorithm, TwentyOctetFingerprint)]
      -> (SignaturePayload, Verification)
      -> [Either String SignaturePayload]
    filterRevs vokers spv =
      case spv of
        (s@(SigV4 SignatureDirectlyOnAKey _ _ _ _ _ _), _) -> [Right s]
        (s@(SigV4 KeyRevocationSig pka _ _ _ _ _), v) ->
          if (v ^. verificationSigner == key ^. tkKey . _1) ||
             any
               (\(p, f) ->
                  p == pka && f == fingerprint (v ^. verificationSigner))
               vokers
            then [Left "Key revoked"]
            else [Right s]
        _ -> []
    vUid :: (Text, SignaturePayload) -> Either String Verification
    vUid (uid, sp) =
      vsf
        (SignaturePkt sp)
        emptyPSC
          { lastPrimaryKey = PublicKeyPkt (key ^. tkKey . _1)
          , lastUIDorUAt = UserIdPkt uid
          }
        mt
    vUAt ::
         ([UserAttrSubPacket], SignaturePayload) -> Either String Verification
    vUAt (uat, sp) =
      vsf
        (SignaturePkt sp)
        emptyPSC
          { lastPrimaryKey = PublicKeyPkt (key ^. tkKey . _1)
          , lastUIDorUAt = UserAttributePkt uat
          }
        mt
    vSig :: SignaturePayload -> Either String Verification
    vSig sp =
      vsf
        (SignaturePkt sp)
        emptyPSC {lastPrimaryKey = PublicKeyPkt (key ^. tkKey . _1)}
        mt
    vSubSig :: Pkt -> SignaturePayload -> Bool
    vSubSig sk sp =
      isRight
        (vsf
           (SignaturePkt sp)
           emptyPSC
             { lastPrimaryKey = PublicKeyPkt (key ^. tkKey . _1)
             , lastSubkey = sk
             }
           mt)
    verifyRevoker ::
         SignaturePayload
      -> Either String [(PubKeyAlgorithm, TwentyOctetFingerprint)]
    verifyRevoker sp = do
      _ <- vSig sp
      return
        (map (\(SigSubPacket _ (RevocationKey _ pka fp)) -> (pka, fp)) .
         filter isRevocationKeySSP $
         getHasheds sp)

verifyAgainstKeyring ::
     Keyring -> Pkt -> Maybe UTCTime -> ByteString -> Either String Verification
verifyAgainstKeyring kr sig mt payload = do
  let ikeys = (kr @=) <$> issuer sig
      ifpkeys = (kr @=) <$> issuerFP sig
  keyset <- maybe (Left "issuer not found") Right (ifpkeys <|> ikeys)
  potentialmatches <-
    if IxSet.null keyset
      then Left "pubkey not found"
      else Right keyset
  verifyAgainstKeys (IxSet.toList potentialmatches) sig mt payload

verifyAgainstKeys ::
     [TK] -> Pkt -> Maybe UTCTime -> ByteString -> Either String Verification
verifyAgainstKeys ks sig mt payload = do
  let allrelevantpkps =
        filter
          (\x ->
             (((fingerprint x ==) <$> issuerFP sig) == Just True) ||
             ((==) <$> issuer sig <*> hush (eightOctetKeyID x)) ==
             Just True)
          (concatMap (\x -> (x ^. tkKey . _1) : map subPKP (_tkSubs x)) ks)
  let results =
        map
          (\pkp ->
             verifyAgainstKey'
               pkp
               sig
               mt
               payload)
          allrelevantpkps
  case rights results of
    [] -> Left (concatMap (++ "/") (lefts results))
    [r] -> do
      _ <- isSignatureExpired sig mt
      return r
    _ -> Left "multiple successes; unexpected condition"
  where
    subPKP (pack, _) = subPKP' pack
    subPKP' (PublicSubkeyPkt p) = p
    subPKP' (SecretSubkeyPkt p _) = p
    subPKP' _ = error "This should never happen (subPKP')"

verifyAgainstKey' ::
     PKPayload -> Pkt -> Maybe UTCTime -> ByteString -> Either String Verification
verifyAgainstKey' pkp sig mt payload = do
--  FIXME: check flags
--  FIXME: check expiration time
      r <- verify'
               sig
               pkp
               (hashalgo sig)
               (BL.toStrict (finalPayload sig payload))
--  FIXME: check signature hash against policy
--  FIXME: check pka against policy
      return (Verification r ((_signaturePayload . fromPkt) sig))
  where
    verify' (SignaturePkt s) pub@(PKPayload V4 _ _ _ pkey) SHA1 pl =
      verify'' (pkaAndMPIs s) CHA.SHA1 pub pkey pl
    verify' (SignaturePkt s) pub@(PKPayload V4 _ _ _ pkey) RIPEMD160 pl =
      verify'' (pkaAndMPIs s) CHA.RIPEMD160 pub pkey pl
    verify' (SignaturePkt s) pub@(PKPayload V4 _ _ _ pkey) SHA256 pl =
      verify'' (pkaAndMPIs s) CHA.SHA256 pub pkey pl
    verify' (SignaturePkt s) pub@(PKPayload V4 _ _ _ pkey) SHA384 pl =
      verify'' (pkaAndMPIs s) CHA.SHA384 pub pkey pl
    verify' (SignaturePkt s) pub@(PKPayload V4 _ _ _ pkey) SHA512 pl =
      verify'' (pkaAndMPIs s) CHA.SHA512 pub pkey pl
    verify' (SignaturePkt s) pub@(PKPayload V4 _ _ _ pkey) SHA224 pl =
      verify'' (pkaAndMPIs s) CHA.SHA224 pub pkey pl
    verify' (SignaturePkt s) pub@(PKPayload V4 _ _ _ pkey) DeprecatedMD5 pl =
      verify'' (pkaAndMPIs s) CHA.MD5 pub pkey pl
    verify' _ _ _ _ = error "This should never happen (verify')."
    verify'' (DSA, mpis) hd pub (DSAPubKey (DSA_PublicKey pkey)) bs =
      dsaVerify pub mpis hd pkey bs
    verify'' (ECDSA, mpis) hd pub (ECDSAPubKey (ECDSA_PublicKey pkey)) bs =
      ecdsaVerify pub mpis hd pkey bs
    verify'' (EdDSA, mpis) hd pub (EdDSAPubKey Ed25519 pkey) bs =
      ed25519Verify pub mpis hd (i2osp (unEPoint pkey)) bs
    verify'' (RSA, mpis) hd pub (RSAPubKey (RSA_PublicKey pkey)) bs =
      rsaVerify pub mpis hd pkey bs
    verify'' _ _ _ _ _ = Left "unimplemented key type"
    dsaVerify pub (r :| [s]) hd pkey bs =
      if DSA.verify hd pkey (dsaMPIsToSig r s) bs
        then Right pub
        else Left ("DSA verification failed: " ++ show (hd, pkey, r, s, bs))
    dsaVerify _ _ _ _ _ = Left "cannot verify DSA signature of wrong shape"
    ecdsaVerify pub (r :| [s]) hd pkey bs =
      if ECDSA.verify hd pkey (ecdsaMPIsToSig r s) bs
        then Right pub
        else Left ("ECDSA verification failed: " ++ show (hd, pkey, r, s, bs))
    ecdsaVerify _ _ _ _ _ = Left "cannot verify ECDSA signature of wrong shape"
    ed25519Verify pub (r :| [s]) hd pkey bs =
      either
        (Left .
         (("Ed25519 verification failed: " ++ show (hd, pkey, r, s, bs) ++ ": ") ++) .
         show)
        return $ do
        ep <- cf2es (Ed25519.publicKey (B.drop 1 pkey)) -- drop the 0x40
        es <- cf2es (Ed25519.signature ((B.append `on` i2osp . unMPI) r s))
        let prehash = crazyHash hd bs :: B.ByteString
        if Ed25519.verify ep prehash es
          then Right pub
          else Left "does not verify"
    ed25519Verify _ _ _ _ _ =
      Left "cannot verify Ed25519 signature of wrong shape"
    cf2es = either (Left . show) return . eitherCryptoError
    rsaVerify pub mpis hd pkey bs =
      if P15.verify (Just hd) pkey bs (rsaMPItoSig mpis)
        then Right pub
        else Left ("DSA verification failed: " ++ show (hd, pkey, mpis, bs))
    dsaMPIsToSig r s = DSA.Signature (unMPI r) (unMPI s)
    ecdsaMPIsToSig r s = ECDSA.Signature (unMPI r) (unMPI s)
    rsaMPItoSig (s :| []) = i2osp (unMPI s)
    hashalgo :: Pkt -> HashAlgorithm
    hashalgo (SignaturePkt (SigV4 _ _ ha _ _ _ _)) = ha
    hashalgo _ = error "This should never happen (hashalgo)."
    pkaAndMPIs (SigV4 _ pka _ _ _ _ mpis) = (pka, mpis)
    pkaAndMPIs _ = error "This should never happen (pkaAndMPIs)."
    crazyHash h = BA.convert . hashWith h

isSignatureExpired :: Pkt -> Maybe UTCTime -> Either String Bool
isSignatureExpired _ Nothing = return False
isSignatureExpired s (Just t) =
      if any
           (expiredBefore t)
           ((\(SigV4 _ _ _ h _ _ _) -> h) . _signaturePayload . fromPkt $ s)
        then Left "signature expired"
        else return True
  where
    expiredBefore :: UTCTime -> SigSubPacket -> Bool
    expiredBefore ct (SigSubPacket _ (SigExpirationTime et)) =
      fromEnum ((posixSecondsToUTCTime . toEnum . fromEnum) et `diffUTCTime` ct) <
      0
    expiredBefore _ _ = False

finalPayload :: Pkt -> ByteString -> ByteString
finalPayload s pl = BL.concat [pl, sigbit, trailer s]
  where
    sigbit = runPut $ putPartialSigforSigning s
    trailer :: Pkt -> ByteString
    trailer (SignaturePkt SigV4 {}) = runPut $ putSigTrailer s
    trailer _ = BL.empty

signUserIDwithRSA ::
     PKPayload -- ^ public key "payload" of user ID being signed
  -> UserId -- ^ user ID being signed
  -> [SigSubPacket] -- ^ hashed signature subpackets
  -> [SigSubPacket] -- ^ unhashed signature subpackets
  -> RSATypes.PrivateKey -- ^ RSA signing key
  -> Either String SignaturePayload
signUserIDwithRSA pkp uid hsigsubs usigsubs prv = do
  uidsig <-
    first
      show
      (P15.sign
         Nothing
         (Just CHA.SHA512)
         prv
         (BL.toStrict (finalPayload (SignaturePkt uidsigp) uidpayload)))
  return (uidsigp' uidsig)
  where
    uidpayload =
      runPut
        (sequence_
           [putKeyforSigning (PublicKeyPkt pkp), putUforSigning (toPkt uid)])
    uidsigp =
      SigV4 PositiveCert RSA SHA512 hsigsubs usigsubs 0 (NE.fromList [MPI 0])
    uidsigp' us =
      SigV4
        PositiveCert
        RSA
        SHA512
        hsigsubs
        usigsubs
        (fromIntegral (os2ip (B.take 2 us)))
        (NE.fromList [MPI (os2ip us)])

crossSignSubkeyWithRSA ::
     PKPayload -- ^ public key "payload" of key being signed
  -> PKPayload -- ^ public subkey "payload" of key being signed
  -> [SigSubPacket] -- ^ hashed signature subpackets for binding sig
  -> [SigSubPacket] -- ^ unhashed signature subpackets for binding sig
  -> [SigSubPacket] -- ^ hashed signature subpackets for embedded sig
  -> [SigSubPacket] -- ^ unhashed signature subpackets for embedded sig
  -> RSATypes.PrivateKey -- ^ RSA signing key
  -> RSATypes.PrivateKey -- ^ RSA signing subkey
  -> Either String SignaturePayload
crossSignSubkeyWithRSA pkp subpkp subhsigsubs subusigsubs embhsigsubs embusigsubs prv ssb = do
  embsig <-
    first
      show
      (P15.sign
         Nothing
         (Just CHA.SHA512)
         ssb
         (BL.toStrict (finalPayload (SignaturePkt embsigp) subkeypayload)))
  subsig <-
    first
      show
      (P15.sign
         Nothing
         (Just CHA.SHA512)
         prv
         (BL.toStrict (finalPayload (SignaturePkt subsigp) subkeypayload)))
  return (subsigp' (embsigp' embsig) subsig)
  where
    subkeypayload =
      runPut
        (sequence_
           [ putKeyforSigning (PublicKeyPkt pkp)
           , putKeyforSigning (PublicSubkeyPkt subpkp)
           ])
    embsigp =
      SigV4
        PrimaryKeyBindingSig
        RSA
        SHA512
        embhsigsubs
        embusigsubs
        0
        (NE.fromList [MPI 0])
    embsigp' es =
      SigV4
        PrimaryKeyBindingSig
        RSA
        SHA512
        embhsigsubs
        embusigsubs
        (fromIntegral (os2ip (B.take 2 es)))
        (NE.fromList [MPI (os2ip es)])
    subsigp =
      SigV4 SubkeyBindingSig RSA SHA512 subhsigsubs [] 0 (NE.fromList [MPI 0])
    sspes es = SigSubPacket False (EmbeddedSignature es)
    subsigp' es ss =
      SigV4
        SubkeyBindingSig
        RSA
        SHA512
        subhsigsubs
        (sspes es : subusigsubs)
        (fromIntegral (os2ip (B.take 2 ss)))
        (NE.fromList [MPI (os2ip ss)])

signDataWithRSA ::
     SigType
  -> RSATypes.PrivateKey
  -> [SigSubPacket]
  -> [SigSubPacket]
  -> ByteString
  -> Either String SignaturePayload
signDataWithRSA st prv has uhas payload =
  sp st <$>
  first
    show
    (P15.sign
       Nothing
       (Just CHA.SHA512)
       prv
       (BL.toStrict (finalPayload (SignaturePkt (sp0 st)) payload)))
  where
    sp0 st = SigV4 st RSA SHA512 has [] 0 (NE.fromList [MPI 0])
    sp st ss =
      SigV4
        st
        RSA
        SHA512
        has
        uhas
        (fromIntegral (os2ip (B.take 2 ss)))
        (NE.fromList [MPI (os2ip ss)])
