{-# LANGUAGE OverloadedStrings #-}

module Zero.SRP.Client
  (
    calcV
  , calcA
  , calcU
  , calcM1K
  ) where

import           Prelude hiding ((^^))
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy.Char8 as L8
import qualified Data.Digest.Pure.SHA as SHA
import           Data.ByteString (ByteString)
import           Data.Binary
import           Data.Bits (xor)
import           Data.Text (Text)
import           GHC.Int (Int64)
import           Text.Printf (printf)

------------------------------------------------------------------------------
-- SRP
------------------------------------------------------------------------------

-- Prime Group (the one here is 2048 bits)

n' :: Integer
n' = 0xEEAF0AB9ADB38DD69C33F80AFA8FC5E86072618775FF3C0B9EA2314C9C256576D674DF7496EA81D3383B4813D692C6E0E0D5D8E250B98BE48E495C1D6089DAD15DC7D7B46154D6B6CE8EF4AD69B15D4982559B297BCF1885C529F566660E57EC68EDBC3C05726CC02FD4CBF4976EAA9AFD5138FE8376435B9FC61D2FC0EB06E3

nsize :: Int64
nsize = L8.length (encode n')

g :: Integer
g = 2

kGlobal = sha1 [encode n', pad g]

sha1 :: [L8.ByteString] -> Integer
sha1 = SHA.integerDigest . SHA.sha1 . L8.concat

hashString :: [String] -> Integer
hashString = SHA.integerDigest . SHA.sha256 . L8.pack . unwords

hashInteger :: [Integer] -> Integer
hashInteger = hashString . map show

------------------------------------------------------------------------------
-- Combine client / host values
------------------------------------------------------------------------------

calcU a' b' = sha1 [pad a', pad b']

pad :: Integer -> L8.ByteString
pad i = L8.append (L8.replicate padsize '\0') iBytes
  where
    iBytes = encode i
    iBytesLen = L8.length iBytes
    padsize = nsize - iBytesLen

(^^) :: Integer -> Integer -> Integer
(^^) = powerWithModulus n'

-- | Client X
calcX :: String -> Integer -> String -> Integer
calcX i s p =
  sha1 [encode s, encode (sha1 [encode (i ++ ":" ++ p)])]

calcv = (g ^^)

-- | The client-side verifier computation
calcV :: String -> String -> String -> String
calcV i s p = printf "%x" (calcv (calcX i s' p))
  where
    s' = hashString [s]

-- | The client-side crypto puzzle
calcA :: Integer -> Integer
calcA a = g ^^ a `mod` n'

-- | Strong shared key
calcK :: Integer -> Integer -> Integer -> Integer -> Integer
calcK x a b' u =
  (b' - (kGlobal * g ^^ x)) ^^ (a + u * x)

-- | Calculates the client side key
calcM1K :: String -> String -> Integer -> Integer -> Integer -> Integer -> String -> Maybe (Integer, Integer)
calcM1K i s a a' b' u p =
  if ((b' `mod` n') == 0) || (u == 0) then
    Nothing
  else
    Just (m1, k)
  where
    m1 = hashInteger [hashInteger [n'] `xor` hashInteger [g],
                      i', s', a', b', k]
    k = hashInteger [calcK x a b' u]
    x = calcX i s' p
    i' = hashString [i]
    s' = hashString [s]

-- | The host challenge
calcM2 :: Integer -> Integer -> Integer -> Integer -> Maybe Integer
calcM2 a m1 m2h k' =
  let m2 = hashInteger [a, m1, k'] in
    if m2 /= m2h then
      Nothing
    else
      Just k'

xorKeys xorKey = B.pack . B.zipWith xor xorKey

powerWithModulus :: Integer -> Integer -> Integer -> Integer
powerWithModulus mo n k
   | mo < 0     = powerWithModulus (-mo) n k
   | mo == 0    = n^k
   | k < 0      = error "powerWithModulus: negative exponent"
   | k == 0     = 1
   | k == 1     = n `mod` mo
   | mo == 1    = 0
   | odd k      = powerMod mo n n (k-1)
   | otherwise  = powerMod mo 1 n k

-- calculate the power with auxiliary value to account for
-- odd exponents on the way, we have @mo >= 2@ and @k >= 1@
powerMod :: Integer -> Integer -> Integer -> Integer -> Integer
powerMod mo aux val k
   | k == 1 = (aux * val) `mod` mo
   | odd k  =
      ((powerMod mo $! ((aux*val) `mod` mo)) $! (val^2 `mod` mo)) $! (k `quot` 2)
   | even k = (powerMod mo aux $! (val^2 `mod` mo)) $! (k `quot` 2)
