get.hs

-- stylistic
{-# LANGUAGE LambdaCase #-}
-- singletons
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE GADTs #-}
-- main
{-# LANGUAGE DataKinds #-}            -- typelevel programming
{-# LANGUAGE TypeFamilies #-}         -- typelevel programming
{-# LANGUAGE ScopedTypeVariables #-}  -- typelevel programming
{-# LANGUAGE KindSignatures #-}       -- explicit mark of kinds in singatures (data family)
{-# LANGUAGE StandaloneDeriving #-}   -- Show and Eq instances for OtrMessage
{-# LANGUAGE FlexibleContexts #-}     -- Show and Eq instances for OtrMessage
{-# LANGUAGE UndecidableInstances #-} -- Show and Eq instances for OtrMessage
{-# LANGUAGE FlexibleInstances #-}  -- generic binary instances for messages
{-# LANGUAGE DeriveGeneric #-}-- generic binary instances for messages
{-# LANGUAGE ImpredicativeTypes #-}   -- nice typesignature of the otrHeader
{-# OPTIONS_GHC -Wall #-}
-- base
import Control.Monad

import Data.Word
import Data.Proxy

import GHC.Generics
import Data.Singletons.TH

-- bytestring
import Data.ByteString hiding ( empty )

-- binary
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put

-- Boilerplate
type BYTE        = Word8
type SHORT       = Word16
type INT         = Word32
type MPI         = ByteString
type DATA        = ByteString
type CTR         = ByteString
type MAC         = ByteString
type InstanceTag = INT

byte :: Get BYTE
byte = get

short :: Get SHORT
short = getWord16be

int :: Get INT
int = getWord32be

mpi :: Get MPI
mpi = getWord32be >>= getByteString . fromIntegral

data_ :: Get DATA
data_ = mpi

-- | Possible types of the OTR messages.
data MessageType = DiffieHellmanCommit
                 | DiffieHellmanKey
                 | RevealSignature
                 | Signature
                 | DataMessage
                 deriving (Eq,Show)

instance Binary MessageType where
  get = flip fmap byte $ \case
    0x03 -> DataMessage
    0x02 -> DiffieHellmanCommit
    0x0a -> DiffieHellmanKey
    0x11 -> RevealSignature
    0x12 -> Signature
    _    -> error "Unknown Tag"
  put DataMessage         = putWord8 0x03
  put DiffieHellmanCommit = putWord8 0x02
  put DiffieHellmanKey    = putWord8 0x0a
  put RevealSignature     = putWord8 0x11
  put Signature           = putWord8 0x12

$(genSingletons [''MessageType])

-- | Type of the 'OtrMessage' payload.
data family OtrMessagePayload (a::MessageType)
data instance OtrMessagePayload 'DiffieHellmanCommit = MkDiffieHellmanCommit { encryptedGX :: !DATA, hashedGX :: !DATA }
    deriving (Eq, Show, Generic)
data instance OtrMessagePayload 'DiffieHellmanKey = MkDiffieHellmanKey { gY :: MPI }
    deriving (Eq, Show, Generic)
data instance OtrMessagePayload 'RevealSignature = MkRevealSignature { revealedKey :: !DATA, signature :: !Signature }
    deriving (Eq, Show)
data instance OtrMessagePayload 'Signature = MkSignature
    { encryptedSignature :: !DATA, macSignature :: !MAC }
    deriving (Eq, Show)
data instance OtrMessagePayload 'DataMessage = MkDataMessage
    { flags                  :: !BYTE
    , senderKeyID            :: !INT
    , recipientKeyID         :: !INT
    , dhY                    :: !MPI
    , topHalfOfCounterInit   :: !CTR
    , encryptedMessage       :: !DATA
    , authenticator          :: !MAC
    , oldMacKeysToBeRevealed :: !DATA }
    deriving (Eq, Show)

-- Binary instances
instance Binary (OtrMessagePayload 'DiffieHellmanCommit)
instance Binary (OtrMessagePayload 'DiffieHellmanKey)

-- Convenient type names
type Signature           = OtrMessagePayload 'Signature
type RevealSignature     = OtrMessagePayload 'RevealSignature
type DiffieHellmanCommit = OtrMessagePayload 'DiffieHellmanCommit
type DiffieHellmanKey    = OtrMessagePayload 'DiffieHellmanKey
type DataMessage         = OtrMessagePayload 'DataMessage

-- | Known 'OtrMessage'
data OtrMessage payload = OtrMessage
    { protovolVersion     :: !SHORT
    , senderInstanceTag   :: !InstanceTag
    , receiverInstanceTag :: !InstanceTag
    , messagePayload      :: !(OtrMessagePayload payload)
    }

deriving instance Eq (OtrMessagePayload payload) => Eq (OtrMessage payload)
deriving instance Show (OtrMessagePayload payload) => Show (OtrMessage payload)

-- | 'OtrMessage' of unknown type
data SomeOtrMessage
       = DiffieHellmanCommitTag (OtrMessage 'DiffieHellmanCommit)
       | DiffieHellmanKeyTag    (OtrMessage 'DiffieHellmanKey)
       | RevealSignatureTag     (OtrMessage 'RevealSignature)
       | SignatureTag           (OtrMessage 'Signature)
       | DataMessageTag         (OtrMessage 'DataMessage)
       deriving (Eq,Show)

instance Binary SomeOtrMessage where
  get = do
    (t, mk) <- otrHeader
    case toSing t of
      SomeSing g@SDiffieHellmanCommit -> DiffieHellmanCommitTag . mk <$> magic g
      SomeSing g@SDiffieHellmanKey    -> DiffieHellmanKeyTag    . mk <$> magic g
      SomeSing _ -> error "Type do not have binary instance."
    where
      magic :: Binary (OtrMessagePayload k) => Sing k -> Get (OtrMessagePayload k)
      magic _ = get
  put (DiffieHellmanCommitTag msg) = put msg
  put (DiffieHellmanKeyTag msg)    = put msg
  put _                            = error "payload do not have binary instance."

instance (Binary (OtrMessagePayload a), SingI a) => Binary (OtrMessage a) where
  get = otrMessage
  put m@(OtrMessage v s r p) = do
    putWord16be v
    put (fromSing $ singByProxy m)
    putWord32be s
    putWord32be r
    put p

otrHeader :: Get (MessageType, forall payload . OtrMessagePayload payload -> OtrMessage payload)
otrHeader = do
     v <- short
     t <- get
     s <- int
     r <- int
     return (t, OtrMessage v s r)

otrMessage :: forall a . (SingI a, Binary (OtrMessagePayload a)) => Get (OtrMessage a)
otrMessage = magic Proxy where
  magic :: SingI a => Proxy (a::MessageType) -> Get (OtrMessage a)
  magic proxy  = do
    (tag, f) <- otrHeader
    let wanted = fromSing (singByProxy proxy) :: MessageType
    guard (tag == wanted)
    f <$> get

otrDiffieHellmanCommit :: Get (OtrMessage 'DiffieHellmanCommit)
otrDiffieHellmanCommit = otrMessage