Source file kuznechik.ml

(* Utils and types *)

let make_random_int64s n = List.init n (fun _ -> Random.bits64 ())

let generate_key =
  Random.self_init () ;
  let k = Bytes.create 32 in
  List.iteri
    (fun i v -> Bytes.set_int64_le k (i * 8) v)
    (make_random_int64s 4) ;
  k

let _gf_mul_step (a, b, c) =
  let a' = if a land 0x80 <> 0 then (a lsl 1) lxor 0xc3 else a lsl 1 in
  let b' = b lsr 1 in
  let c' = if b land 1 = 1 then c lxor a else c in
  (a', b', c')

let gf_mul a b =
  match Seq.drop 8 (Seq.iterate _gf_mul_step (a, b, 0)) () with
  | Seq.Cons ((_, _, c), _) -> c mod 256
  | _ -> -1

(* X Transition *)

let x_trans a b =
  Seq.map2
    (fun a_el b_el -> Char.chr (Char.code a_el lxor Char.code b_el))
    (Bytes.to_seq a) (Bytes.to_seq b)
  |> Bytes.of_seq

(* S Transition and Inverse S Transition*)

let pi_table: (int) list =
  [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) list =
 [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 list) (a : bytes) : bytes =
  a |> Bytes.map (fun x -> List.nth l (Char.code x) |> Char.chr)

let s_trans : bytes -> bytes = s_generic_trans pi_table

let s_inv_trans : bytes -> bytes = s_generic_trans pi_inv_table

(* L Transition and Inverse L Transition*)

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

let r_trans (a : bytes) : bytes =
  let new_val =
    Seq.fold_lefti
      (fun v i c -> gf_mul (Char.code c) (List.nth l_vec i) lxor v)
      0 (Bytes.to_seq a)
  in
  Bytes.sub (Bytes.cat (Bytes.make 1 (Char.chr new_val)) a) 0 16

let r_inv_trans a =
  let new_val =
    Seq.fold_lefti
      (fun v i c -> gf_mul (Char.code c) (List.nth l_vec i) lxor v)
      (Char.code (Bytes.get a 0))
      (Seq.drop 1 (Bytes.to_seq a))
  in
  Bytes.cat (Bytes.sub a 1 15) (Bytes.make 1 (Char.chr new_val))

let l_generic_trans f a =
  match Seq.drop 16 (Seq.iterate f a) () with Seq.Cons (k, _) -> k | _ -> a

let l_trans = l_generic_trans r_trans

let l_inv_trans = l_generic_trans r_inv_trans

(* Round constants computation *)

let generate_round_constants =
  List.init 33 (fun i ->
      Bytes.cat (Bytes.make 15 (Char.chr 0)) (Bytes.make 1 (Char.chr i))
      |> l_trans )
  |> List.tl

(* Key extension *)

let f_trans k1 k2 iter_const =
  (k1 |> x_trans iter_const |> s_trans |> l_trans |> x_trans k2, k1)

let make_iter_keys k1 k2 round_const =
  Seq.iterate
    (fun ((ik1, ik2), i) -> (f_trans ik1 ik2 (List.nth round_const i), i + 1))
    ((k1, k2), 0)
  |> Seq.filter (fun (_, i) -> i mod 8 = 0)
  |> Seq.take 5
  |> Seq.map (fun ((ik1, ik2), _) -> [ik1; ik2])
  |> List.of_seq |> List.flatten

(* Encrypt and Decrypt *)

let encrypt_block key msg =
  let ik =
    make_iter_keys (Bytes.sub key 0 16) (Bytes.sub key 16 16)
      generate_round_constants
  in
  match
    Seq.drop 9
      (Seq.iterate
         (fun (x, i) ->
           (x |> x_trans (List.nth ik i) |> s_trans |> l_trans, i + 1) )
         (msg, 0) )
      ()
  with
  | Seq.Cons ((x, _), _) -> x |> x_trans (List.nth ik 9)
  | _ -> msg

let decrypt_block key msg =
  let ik =
    make_iter_keys (Bytes.sub key 0 16) (Bytes.sub key 16 16)
      generate_round_constants
  in
  match
    Seq.drop 9
      (Seq.iterate
         (fun (x, i) ->
           ( x |> l_inv_trans |> s_inv_trans |> x_trans (List.nth ik (8 - i))
           , i + 1 ) )
         (msg |> x_trans (List.nth ik 9), 0) )
      ()
  with
  | Seq.Cons ((x, _), _) -> x
  | _ -> msg