{-# LANGUAGE OverloadedStrings #-}

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

module Main where

import Types
import Tunables
import qualified CmdLine
import UI
import Encryption
import Entropy
import ExpensiveHash
import Cost
import Shard
import Storage
import Storage.LocalFiles
import qualified Gpg
import Data.Maybe
import Data.Time.Clock
import Data.Time.Calendar
import Control.Monad
import qualified Data.ByteString as B
import qualified Data.ByteString.UTF8 as BU8
import System.Posix.User (userGecos, getUserEntryForID, getEffectiveUserID)

main :: IO ()
main = do
	cmdline <- CmdLine.get
	ui <- selectUI (CmdLine.gui cmdline)
	let tunables = if CmdLine.testMode cmdline
		then testModeTunables
		else defaultTunables
	mode <- CmdLine.selectMode cmdline
	go mode (CmdLine.secretkeysource cmdline) tunables ui
  where
	go CmdLine.Backup (Just secretkeysource@(GpgKey kid)) tunables ui = do
		ok <- Gpg.knownByKeyServer kid
		unless ok $
			error "Your gpg public key has to be stored on the keyservers before you can back it up by keyid. Either use gpg --send-key to store the public key on the keyservers, or omit the --gpgkeyid option"
		backup ui tunables secretkeysource
			=<< getSecretKey secretkeysource
	go CmdLine.Backup (Just secretkeysource) tunables ui =
		backup ui tunables secretkeysource
			=<< getSecretKey secretkeysource
	go CmdLine.Backup Nothing tunables ui =
		backup ui tunables anyGpgKey =<< pickGpgKeyToBackup ui
	go CmdLine.Restore (Just secretkeydest) _ ui =
		restore ui secretkeydest
	go CmdLine.Restore Nothing _ ui =
		restore ui anyGpgKey
	go CmdLine.Benchmark _ tunables _ =
		benchmarkTunables tunables

getSecretKey :: SecretKeySource -> IO SecretKey
getSecretKey (GpgKey kid) = Gpg.getSecretKey kid
getSecretKey (KeyFile f) = SecretKey <$> B.readFile f

-- | Pick gpg secret key to back up.
--
-- If there is only one gpg secret key,
-- the choice is obvious. Otherwise prompt the user with a list.
pickGpgKeyToBackup :: UI -> IO SecretKey
pickGpgKeyToBackup ui = go =<< Gpg.listSecretKeys
  where
	go [] = do
		showError ui "You have no gpg secret keys to back up."
		error "Aborting on no gpg secret keys."
	go [(_, kid)] = Gpg.getSecretKey kid
	go l = maybe (error "Canceled") Gpg.getSecretKey
		=<< promptKeyId ui "Pick gpg secret key"
			"Pick gpg secret key to back up:" l

-- | Use when the gpg keyid will not be known at restore time.
anyGpgKey :: SecretKeySource
anyGpgKey = GpgKey (KeyId "")

backup :: UI -> Tunables -> SecretKeySource -> SecretKey -> IO ()
backup ui tunables secretkeysource secretkey = do
	username <- userName
	name <- fromMaybe (error "Aborting on no name") 
		<$> promptName ui "Enter name"
			namedesc username validateName
	kek <- promptkek name
	let esk = encrypt tunables kek secretkey
	let sis = shardIdents tunables name secretkeysource
	shards <- genShards esk tunables
	print =<< mapM (uncurry (storeShard localFiles)) (zip (getIdents sis) shards)
	print =<< obscureShards localFiles
  where
	promptkek name = do
		password <- fromMaybe (error "Aborting on no password") 
			<$> promptPassword ui True "Enter password" passworddesc
		kek <- genKeyEncryptionKey tunables name password
		username <- userName
		let badwords = concatMap namewords [name, username]
		let crackcost = estimateAttackCost spotAWS $
			estimateBruteforceOf kek $
				passwordEntropy password badwords
		let mincost = Dollars 100000
		if crackcost < mincost
			then do
				showError ui $ "Weak password! It would cost only " ++ show crackcost ++ " to crack the password. Please think of a better one. More words would be good.."
				promptkek name
			else do
				(thisyear, _, _) <- toGregorian . utctDay
					<$> getCurrentTime
				ok <- promptQuestion ui "Password strength estimate"
					(crackdesc crackcost thisyear)
					"Is your password strong enough?"
				if ok
					then return kek
					else promptkek name
	namewords (Name nb) = words (BU8.toString nb)
	namedesc = unlines
		[ "To back up your secret key, you will need to enter a name and a password."
		, ""
		, "Make sure to pick a name you will remember at some point in the future,"
		, "perhaps years from now, when you will need to enter it with the same"
		, "spelling and capitalization in order to restore your secret key."
		, ""
		, "(Your own full name is a pretty good choice for the name to enter here.)"
		]
	passworddesc = unlines
		[ "Pick a password that will be used to protect your secret key."
		, ""
		, "It's very important that this password be hard to guess."
		, ""
		, "And, it needs to be one that you will be able to remember years from now"
		, "in order to restore your secret key."
		]
	crackdesc crackcost thisyear = unlines $
		"Rough estimate of the cost to crack your password: " :
		costOverTimeTable crackcost thisyear

restore :: UI -> SecretKeySource -> IO ()
restore ui secretkeydest = do
	username <- userName
	name <- fromMaybe (error "Aborting on no name") 
		<$> promptName ui "Enter name"
			namedesc username validateName
	password <- fromMaybe (error "Aborting on no password") 
		<$> promptPassword ui True "Enter password" passworddesc
	
	let sis = shardIdents tunables name secretkeydest
	-- we drop 1 to simulate not getting all shards from the servers
	let l = drop 1 $ zip [1..] (getIdents sis)

	shards <- map (\(RetrieveSuccess s) -> s)
		<$> mapM (uncurry (retrieveShard localFiles)) l
	_ <- obscureShards localFiles

	let esk = combineShards tunables shards
	go esk (candidateKeyEncryptionKeys tunables name password)
  where
	go _ [] = error "decryption failed"
	go esk (kek:rest) = case decrypt kek esk of
		Just (SecretKey sk) -> print sk
		Nothing -> go esk rest

	-- TODO: derive by probing to find objects
	tunables = testModeTunables -- defaultTunables
	namedesc = unlines
		[ "When you backed up your secret key, you entered a name and a password."
		, "Now it's time to remember what you entered back then."
		, ""
		, "(If you can't remember the name you used, your own full name is the best guess.)"
		]
	passworddesc = unlines
		[ "Enter the password to unlock your secret key."
		]

validateName :: Name -> Maybe Problem
validateName (Name n)
	| B.length n < 2 = Just "The name should be at least 2 letters long."
	| otherwise = Nothing

userName :: IO Name
userName = do
	u <- getUserEntryForID =<< getEffectiveUserID
	return $ Name $ BU8.fromString $ takeWhile (/= ',') (userGecos u)