Haskell Quiz/The Solitaire Cipher/Solution Mike McClurg

module Main where

import Char (chr, ord, toUpper)
import Data.List (delete, elemIndices, findIndices)
import System.Environment (getArgs, getProgName)

-- Driver for the program
main :: IO ()
main = getArgs >>= parse

parse :: [String] -> IO ()
parse (key:[msg])      = putStrLn $ encodeWithKey key msg
parse ("-d":key:[msg]) = putStrLn $ decodeWithKey key msg
parse ["-h"]           = usage
parse _                = usage

usage :: IO ()
usage = do prog <- getProgName
           putStrLn $ "usage: " ++ prog ++ " [-d] <keyphrase> <message>"

-- Define the deck we'll be using
type Deck = [Card]

data Card = Card !Int
          | JokerA
          | JokerB
            deriving (Eq, Show)


-- Standard deck, in bridge order (0==Ace Spades, 51==King Clubs), jokers at end
mkStdDeck :: Deck
mkStdDeck = map Card [0..51] ++ [JokerA, JokerB]

-- Encode and decode
encode :: Deck -> String -> String
encode d msg = formatOutput $ map (uncurry add) $ zip (formatInput msg) $ keystream d

decode :: Deck -> String -> String
decode d cph = formatOutput $ map (uncurry add) $ zip (formatInput cph) $ map (0-) $ keystream d

encodeWithKey :: String -> String -> String
encodeWithKey key msg = encode (keyDeck key mkStdDeck) msg

decodeWithKey :: String -> String -> String
decodeWithKey key msg = decode (keyDeck key mkStdDeck) msg

-- Initialize the deck with the given key
keyDeck :: String -> Deck -> Deck
keyDeck []     = id
keyDeck (k:ks) = keyDeck ks . countCut' (charVal k) . countCut . tripleCut . moveJokers
    where charVal c = (ord $ toUpper c) - (ord 'A') + 1

-- Generate an infinite keystream
keystream :: Deck -> [Int]
keystream d = let d' = step d in
              case getOutput d' of
                Nothing -> keystream d' -- skip jokers
                Just i  -> i : keystream d'
    where step = countCut . tripleCut . moveJokers

-- Plaintext must have no spaces, and be padded to multiple of five
formatInput :: String -> String
formatInput s = pad $ map toUpper $ filter (/=' ') s
    where pad s = if (length s `mod` 5 == 0)
                  then s
                  else pad $ s ++ "X"

formatOutput :: String -> String
formatOutput [] = []
formatOutput cs = (take 5 cs) ++ " " ++ (formatOutput $ drop 5 cs)

-- Move jokers
moveJokers :: Deck -> Deck
moveJokers = moveJoker JokerB 2 . moveJoker JokerA 1

moveJoker :: Card -> Int -> Deck -> Deck
moveJoker j n d = a ++ [j] ++ b
    where n' = let i = (elemIndices j d)!!0 in
               if (i + n == 54) -- special case where joker would end up as top card; need to move under top one or two
                  then (i + n + 1) `mod` 54
                  else (i + n) `mod` 54
          (a,b) = splitAt n' $ delete j d

-- Triple cut: swap cards above first joker with cards below second joker
tripleCut :: Deck -> Deck
tripleCut d = let a = take j1 d                      -- first 'third'
                  b = take (j2 - j1 + 1) $ drop j1 d -- second 'third' (drop first third, then take up to next joker)
                  c = drop (j2 + 1) d in             -- third 'third'
              c ++ b ++ a
    where is = findIndices (\e -> (e==JokerA) || (e==JokerB)) d
          j1 = is!!0
          j2 = is!!1

-- Count cut: cut deck at n cards, where n is value of last card, leave last card in place
countCut :: Deck -> Deck
countCut d = countCut' i d
    where i = (cardVal $ last d)

countCut' :: Int -> Deck -> Deck
countCut' i d = (drop i d') ++ (take i d') ++ [l]
    where d' = take 53 d
          l  = last d

-- Return value of output card, or Nothing if joker
getOutput :: Deck -> Maybe Int
getOutput []     = Nothing
getOutput (c:cs) = let i = (cardVal c) in
                   case (c:cs)!!i of
                     JokerA   -> Nothing
                     JokerB   -> Nothing
                     (Card a) -> Just $ a+1

-- Int value of Card
cardVal :: Card -> Int
cardVal JokerA = 53
cardVal JokerB = 53
cardVal (Card c) = c + 1

-- Add Chars and Ints, modulo 26
add :: Char -> Int -> Char
add c i = intToChar $ i + charToInt c
    where charToInt c = (ord $ toUpper c) - ord 'A'
          intToChar i = chr $ i `mod` 26 + ord 'A'