module Data.Certificate.X509.Cert
        (
        -- * Data Structure
          SignatureALG(..)
        , HashALG(..)
        , PubKeyALG(..)
        , PubKey(..)
        , ECDSA_Hash(..)
        , ASN1StringType(..)
        , ASN1String
        , Certificate(..)
        , DistinguishedName(..)
        , OID

        -- various OID
        , oidCommonName
        , oidCountry
        , oidOrganization
        , oidOrganizationUnit

        -- signature to/from oid
        , oidSig
        , sigOID

        -- * Parse and encode a single distinguished name
        , parseDN
        , encodeDNinner
        , encodeDN

        -- * extensions
        , module Data.Certificate.X509.Ext
        ) where

import Data.Word
import Data.Monoid
import Data.List (find)
import Data.ASN1.Stream
import Data.ASN1.Encoding
import Data.ASN1.BinaryEncoding
import Data.ASN1.BitArray
import Data.ASN1.Object
import Data.Maybe
import Data.Time.Calendar
import Data.Time.Clock (DiffTime, secondsToDiffTime)
import qualified Data.ByteString.Lazy as L
import Control.Applicative ((<$>))
import Control.Monad.State
import Control.Monad.Error
import Data.Certificate.X509.Internal
import Data.Certificate.X509.Ext
import qualified Crypto.Types.PubKey.RSA as RSA
import qualified Crypto.Types.PubKey.DSA as DSA

import Data.Certificate.KeyRSA (parse_RSA)

data HashALG =
          HashMD2
        | HashMD5
        | HashSHA1
        | HashSHA224
        | HashSHA256
        | HashSHA384
        | HashSHA512
        deriving (Show,Eq)

data PubKeyALG =
          PubKeyALG_RSA
        | PubKeyALG_DSA
        | PubKeyALG_ECDSA
        | PubKeyALG_DH
        | PubKeyALG_Unknown OID
        deriving (Show,Eq)

data SignatureALG =
          SignatureALG HashALG PubKeyALG
        | SignatureALG_Unknown OID
        deriving (Show,Eq)

data ECDSA_Hash = ECDSA_Hash_SHA384
                deriving (Show,Eq)

data PubKey =
          PubKeyRSA RSA.PublicKey -- ^ RSA public key
        | PubKeyDSA DSA.PublicKey -- ^ DSA public key
        | PubKeyDH (Integer,Integer,Integer,Maybe Integer,([Word8], Integer))
                                    -- ^ DH format with (p,g,q,j,(seed,pgenCounter))
        | PubKeyECDSA ECDSA_Hash L.ByteString -- ^ ECDSA format not done yet FIXME
        | PubKeyUnknown OID [Word8] -- ^ unrecognized format
        deriving (Show,Eq)

type Time = (Day, DiffTime, Bool)

data CertKeyUsage =
          CertKeyUsageDigitalSignature
        | CertKeyUsageNonRepudiation
        | CertKeyUsageKeyEncipherment
        | CertKeyUsageDataEncipherment
        | CertKeyUsageKeyAgreement
        | CertKeyUsageKeyCertSign
        | CertKeyUsageCRLSign
        | CertKeyUsageEncipherOnly
        | CertKeyUsageDecipherOnly
        deriving (Show, Eq)

data ASN1StringType = UTF8 | Printable | Univ | BMP | IA5 | T61 deriving (Show,Eq,Ord,Enum)
type ASN1String = (ASN1StringType, String)

newtype DistinguishedName = DistinguishedName { getDistinguishedElements :: [(OID, ASN1String)] }
    deriving (Show,Eq,Ord)

instance Monoid DistinguishedName where
    mempty  = DistinguishedName []
    mappend (DistinguishedName l1) (DistinguishedName l2) = DistinguishedName (l1++l2)

data Certificate = Certificate
        { certVersion      :: Int                    -- ^ Certificate Version
        , certSerial       :: Integer                -- ^ Certificate Serial number
        , certSignatureAlg :: SignatureALG           -- ^ Certificate Signature algorithm
        , certIssuerDN     :: DistinguishedName      -- ^ Certificate Issuer DN
        , certSubjectDN    :: DistinguishedName      -- ^ Certificate Subject DN
        , certValidity     :: (Time, Time)           -- ^ Certificate Validity period
        , certPubKey       :: PubKey                 -- ^ Certificate Public key
        , certExtensions   :: Maybe [ExtensionRaw]   -- ^ Certificate Extensions
        } deriving (Show,Eq)

instance ASN1Object Certificate where
    toASN1   certificate = encodeCertificateHeader certificate
    fromASN1 s           = runParseASN1State parseCertificate s

oidCommonName, oidCountry, oidOrganization, oidOrganizationUnit :: OID
oidCommonName       = [2,5,4,3]
oidCountry          = [2,5,4,6]
oidOrganization     = [2,5,4,10]
oidOrganizationUnit = [2,5,4,11]

parseCertHeaderVersion :: ParseASN1 Int
parseCertHeaderVersion =
    maybe 1 id <$> onNextContainerMaybe (Container Context 0) (getNext >>= getVer)
    where getVer (IntVal v) = return $ fromIntegral v
          getVer _          = throwError "unexpected type for version"

parseCertHeaderSerial :: ParseASN1 Integer
parseCertHeaderSerial = do
        n <- getNext
        case n of
                IntVal v -> return v
                _        -> throwError ("missing serial" ++ show n)

sig_table :: [ (OID, SignatureALG) ]
sig_table =
        [ ([1,2,840,113549,1,1,5], SignatureALG HashSHA1 PubKeyALG_RSA)
        , ([1,2,840,113549,1,1,4], SignatureALG HashMD5 PubKeyALG_RSA)
        , ([1,2,840,113549,1,1,2], SignatureALG HashMD2 PubKeyALG_RSA)
        , ([1,2,840,113549,1,1,11], SignatureALG HashSHA256 PubKeyALG_RSA)
        , ([1,2,840,113549,1,1,12], SignatureALG HashSHA384 PubKeyALG_RSA)
        , ([1,2,840,10040,4,3],    SignatureALG HashSHA1 PubKeyALG_DSA)
        , ([1,2,840,10045,4,3,1],  SignatureALG HashSHA224 PubKeyALG_ECDSA)
        , ([1,2,840,10045,4,3,2],  SignatureALG HashSHA256 PubKeyALG_ECDSA)
        , ([1,2,840,10045,4,3,3],  SignatureALG HashSHA384 PubKeyALG_ECDSA)
        , ([1,2,840,10045,4,3,4],  SignatureALG HashSHA512 PubKeyALG_ECDSA)
        ]

pk_table :: [ (OID, PubKeyALG) ]
pk_table =
        [ ([1,2,840,113549,1,1,1], PubKeyALG_RSA)
        , ([1,2,840,10040,4,1],    PubKeyALG_DSA)
        , ([1,2,840,10045,2,1],    PubKeyALG_ECDSA)
        , ([1,2,840,10046,2,1],    PubKeyALG_DH)
        ]

oidSig :: OID -> SignatureALG
oidSig oid = maybe (SignatureALG_Unknown oid) id $ lookup oid sig_table

oidPubKey :: OID -> PubKeyALG
oidPubKey oid = maybe (PubKeyALG_Unknown oid) id $ lookup oid pk_table

sigOID :: SignatureALG -> OID
sigOID (SignatureALG_Unknown oid) = oid
sigOID sig = maybe [] fst $ find ((==) sig . snd) sig_table

pubkeyalgOID :: PubKeyALG -> OID
pubkeyalgOID (PubKeyALG_Unknown oid) = oid
pubkeyalgOID sig = maybe [] fst $ find ((==) sig . snd) pk_table

pubkeyToAlg :: PubKey -> PubKeyALG
pubkeyToAlg (PubKeyRSA _)         = PubKeyALG_RSA
pubkeyToAlg (PubKeyDSA _)         = PubKeyALG_DSA
pubkeyToAlg (PubKeyDH _)          = PubKeyALG_DH
pubkeyToAlg (PubKeyECDSA _ _)     = PubKeyALG_ECDSA
pubkeyToAlg (PubKeyUnknown oid _) = PubKeyALG_Unknown oid

parseCertHeaderAlgorithmID :: ParseASN1 SignatureALG
parseCertHeaderAlgorithmID = do
        n <- getNextContainer Sequence
        case n of
                [ OID oid, Null ] -> return $ oidSig oid
                [ OID oid ]       -> return $ oidSig oid
                _                 -> throwError ("algorithm ID bad format " ++ show n)

asn1String :: ASN1 -> ASN1String
asn1String (PrintableString x) = (Printable, x)
asn1String (UTF8String x)      = (UTF8, x)
asn1String (UniversalString x) = (Univ, x)
asn1String (BMPString x)       = (BMP, x)
asn1String (IA5String x)       = (IA5, x)
asn1String (T61String x)       = (IA5, x)
asn1String x                   = error ("not a print string " ++ show x)

encodeAsn1String :: ASN1String -> ASN1
encodeAsn1String (Printable, x) = PrintableString x
encodeAsn1String (UTF8, x)      = UTF8String x
encodeAsn1String (Univ, x)      = UniversalString x
encodeAsn1String (BMP, x)       = BMPString x
encodeAsn1String (IA5, x)       = IA5String x
encodeAsn1String (T61, x)       = T61String x

parseDN :: ParseASN1 DistinguishedName
parseDN = DistinguishedName <$> onNextContainer Sequence getDNs where
        getDNs = do
                n <- hasNext
                if n
                        then liftM2 (:) parseOneDN getDNs
                        else return []

parseOneDN :: ParseASN1 (OID, ASN1String)
parseOneDN = onNextContainer Set $ do
        s <- getNextContainer Sequence
        case s of
                [OID oid, val] -> return (oid, asn1String val)
                _              -> throwError "expecting sequence"

parseCertHeaderValidity :: ParseASN1 (Time, Time)
parseCertHeaderValidity = getNextContainer Sequence >>= toTimeBound
    where toTimeBound [ UTCTime t1, UTCTime t2 ]                 = return (convertTime t1, convertTime t2)
          toTimeBound [ GeneralizedTime t1, GeneralizedTime t2 ] = return (convertTime t1, convertTime t2)
          toTimeBound _                                          = throwError "bad validity format"

          convertTime (y,m,d,h,mi,s,u) =
              let day   = fromGregorian (fromIntegral y) m d
                  dtime = secondsToDiffTime (fromIntegral h * 3600 + fromIntegral mi * 60 + fromIntegral s)
               in (day, dtime, u)

parseCertHeaderSubjectPK :: ParseASN1 PubKey
parseCertHeaderSubjectPK = onNextContainer Sequence $ do
    l <- getNextContainer Sequence
    bits <- getNextBitString
    case l of
        (OID pkalg):xs -> toKey (oidPubKey pkalg) xs bits
        _              -> throwError ("subject public unknown key format : " ++ show l)
    where toKey PubKeyALG_RSA _ bits = do
                either (throwError) (return . PubKeyRSA) (parse_RSA bits)
          toKey PubKeyALG_ECDSA xs bits = do
                case xs of
                    [(OID [1,3,132,0,34])] -> return $ PubKeyECDSA ECDSA_Hash_SHA384 bits
                    _                      -> return $ PubKeyUnknown (pubkeyalgOID PubKeyALG_ECDSA) $ L.unpack bits
          toKey PubKeyALG_DSA [Start Sequence,IntVal p,IntVal q,IntVal g,End Sequence] bits = do
                case decodeASN1 BER bits of
                     Right [IntVal dsapub] -> return $ PubKeyDSA $ DSA.PublicKey
                                                                     { DSA.public_params = DSA.Params
                                                                             { DSA.params_p = p
                                                                             , DSA.params_q = q
                                                                             , DSA.params_g = g }
                                                                     , DSA.public_y = dsapub }
                     _                     -> return $ PubKeyUnknown (pubkeyalgOID PubKeyALG_DSA) $ L.unpack bits
          toKey (PubKeyALG_Unknown oid) _ bits = return $ PubKeyUnknown oid $ L.unpack bits
          toKey other _ bits = return $ PubKeyUnknown (pubkeyalgOID other) $ L.unpack bits

          getNextBitString = getNext >>= \bs -> case bs of
                BitString bits -> return $ bitArrayGetData bits
                _              -> throwError "expecting bitstring"

parseCertExtensions :: ParseASN1 (Maybe [ExtensionRaw])
parseCertExtensions = onNextContainerMaybe (Container Context 3) (mapMaybe extractExtension <$> onNextContainer Sequence getSequences)
        where
                getSequences = do
                        n <- hasNext
                        if n
                                then getNextContainer Sequence >>= \sq -> liftM (sq :) getSequences
                                else return []
                extractExtension [OID oid,Boolean True,OctetString obj] = case decodeASN1 BER obj of
                        Left _  -> Nothing
                        Right r -> Just (oid, True, r)
                extractExtension [OID oid,OctetString obj]              = case decodeASN1 BER obj of
                        Left _  -> Nothing
                        Right r -> Just (oid, False, r)
                extractExtension _                                      = Nothing

{- | parse header structure of a x509 certificate. the structure the following:
        Version
        Serial Number
        Algorithm ID
        Issuer
        Validity
                Not Before
                Not After
        Subject
        Subject Public Key Info
                Public Key Algorithm
                Subject Public Key
        Issuer Unique Identifier (Optional)  (>= 2)
        Subject Unique Identifier (Optional) (>= 2)
        Extensions (Optional)   (>= v3)
-}
parseCertificate :: ParseASN1 Certificate
parseCertificate = do
        version  <- parseCertHeaderVersion
        serial   <- parseCertHeaderSerial
        sigalg   <- parseCertHeaderAlgorithmID
        issuer   <- parseDN
        validity <- parseCertHeaderValidity
        subject  <- parseDN
        pk       <- parseCertHeaderSubjectPK
        exts     <- parseCertExtensions
        hnext    <- hasNext
        when hnext $ throwError "expecting End Of Data."
        
        return $ Certificate
                { certVersion      = version
                , certSerial       = serial
                , certSignatureAlg = sigalg
                , certIssuerDN     = issuer
                , certSubjectDN    = subject
                , certValidity     = validity
                , certPubKey       = pk
                , certExtensions   = exts
                }

encodeDNinner :: (ASN1String -> ASN1String) -> DistinguishedName -> [ASN1]
encodeDNinner f (DistinguishedName dn) = concatMap dnSet dn
    where dnSet (oid, stringy) = asn1Container Set (asn1Container Sequence [OID oid, encodeAsn1String (f stringy)])

encodeDN :: DistinguishedName -> [ASN1]
encodeDN dn = asn1Container Sequence $ encodeDNinner id dn

encodePK :: PubKey -> [ASN1]
encodePK k@(PubKeyRSA pubkey) =
        asn1Container Sequence (asn1Container Sequence [pkalg,Null] ++ [BitString $ toBitArray bits 0])
        where
                pkalg = OID $ pubkeyalgOID $ pubkeyToAlg k
                bits  = encodeASN1 DER $ asn1Container Sequence [IntVal (RSA.public_n pubkey), IntVal (RSA.public_e pubkey)]

encodePK k@(PubKeyDSA pubkey) =
        asn1Container Sequence (asn1Container Sequence ([pkalg] ++ dsaseq) ++ [BitString $ toBitArray bits 0])
        where
                pkalg   = OID $ pubkeyalgOID $ pubkeyToAlg k
                dsaseq  = asn1Container Sequence [IntVal (DSA.params_p params)
                                                 ,IntVal (DSA.params_q params)
                                                 ,IntVal (DSA.params_g params)]
                params  = DSA.public_params pubkey
                bits    = encodeASN1 DER [IntVal $ DSA.public_y pubkey]

encodePK k@(PubKeyUnknown _ l) =
        asn1Container Sequence (asn1Container Sequence [pkalg,Null] ++ [BitString $ toBitArray (L.pack l) 0])
        where
                pkalg = OID $ pubkeyalgOID $ pubkeyToAlg k

encodeExts :: Maybe [ExtensionRaw] -> [ASN1]
encodeExts Nothing  = []
encodeExts (Just l) = asn1Container (Container Context 3) $ concatMap encodeExt l
        where encodeExt (oid, critical, asn1) =
                let bs = encodeASN1 DER asn1
                 in asn1Container Sequence ([OID oid] ++ (if critical then [Boolean True] else []) ++ [OctetString bs])

encodeCertificateHeader :: Certificate -> [ASN1]
encodeCertificateHeader cert =
        eVer ++ eSerial ++ eAlgId ++ eIssuer ++ eValidity ++ eSubject ++ epkinfo ++ eexts
        where
                eVer      = asn1Container (Container Context 0) [IntVal (fromIntegral $ certVersion cert)]
                eSerial   = [IntVal $ certSerial cert]
                eAlgId    = asn1Container Sequence [OID (sigOID $ certSignatureAlg cert), Null]
                eIssuer   = encodeDN $ certIssuerDN cert
                (t1, t2)  = certValidity cert
                eValidity = asn1Container Sequence [UTCTime $ unconvertTime t1, UTCTime $ unconvertTime t2]
                eSubject  = encodeDN $ certSubjectDN cert
                epkinfo   = encodePK $ certPubKey cert
                eexts     = encodeExts $ certExtensions cert

                unconvertTime (day, difftime, z) =
                        let (y, m, d) = toGregorian day in
                        let seconds = floor $ toRational difftime in
                        let h = seconds `div` 3600 in
                        let mi = (seconds `div` 60) `mod` 60 in
                        let s  = seconds `mod` 60 in
                        (fromIntegral y,m,d,h,mi,s,z)
