content: -- To compile into an executable:
  -- ghc -main-is OneTimePad OneTimePad.hs
  -- To run:
  -- ./OneTimePad --help
  
  module OneTimePad (main) where
  
  import           Control.Monad
  import           Data.Char
  import           Data.Function         (on)
  import qualified Data.Text             as T
  import qualified Data.Text.IO          as TI
  import           Data.Time
  import           System.Console.GetOpt
  import           System.Environment
  import           System.Exit
  import           System.IO
  
  -- Command-line options parsing
  data Options = Options  { optCommand :: String
                          , optInput   :: IO T.Text
                          , optOutput  :: T.Text -> IO ()
                          , optPad     :: (IO T.Text, T.Text -> IO ())
                          , optLines   :: Int
                          }
  
  startOptions :: Options
  startOptions = Options  { optCommand    = "decrypt"
                          , optInput      = TI.getContents
                          , optOutput     = TI.putStr
                          , optPad        = (TI.getContents, TI.putStr)
                          , optLines      = 0
                          }
  
  options :: [ OptDescr (Options -> IO Options) ]
  options =
      [ Option "e" ["encrypt"]
          (NoArg
              (\opt -> return opt { optCommand = "encrypt" }))
          "Encrypt file"
      , Option "d" ["decrypt"]
          (NoArg
              (\opt -> return opt { optCommand = "decrypt" }))
          "Decrypt file (default)"
      , Option "g" ["generate"]
          (NoArg
              (\opt -> return opt { optCommand = "generate" }))
          "Generate a one-time pad"
      , Option "i" ["input"]
          (ReqArg
              (\arg opt -> return opt { optInput = TI.readFile arg })
              "FILE")
          "Input file (for decryption and encryption)"
      , Option "o" ["output"]
          (ReqArg
              (\arg opt -> return opt { optOutput = TI.writeFile arg })
              "FILE")
          "Output file (for generation, decryption, and encryption)"
      , Option "p" ["pad"]
          (ReqArg
              (\arg opt -> return opt { optPad = (TI.readFile arg,
                                                  TI.writeFile arg) })
              "FILE")
          "One-time pad to use (for decryption and encryption)"
      , Option "l" ["lines"]
          (ReqArg
              (\arg opt -> return opt { optLines = read arg :: Int })
              "LINES")
          "New one-time pad's length (in lines of 48 characters) (for generation)"
      , Option "V" ["version"]
          (NoArg
              (\_ -> do
                  hPutStrLn stderr "Version 0.01"
                  exitWith ExitSuccess))
          "Print version"
      , Option "h" ["help"]
          (NoArg
              (\_ -> do
                  prg <- getProgName
                  putStrLn "usage: OneTimePad [-h] [-V] [--lines LINES] [-i FILE] [-o FILE] [-p FILE] [--encrypt | --decrypt | --generate]"
                  hPutStrLn stderr (usageInfo prg options)
                  exitWith ExitSuccess))
          "Show this help message and exit"
      ]
  
  main :: IO ()
  main = do
    args <- getArgs
    let (actions, nonOptions, errors) = getOpt RequireOrder options args
    opts <- Prelude.foldl (>>=) (return startOptions) actions
    let Options { optCommand = command
                , optInput   = input
                , optOutput  = output
                , optPad     = (inPad, outPad)
                , optLines   = linecnt } = opts
  
    case command of
      "generate" -> generate linecnt output
      "encrypt"  -> do
        inputContents <- clean <$> input
        padContents <- inPad
        output $ format $ encrypt inputContents $ unformat $ T.concat
          $ dropWhile (\t -> T.head t == '-' || T.head t == '#')
          $ T.lines padContents
      "decrypt"  -> do
        inputContents <- unformat <$> input
        padContents <- inPad
        output $ decrypt inputContents $ unformat $ T.concat
          $ dropWhile (\t -> T.head t == '-' || T.head t == '#')
          $ T.lines padContents
        let discardLines = ceiling
              $ ((/) `on` fromIntegral) (T.length inputContents) 48
        outPad $ discard discardLines $ T.lines padContents
  
  {- | Discard used pad lines. Is only called at decryption to enable using the
  same pad file for both encryption and decryption.
  -}
  discard :: Int -> [T.Text] -> T.Text
  discard 0 ts = T.unlines ts
  discard x (t:ts) = if (T.head t == '-' || T.head t == '#')
    then T.unlines [t, (discard x ts)]
    else T.unlines [(T.append (T.pack "- ") t), (discard (x-1) ts)]
  
  {- | Clean the text from symbols that cannot be encrypted.
  -}
  clean :: T.Text -> T.Text
  clean = T.map toUpper . T.filter (\c -> let oc = ord c
                                     in oc >= 65 && oc <= 122
                                     && (not $ oc >=91 && oc <= 96))
  
  {- | Format text (usually encrypted text) for pretty-printing it in a similar
  way to the example from Wikipedia (see Rosetta Code page for this task)
  -}
  format :: T.Text -> T.Text
  format = T.unlines . map (T.intercalate (T.pack " ") . T.chunksOf 6)
    . T.chunksOf 48
  
  {- | Unformat encrypted text, getting rid of characters that are irrelevant for
  decryption.
  -}
  unformat :: T.Text -> T.Text
  unformat = T.filter (\c -> c/='\n' && c/=' ')
  
  {- | Generate a one-time pad and write it to file (specified as second
  parameter). Note: this only works on operating systems that have the
  "/dev/random" file.
  -}
  generate :: Int -> (T.Text -> IO ()) -> IO ()
  generate lines output = do
    withBinaryFile "/dev/random" ReadMode
      (\handle -> do
          contents <- replicateM (48 * lines) $ hGetChar handle
          time <- getCurrentTime
          output
            $ T.unlines [ T.pack
                          $ "# OTP pad, generated by https://github.com/kssytsrk/one-time-pad on "
                          ++ show time
                        , format $ T.pack
                          $ map (chr . (65 +) . flip mod 26 . ord) contents
                        ])
  
  -- Helper function for encryption/decryption.
  crypt :: (Int -> Int -> Int) -> T.Text -> T.Text -> T.Text
  crypt f = T.zipWith ((chr .) . f `on` ord)
  
  -- Encrypt first parameter's contents, using the second parameter as a key.
  encrypt :: T.Text -> T.Text -> T.Text
  encrypt = crypt ((((+65) . flip mod 26 . subtract 130) .) . (+))
  
  -- Decrypt first parameter's contents, using the second parameter as a key.
  decrypt :: T.Text -> T.Text -> T.Text
  decrypt = crypt ((((+65) . flip mod 26) .) . (-))
  