{-# 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 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)
	tunables <- 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 testModeTunables
		else return defaultTunables
	dispatch cmdline ui tunables

dispatch :: CmdLine.CmdLine -> UI -> Tunables -> IO ()
dispatch cmdline ui tunables = 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 secretkeydest
	go CmdLine.Backup Nothing =
		backup storage ui tunables Gpg.anyKey
			=<< Gpg.getKeyToBackup ui
	go CmdLine.Restore Nothing =
		restore storage ui Gpg.anyKey
	go CmdLine.Benchmark _ =
		benchmarkTunables tunables

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

backup :: Storage -> UI -> Tunables -> SecretKeySource -> SecretKey -> IO ()
backup storage 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 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)
	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
	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..."
		]

restore :: Storage -> UI -> SecretKeySource -> IO ()
restore storage 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
	shards <- catMaybes <$> downloadShards storage ui sis
	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 secretkey) -> do
						setpercent 100
						-- TODO save
						print secretkey
						showInfo ui "Success" "Your secret key successfully restored!"
  where
	-- 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."
		]
	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 -> ShardIdents -> IO [Maybe Shard]
downloadShards storage ui sis = bracket_ (return ()) cleanup
	(withProgress ui "Downloading encrypted data" message go)
  where
	go setpercent = do
		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]

		forM (zip percentsteps l) $ \(pct, (n, i)) -> do
			r <- retrieveShard storage n i
			_ <- setpercent pct
			case r of
				RetrieveSuccess s -> do
					return (Just s)
				RetrieveFailure f -> do
					hPutStrLn stderr $
						"warning: retrieval of shard " ++ show n ++ " failed: " ++ f
					return Nothing
	cleanup = obscureShards storage
	message = unlines
		[ "This will probably take around "
			++ showCostMinutes (getCreationCost sis)
		, ""
		, "(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)