Source file kuznechik.ml

(* Utils and types *)
open! Core

let gf_mul (a : int) (b : int) =
  Bigarray.Array1.unsafe_get Gf_mul_table.gf_mul_table ((a * 256) + b)

(* X Transition *)

let x_trans (a : int array) (b : int array) : int array =
  Array.map2_exn a b ~f:(fun a_el b_el -> a_el lxor b_el)

(* S Transition and Inverse S Transition*)

let pi_table: int array =
  [|0xFC; 0xEE; 0xDD; 0x11; 0xCF; 0x6E; 0x31; 0x16;
   0xFB; 0xC4; 0xFA; 0xDA; 0x23; 0xC5; 0x04; 0x4D;
   0xE9; 0x77; 0xF0; 0xDB; 0x93; 0x2E; 0x99; 0xBA;
   0x17; 0x36; 0xF1; 0xBB; 0x14; 0xCD; 0x5F; 0xC1;
   0xF9; 0x18; 0x65; 0x5A; 0xE2; 0x5C; 0xEF; 0x21;
   0x81; 0x1C; 0x3C; 0x42; 0x8B; 0x01; 0x8E; 0x4F;
   0x05; 0x84; 0x02; 0xAE; 0xE3; 0x6A; 0x8F; 0xA0;
   0x06; 0x0B; 0xED; 0x98; 0x7F; 0xD4; 0xD3; 0x1F;
   0xEB; 0x34; 0x2C; 0x51; 0xEA; 0xC8; 0x48; 0xAB;
   0xF2; 0x2A; 0x68; 0xA2; 0xFD; 0x3A; 0xCE; 0xCC;
   0xB5; 0x70; 0x0E; 0x56; 0x08; 0x0C; 0x76; 0x12;
   0xBF; 0x72; 0x13; 0x47; 0x9C; 0xB7; 0x5D; 0x87;
   0x15; 0xA1; 0x96; 0x29; 0x10; 0x7B; 0x9A; 0xC7;
   0xF3; 0x91; 0x78; 0x6F; 0x9D; 0x9E; 0xB2; 0xB1;
   0x32; 0x75; 0x19; 0x3D; 0xFF; 0x35; 0x8A; 0x7E;
   0x6D; 0x54; 0xC6; 0x80; 0xC3; 0xBD; 0x0D; 0x57;
   0xDF; 0xF5; 0x24; 0xA9; 0x3E; 0xA8; 0x43; 0xC9;
   0xD7; 0x79; 0xD6; 0xF6; 0x7C; 0x22; 0xB9; 0x03;
   0xE0; 0x0F; 0xEC; 0xDE; 0x7A; 0x94; 0xB0; 0xBC;
   0xDC; 0xE8; 0x28; 0x50; 0x4E; 0x33; 0x0A; 0x4A;
   0xA7; 0x97; 0x60; 0x73; 0x1E; 0x00; 0x62; 0x44;
   0x1A; 0xB8; 0x38; 0x82; 0x64; 0x9F; 0x26; 0x41;
   0xAD; 0x45; 0x46; 0x92; 0x27; 0x5E; 0x55; 0x2F;
   0x8C; 0xA3; 0xA5; 0x7D; 0x69; 0xD5; 0x95; 0x3B;
   0x07; 0x58; 0xB3; 0x40; 0x86; 0xAC; 0x1D; 0xF7;
   0x30; 0x37; 0x6B; 0xE4; 0x88; 0xD9; 0xE7; 0x89;
   0xE1; 0x1B; 0x83; 0x49; 0x4C; 0x3F; 0xF8; 0xFE;
   0x8D; 0x53; 0xAA; 0x90; 0xCA; 0xD8; 0x85; 0x61;
   0x20; 0x71; 0x67; 0xA4; 0x2D; 0x2B; 0x09; 0x5B;
   0xCB; 0x9B; 0x25; 0xD0; 0xBE; 0xE5; 0x6C; 0x52;
   0x59; 0xA6; 0x74; 0xD2; 0xE6; 0xF4; 0xB4; 0xC0;
   0xD1; 0x66; 0xAF; 0xC2; 0x39; 0x4B; 0x63; 0xB6|] [@@ocamlformat "disable"]

let pi_inv_table: int array =
 [|0xA5; 0x2D; 0x32; 0x8F; 0x0E; 0x30; 0x38; 0xC0;
  0x54; 0xE6; 0x9E; 0x39; 0x55; 0x7E; 0x52; 0x91;
  0x64; 0x03; 0x57; 0x5A; 0x1C; 0x60; 0x07; 0x18;
  0x21; 0x72; 0xA8; 0xD1; 0x29; 0xC6; 0xA4; 0x3F;
  0xE0; 0x27; 0x8D; 0x0C; 0x82; 0xEA; 0xAE; 0xB4;
  0x9A; 0x63; 0x49; 0xE5; 0x42; 0xE4; 0x15; 0xB7;
  0xC8; 0x06; 0x70; 0x9D; 0x41; 0x75; 0x19; 0xC9;
  0xAA; 0xFC; 0x4D; 0xBF; 0x2A; 0x73; 0x84; 0xD5;
  0xC3; 0xAF; 0x2B; 0x86; 0xA7; 0xB1; 0xB2; 0x5B;
  0x46; 0xD3; 0x9F; 0xFD; 0xD4; 0x0F; 0x9C; 0x2F;
  0x9B; 0x43; 0xEF; 0xD9; 0x79; 0xB6; 0x53; 0x7F;
  0xC1; 0xF0; 0x23; 0xE7; 0x25; 0x5E; 0xB5; 0x1E;
  0xA2; 0xDF; 0xA6; 0xFE; 0xAC; 0x22; 0xF9; 0xE2;
  0x4A; 0xBC; 0x35; 0xCA; 0xEE; 0x78; 0x05; 0x6B;
  0x51; 0xE1; 0x59; 0xA3; 0xF2; 0x71; 0x56; 0x11;
  0x6A; 0x89; 0x94; 0x65; 0x8C; 0xBB; 0x77; 0x3C;
  0x7B; 0x28; 0xAB; 0xD2; 0x31; 0xDE; 0xC4; 0x5F;
  0xCC; 0xCF; 0x76; 0x2C; 0xB8; 0xD8; 0x2E; 0x36;
  0xDB; 0x69; 0xB3; 0x14; 0x95; 0xBE; 0x62; 0xA1;
  0x3B; 0x16; 0x66; 0xE9; 0x5C; 0x6C; 0x6D; 0xAD;
  0x37; 0x61; 0x4B; 0xB9; 0xE3; 0xBA; 0xF1; 0xA0;
  0x85; 0x83; 0xDA; 0x47; 0xC5; 0xB0; 0x33; 0xFA;
  0x96; 0x6F; 0x6E; 0xC2; 0xF6; 0x50; 0xFF; 0x5D;
  0xA9; 0x8E; 0x17; 0x1B; 0x97; 0x7D; 0xEC; 0x58;
  0xF7; 0x1F; 0xFB; 0x7C; 0x09; 0x0D; 0x7A; 0x67;
  0x45; 0x87; 0xDC; 0xE8; 0x4F; 0x1D; 0x4E; 0x04;
  0xEB; 0xF8; 0xF3; 0x3E; 0x3D; 0xBD; 0x8A; 0x88;
  0xDD; 0xCD; 0x0B; 0x13; 0x98; 0x02; 0x93; 0x80;
  0x90; 0xD0; 0x24; 0x34; 0xCB; 0xED; 0xF4; 0xCE;
  0x99; 0x10; 0x44; 0x40; 0x92; 0x3A; 0x01; 0x26;
  0x12; 0x1A; 0x48; 0x68; 0xF5; 0x81; 0x8B; 0xC7;
  0xD6; 0x20; 0x0A; 0x08; 0x00; 0x4C; 0xD7; 0x74|] [@@ocamlformat "disable"]

let s_generic_trans (l : int array) (a : int array) : int array =
  a |> Array.map ~f:(fun x -> Array.unsafe_get l x)

let s_trans : int array -> int array = s_generic_trans pi_table

let s_inv_trans : int array -> int array = s_generic_trans pi_inv_table

(* L Transition and Inverse L Transition*)

let l_vec : int array =
  [|148; 32; 133; 16; 194; 192; 1; 251; 1; 192; 194; 16; 133; 32; 148; 1|]

let r_trans (a : int array) : int array =
  let new_val =
    match
      Array.fold a ~init:(0, 0) ~f:(fun (v, i) c ->
          (gf_mul c (Array.unsafe_get l_vec i) lxor v, i + 1) )
    with
    | v, _ -> v
  in
  Array.unsafe_blit ~src:a ~src_pos:0 ~dst:a ~dst_pos:1 ~len:15 ;
  Array.unsafe_set a 0 new_val ;
  a

let r_inv_trans (a : int array) : int array =
  let new_val =
    match
      Array.fold
        (Array.sub a ~pos:1 ~len:15)
        ~init:(Array.unsafe_get a 0, 0)
        ~f:(fun (v, i) c ->
          (gf_mul c (Array.unsafe_get l_vec i) lxor v, i + 1) )
    with
    | v, _ -> v
  in
  Array.unsafe_blit ~src:a ~src_pos:1 ~dst:a ~dst_pos:0 ~len:15 ;
  Array.unsafe_set a 15 new_val ;
  a

let rec _l_generic_trans_inner (f : int array -> int array) (a : int array)
    (i : int) : int array =
  match i with 0 -> a | _ -> _l_generic_trans_inner f (f a) (i - 1)

let l_generic_trans (f : int array -> int array) (a : int array) : int array
    =
  _l_generic_trans_inner f a 16

let l_trans : int array -> int array = l_generic_trans r_trans

let l_inv_trans : int array -> int array = l_generic_trans r_inv_trans

(* Round constants computation *)

let round_constants : int array array =
  Array.sub ~pos:1 ~len:32
    (Array.init 33 ~f:(fun i ->
         Array.append (Array.create ~len:15 0) (Array.create ~len:1 i)
         |> l_trans ) )

(* Key extension *)

let f_trans (k1 : int array) (k2 : int array) (iter_const : int array) :
    int array * int array =
  (k1 |> x_trans iter_const |> s_trans |> l_trans |> x_trans k2, k1)

let make_iter_keys (k1 : int array) (k2 : int array)
    (round_const : int array array) : int array array =
  Sequence.take
    ( Sequence.append
        (Sequence.singleton ((k1, k2), 0))
        (Sequence.unfold
           ~init:((k1, k2), 0)
           ~f:(fun ((ik1, ik2), i) ->
             let new_val =
               (f_trans ik1 ik2 (Array.get round_const i), i + 1)
             in
             Option.some (new_val, new_val) ) )
    |> Sequence.filter ~f:(fun (_, i) -> i mod 8 = 0) )
    5
  |> Sequence.map ~f:(fun ((ik1, ik2), _) -> [ik1; ik2])
  |> Sequence.to_list |> List.concat |> Array.of_list

(* Encrypt and Decrypt *)

module Cipher = struct
  type t = {key: int array; ik: int array array}

  let make (key : int array) : t =
    { key
    ; ik=
        make_iter_keys
          (Array.sub key ~pos:0 ~len:16)
          (Array.sub key ~pos:16 ~len:16)
          round_constants }

  let rec _encrypt_block_inner (ik : int array array) (x : int array)
      (i : int) : int array =
    match i with
    | 9 -> x
    | _ ->
        _encrypt_block_inner ik
          (x |> x_trans ik.(i) |> s_trans |> l_trans)
          (i + 1)

  let encrypt_block (cipher : t) (msg : int array) : int array =
    _encrypt_block_inner cipher.ik msg 0 |> x_trans cipher.ik.(9)

  let rec _decrypt_block_inner (ik : int array array) (x : int array)
      (i : int) : int array =
    match i with
    | 9 -> x
    | _ ->
        _decrypt_block_inner ik
          (x |> l_inv_trans |> s_inv_trans |> x_trans ik.(8 - i))
          (i + 1)

  let decrypt_block (cipher : t) (msg : int array) : int array =
    _decrypt_block_inner cipher.ik (msg |> x_trans cipher.ik.(9)) 0
end