1	# SM4 (GB/T 32907-2016, Information security technology - SM4 block cipher
2	# algorithm; in English, draft-ribose-cfrg-sm4 and RFC 8998, section 2.1):
3	# a 128-bit block cipher with a 128-bit key, as a cipher for cipher-modes'
4	# modes (~LcSm4Key, ~LcSm4.encrypt, ~LcSm4.decrypt).
5	#
6	# The state is four 32-bit words X0..X3, big-endian from the block, which is
7	# exactly a CmBlock. Each of the 32 rounds is
8	#   (X0, X1, X2, X3) -> (X1, X2, X3, X0 xor T(X1 xor X2 xor X3 xor rk_i))
9	# and the output is the last four words in reverse order. Decryption undoes
10	# each step in reverse order, which is the same function as the
11	# specification's "the round keys in reverse order", and makes
12	# decrypt(encrypt(b)) == b provable for every key and block: a round is
13	# undone whatever T is.
14	import Base
15	import bend-osdev-bytes@0.1.0.0/bytes.bend as B
16	import bend-osdev-cipher-modes@0.1.0.0/cipher_modes.bend as M
17	import ./legacy_ciphers.bend as L
18	
19	# The S-box (GB/T 32907-2016, 6.2: row i, column j is entry 16i + j)
20	def sm4.sbox.list() -> List<&2, U32>:
21	  [214, 144, 233, 254, 204, 225, 61, 183, 22, 182, 20, 194, 40, 251, 44, 5,
22	    43, 103, 154, 118, 42, 190, 4, 195, 170, 68, 19, 38, 73, 134, 6, 153,
23	    156, 66, 80, 244, 145, 239, 152, 122, 51, 84, 11, 67, 237, 207, 172, 98,
24	    228, 179, 28, 169, 201, 8, 232, 149, 128, 223, 148, 250, 117, 143, 63, 166,
25	    71, 7, 167, 252, 243, 115, 23, 186, 131, 89, 60, 25, 230, 133, 79, 168,
26	    104, 107, 129, 178, 113, 100, 218, 139, 248, 235, 15, 75, 112, 86, 157, 53,
27	    30, 36, 14, 94, 99, 88, 209, 162, 37, 34, 124, 59, 1, 33, 120, 135,
28	    212, 0, 70, 87, 159, 211, 39, 82, 76, 54, 2, 231, 160, 196, 200, 158,
29	    234, 191, 138, 210, 64, 199, 56, 181, 163, 247, 242, 206, 249, 97, 21, 161,
30	    224, 174, 93, 164, 155, 52, 26, 85, 173, 147, 50, 48, 245, 140, 177, 227,
31	    29, 246, 226, 46, 130, 102, 202, 96, 192, 41, 35, 171, 13, 83, 78, 111,
32	    213, 219, 55, 69, 222, 253, 142, 47, 3, 255, 106, 114, 109, 108, 91, 81,
33	    141, 27, 175, 146, 187, 221, 188, 127, 17, 217, 92, 65, 31, 16, 90, 216,
34	    10, 193, 49, 136, 165, 205, 123, 189, 45, 116, 208, 18, 184, 229, 180, 176,
35	    137, 105, 151, 74, 12, 150, 119, 126, 101, 185, 241, 9, 197, 110, 198, 132,
36	    24, 240, 125, 236, 58, 220, 77, 32, 121, 238, 95, 62, 215, 203, 57, 72]
37	
38	def sm4.sbox() -> L.LcTab:
39	  L.lc.tab(8n, sm4.sbox.list())
40	
41	def sm4.s(+tab: L.LcTab, +a: U32, n: Nat) -> U32:
42	  L.lc.tab.get(8n, tab, L.lc.byte(a, n))
43	
44	# tau: the S-box on each byte of a word
45	def sm4.tau(+tab: L.LcTab, +a: U32) -> U32:
46	  L.lc.be(sm4.s(tab, a, 3n), sm4.s(tab, a, 2n), sm4.s(tab, a, 1n), sm4.s(tab, a, 0n))
47	
48	# T = L(tau(x)), L(B) = B xor B<<<2 xor B<<<10 xor B<<<18 xor B<<<24
49	def sm4.t(+tab: L.LcTab, x: U32) -> U32:
50	  +b = sm4.tau(tab, x)
51	  U32.xor(U32.xor(U32.xor(b, L.lc.rotl(b, 2n)), U32.xor(L.lc.rotl(b, 10n), L.lc.rotl(b, 18n))), L.lc.rotl(b, 24n))
52	
53	# T' of the key schedule: L'(B) = B xor B<<<13 xor B<<<23
54	def sm4.t2(+tab: L.LcTab, x: U32) -> U32:
55	  +b = sm4.tau(tab, x)
56	  U32.xor(U32.xor(b, L.lc.rotl(b, 13n)), L.lc.rotl(b, 23n))
57	
58	# The rounds, on the state as a CmBlock
59	
60	def sm4.step(+tab: L.LcTab, rk: U32, s: M.CmBlock) -> M.CmBlock:
61	  match s:
62	    case M.CmBlock{x0, +x1, +x2, +x3}:
63	      M.CmBlock{x1, x2, x3, U32.xor(x0, sm4.t(tab, U32.xor(U32.xor(U32.xor(x1, x2), x3), rk)))}
64	
65	# the round undone
66	def sm4.unstep(+tab: L.LcTab, rk: U32, s: M.CmBlock) -> M.CmBlock:
67	  match s:
68	    case M.CmBlock{+y0, +y1, +y2, y3}:
69	      M.CmBlock{U32.xor(y3, sm4.t(tab, U32.xor(U32.xor(U32.xor(y0, y1), y2), rk))), y0, y1, y2}
70	
71	def sm4.rounds(+tab: L.LcTab, rks: List<&2, U32>, s: M.CmBlock) -> M.CmBlock:
72	  match rks:
73	    case Nil{}:
74	      s
75	    case Con{k, t}:
76	      sm4.rounds(tab, t, sm4.step(tab, k, s))
77	
78	def sm4.unrounds(+tab: L.LcTab, rks: List<&2, U32>, s: M.CmBlock) -> M.CmBlock:
79	  match rks:
80	    case Nil{}:
81	      s
82	    case Con{k, t}:
83	      sm4.unstep(tab, k, sm4.unrounds(tab, t, s))
84	
85	# the reverse transformation R: (A0, A1, A2, A3) -> (A3, A2, A1, A0)
86	def sm4.rev(s: M.CmBlock) -> M.CmBlock:
87	  match s:
88	    case M.CmBlock{a, b, c, d}:
89	      M.CmBlock{d, c, b, a}
90	
91	# The key schedule (GB/T 32907-2016, 7.3)
92	
93	# FK and CK_i (ck_(i,j) = (4i + j) * 7 mod 256)
94	def sm4.fk(n: Nat) -> U32:
95	  L.lc.at([2746333894, 1453994832, 1736282519, 2993693404], n)
96	
97	def sm4.ck(+i: U32) -> U32:
98	  +b = U32.mul(i, 28)
99	  L.lc.be(b, U32.add(b, 7), U32.add(b, 14), U32.add(b, 21))
100	
101	# rk_i = K_(i+4) = K_i xor T'(K_(i+1) xor K_(i+2) xor K_(i+3) xor CK_i), for
102	# n more rounds from round i, the window K_i .. K_(i+3) in k0 .. k3
103	def sm4.schedule(n: Nat, +tab: L.LcTab, +i: U32, k0: U32, +k1: U32, +k2: U32, +k3: U32) -> List<&2, U32>:
104	  match n:
105	    case 0n:
106	      Nil{}
107	    case 1n+p:
108	      +k4 = U32.xor(k0, sm4.t2(tab, U32.xor(U32.xor(U32.xor(k1, k2), k3), sm4.ck(i))))
109	      Con{k4, sm4.schedule(p, tab, U32.add(i, 1), k1, k2, k3, k4)}
110	
111	# An SM4 key: the S-box table (built once per key) and the 32 round keys
112	# (opaque; built by LcSm4.key)
113	type LcSm4Key is Data:
114	  LcSm4Key{tab: L.LcTab, rks: List<&2, U32>}
115	
116	def sm4.mk(+xs: List<&2, U32>, +i: Nat) -> U32:
117	  +j = Nat.mul(4n, i)
118	  U32.xor(L.lc.be(L.lc.at(xs, j), L.lc.at(xs, Nat.add(j, 1n)), L.lc.at(xs, Nat.add(j, 2n)), L.lc.at(xs, Nat.add(j, 3n))), sm4.fk(i))
119	
120	def sm4.key.pick(ok: Bool, +xs: List<&2, U32>) -> Result<&2, &2, L.LcError, LcSm4Key>:
121	  match ok:
122	    case True{}:
123	      +tab = sm4.sbox()
124	      Done{LcSm4Key{tab, sm4.schedule(32n, tab, 0, sm4.mk(xs, 0n), sm4.mk(xs, 1n), sm4.mk(xs, 2n), sm4.mk(xs, 3n))}}
125	    case False{}:
126	      Fail{L.LcBadKeyLength{B.len(xs)}}
127	
128	# The cipher
129	
130	# a 16-byte key
131	def LcSm4.key(o: B.Octets) -> Result<&2, &2, L.LcError, LcSm4Key>:
132	  +xs = B.Octets.to_list(o)
133	  sm4.key.pick(Nat.is_eq(B.len(xs), 16n), xs)
134	
135	def LcSm4.encrypt(k: LcSm4Key, b: M.CmBlock) -> M.CmBlock:
136	  match k:
137	    case LcSm4Key{+tab, rks}:
138	      sm4.rev(sm4.rounds(tab, rks, b))
139	
140	def LcSm4.decrypt(k: LcSm4Key, b: M.CmBlock) -> M.CmBlock:
141	  match k:
142	    case LcSm4Key{+tab, rks}:
143	      sm4.unrounds(tab, rks, sm4.rev(b))