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

-- | An AES key, which is used to encrypt the 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 AES key.
-- 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
		prg <- Raaz.newPRG () :: IO Raaz.SystemPRG
		ob <- genRandomObstacle tunables
		randomkey <- Raaz.random prg :: IO AesKey
		let k = hashToAESKey randomkey hash
		let strongk = mixinRandomObstacle ob k
		let decryptcost = CombinedCost puzzlecost (castCost hashcost)
		iv <- genIV (Password password)
		-- 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 iv decryptcost bruteforcecalc
  where
	hash@(ExpensiveHash hashcost _) = expensiveHash tunables salt password
	salt = Salt name

-- Use the sha256 of the password (truncated) as the IV.
genIV :: Password -> IO Raaz.IV
genIV (Password password) = do
	prg <- Raaz.newPRG () :: IO Raaz.SystemPRG
	exampleiv <- Raaz.random prg :: IO Raaz.IV
	let ivlen = B.length (Raaz.toByteString exampleiv)
	return $ fromMaybe (error "genIV fromByteString failed") $
		Raaz.fromByteString $ B.take ivlen $
			Raaz.toByteString $ Raaz.sha256 password

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

-- | A stream of all the key encryption keys that need to be tried to 
-- decrypt. 
--
-- The DecryptionPuzzleTunable is used to alter the AES key
-- in some way; this stream includes every possible solution to the puzzle.
candidateKeyEncryptionKeys :: Tunables -> KeyEncryptionKey -> [KeyEncryptionKey]
candidateKeyEncryptionKeys tunables basekek = 
	case decryptionPuzzleTunable tunables of
		KeyBlindingLeftSide _ -> 
			let k = keyEncryptionKey basekek
			    nrand = sizeRandomObstacle tunables
			    mkcandidate b = 
			    	let ob = mkRandomObstacle' k b
				in basekek { keyEncryptionKey = mixinRandomObstacle ob k }
			in map mkcandidate (allByteStringsOfLength nrand)

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)

-- | 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 -> AesKey
hashToAESKey samplekey (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 
		| targetseconds < 1 = 0
		| otherwise = logBase 2 (fromIntegral $ targetseconds * triespersecond) + 1

mkRandomObstacle :: AesKey -> Int -> RandomObstacle
mkRandomObstacle k nbytes = mkRandomObstacle' k $
	B.take nbytes $ Raaz.toByteString k

mkRandomObstacle' :: AesKey -> B.ByteString -> RandomObstacle
mkRandomObstacle' examplek b = RandomObstacle $
	fromMaybe (error "mkRandomObstacle' fromByteString failed") $
		Raaz.fromByteString (b <> padding)
  where
	klen = B.length (Raaz.toByteString examplek)
	padding = B.replicate (klen - B.length b) 0

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

mixinRandomObstacle :: RandomObstacle -> AesKey -> AesKey
mixinRandomObstacle (RandomObstacle r) k = k'
  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)

-- | 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