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

{-# LANGUAGE OverloadedStrings #-}

module Storage (module Storage, module Types.Storage) where

import Types
import Types.Storage
import Types.Server
import Types.Cost
import Output
import Share
import Storage.Network
import Servers
import Tunables
import ByteStrings
import Data.Maybe
import Data.List
import Control.Monad
import Control.Concurrent.Thread.Delay
import Control.Concurrent.Async
import qualified Data.Set as S
import System.Random
import System.Random.Shuffle

networkStorageLocations :: Maybe LocalStorageDirectory -> StorageLocations
networkStorageLocations = StorageLocations . serverList

type UpdateProgress = IO ()

data StorageProblem
	= FatalProblem String
	| OverridableProblem String
	deriving (Show)

-- | Check if there is a problem with storing shares amoung the provided
-- storage locations, assuming that some random set of the storage
-- locations will be used.
--
-- It's always a problem to store anything on an Untrusted server.
--
-- It should not be possible to reconstruct the encrypted
-- secret key using only objects from Alternate servers, so
-- fewer than neededObjects Alternate servers can be used.
problemStoringIn :: StorageLocations -> Tunables -> Maybe StorageProblem
problemStoringIn (StorageLocations locs) tunables
	| not (null (getlevel Untrusted)) || length locs < totalObjects ps = 
		Just $ FatalProblem
			"Not enough servers are available to store your encrypted secret key."
	| length alternates >= neededObjects ps = Just $ OverridableProblem $ unlines $
		[ "Not enough keysafe servers are available that can store"
		, "your encrypted secret key with a recommended level of"
		, "security."
		, "" 
		, "If you continue, some of the following less secure"
		, "servers will be used:"
		, ""
		] ++ map descserver (mapMaybe getServer alternates)
	| otherwise = Nothing
  where
	ps = shareParams tunables
	getlevel sl = filter (\s -> storageLevel s == sl) locs
	alternates = getlevel Alternate
	descserver (Server { serverName = ServerName n, serverDesc = d}) = 
		"* " ++ n ++ " -- " ++ d

-- | Stores the shares amoung the storage locations. Each location
-- gets at most one share from each set.
--
-- If a server is not currently accessible, it will be queued locally.
-- If any uploads are queued, returns True.
storeShares :: StorageLocations -> ShareIdents -> [S.Set Share] -> UpdateProgress -> IO (StoreResult, Bool, [Storage])
storeShares (StorageLocations locs) allsis shares updateprogress = do
	((r, anyqueued), usedlocs) <- go allsis shares [] False
	_ <- mapM_ obscureShares usedlocs
	return (r, anyqueued, usedlocs)
  where
	go sis (s:rest) usedlocs anyqueued = do
		let (is, sis') = nextShareIdents sis
		(r, usedlocs', queued) <- storeset locs [] Nothing (zip (S.toList is) (S.toList s)) False
		case r of
			StoreSuccess -> go sis' rest (usedlocs ++ usedlocs') (anyqueued || queued)
			_ -> return ((r, anyqueued || queued), usedlocs ++ usedlocs')
	go _ [] usedlocs anyqueued = return ((StoreSuccess, anyqueued), usedlocs)

	storeset _ usedlocs _ [] queued = return (StoreSuccess, usedlocs, queued)
	storeset [] usedlocs lasterr _ queued =
		return (fromMaybe (StoreFailure "no storage locations") lasterr, usedlocs, queued)
	storeset (loc:otherlocs) usedlocs _ ((i, s):rest) queued = do
		r <- storeShare loc i s
		case r of
			StoreSuccess -> do
				_ <- updateprogress
				storeset otherlocs (loc:usedlocs) Nothing rest queued
			-- Give up if any location complains a share
			-- already exists, because we have a name conflict.
			StoreAlreadyExists -> return (StoreAlreadyExists, usedlocs, queued)
			-- Queue or try storing it somewhere else on failure.
			StoreFailure _ -> case uploadQueue loc of
				Just q -> do
					r' <- storeShare q i s
					case r' of
						StoreSuccess -> do
							_ <- updateprogress
							storeset otherlocs (loc:usedlocs) Nothing rest True
						StoreAlreadyExists -> return (StoreAlreadyExists, usedlocs, queued)
						StoreFailure _ -> storeset otherlocs usedlocs (Just r) ((i, s):rest) queued
				Nothing -> storeset otherlocs usedlocs (Just r) ((i, s):rest) queued

-- | Retrieves one set of shares from the storage locations.
-- Returns all the shares it can find, which may not be enough,
-- and the remaining Shareidents, to use to get subsequent sets.
--
-- Assumes that each location only contains one share. So, once a
-- share has been found on a location, can avoid asking that location
-- for any other shares.
retrieveShares :: StorageLocations -> ShareIdents -> UpdateProgress -> IO (S.Set Share, ShareIdents, [Server])
retrieveShares (StorageLocations locs) sis updateprogress = do
	let (is, sis') = nextShareIdents sis
	let want = zip [1..] (S.toList is)
	(shares, usedlocs, _unusedlocs) <- go locs [] want []
	_ <- mapM_ obscureShares usedlocs
	let usedservers = mapMaybe getServer usedlocs
	return (S.fromList shares, sis', usedservers)
  where
	go unusedlocs usedlocs [] shares = return (shares, usedlocs, unusedlocs)
	go [] usedlocs _ shares = return (shares, usedlocs, [])
	go (loc:otherlocs) usedlocs ((n, i):rest) shares = do
		r <- retrieveShare loc n i
		case r of
			RetrieveSuccess s -> do
				_ <- updateprogress
				go otherlocs (loc:usedlocs) rest (s:shares)
			RetrieveFailure _ -> do
				-- Try to get the share from other locations.
				(shares', usedlocs', unusedlocs) <-
					go otherlocs usedlocs [(n, i)] shares
				-- May need to ask the location that didn't
				-- have the share for a later share, but
				-- ask it last. This way, the first
				-- location on the list can't deny having
				-- all shares and so learn the idents of
				-- all of them.
				go (unusedlocs++[loc]) usedlocs' rest shares'

-- | Returns descriptions of any failures.
tryUploadQueued :: Maybe LocalStorageDirectory -> IO [String]
tryUploadQueued d = do
	StorageLocations locs <- shuffleStorageLocations $
		networkStorageLocations d
	results <- forM locs $ \loc -> case uploadQueue loc of
		Nothing -> return []
		Just q -> moveShares q loc
	return $ processresults (concat results) []
  where
	processresults [] c = nub c
	processresults (StoreSuccess:rs) c = processresults rs c
	processresults (StoreFailure e:rs) c = processresults rs (e:c)
	processresults (StoreAlreadyExists:rs) c =
		processresults rs ("Unable to upload a share to a server due to a name conflict.":c)

storeChaff :: HostName -> Port -> Maybe Seconds -> IO ()
storeChaff hn port delayseconds = forever $ do
	say $ "Sending chaff to " ++ hn ++ " (press ctrl-c to stop)"
	say "Legend: + = successful upload, ! = upload failure"
	randomname <- randomByteStringOfLength 128
	-- It's ok the use the testModeTunables here because
	-- the randomname is not something that can be feasibly guessed.
	-- Prefix "random chaff" to the name to avoid ever using a name
	-- that a real user might want to use.
	let sis = shareIdents testModeTunables (Name $ "random chaff:" <> randomname) AnyGpgKey
	mapConcurrently (go sis)
		[1..totalObjects (shareParams testModeTunables)]
  where
	server = networkStorage Untrusted Nothing $ 
		Server (ServerName hn) [ServerAddress hn port] "chaff server"
	objsize = objectSize defaultTunables * shareOverhead defaultTunables
	maxmsdelay = ceiling $ 1000000 * fromMaybe 0 delayseconds
	go sis n = do
		msdelay <- getStdRandom (randomR (0, maxmsdelay))
		delay msdelay

		b <- randomByteStringOfLength objsize
		let share = Share 0 (StorableObject b)
		let (is, sis') = nextShareIdents sis
		let i = S.toList is !! (n - 1)
		r <- storeShare server i share
		case r of
			StoreSuccess -> progress "+"
			_ -> progress "!"
		go sis' n

-- | Shuffles the list, keeping Recommended first, then
-- Alternate, and finally Untrusted.
shuffleStorageLocations :: StorageLocations -> IO StorageLocations
shuffleStorageLocations (StorageLocations l) = 
	StorageLocations . concat <$> mapM shuf [minBound..maxBound]
  where
	shuf sl = shuffleM (filter (\s -> storageLevel s == sl) l)