{-# 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 SecretKey
import Shard
import Storage
import qualified Gpg
import Data.Maybe
import Data.Time.Clock
import Data.Time.Calendar
import Data.Monoid
import Control.Monad
import Control.DeepSeq
import Control.Exception
import System.IO
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 mkt = CmdLine.customizeShardParams cmdline
	(tunables, possibletunables) <- if CmdLine.testMode cmdline
		then do
			showInfo ui "Test mode"
				"Keysafe is running in test mode. This is not secure, and should not be used with real secret keys!"
			return (mkt testModeTunables, [mkt testModeTunables])
		else return (mkt defaultTunables, map (mkt . snd) knownTunings)
	dispatch cmdline ui tunables possibletunables

dispatch :: CmdLine.CmdLine -> UI -> Tunables -> [Tunables] -> IO ()
dispatch cmdline ui tunables possibletunables = do
	mode <- CmdLine.selectMode cmdline
	go mode (CmdLine.secretkeysource cmdline)
  where
	storage = CmdLine.storage cmdline
	go CmdLine.Backup (Just secretkeysource) =
		backup storage ui tunables secretkeysource
			=<< getSecretKey secretkeysource
	go CmdLine.Restore (Just secretkeydest) =
		restore storage ui possibletunables secretkeydest
	go CmdLine.Backup Nothing =
		backup storage ui tunables Gpg.anyKey
			=<< Gpg.getKeyToBackup ui
	go CmdLine.Restore Nothing =
		restore storage ui possibletunables Gpg.anyKey
	go CmdLine.Benchmark _ =
		benchmarkTunables tunables

backup :: Storage -> UI -> Tunables -> SecretKeySource -> SecretKey -> IO ()
backup storage ui tunables secretkeysource secretkey = do
	username <- userName
	Name theirname <- fromMaybe (error "Aborting on no username") 
		<$> promptName ui "Enter your name"
			usernamedesc (Just username) validateName
	Name othername <- fromMaybe (error "aborting on no othername")
		<$> promptName ui "Enter other name"
			othernamedesc Nothing validateName
	let name = Name (theirname <> " " <> othername)
	kek <- promptkek name
	let sis = shardIdents tunables name secretkeysource
	let cost = getCreationCost kek <> getCreationCost sis
	r <- withProgress ui "Encrypting and storing data"
		(encryptdesc cost) $ \setpercent -> do
			let esk = encrypt tunables kek secretkey
			shards <- genShards esk tunables
			_ <- esk `deepseq` setpercent 25
			_ <- sis `deepseq` setpercent 50
			let step = 50 `div` length shards
			let percentsteps = map setpercent [50+step, 50+step*2..100]
			storeShards storage sis (zip percentsteps shards)
	case r of
		StoreSuccess -> showInfo ui "Success" "Your secret key successfully encrypted and backed up."
		StoreFailure s -> showError ui ("There was a problem storing your encrypted secret key: " ++ s)
		StoreAlreadyExists -> do
			showError ui $ unlines
				[ "Another secret key is already being stored under the name you entered."
				, "Please try again with a different name."
				]
			backup storage ui tunables secretkeysource secretkey
  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)
	keydesc = case secretkeysource of
		GpgKey _ -> "gpg secret key"
		KeyFile _ -> "secret key"
	usernamedesc = unlines
		[ "Keysafe is going to backup your " ++ keydesc ++ " securely."
		, ""
		, "You will be prompted for some information. To restore your " ++ keydesc
		, "at a later date, you will need to remember and enter the same information."
		, ""
		, "To get started, what is your name?"
		]
	othernamedesc = unlines
		[ "Now think of another name, which not many people know."
		, "This will be used to make it hard for anyone else to find"
		, "the backup of your " ++ keydesc ++ "."
		, ""
		, "Some suggestions:"
		, otherNameSuggestions
		, ""
		, "Make sure to pick a name you will remember later."
		]
	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
	encryptdesc cost = unlines
		[ "This will probably take around " ++ showCostMinutes cost
		, ""
		, "(It's a feature that this takes a while; it makes it hard"
		, "for anyone to find your data, or crack your password.)"
		, ""
		, "Please wait..."
		]

otherNameSuggestions :: String
otherNameSuggestions = unlines $ map ("  * " ++)
	[ "Your high-school sweetheart."
	, "Your first pet."
	, "Your favorite teacher."
	, "Your college roomate."
	]

restore :: Storage -> UI -> [Tunables] -> SecretKeySource -> IO ()
restore storage ui possibletunables secretkeydest = do
	username <- userName
	Name theirname <- fromMaybe (error "Aborting on no username") 
		<$> promptName ui "Enter your name"
			namedesc (Just username) validateName
	Name othername <- fromMaybe (error "aborting on no othername")
		<$> promptName ui "Enter other name"
			othernamedesc Nothing validateName
	let name = Name (theirname <> " " <> othername)
	password <- fromMaybe (error "Aborting on no password") 
		<$> promptPassword ui True "Enter password" passworddesc
	
	let mksis tunables = shardIdents tunables name secretkeydest
	(tunables, shards) <- downloadShards storage ui mksis possibletunables

	let candidatekeys = candidateKeyEncryptionKeys tunables name password
	let cost = getCreationCost candidatekeys 
		<> castCost (getDecryptionCost candidatekeys)
	case combineShards tunables shards of
		Left e -> showError ui e
		Right esk -> withProgress ui "Decrypting"
			(decryptdesc cost) $ \setpercent -> do
				case decrypt candidatekeys esk of
					Nothing -> showError ui "Decryption failed! Unknown why it would fail at this point."
					Just secretkey -> do
						setpercent 100
						writeSecretKey secretkeydest secretkey
						showInfo ui "Success" "Your secret key successfully restored!"
  where
	namedesc = unlines
		[ "When you backed up your secret key, you entered some information."
		, "To restore it, you'll need to remember what you entered back then."
		, ""
		, "To get started, what is your name?"
		]
	othernamedesc = unlines
		[ "What other name did you enter when you backed up your secret key?"
		, ""
		, "Back then, you were given some suggestions, like these:"
		, otherNameSuggestions
		]
	passworddesc = unlines
		[ "Enter the password to unlock your secret key."
		]
	decryptdesc cost = unlines
		[ "This will probably take around " ++ showCostMinutes cost
		, ""
		, "(It's a feature that this takes so long;"
		, "it prevents cracking your password.)"
		, ""
		, "Please wait..."
		]

downloadShards :: Storage -> UI -> (Tunables -> ShardIdents) -> [Tunables] -> IO (Tunables, [Shard])
downloadShards storage ui mksis possibletunables =
	bracket_ setup cleanup download
  where
	setup = return ()
	download = withProgress ui "Downloading encrypted data" message (go possibletunables)
	cleanup = obscureShards storage

	go [] _ = return (defaultTunables, [])
	go (tunables:othertunables) setpercent = do
		let sis = mksis tunables
		let l = zip [1..] (getIdents sis)
		-- Just calculating the idents probably takes
		-- most of the time.
		_ <- l `deepseq` setpercent 50
		let step = 50 `div` length l
		let percentsteps = [50+step, 50+step*2..100]

		mshards <- forM (zip percentsteps l) $ \(pct, (n, i)) -> do
			r <- retrieveShard storage n i
			case r of
				RetrieveSuccess s -> do
					_ <- setpercent pct
					return (Just s)
				RetrieveFailure f -> do
					hPutStrLn stderr $
						"warning: retrieval of shard " ++ show n ++ " failed: " ++ f
					return Nothing
		let shards = catMaybes mshards
		if null shards
			then go othertunables setpercent
			else return (tunables, shards)

	possiblesis = map mksis possibletunables
	message = unlines
		[ "This will probably take around "
			++ showCostMinutes (mconcat $ map getCreationCost possiblesis)
		, ""
		, "(It's a feature that this takes a while; it makes it hard"
		, "for anyone else to find your data.)"
		, ""
		, "Please wait..."
		]

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)