{-# LANGUAGE TemplateHaskell #-}
-- | Main module generating passwords.
module Main where

import           System.IO       (hFlush, stdout)
import           System.Directory(getDirectoryContents)
import qualified Data.Text    as Text
import qualified Data.Text.IO as Text
import           Data.Text(Text)
import qualified Data.Set     as Set
import           Data.Char           (isAlpha, isPunctuation, isSymbol)
import           Data.Random.RVar
import           Data.Random.Sample
import           Data.Random.RVar.Enum
import           Data.Random.Distribution
import           Data.Random.Distribution.Uniform
import           Data.Random.Source.IO
import qualified Data.Vector  as V
import           Control.Applicative
import           Control.Monad       (replicateM, foldM)
import           HFlags

-- | Reads a dict format to get a list of unique words without any special
--   chars.
readDict ::  FilePath -> IO (V.Vector Text)
readDict filename = do
    input <- Text.readFile filename
    return $! V.fromList . Set.toList . Set.fromList . map stripTails . Text.lines $! input
  where
    stripTails = head . Text.split (not . isAlpha)

-- | Find all plausible dictionaries in a given directory
dictFiles dir = postprocess `fmap`
                  getDirectoryContents dir
  where
    postprocess = map ((dir ++ "/") ++) . filter (not . (=='.') . head)

-- | Default directory where to look for the word lists.


-- | Read a set of dictionaries and put the together.
readDicts filenames = do putStr $ "Reading " ++ show (length filenames) ++ " files"
                         result <- (V.fromList . Set.toList) `fmap` foldM action Set.empty filenames
                         putStrLn ""
                         return result
  where
    action currentSet filename = do newSet <- readDict filename
                                    putStr "."
                                    hFlush stdout
                                    return $! Set.fromList (V.toList newSet) `Set.union` currentSet

-- | Read all dictionaries from a given directory.
readDictDir dirname = dictFiles dirname >>= readDicts

-- | Filename for default dictionary (should be command line argument or default glob.)
defaultDictionary ::  FilePath
defaultDictionary = "/usr/share/dict/british-english"

-- | Take a random element of a vector.
randomElement :: V.Vector a -> RVar a
randomElement words = (words V.!) <$> uniform 0 (V.length words - 1)

-- | Pick a random password, given a words list, and a number of words it will contain.
randomPassword :: V.Vector Text -> Int -> RVar Text
randomPassword words numWords = do ws   <- replicateM numWords $ randomElement words
                                   seps <- replicateM numWords randomSeparator
                                   return $ Text.concat $ zipWith Text.append ws seps

-- | Estimate strength of random password with given inputs.
randomPasswordStrength words numWords = fromIntegral numWords * logBase 2 wordStrength
  where
    wordStrength = fromIntegral $ V.length words * (32 + 100)

-- * Random separators
-- | Randomly pick a word separator as a two-digit number, or a symbol
--   character.
randomSeparator ::  RVar Text
randomSeparator = do b <- uniform False True
                     if b then symbolSeparator
                          else numericSeparator

-- | Two-digit number as a separator 10^2 = 6.6 bits of entropy.
numericSeparator ::  RVar Text
numericSeparator = Text.pack <$> show <$> uniform 0 (99 :: Int)

-- | Conjunction of two unary predicates
(|||) ::  (t -> Bool) -> (t -> Bool) -> t -> Bool
(|||) f g x = f x || g x

-- | List of symbol and punctuation characters in ASCII
--   Should be 5 bits of entropy
symbolChars ::  V.Vector Char
symbolChars = V.fromList $ filter (isSymbol ||| isPunctuation) $ map toEnum [0..127]

-- | Text with random symbol character, 5 bits of entropy
symbolSeparator ::  RVar Text
symbolSeparator = Text.singleton <$> randomElement symbolChars

-- * Command-line flags
-- | Number of words per password
defineFlag "w:words" (4 :: Int) "Number of words for each password."

-- | Number of passwords
defineFlag "p:passwords" (10 :: Int) "Number of passwords to generate."

-- | Default word list directory.
defineFlag "d:directory" ("/usr/share/dict" :: FilePath) ("Default directory to search for dictionaries\n (works only if --wordlist options is NOT USED.)")

-- | Pick specific wordlist.
defineFlag "l:wordlist" ("" :: FilePath) "Select particular dictionary (filepath)."

-- | Read wordlist given by explict filepath, or search for all wordlists in a given directory.
selectWordList :: FilePath -> FilePath -> IO (V.Vector Text)
selectWordList ""       dir = readDictDir dir
selectWordList filename _   = readDict    filename

main = do $initHFlags "WordPass - dictionary-based password generator"
          dictWords <- selectWordList flags_directory flags_wordlist
          putStrLn  $ "Read " ++ show (V.length dictWords) ++ " words from dictionaries."
          putStr "Estimated password strength (bits): "
          print $ randomPasswordStrength dictWords flags_words
          replicateM flags_passwords $ do 
            rv <- sample $ randomPassword dictWords flags_words
            Text.putStrLn rv