{-# LANGUAGE LambdaCase, ScopedTypeVariables #-}
{-
    Demonstration of AES as a random number generator seeded by a string.

    Copyright 2015 Ken Takusagawa

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

-}

module Main(main) where {
import Crypto.Cipher.AES;
import qualified Data.ByteString as ByteString;
import Data.ByteString(ByteString);
import Control.Monad.Trans.State(State);
import qualified Control.Monad.Trans.State as State;
import Data.Byteable(toBytes);
import Data.Word(Word8);
import Data.ByteString.Base16 as Base16;
-- prefer Crypto.PBKDF.ByteString over Crypto.PBKDF because the String
-- version returns a hexadecimal string which would need to get
-- parsed to a ByteString.
import Crypto.PBKDF.ByteString(sha512PBKDF2);
import System.Environment(getArgs);
import qualified Crypto.Scrypt as Scrypt;
import Data.Maybe;
import qualified Crypto.Hash.SHA256 as SHA256;
import Data.Text.Encoding(encodeUtf8);
import qualified Data.Text as Text;

main :: IO();
main = getArgs >>= \case {
-- main demonstration: prints out an infinite stream of 8-bit numbers.
[] -> mapM_ print $ aes_stream $ makekey_scrypt mypassphrase;

-- check AES
["runtest1"] -> runtest1;

-- benchmarking kdfs
["pbkdf2"] -> byteStringPutStrLn $ Base16.encode $ makekey_pbkdf2 mypassphrase;
["scrypt"] -> byteStringPutStrLn $ Base16.encode $ makekey_scrypt mypassphrase;
["sha"] -> byteStringPutStrLn $ Base16.encode $  makekey_sha mypassphrase;

_ -> error "args";
};

-- world's most famous password
mypassphrase :: String;
mypassphrase = "ACollectionOfDiplomaticHistorySince_1966_ToThe_PresentDay#";

salt :: String;
salt = "mysalt98754321";

initial_aes_state :: AESIV;
initial_aes_state = aesIV_ $ ByteString.replicate 16 0;  -- set the counter to 0

next :: AES -> State AESIV ByteString;
next aes = do {
iv :: AESIV <- State.get;
let {(out,nextiv) = genCounter aes iv 1};
-- asking for 1 byte should return 1 block
State.put nextiv;
return out;
};

aes_stream :: ByteString -> [Word8];
aes_stream key = concat $ map ByteString.unpack $ State.evalState (sequence $ repeat $ next $ initAES key) $ initial_aes_state;

-- check that AES works
runtest1 :: IO ();
runtest1 = do {
let {(o,ns) = test1};
putStrLn $ toHex o;
print $ ByteString.unpack $ toBytes $ ns;
};

test1 :: (ByteString, AESIV);
test1 = State.runState (next $ initAES zerokey ) initial_aes_state;

zerokey :: ByteString;
zerokey = ByteString.replicate 16 0;  -- sha256

-- workaround for ByteString.putStrLn being deprecated but
-- ByteString.putStr being not.
byteStringPutStrLn :: ByteString -> IO();
byteStringPutStrLn b = do {
 ByteString.putStr b;
 putStrLn "";
};

-- AES256
keylength_bytes :: Integer;
keylength_bytes = div 256 8;

-- 1e6 is about 10 seconds on my machine
iterations_pbkdf2 :: Integer;
iterations_pbkdf2 = 1000000;

makekey_pbkdf2 :: String -> ByteString;
makekey_pbkdf2 s = sha512PBKDF2 (string2bs s) (string2bs salt) (fromIntegral iterations_pbkdf2) (fromIntegral keylength_bytes);

{-
30 indicates that the unit of measurement is 2^30 bytes = 1 GiB
-7 is to compensate for the 128 factor
5 = 5GiB memory
1 = repeat the algorithm 1 times
takes 20 seconds on my machine
-}
scrypt_params :: Scrypt.ScryptParams;
scrypt_params = fromJust $ Scrypt.scryptParamsLen (30-7) 5 1 keylength_bytes;

string2bs :: String -> ByteString;
-- this does not handle UTF8, etc.
-- string2bs = ByteString.pack . map (fromIntegral . fromEnum);
-- but this does
string2bs = encodeUtf8 . Text.pack;

toHex :: ByteString -> String;
toHex =  map (toEnum . fromEnum) . ByteString.unpack . Base16.encode;

makekey_scrypt :: String -> ByteString;
makekey_scrypt s = Scrypt.getHash $ Scrypt.scrypt scrypt_params (Scrypt.Salt $ string2bs salt) (Scrypt.Pass $ string2bs s);

-- This is NOT a good PBKDF!  Too easy to brute force.
makekey_sha :: String -> ByteString;
makekey_sha = SHA256.hash . string2bs;

} -- end