||| [specification](https://csrc.nist.gov/csrc/media/publications/fips/46/3/archive/1999-10-25/documents/fips46-3.pdf)
module Data.Crypto.Encryption.Block.DEA

-- import Control.Isomorphism
import Data.Bits
import Data.Fin
import Data.Vect

import Data.Crypto.Encryption
import Data.Crypto.Encryption.Block
import Data.Crypto.Util

%default total
%access private

-- utility functions

trim : Bits (1 + n) -> Bits n
trim b = truncate (shiftRightLogical (intToBits 1) b)

-- Plenty of places in the 3DES spec use 1-based indexes, where we would like
-- 0-based indexes. So we embed the same numbers from the spec (for easy
-- eyeball-checking), then use this to correct the difference.
offByOne : Fin (S (S n)) -> Fin (S n)
offByOne FZ     = FZ
offByOne (FS y) = y

selectBits : Vect m (Fin n) -> Bits n -> Bits m
selectBits positions input = foldl (\acc, b => shiftLeft acc (intToBits 1) `and` b)
                                   (intToBits 0)
                                   (map (\x => intToBits (if getBit x input then 1 else 0)) positions)

||| The initial permutation
IP : Bits 64 -> Bits 64
IP = selectBits (map offByOne [58, 50, 42, 34, 26, 18, 10,  2,
                               60, 52, 44, 36, 28, 20, 12,  4,
                               62, 54, 46, 38, 30, 22, 14,  6,
                               64, 56, 48, 40, 32, 24, 16,  8,
                               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])

inverseIP : Bits 64 -> Bits 64
inverseIP = selectBits (map offByOne [40,  8, 48, 16, 56, 24, 64, 32,
                                      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])

-- TODO: prove this
-- permutation : Iso (Bits 64) (Bits 64)
-- permutation = MkIso IP IP'

E : Bits 32 -> Bits 48
E = selectBits (map offByOne [32,  1,  2,  3,  4,  5,
                               4,  5,  6,  7,  8,  9,
                               8,  9, 10, 11, 12, 13,
                              12, 13, 14, 15, 16, 17,
                              16, 17, 18, 19, 20, 21,
                              20, 21, 22, 23, 24, 25,
                              24, 25, 26, 27, 28, 29,
                              28, 29, 30, 31, 32,  1])

P : Bits 32 -> Bits 32
P = selectBits (map offByOne [16,  7, 20, 21,
                              29, 12, 28, 17,
                               1, 15, 23, 26,
                               5, 18, 31, 10,
                               2,  8, 24, 14,
                              32, 27,  3,  9,
                              19, 13, 30,  6,
                              22, 11,  4, 25])

-- TODO: Generalize from `16 in `Fin 16` and `nextPow2 16` to `n`.
select : Vect 4 (Vect 16 (Fin 16)) -> Bits 6 -> Bits (nextPow2 16)
select table bits =
  let row = bitsToFin (truncate {m=4} (or (shiftRightLogical (cast 5) (and bits (intToBits 32))) (and bits (intToBits 1))))
      col = bitsToFin (truncate {m=2} (shiftRightLogical (cast 1) (and bits (intToBits 30))))
  in finToBits (index col (index row table))

S : Vect 8 (Bits 6 -> Bits 4)
S = map select
        --  0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15
        [[[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]],
         [[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]],
         [[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]],
         [[ 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]],
         [[ 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]],
         [[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]],
         [[ 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]],
         [[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]]]

f : Bits 32 -> Bits 48 -> Bits 32
f R K = P (concat (zipWith apply S (partition (E R `xor` K))))

iteration : (Bits 32, Bits 32) -> Bits 48 -> (Bits 32, Bits 32)
iteration (L, R) K = (R, L `xor` f R K)

centralDEA : Bits 64 -> Vect 16 (Bits 48) -> Bits 64
centralDEA input keys =
  let [L, R] = the (Vect 2 (Bits 32)) (partition (IP input))
  in inverseIP (uncurry (flip append) (foldl iteration (L, R) keys))

||| permuted choice
PC1 : Bits 64 -> Bits 56
PC1 = selectBits (map offByOne [57, 49, 41, 33, 25, 17,  9,
                                 1, 58, 50, 42, 34, 26, 18,
                                10,  2, 59, 51, 43, 35, 27,
                                19, 11,  3, 60, 52, 44, 36,
                                63, 55, 47, 39, 31, 23, 15,
                                 7, 62, 54, 46, 38, 30, 22,
                                14,  6, 61, 53, 45, 37, 29,
                                21, 13,  5, 28, 20, 12,  4])

||| permuted choice
PC2 : Bits 56 -> Bits 48
PC2 = selectBits (map offByOne [14, 17, 11, 24,  1,  5,
                                 3, 28, 15,  6, 21, 10,
                                23, 19, 12,  4, 26,  8,
                                16,  7, 27, 20, 13,  2,
                                41, 52, 31, 37, 47, 55,
                                30, 40, 51, 45, 33, 48,
                                44, 49, 39, 56, 34, 53,
                                46, 42, 50, 36, 29, 32])

public export
DEAKey : Type
DEAKey = Bits 64

||| INSECURE!
public export
data DataEncryptionAlgorithm : Type where
  DEA : DEAKey -> DataEncryptionAlgorithm


||| key schedule
KS : DEAKey -> Vect 16 (Bits 48)
KS key = map PC2
             (tail (scanl (\prevKey, shift =>
                            concat (map (rotateLeft shift)
                                        (partition {m=28} {n=2} prevKey)))
                          (PC1 key)
                          [1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 1]))


implementation BlockCipher DataEncryptionAlgorithm 64 0 where
  encryptBlock (DEA key) block = centralDEA block (KS key)
  decryptBlock (DEA key) block = centralDEA block (reverse (KS key))
