{-# LANGUAGE FlexibleInstances #-}

-----------------------------------------------------------------------------

-----------------------------------------------------------------------------

-- |
-- Module      :  Crypto.Cipher.DES.Primitive
-- License     :  BSD-style
--
-- This module is copy of DES module from Crypto package.
-- http://hackage.haskell.org/package/Crypto
module Crypto.Cipher.DES.Primitive (
    encrypt,
    decrypt,
    Block (..),
) where

import Data.Bits
import Data.Word

-- | a DES block (64 bits)
newtype Block = Block {unBlock :: Word64}

type Rotation = Int
type Key = Word64

type Bits4 = [Bool]
type Bits6 = [Bool]
type Bits32 = [Bool]
type Bits48 = [Bool]
type Bits56 = [Bool]
type Bits64 = [Bool]

desXor :: [Bool] -> [Bool] -> [Bool]
desXor a b = zipWith (/=) a b

desRotate :: [Bool] -> Int -> [Bool]
desRotate bits rot = drop rot' bits ++ take rot' bits
  where
    rot' = rot `mod` length bits

bitify :: Word64 -> Bits64
bitify w = map (\b -> w .&. (shiftL 1 b) /= 0) [63, 62 .. 0]

unbitify :: Bits64 -> Word64
unbitify bs = foldl (\i b -> if b then 1 + shiftL i 1 else shiftL i 1) 0 bs

initial_permutation :: Bits64 -> Bits64
initial_permutation mb = map ((!!) mb) i
  where
    i =
        [ 57
        , 49
        , 41
        , 33
        , 25
        , 17
        , 9
        , 1
        , 59
        , 51
        , 43
        , 35
        , 27
        , 19
        , 11
        , 3
        , 61
        , 53
        , 45
        , 37
        , 29
        , 21
        , 13
        , 5
        , 63
        , 55
        , 47
        , 39
        , 31
        , 23
        , 15
        , 7
        , 56
        , 48
        , 40
        , 32
        , 24
        , 16
        , 8
        , 0
        , 58
        , 50
        , 42
        , 34
        , 26
        , 18
        , 10
        , 2
        , 60
        , 52
        , 44
        , 36
        , 28
        , 20
        , 12
        , 4
        , 62
        , 54
        , 46
        , 38
        , 30
        , 22
        , 14
        , 6
        ]

{-
"\x39\x31\x29\x21\x19\x11\x09\x01\x3b\x33\x2b\x23\x1b\x13\
\\x0b\x03\x3d\x35\x2d\x25\x1d\x15\x0d\x05\x3f\x37\x2f\x27\
\\x1f\x17\x0f\x07\x38\x30\x28\x20\x18\x10\x08\x00\x3a\x32\
\\x2a\x22\x1a\x12\x0a\x02\x3c\x34\x2c\x24\x1c\x14\x0c\x04\
\\x3e\x36\x2e\x26\x1e\x16\x0e\x06"
-}

key_transformation :: Bits64 -> Bits56
key_transformation kb = map ((!!) kb) i
  where
    i =
        [ 56
        , 48
        , 40
        , 32
        , 24
        , 16
        , 8
        , 0
        , 57
        , 49
        , 41
        , 33
        , 25
        , 17
        , 9
        , 1
        , 58
        , 50
        , 42
        , 34
        , 26
        , 18
        , 10
        , 2
        , 59
        , 51
        , 43
        , 35
        , 62
        , 54
        , 46
        , 38
        , 30
        , 22
        , 14
        , 6
        , 61
        , 53
        , 45
        , 37
        , 29
        , 21
        , 13
        , 5
        , 60
        , 52
        , 44
        , 36
        , 28
        , 20
        , 12
        , 4
        , 27
        , 19
        , 11
        , 3
        ]

{-
"\x38\x30\x28\x20\x18\x10\x08\x00\x39\x31\x29\x21\x19\x11\
\\x09\x01\x3a\x32\x2a\x22\x1a\x12\x0a\x02\x3b\x33\x2b\x23\
\\x3e\x36\x2e\x26\x1e\x16\x0e\x06\x3d\x35\x2d\x25\x1d\x15\
\\x0d\x05\x3c\x34\x2c\x24\x1c\x14\x0c\x04\x1b\x13\x0b\x03"
-}

des_enc :: Block -> Key -> Block
des_enc = do_des [1, 2, 4, 6, 8, 10, 12, 14, 15, 17, 19, 21, 23, 25, 27, 28]

des_dec :: Block -> Key -> Block
des_dec = do_des [28, 27, 25, 23, 21, 19, 17, 15, 14, 12, 10, 8, 6, 4, 2, 1]

do_des :: [Rotation] -> Block -> Key -> Block
do_des rots (Block m) k = Block $ des_work rots (takeDrop 32 mb) kb
  where
    kb = key_transformation $ bitify k
    mb = initial_permutation $ bitify m

des_work :: [Rotation] -> (Bits32, Bits32) -> Bits56 -> Word64
des_work [] (ml, mr) _ = unbitify $ final_perm $ (mr ++ ml)
des_work (r : rs) mb kb = des_work rs mb' kb
  where
    mb' = do_round r mb kb

do_round :: Rotation -> (Bits32, Bits32) -> Bits56 -> (Bits32, Bits32)
do_round r (ml, mr) kb = (mr, m')
  where
    kb' = get_key kb r
    comp_kb = compression_permutation kb'
    expa_mr = expansion_permutation mr
    res = comp_kb `desXor` expa_mr
    res' = drop 1 $ iterate (trans 6) ([], res)
    trans n (_, b) = (take n b, drop n b)
    res_s =
        concat $
            zipWith
                (\f (x, _) -> f x)
                [ s_box_1
                , s_box_2
                , s_box_3
                , s_box_4
                , s_box_5
                , s_box_6
                , s_box_7
                , s_box_8
                ]
                res'
    res_p = p_box res_s
    m' = res_p `desXor` ml

get_key :: Bits56 -> Rotation -> Bits56
get_key kb r = kb'
  where
    (kl, kr) = takeDrop 28 kb
    kb' = desRotate kl r ++ desRotate kr r

compression_permutation :: Bits56 -> Bits48
compression_permutation kb = map ((!!) kb) i
  where
    i =
        [ 13
        , 16
        , 10
        , 23
        , 0
        , 4
        , 2
        , 27
        , 14
        , 5
        , 20
        , 9
        , 22
        , 18
        , 11
        , 3
        , 25
        , 7
        , 15
        , 6
        , 26
        , 19
        , 12
        , 1
        , 40
        , 51
        , 30
        , 36
        , 46
        , 54
        , 29
        , 39
        , 50
        , 44
        , 32
        , 47
        , 43
        , 48
        , 38
        , 55
        , 33
        , 52
        , 45
        , 41
        , 49
        , 35
        , 28
        , 31
        ]

expansion_permutation :: Bits32 -> Bits48
expansion_permutation mb = map ((!!) mb) i
  where
    i =
        [ 31
        , 0
        , 1
        , 2
        , 3
        , 4
        , 3
        , 4
        , 5
        , 6
        , 7
        , 8
        , 7
        , 8
        , 9
        , 10
        , 11
        , 12
        , 11
        , 12
        , 13
        , 14
        , 15
        , 16
        , 15
        , 16
        , 17
        , 18
        , 19
        , 20
        , 19
        , 20
        , 21
        , 22
        , 23
        , 24
        , 23
        , 24
        , 25
        , 26
        , 27
        , 28
        , 27
        , 28
        , 29
        , 30
        , 31
        , 0
        ]

s_box :: [[Word8]] -> Bits6 -> Bits4
s_box s [a, b, c, d, e, f] = to_bool 4 $ (s !! row) !! col
  where
    row = sum $ zipWith numericise [a, f] [1, 0]
    col = sum $ zipWith numericise [b, c, d, e] [3, 2, 1, 0]
    numericise :: Bool -> Int -> Int
    numericise = (\x y -> if x then 2 ^ y else 0)

    to_bool :: Int -> Word8 -> [Bool]
    to_bool 0 _ = []
    to_bool n i = ((i .&. 8) == 8) : to_bool (n - 1) (shiftL i 1)
s_box _ _ = error "DES: internal error bits6 more than 6 elements"

s_box_1 :: Bits6 -> Bits4
s_box_1 = s_box i
  where
    i =
        [ [14, 4, 13, 1, 2, 15, 11, 8, 3, 10, 6, 12, 5, 9, 0, 7]
        , [0, 15, 7, 4, 14, 2, 13, 1, 10, 6, 12, 11, 9, 5, 3, 8]
        , [4, 1, 14, 8, 13, 6, 2, 11, 15, 12, 9, 7, 3, 10, 5, 0]
        , [15, 12, 8, 2, 4, 9, 1, 7, 5, 11, 3, 14, 10, 0, 6, 13]
        ]

s_box_2 :: Bits6 -> Bits4
s_box_2 = s_box i
  where
    i =
        [ [15, 1, 8, 14, 6, 11, 3, 4, 9, 7, 2, 13, 12, 0, 5, 10]
        , [3, 13, 4, 7, 15, 2, 8, 14, 12, 0, 1, 10, 6, 9, 11, 5]
        , [0, 14, 7, 11, 10, 4, 13, 1, 5, 8, 12, 6, 9, 3, 2, 15]
        , [13, 8, 10, 1, 3, 15, 4, 2, 11, 6, 7, 12, 0, 5, 14, 9]
        ]

s_box_3 :: Bits6 -> Bits4
s_box_3 = s_box i
  where
    i =
        [ [10, 0, 9, 14, 6, 3, 15, 5, 1, 13, 12, 7, 11, 4, 2, 8]
        , [13, 7, 0, 9, 3, 4, 6, 10, 2, 8, 5, 14, 12, 11, 15, 1]
        , [13, 6, 4, 9, 8, 15, 3, 0, 11, 1, 2, 12, 5, 10, 14, 7]
        , [1, 10, 13, 0, 6, 9, 8, 7, 4, 15, 14, 3, 11, 5, 2, 12]
        ]

s_box_4 :: Bits6 -> Bits4
s_box_4 = s_box i
  where
    i =
        [ [7, 13, 14, 3, 0, 6, 9, 10, 1, 2, 8, 5, 11, 12, 4, 15]
        , [13, 8, 11, 5, 6, 15, 0, 3, 4, 7, 2, 12, 1, 10, 14, 9]
        , [10, 6, 9, 0, 12, 11, 7, 13, 15, 1, 3, 14, 5, 2, 8, 4]
        , [3, 15, 0, 6, 10, 1, 13, 8, 9, 4, 5, 11, 12, 7, 2, 14]
        ]

s_box_5 :: Bits6 -> Bits4
s_box_5 = s_box i
  where
    i =
        [ [2, 12, 4, 1, 7, 10, 11, 6, 8, 5, 3, 15, 13, 0, 14, 9]
        , [14, 11, 2, 12, 4, 7, 13, 1, 5, 0, 15, 10, 3, 9, 8, 6]
        , [4, 2, 1, 11, 10, 13, 7, 8, 15, 9, 12, 5, 6, 3, 0, 14]
        , [11, 8, 12, 7, 1, 14, 2, 13, 6, 15, 0, 9, 10, 4, 5, 3]
        ]

s_box_6 :: Bits6 -> Bits4
s_box_6 = s_box i
  where
    i =
        [ [12, 1, 10, 15, 9, 2, 6, 8, 0, 13, 3, 4, 14, 7, 5, 11]
        , [10, 15, 4, 2, 7, 12, 9, 5, 6, 1, 13, 14, 0, 11, 3, 8]
        , [9, 14, 15, 5, 2, 8, 12, 3, 7, 0, 4, 10, 1, 13, 11, 6]
        , [4, 3, 2, 12, 9, 5, 15, 10, 11, 14, 1, 7, 6, 0, 8, 13]
        ]

s_box_7 :: Bits6 -> Bits4
s_box_7 = s_box i
  where
    i =
        [ [4, 11, 2, 14, 15, 0, 8, 13, 3, 12, 9, 7, 5, 10, 6, 1]
        , [13, 0, 11, 7, 4, 9, 1, 10, 14, 3, 5, 12, 2, 15, 8, 6]
        , [1, 4, 11, 13, 12, 3, 7, 14, 10, 15, 6, 8, 0, 5, 9, 2]
        , [6, 11, 13, 8, 1, 4, 10, 7, 9, 5, 0, 15, 14, 2, 3, 12]
        ]

s_box_8 :: Bits6 -> Bits4
s_box_8 = s_box i
  where
    i =
        [ [13, 2, 8, 4, 6, 15, 11, 1, 10, 9, 3, 14, 5, 0, 12, 7]
        , [1, 15, 13, 8, 10, 3, 7, 4, 12, 5, 6, 11, 0, 14, 9, 2]
        , [7, 11, 4, 1, 9, 12, 14, 2, 0, 6, 10, 13, 15, 3, 5, 8]
        , [2, 1, 14, 7, 4, 10, 8, 13, 15, 12, 9, 0, 3, 5, 6, 11]
        ]

p_box :: Bits32 -> Bits32
p_box kb = map ((!!) kb) i
  where
    i =
        [ 15
        , 6
        , 19
        , 20
        , 28
        , 11
        , 27
        , 16
        , 0
        , 14
        , 22
        , 25
        , 4
        , 17
        , 30
        , 9
        , 1
        , 7
        , 23
        , 13
        , 31
        , 26
        , 2
        , 8
        , 18
        , 12
        , 29
        , 5
        , 21
        , 10
        , 3
        , 24
        ]

final_perm :: Bits64 -> Bits64
final_perm kb = map ((!!) kb) i
  where
    i =
        [ 39
        , 7
        , 47
        , 15
        , 55
        , 23
        , 63
        , 31
        , 38
        , 6
        , 46
        , 14
        , 54
        , 22
        , 62
        , 30
        , 37
        , 5
        , 45
        , 13
        , 53
        , 21
        , 61
        , 29
        , 36
        , 4
        , 44
        , 12
        , 52
        , 20
        , 60
        , 28
        , 35
        , 3
        , 43
        , 11
        , 51
        , 19
        , 59
        , 27
        , 34
        , 2
        , 42
        , 10
        , 50
        , 18
        , 58
        , 26
        , 33
        , 1
        , 41
        , 9
        , 49
        , 17
        , 57
        , 25
        , 32
        , 0
        , 40
        , 8
        , 48
        , 16
        , 56
        , 24
        ]

takeDrop :: Int -> [a] -> ([a], [a])
takeDrop _ [] = ([], [])
takeDrop 0 xs = ([], xs)
takeDrop n (x : xs) = (x : ys, zs)
  where
    (ys, zs) = takeDrop (n - 1) xs

-- | Basic DES encryption which takes a key and a block of plaintext
-- and returns the encrypted block of ciphertext according to the standard.
encrypt :: Word64 -> Block -> Block
encrypt = flip des_enc

-- | Basic DES decryption which takes a key and a block of ciphertext and
-- returns the decrypted block of plaintext according to the standard.
decrypt :: Word64 -> Block -> Block
decrypt = flip des_dec