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

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

import Codec.Encryption.OpenPGP.Serialize (getPackets)
import Codec.Encryption.OpenPGP.Types
import Control.Applicative (many, (<|>), (<$>))
import Data.Attoparsec.ByteString (Parser, many1, string, inClass, notInClass, satisfy, word8, (<?>), parse, IResult(..))
import Data.Attoparsec.ByteString.Char8 (isDigit_w8)
import Data.Bits (shiftL)
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
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.Serialize.Put (runPut, putWord32be)
import Data.String (IsString, fromString)
import Data.Word (Word32)

decodeArmor :: (Integral a, Read a, Show a, IsString e) => ByteString -> Either e (Armor a)
decodeArmor bs = case parse parseArmor bs of
                     Fail t c e -> Left (fromString e)
                     Partial _ -> Left (fromString "what")
                     Done _ r -> Right r

parseArmor :: (Integral a, Read a, Show a) => Parser (Armor a)
parseArmor = do
    atype <- beginLine <?> "begin line"
    headers <- armorHeaders <?> "headers"
    blankishLine <?> "blank line"
    payload <- base64Data <?> "base64 data"
    endLine atype <?> "end line"
    case runGet getPackets payload of
        Left err -> fail err
        Right packets -> return $ Armor atype headers packets

beginLine :: (Integral a, Read a, Show a) => Parser (ArmorType a)
beginLine = do
    string "-----BEGIN PGP "
    atype <- message <|> pubkey <|> privkey<|> parts <|> signature
    string "-----"
    many (satisfy (inClass " \t"))
    lineEnding
    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 = do
            string "MESSAGE, PART "
            firstnum <- read . BC8.unpack . B.pack <$> many1 (satisfy isDigit_w8)
            return $ ArmorSplitMessageIndefinite firstnum

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

armorHeaders :: Parser [ArmorHeader]
armorHeaders = many armorHeader

armorHeader :: Parser ArmorHeader
armorHeader = do
    key <- many1 (satisfy (inClass "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"))
    string ": "
    val <- many1 (satisfy (notInClass "\n\r"))
    lineEnding
    return (B.pack key, B.pack val)

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

endLine :: (Integral a, Read a, Show a) => ArmorType a -> Parser ByteString
endLine atype = do
    string $ "-----END PGP " `B.append` aType atype `B.append` "-----"
    lineEnding

aType :: (Integral a, Read a, Show a) => ArmorType a -> 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 " ++ show x ++ "/" ++ show y
aType (ArmorSplitMessageIndefinite x) = BC8.pack $ "MESSAGE, PART " ++ show x
aType (ArmorSignature) = BC8.pack "SIGNATURE"

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)