{-# LANGUAGE OverloadedStrings, MultiParamTypeClasses #-}

module Encryption where

import Types
import Cost
import ExpensiveHash
import Data.Bits
import Data.Monoid
import Data.Maybe
import qualified Data.ByteString as B
import Raaz
import qualified Raaz.Cipher.AES as AES

-- | An AES key, which is used to encrypt the key that is stored
-- in keysafe.
data KeyEncryptionKey = KeyEncryptionKey
	{ keyEncryptionKey :: AES.KEY256
	, keyDecryptionCost :: Cost DecryptionOp
	, keyBruteForceCalc :: CostCalc BruteForceOp UnknownPassword
	}

instance Bruteforceable KeyEncryptionKey UnknownPassword where
	getBruteCostCalc = keyBruteForceCalc

-- | The ExpensiveHash of the Password is combined with a 
-- RandomObstacle to form the AES key. Combination method is logical OR.
genKeyEncryptionKey :: Tunables -> KeyIdent -> Password -> IO KeyEncryptionKey
genKeyEncryptionKey tunables keyident password = case decryptionPuzzleTunable tunables of
	KeyBlindingLeftSide puzzlecost -> do
		ob@(RandomObstacle ok) <- genRandomObstacle tunables
		-- Truncate the hash to the AES key length.
		let truncatedhashb = B.take (B.length (toByteString ok)) hashb
		let k = fromMaybe (error "genKeyEncryptionKey fromByteString failed") $
			fromByteString truncatedhashb
		let strongk = mixinRandomObstacle ob k
		let decryptcost = CombinedCost puzzlecost (castCost hashcost)
		-- To brute force data encrypted with this key,
		-- an attacker needs to pay the decryptcost for
		-- each password checked.
		let bruteforcecalc = bruteForceLinearSearch decryptcost
		return $ KeyEncryptionKey strongk decryptcost bruteforcecalc
  where
	(ExpensiveHash hashcost hashb) = expensiveHash tunables salt password
	salt = Salt keyident

-- | A random value which can be mixed into an AES key to 
-- require decrypting it to perform some brute-force work.
--
-- The random value is right-padded with NULL bytes, so ORing it with an AES
-- key varies the initial bytes of the key.
data RandomObstacle = RandomObstacle AES.KEY256

-- | Length of the random obstacle, in bytes, that will satisfy the
-- decryptionPuzzleCost.
--
-- AES can be calculated more efficiently by a GPU, so the cost must be
-- a GPU cost.
--
-- This depends on the objectSize, because to brute force the
-- RandomObstable, AES decryptions must be done repeatedly, and the
-- time needed for an AES decryption depends on the amount of data.
sizeRandomObstacle :: Tunables -> Int
sizeRandomObstacle tunables = ceiling $ nbits / 8
  where
	-- in 2016, a GPU can run AES at 10 GB/s.
	bytespersecond = 10*1024*1024*1024
	triespersecond = bytespersecond `div` fromIntegral (objectSize tunables)
	targetseconds = case decryptionPuzzleTunable tunables of
		KeyBlindingLeftSide cost -> case cost of
			GPUCost (Seconds n) -> n
			_ -> error "decryptionPuzzleCost must be a GPUCost"
	-- Add one bit of entropy, because a brute-force attack will
	-- on average succeed half-way through the search space.
	nbits :: Double
	nbits = logBase 2 (fromIntegral $ targetseconds * triespersecond) + 1

mkRandomObstacle :: AES.KEY256 -> Int -> AES.KEY256
mkRandomObstacle k nbytes = 
	fromMaybe (error "mkRandomObstacle fromByteString failed") $
		fromByteString ob
  where
	kb = toByteString k
	rightnulls = B.replicate (B.length kb - nbytes) 0
	ob = B.take nbytes kb <> rightnulls

genRandomObstacle :: Tunables -> IO RandomObstacle
genRandomObstacle tunables = do
	prg <- newPRG () :: IO SystemPRG
	randomkey <- random prg :: IO AES.KEY256
	let size = sizeRandomObstacle tunables
	return $ RandomObstacle $ mkRandomObstacle randomkey size

mixinRandomObstacle :: RandomObstacle -> AES.KEY256 -> AES.KEY256
mixinRandomObstacle (RandomObstacle r) k = 
	fromMaybe (error "mixinRandomObstacle fromByteString failed") $
		fromByteString $ toByteString r `orBytes` toByteString k

orBytes :: B.ByteString -> B.ByteString -> B.ByteString
orBytes a b = B.pack $ map (uncurry (.|.)) $ zip (B.unpack a) (B.unpack b)