-- ASCIIArmor/Encode.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.Encode (
   encode
 , encodeLazy
) where

import Codec.Encryption.OpenPGP.ASCIIArmor.Types
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.Lazy.Char8 as BLC8
import qualified Data.ByteString.Base64 as Base64
import Data.Digest.CRC24 (crc24Lazy)
import Data.Serialize (put)
import Data.Serialize.Put (runPutLazy, putWord32be)
import Data.String (IsString, fromString)

encode :: [Armor] -> B.ByteString
encode = B.concat . BL.toChunks . encodeLazy

encodeLazy :: [Armor] -> ByteString
encodeLazy = BL.concat . map armor

armor :: Armor -> ByteString
armor (Armor atype ahs bs) = beginLine atype `BL.append` armorHeaders ahs `BL.append` blankLine `BL.append` armorData bs `BL.append` armorChecksum bs `BL.append` endLine atype
armor (ClearSigned chs ctxt csig) = BLC8.pack "-----BEGIN PGP SIGNED MESSAGE-----\n" `BL.append` armorHeaders chs `BL.append` blankLine `BL.append` dashEscape ctxt `BL.append` armor csig

blankLine :: ByteString
blankLine = BLC8.singleton '\n'

beginLine :: ArmorType -> ByteString
beginLine atype = BLC8.pack "-----BEGIN PGP " `BL.append` aType atype `BL.append` BLC8.pack "-----\n"

endLine :: ArmorType -> ByteString
endLine atype = BLC8.pack "-----END PGP " `BL.append` aType atype `BL.append` BLC8.pack "-----\n"

aType :: ArmorType -> ByteString
aType (ArmorMessage) = BLC8.pack "MESSAGE"
aType (ArmorPublicKeyBlock) = BLC8.pack "PUBLIC KEY BLOCK"
aType (ArmorPrivateKeyBlock) = BLC8.pack "PRIVATE KEY BLOCK"
aType (ArmorSplitMessage x y) = BLC8.pack $ "MESSAGE, PART " ++ show x ++ "/" ++ show y
aType (ArmorSplitMessageIndefinite x) = BLC8.pack $ "MESSAGE, PART " ++ show x
aType (ArmorSignature) = BLC8.pack "SIGNATURE"

armorHeaders :: [(String, String)] -> ByteString
armorHeaders ahs = BLC8.unlines . map armorHeader $ ahs
    where
        armorHeader :: (String, String) -> ByteString
        armorHeader (k, v) = BLC8.pack k `BL.append` BLC8.pack ": " `BL.append` BLC8.pack v

armorData :: ByteString -> ByteString
armorData = BLC8.unlines . wordWrap 64 . BL.fromChunks . return . Base64.encode . B.concat . BL.toChunks

wordWrap :: Int -> ByteString -> [ByteString]
wordWrap lw bs
    | BL.null bs = []
    | lw < 1 || lw > 76 = wordWrap 76 bs
    | otherwise = BL.take (fromIntegral lw) bs : wordWrap lw (BL.drop (fromIntegral lw) bs)

armorChecksum :: ByteString -> ByteString
armorChecksum = BLC8.cons '=' . armorData . BL.tail . runPutLazy . putWord32be . crc24Lazy

dashEscape :: ByteString -> ByteString
dashEscape = BLC8.unlines . map escapeLine . BLC8.lines
    where
        escapeLine :: ByteString -> ByteString
        escapeLine l
            | BLC8.singleton '-' `BL.isPrefixOf` l = BLC8.pack "- " `BL.append` l
            | BLC8.pack "From " `BL.isPrefixOf` l = BLC8.pack "- " `BL.append` l
            | otherwise = l