{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program 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 Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Data.ZAuth.Token
  ( -- * Token
    Token (..),
    SerializableToken,

    -- * Header
    Header (..),
    Type (..),
    KnownType (..),
    Tag (..),

    -- * Body
    Body,

    -- * Access body
    Access (..),

    -- * User body
    User (..),

    -- * Bot body
    Bot (..),

    -- * Provider body
    Provider (..),

    -- * Serialization
    writeData,
  )
where

import Control.Error
import Data.Attoparsec.ByteString (takeLazyByteString)
import Data.ByteString.Base64.URL
import Data.ByteString.Builder (Builder, byteString, char8)
import Data.ByteString.Conversion
import Data.ByteString.Lazy (toStrict)
import Data.ByteString.Lazy.Char8 (break, drop, split)
import Data.UUID
import Imports hiding (break, drop)
import Sodium.Crypto.Sign (Signature (..))

data Type
  = -- | Access (Used as short-lived token for Users)
    A
  | -- | User (Used as a cookie for Users to refresh access tokens)
    U
  | -- | Bot
    B
  | -- | Provider
    P
  | -- | LegalHold Access (Used as short-lived token for LegalHold Service)
    LA
  | -- | LegalHold User (Used as a cookie for LegalHold Service to refresh access tokens)
    LU
  deriving (Eq, Show)

class KnownType (t :: Type) where
  typeVal :: Type

instance KnownType A where
  typeVal = A

instance KnownType U where
  typeVal = U

instance KnownType B where
  typeVal = B

instance KnownType P where
  typeVal = P

instance KnownType LA where
  typeVal = LA

instance KnownType LU where
  typeVal = LU

-- | Tag: Tokens for Users with no tag are refreshable themselves and called "UserToken"
-- Tokens for Users with the tag 'S' are non-refreshable themselves and called "SessionToken"
-- FUTUREWORK: rename 'S' to 'SessionTag' for clarity
data Tag = S deriving (Eq, Show)

data Token (t :: Type) = Token
  { signature :: !Signature,
    header :: !(Header t),
    body :: !(Body t)
  }

deriving instance (Eq (Body t)) => Eq (Token t)

deriving instance (Show (Body t)) => Show (Token t)

data Header (t :: Type) = Header
  { version :: !Int,
    key :: !Int,
    time :: !Integer,
    tag :: Maybe Tag
  }
  deriving (Eq, Show)

type family Body (t :: Type) where
  Body A = Access
  Body U = User
  Body B = Bot
  Body P = Provider
  Body LA = Access
  Body LU = User

data Access = Access
  { userId :: !UUID,
    clientId :: Maybe Text,
    -- | 'ConnId' is derived from this.
    connection :: !Word64
  }
  deriving (Eq, Show)

data User = User
  { user :: !UUID,
    client :: Maybe Text,
    rand :: !Word32
  }
  deriving (Eq, Show)

data Bot = Bot
  { prov :: !UUID,
    bot :: !UUID,
    conv :: !UUID
  }
  deriving (Eq, Show)

newtype Provider = Provider
  { provider :: UUID
  }
  deriving (Eq, Show)

type Properties = [(LByteString, LByteString)]

instance (KnownType t, ReadProperties (Body t)) => FromByteString (Token t) where
  parser =
    takeLazyByteString >>= \bs ->
      case readToken bs of
        Nothing -> fail "Invalid zauth token"
        Just t -> pure t

-- | Litte alias to help with reducing constraints. These constraints come
-- together in many places because the 'Header' and 'Body' need to be serialized
-- separately from the 'Token' for signing.
--
-- Perhaps it'd be better to just have a special type for an unsigned token.
type SerializableToken t = (ToByteString (Header t), ToByteString (Body t))

instance (SerializableToken t) => ToByteString (Token t) where
  builder = writeToken

-----------------------------------------------------------------------------
-- Reading

readToken :: (KnownType t, ReadProperties (Body t)) => LByteString -> Maybe (Token t)
readToken b = case split '.' b of
  (s : rest) ->
    let p = map pairwise rest
     in Token
          <$> hush (Signature <$> decode (toStrict s))
          <*> readProperties p
          <*> readProperties p
  _ -> Nothing
  where
    pairwise :: LByteString -> (LByteString, LByteString)
    pairwise x = let (k, v) = break (== '=') x in (k, drop 1 v)

class ReadProperties a where
  readProperties :: Properties -> Maybe a

instance (KnownType t) => ReadProperties (Header t) where
  readProperties p = do
    void $ lookup "t" p >>= readType (typeVal @t)
    Header
      <$> (lookup "v" p >>= fromByteString')
      <*> (lookup "k" p >>= fromByteString')
      <*> (lookup "d" p >>= fromByteString')
      <*> (readTag <$> lookup "l" p)
    where
      readType A "a" = Just A
      readType U "u" = Just U
      readType B "b" = Just B
      readType P "p" = Just P
      readType LA "la" = Just LA
      readType LU "lu" = Just LU
      readType _ _ = Nothing
      readTag "s" = Just S
      readTag _ = Nothing

instance ReadProperties Access where
  readProperties t =
    Access
      <$> (lookup "u" t >>= fromLazyASCIIBytes)
      <*> pure (lookup "i" t >>= fromByteString')
      <*> (lookup "c" t >>= fromByteString')

instance ReadProperties User where
  readProperties t =
    User
      <$> (lookup "u" t >>= fromLazyASCIIBytes)
      <*> pure (lookup "i" t >>= fromByteString')
      <*> (lookup "r" t >>= fmap fromHex . fromByteString')

instance ReadProperties Bot where
  readProperties t =
    Bot
      <$> (lookup "p" t >>= fromLazyASCIIBytes)
      <*> (lookup "b" t >>= fromLazyASCIIBytes)
      <*> (lookup "c" t >>= fromLazyASCIIBytes)

instance ReadProperties Provider where
  readProperties t = Provider <$> (lookup "p" t >>= fromLazyASCIIBytes)

-----------------------------------------------------------------------------
-- Writing

writeToken :: (ToByteString (Body t), ToByteString (Header t)) => Token t -> Builder
writeToken t =
  byteString (encode (sigBytes (t.signature)))
    <> dot
    <> writeData (t.header) (t.body)

writeData :: (SerializableToken t) => Header t -> Body t -> Builder
writeData h a = builder h <> dot <> builder a

instance (KnownType t) => ToByteString (Header t) where
  builder h =
    field "v" h.version
      <> dot
      <> field "k" h.key
      <> dot
      <> field "d" h.time
      <> dot
      <> field "t" (typeVal @t)
      <> dot
      <> field "l" (foldMap builder (h.tag))

instance ToByteString Access where
  builder t =
    field "u" (toLazyASCIIBytes $ t.userId)
      <> foldMap (\c -> dot <> field "i" c) (t.clientId)
      <> dot
      <> field "c" (t.connection)

instance ToByteString User where
  builder t =
    field "u" (toLazyASCIIBytes $ t.user)
      <> dot
      <> field "r" (Hex (t.rand))
      <> foldMap (\c -> dot <> field "i" c) (t.client)

instance ToByteString Bot where
  builder t =
    field "p" (toLazyASCIIBytes $ t.prov)
      <> dot
      <> field "b" (toLazyASCIIBytes $ t.bot)
      <> dot
      <> field "c" (toLazyASCIIBytes $ t.conv)

instance ToByteString Provider where
  builder t = field "p" (toLazyASCIIBytes $ t.provider)

instance ToByteString Type where
  builder A = char8 'a'
  builder U = char8 'u'
  builder B = char8 'b'
  builder P = char8 'p'
  builder LA = char8 'l' <> char8 'a'
  builder LU = char8 'l' <> char8 'u'

instance ToByteString Tag where
  builder S = char8 's'

field :: (ToByteString a) => LByteString -> a -> Builder
field k v = builder k <> eq <> builder v

dot, eq :: Builder
dot = char8 '.'
eq = char8 '='
