{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE OverloadedStrings #-}

module Nats.Nkeys.Codec (KeyPrefix (..), Nats.Nkeys.Codec.encode, encodeSeed, fromByte, toByte, decode, extractSeedPrefix, extractCrc) where

import Data.Binary (encode, putWord8)
import Data.Bits
import Debug.Trace (trace)
import Text.Printf (printf)
import Data.ByteString as B
import Data.ByteString.Base32 (decodeBase32, decodeBase32Unpadded, encodeBase32Unpadded)
import Data.Data
import Data.Text (Text, pack, append)
import Data.Text.Encoding (encodeUtf8)
import Data.Word (Word16, Word8)
import Nats.Nkeys.Crc

-- | Represents the well-known prefixes available for NATS-encoded keys
data KeyPrefix = Seed -- ^ (__S__) Precedes all seed keys, followed by a type prefix
  | Private -- ^ (__P__) Used for private keys
  | Server -- ^ (__N__) Servers and their ilk (nodes, processes, etc)
  | Cluster -- ^ (__C__) Clusters
  | Operator -- ^ (__O__) Operators
  | Account -- ^ (__A__) Accounts
  | User -- ^ (__U__) Users
  | Curve  -- ^ (__X__) Curve keys used for encryption/decryption
  | Unknown -- ^ (__Z__) Catch-all for unknown prefixes
   deriving (KeyPrefix -> KeyPrefix -> Bool
(KeyPrefix -> KeyPrefix -> Bool)
-> (KeyPrefix -> KeyPrefix -> Bool) -> Eq KeyPrefix
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: KeyPrefix -> KeyPrefix -> Bool
== :: KeyPrefix -> KeyPrefix -> Bool
$c/= :: KeyPrefix -> KeyPrefix -> Bool
/= :: KeyPrefix -> KeyPrefix -> Bool
Eq, Int -> KeyPrefix -> ShowS
[KeyPrefix] -> ShowS
KeyPrefix -> String
(Int -> KeyPrefix -> ShowS)
-> (KeyPrefix -> String)
-> ([KeyPrefix] -> ShowS)
-> Show KeyPrefix
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KeyPrefix -> ShowS
showsPrec :: Int -> KeyPrefix -> ShowS
$cshow :: KeyPrefix -> String
show :: KeyPrefix -> String
$cshowList :: [KeyPrefix] -> ShowS
showList :: [KeyPrefix] -> ShowS
Show, Typeable KeyPrefix
Typeable KeyPrefix =>
(forall (c :: * -> *).
 (forall d b. Data d => c (d -> b) -> d -> c b)
 -> (forall g. g -> c g) -> KeyPrefix -> c KeyPrefix)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c KeyPrefix)
-> (KeyPrefix -> Constr)
-> (KeyPrefix -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c KeyPrefix))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c KeyPrefix))
-> ((forall b. Data b => b -> b) -> KeyPrefix -> KeyPrefix)
-> (forall r r'.
    (r -> r' -> r)
    -> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r)
-> (forall r r'.
    (r' -> r -> r)
    -> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r)
-> (forall u. (forall d. Data d => d -> u) -> KeyPrefix -> [u])
-> (forall u.
    Int -> (forall d. Data d => d -> u) -> KeyPrefix -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix)
-> Data KeyPrefix
KeyPrefix -> Constr
KeyPrefix -> DataType
(forall b. Data b => b -> b) -> KeyPrefix -> KeyPrefix
forall a.
Typeable a =>
(forall (c :: * -> *).
 (forall d b. Data d => c (d -> b) -> d -> c b)
 -> (forall g. g -> c g) -> a -> c a)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c a)
-> (a -> Constr)
-> (a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c a))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c a))
-> ((forall b. Data b => b -> b) -> a -> a)
-> (forall r r'.
    (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall r r'.
    (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall u. (forall d. Data d => d -> u) -> a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> Data a
forall u. Int -> (forall d. Data d => d -> u) -> KeyPrefix -> u
forall u. (forall d. Data d => d -> u) -> KeyPrefix -> [u]
forall r r'.
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r
forall r r'.
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r
forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c KeyPrefix
forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> KeyPrefix -> c KeyPrefix
forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c KeyPrefix)
forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c KeyPrefix)
$cgfoldl :: forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> KeyPrefix -> c KeyPrefix
gfoldl :: forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> KeyPrefix -> c KeyPrefix
$cgunfold :: forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c KeyPrefix
gunfold :: forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c KeyPrefix
$ctoConstr :: KeyPrefix -> Constr
toConstr :: KeyPrefix -> Constr
$cdataTypeOf :: KeyPrefix -> DataType
dataTypeOf :: KeyPrefix -> DataType
$cdataCast1 :: forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c KeyPrefix)
dataCast1 :: forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c KeyPrefix)
$cdataCast2 :: forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c KeyPrefix)
dataCast2 :: forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c KeyPrefix)
$cgmapT :: (forall b. Data b => b -> b) -> KeyPrefix -> KeyPrefix
gmapT :: (forall b. Data b => b -> b) -> KeyPrefix -> KeyPrefix
$cgmapQl :: forall r r'.
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r
gmapQl :: forall r r'.
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r
$cgmapQr :: forall r r'.
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r
gmapQr :: forall r r'.
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> KeyPrefix -> r
$cgmapQ :: forall u. (forall d. Data d => d -> u) -> KeyPrefix -> [u]
gmapQ :: forall u. (forall d. Data d => d -> u) -> KeyPrefix -> [u]
$cgmapQi :: forall u. Int -> (forall d. Data d => d -> u) -> KeyPrefix -> u
gmapQi :: forall u. Int -> (forall d. Data d => d -> u) -> KeyPrefix -> u
$cgmapM :: forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
gmapM :: forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
$cgmapMp :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
gmapMp :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
$cgmapMo :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
gmapMo :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> KeyPrefix -> m KeyPrefix
Data, Typeable)

toByte :: KeyPrefix -> Word8
toByte :: KeyPrefix -> Word8
toByte KeyPrefix
prefix = case KeyPrefix
prefix of
  KeyPrefix
Seed -> Word8
18 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
  KeyPrefix
Private -> Word8
15 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
  KeyPrefix
Server -> Word8
13 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
  KeyPrefix
Cluster -> Word8
2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
  KeyPrefix
Operator -> Word8
14 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
  KeyPrefix
Account -> Word8
0
  KeyPrefix
User -> Word8
20 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
  KeyPrefix
Curve -> Word8
23 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
  KeyPrefix
Unknown -> Word8
25 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3

fromByte :: Word8 -> KeyPrefix
fromByte :: Word8 -> KeyPrefix
fromByte Word8
input = case Word8
input of
  Word8
144 -> KeyPrefix
Seed
  Word8
120 -> KeyPrefix
Private
  Word8
104 -> KeyPrefix
Server
  Word8
16 -> KeyPrefix
Cluster
  Word8
112 -> KeyPrefix
Operator
  Word8
0 -> KeyPrefix
Account
  Word8
160 -> KeyPrefix
User
  Word8
184 -> KeyPrefix
Curve
  Word8
200 -> KeyPrefix
Unknown
  Word8
_ -> KeyPrefix
Unknown

encode :: KeyPrefix -> ByteString -> ByteString
encode :: KeyPrefix -> ByteString -> ByteString
encode KeyPrefix
prefix ByteString
input =
  let raw :: ByteString
raw = Word8 -> ByteString -> ByteString
B.cons (KeyPrefix -> Word8
toByte KeyPrefix
prefix) ByteString
input
   in Text -> ByteString
encodeUtf8 (Text -> ByteString)
-> (ByteString -> Text) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
encodeBase32Unpadded (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
appendCrc ByteString
raw

encodeSeed :: KeyPrefix -> ByteString -> ByteString
encodeSeed :: KeyPrefix -> ByteString -> ByteString
encodeSeed KeyPrefix
publicPrefix ByteString
input =
  let input' :: ByteString
input' = Int -> ByteString -> ByteString
B.take Int
32 ByteString
input
      s :: Word8
s = KeyPrefix -> Word8
toByte KeyPrefix
Seed
      p :: Word8
p = KeyPrefix -> Word8
toByte KeyPrefix
publicPrefix
      raw :: ByteString
raw = [Word8] -> ByteString -> ByteString
prefixBytes [Word8
s Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
p Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
5, Word8 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word8) -> Word8 -> Word8
forall a b. (a -> b) -> a -> b
$ (Word8
p Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
31) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3] ByteString
input'
   in Text -> ByteString
encodeUtf8 (Text -> ByteString)
-> (ByteString -> Text) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
encodeBase32Unpadded (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
appendCrc ByteString
raw

decode :: ByteString -> Either Text ByteString
decode :: ByteString -> Either Text ByteString
decode ByteString
input =
  let decoded :: Either Text ByteString
decoded = ByteString -> Either Text ByteString
decodeBase32Unpadded ByteString
input
      trimmed :: Either Text ByteString
trimmed = Int -> ByteString -> ByteString
dropEnd Int
2 (ByteString -> ByteString)
-> Either Text ByteString -> Either Text ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either Text ByteString
decoded
      crc :: Either Text Word16
crc = ByteString -> Word16
crc16 (ByteString -> Word16)
-> Either Text ByteString -> Either Text Word16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either Text ByteString
trimmed
      expectedCrc :: Either Text Word16
expectedCrc = ByteString -> Word16
extractCrc (ByteString -> Word16)
-> Either Text ByteString -> Either Text Word16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either Text ByteString
decoded       
      crcValid :: Bool
crcValid = case (Either Text Word16
expectedCrc, Either Text Word16
crc) of
        (Left Text
_, Either Text Word16
_) -> Bool
False
        (Either Text Word16
_, Left Text
_) -> Bool
False
        (Right Word16
d, Right Word16
c) -> Word16
d Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
c  
  in
    if Bool
crcValid then
      Either Text ByteString
trimmed
    else
      Text -> Either Text ByteString
forall a b. a -> Either a b
Left (Text
"Invalid CRC " :: Text)

extractCrc :: ByteString -> Word16
extractCrc :: ByteString -> Word16
extractCrc ByteString
input =
  let input' :: ByteString
input' = Int -> ByteString -> ByteString
B.takeEnd Int
2 ByteString
input
  in
  case ByteString -> [Word8]
B.unpack ByteString
input' of
    [Word8
a,Word8
b] -> (Word8, Word8) -> Word16
word16FromBytes (Word8
a, Word8
b)
    [Word8]
_ ->
      Word16
0

extractSeedPrefix :: ByteString -> KeyPrefix
extractSeedPrefix :: ByteString -> KeyPrefix
extractSeedPrefix ByteString
input =
  let [Word8
r0, Word8
r1] = ByteString -> [Word8]
B.unpack (ByteString -> [Word8]) -> ByteString -> [Word8]
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
B.take Int
2 ByteString
input 
      b0 :: Word8
b0 = Word8
r0 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
248
      b1 :: Word8
b1 = ((Word8
r0 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
7) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
5) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. ((Word8
r1 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
248) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
3)
      pb0 :: KeyPrefix
pb0 = Word8 -> KeyPrefix
fromByte Word8
b0
      pb1 :: KeyPrefix
pb1 = Word8 -> KeyPrefix
fromByte Word8
b1
  in
    if KeyPrefix
pb0 KeyPrefix -> KeyPrefix -> Bool
forall a. Eq a => a -> a -> Bool
/= KeyPrefix
Seed
    then KeyPrefix
Unknown
    else KeyPrefix
pb1

prefixBytes :: [Word8] -> ByteString -> ByteString
prefixBytes :: [Word8] -> ByteString -> ByteString
prefixBytes [Word8]
bytes ByteString
input =
  [Word8] -> ByteString
B.pack [Word8]
bytes ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
input
  
appendBytes :: [Word8] -> ByteString -> ByteString
appendBytes :: [Word8] -> ByteString -> ByteString
appendBytes [Word8]
bytes ByteString
input =
  ByteString
input ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> [Word8] -> ByteString
B.pack [Word8]
bytes
  
appendCrc :: ByteString -> ByteString
appendCrc :: ByteString -> ByteString
appendCrc ByteString
raw = 
  ByteString
raw ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> [Word8] -> ByteString
B.pack (Word16 -> [Word8]
encodeWord16 (Word16 -> [Word8]) -> Word16 -> [Word8]
forall a b. (a -> b) -> a -> b
$ ByteString -> Word16
crc16 ByteString
raw)

encodeWord16 :: Word16 -> [Word8]
encodeWord16 :: Word16 -> [Word8]
encodeWord16 Word16
x =
  let right_byte :: Word16
right_byte = Word16
x Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. Word16
0xFF
      left_byte :: Word16
left_byte = ( Word16
x Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftR` Int
8 ) Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. Word16
0xFF
  in (Word16 -> Word8) -> [Word16] -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map Word16 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word16
right_byte, Word16
left_byte]

word16FromBytes :: (Word8, Word8) -> Word16 
word16FromBytes :: (Word8, Word8) -> Word16
word16FromBytes (Word8
a, Word8
b) =
  let a' :: Word16
a' = Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
a
      b' :: Word16
b' = Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
shift (Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b) Int
8
  in
    Word16
a' Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16
b'