    -- |
    -- Module      : Data.X509.PublicKey
    -- License     : BSD-style
    -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
    -- Stability   : experimental
    -- Portability : unknown
    --
    -- Private key handling in X.509 infrastructure
    --
    module Data.X509.PrivateKey
        ( PrivKey(..)
        , PrivKeyEC(..)
        , privkeyToAlg
        ) where

    import Control.Applicative ((<$>), pure)
    import Data.Maybe (fromMaybe)
    import Data.Word (Word)

    import Data.ByteArray (ByteArrayAccess, convert)
    import qualified Data.ByteString as B

    import Data.ASN1.Types
    import Data.ASN1.Encoding
    import Data.ASN1.BinaryEncoding
    import Data.ASN1.BitArray
    import Data.ASN1.Stream (getConstructedEnd)

    import Data.X509.AlgorithmIdentifier
    import Data.X509.PublicKey (SerializedPoint(..))
    import Data.X509.OID (lookupByOID, lookupOID, curvesOIDTable)

    import Crypto.Error (CryptoFailable(..))
    import Crypto.Number.Serialize (i2osp, os2ip)
    import qualified Crypto.PubKey.RSA as RSA
    import qualified Crypto.PubKey.DSA as DSA
    import qualified Crypto.PubKey.ECC.Types as ECC
    import qualified Crypto.PubKey.Curve25519 as X25519
    import qualified Crypto.PubKey.Curve448   as X448
    import qualified Crypto.PubKey.Ed25519    as Ed25519
    import qualified Crypto.PubKey.Ed448      as Ed448

    -- | Elliptic Curve Private Key
    --
    -- TODO: missing support for binary curve.
    data PrivKeyEC =
          PrivKeyEC_Prime
            { privkeyEC_priv      :: Integer
            , privkeyEC_a         :: Integer
            , privkeyEC_b         :: Integer
            , privkeyEC_prime     :: Integer
            , privkeyEC_generator :: SerializedPoint
            , privkeyEC_order     :: Integer
            , privkeyEC_cofactor  :: Integer
            , privkeyEC_seed      :: Integer
            }
        | PrivKeyEC_Named
            { privkeyEC_name      :: ECC.CurveName
            , privkeyEC_priv      :: Integer
            }
        deriving (Show,Eq)

    -- | Private key types known and used in X.509
    data PrivKey =
          PrivKeyRSA RSA.PrivateKey -- ^ RSA private key
        | PrivKeyDSA DSA.PrivateKey -- ^ DSA private key
        | PrivKeyEC  PrivKeyEC      -- ^ EC private key
        | PrivKeyX25519 X25519.SecretKey   -- ^ X25519 private key
        | PrivKeyX448 X448.SecretKey       -- ^ X448 private key
        | PrivKeyEd25519 Ed25519.SecretKey -- ^ Ed25519 private key
        | PrivKeyEd448 Ed448.SecretKey     -- ^ Ed448 private key
        deriving (Show,Eq)

    instance ASN1Object PrivKey where
        fromASN1 = privkeyFromASN1
        toASN1 = privkeyToASN1

    privkeyFromASN1 :: [ASN1] -> Either String (PrivKey, [ASN1])
    privkeyFromASN1 asn1 =
      (mapFst PrivKeyRSA <$> rsaFromASN1 asn1) <!>
      (mapFst PrivKeyDSA <$> dsaFromASN1 asn1) <!>
      (mapFst PrivKeyEC <$> ecdsaFromASN1 asn1) <!>
      newcurveFromASN1 asn1
      where
        mapFst f (a, b) = (f a, b)

        Left _ <!> b = b
        a      <!> _ = a

    rsaFromASN1 :: [ASN1] -> Either String (RSA.PrivateKey, [ASN1])
    rsaFromASN1 (Start Sequence : IntVal 0 : IntVal n : IntVal e : IntVal d
        : IntVal p : IntVal q : IntVal dP : IntVal dQ : IntVal qinv
        : End Sequence : as) = pure (key, as)
      where
        key = RSA.PrivateKey (RSA.PublicKey (go n 1) n e) d p q dP dQ qinv
        go m i
            | 2 ^ (i * 8) > m = i
            | otherwise = go m (i + 1)
    rsaFromASN1 (Start Sequence : IntVal 0 : Start Sequence
        : OID [1, 2, 840, 113549, 1, 1, 1] : Null : End Sequence
        : OctetString bytes : End Sequence : as) = do
            asn1 <- mapLeft failure (decodeASN1' BER bytes)
            fmap (const as) <$> rsaFromASN1 asn1
      where
        failure = ("rsaFromASN1: " ++) . show
    rsaFromASN1 _ = Left "rsaFromASN1: unexpected format"

    dsaFromASN1 :: [ASN1] -> Either String (DSA.PrivateKey, [ASN1])
    dsaFromASN1 (Start Sequence : IntVal 0 : IntVal p : IntVal q : IntVal g
        : IntVal _ : IntVal x : End Sequence : as) =
            pure (DSA.PrivateKey (DSA.Params p g q) x, as)
    dsaFromASN1 (Start Sequence : IntVal 0 : Start Sequence
        : OID [1, 2, 840, 10040, 4, 1] : Start Sequence : IntVal p : IntVal q
        : IntVal g : End Sequence : End Sequence : OctetString bytes
        : End Sequence : as) = case decodeASN1' BER bytes of
            Right [IntVal x] -> pure (DSA.PrivateKey (DSA.Params p g q) x, as)
            Right _ -> Left "DSA.PrivateKey.fromASN1: unexpected format"
            Left e -> Left $ "DSA.PrivateKey.fromASN1: " ++ show e
    dsaFromASN1 _ = Left "DSA.PrivateKey.fromASN1: unexpected format"

    ecdsaFromASN1 :: [ASN1] -> Either String (PrivKeyEC, [ASN1])
    ecdsaFromASN1 = go []
      where
        failing = ("ECDSA.PrivateKey.fromASN1: " ++)

        go acc (Start Sequence : IntVal 1 : OctetString bytes : rest) = do
            key <- subgo (oid ++ acc)
            case rest'' of
                End Sequence : rest''' -> pure (key, rest''')
                _ -> Left $ failing "unexpected EC format"
          where
            d = os2ip bytes
            (oid, rest') = spanTag 0 rest
            (_, rest'') = spanTag 1 rest'
            subgo (OID oid_ : _) = maybe failure success mcurve
              where
                failure = Left $ failing $ "unknown curve " ++ show oid_
                success = Right . flip PrivKeyEC_Named d
                mcurve = lookupByOID curvesOIDTable oid_
            subgo (Start Sequence : IntVal 1 : Start Sequence
                : OID [1, 2, 840, 10045, 1, 1] : IntVal p : End Sequence
                : Start Sequence : OctetString a : OctetString b : BitString s
                : End Sequence : OctetString g : IntVal o : IntVal c
                : End Sequence : _) =
                    pure $ PrivKeyEC_Prime d a' b' p g' o c s'
              where
                a' = os2ip a
                b' = os2ip b
                g' = SerializedPoint g
                s' = os2ip $ bitArrayGetData s
            subgo (Null : rest_) = subgo rest_
            subgo [] = Left $ failing "curve is missing"
            subgo _ = Left $ failing "unexpected curve format"
        go acc (Start Sequence : IntVal 0 : Start Sequence
            : OID [1, 2, 840, 10045, 2, 1] : rest) = case rest' of
                (OctetString bytes : rest'') -> do
                    asn1 <- mapLeft (failing . show) (decodeASN1' BER bytes)
                    fmap (const rest'') <$> go (oid ++ acc) asn1
                _ -> Left $ failing "unexpected EC format"
          where
            (oid, rest') = spanEnd 0 rest
        go _ _ = Left $ failing "unexpected EC format"

        spanEnd :: Word -> [ASN1] -> ([ASN1], [ASN1])
        spanEnd = loop id
          where
            loop dlist n (a@(Start _) : as) = loop (dlist . (a :)) (n + 1) as
            loop dlist 0 (End _ : as) = (dlist [], as)
            loop dlist n (a@(End _) : as) = loop (dlist . (a :)) (n - 1) as
            loop dlist n (a : as) = loop (dlist . (a :)) n as
            loop dlist _ [] = (dlist [], [])

        spanTag :: Int -> [ASN1] -> ([ASN1], [ASN1])
        spanTag a (Start (Container _ b) : as) | a == b = spanEnd 0 as
        spanTag _ as = ([], as)

    newcurveFromASN1 :: [ASN1] -> Either String (PrivKey, [ASN1])
    newcurveFromASN1 ( Start Sequence
                      : IntVal v
                      : Start Sequence
                      : OID oid
                      : End Sequence
                      : OctetString bs
                      : xs)
        | isValidVersion v = do
            let (_, ys) = containerWithTag 0 xs
            case primitiveWithTag 1 ys of
                (_, End Sequence : zs) ->
                    case getP oid of
                        Just (name, parse) -> do
                            let err s = Left (name ++ ".SecretKey.fromASN1: " ++ s)
                            case decodeASN1' BER bs of
                                Right [OctetString key] ->
                                    case parse key of
                                        CryptoPassed s -> Right (s, zs)
                                        CryptoFailed e -> err ("invalid secret key: " ++ show e)
                                Right _ -> err "unexpected inner format"
                                Left  e -> err (show e)
                        Nothing -> Left ("newcurveFromASN1: unexpected OID " ++ show oid)
                _ -> Left "newcurveFromASN1: unexpected end format"
        | otherwise = Left ("newcurveFromASN1: unexpected version: " ++ show v)
      where
        getP [1,3,101,110] = Just ("X25519", fmap PrivKeyX25519 . X25519.secretKey)
        getP [1,3,101,111] = Just ("X448", fmap PrivKeyX448 . X448.secretKey)
        getP [1,3,101,112] = Just ("Ed25519", fmap PrivKeyEd25519 . Ed25519.secretKey)
        getP [1,3,101,113] = Just ("Ed448", fmap PrivKeyEd448 . Ed448.secretKey)
        getP _             = Nothing
        isValidVersion version = version >= 0 && version <= 1
    newcurveFromASN1 _ =
        Left "newcurveFromASN1: unexpected format"

    containerWithTag :: ASN1Tag -> [ASN1] -> ([ASN1], [ASN1])
    containerWithTag etag (Start (Container _ atag) : xs)
        | etag == atag = getConstructedEnd 0 xs
    containerWithTag _    xs = ([], xs)

    primitiveWithTag :: ASN1Tag -> [ASN1] -> (Maybe B.ByteString, [ASN1])
    primitiveWithTag etag (Other _ atag bs : xs)
        | etag == atag = (Just bs, xs)
    primitiveWithTag _    xs = (Nothing, xs)

    privkeyToASN1 :: PrivKey -> ASN1S
    privkeyToASN1 (PrivKeyRSA rsa) = rsaToASN1 rsa
    privkeyToASN1 (PrivKeyDSA dsa) = dsaToASN1 dsa
    privkeyToASN1 (PrivKeyEC ecdsa) = ecdsaToASN1 ecdsa
    privkeyToASN1 (PrivKeyX25519 k)  = newcurveToASN1 [1,3,101,110] k
    privkeyToASN1 (PrivKeyX448 k)    = newcurveToASN1 [1,3,101,111] k
    privkeyToASN1 (PrivKeyEd25519 k) = newcurveToASN1 [1,3,101,112] k
    privkeyToASN1 (PrivKeyEd448 k)   = newcurveToASN1 [1,3,101,113] k

    rsaToASN1 :: RSA.PrivateKey -> ASN1S
    rsaToASN1 key = (++)
        [ Start Sequence, IntVal 0, IntVal n, IntVal e, IntVal d, IntVal p
        , IntVal q, IntVal dP, IntVal dQ, IntVal qinv, End Sequence
        ]
      where
        RSA.PrivateKey (RSA.PublicKey _ n e) d p q dP dQ qinv = key

    dsaToASN1 :: DSA.PrivateKey -> ASN1S
    dsaToASN1 (DSA.PrivateKey params@(DSA.Params p g q) y) = (++)
        [ Start Sequence, IntVal 0, IntVal p, IntVal q, IntVal g, IntVal x
        , IntVal y, End Sequence
        ]
      where
        x = DSA.calculatePublic params y

    ecdsaToASN1 :: PrivKeyEC -> ASN1S
    ecdsaToASN1 (PrivKeyEC_Named curveName d) = (++)
        [ Start Sequence, IntVal 1, OctetString (i2osp d)
        , Start (Container Context 0), OID oid, End (Container Context 0)
        , End Sequence
        ]
      where
        err = error . ("ECDSA.PrivateKey.toASN1: " ++)
        oid = fromMaybe (err $ "missing named curve " ++ show curveName)
                        (lookupOID curvesOIDTable curveName)
    ecdsaToASN1 (PrivKeyEC_Prime d a b p g o c s) = (++)
        [ Start Sequence, IntVal 1, OctetString (i2osp d)
        , Start (Container Context 0), Start Sequence, IntVal 1
        , Start Sequence, OID [1, 2, 840, 10045, 1, 1], IntVal p, End Sequence
        , Start Sequence, OctetString a', OctetString b', BitString s'
        , End Sequence, OctetString g' , IntVal o, IntVal c, End Sequence
        , End (Container Context 0), End Sequence
        ]
      where
        a' = i2osp a
        b' = i2osp b
        SerializedPoint g' = g
        s' = BitArray (8 * fromIntegral (B.length bytes)) bytes
          where
            bytes = i2osp s

    newcurveToASN1 :: ByteArrayAccess key => OID -> key -> ASN1S
    newcurveToASN1 oid key = (++)
        [ Start Sequence, IntVal 0, Start Sequence, OID oid, End Sequence
        , OctetString (encodeASN1' DER [OctetString $ convert key])
        , End Sequence
        ]

    mapLeft :: (a0 -> a1) -> Either a0 b -> Either a1 b
    mapLeft f (Left x) = Left (f x)
    mapLeft _ (Right x) = Right x

    -- | Convert a Private key to the Public Key Algorithm type
    privkeyToAlg :: PrivKey -> PubKeyALG
    privkeyToAlg (PrivKeyRSA _)         = PubKeyALG_RSA
    privkeyToAlg (PrivKeyDSA _)         = PubKeyALG_DSA
    privkeyToAlg (PrivKeyEC _)          = PubKeyALG_EC
    privkeyToAlg (PrivKeyX25519 _)      = PubKeyALG_X25519
    privkeyToAlg (PrivKeyX448 _)        = PubKeyALG_X448
    privkeyToAlg (PrivKeyEd25519 _)     = PubKeyALG_Ed25519
    privkeyToAlg (PrivKeyEd448 _)       = PubKeyALG_Ed448