-- Copyright (C) 2013, 2014, 2017  Fraser Tweedale
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings #-}

{-|

Internal utility functions for encoding/decoding JOSE types.

-}
module Crypto.JOSE.Types.Internal
  (
    insertToObject
  , insertManyToObject
  , encodeB64
  , parseB64
  , encodeB64Url
  , parseB64Url
  , bsToInteger
  , integerToBS
  , intBytes
  , sizedIntegerToBS
  , base64url
  ) where

import Data.Bifunctor (first)
import Data.Tuple (swap)
import Data.Word (Word8)

import Control.Lens
import Control.Lens.Cons.Extras
import Crypto.Number.Basic (log2)
import Data.Aeson.Types
import qualified Data.Aeson.KeyMap as M
import qualified Data.ByteString as B
import qualified Data.ByteString.Base64 as B64
import qualified Data.ByteString.Base64.URL as B64U
import qualified Data.Text as T
import qualified Data.Text.Encoding as E

-- | Insert the given key and value into the given @Value@, which
-- is expected to be an @Object@.  If the value is not an @Object@,
-- this is a no-op.
--
insertToObject :: ToJSON v => Key -> v -> Value -> Value
insertToObject :: Key -> v -> Value -> Value
insertToObject Key
k v
v (Object Object
o) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Key -> Value -> Object -> Object
forall v. Key -> v -> KeyMap v -> KeyMap v
M.insert Key
k (v -> Value
forall a. ToJSON a => a -> Value
toJSON v
v) Object
o
insertToObject Key
_ v
_ Value
v          = Value
v

-- | Insert several key/value pairs to the given @Value@, which
-- is expected to be an @Object@.  If the value is not an @Object@,
-- this is a no-op.
--
insertManyToObject :: [Pair] -> Value -> Value
insertManyToObject :: [Pair] -> Value -> Value
insertManyToObject [Pair]
kvs (Object Object
o) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ (Pair -> Object -> Object) -> Object -> [Pair] -> Object
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Key -> Value -> Object -> Object) -> Pair -> Object -> Object
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Key -> Value -> Object -> Object
forall v. Key -> v -> KeyMap v -> KeyMap v
M.insert) Object
o [Pair]
kvs
insertManyToObject [Pair]
_ Value
v            = Value
v

-- | Produce a parser of base64 encoded text from a bytestring parser.
--
parseB64 :: (B.ByteString -> Parser a) -> T.Text -> Parser a
parseB64 :: (ByteString -> Parser a) -> Text -> Parser a
parseB64 ByteString -> Parser a
f = (String -> Parser a)
-> (ByteString -> Parser a) -> Either String ByteString -> Parser a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail ByteString -> Parser a
f (Either String ByteString -> Parser a)
-> (Text -> Either String ByteString) -> Text -> Parser a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Either String ByteString
decodeB64
  where
    decodeB64 :: Text -> Either String ByteString
decodeB64 = ByteString -> Either String ByteString
B64.decode (ByteString -> Either String ByteString)
-> (Text -> ByteString) -> Text -> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
E.encodeUtf8

-- | Convert a bytestring to a base64 encoded JSON 'String'
--
encodeB64 :: B.ByteString -> Value
encodeB64 :: ByteString -> Value
encodeB64 = Text -> Value
String (Text -> Value) -> (ByteString -> Text) -> ByteString -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
E.decodeUtf8 (ByteString -> Text)
-> (ByteString -> ByteString) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
B64.encode


-- | Prism for encoding / decoding base64url.
--
-- To encode, @'review' base64url@.
-- To decode, @'preview' base64url@.
--
-- Works with any combinations of strict/lazy @ByteString@.
--
base64url ::
  ( AsEmpty s1, AsEmpty s2
  , Cons s1 s1 Word8 Word8
  , Cons s2 s2 Word8 Word8
  ) => Prism' s1 s2
base64url :: Prism' s1 s2
base64url = p ByteString (f ByteString) -> p s1 (f s1)
forall a a s s a t (p :: * -> * -> *) (f :: * -> *).
(Cons a a a a, Cons s s a a, Cons s s a a, Cons t t a a, AsEmpty a,
 AsEmpty t, Profunctor p, Functor f) =>
p a (f s) -> p s (f t)
reconsIso (p ByteString (f ByteString) -> p s1 (f s1))
-> (p s2 (f s2) -> p ByteString (f ByteString))
-> p s2 (f s2)
-> p s1 (f s1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. p ByteString (f ByteString) -> p ByteString (f ByteString)
forall (p :: * -> * -> *) (f :: * -> *).
(Choice p, Applicative f) =>
p ByteString (f ByteString) -> p ByteString (f ByteString)
b64u (p ByteString (f ByteString) -> p ByteString (f ByteString))
-> (p s2 (f s2) -> p ByteString (f ByteString))
-> p s2 (f s2)
-> p ByteString (f ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. p s2 (f s2) -> p ByteString (f ByteString)
forall a a s s a t (p :: * -> * -> *) (f :: * -> *).
(Cons a a a a, Cons s s a a, Cons s s a a, Cons t t a a, AsEmpty a,
 AsEmpty t, Profunctor p, Functor f) =>
p a (f s) -> p s (f t)
reconsIso
  where
    b64u :: p ByteString (f ByteString) -> p ByteString (f ByteString)
b64u = (ByteString -> ByteString)
-> (ByteString -> Either ByteString ByteString)
-> forall (p :: * -> * -> *) (f :: * -> *).
   (Choice p, Applicative f) =>
   p ByteString (f ByteString) -> p ByteString (f ByteString)
forall b t s a. (b -> t) -> (s -> Either t a) -> Prism s t a b
prism ByteString -> ByteString
B64U.encodeUnpadded (\ByteString
s -> (String -> ByteString)
-> Either String ByteString -> Either ByteString ByteString
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (ByteString -> String -> ByteString
forall a b. a -> b -> a
const ByteString
s) (ByteString -> Either String ByteString
B64U.decodeUnpadded ByteString
s))
    reconsIso :: p a (f s) -> p s (f t)
reconsIso = (s -> a) -> (s -> t) -> Iso s t a s
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso (Getting a s a -> s -> a
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting a s a
forall s1 a s2.
(Cons s1 s1 a a, Cons s2 s2 a a, AsEmpty s2) =>
Getter s1 s2
recons) (Getting t s t -> s -> t
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting t s t
forall s1 a s2.
(Cons s1 s1 a a, Cons s2 s2 a a, AsEmpty s2) =>
Getter s1 s2
recons)
{-# INLINE base64url #-}


-- | Produce a parser of base64url encoded text from a bytestring parser.
--
parseB64Url :: (B.ByteString -> Parser a) -> T.Text -> Parser a
parseB64Url :: (ByteString -> Parser a) -> Text -> Parser a
parseB64Url ByteString -> Parser a
f = Parser a
-> (ByteString -> Parser a) -> Maybe ByteString -> Parser a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Not valid base64url") ByteString -> Parser a
f (Maybe ByteString -> Parser a)
-> (Text -> Maybe ByteString) -> Text -> Parser a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (First ByteString) ByteString ByteString
-> ByteString -> Maybe ByteString
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview Getting (First ByteString) ByteString ByteString
forall s1 s2.
(AsEmpty s1, AsEmpty s2, Cons s1 s1 Word8 Word8,
 Cons s2 s2 Word8 Word8) =>
Prism' s1 s2
base64url (ByteString -> Maybe ByteString)
-> (Text -> ByteString) -> Text -> Maybe ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
E.encodeUtf8

-- | Convert a bytestring to a base64url encoded JSON 'String'
--
encodeB64Url :: B.ByteString -> Value
encodeB64Url :: ByteString -> Value
encodeB64Url = Text -> Value
String (Text -> Value) -> (ByteString -> Text) -> ByteString -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
E.decodeUtf8 (ByteString -> Text)
-> (ByteString -> ByteString) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AReview ByteString ByteString -> ByteString -> ByteString
forall b (m :: * -> *) t. MonadReader b m => AReview t b -> m t
review AReview ByteString ByteString
forall s1 s2.
(AsEmpty s1, AsEmpty s2, Cons s1 s1 Word8 Word8,
 Cons s2 s2 Word8 Word8) =>
Prism' s1 s2
base64url

-- | Convert an unsigned big endian octet sequence to the integer
-- it represents.
--
bsToInteger :: B.ByteString -> Integer
bsToInteger :: ByteString -> Integer
bsToInteger = (Integer -> Word8 -> Integer) -> Integer -> ByteString -> Integer
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
B.foldl (\Integer
acc Word8
x -> Integer
acc Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
256 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Word8 -> Integer
forall a. Integral a => a -> Integer
toInteger Word8
x) Integer
0

-- | Convert an integer to its unsigned big endian representation as
-- an octet sequence.
--
integerToBS :: Integral a => a -> B.ByteString
integerToBS :: a -> ByteString
integerToBS = ByteString -> ByteString
B.reverse (ByteString -> ByteString) -> (a -> ByteString) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe (Word8, a)) -> a -> ByteString
forall a. (a -> Maybe (Word8, a)) -> a -> ByteString
B.unfoldr (((a, Word8) -> (Word8, a)) -> Maybe (a, Word8) -> Maybe (Word8, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Word8) -> (Word8, a)
forall a b. (a, b) -> (b, a)
swap (Maybe (a, Word8) -> Maybe (Word8, a))
-> (a -> Maybe (a, Word8)) -> a -> Maybe (Word8, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe (a, Word8)
forall a b. (Integral a, Num b) => a -> Maybe (a, b)
f)
  where
    f :: a -> Maybe (a, b)
f a
0 = Maybe (a, b)
forall a. Maybe a
Nothing
    f a
x = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> b) -> (a, a) -> (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
quotRem a
x a
256)

sizedIntegerToBS :: Integral a => Int -> a -> B.ByteString
sizedIntegerToBS :: Int -> a -> ByteString
sizedIntegerToBS Int
w = ByteString -> ByteString
zeroPad (ByteString -> ByteString) -> (a -> ByteString) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. Integral a => a -> ByteString
integerToBS
  where zeroPad :: ByteString -> ByteString
zeroPad ByteString
xs = Int -> Word8 -> ByteString
B.replicate (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
B.length ByteString
xs) Word8
0 ByteString -> ByteString -> ByteString
`B.append` ByteString
xs

intBytes :: Integer -> Int
intBytes :: Integer -> Int
intBytes Integer
n = (Integer -> Int
log2 Integer
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1