{-# LANGUAGE OverloadedStrings, MultiParamTypeClasses, DataKinds #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

{- 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.Monoid
import Data.Maybe
import Data.Word
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.KEY256

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

encrypt :: Tunables -> KeyEncryptionKey -> SecretKey -> EncryptedSecretKey
encrypt tunables kek (SecretKey secret) = EncryptedSecretKey b (keyBruteForceCalc kek)
  where
	-- Raaz does not seem to provide a high-level interface
	-- for AES encryption, so use unsafeEncrypt. The use of 
	-- EncryptableBytes makes sure it's provided with a 
	-- multiple of the AES block size.
	b = Raaz.unsafeEncrypt cipher (keyEncryptionKey kek, keyEncryptionIV kek) $
		getEncryptableBytes $ toEncryptableBytes tunables secret

decrypt :: KeyEncryptionKey -> EncryptedSecretKey -> Maybe SecretKey
decrypt kek (EncryptedSecretKey b _) = SecretKey <$> fromEncryptableBytes pbs
  where
	pbs = EncryptableBytes $ Raaz.unsafeDecrypt cipher (keyEncryptionKey kek, keyEncryptionIV kek) b

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

instance Bruteforceable KeyEncryptionKey UnknownPassword where
	getBruteCostCalc = keyBruteForceCalc

-- | The ExpensiveHash of the Password used as the KeyEncryptionKey
--
-- Name is used as a salt, to prevent rainbow table attacks.
--
-- A random prefix is added to the salt, to force an attacker to
-- run the hash repeatedly.
genKeyEncryptionKey :: Tunables -> Name -> Password -> IO KeyEncryptionKey
genKeyEncryptionKey tunables name password = do
	prg <- Raaz.newPRG () :: IO Raaz.SystemPRG
	saltprefix <- genRandomSaltPrefix prg tunables
	return $ head $
		genKeyEncryptionKeys [saltprefix] tunables name password

-- | A stream of KeyEncryptionKeys, using the specified salt prefixes.
genKeyEncryptionKeys :: [SaltPrefix] -> Tunables -> Name -> Password -> [KeyEncryptionKey]
genKeyEncryptionKeys saltprefixes tunables (Name name) (Password password) =
	map mk saltprefixes
  where
	iv = genIV (Password password)
	-- To brute force data encrypted with a key,
	-- an attacker needs to pay the decryptcost for
	-- each password checked.
	bruteforcecalc = bruteForceLinearSearch decryptcost
	decryptcost = castCost $ randomSaltBytesBruteForceCost kektunables
	kektunables = keyEncryptionKeyTunable tunables

	mk saltprefix = KeyEncryptionKey (hashToAESKey hash) iv decryptcost bruteforcecalc
	  where
		salt = Salt (saltprefix <> name)
		hash = expensiveHash (keyEncryptionKeyHash kektunables) salt password

-- | A stream of all the key encryption keys that need to be tried to 
-- decrypt.
candidateKeyEncryptionKeys :: Tunables -> Name -> Password -> [KeyEncryptionKey]
candidateKeyEncryptionKeys tunables name password =
	genKeyEncryptionKeys saltprefixes tunables name password
  where
	saltprefixes = allByteStringsOfLength $ 
		randomSaltBytes $ keyEncryptionKeyTunable tunables

allByteStringsOfLength :: Int -> [B.ByteString]
allByteStringsOfLength = go []
  where
	go ws n
		| n == 0 = return (B.pack ws)
		| otherwise = do
			w <- [0..255]
			go (w:ws) (n-1)

-- Use the sha256 of the password (truncated) as the IV.
genIV :: Password -> Raaz.IV
genIV (Password password) =
	fromMaybe (error "genIV fromByteString failed") $
		Raaz.fromByteString $ B.take ivlen $
			Raaz.toByteString $ Raaz.sha256 password
  where
	ivlen = fromIntegral $ Raaz.byteSize (undefined :: Raaz.IV)

type SaltPrefix = B.ByteString

genRandomSaltPrefix :: Raaz.SystemPRG -> Tunables -> IO SaltPrefix
genRandomSaltPrefix prg tunables = go [] 
	(randomSaltBytes $ keyEncryptionKeyTunable tunables)
  where
	go ws 0 = return (B.pack ws)
	go ws n = do
		b <- Raaz.random prg :: IO Word8
		go (b:ws) (n-1)

instance Raaz.Random Word8

-- | 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 :: ExpensiveHash -> AesKey
hashToAESKey (ExpensiveHash _ t) = 
	fromMaybe (error "hashToAESKey fromByteString failed") $
		Raaz.fromByteString b
  where
	b = B.take (fromIntegral $ Raaz.byteSize (undefined :: AesKey)) $
		Raaz.toByteString $ Raaz.sha256 (E.encodeUtf8 t)

-- | A bytestring that can be AES enctypted. It includes a checksum,
-- and size, and is padded to the objectSize with NULs.
--
-- This is a multiple of the AES blocksize, as long as objectSize is,
-- which should always be the case.
newtype EncryptableBytes = EncryptableBytes { getEncryptableBytes :: B.ByteString }
	deriving (Show)

toEncryptableBytes :: Tunables -> B.ByteString -> EncryptableBytes
toEncryptableBytes tunables b = EncryptableBytes $ 
	padBytes (objectSize tunables) $ 
		checksum <> sep <> len <> sep <> b
  where
	checksum = Raaz.toByteString $ Raaz.sha256 b
	len = B8.pack (show (B.length b))
	sep = B.singleton 0

padBytes :: Int -> B.ByteString -> B.ByteString
padBytes n b = b <> padding
  where
	len = B.length b
	r = len `rem` n
	padding
		| r == 0 = B.empty
		| otherwise = B.replicate (n - r) 0

fromEncryptableBytes :: EncryptableBytes -> Maybe B.ByteString
fromEncryptableBytes (EncryptableBytes b) = case B.break (== 0) b of
	(checksum, rest)
		| B.null checksum || B.null rest -> Nothing
		| otherwise -> do
			case B.break (== 0) (B.drop 1 rest) of
				(lenb, rest')
					| B.null lenb || B.null rest' -> Nothing
					| otherwise -> do
						len <- readMaybe (B8.unpack lenb)
						let d = B.take len $ B.drop 1 rest'
						if checksum == Raaz.toByteString (Raaz.sha256 d)
							then Just d
							else Nothing