A haskell implementation of a pseudo-random number generator using Salsa20 algorithm

random.hs

#!/usr/bin/env runhaskell  -threaded --

{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}

import           Control.Concurrent  (forkIO)
import           Control.Exception   (IOException, catch)
import           Control.Monad       (forM, forM_, when)
import           Data.Array.IO
import           Data.Array.Storable
import           Data.Bits
import qualified Data.Map.Strict     as Map
import           Data.Word           (Word32, Word64, Word8)
import           Foreign.ForeignPtr  (mallocForeignPtrBytes, withForeignPtr)
import           Foreign.Ptr         (Ptr)
import           Numeric             (showHex)
import           System.Environment  (getArgs)
import           System.Exit         (ExitCode (ExitSuccess), exitWith)
import           System.IO           (BufferMode (..), Handle, IOMode (..),
                                      hClose, hGetBuf, hPutBuf, hPutStr,
                                      hPutStrLn, hSetBuffering, stderr, stdin,
                                      stdout, withBinaryFile)
import           System.Process

-- | Actions this program can do.
data Command = Help
             | Entropy Int
               -- ^Estimates entropy of given number of bytes from stdin
             | CheatMode
               -- ^Simply read from /dev/random to generate random stream...
             | Tcp String
               -- ^Read output from given interface put in prosmicuous mode using tcpdump
             | Salsa Bool
               -- ^Apply salsa20/8 stream cipher to stdin and send it to stdout in 512-bytes long blocks
               -- in verbose mode (dump generated blocks) or not
  deriving (Show, Read)

parseCommandLine :: [ String ] -> Command
parseCommandLine []                = Tcp "en0"
parseCommandLine ("-s":_)          = Salsa False
parseCommandLine ("--salsa":_)     = Salsa False
parseCommandLine ("-S":_)          = Salsa True
parseCommandLine ("-C":_)          = CheatMode
parseCommandLine ("--cheat":_)     = CheatMode
parseCommandLine ("-e":n:_)        = Entropy (read n)
parseCommandLine ("--entropy":n:_) = Entropy (read n)
parseCommandLine ("-h":_)          = Help
parseCommandLine ("--help":_)      = Help
parseCommandLine [iface]           = Tcp iface
parseCommandLine _                 = Help

go :: Command -> IO ()
go Help            = help
go CheatMode       = cheatMode
go (Entropy n)     = entropy n
go (Tcp iface)     = tcpdump iface
go (Salsa verbose) = streamCipher verbose stdin

type Key = [ Word32 ] -- should be 8 words long

-- | Initialise an array for use with `salsa` algorithm.
salsainit :: Key -> Word64 -> Word64 -> IO (StorableArray Int Word32)
salsainit key nonce position = do
  buf <- newArray (0,15) 0
  writeArray buf 0 0x61707865
  forM (zip [1..] (take 4 key)) $ \ (i, k) -> writeArray buf i k
  writeArray buf 5 0x3320646e
  writeArray buf 6 (high nonce)
  writeArray buf 7 (low nonce)
  writeArray buf 8 (high position)
  writeArray buf 9 (low position)
  writeArray buf 10 0x79622d32
  forM (zip [11 ..] (drop 4 key)) $ \ (i, k) -> writeArray buf i k
  writeArray buf 15 0x6b206574
  pure buf

high :: Word64 -> Word32
high w = fromIntegral (w `shiftR` 32)

low :: Word64 -> Word32
low w = fromIntegral w

-- | Applies some number of rounds of Salsa algorithm on an initial block.
--
-- See <Salsa20 https://cr.yp.to/snuffle/salsafamily-20071225.pdf> paper for more details on
-- how it works.
salsa :: (MArray a Word32 IO)
      => Int -> a Int Word32 -> IO (a Int Word32)
salsa 0     init = pure init
salsa count init = do
  buf <- copy init
  op buf  4  0 12  7
  op buf  8  4  0  9
  op buf 12  8  4 13
  op buf  0 12  8 18
  op buf  9  5  1  7
  op buf 13  9  5  9
  op buf  1 13  9 13
  op buf  5  1 13 18
  op buf 14 10  6  7
  op buf  2 14 10  9
  op buf  6  2 14 13
  op buf 10  6  2 18
  op buf  3 15 11  7
  op buf  7  3 15  9
  op buf 11  7  3 13
  op buf 15 11  7 18
  op buf  1  0  3  7
  op buf  2  1  0  9
  op buf  3  2  1 13
  op buf  0  3  2 18
  op buf  6  5  4  7
  op buf  7  6  5  9
  op buf  4  7  6 13
  op buf  5  4  7 18
  op buf 11 10  9  7
  op buf  8 11 10  9
  op buf  9  8 11 13
  op buf 10  9  8 18
  op buf 12 15 14  7
  op buf 13 12 15  9
  op buf 14 13 12 13
  op buf 15 14 13 18
  update (+) init buf
  salsa (count - 1) buf
    where
      copy = mapArray id

      op :: (MArray a Word32 IO)
         => a Int Word32 -> Int -> Int -> Int -> Int -> IO ()
      op array x y z t = do
        [v_x, v_y, v_z] <- forM [x,y,z] (readArray array)
        writeArray array x (v_x `xor` (v_y + v_z `rotateL` t))

-- | Convert a 16 elements array of `Word32` to a 64 bytes array.
toBytesArray :: (MArray a Word32 IO) => a Int Word32 -> IO  (IOUArray Int Word8)
toBytesArray array = do
  bytes <- newArray (0,63) (0 :: Word8) :: IO (IOUArray Int Word8)
  forM_ [0..15] ( \ i -> do
                    word <- readArray array i
                    forM_ [0..3] $ \ j -> do
                      let byte = (word `rotate` (j * 8)) .&. 0x000000ff
                      writeArray bytes (i * 4 + j) (fromIntegral byte)
                )
  return bytes

-- | Apply a binary operation index-wise to 2 mutable arrays.
--
-- The second array contains the result of `op`.
update :: (MArray a Word32 IO)
     => (Word32 -> Word32 -> Word32) -> a Int Word32 -> a Int Word32 -> IO ()
update op a1 a2 = do
  (l1,u1) <- getBounds a1
  (l2,u2) <- getBounds a2
  forM_ [(max l1 l2) .. (min u1 u2)]
    (\ i -> do
        x <- readArray a1 i
        y <- readArray a2 i
        writeArray a2 i (x `op` y))

-- | Encode stdin using <Salsa20/12 https://cr.yp.to/snuffle/salsafamily-20071225.pdf> stream cipher.
-- The first 8 bytes of the input stream are used as the initial `key` for generating blocks of Salsa20.
-- Nonce and starting position are fixed to some arbitrary value, then input is read in blocks of 64 bytes,
-- a new Salsa20/12 block is generated, input is xor-ed with stream cipher block, until EOF is reached.
--
-- When `verbose` is true, the generated block is ouput to `stderr` in hexadecimal form.
streamCipher :: Bool -> Handle -> IO ()
streamCipher verbose h = do
  key <- readKey h
  let nonce = 0x123456789012345
      startBlockNum = 0x234567890123451

      dumpBytes array = getElems array >>= hPutStr stderr . foldr showHex ""

      go key pos buf = do
        numBytesRead <- withStorableArray buf $ \ p -> hGetBuf h p 64
        cipherBlock  <- salsainit key nonce pos >>= salsa 12
        nextKey <- take 8 <$> getElems cipherBlock
        update xor buf cipherBlock
        bytes <- toBytesArray cipherBlock
        when verbose $ dumpBytes bytes
        hPutArray stdout bytes numBytesRead `catch` \ (e :: IOException) -> hPutStrLn stderr (show e) >> exitWith ExitSuccess
        go nextKey (pos + 1) buf

  buf <- newArray (0,15) (0 :: Word32)
  go key startBlockNum buf

  where
      readKey h = do
        buf <- newArray (0,7) (0 :: Word32)
        withStorableArray buf $ \ p -> hGetBuf h p 256
        getElems buf

-- | Run `tcpdump` on en0 interface in promiscuous mode to generate entropy
tcpdump :: String -> IO ()
tcpdump iface = do
  (_, Just hout, Just herr, process) <-
      createProcess (proc "tcpdump" ["-KnOSx", "-vvv",  "-i", iface ]){ std_out = CreatePipe, std_err = CreatePipe }
  forkIO $ devNull herr
  ptr <- mallocForeignPtrBytes 1024
  withForeignPtr ptr $ flip readAndDump hout

-- | Compute Shannon entropy of an input stream
shannon :: [ Word8 ] -> Double
shannon stream = shannon' stream Map.empty
  where
    upsert Nothing  = Just 1
    upsert (Just n) = Just (n + 1)

    shannon' [] map = let count = Map.foldr (+) 0 map
                          freqs = Map.map ((/ fromIntegral count) . fromIntegral) map
                      in - (Map.foldr (\ p acc -> acc + p * logBase 2 p) 0 freqs)
    shannon' (x:xs) map = shannon' xs (Map.alter upsert x map)

-- | Compute a rough approximate of the entropy of some number of bytes of stdin using Shannon's entropy
-- e.g. <shannon formula http://www.bearcave.com/misl/misl_tech/wavelets/compression/shannon.html>
entropy :: Int -> IO ()
entropy numbytes | numbytes < 256 = hPutStrLn stderr $ "Need more than 256 bytes to compute entropy"
                 | otherwise = do
  buf <- newArray (0,256) (0 :: Word8)
  go buf numbytes []
    where
      go buf 0   acc = let s = shannon acc
                       in  putStrLn $ "Shannon entropy: " ++ show s
      go buf num acc = do
        numread <- withStorableArray buf $ \ p -> hGetBuf stdin p 256
        bytes <- take numread <$> getElems buf
        go buf (if numread < 256 || numread < num
                 then 0
                 else num - numread) (acc ++ bytes)


-- | Simulate reading from /dev/random by reading from /dev/random
cheatMode :: IO ()
cheatMode = do
  ptr <- mallocForeignPtrBytes 1024

  withForeignPtr ptr
    (\ p -> withBinaryFile "/dev/random" ReadMode $ readAndDump p)

-- * Helper Functions

-- | Reads data into a buffer from some `Handle` and output it to stdout.
-- Reads at most 1024 bytes so the input buffer should be able to handle that
-- much number of bytes...
--
-- Exits in case an `IOException` is raised while trying to write to `stdout` which probably
-- means pipe is broken and we must stop working.
readAndDump :: Ptr Word8 -> Handle -> IO ()
readAndDump ptr h = do
  numBytesRead <- hGetBuf h ptr 1024
  hPutBuf stdout ptr numBytesRead `catch` \ (_ :: IOException) -> exitWith ExitSuccess
  readAndDump ptr h

-- | Reads data from `Handle` and discards it
devNull :: Handle -> IO ()
devNull h = do
  ptr <- mallocForeignPtrBytes 1024
  withForeignPtr ptr $ go h
    where
      go h p = hGetBuf h p 1024 >> go h p

help :: IO ()
help = putStrLn $ unlines [ "random - A quick and hopefully not too dirty pseudo-random number generator"
                          , "Copyright 2017 - Arnaud Bailly"
                          , "Usage: random [options] iface?"
                          , "By default, uses tcpdump in promiscuous mode on interface en0 to gather 'random' data. Then flags can"
                          , "be used to encode this stream and get pseudorandom bytes. If given a non-option argument, this is used"
                          , "as the name of the interface to read from instead of en0. On linux systems, eth0 could be a good candidate."
                          , "Example:"
                          , "$ ./random.hs | ./random.hs -s | hexdump -bc"
                          , ""
                          , "Options: "
                          , "  -h|--help            : print this help message"
                          , "  -C|--cheat           : cheat mode, output what's read from /dev/random"
                          , "  -s|--salsa           : stream cipher mode. Uses Salsa20/12 to encore standard input to standard output"
                          , "  -S                   : verbose stream cipher mode. Outputs computed blocks in hexadecimal"
                          , "  -e|--entropy <numbytes> : approximate measure of 'entropy' of given number of bytes from stdin. "
                          , "      This is used as a test to assess the randomness of a given string of bytes and simply tries"
                          , "      to compress the stream, returning the compression ratio as an estimate of entropy."
                          ]


main :: IO ()
main = do
  hSetBuffering stdout NoBuffering
  hSetBuffering stderr NoBuffering
  parseCommandLine <$> getArgs >>= go