{-# 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 ByteStrings
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.KEY256

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

encrypt :: Tunables -> KeyEncryptionKey -> SecretKey -> EncryptedSecretKey
encrypt tunables kek (SecretKey secret) = 
	EncryptedSecretKey (chunkByteString (objectSize tunables) b) (keyBruteForceCalc kek)
  where
	-- Raaz does not provide a high-level interface for AES encryption,
	-- so we 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 $ encodeEncryptableBytes tunables secret

data DecryptResult 
	= DecryptSuccess SecretKey
	| DecryptIncomplete KeyEncryptionKey
	-- ^ Returned when the EncryptedSecretKey is truncated.
	| DecryptFailed

instance Show DecryptResult where
	show (DecryptSuccess _) = "DecryptSuccess"
	show (DecryptIncomplete _) = "DecryptIncomplete"
	show DecryptFailed = "DecryptFailed"

decrypt :: KeyEncryptionKey -> EncryptedSecretKey -> DecryptResult
decrypt kek (EncryptedSecretKey cs _) = case decodeEncryptableBytes pbs of
	Nothing -> DecryptFailed
	Just (DecodeSuccess secretkey) -> DecryptSuccess (SecretKey secretkey)
	Just DecodeIncomplete -> DecryptIncomplete kek
  where
	pbs = EncryptableBytes $
		Raaz.unsafeDecrypt cipher (keyEncryptionKey kek, keyEncryptionIV kek) b
	b = B.concat cs

-- | Tries each candidate key in turn until one unlocks the encrypted data.
tryDecrypt :: Candidates KeyEncryptionKey -> EncryptedSecretKey -> DecryptResult
tryDecrypt (Candidates l _ _) esk = go l
  where
	go [] = DecryptFailed
	go (kek:rest) = case decrypt kek esk of
		DecryptFailed -> go rest
		r -> r

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

instance HasCreationCost KeyEncryptionKey where
	getCreationCost = keyCreationCost

instance HasDecryptionCost KeyEncryptionKey where
	getDecryptionCost = keyDecryptionCost

instance Bruteforceable KeyEncryptionKey UnknownPassword where
	getBruteCostCalc = keyBruteForceCalc

data Candidates a = Candidates [a] (Cost CreationOp) (Cost DecryptionOp)

instance HasCreationCost (Candidates a) where
	getCreationCost (Candidates _ c _) = c

instance HasDecryptionCost (Candidates a) where
	getDecryptionCost (Candidates _ _ c) = c

-- | 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
	saltprefix <- genRandomSaltPrefix 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 (Name name)
	-- 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 (getCreationCost hash) 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 -> Candidates KeyEncryptionKey
candidateKeyEncryptionKeys tunables name password =
	let ks@(k:_) = genKeyEncryptionKeys saltprefixes tunables name password
	in Candidates ks (getCreationCost k) (getDecryptionCost k)
  where
	saltprefixes = allByteStringsOfLength $ 
		randomSaltBytes $ keyEncryptionKeyTunable tunables

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

type SaltPrefix = B.ByteString

genRandomSaltPrefix :: Tunables -> IO SaltPrefix
genRandomSaltPrefix tunables = randomByteStringOfLength n
  where
	n = randomSaltBytes $ keyEncryptionKeyTunable tunables

-- | 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.sizeOf (undefined :: AesKey)) $
		Raaz.toByteString $ Raaz.sha256 (E.encodeUtf8 t)

-- | A bytestring that can be AES encrypted.
--
-- It is padded to a multiple of the objectSize with NULs.
-- Since objectSize is a multiple of the AES blocksize, so is this.
--
-- Format is:
--
-- 	sizeNULsizeshaNULdatashaNULdata
--
-- The size gives the length of the data. If the data is shorter
-- than that, we know that the bytestring is truncated.
--
-- The datasha is the sha256 of the data. This is checked when decoding
-- to guard against corruption.
--
-- The sizesha is the sha256 of the size. This is included as a sanity
-- check that the right key was used to decrypt it. It's not unlikely
-- that using the wrong key could result in a bytestring that starts
-- with wrongsizeNUL, but it's astronomically unlikely that the
-- sizesha matches in this case.
newtype EncryptableBytes = EncryptableBytes { getEncryptableBytes :: B.ByteString }
	deriving (Show)

encodeEncryptableBytes :: Tunables -> B.ByteString -> EncryptableBytes
encodeEncryptableBytes tunables content = EncryptableBytes $ 
	padBytes (objectSize tunables) $ B.intercalate sep
		[ size
		, sha size
		, sha content
		, content
		]
  where
	size = B8.pack (show (B.length content))
	sep = B.singleton 0

-- | Encoded, so that it does not contain any NULs.
sha :: B.ByteString -> B.ByteString
sha = Raaz.toByteString . enc . Raaz.sha256
  where
	enc :: Raaz.Encodable a => a -> Raaz.Base16
	enc = Raaz.encode

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

data DecodeResult
	= DecodeSuccess B.ByteString
	| DecodeIncomplete
	deriving (Show)

decodeEncryptableBytes :: EncryptableBytes -> Maybe DecodeResult
decodeEncryptableBytes (EncryptableBytes b) = do
	(sizeb, rest) <- getword b
	(sizesha, rest') <- getword rest
	(contentsha, rest'') <- getword rest'
	if sha sizeb /= sizesha
		then Nothing
		else do
			size <- readMaybe (B8.unpack sizeb)
			let content = B.take size rest''
			if B.length content /= size
				then return DecodeIncomplete
				else if sha content /= contentsha
					then Nothing
					else return (DecodeSuccess content)
  where
	getword d = case B.break (== 0) d of
		(w, rest)
			| B.null w || B.null rest-> Nothing
			| otherwise -> Just (w, B.drop 1 rest)