-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 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 Testlib.Prekeys
  ( getPrekey,
    getPrekeys,
    getLastPrekey,
  )
where

import qualified Codec.CBOR.Encoding as CBOR
import qualified Codec.CBOR.Write as CBOR
import qualified Crypto.PubKey.Ed25519 as Ed25519
import qualified Crypto.Random as Crypto
import Data.Aeson
import qualified Data.ByteArray as ByteArray
import qualified Data.ByteString.Base64 as Base64
import Data.String
import Data.String.Conversions (cs)
import Data.Word
import Prelude

getPrekey :: (Crypto.MonadRandom m) => m Value
getPrekey :: forall (m :: * -> *). MonadRandom m => m Value
getPrekey = Word16 -> m Value
forall (m :: * -> *). MonadRandom m => Word16 -> m Value
mkPrekey Word16
1

getPrekeys :: (Crypto.MonadRandom m) => Word16 -> m [Value]
getPrekeys :: forall (m :: * -> *). MonadRandom m => Word16 -> m [Value]
getPrekeys Word16
n = (Word16 -> m Value) -> [Word16] -> m [Value]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Word16 -> m Value
forall (m :: * -> *). MonadRandom m => Word16 -> m Value
mkPrekey [Word16
0 .. Word16
n]

getLastPrekey :: (Crypto.MonadRandom m) => m Value
getLastPrekey :: forall (m :: * -> *). MonadRandom m => m Value
getLastPrekey = Word16 -> m Value
forall (m :: * -> *). MonadRandom m => Word16 -> m Value
mkPrekey Word16
forall a. Bounded a => a
maxBound

mkPrekey :: (Crypto.MonadRandom m) => Word16 -> m Value
mkPrekey :: forall (m :: * -> *). MonadRandom m => Word16 -> m Value
mkPrekey Word16
prekeyId = do
  pk <- Word16 -> m String
forall (m :: * -> *). MonadRandom m => Word16 -> m String
newPrekey Word16
prekeyId
  pure $ object [fromString "id" .= prekeyId, fromString "key" .= pk]

-- | https://github.com/wireapp/proteus/blob/bb759d762bfde376fa5a8a08b1d1153a345ab28a/src/internal/keys.rs#L305
newPrekey :: (Crypto.MonadRandom m) => Word16 -> m String
newPrekey :: forall (m :: * -> *). MonadRandom m => Word16 -> m String
newPrekey Word16
prekeyId = do
  secretKey <- m SecretKey
forall (m :: * -> *). MonadRandom m => m SecretKey
Ed25519.generateSecretKey
  let publicKey = SecretKey -> PublicKey
Ed25519.toPublic SecretKey
secretKey
  identitySecretKey <- Ed25519.generateSecretKey
  let identityPublicKey = SecretKey -> PublicKey
Ed25519.toPublic SecretKey
identitySecretKey
      encodePublicKey bin
k = Word -> Encoding
CBOR.encodeMapLen Word
1 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Word8 -> Encoding
CBOR.encodeWord8 Word8
0 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> ByteString -> Encoding
CBOR.encodeBytes (bin -> ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
ByteArray.convert bin
k)
      encodedIdentityKey = Word -> Encoding
CBOR.encodeMapLen Word
1 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Word8 -> Encoding
CBOR.encodeWord8 Word8
0 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> PublicKey -> Encoding
forall {bin}. ByteArrayAccess bin => bin -> Encoding
encodePublicKey PublicKey
identityPublicKey
      cbor =
        Encoding -> ByteString
CBOR.toStrictByteString (Encoding -> ByteString) -> Encoding -> ByteString
forall a b. (a -> b) -> a -> b
$
          Word -> Encoding
CBOR.encodeMapLen Word
5
            Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> (Word8 -> Encoding
CBOR.encodeWord8 Word8
0 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Word8 -> Encoding
CBOR.encodeWord8 Word8
1)
            Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> (Word8 -> Encoding
CBOR.encodeWord8 Word8
1 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Word16 -> Encoding
CBOR.encodeWord16 Word16
prekeyId)
            Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> (Word8 -> Encoding
CBOR.encodeWord8 Word8
2 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> PublicKey -> Encoding
forall {bin}. ByteArrayAccess bin => bin -> Encoding
encodePublicKey PublicKey
publicKey)
            Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> (Word8 -> Encoding
CBOR.encodeWord8 Word8
3 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Encoding
encodedIdentityKey)
            Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> (Word8 -> Encoding
CBOR.encodeWord8 Word8
4 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Encoding
CBOR.encodeNull)
  pure . cs $ Base64.encode cbor