{-# LANGUAGE OverloadedStrings #-}

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

module Tests where

import Types
import Output
import Tunables
import Encryption
import Share
import Storage
import Storage.Local
import Control.Exception
import System.Directory
import System.Posix.Temp
import qualified Data.ByteString.UTF8 as BU8
import qualified Data.ByteString as B
import qualified Data.Set as S

type TestDesc = B.ByteString

type TestResult = Either String ()

type Test = (TestDesc, IO TestResult)

testSuccess :: IO TestResult
testSuccess = return $ Right ()

testFailed :: String -> IO TestResult
testFailed = return . Left

runTest :: Test -> IO Bool
runTest (d, t) = do
	progress $ "testing: " ++ show d ++ " ..."
	r <- t
	case r of
		Right () -> do
			say "ok"
			return True
		Left e -> do
			say $ "failed: " ++ show e
			return False

runTests :: IO ()
runTests = do
	r <- mapM runTest tests
	if all (== True) r
		then say "All tests succeeded."
		else error "Tests failed. Report a bug!"

tests :: [Test]
tests = 
	[ stableNamingTest "stable naming"
	, backupRestoreTest "very small" $ 
		testSecretKey 1
	, backupRestoreTest "full size" $ 
		testSecretKey (objectSize testModeTunables)
	, backupRestoreTest "two chunks" $
		testSecretKey (objectSize testModeTunables + 1)
	, backupRestoreTest "many chunks" $
		testSecretKey (objectSize testModeTunables * 10)
	]

testSecretKey :: Int -> SecretKey
testSecretKey sz = SecretKey $ BU8.fromString $ take sz $ concatMap show [1..sz]

withTestStorageLocations :: (StorageLocations -> IO a) -> IO a
withTestStorageLocations a = bracket setup cleanup go
  where
	setup = mkdtemp "keysafe-test"
	cleanup = removeDirectoryRecursive
	go tmpdir = a $ StorageLocations $
		map (localStorage LocallyPreferred (testStorageDir tmpdir) . show)
			[1..100 :: Int]

-- | Test of backup and restore of a SecretKey.
backupRestoreTest :: TestDesc -> SecretKey -> Test
backupRestoreTest testdesc secretkey =
	("backup and restore (" <> testdesc <> ")", runtest)
  where
	runtest = withTestStorageLocations $ \storagelocations -> do
		backup storagelocations
		restore storagelocations

	backup storagelocations = do
		kek <- genKeyEncryptionKey tunables name password
		let esk = encrypt tunables kek secretkey
		shares <- genShares esk tunables
		let sis = shareIdents tunables name AnyGpgKey
		_ <- storeShares storagelocations sis shares (return ())
		return ()

	restore storagelocations = do
		let sis = shareIdents tunables name AnyGpgKey
		(shares, sis', _) <- retrieveShares storagelocations sis (return ())
		let candidatekeys = candidateKeyEncryptionKeys tunables name password
		case combineShares tunables [shares] of
			Left e -> testFailed e
			Right esk -> restorerest storagelocations [shares] sis' $
				tryDecrypt candidatekeys esk
	
	restorerest storagelocations firstshares sis r = case r of
		DecryptFailed -> testFailed "DecryptFailed"
		DecryptSuccess restoredsecretkey ->
			if restoredsecretkey == secretkey
				then testSuccess
				else testFailed "restore yielded different value than was backed up"
		DecryptIncomplete kek -> do
			(nextshares, sis', _) <- retrieveShares storagelocations sis (return ())
			let shares = firstshares ++ [nextshares]
			case combineShares tunables shares of
				Left e -> testFailed e
				Right esk -> restorerest storagelocations shares sis' $
					decrypt kek esk
	
	name = Name testdesc
	password = Password "password"
	-- testModeTunables is used, to avoid this taking a very
	-- long time to run.
	tunables = testModeTunables

-- | It's important that StorableObjectIdent generation be stable;
-- any change to it will cause shards to get lost.
stableNamingTest :: TestDesc -> Test
stableNamingTest testdesc = (testdesc, runtest $ map snd knownTunings)
  where
	runtest [] = testFailed "not stable!"
	runtest (tunables:rest) = do
		let sis = shareIdents tunables name (Distinguisher secretkeysource)
		if S.member knownvalue (head (identsStream sis))
			then testSuccess
			else runtest rest

	name = Name "stable name"
	secretkeysource = GpgKey (KeyId "stable keyid")
	knownvalue = StorableObjectIdent "18b112da9108b4b5e21fa07cfc672e11688110e4c2dc56c8365f0de488cca8cb"