{-# LANGUAGE OverloadedStrings #-}
-- ASCIIArmor/Decode.hs: OpenPGP (RFC9580) ASCII armor implementation
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

module Codec.Encryption.OpenPGP.ASCIIArmor.Decode (
   parseArmor
 , decode
 , decodeLazy
 , decodeWith
 , decodeLazyWith
) where

import Codec.Encryption.OpenPGP.ASCIIArmor.Types
import Codec.Encryption.OpenPGP.ASCIIArmor.Utils
import Control.Applicative (many, (<|>), optional)
import qualified Data.Attoparsec.Combinator as AC
import Data.Attoparsec.ByteString (Parser, many1, string, inClass, notInClass, satisfy, word8, (<?>))
import qualified Data.Attoparsec.ByteString as AS
import qualified Data.Attoparsec.ByteString.Lazy as AL
import Data.Attoparsec.ByteString.Char8 (isDigit_w8, anyChar)
import Data.Bits (shiftL)
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Char8 as BC8
import qualified Data.ByteString.Base64 as Base64
import Data.Char (isAlphaNum, isSpace)
import Data.Digest.CRC24 (crc24)
import Data.Binary.Get (Get, runGetOrFail, getWord8)
import Data.Functor (($>))
import Data.List (dropWhileEnd)
import Data.String (IsString, fromString)
import Data.Word (Word32)

decode :: IsString e => B.ByteString -> Either e [Armor]
decode = decodeWith defaultDecodeOptions

decodeWith :: IsString e => DecodeOptions -> B.ByteString -> Either e [Armor]
decodeWith options bs = go (AS.parse (parseArmors options) bs)
    where
        go (AS.Fail _ _ e) = Left (fromString e)
        go (AS.Partial cont) = go (cont B.empty)
        go (AS.Done _ r) = Right r

decodeLazy :: IsString e => BL.ByteString -> Either e [Armor]
decodeLazy = decodeLazyWith defaultDecodeOptions

decodeLazyWith :: IsString e => DecodeOptions -> BL.ByteString -> Either e [Armor]
decodeLazyWith options bs = go (AL.parse (parseArmors options) bs)
    where
        go (AL.Fail _ _ e) = Left (fromString e)
        go (AL.Done _ r) = Right r

parseArmors :: DecodeOptions -> Parser [Armor]
parseArmors options = many (parseArmorWith options)

parseArmor :: Parser Armor
parseArmor = parseArmorWith defaultDecodeOptions

parseArmorWith :: DecodeOptions -> Parser Armor
parseArmorWith options = prefixed (clearsigned options <|> armor options) <?> "armor"

clearsigned :: DecodeOptions -> Parser Armor
clearsigned options = do
    _ <- string "-----BEGIN PGP SIGNED MESSAGE-----" <?> "clearsign header"
    _ <- lineEnding <?> "line ending"
    headers <- cleartextHeaders options <?> "clearsign headers"
    _ <- blankishLine <?> "blank line"
    cleartext <- dashEscapedCleartext
    sig <- armor options
    return $ ClearSigned headers cleartext sig

armor :: DecodeOptions -> Parser Armor
armor options = do
    atype <- beginLine <?> "begin line"
    headers <- armorHeaders <?> "headers"
    _ <- blankishLine <?> "blank line"
    payload <- base64Data options <?> "base64 data"
    _ <- endLine atype <?> "end line"
    return $ Armor atype headers payload

beginLine :: Parser ArmorType
beginLine = do
    _ <- string "-----BEGIN PGP " <?> "leading minus-hyphens"
    atype <- pubkey <|> privkey <|> parts <|> message <|> signature
    _ <- string "-----" <?> "trailing minus-hyphens"
    _ <- many (satisfy (inClass " \t")) <?> "whitespace"
    _ <- lineEnding <?> "line ending"
    return atype
    where
        message = string "MESSAGE" $> ArmorMessage
        pubkey = string "PUBLIC KEY BLOCK" $> ArmorPublicKeyBlock
        privkey = string "PRIVATE KEY BLOCK" $> ArmorPrivateKeyBlock
        signature = string "SIGNATURE" $> ArmorSignature
        parts = string "MESSAGE, PART " *> (partsdef <|> partsindef)
        partsdef = do
            firstnum <- num
            _ <- word8 (fromIntegral . fromEnum $ '/')
            secondnum <- num
            return $ ArmorSplitMessage (BL.pack firstnum) (BL.pack secondnum)
        partsindef = ArmorSplitMessageIndefinite . BL.pack <$> num
        num = many1 (satisfy isDigit_w8) <?> "number"

lineEnding :: Parser B.ByteString
lineEnding = string "\n" <|> string "\r\n"

cleartextHeaders :: DecodeOptions -> Parser [(String, String)]
cleartextHeaders options
    | shouldRequireHashOnlyHeaders options = many hashHeader
    | otherwise = armorHeaders
    where
        hashHeader = do
            hdr@(k, v) <- armorHeader
            if k == "Hash" && isConformantHashHeader v
                then return hdr
                else fail "Only well-formed Hash headers are allowed before cleartext payload"

isConformantHashHeader :: String -> Bool
isConformantHashHeader v = not (null tokens) && all isValidToken tokens
    where
        tokens = map trim (splitOnComma v)
        splitOnComma [] = [""]
        splitOnComma xs =
            let (a, rest) = break (== ',') xs
             in a : case rest of
                 [] -> []
                 (_:ys) -> splitOnComma ys
        trim = dropWhileEnd isSpace . dropWhile isSpace
        isValidToken [] = False
        isValidToken t = all (\c -> isAlphaNum c || c == '-') t

armorHeaders :: Parser [(String, String)]
armorHeaders = many armorHeader

armorHeader :: Parser (String, String)
armorHeader = do
    key <- many1 (satisfy (inClass "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"))
    _ <- string ": "
    val <- many1 (satisfy (notInClass "\n\r"))
    _ <- lineEnding
    return (w8sToString key, w8sToString val)
    where
        w8sToString = BC8.unpack . B.pack

blankishLine ::  Parser B.ByteString
blankishLine = many (satisfy (inClass " \t")) *> lineEnding

endLine :: ArmorType -> Parser B.ByteString
endLine atype = do
    _ <- string $ "-----END PGP " `B.append` aType atype `B.append` "-----"
    lineEnding

aType :: ArmorType -> B.ByteString
aType ArmorMessage = BC8.pack "MESSAGE"
aType ArmorPublicKeyBlock = BC8.pack "PUBLIC KEY BLOCK"
aType ArmorPrivateKeyBlock = BC8.pack "PRIVATE KEY BLOCK"
aType (ArmorSplitMessage x y) = BC8.pack "MESSAGE, PART " `B.append` l2s x `B.append` BC8.singleton '/' `B.append` l2s y
aType (ArmorSplitMessageIndefinite x) = BC8.pack "MESSAGE, PART " `B.append` l2s x
aType ArmorSignature = BC8.pack "SIGNATURE"

l2s :: BL.ByteString -> B.ByteString
l2s = B.concat . BL.toChunks

base64Data :: DecodeOptions -> Parser ByteString
base64Data options = do
    ls <- many1 base64Line
    let payload = B.concat ls
    case crc24Validation options of
        DecodeCRC24ValidationStrict -> do
            decodedChecksum <- checksumLineStrict
            validateChecksum payload (Just decodedChecksum) True
            return (BL.fromStrict payload)
        DecodeCRC24ValidationPermissive -> do
            decodedChecksum <- optional checksumLinePermissive
            validateChecksum payload decodedChecksum False
            return (BL.fromStrict payload)
        DecodeCRC24ValidationAuto -> fail "Internal error: unresolved CRC24 validation mode"
  where
    base64Line :: Parser B.ByteString
    base64Line = do
        startsChecksum <- isChecksumLineStart
        if startsChecksum
            then fail "checksum line"
            else do
                b64line <- B.pack <$> many1 (satisfy (inClass "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/= \t"))
                _ <- lineEnding
                let cleaned = stripInlineWhitespace b64line
                case Base64.decode cleaned of
                    Left err -> fail err
                    Right bs -> return bs

    isChecksumLineStart :: Parser Bool
    isChecksumLineStart = AC.lookAhead $ do
        _ <- many (satisfy (inClass " \t"))
        c <- satisfy (notInClass "\n\r")
        return (c == fromIntegral (fromEnum '='))

    checksumLineStrict :: Parser B.ByteString
    checksumLineStrict = do
        _ <- word8 (fromIntegral . fromEnum $ '=')
        b64 <- B.pack <$> many1 (satisfy (inClass "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/ \t"))
        _ <- lineEnding
        case Base64.decode (stripInlineWhitespace b64) of
            Left err -> fail err
            Right bs -> return bs

    checksumLinePermissive :: Parser B.ByteString
    checksumLinePermissive = do
        _ <- word8 (fromIntegral . fromEnum $ '=')
        raw <- B.pack <$> many (satisfy (notInClass "\n\r"))
        _ <- lineEnding
        case Base64.decode (stripInlineWhitespace raw) of
            Left _ -> return B.empty
            Right bs -> return bs

validateChecksum :: B.ByteString -> Maybe B.ByteString -> Bool -> Parser ()
validateChecksum payload mDecodedChecksum requireValidCrc =
    case mDecodedChecksum of
        Nothing ->
            if requireValidCrc
                then fail "Missing CRC24 checksum line"
                else return ()
        Just decodedChecksum ->
            if B.length decodedChecksum /= 3
                then if requireValidCrc then fail "Malformed CRC24 checksum line" else return ()
                else
                    case runGetOrFail d24 (BL.fromStrict decodedChecksum) of
                        Left (_,_,err) ->
                            if requireValidCrc then fail err else return ()
                        Right (_,_,theircksum) ->
                            let ourcksum = crc24 payload
                             in if theircksum == ourcksum || not requireValidCrc
                                   then return ()
                                   else fail ("CRC24 mismatch: " ++ show (B.unpack decodedChecksum) ++ "/" ++ show theircksum ++ " vs. " ++ show ourcksum)

stripInlineWhitespace :: B.ByteString -> B.ByteString
stripInlineWhitespace = B.filter (\w -> w /= 0x20 && w /= 0x09)

crc24Validation :: DecodeOptions -> DecodeCRC24Validation
crc24Validation options =
    case decodeCRC24Validation options of
        DecodeCRC24ValidationAuto ->
            case decodeConformanceProfile options of
                Strict_4880_AA -> DecodeCRC24ValidationStrict
                Strict_9580_AA -> DecodeCRC24ValidationPermissive
                Postel_AA -> DecodeCRC24ValidationPermissive
        explicitMode -> explicitMode

shouldRequireHashOnlyHeaders :: DecodeOptions -> Bool
shouldRequireHashOnlyHeaders options =
    case decodeCleartextHeaderPolicy options of
        DecodeCleartextHeaderPolicyAuto ->
            decodeConformanceProfile options == Strict_9580_AA
        DecodeCleartextHeaderPolicyLegacy -> False
        DecodeCleartextHeaderPolicyHashOnly -> True

d24 :: Get Word32
d24 = do
    a <- getWord8
    b <- getWord8
    c <- getWord8
    return $ shiftL (fromIntegral a :: Word32) 16 + shiftL (fromIntegral b :: Word32) 8 + (fromIntegral c :: Word32)

prefixed :: Parser a -> Parser a
prefixed end = end <|> anyChar *> prefixed end

dashEscapedCleartext :: Parser ByteString
dashEscapedCleartext = BL.fromStrict . crlfUnlines <$> go []
    where
        go acc = do
            isSignatureStart <- AC.lookAhead ((string "-----BEGIN PGP " >> return True) <|> return False)
            if isSignatureStart
                then return (reverse acc)
                else do
                    rawLine <- B.pack <$> many (satisfy (notInClass "\n\r"))
                    _ <- lineEnding
                    go (undash rawLine : acc)

        undash :: B.ByteString -> B.ByteString
        undash line
            | BC8.pack "- " `B.isPrefixOf` line = B.drop 2 line
            | otherwise = line
