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

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

import Codec.Encryption.OpenPGP.ASCIIArmor.Types
import Codec.Encryption.OpenPGP.ASCIIArmor.Utils
import Control.Applicative (many, (<|>), (<$>), Alternative, (<*), (<*>), (*>), optional)
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 (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.Digest.CRC24 (crc24)
import Data.Serialize.Get (Get, runGet, getWord8)
import Data.String (IsString, fromString)
import Data.Word (Word32)

decode :: IsString e => ByteString -> Either e [Armor]
decode bs = go (AS.parse parseArmors 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 bs = go (AL.parse parseArmors bs)
    where
        go (AL.Fail _ _ e) = Left (fromString e)
        go (AL.Done _ r) = Right r

parseArmors :: Parser [Armor]
parseArmors = many parseArmor

parseArmor :: Parser Armor
parseArmor = prefixed (clearsigned <|> armor) <?> "armor"

clearsigned :: Parser Armor
clearsigned = do
    _ <- string "-----BEGIN PGP SIGNED MESSAGE-----" <?> "clearsign header"
    _ <- lineEnding <?> "line ending"
    headers <- armorHeaders <?> "clearsign headers"
    _ <- blankishLine <?> "blank line"
    cleartext <- dashEscapedCleartext
    sig <- armor
    return $ ClearSigned headers (BL.fromChunks [cleartext]) sig

armor :: Parser Armor
armor = do
    atype <- beginLine <?> "begin line"
    headers <- armorHeaders <?> "headers"
    _ <- blankishLine <?> "blank line"
    payload <- base64Data <?> "base64 data"
    _ <- endLine atype <?> "end line"
    return $ Armor atype headers (BL.fromChunks [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" *> return ArmorMessage
        pubkey = string "PUBLIC KEY BLOCK" *> return ArmorPublicKeyBlock
        privkey = string "PRIVATE KEY BLOCK" *> return ArmorPrivateKeyBlock
        signature = string "SIGNATURE" *> return 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 ByteString
lineEnding = string "\n" <|> string "\r\n"

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 ByteString
blankishLine = many (satisfy (inClass " \t")) *> lineEnding

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

aType :: ArmorType -> 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 -> ByteString
l2s = B.concat . BL.toChunks

base64Data :: Parser ByteString
base64Data = do
    ls <- many1 base64Line
    cksum <- checksumLine
    let payload = B.concat ls
    let ourcksum = crc24 payload
    case runGet d24 cksum of
        Left err -> fail err
        Right theircksum -> if theircksum == ourcksum then return payload else fail ("CRC24 mismatch: " ++ show (B.unpack cksum) ++  "/" ++ show theircksum ++ " vs. " ++ show ourcksum)
    where
        base64Line :: Parser ByteString
        base64Line = do
                        b64 <- many1 (satisfy (inClass "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"))
                        pad <- many (word8 (fromIntegral . fromEnum $ '='))
                        _ <- lineEnding
                        let line = B.pack b64 `B.append` B.pack pad
                        case Base64.decode line of
                            Left err -> fail err
                            Right bs -> return bs
        checksumLine :: Parser ByteString
        checksumLine = do
                        _ <- word8 (fromIntegral . fromEnum $ '=')
                        b64 <- many1 (satisfy (inClass "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"))
                        _ <- lineEnding
                        let line = B.pack b64
                        case Base64.decode line of
                            Left err -> fail err
                            Right bs -> return bs


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 = do
    ls <- many1 ((deLine <|> unescapedLine) <* lineEnding)
    return $ crlfUnlines ls
    where
        deLine :: Parser ByteString
        deLine = B.pack <$> (string "- " *> many (satisfy (notInClass "\n\r")))
        unescapedLine :: Parser ByteString
        unescapedLine = maybe B.empty B.pack <$> optional ((:) <$> satisfy (notInClass "-\n\r") <*> many (satisfy (notInClass "\n\r")))