Data commit

Task/One-time-pad/Haskell/one-time-pad.hs

@@ -0,0 +1,172 @@
+-- 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) .) . (-))