-- Copyright (c) 2020  Herbert Valerio Riedel <hvr@gnu.org>
--
--  This file is free software: you may copy, redistribute and/or modify it
--  under the terms of the GNU General Public License as published by the
--  Free Software Foundation, either version 2 of the License, or (at your
--  option) any later version.
--
--  This file is distributed in the hope that it will be useful, but
--  WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
--  General Public License for more details.
--
--  You should have received a copy of the GNU General Public License
--  along with this program (see `LICENSE`).  If not, see
--  <https://www.gnu.org/licenses/old-licenses/gpl-2.0.html>.

{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE Trustworthy       #-}
{-# LANGUAGE TypeOperators     #-}

-- |
-- Copyright: © Herbert Valerio Riedel 2020
-- SPDX-License-Identifier: GPL-2.0-or-later
--
-- ASN.1 String Types
--
-- This modules features types and associated functions for encoding/decoding common ASN.1 string types from their ASN.1 BER representation according to their standard /universal/ ASN.1 tag number.
--
-- @since 0.1.1
module LDAPv3.ASN1String
  ( ASN1String(..)

  -- * Convenience Sum-type

  , ASN1StringChoice(..)
  , asn1StringChoice'encode
  , asn1StringChoice'decode

  -- * UTF8String

  , UTF8String(UTF8String, utf8String'toShortText)

  -- * UniversalString

  , UniversalString

  -- * BMPString

  , BMPString
  , bmpString'toUcs2CodePoints
  , bmpString'fromUcs2CodePoints

  -- * IA5String

  , IA5String
  , ia5String'toShortText
  , ia5String'fromShortText

  -- * VisibleString

  , VisibleString
  , visibleString'toShortText
  , visibleString'fromShortText

  -- * PrintableString

  , PrintableString
  , printableString'toShortText
  , printableString'fromShortText

  -- * NumericString

  , NumericString
  , numericString'toShortText
  , numericString'fromShortText
  ) where

import           Common                hiding (Option, many, option, some, (<|>))

import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Lazy  as BL
import qualified Data.ByteString.Short as SBS
import           Data.Char             (chr, ord)
import qualified Data.Text.Short       as TS

import           Data.ASN1
import           Data.ASN1.Prim

import qualified Data.Binary           as Bin
import qualified Data.Binary.Get       as Bin
import qualified Data.Binary.Put       as Bin

-- | Typeclass abstracting over common ASN.1 string operations
--
-- @since 0.1.1
class ASN1String a where
  -- | Decode ASN.1 string type from its ASN.1 BER encoding
  asn1string'decode :: ByteString -> Maybe a
  default asn1string'decode :: ASN1 a => ByteString -> Maybe a
  asn1string'decode = runGetMaybe (toBinaryGet asn1decode)

  -- | Encode ASN.1 string type to its ASN.1 BER encoding
  asn1string'encode :: a -> ByteString
  default asn1string'encode :: ASN1 a => a -> ByteString
  asn1string'encode = BL.toStrict . Bin.runPut . void . toBinaryPut . asn1encode

  -- | Predicate for determining whether given code-point is allowed by the respective ASN.1 string type
  asn1string'supportsCodePoint :: Proxy a -> Char -> Bool

  -- | Convert ASN.1 string type to list of code-points
  asn1string'toCodePoints :: a -> [Char]

  -- | Construct ASN.1 string type from list of code-points
  --
  -- This returns 'Nothing' if a code-point cannot be expressed in the respective ASN.1 string type.
  asn1string'fromCodePoints :: [Char] -> Maybe a

-- | Convenient Sum-type combining a subset of the standard ASN.1 string-like types
--
-- See specific string types in "LDAPv3.ASN1String" for details.
data ASN1StringChoice
  = ASN1String'OCTET_STRING     ShortByteString
  | ASN1String'UniversalString  UniversalString
  | ASN1String'UTF8String       ShortText
  | ASN1String'BMPString        BMPString
  | ASN1String'IA5String        IA5String
  | ASN1String'VisibleString    VisibleString
  | ASN1String'PrintableString  PrintableString
  | ASN1String'NumericString    NumericString
  deriving (Show,Eq)

-- | Encodes as ASN.1 BER
instance Bin.Binary ASN1StringChoice where
    put = void . toBinaryPut . go
      where
        go (ASN1String'BMPString t)       = asn1encode t
        go (ASN1String'IA5String t)       = asn1encode t
        go (ASN1String'NumericString t)   = asn1encode t
        go (ASN1String'OCTET_STRING b)    = asn1encode b
        go (ASN1String'PrintableString t) = asn1encode t
        go (ASN1String'UTF8String t)      = asn1encode (UTF8String t)
        go (ASN1String'UniversalString t) = asn1encode t
        go (ASN1String'VisibleString t)   = asn1encode t

    get = toBinaryGet go
      where
        go = dec'CHOICE
               [ ASN1String'OCTET_STRING <$> asn1decode
               , ASN1String'UTF8String . utf8String'toShortText <$> asn1decode
               , ASN1String'PrintableString <$> asn1decode
               , ASN1String'IA5String <$> asn1decode
               , ASN1String'BMPString <$> asn1decode
               , ASN1String'UniversalString <$> asn1decode
               , ASN1String'VisibleString <$> asn1decode
               , ASN1String'NumericString <$> asn1decode
               ]

-- | Encode ASN.1 string choice to its ASN.1 BER encoding
--
-- @since 0.1.1
asn1StringChoice'encode :: ASN1StringChoice -> ByteString
asn1StringChoice'encode = BL.toStrict . Bin.runPut . Bin.put

-- | Decode ASN.1 string choice from its ASN.1 BER encoding
--
-- @since 0.1.1
asn1StringChoice'decode :: ByteString -> Maybe ASN1StringChoice
asn1StringChoice'decode = runGetMaybe Bin.get

----------------------------------------------------------------------------

-- | ASN.1 UTF8String
--
-- > UTF8String ::= [UNIVERSAL 12] IMPLICIT OCTET STRING
--
-- @since 0.1.1
newtype UTF8String = UTF8String { utf8String'toShortText :: ShortText } deriving (Eq,Ord)

instance ASN1String UTF8String where
    asn1string'supportsCodePoint _ = not . isSurr
    asn1string'toCodePoints (UTF8String t) = TS.unpack t
    asn1string'fromCodePoints cps
      | all (not . isSurr) cps = Just $! UTF8String (TS.pack cps)
      | otherwise              = Nothing

instance Show UTF8String where
    show        (UTF8String s) = show s
    showsPrec p (UTF8String s) = showsPrec p s

instance ASN1 UTF8String where
    asn1defTag _ = Universal 12
    asn1encode (UTF8String t) = asn1encode (IMPLICIT t :: 'UNIVERSAL 12 `IMPLICIT` ShortText)
    asn1decode = unwrap <$> asn1decode
      where
        unwrap :: 'UNIVERSAL 12 `IMPLICIT` ShortText -> UTF8String
        unwrap (IMPLICIT t) = UTF8String t

-- | Encodes as ASN.1 BER
instance Bin.Binary UTF8String where
    get = toBinaryGet asn1decode
    put = void . toBinaryPut . asn1encode

----------------------------------------------------------------------------

-- | ASN.1 PrintableString
--
-- > PrintableString ::= [UNIVERSAL 19] IMPLICIT OCTET STRING
--
-- @since 0.1.1
newtype PrintableString = PrintableString ShortText deriving (Eq,Ord)

instance ASN1String PrintableString where
    asn1string'supportsCodePoint _ = isPrintableChar
    asn1string'toCodePoints (PrintableString t) = TS.unpack t
    asn1string'fromCodePoints cps
      | all isPrintableChar cps = Just $! PrintableString (TS.pack cps)
      | otherwise               = Nothing

instance Show PrintableString where
    show        (PrintableString s) = show s
    showsPrec p (PrintableString s) = showsPrec p s

printableString'fromShortText :: ShortText -> Maybe PrintableString
printableString'fromShortText t
  | TS.all isPrintableChar t = Just $! PrintableString t
  | otherwise = Nothing

printableString'fromByteString :: ByteString -> Maybe PrintableString
printableString'fromByteString bs
  | BSC.all isPrintableChar bs = PrintableString <$> TS.fromByteString bs
  | otherwise = Nothing

printableString'toShortText :: PrintableString -> ShortText
printableString'toShortText (PrintableString t) = t

isPrintableChar :: Char -> Bool
isPrintableChar c = case c of
  ' ' -> True
  '*' -> False
  ':' -> True
  '=' -> True
  '?' -> True
  _ | c `inside` ('A','Z') -> True
    | c `inside` ('a','z') -> True
    | c `inside` ('0','9') -> True
    | c `inside` ('\x27','\x2f') -> True -- "'()*+,-./"
    | otherwise -> False

instance ASN1 PrintableString where
    asn1defTag _ = Universal 19
    asn1encode (PrintableString t) = asn1encode (IMPLICIT t :: 'UNIVERSAL 19 `IMPLICIT` ShortText)
    asn1decode = (unwrap <$> asn1decode) `transformVia`
                 (maybe (Left "Invalid code-point in PrintableString") Right . printableString'fromByteString)
      where
        unwrap :: 'UNIVERSAL 19 `IMPLICIT` OCTET_STRING -> ByteString
        unwrap (IMPLICIT t) = t

-- | Encodes as ASN.1 BER
instance Bin.Binary PrintableString where
    get = toBinaryGet asn1decode
    put = void . toBinaryPut . asn1encode

----------------------------------------------------------------------------

-- | ASN.1 NumericString
--
-- > NumericString ::= [UNIVERSAL 18] IMPLICIT OCTET STRING
--
-- @since 0.1.1
newtype NumericString = NumericString ShortText deriving (Eq,Ord)

instance Show NumericString where
    show        (NumericString s) = show s
    showsPrec p (NumericString s) = showsPrec p s

instance ASN1String NumericString where
    asn1string'supportsCodePoint _ = isNumericChar
    asn1string'toCodePoints (NumericString t) = TS.unpack t
    asn1string'fromCodePoints cps
      | all isNumericChar cps = Just $! NumericString (TS.pack cps)
      | otherwise               = Nothing

numericString'fromShortText :: ShortText -> Maybe NumericString
numericString'fromShortText t
  | TS.all isNumericChar t = Just $! NumericString t
  | otherwise = Nothing

numericString'fromByteString :: ByteString -> Maybe NumericString
numericString'fromByteString bs
  | BSC.all isNumericChar bs = NumericString <$> TS.fromByteString bs
  | otherwise = Nothing

numericString'toShortText :: NumericString -> ShortText
numericString'toShortText (NumericString t) = t

isNumericChar :: Char -> Bool
isNumericChar ' ' = True
isNumericChar c   = c `inside` ('0','9')

instance ASN1 NumericString where
    asn1defTag _ = Universal 18
    asn1encode (NumericString t) = asn1encode (IMPLICIT t :: 'UNIVERSAL 18 `IMPLICIT` ShortText)
    asn1decode = (unwrap <$> asn1decode) `transformVia`
                 (maybe (Left "Invalid code-point in NumericString") Right . numericString'fromByteString)
      where
        unwrap :: 'UNIVERSAL 18 `IMPLICIT` OCTET_STRING -> ByteString
        unwrap (IMPLICIT t) = t

-- | Encodes as ASN.1 BER
instance Bin.Binary NumericString where
    get = toBinaryGet asn1decode
    put = void . toBinaryPut . asn1encode

----------------------------------------------------------------------------

-- | ASN.1 VisibleString
--
-- > VisibleString ::= [UNIVERSAL 26] IMPLICIT OCTET STRING
--
-- @since 0.1.1
newtype VisibleString = VisibleString ShortText deriving (Eq,Ord)

instance ASN1String VisibleString where
    asn1string'supportsCodePoint _ = isVisibleChar
    asn1string'toCodePoints (VisibleString t) = TS.unpack t
    asn1string'fromCodePoints cps
      | all isVisibleChar cps = Just $! VisibleString (TS.pack cps)
      | otherwise             = Nothing

instance Show VisibleString where
    show        (VisibleString s) = show s
    showsPrec p (VisibleString s) = showsPrec p s

visibleString'fromShortText :: ShortText -> Maybe VisibleString
visibleString'fromShortText t
  | TS.all isVisibleChar t = Just $! VisibleString t
  | otherwise = Nothing

visibleString'fromByteString :: ByteString -> Maybe VisibleString
visibleString'fromByteString bs
  | BSC.all isVisibleChar bs = VisibleString <$> TS.fromByteString bs
  | otherwise = Nothing

visibleString'toShortText :: VisibleString -> ShortText
visibleString'toShortText (VisibleString t) = t

isVisibleChar :: Char -> Bool
isVisibleChar c = '\x20' <= c && c <= '\x7e' -- aka 'isPrint && isAscii'

instance ASN1 VisibleString where
    asn1defTag _ = Universal 26
    asn1encode (VisibleString t) = asn1encode (IMPLICIT t :: 'UNIVERSAL 26 `IMPLICIT` ShortText)
    asn1decode = (unwrap <$> asn1decode) `transformVia`
                 (maybe (Left "Invalid code-point in VisibleString") Right . visibleString'fromByteString)
      where
        unwrap :: 'UNIVERSAL 26 `IMPLICIT` OCTET_STRING -> ByteString
        unwrap (IMPLICIT t) = t

-- | Encodes as ASN.1 BER
instance Bin.Binary VisibleString where
    get = toBinaryGet asn1decode
    put = void . toBinaryPut . asn1encode

----------------------------------------------------------------------------

-- | ASN.1 IA5String
--
-- > IA5String ::= [UNIVERSAL 22] IMPLICIT OCTET STRING
--
-- @since 0.1.1
newtype IA5String = IA5String ShortText deriving (Eq,Ord)

instance ASN1String IA5String where
    asn1string'supportsCodePoint _ = isIA5Char
    asn1string'toCodePoints (IA5String t) = TS.unpack t
    asn1string'fromCodePoints cps
      | all isIA5Char cps = Just $! IA5String (TS.pack cps)
      | otherwise         = Nothing

instance Show IA5String where
    show        (IA5String s) = show s
    showsPrec p (IA5String s) = showsPrec p s

ia5String'fromShortText :: ShortText -> Maybe IA5String
ia5String'fromShortText t
  | TS.all isIA5Char t = Just $! IA5String t
  | otherwise = Nothing

ia5String'fromByteString :: ByteString -> Maybe IA5String
ia5String'fromByteString bs
  | BSC.all isIA5Char bs = IA5String <$> TS.fromByteString bs
  | otherwise = Nothing

ia5String'toShortText :: IA5String -> ShortText
ia5String'toShortText (IA5String t) = t

isIA5Char :: Char -> Bool
isIA5Char c = c <= '\x7f'

instance ASN1 IA5String where
    asn1defTag _ = Universal 22
    asn1encode (IA5String t) = asn1encode (IMPLICIT t :: 'UNIVERSAL 22 `IMPLICIT` ShortText)
    asn1decode = (unwrap <$> asn1decode) `transformVia`
                 (maybe (Left "Invalid code-point in IA5String") Right . ia5String'fromByteString)
      where
        unwrap :: 'UNIVERSAL 22 `IMPLICIT` OCTET_STRING -> ByteString
        unwrap (IMPLICIT t) = t

-- | Encodes as ASN.1 BER
instance Bin.Binary IA5String where
    get = toBinaryGet asn1decode
    put = void . toBinaryPut . asn1encode

----------------------------------------------------------------------------

-- | ASN.1 BMPString
--
-- > BMPString ::= [UNIVERSAL 30] IMPLICIT OCTET STRING
--
-- NB: The surrogate-pair range U+D800 through U+DFFF is tolerated and thus the responsibility of code converting to and
-- from 'BMPString'
--
-- @since 0.1.1
newtype BMPString = BMPString SBS.ShortByteString deriving (Eq,Ord)

instance ASN1String BMPString where
    asn1string'supportsCodePoint _ = (<= '\xffff')
    asn1string'toCodePoints   = bmpString'toString
    asn1string'fromCodePoints = bmpString'fromString

instance Show BMPString where
    show = show . bmpString'toString
    showsPrec p = showsPrec p . bmpString'toString

bmpString'toUcs2CodePoints :: BMPString -> [Word16]
bmpString'toUcs2CodePoints (BMPString sbs) = go (SBS.unpack sbs)
  where
    go (h:l:rest) = (fromIntegral h*0x100)+fromIntegral l : go rest
    go []         = []
    go [_]        = impossible -- forbidden by invariant

bmpString'fromUcs2CodePoints :: [Word16] -> BMPString
bmpString'fromUcs2CodePoints cps = BMPString (SBS.pack $ go cps)
  where
    go (cp:rest) = fromIntegral (cp `unsafeShiftR` 8) : fromIntegral (cp .&. 0xff) : go rest
    go []        = []

-- NB: Surrogate pair code-points (U+D800 through U+DFFF) are transparently emitted as surrogate 'Char' code-points
bmpString'toString :: BMPString -> String
bmpString'toString = map (chr . fromIntegral) . bmpString'toUcs2CodePoints

-- NB: Surrogate pair code-points (U+D800 through U+DFFF) are not rejected in order for 'bmpString'toString' to be an inverse operation.
bmpString'fromString :: String -> Maybe BMPString
bmpString'fromString s
  | all (\c -> c <= '\xffff') s = Just $! bmpString'fromUcs2CodePoints $ map (fromIntegral . ord) s
  | otherwise = Nothing

bmpString'fromByteString :: ByteString -> Maybe BMPString
bmpString'fromByteString bs
 | even (BSC.length bs) = Just $! BMPString (SBS.toShort bs)
 | otherwise = Nothing

instance ASN1 BMPString where
    asn1defTag _ = Universal 30
    asn1encode (BMPString t) = asn1encode (IMPLICIT (SBS.fromShort t) :: 'UNIVERSAL 30 `IMPLICIT` OCTET_STRING)
    asn1decode = (unwrap <$> asn1decode) `transformVia`
                 (maybe (Left "Invalid code-point in BMPString") Right . bmpString'fromByteString)
      where
        unwrap :: 'UNIVERSAL 30 `IMPLICIT` OCTET_STRING -> ByteString
        unwrap (IMPLICIT t) = t

-- | Encodes as ASN.1 BER
instance Bin.Binary BMPString where
    get = toBinaryGet asn1decode
    put = void . toBinaryPut . asn1encode

----------------------------------------------------------------------------

-- | ASN.1 UniversalString
--
-- > UniversalString ::= [UNIVERSAL 28] IMPLICIT OCTET STRING
--
-- NB: The surrogate-pair range U+D800 through U+DFFF is tolerated and thus becomes the responsibility of code converting to and from 'UniversalString'
--
-- @since 0.1.1
newtype UniversalString = UniversalString SBS.ShortByteString
  deriving (Eq,Ord)

instance ASN1String UniversalString where
    asn1string'supportsCodePoint _ = const True
    asn1string'toCodePoints = universalString'toString
    asn1string'fromCodePoints = \s -> Just $! universalString'fromString s

instance Show UniversalString where
    show = show . universalString'toString
    showsPrec p = showsPrec p . universalString'toString

-- NB: Surrogate pair code-points (U+D800 through U+DFFF) are transparently emitted as surrogate 'Char' code-points
universalString'toString :: UniversalString -> String
universalString'toString = maybe impossible id . bsToUcs4 . (\(UniversalString x) -> SBS.fromShort x)

-- NB: Surrogate pair code-points (U+D800 through U+DFFF) are not rejected in order for 'universalString'toString' to be an inverse operation.
universalString'fromString :: String -> UniversalString
universalString'fromString = UniversalString . SBS.toShort . ucs4ToBs

universalString'fromByteString :: ByteString -> Maybe UniversalString
universalString'fromByteString bs
 | Just _ <- bsToUcs4 bs = Just (UniversalString $ SBS.toShort bs)
 | otherwise = Nothing

-- internal
bsToUcs4 :: ByteString -> Maybe [Char]
bsToUcs4 bs
  | (n,0) <- BSC.length bs `quotRem` 4 = runGetMaybe (repGet n f) bs
  | otherwise = Nothing
  where
    f = do x <- Bin.getWord32be
           guard (x <= 0x10ffff)
           pure $! chr (fromIntegral x)

ucs4ToBs :: [Char] -> ByteString
ucs4ToBs cs = BL.toStrict $ Bin.runPut (mapM_ Bin.putWord32be cs')
  where
    cs' :: [Word32]
    cs' = map (fromIntegral . ord) cs

instance ASN1 UniversalString where
    asn1defTag _ = Universal 28
    asn1encode (UniversalString t) = asn1encode (IMPLICIT (SBS.fromShort t) :: 'UNIVERSAL 28 `IMPLICIT` OCTET_STRING)
    asn1decode = (unwrap <$> asn1decode) `transformVia`
                 (maybe (Left "Invalid code-point in UniversalString") Right . universalString'fromByteString)
      where
        unwrap :: 'UNIVERSAL 28 `IMPLICIT` OCTET_STRING -> ByteString
        unwrap (IMPLICIT t) = t

-- | Encodes as ASN.1 BER
instance Bin.Binary UniversalString where
    get = toBinaryGet asn1decode
    put = void . toBinaryPut . asn1encode

----------------------------------------------------------------------------
-- helpers

runGetMaybe :: Bin.Get a -> ByteString -> Maybe a
runGetMaybe g bs = case Bin.runGetOrFail g (BL.fromStrict bs) of
  Left _ -> Nothing
  Right (rest,_,x)
    | BL.null rest -> Just $! x
    | otherwise     -> Nothing

repGet :: Int -> Bin.Get a -> Bin.Get [a]
repGet n g = go [] n
 where
   go xs 0 = return $! reverse xs
   go xs i = do { x <- g; x `seq` go (x:xs) (i-1) }

isSurr :: Char -> Bool
isSurr c = c >= '\xd800' && c <= '\xdfff'