{-# LANGUAGE OverloadedStrings, MultiParamTypeClasses #-}

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

module Share where

import Types
import Tunables
import ExpensiveHash
import Cost
import qualified Crypto.SecretSharing.Internal as SS
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Raaz.Core.Encode as Raaz
import qualified Raaz.Hash.Sha256 as Raaz
import qualified Data.Text as T
import qualified Data.Text.Encoding as E
import qualified Data.Set as S
import Data.Word

data ShareIdents = ShareIdents
	{ identsStream :: [S.Set StorableObjectIdent]
	-- ^ Each item in the infinite list is the idents to
	-- use for the shares of a chunk of data.
	, identsCreationCost :: Cost CreationOp
	, identsBruteForceCalc :: CostCalc BruteForceOp UnknownName
	}

nextShareIdents :: ShareIdents -> (S.Set StorableObjectIdent, ShareIdents)
nextShareIdents sis = 
	let (s:rest) = identsStream sis
	in (s, sis { identsStream = rest })

instance HasCreationCost ShareIdents where
	getCreationCost = identsCreationCost

instance Bruteforceable ShareIdents UnknownName where
	getBruteCostCalc = identsBruteForceCalc

data Distinguisher 
	= Distinguisher SecretKeySource
	| AnyGpgKey
	-- ^ Use to avoid the gpg keyid needing to be provided 
	-- at restore time.
	deriving (Eq)

-- | Generates identifiers to use for storing shares.
--
-- This is an expensive operation, to make it difficult for an attacker
-- to brute force known/guessed names and find matching shares.
-- The keyid or filename is used as a salt, to avoid collisions
-- when the same name is chosen for multiple keys.
shareIdents :: Tunables -> Name -> Distinguisher -> ShareIdents
shareIdents tunables (Name name) shareident =
	ShareIdents (segmentbyshare idents) creationcost bruteforcecalc
  where
	(ExpensiveHash creationcost basename) =
		expensiveHash hashtunables salt name
	salt = case shareident of
		Distinguisher sks -> Salt sks
		AnyGpgKey -> Salt (GpgKey (KeyId ""))
	mk n = StorableObjectIdent $ Raaz.toByteString $ mksha $
		E.encodeUtf8 $ basename <> T.pack (show n)
	mksha :: B.ByteString -> Raaz.Base16
	mksha = Raaz.encode . Raaz.sha256
	bruteforcecalc = bruteForceLinearSearch creationcost
	hashtunables = nameGenerationHash $ nameGenerationTunable tunables
	idents = map mk ([1..] :: [Integer])
	m = totalObjects (shareParams tunables)
	segmentbyshare l = 
		let (shareis, l') = splitAt m l
		in S.fromList shareis : segmentbyshare l'

-- | Generates shares of an EncryptedSecretKey.
-- Each chunk of the key creates its own set of shares.
genShares :: EncryptedSecretKey -> Tunables -> IO [S.Set Share]
genShares (EncryptedSecretKey cs _) tunables = do 
	shares <- mapM encode cs
	return $ map (S.fromList . map (uncurry Share) . zip [1..]) shares
  where
	encode :: B.ByteString -> IO [StorableObject]
	encode b = map (StorableObject . encodeShare) 
		<$> SS.encode
			(neededObjects $ shareParams tunables)
			(totalObjects $ shareParams tunables)
			(BL.fromStrict b)

-- | If not enough sets of shares are provided, the EncryptedSecretKey may
-- be incomplete, only containing some chunks of the key
combineShares :: Tunables -> [S.Set Share] -> Either String EncryptedSecretKey
combineShares tunables shares
	| null shares || any null shares || any (\l -> length l < sharesneeded) shares = 
		Left $ "Not enough shares are currently available to reconstruct your data. " ++
			concatMap (\l -> "(Got " ++ show (length l) ++ "/" ++ show sharesneeded ++ ") ") shares
	| otherwise = Right $ mk $
		map (BL.toStrict . SS.decode . map decodeshare . S.toList) shares
  where
	mk cs = EncryptedSecretKey cs unknownCostCalc
	decodeshare (Share sharenum so) = decodeShare sharenum sharesneeded $
		fromStorableObject so
	sharesneeded = neededObjects (shareParams tunables)

-- Note that this does not include the share number in the encoded
-- bytestring. This prevents an attacker from partitioning their shares
-- by share number.
encodeShare :: SS.Share -> B.ByteString
encodeShare = B.pack . concatMap (encodeShare' . SS.shareValue) . SS.theShare

decodeShare :: Int -> Int -> B.ByteString -> SS.Share
decodeShare sharenum sharesneeded = SS.Share . map mk . decodeShare' . B.unpack
  where
	mk v = SS.ByteShare
		{ SS.shareId = sharenum
		, SS.reconstructionThreshold = sharesneeded
		, SS.shareValue = v
		}

-- | Each input byte generates a share in a finite field of size 1021,
-- so encode it as the product of two bytes. This is inneffient; if the
-- finite field was 255 then the encoded share would be the same size as
-- the input. But, the finite-field library used by secret-sharing does
-- not support a non-prime size.
encodeShare' :: Int -> [Word8]
encodeShare' v =
	let (q, r) = quotRem v 255
	in [fromIntegral q, fromIntegral r]

decodeShare' :: [Word8] -> [Int]
decodeShare' = go []
  where
	go c [] = reverse c
	go c (q:r:rest) = go (((255 * fromIntegral q) + fromIntegral r):c) rest
	go _ _ = error "Badly encoded share has odd number of bytes"