{-# LANGUAGE OverloadedStrings, MultiParamTypeClasses, DataKinds #-}

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

module Encryption where

import Types
import Tunables
import Cost
import ExpensiveHash
import Data.Bits
import Data.Monoid
import Data.Maybe
import qualified Raaz
import qualified Raaz.Cipher.AES as Raaz
import qualified Raaz.Cipher.Internal as Raaz
import qualified Data.Text.Encoding as E
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8
import Text.Read

type AesKey = Raaz.Key (Raaz.AES 256 'Raaz.CBC)

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

instance Bruteforceable KeyEncryptionKey UnknownPassword where
	getBruteCostCalc = keyBruteForceCalc

cipher :: Raaz.AES 256 'Raaz.CBC
cipher = Raaz.aes256cbc

blocksize :: Int	
blocksize = fromIntegral $ Raaz.blockSize cipher

encrypt :: KeyEncryptionKey -> SecretKey -> EncryptedSecretKey
encrypt kek (SecretKey secret) = EncryptedSecretKey b (keyBruteForceCalc kek)
  where
	-- Raaz does not seem to provide a high-level interface
	-- for AES encryption, so use unsafeEncrypt, doing our own padding
	-- of the secret key, with NULLs, so that it is a multiple of
	-- the block size.
	b = Raaz.unsafeEncrypt cipher (keyEncryptionKey kek) $ getPaddedBytes $
		toPaddedBytes blocksize secret

decrypt :: KeyEncryptionKey -> EncryptedSecretKey -> Maybe SecretKey
decrypt kek (EncryptedSecretKey b _) = SecretKey <$> fromPaddedBytes pbs
  where
	pbs = PaddedBytes $ Raaz.unsafeDecrypt cipher (keyEncryptionKey kek) b

-- | The ExpensiveHash of the Password is combined with a 
-- RandomObstacle to form the AES key. Combination method is logical OR.
--
-- Name is used as a salt, to prevent rainbow table attacks.
genKeyEncryptionKey :: Tunables -> Name -> Password -> IO KeyEncryptionKey
genKeyEncryptionKey tunables name (Password password) = case decryptionPuzzleTunable tunables of
	KeyBlindingLeftSide puzzlecost -> do
		ob@(RandomObstacle ok) <- genRandomObstacle tunables
		let k = hashToAESKey ok hash
		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
	hash@(ExpensiveHash hashcost _) = expensiveHash tunables salt password
	salt = Salt name

-- | Make an AES key out of a hash value.
--
-- Since the ExpensiveHash value is ascii encoded, and has a common prefix,
-- it does not have a high entropy in every byte, and its length is longer
-- than the AES key length. To deal with this, use the SHA256 of
-- the ExpensiveHash, as a bytestring.
hashToAESKey :: AesKey -> ExpensiveHash -> Raaz.KEY256
hashToAESKey (samplekey, _iv) (ExpensiveHash _ t) = 
	fromMaybe (error "hashToAESKey fromByteString failed") $
		Raaz.fromByteString b
  where
	b = B.take (B.length (Raaz.toByteString samplekey)) $
		Raaz.toByteString $ Raaz.sha256 (E.encodeUtf8 t)

-- | 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 left-padded with NULL bytes, so ORing it with an AES
-- key varies the initial bytes of the key.
--
-- The AesKey also includes a random IV.
data RandomObstacle = RandomObstacle AesKey

-- | 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 :: AesKey -> Int -> AesKey
mkRandomObstacle (k, iv) nbytes = (k', iv)
  where
	k' = fromMaybe (error "mkRandomObstacle fromByteString failed") $
		Raaz.fromByteString ob
	kb = Raaz.toByteString k
	padding = B.replicate (B.length kb - nbytes) 0
	ob = padding <> B.take nbytes kb

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

mixinRandomObstacle :: RandomObstacle -> Raaz.KEY256 -> AesKey
mixinRandomObstacle (RandomObstacle (r, iv)) k = (k', iv)
  where
	k' = fromMaybe (error "mixinRandomObstacle fromByteString failed") $
		Raaz.fromByteString $
			Raaz.toByteString r `orBytes` Raaz.toByteString k

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

newtype PaddedBytes = PaddedBytes { getPaddedBytes :: B.ByteString }
	deriving (Show)

-- Pad with NULs. Since the bytestring can itself include NULs, prefix
-- with the length.
toPaddedBytes :: Int -> B.ByteString -> PaddedBytes
toPaddedBytes n b = PaddedBytes $
	B8.pack (show len) <> B.singleton 0 <> b <> padding
  where
	len = B.length b
	r = len `rem` n
	padding
		| r == 0 = B.empty
		| otherwise = B.replicate (n - r) 0

fromPaddedBytes :: PaddedBytes -> Maybe B.ByteString
fromPaddedBytes (PaddedBytes b) = case B.break (== 0) b of
	(header, rest)
		| B.null header || B.null rest -> Nothing
		| otherwise -> do
			len <- readMaybe (B8.unpack header)
			return $ B.take len $ B.drop 1 rest