{-# LANGUAGE OverloadedStrings #-}

{- Copyright 2016 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

module ExpensiveHash where

import Types
import Tunables
import Cost
import Serialization ()
import qualified Data.Text as T
import qualified Data.ByteString as B
import qualified Crypto.Argon2 as Argon2
import Raaz.Core.Encode
import Data.Time.Clock
import Control.DeepSeq
import Control.Monad
import Data.Monoid
import Data.List
import Data.Maybe
import Text.Read

-- | A hash that is expensive to calculate.
--
-- This is a lynchpin of keysafe's security, because using this hash
-- as an encryption key forces brute force attackers to generate
-- hashes over and over again, taking a very long time.
data ExpensiveHash = ExpensiveHash (Cost CreationOp) T.Text
	deriving (Show)

instance HasCreationCost ExpensiveHash where
	getCreationCost (ExpensiveHash c _) = c

data Salt t = Salt t

expensiveHash :: Encodable t => ExpensiveHashTunable -> Salt t -> B.ByteString -> ExpensiveHash
expensiveHash (UseArgon2 cost opts) (Salt s) b = ExpensiveHash cost $
	-- Using hashEncoded here and not hash,
	-- because of this bug:
	-- https://github.com/ocharles/argon2/issues/3
	Argon2.hashEncoded opts b argonsalt
  where
	-- argon salt cannot be shorter than 8 bytes, so pad with spaces.
	argonsalt = 
		let sb = toByteString s
		in sb <> B.replicate (8 - B.length sb ) 32

benchmarkExpensiveHash :: Int -> ExpensiveHashTunable -> Cost op -> IO (BenchmarkResult (Cost op))
benchmarkExpensiveHash rounds tunables@(UseArgon2 _ hashopts) expected = do
	numcores <- fromIntegral <$> getNumCores
	start <- getCurrentTime
	forM_ [1..rounds] $ \_ -> do
		let ExpensiveHash _ t = expensiveHash tunables
			(Salt (GpgKey (KeyId ("dummy" :: B.ByteString))))
			("himom" :: B.ByteString)
		t `deepseq` return ()
	end <- getCurrentTime
	let diff = floor $ end `diffUTCTime` start
	let actual = CPUCost $ Seconds diff
	-- The expected cost is for a single core, so adjust it
	-- based on the number of cores, up to a maximum of the number
	-- of threads that the hash is configred to use.
	let maxthreads = Argon2.hashParallelism hashopts
	let usedcpus = if numcores > maxthreads then maxthreads else numcores
	let adjustedexpected = mapCost (`div` fromIntegral usedcpus) expected
	return $ BenchmarkResult
		{ expectedBenchmark = adjustedexpected
		, actualBenchmark = actual
		}

-- Number of physical cores. This is not the same as
-- getNumProcessors, which includes hyper-threading.
getNumCores :: IO Integer
getNumCores = getmax . mapMaybe parse . lines <$> readFile "/proc/cpuinfo"
  where
	getmax [] = error "Unknown number of physical cores."
	getmax l = maximum l
	parse l
		| "core id" `isPrefixOf` l = readMaybe $ dropWhile (/= ':') l
		| otherwise = Nothing

benchmarkTunables :: Tunables -> IO ()
benchmarkTunables tunables = do
	putStrLn "/proc/cpuinfo:"
	putStrLn =<< readFile "/proc/cpuinfo"
	
	-- Rather than run all 256 rounds of this hash, which would
	-- probably take on the order of 1 hour, run only 16, and scale
	-- the expected cost accordingly.
	let normalrounds = fromIntegral $ 
		256 * randomSaltBytes (keyEncryptionKeyTunable tunables)
	putStrLn $ "Benchmarking 16/" ++ show normalrounds ++ " rounds of key encryption key hash..."
	r <- benchmarkExpensiveHash 16
		(keyEncryptionKeyHash $ keyEncryptionKeyTunable tunables)
		(mapCost (`div` (normalrounds `div` 16)) $ randomSaltBytesBruteForceCost $ keyEncryptionKeyTunable tunables)
	print r
	putStrLn $ "Estimated time for " ++ show normalrounds ++ " rounds of key encryption key hash..."
	print $ BenchmarkResult
		{ expectedBenchmark = mapCost (* 16) (expectedBenchmark r)
		, actualBenchmark = mapCost (* 16) (actualBenchmark r)
		}

	putStrLn "Benchmarking 1 round of name generation hash..."
	print =<< benchmarkExpensiveHash 1
		(nameGenerationHash $ nameGenerationTunable tunables)
		(getexpected $ nameGenerationHash $ nameGenerationTunable tunables)
  where
	getexpected (UseArgon2 cost _) = cost