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

import Debug.Trace

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

-- | 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
	UsePuzzleIV puzzlecost -> do
		prg <- Raaz.newPRG () :: IO Raaz.SystemPRG
		samplekey <- Raaz.random prg :: IO AesKey
		let basek = hashToAESKey samplekey hash
		(PuzzleIV iv) <- genPuzzleIV tunables
		let strongk = (basek, iv)
		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

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 $ 
		let bb = Raaz.unsafeDecrypt cipher (keyEncryptionKey kek) b
		in traceShow ("padded", bb) bb

-- | 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
		UsePuzzleIV _ -> 
			let (k, randomiv) = keyEncryptionKey basekek
			    nrand = sizePuzzleIV tunables
			    mkcandidate (PuzzleIV iv) = basekek
			    	{ keyEncryptionKey = (k, iv) }
			in map (mkcandidate . mkPuzzleIV' randomiv)
				(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 -> 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 puzzle IV starts with some number of random bytes of IV data,
-- and is right-padded with NULL bytes.
data PuzzleIV = PuzzleIV Raaz.IV

-- | Length of the puzzle IV, 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
-- RandomIV, AES decryptions must be done repeatedly, and the
-- time needed for an AES decryption depends on the amount of data.
sizePuzzleIV :: Tunables -> Int
sizePuzzleIV 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
		UsePuzzleIV cost -> case cost of
			GPUCost (Seconds n) -> n
			_ -> error "decryptionPuzzleTunable 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

mkPuzzleIV :: Raaz.IV -> Int -> PuzzleIV
mkPuzzleIV randomiv nrand = mkPuzzleIV' randomiv 
	(B.take nrand $ Raaz.toByteString randomiv)

mkPuzzleIV' :: Raaz.IV -> B.ByteString -> PuzzleIV
mkPuzzleIV' randomiv b = PuzzleIV $
	fromMaybe (error "mkPuzzleIV fromByteString failed") $
		Raaz.fromByteString (b <> padding)
  where
	ivlen = B.length (Raaz.toByteString randomiv)
	padding = B.replicate (ivlen - B.length b) 0

genPuzzleIV :: Tunables -> IO PuzzleIV
genPuzzleIV tunables = do
	prg <- Raaz.newPRG () :: IO Raaz.SystemPRG
	randomiv <- Raaz.random prg :: IO Raaz.IV
	let size = sizePuzzleIV tunables
	return $ mkPuzzleIV randomiv size

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