-- ============================================================================
-- AES MixColumns Formal Verification — Lean 4
-- Canonical test vector and linearity properties
-- Authors: Ahmad Ali Parr, Jessica L. Williams (SNAPKITTYWEST)
-- ============================================================================

import Mathlib.Data.BitVec.Basic

namespace AESVerification

-- GF(2^8) xtime: multiply by x mod (x^8 + x^4 + x^3 + x + 1)
def xtime (b : BitVec 8) : BitVec 8 :=
  let msb : Bool := b.getMsbD 0
  let shifted : BitVec 8 := b <<< 1
  if msb then shifted ^^^ 27 else shifted

-- Row 0 of MixColumns: y0 = a0 ⊕ t ⊕ xtime(a0⊕a1)
-- where t = a0⊕a1⊕a2⊕a3
def mix_col_r0 (a0 a1 a2 a3 : BitVec 8) : BitVec 8 :=
  let t := a0 ^^^ a1 ^^^ a2 ^^^ a3
  a0 ^^^ t ^^^ xtime (a0 ^^^ a1)

-- Canonical test vector: D4 BF 5D 30 → row0 = 04
def canonical_a0 : BitVec 8 := 0xD4
def canonical_a1 : BitVec 8 := 0xBF
def canonical_a2 : BitVec 8 := 0x5D
def canonical_a3 : BitVec 8 := 0x30

theorem canonical_r0_correct :
    mix_col_r0 canonical_a0 canonical_a1 canonical_a2 canonical_a3 = 0x04 := by
  decide

-- Zero vector: MixColumns(0,0,0,0) = (0,0,0,0)
theorem zero_column :
    mix_col_r0 0 0 0 0 = 0 := by decide

-- Repeated byte: MixColumns(k,k,k,k)[0] = k (linearity: all differences cancel)
-- t = 0, xtime(0) = 0, so y0 = k ⊕ 0 ⊕ 0 = k
theorem repeated_byte (k : BitVec 8) :
    mix_col_r0 k k k k = k := by
  simp [mix_col_r0, xtime]

-- Linearity: xtime(a ⊕ b) = xtime(a) ⊕ xtime(b)
-- (GF(2^8) multiplication is a linear map over GF(2))
theorem xtime_linear (a b : BitVec 8) :
    xtime (a ^^^ b) = xtime a ^^^ xtime b := by
  simp [xtime]
  split <;> split <;> split <;> simp_all [BitVec.xor_assoc]
  all_goals decide

end AESVerification